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