algorithms/tracking/CKFTracking.cc

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