Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
RootAthenaNTupleReader.cpp
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file RootAthenaNTupleReader.cpp
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2017-2022 CERN for the benefit of the Acts project
4 //
5 // This Source Code Form is subject to the terms of the Mozilla Public
6 // License, v. 2.0. If a copy of the MPL was not distributed with this
7 // file, You can obtain one at http://mozilla.org/MPL/2.0/.
8 
9 #include "ActsExamples/Io/Root/RootAthenaNTupleReader.hpp"
10 
11 #include "Acts/Definitions/TrackParametrization.hpp"
12 #include "Acts/EventData/GenericBoundTrackParameters.hpp"
13 #include "Acts/EventData/TrackParameters.hpp"
14 #include "Acts/Surfaces/PerigeeSurface.hpp"
15 #include "Acts/Surfaces/Surface.hpp"
16 #include "Acts/Vertexing/Vertex.hpp"
17 #include "ActsExamples/EventData/Track.hpp"
18 #include "ActsExamples/Framework/AlgorithmContext.hpp"
19 
20 #include <cstdint>
21 #include <iostream>
22 #include <optional>
23 #include <stdexcept>
24 
25 #include <TChain.h>
26 #include <TMathBase.h>
27 
28 ActsExamples::RootAthenaNTupleReader::RootAthenaNTupleReader(
29  const ActsExamples::RootAthenaNTupleReader::Config& config,
30  Acts::Logging::Level level)
31  : ActsExamples::IReader(),
32  m_cfg(config),
33  m_logger(Acts::getDefaultLogger(name(), level)) {
34  if (m_cfg.inputFilePath.empty()) {
35  throw std::invalid_argument("Missing input filename");
36  }
37  if (m_cfg.inputTreeName.empty()) {
38  throw std::invalid_argument("Missing tree name");
39  }
40 
41  m_outputTrackParameters.initialize(m_cfg.outputTrackParameters);
42  m_outputTruthVtxParameters.initialize(m_cfg.outputTruthVtxParameters);
43  m_outputRecoVtxParameters.initialize(m_cfg.outputRecoVtxParameters);
44  m_outputBeamspotConstraint.initialize(m_cfg.outputBeamspotConstraint);
45 
46  m_inputChain = new TChain(m_cfg.inputTreeName.c_str());
47 
48  // unused event identifier
49  std::int32_t eventNumber = 0;
50 
51  // Set the branches
52  m_inputChain->SetBranchAddress("EventNumber", &eventNumber);
53  m_inputChain->SetBranchAddress("track_d0", &m_branches.track_d0);
54  m_inputChain->SetBranchAddress("track_z0", &m_branches.track_z0);
55  m_inputChain->SetBranchAddress("track_theta", &m_branches.track_theta);
56  m_inputChain->SetBranchAddress("track_phi", &m_branches.track_phi);
57  m_inputChain->SetBranchAddress("track_qOverP", &m_branches.track_qOverP);
58  m_inputChain->SetBranchAddress("track_t", &m_branches.track_t);
59  m_inputChain->SetBranchAddress("track_z", &m_branches.track_z);
60 
61  // Covariance stuff
62  m_inputChain->SetBranchAddress("track_var_d0", &m_branches.track_var_d0);
63  m_inputChain->SetBranchAddress("track_var_z0", &m_branches.track_var_z0);
64  m_inputChain->SetBranchAddress("track_var_phi", &m_branches.track_var_phi);
65  m_inputChain->SetBranchAddress("track_var_theta",
66  &m_branches.track_var_theta);
67  m_inputChain->SetBranchAddress("track_var_qOverP",
68  &m_branches.track_var_qOverP);
69  m_inputChain->SetBranchAddress("track_cov_d0z0", &m_branches.track_cov_d0z0);
70  m_inputChain->SetBranchAddress("track_cov_d0phi",
71  &m_branches.track_cov_d0phi);
72  m_inputChain->SetBranchAddress("track_cov_d0theta",
73  &m_branches.track_cov_d0theta);
74  m_inputChain->SetBranchAddress("track_cov_d0qOverP",
75  &m_branches.track_cov_d0qOverP);
76  m_inputChain->SetBranchAddress("track_cov_z0phi",
77  &m_branches.track_cov_z0phi);
78  m_inputChain->SetBranchAddress("track_cov_z0theta",
79  &m_branches.track_cov_z0theta);
80  m_inputChain->SetBranchAddress("track_cov_z0qOverP",
81  &m_branches.track_cov_z0qOverP);
82  m_inputChain->SetBranchAddress("track_cov_phitheta",
83  &m_branches.track_cov_phitheta);
84  m_inputChain->SetBranchAddress("track_cov_phiqOverP",
85  &m_branches.track_cov_phiqOverP);
86  m_inputChain->SetBranchAddress("track_cov_tehtaqOverP",
87  &m_branches.track_cov_tehtaqOverP);
88 
89  // Truth vertex
90  m_inputChain->SetBranchAddress("truthvertex_x", &m_branches.truthvertex_x);
91  m_inputChain->SetBranchAddress("truthvertex_y", &m_branches.truthvertex_y);
92  m_inputChain->SetBranchAddress("truthvertex_z", &m_branches.truthvertex_z);
93  m_inputChain->SetBranchAddress("truthvertex_t", &m_branches.truthvertex_t);
94 
95  m_inputChain->SetBranchAddress("recovertex_x", &m_branches.recovertex_x);
96  m_inputChain->SetBranchAddress("recovertex_y", &m_branches.recovertex_y);
97  m_inputChain->SetBranchAddress("recovertex_z", &m_branches.recovertex_z);
98  m_inputChain->SetBranchAddress("truthvertex_tracks_idx",
99  &m_branches.truthvertex_tracks_idx);
100 
101  m_inputChain->SetBranchAddress("beamspot_x", &m_branches.beamspot_x);
102  m_inputChain->SetBranchAddress("beamspot_y", &m_branches.beamspot_y);
103  m_inputChain->SetBranchAddress("beamspot_z", &m_branches.beamspot_z);
104  m_inputChain->SetBranchAddress("beamspot_sigX", &m_branches.beamspot_sigX);
105  m_inputChain->SetBranchAddress("beamspot_sigY", &m_branches.beamspot_sigY);
106  m_inputChain->SetBranchAddress("beamspot_sigZ", &m_branches.beamspot_sigZ);
107 
108  auto path = m_cfg.inputFilePath;
109 
110  // add file to the input chain
111  m_inputChain->Add(path.c_str());
112  ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.inputTreeName
113  << "'.");
114 
115  m_events = m_inputChain->GetEntries();
116  ACTS_DEBUG("The full chain has " << m_events << " entries.");
117 
118  {
119  // The entry numbers for accessing events in increased order (there could be
120  // multiple entries corresponding to one event number)
121  std::vector<long long> m_entryNumbers;
122  // If the events are not in order, get the entry numbers for ordered events
123  m_entryNumbers.resize(m_events);
124  m_inputChain->Draw("EventNumber", "", "goff");
125  // Sort to get the entry numbers of the ordered events
126  TMath::Sort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
127  m_entryNumbers.data(), false);
128  }
129 }
130 
131 ActsExamples::ProcessCode ActsExamples::RootAthenaNTupleReader::read(
132  const ActsExamples::AlgorithmContext& context) {
133  ACTS_DEBUG("Trying to read track parameters from ntuple.");
134 
135  Acts::Vector3 pos(0, 0, 0);
136  std::shared_ptr<Acts::PerigeeSurface> surface =
137  Acts::Surface::makeShared<Acts::PerigeeSurface>(pos);
138 
139  if (context.eventNumber >= m_events) {
140  ACTS_ERROR("event out of bounds");
141  return ProcessCode::ABORT;
142  }
143 
144  std::lock_guard<std::mutex> lock(m_read_mutex);
145 
146  auto entry = context.eventNumber;
147  m_inputChain->GetEntry(entry);
148  ACTS_INFO("Reading event: " << context.eventNumber
149  << " stored as entry: " << entry);
150 
151  const unsigned int nTracks = m_branches.track_d0.size();
152  const unsigned int nTruthVtx = m_branches.truthvertex_z.size();
153  const unsigned int nRecoVtx = m_branches.recovertex_z.size();
154 
155  ACTS_DEBUG("nTracks = " << nTracks);
156  ACTS_DEBUG("nTruthVtx = " << nTruthVtx);
157  ACTS_DEBUG("nRecoVtx = " << nRecoVtx);
158 
159  TrackParametersContainer trackContainer;
160  trackContainer.reserve(nTracks);
161 
162  for (unsigned int i = 0; i < nTracks; i++) {
163  Acts::BoundVector params;
164 
165  params[Acts::BoundIndices::eBoundLoc0] = m_branches.track_d0[i];
166  params[Acts::BoundIndices::eBoundLoc1] = m_branches.track_z0[i];
167  params[Acts::BoundIndices::eBoundPhi] = m_branches.track_phi[i];
168  params[Acts::BoundIndices::eBoundTheta] = m_branches.track_theta[i];
169  params[Acts::BoundIndices::eBoundQOverP] = m_branches.track_qOverP[i];
170  params[Acts::BoundIndices::eBoundTime] = m_branches.track_t[i];
171 
172  // Construct and fill covariance matrix
173  Acts::BoundSquareMatrix cov;
174 
175  // Variances
176  cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundLoc0) =
177  m_branches.track_var_d0[i];
178  cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundLoc1) =
179  m_branches.track_var_z0[i];
180  cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundPhi) =
181  m_branches.track_var_phi[i];
182  cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundTheta) =
183  m_branches.track_var_theta[i];
184  cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundQOverP) =
185  m_branches.track_var_qOverP[i];
186 
187  cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundLoc1) =
188  m_branches.track_cov_d0z0[i];
189  cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundPhi) =
190  m_branches.track_cov_d0phi[i];
191  cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundTheta) =
192  m_branches.track_cov_d0theta[i];
193  cov(Acts::BoundIndices::eBoundLoc0, Acts::BoundIndices::eBoundQOverP) =
194  m_branches.track_cov_d0qOverP[i];
195  cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundPhi) =
196  m_branches.track_cov_z0phi[i];
197  cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundTheta) =
198  m_branches.track_cov_z0theta[i];
199  cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundQOverP) =
200  m_branches.track_cov_z0qOverP[i];
201  cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundTheta) =
202  m_branches.track_cov_phitheta[i];
203  cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundQOverP) =
204  m_branches.track_cov_phiqOverP[i];
205  cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundQOverP) =
206  m_branches.track_cov_tehtaqOverP[i];
207 
208  cov(Acts::BoundIndices::eBoundLoc1, Acts::BoundIndices::eBoundLoc0) =
209  m_branches.track_cov_d0z0[i];
210  cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundLoc0) =
211  m_branches.track_cov_d0phi[i];
212  cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundLoc0) =
213  m_branches.track_cov_d0theta[i];
214  cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundLoc0) =
215  m_branches.track_cov_d0qOverP[i];
216  cov(Acts::BoundIndices::eBoundPhi, Acts::BoundIndices::eBoundLoc1) =
217  m_branches.track_cov_z0phi[i];
218  cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundLoc1) =
219  m_branches.track_cov_z0theta[i];
220  cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundLoc1) =
221  m_branches.track_cov_z0qOverP[i];
222  cov(Acts::BoundIndices::eBoundTheta, Acts::BoundIndices::eBoundPhi) =
223  m_branches.track_cov_phitheta[i];
224  cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundPhi) =
225  m_branches.track_cov_phiqOverP[i];
226  cov(Acts::BoundIndices::eBoundQOverP, Acts::BoundIndices::eBoundTheta) =
227  m_branches.track_cov_tehtaqOverP[i];
228 
229  // TODO we do not have a hypothesis at hand here. defaulting to pion
230  Acts::BoundTrackParameters tc(surface, params, cov,
231  Acts::ParticleHypothesis::pion());
232  trackContainer.push_back(tc);
233  }
234 
235  std::vector<Acts::Vector4> truthVertexContainer;
236  for (unsigned int i = 0; i < nTruthVtx; i++) {
237  Acts::Vector4 vtx(m_branches.truthvertex_x[i], m_branches.truthvertex_y[i],
238  m_branches.truthvertex_z[i], m_branches.truthvertex_t[i]);
239  truthVertexContainer.push_back(vtx);
240  }
241  std::vector<Acts::Vector4> recoVertexContainer;
242  for (unsigned int i = 0; i < nRecoVtx; i++) {
243  Acts::Vector4 vtx(m_branches.recovertex_x[i], m_branches.recovertex_y[i],
244  m_branches.recovertex_z[i], 0);
245  recoVertexContainer.push_back(vtx);
246  }
247 
248  Acts::Vertex<Acts::BoundTrackParameters> beamspotConstraint;
249  Acts::Vector3 beamspotPos;
250  Acts::SquareMatrix3 beamspotCov;
251 
252  beamspotPos << m_branches.beamspot_x, m_branches.beamspot_y,
253  m_branches.beamspot_z;
254  beamspotCov << m_branches.beamspot_sigX * m_branches.beamspot_sigX, 0, 0, 0,
255  m_branches.beamspot_sigY * m_branches.beamspot_sigY, 0, 0, 0,
256  m_branches.beamspot_sigZ * m_branches.beamspot_sigZ;
257 
258  beamspotConstraint.setPosition(beamspotPos);
259  beamspotConstraint.setCovariance(beamspotCov);
260 
261  m_outputTrackParameters(context, std::move(trackContainer));
262  m_outputTruthVtxParameters(context, std::move(truthVertexContainer));
263  m_outputRecoVtxParameters(context, std::move(recoVertexContainer));
264  m_outputBeamspotConstraint(context, std::move(beamspotConstraint));
265 
266  // Return success flag
267  return ProcessCode::SUCCESS;
268 }