{ "cells": [ { "cell_type": "markdown", "id": "neural-welsh", "metadata": { "tags": [] }, "source": [ "# Interfacing PyTorch models with collections of AnnData objects via `annbatch`" ] }, { "cell_type": "markdown", "id": "480364bc-813f-4834-a5b4-d15ef128171e", "metadata": {}, "source": [ "**Author:** [Sergei Rybakov](https://github.com/koncopd)\n", "**Author:** [Ilan Gold](https://github.com/ilan-gold)" ] }, { "cell_type": "markdown", "id": "given-belarus", "metadata": {}, "source": [ "This tutorial demonstrates how to set up a {class}`annbatch.Loader` for an in-memory {class}`~anndata.AnnData` object to interface with a PyTorch model and lay out complex covariates." ] }, { "cell_type": "markdown", "id": "annbatch-intro", "metadata": {}, "source": [ "Note, however, that when training models on many large `AnnData` objects, loading them all into memory at once is impractical as this notebook does.\n", "\n", "`annbatch.Loader` solves this by allowing users to streams batches formed of contiguous on-memory runs of the training data across multiple AnnData objects without requiring a full in-memory concatenation — compatible with PyTorch, JAX, and other frameworks while being extremely fast.\n", "Preshuffling is a crucial part of making this “contiguous fetching” strategy work – see {class}`annbatch.DatasetCollection` for more information on performing the on-disk, lazy preshuffle." ] }, { "cell_type": "markdown", "id": "60ed6b48-175b-4f67-94d5-60b3376d9c4b", "metadata": {}, "source": [ "To install `annbatch` for in-memory sparse data (which is most common in single-cell RNA-seq and will be used in this tutorial), use `pip install \"annbatch[numba]\"` – note that the `numba` extra is required for in-memory sparse data.\n", "To use the `preload_to_gpu` option in {class}`annbatch.Loader`, install the `cupy-cudaxx` extra as well where `xx` is your `cuda` version.\n", "See the [docs](https://annbatch.readthedocs.io/en/stable/installation.html) for more info on installation." ] }, { "cell_type": "code", "execution_count": null, "id": "possible-possession", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "from typing import TYPE_CHECKING\n", "\n", "import numpy as np\n", "import pandas as pd\n", "import pooch\n", "import pyro\n", "import pyro.distributions as dist\n", "import scanpy as sc\n", "import torch\n", "from annbatch import Loader\n", "from sklearn.preprocessing import OneHotEncoder\n", "from torch import nn\n", "\n", "import anndata as ad\n", "\n", "if TYPE_CHECKING:\n", " from collections.abc import Sequence\n", "\n", "\n", "pyro.clear_param_store()" ] }, { "cell_type": "markdown", "id": "inner-piano", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "source": [ "## VAE model definition" ] }, { "cell_type": "markdown", "id": "semisup-motivation", "metadata": {}, "source": [ "The task is semi-supervised label transfer: `pbmc` cells have annotated cell types, while `covid` cells do not.\n", "The model learns a shared latent space across both datasets and uses the labeled pbmc cells to predict cell types for the unlabeled covid cells.\n", "Batch labels (study of origin) are included as covariates to correct for dataset-specific effects." ] }, { "cell_type": "code", "execution_count": 2, "id": "incorporated-cargo", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "class MLP(nn.Module):\n", " def __init__(self, input_dim, hidden_dims, out_dim):\n", " super().__init__()\n", "\n", " modules = []\n", " for in_size, out_size in zip(\n", " [input_dim, *hidden_dims], hidden_dims, strict=True\n", " ):\n", " modules.append(nn.Linear(in_size, out_size))\n", " modules.append(nn.LayerNorm(out_size))\n", " modules.append(nn.ReLU())\n", " modules.append(nn.Dropout(p=0.05))\n", " modules.append(nn.Linear(hidden_dims[-1], out_dim))\n", " self.fc = nn.Sequential(*modules)\n", "\n", " def forward(self, *inputs):\n", " shape = (*dist.util.broadcast_shape(*[s.shape[:-1] for s in inputs]), -1)\n", " inputs = [s.expand(shape) for s in inputs]\n", "\n", " input_cat = torch.cat(inputs, dim=-1)\n", " return self.fc(input_cat)" ] }, { "cell_type": "code", "execution_count": null, "id": "victorian-brain", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "# The code is based on the scArches trVAE model: https://github.com/theislab/scarches/blob/v0.3.5/scarches/models/trvae/trvae.py\n", "# and on the pyro.ai Variational Autoencoders tutorial: http://pyro.ai/examples/vae.html\n", "# and the Semi-Supervised VAE tutorial: http://pyro.ai/examples/ss-vae.html\n", "class CVAE(nn.Module):\n", " def __init__(\n", " self,\n", " input_dim,\n", " *,\n", " n_conds: int,\n", " n_classes: int,\n", " hidden_dims: Sequence[int],\n", " latent_dim: int,\n", " classifier_dims: Sequence[int] = (128,),\n", " ) -> None:\n", " super().__init__()\n", " # output - mean and logvar of z\n", " self.encoder = MLP(input_dim + n_conds, hidden_dims, 2 * latent_dim)\n", " self.decoder = MLP(\n", " latent_dim + n_conds + n_classes, hidden_dims[::-1], input_dim\n", " )\n", " self.theta = nn.Linear(n_conds, input_dim, bias=False)\n", " self.classifier = MLP(latent_dim, classifier_dims, n_classes)\n", " self.latent_dim = latent_dim\n", "\n", " def model(self, x, batches, classes, size_factors, supervised):\n", " pyro.module(\"cvae\", self)\n", "\n", " batch_size = x.shape[0]\n", "\n", " with pyro.plate(\"data\", batch_size):\n", " z_loc = x.new_zeros((batch_size, self.latent_dim))\n", " z_scale = x.new_ones((batch_size, self.latent_dim))\n", " z = pyro.sample(\"Z\", dist.Normal(z_loc, z_scale).to_event(1))\n", "\n", " classes_probs = self.classifier(z).softmax(dim=-1)\n", " obs = classes if supervised else None\n", " classes = pyro.sample(\n", " \"Class\", dist.OneHotCategorical(probs=classes_probs), obs=obs\n", " )\n", "\n", " dec_mu = (\n", " self.decoder(z, batches, classes).softmax(dim=-1)\n", " * size_factors[:, None]\n", " )\n", " dec_theta = torch.exp(self.theta(batches))\n", "\n", " logits = (dec_mu + 1e-6).log() - (dec_theta + 1e-6).log()\n", "\n", " pyro.sample(\n", " \"X\",\n", " dist.NegativeBinomial(total_count=dec_theta, logits=logits).to_event(1),\n", " obs=x.int(),\n", " )\n", "\n", " def guide(self, x, batches, classes, size_factors, supervised):\n", " batch_size = x.shape[0]\n", "\n", " with pyro.plate(\"data\", batch_size):\n", " z_loc_scale = self.encoder(x, batches)\n", "\n", " z_mu = z_loc_scale[:, : self.latent_dim]\n", " z_var = torch.sqrt(torch.exp(z_loc_scale[:, self.latent_dim :]) + 1e-4)\n", "\n", " z = pyro.sample(\"Z\", dist.Normal(z_mu, z_var).to_event(1))\n", "\n", " if not supervised:\n", " classes_probs = self.classifier(z).softmax(dim=-1)\n", " pyro.sample(\"Class\", dist.OneHotCategorical(probs=classes_probs))" ] }, { "cell_type": "markdown", "id": "floral-relation", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "source": [ "## Create a concatenated `AnnData` from two `AnnData` objects" ] }, { "cell_type": "markdown", "id": "7eb9d012-a6a2-4854-89d2-99251286dc5e", "metadata": { "editable": true, "jp-MarkdownHeadingCollapsed": true, "slideshow": { "slide_type": "" }, "tags": [] }, "source": [ "The data is from this [scvi reproducibility notebook](https://yoseflab.github.io/scvi-tools-reproducibility/scarches_totalvi_seurat_data/)." ] }, { "cell_type": "code", "execution_count": null, "id": "190da95b", "metadata": { "tags": [ "hide-cell" ] }, "outputs": [], "source": [ "def download(url, fname, sha_hash):\n", " pooch.retrieve(\n", " url=url,\n", " known_hash=sha_hash,\n", " fname=fname,\n", " path=\".\",\n", " downloader=pooch.HTTPDownloader(\n", " progressbar=True,\n", " chunk_size=1024,\n", " timeout=120,\n", " headers={\n", " \"User-Agent\": \"an/ndata/1.0.0 (https://github.com/scverse/anndata)\"\n", " },\n", " ),\n", " )" ] }, { "cell_type": "code", "execution_count": null, "id": "interstate-crack", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [ "hide-output" ] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Downloading data from 'https://scverse-exampledata.s3.eu-west-1.amazonaws.com/anndata/cite_covid_full.h5ad' to file '/localscratch/ilan.gold/anndata-tutorials/covid_cite.h5ad'.\n", "100%|████████████████████████████████████████| 304M/304M [00:00<00:00, 491GB/s]\n", "Downloading data from 'https://scverse-exampledata.s3.eu-west-1.amazonaws.com/anndata/pbmc_seurat_v4.h5ad' to file '/localscratch/ilan.gold/anndata-tutorials/pbmc_seurat_v4.h5ad'.\n", "100%|█████████████████████████████████████| 1.00G/1.00G [00:00<00:00, 1.50TB/s]\n" ] } ], "source": [ "download(\n", " \"https://scverse-exampledata.s3.eu-west-1.amazonaws.com/anndata/cite_covid_full.h5ad\",\n", " \"covid_cite.h5ad\",\n", " \"cb9745b2d642459926194961f34110a17c16ef6b11777304a3478e12bd682657\",\n", ")\n", "download(\n", " \"https://scverse-exampledata.s3.eu-west-1.amazonaws.com/anndata/pbmc_seurat_v4.h5ad\",\n", " \"pbmc_seurat_v4.h5ad\",\n", " \"c3b0100a6ce27beb64eff53692e09f98da2a58cfdfea08d15ff204f834b41396\",\n", ")" ] }, { "cell_type": "code", "execution_count": 6, "id": "listed-fairy", "metadata": { "editable": true, "slideshow": { "slide_type": "" }, "tags": [] }, "outputs": [], "source": [ "covid = ad.read_h5ad(\"covid_cite.h5ad\")\n", "pbmc = ad.read_h5ad(\"pbmc_seurat_v4.h5ad\")" ] }, { "cell_type": "code", "execution_count": 9, "id": "e1380244-39e9-45a1-bf30-1acc6224361b", "metadata": {}, "outputs": [], "source": [ "covid.obs[\"size_factors\"] = np.asarray(covid.X.sum(1)).ravel()\n", "pbmc.obs[\"size_factors\"] = np.asarray(pbmc.X.sum(1)).ravel()" ] }, { "cell_type": "markdown", "id": "hvg-explanation", "metadata": {}, "source": [ "We select highly variable genes using pbmc as the reference, since it has rich cell type annotations and multiple batches to compute variability across.\n", "Concatenating with `join='inner'` then restricts covid to the same gene set." ] }, { "cell_type": "code", "execution_count": 10, "id": "closed-greeting", "metadata": {}, "outputs": [], "source": [ "sc.pp.highly_variable_genes(\n", " pbmc,\n", " n_top_genes=4000,\n", " flavor=\"seurat_v3\",\n", " batch_key=\"orig.ident\",\n", " subset=True,\n", ")" ] }, { "cell_type": "code", "execution_count": 11, "id": "2e91bbe4-4cc1-4f82-b2b0-cde78eb17e5f", "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/plain": [ "AnnData object with n_obs × n_vars = 57669 × 33538\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'RNA_snn_res.0.4', 'seurat_clusters', 'set', 'Resp', 'disease', 'subj_code', 'covidpt_orhealth', 'mito', 'ncount', 'nfeat', 'bust_21', 'og_clust', 'severmod_other', 'og_clusts', 'nCount_ADT', 'nFeature_ADT', 'UMAP1', 'UMAP2', 'final_clust', 'final_clust_v2', 'new_pt_id', 'Resp_og', 'final_clust_withnum', 'final_clust_review', 'Age', 'Gender', 'Gender_num', 'size_factors'\n", " obsm: 'pro_exp'" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "covid" ] }, { "cell_type": "code", "execution_count": 12, "id": "1f101ae4-d68a-4e36-8349-4ac867fa8f61", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "AnnData object with n_obs × n_vars = 161764 × 4000\n", " obs: 'nCount_ADT', 'nFeature_ADT', 'nCount_RNA', 'nFeature_RNA', 'orig.ident', 'lane', 'donor', 'time', 'celltype.l1', 'celltype.l2', 'celltype.l3', 'Phase', 'nCount_SCT', 'nFeature_SCT', 'X_index', 'size_factors'\n", " var: 'highly_variable', 'highly_variable_rank', 'means', 'variances', 'variances_norm', 'highly_variable_nbatches'\n", " uns: 'hvg'\n", " obsm: 'protein_counts'" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "pbmc" ] }, { "cell_type": "code", "execution_count": 13, "id": "straight-license", "metadata": {}, "outputs": [], "source": [ "dataset = ad.concat(\n", " {\"covid\": covid, \"pbmc\": pbmc},\n", " join=\"inner\",\n", " label=\"dataset\",\n", ")" ] }, { "cell_type": "code", "execution_count": 14, "id": "chief-heading", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "AnnData object with n_obs × n_vars = 219433 × 4000\n", " obs: 'orig.ident', 'nCount_RNA', 'nFeature_RNA', 'nCount_ADT', 'nFeature_ADT', 'size_factors', 'dataset'" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "dataset" ] }, { "cell_type": "markdown", "id": "b6fddc99-f389-4e33-88aa-a92a0806d42a", "metadata": { "jp-MarkdownHeadingCollapsed": true, "tags": [] }, "source": [ "Setup important covariates in the concatenated object as well as one-hot encoders." ] }, { "cell_type": "code", "execution_count": 15, "id": "friendly-lightweight", "metadata": {}, "outputs": [], "source": [ "dataset.obs[\"orig.ident\"] = dataset.obs[\"orig.ident\"].astype(str)\n", "dataset.obs.loc[covid.obs_names, \"orig.ident\"] = covid.obs[\"set\"]\n", "dataset.obs[\"orig.ident\"] = dataset.obs[\"orig.ident\"].astype(\"category\")" ] }, { "cell_type": "code", "execution_count": 16, "id": "restricted-interstate", "metadata": {}, "outputs": [], "source": [ "dataset.obs[\"size_factors\"] = dataset.X.sum(1)" ] }, { "cell_type": "code", "execution_count": 17, "id": "satisfied-slovenia", "metadata": {}, "outputs": [], "source": [ "labels = \"celltype.l1\"" ] }, { "cell_type": "code", "execution_count": 18, "id": "6e4edca9-493f-488e-9a24-1a29f8778740", "metadata": {}, "outputs": [], "source": [ "dataset.obs[labels] = pd.concat([\n", " pd.Series(index=covid.obs_names, data=[\"not-annotated\"] * covid.shape[0]),\n", " pbmc.obs[\"celltype.l1\"].astype(\"str\"),\n", "]).astype(\"category\")" ] }, { "cell_type": "code", "execution_count": 19, "id": "d0d21b68-4320-4571-88aa-6f3c3aca7211", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "AAACCCACACCAGCGT-1 not-annotated\n", "AAACCCACATCTCAAG-1 not-annotated\n", "AAACGAAAGACCTGGA-1 not-annotated\n", "AAACGCTCAGTGGGTA-1 not-annotated\n", "AAACGCTGTAGCTTGT-1 not-annotated\n", " ... \n", "E2L8_TTTGTTGGTCGTGATT CD8 T\n", "E2L8_TTTGTTGGTGTGCCTG Mono\n", "E2L8_TTTGTTGGTTAGTTCG B\n", "E2L8_TTTGTTGGTTGGCTAT Mono\n", "E2L8_TTTGTTGTCTCATGGA Mono\n", "Name: celltype.l1, Length: 219433, dtype: category\n", "Categories (9, object): ['B', 'CD4 T', 'CD8 T', 'DC', ..., 'NK', 'not-annotated', 'other', 'other T']" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "dataset.obs[labels]" ] }, { "cell_type": "code", "execution_count": 20, "id": "loose-czech", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "use_cuda = torch.cuda.is_available()\n", "use_cuda" ] }, { "cell_type": "code", "execution_count": 21, "id": "06281003-52b1-44f7-841d-673874991945", "metadata": { "tags": [ "hide-output" ] }, "outputs": [ { "data": { "text/html": [ "
OneHotEncoder(dtype=<class 'numpy.float32'>, sparse_output=False)In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
| x0_B | \n", "
| x0_CD4 T | \n", "
| x0_CD8 T | \n", "
| x0_DC | \n", "
| x0_Mono | \n", "
| x0_NK | \n", "
| x0_other | \n", "
| x0_other T | \n", "