graphclust.py 文件源码

def compute_nearest_neighbors(submatrix, balltree, k, row_start):
    """ Compute k nearest neighbors on a submatrix
    Args: submatrix (np.ndarray): Data submatrix
          balltree: Nearest neighbor index (from sklearn)
          k: number of nearest neigbors to compute
          row_start: row offset into larger matrix
    Returns a COO sparse adjacency matrix of nearest neighbor relations as (i,j,x)"""

    nn_dist, nn_idx = balltree.query(submatrix, k=k+1)

    # Remove the self-as-neighbors
    nn_idx = nn_idx[:,1:]
    nn_dist = nn_dist[:,1:]

    # Construct a COO sparse matrix of edges and distances
    i = np.repeat(row_start + np.arange(nn_idx.shape[0]), k)
    j = nn_idx.ravel().astype(int)
    return (i, j, nn_dist.ravel())