test/utils/random_scatterer.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 // Project include(s).
0012 #include "detray/definitions/algebra.hpp"
0013 #include "detray/definitions/detail/qualifiers.hpp"
0014 #include "detray/definitions/pdg_particle.hpp"
0015 #include "detray/definitions/track_parametrization.hpp"
0016 #include "detray/definitions/units.hpp"
0017 #include "detray/material/concepts.hpp"
0018 #include "detray/material/detail/material_accessor.hpp"
0019 #include "detray/material/interaction.hpp"
0020 #include "detray/propagator/actors/parameter_updater.hpp"
0021 #include "detray/propagator/base_actor.hpp"
0022 #include "detray/tracks/bound_track_parameters.hpp"
0023 #include "detray/utils/axis_rotation.hpp"
0024 #include "detray/utils/geometry_utils.hpp"
0025 #include "detray/utils/unit_vectors.hpp"
0026
0027 // Detray test include(s)
0028 #include "detray/test/utils/landau_distribution.hpp"
0029 #include "detray/test/utils/scattering_helper.hpp"
0030
0031 // System include(s).
0032 #include <random>
0033
0034 namespace detray::actor {
0035
0036 template <concepts::algebra algebra_t>
0037 struct random_scatterer : public base_actor {
0038   using scalar_type = dscalar<algebra_t>;
0039   using vector3_type = dvector3D<algebra_t>;
0040   using transform3_type = dtransform3D<algebra_t>;
0041   using interaction_type = interaction<scalar_type>;
0042
0043   struct state {
0044     std::random_device rd{};
0045     std::mt19937_64 generator{rd()};
0046
0047     /// most probable energy loss
0048     scalar_type e_loss_mpv = 0.f;
0049
0050     /// energy loss sigma
0051     scalar_type e_loss_sigma = 0.f;
0052
0053     /// projected scattering angle
0054     scalar_type projected_scattering_angle = 0.f;
0055
0056     // Simulation setup
0057     bool do_energy_loss = true;
0058     bool do_multiple_scattering = true;
0059
0060     /// Constructor with seed
0061     ///
0062     /// @param sd the seed number
0063     explicit state(const uint_fast64_t sd = 0u) { generator.seed(sd); }
0064
0065     void set_seed(const uint_fast64_t sd) { generator.seed(sd); }
0066   };
0067
0068   /// Material store visitor
0069   struct kernel {
0070     template <typename mat_group_t, typename index_t>
0071     DETRAY_HOST_DEVICE inline bool operator()(
0072         [[maybe_unused]] const mat_group_t& material_group,
0073         [[maybe_unused]] const index_t& mat_index,
0074         [[maybe_unused]] typename random_scatterer::state& s,
0075         [[maybe_unused]] const pdg_particle<scalar_type>& ptc,
0076         [[maybe_unused]] const bound_track_parameters<algebra_t>& bound_params,
0077         [[maybe_unused]] const scalar_type cos_inc_angle,
0078         [[maybe_unused]] const scalar_type approach) const {
0079       using material_t = typename mat_group_t::value_type;
0080
0081       if constexpr (concepts::surface_material<material_t>) {
0082         const scalar_type qop = bound_params.qop();
0083
0084         const auto mat = detail::material_accessor::get(
0085             material_group, mat_index, bound_params.bound_local());
0086
0087         // return early in case of zero thickness
0088         if (mat.thickness() <= std::numeric_limits<scalar_type>::epsilon()) {
0089           return false;
0090         }
0091
0092         const scalar_type path_segment{
0093             mat.path_segment(cos_inc_angle, approach)};
0094
0095         const detail::relativistic_quantities rq(ptc, qop);
0096
0097         // Energy Loss
0098         if (s.do_energy_loss) {
0099           s.e_loss_mpv = interaction_type().compute_energy_loss_landau(
0100               path_segment, mat.get_material(), ptc, rq);
0101
0102           s.e_loss_sigma = interaction_type().compute_energy_loss_landau_sigma(
0103               path_segment, mat.get_material(), ptc, rq);
0104         }
0105
0106         // Scattering angle
0107         if (s.do_multiple_scattering) {
0108           // @todo: use momentum before or after energy loss in
0109           // backward mode?
0110           s.projected_scattering_angle =
0111               interaction_type().compute_multiple_scattering_theta0(
0112                   mat.path_segment_in_X0(cos_inc_angle, approach), ptc, rq);
0113         }
0114
0115         return true;
0116       } else {
0117         // For non-pointwise material interactions, do nothing
0118         return false;
0119       }
0120     }
0121   };
0122
0123   template <typename propagator_state_t>
0124   DETRAY_HOST inline void operator()(
0125       state& simulator_state, propagator_state_t& prop_state,
0126       parameter_transporter_result<algebra_t>& res) const {
0127     // @Todo: Make context part of propagation state
0128     using detector_type = typename propagator_state_t::detector_type;
0129     using geo_context_type = typename detector_type::geometry_context;
0130
0131     auto& navigation = prop_state.navigation();
0132
0133     if (!navigation.encountered_sf_material()) {
0134       return;
0135     }
0136
0137     auto& stepping = prop_state.stepping();
0138     const auto& ptc = stepping.particle_hypothesis();
0139     auto& bound_params = res.destination_params();
0140     const auto sf = navigation.current_surface();
0141     const scalar_type cos_inc_angle{cos_angle(geo_context_type{}, sf,
0142                                               bound_params.dir(),
0143                                               bound_params.bound_local())};
0144
0145     const bool success = sf.template visit_material<kernel>(
0146         simulator_state, ptc, bound_params, cos_inc_angle,
0147         bound_params.bound_local()[0]);
0148
0149     if (success) {
0150       // Get the new momentum
0151       const auto new_mom = attenuate(
0152           simulator_state.e_loss_mpv, simulator_state.e_loss_sigma, ptc.mass(),
0153           bound_params.p(ptc.charge()), simulator_state.generator);
0154
0155       // Update Qop
0156       bound_params.set_qop(ptc.charge() / new_mom);
0157
0158       // Get the new direction from random scattering
0159       const auto new_dir = scatter(bound_params.dir(),
0160                                    simulator_state.projected_scattering_angle,
0161                                    simulator_state.generator);
0162
0163       // Update Phi and Theta
0164       bound_params.set_phi(vector::phi(new_dir));
0165       bound_params.set_theta(vector::theta(new_dir));
0166       // Signal parameter update
0167       res.status = actor::status::e_success;
0168
0169       // Flag renavigation of the current candidate
0170       prop_state.navigation().set_high_trust();
0171     }
0172   }
0173
0174   /// @brief Get the new momentum from the landau distribution
0175   template <typename generator_t>
0176   DETRAY_HOST inline scalar_type attenuate(const scalar_type mpv,
0177                                            const scalar_type sigma,
0178                                            const scalar_type m0,
0179                                            const scalar_type p0,
0180                                            generator_t& generator) const {
0181     // Get the random energy loss
0182     // @todo tune the scale parameters (e_loss_mpv and e_loss_sigma)
0183     const auto e_loss =
0184         landau_distribution<scalar_type>{}(generator, mpv, sigma);
0185
0186     // E = sqrt(m^2 + p^2)
0187     const auto energy = math::sqrt(m0 * m0 + p0 * p0);
0188     const auto new_energy = energy - e_loss;
0189
0190     auto p2 = new_energy * new_energy - m0 * m0;
0191
0192     // To avoid divergence
0193     if (p2 < 0) {
0194       p2 = 1.f * unit<scalar_type>::eV * unit<scalar_type>::eV;
0195     }
0196
0197     // p = sqrt(E^2 - m^2)
0198     return math::sqrt(p2);
0199   }
0200
0201   /// @brief Scatter the direction with projected scattering angle
0202   ///
0203   /// @param dir  input direction
0204   /// @param projected_scattering_angle  projected scattering angle
0205   /// @param generator random generator
0206   /// @returns the new direction from random scattering
0207   template <typename generator_t>
0208   DETRAY_HOST inline vector3_type scatter(
0209       const vector3_type& dir, const scalar_type projected_scattering_angle,
0210       generator_t& generator) const {
0211     const auto scattering_angle =
0212         constant<scalar_type>::sqrt2 * projected_scattering_angle;
0213
0214     return scattering_helper<algebra_t>()(dir, scattering_angle, generator);
0215   }
0216 };
0217
0218 }  // namespace detray::actor