generic algebraic graph lib

graph.py

import numpy as np
import pylab as pl

def check_edges(edges):
    '''Check that no edge is a loop or that has a multiple edge.'''
    sorted_edges = []
    out = True
    try:
        for i, j in edges:
            assert i != j, "No loop allowed in Graph"
            i, j = min([i, j]), max([i, j])
            s_edge = np.array([i, j])
            ## check that s_edge is not already in sorted_edges
            if len(sorted_edges) > 0:
                assert min([np.linalg.norm(e - s_edge) for e in sorted_edges]) > 1e-8, 'Duplicate edge not allowed'
            sorted_edges.append(s_edge)
            
    except AssertionError:
        out = False
    return out

class Graph:
    def __init__(self, nodes, edges):
        self.nodes = nodes
        self.edges = edges
        self.n_node = len(nodes)
        self.n_edge = len(edges)
        self.__adj = self.getAdj()
        self.__inc = self.getInc()
        assert check_edges(edges)
                
    def __str__(self):
        return "Graph(%s, %s)" % (self.nodes, self.edges)

    def getAdj(self):
        if hasattr(self, '__adj'):
            return self.__adj ### only compute adj once
        out = np.zeros((self.n_node, self.n_node), bool)
        for i, j in self.edges:
            out[i, j] = True
            # out[j, i] = True
        out += out.T
        return out

    def getInc(self):
        if hasattr(self, '__inc'):
            return self.__inc ### only compute inc once
        out = np.zeros((self.n_node, self.n_edge), bool)
        for idx, edge in enumerate(self.edges):
            out[edge[0], idx] = True
            out[edge[1], idx] = True
        return out

    def getLap(self):
        adj = self.getAdj().astype(np.int16)
        out = np.diag(np.sum(adj, axis=1)) ### row sums!
        out -= adj
        return out

    def getDInc(self):
        out = self.__inc.astype(np.int16)
        for i in range(self.n_edge):
            col_i = out[:,i]
            idx = np.where(col_i)[0][0]
            out[idx, i] *= -1
        return out
    adj = property(getAdj)
    inc = property(getInc)
    dinc = property(getDInc)
    lap = property(getLap)
    
    def plot(self, type='Circle', poss=None):
        if poss is not None:
            type = 'Given'
            
        if type == 'Circle':
            theta = np.arange(0, 2 * np.pi, 2 * np.pi / self.n_node)
            xs = np.sin(theta)
            ys = np.cos(theta)
        elif type == 'Random':
            xs = np.random.uniform(0, 1, self.n_node)
            ys = np.random.uniform(0, 1, self.n_node)
        elif type == 'Given':
            xs = poss[:,0]
            ys = poss[:,1]
        else:
            raise ValueError('Plot type %s not supported.' % type)

        for edge in self.edges:
            i, j = edge
            firstx = xs[i]
            secondx = xs[j]
            firsty = ys[i]
            secondy = ys[j]
            pl.plot([firstx, secondx], [firsty, secondy], 'b-')
        pl.plot(xs, ys, 'ko')

    def is_connected(self):
        pwr_adj = np.zeros(self.adj.shape)
        a_n = np.eye(self.n_node)
        for dummy in range(self.n_node):
            pwr_adj += a_n
            a_n = a_n @ self.adj
        pwr_adj += a_n
        return np.amin(pwr_adj) > 1e-8
        
def gen_random_graph(nodes, n_edge):
    N = len(nodes)
    max_edges = N * (N - 1) / 2
    assert n_edge <= max_edges, 'Too many edges; the maximal number of edges on %d nodes is %d.' % (N, max_edges)
    n_node = len(nodes)
    edges = []
    while len(edges) < n_edge:
        i = np.random.randint(n_node)
        j = np.random.randint(n_node)
        edge = (i, j)
        edges.append(edge)
        if not check_edges(edges):
            edges = edges[:-1]
    return Graph(nodes, edges)

def graph_from_adj(adj):
    shp = np.shape(adj)
    N = shp[0]
    assert len(shp) == 2, 'Only 2D adj allowed'
    assert shp[0] == shp[1], 'Only square adj allowed'
    assert np.linalg.norm(adj - adj.T) < 1e-8, 'Only symmetric adj allowed'
    assert np.linalg.norm([adj[i, i] for i in range(N)]) < 1e-8, 'Zeros required on diagonal of adj.'

    edges = []
    for i in range(N - 1): ### last ele on diag is 0
        for j in range(i + 1, N):
            if adj[i, j]:
                edges.append((i, j))
    nodes = list(range(N))
    return Graph(nodes, edges)

def gen_complete_graph(n_node):
    nodes = range(n_node)
    adj = np.ones((n_node, n_node), bool) - np.eye(n_node)
    return graph_from_adj(adj)

def test():
    g = gen_random_graph('abcd', 3)
    adj = g.adj
    g2 = graph_from_adj(adj)
    
    c = gen_complete_graph(4)
    print(c.lap)

if __name__ == '__main__':
    test()