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 initialize_heaps(data, n_neighbors, leaf_array, dist=dist.euclidean, dist_args=()): graph_heap = make_heap(data.shape[0], 10) search_heap = make_heap(data.shape[0], n_neighbors * 2) tried = set([(-1, -1)]) for n in range(leaf_array.shape[0]): for i in range(leaf_array.shape[1]): if leaf_array[n, i] < 0: break for j in range(i + 1, leaf_array.shape[1]): if leaf_array[n, j] < 0: break if (leaf_array[n, i], leaf_array[n, j]) in tried: continue d = dist(data[leaf_array[n, i]], data[leaf_array[n, j]], *dist_args) unchecked_heap_push(graph_heap, leaf_array[n, i], d, leaf_array[n, j], 1) unchecked_heap_push(graph_heap, leaf_array[n, j], d, leaf_array[n, i], 1) unchecked_heap_push(search_heap, leaf_array[n, i], d, leaf_array[n, j], 1) unchecked_heap_push(search_heap, leaf_array[n, j], d, leaf_array[n, i], 1) tried.add((leaf_array[n, i], leaf_array[n, j])) tried.add((leaf_array[n, j], leaf_array[n, i])) return graph_heap, search_heap
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 search( query_inds, query_indptr, query_data, k, inds, indptr, data, forest, search_indptr, search_indices, epsilon, n_neighbors, tried, sparse_dist, rng_state, ): n_query_points = query_indptr.shape[0] - 1 result = make_heap(n_query_points, k) for i in range(n_query_points): tried[:] = 0 current_query_inds = query_inds[query_indptr[i] : query_indptr[i + 1]] current_query_data = query_data[query_indptr[i] : query_indptr[i + 1]] heap_priorities, heap_indices = search_init( current_query_inds, current_query_data, k, inds, indptr, data, forest, n_neighbors, tried, sparse_dist, rng_state, ) heap_priorities, heap_indices = search_from_init( current_query_inds, current_query_data, inds, indptr, data, search_indptr, search_indices, heap_priorities, heap_indices, epsilon, tried, sparse_dist, ) result[0][i] = heap_indices result[1][i] = heap_priorities return result
def initialise_search(forest, data, query_points, n_neighbors, init_from_random, init_from_tree, rng_state): results = make_heap(query_points.shape[0], n_neighbors) init_from_random(n_neighbors, data, query_points, results, rng_state) if forest is not None: for tree in forest: init_from_tree(tree, data, query_points, results, rng_state) return results
def nn_descent( inds, indptr, data, n_neighbors, rng_state, max_candidates=50, dist=sparse_euclidean, n_iters=10, delta=0.001, rp_tree_init=True, leaf_array=None, low_memory=False, verbose=False, seed_per_row=False, ): n_samples = indptr.shape[0] - 1 current_graph = make_heap(n_samples, n_neighbors) if rp_tree_init: init_rp_tree(inds, indptr, data, dist, current_graph, leaf_array) init_random(n_neighbors, inds, indptr, data, current_graph, dist, rng_state) if low_memory: nn_descent_internal_low_memory_parallel( current_graph, inds, indptr, data, n_neighbors, rng_state, max_candidates=max_candidates, dist=dist, n_iters=n_iters, delta=delta, verbose=verbose, seed_per_row=seed_per_row, ) else: nn_descent_internal_high_memory_parallel( current_graph, inds, indptr, data, n_neighbors, rng_state, max_candidates=max_candidates, dist=dist, n_iters=n_iters, delta=delta, verbose=verbose, seed_per_row=seed_per_row, ) return deheap_sort(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_init_rp_tree(): # Use more graph_data than the other tests since otherwise init_rp_tree has nothing to do np.random.seed(42) N = 100 D = 128 chunk_size = N // 8 n_neighbors = 25 data = np.random.rand(N, D).astype(np.float32) rng_state = new_rng_state() random_state = check_random_state(42) current_graph = utils.make_heap(data.shape[0], n_neighbors) _rp_forest = make_forest( data, n_neighbors, n_trees=8, leaf_size=None, rng_state=rng_state, random_state=random_state, ) leaf_array = rptree_leaf_array(_rp_forest) pynndescent_.init_rp_tree(data, dist, current_graph, leaf_array) rng_state = new_rng_state() random_state = check_random_state(42) current_graph_threaded = utils.make_heap(data.shape[0], n_neighbors) _rp_forest = make_forest( data, n_neighbors, n_trees=8, leaf_size=None, rng_state=rng_state, random_state=random_state, ) leaf_array = rptree_leaf_array(_rp_forest) parallel = joblib.Parallel(n_jobs=2, prefer="threads") threaded.init_rp_tree(data, dist, current_graph_threaded, leaf_array, chunk_size, parallel) assert_allclose(current_graph_threaded, current_graph)
def init_current_graph( data, dist, dist_args, n_neighbors, rng_state, seed_per_row=False ): current_graph = make_heap(data.shape[0], n_neighbors) for i in range(data.shape[0]): if seed_per_row: seed(rng_state, i) indices = rejection_sample(n_neighbors, data.shape[0], rng_state) for j in range(indices.shape[0]): d = dist(data[i], data[indices[j]], *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) return current_graph
def sparse_initialise_search( forest, inds, indptr, data, query_inds, query_indptr, query_data, n_neighbors, rng_state, sparse_dist, dist_args, ): results = make_heap(query_indptr.shape[0] - 1, n_neighbors) sparse_init_from_random( n_neighbors, inds, indptr, data, query_inds, query_indptr, query_data, results, rng_state, sparse_dist, dist_args, ) if forest is not None: for tree in forest: sparse_init_from_tree( tree, inds, indptr, data, query_inds, query_indptr, query_data, results, rng_state, sparse_dist, dist_args, ) return results
def nn_descent( inds, indptr, data, n_vertices, n_neighbors, rng_state, max_candidates=50, n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, verbose=False, ): current_graph = make_heap(n_vertices, n_neighbors) for i in range(n_vertices): indices = rejection_sample(n_neighbors, n_vertices, rng_state) for j in range(indices.shape[0]): from_inds = inds[indptr[i]:indptr[i + 1]] from_data = data[indptr[i]:indptr[i + 1]] to_inds = inds[indptr[indices[j]]:indptr[indices[j] + 1]] to_data = data[indptr[indices[j]]:indptr[indices[j] + 1]] d = sparse_dist(from_inds, from_data, to_inds, to_data, *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) if rp_tree_init: for n in range(leaf_array.shape[0]): for i in range(leaf_array.shape[1]): if leaf_array[n, i] < 0: break for j in range(i + 1, leaf_array.shape[1]): if leaf_array[n, j] < 0: break from_inds = inds[indptr[leaf_array[ n, i]]:indptr[leaf_array[n, i] + 1]] from_data = data[indptr[leaf_array[ n, i]]:indptr[leaf_array[n, i] + 1]] to_inds = inds[indptr[leaf_array[ n, j]]:indptr[leaf_array[n, j] + 1]] to_data = data[indptr[leaf_array[ n, j]]:indptr[leaf_array[n, j] + 1]] d = sparse_dist(from_inds, from_data, to_inds, to_data, *dist_args) heap_push(current_graph, leaf_array[n, i], d, leaf_array[n, j], 1) heap_push(current_graph, leaf_array[n, j], d, leaf_array[n, i], 1) for n in range(n_iters): if verbose: print("\t", n, " / ", n_iters) candidate_neighbors = build_candidates(current_graph, n_vertices, n_neighbors, max_candidates, rng_state) c = 0 for i in range(n_vertices): for j in range(max_candidates): p = int(candidate_neighbors[0, i, j]) if p < 0 or tau_rand(rng_state) < rho: continue for k in range(max_candidates): q = int(candidate_neighbors[0, i, k]) if (q < 0 or not candidate_neighbors[2, i, j] and not candidate_neighbors[2, i, k]): 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, *dist_args) c += heap_push(current_graph, p, d, q, 1) c += heap_push(current_graph, q, d, p, 1) if c <= delta * n_neighbors * n_vertices: break return deheap_sort(current_graph)
def sparse_nn_descent( inds, indptr, data, n_vertices, n_neighbors, rng_state, max_candidates=50, sparse_dist=sparse_euclidean, dist_args=(), n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, verbose=False, ): tried = set([(-1, -1)]) current_graph = make_heap(n_vertices, n_neighbors) for i in range(n_vertices): indices = rejection_sample(n_neighbors, n_vertices, rng_state) for j in range(indices.shape[0]): from_inds = inds[indptr[i]:indptr[i + 1]] from_data = data[indptr[i]:indptr[i + 1]] to_inds = inds[indptr[indices[j]]:indptr[indices[j] + 1]] to_data = data[indptr[indices[j]]:indptr[indices[j] + 1]] d = sparse_dist(from_inds, from_data, to_inds, to_data, *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) tried.add((i, indices[j])) tried.add((indices[j], i)) if rp_tree_init: sparse_init_rp_tree( inds, indptr, data, sparse_dist, dist_args, current_graph, leaf_array, tried=tried, ) 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, False, ) 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 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, *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 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, *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 * n_vertices: break return deheap_sort(current_graph)
def nn_descent( inds, indptr, data, n_neighbors, rng_state, max_candidates=50, dist=sparse_euclidean, n_iters=10, delta=0.001, init_graph=EMPTY_GRAPH, rp_tree_init=True, leaf_array=None, low_memory=False, verbose=False, ): n_samples = indptr.shape[0] - 1 if init_graph[0].shape[0] == 1: # EMPTY_GRAPH current_graph = make_heap(n_samples, n_neighbors) if rp_tree_init: init_rp_tree(inds, indptr, data, dist, current_graph, leaf_array) init_random(n_neighbors, inds, indptr, data, current_graph, dist, rng_state) elif init_graph[0].shape[0] == n_samples and init_graph[0].shape[1] == n_neighbors: current_graph = init_graph else: raise ValueError("Invalid initial graph specified!") if low_memory: nn_descent_internal_low_memory_parallel( current_graph, inds, indptr, data, n_neighbors, rng_state, max_candidates=max_candidates, dist=dist, n_iters=n_iters, delta=delta, verbose=verbose, ) else: nn_descent_internal_high_memory_parallel( current_graph, inds, indptr, data, n_neighbors, rng_state, max_candidates=max_candidates, dist=dist, n_iters=n_iters, delta=delta, verbose=verbose, ) return deheap_sort(current_graph[0], current_graph[1])
from pynndescent.utils import ( tau_rand_int, make_heap, new_build_candidates, deheap_sort, checked_flagged_heap_push, apply_graph_updates_high_memory, apply_graph_updates_low_memory, ) from pynndescent.sparse import sparse_euclidean locale.setlocale(locale.LC_NUMERIC, "C") EMPTY_GRAPH = make_heap(1, 1) @numba.njit(parallel=True, cache=False) def generate_leaf_updates(leaf_block, dist_thresholds, inds, indptr, data, dist): updates = [[(-1, -1, np.inf)] for i in range(leaf_block.shape[0])] for n in numba.prange(leaf_block.shape[0]): for i in range(leaf_block.shape[1]): p = leaf_block[n, i] if p < 0: break for j in range(i + 1, leaf_block.shape[1]): q = leaf_block[n, j]
def nn_descent(data, n_neighbors, rng_state, max_candidates=50, dist=dist.euclidean, dist_args=(), n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, verbose=False): n_vertices = data.shape[0] current_graph = make_heap(data.shape[0], n_neighbors) for i in range(data.shape[0]): indices = rejection_sample(n_neighbors, data.shape[0], rng_state) for j in range(indices.shape[0]): d = dist(data[i], data[indices[j]], *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) if rp_tree_init: for n in range(leaf_array.shape[0]): tried = set([(-1, -1)]) for i in range(leaf_array.shape[1]): if leaf_array[n, i] < 0: break for j in range(i + 1, leaf_array.shape[1]): if leaf_array[n, j] < 0: break if (leaf_array[n, i], leaf_array[n, j]) in tried: continue d = dist(data[leaf_array[n, i]], data[leaf_array[n, j]], *dist_args) heap_push(current_graph, leaf_array[n, i], d, leaf_array[n, j], 1) heap_push(current_graph, leaf_array[n, j], d, leaf_array[n, i], 1) tried.add((leaf_array[n, i], leaf_array[n, j])) tried.add((leaf_array[n, j], leaf_array[n, i])) for n in range(n_iters): (new_candidate_neighbors, old_candidate_neighbors) = 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 d = dist(data[p], data[q], *dist_args) c += heap_push(current_graph, p, d, q, 1) 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 d = dist(data[p], data[q], *dist_args) c += heap_push(current_graph, p, d, q, 1) c += heap_push(current_graph, q, d, p, 1) if c <= delta * n_neighbors * data.shape[0]: break return deheap_sort(current_graph)
def new_build_candidates( current_graph, n_vertices, n_neighbors, max_candidates, chunk_size, rng_state, parallel, seed_per_row=False, ): n_tasks = int(math.ceil(float(n_vertices) / chunk_size)) new_candidate_neighbors = make_heap(n_vertices, max_candidates) old_candidate_neighbors = make_heap(n_vertices, max_candidates) # store the updates in an array max_heap_update_count = chunk_size * n_neighbors * 2 heap_updates = np.zeros((n_tasks, max_heap_update_count, 4), dtype=np.float32) heap_update_counts = np.zeros((n_tasks, ), dtype=np.int64) rng_state_threads = per_thread_rng_state(n_tasks, rng_state) def candidates_map(index): rows = chunk_rows(chunk_size, index, n_vertices) return ( index, candidates_map_jit( rows, n_neighbors, current_graph, heap_updates[index], offset=0, rng_state=rng_state_threads[index], seed_per_row=seed_per_row, ), ) def candidates_reduce(index): return candidates_reduce_jit( n_tasks, current_graph, new_candidate_neighbors, old_candidate_neighbors, heap_updates, offsets, index, ) # run map functions for index, count in parallel(parallel_calls(candidates_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 parallel(parallel_calls(candidates_reduce, n_tasks)) def mark_candidate_results(index): rows = chunk_rows(chunk_size, index, n_vertices) return mark_candidate_results_map(rows, current_graph, n_neighbors, max_candidates, new_candidate_neighbors) # Now mark whether things were used correctly parallel(parallel_calls(mark_candidate_results, n_tasks)) return new_candidate_neighbors, old_candidate_neighbors
def nn_descent( data, n_neighbors, rng_state, max_candidates=50, dist=dist.euclidean, dist_args=(), n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, verbose=False, seed_per_row=False, ): n_vertices = data.shape[0] tried = set([(-1, -1)]) current_graph = make_heap(data.shape[0], n_neighbors) for i in range(data.shape[0]): if seed_per_row: seed(rng_state, i) indices = rejection_sample(n_neighbors, data.shape[0], rng_state) for j in range(indices.shape[0]): d = dist(data[i], data[indices[j]], *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) tried.add((i, indices[j])) tried.add((indices[j], i)) if rp_tree_init: init_rp_tree(data, dist, dist_args, current_graph, leaf_array, tried=tried) 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, seed_per_row, ) 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]: break return deheap_sort(current_graph)
def custom_search_closure(query_points, candidate_indices, k, epsilon, visited): result = make_heap(query_points.shape[0], k) distance_scale = 1.0 + epsilon for i in range(query_points.shape[0]): visited[:] = 0 if dist == alternative_dot or dist == alternative_cosine: norm = np.sqrt((query_points[i]**2).sum()) if norm > 0.0: current_query = query_points[i] / norm else: continue else: current_query = query_points[i] heap_priorities = result[1][i] heap_indices = result[0][i] seed_set = [(np.float32(np.inf), np.int32(-1)) for j in range(0)] ############ Init ################ n_initial_points = candidate_indices.shape[0] for j in range(n_initial_points): candidate = np.int32(candidate_indices[j]) d = dist(data[candidate], current_query) # indices are guaranteed different simple_heap_push(heap_priorities, heap_indices, d, candidate) heapq.heappush(seed_set, (d, candidate)) mark_visited(visited, candidate) ############ Search ############## distance_bound = distance_scale * heap_priorities[0] # Find smallest seed point d_vertex, vertex = heapq.heappop(seed_set) while d_vertex < distance_bound: for j in range(indptr[vertex], indptr[vertex + 1]): candidate = indices[j] if has_been_visited(visited, candidate) == 0: mark_visited(visited, candidate) d = dist(data[candidate], current_query) if d < distance_bound: simple_heap_push(heap_priorities, heap_indices, d, candidate) heapq.heappush(seed_set, (d, candidate)) # Update bound distance_bound = distance_scale * heap_priorities[0] # find new smallest seed point if len(seed_set) == 0: break else: d_vertex, vertex = heapq.heappop(seed_set) return result
def nn_descent( data, n_neighbors, rng_state, max_candidates=50, dist=dist.euclidean, dist_args=(), n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, low_memory=False, verbose=False, seed_per_row=False, ): tried = set([(-1, -1)]) current_graph = make_heap(data.shape[0], n_neighbors) for i in range(data.shape[0]): if seed_per_row: seed(rng_state, i) indices = rejection_sample(n_neighbors, data.shape[0], rng_state) for j in range(indices.shape[0]): d = dist(data[i], data[indices[j]], *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) tried.add((i, indices[j])) tried.add((indices[j], i)) if rp_tree_init: init_rp_tree(data, dist, dist_args, current_graph, leaf_array, tried=tried) if low_memory: nn_descent_internal_low_memory( current_graph, data, n_neighbors, rng_state, max_candidates=max_candidates, dist=dist, dist_args=dist_args, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, seed_per_row=seed_per_row, ) else: nn_descent_internal_high_memory( current_graph, data, n_neighbors, rng_state, tried, max_candidates=max_candidates, dist=dist, dist_args=dist_args, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, seed_per_row=seed_per_row, ) return deheap_sort(current_graph)
def sparse_nn_descent( inds, indptr, data, n_vertices, n_neighbors, rng_state, max_candidates=50, sparse_dist=sparse_euclidean, dist_args=(), n_iters=10, delta=0.001, rho=0.5, low_memory=False, rp_tree_init=True, leaf_array=None, verbose=False, ): tried = set([(-1, -1)]) current_graph = make_heap(n_vertices, n_neighbors) for i in range(n_vertices): indices = rejection_sample(n_neighbors, n_vertices, rng_state) for j in range(indices.shape[0]): from_inds = inds[indptr[i]:indptr[i + 1]] from_data = data[indptr[i]:indptr[i + 1]] to_inds = inds[indptr[indices[j]]:indptr[indices[j] + 1]] to_data = data[indptr[indices[j]]:indptr[indices[j] + 1]] d = sparse_dist(from_inds, from_data, to_inds, to_data, *dist_args) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) tried.add((i, indices[j])) tried.add((indices[j], i)) if rp_tree_init: sparse_init_rp_tree( inds, indptr, data, sparse_dist, dist_args, current_graph, leaf_array, tried=tried, ) if low_memory: sparse_nn_descent_internal_low_memory( current_graph, inds, indptr, data, n_vertices, n_neighbors, rng_state, max_candidates=max_candidates, sparse_dist=sparse_dist, dist_args=dist_args, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, ) else: sparse_nn_descent_internal_high_memory( current_graph, inds, indptr, data, n_vertices, n_neighbors, rng_state, tried, max_candidates=max_candidates, sparse_dist=sparse_dist, dist_args=dist_args, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, ) return deheap_sort(current_graph)
def nn_descent( data, n_neighbors, rng_state, max_candidates=50, dist=dst.euclidean, dist_args=(), 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_vertices = data.shape[0] n_tasks = effective_n_jobs_with_context(n_jobs) chunk_size = int(math.ceil(n_vertices / n_tasks)) current_graph = make_heap(data.shape[0], n_neighbors) if rp_tree_init: init_rp_tree(data, dist, dist_args, current_graph, leaf_array, chunk_size, parallel) init_random( current_graph, data, dist, dist_args, 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, nn_descent_map_jit( rows, max_candidates, data, new_candidate_neighbors, old_candidate_neighbors, heap_updates[index], offset=0, dist=dist, dist_args=dist_args, ), ) 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]
def init_current_graph( data, dist, dist_args, n_neighbors, chunk_size, rng_state, parallel, seed_per_row=False, ): n_vertices = data.shape[0] n_tasks = int(math.ceil(float(n_vertices) / chunk_size)) current_graph = make_heap(n_vertices, n_neighbors) # store the updates in an array max_heap_update_count = chunk_size * n_neighbors * 2 heap_updates = np.zeros((n_tasks, max_heap_update_count, 4), dtype=np.float32) heap_update_counts = np.zeros((n_tasks, ), dtype=np.int64) rng_state_threads = per_thread_rng_state(n_tasks, rng_state) def current_graph_map(index): rows = chunk_rows(chunk_size, index, n_vertices) return ( index, current_graph_map_jit( rows, n_vertices, n_neighbors, data, heap_updates[index], rng_state_threads[index], seed_per_row=seed_per_row, dist=dist, dist_args=dist_args, ), ) def current_graph_reduce(index): return current_graph_reduce_jit(n_tasks, current_graph, heap_updates, offsets, index) # run map functions for index, count in parallel(parallel_calls(current_graph_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 parallel(parallel_calls(current_graph_reduce, n_tasks)) return current_graph