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