Skip to content

Node-link tree layout optimization

Lays out a rooted tree (a NetworkX arborescence) with TreeLayoutOptimizer: nodes at the same depth align along a growth axis, parents are centered over their children, and siblings/subtrees keep a minimum separation in their fixed left-to-right order (so edges never cross).

import networkx as nx
from IPython.display import SVG
from matplotlib import pyplot as plt

from vizopt.animation import SnapshotCallback
from vizopt.base import OptimConfig
from vizopt.templates.trees.tree_layout import TreeLayoutOptimizer

A small tree

An unbalanced tree (varied branching factor) to exercise order_separation and parent_centering together.

tree = nx.DiGraph()
tree.add_edges_from(
    [
        ("root", "a"),
        ("root", "b"),
        ("root", "c"),
        ("a", "a1"),
        ("a", "a2"),
        ("a", "a3"),
        ("b", "b1"),
        ("c", "c1"),
        ("c", "c2"),
        ("a1", "a1x"),
        ("a1", "a1y"),
    ]
)

optimizer = TreeLayoutOptimizer(tree, min_distance=1.0, layer_spacing=1.5)
optimizer.optimize(OptimConfig(n_iters=2000, learning_rate=1e-2))

optimizer.plot()
plt.title("Tree layout (top-down)")
plt.show()

output

Animating the optimization

snapshot_cb = SnapshotCallback(every=25)
optimizer = TreeLayoutOptimizer(tree, min_distance=1.0, layer_spacing=1.5)
optimizer.optimize(OptimConfig(n_iters=2000, learning_rate=1e-2), callback=snapshot_cb)

svg = optimizer.animate_svg(snapshot_cb, fps=12, size=500)
SVG(data=svg)

output

Left-to-right orientation

growth_direction controls the axis depth grows along — (1, 0) gives a left-to-right dendrogram-style layout instead of the default top-down one.

optimizer = TreeLayoutOptimizer(
    tree, min_distance=1.0, layer_spacing=2.0, growth_direction=(1.0, 0.0)
)
optimizer.optimize(OptimConfig(n_iters=2000, learning_rate=1e-2))

optimizer.plot()
plt.title("Tree layout (left-to-right)")
plt.show()

output

A larger random tree

nx.random_labeled_tree produces an undirected tree; orient it away from an arbitrary root with nx.bfs_tree to get the parent → child arborescence TreeLayoutOptimizer expects.

undirected = nx.random_labeled_tree(30, seed=0)
big_tree = nx.bfs_tree(undirected, source=0)

optimizer = TreeLayoutOptimizer(big_tree, min_distance=0.8, layer_spacing=1.2)
optimizer.optimize(OptimConfig(n_iters=3000, learning_rate=1e-2))

optimizer.plot()
plt.title("Random tree (30 nodes)")
plt.show()

output