ECCE @ EIC Software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
RootVertexAndTracksWriter.cpp
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file RootVertexAndTracksWriter.cpp
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2019 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 "ACTFW/Io/Root/RootVertexAndTracksWriter.hpp"
10 
11 #include <Acts/Utilities/Helpers.hpp>
12 #include <TFile.h>
13 #include <TTree.h>
14 #include <ios>
15 #include <stdexcept>
16 
17 using Acts::VectorHelpers::eta;
18 using Acts::VectorHelpers::perp;
19 using Acts::VectorHelpers::phi;
20 
21 FW::RootVertexAndTracksWriter::RootVertexAndTracksWriter(
22  const FW::RootVertexAndTracksWriter::Config& cfg, Acts::Logging::Level lvl)
23  : WriterT(cfg.collection, "RootVertexAndTracksWriter", lvl),
24  m_cfg(cfg),
25  m_outputFile(cfg.rootFile) {
26  // An input collection name and tree name must be specified
27  if (m_cfg.collection.empty()) {
28  throw std::invalid_argument("Missing input collection");
29  } else if (m_cfg.treeName.empty()) {
30  throw std::invalid_argument("Missing tree name");
31  }
32 
33  // Setup ROOT I/O
34  if (m_outputFile == nullptr) {
35  m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
36  if (m_outputFile == nullptr) {
37  throw std::ios_base::failure("Could not open '" + m_cfg.filePath);
38  }
39  }
40  m_outputFile->cd();
41  m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
42  if (m_outputTree == nullptr)
43  throw std::bad_alloc();
44  else {
45  // I/O parameters
46  m_outputTree->Branch("event_nr", &m_eventNr);
47  m_outputTree->Branch("vx", &m_ptrVx);
48  m_outputTree->Branch("vy", &m_ptrVy);
49  m_outputTree->Branch("vz", &m_ptrVz);
50 
51  m_outputTree->Branch("d0", &m_ptrD0);
52  m_outputTree->Branch("z0", &m_ptrZ0);
53  m_outputTree->Branch("phi", &m_ptrPhi);
54  m_outputTree->Branch("theta", &m_ptrTheta);
55  m_outputTree->Branch("qp", &m_ptrQP);
56  m_outputTree->Branch("time", &m_ptrTime);
57  m_outputTree->Branch("vtxID", &m_ptrVtxID);
58 
59  m_outputTree->Branch("trkCov11", &m_ptrCov11);
60  m_outputTree->Branch("trkCov12", &m_ptrCov12);
61  m_outputTree->Branch("trkCov13", &m_ptrCov13);
62  m_outputTree->Branch("trkCov14", &m_ptrCov14);
63  m_outputTree->Branch("trkCov15", &m_ptrCov15);
64  m_outputTree->Branch("trkCov16", &m_ptrCov16);
65 
66  m_outputTree->Branch("trkCov21", &m_ptrCov21);
67  m_outputTree->Branch("trkCov22", &m_ptrCov22);
68  m_outputTree->Branch("trkCov23", &m_ptrCov23);
69  m_outputTree->Branch("trkCov24", &m_ptrCov24);
70  m_outputTree->Branch("trkCov25", &m_ptrCov25);
71  m_outputTree->Branch("trkCov26", &m_ptrCov26);
72 
73  m_outputTree->Branch("trkCov31", &m_ptrCov31);
74  m_outputTree->Branch("trkCov32", &m_ptrCov32);
75  m_outputTree->Branch("trkCov33", &m_ptrCov33);
76  m_outputTree->Branch("trkCov34", &m_ptrCov34);
77  m_outputTree->Branch("trkCov35", &m_ptrCov35);
78  m_outputTree->Branch("trkCov36", &m_ptrCov36);
79 
80  m_outputTree->Branch("trkCov41", &m_ptrCov41);
81  m_outputTree->Branch("trkCov42", &m_ptrCov42);
82  m_outputTree->Branch("trkCov43", &m_ptrCov43);
83  m_outputTree->Branch("trkCov44", &m_ptrCov44);
84  m_outputTree->Branch("trkCov45", &m_ptrCov45);
85  m_outputTree->Branch("trkCov46", &m_ptrCov46);
86 
87  m_outputTree->Branch("trkCov51", &m_ptrCov51);
88  m_outputTree->Branch("trkCov52", &m_ptrCov52);
89  m_outputTree->Branch("trkCov53", &m_ptrCov53);
90  m_outputTree->Branch("trkCov54", &m_ptrCov54);
91  m_outputTree->Branch("trkCov55", &m_ptrCov55);
92  m_outputTree->Branch("trkCov56", &m_ptrCov56);
93 
94  m_outputTree->Branch("trkCov61", &m_ptrCov61);
95  m_outputTree->Branch("trkCov62", &m_ptrCov62);
96  m_outputTree->Branch("trkCov63", &m_ptrCov63);
97  m_outputTree->Branch("trkCov64", &m_ptrCov64);
98  m_outputTree->Branch("trkCov65", &m_ptrCov65);
99  m_outputTree->Branch("trkCov66", &m_ptrCov66);
100  }
101 }
102 
103 FW::RootVertexAndTracksWriter::~RootVertexAndTracksWriter() {
104  if (m_outputFile) {
105  m_outputFile->Close();
106  }
107 }
108 
109 FW::ProcessCode FW::RootVertexAndTracksWriter::endRun() {
110  if (m_outputFile) {
111  m_outputFile->cd();
112  m_outputTree->Write();
113  ACTS_INFO("Wrote event to tree '" << m_cfg.treeName << "' in '"
114  << m_cfg.filePath << "'");
115  }
116  return ProcessCode::SUCCESS;
117 }
118 
119 void FW::RootVertexAndTracksWriter::ClearAll() {
120  m_vx.clear();
121  m_vy.clear();
122  m_vz.clear();
123  m_d0.clear();
124  m_z0.clear();
125  m_phi.clear();
126  m_theta.clear();
127  m_qp.clear();
128  m_time.clear();
129  m_vtxID.clear();
130 
131  m_cov11.clear();
132  m_cov12.clear();
133  m_cov13.clear();
134  m_cov14.clear();
135  m_cov15.clear();
136  m_cov16.clear();
137 
138  m_cov21.clear();
139  m_cov22.clear();
140  m_cov23.clear();
141  m_cov24.clear();
142  m_cov25.clear();
143  m_cov26.clear();
144 
145  m_cov31.clear();
146  m_cov32.clear();
147  m_cov33.clear();
148  m_cov34.clear();
149  m_cov35.clear();
150  m_cov36.clear();
151 
152  m_cov41.clear();
153  m_cov42.clear();
154  m_cov43.clear();
155  m_cov44.clear();
156  m_cov45.clear();
157  m_cov46.clear();
158 
159  m_cov51.clear();
160  m_cov52.clear();
161  m_cov53.clear();
162  m_cov54.clear();
163  m_cov55.clear();
164  m_cov56.clear();
165 
166  m_cov61.clear();
167  m_cov62.clear();
168  m_cov63.clear();
169  m_cov64.clear();
170  m_cov65.clear();
171  m_cov66.clear();
172 }
173 
174 FW::ProcessCode FW::RootVertexAndTracksWriter::writeT(
175  const AlgorithmContext& context,
176  const std::vector<VertexAndTracks>& vertexAndTracksCollection) {
177  if (m_outputFile == nullptr || vertexAndTracksCollection.empty()) {
178  return ProcessCode::SUCCESS;
179  }
180 
181  // Exclusive access to the tree while writing
182  std::lock_guard<std::mutex> lock(m_writeMutex);
183 
184  ClearAll();
185 
186  // Get the event number
187  m_eventNr = context.eventNumber;
188 
189  for (auto& vertexAndTracks : vertexAndTracksCollection) {
190  // Collect the vertex information
191  m_vx.push_back(vertexAndTracks.vertex.position().x());
192  m_vy.push_back(vertexAndTracks.vertex.position().y());
193  m_vz.push_back(vertexAndTracks.vertex.position().z());
194 
195  for (auto& track : vertexAndTracks.tracks) {
196  // Collect the track information
197  m_d0.push_back(track.parameters()[Acts::ParDef::eLOC_D0]);
198  m_z0.push_back(track.parameters()[Acts::ParDef::eLOC_Z0]);
199  m_phi.push_back(track.parameters()[Acts::ParDef::ePHI]);
200  m_theta.push_back(track.parameters()[Acts::ParDef::eTHETA]);
201  m_qp.push_back(track.parameters()[Acts::ParDef::eQOP]);
202  m_time.push_back(track.parameters()[Acts::ParDef::eT]);
203  // Current vertex index as vertex ID
204  m_vtxID.push_back(m_vx.size() - 1);
205 
206  // Save track covariance
207  Acts::BoundSymMatrix cov = *track.covariance();
208 
209  m_cov11.push_back(cov(0, 0));
210  m_cov12.push_back(cov(0, 1));
211  m_cov13.push_back(cov(0, 2));
212  m_cov14.push_back(cov(0, 3));
213  m_cov15.push_back(cov(0, 4));
214  m_cov16.push_back(cov(0, 5));
215 
216  m_cov21.push_back(cov(1, 0));
217  m_cov22.push_back(cov(1, 1));
218  m_cov23.push_back(cov(1, 2));
219  m_cov24.push_back(cov(1, 3));
220  m_cov25.push_back(cov(1, 4));
221  m_cov26.push_back(cov(1, 5));
222 
223  m_cov31.push_back(cov(2, 0));
224  m_cov32.push_back(cov(2, 1));
225  m_cov33.push_back(cov(2, 2));
226  m_cov34.push_back(cov(2, 3));
227  m_cov35.push_back(cov(2, 4));
228  m_cov36.push_back(cov(2, 5));
229 
230  m_cov41.push_back(cov(3, 0));
231  m_cov42.push_back(cov(3, 1));
232  m_cov43.push_back(cov(3, 2));
233  m_cov44.push_back(cov(3, 3));
234  m_cov45.push_back(cov(3, 4));
235  m_cov46.push_back(cov(3, 5));
236 
237  m_cov51.push_back(cov(4, 0));
238  m_cov52.push_back(cov(4, 1));
239  m_cov53.push_back(cov(4, 2));
240  m_cov54.push_back(cov(4, 3));
241  m_cov55.push_back(cov(4, 4));
242  m_cov56.push_back(cov(4, 5));
243 
244  m_cov61.push_back(cov(5, 0));
245  m_cov62.push_back(cov(5, 1));
246  m_cov63.push_back(cov(5, 2));
247  m_cov64.push_back(cov(5, 3));
248  m_cov65.push_back(cov(5, 4));
249  m_cov66.push_back(cov(5, 5));
250  }
251  }
252 
253  m_outputTree->Fill();
254 
255  return ProcessCode::SUCCESS;
256 }