def _nhood_search(umap_object, nhood_size):
if umap_object._small_data:
dmat = sklearn.metrics.pairwise_distances(umap_object._raw_data)
indices = np.argpartition(dmat, nhood_size)[:, :nhood_size]
dmat_shortened = submatrix(dmat, indices, nhood_size)
indices_sorted = np.argsort(dmat_shortened)
indices = submatrix(indices, indices_sorted, nhood_size)
dists = submatrix(dmat_shortened, indices_sorted, nhood_size)
else:
rng_state = np.empty(3, dtype=np.int64)
if len(umap_object._metric_kwds) >= 1:
_dist = umap_object._input_distance_func
_args = tuple(umap_object._metric_kwds.values())
@numba.njit()
def _metric(x, y):
_dist(x, y, *_args)
else:
_metric = umap_object._input_distance_func
init = initialise_search(
umap_object._rp_forest,
umap_object._raw_data,
umap_object._raw_data,
int(nhood_size * umap_object.transform_queue_size),
rng_state,
_metric,
)
result = initialized_nnd_search(
umap_object._raw_data,
umap_object._search_graph.indptr,
umap_object._search_graph.indices,
init,
umap_object._raw_data,
_metric,
)
indices, dists = deheap_sort(result)
indices = indices[:, :nhood_size]
dists = dists[:, :nhood_size]
return indices, dists
def _nhood_search(umap_object, nhood_size):
if umap_object._small_data:
dmat = sklearn.metrics.pairwise_distances(umap_object._raw_data)
indices = np.argpartition(dmat, nhood_size)[:, :nhood_size]
dmat_shortened = submatrix(dmat, indices, nhood_size)
indices_sorted = np.argsort(dmat_shortened)
indices = submatrix(indices, indices_sorted, nhood_size)
dists = submatrix(dmat_shortened, indices_sorted, nhood_size)
else:
rng_state = np.empty(3, dtype=np.int64)
init = initialise_search(
umap_object._rp_forest,
umap_object._raw_data,
umap_object._raw_data,
int(nhood_size * umap_object.transform_queue_size),
rng_state,
umap_object._distance_func,
umap_object._dist_args,
)
result = initialized_nnd_search(
umap_object._raw_data,
umap_object._search_graph.indptr,
umap_object._search_graph.indices,
init,
umap_object._raw_data,
umap_object._distance_func,
umap_object._dist_args,
)
indices, dists = deheap_sort(result)
indices = indices[:, :nhood_size]
dists = dists[:, :nhood_size]
return indices, dists