mcf_sparse.py

import numpy as np
import networkx as nx
from scipy.spatial import Delaunay

def W(x):    
    return (x + np.pi) % (2 * np.pi) - np.pi

def mcf_sparse(x, y, psi, capacity=None):
    points = np.vstack((x, y)).T
    num_points = points.shape[0]

    # Delaunay triangularization
    tri = Delaunay(points)
    simplex = tri.simplices
    simplex_neighbors = tri.neighbors
    num_simplex = simplex.shape[0]

    if capacity is None:
        cap_args = dict()
    else:
        cap_args = {'capacity': capacity}

    # get pseudo estimation of gradients and vertex edges
    psi_diff = W(psi[simplex] - psi[np.roll(simplex, 1, 1)])
    edges = np.stack((np.roll(simplex, 1, 1), simplex), axis=2).reshape(-1, 2)

    # get corresponding simplex's edges orthogonal to original vertex edges
    simplex_edges = np.stack((
        np.broadcast_to(np.arange(num_simplex)[:, None], simplex_neighbors.shape),
        np.roll(simplex_neighbors, -1, 1)), axis=2
    ).reshape(-1, 2)

    # get demands
    demands = np.round(-psi_diff.sum(1) * 0.5 / np.pi).astype(np.int)
    psi_diff = psi_diff.flatten()

    G = nx.Graph()
    G.add_nodes_from(zip(range(num_simplex),
                         [{'demand': d} for d in demands]))

    # set earth node index to -1, and its demand is the negative of the sum of all demands,
    # so the total demands is zero
    G.add_node(-1, demand=-demands.sum())

    # set the edge weight to 1 whenever one of its nodes has zero demand
    demands_dummy = np.concatenate((demands, [-demands.sum()]))
    weights = np.any(demands_dummy[simplex_edges] == 0, 1)
    G.add_weighted_edges_from(zip(simplex_edges[:, 0], simplex_edges[:, 1], weights), **cap_args)

    # make graph to directed graph, so we can distinguish positive and negative flow
    G = G.to_directed()

    # perform MCF
    cost, flowdict = nx.network_simplex(G)

    # construct K matrix with the same shape as the edges
    K = np.empty(edges.shape[0])

    # add the flow to their orthogonal edge
    for i, (u, v) in enumerate(simplex_edges):
        K[i] = -flowdict[u][v] + flowdict[v][u]

    # derive correct gradients
    psi_diff += K * 2 * np.pi

    # construct temporary dict that hold all edge gradients from different direction
    psi_dict = {i: {} for i in range(num_points)}
    for diff, (u, v) in zip(psi_diff, edges):
        psi_dict[u][v] = diff

    # integrate the gradients
    result = psi.copy()

    # choose an integration path, here let's use BFS
    traverse_G = nx.DiGraph(edges.tolist())
    for u, v in nx.algorithms.traversal.breadth_first_search.bfs_edges(traverse_G, 0):
        result[v] = result[u] + psi_dict[u][v]
    return result