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