{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Triallele Model\n",
    "\n",
    "An example of obtaining the sample triallelic frequency spectrum for a simple two epoch demography with selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "time1 = time.time()\n",
    "\n",
    "import dadi\n",
    "import dadi.Triallele\n",
    "import numpy as np, scipy, matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sig1 = 0.0 # selection coefficient for first derived allele\n",
    "sig2 = 0.0 # selection coefficient for second derived allele\n",
    "theta1 = 1.\n",
    "theta2 = 1.\n",
    "misid = 0.0 # no ancestral misidentification\n",
    "dts = [0.01, 0.025, 0.001] # time steps for integration\n",
    "grid_pts = [40,60,80] # evaluate over these grid points, then extrapolate to $\\Delta = 0$\n",
    "ns = 20\n",
    "\n",
    "T = 0.1 # equilibrium\n",
    "#nu = lambda t: 1.0\n",
    "#nuB = nuF = 1.0\n",
    "nu = 2.0\n",
    "\n",
    "fs = {}\n",
    "for dt in dts:\n",
    "    params = [nu,T]\n",
    "    fs0 = dadi.Triallele.demographics.two_epoch(params, ns, grid_pts[0], sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, misid=misid, dt=dt, folded = False)\n",
    "    fs1 = dadi.Triallele.demographics.two_epoch(params, ns, grid_pts[1], sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, misid=misid, dt=dt, folded = False)\n",
    "    fs2 = dadi.Triallele.demographics.two_epoch(params, ns, grid_pts[2], sig1=sig1, sig2=sig2, theta1=theta1, theta2=theta2, misid=misid, dt=dt, folded = False)\n",
    "    fs[dt] = dadi.Numerics.quadratic_extrap((fs0,fs1,fs2),(fs0.extrap_x,fs1.extrap_x,fs2.extrap_x))\n",
    "\n",
    "tri_fs = dadi.Numerics.quadratic_extrap((fs[dts[0]],fs[dts[1]],fs[dts[2]]),(dts[0],dts[1],dts[2]))\n",
    "\n",
    "#tri_fs = tri_fs.fold_major()\n",
    "\n",
    "time2 = time.time()\n",
    "print(\"total runtime = \" + str(time2-time1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## plot the triallele spectrum\n",
    "import matplotlib.pylab as plt\n",
    "dadi.Triallele.plotting.plot_single_trispectrum(tri_fs, folded=True, colorbar=True)\n",
    "plt.show()"
   ]
  }
 ],
 "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
