usearch 2.21.3 → 2.25.0

package/numkong/test/test_maxsim.py

@@ -0,0 +1,209 @@
+#!/usr/bin/env python3
+"""Test MaxSim late-interaction scoring: nk.maxsim, nk.maxsim_pack, nk.maxsim_packed.
+
+Dtypes: float32, bfloat16, float16.
+Baselines: high-precision Decimal angular distances, NumPy reference.
+Matches C++ suite: test_maxsim.cpp.
+"""
+
+import atexit
+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 (
+    assert_allclose,
+    collect_warnings,
+    create_stats,
+    downcast_f32_to_dtype,
+    keep_one_capability,
+    make_nk,
+    nk_seed,  # noqa: F401 — pytest fixture
+    numpy_available,
+    possible_capabilities,
+    precise_decimal,
+    print_stats_report,
+    reduced_repetitions_count,
+    seed_rng,  # noqa: F401 — pytest fixture (autouse)
+    test_depth_dimensions,
+    test_height_dimensions,
+    test_width_dimensions,
+)
+
+stats = create_stats()
+atexit.register(print_stats_report, stats)
+
+_MAXSIM_DTYPE = {"float32": "f32", "bfloat16": "bf16", "float16": "f16"}
+
+
+def baseline_maxsim(queries, documents, dtype=None):
+    """Pure NumPy reference: sum of per-query minimum angular distances."""
+    total = 0.0
+    for q in queries:
+        q64 = q.astype(np.float64)
+        norm_q = np.linalg.norm(q64)
+        min_ang = float("inf")
+        for d in documents:
+            d64 = d.astype(np.float64)
+            norm_d = np.linalg.norm(d64)
+            dot = np.dot(q64, d64)
+            angular = max(0.0, 1.0 - dot / (norm_q * norm_d)) if norm_q > 0 and norm_d > 0 else 1.0
+            min_ang = min(min_ang, angular)
+        total += min_ang
+    return total
+
+
+def precise_maxsim(queries, documents, dtype=None):
+    """High-precision MaxSim via Python Decimal: sum of per-query min angular distances."""
+    with precise_decimal(dtype) as (upcast, sqrt, _ln):
+        total = upcast(0)
+        for query_vector in queries:
+            query_values = [upcast(value) for value in query_vector]
+            norm_query = sqrt(sum(value * value for value in query_values))
+            min_angular = None
+            for document_vector in documents:
+                dot_product = upcast(0)
+                norm_squared = upcast(0)
+                for query_value, document_value in zip(query_values, document_vector):
+                    document_precise = upcast(document_value)
+                    dot_product += query_value * document_precise
+                    norm_squared += document_precise * document_precise
+                norm_document = sqrt(norm_squared)
+                if norm_query > 0 and norm_document > 0:
+                    cosine_similarity = dot_product / (norm_query * norm_document)
+                    angular = max(upcast(0), upcast(1) - cosine_similarity)
+                else:
+                    angular = upcast(1)
+                if min_angular is None or angular < min_angular:
+                    min_angular = angular
+            total += min_angular
+        return float(total)
+
+
+KERNELS_MAXSIM: dict[str, tuple[Callable, Callable, Callable]] = {
+    "maxsim": (baseline_maxsim, nk.maxsim, precise_maxsim),
+}
+
+
+def _make_matrix(rows, cols, dtype):
+    """Create a test matrix in the target dtype."""
+    raw, _ = downcast_f32_to_dtype(np.random.randn(rows, cols).astype(np.float32), dtype)
+    return make_nk(raw, dtype)
+
+
+@pytest.mark.skipif(not numpy_available, reason="NumPy is not installed")
+@pytest.mark.repeat(reduced_repetitions_count)
+@pytest.mark.parametrize("rows", test_height_dimensions)
+@pytest.mark.parametrize("columns", test_width_dimensions)
+@pytest.mark.parametrize("depth", test_depth_dimensions)
+@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"])
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_maxsim_pack_and_packed(rows: int, columns: int, depth: int, dtype: str, capability: str):
+    """Pack + compute vs baseline."""
+    keep_one_capability(capability)
+    baseline_kernel, _, precise_kernel = KERNELS_MAXSIM["maxsim"]
+    dtype_str = _MAXSIM_DTYPE[dtype]
+    queries = _make_matrix(rows, depth, dtype)
+    documents = _make_matrix(columns, depth, dtype)
+
+    qp = nk.maxsim_pack(queries, dtype=dtype_str)
+    dp = nk.maxsim_pack(documents, dtype=dtype_str)
+
+    assert isinstance(qp, nk.MaxSimPackedMatrix)
+    assert qp.vector_count == rows
+    assert qp.depth == depth
+    assert qp.nbytes > 0
+    assert repr(qp)  # smoke-test repr doesn't crash
+
+    result = nk.maxsim_packed(qp, dp)
+
+    # MaxSim is an approximation (coarse screening + fine pass); compound
+    # angular-distance accumulation amplifies f32 error to ~2% relative.
+    # Verify sanity bounds only and log deviations for the stats report.
+    assert np.isfinite(result), f"maxsim_packed({dtype}): got non-finite {result}"
+    assert result >= 0, f"maxsim_packed({dtype}): got negative {result}"
+    if dtype != "bfloat16":
+        q_np = np.array(queries, dtype=dtype)
+        d_np = np.array(documents, dtype=dtype)
+        expected = baseline_kernel(q_np, d_np, dtype=dtype_str)
+        rel_err = abs(result - expected) / max(abs(expected), 1e-12)
+        if rel_err > 0.1:
+            collect_warnings(f"maxsim_packed({dtype}) rel_err={rel_err:.4f}", stats)
+
+
+@pytest.mark.skipif(not numpy_available, reason="NumPy is not installed")
+@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"])
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_maxsim_convenience(dtype: str, capability: str):
+    """Verify maxsim() matches maxsim_pack + maxsim_packed."""
+    keep_one_capability(capability)
+    n_q, n_d, depth = 4, 8, 32
+    dtype_str = _MAXSIM_DTYPE[dtype]
+    queries = _make_matrix(n_q, depth, dtype)
+    documents = _make_matrix(n_d, depth, dtype)
+
+    # Packed path
+    qp = nk.maxsim_pack(queries, dtype=dtype_str)
+    dp = nk.maxsim_pack(documents, dtype=dtype_str)
+    packed_result = nk.maxsim_packed(qp, dp)
+
+    # Convenience path
+    _, simd_kernel, _ = KERNELS_MAXSIM["maxsim"]
+    conv_result = simd_kernel(queries, documents, dtype=dtype_str)
+
+    assert_allclose(packed_result, conv_result)
+
+
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_maxsim_self_zero(capability: str):
+    """Identical queries and documents should give score near zero."""
+    keep_one_capability(capability)
+    vectors = nk.ones((8, 64), dtype="float32")
+    result = nk.maxsim(vectors, vectors, dtype="f32")
+    assert_allclose(result, 0.0, err_msg="Expected near-zero self-distance")
+
+
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_maxsim_type_errors(capability: str, nk_seed: int):
+    """Wrong type and mismatched dtype/depth."""
+    keep_one_capability(capability)
+    query = nk.hash((4, 32), seed=nk_seed, dtype="float32")
+    documents = nk.hash((8, 32), seed=nk_seed + 1, dtype="float32")
+    query_packed = nk.maxsim_pack(query, dtype="f32")
+    documents_packed = nk.maxsim_pack(documents, dtype="f32")
+
+    # Wrong type for maxsim_packed
+    with pytest.raises(TypeError):
+        nk.maxsim_packed(query, documents_packed)
+    with pytest.raises(TypeError):
+        nk.maxsim_packed(query_packed, documents)
+
+    # Mismatched depth
+    documents_wrong = nk.hash((8, 16), seed=nk_seed + 2, dtype="float32")
+    documents_wrong_packed = nk.maxsim_pack(documents_wrong, dtype="f32")
+    with pytest.raises(ValueError):
+        nk.maxsim_packed(query_packed, documents_wrong_packed)
+
+
+@pytest.mark.parametrize("capability", possible_capabilities)
+def test_maxsim_packed_size(capability: str):
+    """Verify classmethod packed_size returns a positive integer."""
+    keep_one_capability(capability)
+    size = nk.MaxSimPackedMatrix.packed_size(8, 64, dtype="bf16")
+    assert isinstance(size, int)
+    assert size > 0
+
+    size_f32 = nk.MaxSimPackedMatrix.packed_size(8, 64, dtype="f32")
+    assert size_f32 > 0

package/numkong/test/test_mesh.cpp

@@ -0,0 +1,203 @@
+/**
+ *  @brief RMSD, Kabsch, and Umeyama alignment tests.
+ *  @file test/test_mesh.cpp
+ *  @author Ash Vardanian
+ *  @date December 28, 2025
+ */
+
+#include "test.hpp"
+#include "numkong/mesh.hpp" // `nk::rmsd`
+
+/**
+ *  @brief Test RMSD kernel.
+ */
+template <typename scalar_type_>
+error_stats_t test_rmsd(typename scalar_type_::mesh_kernel_t kernel) {
+    using scalar_t = scalar_type_;
+    using transform_t = typename scalar_t::mesh_transform_t;
+    using metric_t = typename scalar_t::mesh_metric_t;
+    using reference_t = reference_for<scalar_t>;
+
+    error_stats_t stats(comparison_family_t::mixed_precision_reduction_k);
+    std::mt19937 generator(global_config.seed);
+
+    std::size_t n = global_config.mesh_points;
+    auto a = make_vector<scalar_t>(n * 3), b = make_vector<scalar_t>(n * 3);
+
+    for (auto start = test_start_time(); within_time_budget(start);) {
+        fill_random(generator, a);
+        fill_random(generator, b);
+
+        transform_t a_centroid[3], b_centroid[3], rot[9], scale;
+        metric_t result;
+        kernel(a.raw_values_data(), b.raw_values_data(), n, &a_centroid[0].raw_, &b_centroid[0].raw_, &rot[0].raw_,
+               &scale.raw_, &result.raw_);
+        reference_t a_centroid_ref[3], b_centroid_ref[3], rot_ref[9], scale_ref, reference;
+        nk::rmsd<scalar_t, reference_t, reference_t, nk::no_simd_k>(a.values_data(), b.values_data(), n, a_centroid_ref,
+                                                                    b_centroid_ref, rot_ref, &scale_ref, &reference);
+
+        stats.accumulate(result, reference);
+    }
+    return stats;
+}
+
+/**
+ *  @brief Test Kabsch alignment kernel.
+ */
+template <typename scalar_type_>
+error_stats_t test_kabsch(typename scalar_type_::mesh_kernel_t kernel) {
+    using scalar_t = scalar_type_;
+    using transform_t = typename scalar_t::mesh_transform_t;
+    using metric_t = typename scalar_t::mesh_metric_t;
+    using reference_t = reference_for<scalar_t>;
+
+    error_stats_t stats(comparison_family_t::mixed_precision_reduction_k);
+    std::mt19937 generator(global_config.seed);
+
+    std::size_t n = global_config.mesh_points;
+    auto a = make_vector<scalar_t>(n * 3), b = make_vector<scalar_t>(n * 3);
+
+    for (auto start = test_start_time(); within_time_budget(start);) {
+        fill_random(generator, a);
+        fill_random(generator, b);
+
+        transform_t a_centroid[3], b_centroid[3], rot[9], scale;
+        metric_t result;
+        kernel(a.raw_values_data(), b.raw_values_data(), n, &a_centroid[0].raw_, &b_centroid[0].raw_, &rot[0].raw_,
+               &scale.raw_, &result.raw_);
+        reference_t a_centroid_ref[3], b_centroid_ref[3], rot_ref[9], scale_ref, reference;
+        nk::kabsch<scalar_t, reference_t, reference_t, nk::no_simd_k>(
+            a.values_data(), b.values_data(), n, a_centroid_ref, b_centroid_ref, rot_ref, &scale_ref, &reference);
+
+        stats.accumulate(result, reference);
+    }
+    return stats;
+}
+
+/**
+ *  @brief Test Umeyama alignment kernel.
+ */
+template <typename scalar_type_>
+error_stats_t test_umeyama(typename scalar_type_::mesh_kernel_t kernel) {
+    using scalar_t = scalar_type_;
+    using transform_t = typename scalar_t::mesh_transform_t;
+    using metric_t = typename scalar_t::mesh_metric_t;
+    using reference_t = reference_for<scalar_t>;
+
+    error_stats_t stats(comparison_family_t::mixed_precision_reduction_k);
+    std::mt19937 generator(global_config.seed);
+
+    std::size_t n = global_config.mesh_points;
+    auto a = make_vector<scalar_t>(n * 3), b = make_vector<scalar_t>(n * 3);
+
+    for (auto start = test_start_time(); within_time_budget(start);) {
+        fill_random(generator, a);
+        fill_random(generator, b);
+
+        transform_t a_centroid[3], b_centroid[3], rot[9], scale;
+        metric_t result;
+        kernel(a.raw_values_data(), b.raw_values_data(), n, &a_centroid[0].raw_, &b_centroid[0].raw_, &rot[0].raw_,
+               &scale.raw_, &result.raw_);
+        reference_t a_centroid_ref[3], b_centroid_ref[3], rot_ref[9], scale_ref, reference;
+        nk::umeyama<scalar_t, reference_t, reference_t, nk::no_simd_k>(
+            a.values_data(), b.values_data(), n, a_centroid_ref, b_centroid_ref, rot_ref, &scale_ref, &reference);
+
+        stats.accumulate(result, reference);
+    }
+    return stats;
+}
+
+void test_mesh() {
+    error_stats_section_t check("Mesh Operations");
+
+#if NK_DYNAMIC_DISPATCH
+    check("rmsd_f64", test_rmsd<f64_t>, nk_rmsd_f64);
+    check("rmsd_f32", test_rmsd<f32_t>, nk_rmsd_f32);
+    check("kabsch_f64", test_kabsch<f64_t>, nk_kabsch_f64);
+    check("kabsch_f32", test_kabsch<f32_t>, nk_kabsch_f32);
+    check("umeyama_f64", test_umeyama<f64_t>, nk_umeyama_f64);
+    check("umeyama_f32", test_umeyama<f32_t>, nk_umeyama_f32);
+#else
+
+#if NK_TARGET_NEON
+    check("rmsd_f64_neon", test_rmsd<f64_t>, nk_rmsd_f64_neon);
+    check("rmsd_f32_neon", test_rmsd<f32_t>, nk_rmsd_f32_neon);
+    check("kabsch_f64_neon", test_kabsch<f64_t>, nk_kabsch_f64_neon);
+    check("kabsch_f32_neon", test_kabsch<f32_t>, nk_kabsch_f32_neon);
+    check("umeyama_f64_neon", test_umeyama<f64_t>, nk_umeyama_f64_neon);
+    check("umeyama_f32_neon", test_umeyama<f32_t>, nk_umeyama_f32_neon);
+    check("rmsd_f16_neon", test_rmsd<f16_t>, nk_rmsd_f16_neon);
+    check("kabsch_f16_neon", test_kabsch<f16_t>, nk_kabsch_f16_neon);
+    check("umeyama_f16_neon", test_umeyama<f16_t>, nk_umeyama_f16_neon);
+#endif // NK_TARGET_NEON
+
+#if NK_TARGET_NEONBFDOT
+    check("rmsd_bf16_neonbfdot", test_rmsd<bf16_t>, nk_rmsd_bf16_neonbfdot);
+    check("kabsch_bf16_neonbfdot", test_kabsch<bf16_t>, nk_kabsch_bf16_neonbfdot);
+    check("umeyama_bf16_neonbfdot", test_umeyama<bf16_t>, nk_umeyama_bf16_neonbfdot);
+#endif // NK_TARGET_NEONBFDOT
+
+#if NK_TARGET_HASWELL
+    check("rmsd_f64_haswell", test_rmsd<f64_t>, nk_rmsd_f64_haswell);
+    check("rmsd_f32_haswell", test_rmsd<f32_t>, nk_rmsd_f32_haswell);
+    check("kabsch_f64_haswell", test_kabsch<f64_t>, nk_kabsch_f64_haswell);
+    check("kabsch_f32_haswell", test_kabsch<f32_t>, nk_kabsch_f32_haswell);
+    check("umeyama_f64_haswell", test_umeyama<f64_t>, nk_umeyama_f64_haswell);
+    check("umeyama_f32_haswell", test_umeyama<f32_t>, nk_umeyama_f32_haswell);
+    check("rmsd_f16_haswell", test_rmsd<f16_t>, nk_rmsd_f16_haswell);
+    check("kabsch_f16_haswell", test_kabsch<f16_t>, nk_kabsch_f16_haswell);
+    check("umeyama_f16_haswell", test_umeyama<f16_t>, nk_umeyama_f16_haswell);
+    check("rmsd_bf16_haswell", test_rmsd<bf16_t>, nk_rmsd_bf16_haswell);
+    check("kabsch_bf16_haswell", test_kabsch<bf16_t>, nk_kabsch_bf16_haswell);
+    check("umeyama_bf16_haswell", test_umeyama<bf16_t>, nk_umeyama_bf16_haswell);
+#endif // NK_TARGET_HASWELL
+
+#if NK_TARGET_SKYLAKE
+    check("rmsd_f64_skylake", test_rmsd<f64_t>, nk_rmsd_f64_skylake);
+    check("rmsd_f32_skylake", test_rmsd<f32_t>, nk_rmsd_f32_skylake);
+    check("kabsch_f64_skylake", test_kabsch<f64_t>, nk_kabsch_f64_skylake);
+    check("kabsch_f32_skylake", test_kabsch<f32_t>, nk_kabsch_f32_skylake);
+    check("umeyama_f64_skylake", test_umeyama<f64_t>, nk_umeyama_f64_skylake);
+    check("umeyama_f32_skylake", test_umeyama<f32_t>, nk_umeyama_f32_skylake);
+    check("rmsd_f16_skylake", test_rmsd<f16_t>, nk_rmsd_f16_skylake);
+    check("rmsd_bf16_skylake", test_rmsd<bf16_t>, nk_rmsd_bf16_skylake);
+    check("kabsch_f16_skylake", test_kabsch<f16_t>, nk_kabsch_f16_skylake);
+    check("kabsch_bf16_skylake", test_kabsch<bf16_t>, nk_kabsch_bf16_skylake);
+    check("umeyama_f16_skylake", test_umeyama<f16_t>, nk_umeyama_f16_skylake);
+    check("umeyama_bf16_skylake", test_umeyama<bf16_t>, nk_umeyama_bf16_skylake);
+#endif // NK_TARGET_SKYLAKE
+
+#if NK_TARGET_RVV
+    check("rmsd_f64_rvv", test_rmsd<f64_t>, nk_rmsd_f64_rvv);
+    check("rmsd_f32_rvv", test_rmsd<f32_t>, nk_rmsd_f32_rvv);
+    check("rmsd_f16_rvv", test_rmsd<f16_t>, nk_rmsd_f16_rvv);
+    check("rmsd_bf16_rvv", test_rmsd<bf16_t>, nk_rmsd_bf16_rvv);
+    check("kabsch_f64_rvv", test_kabsch<f64_t>, nk_kabsch_f64_rvv);
+    check("kabsch_f32_rvv", test_kabsch<f32_t>, nk_kabsch_f32_rvv);
+    check("kabsch_f16_rvv", test_kabsch<f16_t>, nk_kabsch_f16_rvv);
+    check("kabsch_bf16_rvv", test_kabsch<bf16_t>, nk_kabsch_bf16_rvv);
+    check("umeyama_f64_rvv", test_umeyama<f64_t>, nk_umeyama_f64_rvv);
+    check("umeyama_f32_rvv", test_umeyama<f32_t>, nk_umeyama_f32_rvv);
+    check("umeyama_f16_rvv", test_umeyama<f16_t>, nk_umeyama_f16_rvv);
+    check("umeyama_bf16_rvv", test_umeyama<bf16_t>, nk_umeyama_bf16_rvv);
+#endif // NK_TARGET_RVV
+
+#if NK_TARGET_V128RELAXED
+    check("rmsd_f32_v128relaxed", test_rmsd<f32_t>, nk_rmsd_f32_v128relaxed);
+    check("rmsd_f64_v128relaxed", test_rmsd<f64_t>, nk_rmsd_f64_v128relaxed);
+    check("kabsch_f32_v128relaxed", test_kabsch<f32_t>, nk_kabsch_f32_v128relaxed);
+    check("kabsch_f64_v128relaxed", test_kabsch<f64_t>, nk_kabsch_f64_v128relaxed);
+    check("umeyama_f32_v128relaxed", test_umeyama<f32_t>, nk_umeyama_f32_v128relaxed);
+    check("umeyama_f64_v128relaxed", test_umeyama<f64_t>, nk_umeyama_f64_v128relaxed);
+#endif // NK_TARGET_V128RELAXED
+
+    // Serial always runs - baseline test
+    check("rmsd_f64_serial", test_rmsd<f64_t>, nk_rmsd_f64_serial);
+    check("rmsd_f32_serial", test_rmsd<f32_t>, nk_rmsd_f32_serial);
+    check("kabsch_f64_serial", test_kabsch<f64_t>, nk_kabsch_f64_serial);
+    check("kabsch_f32_serial", test_kabsch<f32_t>, nk_kabsch_f32_serial);
+    check("umeyama_f64_serial", test_umeyama<f64_t>, nk_umeyama_f64_serial);
+    check("umeyama_f32_serial", test_umeyama<f32_t>, nk_umeyama_f32_serial);
+
+#endif // NK_DYNAMIC_DISPATCH
+}