Add CaloVtxAlgoVit method

offline/packages/calovtxreco/CaloVtxAlgoVit.cc

@@ -0,0 +1,182 @@
+#include "CaloVtxAlgoVit.h"
+
+#include <fun4all/Fun4AllReturnCodes.h>
+
+#include <calobase/TowerInfo.h>
+#include <calobase/TowerInfoContainer.h>
+
+#include <phool/PHCompositeNode.h>
+#include <phool/getClass.h>
+#include <phool/phool.h>
+
+#include <onnxruntime_cxx_api.h>
+
+#include <algorithm>
+#include <cmath>
+#include <iostream>
+#include <limits>
+
+// 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<Ort::Session> 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<Ort::Session>(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<size_t>(kNChan) * kNEta[calo] * kNPhi[calo], 0.);
+  }
+
+  try
+  {
+    m_onnx = std::make_unique<OnnxSession>(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<float>::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<float>::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<TowerInfoContainer>(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<size_t>(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<size_t>(ieta) * nPhi + iphi;
+    eChan[idx] = e;
+    tChan[idx] = tower->get_time();
+    m_etot += e;
+  }
+  return 0;
+}
+
+bool CaloVtxAlgoVit::predict(float &z)
+{
+  try
+  {
+    std::vector<Ort::Value> 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<float>(
+          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<float>()[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;
+  }
+}

offline/packages/calovtxreco/CaloVtxAlgoVit.h

@@ -0,0 +1,75 @@
+#ifndef CALOVTXALGOVIT_H
+#define CALOVTXALGOVIT_H
+
+#include "CaloVtxAlgo.h"
+
+#include <array>
+#include <memory>
+#include <string>
+#include <vector>
+
+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<int, kNCalo> kNEta{{96, 24, 24}};
+  static constexpr std::array<int, kNCalo> 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<std::string, kNCalo> m_towerNode{{"TOWERINFO_CALIB_CEMC", "TOWERINFO_CALIB_HCALIN", "TOWERINFO_CALIB_HCALOUT"}};
+  std::array<std::string, kNCalo> m_inputName{{"emcal", "ihcal", "ohcal"}};
+  std::string m_outputName{"pred_z"};
+  float m_minTotalEnergy{0.};
+  bool m_useGoodTowersOnly{true};
+
+  std::unique_ptr<OnnxSession> m_onnx;
+
+  // flat (1,2,eta,phi) input tensors, one per calorimeter
+  std::array<std::vector<float>, kNCalo> m_input;
+  double m_etot{0.};
+
+  bool m_warnedMissingTowers{false};
+  bool m_warnedPredict{false};
+};
+
+#endif

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 \

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
   };
 

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());

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);

offline/packages/globalvertex/VertexDefs.h

@@ -16,7 +16,8 @@ namespace VertexDefs
       JETSKEW=1,
       AVGZ=2,
       JETMLP=3,
-      CNN=4
+      CNN=4,
+      VIT=5
     };
 };
 #endif