# Copyright (c) Meta Platforms, Inc. and affiliates. # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. """a few tests for graph-based indices (HNSW, nndescent and NSG)""" import numpy as np import unittest import faiss from common_faiss_tests import get_dataset_2, for_all_simd_levels @for_all_simd_levels class TestHNSW(unittest.TestCase): def __init__(self, *args, **kwargs): unittest.TestCase.__init__(self, *args, **kwargs) d = 32 nt = 0 nb = 1500 nq = 500 (_, self.xb, self.xq) = get_dataset_2(d, nt, nb, nq) index = faiss.IndexFlatL2(d) index.add(self.xb) Dref, Iref = index.search(self.xq, 1) self.Iref = Iref def test_hnsw(self): d = self.xq.shape[1] index = faiss.IndexHNSWFlat(d, 16) index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 460) self.io_and_retest(index, Dhnsw, Ihnsw) def test_deterministic_build(self): # The HNSW build is deterministic and lock-free: repeated builds, and # builds at different thread counts, produce byte-identical graphs. # nb must be large enough that the batch-size cap (2% of n) exceeds # the >100 threshold gating the parallel build region; otherwise the # build runs single-threaded and determinism is trivially untested. d = 32 nb = 20000 xb = np.random.RandomState(123).random((nb, d)).astype("float32") # build() mutates the process-global OpenMP thread count; restore it so # it does not leak into subsequent tests. self.addCleanup(faiss.omp_set_num_threads, faiss.omp_get_max_threads()) def build(nthreads): faiss.omp_set_num_threads(nthreads) index = faiss.IndexHNSWFlat(d, 16) index.add(xb) return index # `index.hnsw` returns a fresh (copied) proxy and `.neighbors` / # `.offsets` are views into it, so chaining `index.hnsw.neighbors` # inline reads a freed temporary. Bind the hnsw proxy to keep it alive # across the (copying) vector_to_array calls. def graph_of(index): hnsw = index.hnsw return ( hnsw.max_level, faiss.vector_to_array(hnsw.neighbors), faiss.vector_to_array(hnsw.offsets), ) idx8 = build(8) max_level, g8a, offs = graph_of(idx8) # graph is non-trivial: multi-level, many real edges self.assertGreaterEqual(max_level, 1) self.assertGreater(int((g8a >= 0).sum()), nb) # reproducible across repeated builds and across thread counts g8b = graph_of(build(8))[1] g1 = graph_of(build(1))[1] np.testing.assert_array_equal(g8a, g8b) np.testing.assert_array_equal(g8a, g1) # no self-loops for i in range(nb): seg = g8a[int(offs[i]) : int(offs[i + 1])] self.assertNotIn(i, seg[seg >= 0].tolist()) def test_range_search(self): index_flat = faiss.IndexFlat(self.xb.shape[1]) index_flat.add(self.xb) D, _ = index_flat.search(self.xq, 10) radius = np.median(D[:, -1]) lims_ref, Dref, Iref = index_flat.range_search(self.xq, radius) index = faiss.IndexHNSWFlat(self.xb.shape[1], 16) index.add(self.xb) lims, D, I = index.range_search(self.xq, radius) nmiss = 0 # check if returned results are a subset of the reference results for i in range(len(self.xq)): ref = Iref[lims_ref[i] : lims_ref[i + 1]] new = I[lims[i] : lims[i + 1]] self.assertLessEqual(set(new), set(ref)) nmiss += len(ref) - len(new) # currently we miss 405 / 6019 neighbors self.assertLessEqual(nmiss, lims_ref[-1] * 0.1) def test_hnsw_unbounded_queue(self): d = self.xq.shape[1] index = faiss.IndexHNSWFlat(d, 16) index.add(self.xb) index.hnsw.search_bounded_queue = False Dhnsw, Ihnsw = index.search(self.xq, 1) self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 460) self.io_and_retest(index, Dhnsw, Ihnsw) def test_hnsw_no_init_level0(self): d = self.xq.shape[1] index = faiss.IndexHNSWFlat(d, 16) index.init_level0 = False index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) # This is expected to be smaller because we are not initializing # vectors into level 0. self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 25) self.io_and_retest(index, Dhnsw, Ihnsw) def io_and_retest(self, index, Dhnsw, Ihnsw): index2 = faiss.deserialize_index(faiss.serialize_index(index)) Dhnsw2, Ihnsw2 = index2.search(self.xq, 1) self.assertTrue(np.all(Dhnsw2 == Dhnsw)) self.assertTrue(np.all(Ihnsw2 == Ihnsw)) # Verify deserialized index is serializable again index2b = faiss.deserialize_index(faiss.serialize_index(index2)) Dhnsw2b, Ihnsw2b = index2b.search(self.xq, 1) self.assertTrue(np.all(Dhnsw2b == Dhnsw)) self.assertTrue(np.all(Ihnsw2b == Ihnsw)) # also test clone index3 = faiss.clone_index(index) Dhnsw3, Ihnsw3 = index3.search(self.xq, 1) self.assertTrue(np.all(Dhnsw3 == Dhnsw)) self.assertTrue(np.all(Ihnsw3 == Ihnsw)) def test_hnsw_2level(self): d = self.xq.shape[1] quant = faiss.IndexFlatL2(d) index = faiss.IndexHNSW2Level(quant, 256, 8, 8) index.train(self.xb) index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 307) self.io_and_retest(index, Dhnsw, Ihnsw) def test_hnsw_2level_mixed_search(self): d = self.xq.shape[1] quant = faiss.IndexFlatL2(d) storage = faiss.IndexIVFPQ(quant, d, 32, 8, 8) storage.make_direct_map() index = faiss.IndexHNSW2Level(quant, 32, 8, 8) index.storage = storage index.train(self.xb) index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) # It is expected that the mixed search will perform worse. self.assertGreaterEqual((self.Iref == Ihnsw).sum(), 200) self.io_and_retest(index, Dhnsw, Ihnsw) def test_add_0_vecs(self): index = faiss.IndexHNSWFlat(10, 16) zero_vecs = np.zeros((0, 10), dtype="float32") # infinite loop index.add(zero_vecs) self.assertEqual(index.ntotal, 0) def test_hnsw_IP(self): d = self.xq.shape[1] index_IP = faiss.IndexFlatIP(d) index_IP.add(self.xb) Dref, Iref = index_IP.search(self.xq, 1) index = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT) index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) self.assertGreaterEqual((Iref == Ihnsw).sum(), 470) mask = Iref[:, 0] == Ihnsw[:, 0] self.assertTrue(np.allclose(Dref[mask, 0], Dhnsw[mask, 0])) def test_ndis_stats(self): d = self.xq.shape[1] index = faiss.IndexHNSWFlat(d, 16) index.add(self.xb) stats = faiss.cvar.hnsw_stats self.addCleanup( faiss.set_search_stats_enabled, faiss.get_search_stats_enabled() ) faiss.set_search_stats_enabled(True) stats.reset() Dhnsw, Ihnsw = index.search(self.xq, 1) self.assertGreater(stats.ndis, len(self.xq) * index.hnsw.efSearch) def test_io_no_storage(self): d = self.xq.shape[1] index = faiss.IndexHNSWFlat(d, 16) index.add(self.xb) Dref, Iref = index.search(self.xq, 5) # test writing without storage index2 = faiss.deserialize_index( faiss.serialize_index(index, faiss.IO_FLAG_SKIP_STORAGE) ) self.assertEqual(index2.storage, None) self.assertRaises(RuntimeError, index2.search, self.xb, 1) # make sure we can store an index with empty storage faiss.deserialize_index(faiss.serialize_index(index2)) # add storage afterwards index.storage = faiss.clone_index(index.storage) index.own_fields = True Dnew, Inew = index.search(self.xq, 5) np.testing.assert_array_equal(Dnew, Dref) np.testing.assert_array_equal(Inew, Iref) def test_hnsw_reset(self): d = self.xb.shape[1] index_flat = faiss.IndexFlat(d) index_flat.add(self.xb) self.assertEqual(index_flat.ntotal, self.xb.shape[0]) index_hnsw = faiss.IndexHNSW(index_flat) index_hnsw.add(self.xb) # * 2 because we add to storage twice. This is just for testing # that storage gets cleared correctly. self.assertEqual(index_hnsw.ntotal, self.xb.shape[0] * 2) index_hnsw.reset() self.assertEqual(index_flat.ntotal, 0) self.assertEqual(index_hnsw.ntotal, 0) class TestHNSWSimilarity(unittest.TestCase): """Coverage for HNSW variants built with a similarity metric (METRIC_INNER_PRODUCT). These tests exercise the HNSW::is_similarity dispatch in IndexHNSW{Flat,PQ,SQ,Cagra}: build, search, and (for IndexHNSWFlat) save/load round trips that have to re-derive `hnsw.is_similarity` from the persisted metric type. """ def setUp(self): d = 32 nt = 1500 nb = 1500 nq = 200 (self.xt, self.xb, self.xq) = get_dataset_2(d, nt, nb, nq) # Reference IP search results. ref = faiss.IndexFlatIP(d) ref.add(self.xb) self.Dref_ip, self.Iref_ip = ref.search(self.xq, 1) def test_hnsw_flat_IP_is_similarity_set(self): d = self.xb.shape[1] index = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT) self.assertTrue(index.hnsw.is_similarity) index_l2 = faiss.IndexHNSWFlat(d, 16, faiss.METRIC_L2) self.assertFalse(index_l2.hnsw.is_similarity) def _check_quantized_IP(self, index, min_agreement): """Build (train if needed), search, then save / load and assert that the round trip preserves ordering and `hnsw.is_similarity`. The agreement floors are calibrated against the observed values on the test dataset with a small (~5%) headroom for RNG noise; they catch ordering regressions (which collapse agreement to ~0) and quantizer-quality regressions (which would shift agreement by tens of points).""" self.assertTrue(index.hnsw.is_similarity) if not index.is_trained: index.train(self.xt) index.add(self.xb) Dhnsw, Ihnsw = index.search(self.xq, 1) agreement = (self.Iref_ip[:, 0] == Ihnsw[:, 0]).sum() self.assertGreaterEqual( agreement, min_agreement, f"{type(index).__name__} agreement={agreement} < {min_agreement}", ) index2 = faiss.deserialize_index(faiss.serialize_index(index)) # `hnsw.is_similarity` is not serialized; index_read must # re-derive it from the persisted metric_type. self.assertTrue(index2.hnsw.is_similarity) self.assertEqual(index2.metric_type, faiss.METRIC_INNER_PRODUCT) D2, I2 = index2.search(self.xq, 1) np.testing.assert_array_equal(I2, Ihnsw) np.testing.assert_array_equal(D2, Dhnsw) # Agreement floors below are calibrated against measured values # (Flat: 199, PQ: 26-30, SQ: 195, Cagra: 200 over nq=200) with a # small headroom for RNG noise. def test_hnsw_flat_IP(self): d = self.xb.shape[1] self._check_quantized_IP( faiss.IndexHNSWFlat(d, 16, faiss.METRIC_INNER_PRODUCT), # observed ~199/200; headroom for OpenMP-parallel-build # nondeterminism min_agreement=180, ) def test_hnsw_pq_IP(self): d = self.xb.shape[1] self._check_quantized_IP( faiss.IndexHNSWPQ(d, 8, 16, 8, faiss.METRIC_INNER_PRODUCT), min_agreement=20, ) def test_hnsw_sq_IP(self): d = self.xb.shape[1] self._check_quantized_IP( faiss.IndexHNSWSQ( d, faiss.ScalarQuantizer.QT_8bit, 16, faiss.METRIC_INNER_PRODUCT, ), min_agreement=190, ) def test_hnsw_cagra_IP(self): d = self.xb.shape[1] self._check_quantized_IP( faiss.IndexHNSWCagra(d, 16, faiss.METRIC_INNER_PRODUCT), min_agreement=195, ) def test_hnsw_cagra_IP_base_level_only(self): """IndexHNSWCagra.base_level_only takes a separate code path (random entry-point selection + search_level_0) that also needs to dispatch on is_similarity.""" d = self.xb.shape[1] index = faiss.IndexHNSWCagra(d, 16, faiss.METRIC_INNER_PRODUCT) index.add(self.xb) index.base_level_only = True index.num_base_level_search_entrypoints = 64 _, Ihnsw = index.search(self.xq, 1) agreement = (self.Iref_ip[:, 0] == Ihnsw[:, 0]).sum() # Measured value on this dataset is 200/200; leave headroom for # the RNG-driven entry-point selection (this path is unseeded). self.assertGreaterEqual( agreement, 190, f"Cagra base_only agreement={agreement}", ) class TestHNSWNaN(unittest.TestCase): """Adding a vector with NaN to an IVF+HNSW index used to crash because NaN distances corrupt the MinimaxHeap ordering in HNSW search. The fix converts NaN to +inf in MinimaxHeap::push so the heap stays well-ordered. """ def test_add_nan_vector_to_ivf_hnsw(self): d = 64 nt = 2000 nb = 1000 xt = np.random.default_rng(42).random((nt, d), dtype="float32") xb = np.random.default_rng(43).random((nb, d), dtype="float32") index = faiss.index_factory(d, "IVF256_HNSW32,SQ8") index.train(xt) index.add(xb) # Create a vector with NaN in the first component vec = np.zeros((1, d), dtype="float32") vec[0, 0] = np.nan # This should not crash index.add(vec) self.assertEqual(index.ntotal, nb + 1) class Issue3684(unittest.TestCase): def test_issue3684(self): np.random.seed(1234) # For reproducibility d = 256 # Example dimension nb = 10 # Number of database vectors nq = 2 # Number of query vectors xb = np.random.random((nb, d)).astype("float32") xq = np.random.random((nq, d)).astype("float32") faiss.normalize_L2(xb) # Normalize both query and database vectors faiss.normalize_L2(xq) hnsw_index_ip = faiss.IndexHNSWFlat(256, 16, faiss.METRIC_INNER_PRODUCT) hnsw_index_ip.hnsw.efConstruction = 512 hnsw_index_ip.hnsw.efSearch = 512 hnsw_index_ip.add(xb) # test knn D, I = hnsw_index_ip.search(xq, 10) self.assertTrue(np.all(D[:, :-1] >= D[:, 1:])) # test range search radius = 0.74 # Cosine similarity threshold lims, D, I = hnsw_index_ip.range_search(xq, radius) self.assertTrue(np.all(D >= radius)) class TestNSG(unittest.TestCase): def __init__(self, *args, **kwargs): unittest.TestCase.__init__(self, *args, **kwargs) d = 32 nt = 0 nb = 1500 nq = 500 self.GK = 32 _, self.xb, self.xq = get_dataset_2(d, nt, nb, nq) def make_knn_graph(self, metric): n = self.xb.shape[0] d = self.xb.shape[1] index = faiss.IndexFlat(d, metric) index.add(self.xb) _, I = index.search(self.xb, self.GK + 1) knn_graph = np.zeros((n, self.GK), dtype=np.int64) # For the inner product distance, the distance between a vector and # itself may not be the smallest, so it is not guaranteed that I[:, 0] # is the query itself. for i in range(n): cnt = 0 for j in range(self.GK + 1): if I[i, j] != i: knn_graph[i, cnt] = I[i, j] cnt += 1 if cnt == self.GK: break return knn_graph def subtest_io_and_clone(self, index, Dnsg, Insg): index2 = faiss.deserialize_index(faiss.serialize_index(index)) Dnsg2, Insg2 = index2.search(self.xq, 1) np.testing.assert_array_equal(Dnsg2, Dnsg) np.testing.assert_array_equal(Insg2, Insg) # Verify deserialized index is serializable again index2b = faiss.deserialize_index(faiss.serialize_index(index2)) Dnsg2b, Insg2b = index2b.search(self.xq, 1) np.testing.assert_array_equal(Dnsg2b, Dnsg) np.testing.assert_array_equal(Insg2b, Insg) # also test clone index3 = faiss.clone_index(index) Dnsg3, Insg3 = index3.search(self.xq, 1) np.testing.assert_array_equal(Dnsg3, Dnsg) np.testing.assert_array_equal(Insg3, Insg) def subtest_connectivity(self, index, nb): vt = faiss.VisitedTableVector(nb) count = index.nsg.dfs(vt, index.nsg.enterpoint, 0) self.assertEqual(count, nb) def subtest_add(self, build_type, thresh, metric=faiss.METRIC_L2): d = self.xq.shape[1] flat_index = faiss.IndexFlat(d, metric) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, 1) index = faiss.IndexNSGFlat(d, 16, metric) index.verbose = True index.build_type = build_type index.GK = self.GK index.add(self.xb) Dnsg, Insg = index.search(self.xq, 1) recalls = (Iref == Insg).sum() self.assertGreaterEqual(recalls, thresh) self.subtest_connectivity(index, self.xb.shape[0]) self.subtest_io_and_clone(index, Dnsg, Insg) def subtest_build(self, knn_graph, thresh, metric=faiss.METRIC_L2): d = self.xq.shape[1] flat_index = faiss.IndexFlat(d, metric) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, 1) index = faiss.IndexNSGFlat(d, 16, metric) index.verbose = True index.build(self.xb, knn_graph) Dnsg, Insg = index.search(self.xq, 1) recalls = (Iref == Insg).sum() self.assertGreaterEqual(recalls, thresh) self.subtest_connectivity(index, self.xb.shape[0]) def test_add_bruteforce_L2(self): self.subtest_add(0, 475, faiss.METRIC_L2) def test_add_nndescent_L2(self): self.subtest_add(1, 475, faiss.METRIC_L2) def test_add_bruteforce_IP(self): self.subtest_add(0, 480, faiss.METRIC_INNER_PRODUCT) def test_add_nndescent_IP(self): self.subtest_add(1, 480, faiss.METRIC_INNER_PRODUCT) def test_build_L2(self): knn_graph = self.make_knn_graph(faiss.METRIC_L2) self.subtest_build(knn_graph, 475, faiss.METRIC_L2) def test_build_IP(self): knn_graph = self.make_knn_graph(faiss.METRIC_INNER_PRODUCT) self.subtest_build(knn_graph, 480, faiss.METRIC_INNER_PRODUCT) def test_build_invalid_knng(self): """Make some invalid entries in the input knn graph. It would cause a warning but IndexNSG should be able to handle this. """ knn_graph = self.make_knn_graph(faiss.METRIC_L2) knn_graph[:100, 5] = -111 self.subtest_build(knn_graph, 475, faiss.METRIC_L2) knn_graph = self.make_knn_graph(faiss.METRIC_INNER_PRODUCT) knn_graph[:100, 5] = -111 self.subtest_build(knn_graph, 480, faiss.METRIC_INNER_PRODUCT) def test_reset(self): """test IndexNSG.reset()""" d = self.xq.shape[1] metric = faiss.METRIC_L2 flat_index = faiss.IndexFlat(d, metric) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, 1) index = faiss.IndexNSGFlat(d, 16) index.verbose = True index.GK = 32 index.add(self.xb) Dnsg, Insg = index.search(self.xq, 1) recalls = (Iref == Insg).sum() self.assertGreaterEqual(recalls, 475) self.subtest_connectivity(index, self.xb.shape[0]) index.reset() index.add(self.xb) Dnsg, Insg = index.search(self.xq, 1) recalls = (Iref == Insg).sum() self.assertGreaterEqual(recalls, 475) self.subtest_connectivity(index, self.xb.shape[0]) def test_order(self): """make sure that output results are sorted""" d = self.xq.shape[1] index = faiss.IndexNSGFlat(d, 32) index.train(self.xb) index.add(self.xb) k = 10 nq = self.xq.shape[0] D, _ = index.search(self.xq, k) indices = np.argsort(D, axis=1) gt = np.arange(0, k)[np.newaxis, :] # [1, k] gt = np.repeat(gt, nq, axis=0) # [nq, k] np.testing.assert_array_equal(indices, gt) def test_nsg_supports_pre_built_knn_graph(self): """Test IndexNSGBuild""" knn_graph = self.make_knn_graph(faiss.METRIC_L2) d = self.xq.shape[1] index = faiss.IndexNSGFlat(d, 16) index.build(self.xb, knn_graph) index.search(self.xq, k=1) self.assertTrue(index.is_built) def test_nsg_with_pre_built_knn_graph_throws_when_rebuilding_via_add(self): """Test IndexNSGBuild""" knn_graph = self.make_knn_graph(faiss.METRIC_L2) d = self.xq.shape[1] index = faiss.IndexNSGFlat(d, 16) index.build(self.xb, knn_graph) index.search(self.xq, k=1) self.assertTrue(index.is_built) index.GK = 32 index.train(self.xb) with self.assertRaises(RuntimeError) as context: index.add(self.xb) self.assertIn( "NSG does not support incremental addition", str(context.exception) ) def test_nsg_rebuild_throws_with_pre_built_knn_graph(self): """Test IndexNSGBuild""" knn_graph = self.make_knn_graph(faiss.METRIC_L2) d = self.xq.shape[1] index = faiss.IndexNSGFlat(d, 16) index.build(self.xb, knn_graph) index.search(self.xq, k=1) with self.assertRaises(RuntimeError) as context: index.build(self.xb, knn_graph) self.assertIn("The IndexNSG is already built", str(context.exception)) def test_nsg_pq(self): """Test IndexNSGPQ""" d = self.xq.shape[1] R, pq_M = 32, 4 index = faiss.index_factory(d, f"NSG{R}_PQ{pq_M}np") self.assertIsInstance(index, faiss.IndexNSGPQ) idxpq = faiss.downcast_index(index.storage) self.assertEqual(index.nsg.R, R) self.assertEqual(idxpq.pq.M, pq_M) flat_index = faiss.IndexFlat(d) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, k=1) index.GK = 32 index.train(self.xb) index.add(self.xb) D, I = index.search(self.xq, k=1) # test accuracy recalls = (Iref == I).sum() # observed ~193/500; PQ4 is coarse -- headroom for OpenMP # FP-reduction / tie-break nondeterminism in NSG build + PQ assign self.assertGreaterEqual(recalls, 180) # test I/O self.subtest_io_and_clone(index, D, I) def test_nsg_sq(self): """Test IndexNSGSQ""" d = self.xq.shape[1] R = 32 index = faiss.index_factory(d, f"NSG{R}_SQ8") self.assertIsInstance(index, faiss.IndexNSGSQ) idxsq = faiss.downcast_index(index.storage) self.assertEqual(index.nsg.R, R) self.assertEqual(idxsq.sq.qtype, faiss.ScalarQuantizer.QT_8bit) flat_index = faiss.IndexFlat(d) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, k=1) index.train(self.xb) index.add(self.xb) D, I = index.search(self.xq, k=1) # test accuracy recalls = (Iref == I).sum() # nominal ~411/500; loosened for OpenMP graph-build nondeterminism self.assertGreaterEqual(recalls, 395) # test I/O self.subtest_io_and_clone(index, D, I) class TestNNDescent(unittest.TestCase): def __init__(self, *args, **kwargs): unittest.TestCase.__init__(self, *args, **kwargs) d = 32 nt = 0 nb = 1500 nq = 500 self.GK = 32 _, self.xb, self.xq = get_dataset_2(d, nt, nb, nq) def test_nndescentflat(self): d = self.xq.shape[1] index = faiss.IndexNNDescentFlat(d, 32) index.nndescent.search_L = 8 flat_index = faiss.IndexFlat(d) flat_index.add(self.xb) Dref, Iref = flat_index.search(self.xq, k=1) index.train(self.xb) index.add(self.xb) D, I = index.search(self.xq, k=1) # test accuracy recalls = (Iref == I).sum() self.assertGreaterEqual(recalls, 450) # 462 # do some IO tests index2 = faiss.deserialize_index(faiss.serialize_index(index)) D2, I2 = index2.search(self.xq, 1) np.testing.assert_array_equal(D2, D) np.testing.assert_array_equal(I2, I) # Verify deserialized index is serializable again index2b = faiss.deserialize_index(faiss.serialize_index(index2)) D2b, I2b = index2b.search(self.xq, 1) np.testing.assert_array_equal(D2b, D) np.testing.assert_array_equal(I2b, I) # also test clone index3 = faiss.clone_index(index) D3, I3 = index3.search(self.xq, 1) np.testing.assert_array_equal(D3, D) np.testing.assert_array_equal(I3, I) def test_order(self): """make sure that output results are sorted""" d = self.xq.shape[1] index = faiss.IndexNNDescentFlat(d, 32) index.train(self.xb) index.add(self.xb) k = 10 nq = self.xq.shape[0] D, _ = index.search(self.xq, k) indices = np.argsort(D, axis=1) gt = np.arange(0, k)[np.newaxis, :] # [1, k] gt = np.repeat(gt, nq, axis=0) # [nq, k] np.testing.assert_array_equal(indices, gt) class TestNNDescentGenRandom(unittest.TestCase): """Regression tests for gen_random edge cases in NNDescent.""" def test_search_L_equals_ntotal(self): """gen_random(size, N) crashed with division by zero when size == N. In search(), L_2 = max(search_L, topk). When search_L >= ntotal, gen_random is called with size == N, causing rng() % 0. """ d = 32 nb = 200 # just above NUM_EVAL_POINTS=100 xb = np.random.default_rng(42).random((nb, d)).astype("float32") xq = np.random.default_rng(43).random((10, d)).astype("float32") index = faiss.IndexNNDescentFlat(d, 32) index.nndescent.search_L = nb # triggers gen_random(size=nb, N=nb) index.train(xb) index.add(xb) # This crashed with division by zero before the fix D, I = index.search(xq, k=1) self.assertEqual(I.shape, (xq.shape[0], 1)) self.assertEqual(D.shape, (xq.shape[0], 1)) class TestNNDescentKNNG(unittest.TestCase): def test_knng_L2(self): self.subtest(32, 10, faiss.METRIC_L2) def test_knng_IP(self): self.subtest(32, 10, faiss.METRIC_INNER_PRODUCT) def subtest(self, d, K, metric): nb = 1000 _, xb, _ = get_dataset_2(d, 0, nb, 0) _, knn = faiss.knn(xb, xb, K + 1, metric) knn = knn[:, 1:] index = faiss.IndexNNDescentFlat(d, K, metric) index.nndescent.S = 10 index.nndescent.R = 32 index.nndescent.L = K + 20 index.nndescent.iter = 5 index.verbose = True index.add(xb) graph = index.nndescent.final_graph graph = faiss.vector_to_array(graph) graph = graph.reshape(nb, K) recalls = 0 for i in range(nb): for j in range(K): for k in range(K): if graph[i, j] == knn[i, k]: recalls += 1 break recall = 1.0 * recalls / (nb * K) self.assertGreater(recall, 0.99) def test_small_nndescent(self): """building a too small graph used to crash, make sure it raises an exception instead. TODO: build the exact knn graph for small cases """ d = 32 K = 10 index = faiss.IndexNNDescentFlat(d, K, faiss.METRIC_L2) index.nndescent.S = 10 index.nndescent.R = 32 index.nndescent.L = K + 20 index.nndescent.iter = 5 index.verbose = True xb = np.zeros((78, d), dtype="float32") self.assertRaises(RuntimeError, index.add, xb)