algorithms/pid/IrtCherenkovParticleID.cc

0001 // Copyright 2023, Christopher Dilks, adapted from Alexander Kiselev's Juggler implementation `IRTAlgorithm`
0002 // Subject to the terms in the LICENSE file found in the top-level directory.
0003
0004 #include "IrtCherenkovParticleID.h"
0005
0006 #include <IRT/ChargedParticle.h>
0007 #include <IRT/CherenkovPID.h>
0008 #include <IRT/OpticalPhoton.h>
0009 #include <IRT/RadiatorHistory.h>
0010 #include <IRT/SinglePDF.h>
0011 #include <TString.h>
0012 #include <TVector3.h>
0013 #include <algorithms/logger.h>
0014 #include <edm4eic/CherenkovParticleIDHypothesis.h>
0015 #include <edm4eic/TrackPoint.h>
0016 #include <edm4hep/MCParticleCollection.h>
0017 #include <edm4hep/SimTrackerHitCollection.h>
0018 #include <edm4hep/Vector2f.h>
0019 #include <edm4hep/Vector3d.h>
0020 #include <edm4hep/Vector3f.h>
0021 #include <fmt/format.h>
0022 #include <fmt/ranges.h>
0023 #include <podio/ObjectID.h>
0024 #include <podio/RelationRange.h>
0025 #include <podio/detail/LinkCollectionImpl.h>
0026 #include <podio/detail/LinkCollectionIterator.h>
0027 #include <algorithm>
0028 #include <cmath>
0029 #include <cstddef>
0030 #include <functional>
0031 #include <iterator>
0032 #include <memory>
0033 #include <set>
0034 #include <stdexcept>
0035 #include <tuple>
0036 #include <utility>
0037 #include <vector>
0038
0039 #include "algorithms/pid/IrtCherenkovParticleIDConfig.h"
0040 #include "algorithms/pid/Tools.h"
0041
0042 namespace eicrecon {
0043
0044 void IrtCherenkovParticleID::init(CherenkovDetectorCollection* irt_det_coll) {
0045   // members
0046   m_irt_det_coll = irt_det_coll;
0047
0048   // print the configuration parameters
0049   debug() << m_cfg << endmsg;
0050
0051   // inform the user if a cheat mode is enabled
0052   auto print_cheat = [this](auto name, bool val, auto desc) {
0053     if (val) {
0054       warning("CHEAT MODE '{}' ENABLED: {}", name, desc);
0055     }
0056   };
0057   print_cheat("cheatPhotonVertex", m_cfg.cheatPhotonVertex,
0058               "use MC photon vertex, wavelength, refractive index");
0059   print_cheat("cheatTrueRadiator", m_cfg.cheatTrueRadiator, "use MC truth to obtain true radiator");
0060
0061   // extract the relevant `CherenkovDetector`, set to `m_irt_det`
0062   const auto& detectors = m_irt_det_coll->GetDetectors();
0063   if (detectors.empty()) {
0064     throw std::runtime_error("No CherenkovDetectors found in input collection `irt_det_coll`");
0065   }
0066   if (detectors.size() > 1) {
0067     warning("IrtCherenkovParticleID currently only supports 1 CherenkovDetector at a time; "
0068             "taking the first");
0069   }
0070   auto this_detector = *detectors.begin();
0071   m_det_name         = this_detector.first;
0072   m_irt_det          = this_detector.second;
0073   debug("Initializing IrtCherenkovParticleID algorithm for CherenkovDetector '{}'", m_det_name);
0074
0075   // readout decoding
0076   m_cell_mask = m_irt_det->GetReadoutCellMask();
0077   debug("readout cellMask = {:#X}", m_cell_mask);
0078
0079   // rebin refractive index tables to have `m_cfg.numRIndexBins` bins
0080   trace("Rebinning refractive index tables to have {} bins", m_cfg.numRIndexBins);
0081   for (auto [rad_name, irt_rad] : m_irt_det->Radiators()) {
0082     // FIXME: m_cfg.numRIndexBins should be a service configurable
0083     std::lock_guard<std::mutex> lock(m_irt_det_mutex);
0084     auto ri_lookup_table_orig = irt_rad->m_ri_lookup_table;
0085     if (ri_lookup_table_orig.size() != m_cfg.numRIndexBins) {
0086       irt_rad->m_ri_lookup_table.clear();
0087       irt_rad->m_ri_lookup_table =
0088           Tools::ApplyFineBinning(ri_lookup_table_orig, m_cfg.numRIndexBins);
0089     }
0090   }
0091
0092   // build `m_pid_radiators`, the list of radiators to use for PID
0093   debug("Obtain List of Radiators:");
0094   for (auto [rad_name, irt_rad] : m_irt_det->Radiators()) {
0095     if (rad_name != "Filter") {
0096       m_pid_radiators.insert({std::string(rad_name), irt_rad});
0097       debug("- {}", rad_name.Data());
0098     }
0099   }
0100
0101   // check radiators' configuration, and pass it to `m_irt_det`'s radiators
0102   for (auto [rad_name, irt_rad] : m_pid_radiators) {
0103     std::lock_guard<std::mutex> lock(m_irt_det_mutex);
0104     // find `cfg_rad`, the associated `IrtCherenkovParticleIDConfig` radiator
0105     auto cfg_rad_it = m_cfg.radiators.find(rad_name);
0106     if (cfg_rad_it != m_cfg.radiators.end()) {
0107       auto cfg_rad = cfg_rad_it->second;
0108       // pass `cfg_rad` params to `irt_rad`, the IRT radiator
0109       irt_rad->m_ID                     = Tools::GetRadiatorID(std::string(rad_name));
0110       irt_rad->m_AverageRefractiveIndex = cfg_rad.referenceRIndex;
0111       irt_rad->SetReferenceRefractiveIndex(cfg_rad.referenceRIndex);
0112       if (cfg_rad.attenuation > 0) {
0113         irt_rad->SetReferenceAttenuationLength(cfg_rad.attenuation);
0114       }
0115       if (cfg_rad.smearing > 0) {
0116         if (cfg_rad.smearingMode == "uniform") {
0117           irt_rad->SetUniformSmearing(cfg_rad.smearing);
0118         } else if (cfg_rad.smearingMode == "gaussian") {
0119           irt_rad->SetGaussianSmearing(cfg_rad.smearing);
0120         } else {
0121           error("Unknown smearing mode '{}' for {} radiator", cfg_rad.smearingMode, rad_name);
0122         }
0123       }
0124     } else {
0125       error("Cannot find radiator '{}' in IrtCherenkovParticleIDConfig instance", rad_name);
0126     }
0127   }
0128
0129   // get PDG info for the particles we want to identify in PID
0130   debug("List of particles for PID:");
0131   for (auto pdg : m_cfg.pdgList) {
0132     auto mass = m_particleSvc.particle(pdg).mass;
0133     m_pdg_mass.insert({pdg, mass});
0134     debug("  {:>8}  M={} GeV", pdg, mass);
0135   }
0136 }
0137
0138 void IrtCherenkovParticleID::process(const IrtCherenkovParticleID::Input& input,
0139                                      const IrtCherenkovParticleID::Output& output) const {
0140   const auto [in_aerogel_tracks, in_gas_tracks, in_merged_tracks, in_raw_hits, in_hit_links,
0141               in_hit_assocs]                          = input;
0142   auto [out_aerogel_particleIDs, out_gas_particleIDs] = output;
0143
0144   // logging
0145   trace("{:=^70}", " call IrtCherenkovParticleID::AlgorithmProcess ");
0146   trace("number of raw sensor hits: {}", in_raw_hits->size());
0147   trace("number of raw sensor hits with associated photons: {}", in_hit_links->size());
0148
0149   std::map<std::string, const edm4eic::TrackSegmentCollection*> in_charged_particles{
0150       {"Aerogel", in_aerogel_tracks},
0151       {"Gas", in_gas_tracks},
0152       {"Merged", in_merged_tracks},
0153   };
0154
0155   // start output collections
0156   std::map<std::string, edm4eic::CherenkovParticleIDCollection*> out_cherenkov_pids{
0157       {"Aerogel", out_aerogel_particleIDs}, {"Gas", out_gas_particleIDs}};
0158
0159   // check `in_charged_particles`: each radiator should have the same number of TrackSegments
0160   std::unordered_map<std::size_t, std::size_t> in_charged_particle_size_distribution;
0161   for (const auto& [rad_name, in_charged_particle] : in_charged_particles) {
0162     ++in_charged_particle_size_distribution[in_charged_particle->size()];
0163   }
0164   if (in_charged_particle_size_distribution.size() != 1) {
0165     std::vector<std::size_t> in_charged_particle_sizes;
0166     std::ranges::transform(
0167         in_charged_particles, std::back_inserter(in_charged_particle_sizes),
0168         [](const auto& in_charged_particle) { return in_charged_particle.second->size(); });
0169     error("radiators have differing numbers of TrackSegments {}",
0170           fmt::join(in_charged_particle_sizes, ", "));
0171     return;
0172   }
0173
0174   // loop over charged particles ********************************************
0175   trace("{:#<70}", "### CHARGED PARTICLES ");
0176   std::size_t num_charged_particles = in_charged_particle_size_distribution.begin()->first;
0177   for (std::size_t i_charged_particle = 0; i_charged_particle < num_charged_particles;
0178        i_charged_particle++) {
0179     trace("{:-<70}", fmt::format("--- charged particle #{} ", i_charged_particle));
0180
0181     // start an `irt_particle`, for `IRT`
0182     auto irt_particle = std::make_unique<ChargedParticle>();
0183
0184     // loop over radiators
0185     // note: this must run exclusively since irt_rad points to shared IRT objects that are
0186     // owned by the RichGeo_service; it holds state (e.g. irt_rad->ResetLocation())
0187     std::lock_guard<std::mutex> lock(m_irt_det_mutex);
0188     for (auto [rad_name, irt_rad] : m_pid_radiators) {
0189
0190       // get the `charged_particle` for this radiator
0191       auto charged_particle_list_it = in_charged_particles.find(rad_name);
0192       if (charged_particle_list_it == in_charged_particles.end()) {
0193         error("Cannot find radiator '{}' in `in_charged_particles`", rad_name);
0194         continue;
0195       }
0196       const auto* charged_particle_list = charged_particle_list_it->second;
0197       auto charged_particle             = charged_particle_list->at(i_charged_particle);
0198
0199       // set number of bins for this radiator and charged particle
0200       if (charged_particle.points_size() == 0) {
0201         trace("No propagated track points in radiator '{}'", rad_name);
0202         continue;
0203       }
0204       irt_rad->SetTrajectoryBinCount(charged_particle.points_size() - 1);
0205
0206       // start a new IRT `RadiatorHistory`
0207       // - must be a raw pointer for `irt` compatibility
0208       // - it will be destroyed when `irt_particle` is destroyed
0209       auto* irt_rad_history = new RadiatorHistory();
0210       irt_particle->StartRadiatorHistory({irt_rad, irt_rad_history});
0211
0212       // loop over `TrackPoint`s of this `charged_particle`, adding each to the IRT radiator
0213       irt_rad->ResetLocations();
0214       trace("TrackPoints in '{}' radiator:", rad_name);
0215       for (const auto& point : charged_particle.getPoints()) {
0216         TVector3 position = Tools::PodioVector3_to_TVector3(point.position);
0217         TVector3 momentum = Tools::PodioVector3_to_TVector3(point.momentum);
0218         irt_rad->AddLocation(position, momentum);
0219         trace(Tools::PrintTVector3(" point: x", position));
0220         trace(Tools::PrintTVector3("        p", momentum));
0221       }
0222
0223       // loop over raw hits ***************************************************
0224       trace("{:#<70}", "### SENSOR HITS ");
0225       for (const auto& raw_hit : *in_raw_hits) {
0226
0227         // get MC photon(s), typically only used by cheat modes or trace logging
0228         // - loop over `in_hit_links`, searching for the matching raw-hit ↔ sim-hit link
0229         // - will not exist for noise hits
0230         edm4hep::MCParticle mc_photon;
0231         bool mc_photon_found = false;
0232         if (m_cfg.cheatPhotonVertex || m_cfg.cheatTrueRadiator) {
0233           for (const auto& hit_link : *in_hit_links) {
0234             if (!hit_link.getFrom().isAvailable() || !hit_link.getTo().isAvailable()) {
0235               continue;
0236             }
0237             if (hit_link.getFrom().id() == raw_hit.id()) {
0238               mc_photon       = hit_link.getTo().getParticle();
0239               mc_photon_found = true;
0240               if (mc_photon.getPDG() != -22) {
0241                 warning("non-opticalphoton hit: PDG = {}", mc_photon.getPDG());
0242               }
0243               break;
0244             }
0245           }
0246         }
0247
0248         // cheat mode, for testing only: use MC photon to get the actual radiator
0249         if (m_cfg.cheatTrueRadiator && mc_photon_found) {
0250           auto vtx     = Tools::PodioVector3_to_TVector3(mc_photon.getVertex());
0251           auto* mc_rad = m_irt_det->GuessRadiator(vtx, vtx); // assume IP is at (0,0,0)
0252           if (mc_rad != irt_rad) {
0253             continue; // skip this photon, if not from radiator `irt_rad`
0254           }
0255           trace(Tools::PrintTVector3(
0256               fmt::format("cheat: radiator '{}' determined from photon vertex", rad_name), vtx));
0257         }
0258
0259         // get sensor and pixel info
0260         // FIXME: signal and timing cuts (ADC, TDC, ToT, ...)
0261         auto cell_id       = raw_hit.getCellID();
0262         uint64_t sensor_id = cell_id & m_cell_mask;
0263         TVector3 pixel_pos = m_irt_det->m_ReadoutIDToPosition(cell_id);
0264
0265         // trace logging
0266         if (level() <= algorithms::LogLevel::kTrace) {
0267           trace("cell_id={:#X}  sensor_id={:#X}", cell_id, sensor_id);
0268           trace(Tools::PrintTVector3("pixel position", pixel_pos));
0269           if (mc_photon_found) {
0270             TVector3 mc_endpoint = Tools::PodioVector3_to_TVector3(mc_photon.getEndpoint());
0271             trace(Tools::PrintTVector3("photon endpoint", mc_endpoint));
0272             trace("{:>30} = {}", "dist( pixel,  photon )", (pixel_pos - mc_endpoint).Mag());
0273           } else {
0274             trace("  no MC photon found; probably a noise hit");
0275           }
0276         }
0277
0278         // start new IRT photon
0279         auto* irt_sensor = m_irt_det->m_PhotonDetectors[0]; // NOTE: assumes one sensor type
0280         auto* irt_photon =
0281             new OpticalPhoton(); // new raw pointer; it will also be destroyed when `irt_particle` is destroyed
0282         irt_photon->SetVolumeCopy(sensor_id);
0283         irt_photon->SetDetectionPosition(pixel_pos);
0284         irt_photon->SetPhotonDetector(irt_sensor);
0285         irt_photon->SetDetected(true);
0286
0287         // cheat mode: get photon vertex info from MC truth
0288         if ((m_cfg.cheatPhotonVertex || m_cfg.cheatTrueRadiator) && mc_photon_found) {
0289           irt_photon->SetVertexPosition(Tools::PodioVector3_to_TVector3(mc_photon.getVertex()));
0290           irt_photon->SetVertexMomentum(Tools::PodioVector3_to_TVector3(mc_photon.getMomentum()));
0291         }
0292
0293         // cheat mode: retrieve a refractive index estimate; it is not exactly the one, which
0294         // was used in GEANT, but should be very close
0295         if (m_cfg.cheatPhotonVertex) {
0296           double ri   = NAN;
0297           auto mom    = 1e9 * irt_photon->GetVertexMomentum().Mag();
0298           auto ri_set = Tools::GetFinelyBinnedTableEntry(irt_rad->m_ri_lookup_table, mom, &ri);
0299           if (ri_set) {
0300             irt_photon->SetVertexRefractiveIndex(ri);
0301             trace("{:>30} = {}", "refractive index", ri);
0302           } else {
0303             warning("Tools::GetFinelyBinnedTableEntry failed to lookup refractive index for "
0304                     "momentum {} eV",
0305                     mom);
0306           }
0307         }
0308
0309         // add each `irt_photon` to the radiator history
0310         // - unless cheating, we don't know which photon goes with which
0311         // radiator, thus we add them all to each radiator; the radiators'
0312         // photons are mixed in `ChargedParticle::PIDReconstruction`
0313         irt_rad_history->AddOpticalPhoton(irt_photon);
0314         /* FIXME: this considers ALL of the `irt_photon`s... we can limit this
0315          * once we add the ability to get a fiducial volume for each track, i.e.,
0316          * a region of sensors where we expect to see this `irt_particle`'s
0317          * Cherenkov photons; this should also combat sensor noise
0318          */
0319       } // end `in_raw_hits` loop
0320
0321     } // end radiator loop
0322
0323     // particle identification +++++++++++++++++++++++++++++++++++++++++++++++++++++
0324
0325     // define a mass hypothesis for each particle we want to check
0326     trace("{:+^70}", " PARTICLE IDENTIFICATION ");
0327     CherenkovPID irt_pid;
0328     std::unordered_map<int, MassHypothesis*> pdg_to_hyp; // `pdg` -> hypothesis
0329     for (auto [pdg, mass] : m_pdg_mass) {
0330       irt_pid.AddMassHypothesis(mass);
0331       pdg_to_hyp.insert({pdg, irt_pid.GetHypothesis(irt_pid.GetHypothesesCount() - 1)});
0332     }
0333
0334     // run IRT PID
0335     irt_particle->PIDReconstruction(irt_pid);
0336     trace("{:-^70}", " IRT RESULTS ");
0337
0338     // loop over radiators
0339     for (auto [rad_name, irt_rad] : m_pid_radiators) {
0340       trace("-> {} Radiator (ID={}):", rad_name, irt_rad->m_ID);
0341
0342       // Cherenkov angle (theta) estimate
0343       unsigned npe      = 0;
0344       double rindex_ave = 0.0;
0345       double energy_ave = 0.0;
0346       std::vector<std::pair<double, double>> phot_theta_phi;
0347
0348       // loop over this radiator's photons, and decide which to include in the theta estimate
0349       auto* irt_rad_history = irt_particle->FindRadiatorHistory(irt_rad);
0350       if (irt_rad_history == nullptr) {
0351         trace("  No radiator history; skip");
0352         continue;
0353       }
0354       trace("  Photoelectrons:");
0355       for (auto* irt_photon : irt_rad_history->Photons()) {
0356
0357         // check whether this photon was selected by at least one mass hypothesis
0358         bool photon_selected = false;
0359         for (auto irt_photon_sel : irt_photon->_m_Selected) {
0360           if (irt_photon_sel.second == irt_rad) {
0361             photon_selected = true;
0362             break;
0363           }
0364         }
0365         if (!photon_selected) {
0366           continue;
0367         }
0368
0369         // trace logging
0370         trace(
0371             Tools::PrintTVector3(fmt::format("- sensor_id={:#X}: hit", irt_photon->GetVolumeCopy()),
0372                                  irt_photon->GetDetectionPosition()));
0373         trace(Tools::PrintTVector3("photon vertex", irt_photon->GetVertexPosition()));
0374
0375         // get this photon's theta and phi estimates
0376         auto phot_theta = irt_photon->_m_PDF[irt_rad].GetAverage();
0377         auto phot_phi   = irt_photon->m_Phi[irt_rad];
0378
0379         // add to the total
0380         npe++;
0381         phot_theta_phi.emplace_back(phot_theta, phot_phi);
0382         if (m_cfg.cheatPhotonVertex) {
0383           rindex_ave += irt_photon->GetVertexRefractiveIndex();
0384           energy_ave += irt_photon->GetVertexMomentum().Mag();
0385         }
0386
0387       } // end loop over this radiator's photons
0388
0389       // compute averages
0390       if (npe > 0) {
0391         rindex_ave /= npe;
0392         energy_ave /= npe;
0393       }
0394
0395       // fill photon info
0396       auto out_cherenkov_pid = out_cherenkov_pids.at(rad_name)->create();
0397       out_cherenkov_pid.setNpe(static_cast<decltype(edm4eic::CherenkovParticleIDData::npe)>(npe));
0398       out_cherenkov_pid.setRefractiveIndex(
0399           static_cast<decltype(edm4eic::CherenkovParticleIDData::refractiveIndex)>(rindex_ave));
0400       out_cherenkov_pid.setPhotonEnergy(
0401           static_cast<decltype(edm4eic::CherenkovParticleIDData::photonEnergy)>(energy_ave));
0402       for (auto [phot_theta, phot_phi] : phot_theta_phi) {
0403         out_cherenkov_pid.addToThetaPhiPhotons(
0404             edm4hep::Vector2f{static_cast<float>(phot_theta), static_cast<float>(phot_phi)});
0405       }
0406
0407       // relate mass hypotheses
0408       for (auto [pdg, mass] : m_pdg_mass) {
0409
0410         // get hypothesis results
0411         auto* irt_hypothesis = pdg_to_hyp.at(pdg);
0412         auto hyp_weight      = irt_hypothesis->GetWeight(irt_rad);
0413         auto hyp_npe         = irt_hypothesis->GetNpe(irt_rad);
0414
0415         // Skip hypotheses with nan weight
0416         if (std::isnan(hyp_weight)) {
0417           continue;
0418         }
0419
0420         // fill `ParticleID` output collection
0421         edm4eic::CherenkovParticleIDHypothesis out_hypothesis;
0422         out_hypothesis.PDG =
0423             static_cast<decltype(edm4eic::CherenkovParticleIDHypothesis::PDG)>(pdg);
0424         out_hypothesis.weight =
0425             static_cast<decltype(edm4eic::CherenkovParticleIDHypothesis::weight)>(hyp_weight);
0426         out_hypothesis.npe =
0427             static_cast<decltype(edm4eic::CherenkovParticleIDHypothesis::npe)>(hyp_npe);
0428
0429         // relate
0430         out_cherenkov_pid.addToHypotheses(out_hypothesis);
0431
0432       } // end hypothesis loop
0433
0434       // logging: Cherenkov angle estimate
0435       auto PrintCherenkovEstimate = [this](edm4eic::CherenkovParticleID pid,
0436                                            bool printHypotheses = true, int indent = 2) {
0437         double thetaAve = 0;
0438         if (pid.getNpe() > 0) {
0439           for (const auto& [theta, phi] : pid.getThetaPhiPhotons()) {
0440             thetaAve += theta / pid.getNpe();
0441           }
0442         }
0443         trace("{:{}}Cherenkov Angle Estimate:", "", indent);
0444         trace("{:{}}  {:>16}:  {:>10}", "", indent, "NPE", pid.getNpe());
0445         trace("{:{}}  {:>16}:  {:>10.8} mrad", "", indent, "<theta>",
0446               thetaAve * 1e3); // [rad] -> [mrad]
0447         trace("{:{}}  {:>16}:  {:>10.8}", "", indent, "<rindex>", pid.getRefractiveIndex());
0448         trace("{:{}}  {:>16}:  {:>10.8} eV", "", indent, "<energy>",
0449               pid.getPhotonEnergy() * 1e9); // [GeV] -> [eV]
0450         if (printHypotheses) {
0451           trace("{:{}}Mass Hypotheses:", "", indent);
0452           trace("{}", Tools::HypothesisTableHead(indent + 2));
0453           for (const auto& hyp : pid.getHypotheses()) {
0454             trace("{}", Tools::HypothesisTableLine(hyp, indent + 2));
0455           }
0456         }
0457       };
0458       PrintCherenkovEstimate(out_cherenkov_pid);
0459
0460       // relate charged particle projection
0461       auto charged_particle_list_it = in_charged_particles.find("Merged");
0462       if (charged_particle_list_it != in_charged_particles.end()) {
0463         const auto* charged_particle_list = charged_particle_list_it->second;
0464         auto charged_particle             = charged_particle_list->at(i_charged_particle);
0465         out_cherenkov_pid.setChargedParticle(charged_particle);
0466       } else {
0467         error("Cannot find radiator 'Merged' in `in_charged_particles`");
0468       }
0469
0470       // keep the legacy association relation while consuming links for MC truth lookup
0471       for (const auto& hit_assoc : *in_hit_assocs) {
0472         out_cherenkov_pid.addToRawHitAssociations(hit_assoc);
0473       }
0474
0475     } // end radiator loop
0476
0477     /* NOTE: `unique_ptr irt_particle` goes out of scope and will now be destroyed, and along with it:
0478      * - raw pointer `irt_rad_history` for each radiator
0479      * - all `irt_photon` raw pointers
0480      */
0481
0482   } // end `in_charged_particles` loop
0483 }
0484
0485 } // namespace eicrecon