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