File indexing completed on 2025-02-24 10:15:51
0001
0002
0003
0004 #pragma once
0005
0006 #include <limits>
0007
0008 #include <spdlog/spdlog.h>
0009
0010
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
0036
0037
0038
0039
0040
0041
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
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
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 ¶m: 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
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
0125
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
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
0143
0144
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
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
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
0175
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
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 }