Analysis Software
Documentation for sPHENIX simulation software
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Groups Pages
RootMaterialTrackWriter.cpp
Go to the documentation of this file. Or view the newest version in sPHENIX GitHub for file RootMaterialTrackWriter.cpp
1 // This file is part of the Acts project.
2 //
3 // Copyright (C) 2017-2018 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/RootMaterialTrackWriter.hpp"
10 
11 #include "Acts/Geometry/GeometryIdentifier.hpp"
12 #include "Acts/Geometry/TrackingVolume.hpp"
13 #include "Acts/Geometry/Volume.hpp"
14 #include "Acts/Material/Material.hpp"
15 #include "Acts/Material/MaterialInteraction.hpp"
16 #include "Acts/Material/MaterialSlab.hpp"
17 #include "Acts/Surfaces/CylinderBounds.hpp"
18 #include "Acts/Surfaces/RadialBounds.hpp"
19 #include "Acts/Surfaces/Surface.hpp"
20 #include "Acts/Surfaces/SurfaceBounds.hpp"
21 #include "Acts/Utilities/Intersection.hpp"
22 #include "Acts/Utilities/Logger.hpp"
23 #include "ActsExamples/Framework/AlgorithmContext.hpp"
24 
25 #include <algorithm>
26 #include <cstddef>
27 #include <ios>
28 #include <stdexcept>
29 #include <type_traits>
30 
31 #include <TFile.h>
32 #include <TTree.h>
33 
34 using Acts::VectorHelpers::eta;
35 using Acts::VectorHelpers::perp;
36 using Acts::VectorHelpers::phi;
37 
38 ActsExamples::RootMaterialTrackWriter::RootMaterialTrackWriter(
39  const ActsExamples::RootMaterialTrackWriter::Config& config,
40  Acts::Logging::Level level)
41  : WriterT(config.collection, "RootMaterialTrackWriter", level),
42  m_cfg(config) {
43  // An input collection name and tree name must be specified
44  if (m_cfg.collection.empty()) {
45  throw std::invalid_argument("Missing input collection");
46  } else if (m_cfg.treeName.empty()) {
47  throw std::invalid_argument("Missing tree name");
48  }
49 
50  // Setup ROOT I/O
51  m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
52  if (m_outputFile == nullptr) {
53  throw std::ios_base::failure("Could not open '" + m_cfg.filePath);
54  }
55 
56  m_outputFile->cd();
57  m_outputTree =
58  new TTree(m_cfg.treeName.c_str(), "TTree from RootMaterialTrackWriter");
59  if (m_outputTree == nullptr) {
60  throw std::bad_alloc();
61  }
62 
63  // Set the branches
64  m_outputTree->Branch("event_id", &m_eventId);
65  m_outputTree->Branch("v_x", &m_v_x);
66  m_outputTree->Branch("v_y", &m_v_y);
67  m_outputTree->Branch("v_z", &m_v_z);
68  m_outputTree->Branch("v_px", &m_v_px);
69  m_outputTree->Branch("v_py", &m_v_py);
70  m_outputTree->Branch("v_pz", &m_v_pz);
71  m_outputTree->Branch("v_phi", &m_v_phi);
72  m_outputTree->Branch("v_eta", &m_v_eta);
73  m_outputTree->Branch("t_X0", &m_tX0);
74  m_outputTree->Branch("t_L0", &m_tL0);
75  m_outputTree->Branch("mat_x", &m_step_x);
76  m_outputTree->Branch("mat_y", &m_step_y);
77  m_outputTree->Branch("mat_z", &m_step_z);
78  m_outputTree->Branch("mat_dx", &m_step_dx);
79  m_outputTree->Branch("mat_dy", &m_step_dy);
80  m_outputTree->Branch("mat_dz", &m_step_dz);
81  m_outputTree->Branch("mat_step_length", &m_step_length);
82  m_outputTree->Branch("mat_X0", &m_step_X0);
83  m_outputTree->Branch("mat_L0", &m_step_L0);
84  m_outputTree->Branch("mat_A", &m_step_A);
85  m_outputTree->Branch("mat_Z", &m_step_Z);
86  m_outputTree->Branch("mat_rho", &m_step_rho);
87 
88  if (m_cfg.prePostStep) {
89  m_outputTree->Branch("mat_sx", &m_step_sx);
90  m_outputTree->Branch("mat_sy", &m_step_sy);
91  m_outputTree->Branch("mat_sz", &m_step_sz);
92  m_outputTree->Branch("mat_ex", &m_step_ex);
93  m_outputTree->Branch("mat_ey", &m_step_ey);
94  m_outputTree->Branch("mat_ez", &m_step_ez);
95  }
96  if (m_cfg.storeSurface) {
97  m_outputTree->Branch("sur_id", &m_sur_id);
98  m_outputTree->Branch("sur_type", &m_sur_type);
99  m_outputTree->Branch("sur_x", &m_sur_x);
100  m_outputTree->Branch("sur_y", &m_sur_y);
101  m_outputTree->Branch("sur_z", &m_sur_z);
102  m_outputTree->Branch("sur_pathCorrection", &m_sur_pathCorrection);
103  m_outputTree->Branch("sur_range_min", &m_sur_range_min);
104  m_outputTree->Branch("sur_range_max", &m_sur_range_max);
105  }
106  if (m_cfg.storeVolume) {
107  m_outputTree->Branch("vol_id", &m_vol_id);
108  }
109 }
110 
111 ActsExamples::RootMaterialTrackWriter::~RootMaterialTrackWriter() {
112  if (m_outputFile != nullptr) {
113  m_outputFile->Close();
114  }
115 }
116 
117 ActsExamples::ProcessCode ActsExamples::RootMaterialTrackWriter::finalize() {
118  // write the tree and close the file
119  ACTS_INFO("Writing ROOT output File : " << m_cfg.filePath);
120 
121  m_outputFile->cd();
122  m_outputTree->Write();
123  m_outputFile->Close();
124 
125  return ActsExamples::ProcessCode::SUCCESS;
126 }
127 
128 ActsExamples::ProcessCode ActsExamples::RootMaterialTrackWriter::writeT(
129  const AlgorithmContext& ctx,
130  const std::unordered_map<size_t, Acts::RecordedMaterialTrack>&
131  materialTracks) {
132  // Exclusive access to the tree while writing
133  std::lock_guard<std::mutex> lock(m_writeMutex);
134 
135  m_eventId = ctx.eventNumber;
136  // Loop over the material tracks and write them out
137  for (auto& [idTrack, mtrack] : materialTracks) {
138  // Clearing the vector first
139  m_step_sx.clear();
140  m_step_sy.clear();
141  m_step_sz.clear();
142  m_step_x.clear();
143  m_step_y.clear();
144  m_step_z.clear();
145  m_step_ex.clear();
146  m_step_ey.clear();
147  m_step_ez.clear();
148  m_step_dx.clear();
149  m_step_dy.clear();
150  m_step_dz.clear();
151  m_step_length.clear();
152  m_step_X0.clear();
153  m_step_L0.clear();
154  m_step_A.clear();
155  m_step_Z.clear();
156  m_step_rho.clear();
157 
158  m_sur_id.clear();
159  m_sur_type.clear();
160  m_sur_x.clear();
161  m_sur_y.clear();
162  m_sur_z.clear();
163  m_sur_pathCorrection.clear();
164  m_sur_range_min.clear();
165  m_sur_range_max.clear();
166 
167  m_vol_id.clear();
168 
169  auto materialInteractions = mtrack.second.materialInteractions;
170  if (m_cfg.collapseInteractions) {
171  std::vector<Acts::MaterialInteraction> collapsed;
172 
173  Acts::Vector3 positionSum = Acts::Vector3::Zero();
174  double pathCorrectionSum = 0;
175 
176  for (std::size_t start = 0, end = 0; end < materialInteractions.size();
177  ++end) {
178  const auto& mintStart = materialInteractions[start];
179  const auto& mintEnd = materialInteractions[end];
180 
181  positionSum += mintEnd.position;
182  pathCorrectionSum += mintEnd.pathCorrection;
183 
184  const bool same = mintStart.materialSlab.material() ==
185  mintEnd.materialSlab.material();
186  const bool last = end == materialInteractions.size() - 1;
187 
188  if (!same || last) {
189  auto mint = mintStart;
190  mint.position = positionSum / (end - start);
191  mint.pathCorrection = pathCorrectionSum;
192 
193  collapsed.push_back(mint);
194 
195  start = end;
196  positionSum = Acts::Vector3::Zero();
197  pathCorrectionSum = 0;
198  }
199  }
200 
201  materialInteractions = std::move(collapsed);
202  }
203 
204  // Reserve the vector then
205  size_t mints = materialInteractions.size();
206  m_step_sx.reserve(mints);
207  m_step_sy.reserve(mints);
208  m_step_sz.reserve(mints);
209  m_step_x.reserve(mints);
210  m_step_y.reserve(mints);
211  m_step_z.reserve(mints);
212  m_step_ex.reserve(mints);
213  m_step_ey.reserve(mints);
214  m_step_ez.reserve(mints);
215  m_step_dx.reserve(mints);
216  m_step_dy.reserve(mints);
217  m_step_dz.reserve(mints);
218  m_step_length.reserve(mints);
219  m_step_X0.reserve(mints);
220  m_step_L0.reserve(mints);
221  m_step_A.reserve(mints);
222  m_step_Z.reserve(mints);
223  m_step_rho.reserve(mints);
224 
225  m_sur_id.reserve(mints);
226  m_sur_type.reserve(mints);
227  m_sur_x.reserve(mints);
228  m_sur_y.reserve(mints);
229  m_sur_z.reserve(mints);
230  m_sur_pathCorrection.reserve(mints);
231  m_sur_range_min.reserve(mints);
232  m_sur_range_max.reserve(mints);
233 
234  m_vol_id.reserve(mints);
235 
236  // reset the global counter
237  if (m_cfg.recalculateTotals) {
238  m_tX0 = 0.;
239  m_tL0 = 0.;
240  } else {
241  m_tX0 = mtrack.second.materialInX0;
242  m_tL0 = mtrack.second.materialInL0;
243  }
244 
245  // set the track information at vertex
246  m_v_x = mtrack.first.first.x();
247  m_v_y = mtrack.first.first.y();
248  m_v_z = mtrack.first.first.z();
249  m_v_px = mtrack.first.second.x();
250  m_v_py = mtrack.first.second.y();
251  m_v_pz = mtrack.first.second.z();
252  m_v_phi = phi(mtrack.first.second);
253  m_v_eta = eta(mtrack.first.second);
254 
255  // and now loop over the material
256  for (const auto& mint : materialInteractions) {
257  auto direction = mint.direction.normalized();
258 
259  // The material step position information
260  m_step_x.push_back(mint.position.x());
261  m_step_y.push_back(mint.position.y());
262  m_step_z.push_back(mint.position.z());
263  m_step_dx.push_back(direction.x());
264  m_step_dy.push_back(direction.y());
265  m_step_dz.push_back(direction.z());
266 
267  if (m_cfg.prePostStep) {
268  Acts::Vector3 prePos =
269  mint.position - 0.5 * mint.pathCorrection * direction;
270  Acts::Vector3 posPos =
271  mint.position + 0.5 * mint.pathCorrection * direction;
272 
273  m_step_sx.push_back(prePos.x());
274  m_step_sy.push_back(prePos.y());
275  m_step_sz.push_back(prePos.z());
276  m_step_ex.push_back(posPos.x());
277  m_step_ey.push_back(posPos.y());
278  m_step_ez.push_back(posPos.z());
279  }
280 
281  // Store surface information
282  if (m_cfg.storeSurface) {
283  const Acts::Surface* surface = mint.surface;
284  if (mint.intersectionID.value() != 0) {
285  m_sur_id.push_back(mint.intersectionID.value());
286  m_sur_pathCorrection.push_back(mint.pathCorrection);
287  m_sur_x.push_back(mint.intersection.x());
288  m_sur_y.push_back(mint.intersection.y());
289  m_sur_z.push_back(mint.intersection.z());
290  } else if (surface != nullptr) {
291  auto sfIntersection = surface
292  ->intersect(ctx.geoContext, mint.position,
293  mint.direction, true)
294  .closest();
295  m_sur_id.push_back(surface->geometryId().value());
296  m_sur_pathCorrection.push_back(1.0);
297  m_sur_x.push_back(sfIntersection.position().x());
298  m_sur_y.push_back(sfIntersection.position().y());
299  m_sur_z.push_back(sfIntersection.position().z());
300  } else {
301  m_sur_id.push_back(Acts::GeometryIdentifier().value());
302  m_sur_x.push_back(0);
303  m_sur_y.push_back(0);
304  m_sur_z.push_back(0);
305  m_sur_pathCorrection.push_back(1.0);
306  }
307  if (surface != nullptr) {
308  m_sur_type.push_back(surface->type());
309  const Acts::SurfaceBounds& surfaceBounds = surface->bounds();
310  const Acts::RadialBounds* radialBounds =
311  dynamic_cast<const Acts::RadialBounds*>(&surfaceBounds);
312  const Acts::CylinderBounds* cylinderBounds =
313  dynamic_cast<const Acts::CylinderBounds*>(&surfaceBounds);
314  if (radialBounds != nullptr) {
315  m_sur_range_min.push_back(radialBounds->rMin());
316  m_sur_range_max.push_back(radialBounds->rMax());
317  } else if (cylinderBounds != nullptr) {
318  m_sur_range_min.push_back(
319  -cylinderBounds->get(Acts::CylinderBounds::eHalfLengthZ));
320  m_sur_range_max.push_back(
321  cylinderBounds->get(Acts::CylinderBounds::eHalfLengthZ));
322  } else {
323  m_sur_range_min.push_back(0);
324  m_sur_range_max.push_back(0);
325  }
326  } else {
327  m_sur_type.push_back(-1);
328  m_sur_range_min.push_back(0);
329  m_sur_range_max.push_back(0);
330  }
331  }
332 
333  // store volume information
334  if (m_cfg.storeVolume) {
335  Acts::GeometryIdentifier vlayerID;
336  if (not mint.volume.empty()) {
337  vlayerID = mint.volume.geometryId();
338  m_vol_id.push_back(vlayerID.value());
339  } else {
340  vlayerID.setVolume(0);
341  vlayerID.setBoundary(0);
342  vlayerID.setLayer(0);
343  vlayerID.setApproach(0);
344  vlayerID.setSensitive(0);
345  m_vol_id.push_back(vlayerID.value());
346  }
347  }
348 
349  // the material information
350  const auto& mprops = mint.materialSlab;
351  m_step_length.push_back(mprops.thickness());
352  m_step_X0.push_back(mprops.material().X0());
353  m_step_L0.push_back(mprops.material().L0());
354  m_step_A.push_back(mprops.material().Ar());
355  m_step_Z.push_back(mprops.material().Z());
356  m_step_rho.push_back(mprops.material().massDensity());
357  // re-calculate if defined to do so
358  if (m_cfg.recalculateTotals) {
359  m_tX0 += mprops.thicknessInX0();
360  m_tL0 += mprops.thicknessInL0();
361  }
362  }
363  // write to
364  m_outputTree->Fill();
365  }
366 
367  // return success
368  return ActsExamples::ProcessCode::SUCCESS;
369 }