Root/src/RootAthenaDumpWriter.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/RootAthenaDumpWriter.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "Acts/Utilities/VectorHelpers.hpp"
0013 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0014 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0015
0016 #include <ios>
0017 #include <stdexcept>
0018 #include <unordered_map>
0019
0020 #include <TFile.h>
0021 #include <TTree.h>
0022
0023 namespace ActsExamples {
0024
0025 RootAthenaDumpWriter::RootAthenaDumpWriter(const Config& config,
0026                                            Acts::Logging::Level level)
0027     : IWriter(),
0028       m_cfg(config),
0029       m_logger(Acts::getDefaultLogger(name(), level)) {
0030   if (m_cfg.filePath.empty()) {
0031     throw std::invalid_argument("Missing file path");
0032   }
0033   if (m_cfg.treeName.empty()) {
0034     throw std::invalid_argument("Missing tree name");
0035   }
0036
0037   m_inputParticles.initialize(m_cfg.inputParticles);
0038   m_inputClusters.initialize(m_cfg.inputClusters);
0039   m_inputMeasurements.initialize(m_cfg.inputMeasurements);
0040   m_inputMeasParticleMap.initialize(m_cfg.inputMeasParticleMap);
0041   m_inputSpacePoints.initialize(m_cfg.inputSpacePoints);
0042
0043   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), "RECREATE");
0044   if (m_outputFile == nullptr) {
0045     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0046   }
0047   m_outputFile->cd();
0048   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0049   if (m_outputTree == nullptr) {
0050     throw std::bad_alloc();
0051   }
0052
0053   // Reserve capacity upfront so that data() pointers passed to ROOT branches
0054   // remain stable for the lifetime of the writer. The vectors are cleared and
0055   // refilled each event; clear() preserves capacity, keeping the pointer valid.
0056   m_partEventNumber.reserve(s_maxParticles);
0057   m_partBarcode.reserve(s_maxParticles);
0058   m_partPt.reserve(s_maxParticles);
0059   m_partEta.reserve(s_maxParticles);
0060   m_partPdgId.reserve(s_maxParticles);
0061   m_partVx.reserve(s_maxParticles);
0062   m_partVy.reserve(s_maxParticles);
0063   m_partVz.reserve(s_maxParticles);
0064
0065   m_clIndex.reserve(s_maxClusters);
0066   m_clBarrelEndcap.reserve(s_maxClusters);
0067   m_clEtaModule.reserve(s_maxClusters);
0068   m_clPhiModule.reserve(s_maxClusters);
0069   m_clModuleId.reserve(s_maxClusters);
0070   m_clX.reserve(s_maxClusters);
0071   m_clY.reserve(s_maxClusters);
0072   m_clZ.reserve(s_maxClusters);
0073
0074   m_spIndex.reserve(s_maxSpacePoints);
0075   m_spX.reserve(s_maxSpacePoints);
0076   m_spY.reserve(s_maxSpacePoints);
0077   m_spZ.reserve(s_maxSpacePoints);
0078   m_spCL1Index.reserve(s_maxSpacePoints);
0079   m_spCL2Index.reserve(s_maxSpacePoints);
0080   m_spIsOverlap.reserve(s_maxSpacePoints);
0081
0082   // Event scalar
0083   m_outputTree->Branch("event_number", &m_eventNumber);
0084
0085   // Particle branches – variable-length C-arrays, count driven by nPartEVT
0086   m_outputTree->Branch("nPartEVT", &m_nPartEVT, "nPartEVT/I");
0087   m_outputTree->Branch("Part_event_number", m_partEventNumber.data(),
0088                        "Part_event_number[nPartEVT]/I");
0089   m_outputTree->Branch("Part_barcode", m_partBarcode.data(),
0090                        "Part_barcode[nPartEVT]/I");
0091   m_outputTree->Branch("Part_pt", m_partPt.data(), "Part_pt[nPartEVT]/F");
0092   m_outputTree->Branch("Part_eta", m_partEta.data(), "Part_eta[nPartEVT]/F");
0093   m_outputTree->Branch("Part_pdg_id", m_partPdgId.data(),
0094                        "Part_pdg_id[nPartEVT]/I");
0095   m_outputTree->Branch("Part_vx", m_partVx.data(), "Part_vx[nPartEVT]/F");
0096   m_outputTree->Branch("Part_vy", m_partVy.data(), "Part_vy[nPartEVT]/F");
0097   m_outputTree->Branch("Part_vz", m_partVz.data(), "Part_vz[nPartEVT]/F");
0098
0099   // Cluster branches
0100   m_outputTree->Branch("nCL", &m_nCL, "nCL/I");
0101   m_outputTree->Branch("CLindex", m_clIndex.data(), "CLindex[nCL]/I");
0102   m_outputTree->Branch("CLhardware", &m_clHardware);
0103   m_outputTree->Branch("CLbarrel_endcap", m_clBarrelEndcap.data(),
0104                        "CLbarrel_endcap[nCL]/I");
0105   m_outputTree->Branch("CLeta_module", m_clEtaModule.data(),
0106                        "CLeta_module[nCL]/I");
0107   m_outputTree->Branch("CLphi_module", m_clPhiModule.data(),
0108                        "CLphi_module[nCL]/I");
0109   m_outputTree->Branch("CLmoduleID", m_clModuleId.data(), "CLmoduleID[nCL]/l");
0110   m_outputTree->Branch("CLx", m_clX.data(), "CLx[nCL]/D");
0111   m_outputTree->Branch("CLy", m_clY.data(), "CLy[nCL]/D");
0112   m_outputTree->Branch("CLz", m_clZ.data(), "CLz[nCL]/D");
0113   m_outputTree->Branch("CLlocal_cov", &m_clLocalCov);
0114   m_outputTree->Branch("CLparticleLink_eventIndex",
0115                        &m_clParticleLinkEventIndex);
0116   m_outputTree->Branch("CLparticleLink_barcode", &m_clParticleLinkBarcode);
0117
0118   // Space point branches
0119   m_outputTree->Branch("nSP", &m_nSP, "nSP/I");
0120   m_outputTree->Branch("SPindex", m_spIndex.data(), "SPindex[nSP]/I");
0121   m_outputTree->Branch("SPx", m_spX.data(), "SPx[nSP]/D");
0122   m_outputTree->Branch("SPy", m_spY.data(), "SPy[nSP]/D");
0123   m_outputTree->Branch("SPz", m_spZ.data(), "SPz[nSP]/D");
0124   m_outputTree->Branch("SPCL1_index", m_spCL1Index.data(),
0125                        "SPCL1_index[nSP]/I");
0126   m_outputTree->Branch("SPCL2_index", m_spCL2Index.data(),
0127                        "SPCL2_index[nSP]/I");
0128   m_outputTree->Branch("SPisOverlap", m_spIsOverlap.data(),
0129                        "SPisOverlap[nSP]/I");
0130 }
0131
0132 RootAthenaDumpWriter::~RootAthenaDumpWriter() {
0133   if (m_outputFile != nullptr) {
0134     m_outputFile->Close();
0135   }
0136 }
0137
0138 ProcessCode RootAthenaDumpWriter::finalize() {
0139   m_outputFile->cd();
0140   m_outputTree->Write();
0141   m_outputFile->Close();
0142   ACTS_VERBOSE("Wrote Athena dump to '" << m_cfg.filePath << "'");
0143   return ProcessCode::SUCCESS;
0144 }
0145
0146 ProcessCode RootAthenaDumpWriter::write(const AlgorithmContext& ctx) {
0147   const auto& particles = m_inputParticles(ctx);
0148   const auto& clusters = m_inputClusters(ctx);
0149   const auto& measurements = m_inputMeasurements(ctx);
0150   const auto& measPartMap = m_inputMeasParticleMap(ctx);
0151   const auto& spacePoints = m_inputSpacePoints(ctx);
0152
0153   std::lock_guard<std::mutex> lock(m_writeMutex);
0154
0155   if (particles.size() > s_maxParticles) {
0156     throw std::runtime_error(
0157         "RootAthenaDumpWriter: particle count exceeds maxParticles (" +
0158         std::to_string(particles.size()) + " > " +
0159         std::to_string(s_maxParticles) + ")");
0160   }
0161   if (measurements.size() > s_maxClusters) {
0162     throw std::runtime_error(
0163         "RootAthenaDumpWriter: measurement count exceeds maxClusters (" +
0164         std::to_string(measurements.size()) + " > " +
0165         std::to_string(s_maxClusters) + ")");
0166   }
0167   if (spacePoints.size() > s_maxSpacePoints) {
0168     throw std::runtime_error(
0169         "RootAthenaDumpWriter: space point count exceeds maxSpacePoints (" +
0170         std::to_string(spacePoints.size()) + " > " +
0171         std::to_string(s_maxSpacePoints) + ")");
0172   }
0173
0174   m_eventNumber = ctx.eventNumber;
0175
0176   // Build barcode map: assign sequential Athena barcodes per event.
0177   // Primaries (vertexSecondary == 0) get barcodes [1, s_maxBarcodeForPrimary].
0178   // Secondaries get barcodes starting at s_maxBarcodeForPrimary + 1.
0179   // The same values are used in the particle table and cluster particle links
0180   // so that make_particle_id(subevent, barcode) is consistent throughout.
0181   std::unordered_map<ActsFatras::Barcode, int> barcodeMap;
0182   barcodeMap.reserve(particles.size());
0183   int primaryCount = 1;
0184   int secondaryCount = s_maxBarcodeForPrimary + 1;
0185   for (const auto& particle : particles) {
0186     const bool isPrimary = particle.particleId().vertexSecondary() == 0;
0187     barcodeMap[particle.particleId()] =
0188         isPrimary ? primaryCount++ : secondaryCount++;
0189   }
0190
0191   // --- Particles ---
0192   m_partEventNumber.clear();
0193   m_partBarcode.clear();
0194   m_partPt.clear();
0195   m_partEta.clear();
0196   m_partPdgId.clear();
0197   m_partVx.clear();
0198   m_partVy.clear();
0199   m_partVz.clear();
0200
0201   for (const auto& particle : particles) {
0202     const float pT = static_cast<float>(particle.transverseMomentum() /
0203                                         Acts::UnitConstants::MeV);
0204     const float eta = static_cast<float>(
0205         Acts::VectorHelpers::eta(particle.momentum().normalized()));
0206
0207     m_partEventNumber.push_back(
0208         static_cast<int>(particle.particleId().vertexPrimary()));
0209     m_partBarcode.push_back(barcodeMap.at(particle.particleId()));
0210     m_partPt.push_back(pT);
0211     m_partEta.push_back(eta);
0212     m_partPdgId.push_back(static_cast<int>(particle.pdg()));
0213     m_partVx.push_back(
0214         static_cast<float>(particle.position().x() / Acts::UnitConstants::mm));
0215     m_partVy.push_back(
0216         static_cast<float>(particle.position().y() / Acts::UnitConstants::mm));
0217     m_partVz.push_back(
0218         static_cast<float>(particle.position().z() / Acts::UnitConstants::mm));
0219   }
0220   m_nPartEVT = static_cast<int>(m_partBarcode.size());
0221
0222   // --- Clusters ---
0223   m_clIndex.clear();
0224   m_clHardware.clear();
0225   m_clBarrelEndcap.clear();
0226   m_clEtaModule.clear();
0227   m_clPhiModule.clear();
0228   m_clModuleId.clear();
0229   m_clX.clear();
0230   m_clY.clear();
0231   m_clZ.clear();
0232   m_clLocalCov.clear();
0233   m_clParticleLinkEventIndex.clear();
0234   m_clParticleLinkBarcode.clear();
0235
0236   for (std::size_t im = 0; im < measurements.size(); ++im) {
0237     const auto meas = measurements.at(im);
0238
0239     const bool isPixel = (meas.size() == 2);
0240
0241     m_clIndex.push_back(static_cast<int>(im));
0242     m_clHardware.push_back(isPixel ? "PIXEL" : "STRIP");
0243     m_clModuleId.push_back(meas.geometryId().value());
0244
0245     m_clBarrelEndcap.push_back(s_sentinel);
0246     m_clEtaModule.push_back(s_sentinel);
0247     m_clPhiModule.push_back(s_sentinel);
0248
0249     const Acts::Vector3& pos = clusters.at(im).globalPosition;
0250     m_clX.push_back(pos.x());
0251     m_clY.push_back(pos.y());
0252     m_clZ.push_back(pos.z());
0253
0254     if (isPixel) {
0255       m_clLocalCov.push_back({meas.covariance()(0, 0), meas.covariance()(0, 1),
0256                               meas.covariance()(1, 0),
0257                               meas.covariance()(1, 1)});
0258     } else {
0259       const double var = meas.covariance()(0, 0);
0260       m_clLocalCov.push_back({var, var});
0261     }
0262
0263     std::vector<int> eventIndices;
0264     std::vector<int> barcodes;
0265     const auto [begin, end] = measPartMap.equal_range(im);
0266     for (auto it = begin; it != end; ++it) {
0267       const auto& bc = it->second;
0268       eventIndices.push_back(s_sentinel);
0269       const auto bcIt = barcodeMap.find(bc);
0270       barcodes.push_back(bcIt != barcodeMap.end() ? bcIt->second : 0);
0271     }
0272     m_clParticleLinkEventIndex.push_back(std::move(eventIndices));
0273     m_clParticleLinkBarcode.push_back(std::move(barcodes));
0274   }
0275   m_nCL = static_cast<int>(m_clIndex.size());
0276
0277   // --- Space points ---
0278   m_spIndex.clear();
0279   m_spX.clear();
0280   m_spY.clear();
0281   m_spZ.clear();
0282   m_spCL1Index.clear();
0283   m_spCL2Index.clear();
0284   m_spIsOverlap.clear();
0285
0286   for (const auto& sp : spacePoints) {
0287     const auto sLinks = sp.sourceLinks();
0288     if (sLinks.empty()) {
0289       ACTS_WARNING("Space point with no source links, skipping");
0290       continue;
0291     }
0292
0293     m_spIndex.push_back(static_cast<int>(m_spIndex.size()));
0294     m_spX.push_back(static_cast<double>(sp.x()));
0295     m_spY.push_back(static_cast<double>(sp.y()));
0296     m_spZ.push_back(static_cast<double>(sp.z()));
0297
0298     m_spCL1Index.push_back(
0299         static_cast<int>(sLinks[0].get<IndexSourceLink>().index()));
0300     m_spCL2Index.push_back(
0301         sLinks.size() > 1
0302             ? static_cast<int>(sLinks[1].get<IndexSourceLink>().index())
0303             : -1);
0304
0305     // ACTS simulation does not produce overlapping SPs
0306     m_spIsOverlap.push_back(0);
0307   }
0308   m_nSP = static_cast<int>(m_spIndex.size());
0309
0310   m_outputTree->Fill();
0311
0312   ACTS_DEBUG("Wrote event " << m_eventNumber << " with " << m_nPartEVT
0313                             << " particles, " << m_nCL << " clusters, and "
0314                             << m_nSP << " space points");
0315
0316   return ProcessCode::SUCCESS;
0317 }
0318
0319 }  // namespace ActsExamples