{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Travelling population wave & gene surfing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interactive, IntSlider, FloatSlider, FloatLogSlider, fixed\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### function for growing population, resulting in traveling population wave\n",
    "* two alleles\n",
    "* see https://onlinelibrary.wiley.com/doi/abs/10.1111/ele.13364"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def travpopwave(numdemes = 100, N = 2000, r0 = 0.01, m = 0.01, numsteps = 4000, rs = 42):\n",
    "    np.random.seed(rs)\n",
    "    n1s = np.zeros((numdemes,numsteps))\n",
    "    n2s = np.zeros((numdemes,numsteps))\n",
    "    n1s[int(np.round(numdemes*0.45)):int(np.round(numdemes*0.55)), 0] = N//2\n",
    "    n2s[int(np.round(numdemes*0.45)):int(np.round(numdemes*0.55)), 0] = N//2\n",
    "    for step in range(1, numsteps):\n",
    "        n1x = np.copy(n1s[:, step - 1])\n",
    "        n2x = np.copy(n2s[:, step - 1])\n",
    "        # migration\n",
    "        n1x = (1 - m) * n1x + m / 2 * (np.roll(n1x, 1) + np.roll(n1x, -1))\n",
    "        n2x = (1 - m) * n2x + m / 2 * (np.roll(n2x, 1) + np.roll(n2x, -1))\n",
    "        # population growth\n",
    "        n1xnew = (1 + r0 * (1 - (n1x + n2x) / N)) * n1x\n",
    "        n2xnew = (1 + r0 * (1 - (n1x + n2x) / N)) * n2x\n",
    "        n1x = np.copy(n1xnew)\n",
    "        n2x = np.copy(n2xnew)\n",
    "        for d in range(numdemes):\n",
    "            f1 = n1x[d] / N\n",
    "            f2 = n2x[d] / N\n",
    "            if f1 + f2 > 1:\n",
    "                f1new = f1 / (f1 + f2)\n",
    "                f2new = f2 / (f1 + f2)\n",
    "                f1 = f1new\n",
    "                f2 = f2new\n",
    "            p = [f1, f2, 1 - (f1 + f2)]\n",
    "            aux = np.random.multinomial(N, p)\n",
    "            n1x[d] = aux[0]\n",
    "            n2x[d] = aux[1]\n",
    "        n1s[:, step] = np.copy(n1x)\n",
    "        n2s[:, step] = np.copy(n2x)\n",
    "    fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (20, 4))\n",
    "    X, Y = np.meshgrid(np.arange(numdemes), np.arange(numsteps))\n",
    "    im1 = ax[0].pcolor(X, Y, np.transpose(n1s + n2s), cmap = 'gray', vmin = 0, vmax = N, shading = 'auto')\n",
    "    ax[0].set_ylabel('time (steps)')\n",
    "    ax[0].set_xlabel('deme')\n",
    "    # https://joseph-long.com/writing/colorbars/\n",
    "    cb1 = fig.colorbar(im1, ax=ax[0])\n",
    "    cb1.set_label('population size in deme')\n",
    "    im2 = ax[1].pcolor(X, Y, np.transpose((n1s - n2s) / N), cmap = 'PRGn',\n",
    "                       vmin = -1.0, vmax = 1.0, shading = 'auto')\n",
    "    ax[1].set_ylabel('time (steps)')\n",
    "    ax[1].set_xlabel('deme')\n",
    "    cb2 = fig.colorbar(im2, ax=ax[1])\n",
    "    cb2.set_label('deviation from allele 1 and 2 equally abundant')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "does it work in principle?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "travpopwave()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Travelling population wave"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(travpopwave,\n",
    "            numdemes = IntSlider(value = 100, min = 20, max = 200, continuous_update = False),\n",
    "            N = IntSlider(value = 2000, min = 100, max = 10000, continuous_update = False),\n",
    "            r0 = FloatSlider(value = 0.05, min = 0.0, max = 0.1, step = 0.01, continuous_update = False),\n",
    "            m = FloatSlider(value = 0.05, min = 0.0, max = 0.1, step = 0.01, continuous_update = False),\n",
    "            numsteps = IntSlider(value = 1000, min = 100, max = 2000, continuous_update = False),\n",
    "            rs = IntSlider(value = 42, min = 0, max = 100, continuous_update = False))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}