File indexing completed on 2024-09-27 07:03:00
0001
0002
0003
0004
0005
0006
0007
0008 #include <algorithms/logger.h>
0009 #include <edm4eic/ClusterCollection.h>
0010 #include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
0011 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0012 #include <edm4eic/ReconstructedParticleCollection.h>
0013 #include <edm4hep/MCParticleCollection.h>
0014 #include <edm4hep/Vector3f.h>
0015 #include <edm4hep/utils/vector_utils.h>
0016 #include <fmt/core.h>
0017 #include <podio/ObjectID.h>
0018 #include <cmath>
0019 #include <gsl/pointers>
0020 #include <map>
0021
0022 #include "MatchClusters.h"
0023
0024 namespace eicrecon {
0025
0026 void MatchClusters::process(
0027 const MatchClusters::Input& input,
0028 const MatchClusters::Output& output) const {
0029
0030 const auto [mcparticles, inparts, inpartsassoc, clusters, clustersassoc] = input;
0031 auto [outparts, outpartsassoc] = output;
0032
0033 debug("Processing cluster info for new event");
0034
0035 debug("Step 0/2: Getting indexed list of clusters...");
0036
0037
0038 auto clusterMap = indexedClusters(clusters, clustersassoc);
0039
0040
0041
0042 debug("Step 1/2: Matching clusters to charged particles...");
0043
0044 for (const auto inpart: *inparts) {
0045 debug(" --> Processing charged particle {}, PDG {}, energy {}", inpart.getObjectID().index,
0046 inpart.getPDG(), inpart.getEnergy());
0047
0048 auto outpart = inpart.clone();
0049 outparts->push_back(outpart);
0050
0051 int mcID = -1;
0052
0053
0054 for (const auto &assoc: *inpartsassoc) {
0055 if (assoc.getRec().getObjectID() == inpart.getObjectID()) {
0056 mcID = assoc.getSim().getObjectID().index;
0057 break;
0058 }
0059 }
0060
0061 trace(" --> Found particle with mcID {}", mcID);
0062
0063 if (mcID < 0) {
0064 debug(" --> cannot match track without associated mcID");
0065 continue;
0066 }
0067
0068 if (clusterMap.count(mcID)) {
0069 const auto &clus = clusterMap[mcID];
0070 debug(" --> found matching cluster with energy: {}", clus.getEnergy());
0071 debug(" --> adding cluster to reconstructed particle");
0072 outpart.addToClusters(clus);
0073 clusterMap.erase(mcID);
0074 }
0075
0076
0077 auto assoc = outpartsassoc->create();
0078 assoc.setRecID(outpart.getObjectID().index);
0079 assoc.setSimID(mcID);
0080 assoc.setWeight(1.0);
0081 assoc.setRec(outpart);
0082 assoc.setSim((*mcparticles)[mcID]);
0083 }
0084
0085
0086
0087 debug("Step 2/2: Creating neutrals for remaining clusters...");
0088 for (const auto &[mcID, clus]: clusterMap) {
0089 debug(" --> Processing unmatched cluster with energy: {}", clus.getEnergy());
0090
0091
0092
0093 const auto mc = (*mcparticles)[mcID];
0094 const double mass = (!mc.getCharge()) ? mc.getMass() : 0;
0095 const int32_t pdg = (!mc.getCharge()) ? mc.getPDG() : 0;
0096 if (level() <= algorithms::LogLevel::kDebug) {
0097 if (mc.getCharge()) {
0098 debug(" --> associated mcparticle is not a neutral (PDG: {}), "
0099 "setting the reconstructed particle ID to 0 (unidentified)", mc.getPDG());
0100 }
0101 debug(" --> found matching associated mcparticle with PDG: {}, energy: {}", pdg,
0102 mc.getEnergy());
0103 }
0104
0105
0106 const auto outpart = reconstruct_neutral(&clus, mass, pdg);
0107 debug(" --> Reconstructed neutral particle with PDG: {}, energy: {}", outpart.getPDG(),
0108 outpart.getEnergy());
0109
0110 outparts->push_back(outpart);
0111
0112
0113 auto assoc = outpartsassoc->create();
0114 assoc.setRecID(outpart.getObjectID().index);
0115 assoc.setSimID(mcID);
0116 assoc.setWeight(1.0);
0117 assoc.setRec(outpart);
0118 assoc.setSim((*mcparticles)[mcID]);
0119 }
0120 }
0121
0122
0123
0124 std::map<int, edm4eic::Cluster> MatchClusters::indexedClusters(
0125 const edm4eic::ClusterCollection* clusters,
0126 const edm4eic::MCRecoClusterParticleAssociationCollection* associations) const {
0127
0128 std::map<int, edm4eic::Cluster> matched = {};
0129
0130
0131 for (const auto cluster: *clusters) {
0132
0133 int mcID = -1;
0134
0135
0136 for (const auto assoc: *associations) {
0137 if (assoc.getRec() == cluster) {
0138 mcID = assoc.getSim().getObjectID().index;
0139 break;
0140 }
0141 }
0142
0143 trace(" --> Found cluster with mcID {} and energy {}", mcID, cluster.getEnergy());
0144
0145 if (mcID < 0) {
0146 trace(" --> WARNING: no valid MC truth link found, skipping cluster...");
0147 continue;
0148 }
0149
0150 const bool duplicate = matched.count(mcID);
0151 if (duplicate) {
0152 trace(" --> WARNING: this is a duplicate mcID, keeping the higher energy cluster");
0153
0154 if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0155 continue;
0156 }
0157 }
0158 matched[mcID] = cluster;
0159 }
0160 return matched;
0161 }
0162
0163
0164
0165
0166 edm4eic::MutableReconstructedParticle MatchClusters::reconstruct_neutral(
0167 const edm4eic::Cluster *cluster, const double mass, const int32_t pdg) const {
0168
0169 const float energy = cluster->getEnergy();
0170 const float p = energy < mass ? 0 : std::sqrt(energy * energy - mass * mass);
0171 const auto position = cluster->getPosition();
0172 const auto momentum = p * (position / edm4hep::utils::magnitude(position));
0173
0174 edm4eic::MutableReconstructedParticle part;
0175 part.setMomentum(momentum);
0176 part.setPDG(pdg);
0177 part.setCharge(0);
0178 part.setEnergy(energy);
0179 part.setMass(mass);
0180 part.addToClusters(*cluster);
0181 return part;
0182 }
0183
0184 }