1.Introduction_to_msprime_STUDENT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "# Practical 1. Recap of the neutral coalescent and an intro to coalescence simulations with msprime\n",
    "---\n",
    "\n",
    "At the end of this practical you should:\n",
    "\n",
    "- Understand what a coalescent simulation is and what the relevant data structures are.\n",
    "- Know how to run and analyse coalescent simulations with [`msprime`](https://tskit.dev/msprime/docs/stable/intro.html) and [`tskit`](https://tskit.dev/tskit/docs/stable/python-api.html#trees-and-tree-sequences).\n",
    "- Understand the extent to which the variance of the coalescent depends on the two major axis of sampling: # of loci, # of individuals ([Felsenstein 2004](https://academic.oup.com/mbe/article-lookup/doi/10.1093/molbev/msj079)).\n",
    "- Understand why it is natural (and helpful) to treat mutations separately from ancestry."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.1 Getting started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importing python tools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> We will generally import python tools in the first code cell of a jupyter practical notebook. \n",
    "    \n",
    "There are several tools we will use throughout the course: <code>msprime</code> is a powerful and well documented coalescent simulator. In fact, this notebook is based largely on the \"Running Simulations\" section in the <a href=\"https://tskit.dev/msprime/docs/stable/mutations.html\"><code>msprime</code> manual</a>.\n",
    "\n",
    "<a href=\"https://numpy.org/\"><code>numpy</code></a> is an open source python library that is widely used for scientific computing.  If you are new to `python` or need a refresher about basic <code>python</code> concepts, have a look at practical_0.\n",
    "<a href=\"https://matplotlib.org/\"><code>matplotlib</code></a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import msprime\n",
    "import tskit\n",
    "from IPython.display import SVG, set_matplotlib_formats\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "from collections import Counter"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.2 Build your own simulator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Population genetic models are algorithmic by nature: forwards in time they specify how a population evolves from one generation to the next, backwards in time models are framed in terms of the probabilities of the next event in the ancestry of a sample. It is unsurprising that a large number of population genetic simulators (both forward and backwards) exist. Simulating under the neutral coalescent, the simplest backwards in time model, is straightforward in principle.\n",
    "\n",
    "In fact, if we're only interested in coalescence times, a coalescent simulation is just a set of random draws of coalescence times from an exponential distribuition. We can do this with the `random` method in `numpy`. Let's draw 1,000 pairwise coalescence times for a population (of diploid individuals) with an effective size of $N_e= 10,000$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_replicates = 1_000\n",
    "pair_times_array = np.random.exponential(20_000, num_replicates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the result which we have assigned to a variable `pair_times_array` is a numpy array of floats:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(pair_times_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use python indexing to examine the first 10 simulation replicates:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_times_array[0:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> We can use the <a href=\"https://seaborn.pydata.org/generated/seaborn.displot.html#seaborn.displot/\"><code>displot()</code></a> plot method to plot the distribution of pairwise coalescence times. Unsurprisingly (given that we've drawn from an exponential distribution), the pairwise coalescence times are exponentially distributed:\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = sns.displot(x = pair_times_array, stat = \"density\", binwidth = 5_000)\n",
    "plt.xlabel(\"times (in generations)\");\n",
    "plt.ylabel(\"count\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 1:</b> Find the <a href=\"https://numpy.org/doc/stable/\"><code>numpy<code></a> methods for computing the mean and standard deviation of an array and check how well the mean of the 1_000 pairwise coalescence times simulated above matches the analytic expectation.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen:</b> A basic property of the coalescent is that the rate is proportional to the number of pairs of lineages, and so changes after each coalescent event. When there are  $i$ lineages the rate of coalescence is $\\lambda = \\binom i 2 = \\frac{i (i-1)}{2}$.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> Writing your own function is easy! The code chunck below defines a coalescence time simulator as a function. <code>coaltimes</code> has three arguments: <code>samplesize</code>,  <code>Ne</code> and  <code>replicates</code>, i.e. for a sample of size $n$ and a given $N_e$, <code>coaltimes</code> uses <a href=\"https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html\"><code>np.random.exponential()</code></a> to draw $n-1$ random coalescence times:\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def coaltimes(samplesize, Ne, replicates):\n",
    "    list_of_coal_intervals = [np.random.exponential(2* Ne/int(i*(i-1)/2), replicates) for i in range(2,samplesize+1)]\n",
    "    return np.asarray(list_of_coal_intervals)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 2:</b> Which <code>numpy</code> functions and methods does <code>coaltimes</code> use? The waiting time for most other possible processes that affect the ancestry of a sample (e.g. mutation, migration, recombination events) depends on the number of lineages that are present (rather than the number of possible pairs). Copy <code>coaltimes</code> in the cell below and modify it to define a new function that draws waiting times at a rate proportional to the number of lineages (rather than the number of pairs).\n",
    "    \n",
    "<b>Hint:</b> Read the docs of <a href=\"https://numpy.org/doc/stable/reference/random/generated/numpy.random.exponential.html\"><code>np.random.exponential()</code></a> to understand what input parameters this function takes.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 3:</b>  Use <code>coaltimes</code> to simulate 1,000 replicate sets of coalescence times for a sample of $n=9$ sequences (don't forget to assign the output to a new, helpful variable name).\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 4:</b> Use the funtion <code>type()</code> and the method <code>shape</code> to check what form the output of <code>coaltimes</code> takes.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 5:</b> Compute the $t_{MRCA}$ of each simulation replicate (as the sum of internode intervals) and use <code>sns.displot()</code> to visualize the distribution of $t_{mrca}$ arcoss replicates.\n",
    "    \n",
    "    \n",
    "<b>Hint: </b> Assign the $t_{mrca}$ to a new variable <code>tmrca</code>\n",
    "       </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 6:</b> What is the mean $t_{mrca}$ arcoss replicates?\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 7:</b> How well does this match the theoretical expectation $E[t_{mrca}]$?\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen hint:</b>\n",
    "work out the expected times between coalescent intervals first and use the fact that $E[t_{MRCA}] = \\sum_{i=2}^{n-1} E[t_i]$\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "      <b>Exercise 8:</b> What is missing from this simulation?\n",
    "      </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.3 A simple `msprime` ancestry simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can set up a single simulation using `sim_ancestry()` which is a method of the coalescent simulator `msprime`. The `samples` argument of `sim_ancestry` defines the sampled individuals whose history will be simulated. In the simplest case, we provide a single integer which defines the number of samples."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> Typing the variable assigned to this <code>msprime</code> simulation generates a useful, tabular summary of the simulation output.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "            <div>\n",
       "              <style>\n",
       "                .tskit-table thead tr th {text-align: left;padding: 0.5em 0.5em;}\n",
       "                .tskit-table tbody tr td {padding: 0.5em 0.5em;}\n",
       "                .tskit-table tbody tr td:first-of-type {text-align: left;}\n",
       "                .tskit-details-label {vertical-align: top; padding-right:5px;}\n",
       "                .tskit-table-set {display: inline-flex;flex-wrap: wrap;margin: -12px 0 0 -12px;width: calc(100% + 12px);}\n",
       "                .tskit-table-set-table {margin: 12px 0 0 12px;}\n",
       "                details {display: inline-block;}\n",
       "                summary {cursor: pointer; outline: 0; display: list-item;}\n",
       "              </style>\n",
       "              <div class=\"tskit-table-set\">\n",
       "                <div class=\"tskit-table-set-table\">\n",
       "                  <table class=\"tskit-table\">\n",
       "                    <thead>\n",
       "                      <tr>\n",
       "                        <th style=\"padding:0;line-height:21px;\">\n",
       "                          <img style=\"height: 32px;display: inline-block;padding: 3px 5px 3px 0;\" src=\"https://raw.githubusercontent.com/tskit-dev/administrative/main/tskit_logo.svg\"/>\n",
       "                          <a target=\"_blank\" href=\"https://tskit.dev/tskit/docs/latest/python-api.html#the-treesequence-class\"> Tree Sequence </a>\n",
       "                        </th>\n",
       "                      </tr>\n",
       "                    </thead>\n",
       "                    <tbody>\n",
       "                      <tr><td>Trees</td><td>1</td></tr>\n",
       "                      <tr><td>Sequence Length</td><td>1.0</td></tr>\n",
       "                      <tr><td>Time Units</td><td>generations</td></tr>\n",
       "                      <tr><td>Sample Nodes</td><td>18</td></tr>\n",
       "                      <tr><td>Total Size</td><td>3.8 KiB</td></tr>\n",
       "                      <tr>\n",
       "                        <td>Metadata</td><td style=\"text-align: left;\">No Metadata</td></tr>\n",
       "                    </tbody>\n",
       "                  </table>\n",
       "                </div>\n",
       "                <div class=\"tskit-table-set-table\">\n",
       "                  <table class=\"tskit-table\">\n",
       "                    <thead>\n",
       "                      <tr>\n",
       "                        <th style=\"line-height:21px;\">Table</th>\n",
       "                        <th>Rows</th>\n",
       "                        <th>Size</th>\n",
       "                        <th>Has Metadata</th>\n",
       "                      </tr>\n",
       "                    </thead>\n",
       "                    <tbody>\n",
       "                    \n",
       "                  <tr>\n",
       "                    <td>Edges</td>\n",
       "                      <td>34</td>\n",
       "                      <td>1.1 KiB</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Individuals</td>\n",
       "                      <td>9</td>\n",
       "                      <td>276 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Migrations</td>\n",
       "                      <td>0</td>\n",
       "                      <td>8 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Mutations</td>\n",
       "                      <td>0</td>\n",
       "                      <td>16 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Nodes</td>\n",
       "                      <td>35</td>\n",
       "                      <td>988 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Populations</td>\n",
       "                      <td>1</td>\n",
       "                      <td>224 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        ✅\n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Provenances</td>\n",
       "                      <td>1</td>\n",
       "                      <td>943 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                  <tr>\n",
       "                    <td>Sites</td>\n",
       "                      <td>0</td>\n",
       "                      <td>16 Bytes</td>\n",
       "                      <td style=\"text-align: center;\">\n",
       "                        \n",
       "                      </td>\n",
       "                    </tr>\n",
       "                \n",
       "                    </tbody>\n",
       "                  </table>\n",
       "                </div>\n",
       "              </div>\n",
       "            </div>\n",
       "            "
      ],
      "text/plain": [
       "<tskit.trees.TreeSequence at 0x7fb6b9d739a0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "simple_ts = msprime.sim_ancestry(9);\n",
    "simple_ts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> Strictly speaking, <code>sim_ancestry()</code> always returns a treesequence (more on these practical 3). However, in this case we have only simulated a single tree. We can use the method <code>.first()</code> to extract this tree:\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_tree = simple_ts.first();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> We can plot our simulated tree using the method <code>.draw_svg()</code>:\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SVG(simple_tree.draw_svg(size=(500, 500)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 9:</b>  Why does the tree we simulated contain 18 sample nodes and 35 nodes in total? How would you change the arguments of <code>sim_ancestry</code> to match the simulation we have assumed for our DIY coalescence time simulator above?\n",
    "    \n",
    "<b>Hint: </b> Check what arguments <a href=\"https://tskit.dev/msprime/docs/stable/ancestry.html\"><code>sim_ancestry</code></a> takes.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 10:</b> Use the <a href=\"https://tskit.dev/tskit/docs/stable/python-api.html#the-tree-class\"><code>.root</code></a> attribute to find the rootnode of this tree and the <code>time()</code> method to find the $t_{mrca}$. What are the units of this node time?\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use a `for` loop to loop over all nodes in the tree to find the associated coalescence times. Note the use of the `.nodes()` method in the loop below which returns a `generator` object. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simple_tree.nodes()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "coal_times =[simple_tree.time(node) for node in simple_tree.nodes()]\n",
    "print(sorted(coal_times))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.4  A simple `msprime` mutation simulation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have so far only simulated genealogies, i.e. genetic ancestry without worrying about mutations. \n",
    "\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen:</b> One of the most useful properties of the coalescent is that it allows us to separate the stochatic process of genetic ancestry (i.e. drift backwards) from the stochatic process of mutations. We can sprinkle mutations onto a simulated tree in a second step.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> We use the method <code>sim_mutations</code> to sprinkle mutations onto our simulated ancestry:\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "ts_with_muts_1 = msprime.sim_mutations(simple_ts, rate=0.2)\n",
    "ts_with_muts_2 = msprime.sim_mutations(simple_ts, rate=0.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 11:</b> Visualize the trees for both replicates and check how they differ.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default `sim_ancestry` and `sim_mutations` use different random seeds. This means that every ancestry and mutation simulation will be different. It is useful to be able to replicate the same simulation; setting the same random seed guarantees this. Try this by rerunning the code chunk above!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.5 Using `msprime` to find the distribution of genealogies and coalescent times"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have so far only simulated single genealogies. We are usually interested in using coalescent simulations to understand the **distribution** of genealogies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> The <code>num_replicates</code> argument of <code>sim_ancestry()</code> provides a convenient and efficient way to iterate over a number of replicate simulations. Let's use this to examine the distribution of the time to the most recent common ancestor.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_replicates = 1_000\n",
    "tmrca_msprime = np.zeros(num_replicates)\n",
    "replicates = msprime.sim_ancestry(9, ploidy=1, num_replicates=num_replicates, random_seed=42)\n",
    "for replicate_index, ts in enumerate(replicates):\n",
    "    tree = ts.first()\n",
    "    tmrca_msprime[replicate_index] = tree.time(tree.root)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the output of `sim_ancestry()` is a generator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "type(replicates)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> We will use <code>for</code> loops like the one above for populating numpy arrays with simulation results throughout the course, so it's worth understanding the syntax used: <code>np.zeros()</code> sets up an empty <code>numpy</code> array. It’s important to note that the replicate simulations are generated lazily here on demand — the set of simulations we have specified using <code>sim_ancestry</code> and assigned to the <code>replicates</code> variable is also a <code>generator</code>. This means that simulations are only specified but not actually run when we define <code>replicates</code>. \n",
    "\n",
    "Inside the <code>for</code> loop we use the built-in <code>enumerate()</code> function to easily get the index of each replicate in the list. Each replicate simulation is done sequentially as we go through the loop. We therefore only use the memory that we need to represent the results of one simulation.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Note:</b> the type of  <code>sim_ancestry()</code> changes when we use the  <code>num_replicates </code> argument. If num_replicates is not specified or None, we return an instance of <code>tskit.TreeSequence</code>. If it is specified, we return an iterator over a set of <code>tskit.TreeSequence </code> instances.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 12:</b> What is the mean $t_{mrca}$? How does the distribution compare to the one we obtained using our own coalescence time simulator <code>coaltimes</code> at the start of this practical?\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "##SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.6 Counting trees..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have **still** only considered coalescence times and ignored topologies. How random are tree topologies under the standard neutral coalescent? The number of possible tree topologies grows very quickly with sample size and is **enormous** even for relatively small samples.\n",
    "\n",
    "We can generate all trees using `tskit.all_trees(i, span=1)`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> Use <a href=\"https://seaborn.pydata.org/generated/seaborn.lineplot.html\"><code>sns.lineplot()</code></a> for plotting simple functions and datasets.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_trees = [len(list(enumerate(tskit.all_trees(i, span=1)))) for i in range(2,8)]\n",
    "print(\"For a samples of 7 lineages, the number of unique tree topologies is\", num_trees[-1])\n",
    "n_samp = np.arange(2,8)\n",
    "plt.xlabel('n')\n",
    "sns.lineplot(x = n_samp, y =  num_trees)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 13:</b>  Use the method <a href=\"https://tskit.dev/tskit/docs/stable/topological-analysis.html#sec-tree-ranks\"><code>rank</code></a> to find the shape of each of the 1,000 simulated trees. Set up an empty numpy array and use a for loop over <code>enumerate(replicates)</code> to add the shape integer of each simulated tree to this array (analogous to the loop at the start of this section).\n",
    "\n",
    "<b>Hint: </b> Note that <code>replicates</code> is a generator object, which needs to be re-evaluated if you want to use the same set of simulations again. \n",
    "        </div>\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL rank generates a tuple, the first element is an integer which uniquely records the unlabelled shape of a tree."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> <a href=\"https://realpython.com/python-counter/\"><code>Counter</code></a> is extremely useful. The output is a python dictionary.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 14:</b> Use a <code>Counter</code> to find out how many unique tree shapes you have simulated. How does this compare to the number possible bifurcating tree shapes, which for $n=9$ is 46 (Felsenstein, 2004)?\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Counter` returns a python dictionary. `most_common()` returns the key and value (in this case the count) of a `Counter`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "most_freq_rank, count = count_dict.most_common()[0]\n",
    "print(most_freq_rank)\n",
    "print(count)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 15:</b> If we focus on the root partition and consider a sample of $n=9$, how likely are asymmetric trees, i.e. with a (8,1) compared to very symetric trees (5,4)? Check the frequency of all root partitions in the 1,000 replicates.\n",
    "\n",
    "Use the function <code>get_root_bipartition</code> defined below on every replicate and store the result in an empty array <code>rep_partitions</code>!\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_root_bipartition(tree):\n",
    "    root_children = tree.children(tree.root)\n",
    "    left_clade_n = len(sorted(tree.leaves(root_children[0])))\n",
    "    right_clade_n = len(sorted(tree.leaves(root_children[1])))\n",
    "    return sorted((left_clade_n, right_clade_n))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_root_bipartition(simple_tree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Try to use the code chunk below to count the number of root partitions and print the results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = Counter(rep_partitions)\n",
    "count_array = np.zeros(4)\n",
    "for index, root_partition in enumerate(counts):\n",
    "    count_array[index] =  counts[root_partition]\n",
    "    print(\"There are\", counts[root_partition], \"genealogies with a root partition\", root_partition)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe:</b> Statistical tests with <code>scipy</code>. We can check whether counts are significantly different from each other with a $\\chi^2$ test:\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "#By hand is good for the soul\n",
    "chi_square = ((count_array - 250)**2)/250;\n",
    "print(chi_square.sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "#scipy\n",
    "from scipy.stats import chisquare\n",
    "chisquare(count_array)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}