Ladderizing a tree in a tree sequence

[hyanwong]

Sometimes we will want to draw a tree in a ladderised style, placing the clades with more tips either first or last. This isn't easy to do at the moment. However, if we reorder the leaf nodes in the right way, the standard minlex ordering will do it for us. Here's an implementation:

def reorder_ladderized(ts, pos=None, right=True, map_nodes=False):
    """
    Reorder nodes in a tree sequence such that when plotted
    with the standard "minlex" order, the tree at position
    ``pos`` in the returned tree sequence is shown ladderized,
    with clades containing more tips placed on the right
    (if ``right=True``; default) or on the left (otherwise).

    Note that the returned tree sequence will have changed node IDs
    (as the routine uses {meth}`TreeSequence.subset`). To maintain
    node identities, use metadata, or specify 
    """
    n = ts.num_nodes
    if pos is None:
        pos = min(ts.edges_left)
    tree = ts.at(pos)
    is_leaf = np.zeros(n, dtype=bool)
    num_leaves = np.zeros(n + 1)  # Add an extra on the end to catch `-1`

    # ascend up the tree, tallying number of leaves underneath each node
    for u in tree.nodes(order="postorder"):
        if tree.is_leaf(u):
            is_leaf[u] = True
            num_leaves[u] = 1
        num_leaves[tree.parent(u)] += num_leaves[u]

    # pass back down the tree, summing these up, to get an order
    for u in tree.nodes(order="preorder"):
        num_leaves[u] += tree.parent(u)

    num_leaves = num_leaves[:-1] * (-1 if right else 1)
    node_map = np.arange(n)
    node_map[is_leaf] = node_map[is_leaf][np.argsort(num_leaves[is_leaf])]
    ts = ts.subset(node_map, remove_unreferenced=False)
    if map_nodes:
        return ts, node_map
    else:
        return ts

Demo, reordering the (default) right-ladderised comb tree:

import tskit
ts = tskit.Tree.generate_comb(10).tree_sequence
display(ts.first().draw_svg())
new_ts, node_map = reorder_ladderized(ts, right=False, map_nodes=True)
new_ts.first().draw_svg(node_labels = {u: v for u, v in enumerate(node_map)})

Note that the ordering should stay the same in the case of ties (hence 8 and 9 below are in the same order)