def test_new_build_candidates():
n_vertices = data.shape[0]
current_graph = pynndescent_.init_current_graph(
data, dist, dist_args, n_neighbors, rng_state=new_rng_state(), seed_per_row=True
)
new_candidate_neighbors, old_candidate_neighbors = utils.new_build_candidates(
current_graph,
n_vertices,
n_neighbors,
max_candidates,
rng_state=new_rng_state(),
seed_per_row=True,
)
current_graph = pynndescent_.init_current_graph(
data, dist, dist_args, n_neighbors, rng_state=new_rng_state(), seed_per_row=True
)
parallel = joblib.Parallel(n_jobs=2, prefer="threads")
new_candidate_neighbors_threaded, old_candidate_neighbors_threaded = threaded.new_build_candidates(
current_graph,
n_vertices,
n_neighbors,
max_candidates,
chunk_size=chunk_size,
rng_state=new_rng_state(),
rho=0.5,
parallel=parallel,
seed_per_row=True,
)
assert_allclose(new_candidate_neighbors_threaded, new_candidate_neighbors)
assert_allclose(old_candidate_neighbors_threaded, old_candidate_neighbors)
def test_mark_candidate_results():
np.random.seed(42)
N = 100
D = 128
chunk_size = N // 8
n_neighbors = 25
data = np.random.rand(N, D).astype(np.float32)
n_vertices = data.shape[0]
current_graph = utils.make_heap(data.shape[0], n_neighbors)
pynndescent_.init_random(
n_neighbors,
data,
current_graph,
dist,
new_rng_state(),
seed_per_row=True,
)
pynndescent_.nn_descent_internal_low_memory_parallel(current_graph,
data,
n_neighbors,
new_rng_state(),
n_iters=2,
seed_per_row=True)
current_graph_threaded = utils.Heap(
current_graph[0].copy(),
current_graph[1].copy(),
current_graph[2].copy(),
)
new_candidate_neighbors, old_candidate_neighbors = utils.new_build_candidates(
current_graph,
n_vertices,
n_neighbors,
max_candidates,
rng_state=new_rng_state(),
seed_per_row=True,
)
parallel = joblib.Parallel(n_jobs=2, prefer="threads")
(
new_candidate_neighbors_threaded,
old_candidate_neighbors_threaded,
) = threaded.new_build_candidates(
current_graph_threaded,
n_vertices,
n_neighbors,
max_candidates,
chunk_size=chunk_size,
rng_state=new_rng_state(),
parallel=parallel,
seed_per_row=True,
)
assert_allclose(current_graph_threaded, current_graph)
def test_new_build_candidates():
np.random.seed(42)
N = 100
D = 128
chunk_size = N // 8
n_neighbors = 25
data = np.random.rand(N, D).astype(np.float32)
n_vertices = data.shape[0]
current_graph = utils.make_heap(data.shape[0], n_neighbors)
pynndescent_.init_random(
n_neighbors,
data,
current_graph,
dist,
dist_args,
new_rng_state(),
seed_per_row=True,
)
new_candidate_neighbors, old_candidate_neighbors = utils.new_build_candidates(
current_graph,
n_vertices,
n_neighbors,
max_candidates,
rng_state=new_rng_state(),
seed_per_row=True,
)
current_graph = utils.make_heap(data.shape[0], n_neighbors)
pynndescent_.init_random(
n_neighbors,
data,
current_graph,
dist,
dist_args,
new_rng_state(),
seed_per_row=True,
)
parallel = joblib.Parallel(n_jobs=2, prefer="threads")
(
new_candidate_neighbors_threaded,
old_candidate_neighbors_threaded,
) = threaded.new_build_candidates(
current_graph,
n_vertices,
n_neighbors,
max_candidates,
chunk_size=chunk_size,
rng_state=new_rng_state(),
parallel=parallel,
seed_per_row=True,
)
assert_allclose(new_candidate_neighbors_threaded, new_candidate_neighbors)
assert_allclose(old_candidate_neighbors_threaded, old_candidate_neighbors)
def sparse_nn_descent(
inds,
indptr,
data,
n_vertices,
n_neighbors,
rng_state,
max_candidates=50,
dist=sparse.sparse_euclidean,
n_iters=10,
delta=0.001,
rp_tree_init=False,
leaf_array=None,
verbose=False,
n_jobs=None,
seed_per_row=False,
):
if rng_state is None:
rng_state = new_rng_state()
with joblib.Parallel(prefer="threads", n_jobs=n_jobs) as parallel:
n_tasks = effective_n_jobs_with_context(n_jobs)
chunk_size = int(math.ceil(n_vertices / n_tasks))
current_graph = make_heap(n_vertices, n_neighbors)
if rp_tree_init:
sparse_init_rp_tree(
inds,
indptr,
data,
dist,
current_graph,
leaf_array,
chunk_size,
parallel,
)
sparse_init_random(
current_graph,
inds,
indptr,
data,
dist,
n_neighbors,
chunk_size,
rng_state,
parallel,
seed_per_row=seed_per_row,
)
# store the updates in an array
# note that the factor here is `n_neighbors * n_neighbors`, not `max_candidates * max_candidates`
# since no more than `n_neighbors` candidates are added for each row
max_heap_update_count = chunk_size * n_neighbors * n_neighbors * 4
heap_updates = np.zeros((n_tasks, max_heap_update_count, 4),
dtype=np.float32)
heap_update_counts = np.zeros((n_tasks, ), dtype=np.int64)
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,
chunk_size,
rng_state,
parallel,
seed_per_row=seed_per_row,
)
def nn_descent_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return (
index,
sparse_nn_descent_map_jit(
rows,
max_candidates,
inds,
indptr,
data,
new_candidate_neighbors,
old_candidate_neighbors,
heap_updates[index],
offset=0,
sparse_dist=dist,
),
)
def nn_decent_reduce(index):
return nn_decent_reduce_jit(n_tasks, current_graph,
heap_updates, offsets, index)
# run map functions
for index, count in parallel(
parallel_calls(nn_descent_map, n_tasks)):
heap_update_counts[index] = count
# sort and chunk heap updates so they can be applied in the reduce
max_count = heap_update_counts.max()
offsets = np.zeros((n_tasks, max_count), dtype=np.int64)
def shuffle(index):
return shuffle_jit(
heap_updates,
heap_update_counts,
offsets,
chunk_size,
n_vertices,
index,
)
parallel(parallel_calls(shuffle, n_tasks))
# then run reduce functions
c = 0
for c_part in parallel(parallel_calls(nn_decent_reduce, n_tasks)):
c += c_part
if c <= delta * n_neighbors * data.shape[0]:
break
def deheap_sort_map(index):
rows = chunk_rows(chunk_size, index, n_vertices)
return index, deheap_sort_map_jit(rows, current_graph)
parallel(parallel_calls(deheap_sort_map, n_tasks))
return current_graph[0].astype(np.int64), current_graph[1]
