Advanced forward simulations#


Add further details on building a forward simulator (see issue #14)

In the previous tutorial, we developed a basic basic forward-time Wright-Fisher (WF) simulator (refer back to that tutorial for a detailed run through of the hidden code):

Hide code cell content
import tskit
import numpy as np

random_seed = 6
random = np.random.default_rng(random_seed)  # A random number generator for general use

L = 50_000  # The sequence length: 50 Kb

def add_inheritance_paths(tables, parent_genomes, child_genome, recombination_rate):
    "Add paths from parent genomes to the child genome, with crossover recombination"
    L = tables.sequence_length
    num_recombinations = random.poisson(recombination_rate * L)
    breakpoints = random.integers(0, L - 1, size=num_recombinations)
    break_pos, counts = np.unique(breakpoints, return_counts=True)
    crossovers = break_pos[counts % 2 == 1]  # no crossover if e.g. 2 breaks at same pos
    left_positions = np.insert(crossovers, 0, 0)
    right_positions = np.append(crossovers, L)

    inherit_from = random.integers(2)
    for left, right in zip(left_positions, right_positions):
            left, right, parent_genomes[inherit_from], child_genome)
        inherit_from = 1 - inherit_from  # switch to other parent genome

def make_diploid(tables, time, parent_individuals=None):
    individual_id = tables.individuals.add_row(parents=parent_individuals)
    return individual_id, (
        tables.nodes.add_row(time=time, individual=individual_id),
        tables.nodes.add_row(time=time, individual=individual_id),

def new_population(tables, time, prev_pop, recombination_rate):
    pop = {}
    prev_individuals = np.array([i for i in prev_pop.keys()], dtype=np.int32)
    for _ in range(len(prev_pop)):
        mother_and_father = random.choice(prev_individuals, 2, replace=True)
        child_id, child_genomes = make_diploid(tables, time, mother_and_father)
        pop[child_id] = child_genomes  # store the genome IDs
        for child_genome, parent_individual in zip(child_genomes, mother_and_father):
            parent_genomes = prev_pop[parent_individual]
            add_inheritance_paths(tables, parent_genomes, child_genome, recombination_rate)
    return pop

def initialise_population(tables, time, size) -> dict:
    return dict(make_diploid(tables, time) for _ in range(size))

The main simulation function, as below, returned an unsimplified tree sequence, which we subsequently simplified:

def forward_WF(num_diploids, seq_len, generations, recombination_rate=0, random_seed=7):
    global random
    random = np.random.default_rng(random_seed) 
    tables = tskit.TableCollection(seq_len)
    tables.time_units = "generations"

    pop = initialise_population(tables, generations, num_diploids)
    while generations > 0:
        generations = generations - 1
        pop = new_population(tables, generations, pop, recombination_rate)

    return tables.tree_sequence()

Repeated simplification#

We can perform simplification directly on the tables within the forward_WF() function, using TableCollection.simplify(). More importantly, we can carry this out at repeated intervals. It is helpful to think of this as regular “garbage collection”, as what we’re really doing is getting rid of extinct lineages while also “trimming” extant lineages down to a minimal representation.


Regular garbage collection forces us to reckon with the fact that simplification changes the node IDs. We therefore need to remap any node (and individual) IDs that are used outside of tskit. In the implementation described here, those IDs are stored in the pop variable.

def simplify_tables(tables, samples, pop) -> dict[int, tuple[int, int]]:
    Simplify the tables with respect to the given samples, returning a
    population dict in which individual and nodes have been remapped to their
    new ID numbers
    node_map = tables.simplify(samples, keep_input_roots=True)
    nodes_individual = tables.nodes.individual
    remapped_pop = {}
    for node1, node2 in pop.values():
        node1, node2 = node_map[[node1, node2]]  # remap
        assert nodes_individual[node1] == nodes_individual[node2]  # sanity check
        remapped_pop[nodes_individual[node1]] = (node1, node2)
    return remapped_pop

def forward_WF(
    simplification_interval=None,  # default to simplifying only at end
    global random
    random = np.random.default_rng(random_seed) 
    tables = tskit.TableCollection(seq_len)
    tables.time_units = "generations"  # optional, but helpful when plotting
    if simplification_interval is None:
        simplification_interval = generations
    simplify_mod = generations % simplification_interval

    pop = initialise_population(tables, generations, num_diploids)
    while generations > 0:
        generations = generations - 1
        pop = new_population(tables, generations, pop, recombination_rate)
        if generations > 0 and generations % simplification_interval == simplify_mod:
            current_nodes = [u for nodes in pop.values() for u in nodes]
            pop = simplify_tables(tables, current_nodes, pop)
            if show:
                print("Simplified", generations, "generations before end")

    pop = simplify_tables(tables, [u for nodes in pop.values() for u in nodes], pop)
    if show:
        print("Final simplification")
    return tables.tree_sequence()

ts = forward_WF(6, L, generations=100, simplification_interval=25, show=True)
ts.draw_svg(size=(800, 200))
Simplified 75 generations before end
Simplified 50 generations before end
Simplified 25 generations before end
Final simplification

Invariance to simplification interval#

A critical concept to keep in mind is that the simulation itself is the only random component. The simplification algorithm is deterministic given a set of (nodes, edges) satisfying tskit’s sorting requirements. Therefore, the results of our new forward_WF() function must be the same for all simplification intervals


This invariance property only holds in some cases. We discuss this in more detail below when we add in mutation.

ts = forward_WF(10, L, 500, simplification_interval=1, random_seed=42)

# Iterate over a range of odd and even simplification intervals.
print("Testing invariance to simplification interval")
test_intervals = list(range(2, 500, 33))
for i in test_intervals:
    # Make sure each new sim starts with same random seed!
    ts_test = forward_WF(10, L, 500, simplification_interval=i, show=False, random_seed=42)
    assert ts.equals(ts_test, ignore_provenance=True)
print(f"Intervals {test_intervals} passed")
Testing invariance to simplification interval
Intervals [2, 35, 68, 101, 134, 167, 200, 233, 266, 299, 332, 365, 398, 431, 464, 497] passed


Testing your own code using loops like the one above is a very good way to identify subtle bugs in book-keeping.


  • Simplifying during a simulation changes IDs in the tree sequence tables, so we need to remap entities that store any of these IDs between generations.

  • Our code to carry out simplification gets called both during the simulation and at the end. It’s therefore worth encapsulating it into a class or function for easier code re-use and testing.

Technical notes#

We have found that it is possible to write a simulation where the results differ by simplification interval, but appear correct in distribution. By this we mean that looking at distributions of numbers of mutations, their frequencies, etc., match predictions from analytical theory. However, our experience is that such simulations contain bugs and that the summaries being used for testing are too crude to catch them. For example, they may affect the variance in a subtle way that would require millions of simulations to catch. Often what is going on is that parent/offspring relationships are not being properly recorded, resulting in lineages that either persist too long or not long enough. (In other words, the variance in offspring number per diploid is no longer what it should be, meaning you’ve changed the effective population size.) Thus, please make sure you get the same tskit tables out of a simulation for any simplification interval.


In this section, we will add mutation to our simulation. Mutations will occur according to the infinitely-many sites model, which means that a new mutation cannot arise at a currently-mutated position. \(\theta = 4N\mu\) is the scaled mutation rate, and is equal to twice the expected number of new mutations per generation. The parameter \(\mu\) is the expected number of new mutations per gamete, per generation. Mutation positions will be uniformly distributed along the genome.

Adding mutations changes the complexity of the simulation quite a bit, because now we must add to and simplify site tables and mutation tables instances. We might also want to add metadata to the sites or mutations, recording details such as the selection coefficient of a mutation, or the type of mutation (e.g., synonymous vs. non-synonymous).

We will write a mutation function here which we will re-use in future examples.


We will be treating mutations as neutral. Doing so is odd, as one big selling point of tskit is the ability to skip the tracking of neutral mutations in forward simulations. However, tracking neutral mutations plus metadata is the same as tracking selected mutations and their metadata, and being able to do neat things like put your selected mutations onto a figure of the genealogy is one of several possible use cases.


The rest of this tutorial is still under construction, and needs porting from this workbook. This will primarily deal with sites and mutations (and mutational metadata). We could also include details on selection, if that seems sensible.

The section in that workbook on “Starting with a prior history” should be put in the Completing forwards simulations tutorial.