coal-lab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Human Variation and Diesase Coalescent Computer Lab\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Welcome to the coalescent computer lab! Here we'll explore some fundamentals of the coalescent using exploratory data analysis (i.e. computing summary statistics / visualization) and simulation. We'll learn about & employ a powerful coalescent simulator called `msprime`, which allows us to efficiently generate genealogies under a given demography and subsequently sprinkle mutations on these genealogies to generate sequence data. Coalescent simulation is a very important tool to be equipped with for the modern population geneticist for a number of reasons...\n",
    "\n",
    "* It allows us to develop intuition for a simple generative model that can help to explain patterns in real data\n",
    "* We can use simulations to compare the accuracy of methods and see where they break down or succeed\n",
    "* Simulation can be used in inference, for instance to obtain monte carlo estimates for particular steps of a intractable model or perform approximate Bayesian inference (ABC)\n",
    "\n",
    "A quick note about computing: The environment we are in is called a `jupyter notebook`. A `jupyter notebook` is an interactive computational portal that allows us to make documents that combine simple text formats (markdown / latex) and code. This is quite powerful as we can write up analyses and explorations in the notebook and annotate visualizations & code with text and even math. Jupyter is primarily geared for `python` but one can actually change the \"kernel\" of the notebook to use `R` and even `julia`. Today we'll actually be using a very interesting python package called `rpy2` which allows us to run `R` commands in a `python` “kernel”. We can even pass objects created in python into an `R` code block. How this works will become more clear as we get started. \n",
    "\n",
    "Alright, now let us begin by importing various libraries that we'll be using in this notebook!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports / Configuration"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we load a jupyter extension that allows us to use `rpy2` to pass objects back and forth between `R` and `python` code blocks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext rpy2.ipython"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is how we import packages in python. `numpy` is a python package for creating, manipulating and operating on arrays. As we mentioned previously, `msprime` is a python package for performing coalescent simulations. It expands upon [Dick Hudson's](http://home.uchicago.edu/rhudson1/), a UChicago professor and coalescent pioneer, `ms` software. `tskit` is a complementary package that allows us to work with the output of `msprime` simulations. We will also load the `Ipython` package which will allow us to better visualize the coalescent trees we simulate. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import msprime \n",
    "import IPython\n",
    "import tskit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a quick example of how we can pass python objects into `R`. Lets first create a simple matrix $\\mathbf{A}$ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A = np.array([[1,2], [3, 4]])\n",
    "A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets now import $\\mathbf{A}$ into `R`. To use `R` run the `%%R` command in the top of the code block ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i A\n",
    "A"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Cool! That seems to work. Lets try making simple plot ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "hist(rnorm(1000))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now equipped to run the rest of the notebook ..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Coalescent simulation using `msprime`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "msprime is a widely used software for runnning coalescent simulations. There is a tutorial located here (https://msprime.readthedocs.io/en/stable/tutorial.html). This is a really nice [jupyter notebook](https://github.com/jhmarcus/spg-chapter/) written by [Jerome Kelleher](https://github.com/jeromekelleher), the author of `msprime`. Click the launch binder button at the bottom of the page. We strongly encourage you to change parameters in the notebook and build intuition on how that affects properties of coalescent trees and see how it matches your predictions from the theory we've learned in class."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Writing your own coalescent simulations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's start by running our first simulation! To do this we will use the function msprime.simulate(), passing in the sample size and effective population size $N_e$ as arguments. Here we are simulating a sample size of 4 with an effective population size of 1,000. Since we haven't provided a recombination rate, we are only simulating a single coalescent tree. Because we didn't set a seed and the coalescent is a random process, every time you run the code block you should observe a different tree.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_sequence = msprime.simulate(sample_size=4, Ne=1000)\n",
    "tree = tree_sequence.first()\n",
    "IPython.display.SVG(tree.draw())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have simulated a coalescent tree, we might want to learn some properties of it. For example, let's find the $T_{MRCA}$ for our tree. To do this we can use tree.time() which outputs the time to the most recent common ancestor in generations. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "T_mrca = tree.time(tree.root)\n",
    "T_mrca"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, this is just one simulation. Let's run 10,000 simulations and find the average $T_{MRCA}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_reps = 10000\n",
    "T_mrca = np.zeros(num_reps)\n",
    "for j, ts in enumerate(msprime.simulate(sample_size=4, Ne=1000,num_replicates=num_reps)):\n",
    "    tree = ts.first()\n",
    "    T_mrca[j] = tree.time(tree.root)\n",
    "np.mean(T_mrca)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Compare the result above to our theoretical expectation for $T_{MRCA}$ for $n=4$ and $N_e = 1000$. Then re-run the 10,000 simulations for $N_e = 10000$, do you expect the average $T_{MRCA}$ to increase or decrease, why? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R \n",
    "# your R code goes here\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your python code goes here \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding Recombination  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "While simulating a single coalescent tree is useful, often times we are interested in a locus of a specific length under the influence of recombination. This will create a series of coalescent trees over the length of our locus. In `msprime` we can use the same msprime.simulate() function and add two arguments, length and recombination_rate. The length parameter is the number of base pairs and recombination rate is in units of per base per generation. Here we are simulating trees across over a 10kb region with a recombination rate of $2×10^{−8}$ per base per generation, with a diploid effective population size of 1000.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_sequence = msprime.simulate(sample_size=4, Ne=1000, length=1e4, recombination_rate=2e-8, random_seed=31)\n",
    "# Let's look at the trees we made and the intervals they occur over \n",
    "for tree in tree_sequence.trees():\n",
    "    print(\"-\" * 20)\n",
    "    print(\"tree {}: interval = {}\".format(tree.index, tree.interval))\n",
    "    print(tree.draw(format=\"unicode\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have printed out the series of 4 coalescent trees for our locus and the intervals they occur over."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Adding mutation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So far we have simulated a region of chromosome with recombination, now we want to put mutations down onto our sequence of coalescent trees. Mutations are generated in msprime by throwing mutations down on the branches of trees at a particular rate. The mutations are generated under the infinite sites model, and so each mutation occurs at a unique (floating point) point position along the genomic interval occupied by a tree. The mutation rate for simulations is specified using the mutation_rate parameter of msprime.simulate(). Below we will simulate 50kb or non-recombiing chromosome and add mutations onto the coalescent tree."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_sequence = msprime.simulate(\n",
    "    sample_size=10, Ne=1000, length=50e3, mutation_rate=2e-8, random_seed=30)\n",
    "tree = tree_sequence.first()\n",
    "for site in tree.sites():\n",
    "     for mutation in site.mutations:\n",
    "            print(\"Mutation @ position {:.2f} over node {}\".format(\n",
    "                site.position, mutation.node))\n",
    "IPython.display.SVG(tree.draw())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From here we can create a genotype matrix for our samples. Python and R store integer numbers in different formats. The int8 command will convert python’s format into R’s format so we can work with it later using R code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = tree_sequence.genotype_matrix() \n",
    "# We need to convert data type so we can work with the genotype matrix in R\n",
    "G = np.int8(G)\n",
    "G"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the genotype matrix each individual is a separate column and each mutation is a row. If we look at individual 2, we see that the should have a \"1\" indicating that they have the mutation for mutations 3, 5, 6, and 8. Remember that python indexes from zero individuals are numbered 0,1,..,9. So individual 2 is the third column of the genotype matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question 2** Write a function to calculate the site frequency spectrum from the genotype matrix. Use your function to plot the SFS for G. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i G\n",
    "# R code goes here\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## python code goes here "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also use `tskit` to calculate some familar summary statistics directly from our tree sequence object. tree_sequence.diversity() calculates nucleotide diversity, $\\pi$. tree_sequence.segregating_sites() calculates the number of segregating sites, the argument span_normalise=False tells the function to give us the number of segregating sites rather than the density of segregating sites. Finally, tree_sequence.Tajimas_D() calculates Tajima's D for our sample.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Nucleotide diversity (pi)\n",
    "print(tree_sequence.diversity()) \n",
    "# Number of segregating sites \n",
    "print(tree_sequence.segregating_sites(span_normalise=False))\n",
    "# Tajima's D\n",
    "print(tree_sequence.Tajimas_D())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating under different demographic models "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thus far we have focused on running coalescent simulations under a constant population size model. While these models are very important for building intuition and establishing theoretical expectations, we know that human populations have not had a constant population size. Now we will get familiar with doing simulations under models that have varrying population sizes through time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Below are three different simulations. For each one please write which scenario you think it represents: standard coalescent, recent bottelneck, or recent population growth and say why. Which one has the most rare mutations? Which one the most intermediate frequency ones? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demographic_events = [msprime.PopulationParametersChange(time=0.5, growth_rate=0.05, population=0)]\n",
    "tree_sequence = msprime.simulate(\n",
    "    sample_size=10, Ne=1000, length=50e3, random_seed=31,demographic_events=demographic_events)\n",
    "IPython.display.SVG(tree_sequence.first().draw())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree_sequence = msprime.simulate(\n",
    "    sample_size=10, Ne=1000, length=50e3, random_seed=13)\n",
    "IPython.display.SVG(tree_sequence.first().draw())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demographic_events = [msprime.InstantaneousBottleneck(time=0.5, strength=1, population=0)]\n",
    "tree_sequence = msprime.simulate(\n",
    "    sample_size=10, Ne=1000, length=5e4, random_seed=30,demographic_events=demographic_events)\n",
    "\n",
    "IPython.display.SVG(tree_sequence.first().draw())\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulating under human demographic models using stdpopsim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to do more realistic simulations for human populations, we want to use demographic models inferred from real data. `stdpopsim` is a software that uses `msprime` to do according to pre-defined demographic models. The `stdpopsim` tutorial is located here (https://stdpopsim.readthedocs.io/en/latest/tutorial.html).  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run stdpopsim we are first going to pick the species (humans) and the genetic map. Here we are simulating 5% of chromosome 2 (approximately 12 Mb). The demographic model is the \"OutOfArfica_3G09\" model which allows us to sample individuals from either present day YRI, CEU, and CHB populations. Here we will simulate 50 samples from each population. The output is the same tree sequence object we saw in the msprime examples above.   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import stdpopsim\n",
    "import json\n",
    "\n",
    "species = stdpopsim.get_species(\"HomSap\")\n",
    "contig = species.get_contig(\"chr2\", length_multiplier=0.05)\n",
    "model = species.get_demographic_model('OutOfAfrica_3G09')\n",
    "samples = model.get_samples(50, 50, 50) #(YRI, CEU, CHB)\n",
    "engine = stdpopsim.get_engine('msprime')\n",
    "ts = engine.simulate(model, contig, samples, seed=16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Use the stdpopsim catalog to look up details of the three population out of Africa model (https://stdpopsim.readthedocs.io/en/latest/catalog.html#sec-catalog). What the population size of CHB after the EU/AS split? What is the growth rate of CHB after the split? "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have the tree sequence object for our simulation let's calculate some summary statistics.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Which population do you think will have the most segregating sites and why? How about the lowest number?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer:**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's do simulations to check our intuition. We will first define our three population samples. YRI are samples 1-50, CEU are 51-100, and CHB are 101-150. We can then tell `tskit` to calculate the number of segregating sites separately for each population."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "YRI = ts.samples()[:50]\n",
    "CEU = ts.samples()[50:100]\n",
    "CHB = ts.samples()[100:150]\n",
    "\n",
    "print(ts.segregating_sites(span_normalise=False,sample_sets=[YRI,CEU,CHB]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Which population do you think will have highest pairwise diversity? How about the lowest? Give your prediction and explain your reasoning. Use the `tskit` we used above for caclulating pariwise diversity to check your prediction. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## your code goes here\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Given what you have just calculated and what you know about human population history, which population do you think will have the highest Tajima's D value? Which one the lowest? Explain your thinking then check your prediction. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# your code goes here \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Question** Below is code to extract the genotype matrix for each population. Use these genotype matrices and the function you wrote above to plot the SFS for all populations and comment on differences that you see.  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Answer** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = ts.genotype_matrix()\n",
    "G = np.int8(G)\n",
    "\n",
    "G_YRI = G[:,0:50]\n",
    "G_CEU = G[:,50:100]\n",
    "G_CHB = G[:,100:150]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i G_CHB,G_YRI,G_CEU\n",
    "# SFS G_YRI\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i G_CHB,G_YRI,G_CEU\n",
    "# SFS G_CEU\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i G_CHB,G_YRI,G_CEU\n",
    "# SFS G_CHB\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python code goes here "
   ]
  }
 ],
 "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
}