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