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