models/scGPT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8a4fbd7",
   "metadata": {},
   "source": [
    "# scGPT \n",
    "This is a modified version of the code provided by the scGPT authors."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec9bd10a",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9406b4da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import copy\n",
    "import gc\n",
    "import json\n",
    "import os\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "import sys\n",
    "import time\n",
    "import traceback\n",
    "from typing import List, Tuple, Dict, Union, Optional\n",
    "import warnings\n",
    "import pandas as pd\n",
    "# from . import asyn\n",
    "import pickle\n",
    "import torch\n",
    "from anndata import AnnData\n",
    "import scanpy as sc\n",
    "import scvi\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import wandb\n",
    "from scipy.sparse import issparse\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, precision_recall_fscore_support, classification_report\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score\n",
    "from torchtext.vocab import Vocab\n",
    "from torchtext._torchtext import (\n",
    "    Vocab as VocabPybind,\n",
    ")\n",
    "from sklearn.metrics import confusion_matrix\n",
    "\n",
    "sys.path.insert(0, \"../\")\n",
    "import scgpt as scg\n",
    "from scgpt.model import TransformerModel, AdversarialDiscriminator\n",
    "from scgpt.tokenizer import tokenize_and_pad_batch, random_mask_value\n",
    "from scgpt.loss import (\n",
    "    masked_mse_loss,\n",
    "    masked_relative_error,\n",
    "    criterion_neg_log_bernoulli,\n",
    ")\n",
    "from scgpt.tokenizer.gene_tokenizer import GeneVocab\n",
    "from scgpt.preprocess import Preprocessor\n",
    "from scgpt import SubsetsBatchSampler\n",
    "from scgpt.utils import set_seed, category_str2int, eval_scib_metrics\n",
    "\n",
    "sc.set_figure_params(figsize=(6, 6))\n",
    "os.environ[\"KMP_WARNINGS\"] = \"off\"\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f3b5d67",
   "metadata": {},
   "source": [
    "## Hyperparameter Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07b5257",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "hyperparameter_defaults = dict(\n",
    "    seed=0,\n",
    "    dataset_name=\"ms\",\n",
    "    do_train=True,\n",
    "    load_model=\"../save/scGPT_human\",\n",
    "    mask_ratio=0.0,\n",
    "    epochs=15,\n",
    "    n_bins=51,\n",
    "    MVC=False, # Masked value prediction for cell embedding\n",
    "    ecs_thres=0.0, # Elastic cell similarity objective, 0.0 to 1.0, 0.0 to disable\n",
    "    dab_weight=0.0,\n",
    "    lr=1e-4,\n",
    "    batch_size=32,\n",
    "    layer_size=128,\n",
    "    nlayers=4,  # number of nn.TransformerEncoderLayer in nn.TransformerEncoder\n",
    "    nhead=4,  # number of heads in nn.MultiheadAttention\n",
    "    dropout=0.2,  # dropout probability\n",
    "    schedule_ratio=0.9,  # ratio of epochs for learning rate schedule\n",
    "    save_eval_interval=5,\n",
    "    fast_transformer=True,\n",
    "    pre_norm=False,\n",
    "    amp=True,  # Automatic Mixed Precision\n",
    "    include_zero_gene = False,\n",
    "    freeze = False, #freeze\n",
    "    DSBN = False,  # Domain-spec batchnorm\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94c08ee0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "run = wandb.init(\n",
    "    config=hyperparameter_defaults,\n",
    "    project=\"scGPT\",\n",
    "    reinit=True,\n",
    "    settings=wandb.Settings(start_method=\"fork\"),\n",
    "    entity=\"sabbagh\"\n",
    ")\n",
    "config = wandb.config\n",
    "print(config)\n",
    "\n",
    "set_seed(config.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d7890b7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# settings for input and preprocessing\n",
    "pad_token = \"<pad>\"\n",
    "special_tokens = [pad_token, \"<cls>\", \"<eoc>\"]\n",
    "mask_ratio = config.mask_ratio\n",
    "mask_value = \"auto\"  # for masked values, now it should always be auto\n",
    "\n",
    "include_zero_gene = config.include_zero_gene  # if True, include zero genes among hvgs in the training\n",
    "max_seq_len = 3001\n",
    "n_bins = config.n_bins\n",
    "\n",
    "# input/output representation\n",
    "input_style = \"binned\"  # \"normed_raw\", \"log1p\", or \"binned\"\n",
    "output_style = \"binned\"  # \"normed_raw\", \"log1p\", or \"binned\"\n",
    "\n",
    "# settings for training\n",
    "MLM = False  # whether to use masked language modeling, currently it is always on.\n",
    "CLS = True  # celltype classification objective\n",
    "ADV = False  # Adversarial training for batch correction\n",
    "CCE = False  # Contrastive cell embedding objective\n",
    "MVC = config.MVC  # Masked value prediction for cell embedding\n",
    "ECS = config.ecs_thres > 0  # Elastic cell similarity objective\n",
    "DAB = False  # Domain adaptation by reverse backpropagation, set to 2 for separate optimizer\n",
    "INPUT_BATCH_LABELS = False  # TODO: have these help MLM and MVC, while not to classifier\n",
    "input_emb_style = \"continuous\"  # \"category\" or \"continuous\" or \"scaling\"\n",
    "cell_emb_style = \"cls\"  # \"avg-pool\" or \"w-pool\" or \"cls\"\n",
    "adv_E_delay_epochs = 0  # delay adversarial training on encoder for a few epochs\n",
    "adv_D_delay_epochs = 0\n",
    "mvc_decoder_style = \"inner product\"\n",
    "ecs_threshold = config.ecs_thres\n",
    "dab_weight = config.dab_weight\n",
    "\n",
    "explicit_zero_prob = MLM and include_zero_gene  # whether explicit bernoulli for zeros\n",
    "do_sample_in_train = False and explicit_zero_prob  # sample the bernoulli in training\n",
    "\n",
    "per_seq_batch_sample = False\n",
    "\n",
    "# settings for optimizer\n",
    "lr = config.lr  # TODO: test learning rate ratio between two tasks\n",
    "lr_ADV = 1e-3  # learning rate for discriminator, used when ADV is True\n",
    "batch_size = config.batch_size\n",
    "eval_batch_size = config.batch_size\n",
    "epochs = config.epochs\n",
    "schedule_interval = 1\n",
    "\n",
    "# settings for the model\n",
    "fast_transformer = config.fast_transformer\n",
    "fast_transformer_backend = \"flash\"  # \"linear\" or \"flash\"\n",
    "embsize = config.layer_size  # embedding dimension\n",
    "d_hid = config.layer_size  # dimension of the feedforward network in TransformerEncoder\n",
    "nlayers = config.nlayers  # number of TransformerEncoderLayer in TransformerEncoder\n",
    "nhead = config.nhead  # number of heads in nn.MultiheadAttention\n",
    "dropout = config.dropout  # dropout probability\n",
    "\n",
    "# logging\n",
    "log_interval = 100  # iterations\n",
    "save_eval_interval = config.save_eval_interval  # epochs\n",
    "do_eval_scib_metrics = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17ff2309",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %% validate settings\n",
    "assert input_style in [\"normed_raw\", \"log1p\", \"binned\"]\n",
    "assert output_style in [\"normed_raw\", \"log1p\", \"binned\"]\n",
    "assert input_emb_style in [\"category\", \"continuous\", \"scaling\"]\n",
    "if input_style == \"binned\":\n",
    "    if input_emb_style == \"scaling\":\n",
    "        raise ValueError(\"input_emb_style `scaling` is not supported for binned input.\")\n",
    "elif input_style == \"log1p\" or input_style == \"normed_raw\":\n",
    "    if input_emb_style == \"category\":\n",
    "        raise ValueError(\n",
    "            \"input_emb_style `category` is not supported for log1p or normed_raw input.\"\n",
    "        )\n",
    "\n",
    "if input_emb_style == \"category\":\n",
    "    mask_value = n_bins + 1\n",
    "    pad_value = n_bins  # for padding gene expr values\n",
    "    n_input_bins = n_bins + 2\n",
    "else:\n",
    "    mask_value = -1\n",
    "    pad_value = -2\n",
    "    n_input_bins = n_bins\n",
    "\n",
    "if ADV and DAB:\n",
    "    raise ValueError(\"ADV and DAB cannot be both True.\")\n",
    "DAB_separate_optim = True if DAB > 1 else False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf7112d8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dataset_name = config.dataset_name\n",
    "save_dir = Path(f\"./save/dev_{dataset_name}-{time.strftime('%b%d-%H-%M')}/\")\n",
    "save_dir.mkdir(parents=True, exist_ok=True)\n",
    "print(f\"save to {save_dir}\")\n",
    "logger = scg.logger\n",
    "scg.utils.add_file_handler(logger, save_dir / \"run.log\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37fc7002",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95b50200",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if dataset_name == \"ms\":\n",
    "    data_dir = Path(\"../data/ms\")\n",
    "    adata = sc.read(data_dir / \"c_data.h5ad\")\n",
    "    adata_test = sc.read(data_dir / \"filtered_ms_adata.h5ad\")\n",
    "    adata.obs[\"celltype\"] = adata.obs[\"Factor Value[inferred cell type - authors labels]\"].astype(\"category\")\n",
    "    adata_test.obs[\"celltype\"] = adata_test.obs[\"Factor Value[inferred cell type - authors labels]\"].astype(\"category\")\n",
    "    adata.obs[\"batch_id\"]  = adata.obs[\"str_batch\"] = \"0\"\n",
    "    adata_test.obs[\"batch_id\"]  = adata_test.obs[\"str_batch\"] = \"1\"          \n",
    "    adata.var.set_index(adata.var[\"gene_name\"], inplace=True)\n",
    "    adata_test.var.set_index(adata.var[\"gene_name\"], inplace=True)\n",
    "    data_is_raw = False\n",
    "    filter_gene_by_counts = False\n",
    "    adata_test_raw = adata_test.copy()\n",
    "    adata = adata.concatenate(adata_test, batch_key=\"str_batch\")\n",
    "                \n",
    "# make the batch category column\n",
    "batch_id_labels = adata.obs[\"str_batch\"].astype(\"category\").cat.codes.values\n",
    "adata.obs[\"batch_id\"] = batch_id_labels\n",
    "celltype_id_labels = adata.obs[\"celltype\"].astype(\"category\").cat.codes.values\n",
    "celltypes = adata.obs[\"celltype\"].unique()\n",
    "num_types = len(np.unique(celltype_id_labels))\n",
    "id2type = dict(enumerate(adata.obs[\"celltype\"].astype(\"category\").cat.categories))\n",
    "adata.obs[\"celltype_id\"] = celltype_id_labels\n",
    "adata.var[\"gene_name\"] = adata.var.index.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0dc5a6ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if config.load_model is not None:\n",
    "    model_dir = Path(config.load_model)\n",
    "    model_config_file = model_dir / \"args.json\"\n",
    "    model_file = model_dir / \"best_model.pt\"\n",
    "    vocab_file = model_dir / \"vocab.json\"\n",
    "\n",
    "    vocab = GeneVocab.from_file(vocab_file)\n",
    "    shutil.copy(vocab_file, save_dir / \"vocab.json\")\n",
    "    for s in special_tokens:\n",
    "        if s not in vocab:\n",
    "            vocab.append_token(s)\n",
    "\n",
    "    adata.var[\"id_in_vocab\"] = [\n",
    "        1 if gene in vocab else -1 for gene in adata.var[\"gene_name\"]\n",
    "    ]\n",
    "    gene_ids_in_vocab = np.array(adata.var[\"id_in_vocab\"])\n",
    "    logger.info(\n",
    "        f\"match {np.sum(gene_ids_in_vocab >= 0)}/{len(gene_ids_in_vocab)} genes \"\n",
    "        f\"in vocabulary of size {len(vocab)}.\"\n",
    "    )\n",
    "    adata = adata[:, adata.var[\"id_in_vocab\"] >= 0]\n",
    "\n",
    "    # model\n",
    "    with open(model_config_file, \"r\") as f:\n",
    "        model_configs = json.load(f)\n",
    "    logger.info(\n",
    "        f\"Resume model from {model_file}, the model args will override the \"\n",
    "        f\"config {model_config_file}.\"\n",
    "    )\n",
    "    embsize = model_configs[\"embsize\"]\n",
    "    nhead = model_configs[\"nheads\"]\n",
    "    d_hid = model_configs[\"d_hid\"]\n",
    "    nlayers = model_configs[\"nlayers\"]\n",
    "    n_layers_cls = model_configs[\"n_layers_cls\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d08757ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# set up the preprocessor, use the args to config the workflow\n",
    "preprocessor = Preprocessor(\n",
    "    use_key=\"X\",  # the key in adata.layers to use as raw data\n",
    "    filter_gene_by_counts=filter_gene_by_counts,  # step 1\n",
    "    filter_cell_by_counts=False,  # step 2\n",
    "    normalize_total=1e4,  # 3. whether to normalize the raw data and to what sum\n",
    "    result_normed_key=\"X_normed\",  # the key in adata.layers to store the normalized data\n",
    "    log1p=data_is_raw,  # 4. whether to log1p the normalized data\n",
    "    result_log1p_key=\"X_log1p\",\n",
    "    subset_hvg=False,  # 5. whether to subset the raw data to highly variable genes\n",
    "    hvg_flavor=\"seurat_v3\" if data_is_raw else \"cell_ranger\",\n",
    "    binning=n_bins,  # 6. whether to bin the raw data and to what number of bins\n",
    "    result_binned_key=\"X_binned\",  # the key in adata.layers to store the binned data\n",
    ")\n",
    "\n",
    "\n",
    "adata_test = adata[adata.obs[\"str_batch\"] == \"1\"]\n",
    "adata = adata[adata.obs[\"str_batch\"] == \"0\"]\n",
    "\n",
    "preprocessor(adata, batch_key=None)\n",
    "preprocessor(adata_test, batch_key=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbc1b20f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_layer_key = {  # the values of this map coorespond to the keys in preprocessing\n",
    "    \"normed_raw\": \"X_normed\",\n",
    "    \"log1p\": \"X_normed\",\n",
    "    \"binned\": \"X_binned\",\n",
    "}[input_style]\n",
    "all_counts = (\n",
    "    adata.layers[input_layer_key].A\n",
    "    if issparse(adata.layers[input_layer_key])\n",
    "    else adata.layers[input_layer_key]\n",
    ")\n",
    "genes = adata.var[\"gene_name\"].tolist()\n",
    "\n",
    "celltypes_labels = adata.obs[\"celltype_id\"].tolist()  # make sure count from 0\n",
    "celltypes_labels = np.array(celltypes_labels)\n",
    "\n",
    "batch_ids = adata.obs[\"batch_id\"].tolist()\n",
    "num_batch_types = len(set(batch_ids))\n",
    "batch_ids = np.array(batch_ids)\n",
    "\n",
    "(\n",
    "    train_data,\n",
    "    valid_data,\n",
    "    train_celltype_labels,\n",
    "    valid_celltype_labels,\n",
    "    train_batch_labels,\n",
    "    valid_batch_labels,\n",
    ") = train_test_split(\n",
    "    all_counts, celltypes_labels, batch_ids, test_size=0.1, shuffle=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd701b1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if config.load_model is None:\n",
    "    vocab = Vocab(\n",
    "        VocabPybind(genes + special_tokens, None)\n",
    "    )  # bidirectional lookup [gene <-> int]\n",
    "vocab.set_default_index(vocab[\"<pad>\"])\n",
    "gene_ids = np.array(vocab(genes), dtype=int)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "818bfcc0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenized_train = tokenize_and_pad_batch(\n",
    "    train_data,\n",
    "    gene_ids,\n",
    "    max_len=max_seq_len,\n",
    "    vocab=vocab,\n",
    "    pad_token=pad_token,\n",
    "    pad_value=pad_value,\n",
    "    append_cls=True,  # append <cls> token at the beginning\n",
    "    include_zero_gene=include_zero_gene,\n",
    ")\n",
    "tokenized_valid = tokenize_and_pad_batch(\n",
    "    valid_data,\n",
    "    gene_ids,\n",
    "    max_len=max_seq_len,\n",
    "    vocab=vocab,\n",
    "    pad_token=pad_token,\n",
    "    pad_value=pad_value,\n",
    "    append_cls=True,\n",
    "    include_zero_gene=include_zero_gene,\n",
    ")\n",
    "logger.info(\n",
    "    f\"train set number of samples: {tokenized_train['genes'].shape[0]}, \"\n",
    "    f\"\\n\\t feature length: {tokenized_train['genes'].shape[1]}\"\n",
    ")\n",
    "logger.info(\n",
    "    f\"valid set number of samples: {tokenized_valid['genes'].shape[0]}, \"\n",
    "    f\"\\n\\t feature length: {tokenized_valid['genes'].shape[1]}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37a80818",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def prepare_data(sort_seq_batch=False) -> Tuple[Dict[str, torch.Tensor]]:\n",
    "    masked_values_train = random_mask_value(\n",
    "        tokenized_train[\"values\"],\n",
    "        mask_ratio=mask_ratio,\n",
    "        mask_value=mask_value,\n",
    "        pad_value=pad_value,\n",
    "    )\n",
    "    masked_values_valid = random_mask_value(\n",
    "        tokenized_valid[\"values\"],\n",
    "        mask_ratio=mask_ratio,\n",
    "        mask_value=mask_value,\n",
    "        pad_value=pad_value,\n",
    "    )\n",
    "    print(\n",
    "        f\"random masking at epoch {epoch:3d}, ratio of masked values in train: \",\n",
    "        f\"{(masked_values_train == mask_value).sum() / (masked_values_train - pad_value).count_nonzero():.4f}\",\n",
    "    )\n",
    "\n",
    "    input_gene_ids_train, input_gene_ids_valid = (\n",
    "        tokenized_train[\"genes\"],\n",
    "        tokenized_valid[\"genes\"],\n",
    "    )\n",
    "    input_values_train, input_values_valid = masked_values_train, masked_values_valid\n",
    "    target_values_train, target_values_valid = (\n",
    "        tokenized_train[\"values\"],\n",
    "        tokenized_valid[\"values\"],\n",
    "    )\n",
    "\n",
    "    tensor_batch_labels_train = torch.from_numpy(train_batch_labels).long()\n",
    "    tensor_batch_labels_valid = torch.from_numpy(valid_batch_labels).long()\n",
    "\n",
    "    tensor_celltype_labels_train = torch.from_numpy(train_celltype_labels).long()\n",
    "    tensor_celltype_labels_valid = torch.from_numpy(valid_celltype_labels).long()\n",
    "\n",
    "    if sort_seq_batch:  \n",
    "        train_sort_ids = np.argsort(train_batch_labels)\n",
    "        input_gene_ids_train = input_gene_ids_train[train_sort_ids]\n",
    "        input_values_train = input_values_train[train_sort_ids]\n",
    "        target_values_train = target_values_train[train_sort_ids]\n",
    "        tensor_batch_labels_train = tensor_batch_labels_train[train_sort_ids]\n",
    "        tensor_celltype_labels_train = tensor_celltype_labels_train[train_sort_ids]\n",
    "\n",
    "        valid_sort_ids = np.argsort(valid_batch_labels)\n",
    "        input_gene_ids_valid = input_gene_ids_valid[valid_sort_ids]\n",
    "        input_values_valid = input_values_valid[valid_sort_ids]\n",
    "        target_values_valid = target_values_valid[valid_sort_ids]\n",
    "        tensor_batch_labels_valid = tensor_batch_labels_valid[valid_sort_ids]\n",
    "        tensor_celltype_labels_valid = tensor_celltype_labels_valid[valid_sort_ids]\n",
    "\n",
    "    train_data_pt = {\n",
    "        \"gene_ids\": input_gene_ids_train,\n",
    "        \"values\": input_values_train,\n",
    "        \"target_values\": target_values_train,\n",
    "        \"batch_labels\": tensor_batch_labels_train,\n",
    "        \"celltype_labels\": tensor_celltype_labels_train,\n",
    "    }\n",
    "    valid_data_pt = {\n",
    "        \"gene_ids\": input_gene_ids_valid,\n",
    "        \"values\": input_values_valid,\n",
    "        \"target_values\": target_values_valid,\n",
    "        \"batch_labels\": tensor_batch_labels_valid,\n",
    "        \"celltype_labels\": tensor_celltype_labels_valid,\n",
    "    }\n",
    "\n",
    "    return train_data_pt, valid_data_pt\n",
    "\n",
    "\n",
    "# dataset\n",
    "class SeqDataset(Dataset):\n",
    "    def __init__(self, data: Dict[str, torch.Tensor]):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.data[\"gene_ids\"].shape[0]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return {k: v[idx] for k, v in self.data.items()}\n",
    "\n",
    "\n",
    "# data_loader\n",
    "def prepare_dataloader(\n",
    "    data_pt: Dict[str, torch.Tensor],\n",
    "    batch_size: int,\n",
    "    shuffle: bool = False,\n",
    "    intra_domain_shuffle: bool = False,\n",
    "    drop_last: bool = False,\n",
    "    num_workers: int = 0,\n",
    ") -> DataLoader:\n",
    "    if num_workers == 0:\n",
    "        num_workers = min(len(os.sched_getaffinity(0)), batch_size // 2)\n",
    "\n",
    "    dataset = SeqDataset(data_pt)\n",
    "\n",
    "    if per_seq_batch_sample:\n",
    "        # find the indices of samples in each seq batch\n",
    "        subsets = []\n",
    "        batch_labels_array = data_pt[\"batch_labels\"].numpy()\n",
    "        for batch_label in np.unique(batch_labels_array):\n",
    "            batch_indices = np.where(batch_labels_array == batch_label)[0].tolist()\n",
    "            subsets.append(batch_indices)\n",
    "        data_loader = DataLoader(\n",
    "            dataset=dataset,\n",
    "            batch_sampler=SubsetsBatchSampler(\n",
    "                subsets,\n",
    "                batch_size,\n",
    "                intra_subset_shuffle=intra_domain_shuffle,\n",
    "                inter_subset_shuffle=shuffle,\n",
    "                drop_last=drop_last,\n",
    "            ),\n",
    "            num_workers=num_workers,\n",
    "            pin_memory=True,\n",
    "        )\n",
    "        return data_loader\n",
    "\n",
    "    data_loader = DataLoader(\n",
    "        dataset=dataset,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=shuffle,\n",
    "        drop_last=drop_last,\n",
    "        num_workers=num_workers,\n",
    "        pin_memory=True,\n",
    "    )\n",
    "    return data_loader\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77105fda",
   "metadata": {},
   "source": [
    "## Loading Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "219bb9db",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "ntokens = len(vocab)  # size of vocabulary\n",
    "model = TransformerModel(\n",
    "    ntokens,\n",
    "    embsize,\n",
    "    nhead,\n",
    "    d_hid,\n",
    "    nlayers,\n",
    "    nlayers_cls=3,\n",
    "    n_cls=num_types if CLS else 1,\n",
    "    vocab=vocab,\n",
    "    dropout=dropout,\n",
    "    pad_token=pad_token,\n",
    "    pad_value=pad_value,\n",
    "    do_mvc=MVC,\n",
    "    do_dab=DAB,\n",
    "    use_batch_labels=INPUT_BATCH_LABELS,\n",
    "    num_batch_labels=num_batch_types,\n",
    "    domain_spec_batchnorm=config.DSBN,\n",
    "    input_emb_style=input_emb_style,\n",
    "    n_input_bins=n_input_bins,\n",
    "    cell_emb_style=cell_emb_style,\n",
    "    mvc_decoder_style=mvc_decoder_style,\n",
    "    ecs_threshold=ecs_threshold,\n",
    "    explicit_zero_prob=explicit_zero_prob,\n",
    "    use_fast_transformer=fast_transformer,\n",
    "    fast_transformer_backend=fast_transformer_backend,\n",
    "    pre_norm=config.pre_norm,\n",
    ")\n",
    "if config.load_model is not None:\n",
    "    try:\n",
    "        model.load_state_dict(torch.load(model_file))\n",
    "        logger.info(f\"Loading all model params from {model_file}\")\n",
    "    except:\n",
    "        # only load params that are in the model and match the size\n",
    "        model_dict = model.state_dict()\n",
    "        pretrained_dict = torch.load(model_file)\n",
    "        pretrained_dict = {\n",
    "            k: v\n",
    "            for k, v in pretrained_dict.items()\n",
    "            if k in model_dict and v.shape == model_dict[k].shape\n",
    "        }\n",
    "        for k, v in pretrained_dict.items():\n",
    "            logger.info(f\"Loading params {k} with shape {v.shape}\")\n",
    "        model_dict.update(pretrained_dict)\n",
    "        model.load_state_dict(model_dict)\n",
    "\n",
    "pre_freeze_param_count = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters() if p.requires_grad).values())\n",
    "\n",
    "# Freeze all pre-decoder weights\n",
    "for name, para in model.named_parameters():\n",
    "    print(\"-\"*20)\n",
    "    print(f\"name: {name}\")\n",
    "    if config.freeze and \"encoder\" in name and \"transformer_encoder\" not in name:\n",
    "    # if config.freeze and \"encoder\" in name:\n",
    "        print(f\"freezing weights for: {name}\")\n",
    "        para.requires_grad = False\n",
    "\n",
    "\n",
    "post_freeze_param_count = sum(dict((p.data_ptr(), p.numel()) for p in model.parameters() if p.requires_grad).values())\n",
    "\n",
    "logger.info(f\"Total Pre freeze Params {(pre_freeze_param_count )}\")\n",
    "logger.info(f\"Total Post freeze Params {(post_freeze_param_count )}\")\n",
    "wandb.log(\n",
    "        {\n",
    "            \"info/pre_freeze_param_count\": pre_freeze_param_count,\n",
    "            \"info/post_freeze_param_count\": post_freeze_param_count,\n",
    "        },\n",
    ")\n",
    "\n",
    "model.to(device)\n",
    "wandb.watch(model)\n",
    "\n",
    "if ADV:\n",
    "    discriminator = AdversarialDiscriminator(\n",
    "        d_model=embsize,\n",
    "        n_cls=num_batch_types,\n",
    "    ).to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e4ea79e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "criterion = masked_mse_loss\n",
    "criterion_cls = nn.CrossEntropyLoss()\n",
    "criterion_dab = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(), lr=lr, eps=1e-4 if config.amp else 1e-8\n",
    ")\n",
    "scheduler = torch.optim.lr_scheduler.StepLR(\n",
    "    optimizer, schedule_interval, gamma=config.schedule_ratio\n",
    ")\n",
    "if DAB_separate_optim:\n",
    "    optimizer_dab = torch.optim.Adam(model.parameters(), lr=lr)\n",
    "    scheduler_dab = torch.optim.lr_scheduler.StepLR(\n",
    "        optimizer_dab, schedule_interval, gamma=config.schedu|le_ratio\n",
    "    )\n",
    "if ADV:\n",
    "    criterion_adv = nn.CrossEntropyLoss()  # consider using label smoothing\n",
    "    optimizer_E = torch.optim.Adam(model.parameters(), lr=lr_ADV)\n",
    "    scheduler_E = torch.optim.lr_scheduler.StepLR(\n",
    "        optimizer_E, schedule_interval, gamma=config.schedule_ratio\n",
    "    )\n",
    "    optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=lr_ADV)\n",
    "    scheduler_D = torch.optim.lr_scheduler.StepLR(\n",
    "        optimizer_D, schedule_interval, gamma=config.schedule_ratio\n",
    "    )\n",
    "\n",
    "scaler = torch.cuda.amp.GradScaler(enabled=config.amp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4529c97e-2b6a-4608-a2b1-0b4c63d882f0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "epoch_time_list = []\n",
    "inference_time_list = []\n",
    "train_loss_list = []\n",
    "val_loss_list = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd29493-1126-49f3-9a6a-f7a4e4656f73",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "res_dir = './results'\n",
    "model_name = 'ms_default'\n",
    "number_split = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b734269a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model: nn.Module, loader: DataLoader) -> None:\n",
    "    \"\"\"\n",
    "    Train the model for one epoch.\n",
    "    \"\"\"\n",
    "    model.train()\n",
    "    (\n",
    "        total_loss,\n",
    "        total_mse,\n",
    "        total_cls,\n",
    "        total_cce,\n",
    "        total_mvc,\n",
    "        total_ecs,\n",
    "        total_dab,\n",
    "        total_adv_E,\n",
    "        total_adv_D,\n",
    "        total_zero_log_prob,\n",
    "        total_mvc_zero_log_prob,\n",
    "    ) = (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)\n",
    "    total_error = 0.0\n",
    "    start_time = time.time()\n",
    "\n",
    "    num_batches = len(loader)\n",
    "    for batch, batch_data in enumerate(loader):\n",
    "        input_gene_ids = batch_data[\"gene_ids\"].to(device)\n",
    "        input_values = batch_data[\"values\"].to(device)\n",
    "        target_values = batch_data[\"target_values\"].to(device)\n",
    "        batch_labels = batch_data[\"batch_labels\"].to(device)\n",
    "        celltype_labels = batch_data[\"celltype_labels\"].to(device)\n",
    "\n",
    "        src_key_padding_mask = input_gene_ids.eq(vocab[pad_token])\n",
    "        with torch.cuda.amp.autocast(enabled=config.amp):\n",
    "            output_dict = model(\n",
    "                input_gene_ids,\n",
    "                input_values,\n",
    "                src_key_padding_mask=src_key_padding_mask,\n",
    "                batch_labels=batch_labels if INPUT_BATCH_LABELS or config.DSBN else None,\n",
    "                CLS=CLS,\n",
    "                CCE=CCE,\n",
    "                MVC=MVC,\n",
    "                ECS=ECS,\n",
    "                do_sample=do_sample_in_train,\n",
    "                #generative_training=False\n",
    "            )\n",
    "\n",
    "            masked_positions = input_values.eq(mask_value)  # the postions to predict\n",
    "            loss = 0.0\n",
    "            metrics_to_log = {}\n",
    "            if MLM:\n",
    "                loss_mse = criterion(\n",
    "                    output_dict[\"mlm_output\"], target_values, masked_positions\n",
    "                )\n",
    "                loss = loss + loss_mse\n",
    "                metrics_to_log = {\"train/mse\": loss_mse.item()}\n",
    "            if explicit_zero_prob:\n",
    "                loss_zero_log_prob = criterion_neg_log_bernoulli(\n",
    "                    output_dict[\"mlm_zero_probs\"], target_values, masked_positions\n",
    "                )\n",
    "                loss = loss + loss_zero_log_prob\n",
    "                metrics_to_log.update({\"train/nzlp\": loss_zero_log_prob.item()})\n",
    "            if CLS:\n",
    "                loss_cls = criterion_cls(output_dict[\"cls_output\"], celltype_labels)\n",
    "                loss = loss + loss_cls\n",
    "                metrics_to_log.update({\"train/cls\": loss_cls.item()})\n",
    "\n",
    "                error_rate = 1 - (\n",
    "                    (output_dict[\"cls_output\"].argmax(1) == celltype_labels)\n",
    "                    .sum()\n",
    "                    .item()\n",
    "                ) / celltype_labels.size(0)\n",
    "            if CCE:\n",
    "                loss_cce = 10 * output_dict[\"loss_cce\"]\n",
    "                loss = loss + loss_cce\n",
    "                metrics_to_log.update({\"train/cce\": loss_cce.item()})\n",
    "            if MVC:\n",
    "                loss_mvc = criterion(\n",
    "                    output_dict[\"mvc_output\"], target_values, masked_positions\n",
    "                )\n",
    "                loss = loss + loss_mvc\n",
    "                metrics_to_log.update({\"train/mvc\": loss_mvc.item()})\n",
    "            if MVC and explicit_zero_prob:\n",
    "                loss_mvc_zero_log_prob = criterion_neg_log_bernoulli(\n",
    "                    output_dict[\"mvc_zero_probs\"], target_values, masked_positions\n",
    "                )\n",
    "                loss = loss + loss_mvc_zero_log_prob\n",
    "                metrics_to_log.update({\"train/mvc_nzlp\": loss_mvc_zero_log_prob.item()})\n",
    "            if ECS:\n",
    "                loss_ecs = 10 * output_dict[\"loss_ecs\"]\n",
    "                loss = loss + loss_ecs\n",
    "                metrics_to_log.update({\"train/ecs\": loss_ecs.item()})\n",
    "            if DAB:\n",
    "                # try weighting and separate optimizer\n",
    "                loss_dab = criterion_dab(output_dict[\"dab_output\"], batch_labels)\n",
    "                loss = loss + dab_weight * loss_dab\n",
    "                metrics_to_log.update({\"train/dab\": loss_dab.item()})\n",
    "        \n",
    "        model.zero_grad()\n",
    "        scaler.scale(loss).backward()\n",
    "        scaler.unscale_(optimizer)\n",
    "        with warnings.catch_warnings(record=True) as w:\n",
    "            warnings.filterwarnings(\"always\")\n",
    "            torch.nn.utils.clip_grad_norm_(\n",
    "                model.parameters(),\n",
    "                1.0,\n",
    "                error_if_nonfinite=False if scaler.is_enabled() else True,\n",
    "            )\n",
    "            if len(w) > 0:\n",
    "                logger.warning(\n",
    "                    f\"Found infinite gradient. This may be caused by the gradient \"\n",
    "                    f\"scaler. The current scale is {scaler.get_scale()}. This warning \"\n",
    "                    \"can be ignored if no longer occurs after autoscaling of the scaler.\"\n",
    "                )\n",
    "        scaler.step(optimizer)\n",
    "        scaler.update()\n",
    "\n",
    "        if ADV:\n",
    "            # rerun the model for adversarial training\n",
    "            output_dict = model(\n",
    "                input_gene_ids,\n",
    "                input_values,\n",
    "                src_key_padding_mask=src_key_padding_mask,\n",
    "                batch_labels=batch_labels if INPUT_BATCH_LABELS or config.DSBN else None,\n",
    "                CLS=CLS,\n",
    "                CCE=CCE,\n",
    "                MVC=MVC,\n",
    "                ECS=ECS,\n",
    "                do_sample=do_sample_in_train,\n",
    "                #generative_training=False\n",
    "            )\n",
    "\n",
    "            # TRAINING DISCRIMINATOR\n",
    "            loss_adv_D = criterion_adv(\n",
    "                discriminator(output_dict[\"cell_emb\"].detach()), batch_labels\n",
    "            )\n",
    "            if epoch > adv_D_delay_epochs:\n",
    "                discriminator.zero_grad()\n",
    "                loss_adv_D.backward()\n",
    "                optimizer_D.step()\n",
    "\n",
    "            # TRAINING ENCODER\n",
    "            loss_adv_E = -criterion_adv(\n",
    "                discriminator(output_dict[\"cell_emb\"]), batch_labels\n",
    "            )\n",
    "            # NOTE: the loss is negative here because we want to maximize\n",
    "            # the cross_entropy_loss, in other words, disguise against the discriminator\n",
    "            if epoch > adv_E_delay_epochs:\n",
    "                model.zero_grad()\n",
    "                discriminator.zero_grad()\n",
    "                loss_adv_E.backward()\n",
    "                optimizer_E.step()\n",
    "\n",
    "        wandb.log(metrics_to_log)\n",
    "        \n",
    "        total_loss += loss.item()\n",
    "        total_mse += loss_mse.item() if MLM else 0.0\n",
    "        total_cls += loss_cls.item() if CLS else 0.0\n",
    "        total_cce += loss_cce.item() if CCE else 0.0\n",
    "        total_mvc += loss_mvc.item() if MVC else 0.0\n",
    "        total_ecs += loss_ecs.item() if ECS else 0.0\n",
    "        total_dab += loss_dab.item() if DAB else 0.0\n",
    "        total_adv_E += loss_adv_E.item() if ADV else 0.0\n",
    "        total_adv_D += loss_adv_D.item() if ADV else 0.0\n",
    "        total_zero_log_prob += loss_zero_log_prob.item() if explicit_zero_prob else 0.0\n",
    "        total_mvc_zero_log_prob += (\n",
    "            loss_mvc_zero_log_prob.item() if MVC and explicit_zero_prob else 0.0\n",
    "        )\n",
    "        total_error += error_rate\n",
    "        if batch % log_interval == 0 and batch > 0:\n",
    "            lr = scheduler.get_last_lr()[0]\n",
    "            ms_per_batch = (time.time() - start_time) * 1000 / log_interval\n",
    "            cur_loss = total_loss / log_interval\n",
    "            cur_mse = total_mse / log_interval\n",
    "            cur_cls = total_cls / log_interval if CLS else 0.0\n",
    "            cur_cce = total_cce / log_interval if CCE else 0.0\n",
    "            cur_mvc = total_mvc / log_interval if MVC else 0.0\n",
    "            cur_ecs = total_ecs / log_interval if ECS else 0.0\n",
    "            cur_dab = total_dab / log_interval if DAB else 0.0\n",
    "            cur_adv_E = total_adv_E / log_interval if ADV else 0.0\n",
    "            cur_adv_D = total_adv_D / log_interval if ADV else 0.0\n",
    "            cur_zero_log_prob = (\n",
    "                total_zero_log_prob / log_interval if explicit_zero_prob else 0.0\n",
    "            )\n",
    "            cur_mvc_zero_log_prob = (\n",
    "                total_mvc_zero_log_prob / log_interval\n",
    "                if MVC and explicit_zero_prob\n",
    "                else 0.0\n",
    "            )\n",
    "            cur_error = total_error / log_interval\n",
    "            # ppl = math.exp(cur_loss)\n",
    "            logger.info(\n",
    "                f\"| epoch {epoch:3d} | {batch:3d}/{num_batches:3d} batches | \"\n",
    "                f\"lr {lr:05.4f} | ms/batch {ms_per_batch:5.2f} | \"\n",
    "                f\"loss {cur_loss:5.2f} | \"\n",
    "                + (f\"mse {cur_mse:5.2f} | mre {cur_error:5.2f} |\" if MLM else \"\")\n",
    "                + (f\"cls {cur_cls:5.2f} | \" if CLS else \"\")\n",
    "                + (f\"err {cur_error:5.2f} | \" if CLS else \"\")\n",
    "                + (f\"cce {cur_cce:5.2f} |\" if CCE else \"\")\n",
    "                + (f\"mvc {cur_mvc:5.2f} |\" if MVC else \"\")\n",
    "                + (f\"ecs {cur_ecs:5.2f} |\" if ECS else \"\")\n",
    "                + (f\"dab {cur_dab:5.2f} |\" if DAB else \"\")\n",
    "                + (f\"adv_E {cur_adv_E:5.2f} |\" if ADV else \"\")\n",
    "                + (f\"adv_D {cur_adv_D:5.2f} |\" if ADV else \"\")\n",
    "                + (f\"nzlp {cur_zero_log_prob:5.2f} |\" if explicit_zero_prob else \"\")\n",
    "                + (\n",
    "                    f\"mvc_nzlp {cur_mvc_zero_log_prob:5.2f} |\"\n",
    "                    if MVC and explicit_zero_prob\n",
    "                    else \"\"\n",
    "                )\n",
    "            )\n",
    "            saving_dir = os.path.join(res_dir, model_name, 'train', f'split_{number_split}')\n",
    "            if not os.path.exists(saving_dir):\n",
    "                os.makedirs(saving_dir)\n",
    "\n",
    "            resources_usage_df_path = os.path.join(saving_dir, f'train_batch_loss.csv')\n",
    "            if os.path.exists(resources_usage_df_path):\n",
    "                resources_usage_df = pd.read_csv(resources_usage_df_path)\n",
    "            else:\n",
    "                resources_usage_df = pd.DataFrame(columns=['train_loss'])\n",
    "\n",
    "\n",
    "            resources_usage_df = pd.concat([resources_usage_df, pd.DataFrame({'train_loss':[cur_loss]})])\n",
    "            resources_usage_df.to_csv(resources_usage_df_path, index=False)\n",
    "            \n",
    "            total_loss = 0\n",
    "            total_mse = 0\n",
    "            total_cls = 0\n",
    "            total_cce = 0\n",
    "            total_mvc = 0\n",
    "            total_ecs = 0\n",
    "            total_dab = 0\n",
    "            total_adv_E = 0\n",
    "            total_adv_D = 0\n",
    "            total_zero_log_prob = 0\n",
    "            total_mvc_zero_log_prob = 0\n",
    "            total_error = 0\n",
    "            start_time = time.time()\n",
    "\n",
    "\n",
    "def define_wandb_metrcis():\n",
    "    wandb.define_metric(\"valid/mse\", summary=\"min\", step_metric=\"epoch\")\n",
    "    wandb.define_metric(\"valid/mre\", summary=\"min\", step_metric=\"epoch\")\n",
    "    wandb.define_metric(\"valid/dab\", summary=\"min\", step_metric=\"epoch\")\n",
    "    wandb.define_metric(\"valid/sum_mse_dab\", summary=\"min\", step_metric=\"epoch\")\n",
    "    wandb.define_metric(\"test/avg_bio\", summary=\"max\")\n",
    "\n",
    "\n",
    "def evaluate(model: nn.Module, loader: DataLoader, return_raw: bool = False, is_test: bool = False) -> float:\n",
    "    \"\"\"\n",
    "    Evaluate the model on the evaluation data.\n",
    "    \"\"\"\n",
    "    if is_test:\n",
    "        all_counts = (\n",
    "        adata.layers[input_layer_key].A\n",
    "        if issparse(adata.layers[input_layer_key])\n",
    "        else adata.layers[input_layer_key])\n",
    "\n",
    "        celltypes_labels = adata.obs[\"celltype_id\"].tolist()  # make sure count from 0\n",
    "        celltypes_labels = np.array(celltypes_labels)\n",
    "\n",
    "        batch_ids = adata.obs[\"batch_id\"].tolist()\n",
    "        batch_ids = np.array(batch_ids)\n",
    "\n",
    "        tokenized_test = tokenize_and_pad_batch(\n",
    "            all_counts,\n",
    "            gene_ids,\n",
    "            max_len=max_seq_len,\n",
    "            vocab=vocab,\n",
    "            pad_token=pad_token,\n",
    "            pad_value=pad_value,\n",
    "            append_cls=True,  # append <cls> token at the beginning\n",
    "            include_zero_gene=include_zero_gene,\n",
    "        )\n",
    "\n",
    "        input_values_test = random_mask_value(\n",
    "            tokenized_test[\"values\"],\n",
    "            mask_ratio=mask_ratio,\n",
    "            mask_value=mask_value,\n",
    "            pad_value=pad_value,\n",
    "        )\n",
    "\n",
    "        test_data_pt = {\n",
    "            \"gene_ids\": tokenized_test[\"genes\"],\n",
    "            \"values\": input_values_test,\n",
    "            \"target_values\": tokenized_test[\"values\"],\n",
    "            \"batch_labels\": torch.from_numpy(batch_ids).long(),\n",
    "            \"celltype_labels\": torch.from_numpy(celltypes_labels).long(),\n",
    "        }\n",
    "\n",
    "        loader = DataLoader(\n",
    "            dataset=SeqDataset(test_data_pt),\n",
    "            batch_size=eval_batch_size,\n",
    "            shuffle=False,\n",
    "            drop_last=False,\n",
    "            num_workers=min(len(os.sched_getaffinity(0)), eval_batch_size // 2),\n",
    "            pin_memory=True,\n",
    "        )\n",
    "\n",
    "    model.eval()\n",
    "    total_loss = 0.0\n",
    "    total_error = 0.0\n",
    "    total_dab = 0.0\n",
    "    total_num = 0\n",
    "    predictions = []\n",
    "    truths = []\n",
    "    with torch.no_grad():\n",
    "        for batch_data in loader:\n",
    "            input_gene_ids = batch_data[\"gene_ids\"].to(device)\n",
    "            input_values = batch_data[\"values\"].to(device)\n",
    "            target_values = batch_data[\"target_values\"].to(device)\n",
    "            batch_labels = batch_data[\"batch_labels\"].to(device)\n",
    "            celltype_labels = batch_data[\"celltype_labels\"].to(device)\n",
    "\n",
    "            src_key_padding_mask = input_gene_ids.eq(vocab[pad_token])\n",
    "            with torch.cuda.amp.autocast(enabled=config.amp):\n",
    "                output_dict = model(\n",
    "                    input_gene_ids,\n",
    "                    input_values,\n",
    "                    src_key_padding_mask=src_key_padding_mask,\n",
    "                    batch_labels=batch_labels if INPUT_BATCH_LABELS or config.DSBN else None,\n",
    "                    CLS=CLS,  # evaluation does not need CLS or CCE\n",
    "                    CCE=False,\n",
    "                    MVC=False,\n",
    "                    ECS=False,\n",
    "                    do_sample=do_sample_in_train,\n",
    "                    #generative_training = False,\n",
    "                )\n",
    "                output_values = output_dict[\"cls_output\"]\n",
    "                loss = criterion_cls(output_values, celltype_labels)\n",
    "\n",
    "                if DAB:\n",
    "                    loss_dab = criterion_dab(output_dict[\"dab_output\"], batch_labels)\n",
    "\n",
    "            total_loss += loss.item() * len(input_gene_ids)\n",
    "            accuracy = (output_values.argmax(1) == celltype_labels).sum().item()\n",
    "            total_error += (1 - accuracy / len(input_gene_ids)) * len(input_gene_ids)\n",
    "            total_dab += loss_dab.item() * len(input_gene_ids) if DAB else 0.0\n",
    "            total_num += len(input_gene_ids)\n",
    "            preds = output_values.argmax(1).cpu().numpy()\n",
    "            predictions.append(list(preds))\n",
    "            truths.append(list(celltype_labels.cpu().numpy()))\n",
    "\n",
    "\n",
    "    predictions = [item for sub_list in predictions for item in sub_list]\n",
    "    truths = [item for sub_list in truths for item in sub_list]\n",
    "    wandb.log(\n",
    "        {\n",
    "            \"valid/mse\": total_loss / total_num,\n",
    "            \"valid/err\": total_error / total_num,\n",
    "            \"valid/dab\": total_dab / total_num,\n",
    "            \"valid/sum_mse_dab\": (total_loss + dab_weight * total_dab) / total_num,\n",
    "            \"epoch\": epoch,\n",
    "        },\n",
    "    )\n",
    "\n",
    "    if return_raw:\n",
    "        return np.concatenate(predictions, axis=0)\n",
    "    \n",
    "    \n",
    "    \n",
    "    return total_loss / total_num, total_error / total_num, truths, predictions\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "734a503c",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e48b893",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "best_val_loss = float(\"inf\")\n",
    "best_avg_bio = 0.0\n",
    "best_model = None\n",
    "define_wandb_metrcis()\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    i = epoch\n",
    "    epoch_start_time = time.time()\n",
    "    train_data_pt, valid_data_pt = prepare_data(sort_seq_batch=per_seq_batch_sample)\n",
    "    train_loader = prepare_dataloader(\n",
    "        train_data_pt,\n",
    "        batch_size=batch_size,\n",
    "        shuffle=False,\n",
    "        intra_domain_shuffle=True,\n",
    "        drop_last=False,\n",
    "    )\n",
    "    valid_loader = prepare_dataloader(\n",
    "        valid_data_pt,\n",
    "        batch_size=eval_batch_size,\n",
    "        shuffle=False,\n",
    "        intra_domain_shuffle=False,\n",
    "        drop_last=False,\n",
    "    )\n",
    "\n",
    "    if config.do_train:\n",
    "        train(\n",
    "            model,\n",
    "            loader=train_loader,\n",
    "        )\n",
    "    val_loss, val_err, truths, predictions = evaluate(\n",
    "        model,\n",
    "        loader=valid_loader,)\n",
    "    test_loss, test_err, ttruths, tpredictions = evaluate(\n",
    "        model,\n",
    "        loader=adata_test,\n",
    "        is_test=True\n",
    "    )\n",
    "    elapsed = time.time() - epoch_start_time\n",
    "    max_memory_reserved = torch.cuda.max_memory_reserved() / 1e9\n",
    "    \n",
    "    \n",
    "    saving_dir = os.path.join(res_dir, model_name, 'train', f'split_{number_split}')\n",
    "    if not os.path.exists(saving_dir):\n",
    "        os.makedirs(saving_dir)\n",
    "\n",
    "    resources_usage_df_path = os.path.join(saving_dir, f'resources_usage.csv')\n",
    "    if os.path.exists(resources_usage_df_path):\n",
    "        resources_usage_df = pd.read_csv(resources_usage_df_path)\n",
    "    else:\n",
    "        resources_usage_df = pd.DataFrame(columns=['time_per_epoch', 'memory_reserved'])\n",
    "\n",
    "\n",
    "    resources_usage_df = pd.concat([resources_usage_df, pd.DataFrame({'time_per_epoch':[elapsed], 'memory_reserved':[max_memory_reserved]})])\n",
    "    resources_usage_df.to_csv(resources_usage_df_path, index=False)\n",
    "    \n",
    "\n",
    "    \n",
    "    conf_mat = confusion_matrix(truths, predictions)\n",
    "    report = classification_report(truths, predictions, digits=4) #, target_names=label_dict.tolist()\n",
    "    report_dict = classification_report(truths, predictions, digits=4, output_dict=True) #, target_names=label_dict.tolist()\n",
    "    saving_dir = os.path.join(res_dir, model_name, 'val', f'split_{number_split}')\n",
    "\n",
    "    if not os.path.exists(saving_dir):\n",
    "        os.makedirs(saving_dir)\n",
    "\n",
    "    resources_usage_df_path = os.path.join(saving_dir, f'val_loss.csv')\n",
    "    if os.path.exists(resources_usage_df_path):\n",
    "        resources_usage_df = pd.read_csv(resources_usage_df_path)\n",
    "    else:\n",
    "        resources_usage_df = pd.DataFrame(columns=['val_loss'])\n",
    "\n",
    "    resources_usage_df = pd.concat([resources_usage_df, pd.DataFrame({'val_loss':[val_loss]})])\n",
    "    resources_usage_df.to_csv(resources_usage_df_path, index=False)\n",
    "    pd.DataFrame(conf_mat).to_csv(os.path.join(saving_dir, f'conf_mat_E{i}.csv'), index=False)\n",
    "    pd.DataFrame(report_dict).transpose().to_csv(os.path.join(saving_dir, f'report_E{i}.csv'), index=False)\n",
    "    \n",
    "    \n",
    "    conf_mat = confusion_matrix(ttruths, tpredictions)\n",
    "    report = classification_report(ttruths, tpredictions, digits=4) #, target_names=label_dict.tolist()\n",
    "    report_dict = classification_report(ttruths, tpredictions, digits=4, output_dict=True) #, target_names=label_dict.tolist()\n",
    "    saving_dir = os.path.join(res_dir, model_name, 'test', f'split_{number_split}')\n",
    "\n",
    "    if not os.path.exists(saving_dir):\n",
    "        os.makedirs(saving_dir)\n",
    "\n",
    "    resources_usage_df_path = os.path.join(saving_dir, f'test_loss.csv')\n",
    "    if os.path.exists(resources_usage_df_path):\n",
    "        resources_usage_df = pd.read_csv(resources_usage_df_path)\n",
    "    else:\n",
    "        resources_usage_df = pd.DataFrame(columns=['test_loss'])\n",
    "\n",
    "    resources_usage_df = pd.concat([resources_usage_df, pd.DataFrame({'test_loss':[test_loss]})])\n",
    "    resources_usage_df.to_csv(resources_usage_df_path, index=False)\n",
    "\n",
    "    pd.DataFrame(conf_mat).to_csv(os.path.join(saving_dir, f'conf_mat_E{i}.csv'), index=False)\n",
    "    pd.DataFrame(report_dict).transpose().to_csv(os.path.join(saving_dir, f'report_E{i}.csv'), index=False)\n",
    "        \n",
    "        \n",
    "    logger.info(\"-\" * 89)\n",
    "    logger.info(\n",
    "        f\"| end of epoch {epoch:3d} | time: {elapsed:5.2f}s | \"\n",
    "        f\"valid loss/mse {val_loss:5.4f} | err {val_err:5.4f}\"\n",
    "    )\n",
    "    logger.info(\"-\" * 89)\n",
    "\n",
    "    if val_loss < best_val_loss:\n",
    "        best_val_loss = val_loss\n",
    "        best_model = copy.deepcopy(model)\n",
    "        best_model_epoch = epoch\n",
    "        logger.info(f\"Best model with score {best_val_loss:5.4f}\")\n",
    "        torch.save(best_model.state_dict(), f\"{saving_dir}/best_model.pt\")\n",
    "\n",
    "    \n",
    "    scheduler.step()\n",
    "    if DAB_separate_optim:\n",
    "        scheduler_dab.step()\n",
    "    if ADV:\n",
    "        scheduler_D.step()\n",
    "        scheduler_E.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a6ce176",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# %% inference\n",
    "def test(model: nn.Module, adata: DataLoader) -> float:\n",
    "    all_counts = (\n",
    "        adata.layers[input_layer_key].A\n",
    "        if issparse(adata.layers[input_layer_key])\n",
    "        else adata.layers[input_layer_key]\n",
    "    )\n",
    "\n",
    "    celltypes_labels = adata.obs[\"celltype_id\"].tolist()  # make sure count from 0\n",
    "    celltypes_labels = np.array(celltypes_labels)\n",
    "\n",
    "    batch_ids = adata.obs[\"batch_id\"].tolist()\n",
    "    batch_ids = np.array(batch_ids)\n",
    "\n",
    "    tokenized_test = tokenize_and_pad_batch(\n",
    "        all_counts,\n",
    "        gene_ids,\n",
    "        max_len=max_seq_len,\n",
    "        vocab=vocab,\n",
    "        pad_token=pad_token,\n",
    "        pad_value=pad_value,\n",
    "        append_cls=True,  # append <cls> token at the beginning\n",
    "        include_zero_gene=include_zero_gene,\n",
    "    )\n",
    "\n",
    "    input_values_test = random_mask_value(\n",
    "        tokenized_test[\"values\"],\n",
    "        mask_ratio=mask_ratio,\n",
    "        mask_value=mask_value,\n",
    "        pad_value=pad_value,\n",
    "    )\n",
    "\n",
    "    test_data_pt = {\n",
    "        \"gene_ids\": tokenized_test[\"genes\"],\n",
    "        \"values\": input_values_test,\n",
    "        \"target_values\": tokenized_test[\"values\"],\n",
    "        \"batch_labels\": torch.from_numpy(batch_ids).long(),\n",
    "        \"celltype_labels\": torch.from_numpy(celltypes_labels).long(),\n",
    "    }\n",
    "\n",
    "    test_loader = DataLoader(\n",
    "        dataset=SeqDataset(test_data_pt),\n",
    "        batch_size=eval_batch_size,\n",
    "        shuffle=False,\n",
    "        drop_last=False,\n",
    "        num_workers=min(len(os.sched_getaffinity(0)), eval_batch_size // 2),\n",
    "        pin_memory=True,\n",
    "    )\n",
    "\n",
    "    model.eval()\n",
    "    predictions = evaluate(\n",
    "        model,\n",
    "        loader=test_loader,\n",
    "        return_raw=True,\n",
    "    )\n",
    "\n",
    "    # compute accuracy, precision, recall, f1\n",
    "    from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "\n",
    "    accuracy = accuracy_score(celltypes_labels, predictions)\n",
    "    precision = precision_score(celltypes_labels, predictions, average=\"macro\")\n",
    "    recall = recall_score(celltypes_labels, predictions, average=\"macro\")\n",
    "    macro_f1 = f1_score(celltypes_labels, predictions, average=\"macro\")\n",
    "\n",
    "    logger.info(\n",
    "        f\"Accuracy: {accuracy:.3f}, Precision: {precision:.3f}, Recall: {recall:.3f}, \"\n",
    "        f\"Macro F1: {macro_f1:.3f}\"\n",
    "    )\n",
    "\n",
    "    results = {\n",
    "        \"test/accuracy\": accuracy,\n",
    "        \"test/precision\": precision,\n",
    "        \"test/recall\": recall,\n",
    "        \"test/macro_f1\": macro_f1,\n",
    "    }\n",
    "\n",
    "    return predictions, celltypes_labels, results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16236bf2",
   "metadata": {},
   "source": [
    "## Inferece"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79730e19",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictions, labels, results = test(best_model, adata_test)\n",
    "adata_test_raw.obs[\"predictions\"] = [id2type[p] for p in predictions]\n",
    "\n",
    "# plot\n",
    "palette_ = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"] \n",
    "palette_ = plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"] + plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"] + plt.rcParams[\"axes.prop_cycle\"].by_key()[\"color\"]\n",
    "palette_ = {c: palette_[i] for i, c in enumerate(celltypes)}\n",
    "\n",
    "with plt.rc_context({\"figure.figsize\": (6, 4), \"figure.dpi\": (300)}):\n",
    "    sc.pl.umap(\n",
    "        adata_test_raw,\n",
    "        color=[\"celltype\", \"predictions\"],\n",
    "        palette=palette_,\n",
    "        show=False,\n",
    "    )\n",
    "    plt.savefig(save_dir / \"results.png\", dpi=300)\n",
    "\n",
    "save_dict = {\n",
    "    \"predictions\": predictions,\n",
    "    \"labels\": labels,\n",
    "    \"results\": results,\n",
    "    \"id_maps\": id2type\n",
    "}\n",
    "with open(save_dir / \"results.pkl\", \"wb\") as f:\n",
    "    pickle.dump(save_dict, f)\n",
    "\n",
    "results[\"test/cell_umap\"] = wandb.Image(\n",
    "    str(save_dir / \"results.png\"),\n",
    "    caption=f\"predictions macro f1 {results['test/macro_f1']:.3f}\",\n",
    ")\n",
    "wandb.log(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b72676ab",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.metrics import confusion_matrix\n",
    "celltypes = list(celltypes)\n",
    "for i in set([id2type[p] for p in predictions]):\n",
    "    if i not in celltypes:\n",
    "        celltypes.remove(i)\n",
    "cm = confusion_matrix(labels, predictions)\n",
    "cm = cm.astype(\"float\") / cm.sum(axis=1)[:, np.newaxis]\n",
    "cm = pd.DataFrame(cm, index=celltypes[:cm.shape[0]], columns=celltypes[:cm.shape[1]])\n",
    "plt.figure(figsize=(10, 10))\n",
    "sns.heatmap(cm, annot=True, fmt=\".1f\", cmap=\"Blues\")\n",
    "plt.savefig(save_dir / \"confusion_matrix.png\", dpi=300)\n",
    "\n",
    "results[\"test/confusion_matrix\"] = wandb.Image(\n",
    "    str(save_dir / \"confusion_matrix.png\"),\n",
    "    caption=f\"confusion matrix\",\n",
    ")"
   ]
  }
 ],
 "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.8.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}