5.Detecting_selective_sweeps_2024_STUDENT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a963fed",
   "metadata": {},
   "outputs": [],
   "source": [
    "import allel\n",
    "import msprime\n",
    "import tskit\n",
    "import math\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "from IPython.display import SVG, set_matplotlib_formats\n",
    "from matplotlib import pyplot as plt\n",
    "#jupyter magic\n",
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "615357c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget --no-check-certificate 'https://www.dropbox.com/scl/fi/c677k4cwysifazgyh4vm3/p5_data_2024.zip?rlkey=gs48i6xckqntscesi7upglto5&dl=0' -q -O p5_data_2024.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74515614",
   "metadata": {},
   "outputs": [],
   "source": [
    "! unzip p5_data_2024.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "a14a5f1c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys; sys.path.insert(0,'/home/jovyan/PGA_course_2024/p5_data_2024/')\n",
    "import sf_functions as sf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d3cab18",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Practical 5. Detecting selective sweeps in population genomic data\n",
    "\n",
    "At the end of this practical you should:\n",
    "\n",
    "- Understand how positive selection acts to favor a beneficial mutation  \n",
    "- Understand the role of drift/randomness on allele frequency trajectories and fixation probability\n",
    "- Understand the effect of positive selection on linked neutral variation\n",
    "- Understand how sweepfinder works using simulation data\n",
    "- Be able to perform a Selective sweep scan on real data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "private-vegetarian",
   "metadata": {},
   "source": [
    "# 5.1 The direct effect of selection on a new mutation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "oriented-congress",
   "metadata": {},
   "source": [
    "We have so far considered the ancestry of samples backwards in time, and, in practical 2, have writen a simple coalescence time simulator in python. In this practical, we study the effect of positive selection on a focal *de novo* mutation. It is easiest to first think about the allele frequency trajectory of a new mutation forwards in time. \n",
    "\n",
    "Let's start by implementing a simple forwards in time simulation for the frequency trajectory of a new mutation under the neutral Wright Fisher (WF) model. We can then update this code to model positive selection acting on the new allele, and compare the dynamics of neutral and beneficial alleles."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eaca814a",
   "metadata": {},
   "source": [
    "### 5.1.1 Simulating drift fowards in time "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "finnish-deviation",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen: Neutral Wright Fisher model</b> \n",
    "    \n",
    "You should already be familiar with the neutral WF model (see Peter's population genetics lectures and PGA lecture 5). Consider a haploid population of $N_e$ individuals reproducing according to the WF model. We want to simulate the frequency trajectory of a *de novo* neutral mutation that appears in an individual.\n",
    "\n",
    "In this model, each individual contributes on-average one offspring to the next generation. However, there is variance in offspring number: some individuals have 2, 3, or more children. And many will have no offspring at all.\n",
    "\n",
    "From a coalescent perspective, we can think of this as each individual giving birth at the same _rate_, $1$. The next generation is formed by randomly assigning a parent to each offspring with probability $p = 1/N_e$, i.e. the rate of birth of the focal individual $1$ relative to the total birth rate in the population $N_e$.\n",
    "\n",
    "Therefore, the frequency $k$ of the new mutation in the next generation, i.e. the probability that the focal individual has $k$ offspring, can be obtained simply by binomial sampling with \"success probability\" $p = 1/N_e$. \n",
    "\n",
    "$$P[k] = \\binom{N_e}{k} p^k ( 1- p)^{N_e - k}$$\n",
    "\n",
    "Furthermore, when $X$ individuals currently carry the new mutation, mutant individuals give birth with a total rate $X$. The probability of observing $k$ offspring is still binomially distributed but with $p=\\frac{X}{N_e}$.  \n",
    "\n",
    "By repeatedly drawing a random number of offspring and updating the probability $p$ of inheriting the new mutation, we can simulate the frequency trajectory of this mutation from generation to generation until it is eventually lost or fixed by drift.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ultimate-julian",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python recipe: <code>while</code> loops </b>\n",
    "\n",
    "We want to track the frequency trajectory of a new neutral mutation forwards in time, stopping the simulation when the mutation is lost or fixed. \n",
    "\n",
    "To do this, we will use <code>while</code> loop.  You've already used <code>for</code> loops to iterate over arrays, e.g.<code>range(0,10)</code>. This loop stops once the end of the array is reached. A <code>while</code> loop, on the other hand, repeats until a particular condition is met. This is useful when we don't know beforehand how many iterations are needed to get the desired result.\n",
    "    \n",
    "Here's a <code>while</code> loop that generates a sequence of 20 random integers drawn uniformly from (0, 1, ... 9). In each iteration of the loop, the sequence gets longer, so we can use the <b>condition</b> <code> len(seq)&lt;20 </code> to stop the loop. \n",
    "\n",
    "E.g., when <code>len(seq)=17</code>, <code>len(seq)&lt;20 </code> evaluates as <code>True</code>. When <code> len(seq)=20</code>, the condition evaluates as <code>False</code> and the loop stops.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "unlimited-death",
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = []\n",
    "while len(seq)<20:\n",
    "    a = np.random.randint(0,10)\n",
    "    seq.append(a)\n",
    "print(seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "documented-spank",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "When there are multiple conditions that can stop the loop, it is often easier to write them within the loop itself. To do this we use <code>while True:</code> (meaning always proceed) and use <code>if</code> statements to define the conditions under which we halt the loop by calling <code>break</code>.\n",
    "\n",
    "Below, we generate a sequence of 20 random numbers drawn randomly from (0,1,...,9), but here we stop adding values to the sequence if a \"5\" is drawn.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "statewide-airfare",
   "metadata": {},
   "outputs": [],
   "source": [
    "seq = []\n",
    "while True:\n",
    "    a = np.random.randint(0,10)\n",
    "    seq.append(a)\n",
    "    if len(seq)>= 20:  ## the same condition as above, stop if seq length is 20, writing inequality the other way##\n",
    "        break\n",
    "    if a==5:  ## an additional condition to stop adding values when a '5' is drawn\n",
    "        break\n",
    "print(seq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wound-hawaii",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\"> <b>Warning!</b> This type of loop can be dangerous. If a stop condition can't be met, the computer will run forever, and it is a good idea to define a limit to the number of iterations, just in case.\n",
    "        </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "continuing-conspiracy",
   "metadata": {},
   "source": [
    "Now that we know how to use a `while` loop, we can implement the neutral WF simulation as a function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "curious-chile",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sim_neutral(i,Ne): \n",
    "    # i is any unique ID for the simulation. is converted to string. Usually pass an integer. \n",
    "    # N_e is the (haploid) population size \n",
    "    \n",
    "    maxGens = 4000 #largest number of generations we will simulate.\n",
    "    bCount = Ne -1 #initial freq of ancestral allele\n",
    "    BCount = 1 #initial freq of new mutation  \n",
    "    traj = [] # will be populated by the frequency trajectory of neutral mutation. \n",
    "    \n",
    "    while True: #starts the simulation\n",
    "        \n",
    "        traj.append(BCount) #append current frequency to trajectory\n",
    "        \n",
    "        \n",
    "        if(BCount >= Ne):\n",
    "            break\n",
    "        if(BCount <= 0):\n",
    "            break\n",
    "        if(len(traj)>maxGens): \n",
    "            break\n",
    "        \n",
    "        \n",
    "        ### sample to make the next generation ###\n",
    "        \n",
    "        pB = (BCount*(1))/(BCount*(1)+ bCount*(1)) #prob of sampling the focal allele\n",
    "        \n",
    "        BCount = np.random.binomial(Ne,pB)  # random number of derived mutations drawn \n",
    "                                            # from binomial distribution to produce next generations \n",
    "        bCount = Ne - BCount \n",
    "        \n",
    "    # create and return a pandas data frame for the replicate smulation \n",
    "    gens = list(range(len(traj)))\n",
    "    sim = [str(i) for _ in range(len(traj))]\n",
    "    df = pd.DataFrame({\"sim\":sim,\"gens\":gens,\"freq\":traj})\n",
    "    return(df)    \n",
    "\n",
    "sim_neutral(1,10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "365b9d2d",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 1:</b> In the code above, what are the three conditions that stop the <code>while</code> loop and what do they mean in the context of simulating the allele frequency trajectory?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "202cc606",
   "metadata": {},
   "source": [
    "SOL:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "provincial-demonstration",
   "metadata": {},
   "source": [
    "The wrapper function below returns a single data frame that combines data across replicate simulations. We can then use `sns.relplot` to plot the trajectories and distinguish replicates by their unique IDs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "furnished-memorial",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_neutral_replicates(num_reps,Ne):\n",
    "    replicates = [sim_neutral(i,Ne) for i in range(num_reps)]\n",
    "    trajectories = pd.concat(replicates,ignore_index=True)\n",
    "    return(trajectories)\n",
    "\n",
    "get_neutral_replicates(num_reps=3,Ne=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "perceived-bridge",
   "metadata": {},
   "outputs": [],
   "source": [
    "neu_trajectories= get_neutral_replicates(num_reps=50,Ne=10)\n",
    "sns.relplot(x = \"gens\", y = \"freq\",hue=\"sim\",kind = \"line\",data = neu_trajectories,legend=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "secondary-emerald",
   "metadata": {},
   "source": [
    "### 5.1.2 Simulating selection **AND** drift"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "rocky-modem",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen:</b> \n",
    "\n",
    "Suppose now that the *de novo* mutation confers a fitness advantage of $s$. In the WF model, carriers of the beneficial mutation leave on average a greater number of offspring, that is, reproduce at rate $(1 + s)$, while the remaining $N_e - 1$ individuals carrying the ancestral allele still give birth at rate $1$. The probability that any offspring inherits the new beneficial mutation is still binomially distributed, but now  $p = \\frac{(1)(1 + s)}{(1)(1+s) + (N_e - 1)(1)}$. \n",
    "\n",
    "As for the neutral case above, when there are $ X > 1$ copies of the beneficial mutation, the probability observing $k$ copies in the next generation is binomially distributed with $p= \\frac{X(1 + s)}{X(1 + s) + ( N_e  - X)( 1)}$.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "political-detroit",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 2:</b>\n",
    " <ul>\n",
    "    <li> copy the code for running the neutral simulation below and modify it to include positive selection for simulating the frequency trajectory of a new beneficial mutation. The new function <code>sim_selection(i,Ne,s)</code> will take a third parameter for the strength of selection. </li>\n",
    "        <li>check that it returns a data frame with the same formatting as above.</li>\n",
    "        <li>once your selection simulations work, use the wrapper function <code>get_selection_replicates()</code> below to run 50 replicate simulations.</li>\n",
    "        <li> plot the results.</li>\n",
    "      </ul>\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "educated-testimony",
   "metadata": {},
   "outputs": [],
   "source": [
    "## complete this function using the neutral simulation code to include positive selection in the simulation ##\n",
    "def sim_selection(i,Ne,s):\n",
    "    # i is any unique ID for the simulation. is converted to string. Usually pass an integer. \n",
    "    # N_e is the (haploid) population size \n",
    "    # s is is selection coefficient\n",
    "    \n",
    "\n",
    "sim_selection(i=1,Ne=50,s=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "czech-mixture",
   "metadata": {},
   "outputs": [],
   "source": [
    "#a wrapper function to get replicate simulations for selection\n",
    "def get_selection_replicates(num_reps,Ne,sb):\n",
    "    replicates = [sim_selection(i,Ne,sb) for i in range(num_reps)]\n",
    "    trajectories = pd.concat(replicates,ignore_index=True)\n",
    "    return(trajectories)\n",
    "\n",
    "sel_trajectories= get_selection_replicates(num_reps=50,Ne=1000,sb=0.1)\n",
    "sns.relplot(x = \"gens\", y = \"freq\",hue=\"sim\",kind = \"line\",data = sel_trajectories,legend=False)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "empirical-settle",
   "metadata": {},
   "source": [
    "### 5.1.3 The effect of population size and selection strength on trajectories"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "restricted-metadata",
   "metadata": {},
   "source": [
    "Using the selection simulations you have implemented above, we can compare the fate of new mutations for different population sizes and strengths of selection. For each parameter combination, we plot 200 replicate trajectories of the new mutation.\n",
    "\n",
    "We plot neutral simulations (first column) weak selection simulations (second column) and strong selection simulations (third column). Each row corresponds to a different population size (100, 1000, 10000), increasing from top to bottom.\n",
    "\n",
    "For the selection scenarios, we also plot the expected deterministic allele frequency trajectory for a beneficial mutation (see lecture 5) using the function <code>calc_expected_traj</code> defined below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "occasional-testing",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calc_expected_traj(Ne,sb):\n",
    "    gensTotal = 10*Ne\n",
    "    pB = 1\n",
    "    traj = []\n",
    "    for t in range(gensTotal):\n",
    "        thisPB =  pB*math.exp(sb*t)/((pB*math.exp(sb*t) + Ne - 1))\n",
    "        traj.append(thisPB)\n",
    "        if thisPB >= (Ne-1)/(Ne):\n",
    "            traj = np.array(traj)\n",
    "            traj = traj*(Ne)\n",
    "            return(traj)\n",
    "            break\n",
    "    traj = np.array(traj)\n",
    "    traj = traj*(Ne)\n",
    "    return(traj)\n",
    "\n",
    "calc_expected_traj(10,0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "surprising-dodge",
   "metadata": {},
   "outputs": [],
   "source": [
    "pop_sizes = [100,1000,10000]\n",
    "sel_strengths = [0.01,0.05]\n",
    "num_reps = 200\n",
    "\n",
    "fig, axs = plt.subplots(3,3,figsize=(16,12))\n",
    "\n",
    "for rowidx, Ne in enumerate(pop_sizes):\n",
    "    trajectories = get_neutral_replicates(num_reps,Ne)\n",
    "    sns.lineplot(x = \"gens\", y = \"freq\",hue=\"sim\",data = trajectories,legend=False,ax = axs[rowidx,0])\n",
    "\n",
    "for rowidx, Ne in enumerate(pop_sizes):\n",
    "    for colidx,sb in enumerate(sel_strengths):\n",
    "        trajectories = get_selection_replicates(num_reps,Ne,sb)\n",
    "        sns.lineplot(x = \"gens\", y = \"freq\",hue=\"sim\",data = trajectories,legend=False,ax = axs[rowidx,colidx +1])\n",
    "        axs[rowidx,colidx+1].plot(calc_expected_traj(Ne,sb),\":k\",linewidth=4.0)\n",
    "        \n",
    "axs[0,0].set_title(\"neutral\")\n",
    "axs[0,1].set_title(\"weak selection\")\n",
    "axs[0,2].set_title(\"strong selection\")\n",
    "\n",
    "for ax,size in list(zip(axs[:,0],pop_sizes)):\n",
    "    ax.set_ylabel(size)\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "armed-retrieval",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 3:</b> Execute the code cell above! \n",
    "     <ul>\n",
    "    <li> How do trajectories differ for neutral vs. weak vs strong selection? How does this depend on pop size?</li> \n",
    "    <li> What is the fate of most new mutations, and why? Write down your conclusions?</li> \n",
    "    </ul\n",
    "         \n",
    "         \n",
    "<b>Note:</b> the different scale of the x-axis (number of generations) in each plot.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "great-application",
   "metadata": {},
   "source": [
    "SOLUTION:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "720e1c5a",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Class discussion:</b> \n",
    "    <ul>\n",
    "    <li> How do trajectories differ when the allele frequency is low vs high? What dynamics of the model can explain these patterns?</li> \n",
    "    <li> How do sim trajectories differ from the expectation (dashed line above) and why?</li> \n",
    "        </ul>\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0700dada",
   "metadata": {},
   "source": [
    "SOL: "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad2c7844",
   "metadata": {},
   "source": [
    "SOL: "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "patent-cleveland",
   "metadata": {},
   "source": [
    "# 5.2 The effect of selection on linked neutral sites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "entitled-conspiracy",
   "metadata": {},
   "source": [
    "Now that we have studied the direct effects of selection on the evolutionary trajectory of a new beneficial mutation, we want to investigate how fixation of an adaptive allele affects linked neutral sites in the genome.\n",
    "\n",
    "Here we will use simulations in `msprime` to investigate the effect of positive selection on linked neutral variants."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "roman-cement",
   "metadata": {},
   "source": [
    "### 5.2.0 Neutral Simulations\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "early-provider",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-warning\">\n",
    "    <b>Python refresher: simulating under the coalescent with <code>msprime</code> </b>  \n",
    "\n",
    "Below, we simulate 50Mb of sequence assuming $N_e = 10,000$ (see <i>Homo</i> like sims in practical 3) under the standard neutral coalescent and plot the total branch length for each tree along the sequence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "allied-petroleum",
   "metadata": {},
   "outputs": [],
   "source": [
    "#This specifies parameters we will use for our main  simulations of a neutral chromsome and a chromosome with a sweep\n",
    "\n",
    "Ne = 10_000\n",
    "L = 50e6\n",
    "rbp = 2e-9\n",
    "n=20\n",
    "sb=0.05\n",
    "\n",
    "#We are specifying a mutation rate, but will use branch-based statistics as much as possible (for efficiency):\n",
    "mu = 1e-8 \n",
    "#We are setting up an empty data frame which we will fill as we go along:\n",
    "divdf = pd.DataFrame() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "separate-physics",
   "metadata": {},
   "outputs": [],
   "source": [
    "#This sets up a single neutral simulation given the parameters above: \n",
    "ts_neutral = msprime.sim_ancestry(\n",
    "    samples=n,\n",
    "    model=[msprime.StandardCoalescent()],\n",
    "    population_size=Ne,\n",
    "    recombination_rate=rbp,\n",
    "    sequence_length=L\n",
    "    #,\n",
    "    #random_seed=seed_neu\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "located-integral",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Making windows of 10kb \n",
    "winStartPositions = np.arange(0,L,10_000)\n",
    "slidingWindows=np.append(winStartPositions,L)\n",
    "\n",
    "divdf[\"position\"]=winStartPositions\n",
    "#computing diversity in windows\n",
    "div_neu = ts_neutral.diversity(windows=slidingWindows, mode='branch')\n",
    "divdf[\"neutral\"]=div_neu\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(18,3))\n",
    "sns.lineplot(data =divdf, x=\"position\",y=\"neutral\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bc26338",
   "metadata": {},
   "source": [
    "### 5.2.1 The effect of selection on genetic diversity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bearing-hobby",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 4 - msprime simulation for positive selection:</b> Modify the code below to implement a sweep model. Calculate diversity in windows, add this to the data frame \"divdf\" as a new column called \"sweep\", and produce a scatter plot (as above). For the sweep model: Place the beneficial mutation at the center of the chromosome. For the initial frequency, assume that the beneficial allele is a <i>de novo</i> mutation. For the final frequency, stop the sweep when there is only one copy of the ancestral mutation remaining.\n",
    "    \n",
    "<b>Hints: </b> Check out the msprime documentation of [<code>SweepGenicSelection</code>](https://tskit.dev/msprime/docs/latest/api.html#msprime.SweepGenicSelection) to:\n",
    "<ul>\n",
    "    <li> Define a positive selection model placing the beneficial mutation at the center of the chromosome.</li> \n",
    "    <li> Combine the selection phase with a neutral coalescent to simulate selective sweep.</li>\n",
    "    <li> Compute and plot genetic diversity in sliding windows as before.</li>\n",
    "    </ul>\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "published-coffee",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Insert parameter values for the position, start and end frequency in the code below to define the \n",
    "# sweep model.\n",
    "\n",
    "sweepModel = msprime.SweepGenicSelection(\n",
    "    position =   ,\n",
    "    #ben mutation at the centre\n",
    "    start_frequency=   ,# initial freq of de-novo mutation\n",
    "    end_frequency=    ,#final frequency of de-novo ben mutation\n",
    "    s=     ,#strength of selection\n",
    "    dt = 1/(40*Ne)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17c70965",
   "metadata": {},
   "outputs": [],
   "source": [
    "## define the correct model below\n",
    "\n",
    "ts_sel = msprime.sim_ancestry(\n",
    "    samples=n,\n",
    "    model=    , #simulate the sweep first, followed by standard coal\n",
    "    population_size=Ne,\n",
    "    recombination_rate=rbp,\n",
    "    sequence_length=L\n",
    "    #,\n",
    "    #random_seed=seed_sel\n",
    ")\n",
    "print(ts_sel.num_trees)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "potential-brother",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Plot genetic diversity\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "tutorial-discipline",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 5:</b> How does selection affect the distribution of window-wise genetic diversity compared to the neutral distribution?\n",
    "\n",
    "<b> Hint:</b> Use <code>sns.histplot</code> to plot both distributions.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "treated-austin",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOLUTION:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mathematical-wrong",
   "metadata": {},
   "source": [
    "### 5.2.2 The effect of selective sweeps on the SFS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8d12f4",
   "metadata": {},
   "source": [
    "We will consider a smaller sample of $n=5$ (diploid) individuals, a shorter length of sequence (4Mb) to investigate the effect of sweeps on the SFS at nearby neutral sites. As in practical 2 we will compute the expected SFS by averaging across replicates (in this case 500)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "polish-semester",
   "metadata": {},
   "outputs": [],
   "source": [
    "# this defines a few local parameters for use in this section\n",
    "popSize = 10_000\n",
    "benStrength = 0.05\n",
    "chrLength = 4e6\n",
    "recRate = 1e-9\n",
    "sampleSize = 5 #diploids\n",
    "numReplicates = 500\n",
    "numWindows=4\n",
    "winBreaks = [chrLength/numWindows*i for i in range(0,numWindows+1)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "considered-belize",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>PopGen reminder: the SFS</b>\n",
    "    \n",
    "You have already enountered the analytic predictions for the SFS for a sample from a neutral WF population in practical 2. Given a sequence with a scaled mutation rate $\\theta$, the SFS (in terms of the expected number of iton mutations) is $S_i = \\theta/i$ for $i = 1,2,...,n-1$. See lecture 2.\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "nasty-conducting",
   "metadata": {},
   "source": [
    "The code chunk below obtains the expected neutral SFS for polarized `neutSFSUnfolded` and non-polarized data `neutSFSFolded` for a sample of $n$ individuals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad281bb8-0ed4-4499-9c5f-4e550d03d20b",
   "metadata": {},
   "outputs": [],
   "source": [
    "nHaplotypes = 2*sampleSize\n",
    "neutExp = np.array([1/i for i in range(1,2*sampleSize)])\n",
    "neutSFSUnfolded = neutExp\n",
    "neutSFSUnfolded /= neutSFSUnfolded.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "chicken-newton",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 6:</b> \n",
    "    <ul>\n",
    "    <li>use the unfolded SFS (<code>neutSFSUnfolded</code>, above) to derive the folded site frequency spectrum and save it as <code>neutSFSFolded</code> </li>\n",
    "\n",
    "        \n",
    "Hint: you can reverse the order of an array using np.flip(arr)\n",
    "        \n",
    "Check that  <code>neutSFSFolded</code> sums to 1. You do not need to normalize it.      \n",
    "    </ul>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b022b765-4d43-4873-ac53-5d34f90b25ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Sol:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "gothic-alexander",
   "metadata": {},
   "source": [
    "Let's investigate the effect of sweeps on the SFS at linked neutral sites. The code chunk below runs 500 replicate sweep simulations, assuming a single beneficial allele at `position = 1`\n",
    "- for each replicate, we get the SFS in 1Mb windows (non-overlapping) increasingly far from the center of the sweep, both for the folded and unfolded SFS\n",
    "- for each distance, we average the SFS across replicates.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e6b8839-630d-4fb6-835f-7a9f2954406a",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 7:</b> \n",
    "    Execute the code cell below.  It will plot the SFS for four windows at increasing distance from the sweep center as well as the corresponding site frequency spectrum. \n",
    "\n",
    "- what affect does the sweep have on the folded SFS?\n",
    "- how does this comopare to the unfolded SFS?\n",
    "- Does polarising the data provide more information?\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87586592-a575-45ee-852a-2694329061c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "popSize = 10_000\n",
    "benStrength = 0.05\n",
    "chrLength = 4e6\n",
    "recRate = 1e-9\n",
    "sampleSize = 5 #diploids\n",
    "numReplicates = 500\n",
    "numWindows=4\n",
    "winBreaks = [chrLength/numWindows*i for i in range(0,numWindows+1)]\n",
    "\n",
    "\n",
    "nHaplotypes = 2*sampleSize\n",
    "neutExp = np.array([1/i for i in range(1,2*sampleSize)])\n",
    "neutSFSUnfolded = neutExp\n",
    "neutSFSUnfolded /= neutSFSUnfolded.sum()\n",
    "\n",
    "flipped = np.flip(neutExp)\n",
    "idx = 0\n",
    "neutSFSFolded = []\n",
    "while idx<nHaplotypes/2:\n",
    "    if idx == nHaplotypes/2 -1:\n",
    "        neutSFSFolded.append(neutSFSUnfolded[idx]) # we only want to count this entry once\n",
    "    else: \n",
    "        neutSFSFolded.append(neutSFSUnfolded[idx] + flipped[idx]) # otherwise, we sum them to get the folded counts\n",
    "    idx +=1\n",
    "\n",
    "\n",
    "\n",
    "sweepModel = msprime.SweepGenicSelection(\n",
    "    position = 1,\n",
    "    start_frequency=1.0/(2*popSize),\n",
    "    end_frequency= 1 - 1.0/(2*popSize),\n",
    "    s=benStrength,\n",
    "    dt = 1/(40*popSize)\n",
    ")\n",
    "\n",
    "replicates = msprime.sim_ancestry(\n",
    "    samples=sampleSize,\n",
    "    model=[sweepModel,msprime.StandardCoalescent()],\n",
    "    population_size=popSize,\n",
    "    recombination_rate=recRate,\n",
    "    sequence_length=chrLength,\n",
    "    num_replicates = numReplicates\n",
    ")\n",
    "\n",
    "sfs_rep_table_folded = np.zeros((numReplicates,numWindows,2*sampleSize+1))\n",
    "sfs_rep_table_unfolded = np.zeros((numReplicates,numWindows,2*sampleSize+1))\n",
    "for replicate_index, ts in enumerate(replicates):\n",
    "    sfsUnfolded = ts.allele_frequency_spectrum(mode = \"branch\",polarised = True,windows=winBreaks)\n",
    "    sfsFolded = ts.allele_frequency_spectrum(mode = \"branch\",polarised = False,windows=winBreaks)\n",
    "    sfs_rep_table_folded[replicate_index] = sfsFolded\n",
    "    sfs_rep_table_unfolded[replicate_index] = sfsUnfolded\n",
    "\n",
    "#win_sfs_table: each row is the SFS for a window of 1Mb increasingly far from the sweep centre.\n",
    "win_sfs_table_folded = np.apply_along_axis(np.mean,0,sfs_rep_table_folded)\n",
    "win_sfs_table_folded = np.apply_along_axis(lambda x: x/sum(x),1,win_sfs_table_folded)\n",
    "win_sfs_table_folded = np.apply_along_axis(lambda x: x[1:-1],1,win_sfs_table_folded)\n",
    "#print(win_sfs_table_folded)\n",
    "\n",
    "#win_sfs_table: each row is the SFS for a window of 1Mb increasingly far from the sweep centre.\n",
    "win_sfs_table_unfolded = np.apply_along_axis(np.mean,0,sfs_rep_table_unfolded)\n",
    "win_sfs_table_unfolded = np.apply_along_axis(lambda x: x/sum(x),1,win_sfs_table_unfolded)\n",
    "win_sfs_table_unfolded = np.apply_along_axis(lambda x: x[1:-1],1,win_sfs_table_unfolded)\n",
    "#print(win_sfs_table_unfolded)\n",
    "\n",
    "\n",
    "##dataframe and plot for folded\n",
    "iTons = list(range(1,10))\n",
    "windows = ['first','second','third','fourth']\n",
    "dfWST_folded = pd.DataFrame()\n",
    "for name,vals in list(zip(windows,win_sfs_table_folded)):\n",
    "# This creates a temporary dataframe for each distance from the sweep \n",
    "    tempdf = pd.DataFrame({'itons':iTons,'freq':vals,'window':[name]*(2*sampleSize-1)})\n",
    "    dfWST_folded= pd.concat([dfWST_folded,tempdf])\n",
    "\n",
    "nSFS_folded = np.zeros(len(neutSFSUnfolded)) #add zero entries to unfolded SFS for plotting and store as \"a\"\n",
    "for idx, val in enumerate(neutSFSFolded):\n",
    "    nSFS_folded[idx]=val\n",
    "dfNeut_folded = pd.DataFrame({'itons':iTons,'freq':nSFS_folded,'window':['neutral']*(2*sampleSize-1)})\n",
    "\n",
    "dfWST_folded = pd.concat([dfWST_folded,dfNeut_folded])\n",
    "\n",
    "pFolded = sns.barplot(x=\"itons\",y=\"freq\",hue=\"window\",data=dfWST_folded)\n",
    "plt.show()\n",
    "\n",
    "\n",
    "##dataframe and plot for unfolded\n",
    "iTons = list(range(1,10))\n",
    "windows = ['first','second','third','fourth']\n",
    "dfWST_unfolded = pd.DataFrame()\n",
    "for name,vals in list(zip(windows,win_sfs_table_unfolded)):\n",
    "# This creates a temporary dataframe for each distance from the sweep \n",
    "    tempdf = pd.DataFrame({'itons':iTons,'freq':vals,'window':[name]*(2*sampleSize-1)})\n",
    "    dfWST_unfolded= pd.concat([dfWST_unfolded,tempdf])\n",
    "\n",
    "nSFS_unfolded = np.zeros(len(neutSFSUnfolded)) #add zero entries to unfolded SFS for plotting and store as \"a\"\n",
    "for idx, val in enumerate(neutSFSUnfolded):\n",
    "    nSFS_unfolded[idx]=val\n",
    "dfNeut_unfolded = pd.DataFrame({'itons':iTons,'freq':nSFS_unfolded,'window':['neutral']*(2*sampleSize-1)})\n",
    "\n",
    "dfWST_unfolded = pd.concat([dfWST_unfolded,dfNeut_unfolded])\n",
    "\n",
    "pUnfolded = sns.barplot(x=\"itons\",y=\"freq\",hue=\"window\",data=dfWST_unfolded)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "constant-university",
   "metadata": {},
   "source": [
    "SOLUTION:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "correct-adelaide",
   "metadata": {},
   "source": [
    "# 5.3 Genomic scans for adaptive evolution using SweepFinder\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7eea549d",
   "metadata": {},
   "source": [
    "## 5.3.0 Sweepfinder Background information"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "inner-grain",
   "metadata": {},
   "source": [
    "Here we will perform scans for selection using SweepFinder using the `ts_neutral` and `ts_sel` simulations obtained in **5.2.2** above (the ones with long chromosomes, $n=20$ and the beneficial mutation in the middle). Details of the method are found [Nielsen et al 2005](https://genome.cshlp.org/content/15/11/1566). We first learn to format the data correctly for use in SweepFinder and then think about how to interpret the results.\n",
    "\n",
    "\n",
    "We have re-implemented the sweepfinder approach in python for use in this tutorial.\n",
    "\n",
    "The two main ingredients we need are the mutation positions and counts as well as the site frequency spectrum.\n",
    "The site frequency spectrum can be given separately or it can be calculated from the list of mutation counts\n",
    "\n",
    "\n",
    "For now, we will restrict ourselves to variable sites that are covered across all individuals.  \n",
    "This is always true for the simulations, but we will need to be careful about this with the real data.\n",
    "\n",
    "Our version of  SweepFinder takes as input:\n",
    "\n",
    "    - a list of test sites\n",
    "    - the (haploid) number of samples n\n",
    "    - a list of mutations with frequencies 1 < x < n\n",
    "    - the positions of these mutations\n",
    "    - the neutral (background) SFS for x = [ 0, 1, ..., (n-1), n] mutations. we assume monormphic sites are absent, i.e. x[0].\n",
    "    - whether the sfs is folded \n",
    "    - whether fixed differences (to the outgroup used for polarising) are included\n",
    "\n",
    "\n",
    "\n",
    "SweepFinder returns a pandas data frame with columns corresponding to:\n",
    "\n",
    "    - position: location of the test site\n",
    "    - lnCL: the log likelihood for the best-fitting sweep model\n",
    "    - alpha: the relative strength of selection for the best-fitting sweep model\n",
    "    - log10 alpha:  (-1)*Log10(alpha)  for visualizing better, so higher values correspond to stronger selection\n",
    "    - lnCL diff:  the difference in log likelihood between the best-fitting sweep model and the neutral background model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "initial-religious",
   "metadata": {},
   "source": [
    "## 5.3.1 Sweep scans on simulations\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "986b7e74",
   "metadata": {},
   "source": [
    "### 5.3.1.1 The neutral background SFS and Mutation Counts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "intimate-intelligence",
   "metadata": {},
   "source": [
    "Below, We use the neutral simulation ts_neutral to obtain the background sitefrequency spectrum, here using branched statistics.\n",
    "\n",
    "Note that the SFS includes entries for mutations with counts i=0,1,2,...,n.  i=0 corresponds to monomorphic sites and should always have frequency 0. \n",
    "\n",
    "For the unfolded SFS, i=n may be non-zero if fixed-differences from an outgroup are included (here, theya re not). \n",
    "\n",
    "The folded SFS has frequencies of 0.0 for all 'major' allele counts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da9a5ff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "backgroundSFSFolded = ts_neutral.allele_frequency_spectrum(mode='branch',polarised=False)\n",
    "backgroundSFSFolded = backgroundSFSFolded/sum(backgroundSFSFolded)\n",
    "print('folded')\n",
    "print(backgroundSFSFolded)\n",
    "\n",
    "backgroundSFSUnfolded = ts_neutral.allele_frequency_spectrum(mode='branch',polarised=True)\n",
    "backgroundSFSUnfolded = backgroundSFSUnfolded/sum(backgroundSFSUnfolded)\n",
    "print('unfolded')\n",
    "print(backgroundSFSUnfolded)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e73cd5",
   "metadata": {},
   "source": [
    "So far, we have used branch-based statistics provided by msprime/tskit. However, when running sweepfinder scans on population genomic data (simulated or real), we can only work with polymorphisms.\n",
    "\n",
    "Below, we add mutations to bot the _neutral_ and _selection_ simulations, obtain mutation counts along the genome, both for a polarised and nonpolarised data set, and format this as an input table for SweepFinder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15392317",
   "metadata": {},
   "outputs": [],
   "source": [
    "mts_neutral = msprime.sim_mutations(ts_neutral, rate=mu)\n",
    "print(\"theta per base= \", 4*Ne*mu)\n",
    "print(\"observed num mutations:\", mts_neutral.num_mutations)\n",
    "print(\"expected num mutations\",mu*4*Ne*L* sum([1/i for i in range(1,2*n)])) #expected number\n",
    "\n",
    "\n",
    "biallelicSNPsNeutralFolded =  [[int(var.site.position), min([sum(var.genotypes),2*n - sum(var.genotypes)]),2*n,1] for var in mts_neutral.variants() if len(var.alleles)==2]\n",
    "biallelicSNPsNeutralFolded = pd.DataFrame(biallelicSNPsNeutralFolded,columns=['position','x','n','folded'])\n",
    "biallelicSNPsNeutralFolded = biallelicSNPsNeutralFolded[biallelicSNPsNeutralFolded[\"x\"]>0]\n",
    "biallelicSNPsNeutralFolded = biallelicSNPsNeutralFolded[biallelicSNPsNeutralFolded[\"x\"]<40]\n",
    "\n",
    "\n",
    "biallelicSNPsNeutralUnfolded =  [[int(var.site.position), sum(var.genotypes),2*n,0] for var in mts_neutral.variants() if len(var.alleles)==2]\n",
    "biallelicSNPsNeutralUnfolded = pd.DataFrame(biallelicSNPsNeutralUnfolded,columns=['position','x','n','folded'])\n",
    "biallelicSNPsNeutralUnfolded = biallelicSNPsNeutralUnfolded[biallelicSNPsNeutralUnfolded[\"x\"]>0]\n",
    "biallelicSNPsNeutralUnfolded = biallelicSNPsNeutralUnfolded[biallelicSNPsNeutralUnfolded[\"x\"]<40]\n",
    "\n",
    "biallelicSNPsNeutralUnfolded[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cd41538",
   "metadata": {},
   "outputs": [],
   "source": [
    "mts_sel = msprime.sim_mutations(ts_sel, rate=mu)\n",
    "print(\"theta per base= \", 4*Ne*mu)\n",
    "print(\"observed num mutations:\", mts_sel.num_mutations)\n",
    "print(\"expected num mutations\",mu*4*Ne*L* sum([1/i for i in range(1,2*n)])) #expected number\n",
    "\n",
    "\n",
    "biallelicSNPsSweepFolded =  [[int(var.site.position), min([sum(var.genotypes),2*n - sum(var.genotypes)]),2*n,1] for var in mts_sel.variants() if len(var.alleles)==2]\n",
    "biallelicSNPsSweepFolded = pd.DataFrame(biallelicSNPsSweepFolded,columns=['position','x','n','folded'])\n",
    "biallelicSNPsSweepFolded = biallelicSNPsSweepFolded[biallelicSNPsSweepFolded[\"x\"]>0]\n",
    "biallelicSNPsSweepFolded = biallelicSNPsSweepFolded[biallelicSNPsSweepFolded[\"x\"]<40]\n",
    "\n",
    "biallelicSNPsSweepFolded[0:5]\n",
    "\n",
    "biallelicSNPsSweepUnfolded =  [[int(var.site.position), sum(var.genotypes),2*n,0] for var in mts_sel.variants() if len(var.alleles)==2]\n",
    "biallelicSNPsSweepUnfolded = pd.DataFrame(biallelicSNPsSweepUnfolded,columns=['position','x','n','folded'])\n",
    "biallelicSNPsSweepUnfolded = biallelicSNPsSweepUnfolded[biallelicSNPsSweepUnfolded[\"x\"]>0]\n",
    "biallelicSNPsSweepUnfolded = biallelicSNPsSweepUnfolded[biallelicSNPsSweepUnfolded[\"x\"]<40]\n",
    "\n",
    "\n",
    "biallelicSNPsSweepUnfolded[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "paperback-soundtrack",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 8:</b> Check the output above. How does the total number of mutations ($S$) differ between the neutral and the sweep simulation? How close are either to the analytic expectation $E[S]$ under neutrality?\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94fa2d20-5d36-43e2-84ea-136aea576214",
   "metadata": {},
   "source": [
    "SOL:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "vulnerable-behalf",
   "metadata": {},
   "source": [
    "### 5.3.1.2 Running sweepfinder on simulation data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "correct-depression",
   "metadata": {},
   "source": [
    "Here, we will run a scan on the simulation data with the selective sweep using the __unfolded__ SFS.\n",
    "\n",
    "We already have the neutral background site frequency spectrum and mutation frequencies obtained above.\n",
    "\n",
    "Here, we define all the variables we pass to our version of sweepfinder:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87b6d48b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutPositions = np.array(biallelicSNPsSweepUnfolded[\"position\"])\n",
    "mutFreqs = np.array(biallelicSNPsSweepUnfolded[\"x\"])\n",
    "\n",
    "# mutPositions = np.array(biallelicSNPsNeutralUnfolded[\"position\"])\n",
    "# mutFreqs = np.array(biallelicSNPsNeutralUnfolded[\"x\"])\n",
    "\n",
    "L = 50e6\n",
    "nSamples = 40\n",
    "foldedBool = False\n",
    "fixedDiffBool = False\n",
    "\n",
    "testSites = np.arange(1,L,100_000)\n",
    "print(len(testSites))\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0df68c4c-82ea-49d4-851a-14e5d6842eac",
   "metadata": {},
   "source": [
    "**note** to avoid errors, reload the sweepfinder module each time before running the scan, as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e121ff70",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "reload(sf)\n",
    "\n",
    "start = time.time()\n",
    "resultsSimUnfolded = sf.run_sweepfinder(testSites,nSamples,mutFreqs,mutPositions,backgroundSFSUnfolded,foldedBool,fixedDiffBool)\n",
    "print(f\"time it took: {time.time()-start}\")\n",
    "resultsSimUnfolded"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "entertaining-branch",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "    <b>PopGen: Estimating $s$ from SweepFinder's $\\alpha$ </b>\n",
    "    \n",
    "$\\alpha$ is a measure of the rate of recombination relative to the effective strength of selection in the population.  $\\alpha = r/s* Log( 2 N_e)$. If we have an estimate of $r$ and $N_e$ (or know them, as in these simulations) we can convert this estimate of $\\alpha$ to an estimate of $s$, the fitness effect of the beneficial mutation.\n",
    "    \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wrong-fellow",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 9:</b> For both the neutral and selection scans, use the formula above to calculate the strength of selection from the estimated 'Alpha' values and add this information as a new column <code>sb</code> in the respective data frames. \n",
    "    \n",
    "<b>Hint</b>: Use <code>math.log(x)</code> to calculate the natural logarithm of \"x\".\n",
    "    </div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "limited-expert",
   "metadata": {},
   "outputs": [],
   "source": [
    "# \"lambda x: x^2\" creates a one-time-use function \"lambda\" that accepts some parameter 'x' and\n",
    "# returns that value squared, x^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "surprised-clerk",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83fc30a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(resultsSimUnfolded[\"position\"],resultsSimUnfolded[\"lnCL diff\"],color='b')\n",
    "plt.show()\n",
    "plt.plot(resultsSimUnfolded[\"position\"],resultsSimUnfolded[\"sb\"],color='b')\n",
    "plt.show()\n",
    "# plt.plot(resultsSimUnfolded[\"position\"],resultsSimUnfolded[\"log10 alpha\"],color='b')\n",
    "# plt.show()\n",
    "\n",
    "mostLikelyIdx = resultsSimUnfolded[\"lnCL diff\"].idxmax()\n",
    "print(resultsSimUnfolded.iloc[mostLikelyIdx])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2389c0f-cae4-4663-a047-c79bc68ac5a4",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 10:</b> Sweep scan using the <b>folded</b> SFS\n",
    "    \n",
    "Reload the sweepfinder module, and modify the code above to run the scan but this time using the folded neutral SFS and the folded mutation counts from the sweep simulation.\n",
    "    \n",
    "save the resulting data frame as \"resultsSimFolded\"\n",
    "    \n",
    "plot the results alongside those for the unfolded site frequency spectrum (that we saved as \"resultsSimUnfolded\" above)\n",
    "\n",
    "Where is the most likely position for the center of the sweep?  How do these results compare to those from the unfolded SFS?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66dd9d31-165c-4424-845c-2ab7d568fd62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solution:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "spare-island",
   "metadata": {},
   "source": [
    "# 5.4 Signatures of selection in _Heliconius_ butterflies\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "wound-colorado",
   "metadata": {},
   "source": [
    "Let's look at a SweepFinder scan on a real data set! We will first import a vcf file containing variation data from several individuals of *Heliconius* butterflies belonging to different closely related species and mimicry races. These data were obtained by mapping WGS Illumina reads against the *Heliconius melpomene* reference genome and are from [(Martin et al 2020)](https://datadryad.org/stash/dataset/doi:10.5061/dryad.sk2pd88).\n",
    "\n",
    "You will analyze this data set much more closely in practical 6. For simplicity, we focus on chromosome 18 and restrict our analysis to 10 individuals from a single sub-species: _H. melpomene malleti_."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "capable-cradle",
   "metadata": {},
   "source": [
    "### 5.4.1 Importing and Subsetting the _Heliconius_ data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "architectural-therapy",
   "metadata": {},
   "source": [
    "The cell below imports the variant data from the vcf file and and subsets it using scikit allel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "cfca7d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "## import the vcf file retaining the genotype calls, variant positions, and sample names\n",
    "samples_mal = ['mal.CS1002', 'mal.CS1011', 'mal.CS1815', 'mal.CS21', 'mal.CS22', 'mal.CS24', 'mal.CS586', 'mal.CS594', 'mal.CS604', 'mal.CS615']\n",
    "gt_key, pos_key, sample_key = 'calldata/GT', 'variants/POS', 'samples'\n",
    "vcf_f ='p5_data_2024/heli_chrom18.vcf.gz'\n",
    "helidata = allel.read_vcf(vcf_f, fields=[gt_key, pos_key, sample_key], samples=samples_mal)\n",
    "L = 16802709\n",
    "nSamples = 10*2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7143e7f",
   "metadata": {},
   "source": [
    "As in the last practical, we store the variant positions and the genotypes as separate arrays and turn the genotype array into an array of allele counts:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "possible-general",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"allel allel-DisplayAs2D\"><span>&lt;AlleleCountsArray shape=(1441823, 3) dtype=int32&gt;</span><table><thead><tr><th></th><th style=\"text-align: center\">0</th><th style=\"text-align: center\">1</th><th style=\"text-align: center\">2</th></tr></thead><tbody><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">0</th><td style=\"text-align: center\">18</td><td style=\"text-align: center\"> 0</td><td style=\"text-align: center\"> 0</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">1</th><td style=\"text-align: center\">17</td><td style=\"text-align: center\"> 1</td><td style=\"text-align: center\"> 0</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">2</th><td style=\"text-align: center\">15</td><td style=\"text-align: center\"> 3</td><td style=\"text-align: center\"> 0</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">...</th><td style=\"text-align: center\" colspan=\"4\">...</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">1441820</th><td style=\"text-align: center\">20</td><td style=\"text-align: center\"> 0</td><td style=\"text-align: center\"> 0</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">1441821</th><td style=\"text-align: center\">12</td><td style=\"text-align: center\"> 8</td><td style=\"text-align: center\"> 0</td></tr><tr><th style=\"text-align: center; background-color: white; border-right: 1px solid black; \">1441822</th><td style=\"text-align: center\">10</td><td style=\"text-align: center\">10</td><td style=\"text-align: center\"> 0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "<AlleleCountsArray shape=(1441823, 3) dtype=int32>\n",
       "18  0  0\n",
       "17  1  0\n",
       "15  3  0\n",
       "...\n",
       "20  0  0\n",
       "12  8  0\n",
       "10 10  0"
      ]
     },
     "execution_count": 69,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "positions = helidata['variants/POS']\n",
    "heliGTs = allel.GenotypeArray(helidata['calldata/GT']);\n",
    "\n",
    "mal_counts = heliGTs.count_alleles()\n",
    "mal_counts"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "honest-identification",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 11:</b>\n",
    "    The data from this VCF file only contain polymorphic sites, so why does the first entry only have a single allele? Use scikit allel's <code>is_biallelic()</code> function to restrict the counts and position arrays to biallelic sites (see practical 4). Store the new counts and positions as <code>ba_mal_counts</code> and <code>ba_positions</code> (check how many sits are filtered out).\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "stainless-gregory",
   "metadata": {},
   "outputs": [],
   "source": [
    "#SOL\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "302ffaf0-f222-43b4-a11d-fa5f47c919ca",
   "metadata": {},
   "source": [
    "We now use the <code>allel.windowed_diversity()</code> method (see pratical 4, exercise 8) to plot genetic diversity along the chromosome in windows of 100_000 bases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "855c98e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "window_size = 100_000\n",
    "pi, windows, bases, counts = allel.windowed_diversity(ba_positions, ba_mal_counts, size=window_size, start=0, stop=ba_positions[-1], step=int(window_size/2))\n",
    "win_coord = windows[:,1] - (window_size / 2)\n",
    "win_pi_df = pd.DataFrame({'win_coord':win_coord,'pi':pi})\n",
    "fig, ax = plt.subplots(figsize=(18,6))\n",
    "sns.lineplot(x='win_coord', y='pi', data = win_pi_df)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c3fb7d6-cf3d-49a5-8f74-a3c074036d95",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 12:</b>\n",
    "    Complete the function __getMajorMinor(countArr)__ below to calculate the major and minor allele counts from the individual allele counts in the array above. Hint: use max() and sum() functions.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1aa112d-302c-4b82-a4e6-0709e7b79907",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def getMajorMinor(countArr):#countArr is the list of allele counts [a0,a1,a2] for some position\n",
    "    \n",
    "    maxCount = \n",
    "    \n",
    "    minorCount = \n",
    "    \n",
    "    return([maxCount,minorCount])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4791276b-b664-4180-8a0d-d88d7b9d50f7",
   "metadata": {},
   "source": [
    "We now use this function to get the major and minor allele count for each (bi-allelic) mutant site in our data set "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ab1e439-1062-46c8-85d7-efb21068bd5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ba_maj_min = np.apply_along_axis(getMajorMinor,1,ba_mal_counts)\n",
    "ba_maj_min"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a628b01f-bd4a-41b0-9450-4893382b9a2f",
   "metadata": {},
   "source": [
    "Now that we have the list of major and minor allele counts, we can create a pandas data frame formatted as input for sweepfinder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28edbcd6-f22f-42e6-bd6b-83696e5cae2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "biallelMutCounts = np.array([[pos,min(ac[0],ac[1]),sum(ac),1] for pos,ac in list(zip(ba_positions,ba_maj_min))])\n",
    "\n",
    "mutCountsDF =  pd.DataFrame(biallelMutCounts,columns=['position','x','n','folded'])\n",
    "mutCountsDF"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63fea4d2-1360-466a-a656-1adc0e15d7a9",
   "metadata": {},
   "source": [
    "We want to restrict our data set to sites that have full coverage across all individuals and obtain the corresonding background SFS using our mutation counts data frame.  The function <code>sf.get_sfs()</code> can do this for us. \n",
    "\n",
    "All we need is the total number of (unmasked) sites in the data set, i.e. both monomorphic and polymorphic sites, the mutCountsDF we obtained above, whether fixed diffs are included in the data set, and whether any sites in the data frame are monomorphic, i.e. have x=0.\n",
    "\n",
    "The function returns to us the updated mutation frequency data frame and the corresponding neutral SFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "060f1d80-c300-40dc-b374-fbd9ac79f091",
   "metadata": {},
   "outputs": [],
   "source": [
    "foldedBool = True\n",
    "fixedDiffBool = False\n",
    "invariantBool=False\n",
    "\n",
    "fullCovDF,neutralSFS = sf.get_sfs(mutCountsDF,L,fixedDiffBool,invariantBool)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1cc72be-b820-42ba-bfff-280fa4e15f5c",
   "metadata": {},
   "source": [
    "We can now define the remaining parameters that we need to run the scan.  \n",
    "\n",
    "Here, we will run the scan on 500 test sites on a region between 1.5e7 and 1.6e7 near the end of the chromosome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d23311ad-eb05-40fe-b258-f982570cc39c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mutFreqs = np.array(fullCovDF[\"x\"])\n",
    "mutPositions = np.array(fullCovDF[\"position\"])\n",
    "nSamples = 20\n",
    "\n",
    "\n",
    "leftLimTestSites = 1.5e7\n",
    "rightLimTestSites = 1.6e7\n",
    "testSites = np.arange(leftLimTestSites,rightLimTestSites,2_000)\n",
    "print(len(testSites))\n",
    "print(testSites[1],testSites[-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01227b7e-c06d-4a23-b175-565f4bc0aad2",
   "metadata": {},
   "source": [
    "Finally, we reload sweepfinder and run the scan "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8728bd43-a941-4a63-8d75-9646098c2b75",
   "metadata": {},
   "outputs": [],
   "source": [
    "from importlib import reload\n",
    "reload(sf)\n",
    "\n",
    "start = time.time()\n",
    "malResultsdf = sf.run_sweepfinder(testSites,nSamples,mutFreqs,mutPositions,neutralSFS,foldedBool,fixedDiffBool)\n",
    "print(f\"time it took: {time.time()-start}\")\n",
    "malResultsdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c592186-3bd1-436b-b3da-9dccd3abf3b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(malResultsdf[\"position\"],malResultsdf[\"lnCL diff\"],color='b')\n",
    "plt.show()\n",
    "\n",
    "plt.plot(malResultsdf[\"position\"],malResultsdf[\"log10 alpha\"],color='b')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f2688b8-b1ed-4bbf-965b-ac753d41d373",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 13:</b> What additional information would we need to convert the estimate of alpha to an estimate of the strength of selection sb?\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5fb1821-8b5b-4d97-88ca-43fd9554cfa2",
   "metadata": {},
   "source": [
    "Solution:\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2e15de2-a563-49d4-a9f0-28ab5038b5ab",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-success\">    \n",
    "    <b>Exercise 14:</b> use <code>allel.windowed_diversity</code> again to plot genetic diversity for 50_000 bp windows in the genome region that we tested. Does this confirm the results we obtained from sweepfinder? \n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7660290-8396-441b-a41e-c4364e9e8b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sol:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0449ed79-b6ca-4d0f-a3a8-2499ed780593",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}