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 DataHandle<edm4hep::MCParticleCollection> m_inputMCParticles{"MCParticles", Gaudi::DataHandle::Reader, this};
0033   mutable DataHandle<edm4eic::ReconstructedParticleCollection> m_inputParticles{"ReconstructedChargedParticles",
0034                                                                     Gaudi::DataHandle::Reader, this};
0035   mutable 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<DataHandle<edm4eic::ClusterCollection>*> m_inputClustersCollections;
0040   std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> m_inputClustersAssocCollections;
0041
0042   // output data
0043   mutable DataHandle<edm4eic::ReconstructedParticleCollection> m_outputParticles{"ReconstructedParticles",
0044                                                                      Gaudi::DataHandle::Writer, this};
0045   mutable 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
0102       // find associated particle
0103       for (const auto& assoc: inpartsassoc) {
0104         if (assoc.getRec() == inpart) {
0105           mcID = assoc.getSimID();
0106           break;
0107         }
0108       }
0109
0110       if (msgLevel(MSG::VERBOSE)) {
0111         verbose() << "    --> Found particle with mcID " << mcID << endmsg;
0112       }
0113
0114       if (mcID < 0) {
0115         if (msgLevel(MSG::DEBUG)) {
0116           debug() << "    --> cannot match track without associated mcID" << endmsg;
0117         }
0118         continue;
0119       }
0120
0121       if (clusterMap.count(mcID)) {
0122         const auto& clus = clusterMap[mcID];
0123         if (msgLevel(MSG::DEBUG)) {
0124           debug() << "    --> found matching cluster with energy: " << clus.getEnergy() << endmsg;
0125         }
0126         clusterMap.erase(mcID);
0127       }
0128
0129       // create truth associations
0130       auto assoc = outpartsassoc.create();
0131       assoc.setRecID(outpart.getObjectID().index);
0132       assoc.setSimID(mcID);
0133       assoc.setWeight(1.0);
0134       assoc.setRec(outpart);
0135       //assoc.setSim(mcparticles[mcID]);
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.setRecID(outpart.getObjectID().index);
0174       assoc.setSimID(mcID);
0175       assoc.setWeight(1.0);
0176       assoc.setRec(outpart);
0177       //assoc.setSim(mcparticles[mcID]);
0178     }
0179     return StatusCode::SUCCESS;
0180   }
0181
0182 private:
0183   std::vector<DataHandle<edm4eic::ClusterCollection>*> getClusterCollections(const std::vector<std::string>& cols) {
0184     std::vector<DataHandle<edm4eic::ClusterCollection>*> ret;
0185     for (const auto& colname : cols) {
0186       debug() << "initializing cluster collection: " << colname << endmsg;
0187       ret.push_back(new DataHandle<edm4eic::ClusterCollection>{colname, Gaudi::DataHandle::Reader, this});
0188     }
0189     return ret;
0190   }
0191
0192   std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> getClusterAssociations(const std::vector<std::string>& cols) {
0193     std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> ret;
0194     for (const auto& colname : cols) {
0195       debug() << "initializing cluster association collection: " << colname << endmsg;
0196       ret.push_back(new DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>{colname, Gaudi::DataHandle::Reader, this});
0197     }
0198     return ret;
0199   }
0200
0201   // get a map of mcID --> cluster
0202   // input: cluster_collections --> list of handles to all cluster collections
0203   std::map<int, edm4eic::Cluster>
0204   indexedClusters(
0205       const std::vector<DataHandle<edm4eic::ClusterCollection>*>& cluster_collections,
0206       const std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*>& associations_collections
0207   ) const {
0208     std::map<int, edm4eic::Cluster> matched = {};
0209
0210     // loop over cluster collections
0211     for (const auto& cluster_handle : cluster_collections) {
0212       const auto& clusters = *(cluster_handle->get());
0213
0214       // loop over clusters
0215       for (const auto& cluster : clusters) {
0216
0217         int mcID = -1;
0218
0219         // loop over association collections
0220         for (const auto& associations_handle : associations_collections) {
0221           const auto& associations = *(associations_handle->get());
0222
0223           // find associated particle
0224           for (const auto& assoc : associations) {
0225             if (assoc.getRec() == cluster) {
0226               mcID = assoc.getSimID();
0227               break;
0228             }
0229           }
0230
0231           // found associated particle
0232           if (mcID != -1) {
0233             break;
0234           }
0235         }
0236
0237         if (msgLevel(MSG::VERBOSE)) {
0238           verbose() << " --> Found cluster with mcID " << mcID << " and energy "
0239                     << cluster.getEnergy() << endmsg;
0240         }
0241
0242         if (mcID < 0) {
0243           if (msgLevel(MSG::VERBOSE)) {
0244             verbose() << "   --> WARNING: no valid MC truth link found, skipping cluster..." << endmsg;
0245           }
0246           continue;
0247         }
0248
0249         const bool duplicate = matched.count(mcID);
0250         if (duplicate) {
0251           if (msgLevel(MSG::VERBOSE)) {
0252             verbose() << "   --> WARNING: this is a duplicate mcID, keeping the higher energy cluster" << endmsg;
0253           }
0254           if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0255             continue;
0256           }
0257         }
0258         matched[mcID] = cluster;
0259       }
0260     }
0261     return matched;
0262   }
0263
0264   // reconstruct a neutral cluster
0265   // (for now assuming the vertex is at (0,0,0))
0266   edm4eic::MutableReconstructedParticle reconstruct_neutral(const edm4eic::Cluster& clus, const double mass,
0267                                                  const int32_t pdg) const {
0268     const float energy = clus.getEnergy();
0269     const float p = energy < mass ? 0 : std::sqrt(energy * energy - mass * mass);
0270     const auto position = clus.getPosition();
0271     const auto momentum = p * (position / edm4hep::utils::magnitude(position));
0272     // setup our particle
0273     edm4eic::MutableReconstructedParticle part;
0274     part.setMomentum(momentum);
0275     part.setPDG(pdg);
0276     part.setCharge(0);
0277     part.setEnergy(energy);
0278     part.setMass(mass);
0279     return part;
0280   }
0281 }; // namespace Jug::Fast
0282
0283 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0284 DECLARE_COMPONENT(MatchClusters)
0285
0286 } // namespace Jug::Fast