def do_loop5_kernel(self, nq, bb): """ unit test for the accumulation kernel """ nb = bb * 32 # databse size nsp = 24 # number of sub-quantizers rs = np.random.RandomState(123) codes = rs.randint(256, size=(nb, nsp // 2)).astype('uint8') LUT = rs.randint(256, size=(nq, nsp, 16)).astype('uint8') accu_ref = reference_accu(codes, LUT) def to_A(x): return faiss.array_to_AlignedTable(x.ravel()) sp = faiss.swig_ptr LUT_a = faiss.AlignedTableUint8(LUT.size) faiss.pq4_pack_LUT( nq, nsp, sp(LUT), LUT_a.get() ) codes_a = faiss.AlignedTableUint8(codes.size) faiss.pq4_pack_codes( sp(codes), nb, nsp, nb, nb, nsp, codes_a.get() ) accu_a = faiss.AlignedTableUint16(nq * nb) accu_a.clear() faiss.loop5_kernel_accumulate_1_block_to_mem( nq, nb, nsp, codes_a.get(), LUT_a.get(), accu_a.get() ) accu = faiss.AlignedTable_to_array(accu_a).reshape(nq, nb) np.testing.assert_array_equal(accu_ref, accu)
def do_test_add(self, d, bbs): ds = datasets.SyntheticDataset(d, 2000, 5000, 200) index = faiss.index_factory(d, f'PQ{d//2}x4np') index.train(ds.get_train()) xb = ds.get_database() index.add(xb[:1235]) index2 = faiss.IndexPQFastScan(index, bbs) index2.add(xb[1235:]) new_codes = faiss.AlignedTable_to_array(index2.codes) index.add(xb[1235:]) index3 = faiss.IndexPQFastScan(index, bbs) ref_codes = faiss.AlignedTable_to_array(index3.codes) self.assertEqual(index3.ntotal, index2.ntotal) np.testing.assert_array_equal(ref_codes, new_codes)
def do_partition(self, n, q, maxval=65536, seed=None): #seed = 1235 if seed is None: for i in range(50): self.do_partition(n, q, maxval, i + 1234) # print("seed=", seed) rs = np.random.RandomState(seed) vals = rs.randint(maxval, size=n).astype('uint16') ids = (rs.permutation(n) + 12345).astype('int64') dic = dict(zip(ids, vals)) sp = faiss.swig_ptr vals_orig = vals.copy() tab_a = faiss.AlignedTableUint16() vals_inv = (65535 - vals).astype('uint16') faiss.copy_array_to_AlignedTable(vals_inv, tab_a) # print("tab a type", tab_a.get()) if type(q) == int: faiss.CMin_uint16_partition_fuzzy(tab_a.get(), sp(ids), n, q, q, None) else: q_min, q_max = q q = np.array([-1], dtype='uint64') thresh2 = faiss.CMin_uint16_partition_fuzzy( tab_a.get(), sp(ids), n, q_min, q_max, sp(q)) q = q[0] assert q_min <= q <= q_max vals_inv = faiss.AlignedTable_to_array(tab_a) vals = 65535 - vals_inv o = vals_orig.argsort() thresh = vals_orig[o[q]] n_eq = (vals_orig[o[:q]] == thresh).sum() for i in range(q): self.assertEqual(vals[i], dic[ids[i]]) self.assertLessEqual(vals[i], thresh) if vals[i] == thresh: n_eq -= 1 self.assertEqual(n_eq, 0)