algorithms/calorimetry/CalorimeterClusterRecoCoG.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2024 Sylvester Joosten, Chao, Chao Peng, Whitney Armstrong, Dhevan Gangadharan, Derek Anderson
0003
0004 /*
0005  *  Reconstruct the cluster with Center of Gravity method
0006  *  Logarithmic weighting is used for mimicking energy deposit in transverse direction
0007  *
0008  *  Author: Chao Peng (ANL), 09/27/2020
0009  */
0010
0011 #include <Evaluator/DD4hepUnits.h>
0012 #include <boost/algorithm/string/join.hpp>
0013 #include <boost/range/adaptor/map.hpp>
0014 #include <edm4eic/CalorimeterHitCollection.h>
0015 #include <edm4eic/Cov3f.h>
0016 #include <edm4eic/EDM4eicVersion.h>
0017 #include <edm4hep/RawCalorimeterHit.h>
0018 #include <edm4hep/SimCalorimeterHitCollection.h>
0019 #include <edm4hep/Vector3f.h>
0020 #include <edm4hep/utils/vector_utils.h>
0021 #include <podio/LinkNavigator.h>
0022 #include <podio/ObjectID.h>
0023 #include <podio/RelationRange.h>
0024 #include <podio/detail/Link.h>
0025 #include <algorithm>
0026 #include <cctype>
0027 #include <cstddef>
0028 #include <gsl/pointers>
0029 #include <limits>
0030 #include <map>
0031 #include <optional>
0032 #include <vector>
0033
0034 #include "CalorimeterClusterRecoCoG.h"
0035 #include "algorithms/calorimetry/CalorimeterClusterRecoCoGConfig.h"
0036
0037 namespace eicrecon {
0038
0039 using namespace dd4hep;
0040
0041 void CalorimeterClusterRecoCoG::init() {
0042   // select weighting method
0043   std::string ew = m_cfg.energyWeight;
0044   // make it case-insensitive
0045   std::ranges::transform(ew, ew.begin(), [](char s) { return std::tolower(s); });
0046   auto it = weightMethods.find(ew);
0047   if (it == weightMethods.end()) {
0048     error("Cannot find energy weighting method {}, choose one from [{}]", m_cfg.energyWeight,
0049           boost::algorithm::join(weightMethods | boost::adaptors::map_keys, ", "));
0050     return;
0051   }
0052   weightFunc = it->second;
0053 }
0054
0055 void CalorimeterClusterRecoCoG::process(const CalorimeterClusterRecoCoG::Input& input,
0056                                         const CalorimeterClusterRecoCoG::Output& output) const {
0057 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0058   const auto [proto, mchitlinks, mchitassociations] = input;
0059   auto [clusters, links, associations]              = output;
0060 #else
0061   const auto [proto, mchitassociations] = input;
0062   auto [clusters, associations]         = output;
0063 #endif
0064
0065 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0066   // Check if truth associations are possible
0067   const bool do_assoc = mchitlinks != nullptr && !mchitlinks->empty();
0068   if (!do_assoc) {
0069     debug("Provided MCRecoCalorimeterHitLink collection is empty. No truth associations "
0070           "will be performed.");
0071   }
0072   // Build fast lookup once per event using podio::LinkNavigator
0073   std::optional<podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>> link_nav;
0074   if (do_assoc) {
0075     link_nav.emplace(*mchitlinks);
0076   }
0077 #else
0078   const bool do_assoc = mchitassociations != nullptr && !mchitassociations->empty();
0079   if (!do_assoc) {
0080     debug("Provided MCRecoCalorimeterHitAssociation collection is empty. No truth associations "
0081           "will be performed.");
0082   }
0083 #endif
0084
0085   for (const auto& pcl : *proto) {
0086     // skip protoclusters with no hits
0087     if (pcl.hits_size() == 0) {
0088       continue;
0089     }
0090
0091     auto cl_opt = reconstruct(pcl);
0092     if (!cl_opt.has_value()) {
0093       continue;
0094     }
0095     auto cl = *std::move(cl_opt);
0096
0097     debug("{} hits: {} GeV, ({}, {}, {})", cl.getNhits(), cl.getEnergy() / dd4hep::GeV,
0098           cl.getPosition().x / dd4hep::mm, cl.getPosition().y / dd4hep::mm,
0099           cl.getPosition().z / dd4hep::mm);
0100     clusters->push_back(cl);
0101
0102     // If sim hits are available, associate cluster with MCParticle
0103 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0104     if (do_assoc) {
0105       associate(cl, mchitassociations, *link_nav, links, associations);
0106     }
0107 #else
0108     if (do_assoc) {
0109       associate(cl, mchitassociations, associations);
0110     }
0111 #endif
0112   }
0113 }
0114
0115 std::optional<edm4eic::MutableCluster>
0116 CalorimeterClusterRecoCoG::reconstruct(const edm4eic::ProtoCluster& pcl) const {
0117   edm4eic::MutableCluster cl;
0118   cl.setNhits(pcl.hits_size());
0119
0120   debug("hit size = {}", pcl.hits_size());
0121
0122   // no hits
0123   if (pcl.hits_size() == 0) {
0124     return {};
0125   }
0126
0127   // calculate total energy, find the cell with the maximum energy deposit
0128   float totalE = 0.;
0129   // Used to optionally constrain the cluster eta to those of the contributing hits
0130   float minHitEta = std::numeric_limits<float>::max();
0131   float maxHitEta = std::numeric_limits<float>::min();
0132   auto time       = 0;
0133   auto timeError  = 0;
0134   for (unsigned i = 0; i < pcl.getHits().size(); ++i) {
0135     const auto& hit   = pcl.getHits()[i];
0136     const auto weight = pcl.getWeights()[i];
0137     debug("hit energy = {} hit weight: {}", hit.getEnergy(), weight);
0138     auto energy = hit.getEnergy() * weight;
0139     totalE += energy;
0140     time += (hit.getTime() - time) * energy / totalE;
0141     cl.addToHits(hit);
0142     cl.addToHitContributions(energy);
0143     const float eta = edm4hep::utils::eta(hit.getPosition());
0144     minHitEta       = std::min(eta, minHitEta);
0145     maxHitEta       = std::max(eta, maxHitEta);
0146   }
0147   cl.setEnergy(totalE / m_cfg.sampFrac);
0148   cl.setEnergyError(0.);
0149   cl.setTime(time);
0150   cl.setTimeError(timeError);
0151
0152   // center of gravity with logarithmic weighting
0153   float tw = 0.;
0154   auto v   = cl.getPosition();
0155
0156   double logWeightBase = m_cfg.logWeightBase;
0157   if (!m_cfg.logWeightBaseCoeffs.empty()) {
0158     double l      = std::log(cl.getEnergy() / m_cfg.logWeightBase_Eref);
0159     logWeightBase = 0;
0160     for (std::size_t i = 0; i < m_cfg.logWeightBaseCoeffs.size(); i++) {
0161       logWeightBase += m_cfg.logWeightBaseCoeffs[i] * pow(l, i);
0162     }
0163   }
0164
0165   for (unsigned i = 0; i < pcl.getHits().size(); ++i) {
0166     const auto& hit   = pcl.getHits()[i];
0167     const auto weight = pcl.getWeights()[i];
0168     //      _DBG_<<" -- weight = " << weight << "  E=" << hit.getEnergy() << " totalE=" <<totalE << " log(E/totalE)=" << std::log(hit.getEnergy()/totalE) << std::endl;
0169     float w = weightFunc(hit.getEnergy() * weight, totalE, logWeightBase, 0);
0170     tw += w;
0171     v = v + (hit.getPosition() * w);
0172   }
0173   if (tw == 0.) {
0174     warning("zero total weights encountered, you may want to adjust your weighting parameter.");
0175     return {};
0176   }
0177   cl.setPosition(v / tw);
0178   cl.setPositionError({}); // @TODO: Covariance matrix
0179
0180   // Optionally constrain the cluster to the hit eta values
0181   if (m_cfg.enableEtaBounds) {
0182     const bool overflow  = (edm4hep::utils::eta(cl.getPosition()) > maxHitEta);
0183     const bool underflow = (edm4hep::utils::eta(cl.getPosition()) < minHitEta);
0184     if (overflow || underflow) {
0185       const double newEta   = overflow ? maxHitEta : minHitEta;
0186       const double newTheta = edm4hep::utils::etaToAngle(newEta);
0187       const double newR     = edm4hep::utils::magnitude(cl.getPosition());
0188       const double newPhi   = edm4hep::utils::angleAzimuthal(cl.getPosition());
0189       cl.setPosition(edm4hep::utils::sphericalToVector(newR, newTheta, newPhi));
0190       debug("Bound cluster position to contributing hits due to {}",
0191             (overflow ? "overflow" : "underflow"));
0192     }
0193   }
0194   return cl;
0195 }
0196
0197 void CalorimeterClusterRecoCoG::associate(
0198     const edm4eic::Cluster& cl,
0199     [[maybe_unused]] const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
0200 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0201     const podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
0202     edm4eic::MCRecoClusterParticleLinkCollection* links,
0203 #endif
0204     edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const {
0205   // --------------------------------------------------------------------------
0206   // Association Logic
0207   // --------------------------------------------------------------------------
0208   /*  1. identify all sim hits associated with a given protocluster, and sum
0209    *     the energy of the sim hits.
0210    *  2. for each sim hit
0211    *     - identify parents of each contributing particles; and
0212    *     - if parent is a primary particle, add to list of contributors
0213    *       and sum the energy contributed by the parent.
0214    *  3. create an association for each contributing primary with a weight
0215    *     of contributed energy over total sim hit energy.
0216    */
0217
0218   // lambda to compare MCParticles
0219   auto compare = [](const edm4hep::MCParticle& lhs, const edm4hep::MCParticle& rhs) {
0220     if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
0221       return (lhs.getObjectID().index < rhs.getObjectID().index);
0222     }
0223     return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
0224   };
0225
0226   // bookkeeping maps for associated primaries
0227   std::map<edm4hep::MCParticle, double, decltype(compare)> mapMCParToContrib(compare);
0228
0229   // --------------------------------------------------------------------------
0230   // 1. get associated sim hits and sum energy
0231   // --------------------------------------------------------------------------
0232   double eSimHitSum = 0.;
0233   for (auto clhit : cl.getHits()) {
0234
0235 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0236     // Get linked sim hits using LinkNavigator
0237     const auto vecAssocSimHits = link_nav.getLinked(clhit.getRawHit());
0238 #else
0239     // Fallback: linear search through associations
0240     std::vector<std::pair<edm4hep::SimCalorimeterHit, double>> vecAssocSimHits;
0241     for (const auto& assoc : *mchitassociations) {
0242       if (assoc.getRawHit() == clhit.getRawHit()) {
0243         vecAssocSimHits.push_back(std::make_pair(assoc.getSimHit(), assoc.getWeight()));
0244       }
0245     }
0246 #endif
0247
0248     for (const auto& [simHit, weight] : vecAssocSimHits) {
0249       eSimHitSum += simHit.getEnergy();
0250     }
0251
0252     debug("{} associated sim hits found for reco hit (cell ID = {})", vecAssocSimHits.size(),
0253           clhit.getCellID());
0254
0255     // ------------------------------------------------------------------------
0256     // 2. loop through associated sim hits
0257     // ------------------------------------------------------------------------
0258     for (const auto& [simHit, weight] : vecAssocSimHits) {
0259       for (const auto& contrib : simHit.getContributions()) {
0260         // --------------------------------------------------------------------
0261         // grab primary responsible for contribution & increment relevant sum
0262         // --------------------------------------------------------------------
0263         edm4hep::MCParticle primary = get_primary(contrib);
0264         mapMCParToContrib[primary] += contrib.getEnergy();
0265
0266         trace("Identified primary: id = {}, pid = {}, total energy = {}, contributed = {}",
0267               primary.getObjectID().index, primary.getPDG(), primary.getEnergy(),
0268               mapMCParToContrib[primary]);
0269       }
0270     }
0271   }
0272   debug("Found {} primaries contributing a total of {} GeV", mapMCParToContrib.size(), eSimHitSum);
0273
0274   // --------------------------------------------------------------------------
0275   // 3. create association for each contributing primary
0276   // --------------------------------------------------------------------------
0277   for (auto [part, contribution] : mapMCParToContrib) {
0278     // calculate weight
0279     const double weight = contribution / eSimHitSum;
0280
0281 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 7, 0)
0282     // create link
0283     auto link = links->create();
0284     link.setWeight(weight);
0285     link.setFrom(cl);
0286     link.setTo(part);
0287 #endif
0288
0289     // set association
0290     auto assoc = assocs->create();
0291     assoc.setWeight(weight);
0292     assoc.setRec(cl);
0293     assoc.setSim(part);
0294
0295     debug("Associated cluster #{} to MC Particle #{} (pid = {}, status = {}, energy = {}) with "
0296           "weight ({})",
0297           cl.getObjectID().index, part.getObjectID().index, part.getPDG(),
0298           part.getGeneratorStatus(), part.getEnergy(), weight);
0299   }
0300 }
0301
0302 edm4hep::MCParticle
0303 CalorimeterClusterRecoCoG::get_primary(const edm4hep::CaloHitContribution& contrib) {
0304   // get contributing particle
0305   const auto contributor = contrib.getParticle();
0306
0307   // walk back through parents to find primary
0308   //   - TODO finalize primary selection. This
0309   //     can be improved!!
0310   edm4hep::MCParticle primary = contributor;
0311   while (primary.parents_size() > 0) {
0312     if (primary.getGeneratorStatus() != 0) {
0313       break;
0314     }
0315     primary = primary.getParents(0);
0316   }
0317   return primary;
0318 }
0319
0320 } // namespace eicrecon