TomDLT
commented
3 years ago From the documentation:
- Brute force query time grows as O(DN)
- Ball tree query time grows as approximately O(DlogN)
- KD tree query time changes with D in a way that is difficult to precisely characterise. For small D (less than 20 or so) the cost is approximately O(DlogN), and the KD tree query can be very efficient. For larger D, the cost increases to nearly O(DN), and the overhead due to the tree structure can lead to queries which are slower than brute force.
Playing a bit with it, it seems that the description of nearly O(DN) for large D is also valid for Ball tree, which seems a bit weird to me.
For D = 5
For D = 50
For comparison, annoy leads to the expected O(logN):
I don't know if this is expected behavior (then we should update the doc), or if it is a bug (then we should fix it). Here is the script I used for reference:
```py
import time
import numpy as np
import matplotlib.pyplot as plt
from sklearn.neighbors import KNeighborsTransformer
import annoy
from sklearn.base import BaseEstimator, TransformerMixin
from scipy.sparse import csr_matrix
class AnnoyTransformer(TransformerMixin, BaseEstimator):
"""Wrapper for using annoy.AnnoyIndex as sklearn's KNeighborsTransformer"""
def __init__(self, n_neighbors=5, metric='euclidean', n_trees=10,
search_k=-1):
self.n_neighbors = n_neighbors
self.n_trees = n_trees
self.search_k = search_k
self.metric = metric
def fit(self, X):
self.n_samples_fit_ = X.shape[0]
metric = self.metric if self.metric != 'sqeuclidean' else 'euclidean'
self.annoy_ = annoy.AnnoyIndex(X.shape[1], metric=metric)
for i, x in enumerate(X):
self.annoy_.add_item(i, x.tolist())
self.annoy_.build(self.n_trees)
return self
def transform(self, X):
return self._transform(X)
def fit_transform(self, X, y=None):
return self.fit(X)._transform(X=None)
def _transform(self, X):
"""As `transform`, but handles X is None for faster `fit_transform`."""
n_samples_transform = self.n_samples_fit_ if X is None else X.shape[0]
# For compatibility reasons, as each sample is considered as its own
# neighbor, one extra neighbor will be computed.
n_neighbors = self.n_neighbors + 1
indices = np.empty((n_samples_transform, n_neighbors), dtype=int)
distances = np.empty((n_samples_transform, n_neighbors))
if X is None:
for i in range(self.annoy_.get_n_items()):
ind, dist = self.annoy_.get_nns_by_item(
i, n_neighbors, self.search_k, include_distances=True)
indices[i], distances[i] = ind, dist
else:
for i, x in enumerate(X):
indices[i], distances[i] = self.annoy_.get_nns_by_vector(
x.tolist(), n_neighbors, self.search_k,
include_distances=True)
if self.metric == 'sqeuclidean':
distances **= 2
indptr = np.arange(0, n_samples_transform * n_neighbors + 1,
n_neighbors)
kneighbors_graph = csr_matrix(
(distances.ravel(), indices.ravel(), indptr),
shape=(n_samples_transform, self.n_samples_fit_))
return kneighbors_graph
n_neighbors = 5
n_samples_query = 1
n_features = 5
for model in [
KNeighborsTransformer(n_neighbors=n_neighbors, algorithm='ball_tree'),
KNeighborsTransformer(n_neighbors=n_neighbors, algorithm='kd_tree'),
KNeighborsTransformer(n_neighbors=n_neighbors, algorithm='brute'),
AnnoyTransformer(n_neighbors=n_neighbors),
]:
def bench(n_samples_fit):
X = np.random.rand(n_samples_fit, n_features)
Y = np.random.rand(n_samples_query, n_features)
start = time.time()
kdt = model.fit(X)
fit_time = time.time() - start
start = time.time()
_ = kdt.transform(Y)
query_time = time.time() - start
return n_samples_fit, fit_time, query_time
results = []
for n_samples_fit in np.int_(np.logspace(2, 5, 10)):
results.append(bench(n_samples_fit))
results = np.array(results)
plt.loglog(results[:, 0], results[:, 1], 'o-', label='fit')
plt.loglog(results[:, 0], results[:, 2], 'o-', label='query')
plt.loglog(results[:, 0], results[:, 0] / 1e5, '--', label='linear')
plt.legend()
plt.ylabel('Time (sec)')
plt.xlabel('n_samples_fit')
plt.title(str(model))
plt.show()
```
simsim314
jnothman
thomasjpfan
gkaranikas
ogrisel
I've run performance analysis on matching NN with BallTree (same with KDTree), and the matching time is linear to number of elements, and should be O(log(n)).
Here are the result of benchmark: num_elements, match_time 10000 0.09097146987915039 20000 0.18194293975830078 40000 0.3668830394744873 80000 0.7527577877044678
Here is my code: