Root/src/RootTrackParameterWriter.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2017-2021 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 http://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/Utilities/Range.hpp"
0017 #include "ActsExamples/Validation/TrackClassification.hpp"
0018 #include "ActsFatras/EventData/Barcode.hpp"
0019 #include "ActsFatras/EventData/Hit.hpp"
0020 #include "ActsFatras/EventData/Particle.hpp"
0021
0022 #include <cmath>
0023 #include <cstddef>
0024 #include <ios>
0025 #include <iostream>
0026 #include <memory>
0027 #include <stdexcept>
0028 #include <tuple>
0029 #include <utility>
0030 #include <vector>
0031
0032 #include <TFile.h>
0033 #include <TTree.h>
0034
0035 using Acts::VectorHelpers::eta;
0036 using Acts::VectorHelpers::phi;
0037 using Acts::VectorHelpers::theta;
0038
0039 ActsExamples::RootTrackParameterWriter::RootTrackParameterWriter(
0040     const ActsExamples::RootTrackParameterWriter::Config& config,
0041     Acts::Logging::Level level)
0042     : TrackParameterWriter(config.inputTrackParameters,
0043                            "RootTrackParameterWriter", level),
0044       m_cfg(config) {
0045   if (m_cfg.inputProtoTracks.empty()) {
0046     throw std::invalid_argument("Missing proto tracks input collection");
0047   }
0048   if (m_cfg.inputParticles.empty()) {
0049     throw std::invalid_argument("Missing particles input collection");
0050   }
0051   if (m_cfg.inputSimHits.empty()) {
0052     throw std::invalid_argument("Missing simulated hits input collection");
0053   }
0054   if (m_cfg.inputMeasurementParticlesMap.empty()) {
0055     throw std::invalid_argument("Missing hit-particles map input collection");
0056   }
0057   if (m_cfg.inputMeasurementSimHitsMap.empty()) {
0058     throw std::invalid_argument(
0059         "Missing hit-simulated-hits map input collection");
0060   }
0061   if (m_cfg.filePath.empty()) {
0062     throw std::invalid_argument("Missing output filename");
0063   }
0064   if (m_cfg.treeName.empty()) {
0065     throw std::invalid_argument("Missing tree name");
0066   }
0067
0068   m_inputProtoTracks.initialize(m_cfg.inputProtoTracks);
0069   m_inputParticles.initialize(m_cfg.inputParticles);
0070   m_inputSimHits.initialize(m_cfg.inputSimHits);
0071   m_inputMeasurementParticlesMap.initialize(m_cfg.inputMeasurementParticlesMap);
0072   m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
0073
0074   // Setup ROOT I/O
0075   if (m_outputFile == nullptr) {
0076     auto path = m_cfg.filePath;
0077     m_outputFile = TFile::Open(path.c_str(), m_cfg.fileMode.c_str());
0078     if (m_outputFile == nullptr) {
0079       throw std::ios_base::failure("Could not open '" + path + "'");
0080     }
0081   }
0082   m_outputFile->cd();
0083   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0084   if (m_outputTree == nullptr) {
0085     throw std::bad_alloc();
0086   } else {
0087     // The estimated track parameters
0088     m_outputTree->Branch("event_nr", &m_eventNr);
0089     m_outputTree->Branch("loc0", &m_loc0);
0090     m_outputTree->Branch("loc1", &m_loc1);
0091     m_outputTree->Branch("phi", &m_phi);
0092     m_outputTree->Branch("theta", &m_theta);
0093     m_outputTree->Branch("qop", &m_qop);
0094     m_outputTree->Branch("time", &m_time);
0095     m_outputTree->Branch("p", &m_p);
0096     m_outputTree->Branch("pt", &m_pt);
0097     m_outputTree->Branch("eta", &m_eta);
0098     // The truth track parameters
0099     m_outputTree->Branch("eventNr", &m_eventNr);
0100     m_outputTree->Branch("t_loc0", &m_t_loc0);
0101     m_outputTree->Branch("t_loc1", &m_t_loc1);
0102     m_outputTree->Branch("t_phi", &m_t_phi);
0103     m_outputTree->Branch("t_theta", &m_t_theta);
0104     m_outputTree->Branch("t_qop", &m_t_qop);
0105     m_outputTree->Branch("t_time", &m_t_time);
0106     m_outputTree->Branch("truthMatched", &m_truthMatched);
0107   }
0108 }
0109
0110 ActsExamples::RootTrackParameterWriter::~RootTrackParameterWriter() {
0111   if (m_outputFile != nullptr) {
0112     m_outputFile->Close();
0113   }
0114 }
0115
0116 ActsExamples::ProcessCode ActsExamples::RootTrackParameterWriter::finalize() {
0117   m_outputFile->cd();
0118   m_outputTree->Write();
0119   m_outputFile->Close();
0120
0121   ACTS_INFO("Wrote estimated parameters from seed to tree '"
0122             << m_cfg.treeName << "' in '" << m_cfg.filePath << "'");
0123
0124   return ProcessCode::SUCCESS;
0125 }
0126
0127 ActsExamples::ProcessCode ActsExamples::RootTrackParameterWriter::writeT(
0128     const ActsExamples::AlgorithmContext& ctx,
0129     const TrackParametersContainer& trackParams) {
0130   // Read additional input collections
0131   const auto& protoTracks = m_inputProtoTracks(ctx);
0132   const auto& particles = m_inputParticles(ctx);
0133   const auto& simHits = m_inputSimHits(ctx);
0134   const auto& hitParticlesMap = m_inputMeasurementParticlesMap(ctx);
0135   const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
0136
0137   // Exclusive access to the tree while writing
0138   std::lock_guard<std::mutex> lock(m_writeMutex);
0139
0140   // Get the event number
0141   m_eventNr = ctx.eventNumber;
0142
0143   ACTS_VERBOSE("Writing " << trackParams.size() << " track parameters");
0144
0145   // Loop over the estimated track parameters
0146   for (std::size_t iparams = 0; iparams < trackParams.size(); ++iparams) {
0147     // The reference surface of the parameters, i.e. also the reference surface
0148     // of the first space point
0149     const auto& surface = trackParams[iparams].referenceSurface();
0150     // The estimated bound parameters vector
0151     const auto params = trackParams[iparams].parameters();
0152     m_loc0 = params[Acts::eBoundLoc0];
0153     m_loc1 = params[Acts::eBoundLoc1];
0154     m_phi = params[Acts::eBoundPhi];
0155     m_theta = params[Acts::eBoundTheta];
0156     m_qop = params[Acts::eBoundQOverP];
0157     m_time = params[Acts::eBoundTime];
0158     m_p = std::abs(1.0 / m_qop);
0159     m_pt = m_p * std::sin(m_theta);
0160     m_eta = std::atanh(std::cos(m_theta));
0161
0162     // Get the proto track from which the track parameters are estimated
0163     const auto& ptrack = protoTracks[iparams];
0164     std::vector<ParticleHitCount> particleHitCounts;
0165     identifyContributingParticles(hitParticlesMap, ptrack, particleHitCounts);
0166     m_truthMatched = false;
0167     if (particleHitCounts.size() == 1) {
0168       m_truthMatched = true;
0169     }
0170     // Get the index of the first space point
0171     const auto& hitIdx = ptrack.front();
0172     // Get the sim hits via the measurement to sim hits map
0173     auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
0174     auto [truthLocal, truthPos4, truthUnitDir] =
0175         averageSimHits(ctx.geoContext, surface, simHits, indices, logger());
0176     // Get the truth track parameter at the first space point
0177     m_t_loc0 = truthLocal[Acts::ePos0];
0178     m_t_loc1 = truthLocal[Acts::ePos1];
0179     m_t_phi = phi(truthUnitDir);
0180     m_t_theta = theta(truthUnitDir);
0181     m_t_time = truthPos4[Acts::eTime];
0182     // momentum averaging makes even less sense than averaging position and
0183     // direction. use the first momentum or set q/p to zero
0184     if (!indices.empty()) {
0185       // we assume that the indices are within valid ranges so we do not
0186       // need to check their validity again.
0187       const auto simHitIdx0 = indices.begin()->second;
0188       const auto& simHit0 = *simHits.nth(simHitIdx0);
0189       const auto p =
0190           simHit0.momentum4Before().template segment<3>(Acts::eMom0).norm();
0191       const auto& particleId = simHit0.particleId();
0192       // The truth charge has to be retrieved from the sim particle
0193       auto ip = particles.find(particleId);
0194       if (ip != particles.end()) {
0195         const auto& particle = *ip;
0196         m_t_charge = static_cast<int>(particle.charge());
0197         m_t_qop = m_t_charge / p;
0198       } else {
0199         ACTS_DEBUG("Truth particle with barcode "
0200                    << particleId << "=" << particleId.value() << " not found!");
0201       }
0202     }
0203
0204     m_outputTree->Fill();
0205   }
0206
0207   return ProcessCode::SUCCESS;
0208 }