truth/todo/TruthEnergyPositionClusterMerger.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 "edm4hep/MCParticleCollection.h"
0019 #include "edm4eic/ClusterCollection.h"
0020 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0021 #include <edm4hep/utils/vector_utils.h>
0022
0023 using namespace Gaudi::Units;
0024
0025 namespace Jug::Fast {
0026
0027 /** Simple algorithm to merge the energy measurement from cluster1 with the position
0028  * measurement of cluster2 (in case matching clusters are found). If not, it will
0029  * propagate the raw cluster from cluster1 or cluster2
0030  *
0031  * Matching occurs based on the mc truth information of the clusters.
0032  *
0033  * \ingroup reco
0034  */
0035 class TruthEnergyPositionClusterMerger : public GaudiAlgorithm {
0036 private:
0037   // Input
0038   DataHandle<edm4hep::MCParticleCollection> m_inputMCParticles{"MCParticles", Gaudi::DataHandle::Reader, this};
0039   DataHandle<edm4eic::ClusterCollection> m_energyClusters{"EnergyClusters", Gaudi::DataHandle::Reader, this};
0040   DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection> m_energyAssociations{"EnergyAssociations", Gaudi::DataHandle::Reader, this};
0041   DataHandle<edm4eic::ClusterCollection> m_positionClusters{"PositionClusters", Gaudi::DataHandle::Reader, this};
0042   DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection> m_positionAssociations{"PositionAssociations", Gaudi::DataHandle::Reader, this};
0043   // Output
0044   DataHandle<edm4eic::ClusterCollection> m_outputClusters{"OutputClusters", Gaudi::DataHandle::Writer, this};
0045   DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection> m_outputAssociations{"OutputAssociations", Gaudi::DataHandle::Writer, this};
0046
0047 public:
0048   TruthEnergyPositionClusterMerger(const std::string& name, ISvcLocator* svcLoc)
0049       : GaudiAlgorithm(name, svcLoc) {
0050     declareProperty("inputMCParticles", m_inputMCParticles, "MCParticles");
0051     declareProperty("inputEnergyClusters", m_energyClusters, "Cluster collection with good energy precision");
0052     declareProperty("inputEnergyAssociations", m_energyAssociations, "Cluster association with good energy precision");
0053     declareProperty("inputPositionClusters", m_positionClusters, "Cluster collection with good position precision");
0054     declareProperty("inputPositionAssociations", m_positionAssociations, "Cluster association with good position precision");
0055     declareProperty("outputClusters", m_outputClusters, "Cluster collection with good energy precision");
0056     declareProperty("outputAssociations", m_outputAssociations, "Cluster association with good energy precision");
0057   }
0058
0059   StatusCode initialize() override { return StatusCode::SUCCESS; }
0060
0061   StatusCode execute() override {
0062     if (msgLevel(MSG::DEBUG)) {
0063       debug() << "Merging energy and position clusters for new event" << endmsg;
0064     }
0065     // input
0066     const auto& mcparticles  = *(m_inputMCParticles.get());
0067     const auto& energy_clus  = *(m_energyClusters.get());
0068     const auto& energy_assoc = *(m_energyAssociations.get());
0069     const auto& pos_clus     = *(m_positionClusters.get());
0070     const auto& pos_assoc    = *(m_positionAssociations.get());
0071     // output
0072     auto& merged_clus  = *(m_outputClusters.createAndPut());
0073     auto& merged_assoc = *(m_outputAssociations.createAndPut());
0074
0075     if (!energy_clus.size() && !pos_clus.size()) {
0076       if (msgLevel(MSG::DEBUG)) {
0077         debug() << "Nothing to do for this event, returning..." << endmsg;
0078       }
0079       return StatusCode::SUCCESS;
0080     }
0081
0082     if (msgLevel(MSG::DEBUG)) {
0083       debug() << "Step 0/2: Getting indexed list of clusters..." << endmsg;
0084     }
0085     // get an indexed map of all clusters
0086     auto energyMap = indexedClusters(energy_clus, energy_assoc);
0087     auto posMap    = indexedClusters(pos_clus, pos_assoc);
0088
0089     // loop over all position clusters and match with energy clusters
0090     if (msgLevel(MSG::DEBUG)) {
0091       debug() << "Step 1/2: Matching all position clusters to the available energy clusters..." << endmsg;
0092     }
0093     for (const auto& [mcID, pclus] : posMap) {
0094       if (msgLevel(MSG::DEBUG)) {
0095         debug() << " --> Processing position cluster " << pclus.id() << ", mcID: " << mcID << ", energy: " << pclus.getEnergy()
0096                 << endmsg;
0097       }
0098       if (energyMap.count(mcID)) {
0099         const auto& eclus = energyMap[mcID];
0100         auto new_clus = merged_clus.create();
0101         new_clus.setEnergy(eclus.getEnergy());
0102         new_clus.setEnergyError(eclus.getEnergyError());
0103         new_clus.setTime(pclus.getTime());
0104         new_clus.setNhits(pclus.getNhits() + eclus.getNhits());
0105         new_clus.setPosition(pclus.getPosition());
0106         new_clus.setPositionError(pclus.getPositionError());
0107         new_clus.addToClusters(pclus);
0108         new_clus.addToClusters(eclus);
0109         for (const auto& cl : {pclus, eclus}) {
0110           for (const auto& hit : cl.getHits()) {
0111             new_clus.addToHits(hit);
0112           }
0113           new_clus.addToSubdetectorEnergies(cl.getEnergy());
0114         }
0115         for (const auto& param : pclus.getShapeParameters()) {
0116           new_clus.addToShapeParameters(param);
0117         }
0118         if (msgLevel(MSG::DEBUG)) {
0119           debug() << "   --> Found matching energy cluster " << eclus.id() << ", energy: " << eclus.getEnergy() << endmsg;
0120           debug() << "   --> Created new combined cluster " << new_clus.id() << ", energy: " << new_clus.getEnergy() << endmsg;
0121         }
0122
0123         // set association
0124         edm4eic::MutableMCRecoClusterParticleAssociation clusterassoc;
0125         clusterassoc.setRecID(new_clus.getObjectID().index);
0126         clusterassoc.setSimID(mcID);
0127         clusterassoc.setWeight(1.0);
0128         clusterassoc.setRec(new_clus);
0129         //clusterassoc.setSim(mcparticles[mcID]);
0130         merged_assoc.push_back(clusterassoc);
0131
0132         // erase the energy cluster from the map, so we can in the end account for all
0133         // remaining clusters
0134         energyMap.erase(mcID);
0135       } else {
0136         if (msgLevel(MSG::DEBUG)) {
0137           debug() << "   --> No matching energy cluster found, copying over position cluster" << endmsg;
0138         }
0139         auto new_clus = pclus.clone();
0140         new_clus.addToClusters(pclus);
0141         merged_clus.push_back(new_clus);
0142
0143         // set association
0144         edm4eic::MutableMCRecoClusterParticleAssociation clusterassoc;
0145         clusterassoc.setRecID(new_clus.getObjectID().index);
0146         clusterassoc.setSimID(mcID);
0147         clusterassoc.setWeight(1.0);
0148         clusterassoc.setRec(new_clus);
0149         //clusterassoc.setSim(mcparticles[mcID]);
0150         merged_assoc.push_back(clusterassoc);
0151       }
0152     }
0153     // Collect remaining energy clusters. Use mc truth position for these clusters, as
0154     // they should really have a match in the position clusters (and if they don't it due
0155     // to a clustering error).
0156     if (msgLevel(MSG::DEBUG)) {
0157       debug() << "Step 2/2: Collecting remaining energy clusters..." << endmsg;
0158     }
0159     for (const auto& [mcID, eclus] : energyMap) {
0160       const auto& mc = mcparticles[mcID];
0161       const auto& p = mc.getMomentum();
0162       const auto phi = std::atan2(p.y, p.x);
0163       const auto theta = std::atan2(std::hypot(p.x, p.y), p.z);
0164       auto new_clus = merged_clus.create();
0165       new_clus.setEnergy(eclus.getEnergy());
0166       new_clus.setEnergyError(eclus.getEnergyError());
0167       new_clus.setTime(eclus.getTime());
0168       new_clus.setNhits(eclus.getNhits());
0169       // use nominal radius of 110cm, and use start vertex theta and phi
0170       new_clus.setPosition(edm4hep::utils::sphericalToVector(110.*cm, theta, phi));
0171       new_clus.addToClusters(eclus);
0172       if (msgLevel(MSG::DEBUG)) {
0173         debug() << " --> Processing energy cluster " << eclus.id() << ", mcID: " << mcID << ", energy: " << eclus.getEnergy()
0174                 << endmsg;
0175         debug() << "   --> Created new 'combined' cluster " << new_clus.id() << ", energy: " << new_clus.getEnergy() << endmsg;
0176       }
0177
0178       // set association
0179       edm4eic::MutableMCRecoClusterParticleAssociation clusterassoc;
0180       clusterassoc.setRecID(new_clus.getObjectID().index);
0181       clusterassoc.setSimID(mcID);
0182       clusterassoc.setWeight(1.0);
0183       clusterassoc.setRec(new_clus);
0184       //clusterassoc.setSim(mc);
0185       merged_assoc.push_back(clusterassoc);
0186     }
0187
0188     // That's all!
0189     return StatusCode::SUCCESS;
0190   }
0191
0192   // get a map of MCParticle index --> cluster
0193   // input: cluster_collections --> list of handles to all cluster collections
0194   std::map<int, edm4eic::Cluster> indexedClusters(
0195       const edm4eic::ClusterCollection& clusters,
0196       const edm4eic::MCRecoClusterParticleAssociationCollection& associations
0197   ) const {
0198
0199     std::map<int, edm4eic::Cluster> matched = {};
0200
0201     for (const auto& cluster : clusters) {
0202       int mcID = -1;
0203
0204       // find associated particle
0205       for (const auto& assoc : associations) {
0206         if (assoc.getRec() == cluster) {
0207           mcID = assoc.getSimID();
0208           break;
0209         }
0210       }
0211
0212       if (msgLevel(MSG::VERBOSE)) {
0213         verbose() << " --> Found cluster: " << cluster.getObjectID().index << " with mcID " << mcID << " and energy "
0214                   << cluster.getEnergy() << endmsg;
0215       }
0216
0217       if (mcID < 0) {
0218         if (msgLevel(MSG::VERBOSE)) {
0219           verbose() << "   --> WARNING: no valid MC truth link found, skipping cluster..." << endmsg;
0220         }
0221         continue;
0222       }
0223
0224       const bool duplicate = matched.count(mcID);
0225       if (duplicate) {
0226         if (msgLevel(MSG::VERBOSE)) {
0227           verbose() << "   --> WARNING: this is a duplicate mcID, keeping the higher energy cluster" << endmsg;
0228         }
0229         if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0230           continue;
0231         }
0232       }
0233
0234       matched[mcID] = cluster;
0235     }
0236     return matched;
0237   }
0238 };
0239
0240 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0241 DECLARE_COMPONENT(TruthEnergyPositionClusterMerger)
0242
0243 } // namespace Jug::Fast