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