vc_soa_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_soa_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
 
namespace detray::algebra::vc_soa::math {

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
  requires(N >= 2)
DETRAY_HOST_DEVICE constexpr auto phi(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  return Vc::atan2(v[1], v[0]);
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
  requires(N >= 2)
DETRAY_HOST_DEVICE constexpr auto perp(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  return Vc::sqrt(Vc::fma(v[0], v[0], v[1] * v[1]));
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
  requires(N >= 3)
DETRAY_HOST_DEVICE constexpr auto theta(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  return Vc::atan2(perp(v), v[2]);
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
  requires(N == 3)
DETRAY_HOST_DEVICE constexpr algebra::storage::vector<N, Vc::Vector<value_t>,
                                                      array_t>
cross(const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &a,
      const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &b) {
  return {Vc::fma(a[1], b[2], -b[1] * a[2]), Vc::fma(a[2], b[0], -b[2] * a[0]),
          Vc::fma(a[0], b[1], -b[0] * a[1])};
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr Vc::Vector<value_t> dot(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &a,
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &b) {
  auto ret = a[0] * b[0];
 
  for (unsigned int i{1u}; i < N; i++) {
    ret = Vc::fma(a[i], b[i], ret);
  }
 
  return ret;
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr auto norm(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  return Vc::sqrt(dot(v, v));
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr algebra::storage::vector<N, Vc::Vector<value_t>,
                                                      array_t>
normalize(const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  return (Vc::Vector<value_t>::One() / norm(v)) * v;
 
  // Less accurate, but faster
  // return Vc::reciprocal(norm(v)) * v;
 
  // Even less accurate, but even faster
  // return Vc::rsqrt(dot(v, v)) * v;
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
  requires(N >= 3)
DETRAY_HOST_DEVICE constexpr auto eta(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  // atanh does not exist in Vc
  auto atanh_func = [](value_t e) { return std::atanh(e); };
 
  return (v[2] / norm(v)).apply(atanh_func);
 
  // Faster, but less accurate
  // return (Vc::reciprocal(norm(v)) * v[2]).apply(atanh_func);
 
  // Even faster, but even less accurate
  // return (Vc::rsqrt(dot(v, v)) * v[2]).apply(atanh_func);
}

template <std::size_t N, concepts::value value_t,
          template <typename, std::size_t> class array_t>
DETRAY_HOST_DEVICE constexpr Vc::Vector<value_t> sum(
    const algebra::storage::vector<N, Vc::Vector<value_t>, array_t> &v) {
  Vc::Vector<value_t> res{v[0]};
 
  for (std::size_t i = 1u; i < Vc::Vector<value_t>::size(); ++i) {
    res = res + v[i];
  }
 
  return res;
}
 
}  // namespace detray::algebra::vc_soa::math