algorithms/reco/MatchClusters.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022, 2024 Sylvester Joosten, Dmitry Romanov, Wouter Deconinck
0003
0004 // Takes a list of particles (presumed to be from tracking), and all available clusters.
0005 // 1. Match clusters to their tracks using the mcID field
0006 // 2. For unmatched clusters create neutrals and add to the particle list
0007
0008 #include <algorithms/logger.h>
0009 #include <edm4eic/ClusterCollection.h>
0010 #include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
0011 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0012 #include <edm4eic/ReconstructedParticleCollection.h>
0013 #include <edm4hep/MCParticleCollection.h>
0014 #include <edm4hep/Vector3f.h>
0015 #include <edm4hep/utils/vector_utils.h>
0016 #include <podio/ObjectID.h>
0017 #include <podio/detail/Link.h>
0018 #include <podio/detail/LinkCollectionImpl.h>
0019 #include <cmath>
0020 #include <iterator>
0021 #include <map>
0022 #include <memory>
0023 #include <tuple>
0024 #include <utility>
0025
0026 #include "MatchClusters.h"
0027
0028 namespace eicrecon {
0029
0030 void MatchClusters::process(const MatchClusters::Input& input,
0031                             const MatchClusters::Output& output) const {
0032
0033   const auto [mcparticles, inparts, inpartsassoc, clusters, clustersassoc] = input;
0034   auto [outparts, outlinks, outpartsassoc]                                 = output;
0035
0036   debug("Processing cluster info for new event");
0037
0038   debug("Step 0/2: Getting indexed list of clusters...");
0039
0040   // get an indexed map of all clusters
0041   auto clusterMap = indexedClusters(clusters, clustersassoc);
0042
0043   // 1. Loop over all tracks and link matched clusters where applicable
0044   // (removing matched clusters from the cluster maps)
0045   debug("Step 1/2: Matching clusters to charged particles...");
0046
0047   for (const auto inpart : *inparts) {
0048     debug(" --> Processing charged particle {}, PDG {}, energy {}", inpart.getObjectID().index,
0049           inpart.getPDG(), inpart.getEnergy());
0050
0051     auto outpart = inpart.clone();
0052     outparts->push_back(outpart);
0053
0054     // find the best associated MC particle (largest weight) for cluster matching
0055     int bestMcID      = -1;
0056     double bestWeight = -1.;
0057
0058     for (const auto& assoc : *inpartsassoc) {
0059       if (assoc.getRec().getObjectID() == inpart.getObjectID()) {
0060         const double w = assoc.getWeight();
0061         if (w > bestWeight) {
0062           bestWeight = w;
0063           bestMcID   = assoc.getSim().getObjectID().index;
0064         }
0065       }
0066     }
0067
0068     trace("    --> Found particle with best mcID {} weight {}", bestMcID, bestWeight);
0069
0070     if (bestMcID < 0) {
0071       debug("    --> cannot match track without associated mcID");
0072     } else if (clusterMap.contains(bestMcID)) {
0073       const auto& clus = clusterMap[bestMcID];
0074       debug("    --> found matching cluster with energy: {}", clus.getEnergy());
0075       debug("    --> adding cluster to reconstructed particle");
0076       outpart.addToClusters(clus);
0077       clusterMap.erase(bestMcID);
0078     }
0079
0080     // propagate all original associations, remapped to the cloned output particle
0081     for (const auto& assoc : *inpartsassoc) {
0082       if (assoc.getRec().getObjectID() == inpart.getObjectID()) {
0083         auto link = outlinks->create();
0084         link.setWeight(assoc.getWeight());
0085         link.setFrom(outpart);
0086         link.setTo(assoc.getSim());
0087         auto outassoc = outpartsassoc->create();
0088         outassoc.setWeight(assoc.getWeight());
0089         outassoc.setRec(outpart);
0090         outassoc.setSim(assoc.getSim());
0091       }
0092     }
0093   }
0094
0095   // 2. Now loop over all remaining clusters and add neutrals. Also add in Hcal energy
0096   // if a matching cluster is available
0097   debug("Step 2/2: Creating neutrals for remaining clusters...");
0098   for (const auto& [mcID, clus] : clusterMap) {
0099     debug(" --> Processing unmatched cluster with energy: {}", clus.getEnergy());
0100
0101     // get mass/PDG from mcparticles, 0 (unidentified) in case the matched particle is charged.
0102     const auto mc     = (*mcparticles)[mcID];
0103     const double mass = 0.;
0104     const int32_t pdg = 0;
0105     if (level() <= algorithms::LogLevel::kDebug) {
0106       if (mc.getCharge() != 0.0F) {
0107         debug("   --> associated mcparticle is not a neutral (PDG: {}), "
0108               "setting the reconstructed particle ID to 0 (unidentified)",
0109               mc.getPDG());
0110       }
0111       debug("   --> found matching associated mcparticle with PDG: {}, energy: {}", pdg,
0112             mc.getEnergy());
0113     }
0114
0115     // Reconstruct our neutrals and add them to the list
0116     const auto outpart = reconstruct_neutral(&clus, mass, pdg);
0117     debug(" --> Reconstructed neutral particle with PDG: {}, energy: {}", outpart.getPDG(),
0118           outpart.getEnergy());
0119
0120     outparts->push_back(outpart);
0121
0122     // Create truth associations
0123     auto link = outlinks->create();
0124     link.setWeight(1.0);
0125     link.setFrom(outpart);
0126     link.setTo((*mcparticles)[mcID]);
0127     auto assoc = outpartsassoc->create();
0128     assoc.setWeight(1.0);
0129     assoc.setRec(outpart);
0130     assoc.setSim((*mcparticles)[mcID]);
0131   }
0132 }
0133
0134 // get a map of mcID --> cluster
0135 // For each cluster, pick the best associated MC particle by association weight.
0136 // Returns a map keyed by mcID and valued with the selected cluster.
0137 std::map<int, edm4eic::Cluster> MatchClusters::indexedClusters(
0138     const edm4eic::ClusterCollection* clusters,
0139     const edm4eic::MCRecoClusterParticleAssociationCollection* associations) const {
0140
0141   // temporary map: mcID -> (cluster, weight) so we can choose the cluster with highest weight per mcID
0142   std::map<int, std::pair<edm4eic::Cluster, float>> bestForMc;
0143
0144   // loop over clusters and pick their best MC association by weight
0145   for (const auto cluster : *clusters) {
0146
0147     int bestMcID     = -1;
0148     float bestWeight = -1.F;
0149
0150     // find best associated MC particle for this cluster (largest association weight)
0151     for (const auto assoc : *associations) {
0152       if (assoc.getRec() == cluster) {
0153         const int candMcID = assoc.getSim().getObjectID().index;
0154         const float w      = assoc.getWeight();
0155         if (w > bestWeight) {
0156           bestWeight = w;
0157           bestMcID   = candMcID;
0158         }
0159       }
0160     }
0161
0162     trace(" --> Found cluster with best mcID {} weight {} and energy {}", bestMcID, bestWeight,
0163           cluster.getEnergy());
0164
0165     if (bestMcID < 0) {
0166       trace("   --> WARNING: no valid MC truth link found, skipping cluster...");
0167       continue;
0168     }
0169
0170     // For this mcID, keep the cluster with the highest association weight (tie-break by energy).
0171     auto it = bestForMc.find(bestMcID);
0172     if (it == bestForMc.end()) {
0173       bestForMc.emplace(bestMcID, std::make_pair(cluster, bestWeight));
0174     } else {
0175       const float existingWeight = it->second.second;
0176       if (bestWeight > existingWeight ||
0177           (bestWeight == existingWeight && cluster.getEnergy() > it->second.first.getEnergy())) {
0178         it->second = std::make_pair(cluster, bestWeight);
0179       }
0180     }
0181   }
0182
0183   // Convert to the old API: map<int, edm4eic::Cluster>
0184   std::map<int, edm4eic::Cluster> matched;
0185   for (const auto& kv : bestForMc) {
0186     matched.emplace(kv.first, kv.second.first);
0187   }
0188
0189   return matched;
0190 }
0191
0192 // reconstruct a neutral cluster
0193 // (for now assuming the vertex is at (0,0,0))
0194 edm4eic::MutableReconstructedParticle
0195 MatchClusters::reconstruct_neutral(const edm4eic::Cluster* cluster, const double mass,
0196                                    const int32_t pdg) {
0197
0198   const float energy  = cluster->getEnergy();
0199   const float p       = energy < mass ? 0 : std::sqrt(energy * energy - mass * mass);
0200   const auto position = cluster->getPosition();
0201   const auto momentum = p * (position / edm4hep::utils::magnitude(position));
0202   // setup our particle
0203   edm4eic::MutableReconstructedParticle part;
0204   part.setMomentum(momentum);
0205   part.setPDG(pdg);
0206   part.setCharge(0);
0207   part.setEnergy(energy);
0208   part.setMass(mass);
0209   part.addToClusters(*cluster);
0210   return part;
0211 }
0212
0213 } // namespace eicrecon