algorithms/pid_lut/PIDLookup.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024, Nathan Brei, Dmitry Kalinkin
0003
0004 #include <algorithms/service.h>
0005 #include <edm4eic/MCRecoParticleAssociationCollection.h>
0006 #include <edm4eic/ReconstructedParticleCollection.h>
0007 #include <edm4hep/MCParticleCollection.h>
0008 #include <edm4hep/Vector3f.h>
0009 #include <edm4hep/utils/vector_utils.h>
0010 #include <podio/LinkNavigator.h>
0011 #include <podio/detail/Link.h>
0012 #include <podio/detail/LinkCollectionImpl.h>
0013 #include <cmath>
0014 #include <exception>
0015 #include <limits>
0016 #include <memory>
0017 #include <random>
0018 #include <stdexcept>
0019 #include <tuple>
0020 #include <vector>
0021
0022 #include "algorithms/interfaces/LinkTruthUtils.h"
0023 #include "algorithms/pid_lut/PIDLookup.h"
0024 #include "algorithms/pid_lut/PIDLookupConfig.h"
0025 #include "services/pid_lut/PIDLookupTableSvc.h"
0026
0027 namespace eicrecon {
0028
0029 void PIDLookup::init() {
0030
0031   try {
0032     m_system = m_detector->constant<int32_t>(m_cfg.system);
0033   } catch (const std::exception& e) {
0034     error("Failed to get {} from the detector: {}", m_cfg.system, e.what());
0035     throw std::runtime_error("Failed to get requested ID from the detector");
0036   }
0037
0038   auto& serviceSvc = algorithms::ServiceSvc::instance();
0039   auto* lut_svc    = serviceSvc.service<PIDLookupTableSvc>("PIDLookupTableSvc");
0040
0041   m_lut = lut_svc->load(m_cfg.filename,
0042                         {
0043                             .pdg_values                   = m_cfg.pdg_values,
0044                             .charge_values                = m_cfg.charge_values,
0045                             .momentum_edges               = m_cfg.momentum_edges,
0046                             .polar_edges                  = m_cfg.polar_edges,
0047                             .azimuthal_binning            = m_cfg.azimuthal_binning,
0048                             .azimuthal_bin_centers_in_lut = m_cfg.azimuthal_bin_centers_in_lut,
0049                             .momentum_bin_centers_in_lut  = m_cfg.momentum_bin_centers_in_lut,
0050                             .polar_bin_centers_in_lut     = m_cfg.polar_bin_centers_in_lut,
0051                             .use_radians                  = m_cfg.use_radians,
0052                             .missing_electron_prob        = m_cfg.missing_electron_prob,
0053                         });
0054   if (m_lut == nullptr) {
0055     throw std::runtime_error("LUT not available");
0056   }
0057 }
0058
0059 void PIDLookup::process(const Input& input, const Output& output) const {
0060   const auto [headers, recoparts_in, partlinks_in]                 = input;
0061   auto [recoparts_out, partlinks_out, partassocs_out, partids_out] = output;
0062   const truth::EventLinkNavigator<edm4eic::MCRecoParticleLinkCollection> link_nav(partlinks_in);
0063
0064   // local random generator
0065   auto seed = m_uid.getUniqueID(*headers, name());
0066   std::default_random_engine generator(seed);
0067   std::uniform_real_distribution<double> uniform;
0068
0069   for (const auto& recopart_without_pid : *recoparts_in) {
0070     auto recopart = recopart_without_pid.clone();
0071
0072     // Find MCParticle from links and propagate the relevant ones further
0073     edm4hep::MCParticle best_sim;
0074     float best_weight = std::numeric_limits<float>::lowest();
0075     bool has_best     = false;
0076     for (const auto& [sim_particle, weight] : link_nav.linked(recopart_without_pid)) {
0077       if (!has_best || best_weight < weight) {
0078         best_sim    = sim_particle;
0079         best_weight = weight;
0080         has_best    = true;
0081       }
0082       auto link_out = partlinks_out->create();
0083       link_out.setFrom(recopart);
0084       link_out.setTo(sim_particle);
0085       link_out.setWeight(weight);
0086       auto assoc_out = partassocs_out->create();
0087       assoc_out.setRec(recopart);
0088       assoc_out.setSim(sim_particle);
0089       assoc_out.setWeight(weight);
0090     }
0091     if (!has_best) {
0092       recoparts_out->push_back(recopart);
0093       continue;
0094     }
0095
0096     edm4hep::MCParticle mcpart = best_sim;
0097
0098     int true_pdg    = mcpart.getPDG();
0099     int true_charge = mcpart.getCharge();
0100     int charge      = recopart.getCharge();
0101     double momentum = edm4hep::utils::magnitude(recopart.getMomentum());
0102
0103     double theta = edm4hep::utils::anglePolar(recopart.getMomentum()) / M_PI * 180.;
0104     double phi   = edm4hep::utils::angleAzimuthal(recopart.getMomentum()) / M_PI * 180.;
0105
0106     trace("lookup for true_pdg={}, true_charge={}, momentum={:.2f} GeV, polar={:.2f}, "
0107           "aziumthal={:.2f}",
0108           true_pdg, true_charge, momentum, theta, phi);
0109     const auto* entry = m_lut->Lookup(true_pdg, true_charge, momentum, theta, phi);
0110
0111     int identified_pdg = 0; // unknown
0112
0113     if ((entry != nullptr) && ((entry->prob_electron != 0.) || (entry->prob_pion != 0.) ||
0114                                (entry->prob_kaon != 0.) || (entry->prob_proton != 0.))) {
0115       double random_unit_interval = uniform(generator);
0116
0117       trace("entry with e:pi:K:P={}:{}:{}:{}", entry->prob_electron, entry->prob_pion,
0118             entry->prob_kaon, entry->prob_proton);
0119
0120       recopart.addToParticleIDs(
0121           partids_out->create(m_system,                                // std::int32_t type
0122                               std::copysign(11, -charge),              // std::int32_t PDG
0123                               0,                                       // std::int32_t algorithmType
0124                               static_cast<float>(entry->prob_electron) // float likelihood
0125                               ));
0126       recopart.addToParticleIDs(
0127           partids_out->create(m_system,                            // std::int32_t type
0128                               std::copysign(211, charge),          // std::int32_t PDG
0129                               0,                                   // std::int32_t algorithmType
0130                               static_cast<float>(entry->prob_pion) // float likelihood
0131                               ));
0132       recopart.addToParticleIDs(
0133           partids_out->create(m_system,                            // std::int32_t type
0134                               std::copysign(321, charge),          // std::int32_t PDG
0135                               0,                                   // std::int32_t algorithmType
0136                               static_cast<float>(entry->prob_kaon) // float likelihood
0137                               ));
0138       recopart.addToParticleIDs(
0139           partids_out->create(m_system,                              // std::int32_t type
0140                               std::copysign(2212, charge),           // std::int32_t PDG
0141                               0,                                     // std::int32_t algorithmType
0142                               static_cast<float>(entry->prob_proton) // float likelihood
0143                               ));
0144
0145       if (random_unit_interval < entry->prob_electron) {
0146         identified_pdg = 11; // electron
0147         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 4]);
0148       } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion)) {
0149         identified_pdg = 211; // pion
0150         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 3]);
0151       } else if (random_unit_interval <
0152                  (entry->prob_electron + entry->prob_pion + entry->prob_kaon)) {
0153         identified_pdg = 321; // kaon
0154         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 2]);
0155       } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion +
0156                                          entry->prob_kaon + entry->prob_proton)) {
0157         identified_pdg = 2212; // proton
0158         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 1]);
0159       }
0160     }
0161
0162     if (identified_pdg != 0) {
0163       recopart.setPDG(std::copysign(identified_pdg, (identified_pdg == 11) ? -charge : charge));
0164       recopart.setMass(m_particleSvc.particle(identified_pdg).mass);
0165       recopart.setEnergy(std::hypot(momentum, m_particleSvc.particle(identified_pdg).mass));
0166     }
0167
0168     if (identified_pdg != 0) {
0169       trace("randomized PDG is {}", recopart.getPDG());
0170     }
0171
0172     recoparts_out->push_back(recopart);
0173   }
0174 }
0175
0176 } // namespace eicrecon