{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Moran model with two alleles"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "function for Moran model (two alleles, mutations, with selection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def moran(N = 100, mu12 = 0.0, mu21 = 0.0, relw = 1.0, fstart = 0.5, numsteps = 1000, numtrajs = 10):\n",
    "    w1 = relw * 1.0\n",
    "    w2 = 1.0\n",
    "    fs = np.zeros((numtrajs,numsteps))\n",
    "    fig, ax = plt.subplots(nrows = 1, ncols = 3, figsize = (20, 4))\n",
    "    for traj in range(numtrajs):\n",
    "        i = int(np.round(fstart*N))\n",
    "        for step in range(numsteps):\n",
    "            f = i/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 = [i/N, (N-i)/N])[0]\n",
    "            if (allelebirth == 1 and np.random.rand() >= mu12) or (allelebirth == 2 and np.random.rand() < mu21):\n",
    "                i = i + 1\n",
    "            if alleledeath == 1:\n",
    "                i = i - 1\n",
    "            fs[traj,step] = i/N\n",
    "        ax[0].plot(np.arange(numsteps) / N, fs[traj,:])\n",
    "    ax[0].set_ylim([0.0, 1.0])\n",
    "    ax[0].set_xlabel('time (generations)')\n",
    "    ax[0].set_ylabel('frequency of allele 1')\n",
    "    ax[1].plot(np.arange(numsteps) / N, np.mean(fs, axis = 0), color = 'black')\n",
    "    ax[1].set_ylim([0.0, 1.0])\n",
    "    ax[1].set_xlabel('time (generations)')\n",
    "    ax[1].set_ylabel('mean frequency of allele 1')\n",
    "    ax[2].plot(np.arange(numsteps) / N, np.var(fs, axis = 0), color = 'black')\n",
    "    ax[2].set_ylim(ymin = 0)\n",
    "    ax[2].set_xlabel('time (generations)')\n",
    "    ax[2].set_ylabel('variance of frequency of allele 1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "does it work in principle?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "moran()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moran model: neutral, no mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(moran,\n",
    "            N = IntSlider(value = 100, min = 2, max = 200, continuous_update = False),\n",
    "            mu12 = fixed(0.0), mu21 = fixed(0.0), relw = fixed(1.0),\n",
    "            fstart = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            numsteps = IntSlider(value = 1000, min = 10, max = 10000, continuous_update = False),\n",
    "            numtrajs = IntSlider(valud = 1, min = 1, max = 20, continuous_update = False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moran model: neutral, with mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(moran,\n",
    "            N = IntSlider(value = 100, min = 2, max = 200, continuous_update = False),\n",
    "            mu12 = FloatSlider(value = 0.0, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            mu21 = FloatSlider(value = 0.0, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            relw = fixed(1.0),\n",
    "            fstart = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            numsteps = IntSlider(value = 1000, min = 10, max = 10000, continuous_update = False),\n",
    "            numtrajs = IntSlider(valud = 1, min = 1, max = 20, continuous_update = False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Moran model: with selection, no mutations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactive(moran,\n",
    "            N = IntSlider(value = 100, min = 2, max = 200, continuous_update = False),\n",
    "            mu12 = fixed(0.0), mu21 = fixed(0.0),\n",
    "            relw = FloatLogSlider(value = 1.0, base = 10, min = -1, max = 1, continuous_update = False),\n",
    "            fstart = FloatSlider(value = 0.5, min = 0.0, max = 1.0, continuous_update = False),\n",
    "            numsteps = IntSlider(value = 1000, min = 10, max = 10000, continuous_update = False),\n",
    "            numtrajs = IntSlider(valud = 1, min = 1, max = 20, 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
}