truth/todo/ClusterMerger.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten
0003
0004 #include <limits>
0005 #include <numbers>
0006
0007 #include <fmt/format.h>
0008 // Gaudi
0009 #include "Gaudi/Property.h"
0010 #include "GaudiAlg/GaudiAlgorithm.h"
0011 #include "GaudiAlg/GaudiTool.h"
0012 #include "GaudiAlg/Transformer.h"
0013 #include "GaudiKernel/PhysicalConstants.h"
0014
0015 #include "JugBase/DataHandle.h"
0016
0017 // Event Model related classes
0018 #include "edm4eic/ClusterCollection.h"
0019 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0020 #include "edm4hep/utils/vector_utils.h"
0021
0022 using namespace Gaudi::Units;
0023
0024 namespace Jug::Fast {
0025
0026 /** Simple algorithm to merge clusters orinating from the same particle together,
0027  *  based on the MC truth.
0028  *
0029  * \ingroup fast
0030  */
0031 class ClusterMerger : public GaudiAlgorithm {
0032 private:
0033   // Input
0034   DataHandle<edm4eic::ClusterCollection> m_inputClusters{"InputClusters", Gaudi::DataHandle::Reader, this};
0035   DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection> m_inputAssociations{"InputAssociations", Gaudi::DataHandle::Reader, this};
0036   // Output
0037   DataHandle<edm4eic::ClusterCollection> m_outputClusters{"OutputClusters", Gaudi::DataHandle::Writer, this};
0038   DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection> m_outputAssociations{"OutputAssociations", Gaudi::DataHandle::Writer, this};
0039 public:
0040   ClusterMerger(const std::string& name, ISvcLocator* svcLoc)
0041       : GaudiAlgorithm(name, svcLoc) {
0042     declareProperty("inputClusters", m_inputClusters, "Input cluster collection");
0043     declareProperty("inputAssociations", m_inputAssociations, "Input cluster association");
0044     declareProperty("outputClusters", m_outputClusters, "Cluster collection with good energy precision");
0045     declareProperty("outputAssociations", m_outputAssociations, "Cluster associations with good energy precision");
0046   }
0047
0048   StatusCode initialize() override {
0049     if (GaudiAlgorithm::initialize().isFailure()) {
0050       return StatusCode::FAILURE;
0051     }
0052
0053     return StatusCode::SUCCESS;
0054   }
0055
0056   StatusCode execute() override {
0057     if (msgLevel(MSG::DEBUG)) {
0058       debug() << "Merging cluster that belong to the same primary particle" << endmsg;
0059     }
0060     // input
0061     const auto& split = *(m_inputClusters.get());
0062     const auto& assoc = *(m_inputAssociations.get());
0063     // output
0064     auto& merged = *(m_outputClusters.createAndPut());
0065     auto& assoc2 = *(m_outputAssociations.createAndPut());
0066
0067     if (!split.size()) {
0068       if (msgLevel(MSG::DEBUG)) {
0069         debug() << "Nothing to do for this event, returning..." << endmsg;
0070       }
0071       return StatusCode::SUCCESS;
0072     }
0073
0074     if (msgLevel(MSG::DEBUG)) {
0075       debug() << "Step 0/1: Getting indexed list of clusters..." << endmsg;
0076     }
0077     // get an indexed map of all vectors of clusters, indexed by mcID
0078     auto clusterMap = indexedClusterLists(split, assoc);
0079
0080     // loop over all position clusters and match with energy clusters
0081     if (msgLevel(MSG::DEBUG)) {
0082       debug() << "Step 1/1: Merging clusters where needed" << endmsg;
0083     }
0084     for (const auto& [mcID, clusters] : clusterMap) {
0085       if (msgLevel(MSG::DEBUG)) {
0086         debug() << " --> Processing " << clusters.size() << " clusters for mcID " << mcID << endmsg;
0087       }
0088       if (clusters.size() == 1) {
0089         const auto& clus = clusters[0];
0090         if (msgLevel(MSG::DEBUG)) {
0091           debug() << "   --> Only a single cluster, energy: " << clus.getEnergy()
0092                   << " for this particle, copying" << endmsg;
0093         }
0094         auto new_clus = clus.clone();
0095         merged.push_back(new_clus);
0096         auto ca = assoc2.create();
0097         ca.setRecID(new_clus.getObjectID().index);
0098         ca.setSimID(mcID);
0099         ca.setWeight(1.0);
0100         ca.setRec(new_clus);
0101         //ca.setSim(//FIXME);
0102       } else {
0103         auto new_clus = merged.create();
0104         // calculate aggregate info
0105         float energy      = 0;
0106         float energyError = 0;
0107         float time        = 0;
0108         int nhits = 0;
0109         auto position = new_clus.getPosition();
0110         for (const auto& clus : clusters) {
0111           if (msgLevel(MSG::DEBUG)) {
0112             debug() << "   --> Adding cluster with energy: " << clus.getEnergy() << endmsg;
0113           }
0114           energy += clus.getEnergy();
0115           energyError += clus.getEnergyError() * clus.getEnergyError();
0116           time += clus.getTime() * clus.getEnergy();
0117           nhits += clus.getNhits();
0118           position = position + energy * clus.getPosition();
0119           new_clus.addToClusters(clus);
0120           for (const auto& hit : clus.getHits()) {
0121             new_clus.addToHits(hit);
0122           }
0123         }
0124         new_clus.setEnergy(energy);
0125         new_clus.setEnergyError(sqrt(energyError));
0126         new_clus.setTime(time / energy);
0127         new_clus.setNhits(nhits);
0128         new_clus.setPosition(position / energy);
0129         if (msgLevel(MSG::DEBUG)) {
0130           debug() << "   --> Merged cluster with energy: " << new_clus.getEnergy() << endmsg;
0131         }
0132         auto ca = assoc2.create();
0133         ca.setSimID(mcID);
0134         ca.setWeight(1.0);
0135         ca.setRec(new_clus);
0136       }
0137     }
0138
0139     // That's all!
0140
0141     return StatusCode::SUCCESS;
0142   }
0143
0144   // get a map of MCParticle index--> std::vector<Cluster> for clusters that belong together
0145   std::map<int, std::vector<edm4eic::Cluster>> indexedClusterLists(
0146       const edm4eic::ClusterCollection& clusters,
0147       const edm4eic::MCRecoClusterParticleAssociationCollection& associations
0148   ) const {
0149
0150     std::map<int, std::vector<edm4eic::Cluster>> matched = {};
0151
0152     // loop over clusters
0153     for (const auto& cluster : clusters) {
0154
0155       int mcID = -1;
0156
0157       // find associated particle
0158       for (const auto& assoc : associations) {
0159         if (assoc.getRec() == cluster) {
0160           mcID = assoc.getSimID();
0161           break;
0162         }
0163       }
0164
0165       if (msgLevel(MSG::VERBOSE)) {
0166         verbose() << " --> Found cluster with mcID " << mcID << " and energy "
0167                   << cluster.getEnergy() << endmsg;
0168       }
0169
0170       if (mcID < 0) {
0171         if (msgLevel(MSG::VERBOSE)) {
0172           verbose() << "   --> WARNING: no valid MC truth link found, skipping cluster..." << endmsg;
0173         }
0174         continue;
0175       }
0176
0177       if (!matched.count(mcID)) {
0178         matched[mcID] = {};
0179       }
0180       matched[mcID].push_back(cluster);
0181     }
0182     return matched;
0183   }
0184 };
0185
0186 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0187 DECLARE_COMPONENT(ClusterMerger)
0188
0189 } // namespace Jug::Fast