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 "GaudiAlg/GaudiAlgorithm.h"
0008 #include "GaudiKernel/ToolHandle.h"
0009 #include "GaudiAlg/Transformer.h"
0010 #include "GaudiAlg/GaudiTool.h"
0011 #include "GaudiKernel/RndmGenerators.h"
0012 #include "Gaudi/Property.h"
0013
0014 #include "DDRec/CellIDPositionConverter.h"
0015 #include "DDRec/SurfaceManager.h"
0016 #include "DDRec/Surface.h"
0017
0018 #include "Acts/Geometry/TrackingGeometry.hpp"
0019 #include "Acts/Plugins/DD4hep/DD4hepDetectorElement.hpp"
0020 #include "Acts/Surfaces/PerigeeSurface.hpp"
0021
0022 #include "Acts/TrackFinding/SourceLinkAccessorConcept.hpp"
0023 #include "Acts/TrackFitting/GainMatrixSmoother.hpp"
0024 #include "Acts/TrackFitting/GainMatrixUpdater.hpp"
0025 #include "Acts/Propagator/EigenStepper.hpp"
0026 #include "Acts/Propagator/Navigator.hpp"
0027 #include "Acts/Propagator/Propagator.hpp"
0028 #include "Acts/Definitions/Common.hpp"
0029 #include "Acts/Utilities/Helpers.hpp"
0030 #include "Acts/Utilities/Logger.hpp"
0031 #include "Acts/Definitions/Units.hpp"
0032
0033 #include <k4FWCore/DataHandle.h>
0034 #include <k4Interface/IGeoSvc.h>
0035 #include "JugTrack/IActsGeoSvc.h"
0036 #include "JugTrack/DD4hepBField.h"
0037
0038 #include "ActsExamples/EventData/GeometryContainers.hpp"
0039 #include "ActsExamples/EventData/Measurement.hpp"
0040 #include "ActsExamples/EventData/MeasurementCalibration.hpp"
0041 #include "ActsExamples/EventData/Index.hpp"
0042 #include "ActsExamples/EventData/IndexSourceLink.hpp"
0043 #include "ActsExamples/EventData/Track.hpp"
0044
0045 #include <fmt/core.h>
0046 #include <fmt/ostream.h>
0047
0048 #include "edm4eic/TrackerHitCollection.h"
0049
0050 #include <functional>
0051 #include <stdexcept>
0052 #include <system_error>
0053 #include <type_traits>
0054 #include <vector>
0055 #include <random>
0056 #include <stdexcept>
0057
0058 template<> struct fmt::formatter<std::error_code> : fmt::ostream_formatter {};
0059
0060 static const std::map<int, Acts::Logging::Level> s_msgMap = {
0061     {MSG::DEBUG, Acts::Logging::DEBUG},
0062     {MSG::VERBOSE, Acts::Logging::VERBOSE},
0063     {MSG::INFO, Acts::Logging::INFO},
0064     {MSG::WARNING, Acts::Logging::WARNING},
0065     {MSG::FATAL, Acts::Logging::FATAL},
0066     {MSG::ERROR, Acts::Logging::ERROR},
0067 };
0068
0069 namespace Jug::Reco {
0070
0071   using namespace Acts::UnitLiterals;
0072
0073   CKFTracking::CKFTracking(const std::string& name, ISvcLocator* svcLoc)
0074       : GaudiAlgorithm(name, svcLoc)
0075   {
0076     declareProperty("inputSourceLinks", m_inputSourceLinks, "");
0077     declareProperty("inputMeasurements", m_inputMeasurements, "");
0078     declareProperty("inputInitialTrackParameters", m_inputInitialTrackParameters, "");
0079     declareProperty("outputTracks", m_outputTracks, "");
0080     declareProperty("outputTrajectories", m_outputTrajectories, "");
0081   }
0082
0083   StatusCode CKFTracking::initialize()
0084   {
0085     if (GaudiAlgorithm::initialize().isFailure()) {
0086       return StatusCode::FAILURE;
0087     }
0088     m_geoSvc = service("GeoSvc");
0089     if (!m_geoSvc) {
0090       error() << "Unable to locate Geometry Service. "
0091               << "Make sure you have GeoSvc and SimSvc in the right order in the configuration." << endmsg;
0092       return StatusCode::FAILURE;
0093     }
0094     m_actsGeoSvc = service("ActsGeoSvc");
0095     if (!m_actsGeoSvc) {
0096       error() << "Unable to locate ACTS Geometry Service. "
0097               << "Make sure you have ActsGeoSvc in the right place in the configuration." << endmsg;
0098       return StatusCode::FAILURE;
0099     }
0100
0101     m_BField   = std::dynamic_pointer_cast<const Jug::BField::DD4hepBField>(m_actsGeoSvc->getFieldProvider());
0102     m_fieldctx = Jug::BField::BFieldVariant(m_BField);
0103
0104     // eta bins, chi2 and #sourclinks per surface cutoffs
0105     m_sourcelinkSelectorCfg = {
0106         {Acts::GeometryIdentifier(),
0107             {m_etaBins, m_chi2CutOff,
0108                 {m_numMeasurementsCutOff.begin(), m_numMeasurementsCutOff.end()}
0109             }
0110         },
0111     };
0112     m_trackFinderFunc = CKFTracking::makeCKFTrackingFunction(m_actsGeoSvc->trackingGeometry(), m_BField);
0113     auto im = s_msgMap.find(msgLevel());
0114     if (im != s_msgMap.end()) {
0115         m_actsLoggingLevel = im->second;
0116     }
0117     return StatusCode::SUCCESS;
0118   }
0119
0120   StatusCode CKFTracking::execute()
0121   {
0122     // Read input data
0123     const auto* const src_links       = m_inputSourceLinks.get();
0124     const auto* const init_trk_params = m_inputInitialTrackParameters.get();
0125     const auto* const measurements    = m_inputMeasurements.get();
0126
0127     //// Prepare the output data with MultiTrajectory
0128     auto* acts_trajectories = m_outputTrajectories.createAndPut();
0129     acts_trajectories->reserve(init_trk_params->size());
0130
0131     //// Construct a perigee surface as the target surface
0132     auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0133
0134     ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("CKFTracking Logger", m_actsLoggingLevel));
0135
0136     Acts::PropagatorPlainOptions pOptions;
0137     pOptions.maxSteps = 10000;
0138
0139     ActsExamples::PassThroughCalibrator pcalibrator;
0140     ActsExamples::MeasurementCalibratorAdapter calibrator(pcalibrator, *measurements);
0141     Acts::GainMatrixUpdater kfUpdater;
0142     Acts::GainMatrixSmoother kfSmoother;
0143     Acts::MeasurementSelector measSel{m_sourcelinkSelectorCfg};
0144
0145     Acts::CombinatorialKalmanFilterExtensions<Acts::VectorMultiTrajectory>
0146         extensions;
0147     extensions.calibrator.connect<
0148         &ActsExamples::MeasurementCalibratorAdapter::calibrate>(
0149         &calibrator);
0150     extensions.updater.connect<
0151         &Acts::GainMatrixUpdater::operator()<Acts::VectorMultiTrajectory>>(
0152         &kfUpdater);
0153     extensions.smoother.connect<
0154         &Acts::GainMatrixSmoother::operator()<Acts::VectorMultiTrajectory>>(
0155         &kfSmoother);
0156     extensions.measurementSelector
0157         .connect<&Acts::MeasurementSelector::select<Acts::VectorMultiTrajectory>>(
0158             &measSel);
0159
0160     ActsExamples::IndexSourceLinkAccessor slAccessor;
0161     slAccessor.container = src_links;
0162     Acts::SourceLinkAccessorDelegate<ActsExamples::IndexSourceLinkAccessor::Iterator>
0163         slAccessorDelegate;
0164     slAccessorDelegate.connect<&ActsExamples::IndexSourceLinkAccessor::range>(&slAccessor);
0165
0166     // Set the CombinatorialKalmanFilter options
0167     CKFTracking::TrackFinderOptions options(
0168         m_geoctx, m_fieldctx, m_calibctx, slAccessorDelegate,
0169         extensions, pOptions, &(*pSurface));
0170
0171     // Create track container
0172     auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
0173     auto trackStateContainer = std::make_shared<Acts::VectorMultiTrajectory>();
0174     ActsExamples::TrackContainer tracks(trackContainer, trackStateContainer);
0175
0176     // Add seed number column
0177     tracks.addColumn<unsigned int>("seed");
0178     Acts::TrackAccessor<unsigned int> seedNumber("seed");
0179
0180     // Loop over seeds
0181     for (std::size_t iseed = 0; iseed < init_trk_params->size(); ++iseed) {
0182         auto result =
0183             (*m_trackFinderFunc)(init_trk_params->at(iseed), options, tracks);
0184
0185         if (!result.ok()) {
0186             debug() << fmt::format("Track finding failed for seed {} with error {}", iseed, result.error()) << endmsg;
0187             continue;
0188         }
0189
0190         // Set seed number for all found tracks
0191         auto& tracksForSeed = result.value();
0192         for (auto& track : tracksForSeed) {
0193             seedNumber(track) = iseed;
0194         }
0195     }
0196
0197     // Move track states and track container to const containers
0198     // NOTE Using the non-const containers leads to references to
0199     // implicitly converted temporaries inside the Trajectories.
0200     auto constTrackStateContainer =
0201         std::make_shared<Acts::ConstVectorMultiTrajectory>(
0202             std::move(*trackStateContainer));
0203
0204     auto constTrackContainer =
0205         std::make_shared<Acts::ConstVectorTrackContainer>(
0206             std::move(*trackContainer));
0207
0208     auto* constTracks = m_outputTracks.put(
0209         std::make_unique<ActsExamples::ConstTrackContainer>(constTrackContainer, constTrackStateContainer)
0210     );
0211
0212     // Seed number column accessor
0213     const Acts::ConstTrackAccessor<unsigned int> constSeedNumber("seed");
0214
0215
0216     // Prepare the output data with MultiTrajectory, per seed
0217     ActsExamples::Trajectories::IndexedParameters parameters;
0218     std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
0219
0220     std::optional<unsigned int> lastSeed;
0221     for (const auto& track : *constTracks) {
0222       if (!lastSeed) {
0223         lastSeed = constSeedNumber(track);
0224       }
0225
0226       if (constSeedNumber(track) != lastSeed.value()) {
0227         // make copies and clear vectors
0228         acts_trajectories->emplace_back(
0229           constTracks->trackStateContainer(),
0230           tips, parameters);
0231
0232         tips.clear();
0233         parameters.clear();
0234       }
0235
0236       lastSeed = constSeedNumber(track);
0237
0238       tips.push_back(track.tipIndex());
0239       parameters.emplace(
0240           std::pair{track.tipIndex(),
0241                     ActsExamples::TrackParameters{track.referenceSurface().getSharedPtr(),
0242                                                   track.parameters(), track.covariance(),
0243                                                   track.particleHypothesis()}});
0244     }
0245
0246     if (tips.empty()) {
0247       info() << fmt::format("Last trajectory is empty") << endmsg;
0248     }
0249
0250     // last entry: move vectors
0251     acts_trajectories->emplace_back(
0252       constTracks->trackStateContainer(),
0253       std::move(tips), std::move(parameters));
0254
0255     return StatusCode::SUCCESS;
0256   }
0257
0258   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0259   DECLARE_COMPONENT(CKFTracking)
0260 } // namespace Jug::Reco