4-analyses-with-tree-sequences.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import msprime, tskit\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "from IPython.display import SVG\n",
    "from matplotlib.patches import Polygon"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Analyses with tskit\n",
    "\n",
    "At the end of the previous worksheet, we learnt how to convert a tree sequence's mutational information back into a VCF sequence file. This allows us to use software like `scikit-allel`, `plink`, `bcftools` and `vcftools` on our simulated datasets.\n",
    "\n",
    "However, there's a lot we can do entirely within `tskit`, and these native methods are often quicker and more efficient than methods that work with the sequence data.\n",
    "Since tree sequences also contain information about ancestry and times, there are also many analyses we can *only* perform on tree sequences.\n",
    "We'll learn about some of these methods in this worksheet.\n",
    "\n",
    " - [4.1: Site statistics](#4.1SiteStatistics)\n",
    " - [4.2: Branch statistics](#4.2BranchStatistics)\n",
    " - [4.3: Simplify](#4.3Simplify)\n",
    " - [4.4: Population and individual-based ancestry](#4.4PopulationBasedAncestry)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  A simulated tree sequence to use\n",
    "\n",
    "First, we'll simulate some tree sequence datasets with mutations to use in the rest of this notebook.\n",
    "\n",
    " <img src=\"pics/worksheet4-demography.png\" width=\"500\" height=\"500\">\n",
    "\n",
    "Here's a simple model of admixture in which two populations split 1000 generations ago from an ancestral population, and the third is formed after an admixture event 100 generations ago.\n",
    "One of the modern-day populations is larger than the other two. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demography = msprime.Demography()\n",
    "demography.add_population(name=\"SMALL\", initial_size=2000)\n",
    "demography.add_population(name=\"BIG\", initial_size=5000)\n",
    "demography.add_population(name=\"ADMIX\", initial_size=2000)\n",
    "demography.add_population(name=\"ANC\", initial_size=5000)\n",
    "demography.add_admixture(\n",
    "    time=100, derived=\"ADMIX\", ancestral=[\"SMALL\", \"BIG\"],\n",
    "    proportions=[0.5, 0.5])\n",
    "demography.add_population_split(\n",
    "    time=1000, derived=[\"SMALL\", \"BIG\"], ancestral=\"ANC\")\n",
    "demography.debug()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples={\"SMALL\": 100, \"BIG\": 100, \"ADMIX\" : 100},\n",
    "    demography=demography,\n",
    "    random_seed=2432,\n",
    "    sequence_length=5e7,\n",
    "    recombination_rate=1e-8\n",
    ")\n",
    "ts = msprime.sim_mutations(ts, rate=1e-8, random_seed=6151)\n",
    "ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll also use a tree sequence simulated using a model of a simple selective sweep."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_size=5000\n",
    "seq_length=50_000_000\n",
    "\n",
    "# Specify sweep parameters.\n",
    "sweep_model = msprime.SweepGenicSelection(\n",
    "    position=seq_length/2,\n",
    "    start_frequency=0.0001,\n",
    "    end_frequency=0.9999,\n",
    "    s=0.25,\n",
    "    dt=1e-6\n",
    ")\n",
    "\n",
    "# Simulate sweep genealogies.\n",
    "ts_sweep = msprime.sim_ancestry(\n",
    "    100,\n",
    "    model=[sweep_model, msprime.StandardCoalescent()],\n",
    "    population_size=pop_size,\n",
    "    sequence_length=seq_length,\n",
    "    recombination_rate=1e-8,\n",
    "    random_seed=1234,\n",
    "    )\n",
    "\n",
    "# Add finite-site mutations.\n",
    "ts_sweep = msprime.sim_mutations(\n",
    "    ts_sweep, rate=1e-8, random_seed=4321)\n",
    "ts_sweep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4.1SiteStatistics'></a>\n",
    "## 4.1 Site statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose you have $n$ sequences typed at $m$ different sites...\n",
    "\n",
    "```\n",
    "   ...GTAACGCGATAAGAGATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGAGATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGAAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGTAATAGCGTA...\n",
    "   ...GTAACGCGATAAGATATTAGCCCAAAAACACAGACATGGTAATAGCGTA...\n",
    "```\n",
    "\n",
    "...and you want to calculate mean pairwise diversity on these samples, i.e.\n",
    "\n",
    "$$ \\pi = \\dfrac{1}{n(n-1)/2}\\sum_{i=1}^{n-1} \\sum_{j=i+1}^n k_{ij}, $$\n",
    "\n",
    "where $k_{ij}$ is the number of sites at which sequences $i$ and $j$ carry a different allele.\n",
    "The scaling of this procedure is\n",
    "\n",
    "$$ O\\left( n^2 m \\right) $$\n",
    "\n",
    "ie. quadratic in the number of samples $n$, and linear in the number of sites $m$.\n",
    "However, there is an equivalent way of performing this calculation by assigning weights to the sample nodes, and propagating these values further up the tree using a 'summary function' at each mutation.\n",
    "\n",
    " <img src=\"pics/worksheet4-node-weights.jpeg\" width=\"350\" height=\"350\">\n",
    " \n",
    " This is what `tskit` calls a *site statistic* calculation, and because the operation is of order\n",
    " \n",
    " $$ O\\left( n + \\rho m (\\log(n))^2 \\right)  << O\\left( n^2 m \\right), $$\n",
    " \n",
    "the calculation is quick to run, especially on large datasets:\n",
    "\n",
    " <img src=\"pics/worksheet4-stat-speed.jpeg\" width=\"500\" height=\"500\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See the following paper for more details, especially if you are interested in coming up with tree sequence statistics of your own.\n",
    "\n",
    "Peter Ralph, Kevin Thornton, Jerome Kelleher, Efficiently Summarizing Relationships in Large Samples: A General Duality Between Statistics of Genealogies and Genomes, Genetics, Volume 215, Issue 3, 1 July 2020, Pages 779–797, https://doi.org/10.1534/genetics.120.303253"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The basic syntax: (nucleotide diversity)\n",
    "\n",
    "`tskit` uses very similar syntax for all of its inbuilt statistics, so we'll explore the options using `diversity()` as an example.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.diversity()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, `tskit` presents a normalised version of the statistic  scaled by the length of the region represented in `ts`. This allows you to make comparisons between different tree sequences that may be of different lengths. However, this isn't how all other genetic software computes diversity -- if you wish to disable `tskit`'s default behaviour, use the `span_normalise` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div = ts.diversity(span_normalise=False)\n",
    "div[()]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating statistics on subsets of the samples\n",
    "\n",
    "Remember that this dataset consists of samples from three contemporary populations here, of different sizes. We’d expect them to each have different diversity levels, and for these to differ from the overall (sample-wide) diversity rate. We can get this information out by specifying each of these with the `sample_nodes` argument.\n",
    "\n",
    "A quick and simple way to get all of the sample node IDs from a particular population is to use the `samples()` method. For instance, the following code returns a numpy array holding all of the sample node IDs from 'population 0':"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.samples(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By cross-checking against the population table, we can see the population IDs of all the populations in our simulated demography."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.tables.populations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we wanted to calculate the value of the diversity statistic for each of the three contemporary populations in our dataset ('SMALL', 'BIG' and 'ADMIX'), as well as for the pooled set of samples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples_of_interest=[ts.samples(population=0),\n",
    "                          ts.samples(population=1),\n",
    "                          ts.samples(population=2),\n",
    "                          ts.samples()]\n",
    "ts.diversity(sample_sets=samples_of_interest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is a 1-dimensional numpy array, where each element is a diversity statistic value for one of the sample sets we specified.\n",
    "Nucleotide diversity is lowest in the set of samples from the small population, and largest in the pooled set of samples, as you'd expect.\n",
    "\n",
    "Note that you can use any list of node IDs as inputs to `sample_sets`. This may be useful if your samples of interest correspond to something other than populations (for instance, samples that hold some phenotype of interest)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts.diversity(sample_sets=[[12, 14, 18, 7, 8],\n",
    "                          [23, 19, 10]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Genome scans\n",
    "\n",
    "So far, we’ve just been calculating statistics summarising diversity values along the entire simulated genome. However, in many cases, we might be more interested in how diversity varies along the genome. We can do this using the `windows` argument.\n",
    "\n",
    "We specify the start and end points of the sequence, and the locations of the breakpoints between each window. For instance, suppose we wanted to specify some windows of length 1Mb covering our 50Mb chromosome:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "breakpoints = [i*1e6 for i in range(0, 50 + 1)]\n",
    "print(breakpoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div = ts.diversity(sample_sets=samples_of_interest, windows=breakpoints)\n",
    "print(\"Dimensions of the output:\", div.shape, \"\\n\")\n",
    "print(\"Diversity values over the first 10 windows in each sample set:\")\n",
    "print(div[:10,:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output is a 50 by 4 array, one row for each of our specified windows. Each element of the row holds diversity value in some particular window amongst one of our sample sets. Let’s plot these:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names_to_plot = ['SMALL', 'BIG', 'ADMIX', 'ALL']\n",
    "lines = plt.plot(breakpoints[:-1], div)\n",
    "plt.legend(lines, names_to_plot);\n",
    "plt.xticks()\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.ylim(0,0.0003)\n",
    "plt.title(\"Windowed diversity values\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because `ts` holds a sample simulated under a selectively neutral simulation, any variation we do see along the genome is just statistical noise. \n",
    "What if we look at `ts_sweep` instead?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names_to_plot = ['SMALL', 'BIG', 'ADMIX', 'ALL_TS', 'TS_SWEEP']\n",
    "div_sweep = ts_sweep.diversity(windows=breakpoints)\n",
    "all_divs = np.append(div, np.reshape(div_sweep, newshape=(-1,1)), axis=1)\n",
    "lines = plt.plot(breakpoints[:-1], all_divs)\n",
    "plt.legend(lines, names_to_plot);\n",
    "plt.xticks()\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.ylim(0,0.0003)\n",
    "plt.title(\"Windowed diversity values\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-way statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Diversity is an example of a 'one-way' statistic. All of the samples are exchangeable in the calculation: if we were to re-label one of the samples with another label, we’d still get the same statistic.\n",
    "\n",
    "Other examples of one-way statistics are:\n",
    "\n",
    " - [allele_frequency_spectrum](https://tskit.dev/tskit/docs/stable/python-api.html#tskit.TreeSequence.allele_frequency_spectrum)\n",
    " - [segregating_sites](https://tskit.dev/tskit/docs/stable/python-api.html#tskit.TreeSequence.segregating_sites)\n",
    " - [Tajimas_D](https://tskit.dev/tskit/docs/stable/python-api.html#tskit.TreeSequence.Tajimas_D)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is *not* the case for certain other population genetic statistics that we call ‘multi-way statistics’.\n",
    "Consider Patterson's $f_3$ statistic:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f3_012 = ts.f3(sample_sets=[ts.samples(0), ts.samples(1), ts.samples(2)],\n",
    "               windows=breakpoints)\n",
    "f3_102 = ts.f3(sample_sets=[ts.samples(1), ts.samples(0), ts.samples(2)],\n",
    "               windows=breakpoints)\n",
    "f3_201 = ts.f3(sample_sets=[ts.samples(2), ts.samples(0), ts.samples(1)],\n",
    "               windows=breakpoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "names_to_plot = ['SMALL;BIG,ADMIX', 'BIG;SMALL,ADMIX', 'ADMIX;BIG,SMALL']\n",
    "lines = plt.plot(breakpoints[:-1], np.transpose([f3_012, f3_102, f3_201]))\n",
    "plt.legend(lines, names_to_plot)\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.title(\"Windowed F3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The green line `f3(ADMIX;BIG,SMALL)` is consistently below 0. This suggests that `ADMIX` might be an admixed population with gene flow from `BIG` and `SMALL` (which is indeed the case, since this is what we simulated)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are some other examples of multi-way statistics provided by `tskit`:\n",
    "\n",
    " - divergence\n",
    " - genetic_relatedness\n",
    " - Patterson's f2, f3, f4\n",
    " - genealogical_nearest_neighbours\n",
    " - Fst"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Exercise*. Modify the following code to compute $f_2$ statistics in 5Mb windows along the genome between the samples from populations SMALL and BIG, as well as between populations SMALL and ADMIX."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "breakpoints = [i*1e6 for i in range(0, 50 + 1)]\n",
    "f2_small_big = ts.f3(sample_sets=[ts.samples(0),\n",
    "                                  ts.samples(1),\n",
    "                                  ts.samples(2)],\n",
    "               windows=breakpoints)\n",
    "f2_small_admix = ts.f3(sample_sets=[ts.samples(2),\n",
    "                                  ts.samples(1),\n",
    "                                  ts.samples(0)],\n",
    "               windows=breakpoints)\n",
    "\n",
    "# Plot.\n",
    "names_to_plot = ['f2(SMALL,BIG)', 'f2(SMALL,ADMIX)']\n",
    "lines = plt.plot(\n",
    "    breakpoints[:-1], np.transpose([f2_small_big, f2_small_admix]))\n",
    "plt.legend(lines, names_to_plot)\n",
    "plt.xticks()\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.title(\"Windowed F2 values\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4.2BranchStatistics'></a>\n",
    "## 4.2 Branch statistics and the 'duality' of  tree-based statistics\n",
    "\n",
    "There are several different types of randomness in genetic models that interact with each other in complex ways.\n",
    "Firstly, there is randomness in the genealogical trees that are produced in a given demographic scenario.\n",
    "Consider how different one of our trees looks..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sts = ts.simplify(samples=[1, 11, 21, 31, 41, 51])\n",
    "tree = sts.at(0)\n",
    "display(SVG(tree.draw(width=300)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "from one much further along the genome:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = sts.at(sts.sequence_length - 1)\n",
    "display(SVG(tree.draw(width=300)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In addition, there is randomness caused by the mutational process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overlaid different mutations onto the same tree\n",
    "tss = ts.simplify(samples=[1, 11, 21, 31, 41, 51])\n",
    "mts1 = msprime.sim_mutations(tss, rate=1e-8, random_seed=1812)\n",
    "mts2 = msprime.sim_mutations(tss, rate=1e-8, random_seed=1903)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(SVG(mts1.last().draw(width=300)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(SVG(mts2.last().draw(width=300)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you calculate site statistics, or anything based on allele frequencies, *both* of these processes contribute to the statistical noisiness you see.\n",
    "\n",
    "However in tree sequences, you have information about branches which allows you to bypass this latter type of mutation.\n",
    "Instead of moving upwards along the trees and updating the statistic every time you come across a mutation, you can update the statistic based on the lengths of the branches.\n",
    "This should have some correspondence with the *number of mutations we may expect*. (This should certainly be true in simulated datasets, where we are certain of the correctness of the underlying trees).\n",
    "\n",
    "This is the basic idea behind the *branch statistics* in `tskit`.\n",
    "\n",
    "Here are the diversity stats we looked at before, this time with the branch versions included. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div_branch = ts.diversity(\n",
    "    sample_sets=samples_of_interest,\n",
    "    windows=breakpoints,\n",
    "    mode='branch')\n",
    "names_to_plot = ['SMALL', 'BIG', 'ADMIX', 'ALL']\n",
    "lines = plt.plot(breakpoints[:-1], div_branch)\n",
    "plt.legend(lines, names_to_plot);\n",
    "plt.ylim((0,30000))\n",
    "plt.title(\"Windowed diversity (branch)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "div_site = ts.diversity(\n",
    "    sample_sets=samples_of_interest,\n",
    "    windows=breakpoints,\n",
    "    mode='site')\n",
    "names_to_plot = ['SMALL', 'BIG', 'ADMIX', 'ALL']\n",
    "lines = plt.plot(breakpoints[:-1], div_site)\n",
    "plt.xlabel(\"Position on chromosome\")\n",
    "plt.legend(lines, names_to_plot)\n",
    "plt.ylim((0,0.0003))\n",
    "plt.title(\"Windowed diversity (site)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4.3Simplify'></a>\n",
    "## 4.3 Simplify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider two of the most common 'subset'-type operations we want to do with genomic datasets:\n",
    "\n",
    "- Look at the data for a subset of the samples in the original dataset.\n",
    "- Look at sequence/variant information at the specific sites that vary within that subsample.\n",
    "\n",
    "`simplify` is the tree sequence version of these operations.\n",
    "\n",
    "### The basic syntax\n",
    "\n",
    "`simplify` requires a list of sample IDs that you wish to include in the new, smaller tree sequence.\n",
    "Here, we return the tree sequence for all of the samples from population 'ADMIX' (which has a population label of `2`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tss = ts.simplify(ts.samples(2))\n",
    "tss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have a smaller tree sequence holding just those 20 sample chromosomes from population 'ADMIX'.\n",
    "Note that the number of sites and mutations is also reduced. This is because `simplify` has also removed all the mutations on the edges that were pruned away. Alleles arising from these mutations were only inherited by other samples, or are ancestral to everyone in this smaller subset.\n",
    "\n",
    "Although there are fewer edges, nodes and mutations in this newer, simplified tree sequence, it's not a drastic difference. This shows the sub-linearity and efficiency of tree sequence structures. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It turns out that operations like `simplify` are essential to generate tree sequences from forward-time simulation programs like SLiM — see [this](https://journals.plos.org/ploscompbiol/article?id=10.1371/journal.pcbi.1006581) if you are interested."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='4.4PopulationBasedAncestry'></a>\n",
    "## 4.4 Population and individual-based ancestry\n",
    "\n",
    "Why is understanding ancestry important?\n",
    "\n",
    " - **Demography and history:** Inference about the dates and composition of evolutionary changes and historical events.\n",
    " - **Medicine:** GWAS and risk prediction studies, admixture mapping studies.\n",
    " - **Genetic pipelines:** Phasing, imputation, genotyping errors, SNP ascertainment.\n",
    " \n",
    "Suppose your genealogical ancestors can be partitioned into distinct *populations*.\n",
    "\n",
    " <img src=\"pics/worksheet4-LA.png\" width=\"500\" height=\"500\">\n",
    " \n",
    " This is typically reported as global and local ancestry:\n",
    " \n",
    " <img src=\"pics/worksheet4-LA2.png\" width=\"500\" height=\"500\">\n",
    " \n",
    " However, there are limitations to population-based notions of ancestry:\n",
    " \n",
    "  - 'Populations' are often labels chosen by the researcher to make models more interpretable.\n",
    "  - There is always variation within populations.\n",
    "  - Different 'populations' may be relevant in different epochs.\n",
    "  - Gene flow is often gradual and spatial in nature.\n",
    "  - (Consider the history of population genetics in humans:) eugenics happened.\n",
    "  \n",
    "The alternative is to think about ancestry on the level of *individuals*.\n",
    " \n",
    " <img src=\"pics/worksheet4-ibd.png\" width=\"600\" height=\"600\">\n",
    " \n",
    "Both of these types of ancestry can be extracted using methods in `tskit`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tss = ts.simplify(samples=[1, 11, 21, 31, 41, 51])\n",
    "tss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Population-based ancestry with `link-ancestors`\n",
    "\n",
    "Tree sequences thus contain richly detailed information about local ancestry: any sample node that descends from a node in a given population at some location will have ancestry with the population at the location."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " <img src=\"pics/worksheet4-link-ancestors.png\" width=\"600\" height=\"600\">\n",
    " \n",
    "However, for realistically large and complicated simulations, recovering this information from the overall genealogies is challenging. A visually intuitive way to do this is to locate each sample haplotype on each tree and trace a path up the tree until an ancestral node from one of the populations of interest is reached.\n",
    "\n",
    "However, this approach will be quite inefficient. Any genealogical feature that is shared between different haplotypes, or across different regions of the genome, will be processed separately for each sample and each tree. Given the substantial correlations in genealogy that typically exist between individuals, and across genomes, this approach would require many repetitive operations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Run a simulation with populations and `add_census`\n",
    "\n",
    "By default, the tree sequences generated by `msprime` may not contain all of the ancestral nodes that are needed to extract information about local ancestry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demography = msprime.Demography()\n",
    "demography.add_population(\n",
    "    name=\"SMALL\", initial_size=200)\n",
    "demography.add_population(\n",
    "    name=\"BIG\", initial_size=500)\n",
    "demography.add_population(\n",
    "    name=\"ADMIX\", initial_size=200)\n",
    "demography.add_population(\n",
    "    name=\"ANC\", initial_size=500)\n",
    "demography.add_admixture(\n",
    "    time=20,\n",
    "    derived=\"ADMIX\",\n",
    "    ancestral=[\"SMALL\", \"BIG\"],\n",
    "    proportions=[0.5, 0.5]\n",
    ")\n",
    "# demography.add_census(time=599.5)\n",
    "demography.add_population_split(\n",
    "    time=600,\n",
    "    derived=[\"SMALL\", \"BIG\"],\n",
    "    ancestral=\"ANC\"\n",
    ")\n",
    "# demography.debug()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts = msprime.sim_ancestry(\n",
    "    samples={\"SMALL\": 0, \"BIG\": 0, \"ADMIX\" : 2},\n",
    "    demography=demography,\n",
    "    random_seed=1008,\n",
    "    sequence_length=10000,\n",
    "    recombination_rate=3e-8\n",
    ")\n",
    "# ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colour_map = {0:\"red\", 1:\"blue\", 2:\"purple\", 3:\"gray\"}\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": [
    "As we discussed above, 'ancestry' only makes sense with reference to a particular time. Generally what's meant by 'ancestry' is 'where the ancestors were at some particular point in the past'.\n",
    "\n",
    "The `add_census` method is a special demographic event that records a node on all lineages that are extant at a given time.\n",
    "This is needed to simulate complete information about local ancestry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demography = msprime.Demography()\n",
    "demography.add_population(\n",
    "    name=\"SMALL\", initial_size=200)\n",
    "demography.add_population(\n",
    "    name=\"BIG\", initial_size=500)\n",
    "demography.add_population(\n",
    "    name=\"ADMIX\", initial_size=200)\n",
    "demography.add_population(\n",
    "    name=\"ANC\", initial_size=500)\n",
    "demography.add_admixture(\n",
    "    time=20,\n",
    "    derived=\"ADMIX\",\n",
    "    ancestral=[\"SMALL\", \"BIG\"],\n",
    "    proportions=[0.5, 0.5]\n",
    ")\n",
    "demography.add_census(time=599.5)\n",
    "demography.add_population_split(\n",
    "    time=600,\n",
    "    derived=[\"SMALL\", \"BIG\"],\n",
    "    ancestral=\"ANC\"\n",
    ")\n",
    "# demography.debug()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq_length = 1e4\n",
    "ts = msprime.sim_ancestry(\n",
    "    samples={\"SMALL\": 0, \"BIG\": 0, \"ADMIX\" : 2},\n",
    "    demography=demography,\n",
    "    random_seed=1008,\n",
    "    sequence_length=seq_length,\n",
    "    recombination_rate=3e-8\n",
    ")\n",
    "# ts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "colour_map = {0:\"red\", 1:\"blue\", 2:\"purple\", 3:\"gray\"}\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": [
    "#### Step 2: Get the census nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "census_nodes = []\n",
    "for n in ts.nodes():\n",
    "    if n.time == 599.5:\n",
    "        census_nodes.append(n.id)\n",
    "        \n",
    "# census_nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3: Apply link-ancestors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ancestor_table = ts.tables.link_ancestors(\n",
    "    samples=ts.samples(), \n",
    "    ancestors=census_nodes\n",
    ")\n",
    "# ancestor_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that each node in the `parent` table above is one of the `census_nodes`. This is essentially what `link_ancestors` does -- it shows you directly which of your samples descend from which of the census ancestors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 3. Replace the nodes with their population label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_population_id(node, ts):\n",
    "    return ts.tables.nodes.population[node]\n",
    "\n",
    "local_ancestry = pd.DataFrame(\n",
    "    data = {\n",
    "        'left': ancestor_table.left,\n",
    "        'right': ancestor_table.right,\n",
    "        'population' : [ts.node(n).population for n in ancestor_table.parent],\n",
    "        'child' : ancestor_table.child\n",
    "    }\n",
    ")\n",
    "\n",
    "print(local_ancestry)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "(If you want: here's some more advanced code to sort and squash the table...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_sample = []\n",
    "new_left = []\n",
    "new_right = []\n",
    "new_population = []\n",
    "\n",
    "local_ancestry.sort_values(\n",
    "    by=['child','left'], inplace=True, ignore_index=True)\n",
    "\n",
    "for ind, row in local_ancestry.iterrows():\n",
    "    if ind > 0 and row['left']==new_right[-1] and row['population'] == new_population[-1] and row['child'] == new_sample[-1]:\n",
    "        new_right[-1] = row['right']\n",
    "    else:\n",
    "        new_sample.append(row['child'])\n",
    "        new_left.append(row['left'])\n",
    "        new_right.append(row['right'])\n",
    "        new_population.append(row['population'])\n",
    "        \n",
    "squashed_ancestry_table = pd.DataFrame({\n",
    "    'child': [int(i) for i in new_sample],\n",
    "    'left' : new_left,\n",
    "    'right': new_right,\n",
    "    'population' : [int(p) for p in new_population]\n",
    "})\n",
    "\n",
    "squashed_ancestry_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this information to plot local ancestry in each of our samples, for instance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_ancestry_chunk(row, chrom):\n",
    "    l = row.left/seq_length\n",
    "    r = row.right/seq_length\n",
    "    p = row.population\n",
    "    if int(p) == 0:\n",
    "        c = 'blue'\n",
    "    elif int(p) == 1:\n",
    "        c = 'red'\n",
    "    print('p is', p)\n",
    "    chunk = np.array([[l, 0], [r, 0], [r, 1], [l, 1]])\n",
    "    chrom.add_patch(Polygon(xy=chunk, color = c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chrom_labels = {0: 'chr0', 1: 'chr1'}\n",
    "colors = ['red', 'blue']\n",
    "length = ts.sequence_length\n",
    "\n",
    "fig, (chr0, chr1) = plt.subplots(2, figsize=(10,2))\n",
    "fig.suptitle('Ancestry in admixed individual')\n",
    "fig.frameon=False\n",
    "fig.legend(\n",
    "    handles = [Polygon(xy = np.array([[0,0],[0,1],[1,1],[1,0]]), color = 'blue'),\n",
    "              Polygon(xy = np.array([[0,0],[0,1],[1,1],[1,0]]), color = 'red')],\n",
    "    labels = ['SMALL', 'BIG'],\n",
    "    loc = 'right'\n",
    ")\n",
    "for ind, row in squashed_ancestry_table.iterrows():\n",
    "    if row.child > 1:\n",
    "        break\n",
    "    chunk = np.array([[row.left/length, 0], [row.right/length, 0],\n",
    "                      [row.right/length, 1], [row.left/length, 1]])\n",
    "    if chrom_labels[row.child] == 'chr0':\n",
    "        chr0.add_patch(Polygon(xy=chunk, color = colors[int(row.population)]))\n",
    "    elif chrom_labels[row.child] == 'chr1':\n",
    "        chr1.add_patch(Polygon(xy=chunk, color = colors[int(row.population)]))\n",
    "\n",
    "chr0.set_ylabel('Chrom 0')\n",
    "chr0.tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)\n",
    "chr1.set_xticks(ticks= [0.2, 0.4, 0.6, 0.8, 1.0], labels=[i*seq_length/5 for i in range(0, 5)])\n",
    "chr1.set_xlabel('Chromosomal position (bases)')\n",
    "chr1.set_ylabel('Chrom 1')\n",
    "chr1.tick_params(left=False, labelleft=False)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Consider re-running the code in this section with a larger sequence length to see a more interesting pattern of ancestry. (You might want to comment out the plotting code...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Individual-based ancestry with `ibd-segments`\n",
    "\n",
    "See [this](https://tskit.dev/tskit/docs/stable/ibd.html?highlight=ibd%20segments) preliminary tutorial and documentation [here](https://tskit.dev/tskit/docs/stable/ibd.html?highlight=ibd%20segments)."
   ]
  },
  {
   "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
}