Programmatically specifying nodes of the same rank within networkx's wrapper for pygraphviz/dot

Question

Is it possible to alter the following code to put Child_4 at the same horizontal level as Grandchild_0 (thereby pushing Grandchild_4 to its own level)?

import networkx as nx
import matplotlib.pyplot as plt

G = nx.DiGraph()
G.add_node("ROOT")
for i in xrange(5):
    G.add_node("Child_%i" % i)
    G.add_node("Grandchild_%i" % i)
    G.add_edge("ROOT", "Child_%i" % i)
    G.add_edge("Child_%i" % i, "Grandchild_%i" % i)

pos=nx.graphviz_layout(G,prog='dot')
nx.draw(G,pos,arrows=False)
plt.show()

The above code produces the following layout, which I'd like to alter by shifting a child down one level to be horizontally aligned with the grandchildren:

---

Within the Python network library networkx, I'm using graphviz's dot engine to render a tree (following this recommendation). I would like to control the y-position of the nodes by specifying which nodes should have the same height. The nodes might be at different depths in the tree.

I know I could control the node height if I wrote my own graphviz code through using the rank=same command (e.g., {rank=same; n4 -> p2;} [ex.]). However, I am relying on networkx.graphviz_layout() [doc | source] to generate the node positions, and graphviz_layout can send only command line arguments to pygraphviz. My attempts to use variants of nx.graphviz_layout(G, prog='dot', args="-Grank=same; n4 -> p2;") have failed. Is it possible to describe the desired node heights within the NetworkX wrapper for pygraphviz, or do I need to write my own wrapper around pygraphviz? Edit: The answer provides a new wrapper around pygraphviz. It would significantly simplify things to send the rank information within the existing NetworkX wrapper for pygraphviz. I'll change my accepted answer if someone can tell me how that might be possible.

Answer 1

I can't find a way to achieve this through the original networkx wrapper.

Instead, I've written a new wrapper for pygraphviz, with most lines copied from the source code. It adds a parameter sameRank = [] for a list of nodes-of-the-same-rank lists and a for loop around an invocation of pygraphviz.add_subgraph(listOfNodes,rank="same").

def graphviz_layout_with_rank(G, prog = "neato", root = None, sameRank = [], args = ""):
    ## See original import of pygraphviz in try-except block
    ## See original identification of root through command line
    A = nx.to_agraph(G)
    for sameNodeHeight in sameRank:
        if type(sameNodeHeight) == str:
            print("node \"%s\" has no peers in its rank group" %sameNodeHeight)
        A.add_subgraph(sameNodeHeight, rank="same")
    A.layout(prog=prog, args=args)
    ## See original saving of each node location to node_pos 
    return node_pos

In the question example, Child_4 can be pushed to the same horizontal level as Grandchild_0 through the line:

pos=graphviz_layout_with_rank(G, prog='dot',sameRank=[["Child_4","Grandchild_0"]])