{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DFE inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle, random\n",
    "import numpy as np\n",
    "import dadi\n",
    "import dadi.DFE as DFE\n",
    "import matplotlib.pyplot as plt\n",
    "from dadi.DFE import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### One-population example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For multiprocessing to work on Windows, all script code must be wrapped\n",
    "# in this block. If you're not on Windows, feel free to remove this if statement.\n",
    "if __name__ == '__main__':\n",
    "    # Set demographic parameters and theta. This is usually inferred from\n",
    "    # synonymous sites. In this case, we'll be using a two-epoch model.\n",
    "    demog_params = [2, 0.05]\n",
    "    theta_ns = 4000.\n",
    "    ns = [250]\n",
    "\n",
    "    # Integrate over a range of gammas\n",
    "    pts_l = [600, 800, 1000]\n",
    "    spectra = DFE.Cache1D(demog_params, ns, DFE.DemogSelModels.two_epoch, pts_l=pts_l, \n",
    "                          gamma_bounds=(1e-5, 500), gamma_pts=100, verbose=True,\n",
    "                          mp=True)\n",
    "    # The spectra can be pickled for usage later. This is especially convenient\n",
    "    # if the process of generating the spectra takes a long time.\n",
    "    pickle.dump(spectra, open('example_spectra.bpkl','wb'))\n",
    "    # To load them, use this code\n",
    "    #spectra = pickle.load(open('example_spectra.bpkl','rb'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#load sample data\n",
    "data = dadi.Spectrum.from_file('example.fs')\n",
    "\n",
    "# Fit a DFE to the data\n",
    "# Initial guess and bounds\n",
    "sel_params = [0.2, 1000.]\n",
    "lower_bound, upper_bound = [1e-3, 1e-2], [1, 50000.]\n",
    "p0 = dadi.Misc.perturb_params(sel_params, lower_bound=lower_bound,\n",
    "                              upper_bound=upper_bound)\n",
    "popt = dadi.Inference.optimize_log(p0, data, spectra.integrate, pts=None,\n",
    "                                   func_args=[DFE.PDFs.gamma, theta_ns],\n",
    "                                   lower_bound=lower_bound, upper_bound=upper_bound, \n",
    "                                   verbose=len(sel_params), maxiter=10, multinom=False)\n",
    "\n",
    "# Get expected SFS for MLE\n",
    "model_sfs = spectra.integrate(popt, None, DFE.PDFs.gamma, theta_ns, None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# One possible characterization of the neutral+gamma DFE\n",
    "# Written using numpy tricks to work with both scalar and array arguments\n",
    "def neugamma(xx, params):\n",
    "    pneu, alpha, beta = params\n",
    "    # Convert xx to an array\n",
    "    xx = np.atleast_1d(xx)\n",
    "    out = (1-pneu)*DFE.PDFs.gamma(xx, (alpha, beta))\n",
    "    # Assume gamma < 1e-4 is essentially neutral\n",
    "    out[np.logical_and(0 <= xx, xx < 1e-4)] += pneu/1e-4\n",
    "    # Reduce xx back to scalar if it's possible\n",
    "    return np.squeeze(out)\n",
    "\n",
    "sel_params = [0.2, 0.2, 1000.]\n",
    "lower_bound, upper_bound = [1e-3, 1e-3, 1e-2], [1, 1, 50000.]\n",
    "p0 = dadi.Misc.perturb_params(sel_params, lower_bound=lower_bound,\n",
    "                              upper_bound=upper_bound)\n",
    "popt = dadi.Inference.optimize_log(p0, data, spectra.integrate, pts=None,\n",
    "                                   func_args=[neugamma, theta_ns],\n",
    "                                   lower_bound=lower_bound, upper_bound=upper_bound, \n",
    "                                   verbose=len(sel_params),\n",
    "                                   maxiter=10, multinom=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Modeling ancestral state misidentification, using dadi's built-in function to \n",
    "# wrap fitdadi's integrate method.\n",
    "#\n",
    "p_misid = 0.05\n",
    "data = dadi.Numerics.apply_anc_state_misid(data, p_misid)\n",
    "misid_func = dadi.Numerics.make_anc_state_misid_func(spectra.integrate)\n",
    "sel_params = [0.2, 1000., 0.2]\n",
    "lower_bound, upper_bound = [1e-3, 1e-2, 0], [1, 50000., 1]\n",
    "p0 = dadi.Misc.perturb_params(sel_params, lower_bound=lower_bound,\n",
    "                              upper_bound=upper_bound)\n",
    "popt = dadi.Inference.optimize_log(p0, data, misid_func, pts=None,\n",
    "                                   func_args=[DFE.PDFs.gamma, theta_ns],\n",
    "                                   lower_bound=lower_bound, upper_bound=upper_bound,\n",
    "                                   verbose=len(sel_params), maxiter=10,\n",
    "                                   multinom=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Including a point mass of positive selection\n",
    "#\n",
    "data = dadi.Spectrum.from_file('example.fs')\n",
    "ppos = 0.1\n",
    "sel_data = theta_ns*DFE.DemogSelModels.two_epoch(tuple(demog_params) + (5,), ns, pts_l[-1])\n",
    "data_pos = (1-ppos)*data + ppos*sel_data\n",
    "\n",
    "sel_params = [0.2, 1000., 0.2, 2]\n",
    "lower_bound, upper_bound = [1e-3, 1e-2, 0, 0], [1, 50000., 1, 50]\n",
    "p0 = dadi.Misc.perturb_params(sel_params, lower_bound=lower_bound,\n",
    "                              upper_bound=upper_bound)\n",
    "popt = dadi.Inference.optimize_log(p0, data_pos, spectra.integrate_point_pos, pts=None,\n",
    "                                   func_args=[DFE.PDFs.gamma, theta_ns, DFE.DemogSelModels.two_epoch], \n",
    "                                   lower_bound=lower_bound, upper_bound=upper_bound, \n",
    "                                   verbose=len(sel_params), maxiter=10, multinom=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Multiple point masses of positive selection\n",
    "#\n",
    "# Parameters are mu, sigma, ppops1, gammapos1, ppos2, gammapos2\n",
    "sel_params = [3,2,0.1,2,0.3,6]\n",
    "input_fs = spectra.integrate_point_pos(sel_params,None,DFE.PDFs.lognormal,theta_ns,\n",
    "                                       DFE.DemogSelModels.two_epoch, 2)\n",
    "data = input_fs.sample()\n",
    "lower_bound, upper_bound = [-1,0.1,0,0,0,0], [5,5,1,10,1,20]\n",
    "p0 = dadi.Misc.perturb_params(sel_params, lower_bound=lower_bound,\n",
    "                              upper_bound=upper_bound)\n",
    "\n",
    "def ieq_constraint(p,*args):\n",
    "    # Our constraint is that ppop1+ppos2 must be less than 1.\n",
    "    return [1-(p[2]-p[4])]\n",
    "\n",
    "popt = dadi.Inference.optimize_cons(p0, data, spectra.integrate_point_pos, pts=None,\n",
    "                                    func_args=[DFE.PDFs.lognormal, theta_ns,\n",
    "                                               DFE.DemogSelModels.two_epoch, 2],\n",
    "                                    lower_bound=lower_bound, upper_bound=upper_bound,\n",
    "                                    ieq_constraint=ieq_constraint,\n",
    "                                    # Fix gammapos1\n",
    "                                    fixed_params=[None,None,None,2,None,None],\n",
    "                                    verbose=len(sel_params),\n",
    "                                    maxiter=10, multinom=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Two-population example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Seed random number generator, so example is reproducible\n",
    "np.random.seed(1398238)\n",
    "\n",
    "#\n",
    "# Plotting a joint DFE\n",
    "#\n",
    "sel_dist = PDFs.biv_lognormal\n",
    "params = [0.5,-0.5,0.5,1,-0.8]\n",
    "gammax, gammay = -np.logspace(-2, 1, 20), -np.logspace(-1, 2, 30)\n",
    "\n",
    "fig = plt.figure(137, figsize=(4,3), dpi=150)\n",
    "fig.clear()\n",
    "ax = fig.add_subplot(1,1,1)\n",
    "Plotting.plot_biv_dfe(gammax, gammay, sel_dist, params, logweight=True, ax=ax)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# With positive selection\n",
    "params = [0.5,-0.5,0.5,1,0.0,0.3,3,0.3,4]\n",
    "fig = Plotting.plot_biv_point_pos_dfe(gammax, gammay, sel_dist, params,\n",
    "                                      fignum=23, rho=params[4])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For multiprocessing to work on Windows, all script code must be wrapped\n",
    "# in this block. If you're not on Windows, feel free to remove this if statement.\n",
    "if __name__ == '__main__':\n",
    "    #\n",
    "    # Full test of optimization machinery. \n",
    "    # Considering only a narrow range of gammas, so integration is faster.\n",
    "    #\n",
    "    demo_params = [0.5,2,0.5,0.1,0,0]\n",
    "    ns = [8, 12]\n",
    "    pts_l = [60, 80, 100]\n",
    "    func_ex = DemogSelModels.IM\n",
    "    # Check whether we already have a chached set of 2d spectra. If not\n",
    "    # generate them.\n",
    "    try:\n",
    "        s2 = pickle.load(open('test.spectra2d.bpkl', 'rb'))\n",
    "    except IOError:\n",
    "        s2 = Cache2D(demo_params, ns, func_ex, pts=pts_l, gamma_pts=100,\n",
    "                     gamma_bounds=(1e-2, 10), verbose=True, mp=True,\n",
    "                     additional_gammas=[1.2, 4.3])\n",
    "        # Save spectra2d object\n",
    "        fid = open('test.spectra2d.bpkl', 'wb')\n",
    "        pickle.dump(s2, fid, protocol=2)\n",
    "        fid.close()\n",
    "\n",
    "        ## Cache generation can be very computationally expensive, even\n",
    "        ## with multiprocessing. If you need to split the work across multiple\n",
    "        ## compute nodes, use the split_jobs and this_job_id arguments, then\n",
    "        ## merge the partial caches.\n",
    "        ## In this example, the 3 partial caches could be generate on independent\n",
    "        ## nodes, then s2a,s2b,s2c would be saved to separate files, then loaded\n",
    "        ## and combined later.\n",
    "        #s2a = Cache2D(demo_params, ns, func_ex, pts=pts_l, gamma_pts=100,\n",
    "        #        gamma_bounds=(1e-2, 10), verbose=True, mp=True,\n",
    "        #        additional_gammas=[1.2, 4.3], split_jobs=3, this_job_id=0)\n",
    "        #s2b = Cache2D(demo_params, ns, func_ex, pts=pts_l, gamma_pts=100,\n",
    "        #        gamma_bounds=(1e-2, 10), verbose=True, mp=True,\n",
    "        #        additional_gammas=[1.2, 4.3], split_jobs=3, this_job_id=1)\n",
    "        #s2c = Cache2D(demo_params, ns, func_ex, pts=pts_l, gamma_pts=100,\n",
    "        #        gamma_bounds=(1e-2, 10), verbose=True, mp=True,\n",
    "        #        additional_gammas=[1.2, 4.3], split_jobs=3, this_job_id=2)\n",
    "        #s2 = Cache2D.merge([s2a, s2b, s2c]) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate test data set to fit\n",
    "input_params, theta = [0.5,0.5,-0.8], 1e5\n",
    "sel_dist = PDFs.biv_lognormal\n",
    "# Expected sfs\n",
    "target = s2.integrate(input_params, None, sel_dist, theta, None)\n",
    "# Get data with Poisson variance around expectation\n",
    "data = target.sample()\n",
    "\n",
    "# Parameters are mean, variance, and correlation coefficient\n",
    "p0 = [0,1.,0.8]\n",
    "popt = dadi.Inference.optimize(p0, data, s2.integrate, pts=None,\n",
    "                               func_args=[sel_dist, theta],\n",
    "                               lower_bound=[None,0,-1],\n",
    "                               upper_bound=[None,None,1],\n",
    "                               verbose=30, multinom=False)\n",
    "print('Input parameters: {0}'.format(input_params))\n",
    "print('Optimized parameters: {0}'.format(popt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot inferred DFE. Note that this will render slowly, because grid of\n",
    "# gammas is fairly dense.\n",
    "fig = plt.figure(231, figsize=(4,3), dpi=150)\n",
    "fig.clear()\n",
    "ax = fig.add_subplot(1,1,1)\n",
    "Plotting.plot_biv_dfe(s2.gammas, s2.gammas, sel_dist, popt, ax=ax)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Test point mass of positive selection. To do so, we test against\n",
    "# the single-population case using very high correlation.\n",
    "#\n",
    "params = [-0.5,0.5,0.99, 0.1, 4.3]\n",
    "fs_biv = s2.integrate_symmetric_point_pos(params, None, PDFs.biv_lognormal, theta,\n",
    "                                          pts=None)\n",
    "\n",
    "func_single_ex = DemogSelModels.IM_single_gamma\n",
    "try:\n",
    "    s1 = pickle.load(open('test.spectra1d.bpkl', 'rb'))\n",
    "except IOError:\n",
    "    s1 = Cache1D(demo_params, ns, func_single_ex, pts_l=pts_l,\n",
    "                 gamma_pts=100, gamma_bounds=(1e-2, 10), mp=True,\n",
    "                 additional_gammas = [1.2, 4.3], verbose=False)\n",
    "    fid = open('test.spectra1d.bpkl', 'wb')\n",
    "    pickle.dump(s1, fid, protocol=2)\n",
    "    fid.close()\n",
    "\n",
    "fs1 = s1.integrate_point_pos([-0.5,0.5,0.1,4.3], None, PDFs.lognormal,\n",
    "                             1e5, func_single_ex)\n",
    "\n",
    "fig = dadi.Plotting.pylab.figure(figsize=(8,6))\n",
    "fig.clear()\n",
    "dadi.Plotting.plot_2d_comp_Poisson(fs1, fs_biv, show=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Test optimization of point mass positive selection.\n",
    "#\n",
    "\n",
    "# Generate test data set to fit\n",
    "# This is a symmetric case, with mu1=mu2=0.5, sigma1=sigma2=0.3, rho=-0.5,\n",
    "# ppos1=ppos2=0.2, gammapos1=gammapos2=1.2.\n",
    "input_params, theta = [0.5,0.3,-0.5,0.2,1.2], 1e5\n",
    "# Expected sfs\n",
    "target = s2.integrate_symmetric_point_pos(input_params, None, sel_dist, theta,\n",
    "                                          pts=None)\n",
    "# Get data with Poisson variance around expectation\n",
    "data = target.sample()\n",
    "\n",
    "# We'll fit using our special-case symmetric function. The last\n",
    "# two arguments are ppos and gammapos. The first three are thus for the\n",
    "# lognormal distribution. Note that our lognormal distribution assumes\n",
    "# symmetry if the length of the arguments is only three. If we wanted\n",
    "# asymmetric lognormal, we would pass in a p0 of total length 7.\n",
    "p0 = [0.3,0.3,0.1,0.2,1.2]\n",
    "popt = dadi.Inference.optimize(p0, data,\n",
    "                               s2.integrate_symmetric_point_pos,\n",
    "                               pts=None, func_args=[sel_dist, theta],\n",
    "                               # Note that mu in principle has no lower or\n",
    "                               # upper bound, sigma has only a lower bound\n",
    "                               # of 0, ppos is bounded between 0 and 1, and\n",
    "                               # gamma pos is bounded from below by 0.\n",
    "                               lower_bound=[-1,0.1,-1,0,0],\n",
    "                               upper_bound=[1,1,1,1,None],\n",
    "                               # We fix the gammapos to be 1.2, because\n",
    "                               # we can't do this integration effectively\n",
    "                               # if gammapos is allowed to vary.\n",
    "                               fixed_params=[None,None,None,None,1.2],\n",
    "                               verbose=30, multinom=False)\n",
    "print('Symmetric test fit')\n",
    "print('  Input parameters: {0}'.format(input_params))\n",
    "print('  Optimized parameters: {0}'.format(popt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Mixture model\n",
    "#\n",
    "\n",
    "# Now a mixture model, which adds together a 2D distribution and a \n",
    "# perfectly correlated 1D distribution.\n",
    "# Input parameters here a mu, sigma, rho for 2D (fixed to zero), \n",
    "#   proportion positive, gamma positive, proportion 2D, \n",
    "input_params, theta = [0.5,0.3,0,0.2,1.2,0.2], 1e5\n",
    "# Expected sfs\n",
    "target = mixture_symmetric_point_pos(input_params,None,s1,s2,PDFs.lognormal,\n",
    "                                     PDFs.biv_lognormal, theta)\n",
    "p0 = [0.3,0.3,0,0.2,1.2,0.3]\n",
    "popt = dadi.Inference.optimize(p0, data, mixture_symmetric_point_pos, pts=None, \n",
    "            func_args=[s1, s2, PDFs.lognormal,\n",
    "            PDFs.biv_lognormal, theta],\n",
    "            lower_bound=[None, 0.1,-1,0,None, 0],\n",
    "            upper_bound=[None,None, 1,1,None, 1],\n",
    "            # We fix both the rho assumed for the 2D distribution,\n",
    "            # and the assumed value of positive selection.\n",
    "            fixed_params=[None,None,0,None,1.2,None],\n",
    "            verbose=30, multinom=False, maxiter=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Test Godambe code for estimating uncertainties\n",
    "#\n",
    "input_params = [0.3,0.3,0.1,0.2,1.2]\n",
    "# Generate data in segments for future bootstrapping\n",
    "fs0 = s2.integrate_symmetric_point_pos(input_params, None, sel_dist,\n",
    "                                       theta/100., pts=None)\n",
    "# The multiplication of fs0 is to create a range of data size among\n",
    "# bootstrap chunks, which creates a range of thetas in the bootstrap\n",
    "# data sets.\n",
    "data_pieces = [(fs0*(0.5 + (1.5-0.5)/99*ii)).sample() for ii in range(100)]\n",
    "# Add up those segments to get our data spectrum\n",
    "data = dadi.Spectrum(np.sum(data_pieces, axis=0))\n",
    "# Do the optimization\n",
    "popt = dadi.Inference.optimize([0.2,0.2,0.15,0.3,1.2], data,\n",
    "                                s2.integrate_symmetric_point_pos,\n",
    "                                pts=None, func_args=[sel_dist, theta],\n",
    "                                lower_bound=[-1,0.1,-1,0,0],\n",
    "                                upper_bound=[1,1,1,1,None],\n",
    "                                fixed_params=[None,None,None,None,1.2],\n",
    "                                verbose=30, multinom=False)\n",
    "\n",
    "print('Symmetric test fit')\n",
    "print('  Input parameters: {0}'.format(input_params))\n",
    "print('  Optimized parameters: {0}'.format(popt))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate bootstraps for Godambe\n",
    "all_boot = []\n",
    "for boot_ii in range(100):\n",
    "    # Each bootstrap is made by sampling, with replacement, from our data\n",
    "    # pieces\n",
    "    this_pieces = [random.choice(data_pieces) for _ in range(100)]\n",
    "    all_boot.append(dadi.Spectrum(np.sum(this_pieces, axis=0)))\n",
    "\n",
    "# The Godambe methods only accept a basic dadi function that takes in\n",
    "# parameters, ns, and pts. Moreover, we can't allow gammapos to vary.\n",
    "# (We don't have variations around those cached. We could work around this\n",
    "# by pre-caching the necessary values, but it would be inelegant.) So\n",
    "# we create this simple function that holds gammapos constant and removes\n",
    "# it from the argument list.\n",
    "def temp_func(pin, ns, pts):\n",
    "    # Add in gammapos parameter\n",
    "    params = np.concatenate([pin, [1.2]])\n",
    "    return s2.integrate_symmetric_point_pos(params, None, sel_dist,\n",
    "                                               theta, pts=None)\n",
    "\n",
    "# Run the uncertainty analysis. Note that each bootstrap data set\n",
    "# needs a different assumed theta. We estimate theta for each bootstrap\n",
    "# data set simply by scaling theta from the orignal data.\n",
    "import dadi.Godambe\n",
    "boot_theta_adjusts = [b.sum()/data.sum() for b in all_boot]\n",
    "uncerts_adj = dadi.Godambe.GIM_uncert(temp_func, [], all_boot, popt[:-1],\n",
    "                                      data, multinom=False,\n",
    "                                      boot_theta_adjusts=boot_theta_adjusts)\n",
    "print('Godambe uncertainty test')\n",
    "print('  Input parameters: {0}'.format(input_params))\n",
    "print('  Optimized parameters: {0}'.format(popt))\n",
    "print('  Estimated 95% uncerts (theta adj): {0}'.format(1.96*uncerts_adj))\n",
    "\n",
    "# Ensure plots show up on screen.\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
}
