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