usearch 2.21.4 → 2.25.0

package/numkong/test/test_spatial.py

@@ -0,0 +1,269 @@
+#!/usr/bin/env python3
+"""Test spatial distances: nk.euclidean, nk.sqeuclidean, nk.angular.
+
+Dtypes: float64, float32, float16, bfloat16, e4m3, e5m2, e2m3, e3m2, int8, uint8.
+Baselines: high-precision Decimal accumulation, SciPy spatial.distance.
+Matches C++ suite: test_spatial.cpp.
+"""
+
+import atexit
+import math
+import warnings
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import pytest
+
+if TYPE_CHECKING:
+    import numpy as np  # static-analysis-only; the runtime try/except below is authoritative
+
+try:
+    import numpy as np
+
+    numpy_available = True
+except Exception:
+    numpy_available = False
+
+import numkong as nk
+from test_base import (
+    NATIVE_COMPUTE_DTYPE,
+    NK_ATOL,
+    NK_RTOL,
+    LazyFormat,
+    assert_allclose,
+    collect_errors,
+    create_stats,
+    dense_dimensions,
+    keep_one_capability,
+    make_random,
+    make_random_buffer,
+    nk_seed,  # noqa: F401 — pytest fixture
+    numpy_available,
+    possible_capabilities,
+    precise_decimal,
+    print_stats_report,
+    profile,
+    randomized_repetitions_count,
+    seed_rng,  # noqa: F401 — pytest fixture (autouse)
+    tolerances_for_dtype,
+)
+
+algebraic_dtypes = ["float32", "float64"]
+algebraic_ndims = [7, 97]
+stats = create_stats()
+atexit.register(print_stats_report, stats)
+
+try:
+    import scipy.spatial.distance as spd
+
+    def baseline_euclidean(x, y, dtype=None):
+        return np.array(spd.euclidean(x, y))
+
+    def baseline_sqeuclidean(x, y, dtype=None):
+        return spd.sqeuclidean(x, y)
+
+    def baseline_angular(x, y, dtype=None):
+        return spd.cosine(x, y)
+
+except ImportError:
+
+    def baseline_angular(x, y, dtype=None):
+        return 1.0 - np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
+
+    def baseline_euclidean(x, y, dtype=None):
+        return np.array([np.sqrt(np.sum((x - y) ** 2))])
+
+    def baseline_sqeuclidean(x, y, dtype=None):
+        return np.sum((x - y) ** 2)
+
+
+def precise_sqeuclidean(a, b, dtype=None):
+    """High-precision squared Euclidean distance via Python Decimal."""
+    with precise_decimal(dtype) as (upcast, _sqrt, _ln):
+        return float(sum((upcast(x) - upcast(y)) ** 2 for x, y in zip(a, b)))
+
+
+def precise_euclidean(a, b, dtype=None):
+    return math.sqrt(precise_sqeuclidean(a, b, dtype=dtype))
+
+
+def precise_angular(a, b, dtype=None):
+    """High-precision angular/cosine distance via Python Decimal."""
+    with precise_decimal(dtype) as (upcast, sqrt, _ln):
+        dot_product = upcast(0)
+        norm_left_squared = upcast(0)
+        norm_right_squared = upcast(0)
+        for x, y in zip(a, b):
+            left_value = upcast(x)
+            right_value = upcast(y)
+            dot_product += left_value * right_value
+            norm_left_squared += left_value * left_value
+            norm_right_squared += right_value * right_value
+        denominator = sqrt(norm_left_squared * norm_right_squared)
+        if norm_left_squared == 0 and norm_right_squared == 0:
+            return 0.0
+        if denominator == 0:
+            return 1.0
+        return float(1 - dot_product / denominator)
+
+
+KERNELS_SPATIAL: dict[str, tuple[Callable | None, Callable, Callable]] = {
+    "euclidean": (baseline_euclidean if numpy_available else None, nk.euclidean, precise_euclidean),
+    "sqeuclidean": (baseline_sqeuclidean if numpy_available else None, nk.sqeuclidean, precise_sqeuclidean),
+    "angular": (baseline_angular if numpy_available else None, nk.angular, precise_angular),
+}
+
+
+@pytest.mark.repeat(randomized_repetitions_count)
+@pytest.mark.parametrize("ndim", dense_dimensions)
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        "float64",
+        "float32",
+        "float16",
+        "bfloat16",
+        "e4m3",
+        "e5m2",
+        "e2m3",
+        "e3m2",
+        "int8",
+        "uint8",
+    ],
+)
+@pytest.mark.parametrize("metric", ["euclidean", "sqeuclidean", "angular"])
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_spatial_random_accuracy(ndim: int, dtype: str, metric: str, capability: str, nk_seed: int):
+    """Spatial distances across all numeric dtypes against high-precision Decimal baselines."""
+    a_raw, a_baseline = make_random((ndim,), dtype, seed=nk_seed)
+    b_raw, b_baseline = make_random((ndim,), dtype, seed=nk_seed + 1)
+    atol, rtol = tolerances_for_dtype(dtype)
+
+    keep_one_capability(capability)
+    baseline_kernel, simd_kernel, precise_kernel = KERNELS_SPATIAL[metric]
+
+    # High-precision baseline
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", RuntimeWarning)
+        accurate_dt, accurate = profile(precise_kernel or baseline_kernel, a_baseline, b_baseline, dtype=dtype)
+
+    # Baseline at native precision (for error stats)
+    if baseline_kernel is not None:
+        native_dt = NATIVE_COMPUTE_DTYPE.get(dtype, np.float64)
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", RuntimeWarning)
+            expected_dt, expected = profile(baseline_kernel, a_baseline.astype(native_dt), b_baseline.astype(native_dt))
+    else:
+        expected_dt, expected = 0, None
+
+    # SIMD result
+    result_dt, result = profile(simd_kernel, a_raw, b_raw, dtype)
+
+    err_msg = LazyFormat(
+        lambda: (f"\n{metric}({dtype}, ndim={ndim}):" f"\n  Accurate:  {accurate}" f"\n  Got:       {result}")
+    )
+
+    assert_allclose(result, accurate, atol=atol, rtol=rtol, err_msg=err_msg)
+    collect_errors(metric, ndim, dtype, accurate, accurate_dt, expected, expected_dt, result, result_dt, stats)
+
+
+@pytest.mark.skipif(not numpy_available, reason="NumPy is not installed")
+@pytest.mark.parametrize("ndim", dense_dimensions)
+@pytest.mark.parametrize("dtype", ["float32", "float16"])
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_angular_zero_vector(ndim: int, dtype: str, capability: str):
+    """Tests the nk.angular() function with zero vectors, to catch division by zero errors."""
+    a = np.zeros(ndim, dtype=dtype)
+    b = (np.random.randn(ndim) + 1).astype(dtype)
+    keep_one_capability(capability)
+
+    result = nk.angular(a, b)
+    assert result == 1, f"Expected 1, but got {result}"
+
+    result = nk.angular(a, a)
+    assert result == 0, f"Expected 0 distance from itself, but got {result}"
+
+    result = nk.angular(b, b)
+    assert abs(result) < NK_ATOL, f"Expected 0 distance from itself, but got {result}"
+
+    assert np.all(result >= 0), "Negative result for angular distance"
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_spatial_self_distance_zero(ndim: int, dtype: str, capability: str):
+    """d(v, v) should be 0 for euclidean, sqeuclidean, and angular."""
+    keep_one_capability(capability)
+    v = nk.full((ndim,), 1.5, dtype=dtype)
+    atol = NK_ATOL
+    assert abs(nk.euclidean(v, v)) < atol
+    assert abs(nk.sqeuclidean(v, v)) < atol
+    assert abs(nk.angular(v, v)) < atol
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_euclidean_known(ndim: int, dtype: str, capability: str):
+    """euclidean(ones, zeros) = sqrt(n)."""
+    keep_one_capability(capability)
+    ones_vector = nk.ones((ndim,), dtype=dtype)
+    zeros_vector = nk.zeros((ndim,), dtype=dtype)
+    result = nk.euclidean(ones_vector, zeros_vector)
+    expected = math.sqrt(ndim)
+    assert abs(result - expected) < NK_ATOL + NK_RTOL * expected
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_sqeuclidean_known(ndim: int, dtype: str, capability: str):
+    """sqeuclidean(ones, zeros) = n."""
+    keep_one_capability(capability)
+    ones_vector = nk.ones((ndim,), dtype=dtype)
+    zeros_vector = nk.zeros((ndim,), dtype=dtype)
+    result = nk.sqeuclidean(ones_vector, zeros_vector)
+    assert abs(result - ndim) < NK_ATOL + NK_RTOL * ndim
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_spatial_symmetry(ndim: int, dtype: str, capability: str):
+    """Commutativity: d(a, b) = d(b, a) for all spatial metrics."""
+    keep_one_capability(capability)
+    a = make_random_buffer(ndim, dtype)
+    b = make_random_buffer(ndim, dtype)
+    for metric_fn in [nk.euclidean, nk.sqeuclidean, nk.angular]:
+        d_ab = metric_fn(a, b)
+        d_ba = metric_fn(b, a)
+        assert abs(d_ab - d_ba) < NK_ATOL, f"{metric_fn.__name__}: d(a,b)={d_ab} != d(b,a)={d_ba}"
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_spatial_non_negative(ndim: int, dtype: str, capability: str):
+    """Non-negativity: d(a, b) >= 0 for all spatial metrics."""
+    keep_one_capability(capability)
+    a = make_random_buffer(ndim, dtype)
+    b = make_random_buffer(ndim, dtype)
+    for metric_fn in [nk.euclidean, nk.sqeuclidean, nk.angular]:
+        distance = metric_fn(a, b)
+        assert distance >= -NK_ATOL, f"{metric_fn.__name__}: d(a,b)={distance} is negative"
+
+
+@pytest.mark.parametrize("ndim", algebraic_ndims)
+@pytest.mark.parametrize("dtype", algebraic_dtypes)
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_euclidean_triangle_inequality(ndim: int, dtype: str, capability: str):
+    """Triangle inequality: euclidean(a, c) <= euclidean(a, b) + euclidean(b, c)."""
+    keep_one_capability(capability)
+    a = make_random_buffer(ndim, dtype)
+    b = make_random_buffer(ndim, dtype)
+    c = make_random_buffer(ndim, dtype)
+    d_ac = nk.euclidean(a, c)
+    d_ab = nk.euclidean(a, b)
+    d_bc = nk.euclidean(b, c)
+    assert d_ac <= d_ab + d_bc + NK_ATOL, f"Triangle inequality violated: d(a,c)={d_ac} > d(a,b)+d(b,c)={d_ab + d_bc}"