reconstruct_cytokines_fromLSmodel_pvalues.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Latent space dynamical model plus reconstruction give a model for cytokine dynamics\n",
    "To run this notebook, you need:\n",
    "- Raw cytokine time series (to estimate noise in the data) in the `data/final/` folder. \n",
    "- Pre-processed cytokine time series in the `data/processed/` folder\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",
    "- For plotting, axes tick properties stored in JSON files in `data/misc/`. \n",
    "Those files are available in the data repository hosted online, or you can generate them yourself from raw cytokine data using [`cytokine-pipeline`](https://github.com/tjrademaker/cytokine-pipeline). \n",
    "\n",
    "To reproduce the results reported in the paper, the following datasets must be in `data/final`:\n",
    "- `\"cytokineConcentrationPickleFile-20190404-PeptideComparison_1-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20190412-PeptideComparison_2-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20190608-PeptideComparison_3-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20190718-PeptideComparison_4-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20190802-TCellNumber_1-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20191022-PeptideComparison_7-final.hdf`\n",
    "- `\"cytokineConcentrationPickleFile-20200220-TCellNumber_3-final.hdf\"` \n",
    "- `\"cytokineConcentrationPickleFile-20200624-HighMI_1-1-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20200624-HighMI_1-2-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20200624-HighMI_1-3-final.hdf\"`\n",
    "- `\"cytokineConcentrationPickleFile-20200624-HighMI_1-4-final.hdf\"` \n",
    "- `\"cytokineConcentrationPickleFile-20200627-TCellNumber_4-final.hdf\"` \n",
    "\n",
    "Using a different selection of datasets will produce similar, but not identical, p-values for the model fits. \n",
    "\n",
    "## Main steps in the model fitting\n",
    "We can use our latent space model, combined with our cytokine reconstruction method (linear regression with quadratic terms), i.e. our \"latent space decoder\", as a dynamical model that accurately fits 5D cytokine concentration time series (in log scale). \n",
    "\n",
    "The first main part consists in fitting a dynamical model to data projected in latent space, and reconstructing cytokine trajectories from model trajectories: \n",
    "1. Import naive OT-1 data to be fitted with the model\n",
    "    1. Select datasets to optimize the reconstruction algorithm (\"training\" data)\n",
    "    2. Select a dataset for the plot (\"test\" data)\n",
    "2. Optimize the reconstruction algorithm\n",
    "3. Fit the force model with matching to that dataset in the latent space\n",
    "4. Compute the model concentration curves corresponding to the fitted parameters values\n",
    "5. Project back those curves to cytokine concentration space\n",
    "\n",
    "The second main part consists in comparing the original cytokine time courses and the ones generated from the ballistic parameters:\n",
    "1. Use more naive OT-1 data to estimate noise (error bars and covariance elements at each time point) on the cytokine data\n",
    "2. Compute residuals between the original time courses and the model-generated, reconstructed cytokine trajectories\n",
    "3. Compute a $\\chi^2$ for each time course, summing the multivariate chi-squared at each 5D time point: $\\chi^2 = \\sum_t \\sum_{ij} c_i(t) {C^{-1}}_{ij}(t) c_j(t)$ where $c_i(t)$ is cytokine i at time $t$. \n",
    "4. The number of degrees of freedom $\\nu$ is (nb time points) $\\times$ (nb cytokines) - (nb model parameters fitted on a time course) $= 12 \\times 5 - 6 = 54$. This takes into account correlations between cytokines but neglects correlations between time points, because we don't have enough data to estimate one large covariance matrix between all time points of all cytokines. In other words, we take this hypothetical big covariance matrix and assume it is block-diagonal, with one block per time point. \n",
    "5. Compute p-value for each time course fit with the chi-squared distribution. This p-value gives the probability that a correct model of the data would give a equal or larger $\\chi^2$, taking into account the amount of noise (i.e. the covariance matrix) of the data. \n",
    "\n",
    "\n",
    "\n",
    "## Code structure\n",
    "\n",
    "The following useful functions are in separate Python scripts, for clarity of the notebook. \n",
    "\n",
    "- Scripts to import and process data: `ltspcyt.scripts.neural_network`\n",
    "- Scripts to optimize a reconstruction model (on \"reconstruction training\" data): `ltspcyt.scripts.reconstruction`\n",
    "- Scripts to fit the sigmoid-ballistic model (on separate \"reconstruction test\" data): `ltspcyt.scripts.sigmoid_ballistic` and `fitting_functions`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "yrlPiV29ae9J"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import scipy as sp\n",
    "import scipy.stats\n",
    "import json, pickle\n",
    "from time import perf_counter\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib as mpl\n",
    "import seaborn as sns\n",
    "\n",
    "# Scripts for data importation\n",
    "from ltspcyt.scripts.neural_network import import_WT_output\n",
    "\n",
    "# Scripts for reconstruction, using distinct functions for distinct methods. \n",
    "from ltspcyt.scripts.reconstruction import train_reconstruction, plot_recon_true, compute_latent_curves\n",
    "\n",
    "# Scripts for curve fitting\n",
    "from ltspcyt.scripts.sigmoid_ballistic import (\n",
    "    return_param_and_fitted_latentspace_dfs, sigmoid_conc_full_freealpha, \n",
    "    ballistic_sigmoid_freealpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part I: Model fitting and reconstruction\n",
    "\n",
    "# I.1 Import data\n",
    "\n",
    "### I.1.1 Import all data\n",
    "Remove unwanted levels, normalize the values with the neural network training data's min and max. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_wt = import_WT_output()\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",
    "df_min, df_max = df_min.xs(\"integral\", level=\"Feature\"), df_max.xs(\"integral\", level=\"Feature\")\n",
    "\n",
    "# Projection matrix\n",
    "P = np.load(os.path.join(\"data\", \"trained-networks\", \"mlp_input_weights-thomasRecommendedTraining.npy\")).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "peptides = [\"N4\", \"Q4\", \"T4\", \"V4\", \"G4\", \"E1\", \"A2\", \"Y3\", \"A8\", \"Q7\"]\n",
    "concentrations = [\"1uM\", \"100nM\", \"10nM\", \"1nM\"]\n",
    "cytokines = df_min.index.get_level_values(\"Cytokine\")\n",
    "times = np.arange(0, 73)\n",
    "\n",
    "# Select only the desired cytokines, times, and T cell number\n",
    "df_wt = df_wt.unstack(\"Time\").loc[:, (slice(None), cytokines, times)].stack(\"Time\")\n",
    "\n",
    "# Rescale and project each feature, but do not offset (don't use MLP's intercepts)\n",
    "proj_dfs = []\n",
    "feat_keys = [\"integral\", \"concentration\", \"derivative\"]\n",
    "cols = pd.Index([\"Node 1\", \"Node 2\"], name=\"Node\", copy=True)\n",
    "\n",
    "for typ in feat_keys:\n",
    "    # Rescale with the training min and max\n",
    "    if typ == \"integral\":\n",
    "        df_wt[typ] = (df_wt[typ] - df_min)/(df_max - df_min)\n",
    "    else:   # for conc and deriv, the constant rescaling term disappears. \n",
    "        df_wt[typ] = df_wt[typ]/(df_max - df_min)\n",
    "    df_temp = pd.DataFrame(np.dot(df_wt[typ], P.T), index=df_wt[typ].index, columns=cols)\n",
    "    proj_dfs.append(df_temp)\n",
    "df_proj = pd.concat(proj_dfs, axis=1, names=[\"Feature\"], keys=feat_keys)\n",
    "del proj_dfs, cols, feat_keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove different T cell numbers\n",
    "tcellnum = \"100k\"\n",
    "df_wt = df_wt.xs(tcellnum, level=\"TCellNumber\", axis=0, drop_level=True)\n",
    "df_proj = df_proj.xs(tcellnum, level=\"TCellNumber\", axis=0, drop_level=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### I.1.2 Select training and testing datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "subset_train = [\"HighMI_1-1\", \"HighMI_1-3\"]\n",
    "\n",
    "df_wt_train = df_wt.loc[subset_train]\n",
    "df_proj_train = df_proj.loc[subset_train]\n",
    "\n",
    "subset_test = [\"HighMI_1-2\", \"HighMI_1-4\"]\n",
    "\n",
    "df_wt_test = df_wt.loc[subset_test]\n",
    "df_proj_test = df_proj.loc[subset_test]\n",
    "\n",
    "# Remove A2 and Y3 from the training\n",
    "df_wt_train = df_wt_train.drop([\"A2\", \"Y3\"], level=\"Peptide\", axis=0)\n",
    "df_proj_train = df_proj_train.drop([\"A2\", \"Y3\"], level=\"Peptide\", axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I.2 Train the reconstruction function\n",
    "Also compute the reconstructed cytokines for both the test and training data sets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the reconstruction matrix, based on reconstructing integrals\n",
    "feature = \"concentration\"\n",
    "model_type = \"mixed_quad\"\n",
    "\n",
    "modelargs = {\"which_to_square\":[0, 1]}\n",
    "\n",
    "# Add some arbitrary features. \n",
    "# Try exponentials\n",
    "tanh_norm_factors = df_proj_train[\"integral\"].mean(axis=0)\n",
    "print(tanh_norm_factors)\n",
    "\n",
    "df_proj_train = pd.concat([df_proj_train[\"concentration\"], np.tanh(df_proj_train[\"integral\"] / tanh_norm_factors)], \n",
    "                           keys=[\"concentration\", \"tanh integral\"], names=[\"Feature\"], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pipe, score = train_reconstruction(df_proj_train, df_wt_train, feature=feature, \n",
    "                                   method=model_type, model_args=modelargs, do_scale_out=False)\n",
    "print(\"Reconstruction training R^2 score:\", score)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I.3 Fit the ballistic model to test data\n",
    "We fit the latent space integrals $LS_1$ and $LS_2$, rescaling time by $\\tilde{t} = 20 h$, as we usually do. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Choice of fitting hyperparameters\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",
    "fit = \"Sigmoid_freealpha\"\n",
    "regul_rate = 0.4\n",
    "tscale = 20.\n",
    "\n",
    "# Fit the integrals\n",
    "start_time = perf_counter()\n",
    "\n",
    "ret = return_param_and_fitted_latentspace_dfs(\n",
    "    df_proj_test.xs(\"integral\", level=\"Feature\", axis=1), \n",
    "    fit, reg_rate=regul_rate, time_scale=tscale)\n",
    "df_params, df_compare, df_hess, ser_v2v1 = ret\n",
    "\n",
    "end_t = perf_counter()\n",
    "print(\"Time to fit: \", perf_counter() - start_time)\n",
    "del start_time\n",
    "\n",
    "nparameters = len(fit_vars[fit])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I.4 Compute ballistic curves for fitted parameters\n",
    "The normalization is a bit tricky here. We fitted $LS_i(t')$, where $t' = t/\\tilde{t}$, $\\tilde{t} = 20 $ h (the time scale). Now, we want $ls_i(t) = \\frac{d N_i}{dt} = \\frac{d t'}{dt} \\frac{d N_i(t')}{d t'} = \\frac{1}{\\tilde{t}} n_i(t', a_0', \\ldots)$, where $ls_i(t', a_0', \\ldots)$ is the function $ls_i(t)$ called with $t'$ and parameters fitted for $LS_i(t')$, instead of $t$: same functional form, different scale of variables. We need to compensate this by dividing by $\\tilde{t}$, because when we fitted $LS_i(t')$, we had the following things happening:\n",
    " - To preserve $\\alpha t = \\alpha' t'$ and $\\beta t = \\beta' t'$ in the exponentials with $t$ replaced by $t'$, $\\alpha' = \\tilde{t} \\alpha$\n",
    " - Because the magnitude of $LS_i$ is proportional to $a_0 / \\alpha$ or $v_i / \\alpha$, then $a_0' = \\tilde{t} a_0$, i.e. the fitted value for $a_0$ is too large in reality by a factor of $t_0$\n",
    " - Since the functional form of $ls_i$ has a magnitude proportional to $a_0$ only, calling that function with the fitted value of $a_0'$ would give a concentration too large by a factor $\\tilde{t}$ for being truly $\\frac{d N_i}{dt}$. \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extend ser_v2v1 to have one entry per entry in df_params\n",
    "ser_v2v1_synth = pd.Series(np.zeros(len(df_params.index)), index=df_params.index)\n",
    "for dset in ser_v2v1.index:\n",
    "    ser_v2v1_synth[dset] = ser_v2v1[dset]\n",
    "\n",
    "# Create a new df_compare, by concatenation.\n",
    "df_latent_fit = compute_latent_curves(df_params.loc[:, :\"beta\"], ser_v2v1_synth, tanh_norm_factors, times,\n",
    "    model=\"Sigmoid_freealpha\", tsc=tscale)\n",
    "\n",
    "# Replace the values in df_compare\n",
    "df_compare2 = df_compare.unstack(\"Processing type\").unstack(\"Feature\").stack(\"Node\").unstack(\"Node\")\n",
    "\n",
    "df_compare2[(\"Fit\", \"concentration\")] = df_latent_fit[\"concentration\"]  # each row is a time\n",
    "df_compare2[(\"Fit\", \"tanh integral\", \"Node 1\")] = df_latent_fit[(\"tanh integral\", \"Node 1\")]\n",
    "df_compare2[(\"Fit\", \"tanh integral\", \"Node 2\")] = df_latent_fit[(\"tanh integral\", \"Node 2\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I.5 Reconstruct cytokines from generated curves"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_recon_test = pd.DataFrame(pipe.predict(df_latent_fit), index=df_latent_fit.index, \n",
    "                             columns=df_wt_test.xs(feature, axis=1, level=\"Feature\", drop_level=False).columns)\n",
    "\n",
    "# Concentrations can't be negative but sometimes this reconstruction gives slightly negative values\n",
    "# So the last part of our reconstructino algorithm is to clip values to zero\n",
    "df_recon_test.clip(lower=0.0, inplace=True)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# Save result if desired; convert cell to code cell to execute\n",
    "folder_file = os.path.join(\"results\", \"reconstruction\", \"df_compare_recon-fit_HighMI-1.hdf\")\n",
    "df_wt_test.to_hdf(folder_file, key=\"df_wt\")\n",
    "df_recon_test.loc[df_recon_test.index.get_level_values(\"Time\") > 0].to_hdf(folder_file, key=\"df_recon\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I.6 Compare visually the generated cytokines to the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "figlist = plot_recon_true(df_wt_test, df_recon_test.loc[df_recon_test.index.get_level_values(\"Time\") > 0], \n",
    "                          feature=feature, sharey=True, do_legend=True, \n",
    "                          palette=sns.color_palette(), pept=peptides)\n",
    "for xp in figlist.keys():\n",
    "    print(xp)\n",
    "    legend = figlist[xp].axes[-1].get_legend()\n",
    "    plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RIsX8cc7lLdM"
   },
   "source": [
    "### Adding E1 from another dataset\n",
    "It is not available in the `HighMI_1` experiment, but we would like to see its reconstruction anyways, so we reconstruct the model fits on another experiment (`PeptideComparison_1`) using the same reconstruction coefficients as for `HighMI_1`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select data for E1\n",
    "dset_with_e1 = [\"PeptideComparison_1\"]\n",
    "df_wt_e1 = df_wt.loc[dset_with_e1]\n",
    "df_proj_e1 = df_proj.loc[dset_with_e1]\n",
    "\n",
    "# Fit the model on N_1 and N_2\n",
    "fit = \"Sigmoid_freealpha\"\n",
    "regul_rate = 0.4\n",
    "tscale = 20.\n",
    "nparameters = len(fit_vars[fit])\n",
    "\n",
    "start_time = perf_counter()\n",
    "\n",
    "ret = return_param_and_fitted_latentspace_dfs(\n",
    "    df_proj_e1.xs(\"integral\", level=\"Feature\", axis=1), \n",
    "    fit, reg_rate=regul_rate, time_scale=tscale)\n",
    "df_params_e1, df_compare_e1, df_hess_e1, ser_v2v1_e1 = ret\n",
    "\n",
    "end_t = perf_counter()\n",
    "print(\"Time to fit: \", perf_counter() - start_time)\n",
    "del start_time, end_t\n",
    "\n",
    "\n",
    "# Extend ser_v2v1 to have one entry per entry in df_params\n",
    "ser_v2v1_synth_e1 = pd.Series(np.zeros(len(df_params_e1.index)), index=df_params_e1.index)\n",
    "for dset in ser_v2v1_e1.index:\n",
    "    ser_v2v1_synth_e1[dset] = ser_v2v1_e1[dset]\n",
    "\n",
    "# Create a new df_compare, by concatenation.\n",
    "df_latent_fit_e1 = compute_latent_curves(df_params_e1.loc[:, :\"beta\"], ser_v2v1_synth_e1, tanh_norm_factors, times,\n",
    "    model=\"Sigmoid_freealpha\", tsc=tscale)\n",
    "\n",
    "# Replace the values in df_compare\n",
    "df_compare_e1 = df_compare_e1.unstack(\"Processing type\").unstack(\"Feature\").stack(\"Node\").unstack(\"Node\")\n",
    "\n",
    "df_compare_e1[(\"Fit\", \"concentration\")] = df_latent_fit_e1[\"concentration\"]  # each row is a time\n",
    "df_compare_e1[(\"Fit\", \"tanh integral\", \"Node 1\")] = df_latent_fit_e1[(\"tanh integral\", \"Node 1\")]\n",
    "df_compare_e1[(\"Fit\", \"tanh integral\", \"Node 2\")] = df_latent_fit_e1[(\"tanh integral\", \"Node 2\")]\n",
    "\n",
    "# REconstruct cytokines for the model curves fitted on this dataset\n",
    "df_recon_e1 = pd.DataFrame(pipe.predict(df_latent_fit_e1), index=df_latent_fit_e1.index, \n",
    "                             columns=df_wt_e1.xs(feature, axis=1, level=\"Feature\", drop_level=False).columns)\n",
    "\n",
    "# Concentrations can't be negative but sometimes this reconstruction gives slightly negative values\n",
    "# So the last part of our reconstructino algorithm is to clip values to zero\n",
    "df_recon_e1.clip(lower=0.0, inplace=True)"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "# optional: save the E1 reconstruction results; convert cell to code to execute\n",
    "df_recon_e1.to_hdf(os.path.join(\"results\", \"reconstruction\", \n",
    "    \"df_compare_recon-fit_PeptideComparison_1.hdf\"), key=\"df_recon\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SsxMWM3wae9P"
   },
   "source": [
    "# Part II: Quality of model fits ($\\chi^2$ p-values)\n",
    "\n",
    "## II.1 Importing a few functions to estimate noise from raw data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Functions to estimate noise from data and put back absolute scale of reconstructions\n",
    "from utils.recon_scaling import extract_process_naive_part, import_folder_naive_data, scale_back"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UmxF7seiae9Q"
   },
   "source": [
    "## II.2 Import reconstruction and original data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EWMmFnJ2ae9R"
   },
   "outputs": [],
   "source": [
    "df_recon = scale_back(df_recon_test, df_min, df_max)\n",
    "# Or, import the result saved to disk\n",
    "#df_recon = scale_back(pd.read_hdf(os.path.join(\n",
    "#    \"results\", \"reconstruction\"\", df_compare_recon-fit_HighMI-1.hdf\"), key=\"df_recon\"), df_min, df_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dHSKqGsNae9R"
   },
   "outputs": [],
   "source": [
    "# Load the original HighMI_1 data, log-scale it, select TCellNumber of interest\n",
    "df_orig_full = {}\n",
    "for i in range(1, 5):\n",
    "    df_orig_full[\"HighMI_1-{}\".format(i)] = pd.read_hdf(os.path.join(\"data\", \"final\", \n",
    "                    \"cytokineConcentrationPickleFile-20200624-HighMI_1-{}-final.hdf\".format(i)))\n",
    "\n",
    "df_orig_full = pd.concat(df_orig_full, names=[\"Data\"])\n",
    "df_orig_full = (df_orig_full.loc[df_orig_full.index.isin(cytokines, level=\"Cytokine\")]\n",
    "                .unstack(\"Cytokine\").stack(\"Time\"))\n",
    "\n",
    "# Rescale by the minimum concentration (lower LOD) and take the log10\n",
    "dset_choice = \"HighMI_1-2\"\n",
    "df_orig_full = np.log10(df_orig_full / df_orig_full.min(axis=0))\n",
    "df_orig_full = df_orig_full.xs(\"100k\", level=\"TCellNumber\", axis=0)\n",
    "df_orig = df_orig_full.xs(dset_choice, level=\"Data\", axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lsZsI9czae9S"
   },
   "source": [
    "## II.3 Prepare more raw data to estimate noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "twX6dFaBae9S"
   },
   "outputs": [],
   "source": [
    "# Similar datasets that we use to get a reasonable noise estimate:\n",
    "# P-values go up it PeptideComparison_8 and 9 are included, because\n",
    "# they increase the estimate of cytokine noise, because they have some\n",
    "# experimental variability compared to other datasets. \n",
    "raw_data_list = [\n",
    "    \"cytokineConcentrationPickleFile-20190404-PeptideComparison_1-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20190412-PeptideComparison_2-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20190608-PeptideComparison_3-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20190718-PeptideComparison_4-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20190802-TCellNumber_1-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20191022-PeptideComparison_7-final.hdf\",\n",
    "    #\"cytokineConcentrationPickleFile-20191029-PeptideComparison_8-final.hdf\",\n",
    "    #\"cytokineConcentrationPickleFile-20191106-PeptideComparison_9-final.hdf\", \n",
    "    \"cytokineConcentrationPickleFile-20200220-TCellNumber_3-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20200624-HighMI_1-1-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20200624-HighMI_1-2-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20200624-HighMI_1-3-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20200624-HighMI_1-4-final.hdf\",\n",
    "    \"cytokineConcentrationPickleFile-20200627-TCellNumber_4-final.hdf\"\n",
    "]\n",
    "\n",
    "df_raw_data, df_nM_min_conc =  import_folder_naive_data(os.path.join(\"data\", \"final\"), raw_data_list)\n",
    "df_raw_data = df_raw_data.loc[:, df_raw_data.columns.isin(cytokines, level=\"Cytokine\")]\n",
    "df_raw_data = df_raw_data.xs(\"100k\", level=\"TCellNumber\", axis=0)\n",
    "\n",
    "# Keep only the time points and peptides of interest\n",
    "df_raw_data = df_raw_data.loc[(slice(None), *df_orig.index.levels)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DurgGANVae9S"
   },
   "outputs": [],
   "source": [
    "# Computing the deviations from the average across the selected raw data sets\n",
    "all_lvl_names = list(df_raw_data.index.names)\n",
    "all_lvl_names.remove(\"Data\")\n",
    "\n",
    "# More fair noise estimation: per time point, peptide, concentration\n",
    "#df_std = np.sqrt(df_raw_data.groupby(all_lvl_names).var())\n",
    "# Equivalent way:\n",
    "df_diff = df_raw_data - df_raw_data.groupby(all_lvl_names).mean()\n",
    "df_std = np.sqrt((df_diff**2).groupby(all_lvl_names).sum() / (df_diff.groupby(all_lvl_names).count() - 1))\n",
    "\n",
    "# Combining all residuals to get one single variance estimate per cytokine\n",
    "ser_std = np.sqrt((df_diff**2).sum(axis=0) / (df_raw_data.shape[0] - 1))\n",
    "# Some points in some peptide conditions have zero variance, e.g. IL-2 at 72hrs for V4 is exactly 0, always. \n",
    "# In this case, replace the variance with the ser_std value below (average variance for the cytokine overall)\n",
    "# Actually, replace all points where the variance is smaller than ser_std; that's an artifact\n",
    "# ser_std is already low because it averages including all points where the variance is zero. \n",
    "for cyto in df_std.columns:\n",
    "    where_wrong = df_std[cyto] < ser_std[cyto]\n",
    "    df_std.loc[where_wrong, cyto] = ser_std[cyto]\n",
    "df_std"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zh2yOUEpae9T"
   },
   "source": [
    "## II.4 Compute residuals and variance per cytokine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XbRtvSK9ae9T"
   },
   "outputs": [],
   "source": [
    "# Compute the chi-square of the reconstruction as a fit to the original data, \n",
    "# one per cytokine (not multivariate just yet). Then, obtain a p-value \n",
    "\n",
    "# Select the conditions we will plot\n",
    "conc_choice = [\"1uM\"]\n",
    "dset_choice = \"HighMI_1-2\"\n",
    "peps_to_plot = [\"N4\", \"Q4\", \"T4\", \"V4\"]\n",
    "\n",
    "# Remove all unwanted conditions\n",
    "df_orig_select = df_orig.loc[(peps_to_plot, conc_choice), :]\n",
    "df_recon_select = df_recon.loc[(dset_choice, peps_to_plot, conc_choice), \"concentration\"].droplevel(\"Data\")\n",
    "df_std_select = df_std.loc[(peps_to_plot, conc_choice), :]\n",
    "\n",
    "# We should get one p-value per time course to plot. \n",
    "traject_indices = list(df_orig.index.names)\n",
    "traject_indices.remove(\"Time\")\n",
    "\n",
    "# Find the time points available in both original data and recon\n",
    "# for the conditions we already selected above. \n",
    "common_times = df_orig_select.index.get_level_values(\"Time\").unique()\n",
    "common_times = [a for a in common_times if a in df_recon_select.index.get_level_values(\"Time\").unique()]\n",
    "df_recon_select = df_recon_select.loc[(slice(None), slice(None), common_times)]\n",
    "df_orig_select = df_orig_select.loc[(slice(None), slice(None), common_times)]\n",
    "df_std_select = df_std_select.loc[(slice(None), slice(None), common_times)]\n",
    "\n",
    "# Now, compute the chisquare, grouped by traject_indices\n",
    "df_resids_select = (df_orig_select - df_recon_select)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tPndqtrUae9U"
   },
   "source": [
    "## II.5 Multivariate $\\chi^2$ distribution and p-value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oNSsMK9Hae9V"
   },
   "outputs": [],
   "source": [
    "# Multivariate chisq: each cytokine at a given time point is a different random variable\n",
    "# So we have 5 times more points (still fitted with 6 d.o.f total), but we must\n",
    "# take the covariance matrix between cytokine residuals into account\n",
    "df_cov_mat = df_diff.groupby(all_lvl_names).cov()\n",
    "# Remove extremely small eigenvalues and invert each matrix\n",
    "min_eigenval = float(ser_std.min())**2\n",
    "df_invcov_mat = df_cov_mat.copy()\n",
    "df_cov_mat_corrected = df_cov_mat.copy()\n",
    "for lvl in df_std.index:\n",
    "    mat = df_cov_mat.loc[lvl].values\n",
    "    umat, sigmat, vmat = np.linalg.svd(mat)\n",
    "    # Remove extremely small eigenvalues\n",
    "    sigmat[sigmat < min_eigenval/2] = min_eigenval/2\n",
    "    inv_eigenvals = 1.0 / sigmat\n",
    "    # Rebuild covmat and inverse cov mat without small eigenvalues\n",
    "    covmat = umat.dot(np.diagflat(sigmat)).dot(vmat)\n",
    "    df_cov_mat_corrected.loc[lvl] = covmat\n",
    "    invcovmat = umat.dot(np.diagflat(inv_eigenvals)).dot(vmat)\n",
    "    df_invcov_mat.loc[lvl] = invcovmat\n",
    "    \n",
    "print(df_invcov_mat.loc[(\"N4\", \"1uM\", 24.0)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lhfzAbNeae9V"
   },
   "outputs": [],
   "source": [
    "# Compute the multivariate chisquare for each time series\n",
    "df_invcov_mat_select = df_invcov_mat.loc[(peps_to_plot, conc_choice, common_times), :]\n",
    "ser_mv_chisq = {}\n",
    "for key in df_orig_select.index.droplevel(\"Time\").unique():\n",
    "    chisq = 0.0\n",
    "    df_mat = df_invcov_mat_select.loc[key]\n",
    "    resids = df_resids_select.loc[key]\n",
    "    for t in common_times:\n",
    "        mat = df_mat.xs(t, level=\"Time\")\n",
    "        chisq += resids.loc[t].values.dot(mat).dot(resids.loc[t].values.T)\n",
    "    ser_mv_chisq[key] = chisq\n",
    "\n",
    "# Convert to Series\n",
    "ser_mv_chisq = pd.Series(ser_mv_chisq, index=df_orig_select.index.droplevel(\"Time\").unique())\n",
    "\n",
    "# Normalize chi squares. Number of dof: timepoints * nb_cytokines - nb_params\n",
    "nu_dof5 = len(common_times)*len(cytokines) - 6\n",
    "ser_mv_chisq_norm = ser_mv_chisq / nu_dof5\n",
    "print(ser_mv_chisq_norm)\n",
    "\n",
    "chi_distrib5 = sp.stats.chi2(df=nu_dof5)\n",
    "ser_pval_mv = 1.0 - chi_distrib5.cdf(ser_mv_chisq)\n",
    "ser_pval_mv = pd.Series(ser_pval_mv, index=ser_mv_chisq.index)\n",
    "print(\"Multivariate p-values (high is good):\")\n",
    "print(ser_pval_mv)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RIsX8cc7lLdM"
   },
   "source": [
    "### Adding E1 from another dataset\n",
    "The estimation of the noise is set to the minimum value in ser_std for each cytokine. It is not a good idea to include E1 in the covariance matrix above, as it would bias all values to artificially low noise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "szaBN9a7lv04",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Put back actual cytokine scale and keep E1 only\n",
    "df_recon_e1 = scale_back(df_recon_e1, df_min, df_max)\n",
    "df_recon_e1 = df_recon_e1.xs(\"E1\", level=\"Peptide\", drop_level=False)\n",
    "df_recon_e1 = df_recon_e1.droplevel(\"Data\").loc[:, \"concentration\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nouWy3H4mNhp"
   },
   "outputs": [],
   "source": [
    "df_orig_e1 = pd.read_hdf(os.path.join(\"data\", \"final\", \n",
    "    \"cytokineConcentrationPickleFile-20190404-PeptideComparison_1-final.hdf\"))\n",
    "cytokines = [\"IFNg\", \"IL-17A\", \"IL-2\", \"IL-6\", \"TNFa\"]\n",
    "df_orig_e1 = df_orig_e1.loc[df_orig_e1.index.isin(cytokines, level=\"Cytokine\")].unstack(\"Cytokine\").stack(\"Time\")\n",
    "\n",
    "# Rescale by the minimum concentration (lower LOD) and take the log10\n",
    "df_orig_e1 = np.log10(df_orig_e1 / df_orig_e1.min(axis=0))\n",
    "\n",
    "# Select TCellNumber and E1 only\n",
    "df_orig_e1 = df_orig_e1.xs(\"100k\", level=\"TCellNumber\", axis=0)\n",
    "df_orig_e1 = df_orig_e1.xs(\"E1\", level=\"Peptide\", drop_level=False)\n",
    "df_orig_e1 = df_orig_e1.reset_index()\n",
    "df_orig_e1 = df_orig_e1.set_index([\"Peptide\", \"Concentration\", \"Time\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6MBDEyUkoWZE"
   },
   "outputs": [],
   "source": [
    "# Add a covariance matrix for each time point of the E1 time series\n",
    "# Just use ser_std on the diagonal for the cytokines, assume no correlation. \n",
    "df_cov_e1 = pd.DataFrame(np.zeros([len(df_orig_e1.index)*len(df_orig_e1.columns), len(df_orig_e1.columns)]), \n",
    "    index=pd.MultiIndex.from_product([*df_orig_e1.index.levels, df_orig_e1.columns]), \n",
    "    columns=df_orig_e1.columns)\n",
    "df_invcov_e1 = df_cov_e1.copy()\n",
    "for key in df_orig_e1.index:\n",
    "    for i in range(len(df_orig_e1.columns)):\n",
    "        df_cov_e1.loc[key].iloc[i, i] = ser_std.iloc[i]**2\n",
    "        df_invcov_e1.loc[key].iloc[i, i] = 1 / ser_std.iloc[i]**2\n",
    "print(df_invcov_e1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JkXNhdHxqgfY"
   },
   "outputs": [],
   "source": [
    "e1_times = df_orig_e1.index.get_level_values(\"Time\").unique()\n",
    "df_resids_e1 = df_orig_e1 - df_recon_e1.loc[df_recon_e1.index.isin(e1_times, level=\"Time\")]\n",
    "df_resids_e1 = df_resids_e1.dropna()  # Times not available\n",
    "e1_times = df_resids_e1.index.get_level_values(\"Time\").unique()\n",
    "chisq_e1 = 0.0\n",
    "for t in e1_times:\n",
    "    mat = df_invcov_e1.xs(t, level=\"Time\")\n",
    "    resids = df_resids_e1.xs(t, level=\"Time\")\n",
    "    chisq_e1 += float(resids.values.dot(mat).dot(resids.values.T))\n",
    "print(\"chi^2:\", chisq_e1)\n",
    "# Compute normalized chi-squared and p-value\n",
    "chisq_e1_norm = chisq_e1 / nu_dof5\n",
    "print(\"Normalized chi^2:\", chisq_e1_norm)\n",
    "pvalue_e1 = 1.0 - chi_distrib5.cdf(chisq_e1)\n",
    "print(\"p-value:\", pvalue_e1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ufEgCHrsae9V"
   },
   "source": [
    "# III. Plotting reconstructed time courses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dG03ui7uhUH2"
   },
   "outputs": [],
   "source": [
    "# Retrive the minimum concentration in HighMI_1 for proper scaling back\n",
    "# This is a constant: log(cyto/min) = log(cyto/units) - log(min/units)\n",
    "# that we will add back to all data before plotting\n",
    "pM_offset = 1000 * df_nM_min_conc.loc[dset_choice]\n",
    "print(pM_offset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8waX2j5GzIWN"
   },
   "outputs": [],
   "source": [
    "# Logarithmic minor ticks (we plotted the real log so need to put log ticks manually)\n",
    "# Find the linear scale limiting ticks\n",
    "def compute_log_minor_ticks(logylims, stp=2, base=10.0):\n",
    "    smallest_major = int(np.floor(logylims[0]))\n",
    "    largest_major = int(np.ceil(logylims[1]))\n",
    "    n_decades = largest_major - smallest_major\n",
    "\n",
    "    # Generate linear ranges with the exponents found\n",
    "    tiles = []\n",
    "    for i in range(n_decades):\n",
    "        tiles.append(np.arange(stp*base**(smallest_major+i), \n",
    "                    base**(smallest_major+i+1), stp*base**(smallest_major+i)))\n",
    "    minorticks = np.concatenate(tiles, axis=0)\n",
    "    minorticks = np.log(minorticks) / np.log(base)\n",
    "    minorticks = minorticks[(minorticks > logylims[0]) * (minorticks < logylims[1])]\n",
    "    return minorticks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "4muNGW6xA-mU"
   },
   "outputs": [],
   "source": [
    "### Plot comparing the data to the reconstruction, with error bars\n",
    "## FIRST ROW: data vs reconstructed cytokine trajectories\n",
    "# For each peptide, plot one quantity\n",
    "fig, allaxes = plt.subplots(2, 5, sharex=False, sharey=\"row\")\n",
    "fig.set_size_inches(4.75, 2*1.5)\n",
    "axes = allaxes[0]\n",
    "\n",
    "cytokines_plot = [\"IFNg\", \"IL-2\", \"IL-17A\", \"IL-6\", \"TNFa\"]\n",
    "cytokines_nice = [r\"IFN-$\\gamma$\", \"IL-2\", \"IL-17A\", \"IL-6\", \"TNF\"]\n",
    "\n",
    "times_recon = df_recon.index.get_level_values(\"Time\").unique()\n",
    "times_orig = df_orig.index.get_level_values(\"Time\").unique()\n",
    "\n",
    "peps_to_plot = [\"N4\", \"Q4\", \"T4\", \"V4\", \"E1\"]\n",
    "pep_palette = sns.color_palette(n_colors=len(peps_to_plot)+1)\n",
    "pep_palette.pop(4)\n",
    "pep_palette = {peps_to_plot[i]:pep_palette[i] for i in range(len(pep_palette))}\n",
    "lw_choice = 1.5\n",
    "for j, pep in enumerate(peps_to_plot[:4]):\n",
    "    for i in range(len(cytokines_plot)):\n",
    "        y = df_recon.loc[(dset_choice, pep, conc_choice, times_recon), (\"concentration\", cytokines_plot[i])].values\n",
    "        y2 = df_orig.loc[(pep, conc_choice, times_orig), (cytokines_plot[i],)].values.flatten()\n",
    "        # Offsets for proper absolute concentration plotting\n",
    "        y = y + np.log10(pM_offset.loc[cytokines_plot[i]])\n",
    "        y2 = y2 + np.log10(pM_offset.loc[cytokines_plot[i]])\n",
    "        yerr = df_std.loc[(pep, conc_choice, times_orig), (cytokines_plot[i],)].values.flatten()\n",
    "        # Display p-value in legend. Optional, currently off. \n",
    "        #pvalue = float(ser_pval_mv.loc[(pep, conc_choice)])\n",
    "        #chi_value = float(ser_mv_chisq_norm.loc[(pep, conc_choice)])\n",
    "        axes[i].plot(times_recon, y, color=pep_palette[pep], lw=lw_choice, zorder=2*(4-j)+5,\n",
    "            label=pep)  #+\" (p = {:.2f},\\n\".format(pvalue)+r\"$\\chi^2/\\nu$ = {:.2f})\".format(chi_value))\n",
    "        axes[i].errorbar(times_orig, y2, yerr=yerr, color=pep_palette[pep], lw=lw_choice, zorder=2*(4-j)+4,\n",
    "                         ls=\"none\", marker=\"o\", ms=2, elinewidth=lw_choice*0.5)\n",
    "\n",
    "# Add E1 manually to the plot\n",
    "for i in range(len(cytokines_plot)):\n",
    "    recon_times = df_recon_e1.index.get_level_values(\"Time\").unique()\n",
    "    y = df_recon_e1.loc[(\"E1\", \"1uM\"), cytokines_plot[i]].values\n",
    "    y2 = df_orig_e1.loc[(\"E1\", \"1uM\", e1_times), cytokines_plot[i]].values\n",
    "    y = y + np.log10(pM_offset.loc[cytokines_plot[i]])\n",
    "    y2 = y2 + np.log10(pM_offset.loc[cytokines_plot[i]])\n",
    "    # Variance: we used ser_std anyways\n",
    "    yerr = ser_std[cytokines_plot[i]]\n",
    "    axes[i].plot(recon_times, y, color=pep_palette[\"E1\"], lw=lw_choice, zorder=1,\n",
    "            label=\"E1\")  #+\" (p = {:.2f},\\n\".format(pvalue_e1)+r\"$\\chi^2/\\nu$ = {:.2f})\".format(chisq_e1_norm))\n",
    "    axes[i].errorbar(e1_times, y2, yerr=yerr, color=pep_palette[\"E1\"], lw=lw_choice, zorder=2,\n",
    "                         ls=\"none\", marker=\"o\", ms=2, elinewidth=lw_choice*0.5)\n",
    "\n",
    "# Add logarithmic ticks manually\n",
    "#ylims_log100 = (np.log(10.0**axes[0].get_ylim()[0]) / np.log(100), \n",
    "#                np.log(10.0**axes[0].get_ylim()[1]) / np.log(100))\n",
    "#print(ylims_log100)\n",
    "#minorticks = compute_log_minor_ticks(ylims_log100, stp=1, base=100.0)\n",
    "#Convert back to log10, which is the scale of the axis\n",
    "#minorticks = np.log10(100**minorticks)\n",
    "\n",
    "with open(os.path.join(\"data\", \"misc\", \"minor_ticks_props.json\"), \"r\") as hd:\n",
    "    props_minorticks = json.load(hd)\n",
    "with open(os.path.join(\"data\", \"misc\", \"major_ticks_props.json\"), \"r\") as hd:\n",
    "    props_majorticks = json.load(hd)\n",
    "ylims_log10 = axes[0].get_ylim()\n",
    "minorticks = compute_log_minor_ticks(ylims_log10, stp=1, base=10.0)\n",
    "\n",
    "def powten_format(x, pos):\n",
    "    return r\"$10^{\" + \"{}\".format(int(x)) + r\"}$\"\n",
    "for i, ylbl in enumerate(cytokines_nice):\n",
    "    axes[i].set_title(ylbl, size=7, pad=4., va=\"top\", y=0.91)\n",
    "    axes[i].tick_params(axis=\"both\", **props_majorticks)\n",
    "    axes[i].tick_params(axis=\"y\", which=\"minor\", **props_minorticks)\n",
    "    axes[i].set_yticks(np.arange(np.ceil(ylims_log10[0]), np.floor(ylims_log10[1])+1, 1))\n",
    "    axes[i].set_yticks(minorticks, minor=True)\n",
    "    axes[i].yaxis.set_major_formatter(mpl.ticker.FuncFormatter(powten_format))\n",
    "\n",
    "# Label the y axes\n",
    "axes[0].set_ylabel(\"[cytokine] (pM)\", size=7, labelpad=0.5)\n",
    "for i in range(1, len(cytokines_plot)):\n",
    "    axes[i].set_ylabel(\"\")\n",
    "    \n",
    "# Label the x axes\n",
    "for i in range(len(cytokines_plot)):\n",
    "    axes[i].set_xlabel(\"Time (h)\", size=7, labelpad=0.5)\n",
    "    axes[i].set_xticks([0, 30, 60])\n",
    "    axes[i].set_xticklabels([0, 30, 60])\n",
    "    for axis in ['bottom', 'left', \"top\", \"right\"]:\n",
    "        axes[i].spines[axis].set_linewidth(0.8)\n",
    "    \n",
    "# Legend for peptides on the side\n",
    "handles, labels = axes[0].get_legend_handles_labels()\n",
    "leg = fig.legend(handles, labels, loc=\"upper left\", bbox_to_anchor=(0.96, 0.95), \n",
    "                 handlelength=0.8, fontsize=7, labelspacing=0.3, frameon=False, \n",
    "                 borderaxespad=0.5, handletextpad=0.3)\n",
    "\n",
    "# Legend on top to emphasize reconstruction vs data\n",
    "handles2 = [\n",
    "    mpl.lines.Line2D([0, 1], [0, 1], color=\"k\", ls=\"none\", marker=\"o\", ms=2), \n",
    "    mpl.lines.Line2D([0, 1], [0, 1], color=\"k\", lw=lw_choice)]\n",
    "labels2 = [\"Data\", \"Reconstruction from latent space model\"]\n",
    "second_legend = mpl.legend.Legend(parent=fig, handles=handles2, labels=labels2, ncol=2, \n",
    "                    loc='upper center', bbox_to_anchor=(0.5, 1.02), frameon=False, \n",
    "                    handlelength=0.8, fontsize=8, labelspacing=0.3, \n",
    "                    borderaxespad=0.5, handletextpad=0.3)\n",
    "fig.add_artist(second_legend)\n",
    "\n",
    "\n",
    "## SECOND ROW: residuals\n",
    "axes = allaxes[1]\n",
    "\n",
    "times_recon = df_recon.index.get_level_values(\"Time\").unique()\n",
    "times_orig = df_orig.index.get_level_values(\"Time\").unique()\n",
    "\n",
    "peps_to_plot = [\"N4\", \"Q4\", \"T4\", \"V4\", \"E1\"]\n",
    "pep_palette = sns.color_palette(n_colors=len(peps_to_plot)+1)\n",
    "pep_palette.pop(4)\n",
    "pep_palette = {peps_to_plot[i]:pep_palette[i] for i in range(len(pep_palette))}\n",
    "lw_choice = 2.\n",
    "for j, pep in enumerate(peps_to_plot[:4]):\n",
    "    for i in range(len(cytokines_plot)):\n",
    "        y = df_resids_select.loc[(pep, conc_choice, times_orig), (cytokines_plot[i],)].values.flatten()\n",
    "        yerr = df_std.loc[(pep, conc_choice, times_orig), (cytokines_plot[i],)].values.flatten()\n",
    "        #pvalue = float(ser_pval_mv.loc[(pep, conc_choice)])\n",
    "        #chi_value = float(ser_mv_chisq_norm.loc[(pep, conc_choice)])\n",
    "        zord = (5-j)\n",
    "        axes[i].plot(times_orig, y, color=pep_palette[pep], lw=lw_choice, \n",
    "                         zorder=zord+2*5, label=pep)\n",
    "        axes[i].plot(times_orig, yerr, color=pep_palette[pep], zorder=zord+5,\n",
    "                         ls=\"-.\", lw=lw_choice*0.5, alpha=0.7, label=r\"$\\sigma_{\"+pep+\"}$\")\n",
    "        axes[i].plot(times_orig, -yerr, color=pep_palette[pep], zorder=zord+5,\n",
    "                         ls=\"-.\", lw=lw_choice*0.5, alpha=0.7)\n",
    "        # Fill between?\n",
    "        #axes[i].fill_between(times_orig, -yerr, yerr, color=pep_palette[pep], \n",
    "        #                     zorder=zord, alpha=0.1)\n",
    "\n",
    "# Add E1 manually to the plot\n",
    "for i in range(len(cytokines_plot)):\n",
    "    y = df_resids_e1.loc[(\"E1\", \"1uM\", e1_times), cytokines_plot[i]].values\n",
    "    # Variance: we used ser_std anyways\n",
    "    yerr = np.asarray([ser_std[cytokines_plot[i]]]*len(e1_times))\n",
    "    axes[i].plot(e1_times, y, color=pep_palette[\"E1\"], lw=lw_choice, zorder=1+10,\n",
    "            label=\"E1\")\n",
    "    axes[i].plot(e1_times, yerr, color=pep_palette[\"E1\"], zorder=1+5, ls=\"-.\", \n",
    "                 lw=lw_choice/1.5, alpha=0.65, label=r\"$\\sigma_{E1}$\")\n",
    "    axes[i].plot(e1_times, -yerr, color=pep_palette[\"E1\"], zorder=1+5, ls=\"-.\", \n",
    "                 lw=lw_choice/1.5, alpha=0.65)\n",
    "    #axes[i].fill_between(e1_times, -yerr, yerr, color=pep_palette[\"E1\"], \n",
    "    #              zorder=1, alpha=0.1)\n",
    "\n",
    "\n",
    "for i, ylbl in enumerate(cytokines_nice):\n",
    "    axes[i].tick_params(axis=\"both\", labelsize=7, length=2., width=0.8)\n",
    "\n",
    "# Label the y axes\n",
    "axes[0].set_ylabel(r\"$\\Delta \\log_{10}$([cytokine])\", size=7, labelpad=1.5)\n",
    "for i, ylbl in enumerate(cytokines_nice):\n",
    "    axes[i].set_title(ylbl, size=7, pad=4., va=\"top\", y=0.91)\n",
    "    for axis in ['bottom', 'left', \"top\", \"right\"]:\n",
    "        axes[i].spines[axis].set_linewidth(0.8)\n",
    "\n",
    "# Label the x axes\n",
    "for i in range(len(cytokines_plot)):\n",
    "    axes[i].set_xlabel(\"Time (h)\", size=7, labelpad=0.5)\n",
    "    axes[i].set_xticks([0, 30, 60])\n",
    "    axes[i].set_xticklabels([0, 30, 60])\n",
    "    for axis in ['bottom', 'left', \"top\", \"right\"]:\n",
    "        axes[i].spines[axis].set_linewidth(0.8)\n",
    "    \n",
    "# Legend for peptides below\n",
    "handles3, labels3 = axes[0].get_legend_handles_labels()\n",
    "handles3 = handles3[1::2]\n",
    "labels3 = labels3[1::2]\n",
    "\n",
    "# Third legend for residuals\n",
    "third_legend = mpl.legend.Legend(parent=fig, handles=handles3, labels=labels3, \n",
    "                  bbox_to_anchor=(0.95, 0.45), loc=\"upper left\", ncol=1, \n",
    "                  handlelength=0.8, fontsize=7, labelspacing=0.3, frameon=False, \n",
    "                 borderaxespad=0.5, handletextpad=0.3)\n",
    "fig.add_artist(third_legend)\n",
    "\n",
    "\n",
    "# Fourth legend for residuals and standard deviation\n",
    "handles4 = [\n",
    "    mpl.lines.Line2D([0, 1], [0, 1], color=\"k\", ls=\"-\", lw=lw_choice), \n",
    "    mpl.lines.Line2D([0, 1], [0, 1], color=\"k\", ls=\"-.\", lw=lw_choice/1.5)]\n",
    "labels4 = [\"Residual reconstruction-data\", \"Standard deviation\"]\n",
    "fourth_legend = mpl.legend.Legend(parent=fig, handles=handles4, labels=labels4, ncol=2, \n",
    "                    loc='center', bbox_to_anchor=(0.5, 0.47), frameon=False, \n",
    "                    handlelength=1.85, fontsize=8, labelspacing=0.3, \n",
    "                    borderaxespad=0.5, handletextpad=0.3)\n",
    "fig.add_artist(fourth_legend)\n",
    "\n",
    "fig.tight_layout(w_pad=0.3, h_pad=2.5)\n",
    "\n",
    "fig.savefig(os.path.join(\"figures\", \"reconstruction\", \"supp_figure_reconstruction_model_pvalues_residuals.pdf\"), \n",
    "           bbox_inches=\"tight\", bbox_extra_artists=(leg, second_legend, third_legend, fourth_legend), transparent=True)\n",
    "plt.show()\n",
    "plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_r_a0guae9W"
   },
   "outputs": [],
   "source": [
    "# Output the p-values and chi^2/nu (and chi^2 and nu separately) to a text file\n",
    "df_fit_stats = pd.DataFrame(\n",
    "    {\"p-value\": ser_pval_mv.loc[(peps_to_plot[:4], conc_choice)].tolist()+[pvalue_e1], \n",
    "    \"chi^2/nu\": ser_mv_chisq_norm.loc[(peps_to_plot[:4], conc_choice)].tolist()+[chisq_e1_norm], \n",
    "    \"chi^2\": ser_mv_chisq.loc[(peps_to_plot[:4], conc_choice)].tolist() + [chisq_e1], \n",
    "    \"nu\": [nu_dof5]*len(peps_to_plot)\n",
    "    }, index=pd.MultiIndex.from_product(\n",
    "        (peps_to_plot, conc_choice), names=[\"Peptide\", \"Concentration\"])\n",
    "    )\n",
    "print(df_fit_stats)\n",
    "# df_fit_stats.to_json(os.path.join(\"results\", \"reconstruction\", \"reconstruction_model_pvalues.json\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "reconstruction_pvalues_panels.ipynb",
   "provenance": []
  },
  "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": 1
}