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/Common.hpp>
0008 #include <Acts/Definitions/Direction.hpp>
0009 #include <Acts/Definitions/TrackParametrization.hpp>
0010 #include <Acts/Definitions/Units.hpp>
0011 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
0012 #include <Acts/EventData/MeasurementHelpers.hpp>
0013 #include <Acts/EventData/TrackStatePropMask.hpp>
0014 #include <Acts/EventData/Types.hpp>
0015 #include <Acts/Geometry/GeometryHierarchyMap.hpp>
0016 #include <algorithm>
0017 #include <any>
0018 #include <array>
0019 #include <cstddef>
0020 #include <functional>
0021 #include <optional>
0022 #include <ostream>
0023 #include <set>
0024 #include <stdexcept>
0025 #include <string>
0026 #include <system_error>
0027 #include <utility>
0028 #if Acts_VERSION_MAJOR >= 39
0029 #include <Acts/TrackFinding/CombinatorialKalmanFilterExtensions.hpp>
0030 #endif
0031 #if (Acts_VERSION_MAJOR >= 37) && (Acts_VERSION_MAJOR < 43)
0032 #include <Acts/Utilities/Iterator.hpp>
0033 #endif
0034 #include <Acts/EventData/MultiTrajectory.hpp>
0035 #include <Acts/EventData/ParticleHypothesis.hpp>
0036 #include <Acts/EventData/ProxyAccessor.hpp>
0037 #include <Acts/EventData/SourceLink.hpp>
0038 #include <Acts/EventData/TrackContainer.hpp>
0039 #include <Acts/EventData/TrackProxy.hpp>
0040 #include <Acts/EventData/VectorMultiTrajectory.hpp>
0041 #include <Acts/EventData/VectorTrackContainer.hpp>
0042 #include <Acts/Geometry/GeometryIdentifier.hpp>
0043 #if Acts_VERSION_MAJOR >= 37
0044 #include <Acts/Propagator/ActorList.hpp>
0045 #else
0046 #include <Acts/Propagator/AbortList.hpp>
0047 #include <Acts/Propagator/ActionList.hpp>
0048 #endif
0049 #include <Acts/Propagator/EigenStepper.hpp>
0050 #include <Acts/Propagator/MaterialInteractor.hpp>
0051 #include <Acts/Propagator/Navigator.hpp>
0052 #include <Acts/Propagator/Propagator.hpp>
0053 #include <Acts/Propagator/PropagatorOptions.hpp>
0054 #include <Acts/Propagator/StandardAborters.hpp>
0055 #include <Acts/Surfaces/PerigeeSurface.hpp>
0056 #include <Acts/Surfaces/Surface.hpp>
0057 #if Acts_VERSION_MAJOR >= 39
0058 #include <Acts/TrackFinding/TrackStateCreator.hpp>
0059 #endif
0060 #include <Acts/TrackFitting/GainMatrixUpdater.hpp>
0061 #include <Acts/Utilities/Logger.hpp>
0062 #include <Acts/Utilities/TrackHelpers.hpp>
0063 #include <ActsExamples/EventData/IndexSourceLink.hpp>
0064 #include <ActsExamples/EventData/Measurement.hpp>
0065 #include <ActsExamples/EventData/MeasurementCalibration.hpp>
0066 #include <ActsExamples/EventData/Track.hpp>
0067 #include <boost/container/vector.hpp>
0068 #include <edm4eic/Cov3f.h>
0069 #include <edm4eic/Cov6f.h>
0070 #include <edm4eic/Measurement2DCollection.h>
0071 #include <edm4eic/TrackParametersCollection.h>
0072 #include <edm4hep/Vector2f.h>
0073 #include <fmt/core.h>
0074 #include <fmt/format.h>
0075 #include <Eigen/Core>
0076 #include <Eigen/Geometry>
0077 // IWYU pragma: no_include <Acts/Utilities/detail/ContextType.hpp>
0078 // IWYU pragma: no_include <Acts/Utilities/detail/ContainerIterator.hpp>
0079
0080 #include "ActsGeometryProvider.h"
0081 #include "DD4hepBField.h"
0082 #include "extensions/edm4eic/EDM4eicToActs.h"
0083 #include "extensions/spdlog/SpdlogFormatters.h" // IWYU pragma: keep
0084 #include "extensions/spdlog/SpdlogToActs.h"
0085
0086 namespace eicrecon {
0087
0088 using namespace Acts::UnitLiterals;
0089
0090 CKFTracking::CKFTracking() = default;
0091
0092 void CKFTracking::init(std::shared_ptr<const ActsGeometryProvider> geo_svc,
0093                        std::shared_ptr<spdlog::logger> log) {
0094   m_log         = log;
0095   m_acts_logger = eicrecon::getSpdlogLogger("CKF", m_log);
0096
0097   m_geoSvc = geo_svc;
0098
0099   m_BField =
0100       std::dynamic_pointer_cast<const eicrecon::BField::DD4hepBField>(m_geoSvc->getFieldProvider());
0101   m_fieldctx = eicrecon::BField::BFieldVariant(m_BField);
0102
0103   // eta bins, chi2 and #sourclinks per surface cutoffs
0104   m_sourcelinkSelectorCfg = {
0105       {Acts::GeometryIdentifier(),
0106        {.etaBins               = m_cfg.etaBins,
0107         .chi2CutOff            = m_cfg.chi2CutOff,
0108         .numMeasurementsCutOff = {m_cfg.numMeasurementsCutOff.begin(),
0109                                   m_cfg.numMeasurementsCutOff.end()}}},
0110   };
0111   m_trackFinderFunc =
0112       CKFTracking::makeCKFTrackingFunction(m_geoSvc->trackingGeometry(), m_BField, logger());
0113 }
0114
0115 std::tuple<std::vector<ActsExamples::Trajectories*>,
0116            std::vector<ActsExamples::ConstTrackContainer*>>
0117 CKFTracking::process(const edm4eic::TrackParametersCollection& init_trk_params,
0118                      const edm4eic::Measurement2DCollection& meas2Ds) {
0119
0120   // Create output collections
0121   std::vector<ActsExamples::Trajectories*> acts_trajectories;
0122   // Prepare the output data with MultiTrajectory, per seed
0123   acts_trajectories.reserve(init_trk_params.size());
0124   // FIXME JANA2 std::vector<T*> requires wrapping ConstTrackContainer, instead of:
0125   //ConstTrackContainer constTracks(constTrackContainer, constTrackStateContainer);
0126   std::vector<ActsExamples::ConstTrackContainer*> constTracks_v;
0127
0128   // If measurements or initial track parameters are empty, return early
0129   if (meas2Ds.empty() || init_trk_params.empty()) {
0130     return std::make_tuple(std::move(acts_trajectories), std::move(constTracks_v));
0131   }
0132
0133   // create sourcelink and measurement containers
0134   auto measurements = std::make_shared<ActsExamples::MeasurementContainer>();
0135
0136   // need list here for stable addresses
0137 #if Acts_VERSION_MAJOR < 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR < 1)
0138   std::list<ActsExamples::IndexSourceLink> sourceLinkStorage;
0139   ActsExamples::IndexSourceLinkContainer src_links;
0140   src_links.reserve(meas2Ds.size());
0141   std::size_t hit_index = 0;
0142 #endif
0143
0144   for (const auto& meas2D : meas2Ds) {
0145
0146     Acts::GeometryIdentifier geoId{meas2D.getSurface()};
0147
0148 #if Acts_VERSION_MAJOR < 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR < 1)
0149     // --follow example from ACTS to create source links
0150     sourceLinkStorage.emplace_back(geoId, hit_index);
0151     ActsExamples::IndexSourceLink& sourceLink = sourceLinkStorage.back();
0152     // Add to output containers:
0153     // index map and source link container are geometry-ordered.
0154     // since the input is also geometry-ordered, new items can
0155     // be added at the end.
0156     src_links.insert(src_links.end(), sourceLink);
0157 #endif
0158     // ---
0159     // Create ACTS measurements
0160
0161     Acts::ActsVector<2> loc = Acts::Vector2::Zero();
0162     loc[Acts::eBoundLoc0]   = meas2D.getLoc().a;
0163     loc[Acts::eBoundLoc1]   = meas2D.getLoc().b;
0164
0165     Acts::ActsSquareMatrix<2> cov = Acts::ActsSquareMatrix<2>::Zero();
0166     cov(0, 0)                     = meas2D.getCovariance().xx;
0167     cov(1, 1)                     = meas2D.getCovariance().yy;
0168     cov(0, 1)                     = meas2D.getCovariance().xy;
0169     cov(1, 0)                     = meas2D.getCovariance().xy;
0170
0171 #if Acts_VERSION_MAJOR > 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR >= 1)
0172     std::array<Acts::BoundIndices, 2> indices{Acts::eBoundLoc0, Acts::eBoundLoc1};
0173     Acts::visit_measurement(
0174         indices.size(), [&](auto dim) -> ActsExamples::VariableBoundMeasurementProxy {
0175           if constexpr (dim == indices.size()) {
0176             return ActsExamples::VariableBoundMeasurementProxy{
0177                 measurements->emplaceMeasurement<dim>(geoId, indices, loc, cov)};
0178           } else {
0179             throw std::runtime_error("Dimension not supported in measurement creation");
0180           }
0181         });
0182 #elif Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR == 0
0183     std::array<Acts::BoundIndices, 2> indices{Acts::eBoundLoc0, Acts::eBoundLoc1};
0184     Acts::visit_measurement(
0185         indices.size(), [&](auto dim) -> ActsExamples::VariableBoundMeasurementProxy {
0186           if constexpr (dim == indices.size()) {
0187             return ActsExamples::VariableBoundMeasurementProxy{
0188                 measurements->emplaceMeasurement<dim>(Acts::SourceLink{sourceLink}, indices, loc,
0189                                                       cov)};
0190           } else {
0191             throw std::runtime_error("Dimension not supported in measurement creation");
0192           }
0193         });
0194 #elif Acts_VERSION_MAJOR == 36 && Acts_VERSION_MINOR >= 1
0195     auto measurement = ActsExamples::makeVariableSizeMeasurement(
0196         Acts::SourceLink{sourceLink}, loc, cov, Acts::eBoundLoc0, Acts::eBoundLoc1);
0197     measurements->emplace_back(std::move(measurement));
0198 #elif Acts_VERSION_MAJOR == 36 && Acts_VERSION_MINOR == 0
0199     auto measurement = ActsExamples::makeFixedSizeMeasurement(
0200         Acts::SourceLink{sourceLink}, loc, cov, Acts::eBoundLoc0, Acts::eBoundLoc1);
0201     measurements->emplace_back(std::move(measurement));
0202 #endif
0203
0204 #if Acts_VERSION_MAJOR < 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR < 1)
0205     hit_index++;
0206 #endif
0207   }
0208
0209   ActsExamples::TrackParametersContainer acts_init_trk_params;
0210   for (const auto& track_parameter : init_trk_params) {
0211
0212     Acts::BoundVector params;
0213     params(Acts::eBoundLoc0) =
0214         track_parameter.getLoc().a * Acts::UnitConstants::mm; // cylinder radius
0215     params(Acts::eBoundLoc1) =
0216         track_parameter.getLoc().b * Acts::UnitConstants::mm; // cylinder length
0217     params(Acts::eBoundPhi)    = track_parameter.getPhi();
0218     params(Acts::eBoundTheta)  = track_parameter.getTheta();
0219     params(Acts::eBoundQOverP) = track_parameter.getQOverP() / Acts::UnitConstants::GeV;
0220     params(Acts::eBoundTime)   = track_parameter.getTime() * Acts::UnitConstants::ns;
0221
0222     Acts::BoundSquareMatrix cov = Acts::BoundSquareMatrix::Zero();
0223     for (std::size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0224       for (std::size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0225         cov(a, b) = track_parameter.getCovariance()(i, j) * x * y;
0226         ++j;
0227       }
0228       ++i;
0229     }
0230
0231     // Construct a perigee surface as the target surface
0232     auto pSurface = Acts::Surface::makeShared<const Acts::PerigeeSurface>(Acts::Vector3(0, 0, 0));
0233
0234     // Create parameters
0235     acts_init_trk_params.emplace_back(pSurface, params, cov, Acts::ParticleHypothesis::pion());
0236   }
0237
0238   //// Construct a perigee surface as the target surface
0239   auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0240
0241   ACTS_LOCAL_LOGGER(eicrecon::getSpdlogLogger("CKF", m_log, {"^No tracks found$"}));
0242
0243   Acts::PropagatorPlainOptions pOptions(m_geoctx, m_fieldctx);
0244   pOptions.maxSteps = 10000;
0245
0246   ActsExamples::PassThroughCalibrator pcalibrator;
0247   ActsExamples::MeasurementCalibratorAdapter calibrator(pcalibrator, *measurements);
0248   Acts::GainMatrixUpdater kfUpdater;
0249   Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0250
0251   Acts::CombinatorialKalmanFilterExtensions<ActsExamples::TrackContainer> extensions;
0252 #if Acts_VERSION_MAJOR < 39
0253   extensions.calibrator.connect<&ActsExamples::MeasurementCalibratorAdapter::calibrate>(
0254       &calibrator);
0255 #endif
0256   extensions.updater.connect<&Acts::GainMatrixUpdater::operator()<
0257       typename ActsExamples::TrackContainer::TrackStateContainerBackend>>(&kfUpdater);
0258 #if Acts_VERSION_MAJOR < 39
0259   extensions.measurementSelector.connect<&Acts::MeasurementSelector::select<
0260       typename ActsExamples::TrackContainer::TrackStateContainerBackend>>(&measSel);
0261 #elif Acts_VERSION_MAJOR < 39
0262   extensions.measurementSelector
0263       .connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(&measSel);
0264 #endif
0265
0266   ActsExamples::IndexSourceLinkAccessor slAccessor;
0267 #if Acts_VERSION_MAJOR > 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR >= 1)
0268   slAccessor.container = &measurements->orderedIndices();
0269 #else
0270   slAccessor.container = &src_links;
0271 #endif
0272 #if Acts_VERSION_MAJOR >= 39
0273   using TrackStateCreatorType =
0274       Acts::TrackStateCreator<ActsExamples::IndexSourceLinkAccessor::Iterator,
0275                               ActsExamples::TrackContainer>;
0276   TrackStateCreatorType trackStateCreator;
0277   trackStateCreator.sourceLinkAccessor
0278       .template connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0279   trackStateCreator.calibrator
0280       .template connect<&ActsExamples::MeasurementCalibratorAdapter::calibrate>(&calibrator);
0281   trackStateCreator.measurementSelector
0282       .template connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(&measSel);
0283
0284   extensions.createTrackStates.template connect<&TrackStateCreatorType::createTrackStates>(
0285       &trackStateCreator);
0286 #else
0287   Acts::SourceLinkAccessorDelegate<ActsExamples::IndexSourceLinkAccessor::Iterator>
0288       slAccessorDelegate;
0289   slAccessorDelegate.connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0290 #endif
0291
0292   // Set the CombinatorialKalmanFilter options
0293 #if Acts_VERSION_MAJOR >= 39
0294   CKFTracking::TrackFinderOptions options(m_geoctx, m_fieldctx, m_calibctx, extensions, pOptions);
0295 #else
0296   CKFTracking::TrackFinderOptions options(m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0297                                           extensions, pOptions);
0298 #endif
0299
0300   using Extrapolator = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
0301 #if Acts_VERSION_MAJOR >= 37
0302   using ExtrapolatorOptions = Extrapolator::template Options<
0303       Acts::ActorList<Acts::MaterialInteractor, Acts::EndOfWorldReached>>;
0304 #else
0305   using ExtrapolatorOptions =
0306       Extrapolator::template Options<Acts::ActionList<Acts::MaterialInteractor>,
0307                                      Acts::AbortList<Acts::EndOfWorldReached>>;
0308 #endif
0309   Extrapolator extrapolator(Acts::EigenStepper<>(m_BField),
0310                             Acts::Navigator({.trackingGeometry = m_geoSvc->trackingGeometry()},
0311                                             logger().cloneWithSuffix("Navigator")),
0312                             logger().cloneWithSuffix("Propagator"));
0313   ExtrapolatorOptions extrapolationOptions(m_geoctx, m_fieldctx);
0314
0315   // Create track container
0316   auto trackContainer      = std::make_shared<Acts::VectorTrackContainer>();
0317   auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0318   ActsExamples::TrackContainer acts_tracks(trackContainer, trackStateContainer);
0319
0320   // Add seed number column
0321   acts_tracks.addColumn<unsigned int>("seed");
0322   Acts::ProxyAccessor<unsigned int> seedNumber("seed");
0323   std::set<Acts::TrackIndexType> passed_tracks;
0324
0325   // Loop over seeds
0326   for (std::size_t iseed = 0; iseed < acts_init_trk_params.size(); ++iseed) {
0327     auto result = (*m_trackFinderFunc)(acts_init_trk_params.at(iseed), options, acts_tracks);
0328
0329     if (!result.ok()) {
0330       m_log->debug("Track finding failed for seed {} with error {}", iseed, result.error());
0331       continue;
0332     }
0333
0334     // Set seed number for all found tracks
0335     auto& tracksForSeed = result.value();
0336     for (auto& track : tracksForSeed) {
0337       auto smoothingResult = Acts::smoothTrack(m_geoctx, track, logger());
0338       if (!smoothingResult.ok()) {
0339         ACTS_ERROR("Smoothing for seed " << iseed << " and track " << track.index()
0340                                          << " failed with error " << smoothingResult.error());
0341         continue;
0342       }
0343
0344       auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
0345           track, *pSurface, extrapolator, extrapolationOptions,
0346           Acts::TrackExtrapolationStrategy::firstOrLast, logger());
0347
0348       if (!extrapolationResult.ok()) {
0349         ACTS_ERROR("Extrapolation for seed " << iseed << " and track " << track.index()
0350                                              << " failed with error "
0351                                              << extrapolationResult.error());
0352         continue;
0353       }
0354
0355       passed_tracks.insert(track.index());
0356       seedNumber(track) = iseed;
0357     }
0358   }
0359
0360   for (std::size_t track_index = acts_tracks.size(); (track_index--) != 0U;) {
0361     if (!passed_tracks.contains(track_index)) {
0362       // NOTE This does not remove track states corresponding to the
0363       // removed tracks. Doing so would require implementing some garbage
0364       // collection. We'll just assume no algorithm will access them
0365       // directly.
0366       acts_tracks.removeTrack(track_index);
0367 #if Acts_VERSION_MAJOR == 36 && Acts_VERSION_MINOR < 1
0368       // Workaround an upstream bug in Acts::VectorTrackContainer::removeTrack_impl()
0369       // https://github.com/acts-project/acts/commit/94cf81f3f1109210b963977e0904516b949b1154
0370       trackContainer->m_particleHypothesis.erase(trackContainer->m_particleHypothesis.begin() +
0371                                                  track_index);
0372 #endif
0373     }
0374   }
0375
0376   // Move track states and track container to const containers
0377   // NOTE Using the non-const containers leads to references to
0378   // implicitly converted temporaries inside the Trajectories.
0379   auto constTrackStateContainer =
0380       std::make_shared<Acts::ConstVectorMultiTrajectory>(std::move(*trackStateContainer));
0381
0382   auto constTrackContainer =
0383       std::make_shared<Acts::ConstVectorTrackContainer>(std::move(*trackContainer));
0384
0385   constTracks_v.push_back(
0386       new ActsExamples::ConstTrackContainer(constTrackContainer, constTrackStateContainer));
0387   auto& constTracks = *(constTracks_v.front());
0388
0389   // Seed number column accessor
0390   const Acts::ConstProxyAccessor<unsigned int> constSeedNumber("seed");
0391
0392   ActsExamples::Trajectories::IndexedParameters parameters;
0393   std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
0394
0395   std::optional<unsigned int> lastSeed;
0396   for (const auto& track : constTracks) {
0397     if (!lastSeed) {
0398       lastSeed = constSeedNumber(track);
0399     }
0400
0401     if (constSeedNumber(track) != lastSeed.value()) {
0402       // make copies and clear vectors
0403       acts_trajectories.push_back(
0404           new ActsExamples::Trajectories(constTracks.trackStateContainer(), tips, parameters));
0405
0406       tips.clear();
0407       parameters.clear();
0408     }
0409
0410     lastSeed = constSeedNumber(track);
0411
0412     tips.push_back(track.tipIndex());
0413     parameters.emplace(std::pair{
0414         track.tipIndex(),
0415         ActsExamples::TrackParameters{track.referenceSurface().getSharedPtr(), track.parameters(),
0416                                       track.covariance(), track.particleHypothesis()}});
0417   }
0418
0419   if (tips.empty()) {
0420     m_log->info("Last trajectory is empty");
0421   }
0422
0423   // last entry: move vectors
0424   acts_trajectories.push_back(new ActsExamples::Trajectories(
0425       constTracks.trackStateContainer(), std::move(tips), std::move(parameters)));
0426
0427   return std::make_tuple(std::move(acts_trajectories), std::move(constTracks_v));
0428 }
0429
0430 } // namespace eicrecon