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