algorithms/pid/MergeParticleID.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2023, Christopher Dilks
0003
0004 #include "MergeParticleID.h"
0005
0006 #include <algorithms/logger.h>
0007 #include <edm4eic/CherenkovParticleIDHypothesis.h>
0008 #include <edm4eic/TrackSegmentCollection.h>
0009 #include <edm4hep/Vector2f.h>
0010 #include <fmt/core.h>
0011 #include <podio/ObjectID.h>
0012 #include <podio/RelationRange.h>
0013 #include <cstddef>
0014 #include <gsl/pointers>
0015 #include <iterator>
0016 #include <stdexcept>
0017 #include <unordered_map>
0018 #include <utility>
0019
0020 #include "algorithms/pid/MergeParticleIDConfig.h"
0021 #include "algorithms/pid/Tools.h"
0022
0023 namespace eicrecon {
0024
0025 void MergeParticleID::init() { debug() << m_cfg << endmsg; }
0026
0027 void MergeParticleID::process(const MergeParticleID::Input& input,
0028                               const MergeParticleID::Output& output) const {
0029
0030   const auto [in_pid_collections_list] = input;
0031   auto [out_pids]                      = output;
0032
0033   /* match input `CherenkovParticleIDCollection` elements from each list of
0034    * collections in `in_pid_collections_list`
0035    * - the matching is done by the ID of the charged particle object
0036    * - each unique charged particle object ID should correspond to a unique
0037    *   charged particle object, so any "merging" of the charged particle
0038    *   objects needs to be done upstream of this algorithm
0039    *
0040    * PROCEDURE:
0041    * - build data structure `particle_pids`, which maps a charged particle object
0042    *   ID to a list of index pairs referring to the `CherenkovParticleID`s
0043    *   associated to it
0044    *   - the 1st index of the pair is that of the input std::vector `in_pid_collections_list`
0045    *   - the 2nd index is that of the corresponding CherenkovParticleIDCollection
0046    *
0047    * EXAMPLE for merging dRICH aerogel and gas PID objects:
0048    *
0049    * - INPUT std::vector: `in_pid_collections_list`
0050    *   0. CherenkovParticleIDCollection: aerogel pids
0051    *      0. aerogel PID for charged particle A
0052    *      1. aerogel PID for charged particle B
0053    *   1. CherenkovParticleIDCollection: gas pids
0054    *      0. gas PID for charged particle A
0055    *      1. gas PID for charged particle B
0056    *      2. gas PID for charged particle C  // outside aerogel acceptance, but within gas acceptance
0057    *
0058    * - OUTPUT std::unordered_map: `particle_pids` (integer => std::vector<pair of integers>)
0059    *   - ID of charged particle A => { (0, 0), (1, 0) }
0060    *   - ID of charged particle B => { (0, 1), (1, 1) }
0061    *   - ID of charged particle C => { (1, 2) }
0062    */
0063
0064   // fill `particle_pids`
0065   // -------------------------------------------------------------------------------
0066   std::unordered_map<decltype(podio::ObjectID::index),
0067                      std::vector<std::pair<std::size_t, std::size_t>>>
0068       particle_pids;
0069   trace("{:-<70}", "Build `particle_pids` indexing data structure ");
0070
0071   // loop over list of PID collections
0072   for (std::size_t idx_coll = 0; idx_coll < in_pid_collections_list.size(); idx_coll++) {
0073     const auto& in_pid_collection = in_pid_collections_list.at(idx_coll);
0074     trace("idx_col={}", idx_coll);
0075
0076     // loop over this PID collection
0077     for (std::size_t idx_pid = 0; idx_pid < in_pid_collection->size(); idx_pid++) {
0078       // make the index pair
0079       const auto& in_pid                         = in_pid_collection->at(idx_pid);
0080       const auto& charged_particle_track_segment = in_pid.getChargedParticle();
0081       if (!charged_particle_track_segment.isAvailable()) {
0082         error("PID object found with no charged particle");
0083         continue;
0084       }
0085       auto id_particle = charged_particle_track_segment.getObjectID().index;
0086       auto idx_paired  = std::make_pair(idx_coll, idx_pid);
0087       trace("  idx_pid={}  id_particle={}", idx_pid, id_particle);
0088
0089       // insert in `particle_pids`
0090       auto it = particle_pids.find(id_particle);
0091       if (it == particle_pids.end()) {
0092         particle_pids.insert({id_particle, {idx_paired}});
0093       } else {
0094         it->second.push_back(idx_paired);
0095       }
0096     }
0097   }
0098
0099   // trace logging
0100   if (level() <= algorithms::LogLevel::kTrace) {
0101     trace("{:-<70}", "Resulting `particle_pids` ");
0102     for (auto& [id_particle, idx_paired_list] : particle_pids) {
0103       trace("id_particle={}", id_particle);
0104       for (auto& [idx_coll, idx_pid] : idx_paired_list) {
0105         trace("  (idx_coll, idx_pid) = ({}, {})", idx_coll, idx_pid);
0106       }
0107     }
0108   }
0109
0110   // --------------------------------------------------------------------------------
0111
0112   // loop over charged particles, combine weights from the associated `CherenkovParticleID` objects
0113   // and create a merged output `CherenkovParticleID` object
0114   trace("{:-<70}", "Begin Merging PID Objects ");
0115   for (auto& [id_particle, idx_paired_list] : particle_pids) {
0116
0117     // trace logging
0118     trace("Charged Particle:");
0119     trace("  id = {}", id_particle);
0120     trace("  PID Hypotheses:");
0121
0122     // create mutable output `CherenkovParticleID` object `out_pid`
0123     auto out_pid                                            = out_pids->create();
0124     decltype(edm4eic::CherenkovParticleIDData::npe) out_npe = 0.0;
0125     decltype(edm4eic::CherenkovParticleIDData::refractiveIndex) out_refractiveIndex = 0.0;
0126     decltype(edm4eic::CherenkovParticleIDData::photonEnergy) out_photonEnergy       = 0.0;
0127
0128     // define `pdg_2_out_hyp`: map of PDG => merged output hypothesis
0129     std::unordered_map<decltype(edm4eic::CherenkovParticleIDHypothesis::PDG),
0130                        edm4eic::CherenkovParticleIDHypothesis>
0131         pdg_2_out_hyp;
0132
0133     // merge each input `CherenkovParticleID` objects associated with this charged particle
0134     for (auto& [idx_coll, idx_pid] : idx_paired_list) {
0135       const auto& in_pid = in_pid_collections_list.at(idx_coll)->at(idx_pid);
0136
0137       // logging
0138       trace("    Hypotheses for PID result (idx_coll, idx_pid) = ({}, {}):", idx_coll, idx_pid);
0139       trace(Tools::HypothesisTableHead(6));
0140
0141       // merge scalar members
0142       out_npe += in_pid.getNpe();                                           // sum
0143       out_refractiveIndex += in_pid.getNpe() * in_pid.getRefractiveIndex(); // NPE-weighted average
0144       out_photonEnergy += in_pid.getNpe() * in_pid.getPhotonEnergy();       // NPE-weighted average
0145
0146       // merge photon Cherenkov angles
0147       for (auto in_photon_vec : in_pid.getThetaPhiPhotons()) {
0148         out_pid.addToThetaPhiPhotons(in_photon_vec);
0149       }
0150
0151       // relate the charged particle
0152       if (!out_pid.getChargedParticle().isAvailable()) { // only needs to be done once
0153         out_pid.setChargedParticle(in_pid.getChargedParticle());
0154       }
0155
0156       // merge PDG hypotheses, combining their weights and other members
0157       for (auto in_hyp : in_pid.getHypotheses()) {
0158         trace(Tools::HypothesisTableLine(in_hyp, 6));
0159         auto out_hyp_it = pdg_2_out_hyp.find(in_hyp.PDG);
0160         if (out_hyp_it == pdg_2_out_hyp.end()) {
0161           edm4eic::CherenkovParticleIDHypothesis out_hyp;
0162           out_hyp.PDG    = in_hyp.PDG; // FIXME: no copy constructor?
0163           out_hyp.npe    = in_hyp.npe;
0164           out_hyp.weight = in_hyp.weight;
0165           pdg_2_out_hyp.insert({out_hyp.PDG, out_hyp});
0166         } else {
0167           auto& out_hyp = out_hyp_it->second;
0168           out_hyp.npe += in_hyp.npe;
0169           // combine hypotheses' weights
0170           switch (m_cfg.mergeMode) {
0171           case MergeParticleIDConfig::kAddWeights:
0172             out_hyp.weight += in_hyp.weight;
0173             break;
0174           case MergeParticleIDConfig::kMultiplyWeights:
0175             out_hyp.weight *= in_hyp.weight;
0176             break;
0177           default:
0178             throw std::runtime_error(
0179                 "unknown MergeParticleIDConfig::mergeMode setting; weights not combined");
0180           }
0181         }
0182       } // end `in_pid.getHypotheses()` loop, for this charged particle
0183
0184     } // end `in_pid` loop, for this charged particle
0185
0186     // finish computing averages of scalar members
0187     out_pid.setNpe(out_npe);
0188     if (out_npe > 0) {
0189       out_pid.setRefractiveIndex(out_refractiveIndex / out_npe);
0190       out_pid.setPhotonEnergy(out_photonEnergy / out_npe);
0191     }
0192
0193     // append hypotheses
0194     for (auto [pdg, out_hyp] : pdg_2_out_hyp) {
0195       out_pid.addToHypotheses(out_hyp);
0196     }
0197
0198     // logging: print merged hypothesis table
0199     trace("    => merged hypothesis weights:");
0200     trace(Tools::HypothesisTableHead(6));
0201     for (auto out_hyp : out_pid.getHypotheses()) {
0202       trace(Tools::HypothesisTableLine(out_hyp, 6));
0203     }
0204
0205   } // end `particle_pids` loop over charged particles
0206 }
0207
0208 } // namespace eicrecon