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