distance.h

Go to the documentation of this file.

#pragma once
 
#include <cmath>
#include <cstddef>
#include <type_traits>
#include <utility>
 
#include "clustering/ndarray.h"
 
#ifdef CLUSTERING_USE_AVX2
#include "clustering/math/detail/distance_avx2.h"
#endif
 
namespace clustering::math::distance {

struct SqEuclideanTag {};

struct CosineTag {};

struct ManhattanTag {};
 
namespace detail {

struct PointwiseSqFn {
  template <class Tag, class A, class B>
  constexpr auto operator()(Tag tag, A &&a, B &&b) const
      noexcept(noexcept(tag_invoke(std::declval<const PointwiseSqFn &>(), tag, std::forward<A>(a),
                                   std::forward<B>(b))))
          -> decltype(tag_invoke(std::declval<const PointwiseSqFn &>(), tag, std::forward<A>(a),
                                 std::forward<B>(b))) {
    return tag_invoke(*this, tag, std::forward<A>(a), std::forward<B>(b));
  }
};
 
} // namespace detail

inline constexpr detail::PointwiseSqFn pointwiseSq{};

template <class T, Layout LA, Layout LB>
T tag_invoke(const detail::PointwiseSqFn & /*cpo*/, SqEuclideanTag, const NDArray<T, 1, LA> &a,
             const NDArray<T, 1, LB> &b) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "pointwiseSq: integer element types not supported");
  const std::size_t n = a.dim(0);
#ifdef CLUSTERING_USE_AVX2
  if constexpr (std::is_same_v<T, float>) {
    // _mm256_load_ps issues an 8-lane contiguous aligned read. Gate dispatch on contiguity
    // (a strided view shares element type but not memory layout), 32-byte alignment (misaligned
    // load is UB), and the 8-dim floor (below 8 the horizontal-sum epilogue is pure tax; matches
    // the kdtree dispatch rationale).
    if (n >= 8 && a.isContiguous() && b.isContiguous() && a.template isAligned<32>() &&
        b.template isAligned<32>()) {
      return detail::sqEuclideanAvx2F32(a.template alignedData<32>(), b.template alignedData<32>(),
                                        n);
    }
  }
#endif
  T sum = T{0};
  for (std::size_t i = 0; i < n; ++i) {
    const T diff = a(i) - b(i);
    sum += diff * diff;
  }
  return sum;
}

template <class T, Layout LA, Layout LB>
T tag_invoke(const detail::PointwiseSqFn & /*cpo*/, ManhattanTag, const NDArray<T, 1, LA> &a,
             const NDArray<T, 1, LB> &b) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "pointwiseSq: integer element types not supported");
  const std::size_t n = a.dim(0);
  T sum = T{0};
  for (std::size_t i = 0; i < n; ++i) {
    const T diff = a(i) - b(i);
    sum += diff < T{0} ? -diff : diff;
  }
  return sum;
}

template <class T, Layout LA, Layout LB>
T tag_invoke(const detail::PointwiseSqFn & /*cpo*/, CosineTag, const NDArray<T, 1, LA> &a,
             const NDArray<T, 1, LB> &b) noexcept {
  static_assert(std::is_same_v<T, float> || std::is_same_v<T, double>,
                "pointwiseSq: integer element types not supported");
  const std::size_t n = a.dim(0);
  T dot = T{0};
  T na = T{0};
  T nb = T{0};
  for (std::size_t i = 0; i < n; ++i) {
    const T ai = a(i);
    const T bi = b(i);
    dot += ai * bi;
    na += ai * ai;
    nb += bi * bi;
  }
  if (na == T{0} || nb == T{0}) {
    return T{1};
  }
  return T{1} - (dot / std::sqrt(na * nb));
}
 
} // namespace clustering::math::distance