Root/src/RootPropagationStepsWriter.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/RootPropagationStepsWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Geometry/GeometryIdentifier.hpp"
0013 #include "Acts/Propagator/ConstrainedStep.hpp"
0014 #include "Acts/Utilities/Helpers.hpp"
0015 #include "Acts/Utilities/VectorHelpers.hpp"
0016 #include "ActsExamples/EventData/PropagationSummary.hpp"
0017 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0018
0019 #include <ios>
0020 #include <memory>
0021 #include <ostream>
0022 #include <stdexcept>
0023
0024 #include <TFile.h>
0025 #include <TTree.h>
0026
0027 ActsExamples::RootPropagationStepsWriter::RootPropagationStepsWriter(
0028     const ActsExamples::RootPropagationStepsWriter::Config& cfg,
0029     Acts::Logging::Level level)
0030     : WriterT(cfg.collection, "RootPropagationStepsWriter", level),
0031       m_cfg(cfg),
0032       m_outputFile(cfg.rootFile) {
0033   // An input collection name and tree name must be specified
0034   if (m_cfg.collection.empty()) {
0035     throw std::invalid_argument("Missing input collection");
0036   } else if (m_cfg.treeName.empty()) {
0037     throw std::invalid_argument("Missing tree name");
0038   }
0039
0040   // Setup ROOT I/O
0041   if (m_outputFile == nullptr) {
0042     m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0043     if (m_outputFile == nullptr) {
0044       throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0045     }
0046   }
0047   m_outputFile->cd();
0048
0049   m_outputTree = new TTree(m_cfg.treeName.c_str(),
0050                            "TTree from RootPropagationStepsWriter");
0051   if (m_outputTree == nullptr) {
0052     throw std::bad_alloc();
0053   }
0054
0055   // Set the branches
0056   m_outputTree->Branch("event_nr", &m_eventNr);
0057   m_outputTree->Branch("track_nr", &m_trackNr);
0058   m_outputTree->Branch("d0", &m_d0);
0059   m_outputTree->Branch("z0", &m_z0);
0060   m_outputTree->Branch("phi", &m_phi);
0061   m_outputTree->Branch("theta", &m_theta);
0062   m_outputTree->Branch("qOverP", &m_qOverP);
0063   m_outputTree->Branch("t", &m_t);
0064   m_outputTree->Branch("eta", &m_eta);
0065   m_outputTree->Branch("pt", &m_pt);
0066   m_outputTree->Branch("p", &m_p);
0067   m_outputTree->Branch("volume_id", &m_volumeID);
0068   m_outputTree->Branch("boundary_id", &m_boundaryID);
0069   m_outputTree->Branch("layer_id", &m_layerID);
0070   m_outputTree->Branch("approach_id", &m_approachID);
0071   m_outputTree->Branch("sensitive_id", &m_sensitiveID);
0072   m_outputTree->Branch("extra_id", &m_extraID);
0073   m_outputTree->Branch("material", &m_material);
0074   m_outputTree->Branch("g_x", &m_x);
0075   m_outputTree->Branch("g_y", &m_y);
0076   m_outputTree->Branch("g_z", &m_z);
0077   m_outputTree->Branch("g_r", &m_r);
0078   m_outputTree->Branch("d_x", &m_dx);
0079   m_outputTree->Branch("d_y", &m_dy);
0080   m_outputTree->Branch("d_z", &m_dz);
0081   m_outputTree->Branch("type", &m_step_type);
0082   m_outputTree->Branch("step_acc", &m_step_acc);
0083   m_outputTree->Branch("step_act", &m_step_act);
0084   m_outputTree->Branch("step_abt", &m_step_abt);
0085   m_outputTree->Branch("step_usr", &m_step_usr);
0086   m_outputTree->Branch("nStepTrials", &m_nStepTrials);
0087 }
0088
0089 ActsExamples::RootPropagationStepsWriter::~RootPropagationStepsWriter() {
0090   if (m_outputFile != nullptr) {
0091     m_outputFile->Close();
0092   }
0093 }
0094
0095 ActsExamples::ProcessCode ActsExamples::RootPropagationStepsWriter::finalize() {
0096   // Write the tree
0097   m_outputFile->cd();
0098   m_outputTree->Write();
0099   /// Close the file if it's yours
0100   if (m_cfg.rootFile == nullptr) {
0101     m_outputFile->Close();
0102   }
0103
0104   ACTS_VERBOSE("Wrote particles to tree '" << m_cfg.treeName << "' in '"
0105                                            << m_cfg.filePath << "'");
0106
0107   return ProcessCode::SUCCESS;
0108 }
0109
0110 ActsExamples::ProcessCode ActsExamples::RootPropagationStepsWriter::writeT(
0111     const AlgorithmContext& context, const PropagationSummaries& summaries) {
0112   // Exclusive access to the tree while writing
0113   std::lock_guard<std::mutex> lock(m_writeMutex);
0114
0115   // Get the event number
0116   m_eventNr = context.eventNumber;
0117
0118   // Initialize the last total trials
0119   // This is used to calculate the number of trials per step
0120   std::size_t lastTotalTrials = 0;
0121
0122   // Loop over the step vector of each test propagation in this
0123   for (const auto& [trackNr, summary] : Acts::enumerate(summaries)) {
0124     // Clear the vectors for each collection
0125     m_volumeID.clear();
0126     m_boundaryID.clear();
0127     m_layerID.clear();
0128     m_approachID.clear();
0129     m_sensitiveID.clear();
0130     m_extraID.clear();
0131     m_material.clear();
0132     m_x.clear();
0133     m_y.clear();
0134     m_z.clear();
0135     m_r.clear();
0136     m_dx.clear();
0137     m_dy.clear();
0138     m_dz.clear();
0139     m_step_type.clear();
0140     m_step_acc.clear();
0141     m_step_act.clear();
0142     m_step_abt.clear();
0143     m_step_usr.clear();
0144     m_nStepTrials.clear();
0145
0146     // Set the track number
0147     m_trackNr = static_cast<int>(trackNr);
0148
0149     // Set the initial trajectory parameters
0150     const auto& startParameters = summary.startParameters;
0151     m_d0 = static_cast<float>(startParameters.parameters()[Acts::eBoundLoc0]);
0152     m_z0 = static_cast<float>(startParameters.parameters()[Acts::eBoundLoc1]);
0153     m_phi = static_cast<float>(startParameters.parameters()[Acts::eBoundPhi]);
0154     m_theta =
0155         static_cast<float>(startParameters.parameters()[Acts::eBoundTheta]);
0156     m_qOverP =
0157         static_cast<float>(startParameters.parameters()[Acts::eBoundQOverP]);
0158     m_t = static_cast<float>(startParameters.parameters()[Acts::eBoundTime]);
0159
0160     // Derived initial trajectory parameters
0161     m_eta = static_cast<float>(
0162         Acts::VectorHelpers::eta(startParameters.direction()));
0163     m_pt = static_cast<float>(startParameters.transverseMomentum());
0164     m_p = static_cast<float>(startParameters.absoluteMomentum());
0165
0166     // Loop over single steps
0167     for (auto& step : summary.steps) {
0168       const auto& geoID = step.geoID;
0169       m_sensitiveID.push_back(geoID.sensitive());
0170       m_approachID.push_back(geoID.approach());
0171       m_layerID.push_back(geoID.layer());
0172       m_boundaryID.push_back(geoID.boundary());
0173       m_volumeID.push_back(geoID.volume());
0174       m_extraID.push_back(geoID.extra());
0175
0176       int material = 0;
0177       if (step.surface != nullptr &&
0178           step.surface->surfaceMaterial() != nullptr) {
0179         material = 1;
0180       }
0181       m_material.push_back(material);
0182
0183       // kinematic information
0184       m_x.push_back(step.position.x());
0185       m_y.push_back(step.position.y());
0186       m_z.push_back(step.position.z());
0187       m_r.push_back(Acts::VectorHelpers::perp(step.position));
0188       auto direction = step.momentum.normalized();
0189       m_dx.push_back(direction.x());
0190       m_dy.push_back(direction.y());
0191       m_dz.push_back(direction.z());
0192
0193       double accuracy = step.stepSize.accuracy();
0194       double actor =
0195           step.stepSize.value(Acts::ConstrainedStep::Type::Navigator);
0196       double aborter = step.stepSize.value(Acts::ConstrainedStep::Type::Actor);
0197       double user = step.stepSize.value(Acts::ConstrainedStep::Type::User);
0198       double actAbs = std::abs(actor);
0199       double accAbs = std::abs(accuracy);
0200       double aboAbs = std::abs(aborter);
0201       double usrAbs = std::abs(user);
0202
0203       if (actAbs < accAbs && actAbs < aboAbs && actAbs < usrAbs) {
0204         m_step_type.push_back(0);
0205       } else if (accAbs < aboAbs && accAbs < usrAbs) {
0206         m_step_type.push_back(1);
0207       } else if (aboAbs < usrAbs) {
0208         m_step_type.push_back(2);
0209       } else {
0210         m_step_type.push_back(3);
0211       }
0212
0213       // Step size information
0214       m_step_acc.push_back(Acts::clampValue<float>(accuracy));
0215       m_step_act.push_back(Acts::clampValue<float>(actor));
0216       m_step_abt.push_back(Acts::clampValue<float>(aborter));
0217       m_step_usr.push_back(Acts::clampValue<float>(user));
0218
0219       // Stepper efficiency
0220       m_nStepTrials.push_back(step.nTotalTrials - lastTotalTrials);
0221       lastTotalTrials = step.nTotalTrials;
0222     }
0223     m_outputTree->Fill();
0224   }
0225
0226   return ActsExamples::ProcessCode::SUCCESS;
0227 }