Csv/src/CsvTrackWriter.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/Csv/CsvTrackWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/EventData/MultiTrajectory.hpp"
0013 #include "Acts/EventData/ProxyAccessor.hpp"
0014 #include "Acts/EventData/VectorMultiTrajectory.hpp"
0015 #include "Acts/Utilities/Helpers.hpp"
0016 #include "Acts/Utilities/MultiIndex.hpp"
0017 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0018 #include "ActsExamples/EventData/Track.hpp"
0019 #include "ActsExamples/Framework/AlgorithmContext.hpp"
0020 #include "ActsExamples/Utilities/Paths.hpp"
0021 #include "ActsExamples/Utilities/Range.hpp"
0022 #include "ActsExamples/Validation/TrackClassification.hpp"
0023
0024 #include <algorithm>
0025 #include <fstream>
0026 #include <iomanip>
0027 #include <map>
0028 #include <memory>
0029 #include <stdexcept>
0030 #include <string>
0031 #include <tuple>
0032 #include <unordered_map>
0033 #include <unordered_set>
0034 #include <utility>
0035
0036 using namespace ActsExamples;
0037
0038 CsvTrackWriter::CsvTrackWriter(const CsvTrackWriter::Config& config,
0039                                Acts::Logging::Level level)
0040     : WriterT<ConstTrackContainer>(config.inputTracks, "CsvTrackWriter", level),
0041       m_cfg(config) {
0042   if (m_cfg.inputTracks.empty()) {
0043     throw std::invalid_argument("Missing input tracks collection");
0044   }
0045
0046   m_inputMeasurementParticlesMap.initialize(m_cfg.inputMeasurementParticlesMap);
0047 }
0048
0049 ProcessCode CsvTrackWriter::writeT(const AlgorithmContext& context,
0050                                    const ConstTrackContainer& tracks) {
0051   // open per-event file
0052   std::string path =
0053       perEventFilepath(m_cfg.outputDir, m_cfg.fileName, context.eventNumber);
0054   std::ofstream mos(path, std::ofstream::out | std::ofstream::trunc);
0055   if (!mos) {
0056     throw std::ios_base::failure("Could not open '" + path + "' to write");
0057   }
0058
0059   const auto& hitParticlesMap = m_inputMeasurementParticlesMap(context);
0060
0061   std::unordered_map<Acts::MultiTrajectoryTraits::IndexType, TrackInfo> infoMap;
0062
0063   // Counter of truth-matched reco tracks
0064   using RecoTrackInfo = std::pair<TrackInfo, std::size_t>;
0065   std::map<ActsFatras::Barcode, std::vector<RecoTrackInfo>> matched;
0066
0067   std::size_t trackId = 0;
0068   for (const auto& track : tracks) {
0069     // Reco track selection
0070     //@TODO: add interface for applying others cuts on reco tracks:
0071     // -> pT, d0, z0, detector-specific hits/holes number cut
0072     if (track.nMeasurements() < m_cfg.nMeasurementsMin) {
0073       continue;
0074     }
0075
0076     // Check if the reco track has fitted track parameters
0077     if (!track.hasReferenceSurface()) {
0078       ACTS_WARNING(
0079           "No fitted track parameters for trajectory with entry index = "
0080           << track.tipIndex());
0081       continue;
0082     }
0083
0084     // Get the majority truth particle to this track
0085     std::vector<ParticleHitCount> particleHitCount;
0086     identifyContributingParticles(hitParticlesMap, track, particleHitCount);
0087     if (m_cfg.onlyTruthMatched && particleHitCount.empty()) {
0088       ACTS_WARNING(
0089           "No truth particle associated with this trajectory with entry "
0090           "index = "
0091           << track.tipIndex());
0092       continue;
0093     }
0094
0095     // Requirement on the pT of the track
0096     auto params = track.createParametersAtReference();
0097     const auto momentum = params.momentum();
0098     const auto pT = Acts::VectorHelpers::perp(momentum);
0099     if (pT < m_cfg.ptMin) {
0100       continue;
0101     }
0102     std::size_t nMajorityHits = 0;
0103     ActsFatras::Barcode majorityParticleId;
0104     if (!particleHitCount.empty()) {
0105       // Get the majority particle counts
0106       majorityParticleId = particleHitCount.front().particleId;
0107       // n Majority hits
0108       nMajorityHits = particleHitCount.front().hitCount;
0109     }
0110
0111     static const Acts::ConstProxyAccessor<unsigned int> seedNumber(
0112         "trackGroup");
0113
0114     // track info
0115     TrackInfo toAdd;
0116     toAdd.trackId = trackId;
0117     if (tracks.hasColumn(Acts::hashString("trackGroup"))) {
0118       toAdd.seedID = seedNumber(track);
0119     } else {
0120       toAdd.seedID = 0;
0121     }
0122     toAdd.particleId = majorityParticleId;
0123     toAdd.nStates = track.nTrackStates();
0124     toAdd.nMajorityHits = nMajorityHits;
0125     toAdd.nMeasurements = track.nMeasurements();
0126     toAdd.nOutliers = track.nOutliers();
0127     toAdd.nHoles = track.nHoles();
0128     toAdd.nSharedHits = track.nSharedHits();
0129     toAdd.chi2Sum = track.chi2();
0130     toAdd.NDF = track.nDoF();
0131     toAdd.truthMatchProb = toAdd.nMajorityHits * 1. / track.nMeasurements();
0132     toAdd.fittedParameters = params;
0133     toAdd.trackType = "unknown";
0134
0135     for (const auto& state : track.trackStatesReversed()) {
0136       if (state.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
0137         auto sl =
0138             state.getUncalibratedSourceLink().template get<IndexSourceLink>();
0139         auto hitIndex = sl.index();
0140         toAdd.measurementsID.insert(toAdd.measurementsID.begin(), hitIndex);
0141       }
0142     }
0143
0144     // Check if the trajectory is matched with truth.
0145     if (toAdd.truthMatchProb >= m_cfg.truthMatchProbMin) {
0146       matched[toAdd.particleId].push_back({toAdd, toAdd.trackId});
0147     } else {
0148       toAdd.trackType = "fake";
0149     }
0150
0151     infoMap[toAdd.trackId] = toAdd;
0152
0153     trackId++;
0154   }
0155
0156   // Find duplicates
0157   std::unordered_set<std::size_t> listGoodTracks;
0158   for (auto& [particleId, matchedTracks] : matched) {
0159     std::ranges::sort(matchedTracks, [](const auto& lhs, const auto& rhs) {
0160       const auto& t1 = lhs.first;
0161       const auto& t2 = rhs.first;
0162       // nMajorityHits are sorted descending, others ascending
0163       return std::tie(t2.nMajorityHits, t1.nOutliers, t1.chi2Sum) <
0164              std::tie(t1.nMajorityHits, t2.nOutliers, t2.chi2Sum);
0165     });
0166
0167     listGoodTracks.insert(matchedTracks.front().first.trackId);
0168   }
0169
0170   // write csv header
0171   mos << "track_id,seed_id,particleId,"
0172       << "nStates,nMajorityHits,nMeasurements,nOutliers,nHoles,nSharedHits,"
0173       << "chi2,ndf,chi2/ndf,"
0174       << "pT,eta,phi,"
0175       << "truthMatchProbability,"
0176       << "good/duplicate/fake,"
0177       << "Hits_ID";
0178
0179   mos << '\n';
0180   mos << std::setprecision(m_cfg.outputPrecision);
0181
0182   // good/duplicate/fake = 0/1/2
0183   for (auto& [id, trajState] : infoMap) {
0184     if (listGoodTracks.contains(id)) {
0185       trajState.trackType = "good";
0186     } else if (trajState.trackType != "fake") {
0187       trajState.trackType = "duplicate";
0188     }
0189
0190     const auto& params = *trajState.fittedParameters;
0191     const auto momentum = params.momentum();
0192
0193     // write the track info
0194     mos << trajState.trackId << ",";
0195     mos << trajState.seedID << ",";
0196     mos << trajState.particleId << ",";
0197     mos << trajState.nStates << ",";
0198     mos << trajState.nMajorityHits << ",";
0199     mos << trajState.nMeasurements << ",";
0200     mos << trajState.nOutliers << ",";
0201     mos << trajState.nHoles << ",";
0202     mos << trajState.nSharedHits << ",";
0203     mos << trajState.chi2Sum << ",";
0204     mos << trajState.NDF << ",";
0205     mos << trajState.chi2Sum * 1.0 / trajState.NDF << ",";
0206     mos << Acts::VectorHelpers::perp(momentum) << ",";
0207     mos << Acts::VectorHelpers::eta(momentum) << ",";
0208     mos << Acts::VectorHelpers::phi(momentum) << ",";
0209     mos << trajState.truthMatchProb << ",";
0210     mos << trajState.trackType << ",";
0211     mos << "\"[";
0212     for (auto& ID : trajState.measurementsID) {
0213       mos << ID << ",";
0214     }
0215     mos << "]\"";
0216     mos << '\n';
0217   }
0218
0219   return ProcessCode::SUCCESS;
0220 }