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