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 initialized_nnd_search(data, indptr, indices, initialization, query_points, dist): 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]) unchecked_heap_push(initialization, i, d, candidates[j], 1) tried.add(candidates[j]) return initialization
def nn_descent( data, n_neighbors, rng_state, max_candidates=50, dist=dist.euclidean, n_iters=10, delta=0.001, rho=0.5, rp_tree_init=True, leaf_array=None, low_memory=False, verbose=False, ): tried = set([(-1, -1)]) 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]]) 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, 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, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, ) else: nn_descent_internal_high_memory( current_graph, data, n_neighbors, rng_state, tried, max_candidates=max_candidates, dist=dist, n_iters=n_iters, delta=delta, rho=rho, verbose=verbose, ) return deheap_sort(current_graph)
def init_current_graph(data, dist, n_neighbors, rng_state): 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]]) heap_push(current_graph, i, d, indices[j], 1) heap_push(current_graph, indices[j], d, i, 1) return current_graph
def init_from_random(n_neighbors, data, query_points, heap, rng_state, dist): for i in range(query_points.shape[0]): indices = rejection_sample(n_neighbors, data.shape[0], rng_state) for j in range(indices.shape[0]): if indices[j] < 0: continue d = dist(data[indices[j]], query_points[i]) heap_push(heap, i, d, indices[j], 1) return
def init_from_tree(tree, data, query_points, heap, rng_state, dist): for i in range(query_points.shape[0]): indices = search_flat_tree( query_points[i], tree.hyperplanes, tree.offsets, tree.children, tree.indices, rng_state, ) for j in range(indices.shape[0]): if indices[j] < 0: continue d = dist(data[indices[j]], query_points[i]) heap_push(heap, i, d, indices[j], 1) return
def init_rp_tree(data, dist, current_graph, leaf_array, tried=None): if tried is None: tried = set([(-1, -1)]) for n in range(leaf_array.shape[0]): for i in range(leaf_array.shape[1]): p = leaf_array[n, i] if p < 0: break for j in range(i + 1, leaf_array.shape[1]): q = leaf_array[n, j] if q < 0: break if (p, q) in tried: continue d = dist(data[p], data[q]) heap_push(current_graph, p, d, q, 1) tried.add((p, q)) if p != q: heap_push(current_graph, q, d, p, 1) tried.add((q, p))
def nn_descent(data, 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): 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]): 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 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) 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 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 current_graph[:2]
def nn_descent(data, n_neighbors, max_candidates=50, n_iters=10, delta=0.001, rho=0.5, leaf_array=None): n_vertices = data.shape[0] rng_state = np.empty(3, dtype=np.int64) current_graph = make_heap(data.shape[0], n_neighbors) if leaf_array is not None: 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 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) else: for i in range(data.shape[0]): indices = np.random.choice(data.shape[0], size=n_neighbors, replace=False) 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) for n in range(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 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 current_graph[:2]