usearch 2.8.7 → 2.8.9

package/python/usearch/eval.py

@@ -1,470 +0,0 @@
-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
-from math import ceil
-
-import numpy as np
-
-from usearch.io import load_matrix
-from usearch.index import (
-    Index,
-    Matches,
-    BatchMatches,
-    MetricKind,
-    MetricKindBitwise,
-    Key,
-)
-
-
-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
-
-
-@dataclass
-class SearchStats:
-    """
-    Contains statistics for one or more search runs, including the number of
-    internal nodes that were fetched (`visited_members`) and the number
-    of times the distance metric was invoked (`computed_distances`).
-
-    Other derivative metrics include the `mean_recall` and `mean_efficiency`.
-    Recall is the share of queried vectors, that were successfully found.
-    Efficiency describes the number of distances that had to be computed for
-    each query, normalized to size of the `index`. Highest efficiency is 0.(9),
-    lowest is zero. Highest is achieved, when the distance metric was computed
-    just once per query. Lowest happens during exact search, when every distance
-    to every present vector had to be computed.
-    """
-
-    index_size: int
-    count_queries: int
-    count_matches: int
-
-    visited_members: int
-    computed_distances: int
-
-    @property
-    def mean_efficiency(self) -> float:
-        return 1 - float(self.computed_distances) / (
-            self.count_queries * self.index_size
-        )
-
-    @property
-    def mean_recall(self) -> float:
-        return self.count_matches / self.count_queries
-
-
-def self_recall(index: Index, sample: float = 1, **kwargs) -> SearchStats:
-    """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
-    :param sample: Share of vectors to search, defaults to 1
-    :type sample: float
-    :return: Evaluation report with key metrics
-    :rtype: SearchStats
-    """
-    if len(index) == 0:
-        return 0
-    if "count" not in kwargs:
-        kwargs["count"] = 1
-
-    if "keys" in kwargs:
-        keys = kwargs["keys"]
-    else:
-        keys = np.array(index.keys)
-
-    if sample != 1:
-        keys = np.random.choice(keys, int(ceil(len(keys) * sample)))
-
-    queries = index.get(keys, index.dtype)
-    matches: BatchMatches = index.search(queries, **kwargs)
-    count_matches: float = matches.count_matches(keys)
-    return SearchStats(
-        index_size=len(index),
-        count_queries=len(keys),
-        count_matches=count_matches,
-        visited_members=matches.visited_members,
-        computed_distances=matches.computed_distances,
-    )
-
-
-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
-
-
-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 keys
-    :type expected: np.ndarray
-    :param predicted: predicted keys
-    :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:
-    keys: 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.keys = np.arange(count, dtype=Key)
-
-            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.keys, (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.keys = np.arange(count, dtype=Key)
-            d.neighbors = np.reshape(d.keys, (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:
-    keys: 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.keys, 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 `keys`. Often used for robustness benchmarks."""
-
-        new_order = np.arange(self.count)
-        np.random.shuffle(new_order)
-        self.keys = self.keys[new_order]
-        self.vectors = self.vectors[new_order, :]
-
-    def slices(self, batch_size: int) -> List[AddTask]:
-        """Splits this dataset into smaller chunks."""
-
-        return [
-            AddTask(
-                keys=self.keys[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(
-                keys=self.keys[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.mean_recall(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, keys=dataset.keys)
-        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),
-        }