Algorithm Implementation/Graphs/Maximum flow/Simulation s-t Cut

Python

import networkx
import numpy
import scipy

#   This software is an implementation of the invention in US Patent 8929363
#   "Method and System for Image Segmentation".
#   The software computes an approximation to the minimum s-t cut using the
#   Simulation s-t cut algorithm.
#   The software applies to the partitioning of undirected graphs provided
#   two seed nodes "s" and "t".

class Circuit(networkx.Graph):
    def matrix(self):
        edges = networkx.get_edge_attributes(self,'weight')
        size = len(self.nodes())
        rows = numpy.array(edges.keys())[:,0]
        columns = numpy.array(edges.keys())[:,1]
        weights = numpy.array(edges.values())    
        matrix = scipy.sparse.coo_matrix((weights,(rows,columns)),shape=(size, size)).tocsr()
        return matrix + matrix.transpose()

    def flow(self):
        vector = numpy.zeros(len(self.nodes()))
        nodes = networkx.get_node_attributes(G, 'flow')
        for (node, flow) in nodes.iteritems():
            if (flow == 's'):
                vector[node] = 1
            if (flow == 't'):
                vector[node] = -1
        return vector

    def s(self):
        nodes = networkx.get_node_attributes(G, 'flow')
        for (node, flow) in nodes.iteritems():
            if (flow == 's'):
                s = node
        return s

    def t(self):
        nodes = networkx.get_node_attributes(G, 'flow')
        for (node, flow) in nodes.iteritems():
            if (flow == 't'):
                t = node
        return t

class Nonlinear:
    def __init__(self, weights):
        self.weights = weights

    def set_voltage(self, voltage):
        self.voltage = voltage

    def linearize(self):
        nodes = self.weights.shape[0]
        entries = self.weights.nnz
        data = self.weights.tocoo().data
        rows = self.weights.tocoo().row
        columns = self.weights.tocoo().col
        ones = numpy.ones(entries)
        negative_ones = -1.0 * numpy.ones(entries)
        positive = scipy.sparse.coo_matrix((ones,(range(0,entries),rows)),shape=(entries, nodes)).tocsr()
        negative = scipy.sparse.coo_matrix((ones,(range(0,entries),columns)),shape=(entries, nodes)).tocsr()
        subtract = positive - negative
        difference = subtract * self.voltage
        C = numpy.divide(self.weights.data, ones + numpy.absolute(difference))
        matrix = scipy.sparse.coo_matrix((C,(rows,columns)),shape=(nodes, nodes)).tocsr()
        return matrix

G = Circuit()
G.add_node(0, flow = 's')
G.add_node(1)
G.add_node(2)
G.add_node(3)
G.add_node(4, flow = 't')
G.add_edge(0, 1, weight= 5)
G.add_edge(0, 2, weight= 3)
G.add_edge(0, 3, weight= 4)
G.add_edge(0, 4, weight= 7)
G.add_edge(1, 2, weight= 10)
G.add_edge(1, 3, weight= 1)
G.add_edge(1, 4, weight= 6)
G.add_edge(2, 3, weight= 9)
G.add_edge(2, 4, weight= 8)
G.add_edge(3, 4, weight= 2)
x = numpy.zeros(len(G.nodes()))
ones = numpy.ones(len(G.nodes()))
W = G.matrix()
rowsum = W * ones
D = scipy.sparse.diags(rowsum,0)
F = D - W
f = G.flow()
s = G.s()
t = G.t()
N = Nonlinear(W)
for iteration in range(0,20):
    N.set_voltage(x)
    L = N.linearize()
    rowsum = L * ones
    D = scipy.sparse.diags(rowsum,0)
    A = D - L
    A.data[A.indptr[s]:A.indptr[s+1]] = 0
    A.data[A.indptr[t]:A.indptr[t+1]] = 0
    A = A + scipy.sparse.diags(numpy.abs(f), 0)
    x = scipy.sparse.linalg.spsolve(A, 1000000.0 * f)
    segmentation = numpy.sign(x)
    cut = F * segmentation
    cut = numpy.multiply(cut, segmentation)
    flow = numpy.sum(cut)
    print 0.25 * flow