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)