Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-09-27 07:02:59

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2024, Sylvester Joosten, Wouter Deconinck, Dmitry Romanov, Christopher Dilks, Dmitry Kalinkin
0003 
0004 #include "MatchToRICHPID.h"
0005 
0006 #include <edm4eic/TrackPoint.h>
0007 #include <edm4eic/TrackSegmentCollection.h>
0008 #include <edm4hep/Vector3f.h>
0009 #include <edm4hep/utils/vector_utils.h>
0010 #include <fmt/core.h>
0011 #include <podio/ObjectID.h>
0012 #include <podio/RelationRange.h>
0013 #include <algorithm>
0014 #include <cmath>
0015 #include <cstddef>
0016 #include <gsl/pointers>
0017 #include <map>
0018 #include <vector>
0019 
0020 #include "algorithms/pid/ConvertParticleID.h"
0021 #include "algorithms/pid/MatchToRICHPIDConfig.h"
0022 
0023 
0024 
0025 namespace eicrecon {
0026 
0027     void MatchToRICHPID::init() {}
0028 
0029     void MatchToRICHPID::process(
0030       const MatchToRICHPID::Input& input, const MatchToRICHPID::Output& output
0031     ) const {
0032         const auto [parts_in, assocs_in, drich_cherenkov_pid] = input;
0033         auto [parts_out, assocs_out, pids]                     = output;
0034 
0035         for (auto part_in : *parts_in) {
0036             auto part_out = part_in.clone();
0037 
0038             // link Cherenkov PID objects
0039             auto success = linkCherenkovPID(part_out, *drich_cherenkov_pid, *pids);
0040             if (success)
0041                 trace("Previous PDG vs. CherenkovPID PDG: {:>10} vs. {:<10}",
0042                         part_in.getPDG(),
0043                         part_out.getParticleIDUsed().isAvailable() ? part_out.getParticleIDUsed().getPDG() : 0
0044                         );
0045 
0046             for (auto assoc_in : *assocs_in) {
0047               if (assoc_in.getRec() == part_in) {
0048                 auto assoc_out = assoc_in.clone();
0049                 assoc_out.setRec(part_out);
0050                 assocs_out->push_back(assoc_out);
0051               }
0052             }
0053 
0054             parts_out->push_back(part_out);
0055         }
0056     }
0057 
0058 
0059     /* link PID objects to input particle
0060      * - finds `CherenkovParticleID` object in `in_pids` associated to particle `in_part`
0061      *   by proximity matching to the associated track
0062      * - converts this `CherenkovParticleID` object's PID hypotheses to `ParticleID` objects,
0063      *   relates them to `in_part`, and adds them to the collection `out_pids` for persistency
0064      * - returns `true` iff PID objects were found and linked
0065      */
0066     bool MatchToRICHPID::linkCherenkovPID(
0067             edm4eic::MutableReconstructedParticle& in_part,
0068             const edm4eic::CherenkovParticleIDCollection& in_pids,
0069             edm4hep::ParticleIDCollection& out_pids
0070             ) const
0071     {
0072 
0073         // skip this particle, if neutral
0074         if (std::abs(in_part.getCharge()) < 0.001)
0075             return false;
0076 
0077         // structure to store list of candidate matches
0078         struct ProxMatch {
0079             double      match_dist;
0080             std::size_t pid_idx;
0081         };
0082         std::vector<ProxMatch> prox_match_list;
0083 
0084         // get input reconstructed particle momentum angles
0085         auto in_part_p   = in_part.getMomentum();
0086         auto in_part_eta = edm4hep::utils::eta(in_part_p);
0087         auto in_part_phi = edm4hep::utils::angleAzimuthal(in_part_p);
0088         trace("Input particle: (eta,phi) = ( {:>5.4}, {:>5.4} deg )",
0089                 in_part_eta,
0090                 in_part_phi * 180.0 / M_PI
0091                 );
0092 
0093         // loop over input CherenkovParticleID objects
0094         for (std::size_t in_pid_idx = 0; in_pid_idx < in_pids.size(); in_pid_idx++) {
0095             auto in_pid = in_pids.at(in_pid_idx);
0096 
0097             // get charged particle track associated to this CherenkovParticleID object
0098             auto in_track = in_pid.getChargedParticle();
0099             if (!in_track.isAvailable()) {
0100                 error("found CherenkovParticleID object with no chargedParticle");
0101                 return false;
0102             }
0103             if (in_track.points_size() == 0) {
0104                 error("found chargedParticle for CherenkovParticleID, but it has no TrackPoints");
0105                 return false;
0106             }
0107 
0108             // get averge momentum direction of the track's TrackPoints
0109             decltype(edm4eic::TrackPoint::momentum) in_track_p{0.0, 0.0, 0.0};
0110             for (const auto& in_track_point : in_track.getPoints())
0111                 in_track_p = in_track_p + ( in_track_point.momentum / in_track.points_size() );
0112             auto in_track_eta = edm4hep::utils::eta(in_track_p);
0113             auto in_track_phi = edm4hep::utils::angleAzimuthal(in_track_p);
0114 
0115             // calculate dist(eta,phi)
0116             auto match_dist = std::hypot(
0117                     in_part_eta - in_track_eta,
0118                     in_part_phi - in_track_phi
0119                     );
0120 
0121             // check if the match is close enough: within user-specified tolerances
0122             auto match_is_close =
0123                 std::abs(in_part_eta - in_track_eta) < m_cfg.etaTolerance &&
0124                 std::abs(in_part_phi - in_track_phi) < m_cfg.phiTolerance;
0125             if (match_is_close)
0126                 prox_match_list.push_back(ProxMatch{match_dist, in_pid_idx});
0127 
0128             // logging
0129             trace("  - (eta,phi) = ( {:>5.4}, {:>5.4} deg ),  match_dist = {:<5.4}{}",
0130                     in_track_eta,
0131                     in_track_phi * 180.0 / M_PI,
0132                     match_dist,
0133                     match_is_close ? " => CLOSE!" : ""
0134                     );
0135 
0136         } // end loop over input CherenkovParticleID objects
0137 
0138         // check if at least one match was found
0139         if (prox_match_list.size() == 0) {
0140             trace("  => no matching CherenkovParticleID found for this particle");
0141             return false;
0142         }
0143 
0144         // choose the closest matching CherenkovParticleID object corresponding to this input reconstructed particle
0145         auto closest_prox_match = *std::min_element(
0146                 prox_match_list.begin(),
0147                 prox_match_list.end(),
0148                 [] (ProxMatch a, ProxMatch b) { return a.match_dist < b.match_dist; }
0149                 );
0150         auto in_pid_matched = in_pids.at(closest_prox_match.pid_idx);
0151         trace("  => best match: match_dist = {:<5.4} at idx = {}",
0152                 closest_prox_match.match_dist,
0153                 closest_prox_match.pid_idx
0154                 );
0155 
0156         // convert `CherenkovParticleID` object's hypotheses => set of `ParticleID` objects
0157         auto out_pid_index_map = ConvertParticleID::ConvertToParticleIDs(in_pid_matched, out_pids, true);
0158         if (out_pid_index_map.size() == 0) {
0159             error("found CherenkovParticleID object with no hypotheses");
0160             return false;
0161         }
0162 
0163         // relate matched ParticleID objects to output particle
0164         for (const auto& [out_pids_index, out_pids_id] : out_pid_index_map) {
0165             const auto& out_pid = out_pids->at(out_pids_index);
0166             if (out_pid.getObjectID().index != out_pids_id) { // sanity check
0167                 error("indexing error in `edm4eic::ParticleID` collection");
0168                 return false;
0169             }
0170             in_part.addToParticleIDs(out_pid);
0171         }
0172         in_part.setParticleIDUsed(in_part.getParticleIDs().at(0)); // highest likelihood is the first
0173         in_part.setGoodnessOfPID(1); // FIXME: not used yet, aside from 0=noPID vs 1=hasPID
0174 
0175         // trace logging
0176         trace("    {:.^50}"," PID results ");
0177         trace("      Hypotheses (sorted):");
0178         for (auto hyp : in_part.getParticleIDs()) {
0179             float npe = hyp.parameters_size() > 0 ? hyp.getParameters(0) : -1; // assume NPE is the first parameter
0180             trace("{:{}}{:>6}  {:>10.8}  {:>10.8}", "", 8, hyp.getPDG(), hyp.getLikelihood(), npe);
0181         }
0182         trace("    {:'^50}","");
0183 
0184         return true;
0185     }
0186 }