usearch 0.18.8 → 0.19.0

package/python/usearch/eval.py

@@ -24,7 +24,7 @@ def random_vectors(
     # Infer default parameters from the `index`, if passed
     if index is not None:
         if not isinstance(index, Index):
-            raise ValueError('Unsupported `index` type')
+            raise ValueError("Unsupported `index` type")
 
         ndim = index.ndim
         dtype = index.numpy_dtype
@@ -37,7 +37,6 @@ def random_vectors(
         return bit_vectors
 
     else:
-
         x = np.random.rand(count, ndim)
         if dtype == np.int8:
             x = (x * 100).astype(np.int8)
@@ -77,13 +76,65 @@ def measure_seconds(f: Callable) -> Tuple[float, Any]:
     result = f()
     b = time_ns()
     c = b - a
-    secs = c / (10 ** 9)
+    secs = c / (10**9)
     return secs, result
 
 
+def dcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray:
+    """Calculate DCG (Discounted Cumulative Gain) up to position k.
+
+    :param relevances: List of true relevance scores (in the order as they are ranked)
+    :type relevances: list
+    :param k: Position up to which DCG is computed
+    :type k: int
+    :return: The DCG score at position k
+    :rtype: float
+    """
+    if k:
+        relevances = np.asarray(relevances)[:k]
+
+    n_relevances = len(relevances)
+    if n_relevances == 0:
+        return 0.0
+
+    discounts = np.log2(np.arange(n_relevances) + 2)
+    return np.sum(relevances / discounts)
+
+
+def ndcg(relevances: np.ndarray, k: Optional[int] = None) -> np.ndarray:
+    """Calculate NDCG (Normalized Discounted Cumulative Gain) at position k.
+
+    :param relevances: List of true relevance scores (in the order as they are ranked)
+    :type relevances: list
+    :param k: Position up to which NDCG is computed
+    :type k: int
+    :return: The NDCG score at position k
+    :rtype: float
+    """
+    best_dcg = dcg(sorted(relevances, reverse=True), k)
+    if best_dcg == 0:
+        return 0.0
+
+    return dcg(relevances, k) / best_dcg
+
+
+def relevance(
+    expected: np.ndarray, predicted: np.ndarray, k: Optional[int] = None
+) -> np.ndarray:
+    """Calculate relevance scores. Binary relevance scores
+
+    :param expected: ground-truth labels
+    :type expected: np.ndarray
+    :param predicted: predicted labels
+    :type predicted: np.ndarray
+    """
+    expected = expected[:k]
+    predicted = predicted[:k]
+    return [1 if i in expected else 0 for i in predicted]
+
+
 @dataclass
 class Dataset:
-
     labels: np.ndarray
     vectors: np.ndarray
     queries: np.ndarray
@@ -124,8 +175,11 @@ class Dataset:
 
             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]
+            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=Label)
 
@@ -139,14 +193,13 @@ class Dataset:
                 if k is not None:
                     d.neighbors = d.neighbors[:, :k]
             else:
-                assert k is None, 'Cant ovveride `k`, will retrieve one neighbor'
+                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'
+            assert k is None, "Cant ovveride `k`, will retrieve one neighbor"
 
             d.vectors = random_vectors(count=count, ndim=ndim)
             d.queries = d.vectors
@@ -158,7 +211,6 @@ class Dataset:
 
 @dataclass
 class TaskResult:
-
     add_operations: Optional[int] = None
     add_per_second: Optional[float] = None
 
@@ -169,12 +221,12 @@ class TaskResult:
     def __repr__(self) -> str:
         parts = []
         if self.add_per_second:
-            parts.append(f'{self.add_per_second:.2f} add/s')
+            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')
+            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)
+            parts.append(f"{self.recall_at_one * 100:.2f}% recall@1")
+        return ", ".join(parts)
 
     @property
     def add_seconds(self) -> float:
@@ -188,8 +240,9 @@ class 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)
+            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
@@ -198,10 +251,12 @@ class TaskResult:
         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)
+                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
@@ -213,21 +268,18 @@ class TaskResult:
 
 @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))
+        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,
+            add_per_second=batch_size / dt,
         )
 
     @property
@@ -251,17 +303,21 @@ class AddTask:
 
         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)]
+                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',
+            random_state=0,
+            n_init="auto",
         ).fit(self.vectors)
 
         partitioning = defaultdict(list)
@@ -272,22 +328,23 @@ class AddTask:
             AddTask(
                 labels=self.labels[rows],
                 vectors=self.vectors[rows, :],
-            ) for rows in partitioning.values()]
+            )
+            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]))
+            lambda: index.search(self.queries, self.neighbors.shape[1])
+        )
 
         return TaskResult(
-            search_per_second=self.queries.shape[0]/dt,
+            search_per_second=self.queries.shape[0] / dt,
             recall_at_one=results.recall_first(self.neighbors[:, 0].flatten()),
         )
 
@@ -296,31 +353,26 @@ class SearchTask:
 
         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)]
+                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:
-
+        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)
+        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))
@@ -340,7 +392,6 @@ class Evaluation:
         )
 
     def __call__(self, index: Index, post_clean: bool = True) -> dict:
-
         task_result = TaskResult()
 
         try:

package/python/usearch/index.py

@@ -1,22 +1,22 @@
+from __future__ import annotations
+
 # The purpose of this file is to provide Pythonic wrapper on top
 # the native precompiled CPython module. It improves compatibility
 # Python tooling, linters, and static analyzers. It also embeds JIT
 # into the primary `Index` class, connecting USearch with Numba.
 import os
-from math import sqrt
-from typing import Optional, Callable, Union, NamedTuple, List, Iterable
+from typing import Optional, Union, NamedTuple, List, Iterable
 
 import numpy as np
 
-from usearch.compiled import Index as _CompiledIndex
+from usearch.compiled import Index as _CompiledIndex, IndexMetadata
 from usearch.compiled import SparseIndex as _CompiledSetsIndex
 
-from usearch.compiled import MetricKind, ScalarKind
+from usearch.compiled import MetricKind, ScalarKind, MetricSignature
 from usearch.compiled import (
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
     DEFAULT_EXPANSION_SEARCH,
-
     USES_OPENMP,
     USES_SIMSIMD,
     USES_NATIVE_F16,
@@ -41,11 +41,11 @@ def _normalize_dtype(dtype) -> ScalarKind:
 
     if isinstance(dtype, str):
         _normalize = {
-            'f64': ScalarKind.F64,
-            'f32': ScalarKind.F32,
-            'f16': ScalarKind.F16,
-            'f8': ScalarKind.F8,
-            'b1': ScalarKind.B1,
+            "f64": ScalarKind.F64,
+            "f32": ScalarKind.F32,
+            "f16": ScalarKind.F16,
+            "f8": ScalarKind.F8,
+            "b1": ScalarKind.B1,
         }
         return _normalize[dtype.lower()]
 
@@ -79,9 +79,8 @@ class Matches(NamedTuple):
         return np.sum(self.counts)
 
     def to_list(self, row: Optional[int] = None) -> Union[List[dict], List[List[dict]]]:
-
         if not self.is_batch:
-            assert row is None, 'Exporting a single sequence is only for batch requests'
+            assert row is None, "Exporting a single sequence is only for batch requests"
             labels = self.labels
             distances = self.distances
 
@@ -94,7 +93,7 @@ class Matches(NamedTuple):
             distances = self.distances[row, :count]
 
         return [
-            {'label': int(label), 'distance': float(distance)}
+            {"label": int(label), "distance": float(distance)}
             for label, distance in zip(labels, distances)
         ]
 
@@ -103,8 +102,17 @@ class Matches(NamedTuple):
         return np.sum(best_matches == expected) / len(expected)
 
     def __repr__(self) -> str:
-        return f'usearch.Matches({self.total_matches})' if self.is_batch else \
-            f'usearch.Matches({self.total_matches} across {self.batch_size} queries)'
+        return (
+            f"usearch.Matches({self.total_matches})"
+            if self.is_batch
+            else f"usearch.Matches({self.total_matches} across {self.batch_size} queries)"
+        )
+
+
+class CompiledMetric(NamedTuple):
+    pointer: int
+    kind: MetricKind
+    signature: MetricSignature
 
 
 class Index:
@@ -120,15 +128,13 @@ class Index:
     def __init__(
         self,
         ndim: int,
-        metric: Union[MetricKind, Callable, str] = MetricKind.IP,
+        metric: Union[str, MetricKind, CompiledMetric] = MetricKind.IP,
         dtype: Optional[str] = None,
         jit: bool = False,
-
         connectivity: int = DEFAULT_CONNECTIVITY,
         expansion_add: int = DEFAULT_EXPANSION_ADD,
         expansion_search: int = DEFAULT_EXPANSION_SEARCH,
         tune: bool = False,
-
         path: Optional[os.PathLike] = None,
         view: bool = False,
     ) -> None:
@@ -171,7 +177,7 @@ class Index:
 
         :param expansion_search: Traversal depth on queries, defaults to None
         :type expansion_search: Optional[int], optional
-            Hyper-parameter for the search depth when querying 
+            Hyper-parameter for the search depth when querying
             nearest neighbors. The original paper calls it "ef".
             Can be changed afterwards, as the `.expansion_search`.
 
@@ -190,58 +196,58 @@ class Index:
         else:
             dtype = _normalize_dtype(dtype)
 
-
         if metric is None:
             metric = MetricKind.IP
-        elif isinstance(metric, str):
+
+        if isinstance(metric, str):
             _normalize = {
-                'cos': MetricKind.Cos,
-                'ip': MetricKind.IP,
-                'l2_sq': MetricKind.L2sq,
-                'haversine': MetricKind.Haversine,
-                'perason': MetricKind.Pearson,
-                'hamming': MetricKind.Hamming,
-                'tanimoto': MetricKind.Tanimoto,
-                'sorensen': MetricKind.Sorensen,
+                "cos": MetricKind.Cos,
+                "ip": MetricKind.IP,
+                "l2_sq": MetricKind.L2sq,
+                "haversine": MetricKind.Haversine,
+                "perason": MetricKind.Pearson,
+                "hamming": MetricKind.Hamming,
+                "tanimoto": MetricKind.Tanimoto,
+                "sorensen": MetricKind.Sorensen,
             }
             metric = _normalize[metric.lower()]
 
-        if isinstance(metric, Callable):
-            self._metric_kind = MetricKind.Unknown
-            self._metric_jit = metric
-            self._metric_pointer = int(metric.address)
-
-        elif isinstance(metric, MetricKind):
-            if jit:
-
-                try:
-                    from usearch.numba import jit
-                except ImportError:
-                    raise ModuleNotFoundError(
-                        'To use JIT install Numba with `pip install numba`.'
-                        'Alternatively, reinstall usearch with `pip install usearch[jit]`')
-
-                self._metric_kind = metric
-                self._metric_jit = jit(
-                    ndim=ndim,
-                    metric=metric,
-                    dtype=dtype,
+        if isinstance(metric, MetricKind) and jit:
+            try:
+                from usearch.numba import jit
+            except ImportError:
+                raise ModuleNotFoundError(
+                    "To use JIT install Numba with `pip install numba`."
+                    "Alternatively, reinstall with `pip install usearch[jit]`"
                 )
-                self._metric_pointer = self._metric_jit.address if \
-                    self._metric_jit else 0
-            else:
-                self._metric_kind = metric
-                self._metric_jit = None
-                self._metric_pointer = 0
+
+            metric = jit(
+                ndim=ndim,
+                metric=metric,
+                dtype=dtype,
+            )
+
+        if isinstance(metric, MetricKind):
+            self._metric_kind = metric
+            self._metric_jit = None
+            self._metric_pointer = 0
+            self._metric_signature = MetricSignature.ArrayArraySize
+        elif isinstance(metric, CompiledMetric):
+            self._metric_jit = metric
+            self._metric_kind = metric.kind
+            self._metric_pointer = metric.pointer
+            self._metric_signature = metric.signature
         else:
             raise ValueError(
-                'The `metric` must be Numba callback or a `MetricKind`')
+                "The `metric` must be a `CompiledMetric` or a `MetricKind`"
+            )
 
         self._compiled = _CompiledIndex(
             ndim=ndim,
+            dtype=dtype,
             metric=self._metric_kind,
             metric_pointer=self._metric_pointer,
-            dtype=dtype,
+            metric_signature=self._metric_signature,
             connectivity=connectivity,
             expansion_add=expansion_add,
             expansion_search=expansion_search,
@@ -255,17 +261,40 @@ class Index:
             else:
                 self._compiled.load(path)
 
+    @staticmethod
+    def metadata(path: os.PathLike) -> IndexMetadata:
+        return IndexMetadata(path)
+
+    @staticmethod
+    def restore(path: os.PathLike, view: bool = False) -> Index:
+        meta = Index.metadata(path)
+        bits_per_scalar = {
+            ScalarKind.F8: 8,
+            ScalarKind.F16: 16,
+            ScalarKind.F32: 32,
+            ScalarKind.F64: 64,
+            ScalarKind.B1: 1,
+        }[meta.scalar_kind]
+        ndim = meta.bytes_for_vectors * 8 // meta.size // bits_per_scalar
+        return Index(
+            ndim=ndim,
+            connectivity=meta.connectivity,
+            metric=meta.metric,
+            path=path,
+            view=view,
+        )
+
     def add(
-            self, labels, vectors, *,
-            copy: bool = True, threads: int = 0) -> Union[int, np.ndarray]:
+        self, labels, vectors, *, copy: bool = True, threads: int = 0
+    ) -> Union[int, np.ndarray]:
         """Inserts one or move vectors into the index.
 
         For maximal performance the `labels` and `vectors`
         should conform to the Python's "buffer protocol" spec.
 
-        To index a single entry: 
+        To index a single entry:
             labels: int, vectors: np.ndarray.
-        To index many entries: 
+        To index many entries:
             labels: np.ndarray, vectors: np.ndarray.
 
         When working with extremely large indexes, you may want to
@@ -284,8 +313,8 @@ class Index:
         :return: Inserted label or labels
         :type: Union[int, np.ndarray]
         """
-        assert isinstance(vectors, np.ndarray), 'Expects a NumPy array'
-        assert vectors.ndim == 1 or vectors.ndim == 2, 'Expects a matrix or vector'
+        assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
+        assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
         is_batch = vectors.ndim == 2
         generate_labels = labels is None
 
@@ -306,8 +335,8 @@ class Index:
         return labels
 
     def search(
-            self, vectors, k: int = 10, *,
-            threads: int = 0, exact: bool = False) -> Matches:
+        self, vectors, k: int = 10, *, threads: int = 0, exact: bool = False
+    ) -> Matches:
         """Performs approximate nearest neighbors search for one or more queries.
 
         :param vectors: Query vector or vectors.
@@ -323,26 +352,28 @@ class Index:
         :rtype: Matches
         """
         tuple_ = self._compiled.search(
-            vectors, k,
-            exact=exact, threads=threads,
+            vectors,
+            k,
+            exact=exact,
+            threads=threads,
         )
         return Matches(*tuple_)
 
     @property
     def specs(self) -> dict:
         return {
-            'Class': 'usearch.Index',
-            'Connectivity': self.connectivity,
-            'Size': self.size,
-            'Dimensions': self.ndim,
-            'Expansion@Add': self.expansion_add,
-            'Expansion@Search': self.expansion_search,
-            'OpenMP': USES_OPENMP,
-            'SimSIMD': USES_SIMSIMD,
-            'NativeF16': USES_NATIVE_F16,
-            'JIT': self.jit,
-            'DType': self.dtype,
-            'Path': self.path,
+            "Class": "usearch.Index",
+            "Connectivity": self.connectivity,
+            "Size": self.size,
+            "Dimensions": self.ndim,
+            "Expansion@Add": self.expansion_add,
+            "Expansion@Search": self.expansion_search,
+            "OpenMP": USES_OPENMP,
+            "SimSIMD": USES_SIMSIMD,
+            "NativeF16": USES_NATIVE_F16,
+            "JIT": self.jit,
+            "DType": self.dtype,
+            "Path": self.path,
         }
 
     def __len__(self) -> int:
@@ -413,10 +444,8 @@ class Index:
         return self._compiled.labels
 
     def get_vectors(
-            self,
-            labels: np.ndarray,
-            dtype: ScalarKind = ScalarKind.F32) -> np.ndarray:
-
+        self, labels: np.ndarray, dtype: ScalarKind = ScalarKind.F32
+    ) -> np.ndarray:
         dtype = _normalize_dtype(dtype)
         return np.vstack([self._compiled.__getitem__(l, dtype) for l in labels])