algorithms/tracking/CKFTracking.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2025 Whitney Armstrong, Wouter Deconinck, Dmitry Romanov, Shujie Li, Dmitry Kalinkin
0003
0004 #include "CKFTracking.h"
0005
0006 #include <Acts/Definitions/Algebra.hpp>
0007 #include <Acts/Definitions/TrackParametrization.hpp>
0008 #include <Acts/Definitions/Units.hpp>
0009 #include <Acts/Utilities/MathHelpers.hpp>
0010 #include <edm4hep/Vector3f.h>
0011 #if Acts_VERSION_MAJOR >= 46
0012 #else
0013 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
0014 #endif
0015 #include <Acts/EventData/MeasurementHelpers.hpp>
0016 #include <Acts/EventData/ParticleHypothesis.hpp>
0017 #include <Acts/EventData/ProxyAccessor.hpp>
0018 #include <Acts/EventData/SourceLink.hpp>
0019 #include <Acts/EventData/TrackContainer.hpp>
0020 #include <Acts/EventData/TrackProxy.hpp>
0021 #include <Acts/EventData/TrackStatePropMask.hpp>
0022 #include <Acts/EventData/VectorMultiTrajectory.hpp>
0023 #include <Acts/EventData/VectorTrackContainer.hpp>
0024 #include <Acts/Geometry/GeometryContext.hpp>
0025 #include <Acts/Geometry/GeometryHierarchyMap.hpp>
0026 #include <Acts/Geometry/GeometryIdentifier.hpp>
0027 #include <Acts/Propagator/ActorList.hpp>
0028 #include <Acts/Propagator/EigenStepper.hpp>
0029 #include <Acts/Propagator/MaterialInteractor.hpp>
0030 #include <Acts/Propagator/Navigator.hpp>
0031 #include <Acts/Propagator/Propagator.hpp>
0032 #include <Acts/Propagator/PropagatorOptions.hpp>
0033 #include <Acts/Propagator/StandardAborters.hpp>
0034 #include <Acts/Surfaces/PerigeeSurface.hpp>
0035 #include <Acts/Surfaces/Surface.hpp>
0036 #include <Acts/TrackFinding/CombinatorialKalmanFilterExtensions.hpp>
0037 #include <Acts/TrackFinding/TrackStateCreator.hpp>
0038 #include <Acts/TrackFitting/GainMatrixUpdater.hpp>
0039 #include <Acts/Utilities/CalibrationContext.hpp>
0040 #include <Acts/Utilities/Logger.hpp>
0041 #include <Acts/Utilities/TrackHelpers.hpp>
0042 #include <ActsExamples/EventData/GeometryContainers.hpp>
0043 #include <ActsExamples/EventData/IndexSourceLink.hpp>
0044 #include <ActsExamples/EventData/Track.hpp>
0045 #include <boost/container/vector.hpp>
0046 #include <edm4eic/Cov3f.h>
0047 #include <edm4eic/Cov6f.h>
0048 #include <edm4eic/Measurement2DCollection.h>
0049 #include <edm4eic/TrackParametersCollection.h>
0050 #include <edm4eic/TrackSeedCollection.h>
0051 #include <edm4eic/unit_system.h>
0052 #include <edm4hep/Vector2f.h>
0053 #include <spdlog/common.h>
0054 #include <Eigen/Core>
0055 #include <Eigen/Geometry>
0056 #include <Eigen/LU> // IWYU pragma: keep
0057 #include <algorithm>
0058 #include <any>
0059 #include <array>
0060 #include <cmath>
0061 #include <cstddef>
0062 #include <cstdint>
0063 #include <format>
0064 #include <functional>
0065 #include <string>
0066 #include <system_error>
0067 #include <tuple>
0068 #include <utility>
0069 // IWYU pragma: no_include <Acts/Utilities/detail/ContextType.hpp>
0070 // IWYU pragma: no_include <Acts/Utilities/detail/ContainerIterator.hpp>
0071
0072 #include "ActsGeometryProvider.h"
0073 #include "extensions/edm4eic/EDM4eicToActs.h"
0074 #include "extensions/spdlog/SpdlogFormatters.h" // IWYU pragma: keep
0075 #include "extensions/spdlog/SpdlogToActs.h"
0076
0077 namespace {
0078
0079 /// Calibrator that reads directly from edm4eic::Measurement2DCollection.
0080 /// Also owns the geometry-ordered IndexSourceLink multiset built from the same collection.
0081 class EDM4eicMeasurementSourceLinkCalibrator {
0082 public:
0083   EDM4eicMeasurementSourceLinkCalibrator(const edm4eic::Measurement2DCollection* meas2Ds,
0084                                          bool useTime)
0085       : m_meas2Ds(meas2Ds), m_useTime(useTime) {
0086     for (std::size_t index = 0; index < meas2Ds->size(); ++index) {
0087       m_orderedSourceLinks.emplace(Acts::GeometryIdentifier{(*meas2Ds)[index].getSurface()}, index);
0088     }
0089   }
0090
0091   const ActsExamples::GeometryIdMultiset<ActsExamples::IndexSourceLink>&
0092   orderedSourceLinks() const {
0093     return m_orderedSourceLinks;
0094   }
0095
0096   void calibrate(const Acts::GeometryContext& /*gctx*/, const Acts::CalibrationContext& /*cctx*/,
0097                  const Acts::SourceLink& sourceLink,
0098                  Acts::VectorMultiTrajectory::TrackStateProxy trackState) const {
0099     trackState.setUncalibratedSourceLink(Acts::SourceLink{sourceLink});
0100     const auto& idxSourceLink = sourceLink.get<ActsExamples::IndexSourceLink>();
0101     const auto& meas2D        = (*m_meas2Ds)[idxSourceLink.index()];
0102
0103     if (m_useTime) {
0104       calibrateImpl<3>(meas2D, trackState);
0105     } else {
0106       calibrateImpl<2>(meas2D, trackState);
0107     }
0108   }
0109
0110 private:
0111   /// Fill the calibrated measurement with loc0, loc1 (N = 2) and additionally time (N = 3)
0112   template <std::size_t N>
0113   static void calibrateImpl(const edm4eic::Measurement2D& meas2D,
0114                             Acts::VectorMultiTrajectory::TrackStateProxy& trackState) {
0115 #if Acts_VERSION_MAJOR > 45 || (Acts_VERSION_MAJOR == 45 && Acts_VERSION_MINOR >= 2)
0116     Acts::Vector<N> loc       = Acts::Vector<N>::Zero();
0117     Acts::SquareMatrix<N> cov = Acts::SquareMatrix<N>::Zero();
0118 #else
0119     Acts::ActsVector<N> loc       = Acts::ActsVector<N>::Zero();
0120     Acts::ActsSquareMatrix<N> cov = Acts::ActsSquareMatrix<N>::Zero();
0121 #endif
0122     // Bound parameters that make up the calibrated measurement, in storage order
0123     constexpr auto indices = []() {
0124       std::array<uint8_t, N> idx{static_cast<uint8_t>(Acts::eBoundLoc0),
0125                                  static_cast<uint8_t>(Acts::eBoundLoc1)};
0126       if constexpr (N == 3) {
0127         idx[2] = static_cast<uint8_t>(Acts::eBoundTime);
0128       }
0129       return idx;
0130     }();
0131     // Position of a bound parameter within the compact calibrated vector
0132     constexpr auto position = [indices](Acts::BoundIndices bound) {
0133       return static_cast<std::size_t>(std::find(indices.begin(), indices.end(), bound) -
0134                                       indices.begin());
0135     };
0136     constexpr std::size_t iLoc0 = position(Acts::eBoundLoc0);
0137     constexpr std::size_t iLoc1 = position(Acts::eBoundLoc1);
0138
0139     constexpr auto mm  = Acts::UnitConstants::mm / edm4eic::unit::mm;
0140     constexpr auto mm2 = mm * mm;
0141     loc[iLoc0]         = meas2D.getLoc().a * mm;
0142     loc[iLoc1]         = meas2D.getLoc().b * mm;
0143     cov(iLoc0, iLoc0)  = meas2D.getCovariance().xx * mm2;
0144     cov(iLoc1, iLoc1)  = meas2D.getCovariance().yy * mm2;
0145     cov(iLoc0, iLoc1)  = meas2D.getCovariance().xy * mm2;
0146     cov(iLoc1, iLoc0)  = meas2D.getCovariance().xy * mm2;
0147
0148     if constexpr (N == 3) {
0149       constexpr std::size_t iTime = position(Acts::eBoundTime);
0150       constexpr auto ns           = Acts::UnitConstants::ns / edm4eic::unit::ns;
0151       constexpr auto ns2          = ns * ns;
0152       // The time variance is stored in the zz component of the measurement covariance
0153       loc[iTime]        = meas2D.getTime() * ns;
0154       cov(iTime, iTime) = meas2D.getCovariance().zz * ns2;
0155       cov(iLoc0, iTime) = meas2D.getCovariance().xz * mm * ns;
0156       cov(iTime, iLoc0) = cov(iLoc0, iTime);
0157       cov(iLoc1, iTime) = meas2D.getCovariance().yz * mm * ns;
0158       cov(iTime, iLoc1) = cov(iLoc1, iTime);
0159     }
0160
0161     trackState.allocateCalibrated(loc, cov);
0162     trackState.setProjectorSubspaceIndices(indices);
0163   }
0164
0165   const edm4eic::Measurement2DCollection* m_meas2Ds;
0166   bool m_useTime;
0167   ActsExamples::GeometryIdMultiset<ActsExamples::IndexSourceLink> m_orderedSourceLinks;
0168 };
0169
0170 /// Per-branch stopper for the CombinatorialKalmanFilter.
0171 ///
0172 /// ACTS's default branch stopper never stops a branch, so a diverging branch is
0173 /// never abandoned. Once a fit goes bad the covariance inflates, the Kalman
0174 /// gain on q/p becomes ill-conditioned, and each further update throws q/p
0175 /// further off. The covariance eventually overflows double precision and the
0176 /// chi2 becomes NaN, which crashes the MeasurementSelector (seen in the
0177 /// far-forward B0 tracker).
0178 ///
0179 /// The divergence is visible in the filtered variance of q/p long before it
0180 /// becomes fatal: a well-behaved fit never exceeds the variance asserted by the
0181 /// seed, while a diverging one exceeds it by many orders of magnitude, and
0182 /// ultimately turns negative or non-finite as positive-definiteness is lost.
0183 /// Testing that variance therefore catches the divergence at the first bad
0184 /// update, while the chi2 still looks harmless, and costs a single matrix
0185 /// element with no history to inspect.
0186 class CKFBranchStopper {
0187 public:
0188   using TrackProxy      = ActsExamples::TrackContainer::TrackProxy;
0189   using TrackStateProxy = ActsExamples::TrackContainer::TrackStateProxy;
0190   using Result          = Acts::CombinatorialKalmanFilterBranchStopperResult;
0191   using LogFunc         = std::function<void(const std::string&)>;
0192
0193   CKFBranchStopper(const eicrecon::CKFTrackingConfig& cfg, LogFunc log)
0194       : m_cfg(cfg), m_log(std::move(log)) {}
0195
0196   Result operator()(const TrackProxy& track, const TrackStateProxy& trackState) const {
0197 #if Acts_VERSION_MAJOR >= 45
0198     if (!trackState.typeFlags().hasMeasurement() || !trackState.hasFiltered()) {
0199 #else
0200     if (!trackState.typeFlags().test(Acts::TrackStateFlag::MeasurementFlag) ||
0201         !trackState.hasFiltered()) {
0202 #endif
0203       return Result::Continue;
0204     }
0205
0206     const double varQOverP =
0207         trackState.filteredCovariance()(Acts::eBoundQOverP, Acts::eBoundQOverP);
0208     // Catches an inflated variance, and also a negative or NaN one, which mean
0209     // the covariance has lost positive-definiteness altogether.
0210     if (!(varQOverP > 0.) || varQOverP > m_cfg.maxQOverPVariance) {
0211       if (m_log) {
0212         m_log(std::format("Dropping diverged CKF branch with {} measurement(s): "
0213                           "filtered var(q/p) = {:g}, limit {:g}",
0214                           track.nMeasurements(), varQOverP, m_cfg.maxQOverPVariance));
0215       }
0216       return Result::StopAndDrop;
0217     }
0218     return Result::Continue;
0219   }
0220
0221 private:
0222   const eicrecon::CKFTrackingConfig& m_cfg;
0223   LogFunc m_log;
0224 };
0225
0226 } // anonymous namespace
0227
0228 namespace eicrecon {
0229
0230 using namespace Acts::UnitLiterals;
0231
0232 void CKFTracking::init() {
0233   m_acts_logger = Acts::getDefaultLogger(
0234       "CKF", eicrecon::SpdlogToActsLevel(static_cast<spdlog::level::level_enum>(this->level())));
0235
0236   // eta bins, chi2 and #sourclinks per surface cutoffs
0237   m_sourcelinkSelectorCfg = {
0238       {Acts::GeometryIdentifier(),
0239        {.etaBins               = m_cfg.etaBins,
0240         .chi2CutOff            = m_cfg.chi2CutOff,
0241         .numMeasurementsCutOff = {m_cfg.numMeasurementsCutOff.begin(),
0242                                   m_cfg.numMeasurementsCutOff.end()}}},
0243   };
0244   m_trackFinderFunc = CKFTracking::makeCKFTrackingFunction(
0245       m_geoSvc->trackingGeometry(), m_geoSvc->getFieldProvider(), acts_logger());
0246 }
0247
0248 void CKFTracking::process(const Input& input, const Output& output) const {
0249   const auto [init_trk_seeds, meas2Ds]      = input;
0250   auto [output_track_states, output_tracks] = output;
0251
0252   // If measurements or initial track parameters are empty, create empty output containers
0253   if (meas2Ds->empty() || init_trk_seeds->empty()) {
0254     debug("No seeds or measurements, creating empty output containers");
0255     *output_track_states = new Acts::ConstVectorMultiTrajectory();
0256     *output_tracks       = new Acts::ConstVectorTrackContainer();
0257     return;
0258   }
0259
0260   // Time-enabled calibration requires a valid time uncertainty on every measurement
0261   if (m_cfg.useTime) {
0262     bool invalid = false;
0263     for (const auto& meas2D : *meas2Ds) {
0264       const float var_t = meas2D.getCovariance().zz;
0265       if (!std::isfinite(var_t) || var_t <= 0) {
0266         error("UseTime is enabled but measurement on surface {} has invalid time variance {}; "
0267               "skipping tracking for this event",
0268               meas2D.getSurface(), var_t);
0269         invalid = true;
0270         break;
0271       }
0272     }
0273     if (invalid) {
0274       *output_track_states = new Acts::ConstVectorMultiTrajectory();
0275       *output_tracks       = new Acts::ConstVectorTrackContainer();
0276       return;
0277     }
0278   }
0279
0280   ActsExamples::TrackParametersContainer acts_init_trk_params;
0281   for (const auto& track_seed : *init_trk_seeds) {
0282
0283     const auto& track_parameter = track_seed.getParams();
0284
0285     Acts::BoundVector params;
0286     params(Acts::eBoundLoc0) =
0287         track_parameter.getLoc().a * Acts::UnitConstants::mm; // cylinder radius
0288     params(Acts::eBoundLoc1) =
0289         track_parameter.getLoc().b * Acts::UnitConstants::mm; // cylinder length
0290     params(Acts::eBoundPhi)    = track_parameter.getPhi();
0291     params(Acts::eBoundTheta)  = track_parameter.getTheta();
0292     params(Acts::eBoundQOverP) = track_parameter.getQOverP() / Acts::UnitConstants::GeV;
0293     params(Acts::eBoundTime)   = track_parameter.getTime() * Acts::UnitConstants::ns;
0294
0295 #if Acts_VERSION_MAJOR > 45 || (Acts_VERSION_MAJOR == 45 && Acts_VERSION_MINOR >= 1)
0296     Acts::BoundMatrix cov = Acts::BoundMatrix::Zero();
0297 #else
0298     Acts::BoundSquareMatrix cov = Acts::BoundSquareMatrix::Zero();
0299 #endif
0300     for (std::size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0301       for (std::size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0302         cov(a, b) = track_parameter.getCovariance()(i, j) * x * y;
0303         ++j;
0304       }
0305       ++i;
0306     }
0307
0308     // Construct the perigee surface at the seed's reference point.
0309     // For IP-originating tracks this is (0,0,0); for displaced seeds (e.g. from
0310     // Lambda decay daughters) the seed carries the actual decay-vertex position
0311     // so the CKF propagates from there instead of from the IP.
0312     const auto& perigee_pos = track_seed.getPerigee();
0313     auto pSurface           = Acts::Surface::makeShared<const Acts::PerigeeSurface>(Acts::Vector3(
0314         perigee_pos.x * Acts::UnitConstants::mm, perigee_pos.y * Acts::UnitConstants::mm,
0315         perigee_pos.z * Acts::UnitConstants::mm));
0316
0317     // Create parameters
0318     acts_init_trk_params.emplace_back(pSurface, params, cov, Acts::ParticleHypothesis::pion());
0319   }
0320
0321   //// Construct a perigee surface as the target surface
0322   auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0323
0324   // Convert algorithm log level to Acts log level for local logger
0325   const auto spdlog_level = static_cast<spdlog::level::level_enum>(this->level());
0326   const auto acts_level   = eicrecon::SpdlogToActsLevel(spdlog_level);
0327   ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("CKF", acts_level));
0328
0329   // Get run-scoped contexts from service
0330   const auto& gctx = m_geoSvc->getActsGeometryContext();
0331   const auto& mctx = m_geoSvc->getActsMagneticFieldContext();
0332   const auto& cctx = m_geoSvc->getActsCalibrationContext();
0333
0334   Acts::PropagatorPlainOptions pOptions(gctx, mctx);
0335   pOptions.maxSteps = 10000;
0336
0337   EDM4eicMeasurementSourceLinkCalibrator calibratorImpl{meas2Ds, m_cfg.useTime};
0338   Acts::GainMatrixUpdater kfUpdater;
0339   Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0340
0341   Acts::CombinatorialKalmanFilterExtensions<ActsExamples::TrackContainer> extensions;
0342   extensions.updater.connect<&Acts::GainMatrixUpdater::operator()<
0343       typename ActsExamples::TrackContainer::TrackStateContainerBackend>>(&kfUpdater);
0344
0345   ActsExamples::IndexSourceLinkAccessor slAccessor;
0346   slAccessor.container = &calibratorImpl.orderedSourceLinks();
0347   using TrackStateCreatorType =
0348       Acts::TrackStateCreator<ActsExamples::IndexSourceLinkAccessor::Iterator,
0349                               ActsExamples::TrackContainer>;
0350   TrackStateCreatorType trackStateCreator;
0351   trackStateCreator.sourceLinkAccessor
0352       .template connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0353   trackStateCreator.calibrator.template connect<&EDM4eicMeasurementSourceLinkCalibrator::calibrate>(
0354       &calibratorImpl);
0355   trackStateCreator.measurementSelector
0356       .template connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(&measSel);
0357
0358   extensions.createTrackStates.template connect<&TrackStateCreatorType::createTrackStates>(
0359       &trackStateCreator);
0360
0361   // Per-branch stopping (ACTS default never stops a branch).
0362   CKFBranchStopper branchStopper{m_cfg, [this](const std::string& msg) { debug("{}", msg); }};
0363   extensions.branchStopper.connect<&CKFBranchStopper::operator()>(&branchStopper);
0364
0365   // Set the CombinatorialKalmanFilter options
0366   CKFTracking::TrackFinderOptions options(gctx, mctx, cctx, extensions, pOptions);
0367
0368   using Extrapolator        = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
0369   using ExtrapolatorOptions = Extrapolator::template Options<
0370       Acts::ActorList<Acts::MaterialInteractor, Acts::EndOfWorldReached>>;
0371   Extrapolator extrapolator(Acts::EigenStepper<>(m_BField),
0372                             Acts::Navigator({.trackingGeometry = m_geoSvc->trackingGeometry()},
0373                                             acts_logger().cloneWithSuffix("Navigator")),
0374                             acts_logger().cloneWithSuffix("Propagator"));
0375   ExtrapolatorOptions extrapolationOptions(gctx, mctx);
0376
0377   // Create track container
0378   auto trackContainer      = std::make_shared<Acts::VectorTrackContainer>();
0379   auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0380   ActsExamples::TrackContainer acts_tracks(trackContainer, trackStateContainer);
0381
0382   // Create temporary track container
0383   auto trackContainerTemp      = std::make_shared<Acts::VectorTrackContainer>();
0384   auto trackStateContainerTemp = std::make_shared<Acts::VectorMultiTrajectory>();
0385   ActsExamples::TrackContainer acts_tracks_temp(trackContainerTemp, trackStateContainerTemp);
0386
0387   // Add seed number column
0388   acts_tracks.addColumn<unsigned int>("seed");
0389   acts_tracks_temp.addColumn<unsigned int>("seed");
0390   Acts::ProxyAccessor<unsigned int> seedNumber("seed");
0391
0392   // Loop over seeds
0393   for (std::size_t iseed = 0; iseed < acts_init_trk_params.size(); ++iseed) {
0394
0395     // Clear trackContainerTemp and trackStateContainerTemp
0396     acts_tracks_temp.clear();
0397
0398     // Run track finding for this seed
0399     auto result = (*m_trackFinderFunc)(acts_init_trk_params.at(iseed), options, acts_tracks_temp);
0400
0401     if (!result.ok()) {
0402       debug("Track finding failed for seed {} with error {}", iseed, result.error().message());
0403       continue;
0404     }
0405
0406     // Set seed number for all found tracks
0407     auto& tracksForSeed = result.value();
0408     for (auto& track : tracksForSeed) {
0409       // Check if track has at least one valid (non-outlier) measurement
0410       // (this check avoids errors inside smoothing and extrapolation)
0411       auto lastMeasurement = Acts::findLastMeasurementState(track);
0412       if (!lastMeasurement.ok()) {
0413         debug("Track {} for seed {} has no valid measurements, skipping", track.index(), iseed);
0414         continue;
0415       }
0416
0417       if (track.nMeasurements() < m_cfg.numMeasurementsMin) {
0418         trace("Track {} for seed {} has fewer measurements than minimum of {}, skipping",
0419               track.index(), iseed, m_cfg.numMeasurementsMin);
0420         continue;
0421       }
0422
0423       auto smoothingResult = Acts::smoothTrack(gctx, track, acts_logger());
0424       if (!smoothingResult.ok()) {
0425         debug("Smoothing for seed {} and track {} failed with error {}", iseed, track.index(),
0426               smoothingResult.error().message());
0427         continue;
0428       }
0429
0430       auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
0431           track, *pSurface, extrapolator, extrapolationOptions,
0432           Acts::TrackExtrapolationStrategy::firstOrLast, acts_logger());
0433
0434       if (!extrapolationResult.ok()) {
0435         debug("Extrapolation for seed {} and track {} failed with error {}", iseed, track.index(),
0436               extrapolationResult.error().message());
0437         continue;
0438       }
0439
0440       seedNumber(track) = iseed;
0441
0442       // Copy accepted track into main track container
0443       auto acts_tracks_proxy = acts_tracks.makeTrack();
0444       acts_tracks_proxy.copyFrom(track);
0445     }
0446   }
0447
0448   // Allocate new const containers and assign pointers to outputs
0449   *output_track_states = new Acts::ConstVectorMultiTrajectory(std::move(*trackStateContainer));
0450   *output_tracks       = new Acts::ConstVectorTrackContainer(std::move(*trackContainer));
0451 }
0452
0453 } // namespace eicrecon