Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-07-09 07:53:25

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2024 - 2025 Whitney Armstrong, Wouter Deconinck, Dmitry Romanov, Shujie Li, Dmitry Kalinkin
0003 
0004 #include <Acts/Definitions/TrackParametrization.hpp>
0005 #include <Acts/Definitions/Units.hpp>
0006 #include <Acts/EventData/MultiTrajectoryHelpers.hpp>
0007 #include <Acts/EventData/ParticleHypothesis.hpp>
0008 #include <Acts/EventData/TrackStateType.hpp>
0009 #include <ActsExamples/EventData/IndexSourceLink.hpp>
0010 #include <edm4eic/Cov6f.h>
0011 #include <edm4eic/RawTrackerHit.h>
0012 #include <edm4eic/TrackerHit.h>
0013 #include <edm4hep/EDM4hepVersion.h>
0014 #include <edm4hep/MCParticleCollection.h>
0015 #include <edm4hep/SimTrackerHit.h>
0016 #include <edm4hep/Vector2f.h>
0017 #include <edm4hep/Vector3f.h>
0018 #include <fmt/core.h>
0019 #include <podio/ObjectID.h>
0020 #include <podio/RelationRange.h>
0021 #include <Eigen/Core>
0022 #include <array>
0023 #include <cmath>
0024 #include <cstddef>
0025 #include <gsl/pointers>
0026 #include <map>
0027 #include <numeric>
0028 #include <optional>
0029 #include <utility>
0030 
0031 #include "ActsToTracks.h"
0032 
0033 namespace eicrecon {
0034 
0035 // This array relates the Acts and EDM4eic covariance matrices, including
0036 // the unit conversion to get from Acts units into EDM4eic units.
0037 //
0038 // Note: std::map is not constexpr, so we use a constexpr std::array
0039 // std::array initialization need double braces since arrays are aggregates
0040 // ref: https://en.cppreference.com/w/cpp/language/aggregate_initialization
0041 static constexpr std::array<std::pair<Acts::BoundIndices, double>, 6> edm4eic_indexed_units{
0042     {{Acts::eBoundLoc0, Acts::UnitConstants::mm},
0043      {Acts::eBoundLoc1, Acts::UnitConstants::mm},
0044      {Acts::eBoundPhi, 1.},
0045      {Acts::eBoundTheta, 1.},
0046      {Acts::eBoundQOverP, 1. / Acts::UnitConstants::GeV},
0047      {Acts::eBoundTime, Acts::UnitConstants::ns}}};
0048 
0049 void ActsToTracks::init() {}
0050 
0051 void ActsToTracks::process(const Input& input, const Output& output) const {
0052   const auto [meas2Ds, acts_trajectories, raw_hit_assocs]     = input;
0053   auto [trajectories, track_parameters, tracks, tracks_assoc] = output;
0054 
0055   // Loop over trajectories
0056   for (const auto traj : acts_trajectories) {
0057     // The trajectory entry indices and the multiTrajectory
0058     const auto& trackTips = traj->tips();
0059     const auto& mj        = traj->multiTrajectory();
0060     if (trackTips.empty()) {
0061       warning("Empty multiTrajectory.");
0062       continue;
0063     }
0064 
0065     // Loop over all trajectories in a multiTrajectory
0066     for (auto trackTip : trackTips) {
0067       // Collect the trajectory summary info
0068       auto trajectoryState = Acts::MultiTrajectoryHelpers::trajectoryState(mj, trackTip);
0069 
0070       // Check if the reco track has fitted track parameters
0071       if (not traj->hasTrackParameters(trackTip)) {
0072         warning("No fitted track parameters for trajectory with entry index = {}", trackTip);
0073         continue;
0074       }
0075 
0076       // Create trajectory
0077       auto trajectory = trajectories->create();
0078       trajectory.setNMeasurements(trajectoryState.nMeasurements);
0079       trajectory.setNStates(trajectoryState.nStates);
0080       trajectory.setNOutliers(trajectoryState.nOutliers);
0081       trajectory.setNHoles(trajectoryState.nHoles);
0082       trajectory.setNSharedHits(trajectoryState.nSharedHits);
0083 
0084       debug("trajectory state, measurement, outlier, hole: {} {} {} {}", trajectoryState.nStates,
0085             trajectoryState.nMeasurements, trajectoryState.nOutliers, trajectoryState.nHoles);
0086 
0087       for (const auto& measurementChi2 : trajectoryState.measurementChi2) {
0088         trajectory.addToMeasurementChi2(measurementChi2);
0089       }
0090 
0091       for (const auto& outlierChi2 : trajectoryState.outlierChi2) {
0092         trajectory.addToOutlierChi2(outlierChi2);
0093       }
0094 
0095       // Get the fitted track parameter
0096       const auto& boundParam = traj->trackParameters(trackTip);
0097       const auto& parameter  = boundParam.parameters();
0098       const auto& covariance = *boundParam.covariance();
0099 
0100       auto pars = track_parameters->create();
0101       pars.setType(0); // type: track head --> 0
0102       pars.setLoc({static_cast<float>(parameter[Acts::eBoundLoc0]),
0103                    static_cast<float>(parameter[Acts::eBoundLoc1])});
0104       pars.setTheta(static_cast<float>(parameter[Acts::eBoundTheta]));
0105       pars.setPhi(static_cast<float>(parameter[Acts::eBoundPhi]));
0106       pars.setQOverP(static_cast<float>(parameter[Acts::eBoundQOverP]));
0107       pars.setTime(static_cast<float>(parameter[Acts::eBoundTime]));
0108       edm4eic::Cov6f cov;
0109       for (std::size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0110         for (std::size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0111           // FIXME why not pars.getCovariance()(i,j) = covariance(a,b) / x / y;
0112           cov(i, j) = covariance(a, b) / x / y;
0113           ++j;
0114         }
0115         ++i;
0116       }
0117       pars.setCovariance(cov);
0118 
0119       trajectory.addToTrackParameters(pars);
0120 
0121       // Fill tracks
0122       auto track = tracks->create();
0123       track.setType( // Flag that defines the type of track
0124           pars.getType());
0125       track.setPosition( // Track 3-position at the vertex
0126           edm4hep::Vector3f());
0127       track.setMomentum( // Track 3-momentum at the vertex [GeV]
0128           edm4hep::Vector3f());
0129       track.setPositionMomentumCovariance( // Covariance matrix in basis [x,y,z,px,py,pz]
0130           edm4eic::Cov6f());
0131       track.setTime( // Track time at the vertex [ns]
0132           static_cast<float>(parameter[Acts::eBoundTime]));
0133       track.setTimeError( // Error on the track vertex time
0134           sqrt(static_cast<float>(covariance(Acts::eBoundTime, Acts::eBoundTime))));
0135       track.setCharge( // Particle charge
0136           std::copysign(1., parameter[Acts::eBoundQOverP]));
0137       track.setChi2(trajectoryState.chi2Sum); // Total chi2
0138       track.setNdf(trajectoryState.NDF);      // Number of degrees of freedom
0139       track.setPdg(                           // PDG particle ID hypothesis
0140           boundParam.particleHypothesis().absolutePdg());
0141       track.setTrajectory(trajectory); // Trajectory of this track
0142 
0143       // Determine track association with MCParticle, weighted by number of used measurements
0144       std::map<edm4hep::MCParticle, double> mcparticle_weight_by_hit_count;
0145 
0146       // save measurement2d to good measurements or outliers according to srclink index
0147       // fix me: ideally, this should be integrated into multitrajectoryhelper
0148       // fix me: should say "OutlierMeasurements" instead of "OutlierHits" etc
0149       mj.visitBackwards(trackTip, [&](const auto& state) {
0150         auto geoID     = state.referenceSurface().geometryId().value();
0151         auto typeFlags = state.typeFlags();
0152 
0153         // find the associated hit (2D measurement) with state sourcelink index
0154         // fix me: calibrated or not?
0155         if (state.hasUncalibratedSourceLink()) {
0156 
0157           std::size_t srclink_index = state.getUncalibratedSourceLink()
0158                                           .template get<ActsExamples::IndexSourceLink>()
0159                                           .index();
0160 
0161           // no hit on this state/surface, skip
0162           if (typeFlags.test(Acts::TrackStateFlag::HoleFlag)) {
0163             debug("No hit found on geo id={}", geoID);
0164 
0165           } else {
0166             auto meas2D = (*meas2Ds)[srclink_index];
0167             if (typeFlags.test(Acts::TrackStateFlag::OutlierFlag)) {
0168               trajectory.addToOutliers_deprecated(meas2D);
0169               debug("Outlier on geo id={}, index={}, loc={},{}", geoID, srclink_index,
0170                     meas2D.getLoc().a, meas2D.getLoc().b);
0171             } else if (typeFlags.test(Acts::TrackStateFlag::MeasurementFlag)) {
0172               track.addToMeasurements(meas2D);
0173               trajectory.addToMeasurements_deprecated(meas2D);
0174               debug("Measurement on geo id={}, index={}, loc={},{}", geoID, srclink_index,
0175                     meas2D.getLoc().a, meas2D.getLoc().b);
0176 
0177               // Determine track associations if hit associations provided
0178               // FIXME: not able to check whether optional inputs were provided
0179               //if (raw_hit_assocs->has_value()) {
0180               for (const auto& hit : meas2D.getHits()) {
0181                 auto raw_hit = hit.getRawHit();
0182                 for (const auto raw_hit_assoc : *raw_hit_assocs) {
0183                   if (raw_hit_assoc.getRawHit() == raw_hit) {
0184                     auto sim_hit = raw_hit_assoc.getSimHit();
0185 #if EDM4HEP_BUILD_VERSION >= EDM4HEP_VERSION(0, 99, 0)
0186                     auto mc_particle = sim_hit.getParticle();
0187 #else
0188                             auto mc_particle = sim_hit.getMCParticle();
0189 #endif
0190                     mcparticle_weight_by_hit_count[mc_particle]++;
0191                   }
0192                 }
0193               }
0194               //}
0195             }
0196           }
0197         }
0198       });
0199 
0200       // Store track associations if hit associations provided
0201       // FIXME: not able to check whether optional inputs were provided
0202       //if (raw_hit_assocs->has_value()) {
0203       double total_weight = std::accumulate(
0204           mcparticle_weight_by_hit_count.begin(), mcparticle_weight_by_hit_count.end(), 0,
0205           [](const double sum, const auto& i) { return sum + i.second; });
0206       for (const auto& [mcparticle, weight] : mcparticle_weight_by_hit_count) {
0207         auto track_assoc = tracks_assoc->create();
0208         track_assoc.setRec(track);
0209         track_assoc.setSim(mcparticle);
0210         double normalized_weight = weight / total_weight;
0211         track_assoc.setWeight(normalized_weight);
0212         debug("track {}: mcparticle {} weight {}", track.id().index, mcparticle.id().index,
0213               normalized_weight);
0214       }
0215       //}
0216     }
0217   }
0218 }
0219 
0220 } // namespace eicrecon