Feat added bk tree benchmark and docs

pablocael · Oct 9, 2023 · a2cbbb3 · a2cbbb3
1 parent 20ae652
commit a2cbbb3
Show file tree

Hide file tree

Showing 5 changed files with 47 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -1,6 +1,6 @@
 # Introduction
 
-PyNear is a python library, internally built in C++, for efficient KNN search using metric distance function such as L2 distance (see VPTreeL2Index) or Hamming distances (VPTreeBinaryIndex) as well as other distance functions. It uses AVX2 instructions to optimize distance functions so to improve search performance.
+PyNear is a python library, internally built in C++, for efficient KNN search using metric distance function such as L2 distance (see VPTreeL2Index) or Hamming distances (VPTreeBinaryIndex and BKTreeBinaryIndex) as well as other distance functions. It uses AVX2 instructions to optimize distance functions so to improve search performance.
 
 PyNear aims providing different efficient algorithms for nearest neighbor search. One of the differentials of PyNear is the adoption of [Vantage Point Tree](./docs/vptrees.md) in order to mitigate course of dimensionality for high dimensional features (see VPTree* indices for more information in [docs](./docs/README.md)).
 

diff --git a/docs/README.md b/docs/README.md
@@ -5,13 +5,22 @@
 PyNear has several available indexes that will use different distance functions or algorithms to perform the search.
 Available indices are:
 
+### Threshold based KNN Indices:
+
+| Index Name                     | Description                                                                                                                                                                                                                                       |
+|--------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| pynear.BKTreeBinaryIndex     | Uses AVX2 optimized Hamming distance function and [BKTree](https://en.wikipedia.org/wiki/BK-tree) algorithm to perform exact searches within a threshold distance.                                                                                                                                         |
+
+### Non Threshold based KNN Indices:
+
 | Index Name                     | Description                                                                                                                                                                                                                                       |
 |--------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
 | pynear.VPTreeL2Index         | Uses AVX2 optimized L2 (euclidean norm) distance function and VPTree algorithm to perform exact searches.                                                                                                                                         |
 | pynear.VPTreeL1Index         | Uses L1 (manhattan) distance function and VPTree algorithm to perform exact searches.                                                                                                                                                             |
-| pynear.VPTreeBinaryIndex     | Uses AVX2 optimized Hamming distances function and VPTree algorithm to perform exact searches. Supports 16, 32, 64, 128 and 256 bit dimensional vectors only.                                                                                                                                                     |
+| pynear.VPTreeBinaryIndex     | Uses AVX2 optimized Hamming distance function and VPTree algorithm to perform exact searches. Supports 16, 32, 64, 128 and 256 bit dimensional vectors only.                                                                                                                                                     |
 | pynear.VPTreeChebyshevIndex  | Uses [Chebyshev](https://en.wikipedia.org/wiki/Chebyshev_distance) distance function and VPTree algorithm to perform exact searches. |
 
+
 ## Usage example
 
 ### Creating the index
@@ -41,6 +50,25 @@ vptree.set(data)
 vptree_indices, vptree_distances = vptree.searchKNN(queries, k)
 ```
 
+#### pynear.BKTreeBinaryIndex
+
+```python
+np.random.seed(seed=42)
+
+dimension = 32 # 32 bytes are 256 bit dimensional vectos
+num_points = 2021
+data = np.random.normal(scale=255, loc=0, size=(num_points, dimension)).astype(dtype=np.uint8)
+
+num_queries = 8
+queries = np.random.normal(scale=255, loc=0, size=(num_queries, dimension)).astype(dtype=np.uint8)
+
+vptree = pynear.BKTreeBinaryIndex()
+vptree.set(data)
+
+# To search using maximum threshold use dimension * 8 (the maximum distance) or set any other threshold
+indices, distances, keys = tree.find_threshold(data, dimensions * 8)
+```
+
 For convenience, apart from `searchKNN` function, vptree also provides `search1NN` for searching the closest nearest neighbor.
 
 #### pynear.VPTreeL2Index

diff --git a/docs/img/binary-index-comparison/result_k=8.png b/docs/img/binary-index-comparison/result_k=8.png
diff --git a/pynear/benchmark/bin_benchmark.yaml b/pynear/benchmark/bin_benchmark.yaml
@@ -6,12 +6,12 @@ benchmark:
   - name: "Binary Index Comparison"
     k: [8]
     num_queries: [16]
-    dimensions: [32, 64, 256]
-    dataset_total_size: 1000000
+    dimensions: [32, 64, 128, 256]
+    dataset_total_size: 200000
     dataset_num_clusters: 50
     dataset_type: "uint8"
     index_types:
     - FaissIndexBinaryFlat
     - AnnoyHamming
     - VPTreeBinaryIndex
-
+    - BKTreeBinaryIndex
diff --git a/pynear/benchmark/index_adapters.py b/pynear/benchmark/index_adapters.py
@@ -19,6 +19,7 @@ def create_index_adapter(index_name: str):
         "AnnoyManhattan": AnnoyManhattanAdapter,
         "AnnoyHamming": AnnoyHammingAdapter,
         "SKLearnL2": SKLearnL2Adapter,
+        "BKTreeBinaryIndex": PyNearBKTreeAdapter,
         "VPTreeL2Index": pynear.VPTreeL2Index,
         "VPTreeL1Index": pynear.VPTreeL1Index,
         "VPTreeBinaryIndex": pynear.VPTreeBinaryIndex,
@@ -28,7 +29,7 @@ def create_index_adapter(index_name: str):
         raise ValueError(f"Index name {index_name} not supported")
 
     if index_name.startswith("VPTree"):
-        return PyNearAdapter(index_name)
+        return PyNearVPAdapter(index_name)
 
     return mapper[index_name]()
 
@@ -52,7 +53,7 @@ def _search_implementation(self, query, k: int):
         pass
 
 
-class PyNearAdapter(IndexAdapter):
+class PyNearVPAdapter(IndexAdapter):
     def __init__(self, pyvp_index_name: str):
         self._index = None
         self._pyvp_index_name = pyvp_index_name
@@ -70,6 +71,17 @@ def build_index(self, data: np.ndarray):
     def _search_implementation(self, query, k: int):
         self._index.searchKNN(query, k)
 
+class PyNearBKTreeAdapter(IndexAdapter):
+    def __init__(self):
+        self._index = pynear.BKTreeBinaryIndex()
+        self._dimensions = 0
+
+    def build_index(self, data: np.ndarray):
+        self._index.set(data)
+        self._dimensions = data.shape[1]
+
+    def _search_implementation(self, query, k: int):
+        self._index.find_threshold(query, self._dimensions)
 
 class FaissIndexFlatL2Adapter(IndexAdapter):
     def __init__(self):