traccc/utils/seed_generator.hpp

0001 /** TRACCC library, part of the ACTS project (R&D line)
0002  *
0003  * (c) 2022-2024 CERN for the benefit of the ACTS project
0004  *
0005  * Mozilla Public License Version 2.0
0006  */
0007
0008 #pragma once
0009
0010 // Library include(s).
0011 #include "traccc/edm/track_parameters.hpp"
0012 #include "traccc/utils/logging.hpp"
0013 #include "traccc/utils/particle.hpp"
0014 #include "traccc/utils/trigonometric_helpers.hpp"
0015
0016 // detray include(s).
0017 #include <detray/definitions/navigation.hpp>  // < navigation::direction
0018 #include <detray/geometry/identifier.hpp>
0019 #include <detray/geometry/tracking_surface.hpp>
0020 #include <detray/propagator/actors/pointwise_material_interactor.hpp>
0021
0022 // System include(s).
0023 #include <random>
0024
0025 namespace traccc {
0026
0027 /// Seed track parameter generator
0028 template <typename detector_t>
0029 struct seed_generator {
0030   using algebra_type = typename detector_t::algebra_type;
0031   using ctx_t = typename detector_t::geometry_context;
0032
0033   /// Configure the seed generator
0034   struct config {
0035     // Smearing parameters
0036     /// Constant term of the loc0 resolution.
0037     double sigma_loc0 = 20 * traccc::unit<scalar>::um;
0038     /// Pt-dependent loc0 resolution of the form sigma_loc0 =
0039     /// A*exp(-1.*abs(B)*pt).
0040     double sigma_loc0_pT_a = 30 * traccc::unit<scalar>::um;
0041     double sigma_loc0_pT_b = 0.3 / traccc::unit<scalar>::GeV;
0042     /// Constant term of the loc1 resolution.
0043     double sigma_loc1 = 20 * traccc::unit<scalar>::um;
0044     /// Pt-dependent loc1 resolution of the form sigma_loc1 =
0045     /// A*exp(-1.*abs(B)*pt).
0046     double sigma_loc1_pT_a = 30 * traccc::unit<scalar>::um;
0047     double sigma_loc1_pT_b = 0.3 / traccc::unit<scalar>::GeV;
0048     /// Time resolution.
0049     double sigma_time = 1 * traccc::unit<scalar>::ns;
0050     /// Phi angular resolution.
0051     double sigma_phi = 1 * traccc::unit<scalar>::degree;
0052     /// Theta angular resolution.
0053     double sigma_theta = 1 * traccc::unit<scalar>::degree;
0054     /// Relative transverse momentum resolution.
0055     double sigma_pT_rel = 0.05;
0056
0057     /// Optional. Initial sigmas for the track parameters which overwrites
0058     /// the smearing params if set.
0059     std::optional<std::array<double, e_bound_size>> initial_sigmas;
0060     /// Initial sigma(q/pt) for the track parameters.
0061     /// @note The resulting q/p sigma is added to the one in `initialSigmas`
0062     double initialsigma_qopt =
0063         0.1 * traccc::unit<scalar>::e / traccc::unit<scalar>::GeV;
0064     /// Initial sigma(pt)/pt for the track parameters.
0065     /// @note The resulting q/p sigma is added to the one in `initialSigmas`
0066     double initialsigma_pT_rel = 0.1;
0067
0068     /// Inflation factors for the variances
0069     std::array<double, e_bound_size> cov_inflation{1., 1., 1., 1., 1., 1.};
0070   };
0071
0072   /// Constructor with detector
0073   ///
0074   /// @param det input detector
0075   /// @param stddevs standard deviations for parameter smearing
0076   explicit seed_generator(const detector_t& det, const config& cfg = {},
0077                           const std::size_t sd = std::mt19937::default_seed,
0078                           ctx_t ctx = {})
0079       : m_detector(det), m_ctx(ctx), m_cfg(cfg) {
0080     m_generator.seed(static_cast<std::mt19937::result_type>(sd));
0081   }
0082
0083   /// Seed generator operation
0084   ///
0085   /// @param vertex vertex of particle
0086   /// @param stddevs standard deviations for track parameter smearing
0087   bound_track_parameters<algebra_type> operator()(
0088       const detray::geometry::identifier surface_link,
0089       const free_track_parameters<algebra_type>& free_param,
0090       const traccc::pdg_particle<scalar>& ptc_type) {
0091     // Get bound parameter
0092     const detray::tracking_surface sf{m_detector, surface_link};
0093
0094     TRACCC_DEBUG_HOST("Surface:\n" << sf);
0095     TRACCC_DEBUG_HOST("Input free param. on surface:\n" << free_param);
0096
0097     auto bound_vec = sf.free_to_bound_vector(m_ctx, free_param);
0098     auto bound_cov = matrix::identity<traccc::bound_matrix<algebra_type>>();
0099
0100     bound_track_parameters<algebra_type> bound_param{surface_link, bound_vec,
0101                                                      bound_cov};
0102
0103     TRACCC_DEBUG_HOST("-> Bound param.:\n" << bound_param);
0104
0105     // Type definitions
0106     using interactor_type =
0107         detray::actor::pointwise_material_interactor<algebra_type>;
0108
0109     assert(ptc_type.charge() * bound_param.qop() > 0.f);
0110
0111     // Apply interactor
0112     TRACCC_DEBUG_HOST("Update material:");
0113     if (sf.has_material()) {
0114       typename interactor_type::state interactor_state;
0115       interactor_state.do_multiple_scattering = false;
0116       interactor_type{}.update(
0117           m_ctx, ptc_type, bound_param, interactor_state,
0118           static_cast<int>(detray::navigation::direction::e_backward), sf);
0119     }
0120
0121     // Call the smearing operator
0122     (*this)(bound_param, ptc_type);
0123
0124     bound_param.set_surface_link(surface_link);
0125
0126     return bound_param;
0127   }
0128
0129   /// Seed generator operation
0130   ///
0131   /// @param vertex vertex of particle
0132   /// @param stddevs standard deviations for track parameter smearing
0133   void operator()(bound_track_parameters<algebra_type>& bound_param,
0134                   const traccc::pdg_particle<scalar>& ptc_type) {
0135     const scalar q{ptc_type.charge()};
0136     const auto pos = bound_param.bound_local();
0137     const auto time = bound_param.time();
0138     const auto phi = bound_param.phi();
0139     const auto theta = bound_param.theta();
0140     const auto pt = bound_param.pT(q);
0141     const auto qop = bound_param.qop();
0142
0143     // Build the track covariance matrix using the smearing sigmas
0144     const auto sigmas = generate_smearing_sigmas(pt, qop, theta);
0145
0146     // Smear the position/time
0147     // Note that we smear d0 and z0 in the perigee frame
0148     const double smeared_loc0 = std::normal_distribution<double>(
0149         pos[0], sigmas[e_bound_loc0])(m_generator);
0150     const double smeared_loc1 = std::normal_distribution<double>(
0151         pos[1], sigmas[e_bound_loc1])(m_generator);
0152     bound_param.set_bound_local(
0153         {static_cast<scalar>(smeared_loc0), static_cast<scalar>(smeared_loc1)});
0154
0155     // Time
0156     bound_param.set_time(static_cast<scalar>(std::normal_distribution<double>(
0157         time, sigmas[e_bound_time])(m_generator)));
0158
0159     // Smear direction angles phi,theta ensuring correct bounds
0160     const double smeared_phi =
0161         std::normal_distribution<double>(phi, sigmas[e_bound_phi])(m_generator);
0162     const double smeared_theta = std::normal_distribution<double>(
0163         theta, sigmas[e_bound_theta])(m_generator);
0164     const auto [new_phi, new_theta] = detail::wrap_phi_theta(
0165         static_cast<scalar>(smeared_phi), static_cast<scalar>(smeared_theta));
0166     bound_param.set_phi(static_cast<scalar>(new_phi));
0167     bound_param.set_theta(static_cast<scalar>(new_theta));
0168
0169     // Compute smeared q/p
0170     bound_param.set_qop(static_cast<scalar>(std::normal_distribution<double>(
0171         qop, sigmas[e_bound_qoverp])(m_generator)));
0172
0173     // Generate the covariances
0174     generate_initial_covariance(bound_param, ptc_type);
0175
0176     // We are not on a surface, but in the curvilinear frame
0177     bound_param.set_surface_link(detray::geometry::identifier{});
0178   }
0179
0180   /// Generate an array of standard deviations
0181   ///
0182   /// @param pt transverse momentum of the particle
0183   /// @param qop q/p of the particle
0184   /// @param theta global theta angle of the particle direction
0185   ///
0186   /// @returns an array of stddeviations corresponding to each track parameter
0187   std::array<double, e_bound_size> generate_smearing_sigmas(
0188       const double pt, const double qop, const double theta) {
0189     if (m_cfg.initial_sigmas.has_value()) {
0190       return *m_cfg.initial_sigmas;
0191     } else {
0192       // Momentum dependent standard deviations
0193       const double sigma_loc_0 =
0194           m_cfg.sigma_loc0 +
0195           m_cfg.sigma_loc0_pT_a *
0196               math::exp(-1.0 * math::fabs(m_cfg.sigma_loc0_pT_b) * pt);
0197
0198       const double sigma_loc_1 =
0199           m_cfg.sigma_loc1 +
0200           m_cfg.sigma_loc1_pT_a *
0201               math::exp(-1.0 * math::fabs(m_cfg.sigma_loc1_pT_b) * pt);
0202
0203       const double sigma_qop = math::sqrt(
0204           math::pow(m_cfg.sigma_pT_rel * qop, 2) +
0205           math::pow(m_cfg.sigma_theta * (qop * math::tan(theta)), 2));
0206
0207       return {sigma_loc_0,       sigma_loc_1, m_cfg.sigma_phi,
0208               m_cfg.sigma_theta, sigma_qop,   m_cfg.sigma_time};
0209     }
0210   }
0211
0212   /// Generate the covariance matrix for the track
0213   ///
0214   /// @param[in, out] bound_param the bound track parameters
0215   /// @param ptc_type the particle hypothesis
0216   void generate_initial_covariance(
0217       bound_track_parameters<algebra_type>& bound_param,
0218       const traccc::pdg_particle<scalar>& ptc_type) {
0219     const scalar q{ptc_type.charge()};
0220
0221     // Build the track covariance matrix using the smearing sigmas
0222     const auto sigmas = generate_smearing_sigmas(
0223         bound_param.pT(q), bound_param.qop(), bound_param.theta());
0224
0225     for (std::size_t i = e_bound_loc0; i < e_bound_size; ++i) {
0226       double sigma = sigmas[i];
0227       double variance = sigma * sigma;
0228
0229       // Inflate the initial covariance
0230       variance *= m_cfg.cov_inflation[i];
0231
0232       getter::element(bound_param.covariance(), i, i) =
0233           static_cast<scalar>(variance);
0234     }
0235
0236     assert(!bound_param.is_invalid());
0237   }
0238
0239  private:
0240   /// Random generator
0241   std::random_device m_rd{};
0242   std::mt19937 m_generator{m_rd()};
0243
0244   /// Detector object
0245   const detector_t& m_detector;
0246   /// Geometry context
0247   ctx_t m_ctx;
0248
0249   /// Seed generator configuration
0250   config m_cfg{};
0251 };
0252
0253 }  // namespace traccc