bound_track_parameters.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/definitions/algebra.hpp"
#include "detray/definitions/detail/qualifiers.hpp"
#include "detray/definitions/indexing.hpp"
#include "detray/definitions/math.hpp"
#include "detray/definitions/track_parametrization.hpp"
#include "detray/definitions/units.hpp"
#include "detray/geometry/identifier.hpp"
 
// System include(s)
#include <limits>
#include <ostream>
 
namespace detray {

template <concepts::algebra algebra_t>
struct bound_parameters_vector {

  using algebra_type = algebra_t;
  using scalar_type = dscalar<algebra_t>;
  using point2_type = dpoint2D<algebra_t>;
  using vector3_type = dvector3D<algebra_t>;
 
  // Underlying vector type related to bound track vector.
  using vector_type = bound_vector<algebra_t>;


  bound_parameters_vector() = default;

  DETRAY_HOST_DEVICE
  explicit bound_parameters_vector(const vector_type& vec) : m_vector(vec) {
    assert(!this->is_invalid());
  }

  DETRAY_HOST_DEVICE
  bound_parameters_vector(const point2_type& loc_p, const scalar_type phi,
                          const scalar_type theta, const scalar_type qop,
                          const scalar_type t) {
    getter::set_block(m_vector, loc_p, e_bound_loc0, 0u);
    getter::element(m_vector, e_bound_phi, 0u) = phi;
    getter::element(m_vector, e_bound_theta, 0u) = theta;
    getter::element(m_vector, e_bound_qoverp, 0u) = qop;
    getter::element(m_vector, e_bound_time, 0u) = t;
 
    assert(!this->is_invalid());
  }

  DETRAY_HOST_DEVICE
  bool operator==(const bound_parameters_vector& rhs) const {
    for (unsigned int i = 0u; i < e_bound_size; i++) {
      const auto lhs_val = getter::element(m_vector, i, 0u);
      const auto rhs_val = getter::element(rhs.vector(), i, 0u);
 
      if (math::fabs(lhs_val - rhs_val) >
          std::numeric_limits<scalar_type>::epsilon()) {
        return false;
      }
    }
 
    return true;
  }

  DETRAY_HOST_DEVICE
  scalar_type operator[](const std::size_t i) const {
    return getter::element(m_vector, static_cast<unsigned int>(i), 0u);
  }

  DETRAY_HOST_DEVICE
  decltype(auto) operator[](const std::size_t i) {
    return getter::element(m_vector, static_cast<unsigned int>(i), 0u);
  }

  DETRAY_HOST_DEVICE
  const vector_type& vector() const { return m_vector; }

  DETRAY_HOST_DEVICE
  vector_type& vector() { return m_vector; }

  DETRAY_HOST_DEVICE
  void set_vector(const vector_type& v,
                  [[maybe_unused]] const bool skip_check = false) {
    m_vector = v;
    assert(skip_check || !this->is_invalid());
  }

  DETRAY_HOST_DEVICE
  point2_type bound_local() const {
    return {getter::element(m_vector, e_bound_loc0, 0u),
            getter::element(m_vector, e_bound_loc1, 0u)};
  }

  DETRAY_HOST_DEVICE
  void set_bound_local(const point2_type& pos) {
    assert(math::isfinite(pos[0]));
    assert(math::isfinite(pos[1]));
    getter::set_block(m_vector, pos, e_bound_loc0, 0u);
  }

  DETRAY_HOST_DEVICE
  scalar_type phi() const { return getter::element(m_vector, e_bound_phi, 0u); }

  DETRAY_HOST_DEVICE
  void set_phi(const scalar_type phi) {
    assert(math::fabs(phi) <= constant<scalar_type>::pi);
    getter::element(m_vector, e_bound_phi, 0u) = phi;
  }

  DETRAY_HOST_DEVICE
  scalar_type theta() const {
    return getter::element(m_vector, e_bound_theta, 0u);
  }

  DETRAY_HOST_DEVICE
  void set_theta(const scalar_type theta) {
    assert(0.f < theta);
    assert(theta <= constant<scalar_type>::pi);
    getter::element(m_vector, e_bound_theta, 0u) = theta;
  }

  DETRAY_HOST_DEVICE
  vector3_type dir() const {
    const scalar_type phi{getter::element(m_vector, e_bound_phi, 0u)};
    const scalar_type theta{getter::element(m_vector, e_bound_theta, 0u)};
    const scalar_type sinTheta{math::sin(theta)};
 
    return {math::cos(phi) * sinTheta, math::sin(phi) * sinTheta,
            math::cos(theta)};
  }

  DETRAY_HOST_DEVICE
  scalar_type time() const {
    return getter::element(m_vector, e_bound_time, 0u);
  }

  DETRAY_HOST_DEVICE
  void set_time(const scalar_type t) {
    assert(math::isfinite(t));
    getter::element(m_vector, e_bound_time, 0u) = t;
  }

  DETRAY_HOST_DEVICE
  scalar_type qop() const {
    return getter::element(m_vector, e_bound_qoverp, 0u);
  }

  DETRAY_HOST_DEVICE
  void set_qop(const scalar_type qop) {
    assert(math::isfinite(qop));
    getter::element(m_vector, e_bound_qoverp, 0u) = qop;
  }

  DETRAY_HOST_DEVICE
  scalar_type qopT() const {
    const scalar_type theta{getter::element(m_vector, e_bound_theta, 0u)};
    const scalar_type sinTheta{math::sin(theta)};
    assert(sinTheta != 0.f);
    return getter::element(m_vector, e_bound_qoverp, 0u) / sinTheta;
  }

  DETRAY_HOST_DEVICE
  scalar_type qopz() const {
    const scalar_type theta{getter::element(m_vector, e_bound_theta, 0u)};
    const scalar_type cosTheta{math::cos(theta)};
    assert(cosTheta != 0.f);
    return getter::element(m_vector, e_bound_qoverp, 0u) / cosTheta;
  }

  DETRAY_HOST_DEVICE
  scalar_type p(const scalar_type q) const {
    assert(qop() != 0.f);
    assert(q * qop() > 0.f);
    return q / qop();
  }

  DETRAY_HOST_DEVICE
  vector3_type mom(const scalar_type q) const { return p(q) * dir(); }

  DETRAY_HOST_DEVICE
  scalar_type pT(const scalar_type q) const {
    assert(qop() != 0.f);
    assert(q * qop() > 0.f);
    return math::fabs(q / qop() * vector::perp(dir()));
  }

  DETRAY_HOST_DEVICE
  scalar_type pz(const scalar_type q) const {
    assert(qop() != 0.f);
    assert(q * qop() > 0.f);
    return math::fabs(q / qop() * dir()[2]);
  }

  DETRAY_HOST_DEVICE
  constexpr bool is_invalid(const bool do_check = true) const {
    if (!do_check) {
      return false;
    }
    bool inv_elem{false};
    bool is_all_zero{true};
    for (std::size_t i = 0u; i < e_bound_size; ++i) {
      inv_elem = inv_elem || !math::isfinite(getter::element(m_vector, i, 0u));
      is_all_zero =
          is_all_zero && (math::fabs(getter::element(m_vector, i, 0u)) == 0.f);
    }
 
    return (inv_elem || is_all_zero);
  }
 
 private:
  DETRAY_HOST
  friend std::ostream& operator<<(std::ostream& out_stream,
                                  const bound_parameters_vector& bparam) {
    out_stream << bparam.m_vector;
    return out_stream;
  }
 
  vector_type m_vector = matrix::zero<vector_type>();
};

template <concepts::algebra algebra_t>
struct bound_track_parameters : public bound_parameters_vector<algebra_t> {
  using base_type = bound_parameters_vector<algebra_t>;


  using algebra_type = algebra_t;
  using scalar_type = dscalar<algebra_t>;
  using point2_type = dpoint2D<algebra_t>;
  using vector3_type = dvector3D<algebra_t>;
 
  // Shorthand vector/matrix types related to bound track parameters.
  using parameter_vector_type = bound_parameters_vector<algebra_t>;
  using covariance_type = bound_matrix<algebra_t>;


  bound_track_parameters() = default;

  DETRAY_HOST_DEVICE
  bound_track_parameters(const geometry::identifier sf_idx,
                         const parameter_vector_type& vec,
                         const covariance_type& cov)
      : base_type(vec), m_covariance(cov), m_identifier(sf_idx) {}

  DETRAY_HOST_DEVICE
  bool operator==(const bound_track_parameters& rhs) const {
    if (m_identifier != rhs.surface_link()) {
      return false;
    }
 
    if (!base_type::operator==(rhs)) {
      return false;
    }
 
    for (unsigned int i = 0u; i < e_bound_size; i++) {
      for (unsigned int j = 0u; j < e_bound_size; j++) {
        const auto lhs_val = getter::element(m_covariance, i, j);
        const auto rhs_val = getter::element(rhs.covariance(), i, j);
 
        if (math::fabs(lhs_val - rhs_val) >
            std::numeric_limits<scalar_type>::epsilon()) {
          return false;
        }
      }
    }
 
    return true;
  }

  DETRAY_HOST_DEVICE
  const geometry::identifier& surface_link() const { return m_identifier; }

  DETRAY_HOST_DEVICE
  void set_surface_link(geometry::identifier link) { m_identifier = link; }

  DETRAY_HOST_DEVICE
  void set_parameter_vector(const parameter_vector_type& v,
                            const bool skip_check = false) {
    this->set_vector(v.vector(), skip_check);
  }

  DETRAY_HOST_DEVICE
  covariance_type& covariance() { return m_covariance; }

  DETRAY_HOST_DEVICE
  const covariance_type& covariance() const { return m_covariance; }

  DETRAY_HOST_DEVICE
  void set_covariance(const covariance_type& c) { m_covariance = c; }

  DETRAY_HOST_DEVICE
  constexpr bool is_invalid(const bool do_check = true) const {
    if (!do_check) {
      return false;
    }
    if (base_type::is_invalid()) {
      return true;
    }
 
    // @TODO: Add tests positive semi-definite, check the determinant etc
    return (m_covariance == matrix::zero<covariance_type>());
  }
 
 private:
  DETRAY_HOST
  friend std::ostream& operator<<(std::ostream& out_stream,
                                  const bound_track_parameters& bparam) {
    out_stream << "Surface: " << bparam.m_identifier << std::endl;
    out_stream << "Param.:\n " << static_cast<parameter_vector_type>(bparam)
               << std::endl;
    out_stream << "Cov.:\n" << bparam.m_covariance;
 
    return out_stream;
  }

  covariance_type m_covariance = matrix::zero<covariance_type>();
  geometry::identifier m_identifier{};
};
 
namespace algebra {

template <concepts::value value_t, concepts::algebra algebra_t>
DETRAY_HOST_DEVICE constexpr decltype(auto) cast_to(
    const bound_parameters_vector<algebra_t>& param_vec) {
  using new_algebra_t = dalgebra<algebra_t, value_t>;
 
  if constexpr (std::same_as<new_algebra_t, algebra_t>) {
    return param_vec;
  } else {
    return bound_parameters_vector<new_algebra_t>{
        ::detray::algebra::cast_to<value_t>(param_vec.to_vector())};
  }
}

template <concepts::value value_t, concepts::algebra algebra_t>
DETRAY_HOST_DEVICE constexpr decltype(auto) cast_to(
    const bound_track_parameters<algebra_t>& param) {
  using new_algebra_t = dalgebra<algebra_t, value_t>;
 
  if constexpr (std::same_as<new_algebra_t, algebra_t>) {
    return param;
  } else {
    const bound_parameters_vector<algebra_t>& param_vec = param;
    return bound_track_parameters<new_algebra_t>{
        param.surface_link(), ::detray::algebra::cast_to<value_t>(param_vec),
        ::detray::algebra::cast_to<value_t>(param.covariance())};
  }
}
 
}  // namespace algebra
 
}  // namespace detray