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/TrackParametrization.hpp>
0008 #include <Acts/Definitions/Units.hpp>
0009 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
0010 #if Acts_VERSION_MAJOR < 36
0011 #include <Acts/EventData/Measurement.hpp>
0012 #endif
0013 #include <Acts/EventData/MultiTrajectory.hpp>
0014 #include <Acts/EventData/ParticleHypothesis.hpp>
0015 #if Acts_VERSION_MAJOR >= 32
0016 #include "Acts/EventData/ProxyAccessor.hpp"
0017 #endif
0018 #include <Acts/EventData/SourceLink.hpp>
0019 #include <Acts/EventData/TrackContainer.hpp>
0020 #include <Acts/EventData/TrackProxy.hpp>
0021 #include <Acts/EventData/VectorMultiTrajectory.hpp>
0022 #include <Acts/EventData/VectorTrackContainer.hpp>
0023 #include <Acts/Geometry/GeometryIdentifier.hpp>
0024 #if Acts_VERSION_MAJOR >= 34
0025 #include "Acts/Propagator/AbortList.hpp"
0026 #include "Acts/Propagator/EigenStepper.hpp"
0027 #include "Acts/Propagator/MaterialInteractor.hpp"
0028 #include "Acts/Propagator/Navigator.hpp"
0029 #endif
0030 #include <Acts/Propagator/Propagator.hpp>
0031 #if Acts_VERSION_MAJOR >= 34
0032 #include "Acts/Propagator/StandardAborters.hpp"
0033 #endif
0034 #include <Acts/Surfaces/PerigeeSurface.hpp>
0035 #include <Acts/Surfaces/Surface.hpp>
0036 #include <Acts/TrackFitting/GainMatrixSmoother.hpp>
0037 #include <Acts/TrackFitting/GainMatrixUpdater.hpp>
0038 #include <Acts/Utilities/Logger.hpp>
0039 #if Acts_VERSION_MAJOR >= 34
0040 #include "Acts/Utilities/TrackHelpers.hpp"
0041 #endif
0042 #include <ActsExamples/EventData/IndexSourceLink.hpp>
0043 #include <ActsExamples/EventData/Measurement.hpp>
0044 #include <ActsExamples/EventData/MeasurementCalibration.hpp>
0045 #include <ActsExamples/EventData/Track.hpp>
0046 #include <edm4eic/Cov3f.h>
0047 #include <edm4eic/Cov6f.h>
0048 #include <edm4eic/Measurement2DCollection.h>
0049 #include <edm4eic/TrackParametersCollection.h>
0050 #include <edm4hep/Vector2f.h>
0051 #include <fmt/core.h>
0052 #include <Eigen/Core>
0053 #include <array>
0054 #include <cmath>
0055 #include <cstddef>
0056 #include <functional>
0057 #include <list>
0058 #include <optional>
0059 #include <utility>
0060
0061 #include "ActsGeometryProvider.h"
0062 #include "DD4hepBField.h"
0063 #include "extensions/spdlog/SpdlogFormatters.h" // IWYU pragma: keep
0064 #include "extensions/spdlog/SpdlogToActs.h"
0065
0066 namespace eicrecon {
0067
0068     using namespace Acts::UnitLiterals;
0069
0070     // This array relates the Acts and EDM4eic covariance matrices, including
0071     // the unit conversion to get from Acts units into EDM4eic units.
0072     //
0073     // Note: std::map is not constexpr, so we use a constexpr std::array
0074     // std::array initialization need double braces since arrays are aggregates
0075     // ref: https://en.cppreference.com/w/cpp/language/aggregate_initialization
0076     static constexpr std::array<std::pair<Acts::BoundIndices, double>, 6> edm4eic_indexed_units{{
0077       {Acts::eBoundLoc0, Acts::UnitConstants::mm},
0078       {Acts::eBoundLoc1, Acts::UnitConstants::mm},
0079       {Acts::eBoundPhi, 1.},
0080       {Acts::eBoundTheta, 1.},
0081       {Acts::eBoundQOverP, 1. / Acts::UnitConstants::GeV},
0082       {Acts::eBoundTime, Acts::UnitConstants::ns}
0083     }};
0084
0085     CKFTracking::CKFTracking() {
0086     }
0087
0088     void CKFTracking::init(std::shared_ptr<const ActsGeometryProvider> geo_svc, std::shared_ptr<spdlog::logger> log) {
0089         m_log = log;
0090         m_acts_logger = eicrecon::getSpdlogLogger("CKF", m_log);
0091
0092         m_geoSvc = geo_svc;
0093
0094         m_BField = std::dynamic_pointer_cast<const eicrecon::BField::DD4hepBField>(m_geoSvc->getFieldProvider());
0095         m_fieldctx = eicrecon::BField::BFieldVariant(m_BField);
0096
0097         // eta bins, chi2 and #sourclinks per surface cutoffs
0098         m_sourcelinkSelectorCfg = {
0099                 {Acts::GeometryIdentifier(),
0100                  {m_cfg.etaBins, m_cfg.chi2CutOff,
0101                   {m_cfg.numMeasurementsCutOff.begin(), m_cfg.numMeasurementsCutOff.end()}
0102                  }
0103                 },
0104         };
0105         m_trackFinderFunc = CKFTracking::makeCKFTrackingFunction(m_geoSvc->trackingGeometry(), m_BField, logger());
0106     }
0107
0108     std::tuple<
0109         std::vector<ActsExamples::Trajectories*>,
0110         std::vector<ActsExamples::ConstTrackContainer*>
0111     >
0112     CKFTracking::process(const edm4eic::Measurement2DCollection& meas2Ds,
0113                          const edm4eic::TrackParametersCollection &init_trk_params) {
0114
0115
0116         // create sourcelink and measurement containers
0117         auto measurements = std::make_shared<ActsExamples::MeasurementContainer>();
0118
0119           // need list here for stable addresses
0120         std::list<ActsExamples::IndexSourceLink> sourceLinkStorage;
0121         ActsExamples::IndexSourceLinkContainer src_links;
0122         src_links.reserve(meas2Ds.size());
0123         std::size_t  hit_index = 0;
0124
0125
0126         for (const auto& meas2D : meas2Ds) {
0127
0128             // --follow example from ACTS to create source links
0129             sourceLinkStorage.emplace_back(meas2D.getSurface(), hit_index);
0130             ActsExamples::IndexSourceLink& sourceLink = sourceLinkStorage.back();
0131             // Add to output containers:
0132             // index map and source link container are geometry-ordered.
0133             // since the input is also geometry-ordered, new items can
0134             // be added at the end.
0135             src_links.insert(src_links.end(), sourceLink);
0136             // ---
0137             // Create ACTS measurements
0138             Acts::Vector2 loc = Acts::Vector2::Zero();
0139             loc[Acts::eBoundLoc0] = meas2D.getLoc().a;
0140             loc[Acts::eBoundLoc1] = meas2D.getLoc().b;
0141
0142
0143             Acts::SquareMatrix2 cov = Acts::SquareMatrix2::Zero();
0144             cov(0, 0) = meas2D.getCovariance().xx;
0145             cov(1, 1) = meas2D.getCovariance().yy;
0146             cov(0, 1) = meas2D.getCovariance().xy;
0147             cov(1, 0) = meas2D.getCovariance().xy;
0148
0149 #if Acts_VERSION_MAJOR >= 36
0150             auto measurement = ActsExamples::makeFixedSizeMeasurement(
0151               Acts::SourceLink{sourceLink}, loc, cov, Acts::eBoundLoc0, Acts::eBoundLoc1);
0152 #else
0153             auto measurement = Acts::makeMeasurement(
0154               Acts::SourceLink{sourceLink}, loc, cov, Acts::eBoundLoc0, Acts::eBoundLoc1);
0155 #endif
0156             measurements->emplace_back(std::move(measurement));
0157
0158             hit_index++;
0159         }
0160
0161         ActsExamples::TrackParametersContainer acts_init_trk_params;
0162         for (const auto& track_parameter: init_trk_params) {
0163
0164             Acts::BoundVector params;
0165             params(Acts::eBoundLoc0)   = track_parameter.getLoc().a * Acts::UnitConstants::mm;  // cylinder radius
0166             params(Acts::eBoundLoc1)   = track_parameter.getLoc().b * Acts::UnitConstants::mm;  // cylinder length
0167             params(Acts::eBoundPhi)    = track_parameter.getPhi();
0168             params(Acts::eBoundTheta)  = track_parameter.getTheta();
0169             params(Acts::eBoundQOverP) = track_parameter.getQOverP() / Acts::UnitConstants::GeV;
0170             params(Acts::eBoundTime)   = track_parameter.getTime() * Acts::UnitConstants::ns;
0171
0172             double charge = std::copysign(1., track_parameter.getQOverP());
0173
0174             Acts::BoundSquareMatrix cov = Acts::BoundSquareMatrix::Zero();
0175             for (size_t i = 0; const auto& [a, x] : edm4eic_indexed_units) {
0176               for (size_t j = 0; const auto& [b, y] : edm4eic_indexed_units) {
0177                 cov(a, b) = track_parameter.getCovariance()(i,j) * x * y;
0178                 ++j;
0179               }
0180               ++i;
0181             }
0182
0183             // Construct a perigee surface as the target surface
0184             auto pSurface = Acts::Surface::makeShared<const Acts::PerigeeSurface>(Acts::Vector3(0,0,0));
0185
0186             // Create parameters
0187             acts_init_trk_params.emplace_back(pSurface, params, cov, Acts::ParticleHypothesis::pion());
0188         }
0189
0190         //// Construct a perigee surface as the target surface
0191         auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0192
0193         ACTS_LOCAL_LOGGER(eicrecon::getSpdlogLogger("CKF", m_log, {"^No tracks found$"}));
0194
0195         Acts::PropagatorPlainOptions pOptions;
0196         pOptions.maxSteps = 10000;
0197
0198         ActsExamples::PassThroughCalibrator pcalibrator;
0199         ActsExamples::MeasurementCalibratorAdapter calibrator(pcalibrator, *measurements);
0200         Acts::GainMatrixUpdater kfUpdater;
0201 #if Acts_VERSION_MAJOR < 34
0202         Acts::GainMatrixSmoother kfSmoother;
0203 #endif
0204         Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0205
0206         Acts::CombinatorialKalmanFilterExtensions<Acts::VectorMultiTrajectory>
0207                 extensions;
0208         extensions.calibrator.connect<&ActsExamples::MeasurementCalibratorAdapter::calibrate>(
0209                 &calibrator);
0210         extensions.updater.connect<
0211                 &Acts::GainMatrixUpdater::operator()<Acts::VectorMultiTrajectory>>(
0212                 &kfUpdater);
0213 #if Acts_VERSION_MAJOR < 34
0214         extensions.smoother.connect<
0215                 &Acts::GainMatrixSmoother::operator()<Acts::VectorMultiTrajectory>>(
0216                 &kfSmoother);
0217 #endif
0218         extensions.measurementSelector.connect<
0219                 &Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(
0220                 &measSel);
0221
0222         ActsExamples::IndexSourceLinkAccessor slAccessor;
0223         slAccessor.container = &src_links;
0224         Acts::SourceLinkAccessorDelegate<ActsExamples::IndexSourceLinkAccessor::Iterator>
0225                 slAccessorDelegate;
0226         slAccessorDelegate.connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0227
0228         // Set the CombinatorialKalmanFilter options
0229 #if Acts_VERSION_MAJOR < 34
0230         CKFTracking::TrackFinderOptions options(
0231                 m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0232                 extensions, pOptions, &(*pSurface));
0233 #else
0234         CKFTracking::TrackFinderOptions options(
0235                 m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0236                 extensions, pOptions);
0237 #endif
0238
0239 #if Acts_VERSION_MAJOR >= 34
0240         Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator> extrapolator(
0241             Acts::EigenStepper<>(m_BField),
0242             Acts::Navigator({m_geoSvc->trackingGeometry()},
0243                             logger().cloneWithSuffix("Navigator")),
0244             logger().cloneWithSuffix("Propagator"));
0245
0246         Acts::PropagatorOptions<Acts::ActionList<Acts::MaterialInteractor>,
0247                                 Acts::AbortList<Acts::EndOfWorldReached>>
0248             extrapolationOptions(m_geoctx, m_fieldctx);
0249 #endif
0250
0251         // Create track container
0252         auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
0253         auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0254         ActsExamples::TrackContainer acts_tracks(trackContainer, trackStateContainer);
0255
0256         // Add seed number column
0257         acts_tracks.addColumn<unsigned int>("seed");
0258 #if Acts_VERSION_MAJOR >= 32
0259         Acts::ProxyAccessor<unsigned int> seedNumber("seed");
0260 #else
0261         Acts::TrackAccessor<unsigned int> seedNumber("seed");
0262 #endif
0263
0264         // Loop over seeds
0265         for (std::size_t iseed = 0; iseed < acts_init_trk_params.size(); ++iseed) {
0266             auto result =
0267                 (*m_trackFinderFunc)(acts_init_trk_params.at(iseed), options, acts_tracks);
0268
0269             if (!result.ok()) {
0270                 m_log->debug("Track finding failed for seed {} with error {}", iseed, result.error());
0271                 continue;
0272             }
0273
0274             // Set seed number for all found tracks
0275             auto& tracksForSeed = result.value();
0276             for (auto& track : tracksForSeed) {
0277
0278 #if Acts_VERSION_MAJOR >=34
0279                 auto smoothingResult = Acts::smoothTrack(m_geoctx, track, logger());
0280                 if (!smoothingResult.ok()) {
0281                     ACTS_ERROR("Smoothing for seed "
0282                         << iseed << " and track " << track.index()
0283                         << " failed with error " << smoothingResult.error());
0284                     continue;
0285                 }
0286
0287                 auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
0288                     track, *pSurface, extrapolator, extrapolationOptions,
0289                     Acts::TrackExtrapolationStrategy::firstOrLast, logger());
0290                 if (!extrapolationResult.ok()) {
0291                     ACTS_ERROR("Extrapolation for seed "
0292                         << iseed << " and track " << track.index()
0293                         << " failed with error " << extrapolationResult.error());
0294                     continue;
0295                 }
0296 #endif
0297
0298                 seedNumber(track) = iseed;
0299             }
0300         }
0301
0302
0303         // Move track states and track container to const containers
0304         // NOTE Using the non-const containers leads to references to
0305         // implicitly converted temporaries inside the Trajectories.
0306         auto constTrackStateContainer =
0307             std::make_shared<Acts::ConstVectorMultiTrajectory>(
0308                 std::move(*trackStateContainer));
0309
0310         auto constTrackContainer =
0311             std::make_shared<Acts::ConstVectorTrackContainer>(
0312                 std::move(*trackContainer));
0313
0314         // FIXME JANA2 std::vector<T*> requires wrapping ConstTrackContainer, instead of:
0315         //ConstTrackContainer constTracks(constTrackContainer, constTrackStateContainer);
0316         std::vector<ActsExamples::ConstTrackContainer*> constTracks_v;
0317         constTracks_v.push_back(
0318           new ActsExamples::ConstTrackContainer(
0319             constTrackContainer,
0320             constTrackStateContainer));
0321         auto& constTracks = *(constTracks_v.front());
0322
0323         // Seed number column accessor
0324 #if Acts_VERSION_MAJOR >= 32
0325         const Acts::ConstProxyAccessor<unsigned int> constSeedNumber("seed");
0326 #else
0327         const Acts::ConstTrackAccessor<unsigned int> constSeedNumber("seed");
0328 #endif
0329
0330         // Prepare the output data with MultiTrajectory, per seed
0331         std::vector<ActsExamples::Trajectories*> acts_trajectories;
0332         acts_trajectories.reserve(init_trk_params.size());
0333
0334         ActsExamples::Trajectories::IndexedParameters parameters;
0335         std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
0336
0337         std::optional<unsigned int> lastSeed;
0338         for (const auto& track : constTracks) {
0339           if (!lastSeed) {
0340             lastSeed = constSeedNumber(track);
0341           }
0342
0343           if (constSeedNumber(track) != lastSeed.value()) {
0344             // make copies and clear vectors
0345             acts_trajectories.push_back(new ActsExamples::Trajectories(
0346               constTracks.trackStateContainer(),
0347               tips, parameters));
0348
0349             tips.clear();
0350             parameters.clear();
0351           }
0352
0353           lastSeed = constSeedNumber(track);
0354
0355           tips.push_back(track.tipIndex());
0356           parameters.emplace(
0357               std::pair{track.tipIndex(),
0358                         ActsExamples::TrackParameters{track.referenceSurface().getSharedPtr(),
0359                                                       track.parameters(), track.covariance(),
0360                                                       track.particleHypothesis()}});
0361         }
0362
0363         if (tips.empty()) {
0364           m_log->info("Last trajectory is empty");
0365         }
0366
0367         // last entry: move vectors
0368         acts_trajectories.push_back(new ActsExamples::Trajectories(
0369           constTracks.trackStateContainer(),
0370           std::move(tips), std::move(parameters)));
0371
0372         return std::make_tuple(std::move(acts_trajectories), std::move(constTracks_v));
0373     }
0374
0375 } // namespace eicrecon