Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-09-18 08:13:06

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008 
0009 #include "ActsExamples/Io/Root/RootTrackParameterWriter.hpp"
0010 
0011 #include "Acts/Definitions/Common.hpp"
0012 #include "Acts/Definitions/TrackParametrization.hpp"
0013 #include "Acts/Utilities/MultiIndex.hpp"
0014 #include "ActsExamples/EventData/AverageSimHits.hpp"
0015 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0016 #include "ActsExamples/Framework/WriterT.hpp"
0017 #include "ActsExamples/Utilities/Range.hpp"
0018 #include "ActsExamples/Validation/TrackClassification.hpp"
0019 #include "ActsFatras/EventData/Barcode.hpp"
0020 #include "ActsFatras/EventData/Hit.hpp"
0021 
0022 #include <cmath>
0023 #include <cstddef>
0024 #include <ios>
0025 #include <iostream>
0026 #include <numbers>
0027 #include <stdexcept>
0028 #include <utility>
0029 #include <vector>
0030 
0031 #include <TFile.h>
0032 #include <TTree.h>
0033 
0034 using Acts::VectorHelpers::eta;
0035 using Acts::VectorHelpers::perp;
0036 using Acts::VectorHelpers::phi;
0037 using Acts::VectorHelpers::theta;
0038 
0039 namespace ActsExamples {
0040 
0041 RootTrackParameterWriter::RootTrackParameterWriter(
0042     const RootTrackParameterWriter::Config& config, Acts::Logging::Level level)
0043     : WriterT<TrackParametersContainer>(config.inputTrackParameters,
0044                                         "RootTrackParameterWriter", level),
0045       m_cfg(config) {
0046   if (m_cfg.inputProtoTracks.empty()) {
0047     throw std::invalid_argument("Missing proto tracks input collection");
0048   }
0049   if (m_cfg.inputParticles.empty()) {
0050     throw std::invalid_argument("Missing particles input collection");
0051   }
0052   if (m_cfg.inputSimHits.empty()) {
0053     throw std::invalid_argument("Missing simulated hits input collection");
0054   }
0055   if (m_cfg.inputMeasurementParticlesMap.empty()) {
0056     throw std::invalid_argument("Missing hit-particles map input collection");
0057   }
0058   if (m_cfg.inputMeasurementSimHitsMap.empty()) {
0059     throw std::invalid_argument(
0060         "Missing hit-simulated-hits map input collection");
0061   }
0062   if (m_cfg.filePath.empty()) {
0063     throw std::invalid_argument("Missing output filename");
0064   }
0065   if (m_cfg.treeName.empty()) {
0066     throw std::invalid_argument("Missing tree name");
0067   }
0068 
0069   m_inputProtoTracks.initialize(m_cfg.inputProtoTracks);
0070   m_inputParticles.initialize(m_cfg.inputParticles);
0071   m_inputSimHits.initialize(m_cfg.inputSimHits);
0072   m_inputMeasurementParticlesMap.initialize(m_cfg.inputMeasurementParticlesMap);
0073   m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
0074 
0075   // Setup ROOT I/O
0076   if (m_outputFile == nullptr) {
0077     auto path = m_cfg.filePath;
0078     m_outputFile = TFile::Open(path.c_str(), m_cfg.fileMode.c_str());
0079     if (m_outputFile == nullptr) {
0080       throw std::ios_base::failure("Could not open '" + path + "'");
0081     }
0082   }
0083   m_outputFile->cd();
0084   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0085   if (m_outputTree == nullptr) {
0086     throw std::bad_alloc();
0087   }
0088 
0089   m_outputTree->Branch("event_nr", &m_eventNr);
0090 
0091   m_outputTree->Branch("volumeId", &m_volumeId);
0092   m_outputTree->Branch("layerId", &m_layerId);
0093   m_outputTree->Branch("surfaceId", &m_surfaceId);
0094 
0095   // The estimated track parameters
0096   m_outputTree->Branch("loc0", &m_loc0);
0097   m_outputTree->Branch("loc1", &m_loc1);
0098   m_outputTree->Branch("phi", &m_phi);
0099   m_outputTree->Branch("theta", &m_theta);
0100   m_outputTree->Branch("qop", &m_qop);
0101   m_outputTree->Branch("time", &m_time);
0102 
0103   // The estimated track parameters errors
0104   m_outputTree->Branch("err_loc0", &m_err_loc0);
0105   m_outputTree->Branch("err_loc1", &m_err_loc1);
0106   m_outputTree->Branch("err_phi", &m_err_phi);
0107   m_outputTree->Branch("err_theta", &m_err_theta);
0108   m_outputTree->Branch("err_qop", &m_err_qop);
0109   m_outputTree->Branch("err_time", &m_err_time);
0110 
0111   // Some derived quantities from the estimated track parameters
0112   m_outputTree->Branch("charge", &m_charge);
0113   m_outputTree->Branch("p", &m_p);
0114   m_outputTree->Branch("pt", &m_pt);
0115   m_outputTree->Branch("eta", &m_eta);
0116 
0117   // The truth meta
0118   m_outputTree->Branch("truthMatched", &m_t_matched);
0119   m_outputTree->Branch("particleId", &m_t_particleId);
0120   m_outputTree->Branch("nMajorityHits", &m_nMajorityHits);
0121 
0122   // The truth track parameters
0123   m_outputTree->Branch("t_loc0", &m_t_loc0);
0124   m_outputTree->Branch("t_loc1", &m_t_loc1);
0125   m_outputTree->Branch("t_phi", &m_t_phi);
0126   m_outputTree->Branch("t_theta", &m_t_theta);
0127   m_outputTree->Branch("t_qop", &m_t_qop);
0128   m_outputTree->Branch("t_time", &m_t_time);
0129   m_outputTree->Branch("t_charge", &m_t_charge);
0130   m_outputTree->Branch("t_p", &m_t_p);
0131   m_outputTree->Branch("t_pt", &m_t_pt);
0132   m_outputTree->Branch("t_eta", &m_t_eta);
0133 
0134   // The residuals
0135   m_outputTree->Branch("res_loc0", &m_res_loc0);
0136   m_outputTree->Branch("res_loc1", &m_res_loc1);
0137   m_outputTree->Branch("res_phi", &m_res_phi);
0138   m_outputTree->Branch("res_theta", &m_res_theta);
0139   m_outputTree->Branch("res_qop", &m_res_qop);
0140   m_outputTree->Branch("res_time", &m_res_time);
0141 
0142   // The pulls
0143   m_outputTree->Branch("pull_loc0", &m_pull_loc0);
0144   m_outputTree->Branch("pull_loc1", &m_pull_loc1);
0145   m_outputTree->Branch("pull_phi", &m_pull_phi);
0146   m_outputTree->Branch("pull_theta", &m_pull_theta);
0147   m_outputTree->Branch("pull_qop", &m_pull_qop);
0148   m_outputTree->Branch("pull_time", &m_pull_time);
0149 }
0150 
0151 RootTrackParameterWriter::~RootTrackParameterWriter() {
0152   if (m_outputFile != nullptr) {
0153     m_outputFile->Close();
0154   }
0155 }
0156 
0157 ProcessCode RootTrackParameterWriter::finalize() {
0158   m_outputFile->cd();
0159   m_outputTree->Write();
0160   m_outputFile->Close();
0161 
0162   ACTS_INFO("Wrote estimated parameters from seed to tree '"
0163             << m_cfg.treeName << "' in '" << m_cfg.filePath << "'");
0164 
0165   return ProcessCode::SUCCESS;
0166 }
0167 
0168 ProcessCode RootTrackParameterWriter::writeT(
0169     const AlgorithmContext& ctx, const TrackParametersContainer& trackParams) {
0170   // Read additional input collections
0171   const auto& protoTracks = m_inputProtoTracks(ctx);
0172   const auto& particles = m_inputParticles(ctx);
0173   const auto& simHits = m_inputSimHits(ctx);
0174   const auto& hitParticlesMap = m_inputMeasurementParticlesMap(ctx);
0175   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0176 
0177   std::vector<ParticleHitCount> particleHitCounts;
0178 
0179   // Exclusive access to the tree while writing
0180   std::lock_guard<std::mutex> lock(m_writeMutex);
0181 
0182   // Get the event number
0183   m_eventNr = ctx.eventNumber;
0184 
0185   ACTS_VERBOSE("Writing " << trackParams.size() << " track parameters");
0186 
0187   // Loop over the estimated track parameters
0188   for (std::size_t iparams = 0; iparams < trackParams.size(); ++iparams) {
0189     const auto& params = trackParams[iparams];
0190     // The reference surface of the parameters, i.e. also the reference surface
0191     // of the first space point
0192     const auto& surface = params.referenceSurface();
0193 
0194     m_volumeId = surface.geometryId().volume();
0195     m_layerId = surface.geometryId().layer();
0196     m_surfaceId = surface.geometryId().sensitive();
0197 
0198     m_loc0 = params.parameters()[Acts::eBoundLoc0];
0199     m_loc1 = params.parameters()[Acts::eBoundLoc1];
0200     m_phi = params.parameters()[Acts::eBoundPhi];
0201     m_theta = params.parameters()[Acts::eBoundTheta];
0202     m_qop = params.parameters()[Acts::eBoundQOverP];
0203     m_time = params.parameters()[Acts::eBoundTime];
0204 
0205     auto getError = [&params](std::size_t idx) -> float {
0206       if (!params.covariance().has_value()) {
0207         return NaNfloat;
0208       }
0209       const auto& cov = *params.covariance();
0210       if (cov(idx, idx) < 0) {
0211         return NaNfloat;
0212       }
0213       return std::sqrt(cov(idx, idx));
0214     };
0215 
0216     m_err_loc0 = getError(Acts::eBoundLoc0);
0217     m_err_loc1 = getError(Acts::eBoundLoc1);
0218     m_err_phi = getError(Acts::eBoundPhi);
0219     m_err_theta = getError(Acts::eBoundTheta);
0220     m_err_qop = getError(Acts::eBoundQOverP);
0221     m_err_time = getError(Acts::eBoundTime);
0222 
0223     m_charge = static_cast<int>(params.charge());
0224     m_p = params.absoluteMomentum();
0225     m_pt = params.transverseMomentum();
0226     m_eta = eta(params.momentum());
0227 
0228     // Get the proto track from which the track parameters are estimated
0229     const auto& ptrack = protoTracks[iparams];
0230     identifyContributingParticles(hitParticlesMap, ptrack, particleHitCounts);
0231 
0232     if (particleHitCounts.size() == 1) {
0233       m_t_matched = true;
0234       m_t_particleId = particleHitCounts.front().particleId.value();
0235       m_nMajorityHits = particleHitCounts.front().hitCount;
0236 
0237       // Get the index of the first space point
0238       const auto& hitIdx = ptrack.front();
0239       // Get the sim hits via the measurement to sim hits map
0240       auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0241       auto [truthLocal, truthPos4, truthUnitDir] =
0242           averageSimHits(ctx.geoContext, surface, simHits, indices, logger());
0243 
0244       // Get the truth track parameter at the first space point
0245       m_t_loc0 = truthLocal[Acts::ePos0];
0246       m_t_loc1 = truthLocal[Acts::ePos1];
0247       m_t_phi = phi(truthUnitDir);
0248       m_t_theta = theta(truthUnitDir);
0249       m_t_qop = NaNfloat;
0250       m_t_time = truthPos4[Acts::eTime];
0251 
0252       m_t_charge = 0;
0253       m_t_p = NaNfloat;
0254       m_t_pt = NaNfloat;
0255       m_t_eta = eta(truthUnitDir);
0256 
0257       // momentum averaging makes even less sense than averaging position and
0258       // direction. use the first momentum or set q/p to zero
0259       if (!indices.empty()) {
0260         // we assume that the indices are within valid ranges so we do not
0261         // need to check their validity again.
0262         const auto simHitIdx0 = indices.begin()->second;
0263         const auto& simHit0 = *simHits.nth(simHitIdx0);
0264         const auto p =
0265             simHit0.momentum4Before().template segment<3>(Acts::eMom0);
0266         const auto& particleId = simHit0.particleId();
0267         // The truth charge has to be retrieved from the sim particle
0268         if (auto ip = particles.find(particleId); ip != particles.end()) {
0269           const auto& particle = *ip;
0270           m_t_charge = static_cast<int>(particle.charge());
0271           m_t_qop = particle.hypothesis().qOverP(p.norm(), particle.charge());
0272           m_t_p = p.norm();
0273           m_t_pt = perp(p);
0274         } else {
0275           ACTS_WARNING("Truth particle with barcode " << particleId << "="
0276                                                       << particleId.value()
0277                                                       << " not found!");
0278         }
0279       }
0280 
0281       m_res_loc0 = m_loc0 - m_t_loc0;
0282       m_res_loc1 = m_loc1 - m_t_loc1;
0283       m_res_phi = Acts::detail::difference_periodic(
0284           m_phi, m_t_phi, static_cast<float>(2 * std::numbers::pi));
0285       m_res_theta = m_theta - m_t_theta;
0286       m_res_qop = m_qop - m_t_qop;
0287       m_res_time = m_time - m_t_time;
0288 
0289       auto getPull = [](float res, float err) -> float {
0290         if (std::isnan(err) || std::abs(err) < 1e-6) {
0291           return NaNfloat;
0292         }
0293         return res / err;
0294       };
0295 
0296       m_pull_loc0 = getPull(m_res_loc0, m_err_loc0);
0297       m_pull_loc1 = getPull(m_res_loc1, m_err_loc1);
0298       m_pull_phi = getPull(m_res_phi, m_err_phi);
0299       m_pull_theta = getPull(m_res_theta, m_err_theta);
0300       m_pull_qop = getPull(m_res_qop, m_err_qop);
0301       m_pull_time = getPull(m_res_time, m_err_time);
0302     } else {
0303       m_t_matched = false;
0304       m_t_particleId = 0;
0305       m_nMajorityHits = 0;
0306 
0307       m_t_loc0 = NaNfloat;
0308       m_t_loc1 = NaNfloat;
0309       m_t_phi = NaNfloat;
0310       m_t_theta = NaNfloat;
0311       m_t_qop = NaNfloat;
0312       m_t_time = NaNfloat;
0313 
0314       m_t_charge = 0;
0315       m_t_p = NaNfloat;
0316       m_t_pt = NaNfloat;
0317       m_t_eta = NaNfloat;
0318     }
0319 
0320     m_outputTree->Fill();
0321   }
0322 
0323   return ProcessCode::SUCCESS;
0324 }
0325 
0326 }  // namespace ActsExamples