2-simulating-ancestry-with-msprime.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Simulating ancestry with msprime\n",
    "\n",
    " - [2.1 Basic syntax](#2.1BasicSyntax)\n",
    " - [2.2 Specifying contig information](#2.2SpecifyingContigInformation)\n",
    " - [2.3 Demographic history](#2.3DemographicHistory)\n",
    " - [2.4 Debugging demography](#2.4DebuggingDemography)\n",
    " - [2.5 Other ancestry models](#2.5OtherAncestryModels)\n",
    " - [2.6 Mixing different ancestry models](#MixingDifferentAncestryModels)\n",
    "\n",
    "This is the first of two sessions about `msprime`, a backwards-time tree sequence simulator. In this first worksheet, we will learn how to design simulations using customisable models of genomic history. In the next session, we will add mutations to these simulated histories to obtain sequence data.\n",
    "\n",
    "#### Some relevant papers:\n",
    " -  [Efficient coalescent simulation and genealogical analysis for large sample sizes](https://doi.org/10.1371/journal.pcbi.1004842)\n",
    " - [Efficient ancestry and mutation simulation with msprime 1.0](https://doi.org/10.1093/genetics/iyab229)\n",
    " - [tskit.dev documentation](https://tskit.dev/)\n",
    " \n",
    "Simulations are important in population genetics for many reasons:\n",
    "\n",
    "**Exploration:**\n",
    "Simulations allow us to explore the influence of various historical scenarios on observed patterns of genetic variation and inheritance.\n",
    "\n",
    "**Benchmarking and evaluating methodologies:**\n",
    "To assess the accuracy of inferential methods, we need test datasets for which the true values of important parameters are known.\n",
    "\n",
    "**Model training:**\n",
    "Some methods for ancestry inference are trained on simulated data (eg. Approximate Bayesian Computation).\n",
    "This is especially important in studies of complex demographies, where there are many potential parameters and models, making it impractical to specify likelihood functions.\n",
    "\n",
    "### A brief history of msprime\n",
    "\n",
    "The first release of `msprime` was an emulation of the popular `ms` coalescent simulator with added support for tree sequences.\n",
    "However, it has since become an expansive and flexible backwards-in-time simulator for various different models of genetic ancestry and mutation, and even for simplified models of selection.\n",
    "\n",
    "#### Forwards and backwards simulation\n",
    "\n",
    "The main characteristic of `msprime` is that it simulates *tree sequences* in *backwards-time*.\n",
    "\n",
    "<img src=\"pics/msprime-1.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/msprime-2.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/msprime-3.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/msprime-4.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/msprime-5.png\" width=\"200\" height=\"200\">\n",
    "\n",
    "The alternative is to use a *forwards-time* simulator like `SLiM`:\n",
    "\n",
    "<img src=\"pics/slim-1.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-2.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-3.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-4.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-5.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-6.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-7.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-8.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-9.png\" width=\"200\" height=\"200\">\n",
    "<img src=\"pics/slim-10.png\" width=\"200\" height=\"200\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, forwards-time simulation with `SLiM` is *detailed* and *more realistic*,\n",
    "while backwards-time simulation with `msprime` is *fast* and *efficient*.\n",
    "`msprime` is designed for *neutral* simulations, while `SLiM` is better suited to complex simulations, including those involving selection.\n",
    "\n",
    "However, `pyslim` allows you to combine these via a process called 'recapitation'.\n",
    "See [this](https://tskit.dev/pyslim/docs/latest/tutorial.html) for details.\n",
    "\n",
    "<img src=\"pics/recapitation-1.png\" width=\"400\" height=\"400\">\n",
    "<img src=\"pics/recapitation-2.png\" width=\"400\" height=\"400\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import msprime\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import SVG"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.1 Basic syntax\n",
    "\n",
    "To simulate a tree sequence in `msprime`, we use the `sim_ancestry()` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(samples=2, random_seed=1)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.2 Specifying contig information\n",
    "Although we have specified 2 samples, our tree sequence contains 4 sample nodes.\n",
    "This is because the `samples` argument specifies the number of *individuals* in the sample,\n",
    "and by default, `sim_ancestry()` assumes diploid organisms.\n",
    "To change this, use the `ploidy` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(samples=2, random_seed=1, ploidy=3)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's easiest to start thinking about genome lengths in units of nucleotides. By default, we are simulating a sequence length that spans just one of these units.\n",
    "We can specify a larger region using the `sequence_length` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(samples=2, random_seed=1, sequence_length=10)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Also, note that our 'tree sequence' consists of just a single tree. This is because we have not yet specified a `recombination_rate`, and the default is 0. \n",
    "This is the probability of a recombination event per genomic unit (base), per generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "                samples=2,\n",
    "                random_seed=82,\n",
    "                sequence_length=10,\n",
    "                recombination_rate=0.1)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may also wish to specify recombination rates that vary across the genome. \n",
    "We do this by creating a `RateMap` object,\n",
    "which lists recombination rates between defined positions in the sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rate_map = msprime.RateMap(position=[0, 10, 20], rate=[0.01, 0.1])\n",
    "ts = msprime.sim_ancestry(3, recombination_rate=rate_map, random_seed=2)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that it's also possible to simulate gene conversion in `msprime`.\n",
    "If this is of interest, have a look at the API documentation for the `gene_conversion_rate` and `gene_conversion_tract_length` arguments,\n",
    "and [this](https://tskit.dev/msprime/docs/stable/ancestry.html?highlight=gene%20conversion#gene-conversion) short tutorial.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, the recombination events will be assigned to integer locations along the sequence. However there may be situations where you want to model the genome using continuous coordinates. In this case, use the `discrete_genome=False` argument:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples=2, random_seed=28, sequence_length=100,\n",
    "    recombination_rate=0.01, discrete_genome=False)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Exercise:* Our study organisms have a ploidy of 3 and a chromosome of length 100.\n",
    "Between the bases at positions 50 and 60, there is a recombination 'hotspot' with a higher recombination rate of 0.1.\n",
    "Modify the following code to simulate chromosomes for 2 such organisms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rate_map = msprime.RateMap(position=[0, 100], rate=[0.01])\n",
    "ts = msprime.sim_ancestry(samples=3,\n",
    "                          recombination_rate=rate_map,\n",
    "                          random_seed=193)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.3 Demographic history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far, we have been simulating samples from a single population of a constant size, which isn’t particularly exciting! One of the strengths of msprime is that it can be used to specify quite complicated models of demography and population history with a simple Python API."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Population size\n",
    "\n",
    "So far, we've assumed our samples and all their ancestors belong to a single population with $2N_e=1$, which scales the simulation to use classical coalescent time units of $\\frac{1}{2N_e}$. (Note that `msprime` scales this for different ploidies, also)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(samples=2, random_seed=234)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = ts.first()\n",
    "print(\"Total branch length:\", tree.total_branch_length)\n",
    "print(\"Time at root:\", ts.tables.nodes.time[tree.root])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By specifying a specific population size using the `population_size` argument, all times and branch lengths will be rescaled into units of generations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples=2,\n",
    "    random_seed=234,\n",
    "    population_size=2)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = ts.first()\n",
    "print(\"Total branch length:\", tree.total_branch_length)\n",
    "print(\"Time at root:\", ts.tables.nodes.time[tree.root])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that this population size should be the number of *individuals*, not the number of genomes. `msprime` uses the supplied `ploidy` to do all the other necessary adjustments under-the-hood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples=1,\n",
    "    ploidy=4,\n",
    "    random_seed=234,\n",
    "    population_size=2)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = ts.first()\n",
    "print(\"Total branch length:\", tree.total_branch_length)\n",
    "print(\"Time at root:\", ts.tables.nodes.time[tree.root])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Population structure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To simulate using more complicated models of demographic history, we will need to create a `msprime.Demography` object.\n",
    "`msprime` supports simulation from multiple discrete populations, each of which is initialized via the `add_population()` method.\n",
    "For each population, you can specify a sample size, an effective population size, an exponential growth rate and a name."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could have specified our previous simulation with the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(\n",
    "    name=\"my_pop\",\n",
    "    initial_size=1\n",
    ")\n",
    "dem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples=2,\n",
    "    random_seed=234,\n",
    "    demography=dem\n",
    ")\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use `Population` objects to specify an exponential growth rate of the population per generation (i.e. forwards in time)\n",
    "\n",
    "<img src=\"pics/worksheet2-dem1.png\" width=\"400\" height=\"400\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(\n",
    "    name=\"my_pop\",\n",
    "    initial_size=1e5,\n",
    "    growth_rate=0.1\n",
    ")\n",
    "dem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples=2,\n",
    "    random_seed=11,\n",
    "    demography=dem\n",
    ")\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the recent branches are much longer than the older ones.\n",
    "This is what we expect to see in a growing population."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can add any number of populations into our `Demography` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(\n",
    "    name=\"R\",\n",
    "    description=\"Red population\",\n",
    "    initial_size=500\n",
    ")\n",
    "dem.add_population(\n",
    "    name=\"B\",\n",
    "    description=\"Blue population\",\n",
    "    initial_size=500,\n",
    "    growth_rate=0.0001\n",
    ")\n",
    "dem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"pics/worksheet2-dem2.png\" width=\"400\" height=\"400\">\n",
    "\n",
    "However, this simulation will run forever unless we also specify some migration between the groups!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # without an end_time, THIS WILL RUN FOREVER!\n",
    "\n",
    "# ts = msprime.sim_ancestry(\n",
    "#   samples={\"R\" : 3, \"B\" : 3}, \n",
    "#   demography=dem,\n",
    "#   random_seed=12,\n",
    "#   sequence_length=1000, \n",
    "#   )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To understand why, recall that msprime is a backwards-time simulator. The simulation will run beginning from the present day working further back in time, simulating until all samples have coalesced to a single common ancestor at each genomic location. However, with no migration between our two populations, samples in one population will never coalesce with samples in another population. To fix this, let’s add some migration events."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Migration\n",
    "\n",
    "With msprime, you can specify continual rates of migrations between populations, as well as admixture events, divergences and one-off mass migrations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Migration rates** are passed to our `msprime.Demography` object individually via the `set_migration_rate()` method.\n",
    "This allows us to specify the expected number of migrants moving from population `dest` to population `source` per generation, divided by the size of population `source`. When this rate is small (close to 0), it is approximately equal to the fraction of population `source` that consists of new migrants from population `dest` in each generation. \n",
    "\n",
    "If both of your populations are exchanging migrants at the same rate, you can save yourself some typing by specifying them with a single `set_symmetric_migration_rate()` call.\n",
    "See the documentation [here](https://tskit.dev/msprime/docs/stable/api.html#msprime.Demography.add_symmetric_migration_rate_change)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(\n",
    "    name=\"R\",\n",
    "    description=\"Red population\",\n",
    "    initial_size=500\n",
    ")\n",
    "dem.add_population(\n",
    "    name=\"B\",\n",
    "    description=\"Blue population\",\n",
    "    initial_size=500,\n",
    "    growth_rate=0.0001\n",
    ")\n",
    "dem\n",
    "\n",
    "# Set migration rates.\n",
    "dem.set_migration_rate(source=0, dest=1, rate=0.05)\n",
    "dem.set_migration_rate(source=1, dest=0, rate=0.02)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For instance, the `Demography` object above specifies that in each generation, approximately 5% of the red population (0) consists of migrants from the blue population (1), and approximately 2% of the blue population consists of migrants from the red population."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate.\n",
    "ts = msprime.sim_ancestry(\n",
    "  samples={\"R\" : 5, \"B\" : 5},\n",
    "  demography=dem,\n",
    "  sequence_length=1e7,\n",
    "  random_seed=145)\n",
    "ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One consequence of explicitly specifying `Population` objects is that each of the simulated nodes will now belong to one of our specified populations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.populations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use this information to draw the tree sequence with nodes coloured by population label using SVG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colour_map = {0:\"red\", 1:\"blue\"}\n",
    "node_colours = {u.id: colour_map[u.population] for u in ts.nodes()}\n",
    "for tree in ts.trees():\n",
    "    print(\"Tree on interval:\", tree.interval)\n",
    "    # The code below will only work in a Jupyter notebook with SVG output enabled.\n",
    "    display(SVG(tree.draw(node_colours=node_colours, width=500, height=400)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More coalescences are happening in the blue population (1) than the red population (0). This makes sense given that population 1 is specifying more migrants to population 0 than vice versa."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Changing migration rates\n",
    "\n",
    "We can change any of the migration rates at any time in the simulation. To do this, we'll use the `add_migration_rate_change()` method on our demography object.\n",
    "This will specify the populations whose migration rates are to be changed, the time of the change and the new migration rate.\n",
    "\n",
    "For instance, say we wanted to specify that in each generation prior to `time=100`, 1% of the red population (0) consisted of migrants from the blue population (1)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem.add_migration_rate_change(time=100, rate=0.01, source=0, dest=1)\n",
    "dem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output above shows that we have successfully added our first demographic event to our `Demography` object, a migration rate change. We are now ready to simulate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "  samples={\"R\" : 4, \"B\" : 4},\n",
    "  demography=dem,\n",
    "  sequence_length=1000,\n",
    "  random_seed=63461)\n",
    "ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colour_map = {0:\"red\", 1:\"blue\"}\n",
    "node_colours = {u.id: colour_map[u.population] for u in ts.nodes()}\n",
    "for tree in ts.trees():\n",
    "    print(\"Tree on interval:\", tree.interval)\n",
    "    # The code below will only work in a Jupyter notebook with SVG output enabled.\n",
    "    display(SVG(tree.draw(node_colours=node_colours, width=500, height=400)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Admixture\n",
    "\n",
    "It is also easy to specify admixture and divergence events with `msprime`. Suppose we wanted to specify our demography so that our admixed population arose 200 generations ago, with 70% of the new population being migrants from the red ancestral population, and 30% being migrants from the blue ancestral population.\n",
    "\n",
    "<img src=\"pics/worksheet2-dem3.png\" width=\"400\" height=\"400\">\n",
    "\n",
    "We can do this by using the `add_admixture()` method on our demography object. We must supply a list of ancestral populations participating in the admixture, and a list of the same size specifying the proportions of migrants from each of these populations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(\n",
    "    name=\"AncestralPop0\", description=\"Plotted in red.\", initial_size=1000, growth_rate=0)\n",
    "dem.add_population(\n",
    "    name=\"AncestralPop1\", description=\"Plotted in blue.\", initial_size=1000, growth_rate=0)\n",
    "dem.add_population(\n",
    "    name=\"AdmixedPop\", description=\"Plotted in purple.\", initial_size=700)\n",
    "dem.set_migration_rate(\n",
    "    source=0, dest=1, rate=0.001)\n",
    "dem.set_migration_rate(\n",
    "    source=1, dest=0, rate=0.001)\n",
    "\n",
    "# Specify admixture event.\n",
    "dem.add_admixture(time=200, derived=\"AdmixedPop\", ancestral=[\"AncestralPop0\", \"AncestralPop1\"], proportions=[0.6, 0.4])\n",
    "dem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this simulated sample, all of the recent ancestral haplotypes belong to the admixed population, and all of those prior to the admixture event belong to one of the ancestral populations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "  samples={\"AncestralPop0\" : 0, \"AncestralPop1\" : 0, \"AdmixedPop\" : 6},\n",
    "  demography=dem,\n",
    "  sequence_length=1000,\n",
    "  random_seed=63\n",
    ")\n",
    "\n",
    "print(\"Populations of nodes from time < 500:\")\n",
    "print([u.population for u in ts.nodes() if u.time < 500])\n",
    "print(\"Populations of nodes from time >= 500:\")\n",
    "print([u.population for u in ts.nodes() if u.time >= 500])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colour_map = {0:\"red\", 1:\"blue\", 2:\"purple\"}\n",
    "node_colours = {u.id: colour_map[u.population] for u in ts.nodes()}\n",
    "for tree in ts.trees():\n",
    "    print(\"Tree on interval:\", tree.interval)\n",
    "    # The code below will only work in a Jupyter notebook with SVG output enabled.\n",
    "    display(SVG(tree.draw(node_colours=node_colours, width=500, height=400)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Admixtures and population splits are special types of demographic events that affect the state of some of the defined populations, in addition to moving lineages between populations. The output below shows that by adding the admixture event, we are triggering a change in the state of `AdmixedPop` at `time = 50`; the population is active at the start of the simulation, but becomes inactive for all steps of the simulation beyond time 50."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that, for example, adding any demographic events that affect the AdmixedPop before this time will produce an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # THIS WON'T WORK\n",
    "# dem.add_migration_rate_change(time=80, rate=0.01, source=\"AncestralPop0\", dest=\"AdmixedPop\");\n",
    "# ts = msprime.sim_ancestry(\n",
    "#   samples={\"AncestralPop0\" : 0, \"AncestralPop1\" : 0, \"AdmixedPop\" : 6},\n",
    "#   demography=dem,\n",
    "#   sequence_length=1000,\n",
    "#   random_seed=63,\n",
    "#   recombination_rate=1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Population splits\n",
    "\n",
    "We can also simulate population divergences with msprime. Suppose we want to model a situation where all lineages from multiple populations are migrants from a single ancestral population at a single point in time.\n",
    "\n",
    "<img src=\"pics/worksheet2-dem4.png\" width=\"400\" height=\"400\">\n",
    "\n",
    "We’ll specify this with the `add_population_split()` method. We need to know the time of the event, and the IDs of the derived and ancestral populations participating in the divergence event.\n",
    "\n",
    "Notice that in this case, we do not need to provide proportions as we did in the case of admixture. This makes sense when you consider that `msprime` simulates backwards-in-time: all lineages in all of the derived populations originate from the ancestral population in a split event. Any differences in ‘quantities’ of migrants must be modelled by sizes of the derived populations at the time of the split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(name=\"R\", description=\"Plotted in red.\", initial_size=500)\n",
    "dem.add_population(name=\"B\", description=\"Plotted in blue.\",initial_size=500)\n",
    "dem.add_population(name=\"AncestralPopulation\", description=\"Plotted in green.\", initial_size=500)\n",
    "\n",
    "# Add the population split.\n",
    "dem.add_population_split(time=100, derived=[\"R\",\"B\"], ancestral=\"AncestralPopulation\")\n",
    "dem"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Population splits will also modify the state of each of the derived populations, changing them from active to inactive at the time of the split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "  samples={\"R\" : 3, \"B\" : 3, \"AncestralPopulation\" : 0},\n",
    "  demography=dem,\n",
    "  sequence_length=1000,\n",
    "  random_seed=63\n",
    ")\n",
    "\n",
    "print(\"Populations of nodes from time < 100:\")\n",
    "print([u.population for u in ts.nodes() if u.time < 100])\n",
    "print(\"Populations of nodes from time >= 100:\")\n",
    "print([u.population for u in ts.nodes() if u.time >= 100])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"pics/worksheet2-ex.png\" width=\"400\" height=\"400\">\n",
    "\n",
    "*Exercise:* The code below specifies a simulation of two populations in which population 1 and population 0 are created by a population split at time 200.\n",
    "\n",
    "Can you add a third contemporary population, population 2, so that populations 1 and 2 are created by a more recent population split  at time 150? (You may need to create some other ancestral populations too.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup.\n",
    "dem = msprime.Demography()\n",
    "dem.add_population(name=\"R\", description=\"Plotted in red.\", initial_size=500)\n",
    "dem.add_population(name=\"B\", description=\"Plotted in blue.\",initial_size=500)\n",
    "dem.add_population(name=\"AncestralPopulation\", description=\"Plotted in green.\", initial_size=500)\n",
    "\n",
    "# Add the population split.\n",
    "dem.add_population_split(time=200, derived=[\"R\",\"B\"], ancestral=\"AncestralPopulation\")\n",
    "dem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate!\n",
    "ts_ex = msprime.sim_ancestry(\n",
    "  samples={\"R\" : 1, \"B\" : 1, \"AncestralPopulation\" : 0},\n",
    "  demography=dem,\n",
    "  sequence_length=1000,\n",
    "  random_seed=63\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot.\n",
    "colour_map = {0:\"red\", 1:\"blue\", 2:\"purple\"}\n",
    "node_colours = {u.id: colour_map[u.population] for u in ts.nodes()}\n",
    "for tree in ts.trees():\n",
    "    print(\"Tree on interval:\", tree.interval)\n",
    "    # The code below will only work in a Jupyter notebook with SVG output enabled.\n",
    "    display(SVG(tree.draw(node_colours=node_colours, width=500, height=400)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Changing population sizes or growth rates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also may wish to specify changes to rates of population growth, or sudden changes in population size at a particular time. Both of these can be specified by applying the `add_population_parameters_change()` method to our `Demography` object.\n",
    "\n",
    "<img src=\"pics/worksheet2-dem5.png\" width=\"400\" height=\"400\">"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dem = msprime.Demography()\n",
    "dem.add_population(name=\"R\", description=\"Plotted in red.\", initial_size=500)\n",
    "dem.add_population(name=\"B\", description=\"Plotted in blue.\",initial_size=500)\n",
    "dem.set_migration_rate(source=0, dest=1, rate=0.05)\n",
    "dem.set_migration_rate(source=1, dest=0, rate=0.02)\n",
    "\n",
    "# Bottleneck in Population 0 between 50 - 150 generations ago.\n",
    "dem.add_population_parameters_change(time=50, initial_size=250, population=0)\n",
    "dem.add_population_parameters_change(time=150, initial_size=500, population=0)\n",
    "\n",
    "# Exponential growth in Population 1 starting 50 generations ago.\n",
    "dem.add_population_parameters_change(time=100, growth_rate=0.01, population=1)\n",
    "\n",
    "# Sort events, since we've added some out of time order.\n",
    "dem.sort_events()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate.\n",
    "ts = msprime.sim_ancestry(\n",
    "    samples={\"R\" : 3, \"B\" : 3},\n",
    "    demography=dem,\n",
    "    sequence_length=1000,\n",
    "    random_seed=63461\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that because `msprime` simulates backwards-in-time, parameter changes must be interpreted backwards-in-time as well. For instance, the population growth event in the example above specifies continual growth in the early history of population 1 up until 100 generations in the past."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.4 Debugging demography\n",
    "\n",
    "To help you spot any mistakes in your specified demography, `msprime` provides a debugger that prints out your population history in a more human-readable form. It’s good to get into the habit of running this debugger before running your simulations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_history = msprime.DemographyDebugger(demography=dem)\n",
    "my_history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also use the debugger to perform some informative calculations based on the given demography."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T = np.linspace(0, 49, 200)\n",
    "RAA, _ = my_history.coalescence_rate_trajectory(T, {\"R\": 2})\n",
    "RBB, _ = my_history.coalescence_rate_trajectory(T, {\"B\": 2})\n",
    "RAB, _ = my_history.coalescence_rate_trajectory(T, {\"R\": 1, \"B\": 1})\n",
    "\n",
    "plt.plot(T, RAA, label=\"two R samples\")\n",
    "plt.plot(T, RBB, label=\"two B samples\")\n",
    "plt.plot(T, RAB, label=\"one of each\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"time ago\")\n",
    "plt.ylabel(\"coalescence rate\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See [the documentation](https://tskit.dev/msprime/docs/stable/demography.html#sec-demography-numerical) for more examples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.5 Other ancestry models\n",
    "\n",
    "When might `msprime`'s default (Hudson coalescent) model of genealogy not be ideal for you? Consider the following assumptions it makes:\n",
    "\n",
    " - Selective neutrality\n",
    " - Random mating and survival outcomes\n",
    " - Sample is small relative to population size\n",
    " \n",
    "However, it is possible to simulate from a selection of other models in `msprime` as well! \n",
    "The full list is [here](https://tskit.dev/msprime/docs/stable/ancestry.html#sec-ancestry-quickref)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discrete-time Wright-Fisher\n",
    "\n",
    "DTWF simulations are less efficient than coalescent simulations, but should perform better\n",
    "\n",
    " - on analyses of very recent history (especially for haplotype-based information like identity-by-descent and local ancestry)\n",
    " - on small populations, or when sample size is a large fraction of the population size\n",
    " - when modelling very long range correlations across chromosomes is important\n",
    " \n",
    "See [this paper](https://doi.org/10.1371/journal.pgen.1008619) for a more comprehensive comparison between DTWF and coalescent simulations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider the following (coalescent) simulation of two individuals in a population of two."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "                samples=2,\n",
    "                population_size=2,\n",
    "                random_seed=82,\n",
    "                sequence_length=20,\n",
    "                recombination_rate=0.3)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we now have over 4 distinct chromosomes within a single generation, and indeed, some of these chromosomes inherit from multiple ancestors within a single generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is no longer a problem when we use a discrete-time Wright-Fisher simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DTWF sim\n",
    "ts = msprime.sim_ancestry(\n",
    "                samples=2,\n",
    "                population_size=2,\n",
    "                model=msprime.DiscreteTimeWrightFisher(),\n",
    "                random_seed=82,\n",
    "                sequence_length=10,\n",
    "                recombination_rate=0.2)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simple selective sweeps\n",
    "\n",
    "We can also run simple simulations of genic selection using a structured coalescent model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_model = msprime.SweepGenicSelection(\n",
    "    position=1,  # start of chrom\n",
    "    start_frequency=0.2,\n",
    "    end_frequency=0.8,\n",
    "    s=0.25,\n",
    "    dt=1e-5,\n",
    ")\n",
    "\n",
    "ts = msprime.sim_ancestry(\n",
    "    samples=3,\n",
    "    sequence_length=5e7,\n",
    "    population_size=100,\n",
    "    model=sweep_model,\n",
    "    recombination_rate=1e-8,\n",
    "    random_seed=123456)\n",
    "SVG(ts.draw_svg())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'd recommend using `SLiM` if selection is an important part of your demographic model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.6 Mixing different ancestry models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can also combine multiple ancestry models inside a single `msprime` simulation. Here's a simple example involving a hybrid DTWF and coalescent simulation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "models=[msprime.DiscreteTimeWrightFisher(duration=50),\n",
    "        msprime.StandardCoalescent()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    10,\n",
    "    population_size=100,\n",
    "    sequence_length=1e6,\n",
    "    recombination_rate=1e-8,\n",
    "    model=models,\n",
    "    random_seed=6789\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}