diff --git a/offline/packages/calovtxreco/CaloVtxAlgo.h b/offline/packages/calovtxreco/CaloVtxAlgo.h new file mode 100644 index 0000000000..6b11d0103b --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgo.h @@ -0,0 +1,34 @@ +#ifndef CALOVTXALGO_H +#define CALOVTXALGO_H + +#include +#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; + + virtual VertexDefs::CALOALGO Algo() const = 0; +}; + +#endif 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/CaloVtxAlgoCaloZ.cc b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc new file mode 100644 index 0000000000..344f698ce8 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoCaloZ.cc @@ -0,0 +1,141 @@ +#include "CaloVtxAlgoCaloZ.h" + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +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..a13680d15e --- /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() = 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; } + + 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..362a2a80ce --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.cc @@ -0,0 +1,182 @@ +#include "CaloVtxAlgoJetSkew.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + + +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); + + 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"); + + 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; + 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 = 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, 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 = 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, m_towers[0], m_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 (std::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..cc05fa2992 --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoJetSkew.h @@ -0,0 +1,55 @@ +#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() = 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; } + 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 *m_towers[3] = {0}; + RawTowerGeomContainer *m_geom[3] = {0}; + + +}; + +#endif diff --git a/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc new file mode 100644 index 0000000000..a2199193da --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.cc @@ -0,0 +1,432 @@ +#include "CaloVtxAlgoMLP.h" +#include "VertexMLP.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +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_IH = tower_geomIH->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; + } + + + 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); + + int ijet = 0; + float dijet_pt[2] = {0}; + int dijet_index[2] = {0}; + + if (jetscon) + { + for (auto *jet : *jetscon) + { + 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; + 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) + { // IHcal + tower = hcalin_towers->get_tower_at_channel(channel); + + if (!tower || !tower_geomIH) + { + 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); + 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; + } + jet_ihcal_eta += new_tower_eta*tower_e; + jet_ihcal += tower_e; + 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); + 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; + } + 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); + 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; + } + 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; + } + + 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_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; + + + 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()); + + b_reco_jet_emcal->push_back(jet_emcal); + b_reco_jet_ohcal->push_back(jet_ohcal); + b_reco_jet_ihcal->push_back(jet_ihcal); + + 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); + + + 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++; + } + + } + + // 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. + 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; + double emcal_sublead_on = 1; + 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; }//<"," + + 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])}; + 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])}; + 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])}; + 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 (double & emcal_lead_z_moment : emcal_lead_z_moments) + { + emcal_lead_z_moment = 0; + } + } + if (!emcal_sublead_on) + { + for (double & emcal_sublead_z_moment : emcal_sublead_z_moments) + { + emcal_sublead_z_moment = 0; + } + } + if (!ohcal_lead_on) + { + for (double & ohcal_lead_z_moment : ohcal_lead_z_moments) + { + ohcal_lead_z_moment = 0; + } + } + if (!ohcal_sublead_on) + { + for (double & ohcal_sublead_z_moment : ohcal_sublead_z_moments) + { + ohcal_sublead_z_moment = 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::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 new file mode 100644 index 0000000000..b42b59479c --- /dev/null +++ b/offline/packages/calovtxreco/CaloVtxAlgoMLP.h @@ -0,0 +1,36 @@ +#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() = 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: + + 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_IH{0}; + float m_radius_OH{0}; +}; + +#endif 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/CaloVtxReco.cc b/offline/packages/calovtxreco/CaloVtxReco.cc index 6e8da31be7..8faa83a42c 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,42 @@ 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) + for (auto &algo : m_algos) { - if(Verbosity() > 0) + status = algo->Init(topNode); + if (status) { - std::cout << "CaloVtxReco::InitRun(): Missing tower geometry node for towergeomEM (address: " << tower_geomEM << ") or towergeomOH (address: " << tower_geomOH << ") - aborting run!" << std::endl; + return status; } - 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)) - { - if(Verbosity() > 0) - { - 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 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) + for (auto &algo : m_algos) { - 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 tempz = std::numeric_limits::quiet_NaN(); + algo->CalculateVertex(topNode, tempz); - if (jpt[0] == 0) - { - if (Verbosity() > 2) - { - std::cout << "NO JETS > 5 GeV!" << std::endl; - } - } + CaloVertex *vertex = new CaloVertexv1(); - 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 ((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; + 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); } - 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; - } - } + return Fun4AllReturnCodes::EVENT_OK; } diff --git a/offline/packages/calovtxreco/CaloVtxReco.h b/offline/packages/calovtxreco/CaloVtxReco.h index e1a7d2b76b..eedceced0f 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..2e15a94254 100644 --- a/offline/packages/calovtxreco/Makefile.am +++ b/offline/packages/calovtxreco/Makefile.am @@ -10,19 +10,32 @@ AM_LDFLAGS = \ -L$(OFFLINE_MAIN)/lib pkginclude_HEADERS = \ - CaloVtxReco.h + CaloVtxReco.h \ + CaloVtxAlgo.h \ + CaloVtxAlgoJetSkew.h \ + CaloVtxAlgoCaloZ.h \ + CaloVtxAlgoMLP.h \ + CaloVtxAlgoCNN.h \ + CaloVtxAlgoVit.h \ + VertexMLP.h lib_LTLIBRARIES = \ libcalovtxreco.la libcalovtxreco_la_SOURCES = \ - CaloVtxReco.cc + CaloVtxReco.cc \ + CaloVtxAlgoJetSkew.cc \ + CaloVtxAlgoCaloZ.cc \ + CaloVtxAlgoMLP.cc \ + CaloVtxAlgoCNN.cc \ + CaloVtxAlgoVit.cc libcalovtxreco_la_LIBADD = \ -lcalo_io \ -lSubsysReco \ -ljetbase \ - -lglobalvertex_io + -lglobalvertex_io \ + -lonnxruntime BUILT_SOURCES = testexternals.cc 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/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 3124cea282..f1469c9beb 100644 --- a/offline/packages/globalvertex/GlobalVertex.h +++ b/offline/packages/globalvertex/GlobalVertex.h @@ -19,11 +19,17 @@ class GlobalVertex : public PHObject { UNDEFINED = 0, TRUTH = 100, + ZERO = 150, SMEARED = 200, MBD = 300, SVTX = 400, SVTX_MBD = 500, CALO = 250, + CALO_JETSKEW = 251, + 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 73681759cc..ab89d473bd 100644 --- a/offline/packages/globalvertex/GlobalVertexReco.cc +++ b/offline/packages/globalvertex/GlobalVertexReco.cc @@ -1,12 +1,13 @@ #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 "CaloVertex.h" +#include "CaloVertexv1.h" #include "CaloVertexMap.h" #include "SvtxVertex.h" #include "SvtxVertexMap.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,22 +283,76 @@ int GlobalVertexReco::process_event(PHCompositeNode *topNode) { continue; } - - GlobalVertex *vertex = new GlobalVertexv3(); + + 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); + } + if (caloalgo == VertexDefs::CALOALGO::CNN) + { + 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()); - - vertex->clone_insert_vtx(GlobalVertex::CALO, calo); + used_calo_vtxids.insert(calo->get_id()); globalmap->insert(vertex); if (Verbosity() > 1) - { + { vertex->identify(); } } } + // 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); + 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); + vertex->set_id(globalmap->size()); + //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)) { @@ -337,7 +393,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); @@ -354,7 +410,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); @@ -393,7 +449,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..e5354b84cb --- /dev/null +++ b/offline/packages/globalvertex/GlobalVertexv4.cc @@ -0,0 +1,272 @@ +#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::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); + } + 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..849e8f96c6 --- /dev/null +++ b/offline/packages/globalvertex/VertexDefs.h @@ -0,0 +1,23 @@ +#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 + { + UNDEFINED=0, + JETSKEW=1, + AVGZ=2, + JETMLP=3, + CNN=4, + VIT=5 + }; +}; +#endif