base_stepper.hpp

// This file is part of the ACTS project.
//
// Copyright (C) 2016 CERN for the benefit of the ACTS project
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
#pragma once
 
// Project include(s).
#include "detray/core/concepts.hpp"
#include "detray/definitions/algebra.hpp"
#include "detray/definitions/detail/qualifiers.hpp"
#include "detray/definitions/pdg_particle.hpp"
#include "detray/definitions/units.hpp"
#include "detray/geometry/identifier.hpp"
#include "detray/geometry/tracking_surface.hpp"
#include "detray/propagator/constrained_step.hpp"
#include "detray/propagator/detail/print_stepper_state.hpp"
#include "detray/propagator/stepping_config.hpp"
#include "detray/propagator/transport_jacobian.hpp"
#include "detray/tracks/tracks.hpp"
#include "detray/utils/curvilinear_frame.hpp"
#include "detray/utils/logging.hpp"
 
namespace detray {
 
namespace stepping {

struct void_inspector {
  template <typename state_t>
  DETRAY_HOST_DEVICE constexpr void operator()(const state_t & /*ignored*/,
                                               const char * /*ignored*/) const {
    /*Do nothing*/
  }
};
 
}  // namespace stepping

template <concepts::algebra algebra_t, typename constraint_t, typename policy_t,
          typename inspector_t = stepping::void_inspector>
class base_stepper {
 public:
  using scalar_type = dscalar<algebra_t>;

  using free_track_parameters_type = free_track_parameters<algebra_t>;
  using bound_track_parameters_type = bound_track_parameters<algebra_t>;
  using free_matrix_type = free_matrix<algebra_t>;
  using bound_matrix_type = bound_matrix<algebra_t>;
  using magnetic_field_type = void;

  using inspector_type = inspector_t;
  using policy_type = policy_t;

  template <typename internal_jacobian_t = free_matrix_type>
  struct state {
    using internal_jacobian_type = internal_jacobian_t;

    DETRAY_HOST_DEVICE
    explicit state(const free_track_parameters_type &free_params)
        : m_track(free_params) {
      assert(!m_track.is_invalid());
 
      // HACK: When the overload resolution for the transport Jacobian
      // type is resolved, turn this into a default member
      // initialization.
      if constexpr (concepts::transport_jacobian<internal_jacobian_type>) {
        m_jac_transport = internal_jacobian_type::identity();
      } else {
        m_jac_transport = matrix::identity<internal_jacobian_type>();
      }
    }

    template <concepts::detector detector_t>
    DETRAY_HOST_DEVICE state(const bound_track_parameters_type &bound_params,
                             const detector_t &det,
                             const typename detector_t::geometry_context &ctx) {
      assert(!bound_params.is_invalid());
      assert(!bound_params.surface_link().is_invalid());
 
      // Departure surface
      const auto sf = tracking_surface{det, bound_params.surface_link()};
 
      // Set free track parameters for stepping/navigation
      m_track = sf.bound_to_free_vector(ctx, bound_params);
 
      assert(!m_track.is_invalid());
 
      // HACK: When the overload resolution for the transport Jacobian
      // type is resolved, turn this into a default member
      // initialization.
      if constexpr (concepts::transport_jacobian<internal_jacobian_type>) {
        m_jac_transport = internal_jacobian_type::identity();
      } else {
        m_jac_transport = matrix::identity<internal_jacobian_type>();
      }
    }

    DETRAY_HOST_DEVICE
    free_track_parameters_type &operator()() { return m_track; }

    DETRAY_HOST_DEVICE
    const free_track_parameters_type &operator()() const { return m_track; }

    DETRAY_HOST_DEVICE
    inline step::direction direction() const {
      return m_step_size >= 0.f ? step::direction::e_forward
                                : step::direction::e_backward;
    }

    DETRAY_HOST_DEVICE inline void count_trials() { ++m_n_total_trials; }

    DETRAY_HOST_DEVICE inline std::size_t n_total_trials() const {
      return m_n_total_trials;
    }

    DETRAY_HOST_DEVICE
    inline void set_step_size(const scalar_type step) {
      assert(math::fabs(step) >= 1e-4f * unit<scalar_type>::mm);
      m_step_size = step;
    }

    DETRAY_HOST_DEVICE
    inline scalar_type step_size() const { return m_step_size; }

    DETRAY_HOST_DEVICE
    inline scalar_type path_length() const { return m_path_length; }

    DETRAY_HOST_DEVICE
    inline scalar_type abs_path_length() const { return m_abs_path_length; }

    DETRAY_HOST_DEVICE inline void update_path_lengths(scalar_type seg) {
      m_path_length += seg;
      m_abs_path_length += math::fabs(seg);
    }

    template <step::constraint type = step::constraint::e_actor>
    DETRAY_HOST_DEVICE inline void set_constraint(scalar_type step_size) {
      m_constraint.template set<type>(step_size);
    }

    DETRAY_HOST_DEVICE
    inline const constraint_t &constraints() const { return m_constraint; }

    template <step::constraint type = step::constraint::e_actor>
    DETRAY_HOST_DEVICE inline void release_step() {
      m_constraint.template release<type>();
    }

    DETRAY_HOST_DEVICE
    const auto &particle_hypothesis() const { return m_ptc; }

    DETRAY_HOST_DEVICE
    inline void set_particle(const pdg_particle<scalar_type> &ptc) {
      m_ptc = ptc;
    }

    DETRAY_HOST_DEVICE
    inline free_matrix_type transport_jacobian() const
      requires(!std::same_as<free_matrix_type, internal_jacobian_type>)
    {
      return m_jac_transport.operator dmatrix<algebra_t, 8, 8>();
    }

    DETRAY_HOST_DEVICE
    inline const free_matrix_type &transport_jacobian() const
      requires std::same_as<free_matrix_type, internal_jacobian_type>
    {
      return m_jac_transport;
    }

    DETRAY_HOST_DEVICE
    inline free_matrix_type &transport_jacobian()
      requires std::same_as<free_matrix_type, internal_jacobian_type>
    {
      return m_jac_transport;
    }

    DETRAY_HOST_DEVICE
    inline internal_jacobian_type &internal_transport_jacobian() {
      return m_jac_transport;
    }

    DETRAY_HOST_DEVICE
    inline const internal_jacobian_type &internal_transport_jacobian() const {
      return m_jac_transport;
    }

    DETRAY_HOST_DEVICE
    inline void reset_transport_jacobian() {
      // HACK: When the overload resolution for the transport Jacobian
      // type is resolved, remove this conditional logic.
      if constexpr (concepts::transport_jacobian<internal_jacobian_type>) {
        m_jac_transport = internal_jacobian_type::identity();
      } else {
        m_jac_transport = matrix::identity<internal_jacobian_type>();
      }
    }

    DETRAY_HOST_DEVICE
    inline typename policy_t::state &policy_state() { return m_policy_state; }

    DETRAY_HOST
    constexpr auto &inspector() { return m_inspector; }

    DETRAY_HOST_DEVICE
    inline void run_inspector([[maybe_unused]] const stepping::config &cfg,
                              [[maybe_unused]] const char *message,
                              [[maybe_unused]] const scalar_type dist) {
      if constexpr (!std::is_same_v<inspector_t, stepping::void_inspector>) {
        m_inspector(*this, cfg, message, dist);
      }
 
      DETRAY_DEBUG_HOST("" << message << "\n"
                           << detray::stepping::print_state(*this, dist));
    }
 
   private:
    internal_jacobian_type m_jac_transport;

    free_track_parameters_type m_track;

    pdg_particle<scalar_type> m_ptc = muon<scalar_type>();

    std::size_t m_n_total_trials{0u};

    scalar_type m_step_size{0.f};

    scalar_type m_path_length{0.f};

    scalar_type m_abs_path_length{0.f};

    DETRAY_NO_UNIQUE_ADDRESS constraint_t m_constraint = {};

    DETRAY_NO_UNIQUE_ADDRESS typename policy_t::state m_policy_state = {};

    DETRAY_NO_UNIQUE_ADDRESS inspector_type m_inspector;
  };
};
 
}  // namespace detray