usearch 0.13.0 → 0.15.0

package/python/scripts/test.py

@@ -3,29 +3,41 @@ import pytest
 import numpy as np
 
 from usearch.io import load_matrix, save_matrix
-from usearch.index import Index, SetsIndex, MetricKind
+from usearch.eval import recall_members, random_vectors
 
+from usearch.index import Index, SparseIndex, MetricKind, ScalarKind, Matches
+from usearch.index import (
+    DEFAULT_CONNECTIVITY,
+    DEFAULT_EXPANSION_ADD,
+    DEFAULT_EXPANSION_SEARCH,
+)
 
-dimensions = [13, 99, 100, 256]
-batch_sizes = [1, 2, 100]
-index_types = ['f32', 'f64', 'f16', 'f8']
-numpy_types = [np.float32, np.float64, np.float16, np.byte]
-connectivity_options = [3, 13, 50]
+
+dimensions = [3, 97, 256]
+batch_sizes = [1, 33]
+index_types = [
+    ScalarKind.F32, ScalarKind.F64,
+    ScalarKind.F16, ScalarKind.F8,
+]
+numpy_types = [np.float32, np.float64, np.float16]
+
+connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
 jit_options = [False]
 continuous_metrics = [
     MetricKind.Cos,
     MetricKind.L2sq,
 ]
 hash_metrics = [
-    MetricKind.BitwiseHamming,
-    MetricKind.BitwiseTanimoto,
-    MetricKind.BitwiseSorensen,
+    MetricKind.Hamming,
+    MetricKind.Tanimoto,
+    MetricKind.Sorensen,
 ]
 
 
 @pytest.mark.parametrize('rows', batch_sizes)
 @pytest.mark.parametrize('cols', dimensions)
 def test_serializing_fbin_matrix(rows: int, cols: int):
+
     original = np.random.rand(rows, cols).astype(np.float32)
     save_matrix(original, 'tmp.fbin')
     reconstructed = load_matrix('tmp.fbin')
@@ -36,6 +48,7 @@ def test_serializing_fbin_matrix(rows: int, cols: int):
 @pytest.mark.parametrize('rows', batch_sizes)
 @pytest.mark.parametrize('cols', dimensions)
 def test_serializing_ibin_matrix(rows: int, cols: int):
+
     original = np.random.randint(0, rows+1, size=(rows, cols)).astype(np.int32)
     save_matrix(original, 'tmp.ibin')
     reconstructed = load_matrix('tmp.ibin')
@@ -59,6 +72,8 @@ def test_index(
         ndim=ndim,
         dtype=index_type,
         connectivity=connectivity,
+        expansion_add=DEFAULT_EXPANSION_ADD,
+        expansion_search=DEFAULT_EXPANSION_SEARCH,
         jit=jit,
     )
     assert index.ndim == ndim
@@ -66,9 +81,15 @@ def test_index(
 
     limit = 0.7 if numpy_type != np.byte else 70
     vector = np.random.uniform(0, limit, (index.ndim)).astype(numpy_type)
+
     index.add(42, vector)
-    matches, distances, count = index.search(vector, 10)
 
+    assert 42 in index
+    assert 42 in index.labels
+    if numpy_type != np.byte:
+        assert np.allclose(index[42], vector, atol=0.1)
+
+    matches, distances, count = index.search(vector, 10)
     assert len(index) == 1
     assert len(matches) == count
     assert len(distances) == count
@@ -95,21 +116,25 @@ def test_index_batch(
 
     index = Index(ndim=ndim, metric=metric, dtype=index_type)
 
-    limit = 0.7 if numpy_type != np.byte else 70
-    vectors = np.random.uniform(
-        0, limit, (batch_size, index.ndim)).astype(numpy_type)
-    labels = np.array(range(batch_size), dtype=np.longlong)
+    labels = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim, dtype=numpy_type)
 
     index.add(labels, vectors)
-    matches, distances, count = index.search(vectors, 10)
+    assert len(index) == batch_size
+    assert np.allclose(index.get_vectors(
+        labels).astype(numpy_type), vectors, atol=0.1)
 
-    assert matches.shape[0] == distances.shape[0]
-    assert count.shape[0] == batch_size
+    matches: Matches = index.search(vectors, 10)
+    assert matches.labels.shape[0] == matches.distances.shape[0]
+    assert matches.counts.shape[0] == batch_size
+    assert np.all(np.sort(index.labels) == np.sort(labels))
+
+    assert recall_members(index, exact=True) == 1
 
 
 @pytest.mark.parametrize('ndim', dimensions)
 @pytest.mark.parametrize('batch_size', batch_sizes)
-def test_index_udf(ndim: int, batch_size: int):
+def test_index_numba(ndim: int, batch_size: int):
 
     try:
         from numba import cfunc, types, carray
@@ -118,37 +143,35 @@ def test_index_udf(ndim: int, batch_size: int):
 
     # Showcases how to use Numba to JIT-compile similarity measures for USearch.
     # https://numba.readthedocs.io/en/stable/reference/jit-compilation.html#c-callbacks
-    signature = types.float32(
-        types.CPointer(types.float32),
-        types.CPointer(types.float32),
-        types.uint64, types.uint64)
+    signature = types.float32(types.CPointer(
+        types.float32), types.CPointer(types.float32))
 
     @cfunc(signature)
-    def python_dot(a, b, n, m):
-        a_array = carray(a, n)
-        b_array = carray(b, n)
+    def python_dot(a, b):
+        a_array = carray(a, ndim)
+        b_array = carray(b, ndim)
         c = 0.0
-        for i in range(n):
+        for i in range(ndim):
             c += a_array[i] * b_array[i]
-        return c
+        return 1 - c
 
     index = Index(ndim=ndim, metric=python_dot)
 
-    vectors = np.random.uniform(
-        0, 0.3, (batch_size, index.ndim)).astype(np.float32)
-    labels = np.array(range(batch_size), dtype=np.longlong)
+    labels = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
 
     index.add(labels, vectors)
     matches, distances, count = index.search(vectors, 10)
-
     assert matches.shape[0] == distances.shape[0]
     assert count.shape[0] == batch_size
 
+    assert recall_members(index, exact=True) == 1
+
 
 @pytest.mark.parametrize('connectivity', connectivity_options)
 def test_sets_index(connectivity: int):
 
-    index = SetsIndex(connectivity=connectivity)
+    index = SparseIndex(connectivity=connectivity)
     index.add(10, np.array([10, 12, 15], dtype=np.uint32))
     index.add(11, np.array([11, 12, 15, 16], dtype=np.uint32))
     results = index.search(np.array([12, 15], dtype=np.uint32), 10)
@@ -159,13 +182,13 @@ def test_sets_index(connectivity: int):
 @pytest.mark.parametrize('metric', hash_metrics)
 @pytest.mark.parametrize('connectivity', connectivity_options)
 @pytest.mark.parametrize('batch_size', batch_sizes)
-def test_hash_index(bits: int, metric: MetricKind, connectivity: int, batch_size: int):
+def test_bitwise_index(bits: int, metric: MetricKind, connectivity: int, batch_size: int):
 
     index = Index(ndim=bits, metric=metric, connectivity=connectivity)
 
+    labels = np.arange(batch_size)
     bit_vectors = np.random.randint(2, size=(batch_size, bits))
     bit_vectors = np.packbits(bit_vectors, axis=1)
-    labels = np.array(range(batch_size), dtype=np.longlong)
 
     index.add(labels, bit_vectors)
     index.search(bit_vectors, 10)

package/python/usearch/eval.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+from time import time_ns
+from typing import Tuple, Any, Callable, Union, Optional, List
+from dataclasses import dataclass, asdict
+from collections import defaultdict
+
+import numpy as np
+
+from usearch.io import load_matrix
+from usearch.index import Index, Matches, MetricKind, MetricKindBitwise
+
+
+def random_vectors(
+    count: int,
+    metric: MetricKind = MetricKind.IP,
+    dtype: np.generic = np.float32,
+    ndim: Optional[int] = None,
+    index: Optional[Index] = None,
+) -> np.ndarray:
+    """Produces a collection of random vectors normalized for the provided `metric`
+    and matching wanted `dtype`, which can both be inferred from an existing `index`.
+    """
+
+    # Infer default parameters from the `index`, if passed
+    if index is not None:
+        if not isinstance(index, Index):
+            raise ValueError('Unsupported `index` type')
+
+        ndim = index.ndim
+        dtype = index.numpy_dtype
+        metric = index.metric
+
+    # Produce data
+    if metric in MetricKindBitwise:
+        bit_vectors = np.random.randint(2, size=(count, ndim))
+        bit_vectors = np.packbits(bit_vectors, axis=1)
+        return bit_vectors
+
+    else:
+
+        x = np.random.rand(count, ndim)
+        if dtype == np.int8:
+            x = (x * 100).astype(np.int8)
+        else:
+            x = x.astype(dtype)
+            if metric == MetricKind.IP:
+                return x / np.linalg.norm(x, axis=1, keepdims=True)
+        return x
+
+
+def recall_members(index: Index, *args, **kwargs) -> float:
+    """Simplest benchmark for a quality of search, which queries every
+    existing member of the index, to make sure approximate search finds
+    the point itself.
+
+    :param index: Non-empty pre-constructed index
+    :type index: Index
+    :return: Value from 0 to 1, for the share of found self-references
+    :rtype: float
+    """
+    if len(index) == 0:
+        return 0
+
+    matches: Matches = index.search(index.vectors, 1, *args, **kwargs)
+    return matches.recall_first(index.labels)
+
+
+def measure_seconds(f: Callable) -> Tuple[float, Any]:
+    """Simple function profiling decorator.
+
+    :param f: Function to be profiled
+    :type f: Callable
+    :return: Time elapsed in seconds and the result of the execution
+    :rtype: Tuple[float, Any]
+    """
+    a = time_ns()
+    result = f()
+    b = time_ns()
+    c = b - a
+    secs = c / (10 ** 9)
+    return secs, result
+
+
+@dataclass
+class Dataset:
+
+    labels: np.ndarray
+    vectors: np.ndarray
+    queries: np.ndarray
+    neighbors: np.ndarray
+
+    def crop_neighbors(self, k: int):
+        self.neighbors = self.neighbors[:, k]
+
+    @staticmethod
+    def build(
+        vectors: Optional[str] = None,
+        queries: Optional[str] = None,
+        neighbors: Optional[str] = None,
+        count: Optional[int] = None,
+        ndim: Optional[int] = None,
+        k: Optional[int] = None,
+    ):
+        """Either loads an existing dataset from disk, or generates one on the fly.
+
+        :param vectors: _description_, defaults to None
+        :type vectors: Optional[str], optional
+        :param queries: _description_, defaults to None
+        :type queries: Optional[str], optional
+        :param neighbors: _description_, defaults to None
+        :type neighbors: Optional[str], optional
+        :param count: _description_, defaults to None
+        :type count: Optional[int], optional
+        :param ndim: _description_, defaults to None
+        :type ndim: Optional[int], optional
+        :param k: _description_, defaults to None
+        :type k: Optional[int], optional
+        """
+
+        d = Dataset(None, None, None, None)
+
+        if vectors is not None:
+            assert ndim is None
+
+            d.vectors = load_matrix(vectors)
+            ndim = d.vectors.shape[1]
+            count = min(
+                d.vectors.shape[0], count) if count is not None else d.vectors.shape[0]
+            d.vectors = d.vectors[:count, :]
+            d.labels = np.arange(count, dtype=np.longlong)
+
+            if queries is not None:
+                d.queries = load_matrix(queries)
+            else:
+                d.queries = d.vectors
+
+            if neighbors is not None:
+                d.neighbors = load_matrix(neighbors)
+                if k is not None:
+                    d.neighbors = d.neighbors[:, :k]
+            else:
+                assert k is None, 'Cant ovveride `k`, will retrieve one neighbor'
+                d.neighbors = np.reshape(d.labels, (count, 1))
+
+        else:
+
+            assert ndim is not None
+            assert count is not None
+            assert k is None, 'Cant ovveride `k`, will retrieve one neighbor'
+
+            d.vectors = random_vectors(count=count, ndim=ndim)
+            d.queries = d.vectors
+            d.labels = np.arange(count, dtype=np.longlong)
+            d.neighbors = np.reshape(d.labels, (count, 1))
+
+        return d
+
+
+@dataclass
+class TaskResult:
+
+    add_operations: Optional[int] = None
+    add_per_second: Optional[float] = None
+
+    search_operations: Optional[int] = None
+    search_per_second: Optional[float] = None
+    recall_at_one: Optional[float] = None
+
+    def __repr__(self) -> str:
+        parts = []
+        if self.add_per_second:
+            parts.append(f'{self.add_per_second:.2f} add/s')
+        if self.search_per_second:
+            parts.append(f'{self.search_per_second:.2f} search/s')
+        if self.recall_at_one:
+            parts.append(f'{self.recall_at_one * 100:.2f}% recall@1')
+        return ', '.join(parts)
+
+    @property
+    def add_seconds(self) -> float:
+        return self.add_operations / self.add_per_second
+
+    @property
+    def search_seconds(self) -> float:
+        return self.search_operations / self.search_per_second
+
+    def __add__(self, other: TaskResult, /):
+        result = TaskResult()
+        if self.add_operations and other.add_operations:
+            result.add_operations = self.add_operations + other.add_operations
+            result.add_per_second = result.add_operations / \
+                (self.add_seconds + other.add_seconds)
+        else:
+            base = self if self.add_operations else other
+            result.add_operations = base.add_operations
+            result.add_per_second = base.add_per_second
+
+        if self.search_operations and other.search_operations:
+            result.search_operations = self.search_operations + other.search_operations
+            result.recall_at_one = (
+                self.recall_at_one * self.search_operations + other.recall_at_one *
+                other.search_operations) / (self.search_operations + other.search_operations)
+            result.search_per_second = result.search_operations / \
+                (self.search_seconds + other.search_seconds)
+        else:
+            base = self if self.search_operations else other
+            result.search_operations = base.search_operations
+            result.search_per_second = base.search_per_second
+            result.recall_at_one = base.recall_at_one
+
+        return result
+
+
+@dataclass
+class AddTask:
+
+    labels: np.ndarray
+    vectors: np.ndarray
+
+    def __call__(self, index: Index) -> TaskResult:
+
+        batch_size: int = self.vectors.shape[0]
+        old_size: int = len(index)
+        dt, _ = measure_seconds(
+            lambda: index.add(self.labels, self.vectors))
+
+        assert len(index) == old_size + batch_size
+        return TaskResult(
+            add_operations=batch_size,
+            add_per_second=batch_size/dt,
+        )
+
+    @property
+    def ndim(self):
+        return self.vectors.shape[1]
+
+    @property
+    def count(self):
+        return self.vectors.shape[0]
+
+    def inplace_shuffle(self):
+        """Rorders the `vectors` and `labels`. Often used for robustness benchmarks."""
+
+        new_order = np.arange(self.count)
+        np.random.shuffle(new_order)
+        self.labels = self.labels[new_order]
+        self.vectors = self.vectors[new_order, :]
+
+    def slices(self, batch_size: int) -> List[AddTask]:
+        """Splits this dataset into smaller chunks."""
+
+        return [
+            AddTask(
+                labels=self.labels[start_row:start_row+batch_size],
+                vectors=self.vectors[start_row:start_row+batch_size, :],
+            ) for start_row in range(0, self.count, batch_size)]
+
+    def clusters(self, number_of_clusters: int) -> List[AddTask]:
+        """Splits this dataset into smaller chunks."""
+
+        from sklearn.cluster import KMeans
+        clustering = KMeans(
+            n_clusters=number_of_clusters,
+            random_state=0, n_init='auto',
+        ).fit(self.vectors)
+
+        partitioning = defaultdict(list)
+        for row, cluster in enumerate(clustering.labels_):
+            partitioning[cluster].append(row)
+
+        return [
+            AddTask(
+                labels=self.labels[rows],
+                vectors=self.vectors[rows, :],
+            ) for rows in partitioning.values()]
+
+
+@dataclass
+class SearchTask:
+
+    queries: np.ndarray
+    neighbors: np.ndarray
+
+    def __call__(self, index: Index) -> TaskResult:
+
+        dt, results = measure_seconds(
+            lambda: index.search(self.queries, self.neighbors.shape[1]))
+
+        return TaskResult(
+            search_per_second=self.queries.shape[0]/dt,
+            recall_at_one=results.recall_first(self.neighbors[:, 0].flatten()),
+        )
+
+    def slices(self, batch_size: int) -> List[SearchTask]:
+        """Splits this dataset into smaller chunks."""
+
+        return [
+            SearchTask(
+                queries=self.queries[start_row:start_row+batch_size, :],
+                neighbors=self.neighbors[start_row:start_row+batch_size, :],
+            ) for start_row in range(0, self.queries.shape[0], batch_size)]
+
+
+@dataclass
+class Evaluation:
+
+    tasks: List[Union[AddTask, SearchTask]]
+    count: int
+    ndim: int
+
+    @staticmethod
+    def for_dataset(dataset: Dataset, batch_size: int = 0, clusters: int = 1) -> Evaluation:
+        tasks = []
+        add = AddTask(
+            vectors=dataset.vectors,
+            labels=dataset.labels)
+        search = SearchTask(
+            queries=dataset.queries,
+            neighbors=dataset.neighbors)
+
+        if batch_size:
+            tasks.extend(add.slices(batch_size))
+            tasks.extend(search.slices(batch_size))
+        elif clusters != 1:
+            tasks.extend(add.clusters(clusters))
+            print(tasks)
+            tasks.append(search)
+        else:
+            tasks.append(add)
+            tasks.append(search)
+
+        return Evaluation(
+            tasks=tasks,
+            count=add.count,
+            ndim=add.ndim,
+        )
+
+    def __call__(self, index: Index, post_clean: bool = True) -> dict:
+
+        task_result = TaskResult()
+
+        try:
+            for task in self.tasks:
+                task_result = task_result + task(index)
+        except KeyboardInterrupt:
+            pass
+
+        if post_clean:
+            index.clear()
+        return {
+            **index.specs,
+            **asdict(task_result),
+        }