73_reproduce_LTEE_fit_to_data_subsets.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Long-term fitness trajectories from the LTEE\n",
    "\n",
    "Data available in the Supplemental Material of Good et al. Nature 2017. \n",
    "Download possible from Ben Good's github repository [here](https://github.com/benjaminhgood/LTEE-metagenomic/blob/master/additional_data/Concatenated.LTEE.data.all.csv)\n",
    "\n",
    "We follow the procedures from Wiser et al. 2013 [here](https://doi.org/10.1126/science.1243357). From the Supplemental Material, we are given the following information. \n",
    "\n",
    "\n",
    "- Summarizing statistical procedures to fit the two models\n",
    "- Models were fit to fitness trajectories using the ‘nls’ package in r. \n",
    "- Model fits were compared using the BIC information criterion scores. These were then converted into an odds ratio. \n",
    "    - Table S1 shows the BIC scores and odds ratios for fits to subsets of the data: a) all 12 populations and all time points, b) excluding 3 populations with incomplete trajectories and c) excluding 6 populations that evolved hypermutability\n",
    "    - Table S2 summarizes BIC scores for fits to individual populations. This also indicates if the population was truncated or a hypermutator \n",
    "    - Table S4 lists the estimated parameters for the power law fit\n",
    "\n",
    "On the bigger picture, there is also the talk from 2013 by Wiser on [Youtube](https://www.youtube.com/watch?v=CmyBn5Cezy4) with 127 views as of September 2022. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from scipy import stats\n",
    "\n",
    "from scipy.optimize import curve_fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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/LTEE_fit/'\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": [
    "df = pd.read_csv('./data/Concatenated.LTEE.data.all.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "### execute script to load modules here\n",
    "exec(open('setup_aesthetics.py').read())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### explanation of column headers and labels\n",
    "\n",
    "    607 is the wild-type strain 'REL607'\n",
    "    D.0 is the dilution factor applied to count colonies in initial inoculum. \n",
    "    D.1 is the dilution factor applied to count colonies in the saturated population. We expect D.1 = 100*D.0. \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### split into two subsets, according to the Ara-marker of evolving population\n",
    "is_Ara_positive = np.array(['+' in v for v in df['Population'].values])\n",
    "\n",
    "df_pos = df[is_Ara_positive]\n",
    "df_neg = df[~is_Ara_positive] ### only use Ara negative lineages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## treat Ara-posative population\n",
    "assert all(df_pos['Red.Pop'] == '606') # wild-type is always the red population\n",
    "\n",
    "### re-construct population sizes\n",
    "df_pos['Nwt.0'] = df_pos['Red.0']*df_pos['D.0']\n",
    "df_pos['Nmut.0'] = df_pos['White.0']*df_pos['D.0']\n",
    "df_pos['Nwt.1'] = df_pos['Red.1']*df_pos['D.1']\n",
    "df_pos['Nmut.1'] = df_pos['White.1']*df_pos['D.1']\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## treat Ara-negative population\n",
    "assert all(df_neg['White.Pop'] == '607') # wild-type is always the white population\n",
    "\n",
    "### re-construct population sizes\n",
    "df_neg['Nwt.0'] = df_neg['White.0']*df_neg['D.0']\n",
    "df_neg['Nmut.0'] = df_neg['Red.0']*df_neg['D.0']\n",
    "df_neg['Nwt.1'] = df_neg['White.1']*df_neg['D.1']\n",
    "df_neg['Nmut.1'] = df_neg['Red.1']*df_neg['D.1']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### join\n",
    "\n",
    "df = df_pos.append(df_neg)\n",
    "\n",
    "## reconstruct frequencies\n",
    "df['xmut.0'] = df['Nmut.0']/(df['Nwt.0'] + df['Nmut.0'])\n",
    "df['xmut.1'] = df['Nmut.1']/(df['Nwt.1']  + df['Nmut.1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### reconstruct fitness statistics\n",
    "df['s'] = np.log(df['Nmut.1']/df['Nwt.1']) - np.log(df['Nmut.0']/df['Nwt.0'])\n",
    "df['W'] = np.divide( np.log(df['Nmut.1']/df['Nmut.0']),\\\n",
    "                            np.log(df['Nwt.1']/df['Nwt.0']))\n",
    "df['delta_log'] = np.log(df['xmut.1']) - np.log(df['xmut.0'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### check that  my number is consistent with existing value for 'Fitness' in the dataset\n",
    "np.allclose(df['W'], df['Fitness'],equal_nan=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## manual check\n",
    "df['gap'] = df['W'] - df['Fitness']\n",
    "print(df['gap'].max())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_s = 'tab:grey'\n",
    "color_W = 'firebrick'\n",
    "color_deltalog = 'navy'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "color_hyper = 'tab:red'\n",
    "def hyperbolic(t, a, b):\n",
    "    ## compare first Equation in paper\n",
    "    return 1 + np.divide(a*t,t+b)\n",
    "\n",
    "color_power = 'tab:blue'\n",
    "def powerlaw(t, a, b):\n",
    "    ## compare second Equation in paper\n",
    "    return np.power(b*t + 1,a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## we drop some superfluous columns\n",
    "columns_auxiliary = [ 'Red.0', 'White.0', 'Red.1', 'White.1', 'D.0', 'D.1', 'gap', 'White.Pop', 'Red.Pop', 'Fitness', 'Nwt.0', 'Nmut.0', 'Nwt.1', 'Nmut.1', 'Nwt.0']\n",
    "\n",
    "df = df.drop(columns_auxiliary, axis = 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## shift data points for alternative statistics, which are based at zero\n",
    "df['s+1'] = df['s'] +1\n",
    "df['delta_log+1'] = df['delta_log'] +1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## reproduce fits to grand mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## set up different subsets\n",
    "\n",
    "truncated_to_remove = ['Ara + 6' , 'Ara - 2', 'Ara - 3']\n",
    "hypermutators_to_remove = ['Ara - 1', 'Ara - 2', 'Ara - 3', 'Ara - 4', 'Ara + 3', 'Ara + 6']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### recreate different datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### set up  dict to later access the dataframe\n",
    "\n",
    "subset2data = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for subset_label, pop_to_exclude in zip(['all', 'truncated_removed', 'hypermutators_removed'],\\\n",
    "                                        [[], truncated_to_remove, hypermutators_to_remove]):\n",
    "\n",
    "    ### remove populations \n",
    "    index_to_include = [v not in pop_to_exclude for v in df['Population'].values ]\n",
    "    df_subset = df[index_to_include].copy(deep = True)\n",
    "\n",
    "    ### we create a new dataframe where each timepoint and population is only represented once\n",
    "    df_subset = df_subset.sort_values(by = ['Population', 'Generation', 'Rep']) # first sort for nice look\n",
    "    df_averaged = df_subset.drop_duplicates([ 'Population', 'Generation']).copy(deep=True)\n",
    "    df_averaged = df_averaged.drop(['Rep'], axis = 1)\n",
    "    df_averaged= df_averaged.reset_index()\n",
    "\n",
    "    ## we average across the number of replicates\n",
    "    df_averaged['no_replicates'] = -1 # as a collateral statistic, we count the number of replicates\n",
    "\n",
    "    for i in df_averaged.index:\n",
    "        row = df_averaged.loc[i]\n",
    "\n",
    "        pop = row['Population']\n",
    "        gen = row['Generation']\n",
    "\n",
    "        is_gen = np.array([v == gen for v in df['Generation'].values])\n",
    "        is_pop = np.array([v == pop for v in df['Population'].values])\n",
    "        df_replicates = df.loc[is_gen & is_pop]\n",
    "        df_averaged.at[i,'no_replicates'] = df_replicates.shape[0]\n",
    "\n",
    "        for v in ['xmut.0', 'xmut.1', 's', 'W', 'delta_log']:\n",
    "            df_averaged.at[i, v] = df_replicates[~df_replicates[v].isna()][v].mean()\n",
    "\n",
    "\n",
    "    ## shift data points for alternative statistics, which are based at zero\n",
    "    df_averaged['s+1'] = df_averaged['s'] +1\n",
    "    df_averaged['delta_log+1'] = df_averaged['delta_log'] +1\n",
    "    \n",
    "    ## store reference\n",
    "    subset2data[subset_label] = df_averaged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot distribution for degree of replication\n",
    "\n",
    "for subset_label in ['all', 'truncated_removed', 'hypermutators_removed']:\n",
    "    df_averaged = subset2data[subset_label]\n",
    "    \n",
    "    assert df_averaged['no_replicates'].min() >= 2, 'expect at least 2 replicates per point'\n",
    "\n",
    "    ### manual check: does everyone have at least 2 replicates\n",
    "    fig, ax =plt.subplots(figsize = (FIGWIDTH_TRIPLET, FIGHEIGHT_TRIPLET))\n",
    "    ax = df_averaged['no_replicates'].hist(bins=np.arange(0.5,11.5),log=True)\n",
    "    ax.set_xlabel('number of replicates for a single timepoint and population')\n",
    "    ax.set_ylabel('count')\n",
    "    ax.set_title('set with ' + subset_label + f' (n={df_averaged.shape[0]})')\n",
    "    fig.savefig(FIG_DIR + f\"histogram_of_replicates_for_set_with_{subset_label}.pdf\", DPI = DPI, bbox_inches = 'tight', pad_inches = PAD_INCHES)\n",
    "    \n",
    "    ## remove cluttering output\n",
    "    if subset_label != 'all':\n",
    "        plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## inspect outliers manually\n",
    "df_averaged = subset2data['all']\n",
    "\n",
    "#df_averaged[df_averaged['no_replicates'] == 10]\n",
    "#df_averaged[df_averaged['no_replicates'] == 4].sort_values('Generation')\n",
    "\n",
    "## compute alternative statistics\n",
    "## we add constant to data values, so we can fit them to the same model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform model fit on each dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### set up Dataframe to store results\n",
    "df_results = pd.DataFrame()\n",
    "for subset_label in ['all', 'truncated_removed', 'hypermutators_removed']:\n",
    "    df_tmp= pd.DataFrame(data = {'subset_label':3*[subset_label], 'target':['W', 's+1', 'delta_log+1']})\n",
    "    df_results = df_results.append(df_tmp)\n",
    "    \n",
    "df_results = df_results.set_index(['subset_label','target'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for subset_label in ['all','truncated_removed', 'hypermutators_removed']:\n",
    "    for target  in ['W', 's+1', 'delta_log+1']:\n",
    "        \n",
    "        df_averaged = subset2data[subset_label]\n",
    "\n",
    "        n_datapoints = df_averaged[~df_averaged[target].isna()].shape[0]\n",
    "        df_results.at[(subset_label, target),'n_datapoints'] = n_datapoints\n",
    "\n",
    "        fig, ax =plt.subplots(figsize = (1.5*FIGWIDTH_TRIPLET, 1.1*FIGHEIGHT_TRIPLET))\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "        t = df_averaged['Generation']\n",
    "        y = df_averaged[target]\n",
    "\n",
    "        ## plot raw data\n",
    "        sns.scatterplot(x=t,y=y, ax=ax, color = 'grey')\n",
    "        sns.lineplot(x=t,y=y, ax =ax, color = 'grey')\n",
    "\n",
    "        ## fit hyperbolic model\n",
    "        popt_hyperbolic,_ = curve_fit(f=hyperbolic, xdata=t, ydata=y)\n",
    "        ## store\n",
    "        df_results.at[(subset_label,target), 'hyper_a'] = popt_hyperbolic[0]\n",
    "        df_results.at[(subset_label,target), 'hyper_b'] = popt_hyperbolic[1]\n",
    "        ## compute trajectory\n",
    "        y_hat = hyperbolic(t, *popt_hyperbolic)\n",
    "        ## store sum of residuals squared\n",
    "        rss = np.power(y_hat - y,2).sum()\n",
    "        rsquared = 1 - (rss/np.power(y - y.mean(),2).sum())\n",
    "        df_results.at[(subset_label, target), 'hyper_rss']  = rss\n",
    "        df_results.at[(subset_label, target), 'hyper_rsquared']  = rsquared\n",
    "        ## plot trajectory\n",
    "        param_label = 'hyperbolic: a=%.2f, b=%.0f' % tuple(popt_hyperbolic)\n",
    "        sns.lineplot(x=t, y =y_hat, color = color_hyper,ax=ax,\n",
    "                     label = param_label + rf\", $R^2={rsquared:.2f}$\")\n",
    "        \n",
    "\n",
    "    \n",
    "\n",
    "\n",
    "        ## fit powerlaw model\n",
    "        popt_powerlaw,_ = curve_fit(f=powerlaw, xdata=t, ydata=y)\n",
    "        ## store\n",
    "        df_results.at[(subset_label,target), 'powerlaw_a'] = popt_powerlaw[0]\n",
    "        df_results.at[(subset_label,target), 'powerlaw_b'] = popt_powerlaw[1]\n",
    "        ## compute trajectory\n",
    "        y_hat = powerlaw(t, *popt_powerlaw)\n",
    "        ## store sum of residuals squared\n",
    "        rss = np.power(y_hat - y,2).sum()\n",
    "        rsquared = 1 - (rss/np.power(y - y.mean(),2).sum())\n",
    "        df_results.at[(subset_label, target), 'powerlaw_rss']  = rss\n",
    "        df_results.at[(subset_label, target), 'powerlaw_rsquared']  = rsquared\n",
    "        ## plot trajectory\n",
    "        param_label = 'powerlaw: a=%5.4f, b=%5.5f' % tuple(popt_powerlaw)\n",
    "        sns.lineplot(x=t, y =y_hat, color = color_power,ax=ax,\\\n",
    "                     label = param_label + rf\", $R^2={rsquared:.2f}$\")\n",
    "        \n",
    "        ## remove cluttering output\n",
    "        if subset_label != 'all':\n",
    "            plt.close(fig)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculate model comparison stats\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## set parameters\n",
    "\n",
    "k = 3 # include one extra parameter for variance of error distribution\n",
    "\n",
    "\n",
    "for subset_label in ['all','truncated_removed', 'hypermutators_removed']:\n",
    "    for target  in ['W', 's+1', 'delta_log+1']:\n",
    "\n",
    "        ## read number of datapoints\n",
    "        n = df_results.at[(subset_label, target),'n_datapoints']\n",
    "\n",
    "        ## read hyperbolic model results\n",
    "        rss = df_results.at[(subset_label, target),'hyper_rss']\n",
    "        ## compute information criteria\n",
    "        aic_hyper = 2*k + n*np.log(rss)\n",
    "        bic_hyper = n*np.log(rss) + k*np.log(n)\n",
    "        likelihood_hyper = -n/2*(np.log(2*np.pi) + 1) -n/2*np.log(rss/n)\n",
    "        ## store\n",
    "        df_results.at[(subset_label, target),'hyper_aic'] = aic_hyper\n",
    "        df_results.at[(subset_label, target),'hyper_bic'] = bic_hyper\n",
    "        df_results.at[(subset_label, target),'hyper_likelihood'] = likelihood_hyper\n",
    "\n",
    "        ## read powerlaw model results\n",
    "        rss = df_results.at[(subset_label, target),'powerlaw_rss']\n",
    "        ## compute information criteria\n",
    "        aic_powerlaw = 2*k + n*np.log(rss)\n",
    "        bic_powerlaw = n*np.log(rss) + k*np.log(n)\n",
    "        likelihood_powerlaw = -n/2*(np.log(2*np.pi) + 1) -n/2*np.log(rss/n)\n",
    "        ## store\n",
    "        df_results.at[(subset_label, target),'powerlaw_aic'] = aic_powerlaw\n",
    "        df_results.at[(subset_label, target),'powerlaw_bic'] = bic_powerlaw\n",
    "        df_results.at[(subset_label, target),'powerlaw_likelihood'] = likelihood_powerlaw\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute model comparison stats based on BIC\n",
    "\n",
    "df_results['delta_bic'] = df_results['hyper_bic'] - df_results['powerlaw_bic']\n",
    "df_results['likelihood_ratio'] = np.exp(df_results['powerlaw_likelihood'] - df_results['hyper_likelihood'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## coompute model comparison stats based on AIC\n",
    "df_results['delta_aic'] = df_results['hyper_aic'] - df_results['powerlaw_aic']\n",
    "\n",
    "for subset_label in ['all','truncated_removed', 'hypermutators_removed']:\n",
    "    for target  in ['W', 's+1', 'delta_log+1']:\n",
    "        aic_hyper = df_results.at[(subset_label, target), 'hyper_aic']\n",
    "        aic_powerlaw = df_results.at[(subset_label, target), 'powerlaw_aic']\n",
    "\n",
    "        ## identify model with smaller aic\n",
    "        aic_min = np.min([aic_hyper,aic_powerlaw])\n",
    "        aic_max = np.max([aic_hyper,aic_powerlaw])\n",
    "\n",
    "        ## compute probability\n",
    "        ### see https://en.wikipedia.org/wiki/Akaike_information_criterion#How_to_use_AIC_in_practice\n",
    "        prob_max_model_is_better = np.exp((aic_min - aic_max)/2)\n",
    "\n",
    "\n",
    "        df_results.at[(subset_label, target), 'akaike_pvalue'] = prob_max_model_is_better"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "col_to_print = ['hyper_bic', 'powerlaw_bic', 'delta_bic', 'likelihood_ratio',\\\n",
    "               'hyper_aic', 'powerlaw_aic', 'delta_aic', 'akaike_pvalue']\n",
    "\n",
    "df_results[col_to_print]"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "### store as excel file\n",
    "## save\n",
    "df_results.to_excel(FIG_DIR + 'stats_for_grand_mean.xlsx', index=True)"
   ]
  },
  {
   "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
}