Root/src/RootVertexReader.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 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/RootVertexReader.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 "ActsFatras/EventData/ProcessType.hpp"
0016
0017 #include <algorithm>
0018 #include <cstdint>
0019 #include <iostream>
0020 #include <stdexcept>
0021
0022 #include <TChain.h>
0023 #include <TMathBase.h>
0024
0025 namespace ActsExamples {
0026
0027 RootVertexReader::RootVertexReader(const RootVertexReader::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_outputVertices.initialize(m_cfg.outputVertices);
0042
0043   // Set the branches
0044   m_inputChain->SetBranchAddress("event_id", &m_eventId);
0045   m_inputChain->SetBranchAddress("vertex_id", &m_vertexId);
0046   m_inputChain->SetBranchAddress("process", &m_process);
0047   m_inputChain->SetBranchAddress("vx", &m_vx);
0048   m_inputChain->SetBranchAddress("vy", &m_vy);
0049   m_inputChain->SetBranchAddress("vz", &m_vz);
0050   m_inputChain->SetBranchAddress("vt", &m_vt);
0051   m_inputChain->SetBranchAddress("outgoing_particles", &m_outgoingParticles);
0052   m_inputChain->SetBranchAddress("vertex_primary", &m_vertexPrimary);
0053   m_inputChain->SetBranchAddress("vertex_secondary", &m_vertexSecondary);
0054   m_inputChain->SetBranchAddress("generation", &m_generation);
0055
0056   auto path = m_cfg.filePath;
0057
0058   // add file to the input chain
0059   m_inputChain->Add(path.c_str());
0060   ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.treeName << "'.");
0061
0062   m_events = m_inputChain->GetEntries();
0063   ACTS_DEBUG("The full chain has " << m_events << " entries.");
0064
0065   // Sort the entry numbers of the events
0066   {
0067     m_entryNumbers.resize(m_events);
0068     m_inputChain->Draw("event_id", "", "goff");
0069     // Sort to get the entry numbers of the ordered events
0070     TMath::Sort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
0071                 m_entryNumbers.data(), false);
0072   }
0073 }
0074
0075 std::pair<std::size_t, std::size_t> RootVertexReader::availableEvents() const {
0076   return {0u, m_events};
0077 }
0078
0079 RootVertexReader::~RootVertexReader() {
0080   delete m_vertexId;
0081   delete m_process;
0082   delete m_vx;
0083   delete m_vy;
0084   delete m_vz;
0085   delete m_vt;
0086   delete m_outgoingParticles;
0087   delete m_vertexPrimary;
0088   delete m_vertexSecondary;
0089   delete m_generation;
0090 }
0091
0092 ProcessCode RootVertexReader::read(const AlgorithmContext& context) {
0093   ACTS_DEBUG("Trying to read recorded vertices.");
0094
0095   if (m_inputChain == nullptr || context.eventNumber >= m_events) {
0096     return ProcessCode::SUCCESS;
0097   }
0098
0099   // lock the mutex
0100   std::lock_guard<std::mutex> lock(m_read_mutex);
0101   // now read
0102
0103   // The vertex collection to be filled
0104   SimVertexContainer vertices;
0105
0106   // Read the correct entry
0107   auto entry = m_entryNumbers.at(context.eventNumber);
0108   m_inputChain->GetEntry(entry);
0109   ACTS_DEBUG("Reading event: " << context.eventNumber
0110                                << " stored as entry: " << entry);
0111
0112   unsigned int nVertices = m_vertexId->size();
0113
0114   for (unsigned int i = 0; i < nVertices; i++) {
0115     SimVertex v;
0116
0117     v.id = (*m_vertexId)[i];
0118     v.process = static_cast<ActsFatras::ProcessType>((*m_process)[i]);
0119     v.position4 = Acts::Vector4((*m_vx)[i] * Acts::UnitConstants::mm,
0120                                 (*m_vy)[i] * Acts::UnitConstants::mm,
0121                                 (*m_vz)[i] * Acts::UnitConstants::mm,
0122                                 (*m_vt)[i] * Acts::UnitConstants::mm);
0123
0124     // TODO ingoing particles
0125
0126     for (auto& id : (*m_outgoingParticles)[i]) {
0127       v.outgoing.insert(static_cast<std::uint64_t>(id));
0128     }
0129
0130     vertices.insert(v);
0131   }
0132
0133   ACTS_DEBUG("Read " << vertices.size() << " vertices for event "
0134                      << context.eventNumber);
0135
0136   // Write the collections to the EventStore
0137   m_outputVertices(context, std::move(vertices));
0138
0139   // Return success flag
0140   return ProcessCode::SUCCESS;
0141 }
0142
0143 }  // namespace ActsExamples