def test_init_random():
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")
current_graph_threaded = utils.make_heap(data.shape[0], n_neighbors)
threaded.init_random(
current_graph_threaded,
data,
dist,
dist_args,
n_neighbors,
chunk_size=chunk_size,
rng_state=new_rng_state(),
parallel=parallel,
seed_per_row=True,
)
assert_allclose(current_graph_threaded, current_graph)
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)