42_simulate_bulk_competition_warringer2003_with_error_analysis.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate bulk competition experiments using empirical traits from Warringer 2003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "\n",
    "from bulk_simulation_code import run_pairwise_experiment, run_bulk_experiment\n",
    "from bulk_simulation_code import CalcRelativeYield,CalcReferenceFrequency\n",
    "from bulk_simulation_code import CalcTotalSelectionCoefficientLogit\n",
    "from m3_model import CalcRelativeSaturationTime as CalcSaturationTimeExact\n",
    "from m3_model import CalcFoldChangeWholePopulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Update dependent parameters according to input\n",
    "import os\n",
    "import os.path\n",
    "from os import path\n",
    "\n",
    "## create export directory if necessary\n",
    "## foldernames for output plots/lists produced in this notebook\n",
    "import os\n",
    "FIG_DIR = f'./figures/bulk_fitness/'\n",
    "os.makedirs(FIG_DIR, exist_ok=True)\n",
    "print(\"All  plots will be stored in: \\n\" + FIG_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### execute script to load modules here\n",
    "# I get some error with this command\n",
    "# exec(open('setup_aesthetics.py').read()) \n",
    "\n",
    "# manual fix\n",
    "\n",
    "FIGSIZE_A4 = (8.27, 11.69) # a4 format in inches\n",
    "\n",
    "FIGWIDTH_TRIPLET = FIGSIZE_A4[0]*0.3*2\n",
    "FIGHEIGHT_TRIPLET = FIGWIDTH_TRIPLET*0.75\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_COLOR = 'darkorange'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SUFFIX_DATASET = 'warringer/'\n",
    "\n",
    "FIG_DIR_DATASET = FIG_DIR + SUFFIX_DATASET\n",
    "os.makedirs(FIG_DIR_DATASET, exist_ok=True)\n",
    "\n",
    "OUTPUT_DIR_DATASET = './output/' + SUFFIX_DATASET\n",
    "os.makedirs(OUTPUT_DIR_DATASET, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load wild-type traits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "INDEX_COL = [0,1,2,3,4]\n",
    "list_na_representations = ['not_present', 'failed_to_compute']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PCWS_TRAITS_WARRINGER = './output/df_M3_traits.csv'\n",
    "df_warringer = pd.read_csv(PCWS_TRAITS_WARRINGER, header = 0, index_col= INDEX_COL,\\\n",
    "                                  float_precision=None, na_values=list_na_representations)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### define default wild_type\n",
    "df_wildtypes = df_warringer[df_warringer['is_wildtype']==True]\n",
    "\n",
    "WILDTYPE = df_wildtypes.median(axis = 0, numeric_only = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load mutant data (averaged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "PCWS_TRAITS_WARRINGER_AVERAGED = './output/df_M3_traits_averaged.csv'\n",
    "df_averaged = pd.read_csv(PCWS_TRAITS_WARRINGER_AVERAGED, header = 0, float_precision=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### assign wild-type label\n",
    "def is_wildtype(row):\n",
    "    genotype = row['genotype']\n",
    "    \n",
    "    if genotype == 'BY4741':\n",
    "        return True\n",
    "    else:\n",
    "        return False\n",
    "    \n",
    "\n",
    "row = df_averaged.iloc[0]\n",
    "is_wildtype(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_averaged['is_wildtype'] = df_averaged.apply(is_wildtype, axis =1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### append mutant values (averaged) to set of individual wild-type strains\n",
    "df_knockouts = df_averaged[~df_averaged['is_wildtype']]\n",
    "df_knockouts = df_knockouts\n",
    "df_input = df_wildtypes.reset_index().append(df_knockouts.reset_index())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### restore index\n",
    "index_col_names = df_warringer.index.names\n",
    "df_input = df_input.set_index(index_col_names)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load trait data into the standard form required by Michaels code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_knockouts = df_knockouts.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### growth rates\n",
    "gs = np.zeros(n_knockouts+1)\n",
    "gs[0] = WILDTYPE['gmax']\n",
    "gs[1:] = df_knockouts['gmax'].values\n",
    "\n",
    "### lag times\n",
    "ls = np.zeros(n_knockouts+1)\n",
    "ls[0] = WILDTYPE['lag']\n",
    "ls[1:] = df_knockouts['lag'].values\n",
    "\n",
    "### adjust units of time\n",
    "gs = gs*60 # change units to growth rate per hour\n",
    "ls = ls/60 # change units to hour\n",
    "\n",
    "### yield\n",
    "Ys = np.zeros(n_knockouts+1)\n",
    "Ys[0] = WILDTYPE['yield']\n",
    "#Ys[1:] = Ys[0] #switch off variation in yield\n",
    "Ys[1:] = df_knockouts['yield'].values\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define initial condition for bulk growth cycle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### set initial resource concentrations\n",
    "\n",
    "CONCENTRATION_GLUCOSE = 20/180 * 1e3 # concentrations are recored  in milliMolar, to match the units of yield\n",
    "print(CONCENTRATION_GLUCOSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### define default initial_OD\n",
    "OD_START = 0.05  #df_warringer['od_start'].median()\n",
    "\n",
    "### compare to initial OD in the monoculture cycles\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "\n",
    "ax = df_warringer['od_start'].hist(bins=41, color = DATASET_COLOR, alpha = 0.6, log = True, rasterized = True)\n",
    "\n",
    "\n",
    "ax.axvline(OD_START, color = 'tab:red', label = f'median value: $N_0={OD_START:.3f}$')\n",
    "ax.legend()\n",
    "ax.set_xlabel('initial OD')\n",
    "ax.set_ylabel('no. growth curves')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate effective yield"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bulk_simulation_code import CalcRelativeYield"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### calculcate effective yields\n",
    "nus = CalcRelativeYield(Ys, R0 = CONCENTRATION_GLUCOSE, N0 = OD_START)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate pairwise competition growth cycles (scenario A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frequencies for scenario A can be summarized as \n",
    "\n",
    "    frequency of the focal mutant strain: x2 = 1/N\n",
    "        frequency of the wildtype strain: x1 = 1 - x2\n",
    "\n",
    "where $N$ is some population size (number of cells or biomass?). Intutively, a spontanteous mutation initially forms only a small fraction $x0 = 1/N$ in the population. The values of the population size in nature are largely unknown, but can be approximated in two ways. \n",
    "\n",
    "- by the effective population size $N_e$, which is inferred from the genomic variation across a set of natural isolates, and leads to estimates of $N\\approx 10^8$ cells [see papers by Howard Ochman]\n",
    "- by the bottleneck size $N$ in laboratory evolution experiments like the LTEE, which leads to an estimate of $N=5\\cdot 10^6$ cells. According to the first paper on the LTEE, there are $5\\cdot10^5$ cells per ml at the starting point of the growth cycle, total volue is 10ml. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs_pair, xs_pair_final, tsats_pair,fcs_both, _,_ = \\\n",
    "run_pairwise_experiment(gs=gs,ls=ls,nus = nus, g1=gs[0],l1=ls[0],nu1=nus[0],x0 = 1/N)\n",
    "\n",
    "si_pair = CalcTotalSelectionCoefficientLogit(xs_pair,xs_pair_final)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate bulk competition with background mutants and added wild-type lineage  (scenario B2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frequencies for this scenario can be summarized as \n",
    "\n",
    "    frequency of the mutant straints:     xi = 1/(k+1)\n",
    "    frequency of the wildtype strain:     x1 = 1/(k+1)\n",
    "\n",
    "where `k` is the number of knockouts strain. Here all lineages, the mutants and the wild-type, have the same initial frequency. This roughly resembles scenario 'Bfull', but with a barcoded wild-type spiked into the culture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = n_knockouts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### set initial frequencies\n",
    "xs = np.zeros_like(gs)\n",
    "xs[1:] = 1/(k+1)           # mutant lineages\n",
    "xs[0] = 1/(k+1)              # wildtype population\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Proportion of mutants: %.8f \" % xs[1:].sum() )\n",
    "print(\"Proportion of wild-type: %.8f \" % xs[0] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## calculate final frequencies\n",
    "xs, xs_final,tsat = run_bulk_experiment(gs=gs, ls = ls, nus =nus, xs=xs)\n",
    "\n",
    "## compute total foldchange\n",
    "fc_bulk = CalcFoldChangeWholePopulation(t=tsat,xs=xs,gs=gs,ls=ls)\n",
    "\n",
    "## calculate total selection coefficient\n",
    "si_bulk_B2 = CalcTotalSelectionCoefficientLogit(xs,xs_final)\n",
    "\n",
    "## compute pairwise selection coefficient in bulk\n",
    "xi1 = CalcReferenceFrequency(xs,ref_strains = [0]) \n",
    "xi1_final = CalcReferenceFrequency(xs_final,ref_strains = [0])\n",
    "si1_bulk_B2 = CalcTotalSelectionCoefficientLogit(xi1,xi1_final)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate error to pairwise competition as ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey=True)\n",
    "\n",
    "\n",
    "ax = axes[0] # pairwise selection coefficient in bulk\n",
    "x = si_pair \n",
    "y = si1_bulk_B2 - si_pair\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey') \n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_21 bulk - s_21 pair')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "ax = axes[1] # total selection coefficient in bulk\n",
    "x = si_pair\n",
    "y = si_bulk_B2 - si_pair\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey') \n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_2 bulk - s_21 pair')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate error between total and pairwise selection coefficient"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate trait components of the selection coefficient"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from m3_model import CalcApproxSijComponentsMultitype, CalcApproxSijComponents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "si1_bulk_growth = np.zeros_like(si1_bulk_B2)\n",
    "si1_bulk_lag = np.zeros_like(si1_bulk_B2)\n",
    "si1_bulk_coupling = np.zeros_like(si1_bulk_B2)\n",
    "\n",
    "for i in range(len(gs)):\n",
    "    si1_bulk_growth[i], si1_bulk_lag[i], si1_bulk_coupling[i] = CalcApproxSijComponentsMultitype(i,0,xs,gs,ls,nus)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "si_pair_growth = np.zeros_like(si_pair)\n",
    "si_pair_lag = np.zeros_like(si_pair)\n",
    "si_pair_coupling = np.zeros_like(si_pair)\n",
    "\n",
    "for i in range(len(gs)):\n",
    "        g1, l1, nu1 = gs[0], ls[0], nus[0]\n",
    "        g2, l2, nu2 = gs[i], ls[i], nus[i] # get traits of the invader\n",
    "        x0 = 1/N\n",
    "        si_pair_growth[i], si_pair_lag[i], si_pair_coupling[i] =CalcApproxSijComponentsMultitype(1,0,\n",
    "                                                xs=[1-x0,x0], gs = [g1,g2], ls= [l1,l2], nus = [nu1,nu2] )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Break down the error by trait components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey = True)\n",
    "\n",
    "\n",
    "ax = axes[0] # error from growth and lag component\n",
    "x = si_pair\n",
    "y = si1_bulk_growth + si1_bulk_lag - si_pair_growth - si_pair_lag\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('delta s_growth + delta s_lag')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "\n",
    "ax = axes[1] # error from coupling component\n",
    "x = si_pair\n",
    "y = si1_bulk_coupling\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_coupling: bulk experiment')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare size of the two error components\n",
    "fig, axes = plt.subplots( figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey = True)\n",
    "\n",
    "\n",
    "ax = axes # error from growth and lag component\n",
    "x = si1_bulk_growth + si1_bulk_lag - si_pair_growth - si_pair_lag\n",
    "y = si1_bulk_coupling\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')\n",
    "\n",
    "## add axis labels\n",
    "ax.set_xlabel('error growth and lag: delta s_growth + delta s_lag')\n",
    "ax.set_ylabel('error from coupling: s_coupling')\n",
    "\n",
    "\n",
    "### make square\n",
    "xmin,xmax = ax.get_xlim()\n",
    "ymin,ymax = ax.get_ylim()\n",
    "xymin = np.min([xmin,ymin])\n",
    "xymax = np.max([xmax,ymax])\n",
    "xyabs = np.max(np.abs([xymin,xymax]))\n",
    "ax.set_xlim(-xyabs,xyabs)\n",
    "ax.set_ylim(-xyabs,xyabs)\n",
    "## add diagonal line\n",
    "ax.plot([-xyabs,xyabs],[-xyabs,xyabs], ls = '--', color = 'black')\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the underlying changes in saturation time and fold-change"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute effective doubling times\n",
    "\n",
    "taus = 1/gs\n",
    "tau_bars_pair = np.zeros_like(taus)\n",
    "x0 = 1/N\n",
    "\n",
    "for i in range(len(gs)):\n",
    "    tau1, tau2 = taus[0], taus[i]\n",
    "    Y1, Y2 = Ys[0], Ys[i]\n",
    "    x1, x2 = 1-x0, x0\n",
    "    tau_bars_pair[i] = (x1/Y1 + x2/Y2)/(x1/Y1/tau1 + x2/Y2/tau2)\n",
    "\n",
    "## for bulk experiment\n",
    "\n",
    "tau_bar_bulk = np.sum(np.divide(xs,Ys)) / np.sum(np.divide(np.divide(xs,Ys),taus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot:  tau_bar vs log fold-change\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey = True)\n",
    "\n",
    "ax = axes[0] # plot tau_bar data\n",
    "x = si_pair\n",
    "y = np.divide(tau_bar_bulk - tau_bars_pair, tau_bars_pair)\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')   # add data\n",
    "ax.axhline(0, color = 'black', ls = '--')               # add horizontal line at zero\n",
    "# center the plot vertically at zero\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = 1.1*np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "# add axis labels\n",
    "ax.set_ylabel('tau_bar: (bulk - pair)/pair')\n",
    "ax.set_xlabel('s_21: pairwise competition')\n",
    "\n",
    "ax = axes[1] # plot log foldchange data\n",
    "x = si_pair\n",
    "y = np.divide(np.log(fc_bulk) - np.log(fcs_pair), np.log(fcs_pair))\n",
    "\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')   # add data\n",
    "ax.axhline(0, color = 'black', ls = '--')               # add horizontal line at zero\n",
    "# center the plot vertically at zero\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = 1.1*np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "# add axis labels\n",
    "ax.set_ylabel('log foldchange: (bulk - pair)/pair')\n",
    "ax.set_xlabel('s_21: pairwise competition')\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show that error is mostly explained by tau_bar\n",
    "\n",
    "### compare size of the two error components\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex=True)\n",
    "\n",
    "\n",
    "ax = axes[0] # error from growth and lag component\n",
    "x = np.divide(tau_bar_bulk - tau_bars_pair, tau_bars_pair)\n",
    "y =  np.divide(si1_bulk_lag - si_pair_lag, si_pair_lag, where = si_pair_lag != 0)\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')\n",
    "## add axis labels\n",
    "ax.set_ylabel('error: delta s_lag')\n",
    "ax.set_xlabel('tau_bar: (bulk - pair)/pair')\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "x = np.divide(tau_bar_bulk - tau_bars_pair, tau_bars_pair)\n",
    "y =  np.divide(si1_bulk_growth - si_pair_growth, si_pair_growth, where = si_pair_growth != 0)\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = 'dimgrey')\n",
    "\n",
    "\n",
    "## add axis labels\n",
    "ax.set_ylabel('error: delta s_growth ')\n",
    "ax.set_xlabel('tau_bar: (bulk - pair)/pair')\n",
    "\n",
    "for ax in axes:\n",
    "    ax.ticklabel_format(useOffset=False, style = 'sci', scilimits = (0,0))\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Finding the optimal fraction of mutants in the bulk experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frequencies for this scenario can be summarized as \n",
    "\n",
    "    frequency of the mutant strain s:     xi = x/k\n",
    "    frequency of the wildtype strain:     x1 = 1-x\n",
    "    \n",
    "where `k` is the number of knockouts strain and `x` is the total proportion of mutant. This scenario is interpolation between the scenario **B1** with dominating wild-type (`x<<1`) and the scenario **Bfull** with no wildtype (`x=1`). \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### choose range of frequencies to test\n",
    "#xrange = np.geomspace(1/N,0.99, num = 30)\n",
    "#xrange = np.linspace(0.01,0.99, num = 20)\n",
    "xrange = np.linspace(0.01, k/(k-1),num = 20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "si_bulk_Bx = np.zeros((len(xrange),n_knockouts+1))\n",
    "si1_bulk_Bx = np.zeros((len(xrange),n_knockouts+1))\n",
    "\n",
    "\n",
    "for i in range(len(xrange)):\n",
    "    x = xrange[i]\n",
    "    \n",
    "    ### set initial frequencies\n",
    "    xs = np.zeros_like(gs)\n",
    "    xs[1:] = x/k         # mutant lineages\n",
    "    xs[0] = 1-x              # wildtype population\n",
    "\n",
    "    ## calculate final frequencies\n",
    "    xs, xs_final,_ = run_bulk_experiment(gs=gs, ls = ls, nus =nus, xs=xs)\n",
    "\n",
    "    ## calculate total selection coefficient\n",
    "    si_bulk_Bx[i,:] = CalcTotalSelectionCoefficientLogit(xs,xs_final)\n",
    "    \n",
    "    ## compute pairwise selection coefficient in bulk\n",
    "    xi1 = CalcReferenceFrequency(xs,ref_strains = [0]) \n",
    "    xi1_final = CalcReferenceFrequency(xs_final,ref_strains = [0])\n",
    "    si1_bulk_Bx[i,:] = CalcTotalSelectionCoefficientLogit(xi1,xi1_final)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot error to pairwise competition as a function of mutant frequency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "truth = np.outer(np.ones(len(xrange)),si_pair) # need the right shape\n",
    "\n",
    "## calculate error for pairwise selection coefficients: from higher order interactions\n",
    "error_sij_abs = si1_bulk_Bx - truth \n",
    "error_sij_rel = np.divide(error_abs,truth, where = truth != 0)\n",
    "\n",
    "## calculate error from total selection coefficients\n",
    "error_si_abs = si_bulk_Bx - truth \n",
    "error_si_rel = np.divide(error_abs,truth, where = truth != 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### use a line plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey = True)\n",
    "\n",
    "ax = axes[0]\n",
    "x = xrange\n",
    "y = error_sij_abs\n",
    "ax.plot(x,y[:,1:], rasterized = True, color = 'dimgrey')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_21 bulk - s_21 pair')\n",
    "\n",
    "ax = axes[1]\n",
    "x = xrange\n",
    "y = error_si_abs\n",
    "ax.plot(x,y[:,1:], rasterized = True, color = 'dimgrey')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_2 bulk - s_21 pair')\n",
    "\n",
    "## center at zero\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "\n",
    "\n",
    "for ax in axes: \n",
    "    ax.set_xlabel('fraction of mutants in bulk competition: x')\n",
    "    ax.axhline(0.01, color = 'red', ls = 'dotted')\n",
    "    ax.axhline(-0.01, color = 'red', ls = 'dotted')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error II: How does error of frame of reference depend on mutant fraction?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### use a line plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET),sharey= True)\n",
    "\n",
    "ax = axes[0]\n",
    "x = xrange\n",
    "y = si_bulk_Bx - si1_bulk_Bx\n",
    "ax.plot(x,y[:,1:], rasterized = True, color = 'dimgrey')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_2 bulk - s_21 bulk')\n",
    "ax.set_xlabel('fraction of mutants in bulk competition: x')\n",
    "\n",
    "\n",
    "## center at zero\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "\n",
    "ax = axes[1]\n",
    "x = si_pair \n",
    "\n",
    "### plot high frequency\n",
    "ihigh,color_high = -1, 'dimgrey'\n",
    "y = si_bulk_Bx[ihigh] - si1_bulk_Bx[ihigh]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_high) \n",
    "### plot lower frequency\n",
    "ilow,color_low = 4, 'red'\n",
    "y = si_bulk_Bx[ilow] - si1_bulk_Bx[ilow]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_low) \n",
    "## visualize on first axis\n",
    "axes[0].axvline(xrange[ilow], color = color_low, ls = '--')\n",
    "## add horizontal line\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_2 bulk - s_21 bulk')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "### center on zero\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "\n",
    "### add prediction\n",
    "xmut = (1-xs[0])\n",
    "error_predicted = -xmut*si_bulk_B2.mean()\n",
    "print(error_predicted)\n",
    "ax.axhline(error_predicted, color = 'red', ls = 'dotted')\n",
    "\n",
    "\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error I: How do trait components  depend on mutant fraction?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## calculate the change in tau_bar\n",
    "\n",
    "tau_bar_range = np.zeros_like(xrange)\n",
    "\n",
    "for i in range(len(xrange)):\n",
    "    x = xrange[i] # read frequency\n",
    "    \n",
    "    ### set initial frequencies\n",
    "    xs = np.zeros_like(gs)\n",
    "    xs[1:] = x/k         # mutant lineages\n",
    "    xs[0] = 1-x              # wildtype population\n",
    "    \n",
    "\n",
    "    tau_bar_range[i] = np.sum(np.divide(xs,Ys)) / np.sum(np.divide(np.divide(xs,Ys),taus))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET,FIGHEIGHT_TRIPLET))\n",
    "\n",
    "tau_wt = taus[0]\n",
    "ax.plot(xrange,tau_bar_range, color = 'dimgrey')\n",
    "ax.axhline(taus[0], color = 'darkorange', label = 'wildtype tau')\n",
    "ax.set_xlabel('fraction of mutants: x')\n",
    "ax.set_ylabel('mean doubling time: tau_bar')\n",
    "\n",
    "ax.set_xlim(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### calculate the change in the coupling component\n",
    "\n",
    "n_subset = 300\n",
    "\n",
    "s_lag=np.zeros((len(xrange), len(gs[:n_subset])))\n",
    "s_growth=np.zeros((len(xrange), len(gs[:n_subset])))\n",
    "s_coupling=np.zeros((len(xrange), len(gs[:n_subset])))\n",
    "\n",
    "for i in range(len(xrange)):\n",
    "    x = xrange[i] # read frequency\n",
    "    \n",
    "    ### set initial frequencies\n",
    "    xs = np.zeros_like(gs)\n",
    "    xs[1:] = x/k         # mutant lineages\n",
    "    xs[0] = 1-x              # wildtype population\n",
    "    \n",
    "    for j in range(n_subset): #range(len(gs)):\n",
    "        s_lag[i,j], s_growth[i,j], s_coupling[i,j] = CalcApproxSijComponentsMultitype(j,0,xs=xs,gs=gs,ls=ls,nus=nus)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,2,figsize = (2*FIGWIDTH_TRIPLET,FIGHEIGHT_TRIPLET), sharey=True, sharex = True)\n",
    "\n",
    "ax = axes[0]\n",
    "y = s_lag + s_growth - si_pair_lag[:n_subset] - si_pair_growth[:n_subset]\n",
    "ax.plot(xrange,y[:,1:], color = 'dimgrey', marker = 'x')\n",
    "ax.set_xlabel('fraction of mutants: x')\n",
    "ax.set_ylabel('delta s_growth + delta s_lag')\n",
    "\n",
    "ax = axes[1]\n",
    "y = s_coupling\n",
    "ax.plot(xrange,y[:,1:], color = 'dimgrey', marker = 'x')\n",
    "ax.set_xlabel('fraction of mutants: x')\n",
    "ax.set_ylabel('s_21_coupling')\n",
    "\n",
    "#ax.set_xscale('log')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the frequency dependence as rotation and shrinking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey = True)\n",
    "\n",
    "\n",
    "ax = axes[0] # error from growth and lag component\n",
    "x = si_pair[:n_subset]\n",
    "\n",
    "### plot at high frequency\n",
    "ihigh, color_high = -1, 'dimgrey'\n",
    "y = s_lag[ihigh] + s_growth[ihigh] - si_pair_lag[:n_subset] - si_pair_growth[:n_subset]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_high)\n",
    "### plot at lower frequency\n",
    "ilow, color_low = 4, 'red'\n",
    "y = s_lag[ilow] + s_growth[ilow] - si_pair_lag[:n_subset] - si_pair_growth[:n_subset]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_low)\n",
    "\n",
    "### add horizontal line\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('delta s_growth + delta s_lag')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "\n",
    "ax = axes[1] # error from coupling component\n",
    "x = si_pair[:n_subset]\n",
    "\n",
    "### plot at high frequency\n",
    "y = s_coupling[ihigh]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_high)\n",
    "### plot at lower frequency\n",
    "y = s_coupling[ilow]\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_low)\n",
    "\n",
    "###\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('s_coupling: bulk experiment')\n",
    "ax.set_xlabel('s_21: pairwise experiment')\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate error bounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## calculate tau3\n",
    "## calculate error bound x_min = (tau3-tau1)/tau3 * \\theta/2\n",
    "## calculate g3 "
   ]
  }
 ],
 "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}