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