File indexing completed on 2024-09-27 07:02:59
0001
0002
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/podioVersion.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 auto recopart = recopart_without_pid.clone();
0048
0049
0050 #if podio_VERSION >= PODIO_VERSION(0, 17, 4)
0051 auto best_assoc = edm4eic::MCRecoParticleAssociation::makeEmpty();
0052 #else
0053 edm4eic::MCRecoParticleAssociationCollection _best_assoc_coll;
0054 edm4eic::MCRecoParticleAssociation best_assoc = _best_assoc_coll.create();
0055 best_assoc.unlink();
0056 #endif
0057 for (auto assoc_in : *partassocs_in) {
0058 if (assoc_in.getRec() == recopart_without_pid) {
0059 if ((not best_assoc.isAvailable()) || (best_assoc.getWeight() < assoc_in.getWeight())) {
0060 best_assoc = assoc_in;
0061 }
0062 auto assoc_out = assoc_in.clone();
0063 assoc_out.setRec(recopart);
0064 partassocs_out->push_back(assoc_out);
0065 }
0066 }
0067 if (not best_assoc.isAvailable()) {
0068 recoparts_out->push_back(recopart);
0069 continue;
0070 }
0071
0072 edm4hep::MCParticle mcpart = best_assoc.getSim();
0073
0074 int true_pdg = mcpart.getPDG();
0075 int true_charge = mcpart.getCharge();
0076 int charge = recopart.getCharge();
0077 double momentum = edm4hep::utils::magnitude(recopart.getMomentum());
0078
0079 double theta = edm4hep::utils::anglePolar(recopart.getMomentum()) / M_PI * 180.;
0080 double phi = edm4hep::utils::angleAzimuthal(recopart.getMomentum()) / M_PI * 180.;
0081
0082 trace("lookup for true_pdg={}, true_charge={}, momentum={:.2f} GeV, polar={:.2f}, aziumthal={:.2f}",
0083 true_pdg, true_charge, momentum, theta, phi);
0084 auto entry = m_lut->Lookup(true_pdg, true_charge, momentum, theta, phi);
0085
0086 int identified_pdg = 0;
0087
0088 if ((entry != nullptr) && ((entry->prob_electron != 0.) || (entry->prob_pion != 0.) || (entry->prob_kaon != 0.) || (entry->prob_electron != 0.))) {
0089 double random_unit_interval = m_dist(m_gen);
0090
0091 trace("entry with e:pi:K:P={}:{}:{}:{}", entry->prob_electron, entry->prob_pion, entry->prob_kaon, entry->prob_proton);
0092
0093 recopart.addToParticleIDs(partids_out->create(
0094 m_cfg.system,
0095 std::copysign(11, -charge),
0096 0,
0097 static_cast<float>(entry->prob_electron)
0098 ));
0099 recopart.addToParticleIDs(partids_out->create(
0100 m_cfg.system,
0101 std::copysign(211, charge),
0102 0,
0103 static_cast<float>(entry->prob_pion)
0104 ));
0105 recopart.addToParticleIDs(partids_out->create(
0106 m_cfg.system,
0107 std::copysign(321, charge),
0108 0,
0109 static_cast<float>(entry->prob_kaon)
0110 ));
0111 recopart.addToParticleIDs(partids_out->create(
0112 m_cfg.system,
0113 std::copysign(2212, charge),
0114 0,
0115 static_cast<float>(entry->prob_proton)
0116 ));
0117
0118 if (random_unit_interval < entry->prob_electron) {
0119 identified_pdg = 11;
0120 recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 4]);
0121 } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion)) {
0122 identified_pdg = 211;
0123 recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 3]);
0124 } else if (random_unit_interval <
0125 (entry->prob_electron + entry->prob_pion + entry->prob_kaon)) {
0126 identified_pdg = 321;
0127 recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 2]);
0128 } else if (random_unit_interval < (entry->prob_electron + entry->prob_pion +
0129 entry->prob_kaon + entry->prob_electron)) {
0130 identified_pdg = 2212;
0131 recopart.setParticleIDUsed((*partids_out)[partids_out->size() - 1]);
0132 }
0133 }
0134
0135 if (identified_pdg != 0) {
0136 recopart.setPDG(std::copysign(identified_pdg, (identified_pdg == 11) ? -charge : charge));
0137 recopart.setMass(m_particleSvc.particle(identified_pdg).mass);
0138 }
0139
0140 if (identified_pdg != 0) {
0141 trace("randomized PDG is {}", recopart.getPDG());
0142 }
0143
0144 recoparts_out->push_back(recopart);
0145 }
0146 }
0147
0148 }