src/components/MatchClusters.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten
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 <algorithm>
0009 #include <cmath>
0010
0011 #include <fmt/format.h>
0012
0013 #include "GaudiAlg/GaudiAlgorithm.h"
0014 #include "GaudiAlg/GaudiTool.h"
0015 #include "GaudiAlg/Producer.h"
0016 #include "GaudiAlg/Transformer.h"
0017 #include "GaudiKernel/RndmGenerators.h"
0018
0019 #include <k4FWCore/DataHandle.h>
0020
0021 // Event Model related classes
0022 #include "edm4hep/MCParticleCollection.h"
0023 #include "edm4eic/ClusterCollection.h"
0024 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0025 #include "edm4eic/MCRecoParticleAssociationCollection.h"
0026 #include "edm4eic/ReconstructedParticleCollection.h"
0027 #include "edm4eic/TrackParametersCollection.h"
0028 #include "edm4hep/utils/vector_utils.h"
0029
0030 namespace Jug::Fast {
0031
0032 class MatchClusters : public GaudiAlgorithm {
0033 private:
0034   // input data
0035   DataHandle<edm4hep::MCParticleCollection> m_inputMCParticles{"MCParticles", Gaudi::DataHandle::Reader, this};
0036   DataHandle<edm4eic::ReconstructedParticleCollection> m_inputParticles{"ReconstructedChargedParticles",
0037                                                                     Gaudi::DataHandle::Reader, this};
0038   DataHandle<edm4eic::MCRecoParticleAssociationCollection> m_inputParticlesAssoc{"ReconstructedChargedParticlesAssoc",
0039                                                                     Gaudi::DataHandle::Reader, this};
0040   Gaudi::Property<std::vector<std::string>> m_inputClusters{this, "inputClusters", {}, "Clusters to be aggregated"};
0041   Gaudi::Property<std::vector<std::string>> m_inputClustersAssoc{this, "inputClustersAssoc", {}, "Cluster associations to be aggregated"};
0042   std::vector<DataHandle<edm4eic::ClusterCollection>*> m_inputClustersCollections;
0043   std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> m_inputClustersAssocCollections;
0044
0045   // output data
0046   DataHandle<edm4eic::ReconstructedParticleCollection> m_outputParticles{"ReconstructedParticles",
0047                                                                      Gaudi::DataHandle::Writer, this};
0048   DataHandle<edm4eic::MCRecoParticleAssociationCollection> m_outputParticlesAssoc{"ReconstructedParticlesAssoc",
0049                                                                      Gaudi::DataHandle::Writer, this};
0050
0051 public:
0052   MatchClusters(const std::string& name, ISvcLocator* svcLoc)
0053       : GaudiAlgorithm(name, svcLoc) {
0054     declareProperty("inputMCParticles", m_inputMCParticles, "MCParticles");
0055     declareProperty("inputParticles", m_inputParticles, "ReconstructedChargedParticles");
0056     declareProperty("inputParticlesAssoc", m_inputParticlesAssoc, "ReconstructedChargedParticlesAssoc");
0057     declareProperty("outputParticles", m_outputParticles, "ReconstructedParticles");
0058     declareProperty("outputParticlesAssoc", m_outputParticlesAssoc, "ReconstructedParticlesAssoc");
0059   }
0060
0061   StatusCode initialize() override {
0062     if (GaudiAlgorithm::initialize().isFailure()) {
0063       return StatusCode::FAILURE;
0064     }
0065     m_inputClustersCollections = getClusterCollections(m_inputClusters);
0066     m_inputClustersAssocCollections = getClusterAssociations(m_inputClustersAssoc);
0067     return StatusCode::SUCCESS;
0068   }
0069   StatusCode execute() override {
0070     if (msgLevel(MSG::DEBUG)) {
0071       debug() << "Processing cluster info for new event" << endmsg;
0072     }
0073     // input collection
0074     const auto& mcparticles  = *(m_inputMCParticles.get());
0075     const auto& inparts      = *(m_inputParticles.get());
0076     const auto& inpartsassoc = *(m_inputParticlesAssoc.get());
0077     auto& outparts           = *(m_outputParticles.createAndPut());
0078     auto& outpartsassoc      = *(m_outputParticlesAssoc.createAndPut());
0079
0080     if (msgLevel(MSG::DEBUG)) {
0081       debug() << "Step 0/2: Getting indexed list of clusters..." << endmsg;
0082     }
0083
0084     // get an indexed map of all clusters
0085     auto clusterMap = indexedClusters(m_inputClustersCollections, m_inputClustersAssocCollections);
0086
0087     // 1. Loop over all tracks and link matched clusters where applicable
0088     // (removing matched clusters from the cluster maps)
0089     if (msgLevel(MSG::DEBUG)) {
0090       debug() << "Step 1/2: Matching clusters to charged particles..." << endmsg;
0091     }
0092     for (const auto& inpart: inparts) {
0093       if (msgLevel(MSG::DEBUG)) {
0094         debug() << " --> Processing charged particle " << inpart.getObjectID().index
0095                 << ", PDG: " << inpart.getPDG()
0096                 << ", energy: " << inpart.getEnergy()
0097                 << endmsg;
0098       }
0099
0100       auto outpart = inpart.clone();
0101       outparts.push_back(outpart);
0102
0103       int mcID = -1;
0104
0105       // find associated particle
0106       for (const auto& assoc: inpartsassoc) {
0107         if (assoc.getRec() == inpart) {
0108           mcID = assoc.getSimID();
0109           break;
0110         }
0111       }
0112
0113       if (msgLevel(MSG::VERBOSE)) {
0114         verbose() << "    --> Found particle with mcID " << mcID << endmsg;
0115       }
0116
0117       if (mcID < 0) {
0118         if (msgLevel(MSG::DEBUG)) {
0119           debug() << "    --> cannot match track without associated mcID" << endmsg;
0120         }
0121         continue;
0122       }
0123
0124       if (clusterMap.count(mcID)) {
0125         const auto& clus = clusterMap[mcID];
0126         if (msgLevel(MSG::DEBUG)) {
0127           debug() << "    --> found matching cluster with energy: " << clus.getEnergy() << endmsg;
0128         }
0129         clusterMap.erase(mcID);
0130       }
0131
0132       // create truth associations
0133       auto assoc = outpartsassoc.create();
0134       assoc.setRecID(outpart.getObjectID().index);
0135       assoc.setSimID(mcID);
0136       assoc.setWeight(1.0);
0137       assoc.setRec(outpart);
0138       //assoc.setSim(mcparticles[mcID]);
0139     }
0140
0141     // 2. Now loop over all remaining clusters and add neutrals. Also add in Hcal energy
0142     // if a matching cluster is available
0143     if (msgLevel(MSG::DEBUG)) {
0144       debug() << "Step 2/2: Creating neutrals for remaining clusters..." << endmsg;
0145     }
0146     for (const auto& [mcID, clus] : clusterMap) {
0147       if (msgLevel(MSG::DEBUG)) {
0148         debug() << " --> Processing unmatched cluster with energy: " << clus.getEnergy()
0149                 << endmsg;
0150       }
0151
0152       // get mass/PDG from mcparticles, 0 (unidentified) in case the matched particle is charged.
0153       const auto& mc    = mcparticles[mcID];
0154       const double mass = (!mc.getCharge()) ? mc.getMass() : 0;
0155       const int32_t pdg = (!mc.getCharge()) ? mc.getPDG() : 0;
0156       if (msgLevel(MSG::DEBUG)) {
0157         if (mc.getCharge()) {
0158           debug() << "   --> associated mcparticle is not a neutral (PDG: " << mc.getPDG()
0159                   << "), setting the reconstructed particle ID to 0 (unidentified)" << endmsg;
0160         }
0161         debug() << "   --> found matching associated mcparticle with PDG: " << pdg << ", energy: " << mc.getEnergy()
0162                 << endmsg;
0163       }
0164
0165       // Reconstruct our neutrals and add them to the list
0166       const auto outpart = reconstruct_neutral(clus, mass, pdg);
0167       if (msgLevel(MSG::DEBUG)) {
0168         debug() << " --> Reconstructed neutral particle with PDG: " << outpart.getPDG()
0169                 << ", energy: " << outpart.getEnergy()
0170                 << endmsg;
0171       }
0172       outparts.push_back(outpart);
0173
0174       // Create truth associations
0175       auto assoc = outpartsassoc.create();
0176       assoc.setRecID(outpart.getObjectID().index);
0177       assoc.setSimID(mcID);
0178       assoc.setWeight(1.0);
0179       assoc.setRec(outpart);
0180       //assoc.setSim(mcparticles[mcID]);
0181     }
0182     return StatusCode::SUCCESS;
0183   }
0184
0185 private:
0186   std::vector<DataHandle<edm4eic::ClusterCollection>*> getClusterCollections(const std::vector<std::string>& cols) {
0187     std::vector<DataHandle<edm4eic::ClusterCollection>*> ret;
0188     for (const auto& colname : cols) {
0189       debug() << "initializing cluster collection: " << colname << endmsg;
0190       ret.push_back(new DataHandle<edm4eic::ClusterCollection>{colname, Gaudi::DataHandle::Reader, this});
0191     }
0192     return ret;
0193   }
0194
0195   std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> getClusterAssociations(const std::vector<std::string>& cols) {
0196     std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> ret;
0197     for (const auto& colname : cols) {
0198       debug() << "initializing cluster association collection: " << colname << endmsg;
0199       ret.push_back(new DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>{colname, Gaudi::DataHandle::Reader, this});
0200     }
0201     return ret;
0202   }
0203
0204   // get a map of mcID --> cluster
0205   // input: cluster_collections --> list of handles to all cluster collections
0206   std::map<int, edm4eic::Cluster>
0207   indexedClusters(
0208       const std::vector<DataHandle<edm4eic::ClusterCollection>*>& cluster_collections,
0209       const std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*>& associations_collections
0210   ) const {
0211     std::map<int, edm4eic::Cluster> matched = {};
0212
0213     // loop over cluster collections
0214     for (const auto& cluster_handle : cluster_collections) {
0215       const auto& clusters = *(cluster_handle->get());
0216
0217       // loop over clusters
0218       for (const auto& cluster : clusters) {
0219
0220         int mcID = -1;
0221
0222         // loop over association collections
0223         for (const auto& associations_handle : associations_collections) {
0224           const auto& associations = *(associations_handle->get());
0225
0226           // find associated particle
0227           for (const auto& assoc : associations) {
0228             if (assoc.getRec() == cluster) {
0229               mcID = assoc.getSimID();
0230               break;
0231             }
0232           }
0233
0234           // found associated particle
0235           if (mcID != -1) {
0236             break;
0237           }
0238         }
0239
0240         if (msgLevel(MSG::VERBOSE)) {
0241           verbose() << " --> Found cluster with mcID " << mcID << " and energy "
0242                     << cluster.getEnergy() << endmsg;
0243         }
0244
0245         if (mcID < 0) {
0246           if (msgLevel(MSG::VERBOSE)) {
0247             verbose() << "   --> WARNING: no valid MC truth link found, skipping cluster..." << endmsg;
0248           }
0249           continue;
0250         }
0251
0252         const bool duplicate = matched.count(mcID);
0253         if (duplicate) {
0254           if (msgLevel(MSG::VERBOSE)) {
0255             verbose() << "   --> WARNING: this is a duplicate mcID, keeping the higher energy cluster" << endmsg;
0256           }
0257           if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0258             continue;
0259           }
0260         }
0261         matched[mcID] = cluster;
0262       }
0263     }
0264     return matched;
0265   }
0266
0267   // reconstruct a neutral cluster
0268   // (for now assuming the vertex is at (0,0,0))
0269   edm4eic::MutableReconstructedParticle reconstruct_neutral(const edm4eic::Cluster& clus, const double mass,
0270                                                  const int32_t pdg) const {
0271     const float energy = clus.getEnergy();
0272     const float p = energy < mass ? 0 : std::sqrt(energy * energy - mass * mass);
0273     const auto position = clus.getPosition();
0274     const auto momentum = p * (position / edm4hep::utils::magnitude(position));
0275     // setup our particle
0276     edm4eic::MutableReconstructedParticle part;
0277     part.setMomentum(momentum);
0278     part.setPDG(pdg);
0279     part.setCharge(0);
0280     part.setEnergy(energy);
0281     part.setMass(mass);
0282     return part;
0283   }
0284 }; // namespace Jug::Fast
0285
0286 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0287 DECLARE_COMPONENT(MatchClusters)
0288
0289 } // namespace Jug::Fast