Root/src/RootSimHitReader.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/RootSimHitReader.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 <stdexcept>
0020
0021 #include <TChain.h>
0022 #include <TMathBase.h>
0023
0024 namespace ActsExamples {
0025
0026 RootSimHitReader::RootSimHitReader(const RootSimHitReader::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_outputSimHits.initialize(m_cfg.outputSimHits);
0041
0042   // add file to the input chain
0043   m_inputChain->Add(m_cfg.filePath.c_str());
0044   m_inputChain->LoadTree(0);
0045   ACTS_DEBUG("Adding File " << m_cfg.filePath << " to tree '" << m_cfg.treeName
0046                             << "'.");
0047
0048   // Set the branches
0049   auto setBranches = [&]<class T>(const auto& keys, T& columns) {
0050     using MappedType = typename std::remove_reference_t<T>::mapped_type;
0051     for (auto key : keys) {
0052       MappedType a{};  // 0 or nullptr
0053       columns.emplace(key, a);
0054     }
0055     for (auto key : keys) {
0056       m_inputChain->SetBranchAddress(key, &columns.at(key));
0057     }
0058   };
0059
0060   setBranches(m_floatKeys, m_floatColumns);
0061   setBranches(m_uint32Keys, m_uint32Columns);
0062   setBranches(m_uint64Keys, m_uint64Columns);
0063   setBranches(m_int32Keys, m_int32Columns);
0064
0065   if (m_inputChain->FindBranch("barcode") != nullptr) {
0066     m_hasBarcodeVector = true;
0067     m_barcodeVector.allocate();
0068     m_inputChain->SetBranchAddress("barcode", &m_barcodeVector.get());
0069   } else {
0070     m_hasBarcodeVector = false;
0071     for (const auto* key : m_barcodeComponentKeys) {
0072       if (!m_uint32Columns.contains(key)) {
0073         m_uint32Columns.emplace(key, std::uint32_t{0});
0074       }
0075       m_inputChain->SetBranchAddress(key, &m_uint32Columns.at(key));
0076     }
0077   }
0078
0079   // Because each hit is stored in a single entry in the root file, we need to
0080   // scan the file first for the positions of the events in the file in order to
0081   // efficiently read the events later on.
0082   // TODO change the file format to store one event per entry
0083
0084   // Disable all branches and only enable event-id for a first scan of the file
0085   m_inputChain->SetBranchStatus("*", false);
0086   m_inputChain->SetBranchStatus("event_id", true);
0087
0088   auto nEntries = static_cast<std::size_t>(m_inputChain->GetEntriesFast());
0089   if (nEntries == 0) {
0090     throw std::runtime_error("Did not find any entries in input file");
0091   }
0092
0093   // Add the first entry
0094   m_inputChain->GetEntry(0);
0095   m_eventMap.push_back({m_uint32Columns.at("event_id"), 0ul, 0ul});
0096
0097   // Go through all entries and store the position of new events
0098   for (auto i = 1ul; i < nEntries; ++i) {
0099     m_inputChain->GetEntry(i);
0100     const auto evtId = m_uint32Columns.at("event_id");
0101
0102     if (evtId != std::get<0>(m_eventMap.back())) {
0103       std::get<2>(m_eventMap.back()) = i;
0104       m_eventMap.push_back({evtId, i, i});
0105     }
0106   }
0107
0108   std::get<2>(m_eventMap.back()) = nEntries;
0109
0110   // Sort by event id
0111   std::ranges::sort(m_eventMap, {},
0112                     [](const auto& m) { return std::get<0>(m); });
0113
0114   // Re-Enable all branches
0115   m_inputChain->SetBranchStatus("*", true);
0116   ACTS_DEBUG("Event range: " << availableEvents().first << " - "
0117                              << availableEvents().second);
0118 }
0119
0120 std::pair<std::size_t, std::size_t> RootSimHitReader::availableEvents() const {
0121   return {std::get<0>(m_eventMap.front()), std::get<0>(m_eventMap.back()) + 1};
0122 }
0123
0124 ProcessCode RootSimHitReader::read(const AlgorithmContext& context) {
0125   auto it = std::ranges::find_if(m_eventMap, [&](const auto& a) {
0126     return std::get<0>(a) == context.eventNumber;
0127   });
0128
0129   if (it == m_eventMap.end()) {
0130     // explicitly warn if it happens for the first or last event as that might
0131     // indicate a human error
0132     if ((context.eventNumber == availableEvents().first) &&
0133         (context.eventNumber == availableEvents().second - 1)) {
0134       ACTS_WARNING("Reading empty event: " << context.eventNumber);
0135     } else {
0136       ACTS_DEBUG("Reading empty event: " << context.eventNumber);
0137     }
0138
0139     m_outputSimHits(context, {});
0140
0141     // Return success flag
0142     return ProcessCode::SUCCESS;
0143   }
0144
0145   // lock the mutex
0146   std::lock_guard<std::mutex> lock(m_read_mutex);
0147
0148   ACTS_DEBUG("Reading event: " << std::get<0>(*it)
0149                                << " stored in entries: " << std::get<1>(*it)
0150                                << " - " << std::get<2>(*it));
0151
0152   SimHitContainer hits;
0153   for (auto entry = std::get<1>(*it); entry < std::get<2>(*it); ++entry) {
0154     m_inputChain->GetEntry(entry);
0155
0156     auto eventId = m_uint32Columns.at("event_id");
0157     if (eventId != context.eventNumber) {
0158       break;
0159     }
0160
0161     const Acts::GeometryIdentifier geoid{m_uint64Columns.at("geometry_id")};
0162     SimBarcode pid = SimBarcode::Invalid();
0163     if (m_hasBarcodeVector && m_barcodeVector.hasValue()) {
0164       pid = SimBarcode().withData(*m_barcodeVector);
0165     } else {
0166       pid = SimBarcode()
0167                 .withVertexPrimary(static_cast<SimBarcode::PrimaryVertexId>(
0168                     m_uint32Columns.at("barcode_vertex_primary")))
0169                 .withVertexSecondary(static_cast<SimBarcode::SecondaryVertexId>(
0170                     m_uint32Columns.at("barcode_vertex_secondary")))
0171                 .withParticle(static_cast<SimBarcode::ParticleId>(
0172                     m_uint32Columns.at("barcode_particle")))
0173                 .withGeneration(static_cast<SimBarcode::GenerationId>(
0174                     m_uint32Columns.at("barcode_generation")))
0175                 .withSubParticle(static_cast<SimBarcode::SubParticleId>(
0176                     m_uint32Columns.at("barcode_sub_particle")));
0177     }
0178     const auto index = m_int32Columns.at("index");
0179
0180     const Acts::Vector4 pos4 = {
0181         m_floatColumns.at("tx") * Acts::UnitConstants::mm,
0182         m_floatColumns.at("ty") * Acts::UnitConstants::mm,
0183         m_floatColumns.at("tz") * Acts::UnitConstants::mm,
0184         m_floatColumns.at("tt") * Acts::UnitConstants::mm,
0185     };
0186
0187     const Acts::Vector4 before4 = {
0188         m_floatColumns.at("tpx") * Acts::UnitConstants::GeV,
0189         m_floatColumns.at("tpy") * Acts::UnitConstants::GeV,
0190         m_floatColumns.at("tpz") * Acts::UnitConstants::GeV,
0191         m_floatColumns.at("te") * Acts::UnitConstants::GeV,
0192     };
0193
0194     const Acts::Vector4 delta = {
0195         m_floatColumns.at("deltapx") * Acts::UnitConstants::GeV,
0196         m_floatColumns.at("deltapy") * Acts::UnitConstants::GeV,
0197         m_floatColumns.at("deltapz") * Acts::UnitConstants::GeV,
0198         m_floatColumns.at("deltae") * Acts::UnitConstants::GeV,
0199     };
0200
0201     SimHit hit(geoid, pid, pos4, before4, before4 + delta, index);
0202
0203     hits.insert(hit);
0204   }
0205
0206   ACTS_DEBUG("Read " << hits.size() << " hits for event "
0207                      << context.eventNumber);
0208
0209   m_outputSimHits(context, std::move(hits));
0210
0211   // Return success flag
0212   return ProcessCode::SUCCESS;
0213 }
0214
0215 RootSimHitReader::~RootSimHitReader() = default;
0216
0217 }  // namespace ActsExamples