Root/src/RootAthenaNTupleReader.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/RootAthenaNTupleReader.hpp"
0010
0011 #include "Acts/Definitions/TrackParametrization.hpp"
0012 #include "Acts/Surfaces/PerigeeSurface.hpp"
0013 #include "Acts/Surfaces/Surface.hpp"
0014 #include "Acts/Vertexing/Vertex.hpp"
0015 #include "ActsExamples/EventData/Track.hpp"
0016 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0017
0018 #include <cstdint>
0019 #include <iostream>
0020 #include <optional>
0021 #include <stdexcept>
0022
0023 #include <TChain.h>
0024
0025 namespace ActsExamples {
0026
0027 RootAthenaNTupleReader::RootAthenaNTupleReader(
0028     const RootAthenaNTupleReader::Config& config, Acts::Logging::Level level)
0029     : IReader(),
0030       m_cfg(config),
0031       m_logger(Acts::getDefaultLogger(name(), level)) {
0032   if (m_cfg.inputFilePath.empty()) {
0033     throw std::invalid_argument("Missing input filename");
0034   }
0035   if (m_cfg.inputTreeName.empty()) {
0036     throw std::invalid_argument("Missing tree name");
0037   }
0038
0039   m_outputTrackParameters.initialize(m_cfg.outputTrackParameters);
0040   m_outputTruthVtxParameters.initialize(m_cfg.outputTruthVtxParameters);
0041   m_outputRecoVtxParameters.initialize(m_cfg.outputRecoVtxParameters);
0042   m_outputBeamspotConstraint.initialize(m_cfg.outputBeamspotConstraint);
0043
0044   m_inputChain = std::make_unique<TChain>(m_cfg.inputTreeName.c_str());
0045
0046   // unused event identifier
0047   std::int32_t eventNumber = 0;
0048
0049   // Set the branches
0050   m_inputChain->SetBranchAddress("EventNumber", &eventNumber);
0051   m_inputChain->SetBranchAddress("track_d0", &m_branches.track_d0);
0052   m_inputChain->SetBranchAddress("track_z0", &m_branches.track_z0);
0053   m_inputChain->SetBranchAddress("track_theta", &m_branches.track_theta);
0054   m_inputChain->SetBranchAddress("track_phi", &m_branches.track_phi);
0055   m_inputChain->SetBranchAddress("track_qOverP", &m_branches.track_qOverP);
0056   m_inputChain->SetBranchAddress("track_t", &m_branches.track_t);
0057   m_inputChain->SetBranchAddress("track_z", &m_branches.track_z);
0058
0059   // Covariance stuff
0060   m_inputChain->SetBranchAddress("track_var_d0", &m_branches.track_var_d0);
0061   m_inputChain->SetBranchAddress("track_var_z0", &m_branches.track_var_z0);
0062   m_inputChain->SetBranchAddress("track_var_phi", &m_branches.track_var_phi);
0063   m_inputChain->SetBranchAddress("track_var_theta",
0064                                  &m_branches.track_var_theta);
0065   m_inputChain->SetBranchAddress("track_var_qOverP",
0066                                  &m_branches.track_var_qOverP);
0067   m_inputChain->SetBranchAddress("track_cov_d0z0", &m_branches.track_cov_d0z0);
0068   m_inputChain->SetBranchAddress("track_cov_d0phi",
0069                                  &m_branches.track_cov_d0phi);
0070   m_inputChain->SetBranchAddress("track_cov_d0theta",
0071                                  &m_branches.track_cov_d0theta);
0072   m_inputChain->SetBranchAddress("track_cov_d0qOverP",
0073                                  &m_branches.track_cov_d0qOverP);
0074   m_inputChain->SetBranchAddress("track_cov_z0phi",
0075                                  &m_branches.track_cov_z0phi);
0076   m_inputChain->SetBranchAddress("track_cov_z0theta",
0077                                  &m_branches.track_cov_z0theta);
0078   m_inputChain->SetBranchAddress("track_cov_z0qOverP",
0079                                  &m_branches.track_cov_z0qOverP);
0080   m_inputChain->SetBranchAddress("track_cov_phitheta",
0081                                  &m_branches.track_cov_phitheta);
0082   m_inputChain->SetBranchAddress("track_cov_phiqOverP",
0083                                  &m_branches.track_cov_phiqOverP);
0084   m_inputChain->SetBranchAddress("track_cov_tehtaqOverP",
0085                                  &m_branches.track_cov_tehtaqOverP);
0086
0087   // Truth vertex
0088   m_inputChain->SetBranchAddress("truthvertex_x", &m_branches.truthvertex_x);
0089   m_inputChain->SetBranchAddress("truthvertex_y", &m_branches.truthvertex_y);
0090   m_inputChain->SetBranchAddress("truthvertex_z", &m_branches.truthvertex_z);
0091   m_inputChain->SetBranchAddress("truthvertex_t", &m_branches.truthvertex_t);
0092
0093   m_inputChain->SetBranchAddress("recovertex_x", &m_branches.recovertex_x);
0094   m_inputChain->SetBranchAddress("recovertex_y", &m_branches.recovertex_y);
0095   m_inputChain->SetBranchAddress("recovertex_z", &m_branches.recovertex_z);
0096   m_inputChain->SetBranchAddress("truthvertex_tracks_idx",
0097                                  &m_branches.truthvertex_tracks_idx);
0098
0099   m_inputChain->SetBranchAddress("beamspot_x", &m_branches.beamspot_x);
0100   m_inputChain->SetBranchAddress("beamspot_y", &m_branches.beamspot_y);
0101   m_inputChain->SetBranchAddress("beamspot_z", &m_branches.beamspot_z);
0102   m_inputChain->SetBranchAddress("beamspot_sigX", &m_branches.beamspot_sigX);
0103   m_inputChain->SetBranchAddress("beamspot_sigY", &m_branches.beamspot_sigY);
0104   m_inputChain->SetBranchAddress("beamspot_sigZ", &m_branches.beamspot_sigZ);
0105
0106   auto path = m_cfg.inputFilePath;
0107
0108   // add file to the input chain
0109   m_inputChain->Add(path.c_str());
0110   ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.inputTreeName
0111                             << "'.");
0112
0113   m_events = m_inputChain->GetEntries();
0114   ACTS_DEBUG("The full chain has " << m_events << " entries.");
0115 }
0116
0117 RootAthenaNTupleReader::~RootAthenaNTupleReader() = default;
0118
0119 ProcessCode RootAthenaNTupleReader::read(const AlgorithmContext& context) {
0120   ACTS_DEBUG("Trying to read track parameters from ntuple.");
0121
0122   Acts::Vector3 pos(0, 0, 0);
0123   std::shared_ptr<Acts::PerigeeSurface> surface =
0124       Acts::Surface::makeShared<Acts::PerigeeSurface>(pos);
0125
0126   if (context.eventNumber >= m_events) {
0127     ACTS_ERROR("event out of bounds");
0128     return ProcessCode::ABORT;
0129   }
0130
0131   std::lock_guard<std::mutex> lock(m_read_mutex);
0132
0133   auto entry = context.eventNumber;
0134   m_inputChain->GetEntry(entry);
0135   ACTS_INFO("Reading event: " << context.eventNumber
0136                               << " stored as entry: " << entry);
0137
0138   const unsigned int nTracks = m_branches.track_d0.size();
0139   const unsigned int nTruthVtx = m_branches.truthvertex_z.size();
0140   const unsigned int nRecoVtx = m_branches.recovertex_z.size();
0141
0142   ACTS_DEBUG("nTracks = " << nTracks);
0143   ACTS_DEBUG("nTruthVtx = " << nTruthVtx);
0144   ACTS_DEBUG("nRecoVtx = " << nRecoVtx);
0145
0146   TrackParametersContainer trackContainer;
0147   trackContainer.reserve(nTracks);
0148
0149   for (unsigned int i = 0; i < nTracks; i++) {
0150     Acts::BoundVector params;
0151
0152     params[Acts::BoundIndices::eBoundLoc0] = m_branches.track_d0[i];
0153     params[Acts::BoundIndices::eBoundLoc1] = m_branches.track_z0[i];
0154     params[Acts::BoundIndices::eBoundPhi] = m_branches.track_phi[i];
0155     params[Acts::BoundIndices::eBoundTheta] = m_branches.track_theta[i];
0156     params[Acts::BoundIndices::eBoundQOverP] = m_branches.track_qOverP[i];
0157     params[Acts::BoundIndices::eBoundTime] = m_branches.track_t[i];
0158
0159     // Construct and fill covariance matrix
0160     Acts::BoundMatrix cov;
0161
0162     // Variances
0163     cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundLoc0) =
0164         m_branches.track_var_d0[i];
0165     cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundLoc1) =
0166         m_branches.track_var_z0[i];
0167     cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundPhi) =
0168         m_branches.track_var_phi[i];
0169     cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundTheta) =
0170         m_branches.track_var_theta[i];
0171     cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundQOverP) =
0172         m_branches.track_var_qOverP[i];
0173
0174     cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundLoc1) =
0175         m_branches.track_cov_d0z0[i];
0176     cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundPhi) =
0177         m_branches.track_cov_d0phi[i];
0178     cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundTheta) =
0179         m_branches.track_cov_d0theta[i];
0180     cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundQOverP) =
0181         m_branches.track_cov_d0qOverP[i];
0182     cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundPhi) =
0183         m_branches.track_cov_z0phi[i];
0184     cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundTheta) =
0185         m_branches.track_cov_z0theta[i];
0186     cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundQOverP) =
0187         m_branches.track_cov_z0qOverP[i];
0188     cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundTheta) =
0189         m_branches.track_cov_phitheta[i];
0190     cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundQOverP) =
0191         m_branches.track_cov_phiqOverP[i];
0192     cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundQOverP) =
0193         m_branches.track_cov_tehtaqOverP[i];
0194
0195     cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundLoc0) =
0196         m_branches.track_cov_d0z0[i];
0197     cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundLoc0) =
0198         m_branches.track_cov_d0phi[i];
0199     cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundLoc0) =
0200         m_branches.track_cov_d0theta[i];
0201     cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundLoc0) =
0202         m_branches.track_cov_d0qOverP[i];
0203     cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundLoc1) =
0204         m_branches.track_cov_z0phi[i];
0205     cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundLoc1) =
0206         m_branches.track_cov_z0theta[i];
0207     cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundLoc1) =
0208         m_branches.track_cov_z0qOverP[i];
0209     cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundPhi) =
0210         m_branches.track_cov_phitheta[i];
0211     cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundPhi) =
0212         m_branches.track_cov_phiqOverP[i];
0213     cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundTheta) =
0214         m_branches.track_cov_tehtaqOverP[i];
0215
0216     // TODO we do not have a hypothesis at hand here. defaulting to pion
0217     Acts::BoundTrackParameters tc(surface, params, cov,
0218                                   Acts::ParticleHypothesis::pion());
0219     trackContainer.push_back(tc);
0220   }
0221
0222   std::vector<Acts::Vector4> truthVertexContainer;
0223   for (unsigned int i = 0; i < nTruthVtx; i++) {
0224     Acts::Vector4 vtx(m_branches.truthvertex_x[i], m_branches.truthvertex_y[i],
0225                       m_branches.truthvertex_z[i], m_branches.truthvertex_t[i]);
0226     truthVertexContainer.push_back(vtx);
0227   }
0228   std::vector<Acts::Vector4> recoVertexContainer;
0229   for (unsigned int i = 0; i < nRecoVtx; i++) {
0230     Acts::Vector4 vtx(m_branches.recovertex_x[i], m_branches.recovertex_y[i],
0231                       m_branches.recovertex_z[i], 0);
0232     recoVertexContainer.push_back(vtx);
0233   }
0234
0235   Acts::Vertex beamspotConstraint;
0236   Acts::Vector3 beamspotPos;
0237   Acts::SquareMatrix3 beamspotCov;
0238
0239   beamspotPos << m_branches.beamspot_x, m_branches.beamspot_y,
0240       m_branches.beamspot_z;
0241   beamspotCov << m_branches.beamspot_sigX * m_branches.beamspot_sigX, 0, 0, 0,
0242       m_branches.beamspot_sigY * m_branches.beamspot_sigY, 0, 0, 0,
0243       m_branches.beamspot_sigZ * m_branches.beamspot_sigZ;
0244
0245   beamspotConstraint.setPosition(beamspotPos);
0246   beamspotConstraint.setCovariance(beamspotCov);
0247
0248   m_outputTrackParameters(context, std::move(trackContainer));
0249   m_outputTruthVtxParameters(context, std::move(truthVertexContainer));
0250   m_outputRecoVtxParameters(context, std::move(recoVertexContainer));
0251   m_outputBeamspotConstraint(context, std::move(beamspotConstraint));
0252
0253   // Return success flag
0254   return ProcessCode::SUCCESS;
0255 }
0256
0257 }  // namespace ActsExamples