Root/src/VertexNTupleWriter.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/VertexNTupleWriter.hpp"
0010
0011 #include "Acts/Definitions/Algebra.hpp"
0012 #include "Acts/Definitions/TrackParametrization.hpp"
0013 #include "Acts/EventData/GenericBoundTrackParameters.hpp"
0014 #include "Acts/EventData/TrackParameters.hpp"
0015 #include "Acts/Propagator/Propagator.hpp"
0016 #include "Acts/Propagator/SympyStepper.hpp"
0017 #include "Acts/Surfaces/PerigeeSurface.hpp"
0018 #include "Acts/Surfaces/Surface.hpp"
0019 #include "Acts/Utilities/Logger.hpp"
0020 #include "Acts/Utilities/MultiIndex.hpp"
0021 #include "Acts/Utilities/UnitVectors.hpp"
0022 #include "Acts/Utilities/Zip.hpp"
0023 #include "Acts/Vertexing/TrackAtVertex.hpp"
0024 #include "ActsExamples/EventData/SimHit.hpp"
0025 #include "ActsExamples/EventData/SimParticle.hpp"
0026 #include "ActsExamples/EventData/SimVertex.hpp"
0027 #include "ActsExamples/EventData/Track.hpp"
0028 #include "ActsExamples/EventData/Trajectories.hpp"
0029 #include "ActsExamples/EventData/TruthMatching.hpp"
0030 #include "ActsExamples/Validation/TrackClassification.hpp"
0031 #include "ActsFatras/EventData/Barcode.hpp"
0032 #include "ActsFatras/EventData/Particle.hpp"
0033
0034 #include <algorithm>
0035 #include <cmath>
0036 #include <cstddef>
0037 #include <cstdint>
0038 #include <ios>
0039 #include <limits>
0040 #include <map>
0041 #include <memory>
0042 #include <numbers>
0043 #include <optional>
0044 #include <ostream>
0045 #include <set>
0046 #include <stdexcept>
0047 #include <utility>
0048
0049 #include <TFile.h>
0050 #include <TTree.h>
0051
0052 using Acts::VectorHelpers::eta;
0053 using Acts::VectorHelpers::perp;
0054 using Acts::VectorHelpers::phi;
0055 using Acts::VectorHelpers::theta;
0056
0057 namespace ActsExamples {
0058
0059 namespace {
0060
0061 std::uint32_t getNumberOfReconstructableVertices(
0062     const SimParticleContainer& collection) {
0063   // map for finding frequency
0064   std::map<std::uint32_t, std::uint32_t> fmap;
0065
0066   std::vector<std::uint32_t> reconstructableTruthVertices;
0067
0068   // traverse the array for frequency
0069   for (const SimParticle& p : collection) {
0070     std::uint32_t generation = p.particleId().generation();
0071     if (generation > 0) {
0072       // truthparticle from secondary vtx
0073       continue;
0074     }
0075     std::uint32_t priVtxId = p.particleId().vertexPrimary();
0076     fmap[priVtxId]++;
0077   }
0078
0079   // iterate over the map
0080   for (const auto& [priVtxId, occurrence] : fmap) {
0081     // Require at least 2 tracks
0082     if (occurrence > 1) {
0083       reconstructableTruthVertices.push_back(priVtxId);
0084     }
0085   }
0086
0087   return reconstructableTruthVertices.size();
0088 }
0089
0090 std::uint32_t getNumberOfTruePriVertices(
0091     const SimParticleContainer& collection) {
0092   // Vector to store indices of all primary vertices
0093   std::set<std::uint32_t> allPriVtxIds;
0094   for (const SimParticle& p : collection) {
0095     std::uint32_t priVtxId = p.particleId().vertexPrimary();
0096     std::uint32_t generation = p.particleId().generation();
0097     if (generation > 0) {
0098       // truthparticle from secondary vtx
0099       continue;
0100     }
0101     // Insert to set, removing duplicates
0102     allPriVtxIds.insert(priVtxId);
0103   }
0104   // Size of set corresponds to total number of primary vertices
0105   return allPriVtxIds.size();
0106 }
0107
0108 }  // namespace
0109
0110 VertexNTupleWriter::VertexNTupleWriter(const VertexNTupleWriter::Config& config,
0111                                        Acts::Logging::Level level)
0112     : WriterT(config.inputVertices, "VertexNTupleWriter", level),
0113       m_cfg(config) {
0114   if (m_cfg.filePath.empty()) {
0115     throw std::invalid_argument("Missing output filename");
0116   }
0117   if (m_cfg.treeName.empty()) {
0118     throw std::invalid_argument("Missing tree name");
0119   }
0120   if (m_cfg.inputTruthVertices.empty()) {
0121     throw std::invalid_argument("Collection with truth vertices missing");
0122   }
0123   if (m_cfg.inputParticles.empty()) {
0124     throw std::invalid_argument("Collection with particles missing");
0125   }
0126   if (m_cfg.inputSelectedParticles.empty()) {
0127     throw std::invalid_argument("Collection with selected particles missing");
0128   }
0129   if (m_cfg.inputTrackParticleMatching.empty()) {
0130     throw std::invalid_argument("Missing input track particles matching");
0131   }
0132
0133   m_inputTruthVertices.initialize(m_cfg.inputTruthVertices);
0134   m_inputParticles.initialize(m_cfg.inputParticles);
0135   m_inputSelectedParticles.initialize(m_cfg.inputSelectedParticles);
0136   m_inputTrackParticleMatching.initialize(m_cfg.inputTrackParticleMatching);
0137
0138   if (m_cfg.useTracks) {
0139     m_inputTracks.initialize(m_cfg.inputTracks);
0140   }
0141
0142   // Setup ROOT I/O
0143   m_outputFile = TFile::Open(m_cfg.filePath.c_str(), m_cfg.fileMode.c_str());
0144   if (m_outputFile == nullptr) {
0145     throw std::ios_base::failure("Could not open '" + m_cfg.filePath + "'");
0146   }
0147   m_outputFile->cd();
0148   m_outputTree = new TTree(m_cfg.treeName.c_str(), m_cfg.treeName.c_str());
0149   if (m_outputTree == nullptr) {
0150     throw std::bad_alloc();
0151   }
0152
0153   m_outputTree->Branch("event_nr", &m_eventNr);
0154
0155   m_outputTree->Branch("nRecoVtx", &m_nRecoVtx);
0156   m_outputTree->Branch("nTrueVtx", &m_nTrueVtx);
0157   m_outputTree->Branch("nCleanVtx", &m_nCleanVtx);
0158   m_outputTree->Branch("nMergedVtx", &m_nMergedVtx);
0159   m_outputTree->Branch("nSplitVtx", &m_nSplitVtx);
0160   m_outputTree->Branch("nVtxDetectorAcceptance", &m_nVtxDetAcceptance);
0161   m_outputTree->Branch("nVtxReconstructable", &m_nVtxReconstructable);
0162
0163   m_outputTree->Branch("nTracksRecoVtx", &m_nTracksOnRecoVertex);
0164
0165   m_outputTree->Branch("recoVertexTrackWeights", &m_recoVertexTrackWeights);
0166
0167   m_outputTree->Branch("sumPt2", &m_sumPt2);
0168
0169   m_outputTree->Branch("recoX", &m_recoX);
0170   m_outputTree->Branch("recoY", &m_recoY);
0171   m_outputTree->Branch("recoZ", &m_recoZ);
0172   m_outputTree->Branch("recoT", &m_recoT);
0173
0174   m_outputTree->Branch("covXX", &m_covXX);
0175   m_outputTree->Branch("covYY", &m_covYY);
0176   m_outputTree->Branch("covZZ", &m_covZZ);
0177   m_outputTree->Branch("covTT", &m_covTT);
0178   m_outputTree->Branch("covXY", &m_covXY);
0179   m_outputTree->Branch("covXZ", &m_covXZ);
0180   m_outputTree->Branch("covXT", &m_covXT);
0181   m_outputTree->Branch("covYZ", &m_covYZ);
0182   m_outputTree->Branch("covYT", &m_covYT);
0183   m_outputTree->Branch("covZT", &m_covZT);
0184
0185   m_outputTree->Branch("seedX", &m_seedX);
0186   m_outputTree->Branch("seedY", &m_seedY);
0187   m_outputTree->Branch("seedZ", &m_seedZ);
0188   m_outputTree->Branch("seedT", &m_seedT);
0189
0190   m_outputTree->Branch("vertex_primary", &m_vertexPrimary);
0191   m_outputTree->Branch("vertex_secondary", &m_vertexSecondary);
0192
0193   m_outputTree->Branch("nTracksTruthVtx", &m_nTracksOnTruthVertex);
0194
0195   m_outputTree->Branch("truthPrimaryVertexDensity",
0196                        &m_truthPrimaryVertexDensity);
0197
0198   m_outputTree->Branch("truthVertexTrackWeights", &m_truthVertexTrackWeights);
0199   m_outputTree->Branch("truthVertexMatchRatio", &m_truthVertexMatchRatio);
0200   m_outputTree->Branch("recoVertexContamination", &m_recoVertexContamination);
0201
0202   m_outputTree->Branch("recoVertexClassification", &m_recoVertexClassification);
0203
0204   m_outputTree->Branch("truthX", &m_truthX);
0205   m_outputTree->Branch("truthY", &m_truthY);
0206   m_outputTree->Branch("truthZ", &m_truthZ);
0207   m_outputTree->Branch("truthT", &m_truthT);
0208
0209   m_outputTree->Branch("resX", &m_resX);
0210   m_outputTree->Branch("resY", &m_resY);
0211   m_outputTree->Branch("resZ", &m_resZ);
0212   m_outputTree->Branch("resT", &m_resT);
0213
0214   m_outputTree->Branch("resSeedZ", &m_resSeedZ);
0215   m_outputTree->Branch("resSeedT", &m_resSeedT);
0216
0217   m_outputTree->Branch("pullX", &m_pullX);
0218   m_outputTree->Branch("pullY", &m_pullY);
0219   m_outputTree->Branch("pullZ", &m_pullZ);
0220   m_outputTree->Branch("pullT", &m_pullT);
0221
0222   m_outputTree->Branch("trk_weight", &m_trkWeight);
0223
0224   m_outputTree->Branch("trk_recoPhi", &m_recoPhi);
0225   m_outputTree->Branch("trk_recoTheta", &m_recoTheta);
0226   m_outputTree->Branch("trk_recoQOverP", &m_recoQOverP);
0227
0228   m_outputTree->Branch("trk_recoPhiFitted", &m_recoPhiFitted);
0229   m_outputTree->Branch("trk_recoThetaFitted", &m_recoThetaFitted);
0230   m_outputTree->Branch("trk_recoQOverPFitted", &m_recoQOverPFitted);
0231
0232   m_outputTree->Branch("trk_particleId", &m_trkParticleId);
0233
0234   m_outputTree->Branch("trk_truthPhi", &m_truthPhi);
0235   m_outputTree->Branch("trk_truthTheta", &m_truthTheta);
0236   m_outputTree->Branch("trk_truthQOverP", &m_truthQOverP);
0237
0238   m_outputTree->Branch("trk_resPhi", &m_resPhi);
0239   m_outputTree->Branch("trk_resTheta", &m_resTheta);
0240   m_outputTree->Branch("trk_resQOverP", &m_resQOverP);
0241   m_outputTree->Branch("trk_momOverlap", &m_momOverlap);
0242
0243   m_outputTree->Branch("trk_resPhiFitted", &m_resPhiFitted);
0244   m_outputTree->Branch("trk_resThetaFitted", &m_resThetaFitted);
0245   m_outputTree->Branch("trk_resQOverPFitted", &m_resQOverPFitted);
0246   m_outputTree->Branch("trk_momOverlapFitted", &m_momOverlapFitted);
0247
0248   m_outputTree->Branch("trk_pullPhi", &m_pullPhi);
0249   m_outputTree->Branch("trk_pullTheta", &m_pullTheta);
0250   m_outputTree->Branch("trk_pullQOverP", &m_pullQOverP);
0251
0252   m_outputTree->Branch("trk_pullPhiFitted", &m_pullPhiFitted);
0253   m_outputTree->Branch("trk_pullThetaFitted", &m_pullThetaFitted);
0254   m_outputTree->Branch("trk_pullQOverPFitted", &m_pullQOverPFitted);
0255 }
0256
0257 VertexNTupleWriter::~VertexNTupleWriter() {
0258   if (m_outputFile != nullptr) {
0259     m_outputFile->Close();
0260   }
0261 }
0262
0263 ProcessCode VertexNTupleWriter::finalize() {
0264   m_outputFile->cd();
0265   m_outputTree->Write();
0266   m_outputFile->Close();
0267
0268   return ProcessCode::SUCCESS;
0269 }
0270
0271 ProcessCode VertexNTupleWriter::writeT(
0272     const AlgorithmContext& ctx, const std::vector<Acts::Vertex>& vertices) {
0273   const double nan = std::numeric_limits<double>::quiet_NaN();
0274
0275   // Read truth vertex input collection
0276   const SimVertexContainer& truthVertices = m_inputTruthVertices(ctx);
0277   // Read truth particle input collection
0278   const SimParticleContainer& particles = m_inputParticles(ctx);
0279   const SimParticleContainer& selectedParticles = m_inputSelectedParticles(ctx);
0280   const TrackParticleMatching& trackParticleMatching =
0281       m_inputTrackParticleMatching(ctx);
0282
0283   const ConstTrackContainer* tracks = nullptr;
0284   SimParticleContainer recoParticles;
0285
0286   auto findParticle = [&](ConstTrackProxy track) -> std::optional<SimParticle> {
0287     // Get the truth-matched particle
0288     auto imatched = trackParticleMatching.find(track.index());
0289     if (imatched == trackParticleMatching.end() ||
0290         !imatched->second.particle.has_value()) {
0291       ACTS_DEBUG("No truth particle associated with this track, index = "
0292                  << track.index() << " tip index = " << track.tipIndex());
0293       return {};
0294     }
0295
0296     const TrackMatchEntry& particleMatch = imatched->second;
0297
0298     auto iparticle = particles.find(particleMatch.particle->value());
0299     if (iparticle == particles.end()) {
0300       ACTS_DEBUG(
0301           "Truth particle found but not monitored with this track, index = "
0302           << track.index() << " tip index = " << track.tipIndex()
0303           << " and this barcode = " << particleMatch.particle->value());
0304       return {};
0305     }
0306
0307     return *iparticle;
0308   };
0309
0310   auto findTrack = [&](const Acts::TrackAtVertex& trkAtVtx)
0311       -> std::optional<ConstTrackProxy> {
0312     // Track parameters before the vertex fit
0313     const Acts::BoundTrackParameters& origTrack =
0314         *trkAtVtx.originalParams.as<Acts::BoundTrackParameters>();
0315
0316     // Finding the matching parameters in the container of all track
0317     // parameters. This allows us to identify the corresponding particle.
0318     // TODO this should not be necessary if the tracks at vertex would keep
0319     // this information
0320     for (ConstTrackProxy inputTrk : *tracks) {
0321       const auto& params = inputTrk.parameters();
0322
0323       if (origTrack.parameters() == params) {
0324         return inputTrk;
0325       }
0326     }
0327
0328     return std::nullopt;
0329   };
0330
0331   auto calcSumPt2 = [this](const Acts::Vertex& vtx) {
0332     double sumPt2 = 0;
0333     for (const auto& trk : vtx.tracks()) {
0334       if (trk.trackWeight > m_cfg.minTrkWeight) {
0335         double pt = trk.originalParams.as<Acts::BoundTrackParameters>()
0336                         ->transverseMomentum();
0337         sumPt2 += pt * pt;
0338       }
0339     }
0340     return sumPt2;
0341   };
0342
0343   auto weightHighEnough = [this](const Acts::TrackAtVertex& trkAtVtx) {
0344     return trkAtVtx.trackWeight > m_cfg.minTrkWeight;
0345   };
0346
0347   // Helper function for computing the pull
0348   auto pull = [this](const double& diff, const double& variance,
0349                      const std::string& variableStr,
0350                      const bool& afterFit = true) {
0351     if (variance <= 0) {
0352       std::string tempStr;
0353       if (afterFit) {
0354         tempStr = "after";
0355       } else {
0356         tempStr = "before";
0357       }
0358       ACTS_WARNING("Nonpositive variance " << tempStr << " vertex fit: Var("
0359                                            << variableStr << ") = " << variance
0360                                            << " <= 0.");
0361       return std::numeric_limits<double>::quiet_NaN();
0362     }
0363     double std = std::sqrt(variance);
0364     return diff / std;
0365   };
0366
0367   auto calculateTruthPrimaryVertexDensity =
0368       [this, truthVertices](const Acts::Vertex& vtx) {
0369         double z = vtx.fullPosition()[Acts::CoordinateIndices::eZ];
0370         int count = 0;
0371         for (const SimVertex& truthVertex : truthVertices) {
0372           if (truthVertex.vertexId().vertexSecondary() != 0) {
0373             continue;
0374           }
0375           double zTruth = truthVertex.position4[Acts::CoordinateIndices::eZ];
0376           if (std::abs(z - zTruth) <= m_cfg.vertexDensityWindow) {
0377             ++count;
0378           }
0379         }
0380         return count / (2 * m_cfg.vertexDensityWindow);
0381       };
0382
0383   if (m_cfg.useTracks) {
0384     tracks = &m_inputTracks(ctx);
0385
0386     for (ConstTrackProxy track : *tracks) {
0387       if (!track.hasReferenceSurface()) {
0388         ACTS_DEBUG("No reference surface on this track, index = "
0389                    << track.index() << " tip index = " << track.tipIndex());
0390         continue;
0391       }
0392
0393       if (std::optional<SimParticle> particle = findParticle(track);
0394           particle.has_value()) {
0395         recoParticles.insert(*particle);
0396       }
0397     }
0398   } else {
0399     // if not using tracks, then all truth particles are associated with the
0400     // vertex
0401     recoParticles = particles;
0402   }
0403
0404   // Exclusive access to the tree while writing
0405   std::lock_guard<std::mutex> lock(m_writeMutex);
0406
0407   m_nRecoVtx = vertices.size();
0408   m_nCleanVtx = 0;
0409   m_nMergedVtx = 0;
0410   m_nSplitVtx = 0;
0411
0412   ACTS_DEBUG("Number of reco vertices in event: " << m_nRecoVtx);
0413
0414   // Get number of generated true primary vertices
0415   m_nTrueVtx = getNumberOfTruePriVertices(particles);
0416   // Get number of detector-accepted true primary vertices
0417   m_nVtxDetAcceptance = getNumberOfTruePriVertices(selectedParticles);
0418
0419   ACTS_DEBUG("Number of truth particles in event : " << particles.size());
0420   ACTS_DEBUG("Number of truth primary vertices : " << m_nTrueVtx);
0421   ACTS_DEBUG("Number of detector-accepted truth primary vertices : "
0422              << m_nVtxDetAcceptance);
0423
0424   // Get the event number
0425   m_eventNr = ctx.eventNumber;
0426
0427   // Get number of track-associated true primary vertices
0428   m_nVtxReconstructable = getNumberOfReconstructableVertices(recoParticles);
0429
0430   ACTS_INFO("Number of reconstructed tracks : "
0431             << ((tracks != nullptr) ? tracks->size() : 0));
0432   ACTS_INFO("Number of reco track-associated truth particles in event : "
0433             << recoParticles.size());
0434   ACTS_INFO("Maximum number of reconstructible primary vertices : "
0435             << m_nVtxReconstructable);
0436
0437   // We compare the reconstructed momenta to the true momenta at the vertex. For
0438   // this, we propagate the reconstructed tracks to the PCA of the true vertex
0439   // position. Setting up propagator:
0440   Acts::SympyStepper stepper(m_cfg.bField);
0441   using Propagator = Acts::Propagator<Acts::SympyStepper>;
0442   auto propagator = std::make_shared<Propagator>(stepper);
0443
0444   struct ToTruthMatching {
0445     std::optional<SimVertexBarcode> vertexId;
0446     double totalTrackWeight{};
0447     double truthMajorityTrackWeights{};
0448     double matchFraction{};
0449
0450     RecoVertexClassification classification{RecoVertexClassification::Unknown};
0451   };
0452   struct ToRecoMatching {
0453     std::size_t recoIndex{};
0454
0455     double recoSumPt2{};
0456   };
0457
0458   std::vector<ToTruthMatching> recoToTruthMatching;
0459   std::map<SimVertexBarcode, ToRecoMatching> truthToRecoMatching;
0460
0461   // Do truth matching for each reconstructed vertex
0462   for (const auto& [vtxIndex, vtx] : Acts::enumerate(vertices)) {
0463     // Reconstructed tracks that contribute to the reconstructed vertex
0464     const auto& tracksAtVtx = vtx.tracks();
0465
0466     // Containers for storing truth particles and truth vertices that
0467     // contribute
0468     // to the reconstructed vertex
0469     std::vector<std::pair<SimVertexBarcode, double>> contributingTruthVertices;
0470
0471     double totalTrackWeight = 0;
0472     for (const Acts::TrackAtVertex& trk : tracksAtVtx) {
0473       if (trk.trackWeight < m_cfg.minTrkWeight) {
0474         continue;
0475       }
0476
0477       totalTrackWeight += trk.trackWeight;
0478
0479       std::optional<ConstTrackProxy> trackOpt = findTrack(trk);
0480       if (!trackOpt.has_value()) {
0481         ACTS_DEBUG("Track has no matching input track.");
0482         continue;
0483       }
0484       const ConstTrackProxy& inputTrk = *trackOpt;
0485       std::optional<SimParticle> particleOpt = findParticle(inputTrk);
0486       if (!particleOpt.has_value()) {
0487         ACTS_VERBOSE("Track has no matching truth particle.");
0488       } else {
0489         contributingTruthVertices.emplace_back(
0490             SimBarcode{particleOpt->particleId()}.vertexId(), trk.trackWeight);
0491       }
0492     }
0493
0494     // Find true vertex that contributes most to the reconstructed vertex
0495     std::map<SimVertexBarcode, std::pair<int, double>> fmap;
0496     for (const auto& [vtxId, weight] : contributingTruthVertices) {
0497       ++fmap[vtxId].first;
0498       fmap[vtxId].second += weight;
0499     }
0500     double truthMajorityVertexTrackWeights = 0;
0501     SimVertexBarcode truthMajorityVertexId{0};
0502     for (const auto& [vtxId, counter] : fmap) {
0503       if (counter.second > truthMajorityVertexTrackWeights) {
0504         truthMajorityVertexId = vtxId;
0505         truthMajorityVertexTrackWeights = counter.second;
0506       }
0507     }
0508
0509     double sumPt2 = calcSumPt2(vtx);
0510
0511     double vertexMatchFraction =
0512         truthMajorityVertexTrackWeights / totalTrackWeight;
0513     RecoVertexClassification recoVertexClassification =
0514         RecoVertexClassification::Unknown;
0515
0516     if (vertexMatchFraction >= m_cfg.vertexMatchThreshold) {
0517       recoVertexClassification = RecoVertexClassification::Clean;
0518     } else {
0519       recoVertexClassification = RecoVertexClassification::Merged;
0520     }
0521
0522     recoToTruthMatching.push_back({truthMajorityVertexId, totalTrackWeight,
0523                                    truthMajorityVertexTrackWeights,
0524                                    vertexMatchFraction,
0525                                    recoVertexClassification});
0526
0527     auto& recoToTruth = recoToTruthMatching.back();
0528
0529     // We have to decide if this reco vertex is a split vertex.
0530     if (auto it = truthToRecoMatching.find(truthMajorityVertexId);
0531         it != truthToRecoMatching.end()) {
0532       // This truth vertex is already matched to a reconstructed vertex so we
0533       // are dealing with a split vertex.
0534
0535       // We have to decide which of the two reconstructed vertices is the
0536       // split vertex.
0537       if (sumPt2 <= it->second.recoSumPt2) {
0538         // Since the sumPt2 is smaller we can simply call this a split vertex
0539
0540         recoToTruth.classification = RecoVertexClassification::Split;
0541
0542         // Keep the existing truth to reco matching
0543       } else {
0544         // The sumPt2 is larger, so we call the other vertex a split vertex.
0545
0546         auto& otherRecoToTruth = recoToTruthMatching.at(it->second.recoIndex);
0547         // Swap the classification
0548         recoToTruth.classification = otherRecoToTruth.classification;
0549         otherRecoToTruth.classification = RecoVertexClassification::Split;
0550
0551         // Overwrite the truth to reco matching
0552         it->second = {vtxIndex, sumPt2};
0553       }
0554     } else {
0555       truthToRecoMatching[truthMajorityVertexId] = {vtxIndex, sumPt2};
0556     }
0557   }
0558
0559   // Loop over reconstructed vertices and see if they can be matched to a true
0560   // vertex.
0561   for (const auto& [vtxIndex, vtx] : Acts::enumerate(vertices)) {
0562     // Reconstructed tracks that contribute to the reconstructed vertex
0563     const auto& tracksAtVtx = vtx.tracks();
0564     // Input tracks matched to `tracksAtVtx`
0565     std::vector<std::uint32_t> trackIndices;
0566
0567     const auto& toTruthMatching = recoToTruthMatching[vtxIndex];
0568
0569     m_recoX.push_back(vtx.fullPosition()[Acts::CoordinateIndices::eX]);
0570     m_recoY.push_back(vtx.fullPosition()[Acts::CoordinateIndices::eY]);
0571     m_recoZ.push_back(vtx.fullPosition()[Acts::CoordinateIndices::eZ]);
0572     m_recoT.push_back(vtx.fullPosition()[Acts::CoordinateIndices::eTime]);
0573
0574     double varX = vtx.fullCovariance()(Acts::CoordinateIndices::eX,
0575                                        Acts::CoordinateIndices::eX);
0576     double varY = vtx.fullCovariance()(Acts::CoordinateIndices::eY,
0577                                        Acts::CoordinateIndices::eY);
0578     double varZ = vtx.fullCovariance()(Acts::CoordinateIndices::eZ,
0579                                        Acts::CoordinateIndices::eZ);
0580     double varTime = vtx.fullCovariance()(Acts::CoordinateIndices::eTime,
0581                                           Acts::CoordinateIndices::eTime);
0582
0583     m_covXX.push_back(varX);
0584     m_covYY.push_back(varY);
0585     m_covZZ.push_back(varZ);
0586     m_covTT.push_back(varTime);
0587     m_covXY.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eX,
0588                                            Acts::CoordinateIndices::eY));
0589     m_covXZ.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eX,
0590                                            Acts::CoordinateIndices::eZ));
0591     m_covXT.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eX,
0592                                            Acts::CoordinateIndices::eTime));
0593     m_covYZ.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eY,
0594                                            Acts::CoordinateIndices::eZ));
0595     m_covYT.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eY,
0596                                            Acts::CoordinateIndices::eTime));
0597     m_covZT.push_back(vtx.fullCovariance()(Acts::CoordinateIndices::eZ,
0598                                            Acts::CoordinateIndices::eTime));
0599
0600     double sumPt2 = calcSumPt2(vtx);
0601     m_sumPt2.push_back(sumPt2);
0602
0603     double recoVertexTrackWeights = 0;
0604     for (const Acts::TrackAtVertex& trk : tracksAtVtx) {
0605       if (trk.trackWeight < m_cfg.minTrkWeight) {
0606         continue;
0607       }
0608       recoVertexTrackWeights += trk.trackWeight;
0609     }
0610     m_recoVertexTrackWeights.push_back(recoVertexTrackWeights);
0611
0612     unsigned int nTracksOnRecoVertex =
0613         std::count_if(tracksAtVtx.begin(), tracksAtVtx.end(), weightHighEnough);
0614     m_nTracksOnRecoVertex.push_back(nTracksOnRecoVertex);
0615
0616     // Saving truth information for the reconstructed vertex
0617     bool truthInfoWritten = false;
0618     std::optional<Acts::Vector4> truthPos;
0619     if (toTruthMatching.vertexId.has_value()) {
0620       auto iTruthVertex = truthVertices.find(toTruthMatching.vertexId.value());
0621       if (iTruthVertex == truthVertices.end()) {
0622         ACTS_ERROR("Truth vertex not found.");
0623         continue;
0624       }
0625       const SimVertex& truthVertex = *iTruthVertex;
0626
0627       m_vertexPrimary.push_back(truthVertex.vertexId().vertexPrimary());
0628       m_vertexSecondary.push_back(truthVertex.vertexId().vertexSecondary());
0629
0630       // Count number of reconstructible tracks on truth vertex
0631       int nTracksOnTruthVertex = 0;
0632       for (const auto& particle : selectedParticles) {
0633         if (static_cast<SimVertexBarcode>(particle.particleId().vertexId()) ==
0634             truthVertex.vertexId()) {
0635           ++nTracksOnTruthVertex;
0636         }
0637       }
0638       m_nTracksOnTruthVertex.push_back(nTracksOnTruthVertex);
0639
0640       double truthPrimaryVertexDensity =
0641           calculateTruthPrimaryVertexDensity(vtx);
0642       m_truthPrimaryVertexDensity.push_back(truthPrimaryVertexDensity);
0643
0644       double truthVertexTrackWeights =
0645           toTruthMatching.truthMajorityTrackWeights;
0646       m_truthVertexTrackWeights.push_back(truthVertexTrackWeights);
0647
0648       double truthVertexMatchRatio = toTruthMatching.matchFraction;
0649       m_truthVertexMatchRatio.push_back(truthVertexMatchRatio);
0650
0651       double recoVertexContamination = 1 - truthVertexMatchRatio;
0652       m_recoVertexContamination.push_back(recoVertexContamination);
0653
0654       RecoVertexClassification recoVertexClassification =
0655           toTruthMatching.classification;
0656       m_recoVertexClassification.push_back(
0657           static_cast<int>(recoVertexClassification));
0658
0659       if (recoVertexClassification == RecoVertexClassification::Clean) {
0660         ++m_nCleanVtx;
0661       } else if (recoVertexClassification == RecoVertexClassification::Merged) {
0662         ++m_nMergedVtx;
0663       } else if (recoVertexClassification == RecoVertexClassification::Split) {
0664         ++m_nSplitVtx;
0665       }
0666
0667       const Acts::Vector4& truePos = truthVertex.position4;
0668       truthPos = truePos;
0669       m_truthX.push_back(truePos[Acts::CoordinateIndices::eX]);
0670       m_truthY.push_back(truePos[Acts::CoordinateIndices::eY]);
0671       m_truthZ.push_back(truePos[Acts::CoordinateIndices::eZ]);
0672       m_truthT.push_back(truePos[Acts::CoordinateIndices::eTime]);
0673
0674       const Acts::Vector4 diffPos = vtx.fullPosition() - truePos;
0675       m_resX.push_back(diffPos[Acts::CoordinateIndices::eX]);
0676       m_resY.push_back(diffPos[Acts::CoordinateIndices::eY]);
0677       m_resZ.push_back(diffPos[Acts::CoordinateIndices::eZ]);
0678       m_resT.push_back(diffPos[Acts::CoordinateIndices::eTime]);
0679
0680       m_pullX.push_back(pull(diffPos[Acts::CoordinateIndices::eX], varX, "X"));
0681       m_pullY.push_back(pull(diffPos[Acts::CoordinateIndices::eY], varY, "Y"));
0682       m_pullZ.push_back(pull(diffPos[Acts::CoordinateIndices::eZ], varZ, "Z"));
0683       m_pullT.push_back(
0684           pull(diffPos[Acts::CoordinateIndices::eTime], varTime, "T"));
0685
0686       truthInfoWritten = true;
0687     }
0688     if (!truthInfoWritten) {
0689       m_vertexPrimary.push_back(-1);
0690       m_vertexSecondary.push_back(-1);
0691
0692       m_nTracksOnTruthVertex.push_back(-1);
0693
0694       m_truthPrimaryVertexDensity.push_back(nan);
0695
0696       m_truthVertexTrackWeights.push_back(nan);
0697       m_truthVertexMatchRatio.push_back(nan);
0698       m_recoVertexContamination.push_back(nan);
0699
0700       m_recoVertexClassification.push_back(
0701           static_cast<int>(RecoVertexClassification::Unknown));
0702
0703       m_truthX.push_back(nan);
0704       m_truthY.push_back(nan);
0705       m_truthZ.push_back(nan);
0706       m_truthT.push_back(nan);
0707
0708       m_resX.push_back(nan);
0709       m_resY.push_back(nan);
0710       m_resZ.push_back(nan);
0711       m_resT.push_back(nan);
0712
0713       m_pullX.push_back(nan);
0714       m_pullY.push_back(nan);
0715       m_pullZ.push_back(nan);
0716       m_pullT.push_back(nan);
0717     }
0718
0719     // Saving the reconstructed/truth momenta. The reconstructed momenta
0720     // are taken at the PCA to the truth vertex position -> we need to
0721     // perform a propagation.
0722
0723     // Get references to inner vectors where all track variables corresponding
0724     // to the current vertex will be saved
0725     auto& innerTrkWeight = m_trkWeight.emplace_back();
0726
0727     auto& innerRecoPhi = m_recoPhi.emplace_back();
0728     auto& innerRecoTheta = m_recoTheta.emplace_back();
0729     auto& innerRecoQOverP = m_recoQOverP.emplace_back();
0730
0731     auto& innerRecoPhiFitted = m_recoPhiFitted.emplace_back();
0732     auto& innerRecoThetaFitted = m_recoThetaFitted.emplace_back();
0733     auto& innerRecoQOverPFitted = m_recoQOverPFitted.emplace_back();
0734
0735     auto& innerTrkParticleId = m_trkParticleId.emplace_back();
0736
0737     auto& innerTruthPhi = m_truthPhi.emplace_back();
0738     auto& innerTruthTheta = m_truthTheta.emplace_back();
0739     auto& innerTruthQOverP = m_truthQOverP.emplace_back();
0740
0741     auto& innerResPhi = m_resPhi.emplace_back();
0742     auto& innerResTheta = m_resTheta.emplace_back();
0743     auto& innerResQOverP = m_resQOverP.emplace_back();
0744
0745     auto& innerResPhiFitted = m_resPhiFitted.emplace_back();
0746     auto& innerResThetaFitted = m_resThetaFitted.emplace_back();
0747     auto& innerResQOverPFitted = m_resQOverPFitted.emplace_back();
0748
0749     auto& innerMomOverlap = m_momOverlap.emplace_back();
0750     auto& innerMomOverlapFitted = m_momOverlapFitted.emplace_back();
0751
0752     auto& innerPullPhi = m_pullPhi.emplace_back();
0753     auto& innerPullTheta = m_pullTheta.emplace_back();
0754     auto& innerPullQOverP = m_pullQOverP.emplace_back();
0755
0756     auto& innerPullPhiFitted = m_pullPhiFitted.emplace_back();
0757     auto& innerPullThetaFitted = m_pullThetaFitted.emplace_back();
0758     auto& innerPullQOverPFitted = m_pullQOverPFitted.emplace_back();
0759
0760     // Perigee at the true vertex position
0761     std::shared_ptr<Acts::PerigeeSurface> perigeeSurface;
0762     if (truthPos.has_value()) {
0763       perigeeSurface =
0764           Acts::Surface::makeShared<Acts::PerigeeSurface>(truthPos->head<3>());
0765     }
0766     // Lambda for propagating the tracks to the PCA
0767     auto propagateToVtx =
0768         [&](const auto& params) -> std::optional<Acts::BoundTrackParameters> {
0769       if (!perigeeSurface) {
0770         return std::nullopt;
0771       }
0772
0773       auto intersection =
0774           perigeeSurface
0775               ->intersect(ctx.geoContext, params.position(ctx.geoContext),
0776                           params.direction(),
0777                           Acts::BoundaryTolerance::Infinite())
0778               .closest();
0779
0780       // Setting the geometry/magnetic field context for the event
0781       using PropagatorOptions = Propagator::Options<>;
0782       PropagatorOptions pOptions(ctx.geoContext, ctx.magFieldContext);
0783       pOptions.direction =
0784           Acts::Direction::fromScalarZeroAsPositive(intersection.pathLength());
0785
0786       auto result = propagator->propagate(params, *perigeeSurface, pOptions);
0787       if (!result.ok()) {
0788         ACTS_ERROR("Propagation to true vertex position failed.");
0789         return std::nullopt;
0790       }
0791       auto& paramsAtVtx = *result->endParameters;
0792       return std::make_optional(paramsAtVtx);
0793     };
0794
0795     for (const Acts::TrackAtVertex& trk : tracksAtVtx) {
0796       if (trk.trackWeight < m_cfg.minTrkWeight) {
0797         continue;
0798       }
0799
0800       innerTrkWeight.push_back(trk.trackWeight);
0801
0802       Acts::Vector3 trueUnitDir = Acts::Vector3::Zero();
0803       Acts::Vector3 trueMom = Acts::Vector3::Zero();
0804
0805       std::optional<SimParticle> particleOpt;
0806       std::optional<ConstTrackProxy> trackOpt = findTrack(trk);
0807       if (trackOpt.has_value()) {
0808         particleOpt = findParticle(*trackOpt);
0809       }
0810
0811       if (particleOpt.has_value()) {
0812         const SimParticle& particle = *particleOpt;
0813
0814         innerTrkParticleId.push_back(particle.particleId().value());
0815
0816         trueUnitDir = particle.direction();
0817         trueMom.head<2>() = Acts::makePhiThetaFromDirection(trueUnitDir);
0818         trueMom[2] = particle.qOverP();
0819
0820         innerTruthPhi.push_back(trueMom[0]);
0821         innerTruthTheta.push_back(trueMom[1]);
0822         innerTruthQOverP.push_back(trueMom[2]);
0823       } else {
0824         ACTS_VERBOSE("Track has no matching truth particle.");
0825
0826         innerTrkParticleId.push_back(-1);
0827
0828         innerTruthPhi.push_back(nan);
0829         innerTruthTheta.push_back(nan);
0830         innerTruthQOverP.push_back(nan);
0831       }
0832
0833       // Save track parameters before the vertex fit
0834       const auto paramsAtVtx = propagateToVtx(
0835           *(trk.originalParams.as<Acts::BoundTrackParameters>()));
0836       if (paramsAtVtx.has_value()) {
0837         Acts::Vector3 recoMom =
0838             paramsAtVtx->parameters().segment(Acts::eBoundPhi, 3);
0839         const Acts::SquareMatrix3& momCov =
0840             paramsAtVtx->covariance()->template block<3, 3>(Acts::eBoundPhi,
0841                                                             Acts::eBoundPhi);
0842         innerRecoPhi.push_back(recoMom[0]);
0843         innerRecoTheta.push_back(recoMom[1]);
0844         innerRecoQOverP.push_back(recoMom[2]);
0845
0846         if (particleOpt.has_value()) {
0847           Acts::Vector3 diffMom = recoMom - trueMom;
0848           // Accounting for the periodicity of phi. We overwrite the
0849           // previously computed value for better readability.
0850           diffMom[0] = Acts::detail::difference_periodic(recoMom(0), trueMom(0),
0851                                                          2 * std::numbers::pi);
0852           innerResPhi.push_back(diffMom[0]);
0853           innerResTheta.push_back(diffMom[1]);
0854           innerResQOverP.push_back(diffMom[2]);
0855
0856           innerPullPhi.push_back(pull(diffMom[0], momCov(0, 0), "phi", false));
0857           innerPullTheta.push_back(
0858               pull(diffMom[1], momCov(1, 1), "theta", false));
0859           innerPullQOverP.push_back(
0860               pull(diffMom[2], momCov(2, 2), "q/p", false));
0861
0862           const auto& recoUnitDir = paramsAtVtx->direction();
0863           double overlap = trueUnitDir.dot(recoUnitDir);
0864           innerMomOverlap.push_back(overlap);
0865         } else {
0866           innerResPhi.push_back(nan);
0867           innerResTheta.push_back(nan);
0868           innerResQOverP.push_back(nan);
0869           innerPullPhi.push_back(nan);
0870           innerPullTheta.push_back(nan);
0871           innerPullQOverP.push_back(nan);
0872           innerMomOverlap.push_back(nan);
0873         }
0874       } else {
0875         innerRecoPhi.push_back(nan);
0876         innerRecoTheta.push_back(nan);
0877         innerRecoQOverP.push_back(nan);
0878         innerResPhi.push_back(nan);
0879         innerResTheta.push_back(nan);
0880         innerResQOverP.push_back(nan);
0881         innerPullPhi.push_back(nan);
0882         innerPullTheta.push_back(nan);
0883         innerPullQOverP.push_back(nan);
0884         innerMomOverlap.push_back(nan);
0885       }
0886
0887       // Save track parameters after the vertex fit
0888       const auto paramsAtVtxFitted = propagateToVtx(trk.fittedParams);
0889       if (paramsAtVtxFitted.has_value()) {
0890         Acts::Vector3 recoMomFitted =
0891             paramsAtVtxFitted->parameters().segment(Acts::eBoundPhi, 3);
0892         const Acts::SquareMatrix3& momCovFitted =
0893             paramsAtVtxFitted->covariance()->block<3, 3>(Acts::eBoundPhi,
0894                                                          Acts::eBoundPhi);
0895         innerRecoPhiFitted.push_back(recoMomFitted[0]);
0896         innerRecoThetaFitted.push_back(recoMomFitted[1]);
0897         innerRecoQOverPFitted.push_back(recoMomFitted[2]);
0898
0899         if (particleOpt.has_value()) {
0900           Acts::Vector3 diffMomFitted = recoMomFitted - trueMom;
0901           // Accounting for the periodicity of phi. We overwrite the
0902           // previously computed value for better readability.
0903           diffMomFitted[0] = Acts::detail::difference_periodic(
0904               recoMomFitted(0), trueMom(0), 2 * std::numbers::pi);
0905           innerResPhiFitted.push_back(diffMomFitted[0]);
0906           innerResThetaFitted.push_back(diffMomFitted[1]);
0907           innerResQOverPFitted.push_back(diffMomFitted[2]);
0908
0909           innerPullPhiFitted.push_back(
0910               pull(diffMomFitted[0], momCovFitted(0, 0), "phi"));
0911           innerPullThetaFitted.push_back(
0912               pull(diffMomFitted[1], momCovFitted(1, 1), "theta"));
0913           innerPullQOverPFitted.push_back(
0914               pull(diffMomFitted[2], momCovFitted(2, 2), "q/p"));
0915
0916           const auto& recoUnitDirFitted = paramsAtVtxFitted->direction();
0917           double overlapFitted = trueUnitDir.dot(recoUnitDirFitted);
0918           innerMomOverlapFitted.push_back(overlapFitted);
0919         } else {
0920           innerResPhiFitted.push_back(nan);
0921           innerResThetaFitted.push_back(nan);
0922           innerResQOverPFitted.push_back(nan);
0923           innerPullPhiFitted.push_back(nan);
0924           innerPullThetaFitted.push_back(nan);
0925           innerPullQOverPFitted.push_back(nan);
0926           innerMomOverlapFitted.push_back(nan);
0927         }
0928       } else {
0929         innerRecoPhiFitted.push_back(nan);
0930         innerRecoThetaFitted.push_back(nan);
0931         innerRecoQOverPFitted.push_back(nan);
0932         innerResPhiFitted.push_back(nan);
0933         innerResThetaFitted.push_back(nan);
0934         innerResQOverPFitted.push_back(nan);
0935         innerPullPhiFitted.push_back(nan);
0936         innerPullThetaFitted.push_back(nan);
0937         innerPullQOverPFitted.push_back(nan);
0938         innerMomOverlapFitted.push_back(nan);
0939       }
0940     }
0941   }
0942
0943   // fill the variables
0944   m_outputTree->Fill();
0945
0946   m_nTracksOnRecoVertex.clear();
0947   m_recoVertexTrackWeights.clear();
0948   m_recoX.clear();
0949   m_recoY.clear();
0950   m_recoZ.clear();
0951   m_recoT.clear();
0952   m_covXX.clear();
0953   m_covYY.clear();
0954   m_covZZ.clear();
0955   m_covTT.clear();
0956   m_covXY.clear();
0957   m_covXZ.clear();
0958   m_covXT.clear();
0959   m_covYZ.clear();
0960   m_covYT.clear();
0961   m_covZT.clear();
0962   m_seedX.clear();
0963   m_seedY.clear();
0964   m_seedZ.clear();
0965   m_seedT.clear();
0966   m_vertexPrimary.clear();
0967   m_vertexSecondary.clear();
0968   m_nTracksOnTruthVertex.clear();
0969   m_truthPrimaryVertexDensity.clear();
0970   m_truthVertexTrackWeights.clear();
0971   m_truthVertexMatchRatio.clear();
0972   m_recoVertexContamination.clear();
0973   m_recoVertexClassification.clear();
0974   m_truthX.clear();
0975   m_truthY.clear();
0976   m_truthZ.clear();
0977   m_truthT.clear();
0978   m_resX.clear();
0979   m_resY.clear();
0980   m_resZ.clear();
0981   m_resT.clear();
0982   m_resSeedZ.clear();
0983   m_resSeedT.clear();
0984   m_pullX.clear();
0985   m_pullY.clear();
0986   m_pullZ.clear();
0987   m_pullT.clear();
0988   m_sumPt2.clear();
0989   m_trkWeight.clear();
0990   m_recoPhi.clear();
0991   m_recoTheta.clear();
0992   m_recoQOverP.clear();
0993   m_recoPhiFitted.clear();
0994   m_recoThetaFitted.clear();
0995   m_recoQOverPFitted.clear();
0996   m_trkParticleId.clear();
0997   m_truthPhi.clear();
0998   m_truthTheta.clear();
0999   m_truthQOverP.clear();
1000   m_resPhi.clear();
1001   m_resTheta.clear();
1002   m_resQOverP.clear();
1003   m_momOverlap.clear();
1004   m_resPhiFitted.clear();
1005   m_resThetaFitted.clear();
1006   m_resQOverPFitted.clear();
1007   m_momOverlapFitted.clear();
1008   m_pullPhi.clear();
1009   m_pullTheta.clear();
1010   m_pullQOverP.clear();
1011   m_pullPhiFitted.clear();
1012   m_pullThetaFitted.clear();
1013   m_pullQOverPFitted.clear();
1014
1015   return ProcessCode::SUCCESS;
1016 }
1017
1018 }  // namespace ActsExamples