def torch_replacement_range_search(self, x, thresh): if type(x) is np.ndarray: # Forward to faiss __init__.py base method return self.range_search_numpy(x, thresh) assert type(x) is torch.Tensor n, d = x.shape assert d == self.d x_ptr = swig_ptr_from_FloatTensor(x) assert not x.is_cuda, 'Range search using GPU tensor not yet implemented' assert not hasattr( self, 'getDevice'), 'Range search on GPU index not yet implemented' res = faiss.RangeSearchResult(n) self.range_search_c(n, x_ptr, thresh, res) # get pointers and copy them # FIXME: no rev_swig_ptr equivalent for torch.Tensor, just convert # np to torch # NOTE: torch does not support np.uint64, just np.int64 lims = torch.from_numpy( faiss.rev_swig_ptr(res.lims, n + 1).copy().astype('int64')) nd = int(lims[-1]) D = torch.from_numpy(faiss.rev_swig_ptr(res.distances, nd).copy()) I = torch.from_numpy(faiss.rev_swig_ptr(res.labels, nd).copy()) return lims, D, I
def range_search_preassigned(index_ivf, x, radius, list_nos, coarse_dis=None): """ Perform a range search in the IVF index, with predefined lists to search into """ n, d = x.shape if isinstance(index_ivf, faiss.IndexBinaryIVF): d *= 8 dis_type = "int32" else: dis_type = "float32" # the coarse distances are used in IVFPQ with L2 distance and by_residual=True # otherwise we provide dummy coarse_dis if coarse_dis is None: coarse_dis = np.empty((n, index_ivf.nprobe), dtype=dis_type) else: assert coarse_dis.shape == (n, index_ivf.nprobe) assert d == index_ivf.d assert list_nos.shape == (n, index_ivf.nprobe) res = faiss.RangeSearchResult(n) sp = faiss.swig_ptr index_ivf.range_search_preassigned( n, sp(x), radius, sp(list_nos), sp(coarse_dis), res ) # get pointers and copy them lims = faiss.rev_swig_ptr(res.lims, n + 1).copy() num_results = int(lims[-1]) dist = faiss.rev_swig_ptr(res.distances, num_results).copy() indices = faiss.rev_swig_ptr(res.labels, num_results).copy() return lims, dist, indices
def do_test_array_type(self, dtype): """ tests swig_ptr and rev_swig_ptr for this type of array """ a = np.arange(12).astype(dtype) ptr = faiss.swig_ptr(a) print(ptr) a2 = faiss.rev_swig_ptr(ptr, 12) np.testing.assert_array_equal(a, a2)
def test_rev_swig_ptr(self): index = faiss.IndexFlatL2(4) xb0 = np.vstack([ i * 10 + np.array([1, 2, 3, 4], dtype='float32') for i in range(5) ]) index.add(xb0) xb = faiss.rev_swig_ptr(index.xb.data(), 4 * 5).reshape(5, 4) self.assertEqual(np.abs(xb0 - xb).sum(), 0)
def test_rev_swig_ptr(self): index = faiss.IndexFlatL2(4) xb0 = np.vstack([ i * 10 + np.array([1, 2, 3, 4], dtype='float32') for i in range(5)]) index.add(xb0) xb = faiss.rev_swig_ptr(index.xb.data(), 4 * 5).reshape(5, 4) self.assertEqual(np.abs(xb0 - xb).sum(), 0)
def ivf_range_search_preassigned(self, xq, list_nos, coarse_dis, radius): index_ivf = faiss.extract_index_ivf(self.index) n, d = xq.shape assert d == index_ivf.d n2, d2 = list_nos.shape assert list_nos.shape == coarse_dis.shape assert n2 == n assert d2 == index_ivf.nprobe res = faiss.RangeSearchResult(n) index_ivf.range_search_preassigned(n, faiss.swig_ptr(xq), radius, faiss.swig_ptr(list_nos), faiss.swig_ptr(coarse_dis), res) lims = faiss.rev_swig_ptr(res.lims, n + 1).copy() nd = int(lims[-1]) D = faiss.rev_swig_ptr(res.distances, nd).copy() I = faiss.rev_swig_ptr(res.labels, nd).copy() return lims, D, I
def get_ids(invlists, list_no: int): invlists = faiss.downcast_InvertedLists(invlists) ls = invlists.list_size(list_no) list_ids = np.zeros(ls, dtype='int64') ids = None origin_ids = None try: ids = invlists.get_ids(list_no) origin_ids = np.array(faiss.rev_swig_ptr(ids, ls)) except Exception as e: print("get_ids failed ", e) finally: if ids is not None: invlists.release_ids(list_no, ids) return origin_ids