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/CaloHitContribution.h>
0017 #include <edm4hep/MCParticle.h>
0018 #include <edm4hep/RawCalorimeterHit.h>
0019 #include <edm4hep/SimCalorimeterHitCollection.h>
0020 #include <edm4hep/Vector3f.h>
0021 #include <edm4hep/utils/vector_utils.h>
0022 #include <gsl/pointers>
0023 #include <podio/LinkNavigator.h>
0024 #include <podio/ObjectID.h>
0025 #include <podio/RelationRange.h>
0026 #include <algorithm>
0027 #include <cctype>
0028 #include <cstddef>
0029 #include <limits>
0030 #include <map>
0031 #include <optional>
0032 #include <tuple>
0033 #include <vector>
0034
0035 #include "CalorimeterClusterRecoCoG.h"
0036 #include "algorithms/calorimetry/CalorimeterClusterRecoCoGConfig.h"
0037 #include "algorithms/interfaces/CompareObjectID.h"
0038 #include "algorithms/interfaces/LinkTruthUtils.h"
0039
0040 namespace eicrecon {
0041
0042 using namespace dd4hep;
0043
0044 void CalorimeterClusterRecoCoG::init() {
0045   // select weighting method
0046   std::string ew = m_cfg.energyWeight;
0047   // make it case-insensitive
0048   std::ranges::transform(ew, ew.begin(), [](char s) { return std::tolower(s); });
0049   auto it = weightMethods.find(ew);
0050   if (it == weightMethods.end()) {
0051     error("Cannot find energy weighting method {}, choose one from [{}]", m_cfg.energyWeight,
0052           boost::algorithm::join(weightMethods | boost::adaptors::map_keys, ", "));
0053     return;
0054   }
0055   weightFunc = it->second;
0056 }
0057
0058 void CalorimeterClusterRecoCoG::process(const CalorimeterClusterRecoCoG::Input& input,
0059                                         const CalorimeterClusterRecoCoG::Output& output) const {
0060   const auto [proto, mchitlinks, mchitassociations] = input;
0061   auto [clusters, links, associations]              = output;
0062
0063   // Check if truth associations are possible
0064   const truth::EventLinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection> link_nav(mchitlinks);
0065   const bool do_assoc = link_nav.enabled();
0066   if (!do_assoc) {
0067     debug("Provided MCRecoCalorimeterHitLink collection is empty. No truth associations "
0068           "will be performed.");
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 truth::EventLinkNavigator<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   // bookkeeping maps for associated primaries
0197   std::map<edm4hep::MCParticle, double, CompareObjectID<edm4hep::MCParticle>> mapMCParToContrib;
0198
0199   // --------------------------------------------------------------------------
0200   // 1. get associated sim hits and sum energy
0201   // --------------------------------------------------------------------------
0202   double eSimHitSum = 0.;
0203   for (auto clhit : cl.getHits()) {
0204
0205     // Get linked sim hits using LinkNavigator
0206     const auto vecAssocSimHits = link_nav.linked(clhit.getRawHit());
0207
0208     for (const auto& [simHit, weight] : vecAssocSimHits) {
0209       eSimHitSum += simHit.getEnergy();
0210     }
0211
0212     debug("{} associated sim hits found for reco hit (cell ID = {})", vecAssocSimHits.size(),
0213           clhit.getCellID());
0214
0215     // ------------------------------------------------------------------------
0216     // 2. loop through associated sim hits
0217     // ------------------------------------------------------------------------
0218     for (const auto& [simHit, weight] : vecAssocSimHits) {
0219       for (const auto& contrib : simHit.getContributions()) {
0220         // --------------------------------------------------------------------
0221         // grab primary responsible for contribution & increment relevant sum
0222         // --------------------------------------------------------------------
0223         edm4hep::MCParticle primary = truth::primaryFrom(contrib, m_cfg.promptDecayPDGs);
0224         mapMCParToContrib[primary] += contrib.getEnergy();
0225
0226         trace("Identified primary: id = {}, pid = {}, total energy = {}, contributed = {}",
0227               primary.getObjectID().index, primary.getPDG(), primary.getEnergy(),
0228               mapMCParToContrib[primary]);
0229       }
0230     }
0231   }
0232   debug("Found {} primaries contributing a total of {} GeV", mapMCParToContrib.size(), eSimHitSum);
0233
0234   // --------------------------------------------------------------------------
0235   // 3. create association for each contributing primary
0236   // --------------------------------------------------------------------------
0237   for (auto [part, contribution] : mapMCParToContrib) {
0238     // calculate weight
0239     const double weight = contribution / eSimHitSum;
0240
0241     truth::addWeightedRelation(
0242         cl, part, static_cast<float>(weight),
0243         gsl::not_null<edm4eic::MCRecoClusterParticleLinkCollection*>{links},
0244         gsl::not_null<edm4eic::MCRecoClusterParticleAssociationCollection*>{assocs});
0245
0246     debug("Associated cluster #{} to MC Particle #{} (pid = {}, status = {}, energy = {}) with "
0247           "weight ({})",
0248           cl.getObjectID().index, part.getObjectID().index, part.getPDG(),
0249           part.getGeneratorStatus(), part.getEnergy(), weight);
0250   }
0251 }
0252
0253 } // namespace eicrecon