Skip to content

Instantly share code, notes, and snippets.

@ahwillia
Last active July 8, 2024 21:32
Show Gist options
  • Save ahwillia/2fc2ce7c42d7d63f071b504cd0b88008 to your computer and use it in GitHub Desktop.
Save ahwillia/2fc2ce7c42d7d63f071b504cd0b88008 to your computer and use it in GitHub Desktop.
symmetric_knn_graph
import numpy as np
from scipy.spatial.distance import pdist, squareform
from scipy.sparse import coo_matrix
from scipy.sparse.csgraph import laplacian
import matplotlib.pyplot as plt
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
"""
# Check that we are asking for a reasonable number of neighbors.
assert k > 0
# 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())), shape=(n, n))
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment