usearch 1.1.1 → 1.2.1

package/python/scripts/index_faiss.py

@@ -4,7 +4,7 @@ from typing import Optional
 import numpy as np
 from faiss import IndexHNSWFlat, IndexIVFPQ, read_index
 
-from usearch.index import Matches
+from usearch.index import Matches, BatchMatches
 from usearch.index import (
     DEFAULT_CONNECTIVITY,
     DEFAULT_EXPANSION_ADD,
@@ -45,9 +45,9 @@ class IndexFAISS:
         # self._faiss.add_with_ids(vectors, keys)
         self._faiss.add(vectors)
 
-    def search(self, queries, k: int) -> Matches:
+    def search(self, queries, k: int) -> BatchMatches:
         distances, keys = self._faiss.search(queries, k)
-        return Matches(keys, distances, np.array([k] * queries.shape[0]))
+        return BatchMatches(keys, distances, np.array([k] * queries.shape[0]))
 
     def __len__(self) -> int:
         return self._faiss.ntotal

package/python/scripts/join.py

@@ -21,7 +21,7 @@ from usearch.index import Index, MetricKind, CompiledMetric, MetricSignature
 from usearch.io import load_matrix
 from usearch.eval import measure_seconds
 
-k = 10
+count = 10
 exact = False
 batch_size = 1024 * 4
 max_elements = 1000000
@@ -104,7 +104,7 @@ if len(run_diagnostics) == 0 or run_diagnostics.lower() == "y":
 
     dt = measure_seconds
     args = dict(
-        k=k,
+        count=count,
         batch_size=batch_size,
         log=True,
         exact=exact,
@@ -113,28 +113,28 @@ if len(run_diagnostics) == 0 or run_diagnostics.lower() == "y":
     secs, a_self_recall = dt(lambda: a.search(a.vectors, **args).recall(a.keys))
     print(
         "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
-            k, a_name, a_self_recall * 100, secs
+            count, a_name, a_self_recall * 100, secs
         )
     )
 
     secs, b_self_recall = dt(lambda: b.search(b.vectors, **args).recall(b.keys))
     print(
         "Self-recall @{} of {} index: {:.2f}%, took {:.2f}s".format(
-            k, b_name, b_self_recall * 100, secs
+            count, b_name, b_self_recall * 100, secs
         )
     )
 
     secs, ab_recall = dt(lambda: b.search(a.vectors, **args).recall(b.keys))
     print(
         "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
-            k, a_name, b_name, ab_recall * 100, secs
+            count, a_name, b_name, ab_recall * 100, secs
         )
     )
 
     secs, ba_recall = dt(lambda: a.search(b.vectors, **args).recall(a.keys))
     print(
         "Cross-recall @{} of {} in {}: {:.2f}%, took {:.2f}s".format(
-            k, b_name, a_name, ba_recall * 100, secs
+            count, b_name, a_name, ba_recall * 100, secs
         )
     )
 

package/python/scripts/test_index.py

@@ -0,0 +1,219 @@
+import os
+
+import pytest
+import numpy as np
+
+from usearch.io import load_matrix, save_matrix
+from usearch.eval import random_vectors
+from usearch.index import search
+
+from usearch.index import (
+    Index,
+    MetricKind,
+    ScalarKind,
+    Match,
+    Matches,
+    BatchMatches,
+    Clustering,
+)
+from usearch.index import (
+    DEFAULT_CONNECTIVITY,
+)
+
+
+ndims = [3, 97, 256]
+batch_sizes = [1, 11, 77]
+quantizations = [
+    ScalarKind.F32,
+    ScalarKind.F64,
+    ScalarKind.F16,
+    ScalarKind.I8,
+]
+dtypes = [np.float32, np.float64, np.float16]
+threads = 2
+
+connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
+continuous_metrics = [MetricKind.Cos, MetricKind.L2sq]
+hash_metrics = [
+    MetricKind.Hamming,
+    MetricKind.Tanimoto,
+    MetricKind.Sorensen,
+]
+
+
+@pytest.mark.parametrize("ndim", [3, 97, 256])
+@pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+@pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
+def test_index_initialization_and_addition(
+    ndim, metric, quantization, dtype, batch_size
+):
+    index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
+    index.add(keys, vectors, threads=threads)
+    assert len(index) == batch_size
+
+
+@pytest.mark.parametrize("ndim", [3, 97, 256])
+@pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+@pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
+def test_index_retrieval(ndim, metric, quantization, dtype, batch_size):
+    index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
+    index.add(keys, vectors, threads=threads)
+    vectors_retrived = np.vstack(index.get(keys, dtype))
+    assert np.allclose(vectors_retrived, vectors, atol=0.1)
+
+
+@pytest.mark.parametrize("ndim", [3, 97, 256])
+@pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+@pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
+def test_index_search(ndim, metric, quantization, dtype, batch_size):
+    index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
+    index.add(keys, vectors, threads=threads)
+
+    if batch_size == 1:
+        matches: Matches = index.search(vectors, 10, threads=threads)
+        assert matches.keys.ndim == 1
+        assert matches.keys.shape[0] == matches.distances.shape[0]
+        assert len(matches) == batch_size
+        assert np.all(np.sort(index.keys) == np.sort(keys))
+
+    else:
+        matches: BatchMatches = index.search(vectors, 10, threads=threads)
+        assert matches.keys.ndim == 2
+        assert matches.keys.shape[0] == matches.distances.shape[0]
+        assert len(matches) == batch_size
+        assert np.all(np.sort(index.keys) == np.sort(keys))
+
+
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+def test_index_duplicates(batch_size):
+    ndim = 8
+    index = Index(ndim=ndim, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
+    index.add(keys, vectors, threads=threads)
+    with pytest.raises(Exception):
+        index.add(keys, vectors, threads=threads)
+
+    index = Index(ndim=ndim, multi=True)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
+    index.add(keys, vectors, threads=threads)
+    index.add(keys, vectors, threads=threads)
+    assert len(index) == batch_size * 2
+
+    two_per_key = index.get(keys)
+    assert np.vstack(two_per_key).shape == (2 * batch_size, ndim)
+
+
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+def test_index_stats(batch_size):
+    ndim = 8
+    index = Index(ndim=ndim, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
+    index.add(keys, vectors, threads=threads)
+
+    assert index.max_level >= 0
+    assert index.levels_stats.nodes >= batch_size
+    assert index.level_stats(0).nodes == batch_size
+
+
+@pytest.mark.parametrize("ndim", [1, 3, 8, 32, 256, 4096])
+@pytest.mark.parametrize("batch_size", [1, 7, 1024])
+@pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
+def test_index_save_load_restore_copy(ndim, quantization, batch_size):
+    index = Index(ndim=ndim, dtype=quantization, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
+    index.add(keys, vectors, threads=threads)
+
+    index.save("tmp.usearch")
+    index.clear()
+    assert len(index) == 0
+
+    index.load("tmp.usearch")
+    assert len(index) == batch_size
+    assert len(index[0].flatten()) == ndim
+
+    index = Index.restore("tmp.usearch", view=True)
+    assert len(index) == batch_size
+    assert len(index[0].flatten()) == ndim
+
+    copied_index = index.copy()
+    assert len(copied_index) == len(index)
+    assert np.allclose(np.vstack(copied_index.get(keys)), np.vstack(index.get(keys)))
+
+    index.reset()
+    os.remove("tmp.usearch")
+
+
+@pytest.mark.parametrize("batch_size", [32])
+def test_index_contains_remove_rename(batch_size):
+    if batch_size <= 1:
+        return
+
+    ndim = 8
+    index = Index(ndim=ndim, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim)
+
+    index.add(keys, vectors, threads=threads)
+    assert np.all(index.contains(keys))
+    assert np.all(index.count(keys) == np.ones(batch_size))
+
+    removed_keys = keys[: batch_size // 2]
+    remaining_keys = keys[batch_size // 2 :]
+    index.remove(removed_keys)
+    assert len(index) == (len(keys) - len(removed_keys))
+    assert np.sum(index.contains(keys)) == len(remaining_keys)
+    assert np.sum(index.count(keys)) == len(remaining_keys)
+    assert np.sum(index.count(removed_keys)) == 0
+
+    assert keys[0] not in index
+    assert keys[-1] in index
+
+    renamed_counts = index.rename(removed_keys, removed_keys)
+    assert np.sum(index.count(renamed_counts)) == 0
+
+    renamed_counts = index.rename(remaining_keys, removed_keys)
+    assert np.sum(index.count(removed_keys)) == len(index)
+
+
+@pytest.mark.parametrize("ndim", [3, 97, 256])
+@pytest.mark.parametrize("metric", [MetricKind.Cos, MetricKind.L2sq])
+@pytest.mark.parametrize("batch_size", [500, 1024])
+@pytest.mark.parametrize("quantization", [ScalarKind.F32, ScalarKind.I8])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64, np.float16])
+def test_index_clustering(ndim, metric, quantization, dtype, batch_size):
+    index = Index(ndim=ndim, metric=metric, dtype=quantization, multi=False)
+    keys = np.arange(batch_size)
+    vectors = random_vectors(count=batch_size, ndim=ndim, dtype=dtype)
+    index.add(keys, vectors, threads=threads)
+
+    clusters: Clustering = index.cluster(vectors=vectors, threads=threads)
+    assert len(clusters.matches.keys) == batch_size
+
+    # If no argument is provided, we cluster the present entries
+    clusters: Clustering = index.cluster(threads=threads)
+    assert len(clusters.matches.keys) == batch_size
+
+    # If no argument is provided, we cluster the present entries
+    clusters: Clustering = index.cluster(keys=keys[:50], threads=threads)
+    assert len(clusters.matches.keys) == 50
+
+    # If no argument is provided, we cluster the present entries
+    clusters: Clustering = index.cluster(min_count=3, max_count=10, threads=threads)
+    unique_clusters = set(clusters.matches.keys.flatten().tolist())
+    assert len(unique_clusters) >= 3 and len(unique_clusters) <= 10

package/python/scripts/test_sparse.py

@@ -1,32 +1,32 @@
-import os
-
 import pytest
 import numpy as np
 
 from usearch.index import (
-    SparseIndex,
-)
-from usearch.index import (
-    DEFAULT_CONNECTIVITY,
-    DEFAULT_EXPANSION_ADD,
-    DEFAULT_EXPANSION_SEARCH,
+    Index,
+    MetricKind,
+    ScalarKind,
 )
 
 
-dimensions = [3, 97, 256]
-batch_sizes = [1, 77]
-connectivity_options = [3, 13, 50, DEFAULT_CONNECTIVITY]
+@pytest.mark.parametrize("bits", [3, 97, 256, 4097])
+@pytest.mark.parametrize("metric", [MetricKind.Tanimoto])
+@pytest.mark.parametrize("connectivity", [3, 13, 50])
+@pytest.mark.parametrize("batch_size", [3, 77])
+def test_bitwise_index(
+    bits: int,
+    metric: MetricKind,
+    connectivity: int,
+    batch_size: int,
+):
+    index = Index(ndim=bits, metric=metric, connectivity=connectivity)
 
+    keys = np.arange(batch_size)
+    byte_vectors = np.random.randint(2, size=(batch_size, bits))
+    bit_vectors = np.packbits(byte_vectors, axis=1)
 
-@pytest.mark.parametrize("connectivity", connectivity_options)
-@pytest.mark.skipif(os.name == "nt", reason="Spurious behaviour on windows")
-def test_sets_index(connectivity: int):
-    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)
-    assert list(results) == [10, 11]
+    index.add(keys, bit_vectors)
 
+    byte_vectors_retrieved = np.vstack(index.get(keys, ScalarKind.B1))
+    assert np.all(byte_vectors_retrieved == bit_vectors)
 
-if __name__ == "__main__":
-    pytest.main(args=["python/scripts/test.py", "-s", "-x", "-v"])
+    index.search(bit_vectors, 10)

package/python/scripts/test_tooling.py

@@ -0,0 +1,111 @@
+import os
+
+import pytest
+import numpy as np
+
+from usearch.io import load_matrix, save_matrix
+from usearch.index import search
+from usearch.eval import random_vectors
+
+from usearch.index import Match, Matches, BatchMatches, Index, Indexes
+
+
+dimensions = [3, 97, 256]
+batch_sizes = [1, 77]
+
+
+@pytest.mark.parametrize("rows", batch_sizes)
+@pytest.mark.parametrize("cols", dimensions)
+def test_serializing_fbin_matrix(rows: int, cols: int):
+    """
+    Test the serialization of floating point binary matrix.
+
+    :param int rows: The number of rows in the matrix.
+    :param int cols: The number of columns in the matrix.
+    """
+    original = np.random.rand(rows, cols).astype(np.float32)
+    save_matrix(original, "tmp.fbin")
+    reconstructed = load_matrix("tmp.fbin")
+    assert np.allclose(original, reconstructed)
+    os.remove("tmp.fbin")
+
+
+@pytest.mark.parametrize("rows", batch_sizes)
+@pytest.mark.parametrize("cols", dimensions)
+def test_serializing_ibin_matrix(rows: int, cols: int):
+    """
+    Test the serialization of integer binary matrix.
+
+    :param int rows: The number of rows in the matrix.
+    :param int cols: The number of columns in the matrix.
+    """
+    original = np.random.randint(0, rows + 1, size=(rows, cols)).astype(np.int32)
+    save_matrix(original, "tmp.ibin")
+    reconstructed = load_matrix("tmp.ibin")
+    assert np.allclose(original, reconstructed)
+    os.remove("tmp.ibin")
+
+
+@pytest.mark.parametrize("rows", batch_sizes)
+@pytest.mark.parametrize("cols", dimensions)
+def test_exact_search(rows: int, cols: int):
+    """
+    Test exact search.
+
+    :param int rows: The number of rows in the matrix.
+    :param int cols: The number of columns in the matrix.
+    """
+    original = np.random.rand(rows, cols)
+    matches: BatchMatches = search(original, original, 10, exact=True)
+    top_matches = (
+        [int(m.keys[0]) for m in matches] if rows > 1 else int(matches.keys[0])
+    )
+    assert np.all(top_matches == np.arange(rows))
+
+    matches: Matches = search(original, original[0], 10, exact=True)
+    top_match = int(matches.keys[0])
+    assert top_match == 0
+
+
+def test_matches_creation_and_methods():
+    matches = Matches(
+        keys=np.array([1, 2]),
+        distances=np.array([0.5, 0.6]),
+        visited_members=2,
+        computed_distances=2,
+    )
+    assert len(matches) == 2
+    assert matches[0] == Match(key=1, distance=0.5)
+    assert matches.to_list() == [(1, 0.5), (2, 0.6)]
+
+
+def test_batch_matches_creation_and_methods():
+    keys = np.array([[1, 2], [3, 4]])
+    distances = np.array([[0.5, 0.6], [0.7, 0.8]])
+    counts = np.array([2, 2])
+    batch_matches = BatchMatches(
+        keys=keys,
+        distances=distances,
+        counts=counts,
+        visited_members=2,
+        computed_distances=2,
+    )
+
+    assert len(batch_matches) == 2
+    assert batch_matches[0].keys.tolist() == [1, 2]
+    assert batch_matches[0].distances.tolist() == [0.5, 0.6]
+    assert batch_matches.to_list() == [(1, 0.5), (2, 0.6), (3, 0.7), (4, 0.8)]
+
+
+def test_multi_index():
+    ndim = 10
+    index_a = Index(ndim=ndim)
+    index_b = Index(ndim=ndim)
+
+    vectors = random_vectors(count=3, ndim=ndim)
+    index_a.add(42, vectors[0])
+    index_b.add(43, vectors[1])
+
+    indexes = Indexes([index_a, index_b])
+    matches = indexes.search(vectors[2], 10)
+    assert len(matches) == 2

package/python/usearch/eval.py

@@ -103,14 +103,18 @@ def self_recall(index: Index, sample: float = 1, **kwargs) -> SearchStats:
     """
     if len(index) == 0:
         return 0
-    if "k" not in kwargs:
-        kwargs["k"] = 1
+    if "count" not in kwargs:
+        kwargs["count"] = 1
+
+    if "keys" in kwargs:
+        keys = kwargs["keys"]
+    else:
+        keys = np.array(index.keys)
 
-    keys = index.keys
     if sample != 1:
         keys = np.random.choice(keys, int(ceil(len(keys) * sample)))
 
-    queries = index.get_vectors(keys, index.dtype)
+    queries = index.get(keys, index.dtype)
     matches: BatchMatches = index.search(queries, **kwargs)
     count_matches: float = matches.count_matches(keys)
     return SearchStats(
@@ -403,7 +407,7 @@ class SearchTask:
 
         return TaskResult(
             search_per_second=self.queries.shape[0] / dt,
-            recall_at_one=results.recall_first(self.neighbors[:, 0].flatten()),
+            recall_at_one=results.mean_recall(self.neighbors[:, 0].flatten()),
         )
 
     def slices(self, batch_size: int) -> List[SearchTask]: