symmetric_knn_graph

knn_graph.py

import numpy as np
from scipy.spatial.distance import pdist, squareform
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import laplacian


def symmetric_knn_graph(X, k):
    """
    Parameters
    ----------
    X : ndarray
        Matrix with dimensions (num_samples x num_features)

    Returns
    -------
    G : sparse matrix
        Specifies the KNN graph using scipy.sparse.csgraph conventions. See:
        >> https://docs.scipy.org/doc/scipy/reference/sparse.csgraph.html#graph-representations
    """

    # Number of nodes (samples)
    n = X.shape[0]

    # (samples x samples) matrix holding all pairwise distances.
    D = squareform(pdist(X, metric='sqeuclidean'))

    # Set diagonal to infinity to prevent each node counting
    # itself as it's nearest neighbor.
    D[np.diag_indices_from(D)] = np.inf

    # Rapidly compute the nearest k neighbors in each row.
    neighbor_idx = np.argpartition(D, k, axis=1)

    # Return sparse matrix defining the graph.
    ii = np.tile(np.arange(n)[:, None], (1, k))
    jj = neighbor_idx[:, :k]

    # Return symmetric matrix for an undirected graph.
    G = coo_matrix((np.ones(ii.size), (ii.ravel(), jj.ravel())))
    return (G + G.T).minimum(1.0)


if __name__ == "__main__":

    # 1000 random points in 10-d space
    X = np.random.RandomState(0).randn(1000, 10) 

    # Make symmetric nearest neighbor graph
    G = symmetric_knn_graph(X, 5)

    # Compute the graph laplacian.
    L = laplacian(G)