paired_integration.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c7e65482",
   "metadata": {},
   "source": [
    "# Paired integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a135aeb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download data\n",
    "# !mkdir data\n",
    "# !wget https://cf.10xgenomics.com/samples/cell-arc/2.0.0/10k_PBMC_Multiome_nextgem_Chromium_X/10k_PBMC_Multiome_nextgem_Chromium_X_filtered_feature_bc_matrix.tar.gz\n",
    "# !cd data; tar -xzf 10k_PBMC_Multiome_nextgem_Chromium_X_filtered_feature_bc_matrix.tar.gz\n",
    "# !mkdir write"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94b381b5",
   "metadata": {},
   "source": [
    "## Environment setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7f41c71c",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'get_num_classes' from 'torchmetrics.utilities.data' (d:\\Anaconda\\envs\\sc_multiomics\\lib\\site-packages\\torchmetrics\\utilities\\data.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[1;32mf:\\sc_multiomics\\sc_multiomics\\paired_integration.ipynb Cell 4\u001b[0m line \u001b[0;36m8\n\u001b[0;32m      <a href='vscode-notebook-cell:/f%3A/sc_multiomics/sc_multiomics/paired_integration.ipynb#W3sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39manndata2ri\u001b[39;00m\n\u001b[0;32m      <a href='vscode-notebook-cell:/f%3A/sc_multiomics/sc_multiomics/paired_integration.ipynb#W3sZmlsZQ%3D%3D?line=6'>7</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mlogging\u001b[39;00m\n\u001b[1;32m----> <a href='vscode-notebook-cell:/f%3A/sc_multiomics/sc_multiomics/paired_integration.ipynb#W3sZmlsZQ%3D%3D?line=7'>8</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtorchmetrics\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mutilities\u001b[39;00m\u001b[39m.\u001b[39;00m\u001b[39mdata\u001b[39;00m \u001b[39mimport\u001b[39;00m get_num_classes\n\u001b[0;32m      <a href='vscode-notebook-cell:/f%3A/sc_multiomics/sc_multiomics/paired_integration.ipynb#W3sZmlsZQ%3D%3D?line=8'>9</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mscvi\u001b[39;00m\n\u001b[0;32m     <a href='vscode-notebook-cell:/f%3A/sc_multiomics/sc_multiomics/paired_integration.ipynb#W3sZmlsZQ%3D%3D?line=9'>10</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39mos\u001b[39;00m\n",
      "\u001b[1;31mImportError\u001b[0m: cannot import name 'get_num_classes' from 'torchmetrics.utilities.data' (d:\\Anaconda\\envs\\sc_multiomics\\lib\\site-packages\\torchmetrics\\utilities\\data.py)"
     ]
    }
   ],
   "source": [
    "import scanpy as sc\n",
    "import anndata as ad\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import muon as mu\n",
    "import anndata2ri\n",
    "import logging\n",
    "from torchmetrics.utilities.data import get_num_classes\n",
    "import scvi\n",
    "import os\n",
    "import scipy\n",
    "import scipy.io\n",
    "import scib\n",
    "import rpy2.rinterface_lib.callbacks\n",
    "\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from rpy2.robjects import r\n",
    "from rpy2.robjects import pandas2ri\n",
    "\n",
    "rpy2.rinterface_lib.callbacks.logger.setLevel(logging.ERROR)\n",
    "\n",
    "pandas2ri.activate()\n",
    "anndata2ri.activate()\n",
    "\n",
    "%load_ext rpy2.ipython\n",
    "\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "480f2531",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- Please select a CRAN mirror for use in this session ---\n",
      "Secure CRAN mirrors \n",
      "\n",
      " 1: 0-Cloud [https]\n",
      " 2: Australia (Canberra) [https]\n",
      " 3: Australia (Melbourne 1) [https]\n",
      " 4: Australia (Melbourne 2) [https]\n",
      " 5: Australia (Perth) [https]\n",
      " 6: Austria [https]\n",
      " 7: Belgium (Brussels) [https]\n",
      " 8: Brazil (PR) [https]\n",
      " 9: Brazil (RJ) [https]\n",
      "10: Brazil (SP 1) [https]\n",
      "11: Brazil (SP 2) [https]\n",
      "12: Bulgaria [https]\n",
      "13: Canada (MB) [https]\n",
      "14: Canada (ON) [https]\n",
      "15: Chile (Santiago) [https]\n",
      "16: China (Beijing 2) [https]\n",
      "17: China (Beijing 3) [https]\n",
      "18: China (Hefei) [https]\n",
      "19: China (Hong Kong) [https]\n",
      "20: China (Guangzhou) [https]\n",
      "21: China (Jinan) [https]\n",
      "22: China (Lanzhou) [https]\n",
      "23: China (Nanjing) [https]\n",
      "24: China (Shanghai 2) [https]\n",
      "25: China (Shenzhen) [https]\n",
      "26: Colombia (Cali) [https]\n",
      "27: Costa Rica [https]\n",
      "28: Cyprus [https]\n",
      "29: Czech Republic [https]\n",
      "30: Denmark [https]\n",
      "31: East Asia [https]\n",
      "32: Ecuador (Cuenca) [https]\n",
      "33: France (Lyon 1) [https]\n",
      "34: France (Lyon 2) [https]\n",
      "35: France (Marseille) [https]\n",
      "36: France (Paris 1) [https]\n",
      "37: Germany (Erlangen) [https]\n",
      "38: Germany (Leipzig) [https]\n",
      "39: Germany (Göttingen) [https]\n",
      "40: Germany (Münster) [https]\n",
      "41: Germany (Regensburg) [https]\n",
      "42: Greece [https]\n",
      "43: Iceland [https]\n",
      "44: India (Bengaluru) [https]\n",
      "45: India (Bhubaneswar) [https]\n",
      "46: Indonesia (Banda Aceh) [https]\n",
      "47: Iran (Mashhad) [https]\n",
      "48: Italy (Milano) [https]\n",
      "49: Italy (Padua) [https]\n",
      "50: Japan (Tokyo) [https]\n",
      "51: Japan (Yonezawa) [https]\n",
      "52: Korea (Gyeongsan-si) [https]\n",
      "53: Mexico (Mexico City) [https]\n",
      "54: Mexico (Texcoco) [https]\n",
      "55: Morocco [https]\n",
      "56: Netherlands (Dronten) [https]\n",
      "57: New Zealand [https]\n",
      "58: Norway [https]\n",
      "59: South Africa (Johannesburg) [https]\n",
      "60: Spain (A Coruña) [https]\n",
      "61: Spain (Madrid) [https]\n",
      "62: Sweden (Umeå) [https]\n",
      "63: Switzerland (Zurich 1) [https]\n",
      "64: Taiwan (Taipei) [https]\n",
      "65: Turkey (Denizli) [https]\n",
      "66: Turkey (Istanbul) [https]\n",
      "67: UK (Bristol) [https]\n",
      "68: UK (London 1) [https]\n",
      "69: USA (IA) [https]\n",
      "70: USA (MI) [https]\n",
      "71: USA (MO) [https]\n",
      "72: USA (OH) [https]\n",
      "73: USA (OR) [https]\n",
      "74: USA (TN) [https]\n",
      "75: United Arab Emirates [https]\n",
      "76: Uruguay [https]\n",
      "77: (other mirrors)\n",
      "\n",
      "Selection: 74\n",
      "\n",
      "The downloaded binary packages are in\n",
      "\t/var/folders/w_/jfhkzbcx6311zqlbnq7j9mm80000gn/T//RtmpZTGQa5/downloaded_packages\n",
      "\n",
      "    WARNING: The R package \"reticulate\" only fixed recently\n",
      "    an issue that caused a segfault when used with rpy2:\n",
      "    https://github.com/rstudio/reticulate/pull/1188\n",
      "    Make sure that you use a version of that package that includes\n",
      "    the fix.\n",
      "    "
     ]
    },
    {
     "data": {
      "text/plain": [
       "trying URL 'https://mirrors.nics.utk.edu/cran/bin/macosx/big-sur-arm64/contrib/4.3/Seurat_4.4.0.tgz'\n",
       "Content type 'application/x-gzip' length 3762575 bytes (3.6 MB)\n",
       "==================================================\n",
       "downloaded 3.6 MB\n",
       "\n",
       "In addition: Warning message:\n",
       "In doTryCatch(return(expr), name, parentenv, handler) :\n",
       "  unable to load shared object '/Library/Frameworks/R.framework/Resources/modules//R_X11.so':\n",
       "  dlopen(/Library/Frameworks/R.framework/Resources/modules//R_X11.so, 0x0006): Library not loaded: /opt/X11/lib/libSM.6.dylib\n",
       "  Referenced from: <B3716E5A-BF4D-3CA3-B8EB-89643DB72A04> /Library/Frameworks/R.framework/Versions/4.3-arm64/Resources/modules/R_X11.so\n",
       "  Reason: tried: '/opt/X11/lib/libSM.6.dylib' (no such file), '/System/Volumes/Preboot/Cryptexes/OS/opt/X11/lib/libSM.6.dylib' (no such file), '/opt/X11/lib/libSM.6.dylib' (no such file), '/usr/local/lib/libSM.6.dylib' (no such file), '/usr/lib/libSM.6.dylib' (no such file, not in dyld cache)\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%R\n",
    "install.packages('Seurat')\n",
    "suppressPackageStartupMessages({\n",
    "    library(Seurat)\n",
    "})\n",
    "set.seed(123)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b515ce6",
   "metadata": {},
   "source": [
    "## CITE-seq data\n",
    "\n",
    "We first show how to integrate a CITE-seq dataset using WNN, MOFA+ and totalVI. CITE-seq data contains raw gene expression counts and counts for surface proteins. The surface protein data is represented as antibody-derived tags (adt) here. We refer to the {ref}`surface-protein:motivation` section of the Surface Proteins chapter for more details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8d8986",
   "metadata": {},
   "source": [
    "### Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af8c3629",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre>MuData object with n_obs × n_vars = 10970 × 148344\n",
       "  var:\t&#x27;gene_ids&#x27;, &#x27;feature_types&#x27;\n",
       "  2 modalities\n",
       "    rna:\t10970 x 36601\n",
       "      var:\t&#x27;gene_ids&#x27;, &#x27;feature_types&#x27;\n",
       "    atac:\t10970 x 111743\n",
       "      var:\t&#x27;gene_ids&#x27;, &#x27;feature_types&#x27;</pre>"
      ],
      "text/plain": [
       "MuData object with n_obs × n_vars = 10970 × 148344\n",
       "  var:\t'gene_ids', 'feature_types'\n",
       "  2 modalities\n",
       "    rna:\t10970 x 36601\n",
       "      var:\t'gene_ids', 'feature_types'\n",
       "    atac:\t10970 x 111743\n",
       "      var:\t'gene_ids', 'feature_types'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# adt = sc.read(\n",
    "#     \"/lustre/groups/ml01/workspace/anastasia.litinetskaya/data/neurips-cite/adt_pp.h5ad\"\n",
    "# )\n",
    "adt = mu.read_10x_mtx(\n",
    "    'filtered_feature_bc_matrix/',  # the directory with the `.mtx` file\n",
    "    var_names='gene_symbols',                # use gene symbols for the variable names (variables-axis index)\n",
    "    cache=True)                              # write a cache file for faster subsequent reading\n",
    "adt.var_names_make_unique()  # this is unnecessary if using `var_names='gene_ids'` in `sc.read_10x_mtx`\n",
    "adt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b09c3c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# cells, # genes before filtering: (10970, 148344)\n"
     ]
    }
   ],
   "source": [
    "print(\"# cells, # genes before filtering:\", adata.shape)\n",
    "\n",
    "sc.pp.filter_genes(adata, min_counts=3)\n",
    "sc.pp.filter_cells(adata, min_counts=3)\n",
    "\n",
    "print(\"# cells, # genes after filtering:\", adata.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2a50fbf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rna = sc.read(\n",
    "    \"/lustre/groups/ml01/workspace/anastasia.litinetskaya/data/neurips-cite/rna_hvg.h5ad\"\n",
    ")\n",
    "rna"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "488145aa",
   "metadata": {},
   "source": [
    "We subset the data to Site 1 and the 3 corresponding donors to reduce the run time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32bd085f",
   "metadata": {},
   "outputs": [],
   "source": [
    "batches_to_keep = [\"s1d1\", \"s1d2\", \"s1d3\"]\n",
    "rna = rna[rna.obs[\"batch\"].isin(batches_to_keep)]\n",
    "adt = adt[adt.obs[\"donor\"].isin(batches_to_keep)]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c664670f",
   "metadata": {},
   "source": [
    "We only keep the cells that are present in both modality objects. First we need to make sure that `.obs_names` of both objects have similar structure and if not clean up a bit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6a35d010",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Index(['AAACAGCCAACAACAA-1', 'AAACAGCCACCGGCTA-1', 'AAACAGCCAGGACACA-1',\n",
       "       'AAACAGCCATCCTAGA-1', 'AAACATGCAAAGGTAC-1', 'AAACATGCAAATTCGT-1',\n",
       "       'AAACATGCAACCGCCA-1', 'AAACATGCACTTGTTC-1', 'AAACATGCAGAAATGC-1',\n",
       "       'AAACATGCAGGACCAA-1',\n",
       "       ...\n",
       "       'TTTGTGGCAGGAACTG-1', 'TTTGTGGCAGTTTCTC-1', 'TTTGTGTTCACATTGA-1',\n",
       "       'TTTGTGTTCGGTACGC-1', 'TTTGTGTTCTAATCCT-1', 'TTTGTGTTCTAGCGTG-1',\n",
       "       'TTTGTTGGTAAGGTTT-1', 'TTTGTTGGTTAGGATT-1', 'TTTGTTGGTTTGAGCA-1',\n",
       "       'TTTGTTGGTTTGGGCG-1'],\n",
       "      dtype='object', length=10970)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "adt.obs_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8fefa107",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'rna' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m rna\u001b[38;5;241m.\u001b[39mobs_names\n",
      "\u001b[0;31mNameError\u001b[0m: name 'rna' is not defined"
     ]
    }
   ],
   "source": [
    "rna.obs_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "985bce9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "adt.obs_names = [\n",
    "    name.split(\"-\")[0] + \"-\" + name.split(\"-\")[1] + \"-\" + batch\n",
    "    for batch, name in zip(adt.obs[\"donor\"], adt.obs_names)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e9bf7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "common_idx = list(set(rna.obs_names).intersection(set(adt.obs_names)))\n",
    "rna = rna[common_idx].copy()\n",
    "adt = adt[common_idx].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e6ee43d",
   "metadata": {},
   "source": [
    "We need to rename the proteins in the `adt` object so that the gene names and protein names do not intersect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82263143",
   "metadata": {},
   "outputs": [],
   "source": [
    "adt.var_names = [\"PROT_\" + name for name in adt.var_names]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2a1e10",
   "metadata": {},
   "source": [
    "Next we create a MuData object where we store data for both modalities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56eefa73",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "mdata = mu.MuData({\"rna\": rna, \"adt\": adt})\n",
    "mdata"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34f0afbb",
   "metadata": {},
   "source": [
    "We copy `batch` and `cell_type` column from one of the modality adatas to `.obs` of the mdata object to later use for visualizations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16a54300",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obs[\"batch\"] = rna.obs[\"batch\"].copy()\n",
    "mdata.obs[\"cell_type\"] = rna.obs[\"cell_type\"].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0321b28",
   "metadata": {},
   "source": [
    "### Weighted Nearest Neighbor (WNN)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42253b78",
   "metadata": {},
   "source": [
    "WNN is a graph-based method that takes neighbor graphs for each modality and constructs a common graph which is a weighted combination of the modality graphs. This constructed WNN graph can later be used together with gene expression matrix to obtain a supervised PCA (sPCA) representation guided by a WNN graph. The sPCA representation can be viewed as an embedding in a latent space.\n",
    "\n",
    "First, we use the `anndata2ri` package (https://github.com/theislab/anndata2ri) to move Python AnnData object to SingleCellExperiment and Seurat R objects. We create slimmer versions of AnnData objects that only contain the information that we need to the analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "546ae5c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_ = ad.AnnData(adt.X.copy())\n",
    "adata_.obs_names = adt.obs_names.copy()\n",
    "adata_.var_names = adt.var_names.copy()\n",
    "adata_.obs[\"batch\"] = adt.obs[\"donor\"].copy()\n",
    "adata_.obsm[\"harmony_pca\"] = adt.obsm[\"X_pcahm\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51285354",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i adata_\n",
    "# indicate that data is stored in .X of AnnData object\n",
    "adt = as.Seurat(adata_, data='X', counts=NULL)\n",
    "# the assay is called \"originalexp\" by default, we rename it to \"ADT\"\n",
    "adt <- RenameAssays(object = adt, originalexp = \"ADT\", verbose=FALSE) \n",
    "adt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a7d280",
   "metadata": {},
   "source": [
    "We repeat the same for RNA data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83168024",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_ = ad.AnnData(rna.X.copy())\n",
    "adata_.obs_names = rna.obs_names.copy()\n",
    "adata_.var_names = rna.var_names.copy()\n",
    "adata_.obs[\"cell_type\"] = rna.obs[\"cell_type\"].copy()\n",
    "adata_.obs[\"batch\"] = rna.obs[\"batch\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b687f27",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -i adata_\n",
    "rna = as.Seurat(adata_, data='X', counts=NULL)\n",
    "rna"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43dfbcbf",
   "metadata": {},
   "source": [
    "Next we create a Seurat object with both assays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "252dee79",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite <- rna\n",
    "cite[[\"ADT\"]] <- CreateAssayObject(data = adt@assays$ADT@data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0936e98",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite <- RenameAssays(object = cite, originalexp = \"RNA\", verbose=FALSE) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "063a5a8b",
   "metadata": {},
   "source": [
    "Since we have several batches in the dataset, we would need to perform batch correction before integrating the modalities. One option would be to batch correct using Seurat's `FindIntegrationAnchors()` and `IntegrateData()` (see https://satijalab.org/seurat/articles/integration_introduction.html) functions separately for each modality. We will use batch corrected embedding from previous analysis done separately for ADT and RNA data, namely Harmony corrected PCA embedding for ADT and scVI batch-corrected latent embedding for RNA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ec2137a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "# TODO need to change after we have the preprocessed data, for now RNA is not batch-corrected\n",
    "DefaultAssay(cite) <- \"RNA\"\n",
    "VariableFeatures(cite) <- rownames(cite)\n",
    "cite <- ScaleData(cite, verbose=FALSE)\n",
    "cite <- RunPCA(cite, verbose=FALSE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "028063ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite@reductions$harmony_pca <- adt@reductions$harmony_pca\n",
    "cite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2a0eabc",
   "metadata": {},
   "source": [
    "Now we follow the WNN vignette to perform the analysis. First, we need to find multimodal neighbors using the specified dimensionality reductions for each of the modalities. This function adds a WNN graph to the Seurat object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2203bf2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite <- FindMultiModalNeighbors(\n",
    "    cite, \n",
    "    reduction.list = list(\"pca\", \"harmony_pca\"), \n",
    "    dims.list = list(1:50, 1:30), \n",
    "    modality.weight.name = \"RNA.weight\",\n",
    "    verbose = FALSE\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53324913",
   "metadata": {},
   "source": [
    "Since we are also interested in finding an embedding for our multimodal data, we additionally run the `RunSPCA()` function that uses RNA gene expression data and the WNN graph for supervised PCA. Supervised PCA is a \"guided\" version of standard PCA run on gene expression data guided by the WNN graph to better preserve relationships between cells learn in WNN graph. We also will need a reference UMAP later for mapping an RNA query onto this multimodal reference as discussed in {ref}`multimodal-integration:advanced-integration`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d14e97f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite <- RunSPCA(cite, assay = \"RNA\", graph = \"wsnn\", npcs = 20)\n",
    "cite <- RunUMAP(cite, nn.name = \"weighted.nn\", reduction.name = \"wnn.umap\", reduction.key = \"wnnUMAP_\", return.model=TRUE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3136a270",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "cite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9726f439",
   "metadata": {},
   "source": [
    "We save the Seurat object as an `.rds` file for the {ref}`multimodal-integration:advanced-integration` section. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f66c27",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "saveRDS(cite, file = \"wnn_ref.rds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c92c6e30",
   "metadata": {},
   "source": [
    "We move the sPCA embedding to Python and store them in `.obsm` of the MuData object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b6161e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -o spca\n",
    "spca = Embeddings(object = cite[[\"spca\"]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcaffd68",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obsm[\"X_spca\"] = spca"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c367f2a",
   "metadata": {},
   "source": [
    "We also need to extract the calculated WNN graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4472b8bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R -o wnn\n",
    "wnn <- as.data.frame(summary(cite@graphs$wknn))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e717f809",
   "metadata": {},
   "source": [
    "The table indicates indices with connections between cells in the WNN graph. Since R starts indexing at 1 but Python at 0, we modify the indices to start with 0."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b662c31e",
   "metadata": {},
   "outputs": [],
   "source": [
    "wnn[:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2e58270",
   "metadata": {},
   "outputs": [],
   "source": [
    "wnn[\"i\"] = wnn[\"i\"] - 1\n",
    "wnn[\"j\"] = wnn[\"j\"] - 1\n",
    "wnn[:5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdde866a",
   "metadata": {},
   "source": [
    "We store the graph in `.obsp` of the MuData object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba5e5543",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obsp[\"wnn_connectivities\"] = scipy.sparse.coo_matrix(\n",
    "    (wnn[\"x\"], (wnn[\"i\"], wnn[\"j\"]))\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aef05a42",
   "metadata": {},
   "source": [
    "Next we use the WWN graph to calculate the UMAP coordinates and save them in `.obsm['X_umap_wnn']`. We could alternatively also just use sPCA coordinates for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6ea3abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# we won't actually need the neighbors\n",
    "# but need to run this anyway as a little trick to make scanpy work with externally-computed neighbors\n",
    "sc.pp.neighbors(mdata, use_rep=\"X_spca\")\n",
    "mdata.obsp[\"connectivities\"] = mdata.obsp[\"wnn_connectivities\"].copy()\n",
    "# delete distances to make sure we are not using anything calculated with sc.pp.neighbors()\n",
    "del mdata.obsp[\"distances\"]\n",
    "sc.tl.umap(mdata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "402f145c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obsm[\"X_umap_wnn\"] = mdata.obsm[\"X_umap\"].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "93c4c4be",
   "metadata": {},
   "source": [
    "Finally we visualize the cell types and batches on a UMAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9303a2d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu.pl.embedding(\n",
    "    mdata, color=[\"cell_type\", \"batch\"], ncols=1, basis=\"umap_wnn\", frameon=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba8e33e8",
   "metadata": {},
   "source": [
    "To be able to quantitatively assess the result of the integration and compare to other methods we compute some of the scIB metrics using the sPCA embedding and WNN graph. More specifically, we calculate the following metrics:\n",
    "- bio conservation: `NMI_cluster/label`, `ARI_cluster/label`, `ASW_label` and `isolated_label_silhouette`;\n",
    "- batch correction: `ASW_label/batch`, `graph_conn`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2cbfa52",
   "metadata": {},
   "outputs": [],
   "source": [
    "scib_anndata = sc.AnnData(mdata.obsm[\"X_spca\"]).copy()\n",
    "scib_anndata.obs = mdata.obs.copy()\n",
    "scib_anndata.obsp[\"connectivities\"] = mdata.obsp[\"connectivities\"].copy()\n",
    "scib_anndata.obsm[\"X_spca\"] = mdata.obsm[\"X_spca\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09217236",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "metrics_wnn = scib.metrics.metrics(\n",
    "    scib_anndata,\n",
    "    scib_anndata,\n",
    "    batch_key=\"batch\",\n",
    "    label_key=\"cell_type\",\n",
    "    embed=\"X_spca\",\n",
    "    ari_=True,\n",
    "    nmi_=True,\n",
    "    silhouette_=True,\n",
    "    graph_conn_=True,\n",
    "    isolated_labels_asw_=True,\n",
    ")\n",
    "metrics_wnn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0971c14b",
   "metadata": {},
   "source": [
    "We note that even though batch correction was performed using Harmony for ADT and scVI for RNA, we still include metrics that assess batch correction here too."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52cd6136",
   "metadata": {},
   "source": [
    "### Multi-Omics Factor Analysis (MOFA+)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6529e9f",
   "metadata": {},
   "source": [
    "MOFA+ is a linear factor model that decomposes the input matrices into the product of low-rank matrices. The low-rank representation can be used as an embedding in a low-dimensional space for visualization and other downstream tasks. The latent dimensions are interpretable with respect to the original input features and represent the leading sources of variation in the data.\n",
    "\n",
    "By default, we are using data from `.X` and the data should be normalized. Since there are some batch effects in the data that MOFA+ can correct for, we also pass the `groups_label` parameter to specify the batch covariate.\n",
    "\n",
    "If you want to run MOFA+ on a GPU, you need to additionally install a version of cuPY (https://cupy.dev) which is compatible with your CUDA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56b5728b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu.tl.mofa(mdata, groups_label=\"batch\", gpu_mode=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "917dd29f",
   "metadata": {},
   "source": [
    "We use the `X_mofa` representation to calculate the neighbors and the UMAP coordinates, and store them in `.obsm['X_umap_mofa']`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b605233",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(mdata, use_rep=\"X_mofa\")\n",
    "sc.tl.umap(mdata)\n",
    "mdata.obsm[\"X_umap_mofa\"] = mdata.obsm[\"X_umap\"].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2bb419",
   "metadata": {},
   "source": [
    "We plot the cell types and batches again on the resulting UMAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40e8f2f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu.pl.embedding(\n",
    "    mdata, color=[\"cell_type\", \"batch\"], ncols=1, basis=\"umap_mofa\", frameon=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ad18992",
   "metadata": {},
   "source": [
    "Finally, we calculate the same scIB metrics as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcb560a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "scib_anndata = sc.AnnData(mdata.obsm[\"X_mofa\"]).copy()\n",
    "scib_anndata.obs = mdata.obs.copy()\n",
    "scib_anndata.obsp[\"connectivities\"] = mdata.obsp[\"connectivities\"].copy()\n",
    "scib_anndata.obsm[\"X_mofa\"] = mdata.obsm[\"X_mofa\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d381d1c9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "metrics_mofa = scib.metrics.metrics(\n",
    "    scib_anndata,\n",
    "    scib_anndata,\n",
    "    batch_key=\"batch\",\n",
    "    label_key=\"cell_type\",\n",
    "    embed=\"X_mofa\",\n",
    "    ari_=True,\n",
    "    nmi_=True,\n",
    "    silhouette_=True,\n",
    "    graph_conn_=True,\n",
    "    isolated_labels_asw_=True,\n",
    ")\n",
    "metrics_mofa"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2203bc54",
   "metadata": {},
   "source": [
    "### Total Variational Inference (totalVI)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99214bc7",
   "metadata": {},
   "source": [
    "TotalVI is variational-inference-based methods for joint analysis of paired gene expression and protein abundance measurements. It takes into account batch effects, protein background noise, which allows the model to learn a join latent representation disentangled form technical factors. TotalVI models transcriptome counts with negative-binomial (NB) distribution and the protein counts as NB mixture of foreground and background signal. Hence, the model takes raw gene expression and raw protein counts as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "764dc473",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata = mdata[\"rna\"].copy()\n",
    "adata.obsm[\"protein_expression\"] = mdata[\"adt\"].layers[\"counts\"].A.copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35ebbb19",
   "metadata": {},
   "source": [
    "We need to specify that raw counts for RNA are stored in `counts` layer of our adata and that we want to correct for batch effect with `batch_key=\"batch\"` parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "daaf51fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "scvi.model.TOTALVI.setup_anndata(\n",
    "    adata,\n",
    "    protein_expression_obsm_key=\"protein_expression\",\n",
    "    layer=\"counts\",\n",
    "    batch_key=\"batch\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8209b161",
   "metadata": {},
   "source": [
    "We initialize the totalVI model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bc584c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae = scvi.model.TOTALVI(adata)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16997023",
   "metadata": {},
   "source": [
    "Next, we train the model with default parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c22d4a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "vae.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c5d28b",
   "metadata": {},
   "source": [
    "Next, we obtain the latent representation and store it in `.obsm['X_totalVI']` and then use it to calculate the UMAP coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3902a59b",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obsm[\"X_totalVI\"] = vae.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b791838",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(mdata, use_rep=\"X_totalVI\")\n",
    "sc.tl.umap(mdata)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd8eb3fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata.obsm[\"X_umap_totalVI\"] = mdata.obsm[\"X_umap\"].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee3bfba3",
   "metadata": {},
   "source": [
    "As above, we plot cell types and batches on a UMAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "718932e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu.pl.embedding(\n",
    "    mdata, color=[\"cell_type\", \"batch\"], ncols=1, basis=\"umap_totalVI\", frameon=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fa67339",
   "metadata": {},
   "source": [
    "And finally, we calculate scIB metrics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db46eb96",
   "metadata": {},
   "outputs": [],
   "source": [
    "scib_anndata = sc.AnnData(mdata.obsm[\"X_totalVI\"]).copy()\n",
    "scib_anndata.obs = mdata.obs.copy()\n",
    "scib_anndata.obsp[\"connectivities\"] = mdata.obsp[\"connectivities\"].copy()\n",
    "scib_anndata.obsm[\"X_totalVI\"] = mdata.obsm[\"X_totalVI\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18dff347",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "metrics_totalvi = scib.metrics.metrics(\n",
    "    scib_anndata,\n",
    "    scib_anndata,\n",
    "    batch_key=\"batch\",\n",
    "    label_key=\"cell_type\",\n",
    "    embed=\"X_totalVI\",\n",
    "    ari_=True,\n",
    "    nmi_=True,\n",
    "    silhouette_=True,\n",
    "    graph_conn_=True,\n",
    "    isolated_labels_asw_=True,\n",
    ")\n",
    "metrics_totalvi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "794f925a",
   "metadata": {},
   "source": [
    "### scIB metrics evaluation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b631c37",
   "metadata": {},
   "source": [
    "To better see the differences in models' performances, we visualize the scIB output for each of the methods. We need to merge the output DataFrames into one and additionally calculate the overall score for each method. We follow the scIB publication and calculate the overall score as `0.4 * batch_correction_metrics + 0.6 * bio_conservation_metrics`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea6d271d",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = pd.DataFrame([metrics_wnn[0], metrics_mofa[0], metrics_totalvi[0]])\n",
    "metrics = metrics.set_index(pd.Index([\"WNN\", \"MOFA+\", \"totalVI\"]))\n",
    "metrics = metrics.dropna(axis=1)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d2021c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics[\"overall\"] = (\n",
    "    0.4 * (metrics[\"ASW_label/batch\"] + metrics[\"graph_conn\"]) / 2\n",
    "    + 0.6\n",
    "    * (\n",
    "        metrics[\"NMI_cluster/label\"]\n",
    "        + metrics[\"ARI_cluster/label\"]\n",
    "        + metrics[\"ASW_label\"]\n",
    "        + metrics[\"isolated_label_silhouette\"]\n",
    "    )\n",
    "    / 4\n",
    ")\n",
    "metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86e962d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.scatterplot(data=metrics)\n",
    "plt.legend(bbox_to_anchor=(1.05, 1), loc=\"upper left\", borderaxespad=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d988f2c5",
   "metadata": {},
   "source": [
    "We observe that totalVI obtained the highest overall score, and therefore we will use totalVI embedding later in the notebook to show how one can annotate the cluster in the latent space using both ADT and RNA markers. Depending on the downstream task and the experimental design, selecting the best performing based on a specific metric is advisable."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9ef7912",
   "metadata": {},
   "source": [
    "## Multiome data\n",
    "To show that integration methods can also work with multiome (i.e. paired RNA-seq and ATAC-seq) data, we demonstrate how multiVI can be used for this task. We note that WNN and MOFA+ can also be run on multiome data with almost exactly the same code as above, so here we only present multiVI where the underlying model differs from totalVI."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86cff6fa",
   "metadata": {},
   "source": [
    "### Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae6304c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "atac = sc.read(\n",
    "    \"/lustre/groups/ml01/workspace/anastasia.litinetskaya/data/neurips-multiome/atac_hvf.h5ad\"\n",
    ")\n",
    "atac"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62bd7ce8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "rna = sc.read(\n",
    "    \"/lustre/groups/ml01/workspace/anastasia.litinetskaya/data/neurips-multiome/rna_hvg.h5ad\"\n",
    ")\n",
    "rna"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a903415",
   "metadata": {},
   "source": [
    "We again subset the data to one site and 3 batches."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baa747f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "batches_to_keep = [\"s1d1\", \"s1d2\", \"s1d3\"]\n",
    "rna = rna[rna.obs[\"batch\"].isin(batches_to_keep)]\n",
    "atac = atac[atac.obs[\"batch\"].isin(batches_to_keep)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddb00d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata_multiome = mu.MuData({\"rna\": rna, \"atac\": atac})\n",
    "mdata_multiome"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94075f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata_multiome.obs[\"batch\"] = mdata_multiome[\"rna\"].obs[\"batch\"].copy()\n",
    "mdata_multiome.obs[\"cell_type\"] = mdata_multiome[\"rna\"].obs[\"cell_type\"].copy()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9d0e3e8",
   "metadata": {},
   "source": [
    "### MultiVI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c582fba",
   "metadata": {},
   "source": [
    "MultiVI is also based on variational inference and conditional variational autoencoders. The gene expression counts are modeled exactly the same way as in totalVI, i.e. using raw counts and NB distribution. Chromatin accessibility on the other hand is modeled using Bernouli distribution modeling how likely a particular region is to be open. Hence, the input data for ATAC assay has to be binary where 0 means a closed region and 1 means an open region."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "241cd43d",
   "metadata": {},
   "outputs": [],
   "source": [
    "adt\n",
    "n_genes = len(adt[\"rna\"].var_names)\n",
    "n_regions = len(adt[\"atac\"].var_names)\n",
    "# n_genes = len(rna.var_names)\n",
    "# n_regions = len(atac.var_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a996faf",
   "metadata": {},
   "source": [
    "MultiVI requires one AnnData object with concatenated genes and peaks as features. Since we start off with two different objects for each modality but have paired measurements, we can use the following trick to concatenate the AnnData objects along the feature axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "dfd1f869",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AnnData object with n_obs × n_vars = 10970 × 36601\n",
       "    obs: 'modality'\n",
       "    var: 'gene_ids', 'feature_types'"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# adata_paired = ad.concat([rna.copy().T, atac.copy().T]).T\n",
    "# adata_paired.obs = adata_paired.obs.join(rna.obs[[\"cell_type\", \"batch\"]])\n",
    "# adata_paired.obs[\"modality\"] = \"paired\"\n",
    "# adata_paired\n",
    "adata_paired = sc.read_10x_mtx(\n",
    "    'filtered_feature_bc_matrix/',  # the directory with the `.mtx` file\n",
    "    var_names='gene_symbols',                # use gene symbols for the variable names (variables-axis index)\n",
    "    cache=True) \n",
    "adata_paired.obs[\"modality\"] = \"paired\"\n",
    "adata_paired"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "373df981",
   "metadata": {},
   "outputs": [],
   "source": [
    "adata_mvi = scvi.data.organize_multiome_anndatas(adata_paired)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c979b93",
   "metadata": {},
   "source": [
    "We also make sure that we pass raw counts as input to the model by specifiying `layer='counts'` in the `setup_anndata` funciton."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "4bcdd263",
   "metadata": {},
   "outputs": [],
   "source": [
    "scvi.model.MULTIVI.setup_anndata(\n",
    "    adata_mvi,\n",
    "    batch_key=\"modality\",\n",
    "#     categorical_covariate_keys=[\"batch\"],\n",
    "#     layer=\"counts\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3494d15f",
   "metadata": {},
   "source": [
    "We initialize the MultiVI model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "602eab67",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">MultiVI Model with INPUTS: n_genes:<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">36601</span>, n_regions:<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">111743</span>, n_proteins:<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span>\n",
       "n_hidden: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">334</span>, n_latent: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">18</span>, n_layers_encoder: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span>, n_layers_decoder: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2</span> , dropout_rate: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0.1</span>, latent_distribution: \n",
       "normal, deep injection: <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>, gene_likelihood: zinb, gene_dispersion:gene, Mod.Weights:equal, \n",
       "Mod.Penalty:Jeffreys, protein_dispersion:protein\n",
       "Training status: Not Trained\n",
       "Model's adata is minified?: <span style=\"color: #ff0000; text-decoration-color: #ff0000; font-style: italic\">False</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "MultiVI Model with INPUTS: n_genes:\u001b[1;36m36601\u001b[0m, n_regions:\u001b[1;36m111743\u001b[0m, n_proteins:\u001b[1;36m0\u001b[0m\n",
       "n_hidden: \u001b[1;36m334\u001b[0m, n_latent: \u001b[1;36m18\u001b[0m, n_layers_encoder: \u001b[1;36m2\u001b[0m, n_layers_decoder: \u001b[1;36m2\u001b[0m , dropout_rate: \u001b[1;36m0.1\u001b[0m, latent_distribution: \n",
       "normal, deep injection: \u001b[3;91mFalse\u001b[0m, gene_likelihood: zinb, gene_dispersion:gene, Mod.Weights:equal, \n",
       "Mod.Penalty:Jeffreys, protein_dispersion:protein\n",
       "Training status: Not Trained\n",
       "Model's adata is minified?: \u001b[3;91mFalse\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mvi = scvi.model.MULTIVI(\n",
    "    adata_mvi,\n",
    "    n_genes=n_genes,\n",
    "    n_regions=n_regions,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f630453",
   "metadata": {},
   "source": [
    "Next, we train the model with the default parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "cde67e69",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "`devices` selected with `CPUAccelerator` should be an int > 0.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[21], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m mvi\u001b[38;5;241m.\u001b[39mtrain()\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/scvi/model/_multivi.py:339\u001b[0m, in \u001b[0;36mMULTIVI.train\u001b[0;34m(self, max_epochs, lr, use_gpu, accelerator, devices, train_size, validation_size, shuffle_set_split, batch_size, weight_decay, eps, early_stopping, save_best, check_val_every_n_epoch, n_steps_kl_warmup, n_epochs_kl_warmup, adversarial_mixing, plan_kwargs, **kwargs)\u001b[0m\n\u001b[1;32m    331\u001b[0m data_splitter \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_data_splitter_cls(\n\u001b[1;32m    332\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39madata_manager,\n\u001b[1;32m    333\u001b[0m     train_size\u001b[38;5;241m=\u001b[39mtrain_size,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    336\u001b[0m     batch_size\u001b[38;5;241m=\u001b[39mbatch_size,\n\u001b[1;32m    337\u001b[0m )\n\u001b[1;32m    338\u001b[0m training_plan \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_training_plan_cls(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmodule, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mplan_kwargs)\n\u001b[0;32m--> 339\u001b[0m runner \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_train_runner_cls(\n\u001b[1;32m    340\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    341\u001b[0m     training_plan\u001b[38;5;241m=\u001b[39mtraining_plan,\n\u001b[1;32m    342\u001b[0m     data_splitter\u001b[38;5;241m=\u001b[39mdata_splitter,\n\u001b[1;32m    343\u001b[0m     max_epochs\u001b[38;5;241m=\u001b[39mmax_epochs,\n\u001b[1;32m    344\u001b[0m     use_gpu\u001b[38;5;241m=\u001b[39muse_gpu,\n\u001b[1;32m    345\u001b[0m     accelerator\u001b[38;5;241m=\u001b[39maccelerator,\n\u001b[1;32m    346\u001b[0m     devices\u001b[38;5;241m=\u001b[39mdevices,\n\u001b[1;32m    347\u001b[0m     early_stopping\u001b[38;5;241m=\u001b[39mearly_stopping,\n\u001b[1;32m    348\u001b[0m     check_val_every_n_epoch\u001b[38;5;241m=\u001b[39mcheck_val_every_n_epoch,\n\u001b[1;32m    349\u001b[0m     early_stopping_monitor\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mreconstruction_loss_validation\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m    350\u001b[0m     early_stopping_patience\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m50\u001b[39m,\n\u001b[1;32m    351\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs,\n\u001b[1;32m    352\u001b[0m )\n\u001b[1;32m    353\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m runner()\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/scvi/train/_trainrunner.py:85\u001b[0m, in \u001b[0;36mTrainRunner.__init__\u001b[0;34m(self, model, training_plan, data_splitter, max_epochs, use_gpu, accelerator, devices, **trainer_kwargs)\u001b[0m\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlightning_devices \u001b[38;5;241m=\u001b[39m lightning_devices\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdevice \u001b[38;5;241m=\u001b[39m device\n\u001b[0;32m---> 85\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtrainer \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_trainer_cls(\n\u001b[1;32m     86\u001b[0m     max_epochs\u001b[38;5;241m=\u001b[39mmax_epochs,\n\u001b[1;32m     87\u001b[0m     accelerator\u001b[38;5;241m=\u001b[39maccelerator,\n\u001b[1;32m     88\u001b[0m     devices\u001b[38;5;241m=\u001b[39mlightning_devices,\n\u001b[1;32m     89\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mtrainer_kwargs,\n\u001b[1;32m     90\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/scvi/train/_trainer.py:139\u001b[0m, in \u001b[0;36mTrainer.__init__\u001b[0;34m(self, accelerator, devices, benchmark, check_val_every_n_epoch, max_epochs, default_root_dir, enable_checkpointing, num_sanity_val_steps, enable_model_summary, early_stopping, early_stopping_monitor, early_stopping_min_delta, early_stopping_patience, early_stopping_mode, enable_progress_bar, progress_bar_refresh_rate, simple_progress_bar, logger, log_every_n_steps, **kwargs)\u001b[0m\n\u001b[1;32m    136\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m logger \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    137\u001b[0m     logger \u001b[38;5;241m=\u001b[39m SimpleLogger()\n\u001b[0;32m--> 139\u001b[0m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39m\u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m    140\u001b[0m     accelerator\u001b[38;5;241m=\u001b[39maccelerator,\n\u001b[1;32m    141\u001b[0m     devices\u001b[38;5;241m=\u001b[39mdevices,\n\u001b[1;32m    142\u001b[0m     benchmark\u001b[38;5;241m=\u001b[39mbenchmark,\n\u001b[1;32m    143\u001b[0m     check_val_every_n_epoch\u001b[38;5;241m=\u001b[39mcheck_val_every_n_epoch,\n\u001b[1;32m    144\u001b[0m     max_epochs\u001b[38;5;241m=\u001b[39mmax_epochs,\n\u001b[1;32m    145\u001b[0m     default_root_dir\u001b[38;5;241m=\u001b[39mdefault_root_dir,\n\u001b[1;32m    146\u001b[0m     enable_checkpointing\u001b[38;5;241m=\u001b[39menable_checkpointing,\n\u001b[1;32m    147\u001b[0m     num_sanity_val_steps\u001b[38;5;241m=\u001b[39mnum_sanity_val_steps,\n\u001b[1;32m    148\u001b[0m     enable_model_summary\u001b[38;5;241m=\u001b[39menable_model_summary,\n\u001b[1;32m    149\u001b[0m     logger\u001b[38;5;241m=\u001b[39mlogger,\n\u001b[1;32m    150\u001b[0m     log_every_n_steps\u001b[38;5;241m=\u001b[39mlog_every_n_steps,\n\u001b[1;32m    151\u001b[0m     enable_progress_bar\u001b[38;5;241m=\u001b[39menable_progress_bar,\n\u001b[1;32m    152\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs,\n\u001b[1;32m    153\u001b[0m )\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/pytorch/utilities/argparse.py:70\u001b[0m, in \u001b[0;36m_defaults_from_env_vars.<locals>.insert_env_defaults\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m     67\u001b[0m kwargs \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mdict\u001b[39m(\u001b[38;5;28mlist\u001b[39m(env_variables\u001b[38;5;241m.\u001b[39mitems()) \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mlist\u001b[39m(kwargs\u001b[38;5;241m.\u001b[39mitems()))\n\u001b[1;32m     69\u001b[0m \u001b[38;5;66;03m# all args were already moved to kwargs\u001b[39;00m\n\u001b[0;32m---> 70\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m fn(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/pytorch/trainer/trainer.py:399\u001b[0m, in \u001b[0;36mTrainer.__init__\u001b[0;34m(self, accelerator, strategy, devices, num_nodes, precision, logger, callbacks, fast_dev_run, max_epochs, min_epochs, max_steps, min_steps, max_time, limit_train_batches, limit_val_batches, limit_test_batches, limit_predict_batches, overfit_batches, val_check_interval, check_val_every_n_epoch, num_sanity_val_steps, log_every_n_steps, enable_checkpointing, enable_progress_bar, enable_model_summary, accumulate_grad_batches, gradient_clip_val, gradient_clip_algorithm, deterministic, benchmark, inference_mode, use_distributed_sampler, profiler, detect_anomaly, barebones, plugins, sync_batchnorm, reload_dataloaders_every_n_epochs, default_root_dir)\u001b[0m\n\u001b[1;32m    396\u001b[0m \u001b[38;5;66;03m# init connectors\u001b[39;00m\n\u001b[1;32m    397\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_data_connector \u001b[38;5;241m=\u001b[39m _DataConnector(\u001b[38;5;28mself\u001b[39m)\n\u001b[0;32m--> 399\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_accelerator_connector \u001b[38;5;241m=\u001b[39m _AcceleratorConnector(\n\u001b[1;32m    400\u001b[0m     devices\u001b[38;5;241m=\u001b[39mdevices,\n\u001b[1;32m    401\u001b[0m     accelerator\u001b[38;5;241m=\u001b[39maccelerator,\n\u001b[1;32m    402\u001b[0m     strategy\u001b[38;5;241m=\u001b[39mstrategy,\n\u001b[1;32m    403\u001b[0m     num_nodes\u001b[38;5;241m=\u001b[39mnum_nodes,\n\u001b[1;32m    404\u001b[0m     sync_batchnorm\u001b[38;5;241m=\u001b[39msync_batchnorm,\n\u001b[1;32m    405\u001b[0m     benchmark\u001b[38;5;241m=\u001b[39mbenchmark,\n\u001b[1;32m    406\u001b[0m     use_distributed_sampler\u001b[38;5;241m=\u001b[39muse_distributed_sampler,\n\u001b[1;32m    407\u001b[0m     deterministic\u001b[38;5;241m=\u001b[39mdeterministic,\n\u001b[1;32m    408\u001b[0m     precision\u001b[38;5;241m=\u001b[39mprecision,\n\u001b[1;32m    409\u001b[0m     plugins\u001b[38;5;241m=\u001b[39mplugins,\n\u001b[1;32m    410\u001b[0m )\n\u001b[1;32m    411\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_logger_connector \u001b[38;5;241m=\u001b[39m _LoggerConnector(\u001b[38;5;28mself\u001b[39m)\n\u001b[1;32m    412\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback_connector \u001b[38;5;241m=\u001b[39m _CallbackConnector(\u001b[38;5;28mself\u001b[39m)\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py:157\u001b[0m, in \u001b[0;36m_AcceleratorConnector.__init__\u001b[0;34m(self, devices, num_nodes, accelerator, strategy, plugins, precision, sync_batchnorm, benchmark, use_distributed_sampler, deterministic)\u001b[0m\n\u001b[1;32m    154\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_accelerator_flag \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_choose_gpu_accelerator_backend()\n\u001b[1;32m    156\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_check_device_config_and_set_final_flags(devices\u001b[38;5;241m=\u001b[39mdevices, num_nodes\u001b[38;5;241m=\u001b[39mnum_nodes)\n\u001b[0;32m--> 157\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_set_parallel_devices_and_init_accelerator()\n\u001b[1;32m    159\u001b[0m \u001b[38;5;66;03m# 3. Instantiate ClusterEnvironment\u001b[39;00m\n\u001b[1;32m    160\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcluster_environment: ClusterEnvironment \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_choose_and_init_cluster_environment()\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/pytorch/trainer/connectors/accelerator_connector.py:390\u001b[0m, in \u001b[0;36m_AcceleratorConnector._set_parallel_devices_and_init_accelerator\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    382\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m MisconfigurationException(\n\u001b[1;32m    383\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`\u001b[39m\u001b[38;5;132;01m{\u001b[39;00maccelerator_cls\u001b[38;5;241m.\u001b[39m\u001b[38;5;18m__qualname__\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m` can not run on your system\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    384\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m since the accelerator is not available. The following accelerator(s)\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    385\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m is available and can be passed into `accelerator` argument of\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    386\u001b[0m         \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m `Trainer`: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mavailable_accelerator\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    387\u001b[0m     )\n\u001b[1;32m    389\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_set_devices_flag_if_auto_passed()\n\u001b[0;32m--> 390\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_devices_flag \u001b[38;5;241m=\u001b[39m accelerator_cls\u001b[38;5;241m.\u001b[39mparse_devices(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_devices_flag)\n\u001b[1;32m    391\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parallel_devices:\n\u001b[1;32m    392\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_parallel_devices \u001b[38;5;241m=\u001b[39m accelerator_cls\u001b[38;5;241m.\u001b[39mget_parallel_devices(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_devices_flag)\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/pytorch/accelerators/cpu.py:48\u001b[0m, in \u001b[0;36mCPUAccelerator.parse_devices\u001b[0;34m(devices)\u001b[0m\n\u001b[1;32m     45\u001b[0m \u001b[38;5;129m@staticmethod\u001b[39m\n\u001b[1;32m     46\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mparse_devices\u001b[39m(devices: Union[\u001b[38;5;28mint\u001b[39m, \u001b[38;5;28mstr\u001b[39m, List[\u001b[38;5;28mint\u001b[39m]]) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28mint\u001b[39m:\n\u001b[1;32m     47\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Accelerator device parsing logic.\"\"\"\u001b[39;00m\n\u001b[0;32m---> 48\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m _parse_cpu_cores(devices)\n",
      "File \u001b[0;32m~/anaconda3/lib/python3.11/site-packages/lightning/fabric/accelerators/cpu.py:85\u001b[0m, in \u001b[0;36m_parse_cpu_cores\u001b[0;34m(cpu_cores)\u001b[0m\n\u001b[1;32m     82\u001b[0m     cpu_cores \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mint\u001b[39m(cpu_cores)\n\u001b[1;32m     84\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(cpu_cores, \u001b[38;5;28mint\u001b[39m) \u001b[38;5;129;01mor\u001b[39;00m cpu_cores \u001b[38;5;241m<\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m---> 85\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m`devices` selected with `CPUAccelerator` should be an int > 0.\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     87\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m cpu_cores\n",
      "\u001b[0;31mTypeError\u001b[0m: `devices` selected with `CPUAccelerator` should be an int > 0."
     ]
    }
   ],
   "source": [
    "mvi.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea8ae04",
   "metadata": {},
   "source": [
    "Finally, we visualize the latent embedding on the UMAP."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be6ad8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata_multiome.obsm[\"X_multiVI\"] = mvi.get_latent_representation()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7842d19c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.pp.neighbors(mdata_multiome, use_rep=\"X_multiVI\")\n",
    "sc.tl.umap(mdata_multiome)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6671d3b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "mdata_multiome.obsm[\"X_umap_multiVI\"] = mdata_multiome.obsm[\"X_umap\"].copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91842bf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu.pl.embedding(\n",
    "    mdata_multiome,\n",
    "    color=[\"cell_type\", \"batch\"],\n",
    "    ncols=1,\n",
    "    basis=\"umap_multiVI\",\n",
    "    frameon=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9c2b194",
   "metadata": {},
   "source": [
    "Session info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "506ecfdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%R\n",
    "sessionInfo()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83da53ba",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb6d5b2c",
   "metadata": {},
   "source": [
    "```{bibliography}\n",
    ":filter: docname in docnames\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e566ddd8",
   "metadata": {},
   "source": [
    "## Contributors\n",
    "\n",
    "We gratefully acknowledge the contributions of:\n",
    "\n",
    "### Authors\n",
    "\n",
    "* Anastasia Litinetskaya\n",
    "\n",
    "### Reviewers\n",
    "\n",
    "* Lukas Heumos"
   ]
  }
 ],
 "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.18"
  },
  "vscode": {
   "interpreter": {
    "hash": "e7370f93d1d0cde622a1f8e1c04877d8463912d04d973331ad4851f04de6915a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}