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