Root/src/RootTrackParameterWriter.cpp

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