145 lines
4.6 KiB
Python
145 lines
4.6 KiB
Python
"""
|
|
Graph utilities and algorithms
|
|
|
|
Graphs are represented with their adjacency matrices, preferably using
|
|
sparse matrices.
|
|
"""
|
|
|
|
# Authors: Aric Hagberg <hagberg@lanl.gov>
|
|
# Gael Varoquaux <gael.varoquaux@normalesup.org>
|
|
# License: BSD
|
|
|
|
import numpy as np
|
|
from scipy import sparse
|
|
|
|
from .graph_shortest_path import graph_shortest_path
|
|
|
|
|
|
###############################################################################
|
|
# Path and connected component analysis.
|
|
# Code adapted from networkx
|
|
|
|
def single_source_shortest_path_length(graph, source, cutoff=None):
|
|
"""Return the shortest path length from source to all reachable nodes.
|
|
|
|
Returns a dictionary of shortest path lengths keyed by target.
|
|
|
|
Parameters
|
|
----------
|
|
graph: sparse matrix or 2D array (preferably LIL matrix)
|
|
Adjacency matrix of the graph
|
|
source : node label
|
|
Starting node for path
|
|
cutoff : integer, optional
|
|
Depth to stop the search - only
|
|
paths of length <= cutoff are returned.
|
|
|
|
Examples
|
|
--------
|
|
>>> from sklearn.utils.graph import single_source_shortest_path_length
|
|
>>> import numpy as np
|
|
>>> graph = np.array([[ 0, 1, 0, 0],
|
|
... [ 1, 0, 1, 0],
|
|
... [ 0, 1, 0, 1],
|
|
... [ 0, 0, 1, 0]])
|
|
>>> single_source_shortest_path_length(graph, 0)
|
|
{0: 0, 1: 1, 2: 2, 3: 3}
|
|
>>> single_source_shortest_path_length(np.ones((6, 6)), 2)
|
|
{0: 1, 1: 1, 2: 0, 3: 1, 4: 1, 5: 1}
|
|
"""
|
|
if sparse.isspmatrix(graph):
|
|
graph = graph.tolil()
|
|
else:
|
|
graph = sparse.lil_matrix(graph)
|
|
seen = {} # level (number of hops) when seen in BFS
|
|
level = 0 # the current level
|
|
next_level = [source] # dict of nodes to check at next level
|
|
while next_level:
|
|
this_level = next_level # advance to next level
|
|
next_level = set() # and start a new list (fringe)
|
|
for v in this_level:
|
|
if v not in seen:
|
|
seen[v] = level # set the level of vertex v
|
|
next_level.update(graph.rows[v])
|
|
if cutoff is not None and cutoff <= level:
|
|
break
|
|
level += 1
|
|
return seen # return all path lengths as dictionary
|
|
|
|
|
|
if hasattr(sparse, 'cs_graph_components'):
|
|
cs_graph_components = sparse.cs_graph_components
|
|
else:
|
|
from ._csgraph import cs_graph_components
|
|
|
|
|
|
###############################################################################
|
|
# Graph laplacian
|
|
def _graph_laplacian_sparse(graph, normed=False, return_diag=False):
|
|
n_nodes = graph.shape[0]
|
|
if not graph.format == 'coo':
|
|
lap = (-graph).tocoo()
|
|
else:
|
|
lap = -graph.copy()
|
|
diag_mask = (lap.row == lap.col)
|
|
if not diag_mask.sum() == n_nodes:
|
|
# The sparsity pattern of the matrix has holes on the diagonal,
|
|
# we need to fix that
|
|
diag_idx = lap.row[diag_mask]
|
|
|
|
lap = lap.tolil()
|
|
|
|
diagonal_holes = list(set(range(n_nodes)).difference(
|
|
diag_idx))
|
|
lap[diagonal_holes, diagonal_holes] = 1
|
|
lap = lap.tocoo()
|
|
diag_mask = (lap.row == lap.col)
|
|
lap.data[diag_mask] = 0
|
|
w = -np.asarray(lap.sum(axis=1)).squeeze()
|
|
if normed:
|
|
w = np.sqrt(w)
|
|
w_zeros = w == 0
|
|
w[w_zeros] = 1
|
|
lap.data /= w[lap.row]
|
|
lap.data /= w[lap.col]
|
|
lap.data[diag_mask] = (1 - w_zeros).astype(lap.data.dtype)
|
|
else:
|
|
lap.data[diag_mask] = w[lap.row[diag_mask]]
|
|
if return_diag:
|
|
return lap, w
|
|
return lap
|
|
|
|
|
|
def _graph_laplacian_dense(graph, normed=False, return_diag=False):
|
|
n_nodes = graph.shape[0]
|
|
lap = -graph.copy()
|
|
lap.flat[::n_nodes + 1] = 0
|
|
w = -lap.sum(axis=0)
|
|
if normed:
|
|
w = np.sqrt(w)
|
|
w_zeros = w == 0
|
|
w[w_zeros] = 1
|
|
lap /= w
|
|
lap /= w[:, np.newaxis]
|
|
lap.flat[::n_nodes + 1] = 1 - w_zeros
|
|
else:
|
|
lap.flat[::n_nodes + 1] = w
|
|
if return_diag:
|
|
return lap, w
|
|
return lap
|
|
|
|
|
|
def graph_laplacian(graph, normed=False, return_diag=False):
|
|
""" Return the Laplacian of the given graph.
|
|
"""
|
|
if normed and (np.issubdtype(graph.dtype, np.int)
|
|
or np.issubdtype(graph.dtype, np.uint)):
|
|
graph = graph.astype(np.float)
|
|
if sparse.isspmatrix(graph):
|
|
return _graph_laplacian_sparse(graph, normed=normed,
|
|
return_diag=return_diag)
|
|
else:
|
|
# We have a numpy array
|
|
return _graph_laplacian_dense(graph, normed=normed,
|
|
return_diag=return_diag)
|