algorithms/tracking/ActsToTracks.cc

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