acts/src/TrackFittingAlgorithm.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Whitney Armstrong, Wouter Deconinck
0003
0004 //
0005 #include "TrackFittingAlgorithm.h"
0006
0007 // Gaudi
0008 #include "GaudiAlg/GaudiAlgorithm.h"
0009 #include "GaudiKernel/ToolHandle.h"
0010 #include "GaudiAlg/Transformer.h"
0011 #include "GaudiAlg/GaudiTool.h"
0012 #include "GaudiKernel/RndmGenerators.h"
0013 #include "Gaudi/Property.h"
0014
0015 #include "DDRec/CellIDPositionConverter.h"
0016 #include "DDRec/SurfaceManager.h"
0017 #include "DDRec/Surface.h"
0018
0019 #include "Acts/Geometry/TrackingGeometry.hpp"
0020 #include "Acts/Plugins/DD4hep/DD4hepDetectorElement.hpp"
0021 #include "Acts/Surfaces/PerigeeSurface.hpp"
0022 #include "Acts/Propagator/EigenStepper.hpp"
0023 #include "Acts/Propagator/Navigator.hpp"
0024 #include "Acts/Propagator/Propagator.hpp"
0025 #include "Acts/Definitions/Common.hpp"
0026 #include "Acts/Utilities/Helpers.hpp"
0027 #include "Acts/Utilities/Logger.hpp"
0028 #include "Acts/Definitions/Units.hpp"
0029
0030 #include "JugBase/DataHandle.h"
0031 #include "JugBase/IGeoSvc.h"
0032 #include "JugBase/BField/DD4hepBField.h"
0033
0034 #include "JugTrack/GeometryContainers.hpp"
0035 #include "JugTrack/IndexSourceLink.hpp"
0036 #include "JugTrack/Track.hpp"
0037 #include "JugTrack/Measurement.hpp"
0038
0039 #include "edm4eic/TrackerHitCollection.h"
0040
0041 #include <functional>
0042 #include <stdexcept>
0043 #include <vector>
0044 #include <random>
0045 #include <stdexcept>
0046
0047 namespace Jug::Reco {
0048
0049   using namespace Acts::UnitLiterals;
0050
0051   TrackFittingAlgorithm::TrackFittingAlgorithm(const std::string& name, ISvcLocator* svcLoc)
0052       : GaudiAlgorithm(name, svcLoc)
0053   {
0054     declareProperty("inputSourceLinks", m_inputSourceLinks, "");
0055     declareProperty("initialTrackParameters", m_initialTrackParameters, "");
0056     declareProperty("inputMeasurements", m_inputMeasurements, "");
0057     declareProperty("inputProtoTracks", m_inputProtoTracks, "");
0058     declareProperty("foundTracks", m_foundTracks, "");
0059     declareProperty("outputTrajectories", m_outputTrajectories, "");
0060   }
0061
0062   StatusCode TrackFittingAlgorithm::initialize()
0063   {
0064     if (GaudiAlgorithm::initialize().isFailure()) {
0065       return StatusCode::FAILURE;
0066     }
0067     m_geoSvc = service("GeoSvc");
0068     if (!m_geoSvc) {
0069       error() << "Unable to locate Geometry Service. "
0070               << "Make sure you have GeoSvc and SimSvc in the right order in the configuration." << endmsg;
0071       return StatusCode::FAILURE;
0072     }
0073     m_BField   = std::dynamic_pointer_cast<const Jug::BField::DD4hepBField>(m_geoSvc->getFieldProvider());
0074     m_fieldctx = Jug::BField::BFieldVariant(m_BField);
0075
0076     m_trackFittingFunc = makeTrackFittingFunction(m_geoSvc->trackingGeometry(), m_BField);
0077     return StatusCode::SUCCESS;
0078   }
0079
0080   StatusCode TrackFittingAlgorithm::execute()
0081   {
0082     // Read input data
0083     const auto* const sourceLinks       = m_inputSourceLinks.get();
0084     const auto* const initialParameters = m_initialTrackParameters.get();
0085     const auto* const measurements      = m_inputMeasurements.get();
0086     const auto* const protoTracks       = m_inputProtoTracks.get();
0087     ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("TrackFittingAlgorithm Logger", Acts::Logging::INFO));
0088
0089     // Consistency cross checks
0090     if (protoTracks->size() != initialParameters->size()) {
0091       ACTS_FATAL("Inconsistent number of proto tracks and initial parameters");
0092       return StatusCode::FAILURE;
0093     }
0094
0095     // TrajectoryContainer trajectories;
0096     auto* trajectories = m_outputTrajectories.createAndPut();
0097     trajectories->reserve(initialParameters->size());
0098
0099     // Construct a perigee surface as the target surface
0100     auto pSurface = Acts::Surface::makeShared<Acts::PerigeeSurface>(Acts::Vector3{0., 0., 0.});
0101
0102     Acts::PropagatorPlainOptions pOptions;
0103     pOptions.maxSteps = 10000;
0104     // kfOptions.multipleScattering = m_cfg.multipleScattering;
0105     // kfOptions.energyLoss         = m_cfg.energyLoss;
0106
0107     Acts::KalmanFitterExtensions<Acts::VectorMultiTrajectory> extensions;
0108     MeasurementCalibrator calibrator{*measurements};
0109     extensions.calibrator.connect<&MeasurementCalibrator::calibrate>(&calibrator);
0110     Acts::GainMatrixUpdater kfUpdater;
0111     Acts::GainMatrixSmoother kfSmoother;
0112     extensions.updater.connect<&Acts::GainMatrixUpdater::operator()<Acts::VectorMultiTrajectory>>(
0113         &kfUpdater);
0114     extensions.smoother.connect<&Acts::GainMatrixSmoother::operator()<Acts::VectorMultiTrajectory>>(
0115         &kfSmoother);
0116
0117     Acts::KalmanFitterOptions kfOptions(
0118         m_geoctx, m_fieldctx, m_calibctx, extensions,
0119         Acts::LoggerWrapper{logger()}, Acts::PropagatorPlainOptions(),
0120         &(*pSurface));
0121
0122     // used for processing the data
0123     std::vector<IndexSourceLink>      trackSourceLinks;
0124     std::vector<const Acts::Surface*> surfSequence;
0125
0126     if (msgLevel(MSG::DEBUG)) {
0127       debug() << "initialParams size:  " << initialParameters->size() << endmsg;
0128       debug() << "measurements size:  " << measurements->size() << endmsg;
0129     }
0130
0131     // Perform the track finding for each starting parameter
0132     // @TODO: use seeds from track seeding algorithm as starting parameter
0133     // initial track params and proto tracks might likely have the same size.
0134     //for (std::size_t iseed = 0; iseed < init_trk_params->size(); ++iseed) {
0135     for (std::size_t itrack = 0; itrack < (*protoTracks).size(); ++itrack) {
0136
0137       const auto& protoTrack    = (*protoTracks)[itrack];
0138       const auto& initialParams = (*initialParameters)[itrack];
0139
0140       if (msgLevel(MSG::DEBUG)) {
0141         debug() << "protoTrack size:  " << protoTrack.size() << endmsg;
0142         debug() << "sourceLinks size:  " << sourceLinks->size() << endmsg;
0143       }
0144
0145       trackSourceLinks.clear();
0146       trackSourceLinks.reserve(protoTrack.size());
0147
0148       for (auto hitIndex : protoTrack) {
0149         if (auto it = sourceLinks->nth(hitIndex); it != sourceLinks->end()) {
0150           const IndexSourceLink& sourceLink = *it;
0151           trackSourceLinks.push_back(std::cref(sourceLink));
0152         } else {
0153           ACTS_FATAL("Proto track " << itrack << " contains invalid hit index"
0154                                     << hitIndex);
0155           return StatusCode::FAILURE;
0156         }
0157       }
0158
0159       if (msgLevel(MSG::DEBUG)) {
0160         debug() << "Invoke track fitting ...  " << itrack << endmsg;
0161       }
0162       auto result = fitTrack(trackSourceLinks, initialParams, kfOptions);
0163       if (msgLevel(MSG::DEBUG)) {
0164         debug() << "fitting done." << endmsg;
0165       }
0166       if (result.ok())
0167       {
0168         // Get the fit output object
0169         auto& fitOutput = result.value();
0170         // The track entry indices container. One element here.
0171         std::vector<Acts::MultiTrajectoryTraits::IndexType> trackTips;
0172         trackTips.reserve(1);
0173         trackTips.emplace_back(fitOutput.lastMeasurementIndex);
0174         // The fitted parameters container. One element (at most) here.
0175         Trajectories::IndexedParameters indexedParams;
0176         if (fitOutput.fittedParameters) {
0177           const auto& params = fitOutput.fittedParameters.value();
0178           ACTS_VERBOSE("Fitted paramemeters for track " << itrack);
0179           ACTS_VERBOSE("  " << params.parameters().transpose());
0180           // Push the fitted parameters to the container
0181           indexedParams.emplace(fitOutput.lastMeasurementIndex, params);
0182         } else {
0183           ACTS_DEBUG("No fitted paramemeters for track " << itrack);
0184         }
0185         // store the result
0186         trajectories->emplace_back(std::move(fitOutput.fittedStates), std::move(trackTips),
0187                                   std::move(indexedParams));
0188       } else {
0189         ACTS_WARNING("Fit failed for track " << itrack << " with error" << result.error());
0190         // Fit failed. Add an empty result so the output container has
0191         // the same number of entries as the input.
0192         trajectories->push_back(Trajectories());
0193       }
0194     }
0195
0196     return StatusCode::SUCCESS;
0197   }
0198
0199   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0200   DECLARE_COMPONENT(TrackFittingAlgorithm)
0201
0202 } // namespace Jug::Reco