src/components/CKFTracking.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Whitney Armstrong, Wouter Deconinck
0003
0004 #include "CKFTracking.h"
0005
0006 // Gaudi
0007 #include "Gaudi/Algorithm.h"
0008 #include "GaudiKernel/ToolHandle.h"
0009 #include "GaudiKernel/RndmGenerators.h"
0010 #include "Gaudi/Property.h"
0011
0012 #include "DDRec/CellIDPositionConverter.h"
0013 #include "DDRec/SurfaceManager.h"
0014 #include "DDRec/Surface.h"
0015
0016 #if Acts_VERSION_MAJOR < 36
0017 #include <Acts/EventData/Measurement.hpp>
0018 #endif
0019 #if Acts_VERSION_MAJOR >= 32
0020 #include "Acts/EventData/ProxyAccessor.hpp"
0021 #endif
0022 #include "Acts/Geometry/TrackingGeometry.hpp"
0023 #include "Acts/Plugins/DD4hep/DD4hepDetectorElement.hpp"
0024 #include "Acts/Surfaces/PerigeeSurface.hpp"
0025
0026 #if Acts_VERSION_MAJOR >= 39
0027 #include <Acts/TrackFinding/TrackStateCreator.hpp>
0028 #endif
0029 #include "Acts/TrackFitting/GainMatrixSmoother.hpp"
0030 #include "Acts/TrackFitting/GainMatrixUpdater.hpp"
0031 #include "Acts/Propagator/EigenStepper.hpp"
0032 #if Acts_VERSION_MAJOR >= 34
0033 #if Acts_VERSION_MAJOR >= 37
0034 #include "Acts/Propagator/ActorList.hpp"
0035 #else
0036 #include "Acts/Propagator/AbortList.hpp"
0037 #endif
0038 #include "Acts/Propagator/EigenStepper.hpp"
0039 #include "Acts/Propagator/MaterialInteractor.hpp"
0040 #include "Acts/Propagator/Navigator.hpp"
0041 #endif
0042 #include "Acts/Propagator/Propagator.hpp"
0043 #if Acts_VERSION_MAJOR >= 34
0044 #include "Acts/Propagator/StandardAborters.hpp"
0045 #endif
0046 #include "Acts/Definitions/Common.hpp"
0047 #include "Acts/Utilities/Helpers.hpp"
0048 #include "Acts/Utilities/Logger.hpp"
0049 #if Acts_VERSION_MAJOR >= 34
0050 #include "Acts/Utilities/TrackHelpers.hpp"
0051 #endif
0052 #include "Acts/Definitions/Units.hpp"
0053
0054 #include <k4FWCore/DataHandle.h>
0055 #include <k4Interface/IGeoSvc.h>
0056 #include "JugTrack/IActsGeoSvc.h"
0057 #include "JugTrack/DD4hepBField.h"
0058
0059 #include "ActsExamples/EventData/GeometryContainers.hpp"
0060 #include "ActsExamples/EventData/Measurement.hpp"
0061 #include "ActsExamples/EventData/MeasurementCalibration.hpp"
0062 #include "ActsExamples/EventData/Index.hpp"
0063 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0064 #include "ActsExamples/EventData/Track.hpp"
0065
0066 #include <fmt/core.h>
0067 #include <fmt/ostream.h>
0068
0069 #include "edm4eic/TrackerHitCollection.h"
0070
0071 #include <functional>
0072 #include <stdexcept>
0073 #include <system_error>
0074 #include <type_traits>
0075 #include <vector>
0076 #include <random>
0077 #include <stdexcept>
0078
0079 template<> struct fmt::formatter<std::error_code> : fmt::ostream_formatter {};
0080
0081 static const std::map<int, Acts::Logging::Level> s_msgMap = {
0082     {MSG::DEBUG, Acts::Logging::DEBUG},
0083     {MSG::VERBOSE, Acts::Logging::VERBOSE},
0084     {MSG::INFO, Acts::Logging::INFO},
0085     {MSG::WARNING, Acts::Logging::WARNING},
0086     {MSG::FATAL, Acts::Logging::FATAL},
0087     {MSG::ERROR, Acts::Logging::ERROR},
0088 };
0089
0090 namespace Jug::Reco {
0091
0092   using namespace Acts::UnitLiterals;
0093
0094   CKFTracking::CKFTracking(const std::string& name, ISvcLocator* svcLoc)
0095       : Gaudi::Algorithm(name, svcLoc)
0096   {
0097 #if Acts_VERSION_MAJOR < 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR < 1)
0098     declareProperty("inputSourceLinks", m_inputSourceLinks, "");
0099 #endif
0100     declareProperty("inputMeasurements", m_inputMeasurements, "");
0101     declareProperty("inputInitialTrackParameters", m_inputInitialTrackParameters, "");
0102     declareProperty("outputTracks", m_outputTracks, "");
0103     declareProperty("outputTrajectories", m_outputTrajectories, "");
0104   }
0105
0106   StatusCode CKFTracking::initialize()
0107   {
0108     if (Gaudi::Algorithm::initialize().isFailure()) {
0109       return StatusCode::FAILURE;
0110     }
0111     m_geoSvc = service("GeoSvc");
0112     if (!m_geoSvc) {
0113       error() << "Unable to locate Geometry Service. "
0114               << "Make sure you have GeoSvc and SimSvc in the right order in the configuration." << endmsg;
0115       return StatusCode::FAILURE;
0116     }
0117     m_actsGeoSvc = service("ActsGeoSvc");
0118     if (!m_actsGeoSvc) {
0119       error() << "Unable to locate ACTS Geometry Service. "
0120               << "Make sure you have ActsGeoSvc in the right place in the configuration." << endmsg;
0121       return StatusCode::FAILURE;
0122     }
0123
0124     m_BField   = std::dynamic_pointer_cast<const Jug::BField::DD4hepBField>(m_actsGeoSvc->getFieldProvider());
0125     m_fieldctx = Jug::BField::BFieldVariant(m_BField);
0126
0127     // eta bins, chi2 and #sourclinks per surface cutoffs
0128     m_sourcelinkSelectorCfg = {
0129         {Acts::GeometryIdentifier(),
0130             {m_etaBins, m_chi2CutOff,
0131                 {m_numMeasurementsCutOff.begin(), m_numMeasurementsCutOff.end()}
0132             }
0133         },
0134     };
0135     m_trackFinderFunc = CKFTracking::makeCKFTrackingFunction(m_actsGeoSvc->trackingGeometry(), m_BField);
0136     auto im = s_msgMap.find(msgLevel());
0137     if (im != s_msgMap.end()) {
0138         m_actsLoggingLevel = im->second;
0139     }
0140     return StatusCode::SUCCESS;
0141   }
0142
0143   StatusCode CKFTracking::execute(const EventContext&) const
0144   {
0145     // Read input data
0146 #if Acts_VERSION_MAJOR < 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR < 1)
0147     const auto* const src_links       = m_inputSourceLinks.get();
0148 #endif
0149     const auto* const init_trk_params = m_inputInitialTrackParameters.get();
0150     const auto* const measurements    = m_inputMeasurements.get();
0151
0152     //// Prepare the output data with MultiTrajectory
0153     auto* acts_trajectories = m_outputTrajectories.createAndPut();
0154     acts_trajectories->reserve(init_trk_params->size());
0155
0156     //// Construct a perigee surface as the target surface
0157     auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0158
0159     ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("CKFTracking Logger", m_actsLoggingLevel));
0160
0161 #if Acts_VERSION_MAJOR >= 36
0162     Acts::PropagatorPlainOptions pOptions(m_geoctx, m_fieldctx);
0163 #else
0164     Acts::PropagatorPlainOptions pOptions;
0165 #endif
0166     pOptions.maxSteps = 10000;
0167
0168     ActsExamples::PassThroughCalibrator pcalibrator;
0169     ActsExamples::MeasurementCalibratorAdapter calibrator(pcalibrator, *measurements);
0170     Acts::GainMatrixUpdater kfUpdater;
0171 #if Acts_VERSION_MAJOR < 34
0172     Acts::GainMatrixSmoother kfSmoother;
0173 #endif
0174     Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0175
0176 #if Acts_VERSION_MAJOR >= 36
0177     Acts::CombinatorialKalmanFilterExtensions<ActsExamples::TrackContainer>
0178         extensions;
0179 #else
0180     Acts::CombinatorialKalmanFilterExtensions<Acts::VectorMultiTrajectory>
0181         extensions;
0182 #endif
0183 #if Acts_VERSION_MAJOR < 39
0184     extensions.calibrator.connect<
0185         &ActsExamples::MeasurementCalibratorAdapter::calibrate>(
0186         &calibrator);
0187 #endif
0188 #if Acts_VERSION_MAJOR >= 36
0189     extensions.updater.connect<
0190         &Acts::GainMatrixUpdater::operator()<
0191         typename ActsExamples::TrackContainer::TrackStateContainerBackend>>(
0192         &kfUpdater);
0193 #else
0194     extensions.updater.connect<
0195         &Acts::GainMatrixUpdater::operator()<Acts::VectorMultiTrajectory>>(
0196         &kfUpdater);
0197 #endif
0198 #if Acts_VERSION_MAJOR < 34
0199     extensions.smoother.connect<
0200         &Acts::GainMatrixSmoother::operator()<Acts::VectorMultiTrajectory>>(
0201         &kfSmoother);
0202 #endif
0203 #if Acts_VERSION_MAJOR < 39
0204     extensions.measurementSelector
0205         .connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(
0206             &measSel);
0207 #endif
0208
0209     ActsExamples::IndexSourceLinkAccessor slAccessor;
0210 #if Acts_VERSION_MAJOR > 37 || (Acts_VERSION_MAJOR == 37 && Acts_VERSION_MINOR >= 1)
0211     slAccessor.container = &measurements->orderedIndices();
0212 #else
0213     slAccessor.container = src_links;
0214 #endif
0215 #if Acts_VERSION_MAJOR >=39
0216         using TrackStateCreatorType =
0217             Acts::TrackStateCreator<ActsExamples::IndexSourceLinkAccessor::Iterator,
0218                                     ActsExamples::TrackContainer>;
0219         TrackStateCreatorType trackStateCreator;
0220         trackStateCreator.sourceLinkAccessor
0221             .template connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0222         trackStateCreator.calibrator
0223             .template connect<&ActsExamples::MeasurementCalibratorAdapter::calibrate>(&calibrator);
0224         trackStateCreator.measurementSelector
0225             .template connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(&measSel);
0226
0227         extensions.createTrackStates
0228             .template connect<&TrackStateCreatorType::createTrackStates>(
0229                 &trackStateCreator);
0230 #else
0231     Acts::SourceLinkAccessorDelegate<ActsExamples::IndexSourceLinkAccessor::Iterator>
0232         slAccessorDelegate;
0233     slAccessorDelegate.connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0234 #endif
0235
0236     // Set the CombinatorialKalmanFilter options
0237 #if Acts_VERSION_MAJOR >= 39
0238         CKFTracking::TrackFinderOptions options(
0239                 m_geoctx, m_fieldctx, m_calibctx,
0240                 extensions, pOptions);
0241 #elif Acts_VERSION_MAJOR >= 34
0242     CKFTracking::TrackFinderOptions options(
0243         m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0244         extensions, pOptions);
0245 #else
0246     CKFTracking::TrackFinderOptions options(
0247         m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0248         extensions, pOptions, &(*pSurface));
0249 #endif
0250
0251 #if Acts_VERSION_MAJOR >= 36
0252     using Extrapolator = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
0253 # if Acts_VERSION_MAJOR >= 37
0254     using ExtrapolatorOptions =
0255         Extrapolator::template Options<Acts::ActorList<Acts::MaterialInteractor,
0256                                                            Acts::EndOfWorldReached>>;
0257 # else
0258     using ExtrapolatorOptions =
0259         Extrapolator::template Options<Acts::ActionList<Acts::MaterialInteractor>,
0260                                            Acts::AbortList<Acts::EndOfWorldReached>>;
0261 # endif
0262     Extrapolator extrapolator(
0263         Acts::EigenStepper<>(m_BField),
0264         Acts::Navigator({m_actsGeoSvc->trackingGeometry()},
0265                         logger().cloneWithSuffix("Navigator")),
0266         logger().cloneWithSuffix("Propagator"));
0267     ExtrapolatorOptions extrapolationOptions(m_geoctx, m_fieldctx);
0268 #elif Acts_VERSION_MAJOR >= 34
0269     Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator> extrapolator(
0270         Acts::EigenStepper<>(m_BField),
0271         Acts::Navigator({m_actsGeoSvc->trackingGeometry()},
0272                         logger().cloneWithSuffix("Navigator")),
0273         logger().cloneWithSuffix("Propagator"));
0274
0275     Acts::PropagatorOptions<Acts::ActionList<Acts::MaterialInteractor>,
0276                             Acts::AbortList<Acts::EndOfWorldReached>>
0277         extrapolationOptions(m_geoctx, m_fieldctx);
0278 #endif
0279
0280     // Create track container
0281     auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
0282     auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0283     ActsExamples::TrackContainer tracks(trackContainer, trackStateContainer);
0284
0285     // Add seed number column
0286     tracks.addColumn<unsigned int>("seed");
0287 #if Acts_VERSION_MAJOR >= 32
0288     Acts::ProxyAccessor<unsigned int> seedNumber("seed");
0289 #else
0290     Acts::TrackAccessor<unsigned int> seedNumber("seed");
0291 #endif
0292     std::set<Acts::TrackIndexType> passed_tracks;
0293
0294     // Loop over seeds
0295     for (std::size_t iseed = 0; iseed < init_trk_params->size(); ++iseed) {
0296         auto result =
0297             (*m_trackFinderFunc)(init_trk_params->at(iseed), options, tracks);
0298
0299         if (!result.ok()) {
0300             debug() << fmt::format("Track finding failed for seed {} with error {}", iseed, result.error()) << endmsg;
0301             continue;
0302         }
0303
0304         // Set seed number for all found tracks
0305         auto& tracksForSeed = result.value();
0306         for (auto& track : tracksForSeed) {
0307
0308 #if Acts_VERSION_MAJOR >=34
0309             auto smoothingResult = Acts::smoothTrack(m_geoctx, track, logger());
0310             if (!smoothingResult.ok()) {
0311                 ACTS_ERROR("Smoothing for seed "
0312                     << iseed << " and track " << track.index()
0313                     << " failed with error " << smoothingResult.error());
0314                 continue;
0315             }
0316
0317             auto extrapolationResult = Acts::extrapolateTrackToReferenceSurface(
0318                 track, *pSurface, extrapolator, extrapolationOptions,
0319                 Acts::TrackExtrapolationStrategy::firstOrLast, logger());
0320             if (!extrapolationResult.ok()) {
0321                 ACTS_ERROR("Extrapolation for seed "
0322                     << iseed << " and track " << track.index()
0323                     << " failed with error " << extrapolationResult.error());
0324                 continue;
0325             }
0326 #endif
0327
0328             passed_tracks.insert(track.index());
0329             seedNumber(track) = iseed;
0330         }
0331     }
0332
0333     for (ssize_t track_index = static_cast<ssize_t>(tracks.size()) - 1; track_index >= 0; track_index--) {
0334         if (not passed_tracks.count(track_index)) {
0335             // NOTE This does not remove track states corresponding to the
0336             // removed tracks. Doing so would require implementing some garbage
0337             // collection. We'll just assume no algorithm will access them
0338             // directly.
0339             tracks.removeTrack(track_index);
0340 #if Acts_VERSION_MAJOR < 36 || (Acts_VERSION_MAJOR == 36 && Acts_VERSION_MINOR < 1)
0341             // Workaround an upstream bug in Acts::VectorTrackContainer::removeTrack_impl()
0342             // https://github.com/acts-project/acts/commit/94cf81f3f1109210b963977e0904516b949b1154
0343             trackContainer->m_particleHypothesis.erase(trackContainer->m_particleHypothesis.begin() + track_index);
0344 #endif
0345         }
0346     }
0347
0348     // Move track states and track container to const containers
0349     // NOTE Using the non-const containers leads to references to
0350     // implicitly converted temporaries inside the Trajectories.
0351     auto constTrackStateContainer =
0352         std::make_shared<Acts::ConstVectorMultiTrajectory>(
0353             std::move(*trackStateContainer));
0354
0355     auto constTrackContainer =
0356         std::make_shared<Acts::ConstVectorTrackContainer>(
0357             std::move(*trackContainer));
0358
0359     auto* constTracks = m_outputTracks.put(
0360         std::make_unique<ActsExamples::ConstTrackContainer>(constTrackContainer, constTrackStateContainer)
0361     );
0362
0363     // Seed number column accessor
0364 #if Acts_VERSION_MAJOR >= 32
0365     const Acts::ConstProxyAccessor<unsigned int> constSeedNumber("seed");
0366 #else
0367     const Acts::ConstTrackAccessor<unsigned int> constSeedNumber("seed");
0368 #endif
0369
0370     // Prepare the output data with MultiTrajectory, per seed
0371     ActsExamples::Trajectories::IndexedParameters parameters;
0372     std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
0373
0374     std::optional<unsigned int> lastSeed;
0375     for (const auto& track : *constTracks) {
0376       if (!lastSeed) {
0377         lastSeed = constSeedNumber(track);
0378       }
0379
0380       if (constSeedNumber(track) != lastSeed.value()) {
0381         // make copies and clear vectors
0382         acts_trajectories->emplace_back(
0383           constTracks->trackStateContainer(),
0384           tips, parameters);
0385
0386         tips.clear();
0387         parameters.clear();
0388       }
0389
0390       lastSeed = constSeedNumber(track);
0391
0392       tips.push_back(track.tipIndex());
0393       parameters.emplace(
0394           std::pair{track.tipIndex(),
0395                     ActsExamples::TrackParameters{track.referenceSurface().getSharedPtr(),
0396                                                   track.parameters(), track.covariance(),
0397                                                   track.particleHypothesis()}});
0398     }
0399
0400     if (tips.empty()) {
0401       info() << fmt::format("Last trajectory is empty") << endmsg;
0402     }
0403
0404     // last entry: move vectors
0405     acts_trajectories->emplace_back(
0406       constTracks->trackStateContainer(),
0407       std::move(tips), std::move(parameters));
0408
0409     return StatusCode::SUCCESS;
0410   }
0411
0412   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0413   DECLARE_COMPONENT(CKFTracking)
0414 } // namespace Jug::Reco