From 4a2a90301689d88919083324001df4e50b811b7e Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Wed, 8 Jul 2026 23:08:26 -0400 Subject: [PATCH 1/8] djl -- refactor to better take in new algorithms --- offline/packages/calovtxreco/CaloVtxAlgo.h | 31 ++ .../packages/calovtxreco/CaloVtxAlgoCaloZ.cc | 146 +++++++ .../packages/calovtxreco/CaloVtxAlgoCaloZ.h | 30 ++ .../calovtxreco/CaloVtxAlgoJetSkew.cc | 185 +++++++++ .../packages/calovtxreco/CaloVtxAlgoJetSkew.h | 54 +++ .../packages/calovtxreco/CaloVtxAlgoMLP.cc | 331 ++++++++++++++++ offline/packages/calovtxreco/CaloVtxAlgoMLP.h | 34 ++ offline/packages/calovtxreco/CaloVtxReco.cc | 369 +----------------- offline/packages/calovtxreco/CaloVtxReco.h | 48 +-- offline/packages/calovtxreco/Makefile.am | 12 +- offline/packages/calovtxreco/VertexMLP.h | 143 +++++++ offline/packages/globalvertex/GlobalVertex.h | 1 + .../packages/globalvertex/GlobalVertexReco.cc | 28 +- 13 files changed, 1028 insertions(+), 384 deletions(-) create mode 100644 offline/packages/calovtxreco/CaloVtxAlgo.h create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoMLP.cc create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoMLP.h create mode 100644 offline/packages/calovtxreco/VertexMLP.h diff --git a/offline/packages/calovtxreco/CaloVtxAlgo.h b/offline/packages/calovtxreco/CaloVtxAlgo.h new file mode 100644 index 0000000000..4b6341ea15 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgo.h @@ -0,0 +1,31 @@ +#ifndef CALOVTXALGO_H +#define CALOVTXALGO_H + +#include + +class PHCompositeNode; + +class CaloVtxAlgo +{ + public: + virtual ~CaloVtxAlgo() = default; + + // Where set-up occurs: + // Calibrations/ML weights etc + // Basic information from geometry + + virtual int Init(PHCompositeNode * /*topNode*/) { return 0; } + + // Called once per event. Read whatever input this algorithm needs from + // topNode and fill vtxz [cm]. Return nonzero (e.g. + // Fun4AllReturnCodes::ABORTEVENT) if no vertex could be found. + + virtual int CalculateVertex(PHCompositeNode *topNode, float &zvtx) = 0; + + // Short identifier used to name this algorithm's output node + // (e.g. "CALOVTXOUT_MLP") and in log/eval output so results from + // different algorithms don't collide and can be compared directly. + virtual std::string Name() const = 0; +}; + +#endif diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc new file mode 100644 index 0000000000..9f33f9f8e7 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc @@ -0,0 +1,146 @@ +#include "CaloVtxAlgoCaloZ.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +CaloVtxAlgoCaloZ::CaloVtxAlgoCaloZ() +{ + +} + +int CaloVtxAlgoCaloZ::Init(PHCompositeNode * /*topNode*/) +{ + + return Fun4AllReturnCodes::EVENT_OK; +} + +int CaloVtxAlgoCaloZ::CalculateVertex(PHCompositeNode *topNode, float &zvtx) +{ + zvtx = std::numeric_limits::quiet_NaN(); + + + TowerInfoContainer *emcal_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC"); + TowerInfoContainer *hcalin_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); + TowerInfoContainer *hcalout_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); + RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + RawTowerGeomContainer *tower_geomIH = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); + RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + + if (!emcal_towers || !hcalin_towers || !hcalout_towers) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + if (!tower_geomEM || !tower_geomIH || !tower_geomOH) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + int size; + + float average_z[3]{0}; + float total_E[3]{0}; + + if (emcal_towers) + { + size = emcal_towers->size(); // online towers should be the same! + for (int channel = 0; channel < size; channel++) + { + TowerInfo *_tower = emcal_towers->get_tower_at_channel(channel); + short good = (_tower->get_isGood() ? 1 : 0); + if (!good) + { + continue; + } + + float energy = _tower->get_energy(); + if (energy < m_energy_cut) + { + continue; + } + + // float time = _tower->get_time_float(); + + unsigned int towerkey = emcal_towers->encode_key(channel); + int ieta = emcal_towers->getTowerEtaBin(towerkey); + int iphi = emcal_towers->getTowerPhiBin(towerkey); + + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::CEMC, ieta, iphi); + float tower_z = tower_geomEM->get_tower_geometry(key)->get_center_z(); + average_z[0] += tower_z * energy; + total_E[0] += energy; + } + } + + if (hcalin_towers) + { + size = hcalin_towers->size(); // online towers should be the same! + for (int channel = 0; channel < size; channel++) + { + TowerInfo *_tower = hcalin_towers->get_tower_at_channel(channel); + float energy = _tower->get_energy(); + if (energy < m_energy_cut) + { + continue; + } + // float time = _tower->get_time_float(); + short good = (_tower->get_isGood() ? 1 : 0); + if (!good) + { + continue; + } + + unsigned int towerkey = hcalin_towers->encode_key(channel); + int ieta = hcalin_towers->getTowerEtaBin(towerkey); + int iphi = hcalin_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); + float tower_z = tower_geomIH->get_tower_geometry(key)->get_center_z(); + average_z[1] += tower_z * energy; + total_E[1] += energy; + } + } + if (hcalout_towers) + { + size = hcalout_towers->size(); // online towers should be the same! + for (int channel = 0; channel < size; channel++) + { + TowerInfo *_tower = hcalout_towers->get_tower_at_channel(channel); + float energy = _tower->get_energy(); + if (energy < m_energy_cut) + { + continue; + } + // float time = _tower->get_time_float(); + unsigned int towerkey = hcalout_towers->encode_key(channel); + int ieta = hcalout_towers->getTowerEtaBin(towerkey); + int iphi = hcalout_towers->getTowerPhiBin(towerkey); + short good = (_tower->get_isGood() ? 1 : 0); + + if (!good) + { + continue; + } + + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); + + float tower_z = tower_geomOH->get_tower_geometry(key)->get_center_z(); + average_z[2] += tower_z * energy; + total_E[2] += energy; + } + } + + double b_calo_vertex_z = (average_z[0] + average_z[1] + average_z[2]) / (total_E[0] + total_E[1] + total_E[2]); + + zvtx = b_calo_vertex_z; + + return Fun4AllReturnCodes::EVENT_OK; +} diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h new file mode 100644 index 0000000000..f4d77e403e --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h @@ -0,0 +1,30 @@ +#ifndef CALOVTXALGOCaloZ_H +#define CALOVTXALGOCaloZ_H + +#include "CaloVtxAlgo.h" +#include + +class PHCompositeNode; + +// Wraps the trained CaloZ (VertexCaloZ.h / vertex_mlp_weights.root) as a +// CaloVtxAlgo so it can be registered alongside other algorithms in +// CaloVtxReco and compared directly. +class CaloVtxAlgoCaloZ : public CaloVtxAlgo +{ + public: + explicit CaloVtxAlgoCaloZ(); + ~CaloVtxAlgoCaloZ() override = default; + + int Init(PHCompositeNode *topNode) override; + int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; + std::string Name() const override { return "CaloZ"; } + + float get_energy_cut() { return m_energy_cut; } + void set_energy_cut(float new_energy) { m_energy_cut = new_energy; } + + private: + + float m_energy_cut{0.1}; +}; + +#endif diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc new file mode 100644 index 0000000000..2c1ed616f7 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc @@ -0,0 +1,185 @@ +#include "CaloVtxAlgoJetSkew.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CaloVtxAlgoJetSkew::CaloVtxAlgoJetSkew() +{ +} + +int CaloVtxAlgoJetSkew::Init(PHCompositeNode *topNode) +{ + RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + if(!tower_geomEM || !tower_geomOH) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + m_radius_EM = tower_geomEM->get_radius(); + m_radius_OH = tower_geomOH->get_radius(); + + if(std::isnan(m_radius_EM) || std::isnan(m_radius_OH)) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + return Fun4AllReturnCodes::EVENT_OK; +} + +int CaloVtxAlgoJetSkew::CalculateVertex(PHCompositeNode *topNode, float &zvtx) +{ + + zvtx = std::numeric_limits::quiet_NaN(); + + JetContainer *jetcon = findNode::getClass(topNode, m_jetnodename); + + towers[0] = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC_RETOWER"); + towers[1] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); + towers[2] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); + + geom[0] = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + geom[1] = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); + geom[2] = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + + const int nz = 601; + const int njet = 2; + Jet *jets[njet]; + float jpt[njet] = {0}; + float jemsum[njet] = {0}; + float johsum[njet] = {0}; + float jemeta[njet] = {0}; + float joheta[njet] = {0}; + + if (jetcon) + { + int tocheck = jetcon->size(); + for (int i = 0; i < tocheck; ++i) + { + Jet *jet = jetcon->get_jet(i); + if (jet) + { + float pt = jet->get_pt(); + if (pt < m_jet_threshold) + { + continue; + } + if (pt > jpt[0]) + { + jpt[1] = jpt[0]; + jets[1] = jets[0]; + jets[0] = jet; + jpt[0] = pt; + } + else if (pt > jpt[1]) + { + jets[1] = jet; + jpt[1] = pt; + } + } + } + } + else + { + return Fun4AllReturnCodes::EVENT_OK; + } + + float metric = std::numeric_limits::max(); + for (int i = 0; i < nz; ++i) + { + float testz = -300 + i; + float testmetric = 0; + for (int j = 0; j < njet; ++j) + { + if (jpt[j] == 0) + { + continue; + } + jemsum[j] = 0; + johsum[j] = 0; + jemeta[j] = 0; + joheta[j] = 0; + + for (auto comp : jets[j]->get_comp_vec()) + { + if (comp.first == 5 || comp.first == 26) + { + continue; + } + unsigned int channel = comp.second; + if (comp.first == 7 || comp.first == 27) + { + TowerInfo *tower = towers[2]->get_tower_at_channel(channel); + if (tower->get_energy() < 0.1) + { + continue; + } + johsum[j] += tower->get_energy(); + float neweta = new_eta(channel, towers[2], geom[2], RawTowerDefs::CalorimeterId::HCALOUT, testz); + joheta[j] += neweta * tower->get_energy(); + } + if (comp.first == 13 || comp.first == 28 || comp.first == 25) + { + TowerInfo *tower = towers[0]->get_tower_at_channel(channel); + if (tower->get_energy() < 0.1) + { + continue; + } + jemsum[j] += tower->get_energy(); + float neweta = new_eta(channel, towers[0], geom[1], RawTowerDefs::CalorimeterId::HCALIN, testz); + jemeta[j] += neweta * tower->get_energy(); + } + } + + jemeta[j] /= jemsum[j]; + joheta[j] /= johsum[j]; + if (!std::isnan(jemeta[j]) && !std::isnan(joheta[j])) + { + testmetric += pow(jemeta[j] - joheta[j], 2); + } + } + if (testmetric < metric && testmetric != 0) + { + metric = testmetric; + zvtx = testz; + } + } + if (fabs(zvtx) >= 305) + { + zvtx = std::numeric_limits::quiet_NaN(); + } + else + { + zvtx *= m_calib_factor; // calibration factor from simulation + } + + return Fun4AllReturnCodes::EVENT_OK; + +} + +float CaloVtxAlgoJetSkew::new_eta(int channel, TowerInfoContainer *towerset, RawTowerGeomContainer *geom, RawTowerDefs::CalorimeterId caloID, float testz) +{ + int key = towerset->encode_key(channel); + const RawTowerDefs::keytype geomkey = RawTowerDefs::encode_towerid(caloID, towerset->getTowerEtaBin(key), towerset->getTowerPhiBin(key)); + + RawTowerGeom *tower_geom = geom->get_tower_geometry(geomkey); + float oldeta = tower_geom->get_eta(); + + float radius = (caloID == RawTowerDefs::CalorimeterId::HCALIN ? m_radius_EM : m_radius_OH); + float towerz = radius / (tanf(2 * std::atanf(std::exp(oldeta)))); + float newz = towerz + testz; + float newTheta = std::atan2(radius, newz); + float neweta = -log(tan(0.5 * newTheta)); + + return neweta; +} diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h new file mode 100644 index 0000000000..819d5d8543 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h @@ -0,0 +1,54 @@ +#ifndef CALOVTXALGOJetSkew_H +#define CALOVTXALGOJetSkew_H + +#include "CaloVtxAlgo.h" +#include +#include +#include "TMath.h" +#include +#include + + +class PHCompositeNode; +class RawTowerGeomContainer; +class TowerInfoContainer; + +// Wraps the trained JetSkew (VertexJetSkew.h / vertex_mlp_weights.root) as a +// CaloVtxAlgo so it can be registered alongside other algorithms in +// CaloVtxReco and compared directly. +class CaloVtxAlgoJetSkew : public CaloVtxAlgo +{ + public: + explicit CaloVtxAlgoJetSkew(); + ~CaloVtxAlgoJetSkew() override = default; + + int Init(PHCompositeNode *topNode) override; + int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; + std::string Name() const override { return "JetSkew"; } + + void setJetNode(std::string jetnode) { m_jetnodename = jetnode; } + float get_jet_threshold() { return m_jet_threshold; } + void set_jet_threshold(float new_thresh) { m_jet_threshold = new_thresh; } + float get_calib_factor() { return m_calib_factor; } + void set_calib_factor(float new_calib) { m_calib_factor = new_calib; } + float get_energy_cut() { return m_energy_cut; } + void set_energy_cut(float new_energy) { m_energy_cut = new_energy; } + + private: + + float new_eta(int channel, TowerInfoContainer *towerset, RawTowerGeomContainer *geom, RawTowerDefs::CalorimeterId caloID, float testz); + + std::string m_jetnodename; + float m_jet_threshold{15}; + float m_calib_factor{1.406}; + float m_energy_cut{0.1}; + float m_radius_EM{std::numeric_limits::quiet_NaN()}; + float m_radius_OH{std::numeric_limits::quiet_NaN()}; + + TowerInfoContainer *towers[3] = {0}; + RawTowerGeomContainer *geom[3] = {0}; + + +}; + +#endif diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc new file mode 100644 index 0000000000..918c319b19 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -0,0 +1,331 @@ +#include "CaloVtxAlgoMLP.h" +#include "VertexMLP.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CaloVtxAlgoMLP::CaloVtxAlgoMLP() +{ +} + +int CaloVtxAlgoMLP::Init(PHCompositeNode *topNode) +{ + if (!VertexMLP::Load(m_weightsFile)) + { + return Fun4AllReturnCodes::ABORTRUN; + } + RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + if(!tower_geomEM || !tower_geomOH) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + m_radius_EM = tower_geomEM->get_radius(); + m_radius_OH = tower_geomOH->get_radius(); + + if(std::isnan(m_radius_EM) || std::isnan(m_radius_OH)) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + return Fun4AllReturnCodes::EVENT_OK; +} + +int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) +{ + TowerInfoContainer *emcalre_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC_RETOWER"); + TowerInfoContainer *hcalin_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); + TowerInfoContainer *hcalout_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); + RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + RawTowerGeomContainer *tower_geomIH = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); + RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + + if (!emcalre_towers || !hcalin_towers || !hcalout_towers) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + if (!tower_geomEM || !tower_geomIH || !tower_geomOH) + { + return Fun4AllReturnCodes::ABORTRUN; + } + + // TODO: fill `features` from topNode in the exact order documented in + // VertexMLP.h (emcal/ohcal lead+sublead flags, zmean/zsig/zskew, + // energy, exj). This is the same per-event extraction CaloVtxReco used + // to do inline before the NN call -- move that block here unchanged. + + JetContainer *jetscon = findNode::getClass(topNode, m_jet_node.c_str()); + + float dijet_pt[2]={0}; + float dijet_eta[2]={0}; + float dijet_phi[2]={0}; + float dijet_E[2]={0}; + float dijet_em_frac[2]={0}; + float dijet_em_mean[2]={0}; + float dijet_em_skew[2]={0}; + float dijet_em_sig[2]={0}; + float dijet_oh_frac[2]={0}; + float dijet_oh_mean[2]={0}; + float dijet_oh_skew[2]={0}; + float dijet_oh_sig[2]={0}; + + if (!jetscon) + { + return Fun4AllReturnCodes::EVENT_OK; + } + for (auto jet : *jetscon) + { + double jet_E = 0; + double jet_emcal = 0; + double jet_ohcal = 0; + double jetpt = jet->get_pt(); + double jetphi = jet->get_phi(); + double jeteta= jet->get_eta(); + + if (jetpt < 0.5) continue; + + double em_S0 = 0.0; + double em_S1 = 0.0; + double em_S2 = 0.0; + double em_S3 = 0.0; + double oh_S0 = 0.0; + double oh_S1 = 0.0; + double oh_S2 = 0.0; + double oh_S3 = 0.0; + + + int itower = 0; + for (auto comp : jet->get_comp_vec()) + { + unsigned int channel = comp.second; + TowerInfo *tower; + float tower_e = 0; + if (comp.first == 26 || comp.first == 30) + { + tower = hcalin_towers->get_tower_at_channel(channel); + + if (!tower || !tower_geomIH) + { + continue; + } + + tower_e = tower->get_energy(); + if (tower_e < 0.005) continue; + jet_E += tower_e; + } + else if (comp.first == 27 || comp.first == 31) + { + // OHCAL + tower = hcalout_towers->get_tower_at_channel(channel); + if (!tower || !tower_geomOH) + { + continue; + } + unsigned int towerkey = hcalout_towers->encode_key(channel); + int ieta = hcalout_towers->getTowerEtaBin(towerkey); + int iphi = hcalout_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); + float tower_eta = tower_geomOH->get_tower_geometry(key)->get_eta(); + float tower_r = m_radius_OH; + float tower_z = tower_r*sinh(tower_eta); + + tower_e = tower->get_energy(); + if (tower_e < 0.035) continue; + jet_ohcal += tower_e; + oh_S0 += tower_e; + oh_S1 += tower_e*tower_z; + oh_S2 += tower_e*tower_z*tower_z; + oh_S3 += tower_e*tower_z*tower_z*tower_z; + jet_E += tower_e; + } + else if (comp.first == 28 || comp.first == 29) + { // EMCAL + tower = emcalre_towers->get_tower_at_channel(channel); + + if (!tower || !tower_geomOH) + { + continue; + } + unsigned int towerkey = emcalre_towers->encode_key(channel); + int ieta = emcalre_towers->getTowerEtaBin(towerkey); + int iphi = emcalre_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); + float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); + float tower_r = m_radius_EM; + float tower_z = tower_r*sinh(tower_eta); + + tower_e = tower->get_energy(); + if (tower_e < 0.068) continue; + jet_emcal += tower_e; + jet_E += tower_e; + em_S0 += tower_e; + em_S1 += tower_e*tower_z; + em_S2 += tower_e*tower_z*tower_z; + em_S3 += tower_e*tower_z*tower_z*tower_z; + } + + itower++; + } + + if (itower == 0) continue; + + double em_mean = em_S1/em_S0; + + double em_variance = em_S2/em_S0 - em_mean*em_mean; + double em_sigma = std::sqrt(std::max(0.0, em_variance)); + + double em_mu3 = + em_S3/em_S0 + - 3.0*em_mean*(em_S2/em_S0) + + 2.0*em_mean*em_mean*em_mean; + + double em_skew = (em_sigma > 0.0) + ? em_mu3/std::pow(em_sigma,3) + : 0.0; + + double oh_mean = oh_S1/oh_S0; + + double oh_variance = oh_S2/oh_S0 - oh_mean*oh_mean; + double oh_sigma = std::sqrt(std::max(0.0, oh_variance)); + + double oh_mu3 = + oh_S3/oh_S0 + - 3.0*oh_mean*(oh_S2/oh_S0) + + 2.0*oh_mean*oh_mean*oh_mean; + + double oh_skew = (oh_sigma > 0.0) + ? oh_mu3/std::pow(oh_sigma,3) + : 0.0; + + jet_emcal /= jet_E; + jet_ohcal /= jet_E; + + if (jetpt > dijet_pt[0]) + { + dijet_pt[1]=dijet_pt[0]; + dijet_eta[1]= dijet_eta[0]; + dijet_phi[1]= dijet_phi[0]; + dijet_E[1]= dijet_E[0]; + dijet_em_frac[1]= dijet_em_frac[0]; + dijet_em_mean[1]= dijet_em_mean[0]; + dijet_em_skew[1]= dijet_em_skew[0]; + dijet_em_sig[1]= dijet_em_sig[0]; + dijet_oh_frac[1]= dijet_oh_frac[0]; + dijet_oh_mean[1]= dijet_oh_mean[0]; + dijet_oh_skew[1]= dijet_oh_skew[0]; + dijet_oh_sig[1]= dijet_oh_sig[0]; + + dijet_pt[0]=jetpt; + dijet_eta[0]= jeteta; + dijet_phi[0]= jetphi; + dijet_E[0]= jet_E; + dijet_em_frac[0]= jet_emcal; + dijet_em_mean[0]= em_mean; + dijet_em_skew[0]= em_skew; + dijet_em_sig[0]= em_sigma; + dijet_oh_frac[0]= jet_ohcal; + dijet_oh_mean[0]= oh_mean; + dijet_oh_skew[0]= oh_skew; + dijet_oh_sig[0]= oh_sigma; + + } + else if (jetpt > dijet_pt[1]) + { + dijet_pt[1]=jetpt; + dijet_eta[1]= jeteta; + dijet_phi[1]= jetphi; + dijet_E[1]= jet_E; + dijet_em_frac[1]= jet_emcal; + dijet_em_mean[1]= em_mean; + dijet_em_skew[1]= em_skew; + dijet_em_sig[1]= em_sigma; + dijet_oh_frac[1]= jet_ohcal; + dijet_oh_mean[1]= oh_mean; + dijet_oh_skew[1]= oh_skew; + dijet_oh_sig[1]= oh_sigma; + } + } + + double emcal_on[2] = {1}; + double ohcal_on[2] = {1}; + + + for (int i = 0 ; i < 2; i++) + { + if (std::isnan(dijet_em_frac[i]) || + std::isnan(dijet_em_mean[i]) || + std::isnan(dijet_em_skew[i]) || + std::isnan(dijet_em_sig[i])) + { + emcal_on[i] = 0; + } + if (std::isnan(dijet_oh_frac[i]) || + std::isnan(dijet_oh_mean[i]) || + std::isnan(dijet_oh_skew[i]) || + std::isnan(dijet_oh_sig[i])) + { + ohcal_on[i] = 0; + } + } + + float emcal_z_moments[2][4] = {0}; + float ohcal_z_moments[2][4] = {0}; + for (int i = 0; i < 2; i++) + { + if (emcal_on[i]) + { + emcal_z_moments[i][0] = dijet_em_frac[i]; + emcal_z_moments[i][1] = dijet_em_mean[i]; + emcal_z_moments[i][2] = dijet_em_skew[i]; + emcal_z_moments[i][3] = dijet_em_sig[i]; + } + if (ohcal_on[i]) + { + ohcal_z_moments[i][0] = dijet_oh_frac[i]; + ohcal_z_moments[i][1] = dijet_oh_mean[i]; + ohcal_z_moments[i][2] = dijet_oh_skew[i]; + ohcal_z_moments[i][3] = dijet_oh_sig[i]; + } + } + + float exj = dijet_E[1]/dijet_E[0]; + + std::array features = { + emcal_on[0], + ohcal_on[0], + emcal_on[1], + ohcal_on[1], + emcal_z_moments[0][0], + emcal_z_moments[0][1], + emcal_z_moments[0][2], + emcal_z_moments[0][3], + ohcal_z_moments[0][0], + ohcal_z_moments[0][1], + ohcal_z_moments[0][2], + ohcal_z_moments[0][3], + emcal_z_moments[1][0], + emcal_z_moments[1][1], + emcal_z_moments[1][2], + emcal_z_moments[1][3], + ohcal_z_moments[1][0], + ohcal_z_moments[1][1], + ohcal_z_moments[1][2], + ohcal_z_moments[1][3], + exj}; + + + zvtx = static_cast(VertexMLP::PredictVertexZ(features)); + return Fun4AllReturnCodes::EVENT_OK; +} diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.h b/offline/packages/calovtxreco/CaloVtxAlgoMLP.h new file mode 100644 index 0000000000..de119e7fcf --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.h @@ -0,0 +1,34 @@ +#ifndef CALOVTXALGOMLP_H +#define CALOVTXALGOMLP_H + +#include "CaloVtxAlgo.h" + +#include + +class PHCompositeNode; + +// Wraps the trained MLP (VertexMLP.h / vertex_mlp_weights.root) as a +// CaloVtxAlgo so it can be registered alongside other algorithms in +// CaloVtxReco and compared directly. +class CaloVtxAlgoMLP : public CaloVtxAlgo +{ + public: + explicit CaloVtxAlgoMLP(); + ~CaloVtxAlgoMLP() override = default; + + int Init(PHCompositeNode *topNode) override; + int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; + std::string Name() const override { return "MLP"; } + + void setWeightsFile(std::string file) { m_weightsFile = file; } + void setJetNode(std::string node) { m_jet_node = node; } +private: + + std::string m_jet_node{"Antikt_TowerInfo_vtx_none_r06"}; + std::string m_weightsFile{"vertex_mlp_weights.root"}; + + float m_radius_EM{0}; + float m_radius_OH{0}; +}; + +#endif diff --git a/offline/packages/calovtxreco/CaloVtxReco.cc b/offline/packages/calovtxreco/CaloVtxReco.cc index 6e8da31be7..0238754f6f 100644 --- a/offline/packages/calovtxreco/CaloVtxReco.cc +++ b/offline/packages/calovtxreco/CaloVtxReco.cc @@ -1,14 +1,5 @@ #include "CaloVtxReco.h" -#include -#include - -#include -#include -#include -#include -#include - #include #include @@ -16,17 +7,15 @@ #include #include - +#include "TMath.h" #include /* float radius_EM = 93.5; float radius_OH = 225.87; */ //____________________________________________________________________________.. -CaloVtxReco::CaloVtxReco(const std::string &name, const std::string &jetnodename, const bool use_z_energy_dep) +CaloVtxReco::CaloVtxReco(const std::string &name) : SubsysReco(name) - , m_use_z_energy_dep(use_z_energy_dep) - , m_jetnodename(jetnodename) { } @@ -74,358 +63,38 @@ int CaloVtxReco::InitRun(PHCompositeNode *topNode) { std::cout << "Initializing!" << std::endl; } - if (createNodes(topNode) == Fun4AllReturnCodes::ABORTRUN) + int status = createNodes(topNode); + if (status) { - return Fun4AllReturnCodes::ABORTRUN; + return status; } - RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); - RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); - if(!tower_geomEM || !tower_geomOH) - { - if(Verbosity() > 0) - { - std::cout << "CaloVtxReco::InitRun(): Missing tower geometry node for towergeomEM (address: " << tower_geomEM << ") or towergeomOH (address: " << tower_geomOH << ") - aborting run!" << std::endl; - } - return Fun4AllReturnCodes::ABORTRUN; - } - - m_radius_EM = tower_geomEM->get_radius(); - m_radius_OH = tower_geomOH->get_radius(); - - if(std::isnan(m_radius_EM) || std::isnan(m_radius_OH)) + for (auto &algo : m_algos) { - if(Verbosity() > 0) + status = algo->Init(topNode); + if (status) { - std::cout << "CaloVtxReco::InitRun(): NaN value for one of radius EM (value: " << m_radius_EM << ") or radius OH (value: " << m_radius_OH << ") after attempting to get - aborting run!" << std::endl; + return status; } - return Fun4AllReturnCodes::ABORTRUN; } - - + return Fun4AllReturnCodes::EVENT_OK; } -float CaloVtxReco::new_eta(int channel, TowerInfoContainer *towers, RawTowerGeomContainer *geom, RawTowerDefs::CalorimeterId caloID, float testz) -{ - int key = towers->encode_key(channel); - const RawTowerDefs::keytype geomkey = RawTowerDefs::encode_towerid(caloID, towers->getTowerEtaBin(key), towers->getTowerPhiBin(key)); - - RawTowerGeom *tower_geom = geom->get_tower_geometry(geomkey); - float oldeta = tower_geom->get_eta(); - - float radius = (caloID == RawTowerDefs::CalorimeterId::HCALIN ? m_radius_EM : m_radius_OH); - float towerz = radius / (tanf(2 * std::atanf(std::exp(oldeta)))); - float newz = towerz + testz; - float newTheta = std::atan2(radius, newz); - float neweta = -log(tan(0.5 * newTheta)); - - return neweta; -} - -float get_dphi(float phi1, float phi2) -{ - float dphi = std::abs(phi1 - phi2); - if (dphi > M_PI) - { - dphi = 2 * M_PI - dphi; - } - return dphi; -} - -int CaloVtxReco::calo_tower_algorithm(PHCompositeNode *topNode) const -{ - TowerInfoContainer *emcal_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC"); - TowerInfoContainer *hcalin_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); - TowerInfoContainer *hcalout_towers = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); - RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); - RawTowerGeomContainer *tower_geomIH = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); - RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); - - int size; - - float average_z[3]{0}; - float total_E[3]{0}; - - if (emcal_towers) - { - size = emcal_towers->size(); // online towers should be the same! - for (int channel = 0; channel < size; channel++) - { - TowerInfo *_tower = emcal_towers->get_tower_at_channel(channel); - short good = (_tower->get_isGood() ? 1 : 0); - if (!good) - { - continue; - } - - float energy = _tower->get_energy(); - if (energy < m_energy_cut) - { - continue; - } - - // float time = _tower->get_time_float(); - - unsigned int towerkey = emcal_towers->encode_key(channel); - int ieta = emcal_towers->getTowerEtaBin(towerkey); - int iphi = emcal_towers->getTowerPhiBin(towerkey); - - const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::CEMC, ieta, iphi); - /* - if (emcal_r < 10) - { - emcal_r = tower_geomEM->get_tower_geometry(key)->get_center_radius(); - } - */ - float tower_z = tower_geomEM->get_tower_geometry(key)->get_center_z(); - average_z[0] += tower_z * energy; - total_E[0] += energy; - } - } - - if (hcalin_towers) - { - size = hcalin_towers->size(); // online towers should be the same! - for (int channel = 0; channel < size; channel++) - { - TowerInfo *_tower = hcalin_towers->get_tower_at_channel(channel); - float energy = _tower->get_energy(); - if (energy < m_energy_cut) - { - continue; - } - // float time = _tower->get_time_float(); - short good = (_tower->get_isGood() ? 1 : 0); - if (!good) - { - continue; - } - - unsigned int towerkey = hcalin_towers->encode_key(channel); - int ieta = hcalin_towers->getTowerEtaBin(towerkey); - int iphi = hcalin_towers->getTowerPhiBin(towerkey); - const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); - float tower_z = tower_geomIH->get_tower_geometry(key)->get_center_z(); - /* - if (hcalin_r < 10) - { - hcalin_r = tower_geomIH->get_tower_geometry(key)->get_center_radius(); - } - */ - average_z[1] += tower_z * energy; - total_E[1] += energy; - } - } - if (hcalout_towers) - { - size = hcalout_towers->size(); // online towers should be the same! - for (int channel = 0; channel < size; channel++) - { - TowerInfo *_tower = hcalout_towers->get_tower_at_channel(channel); - float energy = _tower->get_energy(); - if (energy < m_energy_cut) - { - continue; - } - // float time = _tower->get_time_float(); - unsigned int towerkey = hcalout_towers->encode_key(channel); - int ieta = hcalout_towers->getTowerEtaBin(towerkey); - int iphi = hcalout_towers->getTowerPhiBin(towerkey); - short good = (_tower->get_isGood() ? 1 : 0); - - if (!good) - { - continue; - } - - const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); - /* - if (hcalout_r < 10) - { - hcalout_r = tower_geomOH->get_tower_geometry(key)->get_center_radius(); - } - */ - float tower_z = tower_geomOH->get_tower_geometry(key)->get_center_z(); - average_z[2] += tower_z * energy; - total_E[2] += energy; - } - } - - double b_calo_vertex_z = (average_z[0] + average_z[1] + average_z[2]) / (total_E[0] + total_E[1] + total_E[2]); - - CaloVertex *vertex = new CaloVertexv1(); - vertex->set_z(b_calo_vertex_z); - m_calovtxmap->insert(vertex); - - return 0; -} - int CaloVtxReco::process_event(PHCompositeNode *topNode) { - if (m_use_z_energy_dep) - { - calo_tower_algorithm(topNode); - return Fun4AllReturnCodes::EVENT_OK; - } - - if (Verbosity() > 1) - { - std::cout << std::endl - << std::endl - << std::endl - << "CaloVtxReco: Beginning event processing" << std::endl; - } - m_zvtx = std::numeric_limits::quiet_NaN(); - JetContainer *jetcon = findNode::getClass(topNode, m_jetnodename); - TowerInfoContainer *towers[3]; - towers[0] = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC_RETOWER"); - towers[1] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); - towers[2] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); - - RawTowerGeomContainer *geom[3]; - geom[0] = findNode::getClass(topNode, "TOWERGEOM_CEMC"); - geom[1] = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); - geom[2] = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); - - const int nz = 601; - const int njet = 2; - Jet *jets[njet]; - float jpt[njet] = {0}; - float jemsum[njet] = {0}; - float johsum[njet] = {0}; - float jemeta[njet] = {0}; - float joheta[njet] = {0}; - - if (jetcon) - { - int tocheck = jetcon->size(); - if (Verbosity() > 2) - { - std::cout << "Found " << tocheck << " jets to check..." << std::endl; - } - for (int i = 0; i < tocheck; ++i) - { - Jet *jet = jetcon->get_jet(i); - if (jet) - { - float pt = jet->get_pt(); - if (pt < m_jet_threshold) - { - continue; - } - if (pt > jpt[0]) - { - jpt[1] = jpt[0]; - jets[1] = jets[0]; - jets[0] = jet; - jpt[0] = pt; - } - else if (pt > jpt[1]) - { - jets[1] = jet; - jpt[1] = pt; - } - } - } - } - else - { - if (Verbosity() > 0) - { - std::cout << "no jets" << std::endl; - } - return Fun4AllReturnCodes::ABORTEVENT; - } + + float zvtx = std::numeric_limits::quiet_NaN(); - if (jpt[0] == 0) - { - if (Verbosity() > 2) + for (auto &algo : m_algos) { - std::cout << "NO JETS > 5 GeV!" << std::endl; - } - } + float tempz = std::numeric_limits::quiet_NaN(); + algo->CalculateVertex(topNode, zvtx); + std::cout << algo->Name() <<": "<::max(); - for (int i = 0; i < nz; ++i) - { - float testz = -300 + i; - float testmetric = 0; - for (int j = 0; j < njet; ++j) - { - if (jpt[j] == 0) - { - continue; - } - jemsum[j] = 0; - johsum[j] = 0; - jemeta[j] = 0; - joheta[j] = 0; - for (auto comp : jets[j]->get_comp_vec()) - { - if (comp.first == 5 || comp.first == 26) - { - continue; - } - unsigned int channel = comp.second; - if (comp.first == 7 || comp.first == 27) - { - TowerInfo *tower = towers[2]->get_tower_at_channel(channel); - if (tower->get_energy() < 0.1) - { - continue; - } - johsum[j] += tower->get_energy(); - float neweta = new_eta(channel, towers[2], geom[2], RawTowerDefs::CalorimeterId::HCALOUT, testz); - joheta[j] += neweta * tower->get_energy(); - } - if (comp.first == 13 || comp.first == 28 || comp.first == 25) - { - TowerInfo *tower = towers[0]->get_tower_at_channel(channel); - if (tower->get_energy() < 0.1) - { - continue; - } - jemsum[j] += tower->get_energy(); - float neweta = new_eta(channel, towers[0], geom[1], RawTowerDefs::CalorimeterId::HCALIN, testz); - jemeta[j] += neweta * tower->get_energy(); - } - } - jemeta[j] /= jemsum[j]; - joheta[j] /= johsum[j]; - if ((jemsum[j] == 0 || johsum[j] == 0) && Verbosity() > 1) - { - std::cout << "zero E sum in at least one calo for a jet" << std::endl; - } - if (!std::isnan(jemeta[j]) && !std::isnan(joheta[j])) - { - testmetric += pow(jemeta[j] - joheta[j], 2); - } - } - if (Verbosity() > 3) - { - std::cout << "metric: " << testmetric << std::endl; - } - if (testmetric < metric && testmetric != 0) - { - metric = testmetric; - m_zvtx = testz; } - } - if (std::abs(m_zvtx) < 305) - { - if (Verbosity() > 2) - { - std::cout << "optimal z: " << m_zvtx << std::endl; - } - CaloVertex *vertex = new CaloVertexv1(); - m_zvtx *= m_calib_factor; // calibration factor from simulation - vertex->set_z(m_zvtx); - m_calovtxmap->insert(vertex); - if (Verbosity() > 3) - { - std::cout << "CaloVtxReco: end event" << std::endl; - } - } + CaloVertex *vertex = new CaloVertexv1(); + vertex->set_z(zvtx); + m_calovtxmap->insert(vertex); return Fun4AllReturnCodes::EVENT_OK; } diff --git a/offline/packages/calovtxreco/CaloVtxReco.h b/offline/packages/calovtxreco/CaloVtxReco.h index e1a7d2b76b..31ab2502be 100644 --- a/offline/packages/calovtxreco/CaloVtxReco.h +++ b/offline/packages/calovtxreco/CaloVtxReco.h @@ -1,54 +1,40 @@ #ifndef CALOVTXRECO_CALOVTXRECO_H #define CALOVTXRECO_CALOVTXRECO_H -#include - +#include "CaloVtxAlgo.h" +#include +#include #include class CaloVertexMap; class PHCompositeNode; -class RawTowerGeomContainer; -class TowerInfoContainer; class CaloVtxReco : public SubsysReco { public: - CaloVtxReco(const std::string &name = "CaloVtxReco", const std::string &jetnodename = "zzjets06", const bool use_z_energy_dep = false); + enum ALGOTYPE + { + JETSKEW = 0, + CALOZ = 1, + JET_MLP = 2 + }; + + CaloVtxReco(const std::string &name = "CaloVtxReco"); virtual ~CaloVtxReco() = default; int createNodes(PHCompositeNode *topNode); - float new_eta(int channel, TowerInfoContainer *towers, RawTowerGeomContainer *geom, RawTowerDefs::CalorimeterId caloID, float testz); - int InitRun(PHCompositeNode *topNode) override; - - int calo_tower_algorithm(PHCompositeNode *topNode) const; - + int process_event(PHCompositeNode *topNode) override; - - float get_jet_threshold() { return m_jet_threshold; } - - void set_jet_threshold(float new_thresh) { m_jet_threshold = new_thresh; } - - float get_calib_factor() { return m_calib_factor; } - - void set_calib_factor(float new_calib) { m_calib_factor = new_calib; } - - float get_energy_cut() { return m_energy_cut; } - - void set_energy_cut(float new_energy) { m_energy_cut = new_energy; } - + + void registerAlgo(CaloVtxAlgo* algo) { m_algos.push_back(std::move(algo)); } + private: + std::vector m_algos; + CaloVertexMap *m_calovtxmap{nullptr}; - float m_jet_threshold{15}; - float m_zvtx{std::numeric_limits::quiet_NaN()}; - float m_calib_factor{1.406}; - float m_energy_cut{0.1}; - float m_radius_EM{std::numeric_limits::quiet_NaN()}; - float m_radius_OH{std::numeric_limits::quiet_NaN()}; - bool m_use_z_energy_dep; - std::string m_jetnodename; }; #endif // CALOVTXRECO_CALOVTXRECO_H diff --git a/offline/packages/calovtxreco/Makefile.am b/offline/packages/calovtxreco/Makefile.am index 0d8d7283db..87c4d0f429 100644 --- a/offline/packages/calovtxreco/Makefile.am +++ b/offline/packages/calovtxreco/Makefile.am @@ -10,13 +10,21 @@ AM_LDFLAGS = \ -L$(OFFLINE_MAIN)/lib pkginclude_HEADERS = \ - CaloVtxReco.h + CaloVtxReco.h \ + CaloVtxAlgo.h \ + CaloVtxAlgoJetSkew.h \ + CaloVtxAlgoCaloZ.h \ + CaloVtxAlgoMLP.h \ + VertexMLP.h lib_LTLIBRARIES = \ libcalovtxreco.la libcalovtxreco_la_SOURCES = \ - CaloVtxReco.cc + CaloVtxReco.cc \ + CaloVtxAlgoJetSkew.cc \ + CaloVtxAlgoCaloZ.cc \ + CaloVtxAlgoMLP.cc libcalovtxreco_la_LIBADD = \ -lcalo_io \ diff --git a/offline/packages/calovtxreco/VertexMLP.h b/offline/packages/calovtxreco/VertexMLP.h new file mode 100644 index 0000000000..e52c010f80 --- /dev/null +++ b/offline/packages/calovtxreco/VertexMLP.h @@ -0,0 +1,143 @@ +#ifndef VERTEXMLP_H +#define VERTEXMLP_H + +// Loads and evaluates the calo-vertex-z MLP exported by +// export_nn_weights.py + export_nn_root.C (calovertex/vertex_mlp_weights.root). +// +// Usage in a Fun4All SubsysReco module: +// in InitRun(): VertexMLP::Load("vertex_mlp_weights.root"); +// in process_event(): double z = VertexMLP::PredictVertexZ(features); +// +// Retraining the model requires no recompilation here: rerun +// export_nn_weights.py + export_nn_root.C to regenerate the .root file, +// point Load() at it (same path or a new one), done. +// +// Input feature order (features[0..20] must be filled in exactly this +// order -- it's the order the model was trained on). Features 0-3 are +// binary "this calo had a valid shower" flags (1 = at least one constituent +// above threshold, 0 = none -- fill the corresponding zmean/zsig/zskew with +// 0 in that case, matching how the training CSV imputes them): +// 0 emcal_lead (flag) 7 emcal_lead_energy 14 emcal_sublead_zskew +// 1 ohcal_lead (flag) 8 ohcal_lead_zmean 15 emcal_sublead_energy +// 2 emcal_sublead (flag) 9 ohcal_lead_zsig 16 ohcal_sublead_zmean +// 3 ohcal_sublead (flag) 10 ohcal_lead_zskew 17 ohcal_sublead_zsig +// 4 emcal_lead_zmean 11 ohcal_lead_energy 18 ohcal_sublead_zskew +// 5 emcal_lead_zsig 12 emcal_sublead_zmean 19 ohcal_sublead_energy +// 6 emcal_lead_zskew 13 emcal_sublead_zsig 20 exj +// (also stored, in this order, in the "feature_names" TObjArray inside the +// .root file -- worth checking against if this list and that file ever +// disagree). + +#include +#include +#include + +#include +#include +#include +#include + +namespace VertexMLP { + +constexpr int kNFeatures = 21; + +inline TMatrixD gW1, gW2, gW3, gW4; +inline TVectorD gB1, gB2, gB3, gB4; +inline TVectorD gMedian, gIqr, gLo, gHi; +inline bool gLoaded = false; + +namespace detail { + +inline bool GetMatrix(TFile *f, const char *name, TMatrixD &out) { + TMatrixD *m = static_cast(f->Get(name)); + if (!m) { + std::cerr << "VertexMLP::Load: missing TMatrixD \"" << name << "\"" << std::endl; + return false; + } + out.ResizeTo(*m); + out = *m; + return true; +} + +inline bool GetVector(TFile *f, const char *name, TVectorD &out) { + TVectorD *v = static_cast(f->Get(name)); + if (!v) { + std::cerr << "VertexMLP::Load: missing TVectorD \"" << name << "\"" << std::endl; + return false; + } + out.ResizeTo(*v); + out = *v; + return true; +} + +// out = ReLU(W*in + b); set relu=false for the final (output) layer. +inline TVectorD ApplyLayer(const TMatrixD &W, const TVectorD &b, const TVectorD &in, bool relu) { + TVectorD out = W * in + b; + if (relu) { + for (int i = 0; i < out.GetNrows(); i++) out[i] = std::max(0.0, out[i]); + } + return out; +} + +} // namespace detail + +// Call once, from InitRun(). Returns false (and leaves the model unloaded, +// so PredictVertexZ() will refuse to run) if the file or any expected +// object inside it is missing. +inline bool Load(const std::string &filename = "vertex_mlp_weights.root") { + gLoaded = false; + + TFile *f = TFile::Open(filename.c_str(), "READ"); + if (!f || f->IsZombie()) { + std::cerr << "VertexMLP::Load: cannot open " << filename << std::endl; + return false; + } + + bool ok = detail::GetMatrix(f, "W1", gW1) && detail::GetMatrix(f, "W2", gW2) && + detail::GetMatrix(f, "W3", gW3) && detail::GetMatrix(f, "W4", gW4) && + detail::GetVector(f, "b1", gB1) && detail::GetVector(f, "b2", gB2) && + detail::GetVector(f, "b3", gB3) && detail::GetVector(f, "b4", gB4) && + detail::GetVector(f, "median", gMedian) && detail::GetVector(f, "iqr", gIqr) && + detail::GetVector(f, "lo", gLo) && detail::GetVector(f, "hi", gHi); + + f->Close(); + delete f; + + if (!ok || gMedian.GetNrows() != kNFeatures) { + std::cerr << "VertexMLP::Load: failed to load a complete model from " << filename << std::endl; + return false; + } + + gLoaded = true; + std::cout << "VertexMLP::Load: loaded " << filename << std::endl; + return true; +} + +// Predicts the vertex z [cm] given the 21 input features (see the ordering +// table above). Returns 0 and prints an error if Load() hasn't succeeded. +inline double PredictVertexZ(const std::array &features) { + if (!gLoaded) { + std::cerr << "VertexMLP::PredictVertexZ: model not loaded -- call " + "VertexMLP::Load() in InitRun() first" + << std::endl; + return 0.0; + } + + // Same preprocessing as training: percentile clip, then median/IQR scale. + TVectorD x(kNFeatures); + for (int i = 0; i < kNFeatures; i++) { + double v = std::clamp(features[i], gLo[i], gHi[i]); + x[i] = (v - gMedian[i]) / gIqr[i]; + } + + TVectorD h1 = detail::ApplyLayer(gW1, gB1, x, true); + TVectorD h2 = detail::ApplyLayer(gW2, gB2, h1, true); + TVectorD h3 = detail::ApplyLayer(gW3, gB3, h2, true); + TVectorD out = detail::ApplyLayer(gW4, gB4, h3, false); + + return out[0]; +} + +} // namespace VertexMLP + +#endif diff --git a/offline/packages/globalvertex/GlobalVertex.h b/offline/packages/globalvertex/GlobalVertex.h index 3124cea282..da34dd625a 100644 --- a/offline/packages/globalvertex/GlobalVertex.h +++ b/offline/packages/globalvertex/GlobalVertex.h @@ -19,6 +19,7 @@ class GlobalVertex : public PHObject { UNDEFINED = 0, TRUTH = 100, + ZERO = 150, SMEARED = 200, MBD = 300, SVTX = 400, diff --git a/offline/packages/globalvertex/GlobalVertexReco.cc b/offline/packages/globalvertex/GlobalVertexReco.cc index 73681759cc..c34f4e1098 100644 --- a/offline/packages/globalvertex/GlobalVertexReco.cc +++ b/offline/packages/globalvertex/GlobalVertexReco.cc @@ -6,7 +6,7 @@ #include "GlobalVertexv3.h" #include "MbdVertex.h" #include "MbdVertexMap.h" -#include "CaloVertex.h" +#include "CaloVertexv1.h" #include "CaloVertexMap.h" #include "SvtxVertex.h" #include "SvtxVertexMap.h" @@ -297,6 +297,32 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) } } + // okay now put in zero + if (useVertexType(GlobalVertex::VTXTYPE::ZERO)) + { + if (Verbosity()) + { + std::cout << "GlobalVertexReco::process_event - zero" << std::endl; + } + + CaloVertex *cvertex = new CaloVertexv1(); + cvertex->set_z(0); + + GlobalVertex *vertex = new GlobalVertexv3(); + vertex->set_id(globalmap->size()); + + vertex->clone_insert_vtx(GlobalVertex::ZERO, cvertex); + used_calo_vtxids.insert(cvertex->get_id()); + + globalmap->insert(vertex); + + if (Verbosity() > 1) + { + vertex->identify(); + } + } + + // okay now loop over all unused MBD vertexes (3rd class)... if (mbdmap && calomap && useVertexType(GlobalVertex::VTXTYPE::MBD_CALO)) { From e44a6fc1b3e0ec6f6dcc8fc02fecc83651cc740c Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Thu, 9 Jul 2026 19:02:52 -0400 Subject: [PATCH 2/8] vertex algo matches offline --- offline/packages/calovtxreco/CaloVtxAlgo.h | 3 + .../packages/calovtxreco/CaloVtxAlgoCaloZ.cc | 5 - .../packages/calovtxreco/CaloVtxAlgoCaloZ.h | 4 +- .../calovtxreco/CaloVtxAlgoJetSkew.cc | 3 - .../packages/calovtxreco/CaloVtxAlgoJetSkew.h | 5 +- .../packages/calovtxreco/CaloVtxAlgoMLP.cc | 539 ++++++++++-------- offline/packages/calovtxreco/CaloVtxAlgoMLP.h | 6 +- offline/packages/calovtxreco/CaloVtxReco.cc | 16 +- offline/packages/globalvertex/CaloVertex.h | 4 + offline/packages/globalvertex/CaloVertexv1.cc | 1 + offline/packages/globalvertex/CaloVertexv1.h | 6 + offline/packages/globalvertex/GlobalVertex.h | 3 + .../packages/globalvertex/GlobalVertexReco.cc | 60 +- .../packages/globalvertex/GlobalVertexv4.cc | 261 +++++++++ .../packages/globalvertex/GlobalVertexv4.h | 73 +++ .../globalvertex/GlobalVertexv4LinkDef.h | 5 + offline/packages/globalvertex/Makefile.am | 26 +- offline/packages/globalvertex/Vertex.h | 4 + offline/packages/globalvertex/VertexDefs.h | 14 + 19 files changed, 761 insertions(+), 277 deletions(-) create mode 100644 offline/packages/globalvertex/GlobalVertexv4.cc create mode 100644 offline/packages/globalvertex/GlobalVertexv4.h create mode 100644 offline/packages/globalvertex/GlobalVertexv4LinkDef.h create mode 100644 offline/packages/globalvertex/VertexDefs.h diff --git a/offline/packages/calovtxreco/CaloVtxAlgo.h b/offline/packages/calovtxreco/CaloVtxAlgo.h index 4b6341ea15..6b11d0103b 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgo.h +++ b/offline/packages/calovtxreco/CaloVtxAlgo.h @@ -2,6 +2,7 @@ #define CALOVTXALGO_H #include +#include class PHCompositeNode; @@ -26,6 +27,8 @@ class CaloVtxAlgo // (e.g. "CALOVTXOUT_MLP") and in log/eval output so results from // different algorithms don't collide and can be compared directly. virtual std::string Name() const = 0; + + virtual VertexDefs::CALOALGO Algo() const = 0; }; #endif diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc index 9f33f9f8e7..344f698ce8 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc @@ -12,11 +12,6 @@ #include -CaloVtxAlgoCaloZ::CaloVtxAlgoCaloZ() -{ - -} - int CaloVtxAlgoCaloZ::Init(PHCompositeNode * /*topNode*/) { diff --git a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h index f4d77e403e..a13680d15e 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h +++ b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.h @@ -12,13 +12,13 @@ class PHCompositeNode; class CaloVtxAlgoCaloZ : public CaloVtxAlgo { public: - explicit CaloVtxAlgoCaloZ(); + explicit CaloVtxAlgoCaloZ() = default; ~CaloVtxAlgoCaloZ() override = default; int Init(PHCompositeNode *topNode) override; int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; std::string Name() const override { return "CaloZ"; } - + VertexDefs::CALOALGO Algo() const override { return VertexDefs::CALOALGO::AVGZ; } float get_energy_cut() { return m_energy_cut; } void set_energy_cut(float new_energy) { m_energy_cut = new_energy; } diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc index 2c1ed616f7..6c3a093508 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc @@ -13,9 +13,6 @@ #include -CaloVtxAlgoJetSkew::CaloVtxAlgoJetSkew() -{ -} int CaloVtxAlgoJetSkew::Init(PHCompositeNode *topNode) { diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h index 819d5d8543..d28814251d 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h @@ -19,13 +19,14 @@ class TowerInfoContainer; class CaloVtxAlgoJetSkew : public CaloVtxAlgo { public: - explicit CaloVtxAlgoJetSkew(); + explicit CaloVtxAlgoJetSkew() = default; ~CaloVtxAlgoJetSkew() override = default; int Init(PHCompositeNode *topNode) override; int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; std::string Name() const override { return "JetSkew"; } - + VertexDefs::CALOALGO Algo() const override { return VertexDefs::CALOALGO::JETSKEW; } + void setJetNode(std::string jetnode) { m_jetnodename = jetnode; } float get_jet_threshold() { return m_jet_threshold; } void set_jet_threshold(float new_thresh) { m_jet_threshold = new_thresh; } diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc index 918c319b19..adfe981863 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -14,25 +14,30 @@ #include -CaloVtxAlgoMLP::CaloVtxAlgoMLP() -{ -} - int CaloVtxAlgoMLP::Init(PHCompositeNode *topNode) { if (!VertexMLP::Load(m_weightsFile)) { return Fun4AllReturnCodes::ABORTRUN; } + RawTowerGeomContainer *tower_geomEM = findNode::getClass(topNode, "TOWERGEOM_CEMC"); RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + RawTowerGeomContainer *tower_geomIH = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); if(!tower_geomEM || !tower_geomOH) { return Fun4AllReturnCodes::ABORTRUN; } + unsigned int emkey = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::CEMC, 48, 128); + m_radius_EM = tower_geomEM->get_tower_geometry(emkey)->get_center_radius(); + unsigned int ihkey = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, 12, 0); + m_radius_IH = tower_geomIH->get_tower_geometry(ihkey)->get_center_radius(); + unsigned int ohkey = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, 12, 0); + m_radius_OH = tower_geomOH->get_tower_geometry(ohkey)->get_center_radius(); - m_radius_EM = tower_geomEM->get_radius(); - m_radius_OH = tower_geomOH->get_radius(); + // m_radius_EM = tower_geomEM->get_radius(); + // m_radius_IH = tower_geomIH->get_radius(); + // m_radius_OH = tower_geomOH->get_radius(); if(std::isnan(m_radius_EM) || std::isnan(m_radius_OH)) { @@ -51,6 +56,8 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) RawTowerGeomContainer *tower_geomIH = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); RawTowerGeomContainer *tower_geomOH = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + + if (!emcalre_towers || !hcalin_towers || !hcalout_towers) { return Fun4AllReturnCodes::ABORTRUN; @@ -61,271 +68,349 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) return Fun4AllReturnCodes::ABORTRUN; } - // TODO: fill `features` from topNode in the exact order documented in - // VertexMLP.h (emcal/ohcal lead+sublead flags, zmean/zsig/zskew, - // energy, exj). This is the same per-event extraction CaloVtxReco used - // to do inline before the NN call -- move that block here unchanged. + + std::vector *b_reco_jet_pt = new std::vector(); + std::vector *b_reco_jet_eta = new std::vector(); + std::vector *b_reco_jet_phi = new std::vector(); + std::vector *b_reco_jet_emcal_eta = new std::vector(); + std::vector *b_reco_jet_ihcal_eta = new std::vector(); + std::vector *b_reco_jet_ohcal_eta = new std::vector(); + std::vector *b_reco_jet_emcal_zmean = new std::vector(); + std::vector *b_reco_jet_ohcal_zmean = new std::vector(); + std::vector *b_reco_jet_emcal_zsig = new std::vector(); + std::vector *b_reco_jet_ohcal_zsig = new std::vector(); + std::vector *b_reco_jet_emcal_zskew = new std::vector(); + std::vector *b_reco_jet_ohcal_zskew = new std::vector(); + std::vector *b_reco_jet_emcal = new std::vector(); + std::vector *b_reco_jet_ohcal = new std::vector(); + std::vector *b_reco_jet_ihcal = new std::vector(); JetContainer *jetscon = findNode::getClass(topNode, m_jet_node.c_str()); - float dijet_pt[2]={0}; - float dijet_eta[2]={0}; - float dijet_phi[2]={0}; - float dijet_E[2]={0}; - float dijet_em_frac[2]={0}; - float dijet_em_mean[2]={0}; - float dijet_em_skew[2]={0}; - float dijet_em_sig[2]={0}; - float dijet_oh_frac[2]={0}; - float dijet_oh_mean[2]={0}; - float dijet_oh_skew[2]={0}; - float dijet_oh_sig[2]={0}; - - if (!jetscon) - { - return Fun4AllReturnCodes::EVENT_OK; - } - for (auto jet : *jetscon) + int ijet = 0; + double dijet_pt[2] = {0}; + int dijet_index[2] = {0}; + + if (jetscon) { - double jet_E = 0; - double jet_emcal = 0; - double jet_ohcal = 0; - double jetpt = jet->get_pt(); - double jetphi = jet->get_phi(); - double jeteta= jet->get_eta(); - - if (jetpt < 0.5) continue; - - double em_S0 = 0.0; - double em_S1 = 0.0; - double em_S2 = 0.0; - double em_S3 = 0.0; - double oh_S0 = 0.0; - double oh_S1 = 0.0; - double oh_S2 = 0.0; - double oh_S3 = 0.0; - - - int itower = 0; - for (auto comp : jet->get_comp_vec()) + for (auto jet : *jetscon) { - unsigned int channel = comp.second; - TowerInfo *tower; - float tower_e = 0; - if (comp.first == 26 || comp.first == 30) + double jet_E = 0; + double jet_emcal = 0; + double jet_ihcal = 0; + double jet_ohcal = 0; + double jet_emcal_eta = 0; + double jet_ihcal_eta = 0; + double jet_ohcal_eta = 0; + double jetpt = jet->get_pt(); + if (jetpt < 2) continue; + + double em_S0 = 0.0; + double em_S1 = 0.0; + double em_S2 = 0.0; + double em_S3 = 0.0; + double oh_S0 = 0.0; + double oh_S1 = 0.0; + double oh_S2 = 0.0; + double oh_S3 = 0.0; + + + int itower = 0; + for (auto comp : jet->get_comp_vec()) { - tower = hcalin_towers->get_tower_at_channel(channel); + unsigned int channel = comp.second; + TowerInfo *tower; + float tower_e = 0; + if (comp.first == 26 || comp.first == 30) + { // IHcal + tower = hcalin_towers->get_tower_at_channel(channel); - if (!tower || !tower_geomIH) - { - continue; - } + if (!tower || !tower_geomIH) + { + continue; + } - tower_e = tower->get_energy(); - if (tower_e < 0.005) continue; - jet_E += tower_e; - } - else if (comp.first == 27 || comp.first == 31) - { - // OHCAL - tower = hcalout_towers->get_tower_at_channel(channel); - if (!tower || !tower_geomOH) - { - continue; - } - unsigned int towerkey = hcalout_towers->encode_key(channel); - int ieta = hcalout_towers->getTowerEtaBin(towerkey); - int iphi = hcalout_towers->getTowerPhiBin(towerkey); - const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); - float tower_eta = tower_geomOH->get_tower_geometry(key)->get_eta(); - float tower_r = m_radius_OH; - float tower_z = tower_r*sinh(tower_eta); + unsigned int towerkey = hcalin_towers->encode_key(channel); + int ieta = hcalin_towers->getTowerEtaBin(towerkey); + int iphi = hcalin_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); + float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); + float tower_r = m_radius_IH; + float tower_z = tower_r*sinh(tower_eta); + float new_tower_z = tower_z;// - jet_vertex; + float new_tower_eta = asinh(new_tower_z/tower_r); - tower_e = tower->get_energy(); - if (tower_e < 0.035) continue; - jet_ohcal += tower_e; - oh_S0 += tower_e; - oh_S1 += tower_e*tower_z; - oh_S2 += tower_e*tower_z*tower_z; - oh_S3 += tower_e*tower_z*tower_z*tower_z; - jet_E += tower_e; - } - else if (comp.first == 28 || comp.first == 29) - { // EMCAL - tower = emcalre_towers->get_tower_at_channel(channel); - - if (!tower || !tower_geomOH) - { - continue; + tower_e = tower->get_energy(); + if (tower_e < 0.005) continue; + jet_ihcal_eta += new_tower_eta*tower_e; + jet_ihcal += tower_e; + jet_E += tower_e; } - unsigned int towerkey = emcalre_towers->encode_key(channel); - int ieta = emcalre_towers->getTowerEtaBin(towerkey); - int iphi = emcalre_towers->getTowerPhiBin(towerkey); - const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); - float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); - float tower_r = m_radius_EM; - float tower_z = tower_r*sinh(tower_eta); + else if (comp.first == 27 || comp.first == 31) + { // OHCAL + tower = hcalout_towers->get_tower_at_channel(channel); + if (!tower || !tower_geomOH) + { + continue; + } + unsigned int towerkey = hcalout_towers->encode_key(channel); + int ieta = hcalout_towers->getTowerEtaBin(towerkey); + int iphi = hcalout_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); + float tower_eta = tower_geomOH->get_tower_geometry(key)->get_eta(); + float tower_r = m_radius_OH; + float tower_z = tower_r*sinh(tower_eta); + float new_tower_z = tower_z;// - jet_vertex; + float new_tower_eta = asinh(new_tower_z/tower_r); + + tower_e = tower->get_energy(); + if (tower_e < 0.035) continue; + jet_ohcal_eta += new_tower_eta*tower_e; + jet_ohcal += tower_e; + oh_S0 += tower_e; + oh_S1 += tower_e*tower_z; + oh_S2 += tower_e*tower_z*tower_z; + oh_S3 += tower_e*tower_z*tower_z*tower_z; + + //tower_e = tower->get_energy(); + //jet_ohcal += tower_e; + jet_E += tower_e; + + } + else if (comp.first == 28 || comp.first == 29) + { // EMCAL + tower = emcalre_towers->get_tower_at_channel(channel); + + if (!tower || !tower_geomOH) + { + continue; + } + unsigned int towerkey = emcalre_towers->encode_key(channel); + int ieta = emcalre_towers->getTowerEtaBin(towerkey); + int iphi = emcalre_towers->getTowerPhiBin(towerkey); + const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); + float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); + float tower_r = m_radius_EM; + float tower_z = tower_r*sinh(tower_eta); + float new_tower_z = tower_z;// - jet_vertex; + float new_tower_eta = asinh(new_tower_z/tower_r); - tower_e = tower->get_energy(); - if (tower_e < 0.068) continue; - jet_emcal += tower_e; - jet_E += tower_e; - em_S0 += tower_e; - em_S1 += tower_e*tower_z; - em_S2 += tower_e*tower_z*tower_z; - em_S3 += tower_e*tower_z*tower_z*tower_z; + tower_e = tower->get_energy(); + if (tower_e < 0.068) continue; + jet_emcal_eta += new_tower_eta*tower_e; + jet_emcal += tower_e; + jet_E += tower_e; + em_S0 += tower_e; + em_S1 += tower_e*tower_z; + em_S2 += tower_e*tower_z*tower_z; + em_S3 += tower_e*tower_z*tower_z*tower_z; + + // if (print) + // { + // std::cout << itower << " | e:" << tower_e << " / eta: " << tower_eta << " / " << new_tower_eta << std::endl; + // } + } + + itower++; } + if (itower == 0) continue; - itower++; - } + double em_mean = em_S1/em_S0; + + double em_variance = em_S2/em_S0 - em_mean*em_mean; + double em_sigma = std::sqrt(std::max(0.0, em_variance)); - if (itower == 0) continue; + double em_mu3 = + em_S3/em_S0 + - 3.0*em_mean*(em_S2/em_S0) + + 2.0*em_mean*em_mean*em_mean; - double em_mean = em_S1/em_S0; + double em_skew = (em_sigma > 0.0) + ? em_mu3/std::pow(em_sigma,3) + : 0.0; - double em_variance = em_S2/em_S0 - em_mean*em_mean; - double em_sigma = std::sqrt(std::max(0.0, em_variance)); + double oh_mean = oh_S1/oh_S0; - double em_mu3 = - em_S3/em_S0 - - 3.0*em_mean*(em_S2/em_S0) - + 2.0*em_mean*em_mean*em_mean; + double oh_variance = oh_S2/oh_S0 - oh_mean*oh_mean; + double oh_sigma = std::sqrt(std::max(0.0, oh_variance)); - double em_skew = (em_sigma > 0.0) - ? em_mu3/std::pow(em_sigma,3) - : 0.0; + double oh_mu3 = + oh_S3/oh_S0 + - 3.0*oh_mean*(oh_S2/oh_S0) + + 2.0*oh_mean*oh_mean*oh_mean; - double oh_mean = oh_S1/oh_S0; + double oh_skew = (oh_sigma > 0.0) + ? oh_mu3/std::pow(oh_sigma,3) + : 0.0; + + jet_emcal_eta /= jet_emcal; + jet_ihcal_eta /= jet_ihcal; + jet_ohcal_eta /= jet_ohcal; + jet_emcal /= jet_E; + jet_ihcal /= jet_E; + jet_ohcal /= jet_E; + - double oh_variance = oh_S2/oh_S0 - oh_mean*oh_mean; - double oh_sigma = std::sqrt(std::max(0.0, oh_variance)); + b_reco_jet_pt->push_back(jet->get_pt()); + b_reco_jet_eta->push_back(jet->get_eta()); + b_reco_jet_phi->push_back(jet->get_phi()); - double oh_mu3 = - oh_S3/oh_S0 - - 3.0*oh_mean*(oh_S2/oh_S0) - + 2.0*oh_mean*oh_mean*oh_mean; + b_reco_jet_emcal->push_back(jet_emcal); + b_reco_jet_ohcal->push_back(jet_ohcal); + b_reco_jet_ihcal->push_back(jet_ihcal); - double oh_skew = (oh_sigma > 0.0) - ? oh_mu3/std::pow(oh_sigma,3) - : 0.0; - - jet_emcal /= jet_E; - jet_ohcal /= jet_E; + b_reco_jet_emcal_zmean->push_back(em_mean); + b_reco_jet_ohcal_zmean->push_back(oh_mean); + b_reco_jet_emcal_zsig->push_back(em_sigma); + b_reco_jet_ohcal_zsig->push_back(oh_sigma); + b_reco_jet_emcal_zskew->push_back(em_skew); + b_reco_jet_ohcal_zskew->push_back(oh_skew); - if (jetpt > dijet_pt[0]) - { - dijet_pt[1]=dijet_pt[0]; - dijet_eta[1]= dijet_eta[0]; - dijet_phi[1]= dijet_phi[0]; - dijet_E[1]= dijet_E[0]; - dijet_em_frac[1]= dijet_em_frac[0]; - dijet_em_mean[1]= dijet_em_mean[0]; - dijet_em_skew[1]= dijet_em_skew[0]; - dijet_em_sig[1]= dijet_em_sig[0]; - dijet_oh_frac[1]= dijet_oh_frac[0]; - dijet_oh_mean[1]= dijet_oh_mean[0]; - dijet_oh_skew[1]= dijet_oh_skew[0]; - dijet_oh_sig[1]= dijet_oh_sig[0]; - - dijet_pt[0]=jetpt; - dijet_eta[0]= jeteta; - dijet_phi[0]= jetphi; - dijet_E[0]= jet_E; - dijet_em_frac[0]= jet_emcal; - dijet_em_mean[0]= em_mean; - dijet_em_skew[0]= em_skew; - dijet_em_sig[0]= em_sigma; - dijet_oh_frac[0]= jet_ohcal; - dijet_oh_mean[0]= oh_mean; - dijet_oh_skew[0]= oh_skew; - dijet_oh_sig[0]= oh_sigma; + b_reco_jet_emcal_eta->push_back(jet_emcal_eta); + b_reco_jet_ohcal_eta->push_back(jet_ohcal_eta); + b_reco_jet_ihcal_eta->push_back(jet_ihcal_eta); + + + if (jetpt > dijet_pt[0]) + { + dijet_pt[1] = dijet_pt[0]; + dijet_index[1] = dijet_index[0]; + dijet_pt[0] = jetpt; + dijet_index[0] = ijet; + } + else if (jetpt > dijet_pt[1]) + { + dijet_pt[1] = jetpt; + dijet_index[1] = ijet; + } + ijet++; } - else if (jetpt > dijet_pt[1]) - { - dijet_pt[1]=jetpt; - dijet_eta[1]= jeteta; - dijet_phi[1]= jetphi; - dijet_E[1]= jet_E; - dijet_em_frac[1]= jet_emcal; - dijet_em_mean[1]= em_mean; - dijet_em_skew[1]= em_skew; - dijet_em_sig[1]= em_sigma; - dijet_oh_frac[1]= jet_ohcal; - dijet_oh_mean[1]= oh_mean; - dijet_oh_skew[1]= oh_skew; - dijet_oh_sig[1]= oh_sigma; - } - } - double emcal_on[2] = {1}; - double ohcal_on[2] = {1}; + } + // TODO: fill `features` from topNode in the exact order documented in + // VertexMLP.h (emcal/ohcal lead+sublead flags, zmean/zsig/zskew, + // energy, exj). This is the same per-event extraction CaloVtxReco used + // to do inline before the NN call -- move that block here unchanged. + float e1 = dijet_pt[0] * cosh(b_reco_jet_eta->at(dijet_index[0])); + float e2 = dijet_pt[1] * cosh(b_reco_jet_eta->at(dijet_index[1])); + + double exj = e2/e1; + double emcal_lead_on = 1; + double emcal_sublead_on = 1; + double ohcal_lead_on = 1; + double ohcal_sublead_on = 1; - for (int i = 0 ; i < 2; i++) + if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal->at(dijet_index[0]))) emcal_lead_on = 0;//<"," + if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal->at(dijet_index[0]))) ohcal_lead_on = 0;//<"," + if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal->at(dijet_index[1]))) emcal_sublead_on = 0;//<"," + if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal->at(dijet_index[1]))) ohcal_sublead_on = 0;//<"," + + float emcal_lead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[0]), + b_reco_jet_emcal_zsig->at(dijet_index[0]), + b_reco_jet_emcal_zskew->at(dijet_index[0]), + b_reco_jet_emcal->at(dijet_index[0])}; + float emcal_sublead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[1]), + b_reco_jet_emcal_zsig->at(dijet_index[1]), + b_reco_jet_emcal_zskew->at(dijet_index[1]), + b_reco_jet_emcal->at(dijet_index[1])}; + float ohcal_lead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[0]), + b_reco_jet_ohcal_zsig->at(dijet_index[0]), + b_reco_jet_ohcal_zskew->at(dijet_index[0]), + b_reco_jet_ohcal->at(dijet_index[0])}; + float ohcal_sublead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[1]), + b_reco_jet_ohcal_zsig->at(dijet_index[1]), + b_reco_jet_ohcal_zskew->at(dijet_index[1]), + b_reco_jet_ohcal->at(dijet_index[1])}; + + if (!emcal_lead_on) { - if (std::isnan(dijet_em_frac[i]) || - std::isnan(dijet_em_mean[i]) || - std::isnan(dijet_em_skew[i]) || - std::isnan(dijet_em_sig[i])) + for (int i = 0; i < 4; i++) { - emcal_on[i] = 0; + emcal_lead_z_moments[i] = 0; } - if (std::isnan(dijet_oh_frac[i]) || - std::isnan(dijet_oh_mean[i]) || - std::isnan(dijet_oh_skew[i]) || - std::isnan(dijet_oh_sig[i])) + } + if (!emcal_sublead_on) + { + for (int i = 0; i < 4; i++) { - ohcal_on[i] = 0; + emcal_sublead_z_moments[i] = 0; } } - - float emcal_z_moments[2][4] = {0}; - float ohcal_z_moments[2][4] = {0}; - for (int i = 0; i < 2; i++) + if (!ohcal_lead_on) { - if (emcal_on[i]) + for (int i = 0; i < 4; i++) { - emcal_z_moments[i][0] = dijet_em_frac[i]; - emcal_z_moments[i][1] = dijet_em_mean[i]; - emcal_z_moments[i][2] = dijet_em_skew[i]; - emcal_z_moments[i][3] = dijet_em_sig[i]; + ohcal_lead_z_moments[i] = 0; } - if (ohcal_on[i]) + } + if (!ohcal_sublead_on) + { + for (int i = 0; i < 4; i++) { - ohcal_z_moments[i][0] = dijet_oh_frac[i]; - ohcal_z_moments[i][1] = dijet_oh_mean[i]; - ohcal_z_moments[i][2] = dijet_oh_skew[i]; - ohcal_z_moments[i][3] = dijet_oh_sig[i]; + ohcal_sublead_z_moments[i] = 0; } } - float exj = dijet_E[1]/dijet_E[0]; - - std::array features = { - emcal_on[0], - ohcal_on[0], - emcal_on[1], - ohcal_on[1], - emcal_z_moments[0][0], - emcal_z_moments[0][1], - emcal_z_moments[0][2], - emcal_z_moments[0][3], - ohcal_z_moments[0][0], - ohcal_z_moments[0][1], - ohcal_z_moments[0][2], - ohcal_z_moments[0][3], - emcal_z_moments[1][0], - emcal_z_moments[1][1], - emcal_z_moments[1][2], - emcal_z_moments[1][3], - ohcal_z_moments[1][0], - ohcal_z_moments[1][1], - ohcal_z_moments[1][2], - ohcal_z_moments[1][3], + std::cout << emcal_lead_on << "," + << ohcal_lead_on << "," + << emcal_sublead_on << "," + << ohcal_sublead_on << "," + << emcal_lead_z_moments[0] << "," + << emcal_lead_z_moments[1] << "," + << emcal_lead_z_moments[2] << "," + << emcal_lead_z_moments[3] << "," + << ohcal_lead_z_moments[0] << "," + << ohcal_lead_z_moments[1] << "," + << ohcal_lead_z_moments[2] << "," + << ohcal_lead_z_moments[3] << "," + << emcal_sublead_z_moments[0] << "," + << emcal_sublead_z_moments[1] << "," + << emcal_sublead_z_moments[2] << "," + << emcal_sublead_z_moments[3] << "," + << ohcal_sublead_z_moments[0] << "," + << ohcal_sublead_z_moments[1] << "," + << ohcal_sublead_z_moments[2] << "," + << ohcal_sublead_z_moments[3] << "," + << exj << std::endl; + + std::array features = {emcal_lead_on, + ohcal_lead_on, + emcal_sublead_on, + ohcal_sublead_on, + emcal_lead_z_moments[0], + emcal_lead_z_moments[1], + emcal_lead_z_moments[2], + emcal_lead_z_moments[3], + ohcal_lead_z_moments[0], + ohcal_lead_z_moments[1], + ohcal_lead_z_moments[2], + ohcal_lead_z_moments[3], + emcal_sublead_z_moments[0], + emcal_sublead_z_moments[1], + emcal_sublead_z_moments[2], + emcal_sublead_z_moments[3], + ohcal_sublead_z_moments[0], + ohcal_sublead_z_moments[1], + ohcal_sublead_z_moments[2], + ohcal_sublead_z_moments[3], exj}; zvtx = static_cast(VertexMLP::PredictVertexZ(features)); + return Fun4AllReturnCodes::EVENT_OK; } diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.h b/offline/packages/calovtxreco/CaloVtxAlgoMLP.h index de119e7fcf..b42b59479c 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoMLP.h +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.h @@ -13,13 +13,14 @@ class PHCompositeNode; class CaloVtxAlgoMLP : public CaloVtxAlgo { public: - explicit CaloVtxAlgoMLP(); + explicit CaloVtxAlgoMLP() = default; ~CaloVtxAlgoMLP() override = default; int Init(PHCompositeNode *topNode) override; int CalculateVertex(PHCompositeNode *topNode, float &zvtx) override; std::string Name() const override { return "MLP"; } - + VertexDefs::CALOALGO Algo() const override { return VertexDefs::CALOALGO::JETMLP; } + void setWeightsFile(std::string file) { m_weightsFile = file; } void setJetNode(std::string node) { m_jet_node = node; } private: @@ -28,6 +29,7 @@ class CaloVtxAlgoMLP : public CaloVtxAlgo std::string m_weightsFile{"vertex_mlp_weights.root"}; float m_radius_EM{0}; + float m_radius_IH{0}; float m_radius_OH{0}; }; diff --git a/offline/packages/calovtxreco/CaloVtxReco.cc b/offline/packages/calovtxreco/CaloVtxReco.cc index 0238754f6f..8faa83a42c 100644 --- a/offline/packages/calovtxreco/CaloVtxReco.cc +++ b/offline/packages/calovtxreco/CaloVtxReco.cc @@ -84,17 +84,21 @@ int CaloVtxReco::InitRun(PHCompositeNode *topNode) int CaloVtxReco::process_event(PHCompositeNode *topNode) { - float zvtx = std::numeric_limits::quiet_NaN(); for (auto &algo : m_algos) { float tempz = std::numeric_limits::quiet_NaN(); - algo->CalculateVertex(topNode, zvtx); - std::cout << algo->Name() <<": "<CalculateVertex(topNode, tempz); + CaloVertex *vertex = new CaloVertexv1(); + + vertex->set_z(tempz); + vertex->set_z_err(0); + vertex->set_t(0); + vertex->set_t_err(0); + vertex->set_calo_algo(algo->Algo()); + m_calovtxmap->insert(vertex); } - CaloVertex *vertex = new CaloVertexv1(); - vertex->set_z(zvtx); - m_calovtxmap->insert(vertex); + return Fun4AllReturnCodes::EVENT_OK; } diff --git a/offline/packages/globalvertex/CaloVertex.h b/offline/packages/globalvertex/CaloVertex.h index d226636c4e..1561ab8907 100644 --- a/offline/packages/globalvertex/CaloVertex.h +++ b/offline/packages/globalvertex/CaloVertex.h @@ -11,6 +11,7 @@ class CaloVertex : public Vertex { public: + ~CaloVertex() override {} // PHObject virtual overloads @@ -20,6 +21,9 @@ class CaloVertex : public Vertex int isValid() const override { return 0; } // vertex info + virtual void set_calo_algo(VertexDefs::CALOALGO) override {} + virtual VertexDefs::CALOALGO get_calo_algo() const override { return VertexDefs::CALOALGO::UNDEFINED; } + virtual unsigned int get_id() const override { return std::numeric_limits::max(); } virtual void set_id(unsigned int) override {} diff --git a/offline/packages/globalvertex/CaloVertexv1.cc b/offline/packages/globalvertex/CaloVertexv1.cc index 58feb5f7c7..2270ac138f 100644 --- a/offline/packages/globalvertex/CaloVertexv1.cc +++ b/offline/packages/globalvertex/CaloVertexv1.cc @@ -5,6 +5,7 @@ void CaloVertexv1::identify(std::ostream& os) const { os << "---CaloVertexv1--------------------------------" << std::endl; + os << "algo: " << get_calo_algo() << std::endl; os << "vertexid: " << get_id() << std::endl; os << " t = " << get_t() << " +/- " << get_t_err() << std::endl; os << " z = " << get_z() << " +/- " << get_z_err() << std::endl; diff --git a/offline/packages/globalvertex/CaloVertexv1.h b/offline/packages/globalvertex/CaloVertexv1.h index b86776752d..5ce4623722 100644 --- a/offline/packages/globalvertex/CaloVertexv1.h +++ b/offline/packages/globalvertex/CaloVertexv1.h @@ -23,6 +23,9 @@ class CaloVertexv1 : public CaloVertex // vertex info + void set_calo_algo(VertexDefs::CALOALGO algo) override { _algo = algo; } + VertexDefs::CALOALGO get_calo_algo() const override { return _algo; } + unsigned int get_id() const override { return _id; } void set_id(unsigned int id) override { _id = id; } @@ -41,6 +44,9 @@ class CaloVertexv1 : public CaloVertex float get_position(unsigned int coor) const override; private: + + VertexDefs::CALOALGO _algo{VertexDefs::CALOALGO::UNDEFINED}; + unsigned int _id{std::numeric_limits::max()}; //< unique identifier within container float _t{std::numeric_limits::quiet_NaN()}; //< collision time float _t_err{std::numeric_limits::quiet_NaN()}; //< collision time uncertainty diff --git a/offline/packages/globalvertex/GlobalVertex.h b/offline/packages/globalvertex/GlobalVertex.h index da34dd625a..197ae3a3b3 100644 --- a/offline/packages/globalvertex/GlobalVertex.h +++ b/offline/packages/globalvertex/GlobalVertex.h @@ -25,6 +25,9 @@ class GlobalVertex : public PHObject SVTX = 400, SVTX_MBD = 500, CALO = 250, + CALO_JETSKEW = 251, + CALO_AVGZ = 252, + CALO_JETMLP = 253, MBD_CALO = 350 }; diff --git a/offline/packages/globalvertex/GlobalVertexReco.cc b/offline/packages/globalvertex/GlobalVertexReco.cc index c34f4e1098..d233f0c428 100644 --- a/offline/packages/globalvertex/GlobalVertexReco.cc +++ b/offline/packages/globalvertex/GlobalVertexReco.cc @@ -1,9 +1,10 @@ #include "GlobalVertexReco.h" +#include "VertexDefs.h" //#include "GlobalVertex.h" // for GlobalVertex, GlobalVe... #include "GlobalVertexMap.h" // for GlobalVertexMap #include "GlobalVertexMapv1.h" -#include "GlobalVertexv3.h" +#include "GlobalVertexv4.h" #include "MbdVertex.h" #include "MbdVertexMap.h" #include "CaloVertexv1.h" @@ -140,7 +141,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) } // we have a matching pair - GlobalVertex *vertex = new GlobalVertexv3(); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->set_id(globalmap->size()); vertex->clone_insert_vtx(GlobalVertex::SVTX, svtx); @@ -193,7 +194,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) } // we have a standalone SVTX vertex - GlobalVertex *vertex = new GlobalVertexv3(); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->set_id(globalmap->size()); @@ -243,10 +244,11 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) continue; } - GlobalVertex *vertex = new GlobalVertexv3(); - vertex->set_id(globalmap->size()); + GlobalVertex *vertex = new GlobalVertexv4(); + vertex->clone_insert_vtx(GlobalVertex::MBD, mbd); + vertex->set_id(globalmap->size()); used_mbd_vtxids.insert(mbd->get_id()); globalmap->insert(vertex); @@ -271,7 +273,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) ++caloiter) { const CaloVertex *calo = caloiter->second; - + if (used_calo_vtxids.contains(calo->get_id())) { continue; @@ -281,17 +283,34 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) { continue; } - - GlobalVertex *vertex = new GlobalVertexv3(); - vertex->set_id(globalmap->size()); - - vertex->clone_insert_vtx(GlobalVertex::CALO, calo); + + GlobalVertex *vertex = new GlobalVertexv4(); + + auto caloalgo = calo->get_calo_algo(); + if (caloalgo == VertexDefs::CALOALGO::UNDEFINED) + { + vertex->clone_insert_vtx(GlobalVertex::CALO, calo); + } + if (caloalgo == VertexDefs::CALOALGO::JETSKEW) + { + vertex->clone_insert_vtx(GlobalVertex::CALO_JETSKEW, calo); + } + if (caloalgo == VertexDefs::CALOALGO::AVGZ) + { + vertex->clone_insert_vtx(GlobalVertex::CALO_AVGZ, calo); + } + if (caloalgo == VertexDefs::CALOALGO::JETMLP) + { + vertex->clone_insert_vtx(GlobalVertex::CALO_JETMLP, calo); + } + vertex->set_id(globalmap->size()); + used_calo_vtxids.insert(calo->get_id()); globalmap->insert(vertex); if (Verbosity() > 1) - { + { vertex->identify(); } } @@ -307,12 +326,15 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) CaloVertex *cvertex = new CaloVertexv1(); cvertex->set_z(0); - - GlobalVertex *vertex = new GlobalVertexv3(); - vertex->set_id(globalmap->size()); + cvertex->set_t(0); + cvertex->set_z_err(0); + cvertex->set_t_err(0); + cvertex->set_id(0); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->clone_insert_vtx(GlobalVertex::ZERO, cvertex); - used_calo_vtxids.insert(cvertex->get_id()); + vertex->set_id(globalmap->size()); + //used_calo_vtxids.insert(cvertex->get_id()); globalmap->insert(vertex); @@ -363,7 +385,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) continue; } - GlobalVertex *vertex = new GlobalVertexv3(); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->set_id(globalmap->size()); vertex->clone_insert_vtx(GlobalVertex::CALO, calo); @@ -380,7 +402,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) } else { - GlobalVertex *vertex = new GlobalVertexv3(); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->set_id(globalmap->size()); vertex->clone_insert_vtx(GlobalVertex::MBD, mbd); @@ -419,7 +441,7 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) tvertex->set_t(0); tvertex->set_t_err(0); // 0.1 - GlobalVertex *vertex = new GlobalVertexv3(); + GlobalVertex *vertex = new GlobalVertexv4(); vertex->clone_insert_vtx(GlobalVertex::TRUTH, tvertex); globalmap->insert(vertex); if (truthmap) diff --git a/offline/packages/globalvertex/GlobalVertexv4.cc b/offline/packages/globalvertex/GlobalVertexv4.cc new file mode 100644 index 0000000000..d6e201fe35 --- /dev/null +++ b/offline/packages/globalvertex/GlobalVertexv4.cc @@ -0,0 +1,261 @@ +#include "GlobalVertexv4.h" + +#include + +GlobalVertexv4::GlobalVertexv4(const unsigned int id) + : _id(id) +{ +} + +GlobalVertexv4::~GlobalVertexv4() +{ + GlobalVertexv4::Reset(); +} + +void GlobalVertexv4::Reset() +{ + for (auto& _vtx : _vtxs) + { + for (const auto* vertex : _vtx.second) + { + delete vertex; + } + } + _vtxs.clear(); +} + +void GlobalVertexv4::identify(std::ostream& os) const +{ + os << "---GlobalVertexv4-----------------------" << std::endl; + + os << " list of vtx ids: " << std::endl; + for (ConstVertexIter iter = begin_vertexes(); iter != end_vertexes(); ++iter) + { + os << " Vertex type " << iter->first << " has " << iter->second.size() + << " vertices associated to it" << std::endl; + for (const auto& vertex : iter->second) + { + vertex->identify(); + } + } + + os << "-----------------------------------------------" << std::endl; +} + +int GlobalVertexv4::isValid() const +{ + if (_vtxs.empty()) + { + return 0; + } + return 1; +} + +void GlobalVertexv4::insert_vtx(GlobalVertex::VTXTYPE type, const Vertex* vertex) +{ + auto it = _vtxs.find(type); + if (it == _vtxs.end()) + { + VertexVector vector; + vector.push_back(vertex); + _vtxs.insert(std::make_pair(type, vector)); + return; + } + + it->second.push_back(vertex); +} + +void GlobalVertexv4::clone_insert_vtx(GlobalVertex::VTXTYPE type, const Vertex* vertex) +{ + auto it = _vtxs.find(type); + Vertex* clone = dynamic_cast(vertex->CloneMe()); + if (it == _vtxs.end()) + { + VertexVector vector; + vector.push_back(clone); + _vtxs.insert(std::make_pair(type, vector)); + return; + } + + it->second.push_back(clone); +} + +size_t GlobalVertexv4::count_vtxs(GlobalVertex::VTXTYPE type) const +{ + auto it = _vtxs.find(type); + if (it == _vtxs.end()) + { + return 0; + } + + return it->second.size(); +} + +float GlobalVertexv4::get_t() const +{ + auto it = _vtxs.find(GlobalVertex::VTXTYPE::MBD); + if (it == _vtxs.end()) + { + return std::numeric_limits::quiet_NaN(); + } + + return it->second[0]->get_t(); +} + +float GlobalVertexv4::get_t_err() const +{ + auto it = _vtxs.find(GlobalVertex::VTXTYPE::MBD); + if (it == _vtxs.end()) + { + return std::numeric_limits::quiet_NaN(); + } + + return it->second[0]->get_t_err(); +} + +float GlobalVertexv4::get_x() const { return get_position(0); } +float GlobalVertexv4::get_y() const { return get_position(1); } +float GlobalVertexv4::get_z() const { return get_position(2); } + +float GlobalVertexv4::get_position(unsigned int coor) const +{ + auto svtxit = _vtxs.find(GlobalVertex::VTXTYPE::SVTX); + if (svtxit == _vtxs.end()) + { + auto mbdit = _vtxs.find(GlobalVertex::VTXTYPE::MBD); + if (mbdit == _vtxs.end()) + { + + GlobalVertex::ConstVertexIter caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_JETSKEW); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + + caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_AVGZ); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + + caloit = find_vertexes(GlobalVertex::VTXTYPE::CALO_JETMLP); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + caloit = find_vertexes(GlobalVertex::VTXTYPE::ZERO); + if (caloit != _vtxs.end()) + { + return caloit->second[0]->get_position(coor); + } + + auto truthit = _vtxs.find(GlobalVertex::VTXTYPE::TRUTH); + if (truthit == _vtxs.end()) + { + return std::numeric_limits::quiet_NaN(); + } + return truthit->second[0]->get_position(coor); + } + return mbdit->second[0]->get_position(coor); + } + + GlobalVertex::VertexVector trackvertices = svtxit->second; + size_t mosttracks = 0; + float pos = std::numeric_limits::quiet_NaN(); + for (const auto* vertex : trackvertices) + { + if (vertex->size_tracks() > mosttracks) + { + mosttracks = vertex->size_tracks(); + pos = vertex->get_position(coor); + } + } + + return pos; +} + +float GlobalVertexv4::get_chisq() const +{ + auto svtxit = _vtxs.find(GlobalVertex::VTXTYPE::SVTX); + if (svtxit == _vtxs.end()) + { + return std::numeric_limits::quiet_NaN(); + } + + GlobalVertex::VertexVector trackvertices = svtxit->second; + size_t mosttracks = 0; + float chisq = std::numeric_limits::quiet_NaN(); + for (const auto* vertex : trackvertices) + { + if (vertex->size_tracks() > mosttracks) + { + mosttracks = vertex->size_tracks(); + chisq = vertex->get_chisq(); + } + } + + return chisq; +} + +unsigned int GlobalVertexv4::get_ndof() const +{ + auto svtxit = _vtxs.find(GlobalVertex::VTXTYPE::SVTX); + if (svtxit == _vtxs.end()) + { + return std::numeric_limits::max(); + } + + GlobalVertex::VertexVector trackvertices = svtxit->second; + size_t mosttracks = 0; + unsigned int ndf = std::numeric_limits::max(); + for (const auto* vertex : trackvertices) + { + if (vertex->size_tracks() > mosttracks) + { + mosttracks = vertex->size_tracks(); + ndf = vertex->get_ndof(); + } + } + + return ndf; +} + +float GlobalVertexv4::get_error(unsigned int i, unsigned int j) const +{ + auto svtxit = _vtxs.find(GlobalVertex::VTXTYPE::SVTX); + if (svtxit == _vtxs.end()) + { + auto mbdit = _vtxs.find(GlobalVertex::VTXTYPE::MBD); + if (mbdit == _vtxs.end()) + { + return std::numeric_limits::quiet_NaN(); + } + // MBD only has z error defined + if (i == 2 && j == 2) + { + return mbdit->second[0]->get_z_err(); + } + + return std::numeric_limits::quiet_NaN(); + } + + GlobalVertex::VertexVector trackvertices = svtxit->second; + size_t mosttracks = 0; + float err = std::numeric_limits::quiet_NaN(); + for (const auto* vertex : trackvertices) + { + if (vertex->size_tracks() > mosttracks) + { + mosttracks = vertex->size_tracks(); + err = vertex->get_error(i, j); + } + } + + return err; +} + diff --git a/offline/packages/globalvertex/GlobalVertexv4.h b/offline/packages/globalvertex/GlobalVertexv4.h new file mode 100644 index 0000000000..a1ced23697 --- /dev/null +++ b/offline/packages/globalvertex/GlobalVertexv4.h @@ -0,0 +1,73 @@ +// Tell emacs that this is a C++ source +// -*- C++ -*-. +#ifndef GLOBALVERTEX_GLOBALVERTEXV4_H +#define GLOBALVERTEX_GLOBALVERTEXV4_H + +#include "GlobalVertex.h" + +#include // for size_t +#include +#include +#include + +class PHObject; + +class GlobalVertexv4 : public GlobalVertex +{ + public: + GlobalVertexv4() = default; + GlobalVertexv4(const unsigned int id); + ~GlobalVertexv4() override; + + // PHObject virtual overloads + void identify(std::ostream& os = std::cout) const override; + void Reset() override; + int isValid() const override; + PHObject* CloneMe() const override { return new GlobalVertexv4(*this); } + + unsigned int get_id() const override { return _id; } + void set_id(unsigned int id) override { _id = id; } + + short int get_beam_crossing() const override { return _bco; } + void set_beam_crossing(short int bco) override { _bco = bco; } + + float get_t() const override; + float get_t_err() const override; + float get_x() const override; + float get_y() const override; + float get_z() const override; + float get_chisq() const override; + unsigned int get_ndof() const override; + float get_position(unsigned int coor) const override; + float get_error(unsigned int i, unsigned int j) const override; + + // + // associated vertex methods + // + bool empty_vtxs() const override { return _vtxs.empty(); } + size_t size_vtxs() const override { return _vtxs.size(); } + size_t count_vtxs(GlobalVertex::VTXTYPE type) const override; + + void clear_vtxs() override { _vtxs.clear(); } + void insert_vtx(GlobalVertex::VTXTYPE type, const Vertex* vertex) override; + void clone_insert_vtx(GlobalVertex::VTXTYPE type, const Vertex* vertex) override; + size_t erase_vtxs(GlobalVertex::VTXTYPE type) override { return _vtxs.erase(type); } + void erase_vtxs(GlobalVertex::VertexIter iter) override { _vtxs.erase(iter); } + + GlobalVertex::ConstVertexIter begin_vertexes() const override { return _vtxs.begin(); } + GlobalVertex::ConstVertexIter find_vertexes(GlobalVertex::VTXTYPE type) const override { return _vtxs.find(type); } + GlobalVertex::ConstVertexIter end_vertexes() const override { return _vtxs.end(); } + + GlobalVertex::VertexIter begin_vertexes() override { return _vtxs.begin(); } + GlobalVertex::VertexIter find_vertexes(GlobalVertex::VTXTYPE type) override { return _vtxs.find(type); } + GlobalVertex::VertexIter end_vertexes() override { return _vtxs.end(); } + + private: + unsigned int _id{std::numeric_limits::max()}; + short int _bco{std::numeric_limits::max()}; //< global bco (signed short) + std::map _vtxs; //< list of vtxs + + ClassDefOverride(GlobalVertexv4, 3); +}; + +#endif diff --git a/offline/packages/globalvertex/GlobalVertexv4LinkDef.h b/offline/packages/globalvertex/GlobalVertexv4LinkDef.h new file mode 100644 index 0000000000..1d2c4bfedc --- /dev/null +++ b/offline/packages/globalvertex/GlobalVertexv4LinkDef.h @@ -0,0 +1,5 @@ +#ifdef __CINT__ + +#pragma link C++ class GlobalVertexv4 + ; + +#endif /* __CINT__ */ diff --git a/offline/packages/globalvertex/Makefile.am b/offline/packages/globalvertex/Makefile.am index 0ff0587389..07ecc346c7 100644 --- a/offline/packages/globalvertex/Makefile.am +++ b/offline/packages/globalvertex/Makefile.am @@ -23,17 +23,19 @@ libglobalvertex_la_LIBADD = \ -ltrackbase_historic_io pkginclude_HEADERS = \ - CaloVertex.h \ - CaloVertexv1.h \ - CaloVertexMap.h \ - CaloVertexMapv1.h \ + VertexDefs.h \ GlobalVertex.h \ GlobalVertexv1.h \ GlobalVertexv2.h \ GlobalVertexv3.h \ + GlobalVertexv4.h \ GlobalVertexMap.h \ GlobalVertexMapv1.h \ GlobalVertexReco.h \ + CaloVertex.h \ + CaloVertexv1.h \ + CaloVertexMap.h \ + CaloVertexMapv1.h \ MbdVertex.h \ MbdVertexv1.h \ MbdVertexv2.h \ @@ -53,16 +55,17 @@ pkginclude_HEADERS = \ Vertex.h ROOTDICTS = \ - CaloVertex_Dict.cc \ - CaloVertexv1_Dict.cc \ - CaloVertexMap_Dict.cc \ - CaloVertexMapv1_Dict.cc \ GlobalVertex_Dict.cc \ GlobalVertexv1_Dict.cc \ GlobalVertexv2_Dict.cc \ GlobalVertexv3_Dict.cc \ + GlobalVertexv4_Dict.cc \ GlobalVertexMap_Dict.cc \ GlobalVertexMapv1_Dict.cc \ + CaloVertex_Dict.cc \ + CaloVertexv1_Dict.cc \ + CaloVertexMap_Dict.cc \ + CaloVertexMapv1_Dict.cc \ MbdVertex_Dict.cc \ MbdVertexv1_Dict.cc \ MbdVertexv2_Dict.cc \ @@ -87,15 +90,16 @@ nobase_dist_pcm_DATA = $(ROOTDICTS:.cc=_rdict.pcm) libglobalvertex_io_la_SOURCES = \ $(ROOTDICTS) \ - CaloVertexv1.cc \ - CaloVertexMap.cc \ - CaloVertexMapv1.cc \ GlobalVertex.cc \ GlobalVertexv1.cc \ GlobalVertexv2.cc \ GlobalVertexv3.cc \ + GlobalVertexv4.cc \ GlobalVertexMap.cc \ GlobalVertexMapv1.cc \ + CaloVertexv1.cc \ + CaloVertexMap.cc \ + CaloVertexMapv1.cc \ MbdVertexv1.cc \ MbdVertexv2.cc \ MbdVertexv3.cc \ diff --git a/offline/packages/globalvertex/Vertex.h b/offline/packages/globalvertex/Vertex.h index e475bfb95d..a9a3c0e9aa 100644 --- a/offline/packages/globalvertex/Vertex.h +++ b/offline/packages/globalvertex/Vertex.h @@ -9,6 +9,7 @@ #include #include #include +#include "VertexDefs.h" class Vertex : public PHObject { @@ -74,6 +75,9 @@ class Vertex : public PHObject virtual float get_bbc_q(int) const { return std::numeric_limits::quiet_NaN(); } virtual float get_bbc_t(int) const { return std::numeric_limits::quiet_NaN(); } + virtual void set_calo_algo(VertexDefs::CALOALGO) {} + virtual VertexDefs::CALOALGO get_calo_algo() const { return VertexDefs::CALOALGO::UNDEFINED; } + // svtxvertex methods virtual void clear_tracks() {} virtual bool empty_tracks() { return true; } diff --git a/offline/packages/globalvertex/VertexDefs.h b/offline/packages/globalvertex/VertexDefs.h new file mode 100644 index 0000000000..469013c6e8 --- /dev/null +++ b/offline/packages/globalvertex/VertexDefs.h @@ -0,0 +1,14 @@ +#ifndef GLOBALVERTEX_VERTEXDEFS_H +#define GLOBALVERTEX_VERTEXDEFS_H + +namespace VertexDefs +{ + enum CALOALGO + { + UNDEFINED=0, + JETSKEW=1, + AVGZ=2, + JETMLP=3 + }; +}; +#endif From 7359dbfd38fa0deee8af9c3f4f8e37171dd175dc Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Fri, 10 Jul 2026 09:55:26 -0400 Subject: [PATCH 3/8] djl -- make it EXACTLY like offline --- .../packages/calovtxreco/CaloVtxAlgoMLP.cc | 45 ++++++++++--------- 1 file changed, 23 insertions(+), 22 deletions(-) diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc index adfe981863..fee405ed61 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -364,29 +364,30 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) ohcal_sublead_z_moments[i] = 0; } } + // { + // std::cout << emcal_lead_on << "," + // << ohcal_lead_on << "," + // << emcal_sublead_on << "," + // << ohcal_sublead_on << "," + // << emcal_lead_z_moments[0] << "," + // << emcal_lead_z_moments[1] << "," + // << emcal_lead_z_moments[2] << "," + // << emcal_lead_z_moments[3] << "," + // << ohcal_lead_z_moments[0] << "," + // << ohcal_lead_z_moments[1] << "," + // << ohcal_lead_z_moments[2] << "," + // << ohcal_lead_z_moments[3] << "," + // << emcal_sublead_z_moments[0] << "," + // << emcal_sublead_z_moments[1] << "," + // << emcal_sublead_z_moments[2] << "," + // << emcal_sublead_z_moments[3] << "," + // << ohcal_sublead_z_moments[0] << "," + // << ohcal_sublead_z_moments[1] << "," + // << ohcal_sublead_z_moments[2] << "," + // << ohcal_sublead_z_moments[3] << "," + // << exj << std::endl; + // } - std::cout << emcal_lead_on << "," - << ohcal_lead_on << "," - << emcal_sublead_on << "," - << ohcal_sublead_on << "," - << emcal_lead_z_moments[0] << "," - << emcal_lead_z_moments[1] << "," - << emcal_lead_z_moments[2] << "," - << emcal_lead_z_moments[3] << "," - << ohcal_lead_z_moments[0] << "," - << ohcal_lead_z_moments[1] << "," - << ohcal_lead_z_moments[2] << "," - << ohcal_lead_z_moments[3] << "," - << emcal_sublead_z_moments[0] << "," - << emcal_sublead_z_moments[1] << "," - << emcal_sublead_z_moments[2] << "," - << emcal_sublead_z_moments[3] << "," - << ohcal_sublead_z_moments[0] << "," - << ohcal_sublead_z_moments[1] << "," - << ohcal_sublead_z_moments[2] << "," - << ohcal_sublead_z_moments[3] << "," - << exj << std::endl; - std::array features = {emcal_lead_on, ohcal_lead_on, emcal_sublead_on, From 0224b190fc226eae34642d13fc3b71327a9d350c Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Fri, 10 Jul 2026 10:10:33 -0400 Subject: [PATCH 4/8] djl -- instructions --- offline/packages/globalvertex/VertexDefs.h | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/offline/packages/globalvertex/VertexDefs.h b/offline/packages/globalvertex/VertexDefs.h index 469013c6e8..5143fec3d8 100644 --- a/offline/packages/globalvertex/VertexDefs.h +++ b/offline/packages/globalvertex/VertexDefs.h @@ -1,6 +1,13 @@ #ifndef GLOBALVERTEX_VERTEXDEFS_H #define GLOBALVERTEX_VERTEXDEFS_H +// Namespace is used to make the Global Vertex and Calo Vertex object not depend +//on each other but allows for CaloVtxReco to pass on which algorithm was used to +// generate the z-vertex + +// If you'd like to add an algorithm add it here, then to GlobalVertex::VTXTYPE +// and also GlobalVertexv4::get_position() + namespace VertexDefs { enum CALOALGO From bebb6c109b13d03824e1d9c467e2bf3e5e75a290 Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Fri, 10 Jul 2026 11:08:36 -0400 Subject: [PATCH 5/8] djl -- clang --- .../calovtxreco/CaloVtxAlgoJetSkew.cc | 2 +- .../packages/calovtxreco/CaloVtxAlgoMLP.cc | 48 +++++++++---------- offline/packages/calovtxreco/CaloVtxReco.h | 2 +- 3 files changed, 26 insertions(+), 26 deletions(-) diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc index 6c3a093508..01cc5fdbe7 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc @@ -151,7 +151,7 @@ int CaloVtxAlgoJetSkew::CalculateVertex(PHCompositeNode *topNode, float &zvtx) zvtx = testz; } } - if (fabs(zvtx) >= 305) + if (std::fabs(zvtx) >= 305) { zvtx = std::numeric_limits::quiet_NaN(); } diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc index fee405ed61..499e8bb16d 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -302,22 +302,22 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) double ohcal_lead_on = 1; double ohcal_sublead_on = 1; - if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[0]))) emcal_lead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal->at(dijet_index[0]))) emcal_lead_on = 0;//<"," - if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[0]))) ohcal_lead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal->at(dijet_index[0]))) ohcal_lead_on = 0;//<"," - if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[1]))) emcal_sublead_on = 0;//maxz_EM <<"," - if (std::isnan(b_reco_jet_emcal->at(dijet_index[1]))) emcal_sublead_on = 0;//<"," - if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[1]))) ohcal_sublead_on = 0;//maxz_OH <<"," - if (std::isnan(b_reco_jet_ohcal->at(dijet_index[1]))) ohcal_sublead_on = 0;//<"," + if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[0]))) { emcal_lead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[0]))) { emcal_lead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[0]))) { emcal_lead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal->at(dijet_index[0]))) { emcal_lead_on = 0; }//<"," + if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[0]))) { ohcal_lead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[0]))) { ohcal_lead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[0]))) { ohcal_lead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal->at(dijet_index[0]))) { ohcal_lead_on = 0; }//<"," + if (std::isnan(b_reco_jet_emcal_zmean->at(dijet_index[1]))) { emcal_sublead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zsig->at(dijet_index[1]))) { emcal_sublead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal_zskew->at(dijet_index[1]))) { emcal_sublead_on = 0; }//maxz_EM <<"," + if (std::isnan(b_reco_jet_emcal->at(dijet_index[1]))) { emcal_sublead_on = 0; }//<"," + if (std::isnan(b_reco_jet_ohcal_zmean->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zsig->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//maxz_OH <<"," + if (std::isnan(b_reco_jet_ohcal->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//<"," float emcal_lead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[0]), b_reco_jet_emcal_zsig->at(dijet_index[0]), @@ -338,30 +338,30 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) if (!emcal_lead_on) { - for (int i = 0; i < 4; i++) + for (float & emcal_lead_z_moment : emcal_lead_z_moments) { - emcal_lead_z_moments[i] = 0; + emcal_lead_z_moment = 0; } } if (!emcal_sublead_on) { - for (int i = 0; i < 4; i++) + for (float & emcal_sublead_z_moment : emcal_sublead_z_moments) { - emcal_sublead_z_moments[i] = 0; + emcal_sublead_z_moment = 0; } } if (!ohcal_lead_on) { - for (int i = 0; i < 4; i++) + for (float & ohcal_lead_z_moment : ohcal_lead_z_moments) { - ohcal_lead_z_moments[i] = 0; + ohcal_lead_z_moment = 0; } } if (!ohcal_sublead_on) { - for (int i = 0; i < 4; i++) + for (float & ohcal_sublead_z_moment : ohcal_sublead_z_moments) { - ohcal_sublead_z_moments[i] = 0; + ohcal_sublead_z_moment = 0; } } // { diff --git a/offline/packages/calovtxreco/CaloVtxReco.h b/offline/packages/calovtxreco/CaloVtxReco.h index 31ab2502be..eedceced0f 100644 --- a/offline/packages/calovtxreco/CaloVtxReco.h +++ b/offline/packages/calovtxreco/CaloVtxReco.h @@ -32,7 +32,7 @@ class CaloVtxReco : public SubsysReco void registerAlgo(CaloVtxAlgo* algo) { m_algos.push_back(std::move(algo)); } private: - std::vector m_algos; + std::vector m_algos{}; CaloVertexMap *m_calovtxmap{nullptr}; }; From a1b46f65d373f08094d7f3bced187a80c3a563c1 Mon Sep 17 00:00:00 2001 From: Daniel J Lis Date: Fri, 10 Jul 2026 12:37:17 -0400 Subject: [PATCH 6/8] djl -- clang tidy --- .../calovtxreco/CaloVtxAlgoJetSkew.cc | 20 ++-- .../packages/calovtxreco/CaloVtxAlgoJetSkew.h | 4 +- .../packages/calovtxreco/CaloVtxAlgoMLP.cc | 97 +++++++++++-------- 3 files changed, 68 insertions(+), 53 deletions(-) diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc index 01cc5fdbe7..362a2a80ce 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc @@ -41,13 +41,13 @@ int CaloVtxAlgoJetSkew::CalculateVertex(PHCompositeNode *topNode, float &zvtx) JetContainer *jetcon = findNode::getClass(topNode, m_jetnodename); - towers[0] = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC_RETOWER"); - towers[1] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); - towers[2] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); + m_towers[0] = findNode::getClass(topNode, "TOWERINFO_CALIB_CEMC_RETOWER"); + m_towers[1] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALIN"); + m_towers[2] = findNode::getClass(topNode, "TOWERINFO_CALIB_HCALOUT"); - geom[0] = findNode::getClass(topNode, "TOWERGEOM_CEMC"); - geom[1] = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); - geom[2] = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); + m_geom[0] = findNode::getClass(topNode, "TOWERGEOM_CEMC"); + m_geom[1] = findNode::getClass(topNode, "TOWERGEOM_HCALIN"); + m_geom[2] = findNode::getClass(topNode, "TOWERGEOM_HCALOUT"); const int nz = 601; const int njet = 2; @@ -116,24 +116,24 @@ int CaloVtxAlgoJetSkew::CalculateVertex(PHCompositeNode *topNode, float &zvtx) unsigned int channel = comp.second; if (comp.first == 7 || comp.first == 27) { - TowerInfo *tower = towers[2]->get_tower_at_channel(channel); + TowerInfo *tower = m_towers[2]->get_tower_at_channel(channel); if (tower->get_energy() < 0.1) { continue; } johsum[j] += tower->get_energy(); - float neweta = new_eta(channel, towers[2], geom[2], RawTowerDefs::CalorimeterId::HCALOUT, testz); + float neweta = new_eta(channel, m_towers[2], m_geom[2], RawTowerDefs::CalorimeterId::HCALOUT, testz); joheta[j] += neweta * tower->get_energy(); } if (comp.first == 13 || comp.first == 28 || comp.first == 25) { - TowerInfo *tower = towers[0]->get_tower_at_channel(channel); + TowerInfo *tower = m_towers[0]->get_tower_at_channel(channel); if (tower->get_energy() < 0.1) { continue; } jemsum[j] += tower->get_energy(); - float neweta = new_eta(channel, towers[0], geom[1], RawTowerDefs::CalorimeterId::HCALIN, testz); + float neweta = new_eta(channel, m_towers[0], m_geom[1], RawTowerDefs::CalorimeterId::HCALIN, testz); jemeta[j] += neweta * tower->get_energy(); } } diff --git a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h index d28814251d..cc05fa2992 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h @@ -46,8 +46,8 @@ class CaloVtxAlgoJetSkew : public CaloVtxAlgo float m_radius_EM{std::numeric_limits::quiet_NaN()}; float m_radius_OH{std::numeric_limits::quiet_NaN()}; - TowerInfoContainer *towers[3] = {0}; - RawTowerGeomContainer *geom[3] = {0}; + TowerInfoContainer *m_towers[3] = {0}; + RawTowerGeomContainer *m_geom[3] = {0}; }; diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc index 499e8bb16d..a2199193da 100644 --- a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -85,25 +85,28 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) std::vector *b_reco_jet_ohcal = new std::vector(); std::vector *b_reco_jet_ihcal = new std::vector(); - JetContainer *jetscon = findNode::getClass(topNode, m_jet_node.c_str()); + JetContainer *jetscon = findNode::getClass(topNode, m_jet_node); int ijet = 0; - double dijet_pt[2] = {0}; + float dijet_pt[2] = {0}; int dijet_index[2] = {0}; if (jetscon) { - for (auto jet : *jetscon) + for (auto *jet : *jetscon) { - double jet_E = 0; - double jet_emcal = 0; - double jet_ihcal = 0; - double jet_ohcal = 0; - double jet_emcal_eta = 0; - double jet_ihcal_eta = 0; - double jet_ohcal_eta = 0; - double jetpt = jet->get_pt(); - if (jetpt < 2) continue; + float jet_E = 0; + float jet_emcal = 0; + float jet_ihcal = 0; + float jet_ohcal = 0; + float jet_emcal_eta = 0; + float jet_ihcal_eta = 0; + float jet_ohcal_eta = 0; + float jetpt = jet->get_pt(); + if (jetpt < 2) + { + continue; + } double em_S0 = 0.0; double em_S1 = 0.0; @@ -134,14 +137,17 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) int ieta = hcalin_towers->getTowerEtaBin(towerkey); int iphi = hcalin_towers->getTowerPhiBin(towerkey); const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); - float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); - float tower_r = m_radius_IH; - float tower_z = tower_r*sinh(tower_eta); - float new_tower_z = tower_z;// - jet_vertex; - float new_tower_eta = asinh(new_tower_z/tower_r); + double tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); + double tower_r = m_radius_IH; + double tower_z = tower_r*sinh(tower_eta); + double new_tower_z = tower_z;// - jet_vertex; + double new_tower_eta = asinh(new_tower_z/tower_r); tower_e = tower->get_energy(); - if (tower_e < 0.005) continue; + if (tower_e < 0.005) + { + continue; + } jet_ihcal_eta += new_tower_eta*tower_e; jet_ihcal += tower_e; jet_E += tower_e; @@ -157,14 +163,17 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) int ieta = hcalout_towers->getTowerEtaBin(towerkey); int iphi = hcalout_towers->getTowerPhiBin(towerkey); const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALOUT, ieta, iphi); - float tower_eta = tower_geomOH->get_tower_geometry(key)->get_eta(); - float tower_r = m_radius_OH; - float tower_z = tower_r*sinh(tower_eta); - float new_tower_z = tower_z;// - jet_vertex; - float new_tower_eta = asinh(new_tower_z/tower_r); + double tower_eta = tower_geomOH->get_tower_geometry(key)->get_eta(); + double tower_r = m_radius_OH; + double tower_z = tower_r*sinh(tower_eta); + double new_tower_z = tower_z;// - jet_vertex; + double new_tower_eta = asinh(new_tower_z/tower_r); tower_e = tower->get_energy(); - if (tower_e < 0.035) continue; + if (tower_e < 0.035) + { + continue; + } jet_ohcal_eta += new_tower_eta*tower_e; jet_ohcal += tower_e; oh_S0 += tower_e; @@ -189,14 +198,17 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) int ieta = emcalre_towers->getTowerEtaBin(towerkey); int iphi = emcalre_towers->getTowerPhiBin(towerkey); const RawTowerDefs::keytype key = RawTowerDefs::encode_towerid(RawTowerDefs::CalorimeterId::HCALIN, ieta, iphi); - float tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); - float tower_r = m_radius_EM; - float tower_z = tower_r*sinh(tower_eta); - float new_tower_z = tower_z;// - jet_vertex; - float new_tower_eta = asinh(new_tower_z/tower_r); + double tower_eta = tower_geomIH->get_tower_geometry(key)->get_eta(); + double tower_r = m_radius_EM; + double tower_z = tower_r*sinh(tower_eta); + double new_tower_z = tower_z;// - jet_vertex; + double new_tower_eta = asinh(new_tower_z/tower_r); tower_e = tower->get_energy(); - if (tower_e < 0.068) continue; + if (tower_e < 0.068) + { + continue; + } jet_emcal_eta += new_tower_eta*tower_e; jet_emcal += tower_e; jet_E += tower_e; @@ -213,7 +225,10 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) itower++; } - if (itower == 0) continue; + if (itower == 0) + { + continue; + } double em_mean = em_S1/em_S0; @@ -293,8 +308,8 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) // VertexMLP.h (emcal/ohcal lead+sublead flags, zmean/zsig/zskew, // energy, exj). This is the same per-event extraction CaloVtxReco used // to do inline before the NN call -- move that block here unchanged. - float e1 = dijet_pt[0] * cosh(b_reco_jet_eta->at(dijet_index[0])); - float e2 = dijet_pt[1] * cosh(b_reco_jet_eta->at(dijet_index[1])); + double e1 = dijet_pt[0] * cosh(b_reco_jet_eta->at(dijet_index[0])); + double e2 = dijet_pt[1] * cosh(b_reco_jet_eta->at(dijet_index[1])); double exj = e2/e1; double emcal_lead_on = 1; @@ -319,47 +334,47 @@ int CaloVtxAlgoMLP::CalculateVertex(PHCompositeNode *topNode, float &zvtx) if (std::isnan(b_reco_jet_ohcal_zskew->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//maxz_OH <<"," if (std::isnan(b_reco_jet_ohcal->at(dijet_index[1]))) { ohcal_sublead_on = 0; }//<"," - float emcal_lead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[0]), + double emcal_lead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[0]), b_reco_jet_emcal_zsig->at(dijet_index[0]), b_reco_jet_emcal_zskew->at(dijet_index[0]), b_reco_jet_emcal->at(dijet_index[0])}; - float emcal_sublead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[1]), + double emcal_sublead_z_moments[4] = {b_reco_jet_emcal_zmean->at(dijet_index[1]), b_reco_jet_emcal_zsig->at(dijet_index[1]), b_reco_jet_emcal_zskew->at(dijet_index[1]), b_reco_jet_emcal->at(dijet_index[1])}; - float ohcal_lead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[0]), + double ohcal_lead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[0]), b_reco_jet_ohcal_zsig->at(dijet_index[0]), b_reco_jet_ohcal_zskew->at(dijet_index[0]), b_reco_jet_ohcal->at(dijet_index[0])}; - float ohcal_sublead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[1]), + double ohcal_sublead_z_moments[4] = {b_reco_jet_ohcal_zmean->at(dijet_index[1]), b_reco_jet_ohcal_zsig->at(dijet_index[1]), b_reco_jet_ohcal_zskew->at(dijet_index[1]), b_reco_jet_ohcal->at(dijet_index[1])}; if (!emcal_lead_on) { - for (float & emcal_lead_z_moment : emcal_lead_z_moments) + for (double & emcal_lead_z_moment : emcal_lead_z_moments) { emcal_lead_z_moment = 0; } } if (!emcal_sublead_on) { - for (float & emcal_sublead_z_moment : emcal_sublead_z_moments) + for (double & emcal_sublead_z_moment : emcal_sublead_z_moments) { emcal_sublead_z_moment = 0; } } if (!ohcal_lead_on) { - for (float & ohcal_lead_z_moment : ohcal_lead_z_moments) + for (double & ohcal_lead_z_moment : ohcal_lead_z_moments) { ohcal_lead_z_moment = 0; } } if (!ohcal_sublead_on) { - for (float & ohcal_sublead_z_moment : ohcal_sublead_z_moments) + for (double & ohcal_sublead_z_moment : ohcal_sublead_z_moments) { ohcal_sublead_z_moment = 0; } From a5e9c0cd0197cdc2c2e0c06e79baef12c77efab9 Mon Sep 17 00:00:00 2001 From: Nikhil Kumar Date: Sat, 11 Jul 2026 05:33:13 -0400 Subject: [PATCH 7/8] 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 From fd711ef29c297d6b4176fa265ccfcaeb64bd0b29 Mon Sep 17 00:00:00 2001 From: Hanpu Jiang Date: Sun, 12 Jul 2026 22:39:05 -0400 Subject: [PATCH 8/8] Add CaloVtxAlgoVit method --- .../packages/calovtxreco/CaloVtxAlgoVit.cc | 182 ++++++++++++++++++ offline/packages/calovtxreco/CaloVtxAlgoVit.h | 75 ++++++++ offline/packages/calovtxreco/Makefile.am | 4 +- offline/packages/globalvertex/GlobalVertex.h | 1 + .../packages/globalvertex/GlobalVertexReco.cc | 4 + .../packages/globalvertex/GlobalVertexv4.cc | 5 + offline/packages/globalvertex/VertexDefs.h | 3 +- 7 files changed, 272 insertions(+), 2 deletions(-) create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoVit.cc create mode 100644 offline/packages/calovtxreco/CaloVtxAlgoVit.h 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