Augment your Graph with an Augmenting Path

augment.py

def augment(augmentingPath, H):
    augmentingPath = list(augmentingPath)
    H_as_dict = digraph_to_dict(H)
    new_nodes = frozenset({})
    new_arcs = frozenset({})
    visited_nodes  = frozenset({})
    visited_arcs = frozenset({})

    bottleneck_residualCapacity = min( augmentingPath, key=lambda a: a.datum.residualCapacity ).datum.residualCapacity
    for x in augmentingPath:
        from_node_uid = x.fromNode.uid if x.datum.action == 'push' else x.toNode.uid
        to_node_uid = x.toNode.uid   if x.datum.action == 'push' else x.fromNode.uid
        from_node = find_node_by_uid( from_node_uid, H )
        to_node = find_node_by_uid( to_node_uid, H )
        load = bottleneck_residualCapacity if x.datum.action == 'push' else -bottleneck_residualCapacity

        for a in H_as_dict[from_node]:
            if(a.toNode == to_node):
                test_sum = a.datum.flow + load
                new_flow = a.datum.flow
                new_from_node_flow_out = a.fromNode.datum.flowOut
                new_to_node_flow_in = a.toNode.datum.flowIn

                new_from_look = {n for n in new_nodes if (n.uid == a.fromNode.uid)}
                new_to_look = {n for n in new_nodes if (n.uid == a.toNode.uid)  }
                prev_from_node = new_from_look.pop() if (len(new_from_look)>0) else a.fromNode
                prev_to_node = new_to_look.pop()   if (len(new_to_look)>0)   else a.toNode
                new_nodes = new_nodes.difference(frozenset({prev_from_node}))
                new_nodes = new_nodes.difference(frozenset({prev_to_node}))

                if(test_sum > a.datum.capacity):
                    new_flow = a.datum.capacity
                    new_from_node_flow_out = prev_from_node.datum.flowOut - a.datum.flow + a.datum.capacity
                    new_to_node_flow_in    = prev_to_node.datum.flowIn - a.datum.flow + a.datum.capacity
                    load = test_sum - a.datum.capacity
                elif(test_sum < 0.0):
                    new_flow = 0.0
                    new_from_node_flow_out = prev_from_node.datum.flowOut - a.datum.flow
                    new_to_node_flow_in    = prev_to_node.datum.flowIn - a.datum.flow
                    load = test_sum
                 else:
                    new_flow = test_sum
                    new_from_node_flow_out = prev_from_node.datum.flowOut - a.datum.flow + new_flow
                    new_to_node_flow_in    = prev_to_node.datum.flowIn - a.datum.flow + new_flow
                    load = 0.0

                new_from_node_flow_out = round(new_from_node_flow_out, TOL)
                new_to_node_flow_in    = round(new_to_node_flow_in, TOL)
                new_flow               = round(new_flow, TOL)

                new_from_node = Node(prev_from_node.uid, FlowNodeDatum(prev_from_node.datum.flowIn, new_from_node_flow_out))
                new_to_node   = Node(prev_to_node.uid, FlowNodeDatum(new_to_node_flow_in, prev_to_node.datum.flowOut))
                new_arc       = Arc(new_from_node, new_to_node, FlowArcDatum(a.datum.capacity, new_flow))

                visited_nodes = visited_nodes.union(frozenset({a.fromNode,a.toNode}))
                visited_arcs  = visited_arcs.union(frozenset({a}))
                new_nodes     = new_nodes.union(frozenset({new_from_node, new_to_node}))
                new_arcs      = new_arcs.union(frozenset({new_arc}))

    not_visited_nodes = H.setOfNodes.difference(visited_nodes)
    not_visited_arcs  = H.setOfArcs.difference(visited_arcs)

    full_new_nodes = new_nodes.union(not_visited_nodes)
    full_new_arcs  = new_arcs.union(not_visited_arcs)
    G = DiGraph(full_new_nodes, full_new_arcs)

    full_new_arcs_update = frozenset([])
    for a in full_new_arcs:
        new_from_node = a.fromNode
        new_to_node   = a.toNode
        new_from_node = find_node_by_uid( a.fromNode.uid, G )
        new_to_node   = find_node_by_uid( a.toNode.uid, G )
        full_new_arcs_update = full_new_arcs_update.union( {Arc(new_from_node, new_to_node, FlowArcDatum(a.datum.capacity, a.datum.flow))} )

    G = DiGraph(full_new_nodes, full_new_arcs_update)

    return G