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 <fmt/core.h>
0011 #include <cmath>
0012 #include <exception>
0013 #include <gsl/pointers>
0014 #include <random>
0015 #include <stdexcept>
0016 #include <vector>
0017
0018 #include "algorithms/pid_lut/PIDLookup.h"
0019 #include "algorithms/pid_lut/PIDLookupConfig.h"
0020 #include "services/pid_lut/PIDLookupTableSvc.h"
0021
0022 namespace eicrecon {
0023
0024 void PIDLookup::init() {
0025
0026   try {
0027     m_system = m_detector->constant<int32_t>(m_cfg.system);
0028   } catch (const std::exception& e) {
0029     error("Failed to get {} from the detector: {}", m_cfg.system, e.what());
0030     throw std::runtime_error("Failed to get requested ID from the detector");
0031   }
0032
0033   auto& serviceSvc = algorithms::ServiceSvc::instance();
0034   auto* lut_svc    = serviceSvc.service<PIDLookupTableSvc>("PIDLookupTableSvc");
0035
0036   m_lut = lut_svc->load(m_cfg.filename,
0037                         {
0038                             .pdg_values                   = m_cfg.pdg_values,
0039                             .charge_values                = m_cfg.charge_values,
0040                             .momentum_edges               = m_cfg.momentum_edges,
0041                             .polar_edges                  = m_cfg.polar_edges,
0042                             .azimuthal_binning            = m_cfg.azimuthal_binning,
0043                             .azimuthal_bin_centers_in_lut = m_cfg.azimuthal_bin_centers_in_lut,
0044                             .momentum_bin_centers_in_lut  = m_cfg.momentum_bin_centers_in_lut,
0045                             .polar_bin_centers_in_lut     = m_cfg.polar_bin_centers_in_lut,
0046                             .use_radians                  = m_cfg.use_radians,
0047                             .missing_electron_prob        = m_cfg.missing_electron_prob,
0048                         });
0049   if (m_lut == nullptr) {
0050     throw std::runtime_error("LUT not available");
0051   }
0052 }
0053
0054 void PIDLookup::process(const Input& input, const Output& output) const {
0055   const auto [headers, recoparts_in, partassocs_in] = input;
0056   auto [recoparts_out, partassocs_out, partids_out] = output;
0057
0058   // local random generator
0059   auto seed = m_uid.getUniqueID(*headers, name());
0060   std::default_random_engine generator(seed);
0061   std::uniform_real_distribution<double> uniform;
0062
0063   for (const auto& recopart_without_pid : *recoparts_in) {
0064     auto recopart = recopart_without_pid.clone();
0065
0066     // Find MCParticle from associations and propagate the relevant ones further
0067     auto best_assoc = edm4eic::MCRecoParticleAssociation::makeEmpty();
0068     for (auto assoc_in : *partassocs_in) {
0069       if (assoc_in.getRec() == recopart_without_pid) {
0070         if ((not best_assoc.isAvailable()) || (best_assoc.getWeight() < assoc_in.getWeight())) {
0071           best_assoc = assoc_in;
0072         }
0073         auto assoc_out = assoc_in.clone();
0074         assoc_out.setRec(recopart);
0075         partassocs_out->push_back(assoc_out);
0076       }
0077     }
0078     if (not best_assoc.isAvailable()) {
0079       recoparts_out->push_back(recopart);
0080       continue;
0081     }
0082
0083     edm4hep::MCParticle mcpart = best_assoc.getSim();
0084
0085     int true_pdg    = mcpart.getPDG();
0086     int true_charge = mcpart.getCharge();
0087     int charge      = recopart.getCharge();
0088     double momentum = edm4hep::utils::magnitude(recopart.getMomentum());
0089
0090     double theta = edm4hep::utils::anglePolar(recopart.getMomentum()) / M_PI * 180.;
0091     double phi   = edm4hep::utils::angleAzimuthal(recopart.getMomentum()) / M_PI * 180.;
0092
0093     trace("lookup for true_pdg={}, true_charge={}, momentum={:.2f} GeV, polar={:.2f}, "
0094           "aziumthal={:.2f}",
0095           true_pdg, true_charge, momentum, theta, phi);
0096     const auto* entry = m_lut->Lookup(true_pdg, true_charge, momentum, theta, phi);
0097
0098     int identified_pdg = 0; // unknown
0099
0100     if ((entry != nullptr) && ((entry->prob_electron != 0.) || (entry->prob_pion != 0.) ||
0101                                (entry->prob_kaon != 0.) || (entry->prob_proton != 0.))) {
0102       double random_unit_interval = uniform(generator);
0103
0104       trace("entry with e:pi:K:P={}:{}:{}:{}", entry->prob_electron, entry->prob_pion,
0105             entry->prob_kaon, entry->prob_proton);
0106
0107       recopart.addToParticleIDs(
0108           partids_out->create(m_system,                                // std::int32_t type
0109                               std::copysign(11, -charge),              // std::int32_t PDG
0110                               0,                                       // std::int32_t algorithmType
0111                               static_cast<float>(entry->prob_electron) // float likelihood
0112                               ));
0113       recopart.addToParticleIDs(
0114           partids_out->create(m_system,                            // std::int32_t type
0115                               std::copysign(211, charge),          // std::int32_t PDG
0116                               0,                                   // std::int32_t algorithmType
0117                               static_cast<float>(entry->prob_pion) // float likelihood
0118                               ));
0119       recopart.addToParticleIDs(
0120           partids_out->create(m_system,                            // std::int32_t type
0121                               std::copysign(321, charge),          // std::int32_t PDG
0122                               0,                                   // std::int32_t algorithmType
0123                               static_cast<float>(entry->prob_kaon) // float likelihood
0124                               ));
0125       recopart.addToParticleIDs(
0126           partids_out->create(m_system,                              // std::int32_t type
0127                               std::copysign(2212, charge),           // std::int32_t PDG
0128                               0,                                     // std::int32_t algorithmType
0129                               static_cast<float>(entry->prob_proton) // float likelihood
0130                               ));
0131
0132       if (random_unit_interval < entry->prob_electron) {
0133         identified_pdg = 11; // electron
0134         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 4]);
0135       } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion)) {
0136         identified_pdg = 211; // pion
0137         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 3]);
0138       } else if (random_unit_interval <
0139                  (entry->prob_electron + entry->prob_pion + entry->prob_kaon)) {
0140         identified_pdg = 321; // kaon
0141         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 2]);
0142       } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion +
0143                                          entry->prob_kaon + entry->prob_proton)) {
0144         identified_pdg = 2212; // proton
0145         recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 1]);
0146       }
0147     }
0148
0149     if (identified_pdg != 0) {
0150       recopart.setPDG(std::copysign(identified_pdg, (identified_pdg == 11) ? -charge : charge));
0151       recopart.setMass(m_particleSvc.particle(identified_pdg).mass);
0152       recopart.setEnergy(std::hypot(momentum, m_particleSvc.particle(identified_pdg).mass));
0153     }
0154
0155     if (identified_pdg != 0) {
0156       trace("randomized PDG is {}", recopart.getPDG());
0157     }
0158
0159     recoparts_out->push_back(recopart);
0160   }
0161 }
0162
0163 } // namespace eicrecon