algorithms/pid_lut/PhaseSpacePID.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2025, Nathan Brei, Dmitry Kalinkin, Simon Gardner
0003
0004 #include <DD4hep/Detector.h>
0005 #include <algorithms/geo.h>
0006 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0007 #include <edm4eic/ReconstructedParticleCollection.h>
0008 #include <edm4hep/MCParticleCollection.h>
0009 #include <edm4hep/Vector3f.h>
0010 #include <edm4hep/utils/vector_utils.h>
0011 #include <fmt/core.h>
0012 #include <exception>
0013 #include <gsl/pointers>
0014 #include <stdexcept>
0015 #include <vector>
0016
0017 #include "algorithms/interfaces/ParticleSvc.h"
0018 #include "algorithms/pid_lut/PhaseSpacePID.h"
0019 #include "algorithms/pid_lut/PhaseSpacePIDConfig.h"
0020
0021 namespace eicrecon {
0022
0023 void PhaseSpacePID::init() {
0024
0025   auto detector = algorithms::GeoSvc::instance().detector();
0026
0027   try {
0028     m_system = detector->constant<int32_t>(m_cfg.system);
0029   } catch (const std::exception& e) {
0030     error("Failed to get {} from the detector: {}", m_cfg.system, e.what());
0031     throw std::runtime_error("Failed to get requested ID from the detector");
0032   }
0033
0034   m_direction = edm4hep::Vector3f{m_cfg.direction[0], m_cfg.direction[1], m_cfg.direction[2]};
0035
0036   auto& particleSvc = algorithms::ParticleSvc::instance();
0037   m_mass            = particleSvc.particle(m_cfg.pdg_value).mass;
0038   m_charge          = particleSvc.particle(m_cfg.pdg_value).charge;
0039 }
0040
0041 void PhaseSpacePID::process(const Input& input, const Output& output) const {
0042   const auto [recoparts_in, partassocs_in]          = input;
0043   auto [recoparts_out, partassocs_out, partids_out] = output;
0044
0045   for (const auto& recopart_without_pid : *recoparts_in) {
0046     auto recopart = recopart_without_pid.clone();
0047
0048     // Find MCParticle from associations and propagate the relevant ones further
0049     auto best_assoc = edm4eic::MCRecoParticleAssociation::makeEmpty();
0050     for (auto assoc_in : *partassocs_in) {
0051       if (assoc_in.getRec() == recopart_without_pid) {
0052         if ((not best_assoc.isAvailable()) || (best_assoc.getWeight() < assoc_in.getWeight())) {
0053           best_assoc = assoc_in;
0054         }
0055         auto assoc_out = assoc_in.clone();
0056         assoc_out.setRec(recopart);
0057         partassocs_out->push_back(assoc_out);
0058       }
0059     }
0060     if (not best_assoc.isAvailable()) {
0061       recoparts_out->push_back(recopart);
0062       continue;
0063     }
0064
0065     edm4hep::MCParticle mcpart = best_assoc.getSim();
0066
0067     // Check if particle is within the phase space
0068     edm4hep::Vector3f momentum = recopart.getMomentum();
0069     double angle               = edm4hep::utils::angleBetween(momentum, m_direction);
0070     if (angle < m_cfg.opening_angle) {
0071
0072       auto partID = partids_out->create(m_system,        // std::int32_t type
0073                                         m_cfg.pdg_value, // std::int32_t PDG
0074                                         1,               // std::int32_t algorithmType
0075                                         1.0              // float likelihood
0076       );
0077
0078       recopart.addToParticleIDs(partID);
0079       recopart.setParticleIDUsed(partID);
0080       recopart.setPDG(m_cfg.pdg_value);
0081       recopart.setMass(m_mass);
0082     }
0083
0084     recoparts_out->push_back(recopart);
0085   }
0086 }
0087
0088 } // namespace eicrecon