segments_networkx.py

#!/usr/bin/env python3

import networkx as nx
from networkx.drawing import nx_agraph
import sys

G = nx.read_edgelist(sys.argv[1], create_using = nx.DiGraph, nodetype = int)
nbunch = [] # list of nodes to remove

for n in G: G.nodes[n]["inst_eddies"] = [n]
# "inst_eddies" is the list of instantaneous eddies represented by a
# node. At first, each node represents only itself. At the end of the
# script, each node will be a segment and "inst_eddies" will contain
# the list of eddies in the segment.
        
for n in G:
    if G.in_degree[n] == 1:
        n2 = next(G.predecessors(n))

        if G.out_degree[n2] == 1:
            # n continues a segment, circumvent n:

            nbunch.append(n)
            G.remove_edge(n2, n)

            for n3 in list(G.successors(n)):
                G.remove_edge(n, n3)
                G.add_edge(n2, n3)

            # Add the segments represented by n to the segments
            # represented by n2:
            G.nodes[n2]["inst_eddies"] += G.nodes[n]["inst_eddies"]

G.remove_nodes_from(nbunch)
for k in sorted(G.nodes): print(k, G.nodes[k])
A = nx_agraph.to_agraph(G)
A.write(sys.argv[2])