Root/src/RootPropagationStepsWriter.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2017-2022 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 http://mozilla.org/MPL/2.0/.
0008
0009 #include "ActsExamples/Io/Root/RootPropagationStepsWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Utilities/Helpers.hpp"
0013 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0014 #include <Acts/Geometry/GeometryIdentifier.hpp>
0015 #include <Acts/Geometry/TrackingVolume.hpp>
0016 #include <Acts/Propagator/ConstrainedStep.hpp>
0017 #include <Acts/Surfaces/Surface.hpp>
0018 #include <Acts/Utilities/Helpers.hpp>
0019
0020 #include <ios>
0021 #include <memory>
0022 #include <ostream>
0023 #include <stdexcept>
0024
0025 #include <TFile.h>
0026 #include <TTree.h>
0027
0028 ActsExamples::RootPropagationStepsWriter::RootPropagationStepsWriter(
0029     const ActsExamples::RootPropagationStepsWriter::Config& cfg,
0030     Acts::Logging::Level level)
0031     : WriterT(cfg.collection, "RootPropagationStepsWriter", level),
0032       m_cfg(cfg),
0033       m_outputFile(cfg.rootFile) {
0034   // An input collection name and tree name must be specified
0035   if (m_cfg.collection.empty()) {
0036     throw std::invalid_argument("Missing input collection");
0037   } else if (m_cfg.treeName.empty()) {
0038     throw std::invalid_argument("Missing tree name");
0039   }
0040
0041   // Setup ROOT I/O
0042   if (m_outputFile == nullptr) {
0043     m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0044     if (m_outputFile == nullptr) {
0045       throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0046     }
0047   }
0048   m_outputFile->cd();
0049
0050   m_outputTree = new TTree(m_cfg.treeName.c_str(),
0051                            "TTree from RootPropagationStepsWriter");
0052   if (m_outputTree == nullptr) {
0053     throw std::bad_alloc();
0054   }
0055
0056   // Set the branches
0057   m_outputTree->Branch("event_nr", &m_eventNr);
0058   m_outputTree->Branch("volume_id", &m_volumeID);
0059   m_outputTree->Branch("boundary_id", &m_boundaryID);
0060   m_outputTree->Branch("layer_id", &m_layerID);
0061   m_outputTree->Branch("approach_id", &m_approachID);
0062   m_outputTree->Branch("sensitive_id", &m_sensitiveID);
0063   m_outputTree->Branch("material", &m_material);
0064   m_outputTree->Branch("g_x", &m_x);
0065   m_outputTree->Branch("g_y", &m_y);
0066   m_outputTree->Branch("g_z", &m_z);
0067   m_outputTree->Branch("d_x", &m_dx);
0068   m_outputTree->Branch("d_y", &m_dy);
0069   m_outputTree->Branch("d_z", &m_dz);
0070   m_outputTree->Branch("type", &m_step_type);
0071   m_outputTree->Branch("step_acc", &m_step_acc);
0072   m_outputTree->Branch("step_act", &m_step_act);
0073   m_outputTree->Branch("step_abt", &m_step_abt);
0074   m_outputTree->Branch("step_usr", &m_step_usr);
0075   m_outputTree->Branch("nStepTrials", &m_nStepTrials);
0076 }
0077
0078 ActsExamples::RootPropagationStepsWriter::~RootPropagationStepsWriter() {
0079   if (m_outputFile != nullptr) {
0080     m_outputFile->Close();
0081   }
0082 }
0083
0084 ActsExamples::ProcessCode ActsExamples::RootPropagationStepsWriter::finalize() {
0085   // Write the tree
0086   m_outputFile->cd();
0087   m_outputTree->Write();
0088   /// Close the file if it's yours
0089   if (m_cfg.rootFile == nullptr) {
0090     m_outputFile->Close();
0091   }
0092
0093   ACTS_VERBOSE("Wrote particles to tree '" << m_cfg.treeName << "' in '"
0094                                            << m_cfg.filePath << "'");
0095
0096   return ProcessCode::SUCCESS;
0097 }
0098
0099 ActsExamples::ProcessCode ActsExamples::RootPropagationStepsWriter::writeT(
0100     const AlgorithmContext& context,
0101     const std::vector<PropagationSteps>& stepCollection) {
0102   // Exclusive access to the tree while writing
0103   std::lock_guard<std::mutex> lock(m_writeMutex);
0104
0105   // Get the event number
0106   m_eventNr = context.eventNumber;
0107
0108   // Loop over the step vector of each test propagation in this
0109   for (auto& steps : stepCollection) {
0110     // Clear the vectors for each collection
0111     m_volumeID.clear();
0112     m_boundaryID.clear();
0113     m_layerID.clear();
0114     m_approachID.clear();
0115     m_sensitiveID.clear();
0116     m_material.clear();
0117     m_x.clear();
0118     m_y.clear();
0119     m_z.clear();
0120     m_dx.clear();
0121     m_dy.clear();
0122     m_dz.clear();
0123     m_step_type.clear();
0124     m_step_acc.clear();
0125     m_step_act.clear();
0126     m_step_abt.clear();
0127     m_step_usr.clear();
0128     m_nStepTrials.clear();
0129
0130     // Loop over single steps
0131     for (auto& step : steps) {
0132       const auto& geoID = step.geoID;
0133       m_sensitiveID.push_back(geoID.sensitive());
0134       m_approachID.push_back(geoID.approach());
0135       m_layerID.push_back(geoID.layer());
0136       m_boundaryID.push_back(geoID.boundary());
0137       m_volumeID.push_back(geoID.volume());
0138
0139       int material = 0;
0140       if (step.surface) {
0141         if (step.surface->surfaceMaterial() != nullptr) {
0142           material = 1;
0143         }
0144       }
0145       m_material.push_back(material);
0146
0147       // kinematic information
0148       m_x.push_back(step.position.x());
0149       m_y.push_back(step.position.y());
0150       m_z.push_back(step.position.z());
0151       auto direction = step.momentum.normalized();
0152       m_dx.push_back(direction.x());
0153       m_dy.push_back(direction.y());
0154       m_dz.push_back(direction.z());
0155
0156       double accuracy = step.stepSize.accuracy();
0157       double actor = step.stepSize.value(Acts::ConstrainedStep::actor);
0158       double aborter = step.stepSize.value(Acts::ConstrainedStep::aborter);
0159       double user = step.stepSize.value(Acts::ConstrainedStep::user);
0160       double actAbs = std::abs(actor);
0161       double accAbs = std::abs(accuracy);
0162       double aboAbs = std::abs(aborter);
0163       double usrAbs = std::abs(user);
0164
0165       // todo - fold with direction
0166       if (actAbs < accAbs && actAbs < aboAbs && actAbs < usrAbs) {
0167         m_step_type.push_back(0);
0168       } else if (accAbs < aboAbs && accAbs < usrAbs) {
0169         m_step_type.push_back(1);
0170       } else if (aboAbs < usrAbs) {
0171         m_step_type.push_back(2);
0172       } else {
0173         m_step_type.push_back(3);
0174       }
0175
0176       // Step size information
0177       m_step_acc.push_back(Acts::clampValue<float>(accuracy));
0178       m_step_act.push_back(Acts::clampValue<float>(actor));
0179       m_step_abt.push_back(Acts::clampValue<float>(aborter));
0180       m_step_usr.push_back(Acts::clampValue<float>(user));
0181
0182       // Stepper efficiency
0183       m_nStepTrials.push_back(step.stepSize.nStepTrials);
0184     }
0185     m_outputTree->Fill();
0186   }
0187   return ActsExamples::ProcessCode::SUCCESS;
0188 }