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