algorithms/calorimetry/TruthEnergyPositionClusterMerger.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten
0003
0004 #pragma once
0005
0006 #include <limits>
0007
0008 #include <spdlog/spdlog.h>
0009
0010 // Event Model related classes
0011 #include <edm4hep/MCParticle.h>
0012 #include <edm4eic/Cluster.h>
0013 #include <edm4eic/MCRecoClusterParticleAssociation.h>
0014 #include <edm4hep/utils/vector_utils.h>
0015
0016 #include "algorithms/interfaces/WithPodConfig.h"
0017
0018 namespace eicrecon {
0019
0020   using TruthEnergyPositionClusterMergerAlgorithm = algorithms::Algorithm<
0021     algorithms::Input<
0022       edm4hep::MCParticleCollection,
0023       edm4eic::ClusterCollection,
0024       edm4eic::MCRecoClusterParticleAssociationCollection,
0025       edm4eic::ClusterCollection,
0026       edm4eic::MCRecoClusterParticleAssociationCollection
0027     >,
0028     algorithms::Output<
0029       edm4eic::ClusterCollection,
0030       edm4eic::MCRecoClusterParticleAssociationCollection
0031     >
0032   >;
0033
0034   /** Simple algorithm to merge the energy measurement from cluster1 with the position
0035    * measurement of cluster2 (in case matching clusters are found). If not, it will
0036    * propagate the raw cluster from cluster1 or cluster2
0037    *
0038    * Matching occurs based on the mc truth information of the clusters.
0039    *
0040    * \ingroup reco
0041    */
0042   class TruthEnergyPositionClusterMerger
0043       : public TruthEnergyPositionClusterMergerAlgorithm {
0044
0045   public:
0046     TruthEnergyPositionClusterMerger(std::string_view name)
0047       : TruthEnergyPositionClusterMergerAlgorithm{name,
0048                             {"mcParticles",
0049                              "energyClusterCollection", "energyClusterAssociations",
0050                              "positionClusterCollection", "positionClusterAssociations"},
0051                             {"outputClusterCollection", "outputClusterAssociations"},
0052                             "Merge energy and position clusters based on truth."} {}
0053
0054   public:
0055     void init() { }
0056
0057     void process(const Input& input, const Output& output) const final {
0058
0059         const auto [mcparticles, energy_clus, energy_assoc, pos_clus, pos_assoc] = input;
0060         auto [merged_clus, merged_assoc] = output;
0061
0062         debug( "Merging energy and position clusters for new event" );
0063
0064         if (energy_clus->size() == 0 && pos_clus->size() == 0) {
0065             debug( "Nothing to do for this event, returning..." );
0066             return;
0067         }
0068
0069         debug( "Step 0/2: Getting indexed list of clusters..." );
0070
0071         // get an indexed map of all clusters
0072         debug(" --> Indexing energy clusters");
0073         auto energyMap = indexedClusters(*energy_clus, *energy_assoc);
0074         trace(" --> Found these energy clusters:");
0075         for (const auto &[mcID, eclus]: energyMap) {
0076             trace("   --> energy cluster {}, mcID: {}, energy: {}", eclus.getObjectID().index, mcID, eclus.getEnergy());
0077         }
0078         debug(" --> Indexing position clusters");
0079         auto posMap = indexedClusters(*pos_clus, *pos_assoc);
0080         trace(" --> Found these position clusters:");
0081         for (const auto &[mcID, pclus]: posMap) {
0082             trace("   --> position cluster {}, mcID: {}, energy: {}", pclus.getObjectID().index, mcID, pclus.getEnergy());
0083         }
0084
0085         // loop over all position clusters and match with energy clusters
0086         debug( "Step 1/2: Matching all position clusters to the available energy clusters..." );
0087         for (const auto &[mcID, pclus]: posMap) {
0088
0089             debug(" --> Processing position cluster {}, mcID: {}, energy: {}", pclus.getObjectID().index, mcID, pclus.getEnergy());
0090             if (energyMap.count(mcID)) {
0091
0092                 const auto &eclus = energyMap[mcID];
0093
0094                 auto new_clus = merged_clus->create();
0095                 new_clus.setEnergy(eclus.getEnergy());
0096                 new_clus.setEnergyError(eclus.getEnergyError());
0097                 new_clus.setTime(pclus.getTime());
0098                 new_clus.setNhits(pclus.getNhits() + eclus.getNhits());
0099                 new_clus.setPosition(pclus.getPosition());
0100                 new_clus.setPositionError(pclus.getPositionError());
0101                 new_clus.addToClusters(pclus);
0102                 new_clus.addToClusters(eclus);
0103                 for (const auto &cl: {pclus, eclus}) {
0104                     for (const auto &hit: cl.getHits()) {
0105                         new_clus.addToHits(hit);
0106                     }
0107                     new_clus.addToSubdetectorEnergies(cl.getEnergy());
0108                 }
0109                 for (const auto &param: pclus.getShapeParameters()) {
0110                     new_clus.addToShapeParameters(param);
0111                 }
0112                 debug("   --> Found matching energy cluster {}, energy: {}", eclus.getObjectID().index, eclus.getEnergy() );
0113                 debug("   --> Created new combined cluster {}, energy: {}", new_clus.getObjectID().index, new_clus.getEnergy() );
0114
0115                 // set association
0116                 auto clusterassoc = merged_assoc->create();
0117                 clusterassoc.setRecID(new_clus.getObjectID().index);
0118                 clusterassoc.setSimID(mcID);
0119                 clusterassoc.setWeight(1.0);
0120                 clusterassoc.setRec(new_clus);
0121                 clusterassoc.setSim((*mcparticles)[mcID]);
0122
0123                 // erase the energy cluster from the map, so we can in the end account for all
0124                 // remaining clusters
0125                 energyMap.erase(mcID);
0126             } else {
0127                 debug("   --> No matching energy cluster found, copying over position cluster" );
0128                 auto new_clus = pclus.clone();
0129                 new_clus.addToClusters(pclus);
0130                 merged_clus->push_back(new_clus);
0131
0132                 // set association
0133                 auto clusterassoc = merged_assoc->create();
0134                 clusterassoc.setRecID(new_clus.getObjectID().index);
0135                 clusterassoc.setSimID(mcID);
0136                 clusterassoc.setWeight(1.0);
0137                 clusterassoc.setRec(new_clus);
0138                 clusterassoc.setSim((*mcparticles)[mcID]);
0139             }
0140         }
0141         // Collect remaining energy clusters. Use mc truth position for these clusters, as
0142         // they should really have a match in the position clusters (and if they don't it due
0143         // to a clustering error).
0144         debug( "Step 2/2: Collecting remaining energy clusters..." );
0145         for (const auto &[mcID, eclus]: energyMap) {
0146             const auto &mc = (*mcparticles)[mcID];
0147             const auto &p = mc.getMomentum();
0148             const auto phi = std::atan2(p.y, p.x);
0149             const auto theta = std::atan2(std::hypot(p.x, p.y), p.z);
0150
0151             auto new_clus = merged_clus->create();
0152             new_clus.setEnergy(eclus.getEnergy());
0153             new_clus.setEnergyError(eclus.getEnergyError());
0154             new_clus.setTime(eclus.getTime());
0155             new_clus.setNhits(eclus.getNhits());
0156             // FIXME use nominal dd4hep::radius of 110cm, and use start vertex theta and phi
0157             new_clus.setPosition(edm4hep::utils::sphericalToVector(78.5 * dd4hep::cm / dd4hep::mm, theta, phi));
0158             new_clus.addToClusters(eclus);
0159
0160             debug(" --> Processing energy cluster {}, mcID: {}, energy: {}", eclus.getObjectID().index, mcID, eclus.getEnergy() );
0161             debug("   --> Created new 'combined' cluster {}, energy: {}", new_clus.getObjectID().index, new_clus.getEnergy() );
0162
0163             // set association
0164             auto clusterassoc = merged_assoc->create();
0165             clusterassoc.setRecID(new_clus.getObjectID().index);
0166             clusterassoc.setSimID(mcID);
0167             clusterassoc.setWeight(1.0);
0168             clusterassoc.setRec(new_clus);
0169             clusterassoc.setSim(mc);
0170         }
0171     }
0172
0173     // get a map of MCParticle index --> cluster
0174     // input: cluster_collections --> list of handles to all cluster collections
0175     std::map<int, edm4eic::Cluster> indexedClusters(
0176             const edm4eic::ClusterCollection& clusters,
0177             const edm4eic::MCRecoClusterParticleAssociationCollection& associations
0178     ) const {
0179
0180         std::map<int, edm4eic::Cluster> matched = {};
0181
0182         for (const auto &cluster: clusters) {
0183             int mcID = -1;
0184
0185             // find associated particle
0186             for (const auto &assoc: associations) {
0187                 if (assoc.getRec() == cluster) {
0188                     mcID = assoc.getSimID();
0189                     break;
0190                 }
0191             }
0192
0193             trace(" --> Found cluster: {} with mcID {} and energy {}", cluster.getObjectID().index, mcID, cluster.getEnergy());
0194
0195             if (mcID < 0) {
0196                 trace( "   --> WARNING: no valid MC truth link found, skipping cluster..." );
0197                 continue;
0198             }
0199
0200             const bool duplicate = matched.count(mcID);
0201             if (duplicate) {
0202                 trace( "   --> WARNING: this is a duplicate mcID, keeping the higher energy cluster");
0203                 if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0204                     continue;
0205                 }
0206             }
0207
0208             matched[mcID] = cluster;
0209         }
0210         return matched;
0211     }
0212 };
0213
0214 } // namespace eicrecon