Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-18 09:11:58

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("particleId", &m_t_particleId);
0124   m_outputTree->Branch("t_loc0", &m_t_loc0);
0125   m_outputTree->Branch("t_loc1", &m_t_loc1);
0126   m_outputTree->Branch("t_phi", &m_t_phi);
0127   m_outputTree->Branch("t_theta", &m_t_theta);
0128   m_outputTree->Branch("t_qop", &m_t_qop);
0129   m_outputTree->Branch("t_time", &m_t_time);
0130   m_outputTree->Branch("t_charge", &m_t_charge);
0131   m_outputTree->Branch("t_p", &m_t_p);
0132   m_outputTree->Branch("t_pt", &m_t_pt);
0133   m_outputTree->Branch("t_eta", &m_t_eta);
0134 
0135   // The residuals
0136   m_outputTree->Branch("res_loc0", &m_res_loc0);
0137   m_outputTree->Branch("res_loc1", &m_res_loc1);
0138   m_outputTree->Branch("res_phi", &m_res_phi);
0139   m_outputTree->Branch("res_theta", &m_res_theta);
0140   m_outputTree->Branch("res_qop", &m_res_qop);
0141   m_outputTree->Branch("res_time", &m_res_time);
0142 
0143   // The pulls
0144   m_outputTree->Branch("pull_loc0", &m_pull_loc0);
0145   m_outputTree->Branch("pull_loc1", &m_pull_loc1);
0146   m_outputTree->Branch("pull_phi", &m_pull_phi);
0147   m_outputTree->Branch("pull_theta", &m_pull_theta);
0148   m_outputTree->Branch("pull_qop", &m_pull_qop);
0149   m_outputTree->Branch("pull_time", &m_pull_time);
0150 }
0151 
0152 RootTrackParameterWriter::~RootTrackParameterWriter() {
0153   if (m_outputFile != nullptr) {
0154     m_outputFile->Close();
0155   }
0156 }
0157 
0158 ProcessCode RootTrackParameterWriter::finalize() {
0159   m_outputFile->cd();
0160   m_outputTree->Write();
0161   m_outputFile->Close();
0162 
0163   ACTS_INFO("Wrote estimated parameters from seed to tree '"
0164             << m_cfg.treeName << "' in '" << m_cfg.filePath << "'");
0165 
0166   return ProcessCode::SUCCESS;
0167 }
0168 
0169 ProcessCode RootTrackParameterWriter::writeT(
0170     const AlgorithmContext& ctx, const TrackParametersContainer& trackParams) {
0171   // Read additional input collections
0172   const auto& protoTracks = m_inputProtoTracks(ctx);
0173   const auto& particles = m_inputParticles(ctx);
0174   const auto& simHits = m_inputSimHits(ctx);
0175   const auto& hitParticlesMap = m_inputMeasurementParticlesMap(ctx);
0176   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0177 
0178   std::vector<ParticleHitCount> particleHitCounts;
0179 
0180   // Exclusive access to the tree while writing
0181   std::lock_guard<std::mutex> lock(m_writeMutex);
0182 
0183   // Get the event number
0184   m_eventNr = ctx.eventNumber;
0185 
0186   ACTS_VERBOSE("Writing " << trackParams.size() << " track parameters");
0187 
0188   // Loop over the estimated track parameters
0189   for (std::size_t iparams = 0; iparams < trackParams.size(); ++iparams) {
0190     const auto& params = trackParams[iparams];
0191     // The reference surface of the parameters, i.e. also the reference surface
0192     // of the first space point
0193     const auto& surface = params.referenceSurface();
0194 
0195     m_volumeId = surface.geometryId().volume();
0196     m_layerId = surface.geometryId().layer();
0197     m_surfaceId = surface.geometryId().sensitive();
0198 
0199     m_loc0 = params.parameters()[Acts::eBoundLoc0];
0200     m_loc1 = params.parameters()[Acts::eBoundLoc1];
0201     m_phi = params.parameters()[Acts::eBoundPhi];
0202     m_theta = params.parameters()[Acts::eBoundTheta];
0203     m_qop = params.parameters()[Acts::eBoundQOverP];
0204     m_time = params.parameters()[Acts::eBoundTime];
0205 
0206     auto getError = [&params](std::size_t idx) -> float {
0207       if (!params.covariance().has_value()) {
0208         return NaNfloat;
0209       }
0210       const auto& cov = *params.covariance();
0211       if (cov(idx, idx) < 0) {
0212         return NaNfloat;
0213       }
0214       return std::sqrt(cov(idx, idx));
0215     };
0216 
0217     m_err_loc0 = getError(Acts::eBoundLoc0);
0218     m_err_loc1 = getError(Acts::eBoundLoc1);
0219     m_err_phi = getError(Acts::eBoundPhi);
0220     m_err_theta = getError(Acts::eBoundTheta);
0221     m_err_qop = getError(Acts::eBoundQOverP);
0222     m_err_time = getError(Acts::eBoundTime);
0223 
0224     m_charge = static_cast<int>(params.charge());
0225     m_p = params.absoluteMomentum();
0226     m_pt = params.transverseMomentum();
0227     m_eta = eta(params.momentum());
0228 
0229     // Get the proto track from which the track parameters are estimated
0230     const auto& ptrack = protoTracks[iparams];
0231     identifyContributingParticles(hitParticlesMap, ptrack, particleHitCounts);
0232 
0233     if (particleHitCounts.size() == 1) {
0234       m_t_matched = true;
0235       m_t_particleId = particleHitCounts.front().particleId.value();
0236       m_nMajorityHits = particleHitCounts.front().hitCount;
0237 
0238       // Get the index of the first space point
0239       const auto& hitIdx = ptrack.front();
0240       // Get the sim hits via the measurement to sim hits map
0241       auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0242       auto [truthLocal, truthPos4, truthUnitDir] =
0243           averageSimHits(ctx.geoContext, surface, simHits, indices, logger());
0244 
0245       // Get the truth track parameter at the first space point
0246       m_t_loc0 = truthLocal[Acts::ePos0];
0247       m_t_loc1 = truthLocal[Acts::ePos1];
0248       m_t_phi = phi(truthUnitDir);
0249       m_t_theta = theta(truthUnitDir);
0250       m_t_qop = NaNfloat;
0251       m_t_time = truthPos4[Acts::eTime];
0252 
0253       m_t_charge = 0;
0254       m_t_p = NaNfloat;
0255       m_t_pt = NaNfloat;
0256       m_t_eta = eta(truthUnitDir);
0257 
0258       // momentum averaging makes even less sense than averaging position and
0259       // direction. use the first momentum or set q/p to zero
0260       if (!indices.empty()) {
0261         // we assume that the indices are within valid ranges so we do not
0262         // need to check their validity again.
0263         const auto simHitIdx0 = indices.begin()->second;
0264         const auto& simHit0 = *simHits.nth(simHitIdx0);
0265         const auto p =
0266             simHit0.momentum4Before().template segment<3>(Acts::eMom0);
0267         const auto& particleId = simHit0.particleId();
0268         // The truth charge has to be retrieved from the sim particle
0269         if (auto ip = particles.find(particleId); ip != particles.end()) {
0270           const auto& particle = *ip;
0271           m_t_charge = static_cast<int>(particle.charge());
0272           m_t_qop = particle.hypothesis().qOverP(p.norm(), particle.charge());
0273           m_t_p = p.norm();
0274           m_t_pt = perp(p);
0275         } else {
0276           ACTS_WARNING("Truth particle with barcode " << particleId << "="
0277                                                       << particleId.value()
0278                                                       << " not found!");
0279         }
0280       }
0281 
0282       m_res_loc0 = m_loc0 - m_t_loc0;
0283       m_res_loc1 = m_loc1 - m_t_loc1;
0284       m_res_phi = Acts::detail::difference_periodic(
0285           m_phi, m_t_phi, static_cast<float>(2 * std::numbers::pi));
0286       m_res_theta = m_theta - m_t_theta;
0287       m_res_qop = m_qop - m_t_qop;
0288       m_res_time = m_time - m_t_time;
0289 
0290       auto getPull = [](float res, float err) -> float {
0291         if (std::isnan(err) || std::abs(err) < 1e-6) {
0292           return NaNfloat;
0293         }
0294         return res / err;
0295       };
0296 
0297       m_pull_loc0 = getPull(m_res_loc0, m_err_loc0);
0298       m_pull_loc1 = getPull(m_res_loc1, m_err_loc1);
0299       m_pull_phi = getPull(m_res_phi, m_err_phi);
0300       m_pull_theta = getPull(m_res_theta, m_err_theta);
0301       m_pull_qop = getPull(m_res_qop, m_err_qop);
0302       m_pull_time = getPull(m_res_time, m_err_time);
0303     } else {
0304       m_t_matched = false;
0305       m_t_particleId = 0;
0306       m_nMajorityHits = 0;
0307 
0308       m_t_loc0 = NaNfloat;
0309       m_t_loc1 = NaNfloat;
0310       m_t_phi = NaNfloat;
0311       m_t_theta = NaNfloat;
0312       m_t_qop = NaNfloat;
0313       m_t_time = NaNfloat;
0314 
0315       m_t_charge = 0;
0316       m_t_p = NaNfloat;
0317       m_t_pt = NaNfloat;
0318       m_t_eta = NaNfloat;
0319     }
0320 
0321     m_outputTree->Fill();
0322   }
0323 
0324   return ProcessCode::SUCCESS;
0325 }
0326 
0327 }  // namespace ActsExamples