algorithms/calorimetry/TruthEnergyPositionClusterMerger.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten
0003
0004 #include "algorithms/calorimetry/TruthEnergyPositionClusterMerger.h"
0005
0006 #include <Evaluator/DD4hepUnits.h>
0007 #include <edm4hep/Vector3d.h>
0008 #include <edm4hep/utils/vector_utils.h>
0009 #include <fmt/core.h>
0010 #include <podio/ObjectID.h>
0011 #include <podio/RelationRange.h>
0012 #include <cmath>
0013 #include <gsl/pointers>
0014 #include <initializer_list>
0015 #include <vector>
0016
0017 namespace eicrecon {
0018
0019 void TruthEnergyPositionClusterMerger::process(const Input& input, const Output& output) const {
0020
0021   const auto [mcparticles, energy_clus, energy_assoc, pos_clus, pos_assoc] = input;
0022   auto [merged_clus, merged_assoc]                                         = output;
0023
0024   debug("Merging energy and position clusters for new event");
0025
0026   if (energy_clus->empty() && pos_clus->empty()) {
0027     debug("Nothing to do for this event, returning...");
0028     return;
0029   }
0030
0031   debug("Step 0/2: Getting indexed list of clusters...");
0032
0033   // get an indexed map of all clusters
0034   debug(" --> Indexing energy clusters");
0035   auto energyMap = indexedClusters(*energy_clus, *energy_assoc);
0036   trace(" --> Found these energy clusters:");
0037   for (const auto& [mcID, eclus] : energyMap) {
0038     trace("   --> energy cluster {}, mcID: {}, energy: {}", eclus.getObjectID().index, mcID,
0039           eclus.getEnergy());
0040   }
0041   debug(" --> Indexing position clusters");
0042   auto posMap = indexedClusters(*pos_clus, *pos_assoc);
0043   trace(" --> Found these position clusters:");
0044   for (const auto& [mcID, pclus] : posMap) {
0045     trace("   --> position cluster {}, mcID: {}, energy: {}", pclus.getObjectID().index, mcID,
0046           pclus.getEnergy());
0047   }
0048
0049   // loop over all position clusters and match with energy clusters
0050   debug("Step 1/2: Matching all position clusters to the available energy clusters...");
0051   for (const auto& [mcID, pclus] : posMap) {
0052
0053     debug(" --> Processing position cluster {}, mcID: {}, energy: {}", pclus.getObjectID().index,
0054           mcID, pclus.getEnergy());
0055     if (energyMap.count(mcID)) {
0056
0057       const auto& eclus = energyMap[mcID];
0058
0059       auto new_clus = merged_clus->create();
0060       new_clus.setEnergy(eclus.getEnergy());
0061       new_clus.setEnergyError(eclus.getEnergyError());
0062       new_clus.setTime(pclus.getTime());
0063       new_clus.setNhits(pclus.getNhits() + eclus.getNhits());
0064       new_clus.setPosition(pclus.getPosition());
0065       new_clus.setPositionError(pclus.getPositionError());
0066       new_clus.addToClusters(pclus);
0067       new_clus.addToClusters(eclus);
0068       for (const auto& cl : {pclus, eclus}) {
0069         for (const auto& hit : cl.getHits()) {
0070           new_clus.addToHits(hit);
0071         }
0072         new_clus.addToSubdetectorEnergies(cl.getEnergy());
0073       }
0074       for (const auto& param : pclus.getShapeParameters()) {
0075         new_clus.addToShapeParameters(param);
0076       }
0077       debug("   --> Found matching energy cluster {}, energy: {}", eclus.getObjectID().index,
0078             eclus.getEnergy());
0079       debug("   --> Created new combined cluster {}, energy: {}", new_clus.getObjectID().index,
0080             new_clus.getEnergy());
0081
0082       // set association
0083       auto clusterassoc = merged_assoc->create();
0084       clusterassoc.setRecID(new_clus.getObjectID().index);
0085       clusterassoc.setSimID(mcID);
0086       clusterassoc.setWeight(1.0);
0087       clusterassoc.setRec(new_clus);
0088       clusterassoc.setSim((*mcparticles)[mcID]);
0089
0090       // erase the energy cluster from the map, so we can in the end account for all
0091       // remaining clusters
0092       energyMap.erase(mcID);
0093     } else {
0094       debug("   --> No matching energy cluster found, copying over position cluster");
0095       auto new_clus = pclus.clone();
0096       new_clus.addToClusters(pclus);
0097       merged_clus->push_back(new_clus);
0098
0099       // set association
0100       auto clusterassoc = merged_assoc->create();
0101       clusterassoc.setRecID(new_clus.getObjectID().index);
0102       clusterassoc.setSimID(mcID);
0103       clusterassoc.setWeight(1.0);
0104       clusterassoc.setRec(new_clus);
0105       clusterassoc.setSim((*mcparticles)[mcID]);
0106     }
0107   }
0108   // Collect remaining energy clusters. Use mc truth position for these clusters, as
0109   // they should really have a match in the position clusters (and if they don't it due
0110   // to a clustering error).
0111   debug("Step 2/2: Collecting remaining energy clusters...");
0112   for (const auto& [mcID, eclus] : energyMap) {
0113     const auto& mc   = (*mcparticles)[mcID];
0114     const auto& p    = mc.getMomentum();
0115     const auto phi   = std::atan2(p.y, p.x);
0116     const auto theta = std::atan2(std::hypot(p.x, p.y), p.z);
0117
0118     auto new_clus = merged_clus->create();
0119     new_clus.setEnergy(eclus.getEnergy());
0120     new_clus.setEnergyError(eclus.getEnergyError());
0121     new_clus.setTime(eclus.getTime());
0122     new_clus.setNhits(eclus.getNhits());
0123     // FIXME use nominal dd4hep::radius of 110cm, and use start vertex theta and phi
0124     new_clus.setPosition(
0125         edm4hep::utils::sphericalToVector(78.5 * dd4hep::cm / dd4hep::mm, theta, phi));
0126     new_clus.addToClusters(eclus);
0127
0128     debug(" --> Processing energy cluster {}, mcID: {}, energy: {}", eclus.getObjectID().index,
0129           mcID, eclus.getEnergy());
0130     debug("   --> Created new 'combined' cluster {}, energy: {}", new_clus.getObjectID().index,
0131           new_clus.getEnergy());
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(mc);
0140   }
0141 }
0142
0143 // get a map of MCParticle index --> cluster
0144 // input: cluster_collections --> list of handles to all cluster collections
0145 std::map<int, edm4eic::Cluster> TruthEnergyPositionClusterMerger::indexedClusters(
0146     const edm4eic::ClusterCollection& clusters,
0147     const edm4eic::MCRecoClusterParticleAssociationCollection& associations) const {
0148
0149   std::map<int, edm4eic::Cluster> matched = {};
0150
0151   for (const auto& cluster : clusters) {
0152     int mcID = -1;
0153
0154     // find associated particle
0155     for (const auto& assoc : associations) {
0156       if (assoc.getRec() == cluster) {
0157         mcID = assoc.getSimID();
0158         break;
0159       }
0160     }
0161
0162     trace(" --> Found cluster: {} with mcID {} and energy {}", cluster.getObjectID().index, mcID,
0163           cluster.getEnergy());
0164
0165     if (mcID < 0) {
0166       trace("   --> WARNING: no valid MC truth link found, skipping cluster...");
0167       continue;
0168     }
0169
0170     const bool duplicate = matched.count(mcID);
0171     if (duplicate) {
0172       trace("   --> WARNING: this is a duplicate mcID, keeping the higher energy cluster");
0173       if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0174         continue;
0175       }
0176     }
0177
0178     matched[mcID] = cluster;
0179   }
0180   return matched;
0181 }
0182
0183 } // namespace eicrecon