Root/src/RootTrackFinderNTupleWriter.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/RootTrackFinderNTupleWriter.hpp"
0010
0011 #include "Acts/Definitions/Units.hpp"
0012 #include "ActsExamples/EventData/Measurement.hpp"
0013 #include "ActsExamples/EventData/SimParticle.hpp"
0014 #include "ActsExamples/EventData/TruthMatching.hpp"
0015 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0016 #include "ActsExamples/Framework/DataHandle.hpp"
0017 #include "ActsExamples/Utilities/Range.hpp"
0018 #include "ActsExamples/Validation/TrackClassification.hpp"
0019 #include "ActsFatras/EventData/Barcode.hpp"
0020
0021 #include <cstddef>
0022 #include <cstdint>
0023 #include <mutex>
0024 #include <stdexcept>
0025 #include <unordered_map>
0026 #include <utility>
0027 #include <vector>
0028
0029 #include <RtypesCore.h>
0030 #include <TFile.h>
0031 #include <TTree.h>
0032
0033 namespace ActsExamples {
0034
0035 struct RootTrackFinderNTupleWriter::Impl {
0036   Config cfg;
0037
0038   ReadDataHandle<SimParticleContainer> inputParticles;
0039   ReadDataHandle<ParticleMeasurementsMap> inputParticleMeasurementsMap;
0040   ReadDataHandle<TrackParticleMatching> inputTrackParticleMatching;
0041
0042   TFile* file = nullptr;
0043
0044   // per-track tree
0045   TTree* trkTree = nullptr;
0046   std::mutex trkMutex;
0047   // track identification
0048   ULong64_t trkEventId = 0;
0049   ULong64_t trkTrackId = 0;
0050   // track content
0051   // number of hits on track
0052   UShort_t trkNumHits = 0;
0053   // number of particles contained in the track
0054   UShort_t trkNumParticles = 0;
0055   // track particle content; for each contributing particle, largest first
0056   std::vector<std::uint32_t> trkParticleVertexPrimary;
0057   std::vector<std::uint32_t> trkParticleVertexSecondary;
0058   std::vector<std::uint32_t> trkParticleParticle;
0059   std::vector<std::uint32_t> trkParticleGeneration;
0060   std::vector<std::uint32_t> trkParticleSubParticle;
0061   // total number of hits generated by this particle
0062   std::vector<UShort_t> trkParticleNumHitsTotal;
0063   // number of hits within this track
0064   std::vector<UShort_t> trkParticleNumHitsOnTrack;
0065
0066   // per-particle tree
0067   TTree* prtTree = nullptr;
0068   std::mutex prtMutex;
0069   // particle identification
0070   ULong64_t prtEventId = 0;
0071   std::uint32_t prtParticleVertexPrimary = 0;
0072   std::uint32_t prtParticleVertexSecondary = 0;
0073   std::uint32_t prtParticleParticle = 0;
0074   std::uint32_t prtParticleGeneration = 0;
0075   std::uint32_t prtParticleSubParticle = 0;
0076   Int_t prtParticleType = 0;
0077   // particle kinematics
0078   // vertex position in mm
0079   float prtVx = 0, prtVy = 0, prtVz = 0;
0080   // vertex time in ns
0081   float prtVt = 0;
0082   // particle momentum at production in GeV
0083   float prtPx = 0, prtPy = 0, prtPz = 0;
0084   // particle mass in GeV
0085   float prtM = 0;
0086   // particle charge in e
0087   float prtQ = 0;
0088   // particle reconstruction
0089   UShort_t prtNumMeasurements = 0;  // number of hits for this particle
0090   UShort_t prtNumTracks =
0091       0;  // number of tracks this particle was reconstructed in
0092   UShort_t prtNumTracksMajority =
0093       0;  // number of tracks reconstructed as majority
0094   // extra logger reference for the logging macros
0095   const Acts::Logger& _logger;
0096
0097   Impl(RootTrackFinderNTupleWriter* parent, Config&& c, const Acts::Logger& l)
0098       : cfg(std::move(c)),
0099         inputParticles{parent, "InputParticles"},
0100         inputParticleMeasurementsMap{parent, "InputParticleMeasurementsMap"},
0101         inputTrackParticleMatching{parent, "InputTrackParticleMatching"},
0102         _logger(l) {
0103     if (cfg.inputTracks.empty()) {
0104       throw std::invalid_argument("Missing track input collection");
0105     }
0106     if (cfg.inputParticles.empty()) {
0107       throw std::invalid_argument("Missing particles input collection");
0108     }
0109     if (cfg.inputParticleMeasurementsMap.empty()) {
0110       throw std::invalid_argument(
0111           "Missing particle-measurements map input collection");
0112     }
0113     if (cfg.inputTrackParticleMatching.empty()) {
0114       throw std::invalid_argument(
0115           "Missing proto track-particle matching input collection");
0116     }
0117     if (cfg.filePath.empty()) {
0118       throw std::invalid_argument("Missing output filename");
0119     }
0120
0121     inputParticles.initialize(cfg.inputParticles);
0122     inputParticleMeasurementsMap.initialize(cfg.inputParticleMeasurementsMap);
0123     inputTrackParticleMatching.initialize(cfg.inputTrackParticleMatching);
0124
0125     // the output file can not be given externally since TFile accesses to the
0126     // same file from multiple threads are unsafe.
0127     // must always be opened internally
0128     file = TFile::Open(cfg.filePath.c_str(), cfg.fileMode.c_str());
0129     if (file == nullptr) {
0130       throw std::invalid_argument("Could not open '" + cfg.filePath + "'");
0131     }
0132
0133     // construct trees
0134     trkTree = new TTree(cfg.treeNameTracks.c_str(), cfg.treeNameTracks.c_str());
0135     trkTree->SetDirectory(file);
0136     trkTree->Branch("event_id", &trkEventId);
0137     trkTree->Branch("track_id", &trkTrackId);
0138     trkTree->Branch("size", &trkNumHits);
0139     trkTree->Branch("nparticles", &trkNumParticles);
0140     trkTree->Branch("particle_id_vertex_primary", &trkParticleVertexPrimary);
0141     trkTree->Branch("particle_id_vertex_secondary",
0142                     &trkParticleVertexSecondary);
0143     trkTree->Branch("particle_id_particle", &trkParticleParticle);
0144     trkTree->Branch("particle_id_generation", &trkParticleGeneration);
0145     trkTree->Branch("particle_id_sub_particle", &trkParticleSubParticle);
0146     trkTree->Branch("particle_nhits_total", &trkParticleNumHitsTotal);
0147     trkTree->Branch("particle_nhits_on_track", &trkParticleNumHitsOnTrack);
0148     prtTree =
0149         new TTree(cfg.treeNameParticles.c_str(), cfg.treeNameParticles.c_str());
0150     prtTree->SetDirectory(file);
0151     prtTree->Branch("event_id", &prtEventId);
0152     prtTree->Branch("particle_id_vertex_primary", &prtParticleVertexPrimary);
0153     prtTree->Branch("particle_id_vertex_secondary",
0154                     &prtParticleVertexSecondary);
0155     prtTree->Branch("particle_id_particle", &prtParticleParticle);
0156     prtTree->Branch("particle_id_generation", &prtParticleGeneration);
0157     prtTree->Branch("particle_id_sub_particle", &prtParticleSubParticle);
0158     prtTree->Branch("particle_type", &prtParticleType);
0159     prtTree->Branch("vx", &prtVx);
0160     prtTree->Branch("vy", &prtVy);
0161     prtTree->Branch("vz", &prtVz);
0162     prtTree->Branch("vt", &prtVt);
0163     prtTree->Branch("px", &prtPx);
0164     prtTree->Branch("py", &prtPy);
0165     prtTree->Branch("pz", &prtPz);
0166     prtTree->Branch("m", &prtM);
0167     prtTree->Branch("q", &prtQ);
0168     prtTree->Branch("nhits", &prtNumMeasurements);
0169     prtTree->Branch("ntracks", &prtNumTracks);
0170     prtTree->Branch("ntracks_majority", &prtNumTracksMajority);
0171   }
0172
0173   const Acts::Logger& logger() const { return _logger; }
0174
0175   void write(std::uint64_t eventId, const ConstTrackContainer& tracks,
0176              const SimParticleContainer& particles,
0177              const ParticleMeasurementsMap& particleMeasurementsMap,
0178              const TrackParticleMatching& trackParticleMatching) {
0179     // How often a particle was reconstructed.
0180     std::unordered_map<ActsFatras::Barcode, std::size_t> reconCount;
0181     reconCount.reserve(particles.size());
0182     // How often a particle was reconstructed as the majority particle.
0183     std::unordered_map<ActsFatras::Barcode, std::size_t> majorityCount;
0184     majorityCount.reserve(particles.size());
0185
0186     // write per-track performance measures
0187     {
0188       std::lock_guard<std::mutex> guardTrk(trkMutex);
0189       for (auto track : tracks) {
0190         // Get the truth-matched particle
0191         auto imatched = trackParticleMatching.find(track.index());
0192         if (imatched == trackParticleMatching.end()) {
0193           ACTS_DEBUG(
0194               "No truth particle associated with this proto track, index = "
0195               << track.index());
0196           continue;
0197         }
0198         const auto& particleMatch = imatched->second;
0199
0200         if (particleMatch.particle.has_value()) {
0201           SimBarcode majorityParticleId = particleMatch.particle.value();
0202
0203           auto it = majorityCount.try_emplace(majorityParticleId, 0u).first;
0204           it->second += 1;
0205
0206           // Find the truth particle via the barcode
0207           if (auto ip = particles.find(majorityParticleId);
0208               ip == particles.end()) {
0209             ACTS_WARNING(
0210                 "Majority particle not found in the particles collection.");
0211           }
0212         }
0213
0214         for (const auto& hc : particleMatch.contributingParticles) {
0215           auto it = reconCount.try_emplace(hc.particleId, 0u).first;
0216           it->second += 1;
0217         }
0218
0219         trkEventId = eventId;
0220         trkTrackId = track.index();
0221         trkNumHits = track.nMeasurements();
0222         trkNumParticles = particleMatch.contributingParticles.size();
0223         trkParticleVertexPrimary.clear();
0224         trkParticleVertexSecondary.clear();
0225         trkParticleParticle.clear();
0226         trkParticleGeneration.clear();
0227         trkParticleSubParticle.clear();
0228         trkParticleNumHitsTotal.clear();
0229         trkParticleNumHitsOnTrack.clear();
0230         for (const auto& phc : particleMatch.contributingParticles) {
0231           const auto barcode = phc.particleId;
0232           trkParticleVertexPrimary.push_back(barcode.vertexPrimary());
0233           trkParticleVertexSecondary.push_back(barcode.vertexSecondary());
0234           trkParticleParticle.push_back(barcode.particle());
0235           trkParticleGeneration.push_back(barcode.generation());
0236           trkParticleSubParticle.push_back(barcode.subParticle());
0237           // count total number of hits for this particle
0238           auto trueParticleHits =
0239               makeRange(particleMeasurementsMap.equal_range(phc.particleId));
0240           trkParticleNumHitsTotal.push_back(trueParticleHits.size());
0241           trkParticleNumHitsOnTrack.push_back(phc.hitCount);
0242         }
0243
0244         trkTree->Fill();
0245       }
0246     }
0247
0248     // write per-particle performance measures
0249     {
0250       std::lock_guard<std::mutex> guardPrt(trkMutex);
0251       for (const auto& particle : particles) {
0252         // find all measurements for this particle
0253         auto measurements = makeRange(
0254             particleMeasurementsMap.equal_range(particle.particleId()));
0255
0256         // identification
0257         prtEventId = eventId;
0258         const auto particleBarcode = particle.particleId();
0259         prtParticleVertexPrimary = particleBarcode.vertexPrimary();
0260         prtParticleVertexSecondary = particleBarcode.vertexSecondary();
0261         prtParticleParticle = particleBarcode.particle();
0262         prtParticleGeneration = particleBarcode.generation();
0263         prtParticleSubParticle = particleBarcode.subParticle();
0264         prtParticleType = particle.pdg();
0265         // kinematics
0266         prtVx = particle.position().x() / Acts::UnitConstants::mm;
0267         prtVy = particle.position().y() / Acts::UnitConstants::mm;
0268         prtVz = particle.position().z() / Acts::UnitConstants::mm;
0269         prtVt = particle.time() / Acts::UnitConstants::mm;
0270         const auto p = particle.absoluteMomentum() / Acts::UnitConstants::GeV;
0271         prtPx = p * particle.direction().x();
0272         prtPy = p * particle.direction().y();
0273         prtPz = p * particle.direction().z();
0274         prtM = particle.mass() / Acts::UnitConstants::GeV;
0275         prtQ = particle.charge() / Acts::UnitConstants::e;
0276         // reconstruction
0277         prtNumMeasurements = measurements.size();
0278         auto nt = reconCount.find(particle.particleId());
0279         prtNumTracks = (nt != reconCount.end()) ? nt->second : 0u;
0280         auto nm = majorityCount.find(particle.particleId());
0281         prtNumTracksMajority = (nm != majorityCount.end()) ? nm->second : 0u;
0282
0283         prtTree->Fill();
0284       }
0285     }
0286   }
0287
0288   /// Write everything to disk and close the file.
0289   void close() {
0290     if (file == nullptr) {
0291       ACTS_ERROR("Output file is not available");
0292       return;
0293     }
0294     file->Write();
0295     file->Close();
0296   }
0297 };
0298
0299 RootTrackFinderNTupleWriter::RootTrackFinderNTupleWriter(
0300     RootTrackFinderNTupleWriter::Config config, Acts::Logging::Level level)
0301     : WriterT(config.inputTracks, "RootTrackFinderNTupleWriter", level),
0302       m_impl(std::make_unique<Impl>(this, std::move(config), logger())) {}
0303
0304 RootTrackFinderNTupleWriter::~RootTrackFinderNTupleWriter() = default;
0305
0306 ProcessCode RootTrackFinderNTupleWriter::writeT(
0307     const AlgorithmContext& ctx, const ConstTrackContainer& tracks) {
0308   const auto& particles = m_impl->inputParticles(ctx);
0309   const auto& particleMeasurementsMap =
0310       m_impl->inputParticleMeasurementsMap(ctx);
0311   const auto& trackParticleMatching = m_impl->inputTrackParticleMatching(ctx);
0312   m_impl->write(ctx.eventNumber, tracks, particles, particleMeasurementsMap,
0313                 trackParticleMatching);
0314   return ProcessCode::SUCCESS;
0315 }
0316
0317 ProcessCode RootTrackFinderNTupleWriter::finalize() {
0318   m_impl->close();
0319   return ProcessCode::SUCCESS;
0320 }
0321
0322 const RootTrackFinderNTupleWriter::Config& RootTrackFinderNTupleWriter::config()
0323     const {
0324   return m_impl->cfg;
0325 }
0326
0327 }  // namespace ActsExamples