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 #if Acts_VERSION_MAJOR >= 46
0010 #include <Acts/EventData/BoundTrackParameters.hpp>
0011 #else
0012 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
0013 #endif
0014 #include <Acts/EventData/MeasurementHelpers.hpp>
0015 #include <Acts/EventData/TrackStatePropMask.hpp>
0016 #include <Acts/Geometry/GeometryContext.hpp>
0017 #include <Acts/Geometry/GeometryHierarchyMap.hpp>
0018 #include <Acts/TrackFinding/CombinatorialKalmanFilterExtensions.hpp>
0019 #include <Acts/Utilities/CalibrationContext.hpp>
0020 #include <spdlog/common.h>
0021 #include <algorithm>
0022 #include <any>
0023 #include <array>
0024 #include <cstddef>
0025 #include <cstdint>
0026 #include <functional>
0027 #include <string>
0028 #include <system_error>
0029 #include <tuple>
0030 #include <utility>
0031 #if Acts_VERSION_MAJOR < 43
0032 #include <Acts/Utilities/Iterator.hpp>
0033 #endif
0034 #include <Acts/EventData/ParticleHypothesis.hpp>
0035 #include <Acts/EventData/ProxyAccessor.hpp>
0036 #include <Acts/EventData/SourceLink.hpp>
0037 #include <Acts/EventData/TrackContainer.hpp>
0038 #include <Acts/EventData/TrackProxy.hpp>
0039 #include <Acts/EventData/VectorMultiTrajectory.hpp>
0040 #include <Acts/EventData/VectorTrackContainer.hpp>
0041 #include <Acts/Geometry/GeometryIdentifier.hpp>
0042 #include <Acts/Propagator/ActorList.hpp>
0043 #include <Acts/Propagator/EigenStepper.hpp>
0044 #include <Acts/Propagator/MaterialInteractor.hpp>
0045 #include <Acts/Propagator/Navigator.hpp>
0046 #include <Acts/Propagator/Propagator.hpp>
0047 #include <Acts/Propagator/PropagatorOptions.hpp>
0048 #include <Acts/Propagator/StandardAborters.hpp>
0049 #include <Acts/Surfaces/PerigeeSurface.hpp>
0050 #include <Acts/Surfaces/Surface.hpp>
0051 #include <Acts/TrackFinding/TrackStateCreator.hpp>
0052 #include <Acts/TrackFitting/GainMatrixUpdater.hpp>
0053 #include <Acts/Utilities/Logger.hpp>
0054 #include <Acts/Utilities/TrackHelpers.hpp>
0055 #include <ActsExamples/EventData/GeometryContainers.hpp>
0056 #include <ActsExamples/EventData/IndexSourceLink.hpp>
0057 #include <ActsExamples/EventData/Track.hpp>
0058 #include <boost/container/vector.hpp>
0059 #include <edm4eic/Cov3f.h>
0060 #include <edm4eic/Cov6f.h>
0061 #include <edm4eic/Measurement2DCollection.h>
0062 #include <edm4eic/TrackParametersCollection.h>
0063 #include <edm4eic/TrackSeedCollection.h>
0064 #include <edm4eic/unit_system.h>
0065 #include <edm4hep/Vector2f.h>
0066 #include <Eigen/Core>
0067 #include <Eigen/Geometry>
0068 #include <Eigen/LU> // IWYU pragma: keep
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   explicit EDM4eicMeasurementSourceLinkCalibrator(const edm4eic::Measurement2DCollection* meas2Ds)
0084       : m_meas2Ds(meas2Ds) {
0085     for (std::size_t index = 0; index < meas2Ds->size(); ++index) {
0086       m_orderedSourceLinks.emplace(Acts::GeometryIdentifier{(*meas2Ds)[index].getSurface()}, index);
0087     }
0088   }
0089
0090   const ActsExamples::GeometryIdMultiset<ActsExamples::IndexSourceLink>&
0091   orderedSourceLinks() const {
0092     return m_orderedSourceLinks;
0093   }
0094
0095   void calibrate(const Acts::GeometryContext& /*gctx*/, const Acts::CalibrationContext& /*cctx*/,
0096                  const Acts::SourceLink& sourceLink,
0097                  Acts::VectorMultiTrajectory::TrackStateProxy trackState) const {
0098     trackState.setUncalibratedSourceLink(Acts::SourceLink{sourceLink});
0099     const auto& idxSourceLink = sourceLink.get<ActsExamples::IndexSourceLink>();
0100     const auto& meas2D        = (*m_meas2Ds)[idxSourceLink.index()];
0101
0102 #if Acts_VERSION_MAJOR > 45 || (Acts_VERSION_MAJOR == 45 && Acts_VERSION_MINOR >= 2)
0103     Acts::Vector<2> loc       = Acts::Vector2::Zero();
0104     Acts::SquareMatrix<2> cov = Acts::SquareMatrix<2>::Zero();
0105 #else
0106     Acts::ActsVector<2> loc       = Acts::Vector2::Zero();
0107     Acts::ActsSquareMatrix<2> cov = Acts::ActsSquareMatrix<2>::Zero();
0108 #endif
0109     constexpr auto mm                       = Acts::UnitConstants::mm / edm4eic::unit::mm;
0110     constexpr auto mm2                      = mm * mm;
0111     loc[Acts::eBoundLoc0]                   = meas2D.getLoc().a * mm;
0112     loc[Acts::eBoundLoc1]                   = meas2D.getLoc().b * mm;
0113     cov(Acts::eBoundLoc0, Acts::eBoundLoc0) = meas2D.getCovariance().xx * mm2;
0114     cov(Acts::eBoundLoc1, Acts::eBoundLoc1) = meas2D.getCovariance().yy * mm2;
0115     cov(Acts::eBoundLoc0, Acts::eBoundLoc1) = meas2D.getCovariance().xy * mm2;
0116     cov(Acts::eBoundLoc1, Acts::eBoundLoc0) = meas2D.getCovariance().xy * mm2;
0117
0118     trackState.allocateCalibrated(loc, cov);
0119     std::array<uint8_t, 2> indices{static_cast<uint8_t>(Acts::eBoundLoc0),
0120                                    static_cast<uint8_t>(Acts::eBoundLoc1)};
0121     trackState.setProjectorSubspaceIndices(indices);
0122   }
0123
0124 private:
0125   const edm4eic::Measurement2DCollection* m_meas2Ds;
0126   ActsExamples::GeometryIdMultiset<ActsExamples::IndexSourceLink> m_orderedSourceLinks;
0127 };
0128
0129 } // anonymous namespace
0130
0131 namespace eicrecon {
0132
0133 using namespace Acts::UnitLiterals;
0134
0135 void CKFTracking::init() {
0136   m_acts_logger = Acts::getDefaultLogger(
0137       "CKF", eicrecon::SpdlogToActsLevel(static_cast<spdlog::level::level_enum>(this->level())));
0138
0139   // eta bins, chi2 and #sourclinks per surface cutoffs
0140   m_sourcelinkSelectorCfg = {
0141       {Acts::GeometryIdentifier(),
0142        {.etaBins               = m_cfg.etaBins,
0143         .chi2CutOff            = m_cfg.chi2CutOff,
0144         .numMeasurementsCutOff = {m_cfg.numMeasurementsCutOff.begin(),
0145                                   m_cfg.numMeasurementsCutOff.end()}}},
0146   };
0147   m_trackFinderFunc = CKFTracking::makeCKFTrackingFunction(
0148       m_geoSvc->trackingGeometry(), m_geoSvc->getFieldProvider(), acts_logger());
0149 }
0150
0151 void CKFTracking::process(const Input& input, const Output& output) const {
0152   const auto [init_trk_seeds, meas2Ds]      = input;
0153   auto [output_track_states, output_tracks] = output;
0154
0155   // If measurements or initial track parameters are empty, create empty output containers
0156   if (meas2Ds->empty() || init_trk_seeds->empty()) {
0157     debug("No seeds or measurements, creating empty output containers");
0158     *output_track_states = new Acts::ConstVectorMultiTrajectory();
0159     *output_tracks       = new Acts::ConstVectorTrackContainer();
0160     return;
0161   }
0162
0163   ActsExamples::TrackParametersContainer acts_init_trk_params;
0164   for (const auto& track_seed : *init_trk_seeds) {
0165
0166     const auto& track_parameter = track_seed.getParams();
0167
0168     Acts::BoundVector params;
0169     params(Acts::eBoundLoc0) =
0170         track_parameter.getLoc().a * Acts::UnitConstants::mm; // cylinder radius
0171     params(Acts::eBoundLoc1) =
0172         track_parameter.getLoc().b * Acts::UnitConstants::mm; // cylinder length
0173     params(Acts::eBoundPhi)    = track_parameter.getPhi();
0174     params(Acts::eBoundTheta)  = track_parameter.getTheta();
0175     params(Acts::eBoundQOverP) = track_parameter.getQOverP() / Acts::UnitConstants::GeV;
0176     params(Acts::eBoundTime)   = track_parameter.getTime() * Acts::UnitConstants::ns;
0177
0178 #if Acts_VERSION_MAJOR > 45 || (Acts_VERSION_MAJOR == 45 && Acts_VERSION_MINOR >= 1)
0179     Acts::BoundMatrix cov = Acts::BoundMatrix::Zero();
0180 #else
0181     Acts::BoundSquareMatrix cov = Acts::BoundSquareMatrix::Zero();
0182 #endif
0183     for (std::size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0184       for (std::size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0185         cov(a, b) = track_parameter.getCovariance()(i, j) * x * y;
0186         ++j;
0187       }
0188       ++i;
0189     }
0190
0191     // Construct a perigee surface as the target surface
0192     auto pSurface = Acts::Surface::makeShared<const Acts::PerigeeSurface>(Acts::Vector3(0, 0, 0));
0193
0194     // Create parameters
0195     acts_init_trk_params.emplace_back(pSurface, params, cov, Acts::ParticleHypothesis::pion());
0196   }
0197
0198   //// Construct a perigee surface as the target surface
0199   auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0200
0201   // Convert algorithm log level to Acts log level for local logger
0202   const auto spdlog_level = static_cast<spdlog::level::level_enum>(this->level());
0203   const auto acts_level   = eicrecon::SpdlogToActsLevel(spdlog_level);
0204   ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("CKF", acts_level));
0205
0206   // Get run-scoped contexts from service
0207   const auto& gctx = m_geoSvc->getActsGeometryContext();
0208   const auto& mctx = m_geoSvc->getActsMagneticFieldContext();
0209   const auto& cctx = m_geoSvc->getActsCalibrationContext();
0210
0211   Acts::PropagatorPlainOptions pOptions(gctx, mctx);
0212   pOptions.maxSteps = 10000;
0213
0214   EDM4eicMeasurementSourceLinkCalibrator calibratorImpl{meas2Ds};
0215   Acts::GainMatrixUpdater kfUpdater;
0216   Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0217
0218   Acts::CombinatorialKalmanFilterExtensions<ActsExamples::TrackContainer> extensions;
0219   extensions.updater.connect<&Acts::GainMatrixUpdater::operator()<
0220       typename ActsExamples::TrackContainer::TrackStateContainerBackend>>(&kfUpdater);
0221
0222   ActsExamples::IndexSourceLinkAccessor slAccessor;
0223   slAccessor.container = &calibratorImpl.orderedSourceLinks();
0224   using TrackStateCreatorType =
0225       Acts::TrackStateCreator<ActsExamples::IndexSourceLinkAccessor::Iterator,
0226                               ActsExamples::TrackContainer>;
0227   TrackStateCreatorType trackStateCreator;
0228   trackStateCreator.sourceLinkAccessor
0229       .template connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0230   trackStateCreator.calibrator.template connect<&EDM4eicMeasurementSourceLinkCalibrator::calibrate>(
0231       &calibratorImpl);
0232   trackStateCreator.measurementSelector
0233       .template connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(&measSel);
0234
0235   extensions.createTrackStates.template connect<&TrackStateCreatorType::createTrackStates>(
0236       &trackStateCreator);
0237
0238   // Set the CombinatorialKalmanFilter options
0239   CKFTracking::TrackFinderOptions options(gctx, mctx, cctx, extensions, pOptions);
0240
0241   using Extrapolator        = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
0242   using ExtrapolatorOptions = Extrapolator::template Options<
0243       Acts::ActorList<Acts::MaterialInteractor, Acts::EndOfWorldReached>>;
0244   Extrapolator extrapolator(Acts::EigenStepper<>(m_BField),
0245                             Acts::Navigator({.trackingGeometry = m_geoSvc->trackingGeometry()},
0246                                             acts_logger().cloneWithSuffix("Navigator")),
0247                             acts_logger().cloneWithSuffix("Propagator"));
0248   ExtrapolatorOptions extrapolationOptions(gctx, mctx);
0249
0250   // Create track container
0251   auto trackContainer      = std::make_shared<Acts::VectorTrackContainer>();
0252   auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0253   ActsExamples::TrackContainer acts_tracks(trackContainer, trackStateContainer);
0254
0255   // Create temporary track container
0256   auto trackContainerTemp      = std::make_shared<Acts::VectorTrackContainer>();
0257   auto trackStateContainerTemp = std::make_shared<Acts::VectorMultiTrajectory>();
0258   ActsExamples::TrackContainer acts_tracks_temp(trackContainerTemp, trackStateContainerTemp);
0259
0260   // Add seed number column
0261   acts_tracks.addColumn<unsigned int>("seed");
0262   acts_tracks_temp.addColumn<unsigned int>("seed");
0263   Acts::ProxyAccessor<unsigned int> seedNumber("seed");
0264
0265   // Loop over seeds
0266   for (std::size_t iseed = 0; iseed < acts_init_trk_params.size(); ++iseed) {
0267
0268     // Clear trackContainerTemp and trackStateContainerTemp
0269     acts_tracks_temp.clear();
0270
0271     // Run track finding for this seed
0272     auto result = (*m_trackFinderFunc)(acts_init_trk_params.at(iseed), options, acts_tracks_temp);
0273
0274     if (!result.ok()) {
0275       debug("Track finding failed for seed {} with error {}", iseed, result.error().message());
0276       continue;
0277     }
0278
0279     // Set seed number for all found tracks
0280     auto& tracksForSeed = result.value();
0281     for (auto& track : tracksForSeed) {
0282       // Check if track has at least one valid (non-outlier) measurement
0283       // (this check avoids errors inside smoothing and extrapolation)
0284       auto lastMeasurement = Acts::findLastMeasurementState(track);
0285       if (!lastMeasurement.ok()) {
0286         debug("Track {} for seed {} has no valid measurements, skipping", track.index(), iseed);
0287         continue;
0288       }
0289
0290       if (track.nMeasurements() < m_cfg.numMeasurementsMin) {
0291         trace("Track {} for seed {} has fewer measurements than minimum of {}, skipping",
0292               track.index(), iseed, m_cfg.numMeasurementsMin);
0293         continue;
0294       }
0295
0296       auto smoothingResult = Acts::smoothTrack(gctx, track, acts_logger());
0297       if (!smoothingResult.ok()) {
0298         debug("Smoothing for seed {} and track {} failed with error {}", iseed, track.index(),
0299               smoothingResult.error().message());
0300         continue;
0301       }
0302
0303       auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
0304           track, *pSurface, extrapolator, extrapolationOptions,
0305           Acts::TrackExtrapolationStrategy::firstOrLast, acts_logger());
0306
0307       if (!extrapolationResult.ok()) {
0308         debug("Extrapolation for seed {} and track {} failed with error {}", iseed, track.index(),
0309               extrapolationResult.error().message());
0310         continue;
0311       }
0312
0313       seedNumber(track) = iseed;
0314
0315       // Copy accepted track into main track container
0316       auto acts_tracks_proxy = acts_tracks.makeTrack();
0317       acts_tracks_proxy.copyFrom(track);
0318     }
0319   }
0320
0321   // Allocate new const containers and assign pointers to outputs
0322   *output_track_states = new Acts::ConstVectorMultiTrajectory(std::move(*trackStateContainer));
0323   *output_tracks       = new Acts::ConstVectorTrackContainer(std::move(*trackContainer));
0324 }
0325
0326 } // namespace eicrecon