def do_partition(n, qin, maxval=65536, seed=123, id_type='int64'):
print(f"n={n} qin={qin} maxval={maxval} id_type={id_type} ",
end="\t",
flush=True)
# print("seed=", seed)
rs = np.random.RandomState(seed)
vals = rs.randint(maxval, size=n).astype('uint16')
ids = (rs.permutation(n) + 12345).astype(id_type)
sp = faiss.swig_ptr
tab_a = faiss.AlignedTableUint16()
faiss.copy_array_to_AlignedTable(vals, tab_a)
nrun = 2000
times = []
nerr = 0
stats = faiss.cvar.partition_stats
stats.reset()
for _run in range(nrun):
faiss.copy_array_to_AlignedTable(vals, tab_a)
t0 = time.time()
# print("tab a type", tab_a.get())
if type(qin) == int:
q = qin
faiss.CMax_uint16_partition_fuzzy(tab_a.get(), sp(ids), n, q, q,
None)
else:
q_min, q_max = qin
q = np.array([-1], dtype='uint64')
thresh2 = faiss.CMax_uint16_partition_fuzzy(
tab_a.get(), sp(ids), n, q_min, q_max, sp(q))
q = q[0]
if not (q_min <= q <= q_max):
nerr += 1
t1 = time.time()
times.append(t1 - t0)
times = np.array(times[100:]) * 1000000
print(f"times {times.mean():.3f} µs (± {times.std():.4f} µs) nerr={nerr} "
f"bissect {stats.bissect_cycles / 1e6:.3f} Mcy "
f"compress {stats.compress_cycles / 1e6:.3f} Mcy")
def do_partition(self, n, q, maxval=65536, seed=None):
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()
faiss.copy_array_to_AlignedTable(vals, tab_a)
# print("tab a type", tab_a.get())
if type(q) == int:
faiss.CMax_uint16_partition_fuzzy(tab_a.get(), sp(ids), n, q, q,
None)
else:
q_min, q_max = q
q = np.array([-1], dtype='uint64')
faiss.CMax_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 = faiss.AlignedTable_to_array(tab_a)
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)