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