grouping nodes in networkx

Question

I don't know the proper terminology, but I have a digraph in networkx and I want to merge related nodes using a mapping function, maintaining a link to the original nodes

For example, if I had this graph:

a1 -> b1
a2 -> b2
a1 -> c1
a2 -> d2
b1 -> d2
c1 -> f1
c2 -> f2
b1 -> c2

and this function

def mymap(node):
    return node[0]

then I would want to get this graph as a result:

a -> b
a -> c
a -> d
b -> d
c -> f
b -> c

with node data

a: original_nodes = [a1, a2]
b: original_nodes = [b1, b2]
c: original_nodes = [c1, c2]
d: original_nodes = [d2]
f: original_nodes = [f1, f2]

Is there a way to do this? (without manually iterating over nodes and edges myself; I can do that, but I would think this would be a common task)

Answer 1

hmm, I thought I'd remembered doing this before, just found this in one of my Python scripts:

def supergraph(g1, keyfunc, allow_selfloops=True):
    g2 = nx.DiGraph()
    for (a,b,d) in g1.edges_iter(data=True):
        result = keyfunc(g1,a,b,d)
        if result is not None:
            a2,b2,w = result
            if a2 != b2 or allow_selfloops:
                g2.add_edge(a2,b2)
                try:
                    g2[a2][b2]['weight'] += w
                except:
                    g2[a2][b2]['weight'] = w

            for u2,u in [(a2,a),(b2,b)]:
                if not u2 in g2:
                    g2.add_node(u2, original_nodes=set([u]))
                else:
                    try:
                        g2.node[u2]['original_nodes'].add(u)
                    except:
                        g2.node[u2]['original_nodes'] = set([u])
    return g2

where keyfunc(g,a,b,d) with graph g, nodes a and b, and edge data d, is a mapping function that is supposed to return either None (to ignore that edge) or a tuple (a2, b2, w) with new nodes a2 and b2, and weight w.