{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Frequency spectrum from data\n",
    "\n",
    "Extracting a frequency spectrum from a VCF file and generating bootstrap samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dadi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse the VCF file to generate a data dictionary\n",
    "datafile = '1KG.YRI.CEU.biallelic.synonymous.snps.withanc.strict.subset.vcf.gz'\n",
    "dd = dadi.Misc.make_data_dict_vcf(datafile, '1KG.YRI.CEU.popfile.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract the spectrum for ['YRI','CEU'] from that dictionary, with \n",
    "# YRI projected down to 20 and CEU projected down to 24.\n",
    "# We project down like this just to make fitting faster. For a real analysis\n",
    "# we would not project so severely.\n",
    "pop_ids, ns = ['YRI','CEU'], [20,24]\n",
    "fs = dadi.Spectrum.from_data_dict(dd, pop_ids, ns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can save our extracted spectrum to disk\n",
    "fs.to_file('1KG.YRI.CEU.biallelic.synonymous.snps.withanc.strict.subset.fs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If we didn't have outgroup information, we could fold the fs.\n",
    "fs_folded = dadi.Spectrum.from_data_dict(dd, pop_ids, ns, polarized=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate 100 bootstrap datasets, by dividing the genome into 2 Mb chunks and\n",
    "# resampling from those chunks.\n",
    "Nboot, chunk_size = 100, 2e6\n",
    "chunks = dadi.Misc.fragment_data_dict(dd, chunk_size)\n",
    "boots = dadi.Misc.bootstraps_from_dd_chunks(chunks, Nboot, pop_ids, ns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you're modeling inbreeding, you cannot project your data downward. Instead,\n",
    "# to deal with missing data we subsample individuals in the VCF file.\n",
    "# If we're modeling inbreeding, it is important that we *never* project downward,\n",
    "# as this destroys the genotype information.\n",
    "dd_subsample = dadi.Misc.make_data_dict_vcf(datafile, '1KG.YRI.CEU.popfile.txt',\n",
    "                                            subsample={'YRI': ns[0]//2, 'CEU': ns[1]//2})\n",
    "fs_subsample = dadi.Spectrum.from_data_dict(dd_subsample, pop_ids, ns)\n",
    "# Bootstrapping with subsampling is more computationally expensive, \n",
    "# because we must repeatedly access the VCF file to subsample.\n",
    "boots_subsample = dadi.Misc.bootstraps_subsample_vcf(datafile, '1KG.YRI.CEU.popfile.txt',\n",
    "                                                     subsample={'YRI': ns[0]//2, 'CEU': ns[1]//2}, Nboot=2, \n",
    "                                                     chunk_size=chunk_size, pop_ids=pop_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot comparing the multiple versions of our data spectra.\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "fig = plt.figure(1, figsize=(10,6))\n",
    "fig.clear()\n",
    "\n",
    "# Note that projection creates fractional entries in the spectrum,\n",
    "# so vmin < 1 is sensible.\n",
    "ax = fig.add_subplot(2,3,1)\n",
    "dadi.Plotting.plot_single_2d_sfs(fs, vmin=1e-2, ax=ax)\n",
    "ax.set_title('Orignal data')\n",
    "\n",
    "ax = fig.add_subplot(2,3,2)\n",
    "dadi.Plotting.plot_single_2d_sfs(fs_folded, vmin=1e-2, ax=ax)\n",
    "ax.set_title('Folded original data')\n",
    "\n",
    "ax = fig.add_subplot(2,3,3)\n",
    "dadi.Plotting.plot_single_2d_sfs(boots[0], vmin=1e-2, ax=ax)\n",
    "ax.set_title('Bootstrap from original data')\n",
    "\n",
    "# Subsampling does not create those fractional entries.\n",
    "ax = fig.add_subplot(2,3,4)\n",
    "dadi.Plotting.plot_single_2d_sfs(fs_subsample, vmin=1e-2, ax=ax)\n",
    "ax.set_title('Subsampled original data')\n",
    "\n",
    "ax = fig.add_subplot(2,3,6)\n",
    "dadi.Plotting.plot_single_2d_sfs(boots_subsample[0], vmin=1e-2, ax=ax)\n",
    "ax.set_title('Bootstrap subsampled data')\n",
    "\n",
    "fig.tight_layout()\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
}
