faiss/tests/test_contrib.py
Lucas Hosseini a8e4c5e2d5 Move build to CMake (#1313)
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/1313

Reviewed By: mdouze

Differential Revision: D22948267

Pulled By: beauby

fbshipit-source-id: ec16fa0342f37672d46fb7886ecc55c7996011c4
2020-08-14 15:03:10 -07:00

28 lines
740 B
Python

import faiss
import unittest
import numpy as np
from common import get_dataset_2
from faiss.contrib.exhaustive_search import knn_ground_truth
class TestComputeGT(unittest.TestCase):
def test_compute_GT(self):
d = 64
xt, xb, xq = get_dataset_2(d, 0, 10000, 100)
index = faiss.IndexFlatL2(d)
index.add(xb)
Dref, Iref = index.search(xq, 10)
# iterator function on the matrix
def matrix_iterator(xb, bs):
for i0 in range(0, xb.shape[0], bs):
yield xb[i0:i0 + bs]
Dnew, Inew = knn_ground_truth(xq, matrix_iterator(xb, 1000), 10)
np.testing.assert_array_equal(Iref, Inew)
np.testing.assert_almost_equal(Dref, Dnew, decimal=5)