Root/src/RootSimHitReader.cpp

0001 // This file is part of the Acts project.
0002 //
0003 // Copyright (C) 2023 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/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 = 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_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
0076   // Add the first entry
0077   m_inputChain->GetEntry(0);
0078   m_eventMap.push_back({m_uint32Columns.at("event_id"), 0ul, 0ul});
0079
0080   // Go through all entries and store the position of new events
0081   for (auto i = 1ul; i < nEntries; ++i) {
0082     m_inputChain->GetEntry(i);
0083     const auto evtId = m_uint32Columns.at("event_id");
0084
0085     if (evtId != std::get<0>(m_eventMap.back())) {
0086       std::get<2>(m_eventMap.back()) = i;
0087       m_eventMap.push_back({evtId, i, i});
0088     }
0089   }
0090
0091   std::get<2>(m_eventMap.back()) = nEntries;
0092
0093   // Sort by event id
0094   std::sort(m_eventMap.begin(), m_eventMap.end(),
0095             [](const auto& a, const auto& b) {
0096               return std::get<0>(a) < std::get<0>(b);
0097             });
0098
0099   // Re-Enable all branches
0100   m_inputChain->SetBranchStatus("*", true);
0101   ACTS_DEBUG("Event range: " << availableEvents().first << " - "
0102                              << availableEvents().second);
0103 }
0104
0105 std::pair<std::size_t, std::size_t> RootSimHitReader::availableEvents() const {
0106   return {std::get<0>(m_eventMap.front()), std::get<0>(m_eventMap.back()) + 1};
0107 }
0108
0109 ProcessCode RootSimHitReader::read(const AlgorithmContext& context) {
0110   auto it = std::find_if(
0111       m_eventMap.begin(), m_eventMap.end(),
0112       [&](const auto& a) { return std::get<0>(a) == context.eventNumber; });
0113
0114   if (it == m_eventMap.end()) {
0115     // explicitly warn if it happens for the first or last event as that might
0116     // indicate a human error
0117     if ((context.eventNumber == availableEvents().first) &&
0118         (context.eventNumber == availableEvents().second - 1)) {
0119       ACTS_WARNING("Reading empty event: " << context.eventNumber);
0120     } else {
0121       ACTS_DEBUG("Reading empty event: " << context.eventNumber);
0122     }
0123
0124     m_outputSimHits(context, {});
0125
0126     // Return success flag
0127     return ProcessCode::SUCCESS;
0128   }
0129
0130   // lock the mutex
0131   std::lock_guard<std::mutex> lock(m_read_mutex);
0132
0133   ACTS_DEBUG("Reading event: " << std::get<0>(*it)
0134                                << " stored in entries: " << std::get<1>(*it)
0135                                << " - " << std::get<2>(*it));
0136
0137   SimHitContainer hits;
0138   for (auto entry = std::get<1>(*it); entry < std::get<2>(*it); ++entry) {
0139     m_inputChain->GetEntry(entry);
0140
0141     auto eventId = m_uint32Columns.at("event_id");
0142     if (eventId != context.eventNumber) {
0143       break;
0144     }
0145
0146     const Acts::GeometryIdentifier geoid = m_uint64Columns.at("geometry_id");
0147     const SimBarcode pid = m_uint64Columns.at("particle_id");
0148     const auto index = m_int32Columns.at("index");
0149
0150     const Acts::Vector4 pos4 = {
0151         m_floatColumns.at("tx") * Acts::UnitConstants::mm,
0152         m_floatColumns.at("ty") * Acts::UnitConstants::mm,
0153         m_floatColumns.at("tz") * Acts::UnitConstants::mm,
0154         m_floatColumns.at("tt") * Acts::UnitConstants::mm,
0155     };
0156
0157     const Acts::Vector4 before4 = {
0158         m_floatColumns.at("tpx") * Acts::UnitConstants::GeV,
0159         m_floatColumns.at("tpy") * Acts::UnitConstants::GeV,
0160         m_floatColumns.at("tpz") * Acts::UnitConstants::GeV,
0161         m_floatColumns.at("te") * Acts::UnitConstants::GeV,
0162     };
0163
0164     const Acts::Vector4 delta = {
0165         m_floatColumns.at("deltapx") * Acts::UnitConstants::GeV,
0166         m_floatColumns.at("deltapy") * Acts::UnitConstants::GeV,
0167         m_floatColumns.at("deltapz") * Acts::UnitConstants::GeV,
0168         m_floatColumns.at("deltae") * Acts::UnitConstants::GeV,
0169     };
0170
0171     SimHit hit(geoid, pid, pos4, before4, before4 + delta, index);
0172
0173     hits.insert(hit);
0174   }
0175
0176   ACTS_DEBUG("Read " << hits.size() << " hits for event "
0177                      << context.eventNumber);
0178
0179   m_outputSimHits(context, std::move(hits));
0180
0181   // Return success flag
0182   return ProcessCode::SUCCESS;
0183 }
0184
0185 }  // namespace ActsExamples