Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-09-28 07:02:58

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