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