algorithms/calorimetry/CalorimeterTruthClustering.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2026 Sylvester Joosten, Whitney Armstrong, Wouter Deconinck, Derek Anderson
0003
0004 #include "CalorimeterTruthClustering.h"
0005
0006 #include <DD4hep/config.h>
0007 #include <edm4hep/CaloHitContribution.h>
0008 #include <edm4hep/MCParticle.h>
0009 #include <edm4hep/RawCalorimeterHit.h>
0010 #include <edm4hep/SimCalorimeterHit.h>
0011 #include <podio/LinkNavigator.h>
0012 #include <podio/ObjectID.h>
0013 #include <podio/RelationRange.h>
0014 #include <cstddef>
0015 #include <cstdint>
0016 #include <map>
0017 #include <set>
0018 #include <tuple>
0019 #include <vector>
0020
0021 #include "algorithms/calorimetry/CalorimeterTruthClusteringConfig.h"
0022 #include "algorithms/interfaces/LinkTruthUtils.h"
0023
0024 using namespace dd4hep;
0025
0026 namespace eicrecon {
0027
0028 void CalorimeterTruthClustering::init() {}
0029
0030 void CalorimeterTruthClustering::process(const CalorimeterTruthClustering::Input& input,
0031                                          const CalorimeterTruthClustering::Output& output) const {
0032
0033   const auto [hits, hitLinks] = input;
0034   auto [clusters]             = output;
0035
0036   const auto navigator = podio::LinkNavigator(*hitLinks);
0037
0038   // Map mc track ID to protoCluster index
0039   std::map<int32_t, int32_t> protoIndex;
0040
0041   // Loop over all calorimeter hits and sort per mcparticle
0042   for (const auto& hit : *hits) {
0043
0044     const auto linkedSimHits = navigator.getLinked(hit.getRawHit());
0045
0046     // Ignore hit if no associated sim hits
0047     std::set<std::size_t> mcIndices;
0048     std::map<std::size_t, float> mcContribs;
0049     float mcContribs_total{0.0};
0050     for (const auto& [simHit, weight] : linkedSimHits) {
0051
0052       // Loop through contributions, create a protocluster for each contributing primary
0053       for (const auto& contrib : simHit.getContributions()) {
0054
0055         edm4hep::MCParticle primary = truth::primaryFrom(contrib, m_cfg.promptDecayPDGs);
0056         const auto trackID          = primary.getObjectID().index;
0057         // Create a new protocluster if we don't have one for this primary
0058         if (!protoIndex.contains(trackID)) {
0059           clusters->create();
0060           protoIndex[trackID] = clusters->size() - 1;
0061         }
0062         mcIndices.insert(trackID);
0063         mcContribs[trackID] += contrib.getEnergy();
0064         mcContribs_total += contrib.getEnergy();
0065       }
0066     }
0067
0068     // Add hit to the appropriate protoclusters
0069     for (const auto& mcIndex : mcIndices) {
0070       const float weight = mcContribs[mcIndex] / mcContribs_total;
0071       (*clusters)[protoIndex[mcIndex]].addToHits(hit);
0072       (*clusters)[protoIndex[mcIndex]].addToWeights(weight);
0073     }
0074   }
0075 }
0076
0077 } // namespace eicrecon