Root/src/TrackFinderNTupleWriter.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/TrackFinderNTupleWriter.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 TrackFinderNTupleWriter::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<ULong64_t> trkParticleId;
0057   // total number of hits generated by this particle
0058   std::vector<UShort_t> trkParticleNumHitsTotal;
0059   // number of hits within this track
0060   std::vector<UShort_t> trkParticleNumHitsOnTrack;
0061
0062   // per-particle tree
0063   TTree* prtTree = nullptr;
0064   std::mutex prtMutex;
0065   // particle identification
0066   ULong64_t prtEventId = 0;
0067   ULong64_t prtParticleId = 0;
0068   Int_t prtParticleType = 0;
0069   // particle kinematics
0070   // vertex position in mm
0071   float prtVx = 0, prtVy = 0, prtVz = 0;
0072   // vertex time in ns
0073   float prtVt = 0;
0074   // particle momentum at production in GeV
0075   float prtPx = 0, prtPy = 0, prtPz = 0;
0076   // particle mass in GeV
0077   float prtM = 0;
0078   // particle charge in e
0079   float prtQ = 0;
0080   // particle reconstruction
0081   UShort_t prtNumMeasurements = 0;  // number of hits for this particle
0082   UShort_t prtNumTracks =
0083       0;  // number of tracks this particle was reconstructed in
0084   UShort_t prtNumTracksMajority =
0085       0;  // number of tracks reconstructed as majority
0086   // extra logger reference for the logging macros
0087   const Acts::Logger& _logger;
0088
0089   Impl(TrackFinderNTupleWriter* parent, Config&& c, const Acts::Logger& l)
0090       : cfg(std::move(c)),
0091         inputParticles{parent, "InputParticles"},
0092         inputParticleMeasurementsMap{parent, "InputParticleMeasurementsMap"},
0093         inputTrackParticleMatching{parent, "InputTrackParticleMatching"},
0094         _logger(l) {
0095     if (cfg.inputTracks.empty()) {
0096       throw std::invalid_argument("Missing track input collection");
0097     }
0098     if (cfg.inputParticles.empty()) {
0099       throw std::invalid_argument("Missing particles input collection");
0100     }
0101     if (cfg.inputParticleMeasurementsMap.empty()) {
0102       throw std::invalid_argument(
0103           "Missing particle-measurements map input collection");
0104     }
0105     if (cfg.inputTrackParticleMatching.empty()) {
0106       throw std::invalid_argument(
0107           "Missing proto track-particle matching input collection");
0108     }
0109     if (cfg.filePath.empty()) {
0110       throw std::invalid_argument("Missing output filename");
0111     }
0112
0113     inputParticles.initialize(cfg.inputParticles);
0114     inputParticleMeasurementsMap.initialize(cfg.inputParticleMeasurementsMap);
0115     inputTrackParticleMatching.initialize(cfg.inputTrackParticleMatching);
0116
0117     // the output file can not be given externally since TFile accesses to the
0118     // same file from multiple threads are unsafe.
0119     // must always be opened internally
0120     file = TFile::Open(cfg.filePath.c_str(), cfg.fileMode.c_str());
0121     if (file == nullptr) {
0122       throw std::invalid_argument("Could not open '" + cfg.filePath + "'");
0123     }
0124
0125     // construct trees
0126     trkTree = new TTree(cfg.treeNameTracks.c_str(), cfg.treeNameTracks.c_str());
0127     trkTree->SetDirectory(file);
0128     trkTree->Branch("event_id", &trkEventId);
0129     trkTree->Branch("track_id", &trkTrackId);
0130     trkTree->Branch("size", &trkNumHits);
0131     trkTree->Branch("nparticles", &trkNumParticles);
0132     trkTree->Branch("particle_id", &trkParticleId);
0133     trkTree->Branch("particle_nhits_total", &trkParticleNumHitsTotal);
0134     trkTree->Branch("particle_nhits_on_track", &trkParticleNumHitsOnTrack);
0135     prtTree =
0136         new TTree(cfg.treeNameParticles.c_str(), cfg.treeNameParticles.c_str());
0137     prtTree->SetDirectory(file);
0138     prtTree->Branch("event_id", &prtEventId);
0139     prtTree->Branch("particle_id", &prtParticleId);
0140     prtTree->Branch("particle_type", &prtParticleType);
0141     prtTree->Branch("vx", &prtVx);
0142     prtTree->Branch("vy", &prtVy);
0143     prtTree->Branch("vz", &prtVz);
0144     prtTree->Branch("vt", &prtVt);
0145     prtTree->Branch("px", &prtPx);
0146     prtTree->Branch("py", &prtPy);
0147     prtTree->Branch("pz", &prtPz);
0148     prtTree->Branch("m", &prtM);
0149     prtTree->Branch("q", &prtQ);
0150     prtTree->Branch("nhits", &prtNumMeasurements);
0151     prtTree->Branch("ntracks", &prtNumTracks);
0152     prtTree->Branch("ntracks_majority", &prtNumTracksMajority);
0153   }
0154
0155   const Acts::Logger& logger() const { return _logger; }
0156
0157   void write(std::uint64_t eventId, const ConstTrackContainer& tracks,
0158              const SimParticleContainer& particles,
0159              const ParticleMeasurementsMap& particleMeasurementsMap,
0160              const TrackParticleMatching& trackParticleMatching) {
0161     // How often a particle was reconstructed.
0162     std::unordered_map<ActsFatras::Barcode, std::size_t> reconCount;
0163     reconCount.reserve(particles.size());
0164     // How often a particle was reconstructed as the majority particle.
0165     std::unordered_map<ActsFatras::Barcode, std::size_t> majorityCount;
0166     majorityCount.reserve(particles.size());
0167
0168     // write per-track performance measures
0169     {
0170       std::lock_guard<std::mutex> guardTrk(trkMutex);
0171       for (auto track : tracks) {
0172         // Get the truth-matched particle
0173         auto imatched = trackParticleMatching.find(track.index());
0174         if (imatched == trackParticleMatching.end()) {
0175           ACTS_DEBUG(
0176               "No truth particle associated with this proto track, index = "
0177               << track.index());
0178           continue;
0179         }
0180         const auto& particleMatch = imatched->second;
0181
0182         if (particleMatch.particle.has_value()) {
0183           SimBarcode majorityParticleId = particleMatch.particle.value();
0184
0185           auto it = majorityCount.try_emplace(majorityParticleId, 0u).first;
0186           it->second += 1;
0187
0188           // Find the truth particle via the barcode
0189           if (auto ip = particles.find(majorityParticleId);
0190               ip == particles.end()) {
0191             ACTS_WARNING(
0192                 "Majority particle not found in the particles collection.");
0193           }
0194         }
0195
0196         for (const auto& hc : particleMatch.contributingParticles) {
0197           auto it = reconCount.try_emplace(hc.particleId, 0u).first;
0198           it->second += 1;
0199         }
0200
0201         trkEventId = eventId;
0202         trkTrackId = track.index();
0203         trkNumHits = track.nMeasurements();
0204         trkNumParticles = particleMatch.contributingParticles.size();
0205         trkParticleId.clear();
0206         trkParticleNumHitsTotal.clear();
0207         trkParticleNumHitsOnTrack.clear();
0208         for (const auto& phc : particleMatch.contributingParticles) {
0209           trkParticleId.push_back(phc.particleId.value());
0210           // count total number of hits for this particle
0211           auto trueParticleHits = makeRange(
0212               particleMeasurementsMap.equal_range(phc.particleId.value()));
0213           trkParticleNumHitsTotal.push_back(trueParticleHits.size());
0214           trkParticleNumHitsOnTrack.push_back(phc.hitCount);
0215         }
0216
0217         trkTree->Fill();
0218       }
0219     }
0220
0221     // write per-particle performance measures
0222     {
0223       std::lock_guard<std::mutex> guardPrt(trkMutex);
0224       for (const auto& particle : particles) {
0225         // find all measurements for this particle
0226         auto measurements = makeRange(
0227             particleMeasurementsMap.equal_range(particle.particleId()));
0228
0229         // identification
0230         prtEventId = eventId;
0231         prtParticleId = particle.particleId().value();
0232         prtParticleType = particle.pdg();
0233         // kinematics
0234         prtVx = particle.position().x() / Acts::UnitConstants::mm;
0235         prtVy = particle.position().y() / Acts::UnitConstants::mm;
0236         prtVz = particle.position().z() / Acts::UnitConstants::mm;
0237         prtVt = particle.time() / Acts::UnitConstants::mm;
0238         const auto p = particle.absoluteMomentum() / Acts::UnitConstants::GeV;
0239         prtPx = p * particle.direction().x();
0240         prtPy = p * particle.direction().y();
0241         prtPz = p * particle.direction().z();
0242         prtM = particle.mass() / Acts::UnitConstants::GeV;
0243         prtQ = particle.charge() / Acts::UnitConstants::e;
0244         // reconstruction
0245         prtNumMeasurements = measurements.size();
0246         auto nt = reconCount.find(particle.particleId());
0247         prtNumTracks = (nt != reconCount.end()) ? nt->second : 0u;
0248         auto nm = majorityCount.find(particle.particleId());
0249         prtNumTracksMajority = (nm != majorityCount.end()) ? nm->second : 0u;
0250
0251         prtTree->Fill();
0252       }
0253     }
0254   }
0255
0256   /// Write everything to disk and close the file.
0257   void close() {
0258     if (file == nullptr) {
0259       ACTS_ERROR("Output file is not available");
0260       return;
0261     }
0262     file->Write();
0263     file->Close();
0264   }
0265 };
0266
0267 TrackFinderNTupleWriter::TrackFinderNTupleWriter(
0268     TrackFinderNTupleWriter::Config config, Acts::Logging::Level level)
0269     : WriterT(config.inputTracks, "TrackFinderNTupleWriter", level),
0270       m_impl(std::make_unique<Impl>(this, std::move(config), logger())) {}
0271
0272 TrackFinderNTupleWriter::~TrackFinderNTupleWriter() = default;
0273
0274 ProcessCode TrackFinderNTupleWriter::writeT(const AlgorithmContext& ctx,
0275                                             const ConstTrackContainer& tracks) {
0276   const auto& particles = m_impl->inputParticles(ctx);
0277   const auto& particleMeasurementsMap =
0278       m_impl->inputParticleMeasurementsMap(ctx);
0279   const auto& trackParticleMatching = m_impl->inputTrackParticleMatching(ctx);
0280   m_impl->write(ctx.eventNumber, tracks, particles, particleMeasurementsMap,
0281                 trackParticleMatching);
0282   return ProcessCode::SUCCESS;
0283 }
0284
0285 ProcessCode TrackFinderNTupleWriter::finalize() {
0286   m_impl->close();
0287   return ProcessCode::SUCCESS;
0288 }
0289
0290 const TrackFinderNTupleWriter::Config& TrackFinderNTupleWriter::config() const {
0291   return m_impl->cfg;
0292 }
0293
0294 }  // namespace ActsExamples