Animate trees evolving along a tree sequence

animate_edges.py

import pathlib
import subprocess

import msprime


def draw_partial_tree(tables, edges, filename):
    tables.edges.clear()
    for edge in edges:
        tables.edges.add_row(*edge)
    tables.sort()
    tmp_ts = tables.tree_sequence()
    tree = tmp_ts.first()
    tree.draw(
        filename, format="svg", width=400, height=400,
        # These are not-yet-merged and experimental features of Tree.draw().
        # See https://github.com/tskit-dev/tskit/pull/168
        tree_height_scale="rank", max_tree_height="ts")


def draw_intermediate_trees(ts, path_pattern):
    tables = ts.dump_tables()
    edges = []
    for j, (_, edges_out, edges_in) in enumerate(ts.edge_diffs()):
        k = 0
        for edge in edges_out:
            edges.remove((0, 1, edge.parent, edge.child))
            k += 1
            draw_partial_tree(tables, edges, path_pattern.format(j, k))
        for edge in edges_in:
            edges.append((0, 1, edge.parent, edge.child))
            k += 1
            draw_partial_tree(tables, edges, path_pattern.format(j, k))


ts = msprime.simulate(15, recombination_rate=4, random_seed=1)
print(f"tree sequence with {ts.num_trees} trees")
outdir = pathlib.Path("tmp")
draw_intermediate_trees(ts, str(outdir / "tree_{:03d}_{:03d}.svg"))
for svgfile in outdir.glob("*"):
    giffile = outdir / (svgfile.stem + ".gif")
    subprocess.run(
        f"convert -background white -flatten {svgfile} {giffile}",
        shell=True, check=True)
subprocess.run(
    f"convert -delay 20 -loop 0 -dispose previous {outdir}/*.gif "
    "animation.gif", shell=True, check=True)