detray/propagator/rk_stepper.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 "./codegen/update_rk_transport_jacobian.hpp"
0013 #include "detray/definitions/algebra.hpp"
0014 #include "detray/definitions/detail/qualifiers.hpp"
0015 #include "detray/definitions/units.hpp"
0016 #include "detray/material/interaction.hpp"
0017 #include "detray/material/material.hpp"
0018 #include "detray/navigation/policies.hpp"
0019 #include "detray/propagator/base_stepper.hpp"
0020 #include "detray/propagator/transport_jacobian.hpp"
0021 #include "detray/tracks/tracks.hpp"
0022
0023 namespace detray {
0024 enum class rk_stepper_flags : std::uint32_t {
0025   e_allow_covariance_transport = 1,
0026   e_allow_field_gradient = 2,
0027 };
0028
0029 /// Runge-Kutta-Nystrom 4th order stepper implementation
0030 ///
0031 /// @tparam magnetic_field_t the type of magnetic field
0032 /// @tparam track_t the type of track that is being advanced by the stepper
0033 /// @tparam constraint_ the type of constraints on the stepper
0034 template <typename magnetic_field_t, concepts::algebra algebra_t,
0035           typename constraint_t = unconstrained_step<dscalar<algebra_t>>,
0036           typename policy_t = stepper_rk_policy<dscalar<algebra_t>>,
0037           typename inspector_t = stepping::void_inspector,
0038           std::uint32_t flags_v =
0039               (static_cast<std::uint32_t>(
0040                    rk_stepper_flags::e_allow_covariance_transport) |
0041                static_cast<std::uint32_t>(
0042                    rk_stepper_flags::e_allow_field_gradient))>
0043 class rk_stepper final
0044     : public base_stepper<algebra_t, constraint_t, policy_t, inspector_t> {
0045   using base_type =
0046       base_stepper<algebra_t, constraint_t, policy_t, inspector_t>;
0047
0048  public:
0049   static constexpr bool uses_gradient = static_cast<bool>(
0050       flags_v &
0051       static_cast<std::uint32_t>(rk_stepper_flags::e_allow_field_gradient));
0052
0053   using algebra_type = algebra_t;
0054   using scalar_type = dscalar<algebra_t>;
0055   using point3_type = dpoint3D<algebra_t>;
0056   using vector3_type = dvector3D<algebra_t>;
0057   using transform3_type = dtransform3D<algebra_t>;
0058   using free_track_parameters_type =
0059       typename base_type::free_track_parameters_type;
0060   using bound_track_parameters_type =
0061       typename base_type::bound_track_parameters_type;
0062   using magnetic_field_type = magnetic_field_t;
0063   template <std::size_t ROWS, std::size_t COLS>
0064   using matrix_type = dmatrix<algebra_t, ROWS, COLS>;
0065   using transport_jacobian_type = std::conditional_t<
0066       uses_gradient,
0067       detail::transport_jacobian_matrix_with_gradient<algebra_type>,
0068       detail::transport_jacobian_matrix_without_gradient<algebra_type>>;
0069
0070   rk_stepper() = default;
0071
0072   struct intermediate_state {
0073     vector3_type b_first{0.f, 0.f, 0.f};
0074     vector3_type b_middle{0.f, 0.f, 0.f};
0075     vector3_type b_last{0.f, 0.f, 0.f};
0076     // t = tangential direction = dr/ds
0077     darray<vector3_type, 4u> t;
0078     // q/p
0079     darray<scalar_type, 4u> qop;
0080     // dt/ds = d^2r/ds^2 = q/p ( t X B )
0081     darray<vector3_type, 4u> dtds;
0082     // d(q/p)/ds
0083     darray<scalar_type, 4u> dqopds;
0084   };
0085
0086   struct state : public base_type::template state<transport_jacobian_type> {
0087     static constexpr const stepping::id id = stepping::id::e_rk;
0088
0089     using base_state =
0090         typename base_type::template state<transport_jacobian_type>;
0091
0092     DETRAY_HOST_DEVICE
0093     state(const free_track_parameters_type& t,
0094           const magnetic_field_t& mag_field)
0095         : base_state(t), m_magnetic_field(mag_field) {}
0096
0097     template <typename detector_t>
0098     DETRAY_HOST_DEVICE state(const bound_track_parameters_type& bound_params,
0099                              const magnetic_field_t& mag_field,
0100                              const detector_t& det,
0101                              const typename detector_t::geometry_context& ctx)
0102         : base_state(bound_params, det, ctx), m_magnetic_field(mag_field) {}
0103
0104     /// @returns the B-field view
0105     DETRAY_HOST_DEVICE
0106     const magnetic_field_type& field() const { return m_magnetic_field; }
0107
0108     /// Set the next step size
0109     DETRAY_HOST_DEVICE
0110     inline void set_next_step_size(const scalar_type step) {
0111       assert(math::fabs(step) >= 1e-4f * unit<scalar_type>::mm);
0112       m_next_step_size = step;
0113     }
0114
0115     /// @returns the next step size to be taken on the following step.
0116     DETRAY_HOST_DEVICE
0117     inline scalar_type next_step_size() const { return m_next_step_size; }
0118
0119     /// Evaluate dtds, where t is the unit tangential direction
0120     DETRAY_HOST_DEVICE
0121     vector3_type dtds() const {
0122       // In case there was no step before
0123       if (this->path_length() == 0.f) [[unlikely]] {
0124         const point3_type pos = (*this)().pos();
0125
0126         const auto bvec_tmp = this->field().at(pos[0], pos[1], pos[2]);
0127         vector3_type bvec;
0128         bvec[0u] = bvec_tmp[0u];
0129         bvec[1u] = bvec_tmp[1u];
0130         bvec[2u] = bvec_tmp[2u];
0131
0132         DETRAY_DEBUG_HOST(
0133             "--> dtds: " << (*this)().qop() *
0134                                 vector::cross((*this)().dir(), bvec));
0135
0136         return (*this)().qop() * vector::cross((*this)().dir(), bvec);
0137       }
0138
0139       return m_dtds_3;
0140     }
0141
0142     /// Set dtds, where t is the unit tangential direction
0143     DETRAY_HOST_DEVICE
0144     void dtds(const vector3_type& v) { m_dtds_3 = v; }
0145
0146     /// Evaluate d(qop)/ds
0147     DETRAY_HOST_DEVICE
0148     scalar_type dqopds(const material<scalar_type>* vol_mat_ptr) const {
0149       // In case there was no step before
0150       if (this->path_length() == 0.f) [[unlikely]] {
0151         return this->dqopds((*this)().qop(), vol_mat_ptr);
0152       }
0153
0154       return m_dqopds_3;
0155     }
0156
0157     DETRAY_HOST_DEVICE
0158     scalar_type dqopds(const scalar_type qop,
0159                        const material<scalar_type>* vol_mat_ptr) const {
0160       // d(qop)ds is zero for empty space
0161       if (vol_mat_ptr == nullptr) {
0162         return 0.f;
0163       }
0164
0165       assert(qop != 0.f);
0166       assert(std::isfinite(qop));
0167       const scalar_type q = this->particle_hypothesis().charge();
0168       const scalar_type p = q / qop;
0169       const scalar_type mass = this->particle_hypothesis().mass();
0170       const scalar_type E = math::sqrt(p * p + mass * mass);
0171
0172       // Compute stopping power
0173       const scalar_type stopping_power =
0174           interaction<scalar_type>().compute_stopping_power(
0175               *vol_mat_ptr, this->particle_hypothesis(), {mass, qop, q});
0176
0177       // Assert that a momentum is a positive value
0178       assert(p >= 0.f);
0179       assert(q != 0.f);
0180
0181       DETRAY_DEBUG_HOST_DEVICE("--> dqopds: %f",
0182                                qop * qop * qop * E * stopping_power / (q * q));
0183
0184       // d(qop)ds, which is equal to (qop) * E * (-dE/ds) / p^2
0185       // or equal to (qop)^3 * E * (-dE/ds) / q^2
0186       return qop * qop * qop * E * stopping_power / (q * q);
0187     }
0188
0189     /// Set d(qop)/ds
0190     DETRAY_HOST_DEVICE void dqopds(const scalar_type& s) { m_dqopds_3 = s; }
0191
0192     /// Evaluate d(d(qop)/ds)dqop
0193     DETRAY_HOST_DEVICE
0194     scalar_type d2qopdsdqop(const scalar_type qop,
0195                             const material<scalar_type>* vol_mat_ptr) const {
0196       if (vol_mat_ptr == nullptr) {
0197         return 0.f;
0198       }
0199
0200       const scalar_type q = this->particle_hypothesis().charge();
0201       const scalar_type p = q / qop;
0202       const scalar_type p2 = p * p;
0203
0204       const auto& mass = this->particle_hypothesis().mass();
0205       const scalar_type E2 = p2 + mass * mass;
0206
0207       // Interaction object
0208       interaction<scalar_type> I;
0209
0210       // g = dE/ds = -1 * (-dE/ds) = -1 * stopping power
0211       const detail::relativistic_quantities<scalar_type> rq(mass, qop, q);
0212       const scalar_type g =
0213           -1.f * I.compute_stopping_power(*vol_mat_ptr,
0214                                           this->particle_hypothesis(), rq);
0215
0216       // dg/d(qop) = -1 * derivation of stopping power
0217       const scalar_type dgdqop =
0218           -1.f * I.derive_stopping_power(*vol_mat_ptr,
0219                                          this->particle_hypothesis(), rq);
0220
0221       // d(qop)/ds = - qop^3 * E * g / q^2
0222       const scalar_type dqopds = this->dqopds(qop, vol_mat_ptr);
0223
0224       // Check Eq 3.12 of
0225       // (https://iopscience.iop.org/article/10.1088/1748-0221/4/04/P04016/meta)
0226       assert(E2 != 0.f);
0227       assert(g != 0.f);
0228       return dqopds * (1.f / qop * (3.f - p2 / E2) + 1.f / g * dgdqop);
0229     }
0230
0231     /// Call the stepping inspector
0232     template <typename... Args>
0233     DETRAY_HOST_DEVICE void run_inspector(
0234         [[maybe_unused]] const stepping::config& cfg,
0235         [[maybe_unused]] const char* message,
0236         [[maybe_unused]] const scalar_type dist,
0237         [[maybe_unused]] Args&&... args) {
0238       if constexpr (!std::is_same_v<inspector_t, stepping::void_inspector>) {
0239         this->inspector()(*this, cfg, message, dist,
0240                           std::forward<Args>(args)...);
0241       }
0242
0243       if constexpr (sizeof...(Args) > 0u) {
0244         DETRAY_DEBUG_HOST("" << message << "\n"
0245                              << detray::stepping::print_state(
0246                                     *this, std::forward<Args>(args)...));
0247       }
0248     }
0249
0250    private:
0251     vector3_type m_dtds_3;
0252     scalar_type m_dqopds_3;
0253
0254     /// Next step size after adaptive step size scaling
0255     scalar_type m_next_step_size{0.f};
0256
0257     /// Magnetic field view
0258     const magnetic_field_type& m_magnetic_field;
0259   };
0260
0261   /// Take a step, using an adaptive Runge-Kutta algorithm.
0262   ///
0263   /// @param dist_to_next The straight line distance to the next surface
0264   /// @param stepping The state object of a stepper
0265   /// @param cfg The stepping configuration
0266   /// @param do_reset whether to reset the RKN step size to "dist to next"
0267   ///
0268   /// @return returning the heartbeat, indicating if the stepping is alive
0269   DETRAY_HOST_DEVICE bool step(
0270       const scalar_type dist_to_next, state& stepping,
0271       const stepping::config& cfg, bool do_reset,
0272       const material<scalar_type>* vol_mat_ptr = nullptr) const {
0273     DETRAY_DEBUG_HOST("Before: " << stepping());
0274
0275     // In case of an overlap do nothing
0276     if (math::fabs(dist_to_next) < 1.f * unit<float>::um) [[unlikely]] {
0277       DETRAY_VERBOSE_HOST_DEVICE("Zero stepsize...");
0278
0279       // Don't allow a too small step size on next step
0280       if (math::fabs(stepping.next_step_size()) < cfg.min_stepsize) {
0281         stepping.set_next_step_size(math::copysign(
0282             static_cast<scalar_type>(cfg.min_stepsize), dist_to_next));
0283       }
0284       DETRAY_DEBUG_HOST_DEVICE("Setting next stepsize: %f mm",
0285                                stepping.next_step_size());
0286
0287       stepping.run_inspector(cfg, "Step skipped (Overlap): ", dist_to_next);
0288
0289       DETRAY_DEBUG_HOST("After: " << stepping());
0290       return true;
0291     }
0292
0293     if (!(do_reset ||
0294           math::fabs(stepping.next_step_size()) >= cfg.min_stepsize)) {
0295       DETRAY_INFO_HOST("Next stepsize: " << stepping.next_step_size() << ", "
0296                                          << cfg.min_stepsize);
0297     }
0298
0299     // Check navigator and actor results
0300     assert(math::fabs(dist_to_next) != 0.f);
0301     assert(do_reset ||
0302            math::fabs(stepping.next_step_size()) >= cfg.min_stepsize);
0303     assert(!stepping().is_invalid());
0304
0305     // Get stepper and navigator states
0306     auto& magnetic_field = stepping.field();
0307
0308     if (do_reset) {
0309       stepping.set_step_size(dist_to_next);
0310     } else if (stepping.next_step_size() > 0) {
0311       assert(math::fabs(stepping.next_step_size()) >= cfg.min_stepsize);
0312       stepping.set_step_size(
0313           math::min(stepping.next_step_size(), dist_to_next));
0314     } else {
0315       assert(math::fabs(stepping.next_step_size()) >= cfg.min_stepsize);
0316       stepping.set_step_size(
0317           math::max(stepping.next_step_size(), dist_to_next));
0318     }
0319
0320     DETRAY_VERBOSE_HOST_DEVICE("Distance to nex: %f mm", dist_to_next);
0321     DETRAY_VERBOSE_HOST_DEVICE("Initial stepsize: %f mm", stepping.step_size());
0322
0323     // Don't allow too small stepsizes, unless the navigation needs it
0324     const scalar_type min_stepsize{
0325         math::min(math::fabs(stepping.step_size()),
0326                   static_cast<scalar_type>(cfg.min_stepsize))};
0327     const point3_type pos = stepping().pos();
0328
0329     intermediate_state sd{};
0330
0331     // First Runge-Kutta point
0332     auto bvec = magnetic_field.at(pos[0], pos[1], pos[2]);
0333     assert(math::isfinite(bvec[0]));
0334     assert(math::isfinite(bvec[1]));
0335     assert(math::isfinite(bvec[2]));
0336     sd.b_first[0] = bvec[0];
0337     sd.b_first[1] = bvec[1];
0338     sd.b_first[2] = bvec[2];
0339
0340     DETRAY_DEBUG_HOST_DEVICE("First stage:");
0341     DETRAY_DEBUG_HOST_DEVICE("-> B-field: [%f, %f, %f]", bvec[0], bvec[1],
0342                              bvec[2]);
0343
0344     // qop should be recalculated at every point
0345     // Reference: Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0346     detray::tie(sd.dqopds[0u], sd.qop[0u]) =
0347         evaluate_dqopds(stepping, 0u, 0.f, 0.f, vol_mat_ptr, cfg);
0348     detray::tie(sd.dtds[0u], sd.t[0u]) = evaluate_dtds(
0349         stepping, sd.b_first, 0u, 0.f, vector3_type{0.f, 0.f, 0.f}, sd.qop[0u]);
0350
0351     /// RKN step trial and error estimation
0352     const auto estimate_error = [&](const scalar_type& h) {
0353       assert(h != 0);
0354       // State the square and half of the step size
0355       const scalar_type h2{h * h};
0356       const scalar_type half_h{h * 0.5f};
0357
0358       // Second Runge-Kutta point
0359       // qop should be recalculated at every point
0360       // Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0361       const point3_type pos1 =
0362           pos + half_h * sd.t[0u] + h2 * 0.125f * sd.dtds[0u];
0363       bvec = magnetic_field.at(pos1[0], pos1[1], pos1[2]);
0364       assert(math::isfinite(bvec[0]));
0365       assert(math::isfinite(bvec[1]));
0366       assert(math::isfinite(bvec[2]));
0367       sd.b_middle[0] = bvec[0];
0368       sd.b_middle[1] = bvec[1];
0369       sd.b_middle[2] = bvec[2];
0370       DETRAY_DEBUG_HOST_DEVICE("Second stage:");
0371       DETRAY_DEBUG_HOST_DEVICE("-> B-field: [%f, %f, %f]", bvec[0], bvec[1],
0372                                bvec[2]);
0373
0374       detray::tie(sd.dqopds[1u], sd.qop[1u]) = evaluate_dqopds(
0375           stepping, 1u, half_h, sd.dqopds[0u], vol_mat_ptr, cfg);
0376       detray::tie(sd.dtds[1u], sd.t[1u]) = evaluate_dtds(
0377           stepping, sd.b_middle, 1u, half_h, sd.dtds[0u], sd.qop[1u]);
0378
0379       // Third Runge-Kutta point
0380       // qop should be recalculated at every point
0381       // Reference: Eq (84) of
0382       // https://doi.org/10.1016/0029-554X(81)90063-X
0383       detray::tie(sd.dqopds[2u], sd.qop[2u]) = evaluate_dqopds(
0384           stepping, 2u, half_h, sd.dqopds[1u], vol_mat_ptr, cfg);
0385       detray::tie(sd.dtds[2u], sd.t[2u]) = evaluate_dtds(
0386           stepping, sd.b_middle, 2u, half_h, sd.dtds[1u], sd.qop[2u]);
0387
0388       // Last Runge-Kutta point
0389       // qop should be recalculated at every point
0390       // Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0391       const point3_type pos2 = pos + h * sd.t[0u] + h2 * 0.5f * sd.dtds[2u];
0392       bvec = magnetic_field.at(pos2[0], pos2[1], pos2[2]);
0393       assert(math::isfinite(bvec[0]));
0394       assert(math::isfinite(bvec[1]));
0395       assert(math::isfinite(bvec[2]));
0396       sd.b_last[0] = bvec[0];
0397       sd.b_last[1] = bvec[1];
0398       sd.b_last[2] = bvec[2];
0399
0400       DETRAY_DEBUG_HOST_DEVICE("Third stage:");
0401       DETRAY_DEBUG_HOST_DEVICE("-> B-field: [%f, %f, %f]", bvec[0], bvec[1],
0402                                bvec[2]);
0403
0404       detray::tie(sd.dqopds[3u], sd.qop[3u]) =
0405           evaluate_dqopds(stepping, 3u, h, sd.dqopds[2u], vol_mat_ptr, cfg);
0406       detray::tie(sd.dtds[3u], sd.t[3u]) =
0407           evaluate_dtds(stepping, sd.b_last, 3u, h, sd.dtds[2u], sd.qop[3u]);
0408
0409       // Compute and check the local integration error estimate
0410       // @Todo
0411       constexpr auto one_sixth{static_cast<scalar_type>(1. / 6.)};
0412       const vector3_type err_vec =
0413           one_sixth * h2 *
0414           (sd.dtds[0u] - sd.dtds[1u] - sd.dtds[2u] + sd.dtds[3u]);
0415
0416       DETRAY_DEBUG_HOST("-> Integration error vector: " << err_vec);
0417
0418       return vector::norm(err_vec);
0419     };
0420
0421     /// Calculate the scale factor for the stepsize adjustment using the
0422     /// error estimate @param err
0423     const auto step_size_scaling =
0424         [&cfg](const scalar_type& err) -> scalar_type {
0425       assert(err != 0.f);
0426       return static_cast<scalar_type>(
0427           math::min(math::max(math::sqrt(math::sqrt(cfg.rk_error_tol / err)),
0428                               static_cast<scalar_type>(0.25)),
0429                     static_cast<scalar_type>(4.)));
0430     };
0431
0432     scalar_type error{1e20f};
0433
0434     // If the estimated error is larger than the tolerance with an
0435     // additional margin, reduce the step size and try again
0436     const auto n_trials{cfg.max_rk_updates};
0437     for (unsigned int i = 0u; i < n_trials; i++) {
0438       stepping.count_trials();
0439
0440       error = math::max(estimate_error(stepping.step_size()),
0441                         static_cast<scalar_type>(1e-20));
0442
0443       DETRAY_DEBUG_HOST_DEVICE("-> Integration error: %f", error);
0444       assert(math::isfinite(error));
0445
0446       // Error is small enough
0447       // ---> break and advance track
0448       if (error <= 4.f * cfg.rk_error_tol) {
0449         break;
0450       }
0451       // Error estimate is too big
0452       // ---> Make step size smaller and estimate error again
0453       else {
0454         stepping.set_step_size(stepping.step_size() * step_size_scaling(error));
0455
0456         // Run inspection while the stepsize is getting adjusted
0457         stepping.run_inspector(cfg, "Adjust stepsize: ", dist_to_next, i,
0458                                step_size_scaling(error));
0459       }
0460     }
0461
0462     stepping.dtds(sd.dtds[3u]);
0463     stepping.dqopds(sd.dqopds[3u]);
0464
0465     // Check constraints
0466     if (const scalar_type max_step =
0467             stepping.constraints().template size<>(stepping.direction());
0468         math::fabs(stepping.step_size()) > math::fabs(max_step)) {
0469       // Run inspection before step size is cut
0470       stepping.run_inspector(cfg, "Before constraint: ", dist_to_next);
0471
0472       stepping.set_step_size(max_step);
0473     }
0474
0475     // Adjust the min step size
0476     if (math::fabs(stepping.step_size()) < min_stepsize) {
0477       stepping.set_step_size(
0478           math::copysign(min_stepsize, stepping.step_size()));
0479     }
0480
0481     DETRAY_VERBOSE_HOST_DEVICE("Take step: %f mm", stepping.step_size());
0482
0483     // The step size estimation for the next step
0484     stepping.set_next_step_size(stepping.step_size() *
0485                                 step_size_scaling(error));
0486
0487     // Don't allow a too small step size
0488     if (math::fabs(stepping.next_step_size()) < cfg.min_stepsize) {
0489       stepping.set_next_step_size(
0490           math::copysign(static_cast<scalar_type>(cfg.min_stepsize),
0491                          stepping.next_step_size()));
0492     }
0493
0494     DETRAY_DEBUG_HOST_DEVICE("Estimated next stepsize: %f mm",
0495                              stepping.next_step_size());
0496
0497     assert(math::fabs(stepping.next_step_size()) >= cfg.min_stepsize);
0498     assert(math::fabs(stepping.step_size()) >= cfg.min_stepsize);
0499
0500     // Advance track state
0501     advance_track(stepping, sd, vol_mat_ptr);
0502     assert(!stepping().is_invalid());
0503
0504     // Advance jacobian transport
0505     if constexpr ((flags_v & static_cast<std::uint32_t>(
0506                                  detray::rk_stepper_flags::
0507                                      e_allow_covariance_transport)) != 0u) {
0508       if (cfg.do_covariance_transport) {
0509         advance_jacobian(stepping, cfg, sd, vol_mat_ptr);
0510       }
0511     } else {
0512       assert(!cfg.do_covariance_transport);
0513     }
0514
0515     // Run final inspection
0516     stepping.run_inspector(cfg, "Step complete: ", dist_to_next);
0517
0518     DETRAY_DEBUG_HOST("After: " << stepping());
0519     return true;
0520   }
0521
0522   /// evaluate dqopds for a given step size and material
0523   DETRAY_HOST_DEVICE
0524   detray::pair<scalar_type, scalar_type> evaluate_dqopds(
0525       detray::rk_stepper<magnetic_field_t, algebra_t, constraint_t, policy_t,
0526                          inspector_t, flags_v>::state& stepping,
0527       const std::size_t i, const scalar_type h, const scalar_type dqopds_prev,
0528       const material<scalar_type>* vol_mat_ptr,
0529       const detray::stepping::config& cfg) const {
0530     const auto& track = stepping();
0531
0532     if (vol_mat_ptr == nullptr) {
0533       const scalar_type qop = track.qop();
0534       DETRAY_DEBUG_HOST_DEVICE("-> qop: %f", qop);
0535       DETRAY_DEBUG_HOST_DEVICE("-> dqopds: 0");
0536
0537       return detray::make_pair(scalar_type(0.f), qop);
0538     } else if (cfg.use_mean_loss && i != 0u) {
0539       // qop_n is calculated recursively like the direction of
0540       // evaluate_dtds.
0541       //
0542       // https://doi.org/10.1016/0029-554X(81)90063-X says:
0543       // "For y  we  have  similar  formulae  as  for x, for y' and
0544       // \lambda similar  formulae as for  x'"
0545       const scalar_type qop = track.qop() + h * dqopds_prev;
0546       DETRAY_DEBUG_HOST_DEVICE("-> qop: %f", qop);
0547
0548       return detray::make_pair(stepping.dqopds(qop, vol_mat_ptr), qop);
0549     } else {
0550       const scalar_type qop = track.qop();
0551       DETRAY_DEBUG_HOST_DEVICE("-> qop: %f", qop);
0552
0553       return detray::make_pair(stepping.dqopds(qop, vol_mat_ptr), qop);
0554     }
0555   }
0556
0557   /// Evaluate dtds for Runge-Kutta stepping
0558   DETRAY_HOST_DEVICE
0559   detray::pair<vector3_type, vector3_type> evaluate_dtds(
0560       detray::rk_stepper<magnetic_field_t, algebra_t, constraint_t, policy_t,
0561                          inspector_t, flags_v>::state& stepping,
0562       const vector3_type& b_field, const std::size_t i, const scalar_type h,
0563       const vector3_type& dtds_prev, const scalar_type qop) const {
0564     auto& track = stepping();
0565     const auto dir = track.dir();
0566
0567     assert(std::isfinite(qop));
0568
0569     // Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0570     vector3_type t{(i == 0u || h == 0.f)
0571                        ? dir
0572                        : static_cast<vector3_type>(dir + h * dtds_prev)};
0573
0574     // dtds = qop * (t X B) from Lorentz force
0575     DETRAY_DEBUG_HOST(
0576         "-> evaluate dtds: " << vector3_type{qop * vector::cross(t, b_field)});
0577
0578     return detray::make_pair(vector3_type{qop * vector::cross(t, b_field)}, t);
0579   }
0580
0581   /// Update the track state by Runge-Kutta-Nystrom integration.
0582   DETRAY_HOST_DEVICE
0583   void advance_track(
0584       detray::rk_stepper<magnetic_field_t, algebra_t, constraint_t, policy_t,
0585                          inspector_t, flags_v>::state& stepping,
0586       const intermediate_state& sd,
0587       const material<scalar_type>* vol_mat_ptr) const {
0588     const scalar_type h{stepping.step_size()};
0589     const scalar_type h_6{h * static_cast<scalar_type>(1. / 6.)};
0590     auto& track = stepping();
0591     auto pos = track.pos();
0592     auto dir = track.dir();
0593
0594     // Update the track parameters according to the equations of motion
0595     // Reference: Eq (82) of https://doi.org/10.1016/0029-554X(81)90063-X
0596     pos = pos + h * (sd.t[0u] + h_6 * (sd.dtds[0] + sd.dtds[1] + sd.dtds[2]));
0597     track.set_pos(pos);
0598
0599     // Reference: Eq (82) of https://doi.org/10.1016/0029-554X(81)90063-X
0600     dir =
0601         dir + h_6 * (sd.dtds[0] + 2.f * (sd.dtds[1] + sd.dtds[2]) + sd.dtds[3]);
0602     dir = vector::normalize(dir);
0603     track.set_dir(dir);
0604
0605     auto qop = track.qop();
0606     if (vol_mat_ptr != nullptr) {
0607       // Reference: Eq (82) of
0608       // https://doi.org/10.1016/0029-554X(81)90063-X
0609       qop = qop + h_6 * (sd.dqopds[0u] + 2.f * (sd.dqopds[1u] + sd.dqopds[2u]) +
0610                          sd.dqopds[3u]);
0611     }
0612     track.set_qop(qop);
0613
0614     // Update path lengths
0615     stepping.update_path_lengths(h);
0616   }
0617
0618   /// Update the jacobian transport from free propagation
0619   DETRAY_HOST_DEVICE
0620   void advance_jacobian(
0621       detray::rk_stepper<magnetic_field_t, algebra_t, constraint_t, policy_t,
0622                          inspector_t, flags_v>::state& stepping,
0623       const stepping::config& cfg, const intermediate_state& sd,
0624       const material<scalar_type>* vol_mat_ptr) const {
0625     DETRAY_VERBOSE_HOST_DEVICE("-> Advance Jacobian...");
0626
0627     /// The calculations are based on ATL-SOFT-PUB-2009-002. The update of
0628     /// the Jacobian matrix is requires only the calculation of eq. 17
0629     /// and 18. Since the terms of eq. 18 are currently 0, this matrix is
0630     /// not needed in the calculation. The matrix A from eq. 17 consists out
0631     /// of 3 different parts. The first one is given by the upper left 3x3
0632     /// matrix that are calculated by the derivatives dF/dT (called dFdT)
0633     /// and dG/dT (calls dGdT). The second is given by the top 3 lines of
0634     /// the rightmost column. This is calculated by dFdqop and dGdqop. The
0635     /// remaining non-zero term is calculated directly. The naming of the
0636     /// variables is explained in eq. 11 and are directly related to the
0637     /// initial problem in eq. 7. The evaluation is based by propagating the
0638     /// parameters T and lambda as given in eq. 16 and evaluating the
0639     /// derivations for matrix A.
0640     /// @note The translation for u_{n+1} in eq. 7 is in this case a
0641     /// 3-dimensional vector without a dependency of Lambda or lambda
0642     /// neither in u_n nor in u_n'. The second and fourth eq. in eq. 14 have
0643     /// the constant offset matrices h * Id and Id respectively. This
0644     /// involves that the constant offset does not exist for rectangular
0645     /// matrix dGdu' (due to the missing Lambda part) and only exists for
0646     /// dFdu' in dlambda/dlambda.
0647
0648     const scalar_type h{stepping.step_size()};
0649     const scalar_type h2{h * h};
0650     const scalar_type half_h{h * 0.5f};
0651     const scalar_type h_6{h * (1.f / 6.f)};
0652
0653     /*---------------------------------------------------------------------------
0654      *  dk_n/dt1
0655      *    = qop_n * (dt_n/dt1 X B_n)
0656      *      + qop_n * ( t_n X dB_n/dt1 ),
0657      *  where dB_n/dt1 == dB_n/dr_n * dr_n/dt1.
0658      *
0659      *  The second term is non-zero only for inhomogeneous magnetic fields
0660      *
0661      *  Note that [ t_n = t1 + h * d(t_{n-1})/ds) ] as indicated by
0662      *  Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0663
0664      *  [ Table for dt_n/dt1 ]
0665      *  dt1/dt1 = I
0666      *  dt2/dt1 = d( t1 + h/2 * dt1/ds ) / dt1 = I + h/2 * dk1/dt1
0667      *  dt3/dt1 = d( t1 + h/2 * dt2/ds ) / dt1 = I + h/2 * dk2/dt1
0668      *  dt4/dt1 = d( t1 + h * dt3/ds ) / dt1 = I + h * dk3/dt1
0669      *
0670      *  [ Table for dr_n/dt1 ]
0671      *  dr1/dt1 = 0
0672      *  dr2/dt1 = d(r1 + h/2 * t1 + h^2/8 dt1/ds)/dt1 = h/2 * I + h^2/8
0673     dk1/dt1
0674      *  dr3/dt1 = d(r1 + h/2 * t1 + h^2/8 dt1/ds)/dt1 = h/2 * I + h^2/8
0675     dk1/dt1
0676      *  dr4/dt1 = d(r1 + h * t1 + h^2/2 dt3/ds)/dt1 = h * I + h^2/2 dk3/dt1
0677      *
0678      *  Note that
0679      *  d/dr [ F(T) X B ]  = dF(T)/dr (X) B, where (X) means the column wise
0680      *  cross product
0681     ---------------------------------------------------------------------------*/
0682
0683     /*---------------------------------------------------------------------------
0684      *  dk_n/dqop_1
0685      *    = dqop_n/dqop1 * ( t_n X B_n )
0686      *      + qop_n * ( dt_n/dqop1 X B_n )
0687      *      + qop_n * ( t_n X dB_n/dqop1 ),
0688      *  where dB_n/dqop1 = dB_n/dr_n * dr_n/dqop1
0689      *
0690      *  Note that [ qop_n = qop1 + h * dqop_{n-1}/ds) ] as indicated by
0691      *  Eq (84) of https://doi.org/10.1016/0029-554X(81)90063-X
0692      *
0693      *  [ Table for dqop_n/dqop1 ]
0694      *  dqop1/dqop1 = 1
0695      *  dqop2/dqop1 = 1 + h/2 * d(dqop1/ds)/dqop1
0696      *  dqop3/dqop1 = 1 + h/2 * d(dqop2/ds)/dqop1
0697      *  dqop4/dqop1 = 1 + h * d(dqop3/ds)/dqop1
0698      *
0699      *  [ Table for dt_n/dqop1 ]
0700      *  dt1/dqop1 = 0
0701      *  dt2/dqop1 = d(t1 + h/2 dt1/ds)/dqop1 = h/2 * dk1/dqop1
0702      *  dt3/dqop1 = d(t1 + h/2 dt2/ds)/dqop1 = h/2 * dk2/dqop1
0703      *  dt4/dqop1 = d(t1 + h dt3/ds)/dqop1 = h * dk3/dqop1
0704      *
0705      *  [ Table for dr_n/dqop1 ]
0706      *  dr1/dqop1 = 0
0707      *  dr2/dqop1 = d(r1 + h/2 * t1 + h^2/8 dt1/ds)/dqop1 = h^2/8 *
0708     dk1/dqop1
0709      *  dr3/dqop1 = d(r1 + h/2 * t1 + h^2/8 dt1/ds)/dqop1 = h^2/8 *
0710     dk1/dqop1
0711      *  dr4/dqop1 = d(r1 + h * t1 + h^2/2 dt3/ds)/dqop1 = h^2/2 dk3/dqop1
0712     ---------------------------------------------------------------------------*/
0713
0714     /*---------------------------------------------------------------------------
0715      *  dk_n/dr1
0716      *    = qop_n * ( dt_n/dr1 X B_n )
0717      *      + qop_n * ( t_n X dB_n/dr1 ),
0718      *  where dB_n/dr1 = dB_n/dr_n * dr_n/dr1
0719      *
0720      *  [ Table for dt_n/dr1 ]
0721      *  dt1/dr1 = 0
0722      *  dt2/dr1 = d(t1 + h/2 * dt1/ds)/dr1 = h/2 * dk1/dr1
0723      *  dt2/dr1 = d(t1 + h/2 * dt2/ds)/dr1 = h/2 * dk2/dr1
0724      *  dt3/dr1 = d(t1 + h * dt3/ds)/dr1 = h * dk3/dr1
0725      *
0726      *  [ Table for dr_n/dr1 ]
0727      *  dr1/dr1 = I
0728      *  dr2/dr1 = (r1 + h/2 * t1 + h^2/8 dt1/ds ) / dr1 = I + h^2/8 dk1/dr1
0729      *  dr3/dr1 = (r1 + h/2 * t1 + h^2/8 dt1/ds ) / dr1 = I + h^2/8 dk1/dr1
0730      *  dr4/dr1 = (r1 + h * t1 + h^2/2 dt3/ds ) / dr1 = I + h^2/2 dk3/dr1
0731     ---------------------------------------------------------------------------*/
0732
0733     /*---------------------------------------------------------------------------
0734      *  d(dqop_n/ds)/dqop1
0735      *
0736      *  Useful equation:
0737      *  dqop/ds = qop^3 * E * (-dE/ds) / q^2 = - qop^3 * E g / q^2
0738
0739      *  [ Table for d(dqop_n/ds)/dqop1 ]
0740      *  d(dqop1/ds)/dqop1 = dqop1/ds * (1/qop * (3 - p^2/E^2) + 1/g1 *
0741     dg1dqop1
0742      *  d(dqop2/ds)/dqop1 = d(dqop2/ds)/dqop2 * (1 + h/2 *
0743     d(dqop1/ds)/dqop1)
0744      *  d(dqop3/ds)/dqop1 = d(dqop3/ds)/dqop3 * (1 + h/2 *
0745     d(dqop2/ds)/dqop1)
0746      *  d(dqop4/ds)/dqop1 = d(dqop4/ds)/dqop4 * (1 + h * d(dqop3/ds)/dqop1)
0747     ---------------------------------------------------------------------------*/
0748
0749     scalar_type dqopqop;
0750     vector3_type dFdqop;
0751     vector3_type dGdqop;
0752
0753     {
0754       darray<scalar_type, 4u> dqopn_dqop{1.f, 1.f, 1.f, 1.f};
0755
0756       if (!cfg.use_eloss_gradient) {
0757         dqopqop = 1.f;
0758       } else {
0759         // Pre-calculate dqop_n/dqop1
0760         const scalar_type d2qop1dsdqop1 =
0761             stepping.d2qopdsdqop(sd.qop[0u], vol_mat_ptr);
0762
0763         dqopn_dqop[0u] = 1.f;
0764         dqopn_dqop[1u] = 1.f + half_h * d2qop1dsdqop1;
0765
0766         const scalar_type d2qop2dsdqop1 =
0767             stepping.d2qopdsdqop(sd.qop[1u], vol_mat_ptr) * dqopn_dqop[1u];
0768         dqopn_dqop[2u] = 1.f + half_h * d2qop2dsdqop1;
0769
0770         const scalar_type d2qop3dsdqop1 =
0771             stepping.d2qopdsdqop(sd.qop[2u], vol_mat_ptr) * dqopn_dqop[2u];
0772         dqopn_dqop[3u] = 1.f + h * d2qop3dsdqop1;
0773
0774         const scalar_type d2qop4dsdqop1 =
0775             stepping.d2qopdsdqop(sd.qop[3u], vol_mat_ptr) * dqopn_dqop[3u];
0776
0777         /*-----------------------------------------------------------------
0778          * Calculate the first terms of d(dqop_n/ds)/dqop1
0779         -------------------------------------------------------------------*/
0780
0781         dqopqop =
0782             1.f + h_6 * (d2qop1dsdqop1 + 2.f * (d2qop2dsdqop1 + d2qop3dsdqop1) +
0783                          d2qop4dsdqop1);
0784       }
0785
0786       /*-----------------------------------------------------------------
0787        * Calculate the first and second terms of dk_n/dqop1
0788       -------------------------------------------------------------------*/
0789
0790       {
0791         darray<vector3_type, 4u> dkndqop;
0792
0793         // dk1/dqop1
0794         dkndqop[0u] = dqopn_dqop[0u] * vector::cross(sd.t[0u], sd.b_first);
0795
0796         // dk2/dqop1
0797         dkndqop[1u] =
0798             dqopn_dqop[1u] * vector::cross(sd.t[1u], sd.b_middle) +
0799             sd.qop[1u] * half_h * vector::cross(dkndqop[0u], sd.b_middle);
0800
0801         // dk3/dqop1
0802         dkndqop[2u] =
0803             dqopn_dqop[2u] * vector::cross(sd.t[2u], sd.b_middle) +
0804             sd.qop[2u] * half_h * vector::cross(dkndqop[1u], sd.b_middle);
0805
0806         // dk4/dqop1
0807         dkndqop[3u] = dqopn_dqop[3u] * vector::cross(sd.t[3u], sd.b_last) +
0808                       sd.qop[3u] * h * vector::cross(dkndqop[2u], sd.b_last);
0809
0810         // Set dF/dqop1 and dG/dqop1
0811         dFdqop = h * h_6 * (dkndqop[0u] + dkndqop[1u] + dkndqop[2u]);
0812         dGdqop = h_6 * (dkndqop[0u] + 2.f * (dkndqop[1u] + dkndqop[2u]) +
0813                         dkndqop[3u]);
0814       }
0815     }
0816
0817     /*-----------------------------------------------------------------
0818      * Calculate the first terms of dk_n/dt1
0819     -------------------------------------------------------------------*/
0820     // Set dF/dt1 and dG/dt1
0821     auto dFdt = matrix::identity<matrix_type<3, 3>>();
0822     auto dGdt = matrix::identity<matrix_type<3, 3>>();
0823
0824     {
0825       const auto I33 = matrix::identity<matrix_type<3, 3>>();
0826       darray<matrix_type<3u, 3u>, 4u> dkndt{I33, I33, I33, I33};
0827
0828       // dk1/dt1
0829       dkndt[0u] = sd.qop[0u] * matrix::column_wise_cross(dkndt[0u], sd.b_first);
0830
0831       // dk2/dt1
0832       dkndt[1u] = dkndt[1u] + half_h * dkndt[0u];
0833       dkndt[1u] =
0834           sd.qop[1u] * matrix::column_wise_cross(dkndt[1u], sd.b_middle);
0835
0836       // dk3/dt1
0837       dkndt[2u] = dkndt[2u] + half_h * dkndt[1u];
0838       dkndt[2u] =
0839           sd.qop[2u] * matrix::column_wise_cross(dkndt[2u], sd.b_middle);
0840
0841       // dk4/dt1
0842       dkndt[3u] = dkndt[3u] + h * dkndt[2u];
0843       dkndt[3u] = sd.qop[3u] * matrix::column_wise_cross(dkndt[3u], sd.b_last);
0844
0845       dFdt = dFdt + h_6 * (dkndt[0u] + dkndt[1u] + dkndt[2u]);
0846       dFdt = h * dFdt;
0847       dGdt =
0848           dGdt + h_6 * (dkndt[0u] + 2.f * (dkndt[1u] + dkndt[2u]) + dkndt[3u]);
0849     }
0850
0851     // Calculate dkndr in case of considering B field gradient
0852     auto dFdr = matrix::identity<matrix_type<3, 3>>();
0853     auto dGdr = matrix::zero<matrix_type<3, 3>>();
0854
0855     if constexpr ((flags_v & static_cast<std::uint32_t>(
0856                                  rk_stepper_flags::e_allow_field_gradient)) !=
0857                   0u) {
0858       if (cfg.use_field_gradient) {
0859         darray<matrix_type<3u, 3u>, 4u> dkndr;
0860         auto& track = stepping();
0861
0862         // Positions and field gradients at initial, middle and final
0863         // points of the fourth order RKN
0864         vector3_type r_ini = track.pos();
0865         vector3_type r_mid =
0866             r_ini + half_h * sd.t[0u] + h2 * 0.125f * sd.dtds[0u];
0867         vector3_type r_fin = r_ini + h * sd.t[0u] + h2 * 0.5f * sd.dtds[2u];
0868
0869         matrix_type<3, 3> dBdr_ini = evaluate_field_gradient(stepping, r_ini);
0870         matrix_type<3, 3> dBdr_mid = evaluate_field_gradient(stepping, r_mid);
0871         matrix_type<3, 3> dBdr_fin = evaluate_field_gradient(stepping, r_fin);
0872
0873         /*-----------------------------------------------------------------
0874          * Calculate all terms of dk_n/dr1
0875         -------------------------------------------------------------------*/
0876         const auto I33 = matrix::identity<matrix_type<3, 3>>();
0877
0878         // dk1/dr1
0879         dkndr[0u] = -sd.qop[0u] * matrix::column_wise_cross(dBdr_ini, sd.t[0u]);
0880
0881         const auto dkndr0_tmp = (I33 + h2 * 0.125f * dkndr[0u]);
0882
0883         // dk2/dr1
0884         dkndr[1u] = sd.qop[1u] *
0885                     matrix::column_wise_cross(half_h * dkndr[0u], sd.b_middle);
0886         dkndr[1u] =
0887             dkndr[1u] - sd.qop[1u] * matrix::column_wise_cross(
0888                                          dBdr_mid * dkndr0_tmp, sd.t[1u]);
0889
0890         // dk3/dr1
0891         dkndr[2u] = sd.qop[2u] *
0892                     matrix::column_wise_cross(half_h * dkndr[1u], sd.b_middle);
0893         dkndr[2u] =
0894             dkndr[2u] - sd.qop[2u] * matrix::column_wise_cross(
0895                                          dBdr_mid * dkndr0_tmp, sd.t[2u]);
0896
0897         // dk4/dr1
0898         dkndr[3u] =
0899             sd.qop[3u] * matrix::column_wise_cross(h * dkndr[2u], sd.b_last);
0900         dkndr[3u] = dkndr[3u] -
0901                     sd.qop[3u] *
0902                         matrix::column_wise_cross(
0903                             dBdr_fin * (I33 + h2 * 0.5f * dkndr[2u]), sd.t[3u]);
0904
0905         // Set dF/dr1 and dG/dr1
0906         dFdr = dFdr + h * h_6 * (dkndr[0u] + dkndr[1u] + dkndr[2u]);
0907         dGdr = h_6 * (dkndr[0u] + 2.f * (dkndr[1u] + dkndr[2u]) + dkndr[3u]);
0908       }
0909     } else {
0910       assert(!cfg.use_field_gradient);
0911     }
0912
0913     const auto old_jacobian = stepping.internal_transport_jacobian();
0914
0915     if constexpr ((flags_v & static_cast<std::uint32_t>(
0916                                  rk_stepper_flags::e_allow_field_gradient)) !=
0917                   0u) {
0918       detail::update_transport_jacobian_with_gradient_impl(
0919           old_jacobian, dFdt, dGdt, dFdr, dGdr, dFdqop, dGdqop, dqopqop,
0920           stepping.internal_transport_jacobian());
0921     } else {
0922       assert(!cfg.use_field_gradient);
0923
0924       detail::update_transport_jacobian_without_gradient_impl(
0925           old_jacobian, dFdt, dGdt, dFdqop, dGdqop, dqopqop,
0926           stepping.internal_transport_jacobian());
0927     }
0928   }
0929
0930   DETRAY_HOST_DEVICE
0931   matrix_type<3, 3> evaluate_field_gradient(
0932       detray::rk_stepper<magnetic_field_t, algebra_t, constraint_t, policy_t,
0933                          inspector_t, flags_v>::state& stepping,
0934       const point3_type& pos) const {
0935     auto dBdr = matrix::zero<matrix_type<3, 3>>();
0936
0937     constexpr auto delta{1e-1f * unit<scalar_type>::mm};
0938
0939     for (unsigned int i = 0; i < 3; i++) {
0940       point3_type dpos1 = pos;
0941       dpos1[i] += delta;
0942       const auto bvec1_tmp = stepping.field().at(dpos1[0], dpos1[1], dpos1[2]);
0943       vector3_type bvec1;
0944       bvec1[0u] = bvec1_tmp[0u];
0945       bvec1[1u] = bvec1_tmp[1u];
0946       bvec1[2u] = bvec1_tmp[2u];
0947
0948       point3_type dpos2 = pos;
0949       dpos2[i] -= delta;
0950       const auto bvec2_tmp = stepping.field().at(dpos2[0], dpos2[1], dpos2[2]);
0951       vector3_type bvec2;
0952       bvec2[0u] = bvec2_tmp[0u];
0953       bvec2[1u] = bvec2_tmp[1u];
0954       bvec2[2u] = bvec2_tmp[2u];
0955
0956       assert(delta != 0.f);
0957       const vector3_type gradient = (bvec1 - bvec2) * (1.f / (2.f * delta));
0958
0959       getter::element(dBdr, 0u, i) = gradient[0u];
0960       getter::element(dBdr, 1u, i) = gradient[1u];
0961       getter::element(dBdr, 2u, i) = gradient[2u];
0962     }
0963
0964     return dBdr;
0965   }
0966 };
0967
0968 }  // namespace detray