generic_transform3.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/algebra/generic/algorithms/matrix/inverse/hard_coded.hpp"
#include "detray/algebra/generic/impl/generic_matrix.hpp"
#include "detray/algebra/generic/impl/generic_vector.hpp"
#include "detray/algebra/type_traits.hpp"
#include "detray/algebra/utils/approximately_equal.hpp"
#include "detray/definitions/detail/qualifiers.hpp"
 
// System include(s)
#include <cassert>
#include <concepts>
#include <limits>
 
namespace detray::algebra::generic::math {

template <concepts::index index_t, concepts::scalar scalar_t,
          template <typename, index_t, index_t> class matrix_t,
          template <typename, index_t> class array_t>
struct transform3 {

  using scalar_type = scalar_t;

  template <index_t N>
  using array_type = array_t<scalar_t, N>;
 
  // 4 x 4 Matrix (keep four rows for better alignment)
  using matrix44 = matrix_t<scalar_t, 4, 4>;
  static_assert(concepts::square_matrix<matrix44>);

  using vector2 = detray::traits::get_vector_t<matrix44, 2, scalar_t>;
  using vector3 = detray::traits::get_vector_t<matrix44, 3, scalar_t>;
  using point2 = vector2;
  using point3 = vector3;

  using element_getter = detray::traits::element_getter_t<matrix44>;
  using block_getter = detray::traits::block_getter_t<matrix44>;

  using matrix_inversion = generic::matrix::inverse::hard_coded<matrix44>;

  template <concepts::scalar o_scalar_t>
  using other_type = transform3<index_t, o_scalar_t, matrix_t, array_t>;


  matrix44 _data{
      generic::math::identity<matrix44>()};  

  matrix44 _data_inv{
      generic::math::identity<matrix44>()};  



  constexpr transform3() = default;

  DETRAY_HOST_DEVICE
  transform3(const vector3 &t, const vector3 &x, const vector3 &y,
             const vector3 &z, bool get_inverse = true) {
    element_getter{}(_data, 0, 0) = element_getter{}(x, 0);
    element_getter{}(_data, 1, 0) = element_getter{}(x, 1);
    element_getter{}(_data, 2, 0) = element_getter{}(x, 2);
    element_getter{}(_data, 3, 0) = 0.f;
    element_getter{}(_data, 0, 1) = element_getter{}(y, 0);
    element_getter{}(_data, 1, 1) = element_getter{}(y, 1);
    element_getter{}(_data, 2, 1) = element_getter{}(y, 2);
    element_getter{}(_data, 3, 1) = 0.f;
    element_getter{}(_data, 0, 2) = element_getter{}(z, 0);
    element_getter{}(_data, 1, 2) = element_getter{}(z, 1);
    element_getter{}(_data, 2, 2) = element_getter{}(z, 2);
    element_getter{}(_data, 3, 2) = 0.f;
    element_getter{}(_data, 0, 3) = element_getter{}(t, 0);
    element_getter{}(_data, 1, 3) = element_getter{}(t, 1);
    element_getter{}(_data, 2, 3) = element_getter{}(t, 2);
    element_getter{}(_data, 3, 3) = 1.f;
 
    if (get_inverse) {
      _data_inv = matrix_inversion{}(_data);
    }
  }

  DETRAY_HOST_DEVICE
  transform3(const vector3 &t, const vector3 &z, const vector3 &x,
             bool get_inverse = true)
      : transform3(t, x, cross(z, x), z, get_inverse) {}

  DETRAY_HOST_DEVICE
  explicit transform3(const vector3 &t) {
    element_getter{}(_data, 0, 0) = 1.f;
    element_getter{}(_data, 1, 0) = 0.f;
    element_getter{}(_data, 2, 0) = 0.f;
    element_getter{}(_data, 3, 0) = 0.f;
    element_getter{}(_data, 0, 1) = 0.f;
    element_getter{}(_data, 1, 1) = 1.f;
    element_getter{}(_data, 2, 1) = 0.f;
    element_getter{}(_data, 3, 1) = 0.f;
    element_getter{}(_data, 0, 2) = 0.f;
    element_getter{}(_data, 1, 2) = 0.f;
    element_getter{}(_data, 2, 2) = 1.f;
    element_getter{}(_data, 3, 2) = 0.f;
    element_getter{}(_data, 0, 3) = element_getter{}(t, 0);
    element_getter{}(_data, 1, 3) = element_getter{}(t, 1);
    element_getter{}(_data, 2, 3) = element_getter{}(t, 2);
    element_getter{}(_data, 3, 3) = 1.f;
 
    _data_inv = matrix_inversion{}(_data);
  }

  DETRAY_HOST_DEVICE
  explicit transform3(const matrix44 &m) : _data{m} {
    _data_inv = matrix_inversion{}(_data);
  }

  DETRAY_HOST_DEVICE
  transform3(const matrix44 &m, const matrix44 &m_inv)
      : _data{m}, _data_inv{m_inv} {
    // The assertion will not hold for (casts to) int
    if constexpr (std::floating_point<scalar_type>) {
      // The concrete type of matrix mult is not available at this point
      auto prod = generic::math::zero<matrix44>();
 
      constexpr element_getter elem{};
 
      for (index_t i = 0; i < 4; ++i) {
        for (index_t j = 0; j < 4; ++j) {
          for (index_t k = 0; k < 4; ++k) {
            elem(prod, k, j) += elem(m, k, i) * elem(m_inv, i, j);
          }
        }
      }
 
      [[maybe_unused]] constexpr auto epsilon{
          std::numeric_limits<scalar_type>::epsilon()};
      assert(algebra::approx_equal(prod, generic::math::identity<matrix44>(),
                                   16.f * epsilon, 1e-6f));
    }
  }

  DETRAY_HOST_DEVICE
  explicit transform3(const array_type<16> &ma) {
    element_getter{}(_data, 0, 0) = ma[0];
    element_getter{}(_data, 1, 0) = ma[4];
    element_getter{}(_data, 2, 0) = ma[8];
    element_getter{}(_data, 3, 0) = 0.f;
    element_getter{}(_data, 0, 1) = ma[1];
    element_getter{}(_data, 1, 1) = ma[5];
    element_getter{}(_data, 2, 1) = ma[9];
    element_getter{}(_data, 3, 1) = 0.f;
    element_getter{}(_data, 0, 2) = ma[2];
    element_getter{}(_data, 1, 2) = ma[6];
    element_getter{}(_data, 2, 2) = ma[10];
    element_getter{}(_data, 3, 2) = 0.f;
    element_getter{}(_data, 0, 3) = ma[3];
    element_getter{}(_data, 1, 3) = ma[7];
    element_getter{}(_data, 2, 3) = ma[11];
    element_getter{}(_data, 3, 3) = 1.f;
 
    _data_inv = matrix_inversion{}(_data);
  }

  DETRAY_HOST_DEVICE
  constexpr bool operator==(const transform3 &rhs) const {
    for (index_t j = 0; j < 4; j++) {
      // Check only the rows that can differ
      for (index_t i = 0; i < 3; i++) {
        if (element_getter{}(_data, i, j) !=
            element_getter{}(rhs._data, i, j)) {
          return false;
        }
      }
    }
 
    return true;
  }

  DETRAY_HOST_DEVICE
  static constexpr vector3 rotate(const matrix44 &m, const vector2 &v) {
    vector3 ret{0.f, 0.f, 0.f};
 
    element_getter{}(ret, 0) +=
        element_getter{}(m, 0, 0) * element_getter{}(v, 0);
    element_getter{}(ret, 1) +=
        element_getter{}(m, 1, 0) * element_getter{}(v, 0);
    element_getter{}(ret, 2) +=
        element_getter{}(m, 2, 0) * element_getter{}(v, 0);
 
    element_getter{}(ret, 0) +=
        element_getter{}(m, 0, 1) * element_getter{}(v, 1);
    element_getter{}(ret, 1) +=
        element_getter{}(m, 1, 1) * element_getter{}(v, 1);
    element_getter{}(ret, 2) +=
        element_getter{}(m, 2, 1) * element_getter{}(v, 1);
 
    return ret;
  }

  DETRAY_HOST_DEVICE
  static constexpr vector3 rotate(const matrix44 &m, const vector3 &v) {
    vector3 ret{0.f, 0.f, 0.f};
 
    element_getter{}(ret, 0) +=
        element_getter{}(m, 0, 0) * element_getter{}(v, 0);
    element_getter{}(ret, 1) +=
        element_getter{}(m, 1, 0) * element_getter{}(v, 0);
    element_getter{}(ret, 2) +=
        element_getter{}(m, 2, 0) * element_getter{}(v, 0);
 
    element_getter{}(ret, 0) +=
        element_getter{}(m, 0, 1) * element_getter{}(v, 1);
    element_getter{}(ret, 1) +=
        element_getter{}(m, 1, 1) * element_getter{}(v, 1);
    element_getter{}(ret, 2) +=
        element_getter{}(m, 2, 1) * element_getter{}(v, 1);
 
    element_getter{}(ret, 0) +=
        element_getter{}(m, 0, 2) * element_getter{}(v, 2);
    element_getter{}(ret, 1) +=
        element_getter{}(m, 1, 2) * element_getter{}(v, 2);
    element_getter{}(ret, 2) +=
        element_getter{}(m, 2, 2) * element_getter{}(v, 2);
 
    return ret;
  }

  DETRAY_HOST_DEVICE
  auto constexpr rotation() const {
    return block_getter{}.template operator()<3, 3>(_data, 0, 0);
  }

  DETRAY_HOST_DEVICE
  constexpr point3 x() const {
    return {element_getter{}(_data, 0, 0), element_getter{}(_data, 1, 0),
            element_getter{}(_data, 2, 0)};
  }

  DETRAY_HOST_DEVICE
  constexpr point3 y() const {
    return {element_getter{}(_data, 0, 1), element_getter{}(_data, 1, 1),
            element_getter{}(_data, 2, 1)};
  }

  DETRAY_HOST_DEVICE
  constexpr point3 z() const {
    return {element_getter{}(_data, 0, 2), element_getter{}(_data, 1, 2),
            element_getter{}(_data, 2, 2)};
  }

  DETRAY_HOST_DEVICE
  constexpr point3 translation() const {
    return {element_getter{}(_data, 0, 3), element_getter{}(_data, 1, 3),
            element_getter{}(_data, 2, 3)};
  }

  DETRAY_HOST_DEVICE
  constexpr const matrix44 &matrix() const { return _data; }

  DETRAY_HOST_DEVICE
  constexpr const matrix44 &matrix_inverse() const { return _data_inv; }

  DETRAY_HOST_DEVICE constexpr point3 point_to_global(const point2 &v) const {
    const vector3 rg = rotate(_data, v);
 
    return {element_getter{}(rg, 0) + element_getter{}(_data, 0, 3),
            element_getter{}(rg, 1) + element_getter{}(_data, 1, 3),
            element_getter{}(rg, 2) + element_getter{}(_data, 2, 3)};
  }

  DETRAY_HOST_DEVICE constexpr point3 point_to_global(const point3 &v) const {
    const vector3 rg = rotate(_data, v);
 
    return {element_getter{}(rg, 0) + element_getter{}(_data, 0, 3),
            element_getter{}(rg, 1) + element_getter{}(_data, 1, 3),
            element_getter{}(rg, 2) + element_getter{}(_data, 2, 3)};
  }

  DETRAY_HOST_DEVICE constexpr point3 point_to_local(const point3 &v) const {
    const vector3 rg = rotate(_data_inv, v);
 
    return {element_getter{}(rg, 0) + element_getter{}(_data_inv, 0, 3),
            element_getter{}(rg, 1) + element_getter{}(_data_inv, 1, 3),
            element_getter{}(rg, 2) + element_getter{}(_data_inv, 2, 3)};
  }

  DETRAY_HOST_DEVICE constexpr vector3 vector_to_global(
      const vector2 &v) const {
    return rotate(_data, v);
  }

  DETRAY_HOST_DEVICE constexpr vector3 vector_to_global(
      const vector3 &v) const {
    return rotate(_data, v);
  }

  DETRAY_HOST_DEVICE constexpr vector3 vector_to_local(const vector3 &v) const {
    return rotate(_data_inv, v);
  }
 
};  // struct transform3
 
}  // namespace detray::algebra::generic::math