vc_aos_vector.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 "algebra/impl/vc_aos_types.hpp"
#include "detray/algebra/common/math.hpp"
#include "detray/algebra/common/vector.hpp"
#include "detray/algebra/concepts.hpp"
#include "detray/definitions/detail/qualifiers.hpp"
 
// Vc include(s).
#ifdef _MSC_VER
#pragma warning(push, 0)
#endif  // MSVC
#include <Vc/Vc>
#ifdef _MSC_VER
#pragma warning(pop)
#endif  // MSVC
 
// System include(s).
#include <concepts>
#include <utility>
 
namespace detray::algebra::vc_aos::math {

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto phi(const vector_t &v) {
  return algebra::math::atan2(v[1], v[0]);
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto perp(const vector_t &v) {
  return algebra::math::sqrt(algebra::math::fma(v[0], v[0], v[1] * v[1]));
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto theta(const vector_t &v) {
  return algebra::math::atan2(perp(v), v[2]);
}

template <detray::concepts::vc_aos_vector vector_t1,
          detray::concepts::vc_aos_vector vector_t2>
DETRAY_HOST_DEVICE constexpr auto dot(const vector_t1 &a, const vector_t2 &b) {
  return (a * b).sum();
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto norm(const vector_t &v) {
  return algebra::math::sqrt(dot(v, v));
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto normalize(const vector_t &v) {
  return v / norm(v);
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto eta(const vector_t &v) noexcept {
  return algebra::math::atanh(v[2] / norm(v));
}

template <detray::concepts::vc_aos_vector vector_t1,
          detray::concepts::vc_aos_vector vector_t2>
  requires(std::same_as<detray::traits::value_t<vector_t1>, float> &&
           std::same_as<detray::traits::value_t<vector_t2>, float>)
DETRAY_HOST_DEVICE constexpr auto cross(const vector_t1 &a, const vector_t2 &b)
    -> decltype(a * b - b * a) {
  using T = detray::traits::value_t<vector_t1>;
  using simd_array_t = Vc::SimdArray<T, 4>;
 
  // Mask to write on the last two elements: [0011]
  const typename simd_array_t::mask_type m{simd_array_t::IndexesFromZero() > 1};
 
  // a     = [a1, a2, a3, 0]
  // a_rot = [a2, a3, 0, a1]
  simd_array_t a_rot = static_cast<simd_array_t>(a).rotated(1);
  // a_rot = [a2, a3, a1, a2]
  a_rot(m) = a_rot.rotated(1);
 
  // Same for b
  simd_array_t b_rot = static_cast<simd_array_t>(b).rotated(1);
  b_rot(m) = b_rot.rotated(1);
 
  simd_array_t res = a_rot * b_rot.rotated(1) - a_rot.rotated(1) * b_rot;
  // Restore trailing zero
  res[3] = 0.f;
 
  return res;
}

template <detray::concepts::vc_aos_vector vector_t1,
          detray::concepts::vc_aos_vector vector_t2>
  requires(std::same_as<detray::traits::value_t<vector_t1>, double> ||
           std::same_as<detray::traits::value_t<vector_t2>, double>)
DETRAY_HOST_DEVICE constexpr auto cross(const vector_t1 &a, const vector_t2 &b)
    -> decltype(a * b - b * a) {
  return {algebra::math::fma(a[1], b[2], -b[1] * a[2]),
          algebra::math::fma(a[2], b[0], -b[2] * a[0]),
          algebra::math::fma(a[0], b[1], -b[0] * a[1]), 0.f};
}

template <detray::concepts::vc_aos_vector vector_t>
DETRAY_HOST_DEVICE constexpr auto sum(const vector_t &v) {
  return v.get().sum();
}
 
}  // namespace detray::algebra::vc_aos::math