Vignette: Parallelizing SLiM simulations in a phylogenetic tree

Imagine you want to simulate the evolutionary history of a group. If there is no migration between any of the branches in your tree, any branches stemming from the same node can be simulated in parallel (see Fig. 7).

Fig. 7 Example of phylogeny we might want to simulate. Note how branches with the same color can be simulated in parallel when there is no migration.

To do this, we’ll need to do two things: (1) be able to simulate branches in parallel, and (2) glue the resulting simulations (one per tip) back together.

Simulating the branches

First, we need to write a SLiM script that will be used for simulating the history of each branch in our phylogeny. We will perform a simple simulation, in which each branch can have a different (but fixed) population size and length (number of ticks). Also, we will allow deleterious mutations to happen across the entire chromosome at a fixed rate.

Here is a SLiM script that would do this:

// The following constants need to be defined:
// - outfile: path to save the tree sequence of the simulation.
// - popsize: population size.
// - num_gens: run simulation for num_gens ticks.
// - infile: a string with path to tree sequence to start from; optional.
initialize()
{
    initializeSLiMModelType("WF");
    initializeTreeSeq();
    initializeMutationRate(1e-8);
    initializeMutationType("m1", 0.5, "f", -0.01);
    initializeGenomicElementType("g1", m1, 0.1);
    initializeGenomicElement(g1, 0, 1e6-1);
    initializeRecombinationRate(1e-9);
}

1 late() {
    // if no input tree sequence is provided, then start a subpopulation
    if (infile == "") {
        p = sim.addSubpop("p1", popsize);
    } else {
        // reloading must happen in late()
        sim.readFromPopulationFile(infile);
        parent = sim.subpopulations[0];
        p = sim.addSubpopSplit(max(sim.subpopulations.id) + 1, popsize, parent);
        parent.setSubpopulationSize(0);
    }
    p.name = popname;
}

// schedule the end of the simulation
1 late() {
    finaltick = num_gens + community.tick;
    community.rescheduleScriptBlock(s0, ticks=finaltick);
}

// event that saves the tree sequence
s0 1000 late() {
    sim.treeSeqRememberIndividuals(sim.subpopulations.individuals);
    sim.treeSeqOutput(outfile);
}

For each branch, the presence or absence of infile tells SLiM whether we want to start it from a previous branch or not. If so, SLiM will read the previous tree sequence and change the population size accordingly. Note that when you read a tree sequence into SLiM, the tick counter will be updated with the time encoded in the tree sequence, so we need to set the end of the simulation as the length of the branch (num_gens) plus the current “time” at the end of the loaded tree sequence. At the end of the simulation, we call sim.treeSeqRememberIndividuals right before saving the resulting tree sequence. This is necessary because we need to ensure the individuals in the final generation are never dropped from the tree sequence in future runs of SLiM which are started from the output of the simulation, as they will later be used to glue the tree sequences together.

I encoded the phylogeny we will simulate in a simple table, which we’ll use as df in the code below:

Show code cell source Hide code cell source

df = pd.read_csv("_static/phylo.tsv", sep="\t")
df = df.fillna('')
df["infile"] = df.parent + ".trees"
df["outfile"] = df.child + ".trees"
df.loc[df["infile"]==".trees", "infile"] = ""
df["is_leaf"] = ~df.child.isin(df.parent)
df

	child	parent	popsize	edgelen	infile	outfile	is_leaf
0	root		500	2000		root.trees	False
1	C	root	50	250	root.trees	C.trees	True
2	I	root	100	200	root.trees	I.trees	False
3	B	I	70	50	I.trees	B.trees	True
4	A	I	40	50	I.trees	A.trees	True

With our phylogeny and the simulation parameters, we are ready to run our simulations. One way to parallelize the simulation of sister branches is to use make. You do not need to know much about this tool (though it is totally worthwhile to check it out). The main idea here is that you can specify dependency between files and make works its magic to run the simulations in the right order. Here is python code that will write out a makefile from the information in df:

f = open("sims.make", "w")
print(f"all: {' '.join(df.outfile.to_list())}\n", file=f)
for i, row in df.iterrows():
    print(f"{row.outfile}: {row.infile} phylo_bgs.slim", file=f)
    print(f"\tslim -d \"infile='{row.infile}'\" -d popsize={row.popsize} "
          f"-d \"popname=\'{row.child}\'\" "
          f"-d num_gens={row.edgelen} " f"-d \"outfile='{row.child}.trees'\" "
          "phylo_bgs.slim\n",
          file=f)
f.close()

Here’s the result. Again, don’t worry about the details, but you can see that the file encodes the phylogeny through a bunch of child : parent “rules”:

Show code cell source Hide code cell source

%%bash
cat sims.make

all: root.trees C.trees I.trees B.trees A.trees

root.trees:  phylo_bgs.slim

	slim -d "infile=''" -d popsize=500 -d "popname='root'" -d num_gens=2000 -d "outfile='root.trees'" phylo_bgs.slim

C.trees: root.trees phylo_bgs.slim

	slim -d "infile='root.trees'" -d popsize=50 -d "popname='C'" -d num_gens=250 -d "outfile='C.trees'" phylo_bgs.slim

I.trees: root.trees phylo_bgs.slim

	slim -d "infile='root.trees'" -d popsize=100 -d "popname='I'" -d num_gens=200 -d "outfile='I.trees'" phylo_bgs.slim

B.trees: I.trees phylo_bgs.slim

	slim -d "infile='I.trees'" -d popsize=70 -d "popname='B'" -d num_gens=50 -d "outfile='B.trees'" phylo_bgs.slim

A.trees: I.trees phylo_bgs.slim

	slim -d "infile='I.trees'" -d popsize=40 -d "popname='A'" -d num_gens=50 -d "outfile='A.trees'" phylo_bgs.slim

With the makefile in hand, we can now run make, specifying the maximum number of simulations to be run simultaneously the -j. (Click on the “+” icon to see SLiM’s output.)

%%bash
make -f sims.make -j 3

Show code cell output Hide code cell output

make[1]: Entering directory '/home/runner/work/tskit-site/tskit-site/docs'

slim -d "infile=''" -d popsize=500 -d "popname='root'" -d num_gens=2000 -d "outfile='root.trees'" phylo_bgs.slim

// Initial random seed:

4152012718795596099

// RunInitializeCallbacks():

initializeSLiMModelType(modelType = 'WF');

initializeTreeSeq();

initializeMutationRate(1e-08);

initializeMutationType(1, 0.5, "f", -0.01);

initializeGenomicElementType(1, m1, 0.1);

initializeGenomicElement(g1, 0, 999999);

initializeRecombinationRate(1e-09);

// Starting run at tick <start>:

1 

slim -d "infile='root.trees'" -d popsize=50 -d "popname='C'" -d num_gens=250 -d "outfile='C.trees'" phylo_bgs.slim

slim -d "infile='root.trees'" -d popsize=100 -d "popname='I'" -d num_gens=200 -d "outfile='I.trees'" phylo_bgs.slim

// Initial random seed:

3687999900229068413

// RunInitializeCallbacks():

initializeSLiMModelType(modelType = 'WF');

initializeTreeSeq();

initializeMutationRate(1e-08);

initializeMutationType(1, 0.5, "f", -0.01);

initializeGenomicElementType(1, m1, 0.1);

initializeGenomicElement(g1, 0, 999999);

initializeRecombinationRate(1e-09);

// Starting run at tick <start>:

1 

// Initial random seed:

6963628146934480595

// RunInitializeCallbacks():

initializeSLiMModelType(modelType = 'WF');

initializeTreeSeq();

initializeMutationRate(1e-08);

initializeMutationType(1, 0.5, "f", -0.01);

initializeGenomicElementType(1, m1, 0.1);

initializeGenomicElement(g1, 0, 999999);

initializeRecombinationRate(1e-09);

// Starting run at tick <start>:

1 

slim -d "infile='I.trees'" -d popsize=70 -d "popname='B'" -d num_gens=50 -d "outfile='B.trees'" phylo_bgs.slim

slim -d "infile='I.trees'" -d popsize=40 -d "popname='A'" -d num_gens=50 -d "outfile='A.trees'" phylo_bgs.slim

// Initial random seed:

9064305084210009827

// Initial random seed:

2681050222390141554

// RunInitializeCallbacks():

// RunInitializeCallbacks():

initializeSLiMModelType(modelType = 'WF');

initializeSLiMModelType(modelType = 'WF');

initializeTreeSeq();

initializeTreeSeq();

initializeMutationRate(1e-08);

initializeMutationRate(1e-08);

initializeMutationType(1, 0.5, "f", -0.01);

initializeMutationType(1, 0.5, "f", -0.01);

initializeGenomicElementType(1, m1, 0.1);

initializeGenomicElementType(1, m1, 0.1);

initializeGenomicElement(g1, 0, 999999);

initializeGenomicElement(g1, 0, 999999);

initializeRecombinationRate(1e-09);

initializeRecombinationRate(1e-09);

// Starting run at tick <start>:

1 

// Starting run at tick <start>:

1 

make[1]: Leaving directory '/home/runner/work/tskit-site/tskit-site/docs'

Click here for how to use python instead of make

You would have to write a recursion over the branches in your tree (starting from the root) and then parallelize the runs of sister branches somehow.

def phylo_recursion(parent, df):
    print(parent)
    childs = df[df.parent==parent]
    print(childs)
    if len(childs) == 0:
        return
    # you could parallelize this loop over childs with same parent
    for i, row in childs.iterrows():
        if not os.path.exists(row.outfile):
            os.system(f"slim -d \"infile='{row.infile}'\" -d popsize={row.popsize} -d num_gens={row.edgelen} -d \"outfile='{row.child}.trees'\" phylo_bgs.slim")
            phylo_recursion(row.child, df)

phylo_recursion("", df)

Putting it all together: unioning the tree sequences

With the tree sequences in hand, we now need to glue them together. This can be done using union from tskit. For two tree sequences which share some of its past history is shared, union works by copying the non-shared parts of one of the tree sequence onto the other. The trickiest part of this operation is defining the parts that are equivalent in the two tree sequences. For that, you will have to create an array that serves as a map of node IDs between the two tree sequences.

Here is a function that will construct a map of the node IDs of two SLiM tree sequences that correspond to the same chromosomes in SLiM at any time older than the given time ago at which the two populations split. Given two tree sequences other and ts, the goal here is to find, for each node born before split_time ago in other, the matching node in ts, where we can identify matching using the SLiM ID in metadata. The code could be made easier to read by iterating over nodes, but the following numpy-based version is much faster:

def match_nodes(other, ts, split_time):
    """
    Given SLiM tree sequences `other` and `ts`, builds a numpy array with length
    `other.num_nodes` in which the indexes represent the node id in `other` and the
    entries represent the equivalent node id in `ts`. If a node in `other` has no
    equivalent in `ts`, then the entry takes the value `tskit.NULL` (-1). The
    matching is done by comparing the IDs assigned by SLiM which are kept in
    node metadata. This matching of SLiM IDs is *only* done for nodes with time
    older than the specified `split_time`.
    """
    node_mapping = np.full(other.num_nodes, tskit.NULL)
    sids0 = np.array([n.metadata["slim_id"] for n in ts.nodes()])
    sids1 = np.array([n.metadata["slim_id"] for n in other.nodes()])
    alive_before_split1 = (other.tables.nodes.time >= split_time)
    is_1in0 = np.isin(sids1, sids0)
    both = np.logical_and(alive_before_split1, is_1in0)
    sorted_ids0 = np.argsort(sids0)
    matches = np.searchsorted(
        sids0,
        sids1[both],
        side='left',
        sorter=sorted_ids0
    )
    node_mapping[both] = sorted_ids0[matches]
    return node_mapping

Now we are finally ready to union our tree sequences. For that, I wrote a recursive function that goes through our data frame with the phylogeny and returns a dictionary with the merged tree sequences from the tip to the root.

merged = {
    row.child : {
        "ts": tskit.load(row.outfile),
        "depth": row.edgelen,
        "children": [row.child]
    }
    for i, row in df[df.is_leaf].iterrows()
}

def union_children(parent, df, merged):
    print(f"Going in: {parent}")
    child_rows = df[df.parent == parent]
    assert (len(child_rows) == 2) or (len(childs) == 0)
    if len(child_rows) == 2:
        children = [row.child for _, row in child_rows.iterrows()]
        for child in children:
            if child not in merged:
                union_children(child, df, merged)
        split_time = merged[children[0]]["depth"]
        assert split_time == merged[children[1]]["depth"] # ultrametric
        print(f'Unioning: {children}, Split time: {split_time}')
        ts0 = merged[children[0]]["ts"]
        ts1 = merged[children[1]]["ts"]
        node_map = match_nodes(ts1, ts0, split_time)
        tsu = ts0.union(ts1, node_map, check_shared_equality=True)
        # the time from tip to start of simulation is split_time plus the
        # length of the edge
        parent_edgelength = df[df.child==parent].edgelen.item()
        merged[parent] = {
            "ts": tsu,
            "depth": split_time + parent_edgelength,
            "children": merged[children[0]]["children"] + merged[children[1]]["children"]
        }

union_children("root", df, merged)
# union of all three species tree sequences is in the root.
tsu = merged["root"]["ts"]
pops = merged["root"]["children"]

Going in: root
Going in: I
Unioning: ['B', 'A'], Split time: 50
Unioning: ['C', 'I'], Split time: 250

A slightly tricky thing we had to do there was to make sure we kept track of which population in the union’ed tree sequence corresponds to which population in our phylogeny. Happily, we’ve stored each population’s name in its metadata field, so it’s easy to match populations in the tree sequence up to what they’re supposed to be.

Let’s make sure we have the right number of present-day samples in each of the populations. To do this we need to make sure to get “alive” samples, because recall that we have saved the state of the population at each species split time.

alive = np.where(np.isclose(tsu.tables.nodes.time, 0))[0]
pop_ids = {}
for pop in tsu.populations():
    if pop.metadata is not None:
        pop_ids[pop.metadata['name']] = pop.id

for name in pops:
    pop_samples = tsu.samples(pop_ids[name])
    n_samples = sum(np.isin(pop_samples, alive)) // 2
    print(f"Union-ed tree sequence has {n_samples} samples in population {name},\n"
          f"\tand we specified {df[df.child==name].popsize.item()} individuals in our simulations.")
    assert n_samples == df[df.child==name].popsize.item()

Union-ed tree sequence has 50 samples in population C,
	and we specified 50 individuals in our simulations.
Union-ed tree sequence has 70 samples in population B,
	and we specified 70 individuals in our simulations.
Union-ed tree sequence has 40 samples in population A,
	and we specified 40 individuals in our simulations.

Let’s do an additional consistency check now, to see if we need to recapitate (i.e., if some trees haven’t coalesced), and to make sure that all roots are in the root population, as they should be:

# TODO: fix up
# for t in tsu.trees():
#     for r in t.roots:
#         assert tsu.node(r).population == pop_ids["root"]

print(f"Max number of roots: {max([t.num_roots for t in tsu.trees()])}.")

Max number of roots: 1.

Finally, we will recapitate the result with a small population size of 100, in case some trees on the root branch haven’t coalesced, and write out the result:

tsu = pyslim.recapitate(tsu, recombination_rate=1e-8, ancestral_Ne=100)
tsu.dump("final.trees")

/opt/hostedtoolcache/Python/3.10.15/x64/lib/python3.10/site-packages/msprime/ancestry.py:831: TimeUnitsMismatchWarning: The initial_state has time_units=ticks but time is measured in generations in msprime. This may lead to significant discrepancies between the timescales. If you wish to suppress this warning, you can use, e.g., warnings.simplefilter('ignore', msprime.TimeUnitsMismatchWarning)
  warnings.warn(message, TimeUnitsMismatchWarning)

Now we’re done, and can analyse the final tree sequence! Just for fun, I’ll look at the trees produced by the simulation. For instance, we might be curious how often there are disagreements between the species tree and the simulated gene trees (also called incomplete lineage sorting, or ILS).

To make it possible to look at the trees, I will first simplify the union-ed tree sequence to keep only two diploid samples per population.

rng = np.random.default_rng(seed=123)
ind_alive = pyslim.individuals_alive_at(tsu, 0)
# TODO: this will work in the next tskit
# ind_pops = tsu.individuals_population[ind_alive]
ind_pops = np.array([tsu.node(tsu.individual(i).nodes[0]).population for i in ind_alive])
subsample_indivs = [
    rng.choice(ind_alive[ind_pops == pop_ids[name]], 2)
    for name in pops
]
subsample_nodes = [
    np.concatenate([tsu.individual(i).nodes for i in x])
    for x in subsample_indivs
]
tsus = tsu.simplify(
        np.concatenate(subsample_nodes),
        filter_populations=False,
)
pop_labels = {v: k for k, v in pop_ids.items()}
SVG(tsus.draw_svg(
    node_labels={
        node.id: pop_labels[node.population]
        for node in tsus.nodes()
        if not node.time > 0.0
    },
    x_lim=[0,2200],
    size=(800, 300),
))

Note

A possible gotcha in the code above lies in getting the time units to work out. Note that in the SLiM script we both save and reload .trees files in the late() stage of the SLiM life cycle. This is important: if we had reloaded the files in early(), then each time we did so the “tskit time” and “SLiM time” would become one step out of sync. This leads to errors either in union (since if the time units in the two tree sequences do not match, union will raise an error) or in recapitate (since recapitate assumes that the “top” of the trees are at the number of generations ago recorded by SLiM in metadata).