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