compute_channel_capacity_HighMI_3.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Channel capacity from the HighMI_3 data set\n",
    "\n",
    "To run this notebook, you need:\n",
    "- Raw cytokine time series in `data/final/`, in particular the dataset `\"cytokineConcentrationPickleFile-20210619-HighMI_3-final.pkl\"`. \n",
    "- Processed cytokine time series of another experiment (by default `HighMI_1-i.hdf`, $i \\in {1, 2,3 ,4}$) in `data/processed/`, to get a better estimate of the $v_2/v_1$ final slope. \n",
    "- the input weights of a neural network and the min and max cytokine concentrations used to scale the data, in `data/trained-networks`. \n",
    "- Table of OT-1 antigens' EC$_{50}$s, values from the literature and from our own measurements, in the JSON file `data/misc/potencies_df_2021.json`\n",
    "Those files are available in the data repository hosted online. \n",
    "\n",
    "\n",
    "### HighMI_3 dataset\n",
    " \n",
    "This is a dataset from an experiment designed specifically to estimate the number of antigen categories encoded by the cytokine latent space. It contains 9 experimental replicates of the cytokine time series for 8 OVA peptides (N4 down to E1) and each of 4 concentrations (1 $\\mu$M down to 1 nM). \n",
    "\n",
    "### Another explanation of channel capacity\n",
    "What we want to compute is channel capacity between antigen quality $Q$ (quantified by the EC$_{50}$) and latent space time trajectories, characterized by the values of model parameters fitted on them. Specifically, we take the vector, continuous random variable $\\mathbf{X} = (a_0, \\tau_0, \\theta)$, made of three parameters from the force model with matching, as the description of a trajectory in latent space in response to some antigen $Q$. The mutual information $MI(\\mathbf{X}; Q)$ measures how much information about antigen quality $Q$ we gain by knowing the value of the three model parameters in $\\mathbf{X}$ that describe (in latent space) the resulting cytokine time series. The MI is defined as\n",
    "\n",
    "$$ MI(\\mathbf{X}; Q) = \\sum_q p_Q(q) \\int_{\\mathcal{X}} \\mathrm{d}\\mathbf{x} f_{\\mathbf{X}|Q=q}(\\mathbf{x}) \\log_2{\\left( \\frac{f_{\\mathbf{X}|Q=q}(\\mathbf{x})}{f_{\\mathbf{X}}(\\mathbf{x})} \\right)}$$\n",
    "\n",
    "where $\\mathcal{X}$ is the range of possible values of the random variable $\\mathbf{X}$, $p_Q(q)$ is the probability to have an antigen of quality $Q$, $f_{\\mathbf{X}|Q=q}$ is the conditional probability density of $\\mathbf{X}$ in response to a given antigen quality $Q$, and $f_{\\mathbf{X}}$ is the marginal probability density of $\\mathbf{X}$, given by $\\sum_q p_Q(q) f_{\\mathbf{X|Q=q}}(\\mathbf{x})$. The conditional probability density can be though of as an input-output function that maps a given antigen quality to a distribution of possible latent space trajectories. In other words, it characterizes the latent space as an information channel that encodes antigen quality (hence the notion of antigen encoding). \n",
    "\n",
    "The antigen probability distribution $p_Q$ is somewhat arbitrary: is it the frequency of a given antigen in the experiment? Is it the abundance of antigens of that strength $\\textit{in vivo}$? We can actually characterize the information content of the cytokine latent space independently from $p_Q$ by computing instead channel capacity, $C(\\mathbf{X}; Q)$. This quantity gives the most information one can possibly extract from the channel, whatever $p_Q$ is. It is commonly used for this purpose in information theory.  Mathematically, \n",
    "\n",
    "$$ C(\\mathbf{X}; Q) = \\max_{p_Q} MI(\\mathbf{X}; Q) $$\n",
    "\n",
    "In other words, we maximize mutual information over the space of all possible input (antigen) distributions. Hence, $C(\\mathbf{X}; Q)$ no longer depends on $p_Q$. \n",
    "\n",
    "\n",
    "## Steps to compute channel capacity\n",
    "1. Process cytokine data: smoothing, computation of integrals, \n",
    "2. Projection of cytokine trajectories to latent space and (min max scaling) normalization with training data normalization parameters (min and max of each cytokine in the training data). \n",
    "3. Fitting model parameters on each time series in latent space: gives $\\mathbf{x}$ samples\n",
    "4. Fitting multivariate normal distributions on the distributions of model parameters for each antigen (i.e., the distributions $f_{\\mathbf{X}|Q=q}$)\n",
    "5. Interpolation, as a function of $Q$, of the means and covariance matrices of the multivariate normal distributions. \n",
    "6. Calculation of the channel capacity between $Q$ and $\\mathbf{X}$ with the Blahut-Arimoto and Monte Carlo integrations for the multidimensional integrals over the domain of $\\mathbf{X}$. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import pickle\n",
    "from time import perf_counter\n",
    "import os\n",
    "\n",
    "from ltspcyt.scripts.process_raw_data import process_file, process_file_filter\n",
    "from ltspcyt.scripts.sigmoid_ballistic import return_param_and_fitted_latentspace_dfs, compute_v2_v1\n",
    "from utils.statistics import build_symmetric\n",
    "from utils.distrib_interpolation import (eval_interpolated_means_covs, interpolate_params_vs_logec50, \n",
    "                                         stats_per_levels, compute_cholesky_dataframe)\n",
    "import utils.custom_pandas as custom_pd\n",
    "from utils.extra_pairplots import dual_pairplot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Process raw data, filter out noise\n",
    "The processing steps we needed for this dataset:\n",
    "1. Removing some outlier time series caused by stronger peptides splashing from one well to another during preparation of these plates for the robot. \n",
    "2. Filtering of cytokine trajectories that are purely noise, using the same filtering method based on the Kolmogorov-Smirnov test that was applied to the human TCR dataset (see supplementary information). \n",
    "3. Usual moving average and spline smoothing in log scale (supp. section 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Remove known outlier trajectories\n",
    "\n",
    "Due to a shortage of tips, plate inversion was used to remove supernatant during preparation. This caused stronger peptides to splash in some wells, which caused unnatural high cytokine time series in those wells. \n",
    "\n",
    "Look at the raw data plot below (IFN-$\\gamma$ shown, but other cytokines display the same effect). \n",
    "- Replicates 2, 5, 8 of V4 at 1 nM and E1 at 1 $\\mu$M were contaminated with stronger peptides. This shows up as unrealistically high cytokine responses. These replicates are in homologous positions in each of the 3 plates (plate 1 is reps 1-3, plate 2 is reps 4-6, plate 3 is reps 7-9, and hence 2,5,8 correspond to the second replicate in all 3 plates). \n",
    "- Similarly, replicates 3-6-9 of G4 at 1 nM have been contaminated, as well as replicates 3-6-9 of E1 at 1 $\\mu$M.\n",
    "\n",
    "The faulty V4, G4 replicates can just be dropped, but the E1 replicates have to be replaced, because E1 is the null response reference for further noise filtering (processing step 2). Since correct E1 1$\\mu$M trajectories are just uncorrelated noise, we replace the faulty ones this by randomly sampling points from the three non-faulty replicates to replace the fault replicates. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import raw cytokine dataset. Concatenate back the nine separate files that\n",
    "# were created by the default processing pipeline. \n",
    "df_raw = {}\n",
    "for i in range(1, 10):\n",
    "    df_raw[str(i)] = pd.read_hdf(os.path.join(\"data\", \"final\", \n",
    "            \"cytokineConcentrationPickleFile-20210619-HighMI_3-{}-final.hdf\".format(i)))\n",
    "df_raw = pd.concat(df_raw, names=[\"Replicate\"])\n",
    "order_levels = np.asarray(df_raw.index.names)\n",
    "order_levels[:2] = order_levels[1::-1]\n",
    "df_raw = df_raw.reorder_levels(order_levels)\n",
    "del order_levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot of all IFNg time series (for all peptides [columns] at all concentrations [rows] \n",
    "# and all replicates [different line colors])\n",
    "peptide_order = [\"N4\", \"Q4\", \"T4\", \"V4\", \"G4\", \"E1\", \"A2\", \"Y3\"]\n",
    "conc_order = [\"1uM\", \"100nM\", \"10nM\", \"1nM\"]\n",
    "g = sns.relplot(data=np.log10(df_raw.unstack(\"Cytokine\").stack(\"Time\")).reset_index(), \n",
    "           x=\"Time\", y=\"IFNg\", col=\"Peptide\", row=\"Concentration\", hue=\"Replicate\", palette=\"Set2\", \n",
    "           kind=\"line\", height=1.3, lw=2., col_order=peptide_order, row_order=conc_order)\n",
    "\n",
    "for i in range(g.axes.shape[0]):\n",
    "    for j in range(g.axes.shape[1]):\n",
    "        g.axes[i, j].set_title(peptide_order[j] + \" \" + conc_order[i])\n",
    "\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We simply drop the E1 replicates with spurious activation due to splashed peptide. \n",
    "# A more fancy strategy that gives the same results is to replace faulty E1 replicates \n",
    "# with points randomly sampled from correct remaining replicates\n",
    "# This is also a correct way to treat those outliers because E1 is just background noise normally, \n",
    "# so we can randomly sample from its fluctuations without order. \n",
    "def replace_time_series(df, lbl_to_replace, good_replicates, rndgen):\n",
    "    \"\"\" The df should be indexed (Cytokine, Replicate, TCellNumber, Peptide, Concentration). \n",
    "    Time should be in the columns. \n",
    "    lbl_to_replace specifies slice(None), Replicate, ..., Concentration\n",
    "    Warning: modifies the dataframe in place\n",
    "    \"\"\"\n",
    "    choices = rndgen.choice(good_replicates, size=len(df.columns), replace=True)\n",
    "    for i, tim in enumerate(df.columns):\n",
    "        df.loc[lbl_to_replace, tim] = df.loc[(slice(None), choices[i], *lbl_to_replace[2:]), tim].values\n",
    "    return 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment if you want to replace the replicates instead of dropping them. \n",
    "#randomgen = np.random.default_rng(seed=1349839)\n",
    "#for rep in [\"2\", \"3\", \"5\", \"6\", \"8\", \"9\"]:\n",
    "#    replace_time_series(df_raw, (slice(None), rep, \"30k\", \"E1\", \"1uM\"), [\"1\", \"4\", \"7\"], randomgen)\n",
    "\n",
    "# Drop faulty replicates for other peptides\n",
    "for cyto in df_raw.index.get_level_values(\"Cytokine\").unique():\n",
    "    for rep in [\"2\", \"5\", \"8\"]:\n",
    "        df_raw = df_raw.drop((cyto, rep, \"30k\", \"V4\", \"1nM\"), axis=0)\n",
    "    for rep in [\"3\", \"6\", \"9\"]:\n",
    "        df_raw = df_raw.drop((cyto, rep, \"30k\", \"G4\", \"1nM\"), axis=0)\n",
    "    for rep in [\"2\", \"3\", \"5\", \"6\", \"8\", \"9\"]:\n",
    "        df_raw = df_raw.drop((cyto, rep, \"30k\", \"E1\", \"1uM\"), axis=0)\n",
    "\n",
    "### Write back to a file\n",
    "fname_corrected = \"cytokineConcentrationPickleFile-20210619-HighMI_3_corrected-final.hdf\"\n",
    "df_raw.to_hdf(os.path.join(\"data\", \"final\", fname_corrected), key=\"df\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 and 1.3 Filter noise, smooth and fit splines the corrected raw data\n",
    "Here, E1 is the reference for a null response. \n",
    "Trajectories for E1 and some other non-responding conditions have large fluctuations in the 100 fM to 1 pM cytokine concentration range. These concentrations are in the low end of the calibration curve and are thus amplified in variation. We filter them out, after first log-transforming the raw data. \n",
    "\n",
    "We compare the statistics of each IFN-$\\gamma$ time series to the corresponding E1 trajectories (same replicate, same peptide concentration) with the Kolmogorov-Smirnov test, and if they are found to be similar, IL-2, IL-17A, TNF and IL-6 are set to zero for both trajectories. E1 trajectories have those cytokines set to zero as well, because we know they are experimental background noise. IFN-$\\gamma$ is not set to zero, because we want to keep some signal, otherwise the trajectories are perfectly zero in latent space, which is somewhat artificial. \n",
    "\n",
    "Then, the usual processing is applied on the corrected data: finding and fixing missing time points, moving average smoothing, cubic B-spline fitting. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filtering out null-like cytokine time series. \n",
    "# If the E1 replicates were resampled from background instead of dropped, use \n",
    "# split_filter_levels=[\"Concentration\", \"Replicate\"], i.e. compare each replicate\n",
    "# to its E1 counterpart, instead of concatenating all replicates. \n",
    "_, _, _, df_wt = process_file_filter(folder=os.path.join(\"data\", \"final\"), file=fname_corrected,\n",
    "                        do_filter_null=True, null_reference=\"E1\", filter_pval=0.5, choice_filter_cyto=\"IFNg\", \n",
    "                        choice_remove_cyto=[\"IL-2\", \"IL-17A\", \"TNFa\", \"IL-6\"], remove_il17=False,\n",
    "                        split_filter_levels=[\"Concentration\"], do_self_filter=True)\n",
    "\n",
    "# Write the corrected and processed file to hdf\n",
    "df_wt.to_hdf(os.path.join(\"data\", \"HighMI_3_corrected.hdf\"), key=\"df\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Project to latent space and normalize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peptides = [\"N4\", \"Q4\", \"T4\", \"V4\", \"G4\", \"E1\", \"A2\", \"Y3\"]\n",
    "concentrations = [\"1uM\", \"100nM\", \"10nM\", \"1nM\"]\n",
    "times = np.arange(1, 73)\n",
    "tcellnum = \"30k\"\n",
    "tcn_fit = \"30k\"\n",
    "folder = \"capacity\"\n",
    "suffix = \"_HighMI_3\"\n",
    "\n",
    "\n",
    "minmaxfile = os.path.join(\"data\", \"trained-networks\", \"min_max-thomasRecommendedTraining.hdf\")\n",
    "df_min = pd.read_hdf(minmaxfile, key=\"df_min\")\n",
    "df_max = pd.read_hdf(minmaxfile, key=\"df_max\")\n",
    "mlpcoefs = np.load(os.path.join(\"data\", \"trained-networks\", \"mlp_input_weights-thomasRecommendedTraining.npy\"))\n",
    "\n",
    "cytokines = df_min.index.get_level_values(\"Cytokine\")\n",
    "df_wt = pd.concat({\"HighMI_3\":df_wt}, names=[\"Data\"])\n",
    "\n",
    "# Normalize and projection to latent space\n",
    "df = df_wt.unstack(\"Time\").loc[:, (\"integral\", cytokines, times)].stack(\"Time\")\n",
    "df = (df - df_min)/(df_max - df_min)\n",
    "df_proj_exp = pd.DataFrame(np.dot(df, mlpcoefs), index=df.index, columns=[\"Node 1\", \"Node 2\"])\n",
    "df_proj_exp.columns.name = \"Node\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Fit a model on all selected integral time courses\n",
    "We use the force model with matching. \n",
    "\n",
    "We need a special dictionary of parameter bounds, because in this experiment, some theta values are below $-\\pi/2$, which is not allowed in the code by default. \n",
    "\n",
    "Moreover, since the IL-2 consumption is slower in this dataset, we don't quite reach the final dynamical phase with constant slope. Hence, we can't fit properly that final slope; instead, we import the final slope of another dataset where IL-2 consumption was faster. \n",
    "\n",
    "\n",
    "#### Remark\n",
    "Instead of processing HighMI_3 and fitting a model on it, you could bypass steps 1-3 by importing a dataframe of fitter model parameter values, such as those saved by the `fit_latentspace_model.ipynb`notebook. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Process HighMI_1 to get a v2/v1 slope\n",
    "df_highmi1 = pd.concat(\n",
    "    {\"HighMI_1-{}\".format(i): pd.read_hdf(os.path.join(\"data\", \"processed\", \"HighMI_1-{}.hdf\".format(i))) \n",
    "        for i in range(1, 5)}, names=[\"Data\"])\n",
    "df_highmi1 = df_highmi1.unstack(\"Time\").loc[:, (\"integral\", cytokines, times)].stack(\"Time\")\n",
    "df_highmi1 = (df_highmi1 - df_min)/(df_max - df_min)\n",
    "df_proj_highmi1 = pd.DataFrame(np.dot(df_highmi1, mlpcoefs), index=df_highmi1.index, columns=[\"Node 1\", \"Node 2\"])\n",
    "v2v1_mean_slope = compute_v2_v1(df_proj_highmi1, slope_type=\"mean\", reject_neg_slope=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit model parameters\n",
    "We fit peptides N4 to T4 and V4 to E1 with different regularization terms. The default regularization coefficients (used for N4, ..., T4) are too large for trajectories that remain close to the origin; they cause too many trajectories to be fitted with zero parameter values. \n",
    "\n",
    "Moreover, we include a small regularization term that enforces our prior knowledge of a correlation between $a_0$, $\\tau_0$, and $\\theta$. The linear correlation enforced was determined from a prior fit without regularization terms for correlation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Fitting parameters for N4 down to T4\n",
    "pep_selection = [\"N4\", \"Q4\", \"T4\", \"A2\", \"Y3\"]\n",
    "\n",
    "fit_vars={\"Constant velocity\":[\"v0\",\"t0\",\"theta\",\"vt\"],\"Constant force\":[\"F\",\"t0\",\"theta\",\"vt\"],\n",
    "         \"Sigmoid\":[\"a0\", \"tau0\", \"theta\", \"v1\", \"gamma\"], \n",
    "         \"Sigmoid_freealpha\":[\"a0\", \"tau0\", \"theta\", \"v1\", \"alpha\", \"beta\"]}\n",
    "\n",
    "time_scale = 20.0\n",
    "duration = np.amax(times)\n",
    "\n",
    "# Special parameter boundaries for this dataset. Model \"Constant force\" is not used anymore\n",
    "# but included for backward compatibility with the code. \n",
    "special_bounds_dict = {\n",
    "    'Constant velocity':[(0, 0, -np.pi/3, 0), (6/20*time_scale, 5/20*time_scale, np.pi, 5/20*time_scale)],\n",
    "    'Constant force': [(0, 0, -np.pi/3, 0), (6/20*time_scale, 5/20*time_scale, np.pi, 5/20*time_scale)],\n",
    "    'Sigmoid':[(0, 0, -2*np.pi/3, -2/20*time_scale, time_scale/100),\n",
    "                (8/20*time_scale, (duration + 100)/time_scale, np.pi/3, 2/20*time_scale, time_scale/2)],\n",
    "    'Sigmoid_freealpha':[(0, 0, -2*np.pi/3, 0, time_scale/72, time_scale/72),\n",
    "                        (8/20*time_scale, (duration + 100)/time_scale, np.pi/3, 2/20*time_scale,\n",
    "                        time_scale/5, time_scale/2)]\n",
    "}\n",
    "\n",
    "# Fitting\n",
    "choice_model = \"Sigmoid_freealpha\"\n",
    "nparameters = len(fit_vars[choice_model])\n",
    "\n",
    "# Regularization\n",
    "regul_rate = 0.9*np.ones(len(fit_vars[choice_model]))\n",
    "regul_rate[1] = 0.7  # reduce regularization on tau0 because 0.9 forces too many trajectories to zero\n",
    "\n",
    "# Each row specifies [i, j, a, b, c] for a term of the form c*(p[i] - a*p[j] - b)^2 in the cost function \n",
    "# The array must contain ints (to slice with i, j), so we input floats as numerators over 10,000\n",
    "# Removing this regularization does not significantly change the parameter space. \n",
    "params_correls = np.asarray([\n",
    "    [1, 0, int(1.9*10000), 0, int(0.05*10000)], # tau0 = 1.9*a0\n",
    "    [2, 0, int(0.75*10000), int(-np.pi/2*10000), int(0.1*10000)]  # theta = 0.75*a0 - pi/2\n",
    "], dtype=int).T\n",
    "\n",
    "start_time = perf_counter()\n",
    "\n",
    "ret = return_param_and_fitted_latentspace_dfs(\n",
    "            df_proj_exp.loc[df_proj_exp.index.isin(pep_selection, level=\"Peptide\")], \n",
    "            choice_model, time_scale=time_scale, reg_rate=regul_rate, reject_neg_slope=v2v1_mean_slope, \n",
    "            special_bounds_dict=special_bounds_dict, correls=params_correls)\n",
    "df_params1, df_compare1, df_hess1, df_v2v11 = ret\n",
    "\n",
    "end_t = perf_counter()\n",
    "print(\"Time to fit: \", end_t - start_time, \"s\")\n",
    "del start_time, end_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Fitting parameters for V4 to E1\n",
    "pep_selection = [\"V4\", \"G4\", \"E1\"]\n",
    "\n",
    "time_scale = 20.0\n",
    "duration = np.amax(times)\n",
    "\n",
    "# Regularization: slightl lower for the weaker peptides\n",
    "regul_rate = 0.6*np.ones(len(fit_vars[choice_model]))\n",
    "regul_rate[1] = 0.4\n",
    "\n",
    "# Lightly enforcing parameter correlations\n",
    "params_correls = np.asarray([\n",
    "    [1, 0, int(1.8*10000), 0, int(0.1*10000)], # tau0 = 1.9*a0\n",
    "    [2, 0, int(0.5*10000), int(10000*-np.pi/2), int(0.1*10000)]  # theta = 0.5*a0 - 2\n",
    "], dtype=int).T\n",
    "\n",
    "start_time = perf_counter()\n",
    "\n",
    "# Use a different final slope, because for those peptides we can properly fit it\n",
    "ret2 = return_param_and_fitted_latentspace_dfs(\n",
    "            df_proj_exp.loc[df_proj_exp.index.isin(pep_selection, level=\"Peptide\")], \n",
    "            choice_model, time_scale=time_scale, reg_rate=regul_rate, reject_neg_slope=v2v1_mean_slope, \n",
    "            special_bounds_dict=special_bounds_dict, correls=params_correls)\n",
    "df_params2, df_compare2, df_hess2, df_v2v12 = ret2\n",
    "\n",
    "end_t = perf_counter()\n",
    "print(\"Time to fit: \", end_t - start_time, \"s\")\n",
    "del start_time, end_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine the fit results for the two previous sets of peptides\n",
    "df_params = df_params1.append(df_params2)\n",
    "df_compare = df_compare1.append(df_compare2)\n",
    "df_hess = df_hess1.append(df_hess2)\n",
    "df_v2v1 = df_v2v11.append(df_v2v12)\n",
    "\n",
    "# We don't need the individual dataframes anymore\n",
    "del df_params1, df_compare1, df_hess1, df_v2v11, df_params2, df_compare2, df_hess2, df_v2v12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the result of the fits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "if tcn_fit == \"all\":\n",
    "        df_params_sel = df_params\n",
    "else:  \n",
    "    df_params_sel = df_params.xs(tcellnum, level=\"TCellNumber\", axis=0)\n",
    "hue_ord = [a for a in peptides]\n",
    "h = sns.pairplot(data=df_params_sel.reset_index().query(\"Peptide in @hue_ord\"), \n",
    "                 vars=fit_vars[choice_model][:3], hue=\"Peptide\", hue_order=hue_ord, \n",
    "                 plot_kws={\"s\":16, \"alpha\":0.8}, diag_kind=\"none\")\n",
    "legend = h.legend\n",
    "\n",
    "h.fig.savefig(os.path.join(\"figures\", folder, \"pairplot_{}_HighMI_3.pdf\".format(choice_model)), \n",
    "    transparent=True, bbox_extra_artists=(legend,), bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Removing outliers\n",
    "One can notice on the plots above the small nonzero latent space trajectories for E1 and G4, which are just due to background fluctuations not completely filtered out and thus giving non-zero integrals. \n",
    "\n",
    "As a result of this background noise not completely filtered, some parameter values for E1, G4 and V4 (low concentration) do not make sense, for instance the very high $\\theta$ values that are completely off the cloud of points. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check how many points we are going to filter out. \n",
    "print(np.sum((df_params.index.isin([\"E1\"], level=\"Peptide\")*(df_params[\"theta\"] > -np.pi/3))))\n",
    "print(np.sum((df_params.index.isin([\"T4\"], level=\"Peptide\")*(df_params[\"theta\"] < -np.pi/2))))\n",
    "print(np.sum((df_params.index.isin([\"G4\"], level=\"Peptide\")*(df_params[\"theta\"] > -np.pi/3))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Those G4, V4, E1 fits don't represent the difference with other peptides and are just noise\n",
    "df_params2 = df_params.loc[~(df_params.index.isin([\"E1\"], level=\"Peptide\")*(df_params[\"theta\"] > -np.pi/3))]\n",
    "df_params2 = df_params2.loc[~(df_params2.index.isin([\"T4\"], level=\"Peptide\")*(df_params2[\"theta\"] < -np.pi/2))]\n",
    "df_params2 = df_params2.loc[~(df_params2.index.isin([\"G4\"], level=\"Peptide\")*(df_params2[\"theta\"] > -np.pi/3))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert cell to code to save the parameter fits on disk\n",
    "df_params2.to_hdf(os.path.join(\"results\", \"fits\", \"df_params_{}_{}.hdf\".format(\n",
    "    choice_model.replace(\" \", \"_\"), \"HighMI_3\")), key=\"df_params\")\n",
    "df_v2v1.to_hdf(os.path.join(\"results\", \"fits\", \"df_v2v1_{}_{}.hdf\".format(\n",
    "    choice_model.replace(\" \", \"_\"), \"HighMI_3\")), key=\"df_params\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Fit multivariate normal distributions in model parameter space \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "params_to_keep = [\"a0\", \"tau0\", \"theta\"]\n",
    "levels_group = [\"Peptide\"]\n",
    "\n",
    "# THIS IS WHERE THE GAUSSIANS ARE ACTUALLY FITTED\n",
    "# Use df_params to mix all T cell numbers, use df_params_sel to use only 100k\n",
    "if tcn_fit == \"all\":\n",
    "    ret = stats_per_levels(df_params2, levels_groupby=levels_group, feats_keep=params_to_keep)\n",
    "else:\n",
    "    ret = stats_per_levels(df_params2.xs(tcn_fit, level=\"TCellNumber\", axis=0), \n",
    "                           levels_groupby=levels_group, feats_keep=params_to_keep)\n",
    "df_params_means, df_params_means_estim_vari, df_params_covs, df_params_covs_estim_vari, ser_npts = ret\n",
    "df_params_covs_estim_vari = np.abs(df_params_covs_estim_vari)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check that the results make some sense by sampling the generated gaussians\n",
    "Create a synthetic parameter space, basically. I am sure it will look very different! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nsamples = 40\n",
    "seed = 1357642\n",
    "rnd_gen = np.random.default_rng(seed=seed)\n",
    "if len(levels_group) == 1:\n",
    "    new_index = pd.MultiIndex.from_product([df_params_means.index] + [range(nsamples)], \n",
    "                                      names=[df_params_means.index.name, \"Sample\"])\n",
    "else:\n",
    "    new_index = pd.MultiIndex.from_product([*zip(*df_params_means.index)] + [range(nsamples)], \n",
    "                                      names=[df_params_means.index.names] + [\"Sample\"])\n",
    "df_params_synth = pd.DataFrame(index=new_index, columns=params_to_keep, dtype=np.float64)\n",
    "df_params_synth.columns.name = \"Parameter\"\n",
    "\n",
    "# Sample from the fitted gaussians\n",
    "for key in df_params_means.index:\n",
    "    cov_mat = build_symmetric(df_params_covs.loc[key].values)\n",
    "    mean_vec = df_params_means.loc[key].values\n",
    "    df_params_synth.loc[key] = rnd_gen.multivariate_normal(mean_vec, cov_mat, nsamples)\n",
    "    \n",
    "# Concatenate the data and synthetic sample points, for plotting purposes   \n",
    "if tcn_fit == \"all\":\n",
    "    params_remove = list(set(df_params2.index.names).difference(levels_group))\n",
    "    df_params_both = df_params2.droplevel(params_remove).sort_index()\n",
    "else:\n",
    "    params_remove = list(set(df_params2.index.names).difference(levels_group))\n",
    "    params_remove.remove(\"TCellNumber\")\n",
    "    df_params_both = df_params2.xs(tcn_fit, level=\"TCellNumber\", axis=0).droplevel(params_remove).sort_index()\n",
    "\n",
    "df_params_both = df_params_both.loc[:, params_to_keep[0]:params_to_keep[-1]]\n",
    "print(df_params_both.groupby(levels_group).count().values)\n",
    "\n",
    "idx = np.concatenate([np.arange(n) for n in df_params_both.groupby(levels_group).count().sort_index().values[:, 0]])\n",
    "df_params_both[\"Sample\"] = idx\n",
    "df_params_both = df_params_both.set_index(\"Sample\", append=True)\n",
    "df_params_both = pd.concat([df_params_both, df_params_synth], axis=1, keys=[\"Data\", \"Synth\"], names=[\"Source\", \"Parameter\"])\n",
    "df_params_both = df_params_both.stack(\"Source\")\n",
    "print(df_params_both)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot the synthetic versus real points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peptides = [\"N4\", \"A2\", \"Y3\", \"Q4\", \"T4\", \"V4\", \"G4\", \"E1\"]\n",
    "pep_color_order = [\"N4\", \"Q4\", \"T4\", \"V4\", \"G4\", \"E1\", \"A2\", \"Y3\", \"A8\", \"Q7\"]\n",
    "pep_palette = {pep_color_order[i]:sns.color_palette()[i] for i in range(len(pep_color_order))}\n",
    "palette_order = [pep_palette.get(a) for a in peptides]\n",
    "rename_dict = {\"a0\":r\"$a_0$\", \"tau0\":r\"$\\tau_0$\", \"theta\":r\"$\\theta$\"}\n",
    "params_to_keep2 = [rename_dict.get(a, a) for a in params_to_keep]\n",
    "df_params_plot = custom_pd.xs_slice(df_params_both.rename(rename_dict, axis=1, level=\"Parameter\"), \n",
    "                                    name=\"Peptide\", lvl_slice=peptides, axis=0).reset_index()\n",
    "\n",
    "fig, axes, legend = dual_pairplot(data=df_params_plot, vari=params_to_keep2, dual_lvl=\"Source\", \n",
    "              dual_labels=[\"Data\", \"Synthetic\"], hue=\"Peptide\",\n",
    "              dual_hues = [(0.5, 0.5, 0.5), plt.cm.viridis([206])[0]], palette=palette_order, \n",
    "              hue_order=peptides, alpha=0.8, s=16, edgecolors=None)\n",
    "\n",
    "# Clean up layout\n",
    "fig.set_size_inches(6.5, 6.5)\n",
    "fig.tight_layout(h_pad=0.5, w_pad=0.5)\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Overlay data and synthetic sample points\n",
    "h = sns.pairplot(\n",
    "    data=custom_pd.xs_slice(df_params_both, name=\"Peptide\", \n",
    "                            lvl_slice=peptides, axis=0).reset_index(), \n",
    "    vars=params_to_keep, hue=\"Source\", hue_order=[\"Synth\", \"Data\"], palette=plt.cm.viridis([40, 206]),\n",
    "    plot_kws={\"s\": 16, \"alpha\":0.7})\n",
    "legend = h.legend\n",
    "\n",
    "h.fig.savefig(os.path.join(\n",
    "    \"figures\", folder, \"pairplot_synthreal_{}_HighMI_3.pdf\".format(choice_model.replace(\" \", \"_\"))), \n",
    "    transparent=True, bbox_extra_artists=(legend,), bbox_inches='tight')\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compute the Cholesky decomposition of the covariance matrices\n",
    "To preserve positive definiteness, we will have to interpolate the elements of the covariance matrix's Cholesky decomposition, $\\Sigma = L L^T$. Then, once we have $L(\\mathrm{EC}_{50})$ interpolated as a function of EC$_{50}$, we can ensure positive definiteness by constructing $\\Sigma(\\mathrm{EC}_{50}) = L(\\mathrm{EC}_{50}) L^T(\\mathrm{EC}_{50})$, which is then necessarily symmetric and positive definite because it has a Cholesky decomposition with non-negative diagonal elements. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Variance of the Cholesky decomposition estimator\n",
    "To interpolate elements of the Cholesky matrices, we need error bars on the statistical estimates of those elements. The paper Olkin, 1985, *Estimating a Cholesky Decomposition* studies a statistical estimator of the Cholesky decomposition of the covariance matrix, and the variance of that estimator. \n",
    "\n",
    "Olkin gives an unbiased estimator for $\\Psi$, which is surprisingly not just the Cholesky decomposition of the covariance matrix estimator. The problem with Olkin's estimator is that (because it is unbiased for $\\Psi$ itself) it does not reconstruct an unbiased covariance matrix, i.e. $\\langle \\hat{\\Psi} \\hat{\\Psi}^T \\rangle \\neq \\Sigma$. Because I only care about $\\Sigma$ in the end, I just define instead $\\hat{\\Psi}$ from $\\hat{\\Psi} \\hat{\\Psi}^T = \\hat{\\Sigma}$, where $\\hat{\\Sigma}$ is the covariance matrix estimator. Both estimators are convergent anyways as $K \\rightarrow \\infty$ ($K = $ number of sample points). \n",
    "\n",
    "Olkin's paper gives a correct derivation of the variance of its unbiased estimator of $\\Psi$, which can readily be applied to my biased estimator (they are related by a scaling factor). However, Olkin does not report the final result for that variance properly. Specifically the paper proves that the $u_{ij}^2$ variables follow a $\\chi^2$ distribution, and hence the variance of the $u_{ij}$ is related to the *mean* of the $\\chi^2$ distribution for $u^2$. However, Olkin uses the variance of the $chi^2$ distribution, which leads to variances of the estimator that do not decay to zero as the number of points $K \\rightarrow \\infty$. \n",
    "\n",
    "In any case, following the derivation carefully and correcting the mistake in the paper's final result,the variance of the Cholesky matrix estimator is (replacing the unknown true $\\Psi$ by $\\hat{\\Psi}$):\n",
    "\n",
    "$$ \\mathrm{Var}[\\hat{\\Psi}^{ii}] =  \\frac{K-i - a_i^2}{K-1} (\\Psi^{ii})^2 $$\n",
    "\n",
    "$$\t\\mathrm{Var}[\\hat{\\Psi}^{ij}] = \\frac{1}{K-1} \\sum_{l=j+1}^i (\\Psi^{il})^2 + \\frac{K - j - a_j^2}{K-1} \\Psi^{ij} \\quad (i \\neq j) $$\n",
    "where \n",
    "$$ a_i = \\sqrt{2} \\frac{\\Gamma((K-i+1)/2)}{\\Gamma((K-i)/2)} \\, . $$\n",
    "\n",
    "The variances on the estimator matrix elements go to zero as $K \\rightarrow \\infty$ due to the asymptotic behaviour of the $\\Gamma$ function, which implies $a_j^2 \\sim (K - j) - \\frac{1}{2}$. \n",
    "\n",
    "Therefore, all we need to estimate $L$ and the variance of that estimator is the covariance matrix estimator, $\\hat{\\Sigma}$, which we Cholesky-decompose to get our (biased) $\\hat{L}$, and the number of sample points $N$, because the variance of the estimator (which is biased, not exact like our long formula for the variance of $\\hat{\\Sigma}$ itself) only depends on elements of $\\hat{L}$ and $n$. This is highly convenient. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## See the implementation of the estimator in utils.statistics\n",
    "df_params_chol, df_params_chol_estim_vari = compute_cholesky_dataframe(df_params_covs, ser_npts)\n",
    "print(df_params_chol_estim_vari)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save fitted multivariate normal parameter distributions\n",
    "Convert the cell below to code to save. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "option = \"_highmi3_\"+choice_model.replace(\" \", \"_\")\n",
    "folderpath = os.path.join(\"results\", folder)\n",
    "         \n",
    "df_params_means.to_hdf(os.path.join(folderpath, \"df_params_means{}.hdf\".format(option)), \n",
    "    key=\"df_params_means\", mode=\"w\")\n",
    "df_params_means_estim_vari.to_hdf(os.path.join(folderpath, \"df_params_means_estim_vari{}.hdf\".format(option)), \n",
    "    key=\"df_params_means_estim_vari\", mode=\"w\")\n",
    "df_params_covs.to_hdf(os.path.join(folderpath, \"df_params_covs{}.hdf\".format(option)), \n",
    "    key=\"df_params_covs\", mode=\"w\")\n",
    "df_params_covs_estim_vari.to_hdf(os.path.join(folderpath, \"df_params_covs_estim_vari{}.hdf\".format(option)), \n",
    "    key=\"df_params_covs_estim_vari\", mode=\"w\")\n",
    "ser_npts.to_hdf(os.path.join(folderpath, \"ser_npts{}.hdf\".format(option)), \n",
    "    key=\"ser_npts\", mode=\"w\")\n",
    "\n",
    "# Also, the Cholesky decomposition\n",
    "df_params_chol.to_hdf(os.path.join(folderpath, \"df_params_chol{}.hdf\".format(option)), \n",
    "    key=\"df_params_chol\", mode=\"w\")\n",
    "df_params_chol_estim_vari.to_hdf(os.path.join(folderpath, \"df_params_chol_estim_vari{}.hdf\".format(option)), \n",
    "    key=\"df_params_chol_estim_vari\", mode=\"w\")\n",
    "\n",
    "# Save also an MultiIndex with all levels that are not in levels_group, \n",
    "# to know which conditions were used. Convert to a Series/DataFrame first. \n",
    "params_remove = list(set(df_params2.index.names).difference(levels_group))\n",
    "index_conditions_used = df_params2.index.droplevel(levels_group).unique()\n",
    "\n",
    "if len(params_remove) == 1:\n",
    "    index_conditions_used.to_series().to_json(os.path.join(\n",
    "    folderpath, \"index_data_conditions_used{}.json\".format(option)))\n",
    "else:\n",
    "    index_conditions_used.to_frame().to_json(os.path.join(\n",
    "    folderpath, \"index_data_conditions_used{}.json\".format(option)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Interpolate multivariate normal distributions as a function of EC$_{50}$\n",
    "For each mean vector and covariance Cholesky decomposition matrix element, use a smoothing spline interpolation to get the value of those elements at arbitrary $\\log_{10}{\\mathrm{EC}_{50}}$. \n",
    "\n",
    "For the interpolation of a given distribution parameter $p$ as a function of $g = \\log_{10}{\\mathrm{EC}_{50}}$, we use this procedure:\n",
    "1. Use the $p(g_i)$ values for each $g_i$ corresponding to one of the peptides\n",
    "2. For error bars, use the square root of the statistical variance on this estimated parameter, $\\sqrt{\\sigma_p^2(g_i)}$. \n",
    "3. Fit a smoothing UnivariateSpline (Dierckx 1975) on $(g_i, p(g_i))$ with weighting factors $1/\\sqrt{\\sigma_p^2(g_i)}$ and the smoothing factor set to $0.25$ times the default, that is, $0.25$ times the number of points $g_i$. \n",
    "4. Evaluate the smoothing spline at the $g_i$ to get $\\hat{p}(g_i)$. The spline may have large non-monotonous artifacts between the data points, but it should be as close to the points as the smoothing allows, so any potential large artifact should not be significant at those locations. \n",
    "5. Fit a PCHIP monotonous cubic interpolator through those values, $(g_i, \\hat{p}(g_i))$. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_ec50s_refs = pd.read_json(os.path.join(\"data\", \"misc\", \"potencies_df_2021.json\"))\n",
    "df_ec50s_refs.columns.name = \"Reference\"; df_ec50s_refs.index.name = \"Peptide\"\n",
    "print(df_ec50s_refs)\n",
    "ser_ec50s_avglog = np.log10(df_ec50s_refs).mean(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add error bars; can't use relplot or catplot anymore, \n",
    "# they don't have the option to include known standard deviations\n",
    "def plot_params_vs_logec50(df_estim, df_estim_vari, ser_x, cols_plot=None, \n",
    "                        ser_interp=None, x_name=\"Peptide\", col_wrap=3):\n",
    "    \"\"\" Optional df_interp, which contains interpolating splines for each parameter\n",
    "    as a function of the x variable (log EC50, usually). \n",
    "    \"\"\"\n",
    "    if cols_plot is None:\n",
    "        nplots = len(df_estim.columns)\n",
    "        cols_plot = df_estim.columns\n",
    "    else:\n",
    "        nplots = len(cols_plot)\n",
    "    \n",
    "    nrows = nplots // col_wrap + min(1, nplots % col_wrap)  # Add 1 if there is a remainder. \n",
    "    ncols = min(nplots, col_wrap)\n",
    "    \n",
    "    fig, axes = plt.subplots(nrows, ncols, sharex=True, sharey=False)\n",
    "    fig.set_size_inches(3*ncols, 2.5*nrows)\n",
    "    axes = axes.flatten()\n",
    "    \n",
    "    for i in range(nplots):\n",
    "        estim = df_estim[cols_plot[i]]\n",
    "        stds = np.sqrt(df_estim_vari[cols_plot[i]])\n",
    "        x_labels = estim.index.get_level_values(x_name)\n",
    "        xpoints = ser_x.reindex(x_labels)  # assume ser_x has a single index level?\n",
    "        xmin, xmax = np.amin(xpoints), np.amax(xpoints)\n",
    "        axes[i].errorbar(xpoints, estim, yerr=stds, ls='none', marker=\"o\", ms=5, label=\"Fit\")\n",
    "        axes[i].set_ylabel(cols_plot[i])\n",
    "        xr = xmax - xmin\n",
    "        yr = np.amax(estim) - np.amin(estim)\n",
    "        for j in range(len(x_labels)):\n",
    "            axes[i].annotate(x_labels[j], xy=(xpoints[j]+0.01*xr, estim[j]+0.02*yr), fontsize=8)\n",
    "        if ser_interp is not None:\n",
    "            spl = ser_interp[cols_plot[i]]\n",
    "            xrange = np.linspace(xmin, xmax, 201)\n",
    "            axes[i].plot(xrange, spl(xrange), lw=1.5, label=\"Interpolation\")\n",
    "            axes[i].legend()\n",
    "    return [fig, axes]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interpolate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "ser_splines_means = interpolate_params_vs_logec50(df_params_means, df_params_means_estim_vari, \n",
    "                                      ser_ec50s_avglog, x_name=\"Peptide\")\n",
    "\n",
    "# Plot the interpolation versus the data\n",
    "fig, axes = plot_params_vs_logec50(df_params_means, df_params_means_estim_vari, ser_ec50s_avglog, \n",
    "                             ser_interp=ser_splines_means, cols_plot=None, x_name=\"Peptide\", col_wrap=3)\n",
    "for ax in axes[-3:]:  # 3 is col_wrap\n",
    "    ax.set_xlabel(r\"$\\log_{10}{\\mathrm{EC}_{50}}$ [-]\")\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ser_splines_covs = interpolate_params_vs_logec50(df_params_covs, df_params_covs_estim_vari, \n",
    "                                      ser_ec50s_avglog, x_name=\"Peptide\")\n",
    "\n",
    "ser_splines_chol = interpolate_params_vs_logec50(df_params_chol, df_params_chol_estim_vari, \n",
    "                                      ser_ec50s_avglog, x_name=\"Peptide\")\n",
    "fig, axes = plot_params_vs_logec50(df_params_chol, df_params_chol_estim_vari, ser_ec50s_avglog, \n",
    "                             ser_interp=ser_splines_chol, cols_plot=None, x_name=\"Peptide\", col_wrap=3)\n",
    "for ax in axes[-3:]:  # 3 is col_wrap\n",
    "    ax.set_xlabel(r\"$\\log_{10}{\\mathrm{EC}_{50}}$ [-]\")\n",
    "fig.tight_layout()\n",
    "\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert cell to code to save the interpolation objects for later use\n",
    "option = \"HighMI_3\"\n",
    "with open(os.path.join(\"results\", folder, \"ser_splines_chol_{}.pkl\".format(folder, tcn_fit)), \"wb\") as hd:\n",
    "    pickle.dump(ser_splines_chol.to_dict(), hd)\n",
    "with open(os.path.join(\"results\", folder, \"ser_splines_means_{}.pkl\".format(folder, option)), \"wb\") as hd:\n",
    "    pickle.dump(ser_splines_means.to_dict(), hd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Discretize the EC50 axis\n",
    "Sampling closely spaced but still discrete $Q$ values allows us to use the Blahut-Arimoto algorithm to optimize over a discrete $p_Q$. The only different is that integrals over $\\mathcal{X}$ (for each value of $Q$) are now continuous; we perform them with Monte Carlo integration in our code for better speed. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Discretize the $log{EC_{50}}$ axis\n",
    "# Let's try $n=25$, which seems enough to extract all the information available\n",
    "n_inputs = 25\n",
    "min_logec50 = ser_ec50s_avglog.min()\n",
    "max_logec50 = ser_ec50s_avglog.max()\n",
    "bins_logec50 = np.linspace(min_logec50, max_logec50, n_inputs+1)\n",
    "# Take the midpoints of those bin separators\n",
    "sampled_logec50 = bins_logec50[:n_inputs] + np.diff(bins_logec50)/2\n",
    "print(sampled_logec50)\n",
    "\n",
    "## Compute the means and covariance matrices at each EC50, using the interpolation\n",
    "n_dims = len(df_params_means.columns)  # number of parameters\n",
    "meanmats, covmats, covmats_direct = eval_interpolated_means_covs(\n",
    "    ser_splines_means, ser_splines_covs, ser_splines_chol, sampled_logec50, \n",
    "    n_inputs, n_dims, epsil=1e-5\n",
    ")\n",
    "# Uncomment to save the interpolated channel to the results/ folder\n",
    "#np.save(os.path.join(\"results\", folder, \"meanmats_{}inputs{}.npy\".format(n_inputs, suffix)), meanmats)\n",
    "#np.save(os.path.join(\"results\", folder, \"covmats_{}inputs{}.npy\".format(n_inputs, suffix)), covmats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 6. Compute channel capacity with Blahut-Arimoto"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from wurlitzer import sys_pipes\n",
    "import json\n",
    "from datetime import date\n",
    "\n",
    "import chancapmc.chancapmc as chancapmc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "## Time and run the Blahut-Arimoto algorithm. \n",
    "seed = 42493287\n",
    "reltol = 0.01\n",
    "\n",
    "# With wurlitzer, the C code output should print here at the end of the run only\n",
    "# https://github.com/minrk/wurlitzer\n",
    "# To print it as it goes in the terminal running the notebook, remove the \"with sys_pipes():\" statement. \n",
    "start_t = perf_counter()\n",
    "\n",
    "# Use this wurlitzer function to print C-level stdout and stderr in the notebook\n",
    "with sys_pipes():\n",
    "    res = chancapmc.ba_discretein_gaussout(meanmats, covmats, sampled_logec50, reltol, seed)\n",
    "\n",
    "capacity_bits, optim_input_distrib = res\n",
    "capacity_error = capacity_bits * reltol\n",
    "\n",
    "print(\"Capacity = {} pm {}\".format(capacity_bits, capacity_error))\n",
    "print(\"Optimal input distribution:\", optim_input_distrib)\n",
    "\n",
    "run_duration = perf_counter() - start_t\n",
    "print(\"Time to converge: \", run_duration, \"s\")\n",
    "\n",
    "del start_t"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save all the info on the run in a JSON file\n",
    "today = date.today().strftime(\"%d-%b-%Y\").lower()\n",
    "run_info = {\n",
    "    \"date\": today, \n",
    "    \"capacity_bits\": capacity_bits, \n",
    "    \"input_values\": list(sampled_logec50.astype(float)), \n",
    "    \"optimal_distribution\": list(optim_input_distrib.astype(float)),\n",
    "    \"n_inputs\": n_inputs, \n",
    "    \"run_duration (s)\": run_duration, \n",
    "    \"relative_tolerance\": reltol, \n",
    "    \"absolute_error\": capacity_error, \n",
    "    \"params\": list(set(a for lbl in ser_splines_means.index \n",
    "                          for a in lbl.split(\"*\"))), \n",
    "    \"seed\": seed    \n",
    "}\n",
    "folder = \"capacity\"\n",
    "suffix = \"_HighMI_3\"\n",
    "filename = os.path.join(\"results\", folder, \"capacity_run{}_{}.json\".format(suffix, today))\n",
    "\n",
    "# UNCOMMENT TO SAVE\n",
    "with open(filename, \"w\") as hand:\n",
    "    json.dump(run_info, hand)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "slideshow": {
     "slide_type": "slide"
    }
   },
   "outputs": [],
   "source": [
    "# Make a histogram (bar plot) of the optimal input distribution\n",
    "fig, ax = plt.subplots()\n",
    "ax.bar(np.around(sampled_logec50, 2), optim_input_distrib, width=np.diff(sampled_logec50)[0])\n",
    "ax.set_xlabel(r\"$\\log_{10}{(\\mathrm{EC}_{50})}$ [-]\", size=14)\n",
    "ax.set_ylabel(\"Probability [-]\", size=14)\n",
    "ax.tick_params(which=\"both\", labelsize=12)\n",
    "ax.annotate(r\"C = {:.4f} bits $\\pm$ {:.2f} %\".format(capacity_bits, reltol*100), xy=(0.3, 0.8), \n",
    "            xycoords=\"axes fraction\", size=12)\n",
    "\n",
    "# Annotate peptides\n",
    "maxprob = np.amax(optim_input_distrib)\n",
    "for pep in ser_ec50s_avglog.index:\n",
    "    if pep not in df_params_means.index:\n",
    "        continue\n",
    "    ax.annotate(pep, xy=(ser_ec50s_avglog[pep], maxprob), fontsize=8, ha=\"center\")\n",
    "    ax.axvline(ser_ec50s_avglog[pep], ls=\":\", lw=0.5)\n",
    "\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}