Jupyter Notebooks/Train_Haber.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample Run with trVAE for Haber et. al dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "os.chdir(\"../\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "import scanpy as sc\n",
    "import reptrvae\n",
    "import pandas as pd \n",
    "from scipy import stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.settings.set_figure_params(dpi=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_name = \"haber\"\n",
    "specific_celltype = \"Tuft\"\n",
    "conditions = ['Control', 'Hpoly.Day3', 'Hpoly.Day10', 'Salmonella']\n",
    "target_conditions = [\"Hpoly.Day3\", \"Hpoly.Day10\", \"Salmonella\"]\n",
    "cell_type_key = \"cell_label\"\n",
    "condition_key = \"condition\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "View of AnnData object with n_obs × n_vars = 9842 × 1000 \n",
       "    obs: 'batch', 'barcode', 'condition', 'cell_label', 'n_counts'\n",
       "    var: 'highly_variable', 'means', 'dispersions', 'dispersions_norm'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata = sc.read(f\"./data/{data_name}/{data_name}_normalized.h5ad\")\n",
    "adata = adata[adata.obs[condition_key].isin(conditions)]\n",
    "adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cell_label             condition  \n",
       "Endocrine              Control        112\n",
       "                       Hpoly.Day3     117\n",
       "                       Hpoly.Day10     82\n",
       "                       Salmonella      69\n",
       "Enterocyte             Control        424\n",
       "                       Hpoly.Day3     201\n",
       "                       Hpoly.Day10    128\n",
       "                       Salmonella     705\n",
       "Enterocyte.Progenitor  Control        545\n",
       "                       Hpoly.Day3     462\n",
       "                       Hpoly.Day10    586\n",
       "                       Salmonella     229\n",
       "Goblet                 Control        216\n",
       "                       Hpoly.Day3      99\n",
       "                       Hpoly.Day10    317\n",
       "                       Salmonella     126\n",
       "Stem                   Control        670\n",
       "                       Hpoly.Day3     388\n",
       "                       Hpoly.Day10    592\n",
       "                       Salmonella     207\n",
       "TA                     Control        421\n",
       "                       Hpoly.Day3     302\n",
       "                       Hpoly.Day10    353\n",
       "                       Salmonella     112\n",
       "TA.Early               Control        792\n",
       "                       Hpoly.Day3     442\n",
       "                       Hpoly.Day10    436\n",
       "                       Salmonella     300\n",
       "Tuft                   Control         60\n",
       "                       Hpoly.Day3     110\n",
       "                       Hpoly.Day10    217\n",
       "                       Salmonella      22\n",
       "dtype: int64"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adata.obs.groupby([cell_type_key, condition_key]).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_adata, valid_adata = reptrvae.utils.train_test_split(adata, 0.80)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((7873, 1000), (1969, 1000))"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_adata.shape, valid_adata.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "net_train_adata = train_adata[~((train_adata.obs[cell_type_key] == specific_celltype) & (train_adata.obs[condition_key].isin(target_conditions)))]\n",
    "net_valid_adata = valid_adata[~((valid_adata.obs[cell_type_key] == specific_celltype) & (valid_adata.obs[condition_key].isin(target_conditions)))]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((7598, 1000), (1895, 1000))"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net_train_adata.shape, net_valid_adata.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "cell_label             condition  \n",
       "Endocrine              Control         89\n",
       "                       Hpoly.Day3     103\n",
       "                       Hpoly.Day10     69\n",
       "                       Salmonella      58\n",
       "Enterocyte             Control        318\n",
       "                       Hpoly.Day3     166\n",
       "                       Hpoly.Day10     99\n",
       "                       Salmonella     571\n",
       "Enterocyte.Progenitor  Control        442\n",
       "                       Hpoly.Day3     378\n",
       "                       Hpoly.Day10    451\n",
       "                       Salmonella     189\n",
       "Goblet                 Control        170\n",
       "                       Hpoly.Day3      81\n",
       "                       Hpoly.Day10    255\n",
       "                       Salmonella     108\n",
       "Stem                   Control        530\n",
       "                       Hpoly.Day3     314\n",
       "                       Hpoly.Day10    494\n",
       "                       Salmonella     168\n",
       "TA                     Control        322\n",
       "                       Hpoly.Day3     246\n",
       "                       Hpoly.Day10    284\n",
       "                       Salmonella      88\n",
       "TA.Early               Control        626\n",
       "                       Hpoly.Day3     337\n",
       "                       Hpoly.Day10    355\n",
       "                       Salmonella     239\n",
       "Tuft                   Control         48\n",
       "dtype: int64"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net_train_adata.obs.groupby([cell_type_key, condition_key]).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /home/mohsen/miniconda3/envs/venv/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:263: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Colocations handled automatically by placer.\n",
      "WARNING:tensorflow:From /home/mohsen/miniconda3/envs/venv/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py:3445: calling dropout (from tensorflow.python.ops.nn_ops) with keep_prob is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Please use `rate` instead of `keep_prob`. Rate should be set to `rate = 1 - keep_prob`.\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "data (InputLayer)               (None, 1000)         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "encoder_labels (InputLayer)     (None, 4)            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_1 (Concatenate)     (None, 1004)         0           data[0][0]                       \n",
      "                                                                 encoder_labels[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "dense_1 (Dense)                 (None, 800)          803200      concatenate_1[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_1 (BatchNor (None, 800)          3200        dense_1[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_1 (LeakyReLU)       (None, 800)          0           batch_normalization_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_1 (Dropout)             (None, 800)          0           leaky_re_lu_1[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "dense_2 (Dense)                 (None, 800)          640000      dropout_1[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_2 (BatchNor (None, 800)          3200        dense_2[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_2 (LeakyReLU)       (None, 800)          0           batch_normalization_2[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_2 (Dropout)             (None, 800)          0           leaky_re_lu_2[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "dense_3 (Dense)                 (None, 128)          102400      dropout_2[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_3 (BatchNor (None, 128)          512         dense_3[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_3 (LeakyReLU)       (None, 128)          0           batch_normalization_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_3 (Dropout)             (None, 128)          0           leaky_re_lu_3[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "dense_4 (Dense)                 (None, 60)           7740        dropout_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "dense_5 (Dense)                 (None, 60)           7740        dropout_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "lambda_1 (Lambda)               (None, 60)           0           dense_4[0][0]                    \n",
      "                                                                 dense_5[0][0]                    \n",
      "==================================================================================================\n",
      "Total params: 1,567,992\n",
      "Trainable params: 1,564,536\n",
      "Non-trainable params: 3,456\n",
      "__________________________________________________________________________________________________\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "latent_data (InputLayer)        (None, 60)           0                                            \n",
      "__________________________________________________________________________________________________\n",
      "decoder_labels (InputLayer)     (None, 4)            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "concatenate_2 (Concatenate)     (None, 64)           0           latent_data[0][0]                \n",
      "                                                                 decoder_labels[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "dense_6 (Dense)                 (None, 128)          8192        concatenate_2[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_4 (BatchNor (None, 128)          512         dense_6[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "mmd (LeakyReLU)                 (None, 128)          0           batch_normalization_4[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_4 (Dropout)             (None, 128)          0           mmd[0][0]                        \n",
      "__________________________________________________________________________________________________\n",
      "dense_7 (Dense)                 (None, 800)          102400      dropout_4[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_5 (BatchNor (None, 800)          3200        dense_7[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_4 (LeakyReLU)       (None, 800)          0           batch_normalization_5[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_5 (Dropout)             (None, 800)          0           leaky_re_lu_4[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "dense_8 (Dense)                 (None, 800)          640000      dropout_5[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_6 (BatchNor (None, 800)          3200        dense_8[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "leaky_re_lu_5 (LeakyReLU)       (None, 800)          0           batch_normalization_6[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "dropout_6 (Dropout)             (None, 800)          0           leaky_re_lu_5[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "dense_9 (Dense)                 (None, 1000)         801000      dropout_6[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "reconstruction_output (Activati (None, 1000)         0           dense_9[0][0]                    \n",
      "==================================================================================================\n",
      "Total params: 1,558,504\n",
      "Trainable params: 1,555,048\n",
      "Non-trainable params: 3,456\n",
      "__________________________________________________________________________________________________\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "data (InputLayer)               (None, 1000)         0                                            \n",
      "__________________________________________________________________________________________________\n",
      "encoder_labels (InputLayer)     (None, 4)            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "encoder (Model)                 [(None, 60), (None,  1567992     data[0][0]                       \n",
      "                                                                 encoder_labels[0][0]             \n",
      "                                                                 data[0][0]                       \n",
      "                                                                 encoder_labels[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "decoder_labels (InputLayer)     (None, 4)            0                                            \n",
      "__________________________________________________________________________________________________\n",
      "decoder (Model)                 (None, 1000)         1558504     encoder[1][2]                    \n",
      "                                                                 decoder_labels[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "decoder_mmd (Model)             (None, 128)          8704        encoder[2][2]                    \n",
      "                                                                 decoder_labels[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "kl_mse (Lambda)                 (None, 1000)         0           decoder[1][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "mmd (Lambda)                    (None, 128)          0           decoder_mmd[1][0]                \n",
      "==================================================================================================\n",
      "Total params: 3,126,496\n",
      "Trainable params: 3,119,584\n",
      "Non-trainable params: 6,912\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "network = reptrvae.models.trVAE(x_dimension=net_train_adata.shape[1],\n",
    "                                z_dimension=60,\n",
    "                                mmd_dimension=128, \n",
    "                                n_conditions=len(net_train_adata.obs[condition_key].unique()),\n",
    "                                alpha=1e-6,\n",
    "                                beta=100,\n",
    "                                eta=100,\n",
    "                                clip_value=100,\n",
    "                                lambda_l1=0.0,\n",
    "                                lambda_l2=0.0,\n",
    "                                learning_rate=0.001,\n",
    "                                model_path=f\"./models/trVAEMulti/best/{data_name}-{specific_celltype}/\",\n",
    "                                dropout_rate=0.2,\n",
    "                                output_activation='relu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_encoder = {'Control': 0, 'Hpoly.Day3': 1, 'Hpoly.Day10': 2, 'Salmonella': 3}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# network.restore_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network.train(net_train_adata,\n",
    "              net_valid_adata,\n",
    "              label_encoder,\n",
    "              condition_key,\n",
    "              n_epochs=10000,\n",
    "              batch_size=512,\n",
    "              verbose=2,\n",
    "              early_stop_limit=750,\n",
    "              lr_reducer=0,\n",
    "              shuffle=True,\n",
    "              )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_labels, _ = trvae.tl.label_encoder(net_train_adata, label_encoder, condition_key)\n",
    "latent_with_true_labels = network.to_latent(net_train_adata, train_labels)\n",
    "mmd_with_true_labels = network.to_mmd_layer(net_train_adata, train_labels, feed_fake=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(latent_with_true_labels)\n",
    "sc.tl.umap(latent_with_true_labels)\n",
    "sc.pl.umap(latent_with_true_labels, color=[condition_key, cell_type_key],\n",
    "           show=True,\n",
    "           wspace=0.15,\n",
    "           frameon=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(mmd_with_true_labels)\n",
    "sc.tl.umap(mmd_with_true_labels)\n",
    "sc.pl.umap(mmd_with_true_labels, color=[condition_key, cell_type_key],\n",
    "           show=True,\n",
    "           wspace=0.15,\n",
    "           frameon=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_type_adata = train_adata[train_adata.obs[cell_type_key] == specific_celltype]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_type_adata.var = pd.DataFrame(index=cell_type_adata.var_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cell_type_adata.obs.groupby([cell_type_key, condition_key]).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_transition(adata, source_cond, target_cond):\n",
    "    \n",
    "    source_adata = adata[adata.obs[condition_key] == source_cond]\n",
    "    \n",
    "    source_cond_key = source_cond.split(\"_to_\")[-1]\n",
    "    source_labels = np.zeros(source_adata.shape[0]) + label_encoder[source_cond_key]\n",
    "    target_labels = np.zeros(source_adata.shape[0]) + label_encoder[target_cond]\n",
    "\n",
    "    pred_target = network.predict(source_adata,\n",
    "                                  encoder_labels=source_labels,\n",
    "                                  decoder_labels=target_labels,\n",
    "                                  )\n",
    "\n",
    "    pred_adata = sc.AnnData(X=pred_target)\n",
    "    pred_adata.obs[condition_key] = [source_cond + \"_to_\" + target_cond] * pred_target.shape[0]\n",
    "    pred_adata.obs[cell_type_key] = [specific_celltype] * pred_target.shape[0]\n",
    "    pred_adata.var_names = source_adata.var_names\n",
    "\n",
    "    adata = adata.concatenate(pred_adata)\n",
    "    return adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_adata = predict_transition(cell_type_adata, \"Control\", \"Hpoly.Day10\")\n",
    "recon_adata = recon_adata.concatenate(predict_transition(recon_adata, \"Control\", \"Hpoly.Day3\"))\n",
    "recon_adata = recon_adata.concatenate(predict_transition(recon_adata, \"Control\", \"Salmonella\"))\n",
    "recon_adata = recon_adata.concatenate(predict_transition(recon_adata, \"Hpoly.Day3\", \"Hpoly.Day10\"))\n",
    "recon_adata = recon_adata.concatenate(predict_transition(recon_adata, \"Control_to_Hpoly.Day3\", \"Hpoly.Day10\"))\n",
    "recon_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_adata.obs.groupby([cell_type_key, condition_key]).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path = \"Control_to_Hpoly.Day10\"\n",
    "from_condition = path.split(\"_to_\")[-2]\n",
    "to_condition = path.split(\"_to_\")[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_adata = recon_adata[recon_adata.obs[condition_key] == path]\n",
    "pred_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ctrl_adata = cell_type_adata[cell_type_adata.obs[condition_key] == from_condition]\n",
    "ctrl_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_adata = cell_type_adata[cell_type_adata.obs[condition_key] == to_condition]\n",
    "real_adata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_adata.X.min(), real_adata.X.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_adata.X.min(), pred_adata.X.max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_mean = np.mean(pred_adata.X, axis=0)\n",
    "ctrl_mean = np.mean(ctrl_adata.X, axis=0)\n",
    "real_mean = np.mean(real_adata.X, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_var = np.var(pred_adata.X, axis=0)\n",
    "ctrl_var = np.var(ctrl_adata.X, axis=0)\n",
    "real_var = np.var(real_adata.X, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m, b, r_value_mean, p_value, std_err = stats.linregress(pred_mean, real_mean)\n",
    "r_value_mean = r_value_mean ** 2\n",
    "r_value_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m, b, r_value_var, p_value, std_err = stats.linregress(pred_var, real_var)\n",
    "r_value_var = r_value_var ** 2\n",
    "r_value_var"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_pred = ctrl_adata.concatenate(pred_adata, real_adata)\n",
    "adata_pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_pred.obs[condition_key].unique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.tl.rank_genes_groups(cell_type_adata,\n",
    "                        groupby=condition_key,\n",
    "                        groups=[to_condition],\n",
    "                        reference=from_condition,\n",
    "                        n_genes=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_genes = cell_type_adata.uns['rank_genes_groups']['names'][to_condition]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trvae.pl.reg_mean_plot(adata_pred,\n",
    "                         top_100_genes=top_genes,\n",
    "                         gene_list=top_genes[:10],\n",
    "                         condition_key=condition_key,\n",
    "                         axis_keys={'x': path, 'y': to_condition},\n",
    "                         labels={'x': path, 'y': to_condition},\n",
    "                         path_to_save=None,\n",
    "                         legend=False,\n",
    "                         show=True,\n",
    "                         x_coeff=1.0,\n",
    "                         y_coeff=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trvae.pl.reg_var_plot(adata_pred,\n",
    "                     top_100_genes=top_genes,\n",
    "                     gene_list=top_genes[:10],\n",
    "                     condition_key=condition_key,\n",
    "                     axis_keys={'x': path, 'y': to_condition},\n",
    "                     labels={'x': path, 'y': to_condition},\n",
    "                     path_to_save=None,\n",
    "                     legend=False,\n",
    "                     show=True,\n",
    "                     x_coeff=1.0,\n",
    "                     y_coeff=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pl.violin(adata_pred, groupby=condition_key, keys=top_genes[:10], rotation=90)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "recon_adata.write_h5ad(f\"../trVAE_reproducibility/data/reconstructed/trVAE_Haber/{specific_celltype}.h5ad\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network.model_to_use = f\"./models/trVAEMulti/best/haber-{specific_celltype}/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network.model_to_use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "network.save_model()"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}