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