truth/src/ParticlesWithTruthPID.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Wouter Deconinck
0003
0004 #include <algorithms/truth/ParticlesWithTruthPID.h>
0005
0006 #include <algorithm>
0007 #include <cmath>
0008 #include <fmt/format.h>
0009
0010 #include <edm4hep/utils/vector_utils.h>
0011
0012 namespace algorithms::truth {
0013
0014 void ParticlesWithTruthPID::init() {
0015   ; // do nothing
0016 }
0017 void ParticlesWithTruthPID::process(const ParticlesWithTruthPID::Input& input,
0018                                     const ParticlesWithTruthPID::Output& output) const {
0019   const auto [mc_ptr, tracks_ptr] = input;
0020   auto [part_ptr, assoc_ptr]      = output;
0021
0022   const auto& mc     = *mc_ptr;
0023   const auto& tracks = *tracks_ptr;
0024   auto& part         = *part_ptr;
0025   auto& assoc        = *assoc_ptr;
0026
0027   const double sinPhiOver2Tolerance = sin(0.5 * m_phiTolerance);
0028   std::vector<bool> consumed(mc.size(), false);
0029   for (const auto& trk : tracks) {
0030     const auto mom =
0031         edm4hep::utils::sphericalToVector(1.0 / std::abs(trk.getQOverP()), trk.getTheta(), trk.getPhi());
0032     const auto charge_rec = std::copysign(1., trk.getQOverP());
0033     // utility variables for matching
0034     int best_match    = -1;
0035     double best_delta = std::numeric_limits<double>::max();
0036     for (size_t ip = 0; ip < mc.size(); ++ip) {
0037       const auto& mcpart = mc[ip];
0038       if (consumed[ip] || mcpart.getGeneratorStatus() > 1 || mcpart.getCharge() == 0 ||
0039           mcpart.getCharge() * charge_rec < 0) {
0040         if (aboveDebugThreshold()) {
0041           debug() << "ignoring non-primary/neutral/opposite charge particle" << endmsg;
0042         }
0043         continue;
0044       }
0045       const auto& p       = mcpart.getMomentum();
0046       const auto p_mag    = std::hypot(p.x, p.y, p.z);
0047       const auto p_phi    = std::atan2(p.y, p.x);
0048       const auto p_eta    = std::atanh(p.z / p_mag);
0049       const double dp_rel = std::abs((edm4hep::utils::magnitude(mom) - p_mag) / p_mag);
0050       // check the tolerance for sin(dphi/2) to avoid the hemisphere problem and allow
0051       // for phi rollovers
0052       const double dsphi = std::abs(sin(0.5 * (edm4hep::utils::angleAzimuthal(mom) - p_phi)));
0053       const double deta  = std::abs((edm4hep::utils::eta(mom) - p_eta));
0054
0055       if (dp_rel < m_pRelativeTolerance && deta < m_etaTolerance && dsphi < sinPhiOver2Tolerance) {
0056         const double delta = std::hypot(dp_rel / m_pRelativeTolerance, deta / m_etaTolerance,
0057                                         dsphi / sinPhiOver2Tolerance);
0058         if (delta < best_delta) {
0059           best_match = ip;
0060           best_delta = delta;
0061         }
0062       }
0063     }
0064     auto rec_part       = part.create();
0065     int32_t best_pid    = 0;
0066     auto referencePoint = rec_part.referencePoint();
0067     // float time          = 0;
0068     float mass = 0;
0069     if (best_match >= 0) {
0070       consumed[best_match] = true;
0071       const auto& mcpart   = mc[best_match];
0072       best_pid             = mcpart.getPDG();
0073       referencePoint       = {
0074           static_cast<float>(mcpart.getVertex().x), static_cast<float>(mcpart.getVertex().y),
0075           static_cast<float>(
0076               mcpart.getVertex().z)}; // @TODO: not sure if vertex/reference poitn makes sense here
0077       // time                 = mcpart.getTime();
0078       mass = mcpart.getMass();
0079     }
0080     rec_part.setType(static_cast<int16_t>(best_match >= 0 ? 0 : -1)); // @TODO: determine type codes
0081     rec_part.setEnergy(std::hypot(edm4hep::utils::magnitude(mom), mass));
0082     rec_part.setMomentum(mom);
0083     rec_part.setReferencePoint(referencePoint);
0084     rec_part.setCharge(charge_rec);
0085     rec_part.setMass(mass);
0086     rec_part.setGoodnessOfPID(1); // perfect PID
0087     rec_part.setPDG(best_pid);
0088     // rec_part.covMatrix()  // @TODO: covariance matrix on 4-momentum
0089     // Also write MC <--> truth particle association if match was found
0090     if (best_match >= 0) {
0091       auto rec_assoc = assoc.create();
0092       rec_assoc.setRecID(rec_part.getObjectID().index);
0093       rec_assoc.setSimID(mc[best_match].getObjectID().index);
0094       rec_assoc.setWeight(1);
0095       rec_assoc.setRec(rec_part);
0096       // rec_assoc.setSim(mc[best_match]);
0097     }
0098     if (aboveDebugThreshold()) {
0099       if (best_match > 0) {
0100         const auto& mcpart = mc[best_match];
0101         debug() << fmt::format("Matched track with MC particle {}\n", best_match) << endmsg;
0102         debug() << fmt::format("  - Track: (mom: {}, theta: {}, phi: {}, charge: {})",
0103                                edm4hep::utils::magnitude(mom), edm4hep::utils::anglePolar(mom),
0104                                edm4hep::utils::angleAzimuthal(mom), charge_rec)
0105                 << endmsg;
0106         const auto& p      = mcpart.getMomentum();
0107         const auto p_mag   = edm4hep::utils::magnitude(p);
0108         const auto p_phi   = edm4hep::utils::angleAzimuthal(p);
0109         const auto p_theta = edm4hep::utils::anglePolar(p);
0110         debug() << fmt::format(
0111                        "  - MC particle: (mom: {}, theta: {}, phi: {}, charge: {}, type: {}", p_mag,
0112                        p_theta, p_phi, mcpart.getCharge(), mcpart.getPDG())
0113                 << endmsg;
0114       } else {
0115         debug() << fmt::format("Did not find a good match for track \n") << endmsg;
0116         debug() << fmt::format("  - Track: (mom: {}, theta: {}, phi: {}, charge: {})",
0117                                edm4hep::utils::magnitude(mom), edm4hep::utils::anglePolar(mom),
0118                                edm4hep::utils::angleAzimuthal(mom), charge_rec)
0119                 << endmsg;
0120       }
0121     }
0122   }
0123 }
0124
0125 } // namespace algorithms::truth
0126