{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import wot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook 7: Validation by Geodesic Interpolation\n",
    "\n",
    "In this notebook we test our ability to interpolate the distribution of cells at held-out time-points. \n",
    "Intuitively, the developmental process $\\mathbb{P}_t$ traces out a *curve in the space of distributions*.\n",
    "We sample cells at various time points and obtain noisy samples along this curve.\n",
    "Optimal transport allows us to connect these noisy samples with straight lines (aka *geodesics*) in the space of distributions. This allows us to interpolate the distribution of cells at intermediate time-points. \n",
    "We can therefore test the performance of OT by holding out time-points and  comparing our interpolation estimate to the real held-out data.\n",
    "\n",
    "We begin by initializing an OT Model for the serum time-course using pre-computed growth rates (stored in the file `serum_g.txt` created in Notebook 2). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "VAR_GENE_DS_PATH = 'data/ExprMatrix.var.genes.h5ad'\n",
    "LEARNED_GROWTH_SCORES_PATH = 'tmaps/serum_g.txt'\n",
    "BATCH_PATH = 'data/batches.txt'\n",
    "CELL_DAYS_PATH = 'data/cell_days.txt'\n",
    "SERUM_CELL_IDS_PATH = 'data/serum_cell_ids.txt'\n",
    "\n",
    "adata = wot.io.read_dataset(VAR_GENE_DS_PATH, obs=[CELL_DAYS_PATH, BATCH_PATH, LEARNED_GROWTH_SCORES_PATH], obs_filter=SERUM_CELL_IDS_PATH)\n",
    "\n",
    "ot_model = wot.ot.OTModel(adata, growth_rate_field='g2',growth_iters = 1) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to test interpolation over a single interval from day $17$ to $18$. \n",
    "We set aside the data from day $17.5$, and we estimate the coupling from $17$ to $18$ by computing a transport map. \n",
    "We then compute an interpolating point cloud of simulated day $17.5$ cells and compare this to the real observations using earth mover's distance. \n",
    "We also compute the distance betwen batches of real observations at time $17.5$. \n",
    "This provides an estimate of the baseline noise level. \n",
    "Finally, we also compute distances between the real point cloud at $17.5$ and several null-estimates different from OT. These include\n",
    "* an interoplation using a random coupling\n",
    "* the first point in the interval (day 17.0 cells)\n",
    "* the last point in the interval (day 18.0 cells)\n",
    "\n",
    "The function `compute_validation_summary` returns a table containing all these point-cloud distances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interp_summary = wot.ot.compute_validation_summary(ot_model, day_triplets=[(17, 17.5, 18)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This produces a table `interp_summary` which contains all the distances between point clouds described above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28688 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "wot.graphics.plot_ot_validation_summary_stats(interp_summary.groupby(['interval_mid', 'name'])['distance'].agg([np.mean, np.std]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now use the `compute_validation_summary` function to check the quality of interpolation at *each* triplet of consecutive time-points $(t_i,t_{i+1},t_{i+2})$. For each triple, we\n",
    "* hold out the data from time $t_{i+1}$,\n",
    "* estimate the coupling $\\gamma_{t_i,t_{i+2}}$ connecting time $t_{i}$ to $t_{i+2}$,\n",
    "* compute an interpolating distribution at time $t_{i+1}$,\n",
    "* and compare this to the held-out data $\\hat{\\mathbb{P}}_{t_{i+1}}$.\n",
    "\n",
    "If the held-out data consists of multiple batches (e.g. $\\hat{\\mathbb{P}}_{t_{i+1}}^{(1)}$ and $\\hat{\\mathbb{P}}_{t_{i+1}}^{(2)}$), then we compare those to each other to estimate the baseline noise level.\n",
    "\n",
    "We also compare several null models to the held-out data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Compute validation for all triplets in dataset. As this can take several hours. We provide precomputed results for convenience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# all_triplets_summary = wot.ot.compute_validation_summary(ot_model)\n",
    "\n",
    "# # save results\n",
    "# all_triplets_summary.to_csv('data/serum_validation_summary.txt')\n",
    "# all_triplets_stats = all_triplets_summary.groupby(['interval_mid', 'name'])['distance'].agg([np.mean, np.std])\n",
    "# all_triplets_stats.to_csv('data/serum_validation_summary_stats.txt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 211888 characters]\n",
      "text/plain": [
       "<Figure size 720x720 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# read in and plot results\n",
    "all_triplets_stats = pd.read_csv('data/serum_validation_summary_stats.txt')\n",
    "wot.graphics.plot_ot_validation_summary_stats(all_triplets_stats)"
   ]
  }
 ],
 "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.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "source": [],
    "metadata": {
     "collapsed": false
    }
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}