File indexing completed on 2024-09-27 07:02:31
0001
0002
0003
0004
0005
0006
0007
0008 #include <algorithm>
0009 #include <cmath>
0010
0011 #include <fmt/format.h>
0012
0013 #include "Gaudi/Algorithm.h"
0014 #include "GaudiAlg/GaudiTool.h"
0015 #include "GaudiAlg/Producer.h"
0016 #include "GaudiAlg/Transformer.h"
0017 #include "GaudiKernel/RndmGenerators.h"
0018
0019 #include "JugBase/DataHandle.h"
0020
0021
0022 #include "edm4hep/MCParticleCollection.h"
0023 #include "edm4eic/ClusterCollection.h"
0024 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0025 #include "edm4eic/MCRecoParticleAssociationCollection.h"
0026 #include "edm4eic/ReconstructedParticleCollection.h"
0027 #include "edm4eic/TrackParametersCollection.h"
0028 #include "edm4hep/utils/vector_utils.h"
0029
0030 namespace Jug::Fast {
0031
0032 class MatchClusters : public GaudiAlgorithm {
0033 private:
0034
0035 DataHandle<edm4hep::MCParticleCollection> m_inputMCParticles{"MCParticles", Gaudi::DataHandle::Reader, this};
0036 DataHandle<edm4eic::ReconstructedParticleCollection> m_inputParticles{"ReconstructedChargedParticles",
0037 Gaudi::DataHandle::Reader, this};
0038 DataHandle<edm4eic::MCRecoParticleAssociationCollection> m_inputParticlesAssoc{"ReconstructedChargedParticlesAssoc",
0039 Gaudi::DataHandle::Reader, this};
0040 Gaudi::Property<std::vector<std::string>> m_inputClusters{this, "inputClusters", {}, "Clusters to be aggregated"};
0041 Gaudi::Property<std::vector<std::string>> m_inputClustersAssoc{this, "inputClustersAssoc", {}, "Cluster associations to be aggregated"};
0042 std::vector<DataHandle<edm4eic::ClusterCollection>*> m_inputClustersCollections;
0043 std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> m_inputClustersAssocCollections;
0044
0045
0046 DataHandle<edm4eic::ReconstructedParticleCollection> m_outputParticles{"ReconstructedParticles",
0047 Gaudi::DataHandle::Writer, this};
0048 DataHandle<edm4eic::MCRecoParticleAssociationCollection> m_outputParticlesAssoc{"ReconstructedParticlesAssoc",
0049 Gaudi::DataHandle::Writer, this};
0050
0051 public:
0052 MatchClusters(const std::string& name, ISvcLocator* svcLoc)
0053 : GaudiAlgorithm(name, svcLoc) {
0054 declareProperty("inputMCParticles", m_inputMCParticles, "MCParticles");
0055 declareProperty("inputParticles", m_inputParticles, "ReconstructedChargedParticles");
0056 declareProperty("inputParticlesAssoc", m_inputParticlesAssoc, "ReconstructedChargedParticlesAssoc");
0057 declareProperty("outputParticles", m_outputParticles, "ReconstructedParticles");
0058 declareProperty("outputParticlesAssoc", m_outputParticlesAssoc, "ReconstructedParticlesAssoc");
0059 }
0060
0061 StatusCode initialize() override {
0062 if (GaudiAlgorithm::initialize().isFailure()) {
0063 return StatusCode::FAILURE;
0064 }
0065 m_inputClustersCollections = getClusterCollections(m_inputClusters);
0066 m_inputClustersAssocCollections = getClusterAssociations(m_inputClustersAssoc);
0067 return StatusCode::SUCCESS;
0068 }
0069 StatusCode execute() override {
0070 if (msgLevel(MSG::DEBUG)) {
0071 debug() << "Processing cluster info for new event" << endmsg;
0072 }
0073
0074 const auto& mcparticles = *(m_inputMCParticles.get());
0075 const auto& inparts = *(m_inputParticles.get());
0076 const auto& inpartsassoc = *(m_inputParticlesAssoc.get());
0077 auto& outparts = *(m_outputParticles.createAndPut());
0078 auto& outpartsassoc = *(m_outputParticlesAssoc.createAndPut());
0079
0080 if (msgLevel(MSG::DEBUG)) {
0081 debug() << "Step 0/2: Getting indexed list of clusters..." << endmsg;
0082 }
0083
0084
0085 auto clusterMap = indexedClusters(m_inputClustersCollections, m_inputClustersAssocCollections);
0086
0087
0088
0089 if (msgLevel(MSG::DEBUG)) {
0090 debug() << "Step 1/2: Matching clusters to charged particles..." << endmsg;
0091 }
0092 for (const auto& inpart: inparts) {
0093 if (msgLevel(MSG::DEBUG)) {
0094 debug() << " --> Processing charged particle " << inpart.getObjectID().index
0095 << ", PDG: " << inpart.getPDG()
0096 << ", energy: " << inpart.getEnergy()
0097 << endmsg;
0098 }
0099
0100 auto outpart = inpart.clone();
0101 outparts.push_back(outpart);
0102
0103 int mcID = -1;
0104
0105
0106 for (const auto& assoc: inpartsassoc) {
0107 if (assoc.getRec() == inpart) {
0108 mcID = assoc.getSimID();
0109 break;
0110 }
0111 }
0112
0113 if (msgLevel(MSG::VERBOSE)) {
0114 verbose() << " --> Found particle with mcID " << mcID << endmsg;
0115 }
0116
0117 if (mcID < 0) {
0118 if (msgLevel(MSG::DEBUG)) {
0119 debug() << " --> cannot match track without associated mcID" << endmsg;
0120 }
0121 continue;
0122 }
0123
0124 if (clusterMap.count(mcID)) {
0125 const auto& clus = clusterMap[mcID];
0126 if (msgLevel(MSG::DEBUG)) {
0127 debug() << " --> found matching cluster with energy: " << clus.getEnergy() << endmsg;
0128 }
0129 clusterMap.erase(mcID);
0130 }
0131
0132
0133 auto assoc = outpartsassoc.create();
0134 assoc.setRecID(outpart.getObjectID().index);
0135 assoc.setSimID(mcID);
0136 assoc.setWeight(1.0);
0137 assoc.setRec(outpart);
0138
0139 }
0140
0141
0142
0143 if (msgLevel(MSG::DEBUG)) {
0144 debug() << "Step 2/2: Creating neutrals for remaining clusters..." << endmsg;
0145 }
0146 for (const auto& [mcID, clus] : clusterMap) {
0147 if (msgLevel(MSG::DEBUG)) {
0148 debug() << " --> Processing unmatched cluster with energy: " << clus.getEnergy()
0149 << endmsg;
0150 }
0151
0152
0153 const auto& mc = mcparticles[mcID];
0154 const double mass = (!mc.getCharge()) ? mc.getMass() : 0;
0155 const int32_t pdg = (!mc.getCharge()) ? mc.getPDG() : 0;
0156 if (msgLevel(MSG::DEBUG)) {
0157 if (mc.getCharge()) {
0158 debug() << " --> associated mcparticle is not a neutral (PDG: " << mc.getPDG()
0159 << "), setting the reconstructed particle ID to 0 (unidentified)" << endmsg;
0160 }
0161 debug() << " --> found matching associated mcparticle with PDG: " << pdg << ", energy: " << mc.getEnergy()
0162 << endmsg;
0163 }
0164
0165
0166 const auto outpart = reconstruct_neutral(clus, mass, pdg);
0167 if (msgLevel(MSG::DEBUG)) {
0168 debug() << " --> Reconstructed neutral particle with PDG: " << outpart.getPDG()
0169 << ", energy: " << outpart.getEnergy()
0170 << endmsg;
0171 }
0172 outparts.push_back(outpart);
0173
0174
0175 auto assoc = outpartsassoc.create();
0176 assoc.setRecID(outpart.getObjectID().index);
0177 assoc.setSimID(mcID);
0178 assoc.setWeight(1.0);
0179 assoc.setRec(outpart);
0180
0181 }
0182 return StatusCode::SUCCESS;
0183 }
0184
0185 private:
0186 std::vector<DataHandle<edm4eic::ClusterCollection>*> getClusterCollections(const std::vector<std::string>& cols) {
0187 std::vector<DataHandle<edm4eic::ClusterCollection>*> ret;
0188 for (const auto& colname : cols) {
0189 debug() << "initializing cluster collection: " << colname << endmsg;
0190 ret.push_back(new DataHandle<edm4eic::ClusterCollection>{colname, Gaudi::DataHandle::Reader, this});
0191 }
0192 return ret;
0193 }
0194
0195 std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> getClusterAssociations(const std::vector<std::string>& cols) {
0196 std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*> ret;
0197 for (const auto& colname : cols) {
0198 debug() << "initializing cluster association collection: " << colname << endmsg;
0199 ret.push_back(new DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>{colname, Gaudi::DataHandle::Reader, this});
0200 }
0201 return ret;
0202 }
0203
0204
0205
0206 std::map<int, edm4eic::Cluster>
0207 indexedClusters(
0208 const std::vector<DataHandle<edm4eic::ClusterCollection>*>& cluster_collections,
0209 const std::vector<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>*>& associations_collections
0210 ) const {
0211 std::map<int, edm4eic::Cluster> matched = {};
0212
0213
0214 for (const auto& cluster_handle : cluster_collections) {
0215 const auto& clusters = *(cluster_handle->get());
0216
0217
0218 for (const auto& cluster : clusters) {
0219
0220 int mcID = -1;
0221
0222
0223 for (const auto& associations_handle : associations_collections) {
0224 const auto& associations = *(associations_handle->get());
0225
0226
0227 for (const auto& assoc : associations) {
0228 if (assoc.getRec() == cluster) {
0229 mcID = assoc.getSimID();
0230 break;
0231 }
0232 }
0233
0234
0235 if (mcID != -1) {
0236 break;
0237 }
0238 }
0239
0240 if (msgLevel(MSG::VERBOSE)) {
0241 verbose() << " --> Found cluster with mcID " << mcID << " and energy "
0242 << cluster.getEnergy() << endmsg;
0243 }
0244
0245 if (mcID < 0) {
0246 if (msgLevel(MSG::VERBOSE)) {
0247 verbose() << " --> WARNING: no valid MC truth link found, skipping cluster..." << endmsg;
0248 }
0249 continue;
0250 }
0251
0252 const bool duplicate = matched.count(mcID);
0253 if (duplicate) {
0254 if (msgLevel(MSG::VERBOSE)) {
0255 verbose() << " --> WARNING: this is a duplicate mcID, keeping the higher energy cluster" << endmsg;
0256 }
0257 if (cluster.getEnergy() < matched[mcID].getEnergy()) {
0258 continue;
0259 }
0260 }
0261 matched[mcID] = cluster;
0262 }
0263 }
0264 return matched;
0265 }
0266
0267
0268
0269 edm4eic::ReconstructedParticle reconstruct_neutral(const edm4eic::Cluster& clus, const double mass,
0270 const int32_t pdg) const {
0271 const float energy = clus.getEnergy();
0272 const float p = energy < mass ? 0 : std::sqrt(energy * energy - mass * mass);
0273 const auto position = clus.getPosition();
0274 const auto momentum = p * (position / edm4hep::utils::magnitude(position));
0275
0276 edm4eic::MutableReconstructedParticle part;
0277 part.setMomentum(momentum);
0278 part.setPDG(pdg);
0279 part.setCharge(0);
0280 part.setEnergy(energy);
0281 part.setMass(mass);
0282 return part;
0283 }
0284 };
0285
0286
0287 DECLARE_COMPONENT(MatchClusters)
0288
0289 }