41_simulate_bulk_competition_warringer2003.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",
    "import seaborn as sns\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, CalcTotalSelectionCoefficientLog\n",
    "from m3_model import CalcRelativeSaturationTime as CalcSaturationTimeExact"
   ]
  },
  {
   "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 = 'all_traits_vary/'\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": [
    "### set colorscheme"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cmap = plt.get_cmap('tab20c')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_B1 = cmap(4.5/20)\n",
    "color_error1 = cmap(5.5/20) #cmap(13.5/20)\n",
    "color_error2 = cmap(7.5/20)\n",
    "\n",
    "color_stotal = cmap(16.5/20)\n",
    "color_B2 = cmap(0.5/20)\n",
    "color_ratio = cmap(13.5/20) #cmap(10.5/20)\n",
    "\n",
    "color_wt = cmap(10.5/20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "knockout_cmap = plt.get_cmap('flag')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rel_threshold = 0.01 ### threshold for relative error plots"
   ]
  },
  {
   "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",
    "### lag times\n",
    "ls = np.zeros(n_knockouts+1)\n",
    "ls[0] = WILDTYPE['lag']\n",
    "ls[1:] = df_knockouts['lag'].values\n",
    "\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:] = 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 # default value: 1e6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xs_pair, xs_pair_final, tsats, _,_,_ = run_pairwise_experiment(gs=gs,ls=ls,nus = nus, g1=gs[0],l1=ls[0],nu1=nus[0],\n",
    "                                                               x0 = 1/N)\n",
    "\n",
    "s_pair = CalcTotalSelectionCoefficientLogit(xs_pair,xs_pair_final)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot  frequency trajectory in pairwise competition for example mutant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## set axis limits\n",
    "TMIN,TMAX = 0,15\n",
    "XMIN,XMAX = 5e-8,2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## set timepoints for each trajectory\n",
    "tvec = np.array([0, TMAX])\n",
    "## set frequency values for each trajectory\n",
    "trajectories = np.vstack((xs_pair,xs_pair_final))\n",
    "\n",
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "ax.plot(tvec,trajectories[:,1:], color = 'dimgrey', label = 'mutants')\n",
    "ax.plot(tvec,trajectories[:,0], color = 'orange', label = 'wild-type')\n",
    "\n",
    "ax.set_xlabel('time')\n",
    "ax.set_ylabel('frequency')\n",
    "\n",
    "# fix axis limits\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(XMIN,XMAX)\n",
    "ax.set_xlim(TMIN,TMAX)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot distribution of fitness effects in the pairwise competition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "plt.hist(s_pair, bins = 42, color = 'dimgrey')\n",
    "ax.axvline(s_pair[0], color = 'orange')\n",
    "ax.axvline(s_pair[1:].mean(), color = 'tab:red', ls = '--')\n",
    "ax.set_xlabel('s_21 in pairwise experiment')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate bulk competition with background mutants added at invasion frequency (scenario B1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frequencies for this scenario can be summarized as \n",
    "\n",
    "    frequency of the mutants:             xi = 1/N\n",
    "    frequency of the wildtype strain:     x1 = 1 - x2 - x3\n",
    "\n",
    "where each mutant lineage is added at the frequency used in the pairwise competition. "
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### set initial frequencies\n",
    "xs = np.zeros_like(gs)\n",
    "xs[1:] = 1e-6                  # mutant lineages\n",
    "xs[0] = 1- xs[1:].sum()       # wildtype population\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### set initial frequencies\n",
    "mutant_library_freq = 0.01\n",
    "\n",
    "xs = np.zeros_like(gs)\n",
    "n_knockouts = len(xs[1:])\n",
    "xs[1:] = mutant_library_freq/n_knockouts                  # mutant lineages\n",
    "xs[0] = 1- xs[1:].sum()       # wildtype population\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutant_ratio_B1 = xs[1:].sum()\n",
    "\n",
    "print(\"Proportion of mutants: %.8f \" % mutant_ratio_B1 )\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",
    "## calculate total selection coefficient\n",
    "sigma_bulk_B1 = 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",
    "s_bulk_B1 = CalcTotalSelectionCoefficientLogit(xi1,xi1_final)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot frequency trajectory for all barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## set timepoints for each trajectory\n",
    "tvec = np.array([0, tsat])\n",
    "## set frequency values for each trajectory\n",
    "trajectories = np.vstack((xs,xs_final))\n",
    "\n",
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "ax.plot(tvec,trajectories[:,1:], color = 'dimgrey', label = 'mutants')\n",
    "ax.plot(tvec,trajectories[:,0], color = 'orange', label = 'wild-type')\n",
    "\n",
    "ax.set_xlabel('time')\n",
    "ax.set_ylabel('frequency')\n",
    "\n",
    "# fix axis limits\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(XMIN,XMAX)\n",
    "ax.set_xlim(TMIN,tsat)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Calculate error to pairwise competition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "y = sigma_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_stotal, alpha = 1, \n",
    "           label = r'total scoeff. $\\sigma$', marker = 'o') \n",
    "\n",
    "y = s_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B1, \n",
    "           label = 'pairwise scoeff. $s$', marker = 'o') \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('absolute error')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'lower right', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_bulk_competition_B.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "\n",
    "y = np.abs(np.divide( sigma_bulk_B1 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_stotal, alpha = 1, \n",
    "           label = r'total scoeff. $\\sigma$', marker = 'o') \n",
    "\n",
    "y = np.abs(np.divide( s_bulk_B1 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_B1, \n",
    "           label = 'pairwise scoeff. $s$', marker = 'o') \n",
    "\n",
    "\n",
    "ax.axhline(rel_threshold, ls = '--', color = 'black')\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_yscale('log')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "ax.set_ylabel('relative error')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'upper left', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'relative_error_bulk_competition_B.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare predictions of the ranking"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "Observation for the comparison of total and relative fitness in bulk: \n",
    "    - the rankings are almost identical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.DataFrame()\n",
    "## load data on selection coefficients\n",
    "data['sigma_bulk_B1'] = sigma_bulk_B1\n",
    "data['s_bulk_B1'] = s_bulk_B1\n",
    "data['s_pair'] = s_pair\n",
    "## exclude wild-type, has different offset\n",
    "data= data.iloc[1:]\n",
    "## create rank variables, fittest mutants come first\n",
    "data['sigma_bulk_B1_rank'] = data['sigma_bulk_B1'].rank(ascending = True,method = 'min') \n",
    "data['s_bulk_B1_rank'] = data['s_bulk_B1'].rank(ascending = True, method = 'min') \n",
    "data['s_pair_rank'] = data['s_pair'].rank(ascending = True,method = 'min') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2.5*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey=True)\n",
    "\n",
    "\n",
    "x = s_pair[1:]\n",
    "\n",
    "ax = axes[0] \n",
    "## plot ranking total bulk vs pairwise bulk\n",
    "y = data['sigma_bulk_B1_rank'] - data['s_bulk_B1_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = 'silver', alpha = 1, marker = 'o') \n",
    "ax.set_ylabel('rank difference: total fitness in bulk\\nvs. pairwise fitness in bulk')\n",
    "#ax.set_ylabel('rank difference between: \\n total fit. (bulk) - pairwise comp.')\n",
    "ax = axes[1] \n",
    "## plot ranking total bulk vs pairwise\n",
    "y = data['s_bulk_B1_rank'] - data['s_pair_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = color_B1, alpha = 1, marker = 'o')\n",
    "ax.set_ylabel('rank difference: pairwise fitness in bulk\\nvs. pairwise fitness in pairwise competition')\n",
    "ax.tick_params(labelleft=True)\n",
    "\n",
    "\n",
    "## plot ranking pairwise bulk vs pairwisse\n",
    "#y = data['s_bulk_B1_rank'] - data['s_pair_rank']\n",
    "#ax.scatter(x,y, rasterized = True, color = 'dimgrey', alpha = 1, marker = 'o') \n",
    "#ax.set_ylabel('rank difference between: \\n pairwise fit. (bulk) - pairwise comp.')\n",
    "\n",
    "\n",
    "fig.subplots_adjust( wspace= 0.5)\n",
    "\n",
    "for ax in axes: \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",
    "\n",
    "    ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'ranking_mismatch_bulk_B1.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "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",
    "s_bulk_growth = np.zeros_like(s_bulk_B1)\n",
    "s_bulk_lag = np.zeros_like(s_bulk_B1)\n",
    "s_bulk_coupling = np.zeros_like(s_bulk_B1)\n",
    "\n",
    "for i in range(len(gs)):\n",
    "    s_bulk_growth[i], s_bulk_lag[i], s_bulk_coupling[i] = CalcApproxSijComponentsMultitype(i,0,xs,gs,ls,nus)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "s_pair_growth = np.zeros_like(s_pair)\n",
    "s_pair_lag = np.zeros_like(s_pair)\n",
    "s_pair_coupling = np.zeros_like(s_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",
    "        s_pair_growth[i], s_pair_lag[i], s_pair_coupling[i] =CalcApproxSijComponentsMultitype(1,0,\n",
    "                                                xs=[1-x0,x0], gs = [g1,g2], ls= [l1,l2], nus = [nu1,nu2] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## calculate error components\n",
    "\n",
    "delta_s_lag = s_bulk_lag - s_pair_lag\n",
    "delta_s_growth = s_bulk_growth - s_pair_growth\n",
    "delta_s_one = delta_s_lag + delta_s_growth\n",
    "delta_s_two = s_bulk_coupling"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot the error components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "y = delta_s_one\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_error1, zorder = 2, \n",
    "           label = r'linear error')\n",
    "\n",
    "y = delta_s_two\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_error2, \n",
    "           label = r'quadratic error')\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('absolute error')#\\n(pairwise relative fitness)')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'lower right', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_components_bulk_competition_B.PDF',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "y = np.abs(np.divide(delta_s_one, s_pair, where = s_pair!=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_error1, zorder = 2, \n",
    "           label = r'linear error')\n",
    "\n",
    "y = np.abs(np.divide(delta_s_two, s_pair, where = s_pair!=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_error2, \n",
    "           label = r'quadratic error')\n",
    "\n",
    "\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('relative error')# pairwise fitness\\n(bulk competition B)')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend( frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "ax.set_yscale('log')\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'relative_error_components_bulk_competition_B.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate traits of the coarse-grained population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlib = xs[1:].sum()\n",
    "Ylib = 1/(np.sum(xs[1:]/Ys[1:])/xlib)\n",
    "nulib = CalcRelativeYield(Ylib, R0 = CONCENTRATION_GLUCOSE, N0 = OD_START)\n",
    "\n",
    "taus = 1/gs\n",
    "taulib = 1/(np.sum(xs[1:]/Ys[1:]/taus[1:])/(xlib/Ylib))\n",
    "glib = 1/taulib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute population lag time\n",
    "lambda_lib = np.sum(ls[1:]*xs[1:]/Ys[1:]/taus[1:])*(taulib*Ylib/xlib)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_lib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nus[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculate traits of the leave-one-out population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute the leave-one-out library abundance\n",
    "xlib_no =np.zeros_like(xs)\n",
    "xlib_no[0] = xlib\n",
    "xlib_no[1:] = xlib-xs[1:]\n",
    "\n",
    "### compute the leave-one-out biomas yield\n",
    "Ylib_no = np.zeros_like(Ys)\n",
    "Ylib_no[0] = Ylib\n",
    "n = len(Ylib_no)\n",
    "for i in range(1,n):\n",
    "    tmp = 0\n",
    "    for j in range(1, n):\n",
    "        if j==i: continue # skip\n",
    "        else: tmp += xs[j]/Ys[j]\n",
    "    Ylib_no[i] = 1/(tmp/xlib_no[i])\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute the individual leave-one-out growth rate\n",
    "taulib_no = np.zeros_like(gs)\n",
    "taulib_no[0]= taulib\n",
    "\n",
    "for i in range(1, len(gs)):\n",
    "    tmp = 0\n",
    "    for j in range(1, len(gs)):\n",
    "        if j==i: continue # skip\n",
    "        else: tmp+= xs[j]/Ys[j]/taus[j]\n",
    "    taulib_no[i] = 1/(tmp/(xlib_no[i]/Ylib_no[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate the leave-one-out lag time\n",
    "\n",
    "lambdalib_no = np.zeros_like(gs) # prepare data container\n",
    "\n",
    "n = len(lambdalib_no)\n",
    "\n",
    "for i in range(1,n):\n",
    "    ## calculate prefactor\n",
    "    delta_tau = taus[i]- taus[0]\n",
    "    \n",
    "    ## calculate the sum over all other strains in the librar\n",
    "    tmp = 0\n",
    "    for j in range(1, n):\n",
    "        if j==i: continue # skip\n",
    "        else: tmp += xs[j]/(Ys[j]*taus[j]) * ((taus[i]-taus[j])*(ls[j]-ls[0]) - (ls[i]-ls[j])*(taus[j]-taus[0]) )\n",
    "\n",
    "    ## calculate first term\n",
    "    the_bracket = (Ylib_no[i]*taulib_no[i]/xlib_no[i]) *tmp\n",
    "    first_term = 1/delta_tau * the_bracket\n",
    "    ## calcualte the second termb\n",
    "    second_term = 1/delta_tau * ( (taus[i]-taulib_no[i])*ls[0] + (taulib_no[i]-taus[0])*ls[i] )\n",
    "    lambdalib_no[i]= first_term + second_term\n",
    "\n",
    "### define lib value for lag time \n",
    "lambdalib = lambdalib_no[1:].mean()\n",
    "lambdalib_no[0] = lambdalib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,3, figsize = (4*0.8*FIGWIDTH_TRIPLET, 0.8*FIGHEIGHT_TRIPLET))\n",
    "\n",
    "ax = axes[0]\n",
    "y = Ylib_no[1:] \n",
    "ax.hist(y, bins = 42, color = 'dimgrey', log= True)\n",
    "## add mean\n",
    "mean = Ylib_no[1:].mean()\n",
    "label = r\"mean value: $Y_{\\mathrm{lib}/i}$\" + f\"={mean:.6f}\"\n",
    "ax.axvline(mean, color = 'tab:red',label = label)\n",
    "ax.legend(frameon=False)\n",
    "## add information in title\n",
    "title = r'wild-type $Y_{\\mathrm{wt}}$' + f'={Ys[0]:.6f}'\n",
    "title += ', '+ r'library $Y_{\\mathrm{lib}}$' + f'={Ylib_no[0]:.6f}'\n",
    "ax.set_title(title, loc = 'right')\n",
    "ax.set_xlabel(r'background biomass yield $Y_{\\mathrm{lib}/i}$ [OD/mM glucose]')\n",
    "\n",
    "\n",
    "ax = axes[1]\n",
    "y = taulib_no[1:]\n",
    "ax.hist(y, bins = 42, color = 'dimgrey', log= True)\n",
    "## add mean\n",
    "mean = taulib_no[1:].mean()\n",
    "label = r\"mean value: $\\tau_{\\mathrm{lib}/i}$\" + f\"={mean:.5f}\"\n",
    "ax.axvline(mean, color = 'tab:red', label = label)\n",
    "ax.legend(frameon=False)\n",
    "## add information in title\n",
    "title = r'wild-type $\\tau_{\\mathrm{wt}}$' + f'={taus[0]:.5f}'\n",
    "title += ', '+ r'library $\\tau_{\\mathrm{lib}}$' + f'={taulib_no[0]:.5f}'\n",
    "ax.set_title(title, loc = 'right')\n",
    "ax.set_xlabel(r'background doubling time $\\tau_{\\mathrm{lib}/i}$ [hours]')\n",
    "\n",
    "ax = axes[2]\n",
    "y = lambdalib_no[1:] \n",
    "ax.hist(y, bins = 42, color = 'dimgrey', log= True)\n",
    "## add mean\n",
    "mean = lambdalib_no[1:].mean()\n",
    "label = r\"mean value: $\\lambda_{\\mathrm{lib}/i}$\"+ f\"={mean:.5f}\"\n",
    "ax.axvline(mean, color = 'tab:red', label = label)\n",
    "ax.legend(frameon=False)\n",
    "## add information in title\n",
    "title = r'wild-type $\\lambda_{\\mathrm{wt}}$' + f'={ls[0]:.5f}'\n",
    "ax.set_title(title, loc = 'right')\n",
    "\n",
    "\n",
    "ax.set_xlabel(r'background lag time $\\lambda_{\\mathrm{lib}/i}$ [hours]')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "taulib_no.min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## doubling time of the library\n",
    "taulib*np.log(2)*60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## doubling time of the wild-type\n",
    "taus[0]*np.log(2)*60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### analyze the outliers\n",
    "\n",
    "data = pd.DataFrame()\n",
    "data['taulib_no'] = taulib_no\n",
    "data['lambdalib_no'] = lambdalib_no\n",
    "data['Ylib_no'] = lambdalib_no\n",
    "\n",
    "index = data.loc[1:,'lambdalib_no'].idxmin()\n",
    "\n",
    "print(f\"Outlier with lower background lag time: index = {index}\")\n",
    "print('biomass yield: %.6f' % Ys[index])\n",
    "print(\"lag time: %.6f\" %ls[index])\n",
    "\n",
    "index = data.loc[1:,'taulib_no'].idxmin()\n",
    "\n",
    "print(f\"Outlier with lower background doubling time: index = {index}\")\n",
    "print('biomass yield: %.6f' % Ys[index])\n",
    "print(\"doubling time: %.6f\" %taus[index])\n",
    "\n",
    "\n",
    "index = data.loc[1:,'Ylib_no'].idxmax()\n",
    "\n",
    "print(f\"Outlier with higher background biomass yield: index = {index}\")\n",
    "print('biomass yield: %.6f' % Ys[index])\n",
    "print(\"doubling time: %.6f\" %taus[index])\n",
    "print(\"lag time: %.6f\" %ls[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Compare to wild-type traits\")\n",
    "print(\"biomass yield %.6f OD/mM glucose\" %Ys[0])\n",
    "print(\"doubling time %.6f hours\" %taus[0])\n",
    "print(\"lag time %.6f hours\" %ls[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## compare to bulk competition traits\n",
    "Ybar_bulk  = 1/np.sum(xs/Ys)/xs.sum()\n",
    "taubar_bulk = 1/np.sum(xs/Ys/taus)/Ybar_bulk\n",
    "\n",
    "\n",
    "\n",
    "print(\"Compare to bulk competition traits\")\n",
    "print(\"biomass yield %.6f OD/mM glucose\" %Ybar_bulk)\n",
    "print(\"doubling time %.6f hours\" %taubar_bulk)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare to triplet traits\n",
    "Ybar_triplet = 1/((xs[0]/Ys[0] + xlib/Ylib)/(xs[0]+xlib))\n",
    "taubar_triplet = 1/(xs[0]/Ys[0]/taus[0] + xlib/Ylib/taulib)/Ybar_triplet\n",
    "\n",
    "print(\"Compare to  triplet competition traits\")\n",
    "print(\"biomass yield %.6f OD/mM glucose\" %Ybar_triplet)\n",
    "print(\"doubling time %.6f hours\" %taubar_triplet)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# calculate terms that are independent of the knockout\n",
    "\n",
    "taus = 1/gs\n",
    "\n",
    "lstars = np.zeros_like(gs) # prepare data container\n",
    "lstars[0] = ls[0]\n",
    "\n",
    "for i in range(1,len(gs)):\n",
    "    ## calculate prefactor\n",
    "    delta_tau = taus[i]- taus[0]\n",
    "    \n",
    "    ## calculate first term\n",
    "    the_summands = xs/(Ys*taus) * ( (taus[i]-taus)*(ls-ls[0]) - (ls[i]-ls)*(taus-taus[0]) )\n",
    "    the_bracket = Ystar*taustar/xstar *the_summands.sum()\n",
    "    first_term = 1/delta_tau * the_bracket\n",
    "    second_term = 1/delta_tau * ( (taus[i]-taustar)*ls[0] + (taustar-taus[0])*ls[i] )\n",
    "    lstars[i]= first_term + second_term"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### plot a histogramm of the distribution of effective lag times\n",
    "\n",
    "lstar = lstars[1:].mean()\n",
    "y = lstars[1:] - lstar\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "plt.hist(y, bins = 42, color = 'dimgrey')\n",
    "#ax.axvline(ls[0], color = 'orange')\n",
    "ax.axvline(0, color = 'tab:red', ls = '--')\n",
    "ax.set_xlabel(r'deviation from average $\\lambda_3$')\n",
    "ax.set_yscale('log')\n",
    "ax.set_title(r\"average $\\lambda_3=$\" + f\"{lstar:.4f}h\" + \", \" \n",
    "             + r\"wild-type $\\lambda_1=$\" + f\"{lstars[0]:.4f}h\", loc = 'right')\n",
    "ax.ticklabel_format(axis = 'x', useOffset = False)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### also calculate as relative yield\n",
    "nulib_no = CalcRelativeYield(Ylib_no, R0 = CONCENTRATION_GLUCOSE, N0 = OD_START)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "s_triplet_growth = np.zeros_like(s_bulk_growth)\n",
    "s_triplet_lag = np.zeros_like(s_bulk_growth)\n",
    "s_triplet_coupling = np.zeros_like(s_bulk_growth)\n",
    "\n",
    "\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",
    "        #g3, l3, nu3 = 1/taulib, lambdalib, nulib  # use traits of the library as a whole\n",
    "        g3, l3, nu3 = 1/taulib_no[i], lambdalib_no[i], nulib_no[i] # use leave-one-out traits\n",
    "        xs_triplet =[1-xlib,1/n_knockouts, xlib - 1/n_knockouts]\n",
    "        s_triplet_growth[i], s_triplet_lag[i], s_triplet_coupling[i] =CalcApproxSijComponentsMultitype(1,0,\n",
    "                                                xs=xs_triplet, gs = [g1,g2,g3], ls= [l1,l2,l3], nus = [nu1,nu2,nu3] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1,3, figsize = (3.*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "\n",
    "ax = axes[0]\n",
    "ax.scatter(s_bulk_growth, s_triplet_growth-s_bulk_growth)\n",
    "ax.set_ylabel(\"s_growth: triplet - bulk\")\n",
    "ax = axes[1]\n",
    "ax.scatter(s_bulk_lag, s_triplet_lag-s_bulk_lag)\n",
    "ax.set_ylabel(\"s_lag: triplet - bulk\")\n",
    "ax = axes[2]\n",
    "ax.scatter(s_bulk_coupling, s_triplet_coupling-s_bulk_coupling)\n",
    "ax.set_ylabel(\"s_coupling: triplet - bulk\")\n",
    "\n",
    "for ax in axes:\n",
    "    ax.axhline(0, ls = '--', color = 'black')\n",
    "    ymin,ymax = ax.get_ylim()\n",
    "    ylim = np.max(np.abs([ymin,ymax]))\n",
    "    ax.set_ylim(-ylim,ylim)\n",
    "    \n",
    "    ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare to Michael's formula for the effective lag time\n",
    "\n",
    "\\begin{align}\n",
    "\\lambda_{MM} = - \\frac{\\sum_{j=1}^n \\exp(-\\lambda_j g_j)}{\\sum_{j=1}^n g_j \\exp(-\\lambda_j g_j)}\n",
    "\\log \\left( \\sum_{j=1}^n \\exp(-\\lambda_j g_j) \\right)\n",
    "\\end{align}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## \n",
    "\n",
    "nominator = np.sum(xs[1:]/xlib*np.exp(-ls[1:]*gs[1:]))\n",
    "denominator = np.sum(xs[1:]/xlib*gs[1:]* np.exp(-ls[1:]*gs[1:]))\n",
    "log_term = np.log(nominator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_MM = -nominator/denominator * log_term"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambda_MM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambdalib"
   ]
  },
  {
   "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": "raw",
   "metadata": {},
   "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": [
    "### set initial frequencies\n",
    "mutant_library_freq = 0.9999\n",
    "\n",
    "xs = np.zeros_like(gs)\n",
    "n_knockouts = len(xs[1:])\n",
    "xs[1:] = mutant_library_freq/n_knockouts                  # mutant lineages\n",
    "xs[0] = 1- xs[1:].sum()       # wildtype population\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mutant_ratio_B2 = xs[1:].sum()\n",
    "\n",
    "print(\"Proportion of mutants: %.8f \" % mutant_ratio_B2 )\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",
    "## calculate total selection coefficient\n",
    "sigma_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",
    "s_bulk_B2 = CalcTotalSelectionCoefficientLogit(xi1,xi1_final)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot frequency trajectory for all barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## set timepoints for each trajectory\n",
    "tvec = np.array([0, tsat])\n",
    "## set frequency values for each trajectory\n",
    "trajectories = np.vstack((xs,xs_final))\n",
    "\n",
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "ax.plot(tvec,trajectories[:,1:], color = 'dimgrey', label = 'mutants')\n",
    "ax.plot(tvec,trajectories[:,0], color = 'orange', label = 'wild-type')\n",
    "\n",
    "ax.set_xlabel('time')\n",
    "ax.set_ylabel('frequency')\n",
    "\n",
    "# fix axis limits\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(XMIN,XMAX)\n",
    "ax.set_xlim(TMIN,tsat)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGHEIGHT_TRIPLET, 0.3* FIGHEIGHT_TRIPLET))\n",
    "                       \n",
    "ax.tick_params(left = False, labelleft = False, bottom = False, labelbottom=False)\n",
    "\n",
    "n = len(xs)\n",
    "colors = [knockout_cmap(i/k+1) for i in range(n)]\n",
    "colors[0] = color_wt\n",
    "for i in range(n):\n",
    "    if i == 0: y1 = [0,0]\n",
    "    else: y1 = [xs[i-1], xs_final[i-1]]\n",
    "    \n",
    "    if i == n-1: y2 = [1,1]\n",
    "    else: y2 = [xs[i], xs_final[i]]\n",
    "        \n",
    "    ax.fill_between(x=[0,tsat], y1=y1,y2=y2)\n",
    "    \n",
    "ax.set_xlim(0,tsat)\n",
    "ax.set_ylim(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "y =  s_bulk_B2\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B2, zorder = 2,\n",
    "           label = f'mutant ratio:{mutant_ratio_B2*100:.2f}%', marker = 'o') \n",
    "\n",
    "y =  s_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B1, zorder = 5,\n",
    "           label = f'mutant ratio:{mutant_ratio_B1*100:6.2f}%', marker = 'o') \n",
    "\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('absolute error')# pairwise fitness\\n(different bulk competitions)')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'lower right', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_bulk_competition_comparison.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "\n",
    "y = np.abs(np.divide( s_bulk_B1 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_B1, zorder = 5, \n",
    "           label = f'mutant ratio:{mutant_ratio_B1*100:6.2f}%', marker = 'o')   \n",
    "\n",
    "y = np.abs(np.divide( s_bulk_B2 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_B2, zorder = 2, \n",
    "           label = f'mutant ratio:{mutant_ratio_B2*100:6.2f}%', marker = 'o')      \n",
    "\n",
    "\n",
    "\n",
    "ax.axhline(rel_threshold, ls = '--', color = 'black')\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_yscale('log')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "ax.set_ylabel('relative error') #\\n(bulk competition B)')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'upper left', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'relative_error_bulk_competition_comparison.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = np.abs(s_pair)\n",
    "\n",
    "## plot data\n",
    "\n",
    "y = np.abs(np.divide( s_bulk_B1 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_B1, zorder = 5, \n",
    "           label = f'mutant ratio:{mutant_ratio_B1*100:6.2f}%', marker = 'o')   \n",
    "\n",
    "y = np.abs(np.divide( s_bulk_B2 - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_B2, zorder = 2, \n",
    "           label = f'mutant ratio:{mutant_ratio_B2*100:6.2f}%', marker = 'o')      \n",
    "\n",
    "\n",
    "\n",
    "ax.axhline(rel_threshold, ls = '--', color = 'black')\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_yscale('log')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "ax.set_ylabel('relative error')\n",
    "ax.set_xlabel('mutant fitness magnitude\\n in pairwise competition')\n",
    "#ax.legend(loc = 'upper left', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'relative_error_bulk_competition_comparison_abs.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot with inset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "y = sigma_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_stotal, alpha = 1, \n",
    "           label = r'total scoeff. $\\sigma$', marker = 'o') \n",
    "\n",
    "y = s_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B1, \n",
    "           label = 'pairwise scoeff. $s$', marker = 'o') \n",
    "\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "yabs = 1.5*yabs\n",
    "ax.set_ylim(-yabs,yabs)\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel('absolute error')\n",
    "ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "#ax.legend(loc = 'lower right', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "### add inset\n",
    "\n",
    "inset = ax.inset_axes([0.75, 0.045, 0.38, 0.38])\n",
    "\n",
    "\n",
    "ax = inset.twinx()\n",
    "\n",
    "## plot data\n",
    "y =  s_bulk_B2\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B2, zorder = 2,\n",
    "           label = f'mutant ratio:{mutant_ratio_B2*100:.2f}%', marker = 'o') \n",
    "\n",
    "y =  s_bulk_B1\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_B1, zorder = 5,\n",
    "           label = f'mutant ratio:{mutant_ratio_B1*100:6.2f}%', marker = 'o') \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.tick_params(labelbottom = False, bottom = False)\n",
    "\n",
    "ax = inset\n",
    "ax.tick_params(labelleft =False, left = False, bottom = False, labelbottom = False)\n",
    "## add labels for inset\n",
    "ax.set_ylabel('abs. error')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_bulk_competition_B_with_inset.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compare predictions of the ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.DataFrame()\n",
    "## load data on selection coefficients\n",
    "data['sigma_bulk_B1'] = sigma_bulk_B1\n",
    "data['s_bulk_B1'] = s_bulk_B1\n",
    "data['sigma_bulk_B2'] = sigma_bulk_B2\n",
    "data['s_bulk_B2'] = s_bulk_B2\n",
    "data['s_pair'] = s_pair\n",
    "\n",
    "## exclude wild-type, has different offset\n",
    "data= data.iloc[1:]\n",
    "## create rank variables, fittest mutants come first\n",
    "data['sigma_bulk_B1_rank'] = data['sigma_bulk_B1'].rank(ascending = True, method = 'min') \n",
    "data['s_bulk_B1_rank'] = data['s_bulk_B1'].rank(ascending = True, method = 'min') \n",
    "data['sigma_bulk_B2_rank'] = data['sigma_bulk_B2'].rank(ascending = True, method = 'min') \n",
    "data['s_bulk_B2_rank'] = data['s_bulk_B2'].rank(ascending = True, method = 'min')\n",
    "data['s_pair_rank'] = data['s_pair'].rank(ascending = True, method = 'min') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2.5*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True, sharey=True)\n",
    "\n",
    "\n",
    "x = s_pair[1:]\n",
    "\n",
    "ax = axes[0] \n",
    "## plot ranking total bulk vs pairwise bulk\n",
    "y = data['sigma_bulk_B1_rank'] - data['s_bulk_B1_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = color_B1, alpha = 1, marker = 'o', zorder = 5) \n",
    "y = data['sigma_bulk_B2_rank'] - data['s_bulk_B2_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = color_B2, alpha = 1, marker = 'o', zorder = 2)\n",
    "ax.set_ylabel('rank difference: total fitness in bulk\\nvs. pairwise fitness in bulk')\n",
    "ax = axes[1] \n",
    "## plot ranking total bulk vs pairwise\n",
    "y = data['s_bulk_B1_rank'] - data['s_pair_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = color_B1, alpha = 1, marker = 'o',zorder = 5)\n",
    "y = data['s_bulk_B2_rank'] - data['s_pair_rank']\n",
    "ax.scatter(x,y, rasterized = True, color = color_B2, alpha = 1, marker = 'o', zorder =2)\n",
    "ax.set_ylabel('rank difference: pairwise fitness in bulk\\nvs. pairwise fitness in pairwise')\n",
    "ax.tick_params(labelleft=True)\n",
    "\n",
    "\n",
    "## plot ranking pairwise bulk vs pairwisse\n",
    "#y = data['s_bulk_B1_rank'] - data['s_pair_rank']\n",
    "#ax.scatter(x,y, rasterized = True, color = 'dimgrey', alpha = 1, marker = 'o') \n",
    "#ax.set_ylabel('rank difference between: \\n pairwise fit. (bulk) - pairwise comp.')\n",
    "\n",
    "\n",
    "fig.subplots_adjust( wspace= 0.5)\n",
    "\n",
    "for ax in axes: \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",
    "\n",
    "    ax.set_xlabel('mutant fitness in pairwise competition')\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'ranking_mismatch_bulk_comparison.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Find the optimum frequency for measuring bulk competition"
   ]
  },
  {
   "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.linspace(1/N,0.99, num = 10)\n",
    "xrange = np.geomspace(1e-6,.99999, num = 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time \n",
    "\n",
    "sigma_bulk_Bx = np.zeros((len(xrange),n_knockouts+1))\n",
    "s_bulk_Bx = np.zeros((len(xrange),n_knockouts+1))\n",
    "\n",
    "\n",
    "for i in range(len(xrange)):\n",
    "    xlib = xrange[i]\n",
    "    \n",
    "    ### set initial frequencies\n",
    "    xs = np.zeros_like(gs)\n",
    "    xs[1:] = xlib/n_knockouts         # mutant lineages\n",
    "    xs[0] = 1-xlib              # 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",
    "    sigma_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",
    "    s_bulk_Bx[i,:] = CalcTotalSelectionCoefficientLogit(xi1,xi1_final)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xrange"
   ]
  },
  {
   "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)),s_pair) # need the right shape\n",
    "error_abs = s_bulk_Bx - truth \n",
    "error_rel = np.divide(error_abs,truth, where = truth != 0)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "## calculate simple error bound\n",
    "\n",
    "n_knockouts = len(gs[1:])\n",
    "Y3 = 1/np.sum(1/n_knockouts/Ys[1:])\n",
    "tau3 = 1/(np.sum(1/n_knockouts/(Ys[1:]/Y3)/taus[1:]))\n",
    "g3 = 1/tau3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### define traits of the coarse-grained population\n",
    "xlib = xs[1:].sum()\n",
    "Ylib = 1/(np.sum(xs[1:]/Ys[1:])/xlib)\n",
    "nulib = CalcRelativeYield(Ylib, R0 = CONCENTRATION_GLUCOSE, N0 = OD_START)\n",
    "\n",
    "taus = 1/gs\n",
    "taulib = 1/(np.sum(xs[1:]/Ys[1:]/taus[1:])/(xlib/Ylib))\n",
    "glib = 1/taulib\n",
    "lambda_lib = np.sum(ls[1:]*xs[1:]/Ys[1:]/taus[1:])*(taulib*Ylib/xlib)\n",
    "\n",
    "### Calculate more complex error bound\n",
    "## define wild-type traits\n",
    "Ywt = Ys[0]\n",
    "xwt = xs[0]\n",
    "lwt = ls[0]\n",
    "gwt = gs[0]\n",
    "tauwt = 1/gwt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"lag times\")\n",
    "print(\"library: %.3f\" % lambda_lib)\n",
    "print(\"wild-type: %.3f\" % lwt)\n",
    "\n",
    "print(\"doubling times\")\n",
    "print(\"library: %.3f\" % taulib)\n",
    "print(\"wild-type: %.3f\" % tauwt)\n",
    "\n",
    "print(\"growth rates\")\n",
    "print(\"library: %.3f\" % glib)\n",
    "print(\"wild-type: %.3f\" % gwt)\n",
    "\n",
    "\n",
    "print(\"biomass yield\")\n",
    "print(\"library: %.4f\" % Ylib)\n",
    "print(\"wild-type: %.4f\" % Ywt)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### specific relative error threshold\n",
    "theta_rel = 0.01"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### calculate simple error bound\n",
    "\n",
    "\n",
    "x_lib_max = gwt/np.abs(glib-gwt) * theta_rel\n",
    "print(f\"The maximum tolerated mutant library frequency from our equation is: {x_lib_max}\")\n",
    "print(f\"mutant library growth rate: {glib}\")\n",
    "print(f\"wild-type growth rate: {gwt}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### calculate complex error bound\n",
    "\n",
    "## find maximum Bj value\n",
    "LFCpairs = np.log(nus)\n",
    "Bj = np.abs(LFCpairs[1:] + (ls[1:]-lwt)/(taus[1]-tauwt))*tauwt\n",
    "Bj_max = Bj.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### calculate specific specific maximum library abundnace\n",
    "\n",
    "assert glib < gwt, \"The formula assumes this case!\"\n",
    "\n",
    "prefactor = gwt/((-1)*(glib-gwt))\n",
    "eps = theta_rel\n",
    "spair_lib_j = np.abs(-(lambda_lib-lwt)*glib - (taulib-tauwt)*LFCpairs[1:]*glib)\n",
    "libfactor = spair_lib_j *glib/((-1)*(glib-gwt))\n",
    "xlib_max_j =  prefactor * (eps/ (eps + libfactor * Bj))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlib_max_j.min()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LFC_value = np.log(100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert glib < gwt, \"The formula assumes this case!\"\n",
    "\n",
    "spair_lib = -(lambda_lib-lwt)*glib - (taulib-tauwt)*LFC_value*glib\n",
    "libfactor = np.abs(spair_lib) *glib/np.abs(glib-gwt)\n",
    "xlib_star =  prefactor * (eps/ (eps + libfactor * Bj_max))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xlib_star"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"selection coefficient:\")\n",
    "print(\"spair_lib: %.3f\" % spair_lib)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### collect plot data into a Dataframe\n",
    "# we want rows to be the genotyes, columns to be the timepoints\n",
    "df_abs = pd.DataFrame(data= np.abs(error_rel.T), columns = xrange)\n",
    "df_abs['s_pair'] = s_pair\n",
    "df_abs = df_abs.melt(id_vars = 's_pair', var_name = 'xlib', value_name = 'error_rel', ignore_index = False,)\n",
    "df_abs = df_abs.reset_index(drop = False)\n",
    "df_abs = df_abs.sort_values('s_pair', ascending = True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.colors import TwoSlopeNorm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### interpolate relative abundance that gives 1% from data\n",
    "df_search = pd.DataFrame(data= np.abs(error_rel.T), columns = xrange)\n",
    "high_error_mutants = set(df_search.idxmax(axis = 0).values)\n",
    "high_error_mutants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xp = df_search.columns\n",
    "fp = df_search.loc[3407]\n",
    "def fun(x):\n",
    "    return np.square(np.interp(x, xp =xp, fp = fp) - eps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fun(xlib_star)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize_scalar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "minimize_scalar(fun, [0.0001,0.01])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_abs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_abs['error_rel'].idxmax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### choose graphical settings\n",
    "\n",
    "# possible palettes to try: jet, nipy_spectral, icefire, coolwarm, viridis, crest\n",
    "cmap = sns.color_palette('viridis', as_cmap=True)\n",
    "#norm = plt.Normalize(-1, 1, clip = True) # basic norm, sames lope on both sides of zero, with cutoff at boundary\n",
    "norm = TwoSlopeNorm(vcenter = 0, vmin = s_pair.min(),vmax = s_pair.max())\n",
    "\n",
    "order = np.sort(s_pair)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(action='ignore', category=FutureWarning)\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize = (1.25*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "    sns.lineplot(df_abs, x = 'xlib', y = 'error_rel', units = 'index', estimator = None, ax = ax, \n",
    "                rasterized = True, hue = 's_pair', palette = cmap, hue_norm = norm, hue_order = order[::1],\n",
    "                 legend = False, alpha = 1.0, zorder = -2,) #color = 'silver'\n",
    "\n",
    "    ## add the colorbar\n",
    "    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)\n",
    "    sm.set_array([])\n",
    "\n",
    "    # Remove the legend and add a colorbar\n",
    "    label = 'true mutant fitness\\n(pairwise competition A)'\n",
    "    label = 'mutant fitness in pairwise competition'\n",
    "    ax.figure.colorbar(sm, label = label,use_gridspec=True)\n",
    "\n",
    "\n",
    "    ax.axhline(0, ls = '--', color = 'black')\n",
    "    ax.set_ylabel(r'relative error')\n",
    "\n",
    "    ax.axvline(x_lib_max, color='black', lw = 2)\n",
    "    ax.axvline(xlib_max_j.min(), color='black', lw = 2)\n",
    "    #ax.axvline(x_lib_max/2, color='magenta', lw = 2)\n",
    "    ax.axhline(rel_threshold,color = 'black', ls = '--',lw = 1)\n",
    "    ax.axvline(mutant_ratio_B1, color = color_B1, lw = 3)\n",
    "    ax.axvline(mutant_ratio_B2, color = color_B2, lw = 3)\n",
    "\n",
    "    ax.set_xlabel('relative abundance\\nof the mutant library' + r' $x_{\\mathrm{lib}}$')\n",
    "    #ax.set_xlabel(r'mutant-to-wildtype ratio $x_{\\mathrm{lib}}$')\n",
    "    ax.set_yscale('log')\n",
    "    ax.set_xscale('log')\n",
    "    #ax.set_xlim(xrange[0],xrange[-1])\n",
    "    #ax.set_xlim(xrange[0],xrange[-1]*1.1)\n",
    "    \n",
    "    ax.set_ylim(1e-12,1e2)\n",
    "\n",
    "    ax.spines['right'].set_visible(False)\n",
    "\n",
    "    #fig.tight_layout()\n",
    "    ax.set_rasterization_zorder(-1)\n",
    "\n",
    "    fig.savefig(FIG_DIR_DATASET + 'relative_error_vs_mutant-to-wildtype_ratio_colorbar.pdf',\\\n",
    "                 DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Look at the mutant genotypes with a high relative error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df= pd.DataFrame(error_rel).T\n",
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ranks = df.rank(axis=0, ascending = False,)\n",
    "df_ranks_avg = df_ranks.mean(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### looking for the index with the highest relative error\n",
    "top_ten = df_ranks_avg.sort_values(ascending = True).index[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_ten"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_ranks = pd.DataFrame(gs).rank(ascending = False)\n",
    "l_ranks = pd.DataFrame(ls).rank(ascending = True)\n",
    "Y_ranks = pd.DataFrame(Ys).rank(ascending = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for index in top_ten: \n",
    "    print(\"####\")\n",
    "    print(f\"rank in relative error: {df_ranks_avg.loc[index]}\")\n",
    "    print(f\"mutant genotype index: {index}\")\n",
    "    print(f\"relative fitness s^pair: {s_pair[index]}\")\n",
    "    print(f\"relative fitness s^bulk: {s_bulk_B1[index]}\")\n",
    "    print(f\"growth rate: {gs[index]}\")\n",
    "    print(f\"rank in growth rate: {g_ranks.loc[index, 0]}\")\n",
    "    print(f\"lag time: {ls[index]}\")\n",
    "    print(f\"rank in lag time: {l_ranks.loc[index, 0]}\")\n",
    "    print(f\"biomass yield: {Ys[index]}\")\n",
    "    print(f\"rank in yield: {Y_ranks.loc[index, 0]}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Observations on the mutants with high relative error\n",
    "\n",
    "- The mutants have selection coefficient of small magnitude, so they are close to neutral. \n",
    "- The mutants have always a beneficial lag time, but a deleterious growth rate. This means they are neutal, but in a trade-off between the traits.\n",
    "- It's interesting, that we see only mutants with one direction in this trade-off...those that have a beneficial lag time. I guess, this makes sense since we neglected the lag time variation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare to the wild-type\n",
    "index = 0\n",
    "print(f\"mutant genotype index: {index}\")\n",
    "print(f\"relative fitness s^pair: {s_pair[index]}\")\n",
    "print(f\"relative fitness s^bulk: {s_bulk_B1[index]}\")\n",
    "print(f\"rank in relative error: {df_ranks_avg.loc[index]}\")\n",
    "print(f\"growth rate: {gs[index]}\")\n",
    "print(f\"rank in growth rate: {g_ranks.loc[index, 0]}\")\n",
    "print(f\"lag time: {ls[index]}\")\n",
    "print(f\"rank in lag time: {l_ranks.loc[index, 0]}\")\n",
    "print(f\"biomass yield: {Ys[index]}\")\n",
    "print(f\"rank in yield: {Y_ranks.loc[index, 0]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the discarded terms in error calculation - use d\n",
    "\n",
    "\\begin{equation}\n",
    "        C_{231} = s^\\pair_{21} \\Delta\\tau_{31} \n",
    "                -       \\frac{1}{\\tau_2}\n",
    "        (\\Delta\\tau_{21}\\Delta\\lambda_{31} - \\Delta\\tau_{31}\\Delta\\lambda_{21})\n",
    "\\end{equation}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_laggrowth = np.zeros_like(gs)\n",
    "c_laggrowth[1:] = s_pair[1:]*(tau3-tau1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l1 = ls[0]\n",
    "tau1 = 1/gs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_coupling = np.zeros_like(gs)\n",
    "c_coupling[1:] = -1/taus[1:] * ( (taus[1:]-tau1)*(lstars[1:]-l1) - (tau3-tau1)*(ls[1:]-l1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "## plot data\n",
    "x = s_pair\n",
    "y = c_laggrowth\n",
    "ax.scatter(x,y, rasterized = True, color = color_error1, zorder = 2, \n",
    "           label = r'laggrowth')\n",
    "\n",
    "y = c_coupling\n",
    "ax.scatter(x,y, rasterized = True, color = color_error2, \n",
    "           label = r'coupling')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compute the discarded terms in error calculation - use rewritten decomposition\n",
    "\n",
    "\\begin{equation}\n",
    "        C_{231} = \\frac{-\\Delta \\tau_{21}}{\\tau_2} \\Big[\\Delta \\tau_{31} \\eta^\\pair + \\Delta\\lambda_{31}\\Big]\n",
    "\\end{equation}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_laggrowth = np.zeros_like(gs)\n",
    "#c_laggrowth[1:] = s_pair_growth[1:]*(tau3-tau1)\n",
    "c_laggrowth[1:] = -((taus[1:]-tau1)/taus[1:])*((tau3-tau1)*np.log(1+nus[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "c_coupling = np.zeros_like(gs)\n",
    "c_coupling[1:] = -((taus[1:]-tau1)/taus[1:])* (lstars[1:]-l1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "## plot data\n",
    "x = s_pair\n",
    "y = c_laggrowth\n",
    "ax.scatter(x,y, rasterized = True, color = color_error1, zorder = 2, \n",
    "           label = r'laggrowth')\n",
    "\n",
    "y = c_coupling\n",
    "ax.scatter(x,y, rasterized = True, color = color_error2, zorder = 4,\n",
    "           label = r'coupling')\n",
    "\n",
    "ax.set_ylabel('absolute error pairwise fitness\\n(bulk competition B)')\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "data = pd.DataFrame()\n",
    "data['s_pair'] = s_pair\n",
    "data['lstar'] = lstars\n",
    "data['Delta'] = data['lstar']-l1\n",
    "data['absDelta'] = np.abs(data['Delta'])\n",
    "threshold = 1e-8\n",
    "data['log10Delta'] = np.log10(data['absDelta'], where =data['absDelta']>threshold)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#plot everything\n",
    "ax = data['Delta'].hist(bins = 42, log=True, color = color_error2)\n",
    "ax.axvline( (tau3-tau1)*np.log(1+nus[0]), color = color_error1, lw = 3)\n",
    "ax.set_yscale('log')\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "#zoom in\n",
    "is_zoomed = (data['absDelta']<5)\n",
    "print(f\"No. of points excluded in this zoom: {sum(~is_zoomed)}\")\n",
    "print(f\"Points plotted: {sum(is_zoomed)}\")\n",
    "ax = data.loc[is_zoomed,'Delta'].hist(bins = 42, log=True, color = color_error2)\n",
    "cutoff = (tau3-tau1)*np.log(1+nus[0])\n",
    "ax.axvline(color = color_error1, lw = 3)\n",
    "#ax.set_yscale('log')\n",
    "print()"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "print(\"Total number of points to the right of the cutoff:\")\n",
    "np.sum(data['Delta']> cutoff)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### create a manual symlog\n",
    "#data['symlogDelta'] = np.log10(data['absDelta'], where = data['absDelta'] > 1e-4)\n",
    "data['symlogDelta'] = np.where(data['absDelta']>1e-7, data['log10Delta']*np.sign(data['Delta']), data['Delta'])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "\n",
    "fig, ax = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "## plot data\n",
    "is_zero = c_laggrowth == 0\n",
    "x =s_pair[~is_zero]\n",
    "y =c_coupling[~is_zero]/c_laggrowth[~is_zero]\n",
    "ax.scatter(x,np.abs(y), rasterized = True, color = 'silver', zorder = 2, \n",
    "           label = r'laggrowth')\n",
    "ax.set_yscale('log')\n"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "data['c_laggrowth'] = c_laggrowth\n",
    "data['c_coupling'] = c_coupling\n",
    "is_zero = data['c_laggrowth'] == 0\n",
    "data['c_ratio'] = np.nan\n",
    "data.loc[~is_zero,'c_ratio'] = data.loc[~is_zero,'c_coupling']/data.loc[~is_zero,'c_laggrowth']\n",
    "data['abs_c_ratio'] = np.abs(data['c_ratio'])\n",
    "data['log10abs_c_ratio'] = np.log10(data['abs_c_ratio'])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "ratio = 5\n",
    "grid = sns.jointplot(data=data, x=\"s_pair\", y=\"log10abs_c_ratio\", color = color_error2,\n",
    "                     marginal_kws = {'multiple':'layer', 'fill':True},\n",
    "                     marginal_ticks= False, space = 0, ratio = ratio,\n",
    "                    height = (1+1/ratio)*FIGHEIGHT_TRIPLET, rasterized=True)\n",
    "\n",
    "ax = grid.ax_joint\n",
    "#ax.axhline(np.log10(cutoff), lw = 3, color = color_error1)\n",
    "ax.axhline(np.log10(1), lw = 3, color = color_error1)\n",
    "\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "ax.set_ylabel(r'ratio of error terms: coupling/laggrowth')\n",
    "\n",
    "\n",
    "yticks = ax.get_yticks()\n",
    "\n",
    "ax.set_yticklabels([\"$10^{%d}$\" %v for v in yticks])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "ratio = 5\n",
    "grid = sns.jointplot(data=data, x=\"s_pair\", y=\"log10Delta\", color = color_error2,\n",
    "                     marginal_kws = {'multiple':'layer', 'fill':True},\n",
    "                     marginal_ticks= False, space = 0, ratio = ratio,\n",
    "                    height = (1+1/ratio)*FIGHEIGHT_TRIPLET, rasterized=True)\n",
    "\n",
    "ax = grid.ax_joint\n",
    "ax.axhline(np.log10(cutoff), lw = 3, color = color_error1)\n",
    "\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "ax.set_ylabel(r'background population lag time: $|\\Delta \\lambda_{31}|$')\n",
    "\n",
    "\n",
    "yticks = ax.get_yticks()\n",
    "\n",
    "ax.set_yticklabels([\"$10^{%d}$\" %v for v in yticks])"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### Plot absolute error"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "scaling = 1e6"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "abs_threshold = 1e-3"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "%%time\n",
    "\n",
    "### use a line plot\n",
    "\n",
    "#fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, 0.8*0.5*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes\n",
    "x= xrange\n",
    "y = error_abs[:,1:]*scaling\n",
    "\n",
    "#symlogy = np.log10(np.abs(y), where = np.abs(y) > 0.0001)\n",
    "## plot all knockouts\n",
    "ax.plot(x,y, rasterized = True, color = 'silver', zorder = -2)\n",
    "## plot the mean\n",
    "#ax.plot(x,y.sum(axis=1)/y.shape[1], rasterized = True, color = 'black', lw = 2)\n",
    "\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "ax.set_ylabel(r'absolute error')\n",
    "\n",
    "ax.axvline(mutant_ratio_B1, color = color_B1, lw = 2)\n",
    "ax.axvline(mutant_ratio_B2, color = color_B2, lw = 2)\n",
    "\n",
    "#ax.axhline(abs_threshold, ls = '--', color = 'magenta')\n",
    "#ax.axhline(-abs_threshold, ls = '--', color = 'magenta')\n",
    "\n",
    "ax.set_xlabel('mutant-to-wildtype ratio $x$')\n",
    "ax.set_xscale('log')\n",
    "ax.set_xlim(xrange[0],xrange[-1])\n",
    "\n",
    "ax.set_yscale('symlog')\n",
    "\n",
    "yticks = ax.get_yticks()\n",
    "ax.set_yticks(yticks[1::2])\n",
    "labels= [f\"{v/scaling:.0e}\" for v in yticks[1::2]]\n",
    "ax.set_yticklabels(labels)  \n",
    "\n",
    "ax.set_rasterization_zorder(-1)\n",
    "\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_vs_mutant-to-wildtype_ratio.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "break"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### Plot optimal frequency as a function of selection coefficient - relative error"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "from scipy.optimize import brentq as find_zero"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "threshold = rel_threshold\n",
    "\n",
    "n = s_bulk_Bx.shape[1]\n",
    "ratio_star = np.zeros(n)\n",
    "\n",
    "for i in range(n):\n",
    "    #print(i)\n",
    "    def fun(v):\n",
    "        return np.abs(np.interp(v, xp = xrange, fp = error_rel[:,i],))  - threshold\n",
    "    try:\n",
    "        root = find_zero(fun, xrange[0], xrange[-1])\n",
    "        #print(root)\n",
    "    except:\n",
    "        if fun(1) <= threshold: root = 1\n",
    "        else: root = -1\n",
    "      \n",
    "        \n",
    "    ratio_star[i] = root"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### use a line plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, 0.8*0.5*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes\n",
    "\n",
    "x = s_pair\n",
    "y = ratio_star\n",
    "ax.scatter(x[1:], y[1:], rasterized = True, color = 'silver')\n",
    "\n",
    "#ax.set_ylim(0,1)\n",
    "ax.set_yscale('log')\n",
    "\n",
    "ax.set_ylabel(r'tolerable ratio $x$' +'\\n(bulk competition)')\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "\n",
    "\n",
    "fig.savefig(FIG_DIR_DATASET + 'tolerable_ratio_vs_selection_coefficient.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### Plot optimal frequency as a function of selection coefficient - absolute error"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "abs_threshold = 1e-3"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "threshold = abs_threshold\n",
    "\n",
    "n = s_bulk_Bx.shape[1]\n",
    "ratio_star = np.zeros(n)\n",
    "\n",
    "for i in range(n):\n",
    "    #print(i)\n",
    "    def fun(v):\n",
    "        return np.abs(np.interp(v, xp = xrange, fp = error_abs[:,i],))  - threshold\n",
    "    try:\n",
    "        root = find_zero(fun, xrange[0], xrange[-1])\n",
    "        #print(root)\n",
    "    except:\n",
    "        if fun(1) <= threshold: root = 1\n",
    "        else: root = -1\n",
    "      \n",
    "        \n",
    "    ratio_star[i] = root"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### use a line plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, 0.8*0.5*FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes\n",
    "\n",
    "x = s_pair\n",
    "y = ratio_star\n",
    "ax.scatter(x[1:], y[1:], rasterized = True, color = 'silver')\n",
    "\n",
    "#ax.set_ylim(0,1)\n",
    "ax.set_yscale('log')\n",
    "\n",
    "ax.set_ylabel(r'tolerable ratio $x$' +'\\n(bulk competition)')\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "\n",
    "fig.savefig(FIG_DIR_DATASET + 'tolerable_ratio_vs_selection_coefficient.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### Plot error trend for the total selection coefficient"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "truth = np.outer(np.ones(len(xrange)),si_pair) # need the right shape\n",
    "error_abs = si_bulk_Bx - truth \n",
    "error_rel = np.divide(error_abs,truth, where = truth != 0)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### use a line plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "ax = axes[0]\n",
    "ax.plot(xrange,error_abs, rasterized = True, color = 'dimgrey', marker = 'x')\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('total selection bulk - pairwise selection pair')\n",
    "\n",
    "ax.axhline(0.01,color = 'tab:red', ls = '--')\n",
    "ax.axhline(-0.01,color = 'tab:red', ls = '--')\n",
    "\n",
    "ax = axes[1]\n",
    "ax.plot(xrange,error_rel, rasterized = True, color = 'dimgrey', marker = 'x')\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('relative error: s_2 bulk competition Bx')\n",
    "ax.set_ylim(-0.1,0.1)\n",
    "\n",
    "for ax in axes: \n",
    "    ax.set_xlabel('total proportion of mutants: x')\n",
    "    #ax.set_xscale('log')\n",
    "    #ax.axvline(1/N, color = 'tab:red', ls = '-', label = 'scenario B1: 1/N')\n",
    "\n",
    "    \n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Simulate bulk competition with background mutants only (scenario Bfull)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The frequencies for this scenario can be summarized as \n",
    "\n",
    "    frequency of the mutant strains:      xi = 1/k\n",
    "    frequency of the wildtype strain:     x1 = 0\n",
    "\n",
    "where `k` is the number of knockouts strain. Here the culture only consists of background mutants, and has no wild-type. "
   ]
  },
  {
   "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           # mutant lineages\n",
    "xs[0] = 0              # 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",
    "## calculate total selection coefficient\n",
    "sigma_bulk_Bfull = CalcTotalSelectionCoefficientLogit(xs,xs_final)\n",
    "sigma_bulk_Bfull_log = CalcTotalSelectionCoefficientLog(xs,xs_final)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Plot frequency trajectory for all barcodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "## set timepoints for each trajectory\n",
    "tvec = np.array([0, tsat])\n",
    "## set frequency values for each trajectory\n",
    "trajectories = np.vstack((xs,xs_final))\n",
    "\n",
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "ax.plot(tvec,trajectories[:,1:], color = 'dimgrey', label = 'mutants')\n",
    "ax.plot(tvec,trajectories[:,0], color = 'orange', label = 'wild-type')\n",
    "\n",
    "ax.set_xlabel('time')\n",
    "ax.set_ylabel('frequency')\n",
    "\n",
    "# fix axis limits\n",
    "ax.set_yscale('log')\n",
    "ax.set_ylim(XMIN,XMAX)\n",
    "ax.set_xlim(TMIN,tsat)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## plot\n",
    "fig, ax = plt.subplots(figsize = (FIGHEIGHT_TRIPLET, 0.3* FIGHEIGHT_TRIPLET))\n",
    "                       \n",
    "ax.tick_params(left = False, labelleft = False, bottom = False, labelbottom=False)\n",
    "\n",
    "n = len(xs)\n",
    "colors = [knockout_cmap(i/k+1) for i in range(n)]\n",
    "colors[0] = color_wt\n",
    "for i in range(n):\n",
    "    if i == 0: y1 = [0,0]\n",
    "    else: y1 = [xs[i-1], xs_final[i-1]]\n",
    "    \n",
    "    if i == n-1: y2 = [1,1]\n",
    "    else: y2 = [xs[i], xs_final[i]]\n",
    "        \n",
    "    ax.fill_between(x=[0,tsat], y1=y1,y2=y2)\n",
    "    \n",
    "ax.set_xlim(0,tsat)\n",
    "ax.set_ylim(0,1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,2, figsize = (2*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "x = s_pair\n",
    "\n",
    "ax = axes[0]\n",
    "## plot data for total selection coefficient in logit\n",
    "y = sigma_bulk_Bfull\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_stotal, alpha = 1, marker = 'o') \n",
    "ax.set_title('using logit encoding', loc = 'right')\n",
    "\n",
    "ax = axes[1]\n",
    "## plot data for total selection coefficient in log\n",
    "y = sigma_bulk_Bfull_log\n",
    "ax.scatter(x[1:],y[1:]-x[1:], rasterized = True, color = color_stotal, alpha = 1, marker = 'o') \n",
    "ax.set_title('using log encoding', loc = 'right')\n",
    "\n",
    "\n",
    "for ax in axes: \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('absolute error total fitness\\n(bulk competition C)')\n",
    "    ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "    #ax.legend(loc = 'lower right', frameon=False)\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'absolute_error_bulk_competition_full.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = pd.DataFrame()\n",
    "## load data on selection coefficients\n",
    "data['sigma_bulk_Bfull'] = sigma_bulk_Bfull\n",
    "data['sigma_bulk_Bfull_log'] =  sigma_bulk_Bfull_log\n",
    "data['s_pair'] = s_pair\n",
    "## create rank variables, fittest mutants come first\n",
    "data['sigma_bulk_Bfull_rank'] = data['sigma_bulk_Bfull'].rank(ascending = False)\n",
    "data['sigma_bulk_Bfull_log_rank'] =  data['sigma_bulk_Bfull_log'].rank(ascending = False)\n",
    "data['s_pair_rank'] = data['s_pair'].rank(ascending = False) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,3, figsize = (3.*FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "x = s_pair\n",
    "\n",
    "ax = axes[0] \n",
    "## plot ranking total bulk vs pairwise\n",
    "y = data['sigma_bulk_Bfull_rank'] - data['s_pair_rank']\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_stotal, alpha = 1, marker = 'o',zorder = 5)\n",
    "ax.set_ylabel('rank diff. between statistics: \\n total fit. (logit) - pairwise fit. (pair)')\n",
    "\n",
    "ax = axes[1] \n",
    "## plot ranking pairwise bulk vs pairwisse\n",
    "y = data['sigma_bulk_Bfull_log_rank'] - data['s_pair_rank']\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_stotal, alpha = 1, marker = 'o',zorder = 5) \n",
    "ax.set_ylabel('rank diff. between statistics: \\n total fit. (log) - pairwise fit. (pair)')\n",
    "\n",
    "ax = axes[2] \n",
    "## plot ranking total bulk vs pairwise bulk\n",
    "y = data['sigma_bulk_Bfull_log_rank'] - data['sigma_bulk_Bfull_rank']\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_stotal, alpha = 1, marker = 'o', zorder = 5) \n",
    "ax.set_ylabel('rank diff. between statistics: \\n total fit. (log) - total fit. (logit)')\n",
    "for ax in axes: \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', zorder = 10)\n",
    "\n",
    "    ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "\n",
    "\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'ranking_mismatch_bulk_Bfull.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compare in a plot\n",
    "\n",
    "fig, axes = plt.subplots(1,1, figsize = (FIGHEIGHT_TRIPLET, FIGHEIGHT_TRIPLET), sharex = True)\n",
    "\n",
    "\n",
    "ax = axes \n",
    "x = s_pair\n",
    "\n",
    "## plot data\n",
    "\n",
    "y = np.abs(np.divide( sigma_bulk_Bfull - s_pair, s_pair, where = s_pair !=0))\n",
    "ax.scatter(x[1:],y[1:], rasterized = True, color = color_stotal, alpha = 1, \n",
    "           label = r'total scoeff. $\\sigma$', marker = 'o') \n",
    "\n",
    "\n",
    "ax.axhline(rel_threshold, ls = '--', color = 'black')\n",
    "\n",
    "ymin,ymax = ax.get_ylim()\n",
    "yabs = np.max(np.abs([ymin,ymax]))\n",
    "ax.set_yscale('log')\n",
    "ax.axhline(0, ls = '--', color = 'black')\n",
    "\n",
    "ax.set_ylabel('relative error\\n(bulk competition C)')\n",
    "ax.set_xlabel('true mutant fitness\\n(pairwise competition A)')\n",
    "ax.legend(loc = 'upper left', frameon=False)\n",
    "error1_min,error1_max = ax.get_ylim()\n",
    "\n",
    "\n",
    "#fig.tight_layout()\n",
    "fig.savefig(FIG_DIR_DATASET + 'relative_error_bulk_competition_C.pdf',\\\n",
    "             DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}