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