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