Benchmarking ann-benchmarks bruteforce NN versus sklearn's bruteforce NN implementation

test_bench_brutenn.py

#Python 3.9.0 | packaged by conda-forge | (default, Oct 14 2020, 22:59:50) 
#[GCC 7.5.0] on linux
#sklearn.__version__ 0.23.2
#numpy.__version__ 1.19.2
#pytest.__version__ 6.1.2
#pytest_benchmark.__version__ 3.2.3

from brutenn.bruteforce import BruteForceBLAS
import numpy as np
from numpy.random import default_rng as drng
import pytest
#from pytest_benchmark import benchmark
from sklearn.neighbors import NearestNeighbors

k_neighbors = 10
n_queries = 128

@pytest.fixture
def rndi3d(size=64*1024, dims=12):

    rng = drng(size)
    dataset = rng.integers(-10, 10,
                           (size*dims)).reshape((size, dims))
    return dataset

@pytest.fixture
def rndf3d(size=64*1024, dims=12):

    rng = drng(size)
    dataset = rng.random((size*dims)).reshape((size, dims))
    return dataset


def euclidean_brutenn_query(ds):

    bf = BruteForceBLAS(metric='euclidean')
    bf.fit(ds)

    indices = np.random.choice(ds.shape[0], n_queries)
    value = 0
    for i in indices:
        res = bf.query_with_distances(ds[i], k_neighbors)
        value += len(list(res))

    return value/len(indices)


def test_euclidean_brutenn_64kints(rndi3d, benchmark):
    result = benchmark(euclidean_brutenn_query, rndi3d)
    assert result > 0


def test_euclidean_brutenn_64kfloats(rndf3d, benchmark):
    result = benchmark(euclidean_brutenn_query, rndf3d)
    assert result > 0


def euclidean_sklearn_query(ds):

    bf = NearestNeighbors(algorithm='brute', metric='euclidean')
    bf.fit(ds)

    assert bf

    indices = np.random.choice(ds.shape[0], n_queries)
    neighs = bf.kneighbors(ds[indices, ...], k_neighbors, return_distance=False)

    return len(neighs[0, :])


def test_euclidean_sklearn_64kints(rndi3d, benchmark):
    result = benchmark(euclidean_sklearn_query, rndi3d)
    assert result > 0


def test_euclidean_sklearn_64kfloats(rndf3d, benchmark):
    result = benchmark(euclidean_sklearn_query, rndf3d)
    assert result > 0