Root/src/RootParticleWriter.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2023-2024 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/RootParticleWriter.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Utilities/Helpers.hpp"
0013 #include "Acts/Utilities/VectorHelpers.hpp"
0014 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0015 #include "ActsFatras/EventData/Barcode.hpp"
0016 #include "ActsFatras/EventData/Particle.hpp"
0017
0018 #include <cstdint>
0019 #include <ios>
0020 #include <limits>
0021 #include <ostream>
0022 #include <stdexcept>
0023
0024 #include <TFile.h>
0025 #include <TTree.h>
0026
0027 ActsExamples::RootParticleWriter::RootParticleWriter(
0028     const ActsExamples::RootParticleWriter::Config& cfg,
0029     Acts::Logging::Level lvl)
0030     : WriterT(cfg.inputParticles, "RootParticleWriter", lvl), m_cfg(cfg) {
0031   // inputParticles is already checked by base constructor
0032   if (m_cfg.filePath.empty()) {
0033     throw std::invalid_argument("Missing file path");
0034   }
0035   if (m_cfg.treeName.empty()) {
0036     throw std::invalid_argument("Missing tree name");
0037   }
0038
0039   m_inputFinalParticles.maybeInitialize(m_cfg.inputFinalParticles);
0040
0041   // open root file and create the tree
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   m_outputFile->cd();
0047   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0048   if (m_outputTree == nullptr) {
0049     throw std::bad_alloc();
0050   }
0051
0052   // setup the branches
0053   m_outputTree->Branch("event_id", &m_eventId);
0054   m_outputTree->Branch("particle_id", &m_particleId);
0055   m_outputTree->Branch("particle_type", &m_particleType);
0056   m_outputTree->Branch("process", &m_process);
0057   m_outputTree->Branch("vx", &m_vx);
0058   m_outputTree->Branch("vy", &m_vy);
0059   m_outputTree->Branch("vz", &m_vz);
0060   m_outputTree->Branch("vt", &m_vt);
0061   m_outputTree->Branch("px", &m_px);
0062   m_outputTree->Branch("py", &m_py);
0063   m_outputTree->Branch("pz", &m_pz);
0064   m_outputTree->Branch("m", &m_m);
0065   m_outputTree->Branch("q", &m_q);
0066   m_outputTree->Branch("eta", &m_eta);
0067   m_outputTree->Branch("phi", &m_phi);
0068   m_outputTree->Branch("pt", &m_pt);
0069   m_outputTree->Branch("p", &m_p);
0070   m_outputTree->Branch("vertex_primary", &m_vertexPrimary);
0071   m_outputTree->Branch("vertex_secondary", &m_vertexSecondary);
0072   m_outputTree->Branch("particle", &m_particle);
0073   m_outputTree->Branch("generation", &m_generation);
0074   m_outputTree->Branch("sub_particle", &m_subParticle);
0075
0076   if (m_inputFinalParticles.isInitialized()) {
0077     m_outputTree->Branch("e_loss", &m_eLoss);
0078     m_outputTree->Branch("total_x0", &m_pathInX0);
0079     m_outputTree->Branch("total_l0", &m_pathInL0);
0080     m_outputTree->Branch("number_of_hits", &m_numberOfHits);
0081     m_outputTree->Branch("outcome", &m_outcome);
0082   }
0083 }
0084
0085 ActsExamples::RootParticleWriter::~RootParticleWriter() {
0086   if (m_outputFile != nullptr) {
0087     m_outputFile->Close();
0088   }
0089 }
0090
0091 ActsExamples::ProcessCode ActsExamples::RootParticleWriter::finalize() {
0092   m_outputFile->cd();
0093   m_outputTree->Write();
0094   m_outputFile->Close();
0095
0096   ACTS_INFO("Wrote particles to tree '" << m_cfg.treeName << "' in '"
0097                                         << m_cfg.filePath << "'");
0098
0099   return ProcessCode::SUCCESS;
0100 }
0101
0102 ActsExamples::ProcessCode ActsExamples::RootParticleWriter::writeT(
0103     const AlgorithmContext& ctx, const SimParticleContainer& particles) {
0104   const SimParticleContainer* finalParticles = nullptr;
0105   if (m_inputFinalParticles.isInitialized()) {
0106     finalParticles = &m_inputFinalParticles(ctx);
0107   }
0108
0109   // ensure exclusive access to tree/file while writing
0110   std::lock_guard<std::mutex> lock(m_writeMutex);
0111
0112   auto nan = std::numeric_limits<float>::quiet_NaN();
0113
0114   m_eventId = ctx.eventNumber;
0115   for (const auto& particle : particles) {
0116     m_particleId.push_back(particle.particleId().value());
0117     m_particleType.push_back(particle.pdg());
0118     m_process.push_back(static_cast<uint32_t>(particle.process()));
0119     // position
0120     m_vx.push_back(Acts::clampValue<float>(particle.fourPosition().x() /
0121                                            Acts::UnitConstants::mm));
0122     m_vy.push_back(Acts::clampValue<float>(particle.fourPosition().y() /
0123                                            Acts::UnitConstants::mm));
0124     m_vz.push_back(Acts::clampValue<float>(particle.fourPosition().z() /
0125                                            Acts::UnitConstants::mm));
0126     m_vt.push_back(Acts::clampValue<float>(particle.fourPosition().w() /
0127                                            Acts::UnitConstants::mm));
0128     // momentum
0129     const auto p = particle.absoluteMomentum() / Acts::UnitConstants::GeV;
0130     m_p.push_back(Acts::clampValue<float>(p));
0131     m_px.push_back(Acts::clampValue<float>(p * particle.direction().x()));
0132     m_py.push_back(Acts::clampValue<float>(p * particle.direction().y()));
0133     m_pz.push_back(Acts::clampValue<float>(p * particle.direction().z()));
0134     // particle constants
0135     m_m.push_back(
0136         Acts::clampValue<float>(particle.mass() / Acts::UnitConstants::GeV));
0137     m_q.push_back(
0138         Acts::clampValue<float>(particle.charge() / Acts::UnitConstants::e));
0139     // derived kinematic quantities
0140     m_eta.push_back(Acts::clampValue<float>(
0141         Acts::VectorHelpers::eta(particle.direction())));
0142     m_phi.push_back(Acts::clampValue<float>(
0143         Acts::VectorHelpers::phi(particle.direction())));
0144     m_pt.push_back(Acts::clampValue<float>(
0145         p * Acts::VectorHelpers::perp(particle.direction())));
0146     // decoded barcode components
0147     m_vertexPrimary.push_back(particle.particleId().vertexPrimary());
0148     m_vertexSecondary.push_back(particle.particleId().vertexSecondary());
0149     m_particle.push_back(particle.particleId().particle());
0150     m_generation.push_back(particle.particleId().generation());
0151     m_subParticle.push_back(particle.particleId().subParticle());
0152
0153     bool wroteFinalParticle = false;
0154     if (finalParticles != nullptr) {
0155       // get the final particle
0156       auto it = finalParticles->find(particle);
0157       if (it == finalParticles->end()) {
0158         ACTS_ERROR("Could not find final particle for "
0159                    << particle.particleId() << " in event " << ctx.eventNumber);
0160       } else {
0161         const auto& finalParticle = *it;
0162         // get the energy loss
0163         m_eLoss.push_back(Acts::clampValue<float>(
0164             (particle.energy() - finalParticle.energy()) /
0165             Acts::UnitConstants::GeV));
0166         // get the path in X0
0167         m_pathInX0.push_back(Acts::clampValue<float>(finalParticle.pathInX0() /
0168                                                      Acts::UnitConstants::mm));
0169         // get the path in L0
0170         m_pathInL0.push_back(Acts::clampValue<float>(finalParticle.pathInL0() /
0171                                                      Acts::UnitConstants::mm));
0172         // get the number of hits
0173         m_numberOfHits.push_back(finalParticle.numberOfHits());
0174         // get the particle outcome
0175         m_outcome.push_back(
0176             static_cast<std::uint32_t>(finalParticle.outcome()));
0177
0178         wroteFinalParticle = true;
0179       }
0180     }
0181     if (!wroteFinalParticle) {
0182       m_eLoss.push_back(nan);
0183       m_pathInX0.push_back(nan);
0184       m_pathInL0.push_back(nan);
0185       m_numberOfHits.push_back(-1);
0186       m_outcome.push_back(0);
0187     }
0188   }
0189
0190   m_outputTree->Fill();
0191
0192   m_particleId.clear();
0193   m_particleType.clear();
0194   m_process.clear();
0195   m_vx.clear();
0196   m_vy.clear();
0197   m_vz.clear();
0198   m_vt.clear();
0199   m_p.clear();
0200   m_px.clear();
0201   m_py.clear();
0202   m_pz.clear();
0203   m_m.clear();
0204   m_q.clear();
0205   m_eta.clear();
0206   m_phi.clear();
0207   m_pt.clear();
0208   m_vertexPrimary.clear();
0209   m_vertexSecondary.clear();
0210   m_particle.clear();
0211   m_generation.clear();
0212   m_subParticle.clear();
0213   m_eLoss.clear();
0214   m_numberOfHits.clear();
0215   m_outcome.clear();
0216   m_pathInX0.clear();
0217   m_pathInL0.clear();
0218
0219   return ProcessCode::SUCCESS;
0220 }