tutorial/sim_analysis_batched_experimental_conditions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyze cell state dynamics batched experimental conditions\n",
    "In this tutorial, we will show how to analyze the effect of batched experimental conditions using exdyn. We will apply exdyn to simulated data set where cluster 0 population differentiate into cluster 1 in condition 1 and cluster 2 in condition 2.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model set up and optimization\n",
    "Exdyn accepts AnnData object for dynamics estimation. The model parameters such as the number of neural network layers can be specified through a dictionary of the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install exdyn torchtext==0.6.0\n",
    "! wget   -O sim_adata.h5ad https://www.dropbox.com/scl/fi/lj9ca6bn058g0xopi7zjd/sim_adata_mod.h5ad?rlkey=d44nnplmrqd5059698qrm73b6&dl=0\n",
    "import scanpy as sc\n",
    "from exdyn import workflow, visualization, condiff, modules, commons\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "seed = 42\n",
    "\n",
    "random.seed(seed)\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "torch.backends.cudnn.benchmark = False\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "orig_adata = sc.read_h5ad('data/sim_adata_mod.h5ad')\n",
    "model_params = {\n",
    "            'z_dim': 10,\n",
    "            'enc_z_h_dim': 50, 'enc_d_h_dim': 50, 'dec_z_h_dim': 50,\n",
    "            'num_enc_z_layers': 2, 'num_enc_d_layers': 2,\n",
    "            'num_dec_z_layers': 2, 'use_ambient': False, 'use_vamp': False, 'no_d_kld': False, 'decreasing_temp': False, 'dec_temp_steps': 30000, 'loss_mode': 'nb'\n",
    "}\n",
    "checkpoint_dirname = 'checkpoint'\n",
    "adata, lit_envdyn = workflow.conduct_cvicdyf_inference(orig_adata, model_params, checkpoint_dirname, batch_size=128, two_step=False, dyn_mode=False, epoch=100, patience=50, module=modules.Cvicdyf, use_highly_variable=False, batch_key='sample', condition_key='condition')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visulaize condition wise dynamics\n",
    "We can visualize cell state dynamics of all the cells in each condition as well as the dynamics in the conditions identical with those where cells were obsserved.\n",
    "\n",
    "### Cell state dynamics in original conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Firstly, we visualize the cell state dynamics in the conditions where each cell originally observed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "visualization.plot_mean_flow(adata, cluster_key='condition', legend_loc='right')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Exdyn can provide an opptunity to derive the conterfactual estimation of cell state dynamics in a condition different from that where the cells were observed. \n",
    "Here, we estimate and display the cell state dynamics of all the cells in condition 1 and 2 respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cond1 = '1'\n",
    "cond2 = '2'\n",
    "conds = [cond1, cond2]\n",
    "import importlib;importlib.reload(condiff)\n",
    "adata = condiff.estimate_two_cond_dynamics(adata, cond1, cond2, lit_envdyn)\n",
    "for cond in conds:\n",
    "    visualization.plot_mean_flow(adata, cluster_key=None, legend_loc='right', vel_key=f'norm_cond_vel_{cond}', du_key=f'cond_dumap_{cond}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the difference between conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6 * 1, 5 * 1))\n",
    "visualization.diff_visualize(adata.obsm['cond_dumap_diff'], adata.obsm['X_umap'], ax)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general, exdyn can calculate the dynamics across multple conditions more than 2. However, the analysis lacks several two condition specific results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multi_adata = condiff.estimate_multi_cond_dynamics(adata, [cond1, cond2], lit_envdyn)\n",
    "print('Two condition version version:')\n",
    "print(adata)\n",
    "print('General version:')\n",
    "print(multi_adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Identification analysis of bifurction points between conditions \n",
    "Exdyn enables us to identify a bifurcation point between the conditions, where the identical cell states have dynamics varying between the econditions. Here, we extract cells with large conditional difference from the cell population shared by both of the conditions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata.obs['Condition 2 ratio'] = commons.calculate_neighbor_ratio(adata.obsm['X_vicdyf_zl'], adata.obs.condition.values.astype(str) == '2', nn=30)\n",
    "adata.obs['Condition difference'] = np.linalg.norm(adata.layers['norm_cond_vel_diff'], axis=1)\n",
    "\n",
    "common_adata = adata[np.logical_and(adata.obs['Condition 2 ratio'] > 0.3, adata.obs['Condition 2 ratio'] < 0.7)]\n",
    "top_adata = common_adata[common_adata.obs['Condition difference'] > common_adata.obs['Condition difference'].quantile(0.7)]\n",
    "\n",
    "adata.obs['Top Condition Difference'] = adata.obs_names.isin(top_adata.obs_names).astype(str)\n",
    "sc.pl.umap(adata, color=['Top Condition Difference'], palette={'True': 'black', 'False': 'gray'})\n",
    "\n",
    "adata.obs['Top Condition Difference'] = adata.obs_names.isin(top_adata.obs_names).astype(str)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can quantify the gene-wise dynamics difference at the bifurcation point, which leads to the identification of the regulator explaining the population difference between the conditions. Here we calculated the conditional difference and visualize them for top difference genes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "cdiff_tot_vels = pd.Series(top_adata.layers[f'norm_cond_vel_diff'].mean(axis=0), index=top_adata.var_names)\n",
    "top_cdiff_vals = pd.concat([cdiff_tot_vels.sort_values(ascending=False)[:5], cdiff_tot_vels.sort_values(ascending=False)[-5:]])\n",
    "fig, ax = plt.subplots(1, 1, figsize=(6 * 1, 5 * 1))\n",
    "visualization.annotated_bars(ax, top_cdiff_vals.index, top_cdiff_vals.values)\n",
    "ax.set_ylabel('Conditional dynamics difference')\n",
    "ax.set_xlabel('Gene')\n",
    "plt.subplots_adjust(bottom=0.15, left=0.2)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can conduct clustering on the bifurcation point and analyze subpulation if the identified bifurcation points are heterogeneous."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(top_adata, use_rep='X_vicdyf_zl')\n",
    "sc.tl.leiden(top_adata, key_added='top_leiden')\n",
    "adata.obs['top_leiden'] = None\n",
    "adata.obs.loc[top_adata.obs_names, 'top_leiden'] = top_adata.obs['top_leiden']\n",
    "sc.pl.umap(top_adata, color='top_leiden', legend_loc='on data')\n",
    "sub_top_adata = adata[adata.obs.top_leiden.isin(['0', '1'])]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_adata"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "functorch2",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}