usearch 0.20.0 → 0.21.0

package/python/usearch/index.py

@@ -5,12 +5,13 @@ from __future__ import annotations
 # Python tooling, linters, and static analyzers. It also embeds JIT
 # into the primary `Index` class, connecting USearch with Numba.
 import os
+import math
 from typing import Optional, Union, NamedTuple, List, Iterable
 
 import numpy as np
 from tqdm import tqdm
 
-from usearch.compiled import Index as _CompiledIndex, IndexMetadata
+from usearch.compiled import Index as _CompiledIndex, IndexMetadata, IndexStats
 from usearch.compiled import SparseIndex as _CompiledSetsIndex
 
 from usearch.compiled import MetricKind, ScalarKind, MetricSignature
@@ -141,6 +142,13 @@ class Matches(NamedTuple):
         best_matches = self.labels if not self.is_multiple else self.labels[:, 0]
         return np.sum(best_matches == expected) / len(expected)
 
+    def recall(self, expected: np.ndarray) -> float:
+        assert len(expected) == self.batch_size
+        recall = 0
+        for i in range(self.batch_size):
+            recall += expected[i] in self.labels[i]
+        return recall / len(expected)
+
     def __repr__(self) -> str:
         return (
             f"usearch.Matches({self.total_matches})"
@@ -347,6 +355,10 @@ class Index:
         count_labels = len(labels) if isinstance(labels, Iterable) else 1
         assert count_labels == count_vectors
 
+        # If logging is requested, and batch size is undefined, set it to grow 1% at a time:
+        if log and batch_size == 0:
+            batch_size = int(math.ceil(count_vectors / 100))
+
         # Split into batches and log progress, if needed
         if batch_size:
             labels = [
@@ -401,13 +413,53 @@ class Index:
         :return: Approximate matches for one or more queries
         :rtype: Matches
         """
-        tuple_ = self._compiled.search(
-            vectors,
-            k,
-            exact=exact,
-            threads=threads,
-        )
-        return Matches(*tuple_)
+
+        assert isinstance(vectors, np.ndarray), "Expects a NumPy array"
+        assert vectors.ndim == 1 or vectors.ndim == 2, "Expects a matrix or vector"
+        count_vectors = vectors.shape[0] if vectors.ndim == 2 else 1
+
+        if log and batch_size == 0:
+            batch_size = int(math.ceil(count_vectors / 100))
+
+        if batch_size:
+            tasks = [
+                vectors[start_row : start_row + batch_size, :]
+                for start_row in range(0, count_vectors, batch_size)
+            ]
+            tasks_matches = []
+            name = log if isinstance(log, str) else "Search"
+            pbar = tqdm(
+                tasks,
+                desc=name,
+                total=count_vectors,
+                unit="Vector",
+                disable=log is False,
+            )
+            for vectors in tasks:
+                tuple_ = self._compiled.search(
+                    vectors,
+                    k,
+                    exact=exact,
+                    threads=threads,
+                )
+                tasks_matches.append(Matches(*tuple_))
+                pbar.update(vectors.shape[0])
+
+            pbar.close()
+            return Matches(
+                labels=np.vstack([m.labels for m in tasks_matches]),
+                distances=np.vstack([m.distances for m in tasks_matches]),
+                counts=np.concatenate([m.counts for m in tasks_matches], axis=None),
+            )
+
+        else:
+            tuple_ = self._compiled.search(
+                vectors,
+                k,
+                exact=exact,
+                threads=threads,
+            )
+            return Matches(*tuple_)
 
     @property
     def specs(self) -> dict:
@@ -442,8 +494,8 @@ class Index:
         return self._compiled.ndim
 
     @property
-    def metric(self) -> MetricKind:
-        return self._metric_kind
+    def metric(self) -> Union[MetricKind, CompiledMetric]:
+        return self._metric_jit if self._metric_jit else self._metric_kind
 
     @property
     def dtype(self) -> ScalarKind:
@@ -501,6 +553,26 @@ class Index:
     def close(self):
         self._compiled.close()
 
+    def copy(self) -> Index:
+        result = Index(
+            ndim=self.ndim,
+            metric=self.metric,
+            dtype=self.dtype,
+            connectivity=self.connectivity,
+            expansion_add=self.expansion_add,
+            expansion_search=self.expansion_search,
+            path=self.path,
+        )
+        result._compiled = self._compiled.copy()
+        return result
+
+    def join(self, other: Index, max_proposals: int = 0, exact: bool = False) -> dict:
+        return self._compiled.join(
+            other=other._compiled,
+            max_proposals=max_proposals,
+            exact=exact,
+        )
+
     @property
     def labels(self) -> np.ndarray:
         return self._compiled.labels
@@ -530,3 +602,40 @@ class Index:
         vector = self._compiled.__getitem__(label, get_dtype)
         view_dtype = _to_numpy_dtype(dtype)
         return None if vector is None else vector.view(view_dtype)
+
+    def remove(self, label: int):
+        pass
+
+    @property
+    def max_level(self) -> int:
+        return self._compiled.max_level
+
+    @property
+    def levels_stats(self) -> IndexStats:
+        return self._compiled.levels_stats
+
+    def level_stats(self, level: int) -> IndexStats:
+        return self._compiled.level_stats(level)
+
+    def __repr__(self) -> str:
+        return f"usearch.index.Index({self.dtype} x {self.ndim}, {self.metric}, expansion: {self.expansion_add} & {self.expansion_search}, {len(self)} vectors across {self.max_level+1} levels)"
+
+    def _repr_pretty_(self) -> str:
+        level_stats = [
+            f"--- {i}. {self.level_stats(i)} nodes" for i in range(self.max_level)
+        ]
+        return "\n".join(
+            [
+                "usearch.index.Index",
+                "- config" f"-- data type: {self.dtype}",
+                f"-- dimensions: {self.ndim}",
+                f"-- metric: {self.metric}",
+                f"-- expansion on addition:{self.expansion_add} candidates",
+                f"-- expansion on search: {self.expansion_search} candidates",
+                "- state",
+                f"-- size: {self.size} vectors",
+                f"-- memory usage: {self.memory_usage} bytes",
+                f"-- max level: {self.max_level}",
+                *level_stats,
+            ]
+        )

package/python/usearch/join.py

@@ -1,98 +0,0 @@
-from collections import deque
-
-import numpy as np
-from numba import njit
-
-from index import Index, Matches, Label
-
-
-class bidict(dict):
-    def __init__(self, *args, **kwargs):
-        super(bidict, self).__init__(*args, **kwargs)
-        self.inverse = {}
-        for key, value in self.items():
-            self.inverse.setdefault(value, []).append(key)
-
-    def __setitem__(self, key, value):
-        if key in self:
-            self.inverse[self[key]].remove(key)
-        super(bidict, self).__setitem__(key, value)
-        self.inverse.setdefault(value, []).append(key)
-
-    def __delitem__(self, key):
-        self.inverse.setdefault(self[key], []).remove(key)
-        if self[key] in self.inverse and not self.inverse[self[key]]:
-            del self.inverse[self[key]]
-        super(bidict, self).__delitem__(key)
-
-
-@njit
-def index_in_array(array: np.ndarray, item: Label) -> int:
-    for idx, val in np.ndenumerate(array):
-        if val == item:
-            return idx
-
-    return len(array)
-
-
-def semantic_join(a: Index, b: Index, resolution: int = 10) -> bidict:
-    """Performs a Semantic Join, matching one entry from `a` with
-    one entry from `b`, converging towards a Stable Marriage.
-    Assuming the collections can be different in size, classical solution
-    doesn't provide
-
-    :param resolution: Approximate matches per member to consider, defaults to 10
-    :type resolution: int, optional
-    :return: Bidirectional mapping from men to women and `.inverse`
-    :rtype: bidict
-    """
-
-    men, women = (a, b) if len(a) > len(b) else (b, a)
-    matches = bidict()
-
-    men_vectors = np.vstack(men[i] for i in range(len(men)))
-    women_vectors = np.vstack(women[i] for i in range(len(women)))
-
-    man_count_proposed = np.zeros(len(men))
-    man_to_women_preferences: Matches = men.search(women_vectors, resolution)
-    woman_to_men_preferences: Matches = women.search(men_vectors, resolution)
-
-    # A nice optimization is to prioritize free men by the quality of their
-    # best remaining variant. TODO: Replace with `heapq`
-    free_men = deque(range(len(men)))
-
-    idle_cycles: int = 0
-    hopeless_men_count = len(men) - len(women)
-    while len(free_men) > hopeless_men_count:
-        # In the worst case scenario we may need to match more candidates for every
-        # remaining man. This, however, may drastically increase the runtime.
-        if len(free_men) == idle_cycles:
-            break
-
-        man: Label = free_men.popleft()
-        count_proposals = man_count_proposed[man]
-        if count_proposals == man_to_women_preferences.counts[man]:
-            free_men.append(man)
-            idle_cycles += 1
-            continue
-
-        woman: Label = man_to_women_preferences.labels[man, count_proposals]
-        man_count_proposed[man] += 1
-
-        if woman not in matches.inverse:
-            matches[man] = woman
-
-        else:
-            her_preferences = woman_to_men_preferences.labels[woman, :]
-            husband: Label = matches.inverse[woman]
-            husband_rank = index_in_array(her_preferences, husband)
-            challenger_rank = index_in_array(her_preferences, man)
-            if challenger_rank < husband_rank:
-                del matches[husband]
-                free_men.append(husband)
-                matches[man] = woman
-
-            else:
-                free_men.append(man)
-
-    return matches