How to draw Caterpillar Trees in Networkx beautifully?

Question

I want to draw Caterpillar trees using Networkx. Is there a way to draw in a way that the spine is in a straight line and the leaves are clearly visible?

The image from Wikipedia is a good example of something I want to draw.

Answer 1

We can generate a caterpillar network by using the random_lobster network generation function, with parameter p2 held to zero.

import itertools
import networkx as nx
import matplotlib.pyplot as plt

spine_sz = 5
G = nx.random_lobster(spine_sz, 0.65, 0, seed=7)

The second part, drawing the network in such a way that the spine is drawn in a straight line with the leaves clearly visible requires quite some calculation.

First, we have to determine what the spine is of the generated network. Since the network is a tree, this is not difficult. We'll use the first path that is as long as the diameter.

eccs = nx.eccentricity(G)
diameter = max(eccs.values())
outliers = tuple(node for node in eccs if eccs[node] == diameter)
for combo in itertools.combinations(outliers, 2):
    path = nx.shortest_path(G, combo[0], combo[1])
    if len(path) == diameter+1:
        break

Second, we have to estimate the space needed for the network, assuming that we'll draw the leaves above and below the spine, as evenly distributed as possible.

def node_space(nleaves):
    return 1+int((nleaves-1)/2)

nb_leaves = {}
for node in path:
    neighbors = set(G.neighbors(node))
    nb_leaves[node] = neighbors - set(path)
max_leaves = max(len(v) for v in nb_leaves.values())
space_needed = 1 + sum(node_space(len(v)) for v in nb_leaves.values())

Third, with this spacing, we can draw the spine, and with each node of the spine, draw the leaves around it, distributing them above and below the node, while also distributing them horizontally with in the "spine space" reserved. We'll use the node number being even to determine if it will appear above or below the spine, alternating this for each subsequent node with leaves; the latter by reversing the yhop variable.

def count_leaves(nleaves, even):
    leaf_cnt = int(nleaves/2)
    if even:
        leaf_cnt += nleaves % 2
    return leaf_cnt

def leaf_spacing(nleaves, even):
    leaf_cnt = count_leaves(nleaves, even)
    if leaf_cnt <= 1:
        return 0
    return 1 / (leaf_cnt-1)

xhop = 2 / (space_needed+2)
yhop = 0.7
pos = {}
xcurr = -1 + xhop/4
for node in path:
    pos[node] = (xcurr, 0)
    if len(nb_leaves[node]) > 0:
        leaves_cnt = len(nb_leaves[node])
        extra_cnt = node_space(leaves_cnt) - 1
        extra_space = xhop * extra_cnt / 2
        xcurr += extra_space
        pos[node] = (xcurr, 0)
        l0hop = 2 * extra_space * leaf_spacing(leaves_cnt, True)
        l1hop = 2 * extra_space * leaf_spacing(leaves_cnt, False)
        l0curr = xcurr - extra_space
        l1curr = xcurr - extra_space
        if l1hop == 0:
            l1curr = xcurr
        for j,leaf in enumerate(nb_leaves[node]):
            if j % 2 == 0:
                pos[leaf] = (l0curr, yhop)
                l0curr += l0hop
            else:
                pos[leaf] = (l1curr, -yhop)
                l1curr += l1hop
        yhop = -yhop
        xcurr += xhop * extra_cnt / 2
    prev_leaves = len(nb_leaves[node])
    xcurr += xhop

The final part is the actual plotting of the graph, with some options set for color and size.

options = {
    "font_size": 9,
    "node_size": 300,
    "node_color": "lime",
    "edgecolors": "black",
    "linewidths": 1,
    "width": 1,
}
nx.draw_networkx(G, pos, **options)
ax = plt.gca()
ax.margins(0.10)
plt.axis("off")
plt.show()

The resulting graph looks like this: