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("process", &m_process.get());
0045   m_inputChain->SetBranchAddress("vx", &m_vx.get());
0046   m_inputChain->SetBranchAddress("vy", &m_vy.get());
0047   m_inputChain->SetBranchAddress("vz", &m_vz.get());
0048   m_inputChain->SetBranchAddress("vt", &m_vt.get());
0049   if (m_inputChain->GetBranch("incoming_particles") != nullptr) {
0050     m_hasCombinedIncoming = true;
0051     m_incomingParticles.allocate();
0052     m_inputChain->SetBranchAddress("incoming_particles",
0053                                    &m_incomingParticles.get());
0054   } else {
0055     m_hasCombinedIncoming = false;
0056     m_incomingParticlesVertexPrimary.allocate();
0057     m_incomingParticlesVertexSecondary.allocate();
0058     m_incomingParticlesParticle.allocate();
0059     m_incomingParticlesGeneration.allocate();
0060     m_incomingParticlesSubParticle.allocate();
0061     m_inputChain->SetBranchAddress("incoming_particles_vertex_primary",
0062                                    &m_incomingParticlesVertexPrimary.get());
0063     m_inputChain->SetBranchAddress("incoming_particles_vertex_secondary",
0064                                    &m_incomingParticlesVertexSecondary.get());
0065     m_inputChain->SetBranchAddress("incoming_particles_particle",
0066                                    &m_incomingParticlesParticle.get());
0067     m_inputChain->SetBranchAddress("incoming_particles_generation",
0068                                    &m_incomingParticlesGeneration.get());
0069     m_inputChain->SetBranchAddress("incoming_particles_sub_particle",
0070                                    &m_incomingParticlesSubParticle.get());
0071   }
0072
0073   if (m_inputChain->GetBranch("outgoing_particles") != nullptr) {
0074     m_hasCombinedOutgoing = true;
0075     m_outgoingParticles.allocate();
0076     m_inputChain->SetBranchAddress("outgoing_particles",
0077                                    &m_outgoingParticles.get());
0078   } else {
0079     m_hasCombinedOutgoing = false;
0080     m_outgoingParticlesVertexPrimary.allocate();
0081     m_outgoingParticlesVertexSecondary.allocate();
0082     m_outgoingParticlesParticle.allocate();
0083     m_outgoingParticlesGeneration.allocate();
0084     m_outgoingParticlesSubParticle.allocate();
0085     m_inputChain->SetBranchAddress("outgoing_particles_vertex_primary",
0086                                    &m_outgoingParticlesVertexPrimary.get());
0087     m_inputChain->SetBranchAddress("outgoing_particles_vertex_secondary",
0088                                    &m_outgoingParticlesVertexSecondary.get());
0089     m_inputChain->SetBranchAddress("outgoing_particles_particle",
0090                                    &m_outgoingParticlesParticle.get());
0091     m_inputChain->SetBranchAddress("outgoing_particles_generation",
0092                                    &m_outgoingParticlesGeneration.get());
0093     m_inputChain->SetBranchAddress("outgoing_particles_sub_particle",
0094                                    &m_outgoingParticlesSubParticle.get());
0095   }
0096   m_inputChain->SetBranchAddress("vertex_primary", &m_vertexPrimary.get());
0097   m_inputChain->SetBranchAddress("vertex_secondary", &m_vertexSecondary.get());
0098   m_inputChain->SetBranchAddress("generation", &m_generation.get());
0099
0100   auto path = m_cfg.filePath;
0101
0102   // add file to the input chain
0103   m_inputChain->Add(path.c_str());
0104   ACTS_DEBUG("Adding File " << path << " to tree '" << m_cfg.treeName << "'.");
0105
0106   m_events = m_inputChain->GetEntries();
0107   ACTS_DEBUG("The full chain has " << m_events << " entries.");
0108
0109   // Sort the entry numbers of the events
0110   {
0111     // necessary to guarantee that m_inputChain->GetV1() is valid for the
0112     // entire range
0113     m_inputChain->SetEstimate(m_events + 1);
0114
0115     m_entryNumbers.resize(m_events);
0116     m_inputChain->Draw("event_id", "", "goff");
0117     RootUtility::stableSort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
0118                             m_entryNumbers.data(), false);
0119   }
0120 }
0121
0122 std::pair<std::size_t, std::size_t> RootVertexReader::availableEvents() const {
0123   return {0u, m_events};
0124 }
0125
0126 ProcessCode RootVertexReader::read(const AlgorithmContext& context) {
0127   ACTS_DEBUG("Trying to read recorded vertices.");
0128
0129   if (m_inputChain == nullptr || context.eventNumber >= m_events) {
0130     return ProcessCode::SUCCESS;
0131   }
0132
0133   // lock the mutex
0134   std::lock_guard<std::mutex> lock(m_read_mutex);
0135   // now read
0136
0137   // The vertex collection to be filled
0138   SimVertexContainer vertices;
0139
0140   // Read the correct entry
0141   auto entry = m_entryNumbers.at(context.eventNumber);
0142   m_inputChain->GetEntry(entry);
0143   ACTS_DEBUG("Reading event: " << context.eventNumber
0144                                << " stored as entry: " << entry);
0145
0146   unsigned int nVertices = m_vx->size();
0147
0148   for (unsigned int i = 0; i < nVertices; i++) {
0149     SimVertex v;
0150
0151     v.id = SimVertexBarcode()
0152                .withVertexPrimary((*m_vertexPrimary)[i])
0153                .withVertexSecondary((*m_vertexSecondary)[i])
0154                .withGeneration((*m_generation)[i]);
0155
0156     v.process = static_cast<ActsFatras::ProcessType>((*m_process)[i]);
0157     v.position4 = Acts::Vector4((*m_vx)[i] * Acts::UnitConstants::mm,
0158                                 (*m_vy)[i] * Acts::UnitConstants::mm,
0159                                 (*m_vz)[i] * Acts::UnitConstants::mm,
0160                                 (*m_vt)[i] * Acts::UnitConstants::mm);
0161
0162     // incoming particles
0163     if (m_hasCombinedIncoming && m_incomingParticles.hasValue()) {
0164       for (auto& barcodeComponents : (*m_incomingParticles)[i]) {
0165         v.incoming.insert(SimBarcode().withData(barcodeComponents));
0166       }
0167     } else if (m_incomingParticlesVertexPrimary.hasValue()) {
0168       const auto& incomingPrimaries = (*m_incomingParticlesVertexPrimary)[i];
0169       const auto& incomingSecondaries =
0170           (*m_incomingParticlesVertexSecondary)[i];
0171       const auto& incomingParticles = (*m_incomingParticlesParticle)[i];
0172       const auto& incomingGenerations = (*m_incomingParticlesGeneration)[i];
0173       const auto& incomingSubParticles = (*m_incomingParticlesSubParticle)[i];
0174       for (std::size_t j = 0; j < incomingPrimaries.size(); ++j) {
0175         v.incoming.insert(
0176             SimBarcode()
0177                 .withVertexPrimary(static_cast<SimBarcode::PrimaryVertexId>(
0178                     incomingPrimaries[j]))
0179                 .withVertexSecondary(static_cast<SimBarcode::SecondaryVertexId>(
0180                     incomingSecondaries[j]))
0181                 .withParticle(
0182                     static_cast<SimBarcode::ParticleId>(incomingParticles[j]))
0183                 .withGeneration(static_cast<SimBarcode::GenerationId>(
0184                     incomingGenerations[j]))
0185                 .withSubParticle(static_cast<SimBarcode::SubParticleId>(
0186                     incomingSubParticles[j])));
0187       }
0188     }
0189
0190     // outgoing particles
0191     if (m_hasCombinedOutgoing && m_outgoingParticles.hasValue()) {
0192       for (auto& barcodeComponents : (*m_outgoingParticles)[i]) {
0193         v.outgoing.insert(SimBarcode().withData(barcodeComponents));
0194       }
0195     } else if (m_outgoingParticlesVertexPrimary.hasValue()) {
0196       const auto& outgoingPrimaries = (*m_outgoingParticlesVertexPrimary)[i];
0197       const auto& outgoingSecondaries =
0198           (*m_outgoingParticlesVertexSecondary)[i];
0199       const auto& outgoingParticles = (*m_outgoingParticlesParticle)[i];
0200       const auto& outgoingGenerations = (*m_outgoingParticlesGeneration)[i];
0201       const auto& outgoingSubParticles = (*m_outgoingParticlesSubParticle)[i];
0202       for (std::size_t j = 0; j < outgoingPrimaries.size(); ++j) {
0203         v.outgoing.insert(
0204             SimBarcode()
0205                 .withVertexPrimary(static_cast<SimBarcode::PrimaryVertexId>(
0206                     outgoingPrimaries[j]))
0207                 .withVertexSecondary(static_cast<SimBarcode::SecondaryVertexId>(
0208                     outgoingSecondaries[j]))
0209                 .withParticle(
0210                     static_cast<SimBarcode::ParticleId>(outgoingParticles[j]))
0211                 .withGeneration(static_cast<SimBarcode::GenerationId>(
0212                     outgoingGenerations[j]))
0213                 .withSubParticle(static_cast<SimBarcode::SubParticleId>(
0214                     outgoingSubParticles[j])));
0215       }
0216     }
0217
0218     vertices.insert(v);
0219   }
0220
0221   ACTS_DEBUG("Read " << vertices.size() << " vertices for event "
0222                      << context.eventNumber);
0223
0224   // Write the collections to the EventStore
0225   m_outputVertices(context, std::move(vertices));
0226
0227   // Return success flag
0228   return ProcessCode::SUCCESS;
0229 }
0230
0231 }  // namespace ActsExamples