steppingstone.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Stepping stone model with two alleles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipywidgets import interactive, IntSlider, FloatSlider, FloatLogSlider, Dropdown, fixed\n",
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "import numba"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### function for Moran model with several demes\n",
    "* two alleles\n",
    "* mutation\n",
    "* selection\n",
    "* see https://journals.aps.org/prl/supplemental/10.1103/PhysRevLett.107.088103/supplemental.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# auxiliary function to facilitate periodic boundary conditions\n",
    "@numba.njit\n",
    "def pbc(i, ls):\n",
    "    if i == ls:\n",
    "        return 0\n",
    "    elif i == -1:\n",
    "        return ls-1\n",
    "    else:\n",
    "        return i\n",
    "\n",
    "def steppingstone(numdemes = 20, N = 50, m = 0.5, relw = 1.0, mu12 = 0.0, mu21 = 0.0,\n",
    "                  initialcond = 'uniform', numsteps = 1000):\n",
    "    w1 = relw * 1.0\n",
    "    w2 = 1.0\n",
    "    if initialcond == 'uniform':\n",
    "        n1s = N//2 * np.ones((numdemes,numsteps)) # number individuals of allele 1\n",
    "    elif initialcond == 'nonuniform':\n",
    "        n1s = N * np.zeros((numdemes,numsteps)) # number individuals of allele 1\n",
    "        n1s[int(np.round(numdemes*0.4)):int(np.round(numdemes*0.6)), 0] = int(np.round(N))\n",
    "    for step in range(1, numsteps):\n",
    "        n1x = np.copy(n1s[:, step - 1]) # number of individuals of allele 1, auxiliary array\n",
    "        # migration (in random order)\n",
    "        for d in np.random.permutation(numdemes):\n",
    "            ld = pbc(d - 1, numdemes)\n",
    "            rd = pbc(d + 1, numdemes)\n",
    "            p = m / 2 / N**2 * np.array([1.0, 1.0, 1.0, 1.0])\n",
    "            # 4 cases, see article referenced above\n",
    "            p[0] = p[0] * n1x[d] * (N - n1x[ld])\n",
    "            p[1] = p[1] * n1x[ld] * (N - n1x[d])\n",
    "            p[2] = p[2] * n1x[d] * (N - n1x[rd])\n",
    "            p[3] = p[3] * n1x[rd] * (N - n1x[d])\n",
    "            # choose event\n",
    "            raux = np.random.rand()\n",
    "            if raux < np.cumsum(p)[0]:\n",
    "                n1x[ld] += 1\n",
    "                n1x[d] -= 1\n",
    "            elif raux < np.cumsum(p)[1]:\n",
    "                n1x[ld] -= 1\n",
    "                n1x[d] += 1\n",
    "            elif raux < np.cumsum(p)[2]:\n",
    "                n1x[d] -= 1\n",
    "                n1x[rd] += 1\n",
    "            elif raux < np.cumsum(p)[3]:\n",
    "                n1x[d] += 1\n",
    "                n1x[rd] -= 1\n",
    "        # birth and death\n",
    "        for d in range(numdemes):\n",
    "            f = n1x[d] / N\n",
    "            allelebirth = np.random.choice([1, 2], 1, p = [w1*f/(w1*f+w2*(1-f)), w2*(1-f)/(w1*f+w2*(1-f))])[0]\n",
    "            alleledeath = np.random.choice([1, 2], 1, p = [f, 1-f])[0]\n",
    "            if (allelebirth == 1 and np.random.rand() >= mu12) or (allelebirth == 2 and np.random.rand() < mu21):\n",
    "                n1x[d] += 1\n",
    "            if alleledeath == 1:\n",
    "                n1x[d] -= 1\n",
    "        n1s[:, step] = np.copy(n1x)\n",
    "    fig, ax = plt.subplots(nrows = 1, ncols = 2, figsize = (20, 4))\n",
    "    X, Y = np.meshgrid(np.arange(numdemes), np.arange(numsteps))\n",
    "    im = ax[0].pcolor(X, Y, np.transpose(n1s) / N, cmap = 'PRGn', vmin = 0.0, vmax = 1.0, 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(im, ax=ax[0])\n",
    "    cb1.set_label('frequency allele 1')\n",
    "    cf = ax[1].contour(X, Y, np.transpose(n1s) / N)\n",
    "    cb2 = fig.colorbar(cf, ax=ax[1])\n",
    "    cb2.set_label('frequency allele 1')\n",
    "    ax[1].set_ylabel('time (steps)')\n",
    "    ax[1].set_xlabel('deme')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "does it work in principle?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "steppingstone()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stepping stone model: neutral"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(steppingstone,\n",
    "            numdemes = IntSlider(value = 20, min = 10, max = 100, continuous_update = False),\n",
    "            N = IntSlider(value = 50, min = 10, max = 200, continuous_update = False),\n",
    "            m = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            relw = fixed(1.0),\n",
    "            mu12 = fixed(0.0), mu21 = fixed(0.0),\n",
    "            initialcond = Dropdown(options=['uniform', 'nonuniform'], value = 'uniform', description='initial cond'),\n",
    "            numsteps = IntSlider(value = 100, min = 10, max = 2000, continuous_update = False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stepping stone model: with selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(steppingstone,\n",
    "            numdemes = IntSlider(value = 20, min = 10, max = 100, continuous_update = False),\n",
    "            N = IntSlider(value = 50, min = 10, max = 200, continuous_update = False),\n",
    "            m = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            relw = FloatLogSlider(value = 1.0, base = 10, min = -1, max = 1, continuous_update = False),\n",
    "            mu12 = fixed(0.0), mu21 = fixed(0.0),\n",
    "            initialcond = Dropdown(options=['uniform', 'nonuniform'], value = 'nonuniform', description='initial cond', disabled=False),\n",
    "            numsteps = IntSlider(value = 100, min = 10, max = 2000, 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
}