Animate trees evolving along a chromosome

better_animate_edges.py

import pathlib
import subprocess
import tempfile

import tskit
import msprime


def draw_transitioning_tree(tree, filename, edge_colours):
    tree.draw(
        filename, format="svg",
        edge_colours=edge_colours,
        # These two parameters are experimental and unmerged.
        tree_height_scale="rank", max_tree_height="ts",
        width=400, height=400)


def draw_intermediate_trees(ts, path_pattern):
    tree = tskit.Tree(ts)
    for j, (_, edges_out, edges_in) in enumerate(ts.edge_diffs()):
        k = 0
        edge_colours = {}
        draw_transitioning_tree(tree, path_pattern.format(j, k), edge_colours)
        for edge in edges_out:
            edge_colours[edge.child] = "red"
            k += 1
            draw_transitioning_tree(tree, path_pattern.format(j, k), edge_colours)
            edge_colours[edge.child] = None
        tree.next()
        edge_colours = {edge.child: None for edge in edges_in}
        for edge in edges_in:
            edge_colours[edge.child] = "green"
            k += 1
            draw_transitioning_tree(tree, path_pattern.format(j, k), edge_colours)
            del edge_colours[edge.child]


ts = msprime.simulate(15, recombination_rate=4, random_seed=1)
print(f"tree sequence with {ts.num_trees} trees")
with tempfile.TemporaryDirectory() as tmpdir:
    outdir = pathlib.Path(tmpdir)
    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)