algorithms/tracking/TrackPropagation.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022, 2023 Wenqing Fan, Barak Schmookler, Whitney Armstrong, Sylvester Joosten, Dmitry Romanov, Christopher Dilks, Wouter Deconinck
0003
0004 #include <Acts/Definitions/Algebra.hpp>
0005 #include <Acts/Definitions/Common.hpp>
0006 #include <Acts/Definitions/Direction.hpp>
0007 #include <Acts/Definitions/TrackParametrization.hpp>
0008 #include <Acts/Definitions/Units.hpp>
0009 #include <Acts/EventData/GenericBoundTrackParameters.hpp>
0010 #include <Acts/EventData/MultiTrajectoryHelpers.hpp>
0011 #include <Acts/EventData/VectorMultiTrajectory.hpp>
0012 #include <Acts/Geometry/GeometryIdentifier.hpp>
0013 #include <Acts/Geometry/TrackingGeometry.hpp>
0014 #include <Acts/MagneticField/MagneticFieldProvider.hpp>
0015 #include <Acts/Material/MaterialInteraction.hpp>
0016 #include <Acts/Propagator/ActorList.hpp>
0017 #include <Acts/Propagator/EigenStepper.hpp>
0018 #include <Acts/Propagator/MaterialInteractor.hpp>
0019 #include <Acts/Propagator/Navigator.hpp>
0020 #include <Acts/Propagator/Propagator.hpp>
0021 #include <Acts/Propagator/PropagatorResult.hpp>
0022 #include <Acts/Surfaces/CylinderBounds.hpp>
0023 #include <Acts/Surfaces/CylinderSurface.hpp>
0024 #include <Acts/Surfaces/DiscSurface.hpp>
0025 #include <Acts/Surfaces/RadialBounds.hpp>
0026 #include <Acts/Utilities/Logger.hpp>
0027 #include <DD4hep/Handle.h>
0028 #include <Evaluator/DD4hepUnits.h>
0029 #include <edm4eic/Cov2f.h>
0030 #include <edm4eic/Cov3f.h>
0031 #include <edm4hep/Vector3f.h>
0032 #include <edm4hep/utils/vector_utils.h>
0033 #include <Eigen/Core>
0034 #include <Eigen/Geometry>
0035 #include <algorithm>
0036 #include <cmath>
0037 #include <cstdint>
0038 #include <functional>
0039 #include <iterator>
0040 #include <map>
0041 #include <optional>
0042 #include <spdlog/common.h>
0043 #include <stdexcept>
0044 #include <string>
0045 #include <tuple>
0046 #include <typeinfo>
0047 #include <utility>
0048 #include <variant>
0049
0050 #include "algorithms/tracking/ActsGeometryProvider.h"
0051 #include "algorithms/tracking/TrackPropagation.h"
0052 #include "algorithms/tracking/TrackPropagationConfig.h"
0053 #include "extensions/spdlog/SpdlogToActs.h"
0054
0055 namespace eicrecon {
0056
0057 template <typename... L> struct multilambda : L... {
0058   using L::operator()...;
0059   constexpr multilambda(L... lambda) : L(std::move(lambda))... {}
0060 };
0061
0062 void TrackPropagation::init() {
0063   const auto* detector = m_detector;
0064
0065   std::map<uint32_t, std::size_t> system_id_layers;
0066
0067   multilambda _toDouble = {
0068       [](const std::string& v) { return dd4hep::_toDouble(v); },
0069       [](const double& v) { return v; },
0070   };
0071
0072   auto _toActsSurface =
0073       [&_toDouble, &detector, &system_id_layers](
0074           const std::variant<CylinderSurfaceConfig, DiscSurfaceConfig> surface_variant)
0075       -> std::shared_ptr<Acts::Surface> {
0076     if (std::holds_alternative<CylinderSurfaceConfig>(surface_variant)) {
0077       CylinderSurfaceConfig surface = std::get<CylinderSurfaceConfig>(surface_variant);
0078       const double rmin =
0079           std::visit(_toDouble, surface.rmin) / dd4hep::mm * Acts::UnitConstants::mm;
0080       const double zmin =
0081           std::visit(_toDouble, surface.zmin) / dd4hep::mm * Acts::UnitConstants::mm;
0082       const double zmax =
0083           std::visit(_toDouble, surface.zmax) / dd4hep::mm * Acts::UnitConstants::mm;
0084       const uint32_t system_id = detector->constant<uint32_t>(surface.id);
0085       auto bounds              = std::make_shared<Acts::CylinderBounds>(rmin, (zmax - zmin) / 2);
0086       auto t                   = Acts::Translation3(Acts::Vector3(0, 0, (zmax + zmin) / 2));
0087       auto tf                  = Acts::Transform3(t);
0088       auto acts_surface        = Acts::Surface::makeShared<Acts::CylinderSurface>(tf, bounds);
0089 #if Acts_VERSION_MAJOR >= 40
0090       acts_surface->assignGeometryId(
0091           Acts::GeometryIdentifier().withExtra(system_id).withLayer(++system_id_layers[system_id]));
0092 #else
0093       acts_surface->assignGeometryId(
0094           Acts::GeometryIdentifier().setExtra(system_id).setLayer(++system_id_layers[system_id]));
0095 #endif
0096       return acts_surface;
0097     }
0098     if (std::holds_alternative<DiscSurfaceConfig>(surface_variant)) {
0099       DiscSurfaceConfig surface = std::get<DiscSurfaceConfig>(surface_variant);
0100       const double zmin =
0101           std::visit(_toDouble, surface.zmin) / dd4hep::mm * Acts::UnitConstants::mm;
0102       const double rmin =
0103           std::visit(_toDouble, surface.rmin) / dd4hep::mm * Acts::UnitConstants::mm;
0104       const double rmax =
0105           std::visit(_toDouble, surface.rmax) / dd4hep::mm * Acts::UnitConstants::mm;
0106       const uint32_t system_id = detector->constant<uint32_t>(surface.id);
0107       auto bounds              = std::make_shared<Acts::RadialBounds>(rmin, rmax);
0108       auto t                   = Acts::Translation3(Acts::Vector3(0, 0, zmin));
0109       auto tf                  = Acts::Transform3(t);
0110       auto acts_surface        = Acts::Surface::makeShared<Acts::DiscSurface>(tf, bounds);
0111 #if Acts_VERSION_MAJOR >= 40
0112       acts_surface->assignGeometryId(
0113           Acts::GeometryIdentifier().withExtra(system_id).withLayer(++system_id_layers[system_id]));
0114 #else
0115       acts_surface->assignGeometryId(
0116           Acts::GeometryIdentifier().setExtra(system_id).setLayer(++system_id_layers[system_id]));
0117 #endif
0118       return acts_surface;
0119     }
0120     throw std::domain_error("Unknown surface type");
0121   };
0122   m_target_surfaces.resize(m_cfg.target_surfaces.size());
0123   std::ranges::transform(m_cfg.target_surfaces, m_target_surfaces.begin(), _toActsSurface);
0124   m_filter_surfaces.resize(m_cfg.filter_surfaces.size());
0125   std::ranges::transform(m_cfg.filter_surfaces, m_filter_surfaces.begin(), _toActsSurface);
0126
0127   trace("Initialized");
0128 }
0129
0130 void TrackPropagation::propagateToSurfaceList(const Input& input, const Output& output) const {
0131   const auto [tracks, track_states, tracks_acts] = input;
0132   auto [track_segments]                          = output;
0133
0134   // logging
0135   trace("Propagate tracks: --------------------");
0136   trace("number of tracks: {}", tracks->size());
0137
0138   // Construct ConstTrackContainer from underlying containers
0139   auto trackStateContainer = std::make_shared<Acts::ConstVectorMultiTrajectory>(*track_states);
0140   auto trackContainer      = std::make_shared<Acts::ConstVectorTrackContainer>(*tracks_acts);
0141   ActsExamples::ConstTrackContainer constTracks(trackContainer, trackStateContainer);
0142
0143   // loop over input tracks
0144   std::size_t i = 0;
0145   for (const auto& track : constTracks) {
0146
0147     // check if this track can be propagated to any filter surface
0148     bool track_reaches_filter_surface{false};
0149     for (const auto& filter_surface : m_filter_surfaces) {
0150       auto point = propagate(edm4eic::Track{}, track, constTracks, filter_surface);
0151       if (point) {
0152         track_reaches_filter_surface = true;
0153         break;
0154       }
0155     }
0156     if (!track_reaches_filter_surface) {
0157       ++i;
0158       continue;
0159     }
0160
0161     // start a mutable TrackSegment
0162     auto track_segment = track_segments->create();
0163
0164     // corresponding track
0165     if (tracks->size() == constTracks.size()) {
0166       trace("track segment connected to track {}", i);
0167       track_segment.setTrack((*tracks)[i]);
0168       ++i;
0169     }
0170
0171     // zero measurements of segment length
0172     decltype(edm4eic::TrackSegmentData::length) length            = 0;
0173     decltype(edm4eic::TrackSegmentData::lengthError) length_error = 0;
0174
0175     // loop over projection-target surfaces
0176     for (const auto& target_surface : m_target_surfaces) {
0177
0178       // project the track to this surface
0179       auto point = propagate(edm4eic::Track{}, track, constTracks, target_surface);
0180       if (!point) {
0181         trace("<> Failed to propagate track to this plane");
0182         continue;
0183       }
0184
0185       // logging
0186       trace("<> track: x=( {:>10.2f} {:>10.2f} {:>10.2f} )", point->position.x, point->position.y,
0187             point->position.z);
0188       trace("               p=( {:>10.2f} {:>10.2f} {:>10.2f} )", point->momentum.x,
0189             point->momentum.y, point->momentum.z);
0190
0191       // track point cut
0192       if (!m_cfg.track_point_cut(*point)) {
0193         trace("                 => REJECTED by trackPointCut");
0194         if (m_cfg.skip_track_on_track_point_cut_failure) {
0195           break;
0196         }
0197         continue;
0198       }
0199
0200       // update the `TrackSegment` length
0201       // FIXME: `length` and `length_error` are currently not used by any callers, and may not be correctly calculated here
0202       if (track_segment.points_size() > 0) {
0203         auto pos0 = point->position;
0204         auto pos1 = std::prev(track_segment.points_end())->position;
0205         auto dist = edm4hep::utils::magnitude(pos0 - pos1);
0206         length += dist;
0207         trace("               dist to previous point: {}", dist);
0208       }
0209
0210       // add the `TrackPoint` to the `TrackSegment`
0211       track_segment.addToPoints(*point);
0212
0213     } // end `targetSurfaces` loop
0214
0215     // set final length and length error
0216     track_segment.setLength(length);
0217     track_segment.setLengthError(length_error);
0218
0219   } // end loop over input tracks
0220 }
0221
0222 std::unique_ptr<edm4eic::TrackPoint>
0223 TrackPropagation::propagate(const edm4eic::Track& /* track */,
0224                             const ActsExamples::ConstTrackProxy& acts_track,
0225                             const ActsExamples::ConstTrackContainer& trackContainer,
0226                             const std::shared_ptr<const Acts::Surface>& targetSurf) const {
0227
0228   auto tipIndex = acts_track.tipIndex();
0229
0230   trace("  Propagating track with tip index {}", tipIndex);
0231
0232   // Collect the trajectory summary info
0233   auto trajState =
0234       Acts::MultiTrajectoryHelpers::trajectoryState(trackContainer.trackStateContainer(), tipIndex);
0235   int m_nMeasurements = trajState.nMeasurements;
0236   int m_nStates       = trajState.nStates;
0237
0238   trace("  Num measurement in trajectory: {}", m_nMeasurements);
0239   trace("  Num states in trajectory     : {}", m_nStates);
0240
0241   // Get track state at last measurement surface
0242   // For last measurement surface, filtered and smoothed results are equivalent
0243   auto trackState        = trackContainer.trackStateContainer().getTrackState(tipIndex);
0244   auto initSurface       = trackState.referenceSurface().getSharedPtr();
0245   const auto& initParams = trackState.filtered();
0246   const auto& initCov    = trackState.filteredCovariance();
0247
0248   Acts::BoundTrackParameters initBoundParams(initSurface, initParams, initCov,
0249                                              acts_track.particleHypothesis());
0250
0251   // Get pathlength of last track state with respect to perigee surface
0252   const auto initPathLength = trackState.pathLength();
0253
0254   trace("    TrackPropagation. Propagating to surface # {}", typeid(targetSurf->type()).name());
0255
0256   std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry   = m_geoSvc->trackingGeometry();
0257   std::shared_ptr<const Acts::MagneticFieldProvider> magneticField = m_geoSvc->getFieldProvider();
0258
0259   // Convert algorithm log level to Acts log level
0260   const auto spdlog_level = static_cast<spdlog::level::level_enum>(this->level());
0261   const auto acts_level   = eicrecon::SpdlogToActsLevel(spdlog_level);
0262   ACTS_LOCAL_LOGGER(Acts::getDefaultLogger("PROP", acts_level));
0263
0264   using Propagator        = Acts::Propagator<Acts::EigenStepper<>, Acts::Navigator>;
0265   using PropagatorOptions = Propagator::template Options<Acts::ActorList<Acts::MaterialInteractor>>;
0266   Propagator propagator(Acts::EigenStepper<>(magneticField),
0267                         Acts::Navigator({.trackingGeometry = m_geoSvc->trackingGeometry()},
0268                                         logger().cloneWithSuffix("Navigator")),
0269                         logger().cloneWithSuffix("Propagator"));
0270
0271   // Get run-scoped contexts from service
0272   const auto& gctx = m_geoSvc->getActsGeometryContext();
0273   const auto& mctx = m_geoSvc->getActsMagneticFieldContext();
0274
0275   PropagatorOptions propagationOptions(gctx, mctx);
0276
0277   auto result = propagator.propagate(initBoundParams, *targetSurf, propagationOptions);
0278
0279   // check propagation result
0280   if (!result.ok()) {
0281     trace("    propagation failed (!result.ok())");
0282     return nullptr;
0283   }
0284   trace("    propagation result is OK");
0285
0286   // Pulling results to convenient variables
0287   auto trackStateParams  = *((*result).endParameters);
0288   const auto& parameter  = trackStateParams.parameters();
0289   const auto& covariance = *trackStateParams.covariance();
0290
0291   // Path length
0292   const float pathLength      = initPathLength + (*result).pathLength;
0293   const float pathLengthError = 0;
0294   trace("    path len = {}", pathLength);
0295
0296   // Position:
0297   auto projectionPos = trackStateParams.position(gctx);
0298   const decltype(edm4eic::TrackPoint::position) position{static_cast<float>(projectionPos(0)),
0299                                                          static_cast<float>(projectionPos(1)),
0300                                                          static_cast<float>(projectionPos(2))};
0301   const decltype(edm4eic::TrackPoint::positionError) positionError{0, 0, 0};
0302   trace("    pos x = {}", position.x);
0303   trace("    pos y = {}", position.y);
0304   trace("    pos z = {}", position.z);
0305
0306   // Momentum
0307   const decltype(edm4eic::TrackPoint::momentum) momentum = edm4hep::utils::sphericalToVector(
0308       static_cast<float>(1.0 / std::abs(parameter[Acts::eBoundQOverP])),
0309       static_cast<float>(parameter[Acts::eBoundTheta]),
0310       static_cast<float>(parameter[Acts::eBoundPhi]));
0311   const decltype(edm4eic::TrackPoint::momentumError) momentumError{
0312       static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundTheta)),
0313       static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundPhi)),
0314       static_cast<float>(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)),
0315       static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundPhi)),
0316       static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundQOverP)),
0317       static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundQOverP))};
0318
0319   // time
0320   const float time{static_cast<float>(parameter(Acts::eBoundTime))};
0321   const float timeError{static_cast<float>(sqrt(covariance(Acts::eBoundTime, Acts::eBoundTime)))};
0322
0323   // Direction
0324   const float theta(parameter[Acts::eBoundTheta]);
0325   const float phi(parameter[Acts::eBoundPhi]);
0326   const decltype(edm4eic::TrackPoint::directionError) directionError{
0327       static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundTheta)),
0328       static_cast<float>(covariance(Acts::eBoundPhi, Acts::eBoundPhi)),
0329       static_cast<float>(covariance(Acts::eBoundTheta, Acts::eBoundPhi))};
0330
0331   // >oO debug print
0332   trace("    loc 0   = {:.4f}", parameter[Acts::eBoundLoc0]);
0333   trace("    loc 1   = {:.4f}", parameter[Acts::eBoundLoc1]);
0334   trace("    phi     = {:.4f}", parameter[Acts::eBoundPhi]);
0335   trace("    theta   = {:.4f}", parameter[Acts::eBoundTheta]);
0336   trace("    q/p     = {:.4f}", parameter[Acts::eBoundQOverP]);
0337   trace("    p       = {:.4f}", 1.0 / parameter[Acts::eBoundQOverP]);
0338   trace("    err phi = {:.4f}", sqrt(covariance(Acts::eBoundPhi, Acts::eBoundPhi)));
0339   trace("    err th  = {:.4f}", sqrt(covariance(Acts::eBoundTheta, Acts::eBoundTheta)));
0340   trace("    err q/p = {:.4f}", sqrt(covariance(Acts::eBoundQOverP, Acts::eBoundQOverP)));
0341   trace("    chi2    = {:.4f}", trajState.chi2Sum);
0342   trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc0)));
0343   trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc1, Acts::eBoundLoc1)));
0344   trace("    loc err = {:.4f}", static_cast<float>(covariance(Acts::eBoundLoc0, Acts::eBoundLoc1)));
0345
0346   uint64_t surface = targetSurf->geometryId().value();
0347   uint32_t system  = 0; // default value...will be set in TrackPropagation factory
0348
0349   return std::make_unique<edm4eic::TrackPoint>(
0350       edm4eic::TrackPoint{.surface         = surface,
0351                           .system          = system,
0352                           .position        = position,
0353                           .positionError   = positionError,
0354                           .momentum        = momentum,
0355                           .momentumError   = momentumError,
0356                           .time            = time,
0357                           .timeError       = timeError,
0358                           .theta           = theta,
0359                           .phi             = phi,
0360                           .directionError  = directionError,
0361                           .pathlength      = pathLength,
0362                           .pathlengthError = pathLengthError});
0363 }
0364
0365 } // namespace eicrecon