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