Root/src/RootParticleReader.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2017-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/RootParticleReader.hpp"
0010
0011 #include "Acts/Definitions/PdgParticle.hpp"
0012 #include "Acts/Utilities/Logger.hpp"
0013 #include "ActsExamples/EventData/SimParticle.hpp"
0014 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0015 #include "ActsExamples/Io/Root/RootUtility.hpp"
0016 #include "ActsFatras/EventData/ProcessType.hpp"
0017
0018 #include <algorithm>
0019 #include <cstdint>
0020 #include <iostream>
0021 #include <stdexcept>
0022
0023 #include <TChain.h>
0024
0025 namespace ActsExamples {
0026
0027 RootParticleReader::RootParticleReader(const RootParticleReader::Config& config,
0028                                        Acts::Logging::Level level)
0029     : IReader(),
0030       m_cfg(config),
0031       m_logger(Acts::getDefaultLogger(name(), level)) {
0032   m_inputChain = new TChain(m_cfg.treeName.c_str());
0033
0034   if (m_cfg.filePath.empty()) {
0035     throw std::invalid_argument("Missing input filename");
0036   }
0037   if (m_cfg.treeName.empty()) {
0038     throw std::invalid_argument("Missing tree name");
0039   }
0040
0041   m_outputParticles.initialize(m_cfg.outputParticles);
0042
0043   // Set the branches
0044   m_inputChain->SetBranchAddress("event_id", &m_eventId);
0045   m_inputChain->SetBranchAddress("particle_id", &m_particleId);
0046   m_inputChain->SetBranchAddress("particle_type", &m_particleType);
0047   m_inputChain->SetBranchAddress("process", &m_process);
0048   m_inputChain->SetBranchAddress("vx", &m_vx);
0049   m_inputChain->SetBranchAddress("vy", &m_vy);
0050   m_inputChain->SetBranchAddress("vz", &m_vz);
0051   m_inputChain->SetBranchAddress("vt", &m_vt);
0052   m_inputChain->SetBranchAddress("p", &m_p);
0053   m_inputChain->SetBranchAddress("px", &m_px);
0054   m_inputChain->SetBranchAddress("py", &m_py);
0055   m_inputChain->SetBranchAddress("pz", &m_pz);
0056   m_inputChain->SetBranchAddress("m", &m_m);
0057   m_inputChain->SetBranchAddress("q", &m_q);
0058   m_inputChain->SetBranchAddress("eta", &m_eta);
0059   m_inputChain->SetBranchAddress("phi", &m_phi);
0060   m_inputChain->SetBranchAddress("pt", &m_pt);
0061   m_inputChain->SetBranchAddress("vertex_primary", &m_vertexPrimary);
0062   m_inputChain->SetBranchAddress("vertex_secondary", &m_vertexSecondary);
0063   m_inputChain->SetBranchAddress("particle", &m_particle);
0064   m_inputChain->SetBranchAddress("generation", &m_generation);
0065   m_inputChain->SetBranchAddress("sub_particle", &m_subParticle);
0066
0067   auto path = m_cfg.filePath;
0068
0069   // add file to the input chain
0070   m_inputChain->Add(path.c_str());
0071   ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.treeName << "'.");
0072
0073   m_events = m_inputChain->GetEntries();
0074   ACTS_DEBUG("The full chain has " << m_events << " entries.");
0075
0076   // Sort the entry numbers of the events
0077   {
0078     m_entryNumbers.resize(m_events);
0079     m_inputChain->Draw("event_id", "", "goff");
0080     RootUtility::stableSort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
0081                             m_entryNumbers.data(), false);
0082   }
0083 }
0084
0085 std::pair<std::size_t, std::size_t> RootParticleReader::availableEvents()
0086     const {
0087   return {0u, m_events};
0088 }
0089
0090 RootParticleReader::~RootParticleReader() {
0091   delete m_particleId;
0092   delete m_particleType;
0093   delete m_process;
0094   delete m_vx;
0095   delete m_vy;
0096   delete m_vz;
0097   delete m_vt;
0098   delete m_p;
0099   delete m_px;
0100   delete m_py;
0101   delete m_pz;
0102   delete m_m;
0103   delete m_q;
0104   delete m_eta;
0105   delete m_phi;
0106   delete m_pt;
0107   delete m_vertexPrimary;
0108   delete m_vertexSecondary;
0109   delete m_particle;
0110   delete m_generation;
0111   delete m_subParticle;
0112 }
0113
0114 ProcessCode RootParticleReader::read(const AlgorithmContext& context) {
0115   ACTS_DEBUG("Trying to read recorded particles.");
0116
0117   if (m_inputChain == nullptr || context.eventNumber >= m_events) {
0118     return ProcessCode::SUCCESS;
0119   }
0120
0121   // lock the mutex
0122   std::lock_guard<std::mutex> lock(m_read_mutex);
0123   // now read
0124
0125   // The particle collection to be filled
0126   SimParticleContainer particles;
0127
0128   // Read the correct entry
0129   auto entry = m_entryNumbers.at(context.eventNumber);
0130   m_inputChain->GetEntry(entry);
0131   ACTS_DEBUG("Reading event: " << context.eventNumber
0132                                << " stored as entry: " << entry);
0133
0134   unsigned int nParticles = m_particleId->size();
0135
0136   for (unsigned int i = 0; i < nParticles; i++) {
0137     SimParticle p;
0138
0139     p.setProcess(static_cast<ActsFatras::ProcessType>((*m_process)[i]));
0140     p.setPdg(static_cast<Acts::PdgParticle>((*m_particleType)[i]));
0141     p.setCharge((*m_q)[i] * Acts::UnitConstants::e);
0142     p.setMass((*m_m)[i] * Acts::UnitConstants::GeV);
0143     p.setParticleId((*m_particleId)[i]);
0144     p.setPosition4((*m_vx)[i] * Acts::UnitConstants::mm,
0145                    (*m_vy)[i] * Acts::UnitConstants::mm,
0146                    (*m_vz)[i] * Acts::UnitConstants::mm,
0147                    (*m_vt)[i] * Acts::UnitConstants::mm);
0148     // NOTE: direction is normalized inside `setDirection`
0149     p.setDirection((*m_px)[i], (*m_py)[i], (*m_pz)[i]);
0150     p.setAbsoluteMomentum((*m_p)[i] * Acts::UnitConstants::GeV);
0151
0152     particles.insert(p);
0153   }
0154
0155   ACTS_DEBUG("Read " << particles.size() << " particles for event "
0156                      << context.eventNumber);
0157
0158   // Write the collections to the EventStore
0159   m_outputParticles(context, std::move(particles));
0160
0161   // Return success flag
0162   return ProcessCode::SUCCESS;
0163 }
0164
0165 }  // namespace ActsExamples