autoencoder_citeseq_saturn.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Integrative analysis of single-cell multiomics data using deep learning\n",
    "\n",
    "**Jupyter notebook:** \n",
    "[![View on Github](https://img.shields.io/static/v1.svg?logo=github&label=&message=View%20On%20Github&color=lightgrey)](https://github.com/naity/citeseq_autoencoder/blob/master/autoencoder_citeseq_saturn.ipynb)   \n",
    "**Recording:** \n",
    "[![YouTubb](https://img.shields.io/static/v1?logo=youtube&label=&message=Youtube&color=red)](https://youtu.be/tad9TPCMWbU)   \n",
    "**Author:** Yuan Tian [![Connect](https://img.shields.io/static/v1?label=&logo=linkedin&message=Connect&color=blue)](https://www.linkedin.com/in/ytiancompbio) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<div style=\"font-size:larger;\">\n",
    "    <p><span style=\"font-size:xx-large;\">S</span>ingle-cell RNA sequencing (scRNA-seq) has offered a comprehensive and unbiased approach to profile various type of cells such as immune cells with a single-cell resolution using next‑generation sequencing. More recently, exciting technologies such as cellular indexing of transcriptomes and epitopes by sequencing (CITE-seq) have been developed to extend scRNA-seq by jointly measuring multiple molecular modalities such as proteome and transcriptome from the same cell as illustrated in the figure below. By utilizing antibodies that are conjugated to oligonucleotides, CITE-seq simultaneously generates sequencing-based readouts for surface protein expression along with gene expression.</p>\n",
    "    <p>Since gene and protein expressions convey distinct and complementary information about a cell, CITE-seq offers a unique opportunity to combine both transcriptomic and proteomic data to decipher the biology of individual cells at a considerably higher resolution than using either one alone. This requires computational methods that can effectively integrate single-cell data from both modalities. In this tutorial, we will conduct integrative analysis of CITE-seq data using an unsupervised deep learning method named autoencoder.</p>\n",
    "    <p>In essence:</p>\n",
    "    <ul>\n",
    "        <li>Single-cell technologies offer considerable promise in dissecting the heterogeneity among individual cells and are being utilized in biomedical studies at an astounding pace.</li>\n",
    "        <li>CITE-seq simultaneously measures gene expression and surface protein at a single-cell level.</li>\n",
    "    </ul>\n",
    "</div>\n",
    "\n",
    "<figure>\n",
    "    <center><img src=\"imgs/citeseq.jpg\"/></center>\n",
    "    <center><figcaption>Image source: 10x Genomics</figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Standard libraries\n",
    "import time\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import urllib.request\n",
    "from pathlib import Path\n",
    "from urllib.error import HTTPError\n",
    "from tqdm.notebook import tqdm \n",
    "from sklearn import preprocessing\n",
    "\n",
    "# Pytorch and Pytorch Lightning\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import pytorch_lightning as pl\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint\n",
    "\n",
    "# Visualization and plotting\n",
    "import umap\n",
    "import plotly.express as px\n",
    "\n",
    "# Tensorboard extension\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "%load_ext tensorboard\n",
    "\n",
    "# Path to datasets\n",
    "DATASET_PATH = Path(\"data\")\n",
    "if not DATASET_PATH.exists():\n",
    "    DATASET_PATH.mkdir()\n",
    "    \n",
    "# Path to saved models\n",
    "CHECKPOINT_PATH = Path(\"saved_models\")\n",
    "if not CHECKPOINT_PATH.exists():\n",
    "    CHECKPOINT_PATH.mkdir()\n",
    "\n",
    "# for reproducibility \n",
    "pl.seed_everything(42)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "torch.backends.cudnn.benchmark = False\n",
    "\n",
    "# Use GPU if available, otherwise use cpu instead \n",
    "device = torch.device(\"cuda:0\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
    "print(\"Device:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the CITE-seq dataset published by [Stuart and Butler et al.](https://www.cell.com/cell/fulltext/S0092-8674(19)30559-8) in 2019. The authors measured the single-cell transcriptomics of 30,672 bone marrow cells together with the expression of 25 proteins. I have already preprocessed the data to generate normalized counts and cell type anonations using [Seurat](https://satijalab.org/seurat/index.html), which is a popular R package for analyzing single-cell genomics data. The script used for preprocessing can be found [here](https://github.com/naity/citeseq_autoencoder/blob/master/preprocessing.R). There are three CSV files (RNA, protein, and cell type annotation), which can be downloaded from my [github repo](https://github.com/naity/citeseq_autoencoder) using the code below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# URL for downloading data\n",
    "data_url = \"https://raw.githubusercontent.com/naity/citeseq_autoencoder/master/data/\"\n",
    "\n",
    "# Files to download\n",
    "data_files = [\"rna_scale.csv.gz\", \"protein_scale.csv.gz\", \"metadata.csv.gz\"]\n",
    "\n",
    "# Download datafile if necessary\n",
    "for file_name in data_files:\n",
    "    file_path = Path(DATASET_PATH/file_name)\n",
    "    if not file_path.exists():\n",
    "        file_url = data_url + file_name\n",
    "        print(f\"Downloading {file_url}...\")\n",
    "        try:\n",
    "            urllib.request.urlretrieve(file_url, file_path)\n",
    "        except HTTPError as e:\n",
    "            print(\"Something went wrong. Please try downloading the file from the Google Drive folder\\n\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use Pandas to read the data\n",
    "rna = pd.read_csv(DATASET_PATH/\"rna_scale.csv.gz\", index_col=0).T\n",
    "pro = pd.read_csv(DATASET_PATH/\"protein_scale.csv.gz\", index_col=0).T\n",
    "\n",
    "ncells = rna.shape[0]\n",
    "nfeatures_rna = rna.shape[1]\n",
    "nfeatures_pro = pro.shape[1]\n",
    "\n",
    "print(\"Number of cells:\", ncells)\n",
    "print(\"Number of genes:\", nfeatures_rna)\n",
    "print(\"Number of proteins:\", nfeatures_pro)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, gene and protein expression data are concatenated together, where each column is a gene or protein while each row is a cell (each cell has a unique barcode). The dataset contains the expression levels of 2000 genes and 25 proteins for a total of 30672 cells. We will also import the annotations of each cell for visualizaiton purpose later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# concat rna and pro\n",
    "assert all(rna.index == pro.index), \"RNA and protein data cell barcodes do not match!\"\n",
    "citeseq = pd.concat([rna, pro], axis=1)\n",
    "print(citeseq.shape)\n",
    "citeseq.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cell type annotations\n",
    "metadata = pd.read_csv(DATASET_PATH/\"metadata.csv.gz\", index_col=0)\n",
    "metadata.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert all(citeseq.index == pro.index), \"CITE-seq data and metadata cell barcodes do not match!\"\n",
    "\n",
    "# separate CD4 and CD8 in l1\n",
    "metadata[\"celltype.l1.5\"] = metadata[\"celltype.l1\"].values\n",
    "metadata.loc[metadata[\"celltype.l2\"].str.startswith(\"CD4\"), \"celltype.l1.5\"] = \"CD4 T\"\n",
    "metadata.loc[metadata[\"celltype.l2\"].str.startswith(\"CD8\"), \"celltype.l1.5\"] = \"CD8 T\"\n",
    "metadata.loc[metadata[\"celltype.l2\"]==\"Treg\", \"celltype.l1.5\"] = \"CD4 T\"\n",
    "metadata.loc[metadata[\"celltype.l2\"]==\"MAIT\", \"celltype.l1.5\"] = \"MAIT\"\n",
    "metadata.loc[metadata[\"celltype.l2\"]==\"gdT\", \"celltype.l1.5\"] = \"gdT\"\n",
    "\n",
    "# convert cell type annoations to integers\n",
    "le = preprocessing.LabelEncoder()\n",
    "labels = le.fit_transform(metadata[\"celltype.l1.5\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pytorch datasets and dataloaders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TabularDataset(Dataset):\n",
    "    \"\"\"Custome dataset for tabular data\"\"\"\n",
    "    def __init__(self, df: pd.DataFrame, labels: np.ndarray):\n",
    "        self.data = torch.tensor(df.to_numpy(), dtype=torch.float)\n",
    "        self.labels = torch.tensor(labels, dtype=torch.float)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        x = self.data[idx]\n",
    "        y = self.labels[idx]\n",
    "        return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset = TabularDataset(citeseq, labels)\n",
    "\n",
    "# train, validation, and test split\n",
    "train_size = int(ncells*0.7)\n",
    "val_size = int(ncells*0.15)\n",
    "train_ds, val_ds, test_ds = random_split(dataset, [train_size, val_size, ncells-train_size-val_size],\n",
    "                                         generator=torch.Generator().manual_seed(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Number of cells for training:\", len(train_ds))\n",
    "print(\"Number of cells for validation:\", len(val_ds))\n",
    "print(\"Number of cells for test:\", len(test_ds))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch size\n",
    "bs = 256\n",
    "\n",
    "train_dl = DataLoader(train_ds, batch_size=bs, shuffle=True, drop_last=True, pin_memory=True)\n",
    "val_dl = DataLoader(val_ds, batch_size=bs, shuffle=False, drop_last=False)\n",
    "test_dl = DataLoader(test_ds, batch_size=bs, shuffle=False, drop_last=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s look at one example of the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = train_dl.dataset[0]\n",
    "print(\"Input data:\", x)\n",
    "print(\"Label:     \", y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Use autoencoders for single-cell analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Autoencoder is a type of unsupervised deep learning model or neural network that consists of three major components: an encoder, a bottleneck, and a decoder as shown in the figure below. The encoder compresses the input, and the bottleneck layer stores the compressed representation of the input. In contrast, the decoder tries to reconstruct the input based upon the compressed data.\n",
    "\n",
    "The dimension of the bottleneck layer is normally substantially lower than that of the input. As a result, the encoder will try to learn as much meaningful information about the input as possible while ignoring the noise so that the decoder can do a better job reconstructing the input. Autoencoder can function as a dimensionality reduction algorithm and the low-dimensional representation of the input stored in the bottleneck layer can be used for data visualization and other purposes. Moreover, thanks to its flexible neural network architecture, it offers unlimited ways to incorporate gene and protein expression data as we shall see below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<figure>\n",
    "    <center><img src=\"imgs/autoencoder.png\"/></center>\n",
    "    <center><figcaption>Image source: Eraslan et al. Nat Rev Genet. 2019</figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since gene and protein data have dramatically different dimensions, we will first encode them separately using two different encoders and then concatenate the outputs, which will be passed through another encoder to generate the bottleneck layer. Subsequently, the decoder will try to reconstruct the input based on the bottleneck layer. The overall neural network architecture is illustrated below:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<figure>\n",
    "    <center><img src=\"imgs/autoencoder_arch.png\"/></center>\n",
    "    <center><figcaption><b>Autoencoder architecture for CITE-seq data</b></figcaption></center>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the module below to group linear, batchnorm, and dropout layers together in order to make it easier to implement encoder and decoder later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinBnDrop(nn.Sequential):\n",
    "    \"\"\"Module grouping `BatchNorm1d`, `Dropout` and `Linear` layers, adapted from fastai.\"\"\"\n",
    "    \n",
    "    def __init__(self, n_in, n_out, bn=True, p=0., act=None, lin_first=True):\n",
    "        layers = [nn.BatchNorm1d(n_out if lin_first else n_in)] if bn else []\n",
    "        if p != 0: layers.append(nn.Dropout(p))\n",
    "        lin = [nn.Linear(n_in, n_out, bias=not bn)]\n",
    "        if act is not None: lin.append(act)\n",
    "        layers = lin+layers if lin_first else layers+lin\n",
    "        super().__init__(*layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We start by implementing the encoder, which consists of three fully connected layer groups, one for RNA, one for protein, and one for the concatenated output that generates the latent representation of size `latent_dim` stored in the bottlenect layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    \"\"\"Encoder for CITE-seq data\"\"\"\n",
    "    def __init__(self,\n",
    "                 nfeatures_rna: int,\n",
    "                 nfeatures_pro: int,\n",
    "                 hidden_rna: int,\n",
    "                 hidden_pro: int,\n",
    "                 latent_dim: int,\n",
    "                 p: float = 0):\n",
    "        super().__init__()\n",
    "        self.nfeatures_rna = nfeatures_rna\n",
    "        self.nfeatures_pro = nfeatures_pro\n",
    "        hidden_dim = hidden_rna + hidden_pro\n",
    "        \n",
    "        self.encoder_rna = nn.Sequential(\n",
    "            LinBnDrop(nfeatures_rna, nfeatures_rna // 2, p=p, act=nn.LeakyReLU()),\n",
    "            LinBnDrop(nfeatures_rna // 2, hidden_rna, act=nn.LeakyReLU())\n",
    "        )\n",
    "        self.encoder_protein = LinBnDrop(nfeatures_pro, hidden_pro, p=p, act=nn.LeakyReLU())\n",
    "        self.encoder = LinBnDrop(hidden_dim, latent_dim, act=nn.LeakyReLU())\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_rna = self.encoder_rna(x[:, :self.nfeatures_rna])\n",
    "        x_pro = self.encoder_protein(x[:, self.nfeatures_rna:])\n",
    "        x = torch.cat([x_rna, x_pro], 1)\n",
    "        return self.encoder(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The decoder is a flipped version of the encoder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    \"\"\"Decoder for CITE-seq data\"\"\"\n",
    "    def __init__(self,\n",
    "                 nfeatures_rna: int,\n",
    "                 nfeatures_pro: int,\n",
    "                 hidden_rna: int,\n",
    "                 hidden_pro: int,\n",
    "                 latent_dim: int):\n",
    "        super().__init__()\n",
    "        hidden_dim = hidden_rna + hidden_pro\n",
    "        out_dim = nfeatures_rna + nfeatures_pro\n",
    "        \n",
    "        self.decoder = nn.Sequential(\n",
    "            LinBnDrop(latent_dim, hidden_dim, act=nn.LeakyReLU()),\n",
    "            LinBnDrop(hidden_dim, out_dim // 2, act=nn.LeakyReLU()),\n",
    "            LinBnDrop(out_dim // 2, out_dim, bn=False)\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.decoder(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we assemble the encoder and decoder into an autoencoder, which is defined as a PyTorch Lightning Module to simplify the training process. We will define the following:\n",
    "* `__init__` for creating and saving parameters and model\n",
    "* `forward`: for inference, which we will use to generate latent representations for downstream analysis\n",
    "* `configure_optimizers` for creating the optimizer and learning rate scheduler\n",
    "* `training_step` for calculating the loss (mean squared error (MSE) for our example) of a single batch\n",
    "* `validation_step` similar to `training_step` but on the validation set \n",
    "* `test_step` same as `validation_step` but on a test set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CiteAutoencoder(pl.LightningModule):\n",
    "    def __init__(self,\n",
    "                 nfeatures_rna: int,\n",
    "                 nfeatures_pro: int,\n",
    "                 hidden_rna: int,\n",
    "                 hidden_pro: int,\n",
    "                 latent_dim: int,\n",
    "                 p: float = 0,\n",
    "                 lr: float = 0.1):\n",
    "        \"\"\" Autoencoder for citeseq data \"\"\"\n",
    "        super().__init__()\n",
    "        \n",
    "        # save hyperparameters\n",
    "        self.save_hyperparameters()\n",
    " \n",
    "        self.encoder = Encoder(nfeatures_rna, nfeatures_pro, hidden_rna, hidden_pro, latent_dim, p)\n",
    "        self.decoder = Decoder(nfeatures_rna, nfeatures_pro, hidden_rna, hidden_pro, latent_dim)\n",
    "        \n",
    "        # example input array for visualizing network graph\n",
    "        self.example_input_array = torch.zeros(256, nfeatures_rna + nfeatures_pro)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # extract latent embeddings\n",
    "        z = self.encoder(x)\n",
    "        return z\n",
    "    \n",
    "    def configure_optimizers(self):\n",
    "        optimizer = optim.Adam(self.parameters(), lr=self.hparams.lr)\n",
    "        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer)\n",
    "        return {\"optimizer\": optimizer, \"lr_scheduler\": scheduler, \"monitor\": \"val_loss\"}\n",
    "    \n",
    "    def _get_reconstruction_loss(self, batch):\n",
    "        \"\"\" Calculate MSE loss for a given batch. \"\"\"\n",
    "        x, _ = batch\n",
    "        z = self.encoder(x)\n",
    "        x_hat = self.decoder(z)\n",
    "        # MSE loss\n",
    "        loss = F.mse_loss(x_hat, x)\n",
    "        return loss\n",
    "    \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        loss = self._get_reconstruction_loss(batch)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        loss = self._get_reconstruction_loss(batch)\n",
    "        self.log(\"val_loss\", loss)\n",
    "        \n",
    "    def test_step(self, batch, batch_idx):\n",
    "        loss = self._get_reconstruction_loss(batch)\n",
    "        self.log(\"test_loss\", loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will take advantage of the `Trainer` API from PyTorch Lightning to execute the training process. The two functions that we will be using are:\n",
    "* `fit`: Train a lightning module using the given train dataloader, and validate on the provided validation dataloader.\n",
    "* `test`: Test the given model on the provided dataloader."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train_citeseq(hidden_rna: int = 30, hidden_pro: int = 18,\n",
    "                  latent_dim: int = 24, p: float = 0.1, lr: float = 0.1):\n",
    "    trainer = pl.Trainer(default_root_dir=CHECKPOINT_PATH,\n",
    "                         gpus=1 if \"cuda\" in str(device) else 0,\n",
    "                         max_epochs=50,\n",
    "                         callbacks=[ModelCheckpoint(save_weights_only=True, mode=\"min\", monitor=\"val_loss\"),\n",
    "                                    LearningRateMonitor(\"epoch\")])\n",
    "    trainer.logger._log_graph = True\n",
    "    trainer.logger._default_hp_metric=None\n",
    "    \n",
    "    model = CiteAutoencoder(nfeatures_rna,\n",
    "                            nfeatures_pro,\n",
    "                            hidden_rna=hidden_rna,\n",
    "                            hidden_pro=hidden_pro,\n",
    "                            latent_dim=latent_dim,\n",
    "                            p=p,\n",
    "                            lr=lr)\n",
    "    trainer.fit(model, train_dl, val_dl)\n",
    "    \n",
    "    train_result = trainer.test(model, train_dl, verbose=False)\n",
    "    val_result = trainer.test(model, val_dl, verbose=False)\n",
    "    test_result = trainer.test(model, test_dl, verbose=False)\n",
    "    result = {\"train\": train_result, \"val\": val_result, \"test\": test_result, }\n",
    "    return model, result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model, result = train_citeseq()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(f\"Training loss:  {result['train'][0]['test_loss']:.3f}\")\n",
    "print(f\"Validation loss:  {result['val'][0]['test_loss']:.3f}\")\n",
    "print(f\"Test loss: {result['test'][0]['test_loss']:.3f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTorch Lightning automatically logs the training results into TensorBoard, which we can open like below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%tensorboard --host 0.0.0.0 --port 8000 --logdir saved_models/lightning_logs/version_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# kill tensorboard process\n",
    "!kill $(ps -e | grep 'tensorboard' | awk '{print $1}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize latent representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The latent space in our example has 24 dimensions. In order to visualize and inspect how different types of immune cells cluster in the latent space, we first use the trained model to generate the latent representations of the test dataset and then use UMAP, which is widely used in single-cell analysis, to reduce the dimensions for visualization in 2D. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_encodings = []\n",
    "test_labels = []\n",
    "    \n",
    "model.eval()\n",
    "with torch.no_grad():    \n",
    "    for x, y in tqdm(test_dl, desc=\"Encoding cells\"):\n",
    "        test_encodings.append(model(x.to(model.device)))\n",
    "        test_labels += y.to(torch.int).tolist()\n",
    "        \n",
    "test_embeds = torch.cat(test_encodings, dim=0).cpu().numpy()\n",
    "test_labels = le.inverse_transform(test_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run umap for dimensionality reduction and visualization\n",
    "embeds_umap = umap.UMAP(random_state=0).fit_transform(test_embeds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# visualize umap\n",
    "fig = px.scatter(x=embeds_umap[:, 0], y=embeds_umap[:, 1], color=test_labels, width=800, height=600,\n",
    "                 labels={\n",
    "                     \"x\": \"UMAP1\",\n",
    "                     \"y\": \"UMAP2\",\n",
    "                     \"color\": \"Cell type\"}\n",
    "                )\n",
    "fig.show(renderer=\"colab\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize and explore latent representations using TensorBoard, which provides a convinient interface for popular dimensionality reduction methods such as UMAP, TSNE, and PCA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# visualization with tensorboard\n",
    "writer = SummaryWriter(\"tensorboard/\")\n",
    "writer.add_embedding(test_embeds, metadata=test_labels)\n",
    "\n",
    "# wait for saving files\n",
    "time.sleep(3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%tensorboard --host 0.0.0.0 --port 8000 --logdir  tensorboard/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# kill tensorboard process\n",
    "!kill $(ps -e | grep 'tensorboard' | awk '{print $1}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<p style=\"font-size:larger;\">In this tutorial, we have built an autoencoder-based deep learning model for dimensionality reduction and visualization of single-cell CITE-seq data. We demonstrate that the integrative analysis of both transcriptomic and proteomic data achieves superior resolution in distinguishing between various immune cell types.</p>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}