{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Demography inference for YRI and CEU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Numpy is the numerical library dadi is built upon\n",
    "import numpy as np\n",
    "import dadi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define custom demography models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define two custom demography models\n",
    "def prior_onegrow_mig(params, ns, pts):\n",
    "    \"\"\"\n",
    "    Model with growth, split, bottleneck in pop2, exp recovery, migration\n",
    "\n",
    "    params list is\n",
    "    nu1F: The ancestral population size after growth. (Its initial size is\n",
    "          defined to be 1.)\n",
    "    nu2B: The bottleneck size for pop2\n",
    "    nu2F: The final size for pop2\n",
    "    m: The scaled migration rate\n",
    "    Tp: The scaled time between ancestral population growth and the split.\n",
    "    T: The time between the split and present\n",
    "\n",
    "    ns = (n1,n2): Size of fs to generate.\n",
    "    pts: Number of points to use in grid for evaluation.\n",
    "    \"\"\"\n",
    "    nu1F, nu2B, nu2F, m, Tp, T = params\n",
    "    n1,n2 = ns\n",
    "    # Define the grid we'll use\n",
    "    xx = yy = dadi.Numerics.default_grid(pts)\n",
    "\n",
    "    # phi for the equilibrium ancestral population\n",
    "    phi = dadi.PhiManip.phi_1D(xx)\n",
    "    # Now do the population growth event.\n",
    "    phi = dadi.Integration.one_pop(phi, xx, Tp, nu=nu1F)\n",
    "\n",
    "    # The divergence\n",
    "    phi = dadi.PhiManip.phi_1D_to_2D(xx, phi)\n",
    "    # We need to define a function to describe the non-constant population 2\n",
    "    # size. lambda is a convenient way to do so.\n",
    "    nu2_func = lambda t: nu2B*(nu2F/nu2B)**(t/T)\n",
    "    phi = dadi.Integration.two_pops(phi, xx, T, nu1=nu1F, nu2=nu2_func, \n",
    "                                    m12=m, m21=m)\n",
    "\n",
    "    # Finally, calculate the spectrum.\n",
    "    sfs = dadi.Spectrum.from_phi(phi, (n1,n2), (xx,yy))\n",
    "    return sfs\n",
    "\n",
    "def prior_onegrow_nomig(params, ns, pts):\n",
    "    \"\"\"\n",
    "    Model with growth, split, bottleneck in pop2, exp recovery, no migration\n",
    "\n",
    "    params\n",
    "    nu1F: The ancestral population size after growth. (Its initial size is\n",
    "          defined to be 1.)\n",
    "    nu2B: The bottleneck size for pop2\n",
    "    nu2F: The final size for pop2\n",
    "    Tp: The scaled time between ancestral population growth and the split.\n",
    "    T: The time between the split and present\n",
    "\n",
    "    ns = (n1,n2): Size of fs to generate.\n",
    "    pts: Number of points to use in grid for evaluation.\n",
    "    \"\"\"\n",
    "    nu1F, nu2B, nu2F, Tp, T = params\n",
    "    return prior_onegrow_mig((nu1F, nu2B, nu2F, 0, Tp, T), ns, pts)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Demography inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the data\n",
    "data = dadi.Spectrum.from_file('YRI_CEU.fs')\n",
    "ns = data.sample_sizes\n",
    "\n",
    "# These are the grid point settings will use for extrapolation.\n",
    "pts_l = [40,50,60]\n",
    "\n",
    "# The Demographics1D and Demographics2D modules contain a few simple models,\n",
    "# mostly as examples. We could use one of those.\n",
    "func = dadi.Demographics2D.split_mig\n",
    "# Instead, we'll work with our custom model\n",
    "func = prior_onegrow_mig\n",
    "\n",
    "# Now let's optimize parameters for this model.\n",
    "\n",
    "# The upper_bound and lower_bound lists are for use in optimization.\n",
    "# Occasionally the optimizer will try wacky parameter values. We in particular\n",
    "# want to exclude values with very long times, very small population sizes, or\n",
    "# very high migration rates, as they will take a long time to evaluate.\n",
    "# Parameters are: (nu1F, nu2B, nu2F, m, Tp, T)\n",
    "upper_bound = [100, 100, 100, 10, 3, 3]\n",
    "lower_bound = [1e-2, 1e-2, 1e-2, 0, 0, 0]\n",
    "\n",
    "# This is our initial guess for the parameters, which is somewhat arbitrary.\n",
    "p0 = [2,0.1,2,1,0.2,0.2]\n",
    "# Make the extrapolating version of our demographic model function.\n",
    "func_ex = dadi.Numerics.make_extrap_log_func(func)\n",
    "\n",
    "# Perturb our parameters before optimization. This does so by taking each\n",
    "# parameter a up to a factor of two up or down.\n",
    "p0 = dadi.Misc.perturb_params(p0, fold=1, upper_bound=upper_bound,\n",
    "                              lower_bound=lower_bound)\n",
    "# Do the optimization. By default we assume that theta is a free parameter,\n",
    "# since it's trivial to find given the other parameters. If you want to fix\n",
    "# theta, add a multinom=False to the call.\n",
    "# The maxiter argument restricts how long the optimizer will run. For real \n",
    "# runs, you will want to set this value higher (at least 10), to encourage\n",
    "# better convergence. You will also want to run optimization several times\n",
    "# using multiple sets of intial parameters, to be confident you've actually\n",
    "# found the true maximum likelihood parameters.\n",
    "print('Beginning optimization ************************************************')\n",
    "popt = dadi.Inference.optimize_log(p0, data, func_ex, pts_l, \n",
    "                                   lower_bound=lower_bound,\n",
    "                                   upper_bound=upper_bound,\n",
    "                                   verbose=len(p0), maxiter=3)\n",
    "# The verbose argument controls how often progress of the optimizer should be\n",
    "# printed. It's useful to keep track of optimization process.\n",
    "print('Finshed optimization **************************************************')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These are the actual best-fit model parameters, which we found through\n",
    "# longer optimizations and confirmed by running multiple optimizations.\n",
    "# We'll work with them through the rest of this script.\n",
    "popt = [1.880, 0.0724, 1.764, 0.930, 0.363, 0.112]\n",
    "print('Best-fit parameters: {0}'.format(popt))\n",
    "\n",
    "# Calculate the best-fit model AFS.\n",
    "model = func_ex(popt, ns, pts_l)\n",
    "# Likelihood of the data given the model AFS.\n",
    "ll_model = dadi.Inference.ll_multinom(model, data)\n",
    "print('Maximum log composite likelihood: {0}'.format(ll_model))\n",
    "# The optimal value of theta given the model.\n",
    "theta = dadi.Inference.optimal_sfs_scaling(model, data)\n",
    "print('Optimal value of theta: {0}'.format(theta))\n",
    "\n",
    "# Plot a comparison of the resulting fs with the data.\n",
    "import pylab\n",
    "pylab.figure(figsize=(8,6))\n",
    "dadi.Plotting.plot_2d_comp_multinom(model, data, vmin=1, resid_range=3,\n",
    "                                    pop_ids =('YRI','CEU'), show=False)\n",
    "# Save the figure\n",
    "pylab.savefig('YRI_CEU.png', dpi=250)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Uncertainty analysis with GIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate parameter uncertainties using the Godambe Information Matrix, to\n",
    "# account for linkage in the data. \n",
    "import dadi.Godambe\n",
    "\n",
    "# To use the GIM approach, we need to have spectra from bootstrapping our \n",
    "# data.  Let's load the ones we've provided for the example.  \n",
    "# (We're using Python list comprehension syntax to do this in one line.)\n",
    "all_boot = [dadi.Spectrum.from_file('bootstraps/{0:02d}.fs'.format(ii)) \n",
    "            for ii in range(100)]\n",
    "uncerts = dadi.Godambe.GIM_uncert(func_ex, pts_l, all_boot, popt, data, \n",
    "                                  multinom=True)\n",
    "# uncert contains the estimated standard deviations of each parameter, with\n",
    "# theta as the final entry in the list.\n",
    "print('Estimated parameter standard deviations from GIM: {0}'.format(uncerts))\n",
    "\n",
    "# For comparison, we can estimate uncertainties with the Fisher Information\n",
    "# Matrix, which doesn't account for linkage in the data and thus underestimates\n",
    "# uncertainty. (Although it's a fine approach if you think your data is truly\n",
    "# unlinked.)\n",
    "uncerts_fim = dadi.Godambe.FIM_uncert(func_ex, pts_l, popt, data, multinom=True)\n",
    "print('Estimated parameter standard deviations from FIM: {0}'.format(uncerts_fim))\n",
    "\n",
    "print('Factors by which FIM underestimates parameter uncertainties: {0}'.format(uncerts/uncerts_fim))\n",
    "\n",
    "# What if we fold the data?\n",
    "# These are the optimal parameters when the spectrum is folded. \n",
    "popt_fold = [1.910, 0.074, 1.787, 0.914, 0.439, 0.114]\n",
    "all_boot_fold = [_.fold() for _ in all_boot]\n",
    "uncerts_folded = dadi.Godambe.GIM_uncert(func_ex, pts_l, all_boot_fold,\n",
    "                                         popt_fold, data.fold(), multinom=True)\n",
    "print('Folding increases parameter uncertainties by factors of: {0}'.format(uncerts_folded/uncerts))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Likelihood ratio test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's do a likelihood-ratio test comparing models with and without migration.\n",
    "# The no migration model is implemented as prior_onegrow_nomig\n",
    "func_nomig = prior_onegrow_nomig\n",
    "func_ex_nomig = dadi.Numerics.make_extrap_log_func(func_nomig)\n",
    "# These are the best-fit parameters, which we found by multiple optimizations\n",
    "popt_nomig = np.array([ 1.897,  0.0388,  9.677,  0.395,  0.070])\n",
    "model_nomig = func_ex_nomig(popt_nomig, ns, pts_l)\n",
    "ll_nomig = dadi.Inference.ll_multinom(model_nomig, data)\n",
    "\n",
    "# Since LRT evaluates the complex model using the best-fit parameters from the\n",
    "# simple model, we need to create list of parameters for the complex model\n",
    "# using the simple (no-mig) best-fit params.  Since evalution is done with more\n",
    "# complex model, need to insert zero migration value at corresponding migration\n",
    "# parameter index in complex model. And we need to tell the LRT adjust function\n",
    "# that the 3rd parameter (counting from 0) is the nested one.\n",
    "p_lrt = [1.897,  0.0388,  9.677, 0, 0.395,  0.070]\n",
    "\n",
    "adj = dadi.Godambe.LRT_adjust(func_ex, pts_l, all_boot, p_lrt, data, \n",
    "                              nested_indices=[3], multinom=True)\n",
    "D_adj = adj*2*(ll_model - ll_nomig)\n",
    "print('Adjusted D statistic: {0:.4f}'.format(D_adj))\n",
    "\n",
    "# Because this is test of a parameter on the boundary of parameter space \n",
    "# (m cannot be less than zero), our null distribution is an even proportion \n",
    "# of chi^2 distributions with 0 and 1 d.o.f. To evaluate the p-value, we use the\n",
    "# point percent function for a weighted sum of chi^2 dists.\n",
    "pval = dadi.Godambe.sum_chi2_ppf(D_adj, weights=(0.5,0.5))\n",
    "print('p-value for rejecting no-migration model: {0:.4f}'.format(pval))\n",
    "\n",
    "# This ensures that the figures pop up. It may be unecessary if you are using\n",
    "# ipython.\n",
    "pylab.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
}
