def nn_descent_internal_high_memory(
current_graph,
data,
n_neighbors,
rng_state,
tried,
max_candidates=50,
dist=dist.euclidean,
n_iters=10,
delta=0.001,
rho=0.5,
verbose=False,
):
n_vertices = data.shape[0]
for n in range(n_iters):
with numba.objmode():
# Call into object mode to temporarily sleep (and thus release GIL)
logging.info("(obj mode) high mem nn descent iter.")
time.sleep(0.05)
if verbose:
print("\t", n, " / ", n_iters)
(new_candidate_neighbors, old_candidate_neighbors) = new_build_candidates(
current_graph, n_vertices, n_neighbors, max_candidates, rng_state, rho
)
c = 0
for i in range(n_vertices):
for j in range(max_candidates):
p = int(new_candidate_neighbors[0, i, j])
if p < 0:
continue
for k in range(j, max_candidates):
q = int(new_candidate_neighbors[0, i, k])
if q < 0 or (p, q) in tried:
continue
d = dist(data[p], data[q])
c += unchecked_heap_push(current_graph, p, d, q, 1)
tried.add((p, q))
if p != q:
c += unchecked_heap_push(current_graph, q, d, p, 1)
tried.add((q, p))
for k in range(max_candidates):
q = int(old_candidate_neighbors[0, i, k])
if q < 0 or (p, q) in tried:
continue
d = dist(data[p], data[q])
c += unchecked_heap_push(current_graph, p, d, q, 1)
tried.add((p, q))
if p != q:
c += unchecked_heap_push(current_graph, q, d, p, 1)
tried.add((q, p))
if c <= delta * n_neighbors * data.shape[0]:
return
def nn_descent_internal_high_memory(
current_graph,
data,
n_neighbors,
rng_state,
tried,
max_candidates=50,
dist=dist.euclidean,
dist_args=(),
n_iters=10,
delta=0.001,
rho=0.5,
verbose=False,
):
n_vertices = data.shape[0]
for n in range(n_iters):
if verbose:
print("\t", n, " / ", n_iters)
(new_candidate_neighbors, old_candidate_neighbors) = new_build_candidates(
current_graph, n_vertices, n_neighbors, max_candidates, rng_state, rho
)
c = 0
for i in range(n_vertices):
for j in range(max_candidates):
p = int(new_candidate_neighbors[0, i, j])
if p < 0:
continue
for k in range(j, max_candidates):
q = int(new_candidate_neighbors[0, i, k])
if q < 0 or (p, q) in tried:
continue
d = dist(data[p], data[q], *dist_args)
c += unchecked_heap_push(current_graph, p, d, q, 1)
tried.add((p, q))
if p != q:
c += unchecked_heap_push(current_graph, q, d, p, 1)
tried.add((q, p))
for k in range(max_candidates):
q = int(old_candidate_neighbors[0, i, k])
if q < 0 or (p, q) in tried:
continue
d = dist(data[p], data[q], *dist_args)
c += unchecked_heap_push(current_graph, p, d, q, 1)
tried.add((p, q))
if p != q:
c += unchecked_heap_push(current_graph, q, d, p, 1)
tried.add((q, p))
if c <= delta * n_neighbors * data.shape[0]:
return
def sparse_nn_descent_internal_low_memory(
current_graph,
inds,
indptr,
data,
n_vertices,
n_neighbors,
rng_state,
max_candidates=50,
sparse_dist=umap.sparse.sparse_euclidean,
n_iters=10,
delta=0.001,
rho=0.5,
verbose=False,
):
for n in range(n_iters):
if verbose:
print("\t", n, " / ", n_iters)
(new_candidate_neighbors,
old_candidate_neighbors) = new_build_candidates(
current_graph, n_vertices, n_neighbors, max_candidates, rng_state,
rho)
c = 0
for i in range(n_vertices):
for j in range(max_candidates):
p = int(new_candidate_neighbors[0, i, j])
if p < 0:
continue
for k in range(j, max_candidates):
q = int(new_candidate_neighbors[0, i, k])
if q < 0:
continue
from_inds = inds[indptr[p]:indptr[p + 1]]
from_data = data[indptr[p]:indptr[p + 1]]
to_inds = inds[indptr[q]:indptr[q + 1]]
to_data = data[indptr[q]:indptr[q + 1]]
d = sparse_dist(from_inds, from_data, to_inds, to_data)
c += heap_push(current_graph, p, d, q, 1)
if p != q:
c += heap_push(current_graph, q, d, p, 1)
for k in range(max_candidates):
q = int(old_candidate_neighbors[0, i, k])
if q < 0:
continue
from_inds = inds[indptr[p]:indptr[p + 1]]
from_data = data[indptr[p]:indptr[p + 1]]
to_inds = inds[indptr[q]:indptr[q + 1]]
to_data = data[indptr[q]:indptr[q + 1]]
d = sparse_dist(from_inds, from_data, to_inds, to_data)
c += heap_push(current_graph, p, d, q, 1)
if p != q:
c += heap_push(current_graph, q, d, p, 1)
if c <= delta * n_neighbors * n_vertices:
return