diff --git a/offline/packages/calovtxreco/CaloVtxAlgoVit.cc b/offline/packages/calovtxreco/CaloVtxAlgoVit.cc new file mode 100644 index 0000000000..2e5f0fa7e9 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoVit.cc @@ -0,0 +1,182 @@ +#include "CaloVtxAlgoVit.h" + +#include + +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include + +// onnxruntime session (pImpl, keeps Ort types out of the header) +struct CaloVtxAlgoVit::OnnxSession +{ + Ort::Env env{ORT_LOGGING_LEVEL_WARNING, "CaloVtxAlgoVit"}; + Ort::SessionOptions opts; + std::unique_ptr session; + Ort::MemoryInfo memInfo{ + Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault)}; + + explicit OnnxSession(const std::string &path) + { + // pin to one compute thread: on shared (Condor) nodes onnxruntime must + // not spawn one thread per node core + opts.SetIntraOpNumThreads(1); + opts.SetInterOpNumThreads(1); + opts.SetGraphOptimizationLevel(ORT_ENABLE_ALL); + session = std::make_unique(env, path.c_str(), opts); + } +}; + +CaloVtxAlgoVit::CaloVtxAlgoVit() = default; +CaloVtxAlgoVit::~CaloVtxAlgoVit() = default; + +int CaloVtxAlgoVit::Init(PHCompositeNode * /*topNode*/) +{ + for (int calo = 0; calo < kNCalo; ++calo) + { + m_input[calo].assign(static_cast(kNChan) * kNEta[calo] * kNPhi[calo], 0.); + } + + 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; + } + + const size_t nIn = m_onnx->session->GetInputCount(); + if (nIn != kNCalo) + { + std::cout << PHWHERE << "model " << m_modelFile << " has " << nIn + << " inputs, expected " << kNCalo << " (emcal, ihcal, ohcal)" << std::endl; + return Fun4AllReturnCodes::ABORTRUN; + } + return Fun4AllReturnCodes::EVENT_OK; +} + +int CaloVtxAlgoVit::CalculateVertex(PHCompositeNode *topNode, float &zvtx) +{ + zvtx = std::numeric_limits::quiet_NaN(); + + if (fillInputs(topNode) != 0) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + if (m_etot <= m_minTotalEnergy) + { + // empty event: 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; +} + +int CaloVtxAlgoVit::fillInputs(PHCompositeNode *topNode) +{ + m_etot = 0.; + for (int calo = 0; calo < kNCalo; ++calo) + { + std::fill(m_input[calo].begin(), m_input[calo].end(), 0.); + if (fillCalo(topNode, calo) != 0) + { + return -1; + } + } + return 0; +} + +int CaloVtxAlgoVit::fillCalo(PHCompositeNode *topNode, int calo) +{ + TowerInfoContainer *towers = findNode::getClass(topNode, m_towerNode[calo]); + if (!towers) + { + if (!m_warnedMissingTowers) + { + std::cout << PHWHERE << "tower node missing: " << m_towerNode[calo] << std::endl; + m_warnedMissingTowers = true; + } + return -1; + } + + const int nEta = kNEta[calo]; + const int nPhi = kNPhi[calo]; + float *eChan = m_input[calo].data(); // channel 0: raw energy + float *tChan = m_input[calo].data() + static_cast(nEta) * nPhi; // channel 1: time + + const unsigned int ntow = towers->size(); + for (unsigned int ch = 0; ch < ntow; ++ch) + { + TowerInfo *tower = towers->get_tower_at_channel(ch); + if (!tower) + { + continue; + } + if (m_useGoodTowersOnly && !tower->get_isGood()) + { + 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 >= nEta || iphi < 0 || iphi >= nPhi) + { + continue; + } + // raw energy floored at 0, NO log1p (the graph applies it); keep the time + const float e = std::max(tower->get_energy(), 0.f); + const size_t idx = static_cast(ieta) * nPhi + iphi; + eChan[idx] = e; + tChan[idx] = tower->get_time(); + m_etot += e; + } + return 0; +} + +bool CaloVtxAlgoVit::predict(float &z) +{ + try + { + std::vector inputs; + inputs.reserve(kNCalo); + const char *inNames[kNCalo]; + for (int calo = 0; calo < kNCalo; ++calo) + { + const int64_t shape[4] = {1, kNChan, kNEta[calo], kNPhi[calo]}; + inputs.push_back(Ort::Value::CreateTensor( + m_onnx->memInfo, m_input[calo].data(), m_input[calo].size(), shape, 4)); + inNames[calo] = m_inputName[calo].c_str(); + } + const char *outNames[] = {m_outputName.c_str()}; + auto outs = m_onnx->session->Run(Ort::RunOptions{nullptr}, inNames, inputs.data(), kNCalo, 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; + } +} diff --git a/offline/packages/calovtxreco/CaloVtxAlgoVit.h b/offline/packages/calovtxreco/CaloVtxAlgoVit.h new file mode 100644 index 0000000000..315d6d8d3b --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoVit.h @@ -0,0 +1,75 @@ +#ifndef CALOVTXALGOVIT_H +#define CALOVTXALGOVIT_H + +#include "CaloVtxAlgo.h" + +#include +#include +#include +#include + +class PHCompositeNode; + +// Calo vision-transformer (CaloViTv2_gate, ONNX) as a CaloVtxAlgo. +// Inputs are three 2-channel tower images at native tower granularity: +// emcal (1,2,96,256), ihcal (1,2,24,64), ohcal (1,2,24,64) +// channel 0 = raw tower energy (negative energies floored at 0, NO log1p -- +// the graph applies it), channel 1 = tower time. Output is pred_z [cm]. +// The EMCal is used at fine 96x256 granularity, so no retowering is needed. +class CaloVtxAlgoVit : public CaloVtxAlgo +{ + public: + static constexpr int kNCalo = 3; + static constexpr int kNChan = 2; // 0 = energy, 1 = time + enum Calo + { + kEMC = 0, + kIHC = 1, + kOHC = 2 + }; + static constexpr std::array kNEta{{96, 24, 24}}; + static constexpr std::array kNPhi{{256, 64, 64}}; + + // ctor/dtor defined in the .cc, where OnnxSession is a complete type + // (required for the std::unique_ptr pImpl member, notably under cling) + CaloVtxAlgoVit(); + ~CaloVtxAlgoVit() override; + + int Init(PHCompositeNode *topNode) override; + int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; + std::string Name() const override { return "ViT"; } + VertexDefs::CALOALGO Algo() const override { return VertexDefs::CALOALGO::VIT; } + + void setModelFile(const std::string &path) { m_modelFile = path; } + void setTowerNode(Calo calo, const std::string &node) { m_towerNode.at(calo) = node; } + // tensor names in the exported graph (defaults match the CaloViTv2_gate export) + void setInputName(Calo calo, const std::string &name) { m_inputName.at(calo) = name; } + void setOutputName(const std::string &name) { m_outputName = name; } + void setMinTotalEnergy(float e) { m_minTotalEnergy = e; } // [GeV] + void setUseGoodTowersOnly(bool b) { m_useGoodTowersOnly = b; } + + private: + struct OnnxSession; // pImpl, defined in the .cc + + int fillInputs(PHCompositeNode *topNode); + int fillCalo(PHCompositeNode *topNode, int calo); + bool predict(float &z); + + std::string m_modelFile{"calovit.onnx"}; + std::array m_towerNode{{"TOWERINFO_CALIB_CEMC", "TOWERINFO_CALIB_HCALIN", "TOWERINFO_CALIB_HCALOUT"}}; + std::array m_inputName{{"emcal", "ihcal", "ohcal"}}; + std::string m_outputName{"pred_z"}; + float m_minTotalEnergy{0.}; + bool m_useGoodTowersOnly{true}; + + std::unique_ptr m_onnx; + + // flat (1,2,eta,phi) input tensors, one per calorimeter + std::array, kNCalo> m_input; + double m_etot{0.}; + + bool m_warnedMissingTowers{false}; + bool m_warnedPredict{false}; +}; + +#endif diff --git a/offline/packages/calovtxreco/Makefile.am b/offline/packages/calovtxreco/Makefile.am index 9a7ba72609..2e15a94254 100644 --- a/offline/packages/calovtxreco/Makefile.am +++ b/offline/packages/calovtxreco/Makefile.am @@ -16,6 +16,7 @@ pkginclude_HEADERS = \ CaloVtxAlgoCaloZ.h \ CaloVtxAlgoMLP.h \ CaloVtxAlgoCNN.h \ + CaloVtxAlgoVit.h \ VertexMLP.h lib_LTLIBRARIES = \ @@ -26,7 +27,8 @@ libcalovtxreco_la_SOURCES = \ CaloVtxAlgoJetSkew.cc \ CaloVtxAlgoCaloZ.cc \ CaloVtxAlgoMLP.cc \ - CaloVtxAlgoCNN.cc + CaloVtxAlgoCNN.cc \ + CaloVtxAlgoVit.cc libcalovtxreco_la_LIBADD = \ -lcalo_io \ diff --git a/offline/packages/globalvertex/GlobalVertex.h b/offline/packages/globalvertex/GlobalVertex.h index 81f28450cb..f1469c9beb 100644 --- a/offline/packages/globalvertex/GlobalVertex.h +++ b/offline/packages/globalvertex/GlobalVertex.h @@ -29,6 +29,7 @@ class GlobalVertex : public PHObject CALO_AVGZ = 252, CALO_JETMLP = 253, CALO_CNN = 254, + CALO_VIT = 255, MBD_CALO = 350 }; diff --git a/offline/packages/globalvertex/GlobalVertexReco.cc b/offline/packages/globalvertex/GlobalVertexReco.cc index dd2b74394f..ab89d473bd 100644 --- a/offline/packages/globalvertex/GlobalVertexReco.cc +++ b/offline/packages/globalvertex/GlobalVertexReco.cc @@ -307,6 +307,10 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) { vertex->clone_insert_vtx(GlobalVertex::CALO_CNN, calo); } + if (caloalgo == VertexDefs::CALOALGO::VIT) + { + vertex->clone_insert_vtx(GlobalVertex::CALO_VIT, 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 f304c3818d..e5354b84cb 100644 --- a/offline/packages/globalvertex/GlobalVertexv4.cc +++ b/offline/packages/globalvertex/GlobalVertexv4.cc @@ -150,6 +150,11 @@ float GlobalVertexv4::get_position(unsigned int coor) const } caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_CNN); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_VIT); 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 f312205032..849e8f96c6 100644 --- a/offline/packages/globalvertex/VertexDefs.h +++ b/offline/packages/globalvertex/VertexDefs.h @@ -16,7 +16,8 @@ namespace VertexDefs JETSKEW=1, AVGZ=2, JETMLP=3, - CNN=4 + CNN=4, + VIT=5 }; }; #endif