simulation-analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compilation of test data for performance tests\n",
    "\n",
    "This notebook records the parameters for Wright-Fisher simulations used to generate our test data sets, as well as commands for running covariance estimation algorithms on the test data and compiling the results. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import required libraries and define global variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "This notebook was prepared using:\n",
      "python version 3.6.10 |Anaconda, Inc.| (default, Mar 25 2020, 18:53:43) \n",
      "[GCC 4.2.1 Compatible Clang 4.0.1 (tags/RELEASE_401/final)]\n",
      "numpy version 1.19.1\n",
      "seaborn version 0.11.0\n",
      "matplotlib version 3.3.2\n",
      "scipy version 1.5.2\n"
     ]
    }
   ],
   "source": [
    "# Full library list and version numbers\n",
    "\n",
    "print('This notebook was prepared using:')\n",
    "\n",
    "import sys\n",
    "print('python version %s' % sys.version)\n",
    "\n",
    "import numpy as np\n",
    "print('numpy version %s' % np.__version__)\n",
    "\n",
    "import seaborn as sns\n",
    "print('seaborn version %s' % sns.__version__)\n",
    "\n",
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.gridspec import GridSpec\n",
    "print('matplotlib version %s' % mpl.__version__)\n",
    "\n",
    "import scipy\n",
    "from scipy import stats\n",
    "print('scipy version %s' % scipy.__version__)\n",
    "\n",
    "import importlib\n",
    "\n",
    "sys.path.append('./src')\n",
    "import simulation_setup as SS\n",
    "import subsample as subsample\n",
    "import data_parser as DP\n",
    "\n",
    "# GLOBAL VARIABLES\n",
    "\n",
    "N = 1000        # number of sequences in the population\n",
    "L = 50          # sequence length\n",
    "T = 700         # number of generations\n",
    "MU = 1e-3       # mutation rate\n",
    "NUM_SELECTIONS = 10  # number of sets of selection coefficients that we used\n",
    "NUM_TRIALS = 20  # for each set of selection coefficients, we simulated for 20 times\n",
    "SAMPLE = [1000, 500, 100, 50, 10]  # sampling depth options when subsampling\n",
    "RECORD = [1, 3, 5, 10]  # sampling time interval options when subsampling\n",
    "LINEAR = [1, 5, 10, 20, 50, 100, 300]  # list of linear shrinkage strengths/sampling time intervals, e.g., interval=5, linear=5 ==> strength=25'\n",
    "GAMMA = [1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1]  # non-linear shrinkage parameter choices that we tested\n",
    "TRUNCATE = [200, 300, 400, 500, 600, 700]  # truncate options to test performance using different lengths of data\n",
    "WINDOW = [0, 1, 2, 3, 4, 5, 10, 20, 40, 80, 160]  # window choices that we tested\n",
    "\n",
    "# loss functions for non-linear shrinkage that we tested\n",
    "LOSS = ['Fro | $\\hat{C}-C$', 'Fro | $\\hat{C}^{-1}-C^{-1}$', 'Fro | $C^{-1}\\hat{C}-I$', \n",
    "        'Fro | $\\hat{C}^{-1}C-I$', 'Fro | $\\hat{C}^{-1/2}C\\hat{C}^{-1/2}-I$', \n",
    "        'Nuc | $\\hat{C}-C$', 'Nuc | $\\hat{C}^{-1}-C^{-1}$', 'Nuc | $C^{-1}\\hat{C}-I$', \n",
    "        'Nuc | $\\hat{C}^{-1}C-I$', 'Nuc | $\\hat{C}^{-1/2}C\\hat{C}^{-1/2}-I$'] \n",
    "\n",
    "# GitHub directories\n",
    "DATA_DIR = './data'\n",
    "SELECTION_DIR = './src/selection'\n",
    "INITIAL_DIR = './src/initial'\n",
    "SIMULATION_OUTPUT_DIR = './data/simulation_output'\n",
    "SUBSAMPLE_OUTPUT_DIR = './data/subsample_output'\n",
    "ESTIMATION_OUTPUT_DIR = './data/estimation_output'\n",
    "JOB_DIR = './src/jobs'\n",
    "\n",
    "# relative directories looking from shell scripts in JOB_DIR\n",
    "SRC_DIR_REL = '..'\n",
    "DATA_DIR_REL = '../../data'\n",
    "SELECTION_DIR_REL = '../../src/selection'\n",
    "INITIAL_DIR_REL = '../../src/initial'\n",
    "SIMULATION_OUTPUT_DIR_REL = '../../data/simulation_output'\n",
    "SUBSAMPLE_OUTPUT_DIR_REL = '../../data/subsample_output'\n",
    "ESTIMATION_OUTPUT_DIR_REL = '../../data/estimation_output'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reload():\n",
    "    importlib.reload(SS)\n",
    "    importlib.reload(subsample)\n",
    "    importlib.reload(DP)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 1. Generation of test data through Wright-Fisher simulations¶\n",
    "\n",
    "Wright-Fisher simulations are performed using src/Wright-Fisher.py. The output of these simulations is saved for processing. The code below creates multiple job files for running many simulations in parallel on a computer cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate selection coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10 sets of selection coefficients used for simulations\n",
    "SS.generate_selection_coefficients(overwrite=False)\n",
    "selections = [SS.load_selection(s) for s in range(NUM_SELECTIONS)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate initial distributions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The initial distribution with four genotypes\n",
    "SS.generate_default_initial_distribution(overwrite=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Genotype 1 01111010010011111100101001101111001110101010111110 count=260\n",
      "Genotype 2 11111010000010000100110010011100011110111010001001 count=210\n",
      "Genotype 3 01100111010010100001101101000011111011011010101111 count=280\n",
      "Genotype 4 10011100010100100000011011011011111111111000011000 count=250\n",
      "\n"
     ]
    }
   ],
   "source": [
    "initial = SS.load_default_initial()\n",
    "SS.print_initial_distribution(initial)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 20 initial distributions of genotypes, with 4 founder genotypes\n",
    "SS.generate_initial_distributions_four_founder_genotypes(overwrite=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Genotype 1 10101011010000010010010111011011000011100111110010 count=210\n",
      "Genotype 2 00110001001100001011100010010011100010100010101101 count=300\n",
      "Genotype 3 11001011001001000010101000001101000011010000011010 count=250\n",
      "Genotype 4 10011111011111100010111110100000000001111110011110 count=240\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Check initial distributions, print out an example\n",
    "i = 15\n",
    "initial = SS.load_initial_four_founder_genotypes(i)\n",
    "SS.print_initial_distribution(initial, max_line=30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 20 initial distributions of genotypes, with 5 to 24 founder genotypes\n",
    "SS.generate_initial_distributions(overwrite=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Genotype 1 01111101101011111101110011000000100111110101101011 count=40\n",
      "Genotype 2 01000101110111110010111111101011101111100010100101 count=30\n",
      "Genotype 3 10100011111011111110110101110000001001000001001010 count=70\n",
      "Genotype 4 11101000000000110101100100100000110001110011000001 count=50\n",
      "Genotype 5 01010111001110001101000110010010001111001101001100 count=50\n",
      "Genotype 6 01101100111011111110010111011011101011101000001111 count=80\n",
      "Genotype 7 10000111100100100010101101100101100001111000010110 count=20\n",
      "Genotype 8 00111110000010111010010011101011001001011010111100 count=100\n",
      "Genotype 9 10011111111010000011001100011000001100111001011101 count=20\n",
      "Genotype 10 11010111100000111100011110001100010110001001010101 count=20\n",
      "Genotype 11 10001001101100111100111110110010111101100110101110 count=40\n",
      "Genotype 12 10001010010101111011000100001001000011010010110010 count=60\n",
      "Genotype 13 00000001100100010010111110111000101010101101000100 count=30\n",
      "Genotype 14 10100111010100001001001011011100010000000100110000 count=30\n",
      "Genotype 15 00111000111110001000010101010011011110010001101000 count=70\n",
      "Genotype 16 01110100001101111011100011101100101000011101000010 count=20\n",
      "Genotype 17 00101010111011110100111100011001001110110010011000 count=50\n",
      "Genotype 18 01010001001001101101100011110001101101101111111110 count=70\n",
      "Genotype 19 11010111011001001010001100111010001110100101110000 count=50\n",
      "Genotype 20 11110100010000001011101010111010110111010111101010 count=100\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Check initial distributions, print out an example\n",
    "i = 15\n",
    "initial = SS.load_initial(i)\n",
    "SS.print_initial_distribution(initial, max_line=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate simulations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In total, there are num_selections x num_trials = 200 simulations.\n",
    "# Create jobs to run on cluster\n",
    "SS.generate_simulation_scripts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulations all with four founder genotypes, but different\n",
    "reload()\n",
    "SS.generate_simulation_scripts(varying_founder_genotype=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulations with varying initial distributions (5 to 24 founder genotypes)\n",
    "SS.generate_simulation_scripts(varying_initial=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulations with recombination\n",
    "SS.generate_simulation_scripts(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can run locally as well. Will take several hours.\n",
    "# SS.generate_simulation_scripts(local=True)\n",
    "# SS.generate_simulation_scripts(local=True, varying_initial=True)\n",
    "# SS.generate_simulation_scripts(local=True, recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test if simulations are complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_simulation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "reload()\n",
    "SS.test_load_simulation(varying_founder_genotype=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_simulation(varying_initial=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_simulation(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 2. Subsample the simulated results for later inference under various limited sampling scenarios\n",
    "We subsample the simulated results varying sampling depth and time interval. The output of is saved for investigating performance of inference under different limited sampling effects. The code below creates multiple files for running many subsampling jobs in parallel on a computer cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In total, there are num_selections x num_trials x len(SAMPLE) x len(RECORD) = 4000 \n",
    "# subsampling runs, aggregated as 200 jobs.\n",
    "SS.generate_subsample_scripts()\n",
    "SS.generate_subsample_scripts(varying_founder_genotype=True)\n",
    "SS.generate_subsample_scripts(varying_initial=True)\n",
    "SS.generate_subsample_scripts(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Store time-series covariance for a case for a supplementary figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_subsample_script_with_time_series_covariance()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test if subsamples are complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_subsample()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "reload()\n",
    "SS.test_load_subsample(varying_founder_genotype=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_subsample(varying_initial=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_subsample(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 3. Running the covariance estimation algorithms and compiling output\n",
    "Finally we estimate covariance and infer selection coefficients for all subsampled cases generated above. The output of is saved for plotting. The code below creates multiple files to run in parallel on a computer cluster."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimation runs with minimal output\n",
    "The jobs defined below run the estimation algorithm and store a minimal size of output, which contains covariance matrix only for some cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "reload()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In total, there are len(TRUNCATE) x len(WINDOW) = 66 jobs\n",
    "# Each job (for a particular (tr, win) pair) outputs ~8MB, totaling ~500MB\n",
    "SS.generate_estimate_scripts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_estimate_scripts(varying_founder_genotype=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_estimate_scripts(varying_initial=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_estimate_scripts(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_estimation(varying_founder_genotype=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_estimation()\n",
    "SS.test_load_estimation(varying_initial=True)\n",
    "SS.test_load_estimation(recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Estimation runs with complete output\n",
    "The jobs defined below run the estimation algorithm and store complete output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "PARAMS_COMPLETE_OUTPUT = {\n",
    "    'truncate_list': [700],\n",
    "    'window_list': [20],\n",
    "    'complete_output': True,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Each job (for a particular (tr, win) pair) outputs ~500MB\n",
    "# If running for all (tr, win) pairs, the total output will be ~30GB\n",
    "SS.generate_estimate_scripts(**PARAMS_COMPLETE_OUTPUT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_estimate_scripts(**PARAMS_COMPLETE_OUTPUT, \n",
    "                             varying_initial=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_estimate_scripts(**PARAMS_COMPLETE_OUTPUT, \n",
    "    recombination=True, r=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 4. Running Evoracle\n",
    "We also test Evoracle on our simulated data, as a compare method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.generate_evoracle_scripts()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then run the evoracle scripts. Instructions on installing the *Evoracle* package can be found in https://github.com/maxwshen/evoracle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "SS.test_load_evoracle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Section 5. Running haploSep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First save simulation trajectories as input for haploSep\n",
    "DP.save_simulation_traj_for_haplosep()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then run the R script, *data/haploSep/run_haploSep_for_simulated_data.R*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 32min 40s, sys: 10.6 s, total: 32min 51s\n",
      "Wall time: 32min 14s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# After running haploSep using the R script, \n",
    "# parse the results and save the parsed output\n",
    "_ = DP.parse_simulation_results_for_haplosep(save=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:root]",
   "language": "python",
   "name": "conda-root-py"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}