Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-01-18 09:11:56

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/RootMaterialTrackWriter.hpp"
0010 
0011 #include "Acts/Geometry/GeometryIdentifier.hpp"
0012 #include "Acts/Geometry/TrackingVolume.hpp"
0013 #include "Acts/Material/Material.hpp"
0014 #include "Acts/Material/MaterialInteraction.hpp"
0015 #include "Acts/Material/MaterialSlab.hpp"
0016 #include "Acts/Surfaces/CylinderBounds.hpp"
0017 #include "Acts/Surfaces/RadialBounds.hpp"
0018 #include "Acts/Surfaces/Surface.hpp"
0019 #include "Acts/Surfaces/SurfaceBounds.hpp"
0020 #include "Acts/Utilities/Intersection.hpp"
0021 #include "Acts/Utilities/Logger.hpp"
0022 #include "Acts/Utilities/VectorHelpers.hpp"
0023 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0024 
0025 #include <cstddef>
0026 #include <ios>
0027 #include <stdexcept>
0028 
0029 #include <TFile.h>
0030 #include <TTree.h>
0031 
0032 using Acts::VectorHelpers::eta;
0033 using Acts::VectorHelpers::perp;
0034 using Acts::VectorHelpers::phi;
0035 
0036 namespace ActsExamples {
0037 
0038 RootMaterialTrackWriter::RootMaterialTrackWriter(
0039     const RootMaterialTrackWriter::Config& config, Acts::Logging::Level level)
0040     : WriterT(config.inputMaterialTracks, "RootMaterialTrackWriter", level),
0041       m_cfg(config) {
0042   // An input collection name and tree name must be specified
0043   if (m_cfg.inputMaterialTracks.empty()) {
0044     throw std::invalid_argument("Missing input collection");
0045   } else if (m_cfg.treeName.empty()) {
0046     throw std::invalid_argument("Missing tree name");
0047   }
0048 
0049   // Setup ROOT I/O
0050   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0051   if (m_outputFile == nullptr) {
0052     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0053   }
0054 
0055   m_outputFile->cd();
0056   m_outputTree =
0057       new TTree(m_cfg.treeName.c_str(), "TTree from RootMaterialTrackWriter");
0058   if (m_outputTree == nullptr) {
0059     throw std::bad_alloc();
0060   }
0061 
0062   // Set the branches
0063   m_outputTree->Branch("event_id", &m_eventId);
0064   m_outputTree->Branch("v_x", &m_v_x);
0065   m_outputTree->Branch("v_y", &m_v_y);
0066   m_outputTree->Branch("v_z", &m_v_z);
0067   m_outputTree->Branch("v_px", &m_v_px);
0068   m_outputTree->Branch("v_py", &m_v_py);
0069   m_outputTree->Branch("v_pz", &m_v_pz);
0070   m_outputTree->Branch("v_phi", &m_v_phi);
0071   m_outputTree->Branch("v_eta", &m_v_eta);
0072   m_outputTree->Branch("t_X0", &m_tX0);
0073   m_outputTree->Branch("t_L0", &m_tL0);
0074   m_outputTree->Branch("mat_x", &m_step_x);
0075   m_outputTree->Branch("mat_y", &m_step_y);
0076   m_outputTree->Branch("mat_z", &m_step_z);
0077   m_outputTree->Branch("mat_r", &m_step_r);
0078   m_outputTree->Branch("mat_dx", &m_step_dx);
0079   m_outputTree->Branch("mat_dy", &m_step_dy);
0080   m_outputTree->Branch("mat_dz", &m_step_dz);
0081   m_outputTree->Branch("mat_step_length", &m_step_length);
0082   m_outputTree->Branch("mat_X0", &m_step_X0);
0083   m_outputTree->Branch("mat_L0", &m_step_L0);
0084   m_outputTree->Branch("mat_A", &m_step_A);
0085   m_outputTree->Branch("mat_Z", &m_step_Z);
0086   m_outputTree->Branch("mat_rho", &m_step_rho);
0087 
0088   if (m_cfg.prePostStep) {
0089     m_outputTree->Branch("mat_sx", &m_step_sx);
0090     m_outputTree->Branch("mat_sy", &m_step_sy);
0091     m_outputTree->Branch("mat_sz", &m_step_sz);
0092     m_outputTree->Branch("mat_ex", &m_step_ex);
0093     m_outputTree->Branch("mat_ey", &m_step_ey);
0094     m_outputTree->Branch("mat_ez", &m_step_ez);
0095   }
0096   if (m_cfg.storeSurface) {
0097     m_outputTree->Branch("sur_id", &m_sur_id);
0098     m_outputTree->Branch("sur_type", &m_sur_type);
0099     m_outputTree->Branch("sur_x", &m_sur_x);
0100     m_outputTree->Branch("sur_y", &m_sur_y);
0101     m_outputTree->Branch("sur_z", &m_sur_z);
0102     m_outputTree->Branch("sur_r", &m_sur_r);
0103     m_outputTree->Branch("sur_distance", &m_sur_distance);
0104     m_outputTree->Branch("sur_pathCorrection", &m_sur_pathCorrection);
0105     m_outputTree->Branch("sur_range_min", &m_sur_range_min);
0106     m_outputTree->Branch("sur_range_max", &m_sur_range_max);
0107   }
0108   if (m_cfg.storeVolume) {
0109     m_outputTree->Branch("vol_id", &m_vol_id);
0110   }
0111 }
0112 
0113 RootMaterialTrackWriter::~RootMaterialTrackWriter() {
0114   if (m_outputFile != nullptr) {
0115     m_outputFile->Close();
0116   }
0117 }
0118 
0119 ProcessCode RootMaterialTrackWriter::finalize() {
0120   // write the tree and close the file
0121   ACTS_INFO("Writing ROOT output File : " << m_cfg.filePath);
0122 
0123   m_outputFile->cd();
0124   m_outputTree->Write();
0125   m_outputFile->Close();
0126 
0127   return ProcessCode::SUCCESS;
0128 }
0129 
0130 ProcessCode RootMaterialTrackWriter::writeT(
0131     const AlgorithmContext& ctx,
0132     const std::unordered_map<std::size_t, Acts::RecordedMaterialTrack>&
0133         materialTracks) {
0134   // Exclusive access to the tree while writing
0135   std::lock_guard<std::mutex> lock(m_writeMutex);
0136 
0137   m_eventId = ctx.eventNumber;
0138   // Loop over the material tracks and write them out
0139   for (auto& [idTrack, mtrack] : materialTracks) {
0140     // Clearing the vector first
0141     m_step_sx.clear();
0142     m_step_sy.clear();
0143     m_step_sz.clear();
0144     m_step_x.clear();
0145     m_step_y.clear();
0146     m_step_z.clear();
0147     m_step_r.clear();
0148     m_step_ex.clear();
0149     m_step_ey.clear();
0150     m_step_ez.clear();
0151     m_step_dx.clear();
0152     m_step_dy.clear();
0153     m_step_dz.clear();
0154     m_step_length.clear();
0155     m_step_X0.clear();
0156     m_step_L0.clear();
0157     m_step_A.clear();
0158     m_step_Z.clear();
0159     m_step_rho.clear();
0160 
0161     m_sur_id.clear();
0162     m_sur_type.clear();
0163     m_sur_x.clear();
0164     m_sur_y.clear();
0165     m_sur_z.clear();
0166     m_sur_r.clear();
0167     m_sur_distance.clear();
0168     m_sur_pathCorrection.clear();
0169     m_sur_range_min.clear();
0170     m_sur_range_max.clear();
0171 
0172     m_vol_id.clear();
0173 
0174     auto materialInteractions = mtrack.second.materialInteractions;
0175     if (m_cfg.collapseInteractions) {
0176       std::vector<Acts::MaterialInteraction> collapsed;
0177 
0178       Acts::Vector3 positionSum = Acts::Vector3::Zero();
0179       double pathCorrectionSum = 0;
0180 
0181       for (std::size_t start = 0, end = 0; end < materialInteractions.size();
0182            ++end) {
0183         const auto& mintStart = materialInteractions[start];
0184         const auto& mintEnd = materialInteractions[end];
0185 
0186         positionSum += mintEnd.position;
0187         pathCorrectionSum += mintEnd.pathCorrection;
0188 
0189         const bool same = mintStart.materialSlab.material() ==
0190                           mintEnd.materialSlab.material();
0191         const bool last = end == materialInteractions.size() - 1;
0192 
0193         if (!same || last) {
0194           auto mint = mintStart;
0195           mint.position = positionSum / (end - start);
0196           mint.pathCorrection = pathCorrectionSum;
0197 
0198           collapsed.push_back(mint);
0199 
0200           start = end;
0201           positionSum = Acts::Vector3::Zero();
0202           pathCorrectionSum = 0;
0203         }
0204       }
0205 
0206       materialInteractions = std::move(collapsed);
0207     }
0208 
0209     // Reserve the vector then
0210     std::size_t mints = materialInteractions.size();
0211     m_step_sx.reserve(mints);
0212     m_step_sy.reserve(mints);
0213     m_step_sz.reserve(mints);
0214     m_step_x.reserve(mints);
0215     m_step_y.reserve(mints);
0216     m_step_z.reserve(mints);
0217     m_step_r.reserve(mints);
0218     m_step_ex.reserve(mints);
0219     m_step_ey.reserve(mints);
0220     m_step_ez.reserve(mints);
0221     m_step_dx.reserve(mints);
0222     m_step_dy.reserve(mints);
0223     m_step_dz.reserve(mints);
0224     m_step_length.reserve(mints);
0225     m_step_X0.reserve(mints);
0226     m_step_L0.reserve(mints);
0227     m_step_A.reserve(mints);
0228     m_step_Z.reserve(mints);
0229     m_step_rho.reserve(mints);
0230 
0231     m_sur_id.reserve(mints);
0232     m_sur_type.reserve(mints);
0233     m_sur_x.reserve(mints);
0234     m_sur_y.reserve(mints);
0235     m_sur_z.reserve(mints);
0236     m_sur_r.reserve(mints);
0237     m_sur_distance.reserve(mints);
0238     m_sur_pathCorrection.reserve(mints);
0239     m_sur_range_min.reserve(mints);
0240     m_sur_range_max.reserve(mints);
0241 
0242     m_vol_id.reserve(mints);
0243 
0244     // reset the global counter
0245     if (m_cfg.recalculateTotals) {
0246       m_tX0 = 0.;
0247       m_tL0 = 0.;
0248     } else {
0249       m_tX0 = mtrack.second.materialInX0;
0250       m_tL0 = mtrack.second.materialInL0;
0251     }
0252 
0253     // set the track information at vertex
0254     m_v_x = mtrack.first.first.x();
0255     m_v_y = mtrack.first.first.y();
0256     m_v_z = mtrack.first.first.z();
0257     m_v_px = mtrack.first.second.x();
0258     m_v_py = mtrack.first.second.y();
0259     m_v_pz = mtrack.first.second.z();
0260     m_v_phi = phi(mtrack.first.second);
0261     m_v_eta = eta(mtrack.first.second);
0262 
0263     // and now loop over the material
0264     for (const auto& mint : materialInteractions) {
0265       auto direction = mint.direction.normalized();
0266 
0267       // The material step position information
0268       m_step_x.push_back(mint.position.x());
0269       m_step_y.push_back(mint.position.y());
0270       m_step_z.push_back(mint.position.z());
0271       m_step_r.push_back(perp(mint.position));
0272       m_step_dx.push_back(direction.x());
0273       m_step_dy.push_back(direction.y());
0274       m_step_dz.push_back(direction.z());
0275 
0276       if (m_cfg.prePostStep) {
0277         Acts::Vector3 prePos =
0278             mint.position - 0.5 * mint.pathCorrection * direction;
0279         Acts::Vector3 posPos =
0280             mint.position + 0.5 * mint.pathCorrection * direction;
0281 
0282         m_step_sx.push_back(prePos.x());
0283         m_step_sy.push_back(prePos.y());
0284         m_step_sz.push_back(prePos.z());
0285         m_step_ex.push_back(posPos.x());
0286         m_step_ey.push_back(posPos.y());
0287         m_step_ez.push_back(posPos.z());
0288       }
0289 
0290       // Store surface information
0291       if (m_cfg.storeSurface) {
0292         const Acts::Surface* surface = mint.surface;
0293         if (mint.intersectionID.value() != 0) {
0294           m_sur_id.push_back(mint.intersectionID.value());
0295           m_sur_pathCorrection.push_back(mint.pathCorrection);
0296           m_sur_x.push_back(mint.intersection.x());
0297           m_sur_y.push_back(mint.intersection.y());
0298           m_sur_z.push_back(mint.intersection.z());
0299           m_sur_r.push_back(perp(mint.intersection));
0300           m_sur_distance.push_back((mint.position - mint.intersection).norm());
0301         } else if (surface != nullptr) {
0302           auto sfIntersection =
0303               surface
0304                   ->intersect(ctx.geoContext, mint.position, mint.direction,
0305                               Acts::BoundaryTolerance::None())
0306                   .closest();
0307           m_sur_id.push_back(surface->geometryId().value());
0308           m_sur_pathCorrection.push_back(1.0);
0309           m_sur_x.push_back(sfIntersection.position().x());
0310           m_sur_y.push_back(sfIntersection.position().y());
0311           m_sur_z.push_back(sfIntersection.position().z());
0312         } else {
0313           m_sur_id.push_back(Acts::GeometryIdentifier().value());
0314           m_sur_x.push_back(0);
0315           m_sur_y.push_back(0);
0316           m_sur_z.push_back(0);
0317           m_sur_pathCorrection.push_back(1.0);
0318         }
0319         if (surface != nullptr) {
0320           m_sur_type.push_back(surface->type());
0321           const Acts::SurfaceBounds& surfaceBounds = surface->bounds();
0322           const Acts::RadialBounds* radialBounds =
0323               dynamic_cast<const Acts::RadialBounds*>(&surfaceBounds);
0324           const Acts::CylinderBounds* cylinderBounds =
0325               dynamic_cast<const Acts::CylinderBounds*>(&surfaceBounds);
0326           if (radialBounds != nullptr) {
0327             m_sur_range_min.push_back(radialBounds->rMin());
0328             m_sur_range_max.push_back(radialBounds->rMax());
0329           } else if (cylinderBounds != nullptr) {
0330             m_sur_range_min.push_back(
0331                 -cylinderBounds->get(Acts::CylinderBounds::eHalfLengthZ));
0332             m_sur_range_max.push_back(
0333                 cylinderBounds->get(Acts::CylinderBounds::eHalfLengthZ));
0334           } else {
0335             m_sur_range_min.push_back(0);
0336             m_sur_range_max.push_back(0);
0337           }
0338         } else {
0339           m_sur_type.push_back(-1);
0340           m_sur_range_min.push_back(0);
0341           m_sur_range_max.push_back(0);
0342         }
0343       }
0344 
0345       // store volume information
0346       if (m_cfg.storeVolume) {
0347         Acts::GeometryIdentifier vlayerID;
0348         if (!mint.volume.empty()) {
0349           vlayerID = mint.volume.geometryId();
0350           m_vol_id.push_back(vlayerID.value());
0351         } else {
0352           vlayerID.setVolume(0);
0353           vlayerID.setBoundary(0);
0354           vlayerID.setLayer(0);
0355           vlayerID.setApproach(0);
0356           vlayerID.setSensitive(0);
0357           m_vol_id.push_back(vlayerID.value());
0358         }
0359       }
0360 
0361       // the material information
0362       const auto& mprops = mint.materialSlab;
0363       m_step_length.push_back(mprops.thickness());
0364       m_step_X0.push_back(mprops.material().X0());
0365       m_step_L0.push_back(mprops.material().L0());
0366       m_step_A.push_back(mprops.material().Ar());
0367       m_step_Z.push_back(mprops.material().Z());
0368       m_step_rho.push_back(mprops.material().massDensity());
0369       // re-calculate if defined to do so
0370       if (m_cfg.recalculateTotals) {
0371         m_tX0 += mprops.thicknessInX0();
0372         m_tL0 += mprops.thicknessInL0();
0373       }
0374     }
0375     // write to
0376     m_outputTree->Fill();
0377   }
0378 
0379   // return success
0380   return ProcessCode::SUCCESS;
0381 }
0382 
0383 }  // namespace ActsExamples