Root/src/RootParticleReader.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/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/GenerationProcess.hpp"
0017 #include "ActsFatras/EventData/SimulationOutcome.hpp"
0018
0019 #include <iostream>
0020 #include <stdexcept>
0021
0022 #include <TChain.h>
0023
0024 namespace ActsExamples {
0025
0026 RootParticleReader::RootParticleReader(const RootParticleReader::Config& config,
0027                                        Acts::Logging::Level level)
0028     : IReader(),
0029       m_cfg(config),
0030       m_logger(Acts::getDefaultLogger(name(), level)) {
0031   m_inputChain = std::make_unique<TChain>(m_cfg.treeName.c_str());
0032
0033   if (m_cfg.filePath.empty()) {
0034     throw std::invalid_argument("Missing input filename");
0035   }
0036   if (m_cfg.treeName.empty()) {
0037     throw std::invalid_argument("Missing tree name");
0038   }
0039
0040   m_outputParticles.initialize(m_cfg.outputParticles);
0041
0042   // Set the branches
0043   m_inputChain->SetBranchAddress("event_id", &m_eventId);
0044   m_inputChain->SetBranchAddress("particle_hash", &m_particleHash.get());
0045   m_inputChain->SetBranchAddress("particle_type", &m_particleType.get());
0046   m_inputChain->SetBranchAddress("process", &m_process.get());
0047   m_inputChain->SetBranchAddress("vx", &m_vx.get());
0048   m_inputChain->SetBranchAddress("vy", &m_vy.get());
0049   m_inputChain->SetBranchAddress("vz", &m_vz.get());
0050   m_inputChain->SetBranchAddress("vt", &m_vt.get());
0051   m_inputChain->SetBranchAddress("p", &m_p.get());
0052   m_inputChain->SetBranchAddress("px", &m_px.get());
0053   m_inputChain->SetBranchAddress("py", &m_py.get());
0054   m_inputChain->SetBranchAddress("pz", &m_pz.get());
0055   m_inputChain->SetBranchAddress("m", &m_m.get());
0056   m_inputChain->SetBranchAddress("q", &m_q.get());
0057   m_inputChain->SetBranchAddress("eta", &m_eta.get());
0058   m_inputChain->SetBranchAddress("phi", &m_phi.get());
0059   m_inputChain->SetBranchAddress("pt", &m_pt.get());
0060   m_inputChain->SetBranchAddress("vertex_primary", &m_vertexPrimary.get());
0061   m_inputChain->SetBranchAddress("vertex_secondary", &m_vertexSecondary.get());
0062   m_inputChain->SetBranchAddress("particle", &m_particle.get());
0063   m_inputChain->SetBranchAddress("generation", &m_generation.get());
0064   m_inputChain->SetBranchAddress("sub_particle", &m_subParticle.get());
0065   m_inputChain->SetBranchAddress("orig_part_idx", &m_origParticleIdx.get());
0066   m_inputChain->SetBranchAddress("hf_origin", &m_hfOrigin.get());
0067
0068   m_inputChain->SetBranchAddress("e_loss", &m_eLoss.get());
0069   m_inputChain->SetBranchAddress("total_x0", &m_pathInX0.get());
0070   m_inputChain->SetBranchAddress("total_l0", &m_pathInL0.get());
0071   m_inputChain->SetBranchAddress("number_of_hits", &m_numberOfHits.get());
0072   m_inputChain->SetBranchAddress("outcome", &m_outcome.get());
0073
0074   auto path = m_cfg.filePath;
0075
0076   // add file to the input chain
0077   m_inputChain->Add(path.c_str());
0078   ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.treeName << "'.");
0079
0080   m_events = m_inputChain->GetEntries();
0081   ACTS_DEBUG("The full chain has " << m_events << " entries.");
0082
0083   // Sort the entry numbers of the events
0084   {
0085     // necessary to guarantee that m_inputChain->GetV1() is valid for the
0086     // entire range
0087     m_inputChain->SetEstimate(m_events + 1);
0088
0089     m_entryNumbers.resize(m_events);
0090     m_inputChain->Draw("event_id", "", "goff");
0091     RootUtility::stableSort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
0092                             m_entryNumbers.data(), false);
0093   }
0094 }
0095
0096 std::pair<std::size_t, std::size_t> RootParticleReader::availableEvents()
0097     const {
0098   return {0u, m_events};
0099 }
0100
0101 ProcessCode RootParticleReader::read(const AlgorithmContext& context) {
0102   ACTS_DEBUG("Trying to read recorded particles.");
0103
0104   if (m_inputChain == nullptr || context.eventNumber >= m_events) {
0105     return ProcessCode::SUCCESS;
0106   }
0107
0108   // lock the mutex
0109   std::lock_guard<std::mutex> lock(m_read_mutex);
0110   // now read
0111
0112   // The particle collection to be filled
0113   SimParticleContainer particles;
0114
0115   // Read the correct entry
0116   auto entry = m_entryNumbers.at(context.eventNumber);
0117   m_inputChain->GetEntry(entry);
0118   ACTS_DEBUG("Reading event: " << context.eventNumber
0119                                << " stored as entry: " << entry);
0120
0121   unsigned int nParticles = m_particleType->size();
0122
0123   for (unsigned int i = 0; i < nParticles; i++) {
0124     SimParticle p;
0125
0126     p.setProcess(
0127         static_cast<ActsFatras::GenerationProcess>((*m_process).at(i)));
0128     p.setPdg(static_cast<Acts::PdgParticle>((*m_particleType).at(i)));
0129     p.setCharge((*m_q).at(i) * Acts::UnitConstants::e);
0130     p.setMass((*m_m).at(i) * Acts::UnitConstants::GeV);
0131     p.setParticleId(SimBarcode()
0132                         .withVertexPrimary((*m_vertexPrimary).at(i))
0133                         .withVertexSecondary((*m_vertexSecondary).at(i))
0134                         .withParticle((*m_particle).at(i))
0135                         .withGeneration((*m_generation).at(i))
0136                         .withSubParticle((*m_subParticle).at(i)));
0137
0138     p.setOrigParticleIdx((*m_origParticleIdx).at(i));
0139     p.setHeavyFlavourOrigin(
0140         static_cast<HeavyFlavourOrigin>((*m_hfOrigin).at(i)));
0141
0142     SimParticleState& initialState = p.initialState();
0143
0144     initialState.setPosition4((*m_vx).at(i) * Acts::UnitConstants::mm,
0145                               (*m_vy).at(i) * Acts::UnitConstants::mm,
0146                               (*m_vz).at(i) * Acts::UnitConstants::mm,
0147                               (*m_vt).at(i) * Acts::UnitConstants::mm);
0148     // NOTE: direction is normalized inside `setDirection`
0149     initialState.setDirection((*m_px).at(i), (*m_py).at(i), (*m_pz).at(i));
0150     initialState.setAbsoluteMomentum((*m_p).at(i) * Acts::UnitConstants::GeV);
0151
0152     SimParticleState& finalState = p.finalState();
0153
0154     // TODO eloss cannot be read since we need the final momentum
0155     finalState.setMaterialPassed((*m_pathInX0).at(i) * Acts::UnitConstants::mm,
0156                                  (*m_pathInL0).at(i) * Acts::UnitConstants::mm);
0157     finalState.setNumberOfHits((*m_numberOfHits).at(i));
0158     finalState.setOutcome(
0159         static_cast<ActsFatras::SimulationOutcome>((*m_outcome).at(i)));
0160
0161     particles.insert(p);
0162   }
0163
0164   ACTS_DEBUG("Read " << particles.size() << " particles for event "
0165                      << context.eventNumber);
0166
0167   // Write the collections to the EventStore
0168   m_outputParticles(context, std::move(particles));
0169
0170   // Return success flag
0171   return ProcessCode::SUCCESS;
0172 }
0173
0174 }  // namespace ActsExamples