{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two Locus Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dadi\n",
    "import dadi.TwoLocus\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loading and manipulating two-locus spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 1. We'll first import a saved two-locus spectrum\n",
    "# this was simulated/sampled with rho=1, nu=2.0, T=0.1\n",
    "fs = dadi.TwoLocus.TLSpectrum.from_file('example_observed_fs.fs')\n",
    "# get the sample size of the spectrum\n",
    "ns = fs.sample_size\n",
    "# fold the spectrum (ancestral state unknown)\n",
    "fs_folded = fs.fold()\n",
    "# project to a smaller sample size\n",
    "fs_project = dadi.TwoLocus.numerics.project(fs, ns//2)\n",
    "# compute mean r squared for observed variable sites\n",
    "r2 = fs.mean_r2()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating expected spectrum using demographic history functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 2. Create a two locus spectrum under a demographic history model\n",
    "# two epoch model, with a recent size change of 2.0 for 0.1 (2*N_e) time units\n",
    "nu,T = 2.0, 0.1\n",
    "# sample size\n",
    "ns = 20\n",
    "# population size-scaled recombination rate between loci\n",
    "rho = 1.0\n",
    "# set grid points and time steps to integrate over\n",
    "# we want the number of grid points to be larger than the sample size (here 20)\n",
    "gridpts = [40,50,60]\n",
    "# we also extrapolate over dt values, which improves accuracy\n",
    "dts = [0.005,0.0025,0.001]\n",
    "\n",
    "# note that the first time computing the equilibrium spectrum for a given \n",
    "# dt and pts can take a while, which is why we cache those spectra\n",
    "\n",
    "# Here we copy the precached spectra we have provided to the central cache\n",
    "# location (by default, ~/.dadi/TwoLocus_cache). In general, a user won't need\n",
    "# to do this. But if you are moving between machines, you may want to copy your\n",
    "# cache over as well, to save substantial calculation time.\n",
    "import glob, shutil\n",
    "for fname in glob.glob('TwoLocus_cache/*.npz'):\n",
    "    shutil.copy(fname, dadi.TwoLocus.demographics.cache_path)\n",
    "\n",
    "spectra = {}\n",
    "for dt in dts:\n",
    "    spectra.setdefault(dt,{})\n",
    "    for pts in gridpts:\n",
    "        print('computing for dt, pts = {0}, {1}'.format(dt,pts))\n",
    "        spectra[dt][pts] = dadi.TwoLocus.demographics.two_epoch((nu,T), pts, ns, rho=rho, dt=dt)\n",
    "\n",
    "# extrap_dt_pts takes a dictionary in the form spectra[dts][pts] and\n",
    "# extrapolates the calculated spectra over dt and pts to output the expected\n",
    "# frequency spectrum\n",
    "F = dadi.TwoLocus.numerics.extrap_dt_pts(spectra)\n",
    "\n",
    "# save this spectrum to working directory\n",
    "F.to_file('./two_epoch_example.fs')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference demographic parameters from example spectrum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## 3. Fit a two-epoch model to the spectrum that we loaded in 1.\n",
    "fs = dadi.TwoLocus.TLSpectrum.from_file('example_observed_fs.fs')\n",
    "# The inference/optimize functions takes a list of two-locus frequency spectra,\n",
    "# and a list of the rho values for each of the spectra in the list, in the same\n",
    "# order as the spectra in the list. Here, we only are optimizing over a single\n",
    "# frequency spectrum for rho = 1.\n",
    "data_list = [fs]\n",
    "rho_list = [1.0]\n",
    "# we know that the sample spectrum was generated under a demography with \n",
    "# nu = 2.0, T = 0.1, so for illustration purposes we pick initial values close\n",
    "# to the true values\n",
    "initial_guess = dadi.Misc.perturb_params([2.0,0.1], fold=0.2)\n",
    "\n",
    "gridpts = [40,50,60]\n",
    "dts = [0.005,0.0025,0.001]\n",
    "\n",
    "# do_calculation computes the expected frequency spectrum for a given sample\n",
    "# size and rho\n",
    "def do_calculation(arguments):\n",
    "    params, pts, ns, rho, dt = arguments\n",
    "    # create spectrum for given param set, rho, etc\n",
    "    temps = {}\n",
    "    for dt in dts:\n",
    "        temps.setdefault(dt,{})\n",
    "        for numpts in pts:\n",
    "            temps[dt][numpts] = dadi.TwoLocus.demographics.two_epoch(params,numpts,ns,rho=rho,dt=dt)\n",
    "    return dadi.TwoLocus.TLSpectrum(dadi.TwoLocus.numerics.extrap_dt_pts(temps))\n",
    "\n",
    "# optimize functions needs a function to output a list of spectra over the rho-values in the rhos list\n",
    "def model_func(params, ns, pts, dts, rhos=[0]):\n",
    "    inputs = [(params,pts,ns,rho,dts) for rho in rhos]\n",
    "    output = [do_calculation(input) for input in inputs]\n",
    "    return output\n",
    "\n",
    "opt_params = dadi.TwoLocus.inference.optimize_log_fmin(initial_guess, data_list, model_func, gridpts, dts, rhos=rho_list, verbose=1)"
   ]
  }
 ],
 "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
}
