From a5e9c0cd0197cdc2c2e0c06e79baef12c77efab9 Mon Sep 17 00:00:00 2001 From: Nikhil Kumar Date: Sat, 11 Jul 2026 05:33:13 -0400 Subject: [PATCH] Add CNN calo vertex implementation For now I am only supplying z. I can add z_err (per event sigma) in another PR. --- .../packages/calovtxreco/CaloVtxAlgoCNN.cc | 407 ++++++++++++++++++ offline/packages/calovtxreco/CaloVtxAlgoCNN.h | 70 +++ offline/packages/calovtxreco/Makefile.am | 7 +- offline/packages/globalvertex/GlobalVertex.h | 1 + .../packages/globalvertex/GlobalVertexReco.cc | 6 +- .../packages/globalvertex/GlobalVertexv4.cc | 6 + offline/packages/globalvertex/VertexDefs.h | 3 +- 7 files changed, 496 insertions(+), 4 deletions(-) create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoCNN.cc create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoCNN.h diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCNN.cc b/offline/packages/calovtxreco/CaloVtxAlgoCNN.cc new file mode 100644 index 0000000000..df1a18f54d --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoCNN.cc @@ -0,0 +1,407 @@ +#include "CaloVtxAlgoCNN.h" + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace +{ + const std::string TowerNode[CaloVtxAlgoCNN::kNLayer] = { + "TOWERINFO_CALIB_CEMC", "TOWERINFO_CALIB_HCALIN", "TOWERINFO_CALIB_HCALOUT"}; + const std::string GeomNodeEmc = "TOWERGEOM_CEMC"; + const std::string GeomNodeIhc = "TOWERGEOM_HCALIN"; +} // namespace + +// onnxruntime session (pImpl, keeps Ort types out of the header) +struct CaloVtxAlgoCNN::OnnxSession +{ + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "CaloVtxAlgoCNN"}; + Ort::SessionOptions opts; + std::unique_ptr session; + Ort::MemoryInfo memInfo{ + Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)}; + + explicit OnnxSession(const std::string &path) + { + opts.SetIntraOpNumThreads(1); + opts.SetGraphOptimizationLevel(ORT_ENABLE_ALL); + session = std::make_unique(env, path.c_str(), opts); + } +}; + +CaloVtxAlgoCNN::~CaloVtxAlgoCNN() = default; + +int CaloVtxAlgoCNN::Init(PHCompositeNode * /*topNode*/) +{ + try + { + m_onnx = std::make_unique(m_modelFile); + } + catch (const std::exception &e) + { + std::cout << PHWHERE << "failed to load model " << m_modelFile << ": " << e.what() << std::endl; + return Fun4AllReturnCodes::ABORTRUN; + } + return Fun4AllReturnCodes::EVENT_OK; +} + +int CaloVtxAlgoCNN::CalculateVertex(PHCompositeNode *topNode, float &zvtx) +{ + zvtx = std::numeric_limits::quiet_NaN(); + + if (fillTowerImage(topNode) != 0) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + double etot = 0.; + for (int layer = 0; layer < kNLayer; ++layer) + { + for (int ieta = 0; ieta < kNEtaImg; ++ieta) + { + for (int iphi = 0; iphi < kNPhiImg; ++iphi) + { + etot += m_image[layer][ieta][iphi]; + } + } + } + if (etot <= m_minTotalEnergy) + { + // empty image: leave zvtx NaN, do not run the network + return Fun4AllReturnCodes::EVENT_OK; + } + + float z = std::numeric_limits::quiet_NaN(); + if (!predict(z)) + { + return Fun4AllReturnCodes::ABORTEVENT; + } + zvtx = z; + return Fun4AllReturnCodes::EVENT_OK; +} + +bool CaloVtxAlgoCNN::predict(float &z) +{ + const int64_t shape[4] = {1, kNLayer, kNEtaImg, kNPhiImg}; + try + { + Ort::Value input = Ort::Value::CreateTensor(m_onnx->memInfo, &m_image[0][0][0], static_cast(kNLayer) * kNEtaImg * kNPhiImg, shape, 4); + const char *inNames[] = {"raw_image"}; + const char *outNames[] = {"z_cal_cm"}; + auto outs = m_onnx->session->Run(Ort::RunOptions{nullptr}, inNames, &input, 1, outNames, 1); + z = outs[0].GetTensorData()[0]; + return true; + } + catch (const std::exception &e) + { + if (!m_warnedPredict) + { + std::cout << PHWHERE << "model evaluation failed: " << e.what() << std::endl; + m_warnedPredict = true; + } + return false; + } +} + +// pre-processing: fill the 3-layer tower image from the nodes, including retowering of the EMCAL +int CaloVtxAlgoCNN::fillTowerImage(PHCompositeNode *topNode) +{ + for (int layer = 0; layer < kNLayer; ++layer) + { + for (int ieta = 0; ieta < kNEtaImg; ++ieta) + { + for (int iphi = 0; iphi < kNPhiImg; ++iphi) + { + m_image[layer][ieta][iphi] = 0.; + } + } + } + // HCal layers are already on the common 24x64 grid. + if (fillHcalLayer(topNode, kIHC) != 0) + { + return -1; + } + if (fillHcalLayer(topNode, kOHC) != 0) + { + return -1; + } + return fillEmcRetower(topNode); +} + +int CaloVtxAlgoCNN::fillHcalLayer(PHCompositeNode *topNode, int layer) +{ + TowerInfoContainer *towers = findNode::getClass(topNode, TowerNode[layer]); + if (!towers) + { + if (!m_warnedMissingTowers) + { + std::cout << PHWHERE << "tower node missing: " << TowerNode[layer] << std::endl; + m_warnedMissingTowers = true; + } + return -1; + } + + const unsigned int ntow = towers->size(); + for (unsigned int ch = 0; ch < ntow; ++ch) + { + TowerInfo *tower = towers->get_tower_at_channel(ch); + if (!tower || !tower->get_isGood()) + { + continue; + } + const float e = tower->get_energy(); + if (e < m_towerEMin[layer]) + { + continue; + } + const unsigned int key = towers->encode_key(ch); + const int ieta = towers->getTowerEtaBin(key); + const int iphi = towers->getTowerPhiBin(key); + if (ieta < 0 || ieta >= kNEtaImg || iphi < 0 || iphi >= kNPhiImg) + { + continue; + } + m_image[layer][ieta][iphi] += e; + } + return 0; +} + +int CaloVtxAlgoCNN::fillEmcRetower(PHCompositeNode *topNode) +{ + if (buildEmcRetowerMap(topNode) != 0) + { + return -1; + } + + TowerInfoContainer *towers = findNode::getClass(topNode, TowerNode[kEMC]); + if (!towers) + { + if (!m_warnedMissingTowers) + { + std::cout << PHWHERE << "tower node missing: " << TowerNode[kEMC] << std::endl; + m_warnedMissingTowers = true; + } + return -1; + } + + for (auto &row : m_rawEmcFine) + { + for (double &e : row) + { + e = 0.; + } + } + + const unsigned int ntow = towers->size(); + for (unsigned int ch = 0; ch < ntow; ++ch) + { + TowerInfo *tower = towers->get_tower_at_channel(ch); + if (!tower || !tower->get_isGood()) + { + continue; + } + const float e = tower->get_energy(); + if (e < m_towerEMin[kEMC]) + { + continue; + } + const unsigned int key = towers->encode_key(ch); + const int ieta = towers->getTowerEtaBin(key); + const int iphi = towers->getTowerPhiBin(key); + if (ieta < 0 || ieta >= kNEtaEmcFine || iphi < 0 || iphi >= kNPhiEmcFine) + { + continue; + } + m_rawEmcFine[ieta][iphi] = e; + } + + // eta-fraction / phi-grouping sums, as in RetowerCEMC + for (int ietaHcal = 0; ietaHcal < kNEtaImg; ++ietaHcal) + { + for (int iphiHcal = 0; iphiHcal < kNPhiImg; ++iphiHcal) + { + double retowerE = 0.; + for (int ietaEmc = m_retowerLowerEta[ietaHcal]; ietaEmc <= m_retowerUpperEta[ietaHcal]; ++ietaEmc) + { + double fraction = 1.; + if (ietaEmc == m_retowerLowerEta[ietaHcal]) + { + fraction = m_retowerLowerFrac[ietaHcal]; + } + else if (ietaEmc == m_retowerUpperEta[ietaHcal]) + { + fraction = m_retowerUpperFrac[ietaHcal]; + } + for (int iphiEmc = m_retowerPhiOffset + iphiHcal * 4; iphiEmc < m_retowerPhiOffset + iphiHcal * 4 + 4; ++iphiEmc) + { + int iphiEmcWrap = iphiEmc; + if (iphiEmcWrap > kNPhiEmcFine - 1) + { + iphiEmcWrap -= kNPhiEmcFine; + } + retowerE += m_rawEmcFine[ietaEmc][iphiEmcWrap] * fraction; + } + } + m_image[kEMC][ietaHcal][iphiHcal] = retowerE; + } + } + return 0; +} + +int CaloVtxAlgoCNN::buildEmcRetowerMap(PHCompositeNode *topNode) +{ + if (m_retowerMapReady) + { + return 0; + } + + RawTowerGeomContainer *geomEM = findNode::getClass(topNode, GeomNodeEmc); + RawTowerGeomContainer *geomIH = findNode::getClass(topNode, GeomNodeIhc); + if (!geomEM || !geomIH) + { + if (!m_warnedRetowerMap) + { + std::cout << PHWHERE << "cannot build EMCal retower map, missing " << GeomNodeEmc << " or " << GeomNodeIhc << std::endl; + m_warnedRetowerMap = true; + } + return -1; + } + + // first fine EMCal phi bin belonging to HCal phi 0, cf. RetowerCEMC::get_first_phi_index() + bool foundFirstLowerBound = false; + int iphiEmc = 0; + while (iphiEmc < kNPhiEmcFine) + { + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CEMC, 0, iphiEmc); + RawTowerGeom *towerGeom = geomEM->get_tower_geometry(key); + if (towerGeom && geomIH->get_phibin(towerGeom->get_phi()) == 0) + { + foundFirstLowerBound = true; + break; + } + ++iphiEmc; + } + + if (foundFirstLowerBound && iphiEmc == 0) + { + bool outOfRange = false; + int iphiEmcTemp = kNPhiEmcFine - 1; + while (iphiEmcTemp > iphiEmc) + { + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CEMC, 0, iphiEmcTemp); + RawTowerGeom *towerGeom = geomEM->get_tower_geometry(key); + if (towerGeom && geomIH->get_phibin(towerGeom->get_phi()) == kNPhiImg - 1) + { + outOfRange = true; + break; + } + --iphiEmcTemp; + } + if (!outOfRange) + { + if (!m_warnedRetowerMap) + { + std::cout << PHWHERE << "cannot build EMCal retower map, no wrap-around " << "phi match" << std::endl; + m_warnedRetowerMap = true; + } + return -1; + } + m_retowerPhiOffset = (iphiEmcTemp + 1 == kNPhiEmcFine) ? 0 : iphiEmcTemp + 1; + } + else if (!foundFirstLowerBound) + { + if (!m_warnedRetowerMap) + { + std::cout << PHWHERE << "cannot build EMCal retower map, no EMCal phi bin " << "maps to HCal phi 0" << std::endl; + m_warnedRetowerMap = true; + } + return -1; + } + else + { + m_retowerPhiOffset = iphiEmc; + } + + // eta-bound overlaps (edge bins fractional), cf. RetowerCEMC::get_weighted_fraction() + int ietaEmc = 0; + for (int ietaHcal = 0; ietaHcal < kNEtaImg; ++ietaHcal) + { + const std::pair rangeHcal = geomIH->get_etabounds(ietaHcal); + const double hcalLower = rangeHcal.first; + const double hcalUpper = rangeHcal.second; + bool foundLower = false; + bool foundUpper = false; + + while ((!foundLower || !foundUpper) && ietaEmc < kNEtaEmcFine) + { + const std::pair rangeEmc = geomEM->get_etabounds(ietaEmc); + const double emcLower = rangeEmc.first; + const double emcUpper = rangeEmc.second; + + if (!foundLower) + { + if (emcUpper > hcalLower && emcLower <= hcalLower) + { + m_retowerLowerEta[ietaHcal] = ietaEmc; + m_retowerLowerFrac[ietaHcal] = (emcUpper - hcalLower) / (emcUpper - emcLower); + foundLower = true; + } + if (emcUpper > hcalLower && emcLower > hcalLower) + { + m_retowerLowerEta[ietaHcal] = ietaEmc; + m_retowerLowerFrac[ietaHcal] = 1.; + foundLower = true; + } + } + else + { + if (emcUpper >= hcalUpper && emcLower < hcalUpper) + { + m_retowerUpperEta[ietaHcal] = ietaEmc; + m_retowerUpperFrac[ietaHcal] = (hcalUpper - emcLower) / (emcUpper - emcLower); + foundUpper = true; + } + if (emcUpper > hcalUpper && emcLower > hcalUpper) + { + --ietaEmc; + m_retowerUpperEta[ietaHcal] = ietaEmc; + m_retowerUpperFrac[ietaHcal] = 1.; + foundUpper = true; + } + } + + if (!(foundLower && foundUpper)) + { + ++ietaEmc; + } + } + + if (!foundLower || !foundUpper) + { + if (!m_warnedRetowerMap) + { + std::cout << PHWHERE << "cannot build EMCal retower map, missing " << (foundLower ? "upper" : "lower") << " eta overlap for HCal ieta " << ietaHcal << std::endl; + m_warnedRetowerMap = true; + } + return -1; + } + } + + m_retowerMapReady = true; + return 0; +} diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCNN.h b/offline/packages/calovtxreco/CaloVtxAlgoCNN.h new file mode 100644 index 0000000000..c1770de5b3 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoCNN.h @@ -0,0 +1,70 @@ +#ifndef CALOVTXALGOCNN_H +#define CALOVTXALGOCNN_H + +#include "CaloVtxAlgo.h" + +#include +#include +#include + +class PHCompositeNode; + +// Calo-image CNN (ONNX) as a CaloVtxAlgo. Input is the 3x24x64 tower-energy image; preprocessing and calibration are inside the graph. +class CaloVtxAlgoCNN : public CaloVtxAlgo +{ + public: + static constexpr int kNLayer = 3; + static constexpr int kNEtaImg = 24; + static constexpr int kNPhiImg = 64; + static constexpr int kNEtaEmcFine = 96; + static constexpr int kNPhiEmcFine = 256; + enum Layer + { + kEMC = 0, + kIHC = 1, + kOHC = 2 + }; + + explicit CaloVtxAlgoCNN() = default; + ~CaloVtxAlgoCNN() override; + + int Init(PHCompositeNode *topNode) override; + int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; + std::string Name() const override { return "CNN"; } + VertexDefs::CALOALGO Algo() const override { return VertexDefs::CALOALGO::CNN; } + + void setModelFile(const std::string &path) { m_modelFile = path; } + void setTowerEMin(Layer layer, float e) { m_towerEMin.at(layer) = e; } // [GeV] + void setMinTotalEnergy(float e) { m_minTotalEnergy = e; } // [GeV] + + private: + struct OnnxSession; // pImpl, defined in the .cc + + int fillTowerImage(PHCompositeNode *topNode); + int fillHcalLayer(PHCompositeNode *topNode, int layer); + int fillEmcRetower(PHCompositeNode *topNode); + int buildEmcRetowerMap(PHCompositeNode *topNode); + bool predict(float &z); + + std::string m_modelFile{"vertex_cnn.onnx"}; + std::array m_towerEMin{{0.068, 0.005, 0.035}}; + float m_minTotalEnergy{0.}; + + std::unique_ptr m_onnx; + + float m_image[kNLayer][kNEtaImg][kNPhiImg]{}; + double m_rawEmcFine[kNEtaEmcFine][kNPhiEmcFine]{}; + + bool m_retowerMapReady{false}; + int m_retowerPhiOffset{-1}; + std::array m_retowerLowerEta{}; + std::array m_retowerUpperEta{}; + std::array m_retowerLowerFrac{}; + std::array m_retowerUpperFrac{}; + + bool m_warnedMissingTowers{false}; + bool m_warnedRetowerMap{false}; + bool m_warnedPredict{false}; +}; + +#endif diff --git a/offline/packages/calovtxreco/Makefile.am b/offline/packages/calovtxreco/Makefile.am index 87c4d0f429..9a7ba72609 100644 --- a/offline/packages/calovtxreco/Makefile.am +++ b/offline/packages/calovtxreco/Makefile.am @@ -15,6 +15,7 @@ pkginclude_HEADERS = \ CaloVtxAlgoJetSkew.h \ CaloVtxAlgoCaloZ.h \ CaloVtxAlgoMLP.h \ + CaloVtxAlgoCNN.h \ VertexMLP.h lib_LTLIBRARIES = \ @@ -24,13 +25,15 @@ libcalovtxreco_la_SOURCES = \ CaloVtxReco.cc \ CaloVtxAlgoJetSkew.cc \ CaloVtxAlgoCaloZ.cc \ - CaloVtxAlgoMLP.cc + CaloVtxAlgoMLP.cc \ + CaloVtxAlgoCNN.cc libcalovtxreco_la_LIBADD = \ -lcalo_io \ -lSubsysReco \ -ljetbase \ - -lglobalvertex_io + -lglobalvertex_io \ + -lonnxruntime BUILT_SOURCES = testexternals.cc diff --git a/offline/packages/globalvertex/GlobalVertex.h b/offline/packages/globalvertex/GlobalVertex.h index 197ae3a3b3..81f28450cb 100644 --- a/offline/packages/globalvertex/GlobalVertex.h +++ b/offline/packages/globalvertex/GlobalVertex.h @@ -28,6 +28,7 @@ class GlobalVertex : public PHObject CALO_JETSKEW = 251, CALO_AVGZ = 252, CALO_JETMLP = 253, + CALO_CNN = 254, MBD_CALO = 350 }; diff --git a/offline/packages/globalvertex/GlobalVertexReco.cc b/offline/packages/globalvertex/GlobalVertexReco.cc index d233f0c428..dd2b74394f 100644 --- a/offline/packages/globalvertex/GlobalVertexReco.cc +++ b/offline/packages/globalvertex/GlobalVertexReco.cc @@ -303,7 +303,11 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) { vertex->clone_insert_vtx(GlobalVertex::CALO_JETMLP, calo); } - vertex->set_id(globalmap->size()); + if (caloalgo == VertexDefs::CALOALGO::CNN) + { + vertex->clone_insert_vtx(GlobalVertex::CALO_CNN, calo); + } + vertex->set_id(globalmap->size()); used_calo_vtxids.insert(calo->get_id()); diff --git a/offline/packages/globalvertex/GlobalVertexv4.cc b/offline/packages/globalvertex/GlobalVertexv4.cc index d6e201fe35..f304c3818d 100644 --- a/offline/packages/globalvertex/GlobalVertexv4.cc +++ b/offline/packages/globalvertex/GlobalVertexv4.cc @@ -144,6 +144,12 @@ float GlobalVertexv4::get_position(unsigned int coor) const } caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_JETMLP); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + + caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_CNN); if (caloit != _vtxs.end()) { return caloit->second[0]->get_position(coor); diff --git a/offline/packages/globalvertex/VertexDefs.h b/offline/packages/globalvertex/VertexDefs.h index 5143fec3d8..f312205032 100644 --- a/offline/packages/globalvertex/VertexDefs.h +++ b/offline/packages/globalvertex/VertexDefs.h @@ -15,7 +15,8 @@ namespace VertexDefs UNDEFINED=0, JETSKEW=1, AVGZ=2, - JETMLP=3 + JETMLP=3, + CNN=4 }; }; #endif