def initialized_nnd_search(
data, indptr, indices, initialization, query_points, dist, dist_args
):
for i in numba.prange(query_points.shape[0]):
tried = set(initialization[0, i])
while True:
# Find smallest flagged vertex
vertex = smallest_flagged(initialization, i)
if vertex == -1:
break
candidates = indices[indptr[vertex] : indptr[vertex + 1]]
for j in range(candidates.shape[0]):
if (
candidates[j] == vertex
or candidates[j] == -1
or candidates[j] in tried
):
continue
d = dist(data[candidates[j]], query_points[i], *dist_args)
unchecked_heap_push(initialization, i, d, candidates[j], 1)
tried.add(candidates[j])
return initialization
def sparse_initialized_nnd_search(
inds,
indptr,
data,
search_indptr,
search_inds,
initialization,
query_inds,
query_indptr,
query_data,
sparse_dist,
dist_args,
):
for i in numba.prange(query_indptr.shape[0] - 1):
tried = set(initialization[0, i])
to_inds = query_inds[query_indptr[i]:query_indptr[i + 1]]
to_data = query_data[query_indptr[i]:query_indptr[i + 1]]
while True:
# Find smallest flagged vertex
vertex = smallest_flagged(initialization, i)
if vertex == -1:
break
candidates = search_inds[search_indptr[vertex]:search_indptr[vertex
+ 1]]
for j in range(candidates.shape[0]):
if (candidates[j] == vertex or candidates[j] == -1
or candidates[j] in tried):
continue
from_inds = inds[indptr[candidates[j]]:indptr[candidates[j] +
1]]
from_data = data[indptr[candidates[j]]:indptr[candidates[j] +
1]]
d = sparse_dist(from_inds, from_data, to_inds, to_data,
*dist_args)
unchecked_heap_push(initialization, i, d, candidates[j], 1)
tried.add(candidates[j])
return initialization
