def distance_table(self, q):
q = pad(q, (2 * self.dims_per_block, ))
dpb = self.dims_per_block
n_blocks = q.size / dpb
parts = q.reshape(-1, dpb)
# Center the data in the range [-128, 128]
# TODO: Is this the best scaling formula?
# TODO: Would this be faster if we used the same centers for each block?
dists = self.center_norms_sq - 2 * np.einsum("ijk,ik->ij",
self.centers, parts)
# dists += (parts * parts).sum(axis=1, keepdims=True)
# shift = np.mean(dists)
# print(np.mean(dists), np.median(dists))
# shift = 1
# shift = 128 / n_blocks
# scale = 128 / (np.max(-(dists-shift)) * np.sqrt(n_blocks))
# shift = np.mean(dists) / 2
shift = np.median(dists)
# shift = 0
# scale = 1
scale = 128 / (np.max(np.abs(dists - shift)) * np.sqrt(n_blocks))
table = np.round(
(dists - shift) * scale) # Round to nearest integer towards zero.
# The transformation doesn't care about the sign, so we just use uint
table = table.astype(np.uint8)
trans = transform_tables(table)
return _FastDistanceTable(q, trans, shift, scale, signed=True)
def _slow_pq(data0, tables0):
data = transform_data(data0)
tables = transform_tables(tables0)
out = np.zeros(2 * len(data), dtype=np.uint64)
query_pq_sse(data, tables, out)
res = out.view(np.uint8)
assert res.shape[0] == data0.shape[0]
return res
def udistance_table(self, q):
# Experimental
q = pad(q, (2 * self.dims_per_block, ))
dpb = self.dims_per_block
n_blocks = q.size / dpb
parts = q.reshape(-1, dpb)
dists = self.center_norms_sq - 2 * np.einsum("ijk,ik->ij",
self.centers, parts)
shift = np.min(dists)
dists -= shift
scale = 255 / (np.max(dists) * np.sqrt(n_blocks))
table *= scale
table = table.astype(np.uint8)
trans = transform_tables(table)
return _FastDistanceTable(q, trans, shift, scale, signed=False)
def transform_query(self, q):
dpb = self.dims_per_block
parts = q.reshape(-1, dpb)
# TODO: This can probably be speed up by a better memory layout
ips = np.einsum("ijk,ik->ij", self.centers, parts)
norms = np.linalg.norm(parts, axis=1, keepdims=True) ** 2
table = self.norms - 2 * ips + norms
# TODO: Try if median or something would work in place of max here
blocks = q.size / dpb
# scale = np.mean(table) / 255 * blocks
scale = np.max(table) / 255 * np.sqrt(blocks)
table = np.floor(table / scale).astype(np.uint8)
trans = transform.transform_tables(table)
return trans, scale