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