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   for (const auto& track : tracks) {
0068     // Reco track selection
0069     //@TODO: add interface for applying others cuts on reco tracks:
0070     // -> pT, d0, z0, detector-specific hits/holes number cut
0071     if (track.nMeasurements() < m_cfg.nMeasurementsMin) {
0072       continue;
0073     }
0074
0075     // Check if the reco track has fitted track parameters
0076     if (!track.hasReferenceSurface()) {
0077       ACTS_WARNING(
0078           "No fitted track parameters for trajectory with entry index = "
0079           << track.tipIndex());
0080       continue;
0081     }
0082
0083     // Get the majority truth particle to this track
0084     std::vector<ParticleHitCount> particleHitCount;
0085     identifyContributingParticles(hitParticlesMap, track, particleHitCount);
0086     if (m_cfg.onlyTruthMatched && particleHitCount.empty()) {
0087       ACTS_WARNING(
0088           "No truth particle associated with this trajectory with entry "
0089           "index = "
0090           << track.tipIndex());
0091       continue;
0092     }
0093
0094     // Requirement on the pT of the track
0095     auto params = track.createParametersAtReference();
0096     const auto momentum = params.momentum();
0097     const auto pT = Acts::VectorHelpers::perp(momentum);
0098     if (pT < m_cfg.ptMin) {
0099       continue;
0100     }
0101     std::size_t nMajorityHits = 0;
0102     ActsFatras::Barcode majorityParticleId;
0103     if (!particleHitCount.empty()) {
0104       // Get the majority particle counts
0105       majorityParticleId = particleHitCount.front().particleId;
0106       // n Majority hits
0107       nMajorityHits = particleHitCount.front().hitCount;
0108     }
0109
0110     static const Acts::ConstProxyAccessor<unsigned int> seedNumber(
0111         "trackGroup");
0112
0113     // track info
0114     TrackInfo toAdd;
0115     toAdd.trackId = track.index();
0116     if (tracks.hasColumn(Acts::hashString("trackGroup"))) {
0117       toAdd.seedID = seedNumber(track);
0118     } else {
0119       toAdd.seedID = 0;
0120     }
0121     toAdd.particleId = majorityParticleId;
0122     toAdd.nStates = track.nTrackStates();
0123     toAdd.nMajorityHits = nMajorityHits;
0124     toAdd.nMeasurements = track.nMeasurements();
0125     toAdd.nOutliers = track.nOutliers();
0126     toAdd.nHoles = track.nHoles();
0127     toAdd.nSharedHits = track.nSharedHits();
0128     toAdd.chi2Sum = track.chi2();
0129     toAdd.NDF = track.nDoF();
0130     toAdd.truthMatchProb = toAdd.nMajorityHits * 1. / track.nMeasurements();
0131     toAdd.fittedParameters = params;
0132     toAdd.trackType = "unknown";
0133
0134     for (const auto& state : track.trackStatesReversed()) {
0135       if (state.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag)) {
0136         auto sl =
0137             state.getUncalibratedSourceLink().template get<IndexSourceLink>();
0138         auto hitIndex = sl.index();
0139         toAdd.measurementsID.insert(toAdd.measurementsID.begin(), hitIndex);
0140       }
0141     }
0142
0143     // Check if the trajectory is matched with truth.
0144     if (toAdd.truthMatchProb >= m_cfg.truthMatchProbMin) {
0145       matched[toAdd.particleId].push_back({toAdd, toAdd.trackId});
0146     } else {
0147       toAdd.trackType = "fake";
0148     }
0149
0150     infoMap[toAdd.trackId] = toAdd;
0151   }
0152
0153   // Find duplicates
0154   std::unordered_set<std::size_t> listGoodTracks;
0155   for (auto& [particleId, matchedTracks] : matched) {
0156     std::ranges::sort(matchedTracks, [](const auto& lhs, const auto& rhs) {
0157       const auto& t1 = lhs.first;
0158       const auto& t2 = rhs.first;
0159       // nMajorityHits are sorted descending, others ascending
0160       return std::tie(t2.nMajorityHits, t1.nOutliers, t1.chi2Sum) <
0161              std::tie(t1.nMajorityHits, t2.nOutliers, t2.chi2Sum);
0162     });
0163
0164     listGoodTracks.insert(matchedTracks.front().first.trackId);
0165   }
0166
0167   // write csv header
0168   mos << "track_id,seed_id,particleId,"
0169       << "nStates,nMajorityHits,nMeasurements,nOutliers,nHoles,nSharedHits,"
0170       << "chi2,ndf,chi2/ndf,"
0171       << "pT,eta,phi,"
0172       << "truthMatchProbability,"
0173       << "good/duplicate/fake,"
0174       << "Measurements_ID";
0175
0176   mos << '\n';
0177   mos << std::setprecision(m_cfg.outputPrecision);
0178
0179   // good/duplicate/fake = 0/1/2
0180   for (auto& [id, trajState] : infoMap) {
0181     if (listGoodTracks.contains(id)) {
0182       trajState.trackType = "good";
0183     } else if (trajState.trackType != "fake") {
0184       trajState.trackType = "duplicate";
0185     }
0186
0187     const auto& params = *trajState.fittedParameters;
0188     const auto momentum = params.momentum();
0189
0190     // write the track info
0191     mos << trajState.trackId << ",";
0192     mos << trajState.seedID << ",";
0193     mos << trajState.particleId << ",";
0194     mos << trajState.nStates << ",";
0195     mos << trajState.nMajorityHits << ",";
0196     mos << trajState.nMeasurements << ",";
0197     mos << trajState.nOutliers << ",";
0198     mos << trajState.nHoles << ",";
0199     mos << trajState.nSharedHits << ",";
0200     mos << trajState.chi2Sum << ",";
0201     mos << trajState.NDF << ",";
0202     mos << trajState.chi2Sum * 1.0 / trajState.NDF << ",";
0203     mos << Acts::VectorHelpers::perp(momentum) << ",";
0204     mos << Acts::VectorHelpers::eta(momentum) << ",";
0205     mos << Acts::VectorHelpers::phi(momentum) << ",";
0206     mos << trajState.truthMatchProb << ",";
0207     mos << trajState.trackType << ",";
0208     mos << "\"[";
0209     for (auto& ID : trajState.measurementsID) {
0210       mos << ID << ",";
0211     }
0212     mos << "]\"";
0213     mos << '\n';
0214   }
0215
0216   return ProcessCode::SUCCESS;
0217 }