diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 8917d84c11..c04bfe226c 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -21,6 +21,9 @@ scvi-tools is composed of models that can perform one or many analysis tasks. In * - :doc:`/user_guide/models/scvi` - Dimensionality reduction, removal of unwanted variation, integration across replicates, donors, and technologies, differential expression, imputation, normalization of other cell- and sample-level confounding factors - :cite:p:`Lopez18` + * - :doc:`/user_guide/models/scvix` + - Cross-technology integration and zero-shot query mapping for scRNA-seq datasets + - In preparation * - :doc:`/user_guide/models/scanvi` - scVI tasks with cell type transfer from reference, seed labeling - :cite:p:`Xu21` diff --git a/docs/user_guide/models/index.md b/docs/user_guide/models/index.md index 65e706e31a..a192504dbc 100644 --- a/docs/user_guide/models/index.md +++ b/docs/user_guide/models/index.md @@ -25,6 +25,7 @@ scanvi scar scbasset scvi +scvix scviva solo stereoscope diff --git a/docs/user_guide/models/scvix.md b/docs/user_guide/models/scvix.md new file mode 100644 index 0000000000..d572b1dbdf --- /dev/null +++ b/docs/user_guide/models/scvix.md @@ -0,0 +1,265 @@ +# scVI-X + +**scVI-X** [^ref1] (single-cell Variational Inference — cross-technology; +Python class {class}`~scvi.external.SCVIX`) is a deep generative model for +learning technology-invariant cell-state representations across diverse +single-cell sequencing platforms. It extends scVI [^ref2] with three targeted +architectural changes that together enable strong cross-technology integration +and zero-shot query mapping without fine-tuning. + +The key advantages of scVI-X over scVI are: + +- **Cross-technology integration**: Explicitly models the sequencing assay + (e.g., 10x single-cell vs. single-nucleus) as a distinct source of + variation, yielding substantially better mixing across platforms. +- **Zero-shot query mapping**: Trained models can embed new cells from + previously unseen arrays via a single forward pass, with no retraining + required. +- **Assay-level adversarial training**: Adversarial correction is applied at + the assay level rather than the batch level, reducing the risk of + overintegration when cell-type compositions differ between batches. +- **Assay-aware decoding**: Learned gene-level assay biases in the output + layer capture technology-specific transcript-detection differences and + provide interpretable, reproducible signatures across datasets. +- **Variational batch embeddings**: Optionally replaces one-hot batch + encodings with low-dimensional variational embeddings, reducing parameter + count and improving identifiability through a KL regularisation term. + +## Preliminaries + +scVI-X takes as input a scRNA-seq count matrix $X$ with $N$ cells and $G$ +genes. For each cell $n$ it requires: + +- An **assay label** $a_n$ identifying the sequencing technology or + suspension type (e.g., 10x Chromium single-cell, 10x Chromium + single-nucleus, Smart-seq2). This is the primary covariate that scVI-X is + designed to correct for. +- A **batch label** $b_n$ for experimental nuisance variation within an + assay (e.g., sequencing run, donor, processing day). + +Optionally, scVI-X can also accept: + +- Cell-type labels $y_n$ for supervised integration via a cell-type + biased prior. +- An **adversarial group label** $g_n$ (defaults to $a_n$) that defines + which grouping the adversarial classifier acts on. To reduce + overintegration if cell-type composition varies across assays. +- Additional categorical or continuous covariates. + +## Generative process + +scVI-X posits a conditional VAE generative model. For each cell $n$: + +1. A latent cell-state vector $z_n$ is drawn from the prior: + + $z_n \sim p(z)$ + + where $p(z)$ is one of: a standard Gaussian, a mixture of Gaussians + (MOG), a variational amortized mixture of posteriors (VampPrior). + +2. The normalised expression scale $h_n$ is decoded from $z_n$ conditioned + on the batch $b_n$ and assay $a_n$: + + $h_n = \mathrm{softmax}\!\left(f_\theta(z_n,\, b_n) + \delta_{a_n}\right)$ + + where $f_\theta$ is a fully-connected decoder network, and + $\delta_{a_n} \in \mathbb{R}^G$ is a learned, assay-specific output bias + vector that captures gene-level differences in transcript-detection + efficiency across sequencing platforms. + +3. Observed counts are generated from a negative binomial (or ZINB / Poisson) + likelihood: + + $x_{ng} \mid h_{ng} \sim \mathrm{NegativeBinomial}(l_n\, h_{ng},\; r_{ng})$ + + where $l_n$ is the observed library size and $r_{ng}$ is the + gene-specific (or cell-specific) inverse dispersion. + +The latent variables are summarised below: + +```{eval-rst} +.. list-table:: + :widths: 20 80 15 + :header-rows: 1 + + * - Variable + - Description + - Code name + * - :math:`z_n \in \mathbb{R}^L` + - Technology-invariant cell-state embedding. + - ``z`` + * - :math:`h_n \in \mathbb{R}^G` + - Normalised gene expression scale. + - ``px_scale`` + * - :math:`l_n \in \mathbb{R}^+` + - Observed library size (log-transformed). + - ``library`` + * - :math:`r_{ng} \in \mathbb{R}^+` + - Gene- (and optionally cell-) specific inverse dispersion. + - ``px_r`` + * - :math:`\delta_{a_n} \in \mathbb{R}^G` + - Assay-specific gene-level output bias. + - output layer weights conditioned on ``assay_index`` +``` + +### Variational batch embeddings (optional) + +When `batch_representation="embedding"` is selected, each batch $b$ is +represented by a low-dimensional embedding +$e_b \sim \mathcal{N}(\mu_b, \sigma_b^2 I)$ rather than a one-hot vector. +This reduces the number of decoder parameters. This embedding can be variational +and then adds a KL regularisation +term to the loss that improves identifiability of batch representations: + +$\mathcal{L}_\text{embed} = \mathrm{KL}\!\left(q(e_b) \,\|\, \mathcal{N}(0,I)\right)$ + +## Inference + +scVI-X uses amortised variational inference. The approximate posterior is: + +$q_\phi(z_n \mid x_n, a_n) = \mathcal{N}\!\left(\mu_\phi(x_n, a_n),\; \sigma^2_\phi(x_n, a_n)\, I\right)$ + +where $\mu_\phi$ and $\sigma^2_\phi$ are encoder neural networks. The assay +index $a_n$ is injected into the encoder through **conditional layer +normalisation** (controlled by `conditional_norm=True`, the default): instead +of a single global scale and shift applied after each hidden layer, the model +learns per-assay scale and shift parameters $\gamma_{a_n}$ and $\beta_{a_n}$. +This provides a lightweight but effective mechanism for correcting +assay-specific distributional shifts in the encoder. + +The evidence lower bound (ELBO) trained is: + +$\mathcal{L} = \mathbb{E}_{q_\phi(z_n \mid x_n, a_n)}\!\left[\log p_\theta(x_n \mid z_n, b_n, a_n)\right] - \mathrm{KL}\!\left(q_\phi(z_n \mid x_n, a_n) \,\|\, p(z_n)\right) + \mathcal{L}_\text{embed}$ + +### Adversarial training + +When multiple assays are present, scVI-X adds an adversarial +classifier that acts on the latent space $z_n$. The classifier is trained to +predict the adversarial label $a_n$ (typically the assay), while the +encoder is simultaneously trained to fool it: + +$\mathcal{L}_\text{adv} = -\mathbb{E}\left[\log p_\psi(\hat{a}_n \neq a_n \mid z_n)\right]$ + +By conditioning adversarial correction at the **assay** level rather than the +fine-grained batch level, scVI-X reduces the risk of overintegration in +settings where cell-type compositions differ between donors or conditions +within the same assay. + +Adversarial training is automatically enabled when more than one assay is +present; the adversarial label can be set to any registered categorical +covariate via `adversarial_key`. + +## Prior options + +scVI-X supports four prior distributions for $z_n$, selected via the `prior` +argument at model initialisation. We found Gaussian to perform well in all tested scenarios. + +- **`'gaussian'`** (default): Standard normal $\mathcal{N}(0, I)$. Fastest + to train; suitable for most use cases. +- **`'mog'`**: Mixture of Gaussians with $K$ learnable components. Encourages + a more structured latent space. +- **`'vamp'`**: VampPrior ([Tomczak & Welling, 2018](https://doi.org/10.48550/arXiv.1705.07120)). + Prior modes are anchored to learned pseudoinputs drawn from the data, + making the prior data-adaptive. +- **`'mog_celltype'`**: Mixture of Gaussians with one component per cell-type + label. Requires `labels_key` to be set in `setup_anndata`. Guides + integration by biasing the prior toward annotated cell-type clusters. + +## Tasks + +### Latent representation + +The primary output of scVI-X is a low-dimensional, technology-corrected +embedding of each cell: + +```python +import scvi + +scvi.external.SCVIX.setup_anndata( + adata, + batch_key="batch", + assay_key="assay", # sequencing technology / suspension type +) +model = scvi.external.SCVIX( + adata +) # defaults: embedding, n_latent=20, n_layers=2, n_hidden=512 +model.train() # defaults: batch_size=1024, n_epochs_kl_warmup=5 +adata.obsm["X_scVIX"] = model.get_latent_representation() +``` + +### Normalised expression + +Batch- and assay-corrected normalised expression can be obtained by decoding +from the latent space while fixing the batch and/or assay to a reference +value. This is analogous to `get_normalized_expression` in scVI and supports +the `transform_batch` and `transform_assay` arguments: + +```python +# Expression as it would look in a given reference batch +norm_expr = model.get_normalized_expression(transform_batch="batch_1") +``` + +### Zero-shot query mapping + +scVI-X supports query-to-reference mapping via the scArches framework without +any retraining. A query dataset with new batches (but the same gene set and assay) can +be embedded directly without additional training. + +```python +# Save the reference model +model.save("scvix_reference/") + +# Prepare the query and embed zero-shot +scvi.external.SCVIX.prepare_query_anndata(adata_query, "scvix_reference/") +query_model = scvi.external.SCVIX.load_query_data(adata_query, "scvix_reference/") +query_model.train(max_epochs=200, plan_kwargs={"weight_decay": 0.0}) +adata_query.obsm["X_scVIX"] = query_model.get_latent_representation() +``` + +### Differential expression + +Standard differential expression analysis between cell groups is available +via the inherited `differential_expression` method: + +```python +de_results = model.differential_expression(groupby="cell_type", group1="T cell") +``` + +### Batch embeddings + +When `batch_representation="embedding"` is used, the learned batch embeddings +can be retrieved and used for downstream analysis (e.g., visualising +batch-level variation or clustering batches): + +```python +model = scvi.external.SCVIX(adata, batch_representation="embedding") +model.train() +batch_emb = model.get_batch_representation() # shape: (n_cells, embedding_dim) +``` + +## Key setup_anndata arguments + +```{eval-rst} +.. list-table:: + :widths: 25 75 + :header-rows: 1 + + * - Argument + - Description + * - ``batch_key`` + - Column in ``adata.obs`` for experimental batch (donor, run, etc.). + * - ``assay_key`` + - Column in ``adata.obs`` for sequencing technology / suspension type. + Drives conditional layer normalisation and assay-output biases. +``` + +[^ref1]: + Can Ergen, Ori Kronfeld, Martin Kim, Shiyi Yang, Florian Ingelfinger, + Nir Yosef (in preparation), + _scvi-X: Learning technology-invariant cell states_. + +[^ref2]: + Romain Lopez, Jeffrey Regier, Michael B Cole, Michael I Jordan, Nir Yosef + (2018), + _Deep generative modeling for single-cell transcriptomics_, + [Nature Methods](https://doi.org/10.1038/s41592-018-0229-2). diff --git a/src/scvi/external/__init__.py b/src/scvi/external/__init__.py index 5752804bdd..54808a63a0 100644 --- a/src/scvi/external/__init__.py +++ b/src/scvi/external/__init__.py @@ -12,6 +12,7 @@ from .scar import SCAR from .scbasset import SCBASSET from .scviva import SCVIVA +from .scvix import SCVIX from .solo import SOLO from .stereoscope import RNAStereoscope, SpatialStereoscope from .sysvi import SysVI @@ -20,6 +21,7 @@ from .velovi import VELOVI __all__ = [ + "SCVIX", "SCAR", "SOLO", "GIMVI", diff --git a/src/scvi/external/scvix/__init__.py b/src/scvi/external/scvix/__init__.py new file mode 100644 index 0000000000..aa49a88060 --- /dev/null +++ b/src/scvi/external/scvix/__init__.py @@ -0,0 +1,4 @@ +from ._model import SCVIX +from ._module import VAEX + +__all__ = ["SCVIX", "VAEX"] diff --git a/src/scvi/external/scvix/_components.py b/src/scvi/external/scvix/_components.py new file mode 100644 index 0000000000..b232284e59 --- /dev/null +++ b/src/scvi/external/scvix/_components.py @@ -0,0 +1,348 @@ +from __future__ import annotations + +import collections +from typing import TYPE_CHECKING + +import torch +from torch import nn +from torch.distributions import Normal + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + from typing import Literal + + +def _identity(x): + return x + + +class ConditionalBatchNorm1d(nn.Module): + def __init__(self, num_features: int, num_classes: int, momentum: float, eps: float): + super().__init__() + self.num_features = num_features + self.bn = nn.BatchNorm1d(num_features, momentum=momentum, eps=eps, affine=False) + self.embed = nn.Embedding(num_classes, num_features * 2) + self.embed.weight.data[:, :num_features].normal_(1, 0.02) + self.embed.weight.data[:, num_features:].zero_() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + out = self.bn(x) + gamma, beta = self.embed(y.long().ravel()).chunk(2, 1) + return gamma.view(-1, self.num_features) * out + beta.view(-1, self.num_features) + + +class ConditionalLayerNorm(nn.Module): + def __init__(self, num_features: int, num_classes: int): + super().__init__() + self.num_features = num_features + self.ln = nn.LayerNorm(num_features, elementwise_affine=False) + self.embed = nn.Embedding(num_classes, num_features * 2) + self.embed.weight.data[:, :num_features].normal_(1, 0.02) + self.embed.weight.data[:, num_features:].zero_() + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + out = self.ln(x) + gamma, beta = self.embed(y.long().ravel()).chunk(2, 1) + return gamma.view(-1, self.num_features) * out + beta.view(-1, self.num_features) + + +class FCLayersX(nn.Module): + def __init__( + self, + n_in: int, + n_out: int, + n_continuous: int = 0, + n_cat_list: Iterable[int] | None = None, + n_layers: int = 1, + n_hidden: int = 128, + dropout_rate: float = 0.1, + use_batch_norm: bool = True, + use_layer_norm: bool = False, + use_activation: bool = True, + bias: bool = True, + inject_covariates: bool = True, + activation_fn: Callable[[], nn.Module] = nn.ReLU, + conditional_norm: bool = False, + conditional_category: int = 0, + ): + super().__init__() + self.inject_covariates = inject_covariates + layers_dim = [n_in] + (n_layers - 1) * [n_hidden] + [n_out] + + if n_cat_list is not None: + self.n_cat_list = [n_cat if n_cat > 1 else 0 for n_cat in n_cat_list] + else: + self.n_cat_list = [] + self.n_continuous = n_continuous + self.cond_cat = conditional_category + + if conditional_norm: + invalid_conditional_category = ( + conditional_category >= len(self.n_cat_list) + or self.n_cat_list[conditional_category] == 0 + ) + if invalid_conditional_category: + raise ValueError( + "Conditional normalization requires a categorical covariate with more than " + "one category." + ) + + self.n_cov = n_continuous + sum(self.n_cat_list) + self.fc_layers = nn.Sequential( + collections.OrderedDict( + [ + ( + f"Layer {i}", + nn.Sequential( + nn.Linear( + n_in + self.n_cov * self.inject_into_layer(i), + n_out, + bias=bias, + ), + ConditionalBatchNorm1d( + n_out, + self.n_cat_list[self.cond_cat], + momentum=0.01, + eps=0.001, + ) + if conditional_norm and use_batch_norm + else nn.BatchNorm1d(n_out, momentum=0.01, eps=0.001) + if use_batch_norm + else None, + ConditionalLayerNorm(n_out, self.n_cat_list[self.cond_cat]) + if conditional_norm and use_layer_norm + else nn.LayerNorm(n_out, elementwise_affine=False) + if use_layer_norm + else None, + activation_fn() if use_activation else None, + nn.Dropout(p=dropout_rate) if dropout_rate > 0 else None, + ), + ) + for i, (n_in, n_out) in enumerate( + zip(layers_dim[:-1], layers_dim[1:], strict=True) + ) + ] + ) + ) + + def inject_into_layer(self, layer_num: int) -> bool: + return layer_num == 0 or (layer_num > 0 and self.inject_covariates) + + def set_online_update_hooks(self, hook_first_layer: bool = True): + self.hooks = [] + + def _hook_fn_weight(grad): + categorical_dims = sum(self.n_cat_list) + new_grad = torch.zeros_like(grad) + if categorical_dims > 0: + new_grad[:, -categorical_dims:] = grad[:, -categorical_dims:] + return new_grad + + def _hook_fn_zero_out(grad): + return grad * 0 + + for i, layers in enumerate(self.fc_layers): + for layer in layers: + if i == 0 and not hook_first_layer: + continue + if isinstance(layer, nn.Linear): + if self.inject_into_layer(i): + w = layer.weight.register_hook(_hook_fn_weight) + else: + w = layer.weight.register_hook(_hook_fn_zero_out) + self.hooks.append(w) + b = layer.bias.register_hook(_hook_fn_zero_out) + self.hooks.append(b) + + def forward( + self, + x: torch.Tensor, + *cat_list: int, + cont_input: torch.Tensor | None = None, + ) -> torch.Tensor: + one_hot_cat_list = [] + cont_list = [cont_input] if cont_input is not None else [] + cat_list = cat_list or [] + + if len(self.n_cat_list) > len(cat_list): + raise ValueError("nb. categorical args provided doesn't match init. params.") + if ( + self.n_continuous > 0 + and cont_input is not None + and cont_input.shape[-1] != self.n_continuous + ): + raise ValueError("continuous dims provided doesn't match init. params.") + + for n_cat, cat in zip(self.n_cat_list, cat_list, strict=False): + if n_cat and cat is None: + raise ValueError("cat not provided while n_cat != 0 in init. params.") + if n_cat > 1: + if cat.size(1) != n_cat: + one_hot_cat = nn.functional.one_hot(cat.squeeze(-1), n_cat) + else: + one_hot_cat = cat + one_hot_cat_list.append(one_hot_cat) + cov_list = cont_list + one_hot_cat_list + + for i, layers in enumerate(self.fc_layers): + for layer in layers: + if layer is None: + continue + if isinstance(layer, (ConditionalBatchNorm1d, ConditionalLayerNorm)): + if x.dim() == 3: + x = torch.cat( + [ + layer(slice_x, cat_list[self.cond_cat]).unsqueeze(0) + for slice_x in x + ], + dim=0, + ) + else: + x = layer(x, cat_list[self.cond_cat]) + elif isinstance(layer, nn.BatchNorm1d): + if x.dim() == 3: + if x.device.type == "mps": + x = torch.cat( + [(layer(slice_x.clone())).unsqueeze(0) for slice_x in x], dim=0 + ) + else: + x = torch.cat([layer(slice_x).unsqueeze(0) for slice_x in x], dim=0) + else: + x = layer(x) + else: + if isinstance(layer, nn.Linear) and self.inject_into_layer(i): + if x.dim() == 3: + cov_list_layer = [ + o.unsqueeze(0).expand((x.size(0), o.size(0), o.size(1))) + for o in cov_list + ] + else: + cov_list_layer = cov_list + x = torch.cat((x, *cov_list_layer), dim=-1) + x = layer(x) + return x + + +class EncoderX(nn.Module): + def __init__( + self, + n_input: int, + n_output: int, + n_continuous: int = 0, + n_cat_list: Iterable[int] | None = None, + n_layers: int = 1, + n_hidden: int = 128, + dropout_rate: float = 0.1, + distribution: str = "normal", + var_eps: float = 1e-4, + var_activation: Callable | None = None, + return_dist: bool = False, + **kwargs, + ): + super().__init__() + self.distribution = distribution + self.var_eps = var_eps + self.encoder = FCLayersX( + n_in=n_input, + n_out=n_hidden, + n_continuous=n_continuous, + n_cat_list=n_cat_list, + n_layers=n_layers, + n_hidden=n_hidden, + dropout_rate=dropout_rate, + **kwargs, + ) + self.mean_encoder = nn.Linear(n_hidden, n_output) + self.var_encoder = nn.Linear(n_hidden, n_output) + self.return_dist = return_dist + + self.z_transformation = nn.Softmax(dim=-1) if distribution == "ln" else _identity + self.var_activation = torch.exp if var_activation is None else var_activation + + def forward( + self, + x: torch.Tensor, + *cat_list: int, + cont: torch.Tensor | None = None, + ): + q = self.encoder(x, *cat_list, cont_input=cont) + q_m = self.mean_encoder(q) + q_v = self.var_activation(self.var_encoder(q)) + self.var_eps + dist = Normal(q_m, q_v.sqrt()) + latent = self.z_transformation(dist.rsample()) + if self.return_dist: + return dist, latent + return q_m, q_v, latent + + +class DecoderSCVIX(nn.Module): + def __init__( + self, + n_input: int, + n_output: int, + n_continuous: int = 0, + n_cat_list: Iterable[int] | None = None, + n_layers: int = 1, + n_hidden: int = 128, + n_conditions_output: int = 0, + inject_covariates: bool = True, + use_batch_norm: bool = False, + use_layer_norm: bool = False, + scale_activation: Literal["softmax", "softplus"] = "softmax", + **kwargs, + ): + super().__init__() + self.n_conditions_output = n_conditions_output + self.px_decoder = FCLayersX( + n_in=n_input, + n_out=n_hidden, + n_continuous=n_continuous, + n_cat_list=n_cat_list, + n_layers=n_layers, + n_hidden=n_hidden, + dropout_rate=0, + inject_covariates=inject_covariates, + use_batch_norm=use_batch_norm, + use_layer_norm=use_layer_norm, + **kwargs, + ) + + if scale_activation == "softmax": + px_scale_activation = nn.Softmax(dim=-1) + elif scale_activation == "softplus": + px_scale_activation = nn.Softplus() + else: + raise ValueError("`scale_activation` must be 'softmax' or 'softplus'.") + + self.px_scale_decoder = nn.Sequential( + nn.Linear(n_hidden + n_conditions_output, n_output), + px_scale_activation, + ) + self.px_r_decoder = nn.Linear(n_hidden + n_conditions_output, n_output) + self.px_dropout_decoder = nn.Linear(n_hidden + n_conditions_output, n_output) + + def forward( + self, + dispersion: str, + z: torch.Tensor, + library: torch.Tensor, + *cat_list: int, + cont: torch.Tensor | None = None, + output_condition: torch.Tensor | None = None, + ): + px = self.px_decoder(z, *cat_list, cont_input=cont) + if output_condition is not None and self.n_conditions_output: + one_hot_cat = nn.functional.one_hot( + output_condition.squeeze(-1), self.n_conditions_output + ).float() + else: + one_hot_cat = torch.zeros(px.size(-2), self.n_conditions_output, device=px.device) + if px.dim() == 3: + one_hot_cat = one_hot_cat.unsqueeze(0).expand(px.size(0), -1, -1) + px_cat = torch.cat([px, one_hot_cat], dim=-1) + + px_scale = self.px_scale_decoder(px_cat) + px_dropout = self.px_dropout_decoder(px_cat) + px_rate = torch.exp(library) * px_scale + px_r = self.px_r_decoder(px_cat) if dispersion == "gene-cell" else None + return px_scale, px_r, px_rate, px_dropout diff --git a/src/scvi/external/scvix/_model.py b/src/scvi/external/scvix/_model.py new file mode 100644 index 0000000000..ea63248f36 --- /dev/null +++ b/src/scvi/external/scvix/_model.py @@ -0,0 +1,486 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import numpy as np + +from scvi import REGISTRY_KEYS +from scvi.data import AnnDataManager +from scvi.data._utils import _get_adata_minify_type +from scvi.data.fields import ( + CategoricalJointObsField, + CategoricalObsField, + LabelsWithUnlabeledObsField, + LayerField, + NumericalJointObsField, +) +from scvi.dataloaders import DataSplitter +from scvi.model._utils import ( + get_max_epochs_heuristic, +) +from scvi.model.base import ( + ArchesMixin, + BaseMinifiedModeModelClass, + EmbeddingMixin, + RNASeqMixin, + VAEMixin, +) +from scvi.train import TrainRunner +from scvi.utils._docstrings import devices_dsp, setup_anndata_dsp + +from ._module import VAEX +from ._trainingplans import SCVIXTrainingPlan + +if TYPE_CHECKING: + from typing import Literal + + import pandas as pd + from anndata import AnnData + +ASSAY_KEY = "assay" +ADVERSARIAL_GROUP_KEY = "adversarial_group" + +logger = logging.getLogger(__name__) + + +class SCVIX( + EmbeddingMixin, + RNASeqMixin, + VAEMixin, + ArchesMixin, + BaseMinifiedModeModelClass, +): + """scVI-X: learning technology-invariant cell states. + + Parameters + ---------- + adata + AnnData object that has been registered via :meth:`~scvi.model.SCVI.setup_anndata`. If + ``None``, then the underlying module will not be initialized until training, and a + :class:`~lightning.pytorch.core.LightningDataModule` must be passed in during training + (``EXPERIMENTAL``). + n_hidden + Number of nodes per hidden layer. + n_latent + Dimensionality of the latent space. + n_layers + Number of hidden layers used for encoder and decoder NNs. + dropout_rate + Dropout rate for neural networks. + dispersion + One of the following: + + * ``'gene'`` - dispersion parameter of NB is constant per gene across cells + * ``'gene-batch'`` - dispersion can differ between different batches + * ``'gene-cell'`` - dispersion can differ for every gene in every cell + gene_likelihood + One of: + + * ``'nb'`` - Negative binomial distribution + * ``'zinb'`` - Zero-inflated negative binomial distribution + * ``'poisson'`` - Poisson distribution + * ``'normal'`` - ``EXPERIMENTAL`` Normal distribution + prior + One of: + * ``'gaussian'`` - Gaussian prior + * ``'mog'`` - Mixture of Gaussians prior + * ``'vamp'`` - Variational Amortized Mixture of Posteriors prior + * ``'mog_celltype'`` - Mixture of Gaussians prior with cell-type bias + pseudoinputs_data_indices + Indices of cells to use as pseudoinputs for the VAMP prior. If ``None``, a random sample of + ``n_prior_components`` cells will be used. + n_prior_components + Number of components to use for the VAMP and MOG priors. This is the number of pseudoinputs + used for the VAMP prior, and the number of components in the MOG prior. + Defaults to 50. + **kwargs + Additional keyword arguments for :class:`~scvi.module.VAE`. + + Examples + -------- + >>> adata = anndata.read_h5ad(path_to_anndata) + >>> scvi.model.SCVI.setup_anndata(adata, batch_key="batch") + >>> vae = scvi.model.SCVI(adata) + >>> vae.train() + >>> adata.obsm["X_scVI"] = vae.get_latent_representation() + >>> adata.obsm["X_normalized_scVI"] = vae.get_normalized_expression() + + Notes + ----- + See further usage examples in the following tutorials: + + 1. :doc:`/tutorials/notebooks/quick_start/api_overview` + 2. :doc:`/tutorials/notebooks/scrna/harmonization` + 3. :doc:`/tutorials/notebooks/scrna/scarches_scvi_tools` + 4. :doc:`/tutorials/notebooks/scrna/scvi_in_R` + + See Also + -------- + :class:`~scvi.module.VAE` + """ + + _module_cls = VAEX + _LATENT_QZM_KEY = "scvix_latent_qzm" + _LATENT_QZV_KEY = "scvix_latent_qzv" + _data_splitter_cls = DataSplitter + _training_plan_cls = SCVIXTrainingPlan + _train_runner_cls = TrainRunner + + def __init__( + self, + adata: AnnData | None = None, + n_hidden: int = 512, + n_latent: int = 20, + n_layers: int = 2, + dropout_rate: float = 0.05, + dispersion: Literal["gene", "gene-batch", "gene-assay", "gene-cell"] = "gene-assay", + gene_likelihood: Literal["zinb", "nb", "poisson", "normal"] = "nb", + prior: Literal["gaussian", "mog", "vamp", "mog_celltype"] = "gaussian", + batch_representation: Literal["one-hot", "embedding"] = "embedding", + batch_embedding_kwargs: dict | None = None, + pseudoinputs_data_indices: np.array | None = None, + n_prior_components: int = 50, + **kwargs, + ): + super().__init__(adata) + + if batch_embedding_kwargs is None: + batch_embedding_kwargs = {"variational": True, "embedding_dim": n_latent} + + self._module_kwargs = { + "n_hidden": n_hidden, + "n_latent": n_latent, + "n_layers": n_layers, + "dropout_rate": dropout_rate, + "dispersion": dispersion, + "gene_likelihood": gene_likelihood, + "batch_representation": batch_representation, + "batch_embedding_kwargs": batch_embedding_kwargs, + **kwargs, + } + self._model_summary_string = ( + "SCVI-X model with the following parameters: \n" + f"n_hidden: {n_hidden}, n_latent: {n_latent}, n_layers: {n_layers}, " + f"dropout_rate: {dropout_rate}, dispersion: {dispersion}, " + f"gene_likelihood: {gene_likelihood}, " + f"batch_representation: {batch_representation}." + ) + + if prior == "vamp": + if pseudoinputs_data_indices is None: + pseudoinputs_data_indices = np.random.randint( + 0, self.summary_stats.n_cells, n_prior_components + ) + if pseudoinputs_data_indices.ndim != 1: + raise ValueError("`pseudoinputs_data_indices` must be one-dimensional.") + if pseudoinputs_data_indices.shape[0] != n_prior_components: + raise ValueError( + "`pseudoinputs_data_indices` must contain exactly " + f"{n_prior_components} indices." + ) + pseudoinput_data = next( + iter( + self._make_data_loader( + adata=adata, + indices=pseudoinputs_data_indices, + batch_size=n_prior_components, + shuffle=False, + ) + ) + ) + else: + pseudoinput_data = None + + n_cats_per_cov = ( + self.adata_manager.get_state_registry(REGISTRY_KEYS.CAT_COVS_KEY).n_cats_per_key + if REGISTRY_KEYS.CAT_COVS_KEY in self.adata_manager.data_registry + else None + ) + self.module = self._module_cls( + n_input=self.summary_stats.n_vars, + n_batch=self.summary_stats.n_batch, + n_assay=self.summary_stats.n_assay, + n_adversarial_group=self.summary_stats.get("n_adversarial_group", 1), + n_labels=self.summary_stats.get("n_labels", 1), + n_continuous_cov=self.summary_stats.get("n_extra_continuous_covs", 0), + n_cats_per_cov=n_cats_per_cov, + n_hidden=n_hidden, + n_latent=n_latent, + n_layers=n_layers, + dropout_rate=dropout_rate, + dispersion=dispersion, + gene_likelihood=gene_likelihood, + batch_representation=batch_representation, + batch_embedding_kwargs=batch_embedding_kwargs, + prior=prior, + pseudoinput_data=pseudoinput_data, + n_prior_components=n_prior_components, + **kwargs, + ) + self.module.minified_data_type = self.minified_data_type + + self.init_params_ = self._get_init_params(locals()) + + @devices_dsp.dedent + def train( + self, + max_epochs: int | None = None, + lr: float = 4e-3, + accelerator: str = "auto", + devices: int | list[int] | str = "auto", + train_size: float | None = None, + validation_size: float | None = None, + shuffle_set_split: bool = True, + batch_size: int = 1024, + early_stopping: bool = False, + n_epochs_kl_warmup: int | None = 5, + adversarial_classifier: bool | None = None, + adversarial_key: str = ASSAY_KEY, + scale_adversarial_loss: float | str = 5.0, + adversarial_steps: int = 1, + weight_kl_sample: float = 0.0, + datasplitter_kwargs: dict | None = None, + plan_kwargs: dict | None = None, + external_indexing: list[np.array] = None, + **kwargs, + ): + """Trains the model using amortized variational inference. + + Parameters + ---------- + max_epochs + Number of passes through the dataset. + lr + Learning rate for optimization. + %(param_accelerator)s + %(param_devices)s + train_size + Size of training set in the range [0.0, 1.0]. + validation_size + Size of the test set. If `None`, defaults to 1 - `train_size`. If + `train_size + validation_size < 1`, the remaining cells belong to a test set. + shuffle_set_split + Whether to shuffle indices before splitting. If `False`, the val, train, and test set + are split in the sequential order of the data according to `validation_size` and + `train_size` percentages. + batch_size + Minibatch size to use during training. + early_stopping + Whether to perform early stopping with respect to the validation set. + n_epochs_kl_warmup + Number of epochs to scale weight on KL divergences from 0 to 1. + adversarial_key + Key in `adata.obs` that corresponds to the batch or assay key to use for adversarial + training. Defaults to the assay key. + scale_adversarial_loss + Scaling factor for the adversarial loss component. Higher values enforce stronger + assay mixing. Use ``"auto"`` to scale automatically with the KL warmup weight. + adversarial_steps + Number of adversarial classifier gradient steps per training step. + weight_kl_sample + Weight on the KL divergence of the variational batch embeddings. + datasplitter_kwargs + Additional keyword arguments passed into :class:`~scvi.dataloaders.DataSplitter`. + plan_kwargs + Keyword args for :class:`~scvi.train.AdversarialTrainingPlan`. Keyword arguments passed + to `train()` will overwrite values present in `plan_kwargs`, when appropriate. + external_indexing + A list of data split indices in the order of training, validation, and test sets. + Validation and test set are not required and can be left empty. + **kwargs + Other keyword args for :class:`~scvi.train.Trainer`. + """ + if adversarial_classifier is None: + if self.module.n_assay > 1: + adversarial_classifier = True + else: + adversarial_classifier = False + n_epochs_kl_warmup = ( + n_epochs_kl_warmup if n_epochs_kl_warmup is not None else max_epochs // 2 + ) + + update_dict = { + "lr": lr, + "adversarial_classifier": adversarial_classifier, + "adversarial_key": adversarial_key, + "n_epochs_kl_warmup": n_epochs_kl_warmup, + "scale_adversarial_loss": scale_adversarial_loss, + "adversarial_steps": adversarial_steps, + "weight_kl_sample": weight_kl_sample, + } + if plan_kwargs is not None: + plan_kwargs.update(update_dict) + else: + plan_kwargs = update_dict + + if max_epochs is None: + max_epochs = get_max_epochs_heuristic(self.adata.n_obs) + + plan_kwargs = plan_kwargs if isinstance(plan_kwargs, dict) else {} + datasplitter_kwargs = datasplitter_kwargs or {} + + data_splitter = self._data_splitter_cls( + self.adata_manager, + train_size=train_size, + validation_size=validation_size, + shuffle_set_split=shuffle_set_split, + batch_size=batch_size, + external_indexing=external_indexing, + **datasplitter_kwargs, + ) + training_plan = self._training_plan_cls(self.module, **plan_kwargs) + runner = self._train_runner_cls( + self, + training_plan=training_plan, + data_splitter=data_splitter, + max_epochs=max_epochs, + accelerator=accelerator, + devices=devices, + early_stopping=early_stopping, + **kwargs, + ) + return runner() + + def _get_transform_batch_gen_kwargs(self, batch): + kwargs = super()._get_transform_batch_gen_kwargs(batch) + if hasattr(self, "_transform_assay_override"): + kwargs["transform_assay"] = self._transform_assay_override + return kwargs + + def get_normalized_expression( + self, + adata: AnnData | None = None, + indices: list[int] | None = None, + transform_batch: list[int | str] | None = None, + transform_assay: int | str | None = None, + gene_list: list[str] | None = None, + library_size: float | Literal["latent"] = 1, + n_samples: int = 1, + n_samples_overall: int | None = None, + weights: Literal["uniform", "importance"] | None = None, + batch_size: int | None = None, + return_mean: bool = True, + return_numpy: bool | None = None, + silent: bool = True, + dataloader=None, + data_loader_kwargs: dict | None = None, + **importance_weighting_kwargs, + ) -> np.ndarray | pd.DataFrame: + """Returns the normalized (decoded) gene expression. + + Extends the base ``get_normalized_expression`` with assay conditioning via + ``transform_assay``. + + Parameters + ---------- + transform_assay + Assay to condition on for all cells when decoding. If a string, must be a valid + assay category registered via :meth:`~scvi.external.SCVIX.setup_anndata`. If an + integer, used directly as the assay index. If ``None`` (default), each cell is + decoded using its own observed assay. + transform_batch + Batch to condition on. See base class for details. + + Notes + ----- + All other parameters are identical to + :meth:`~scvi.model.base.RNASeqMixin.get_normalized_expression`. + """ + if transform_assay is not None: + adata_ = self._validate_anndata(adata) + assay_mappings = ( + self.get_anndata_manager(adata_, required=True) + .get_state_registry(ASSAY_KEY) + .categorical_mapping + ) + if isinstance(transform_assay, str): + if transform_assay not in assay_mappings: + raise ValueError(f'"{transform_assay}" is not a valid assay category.') + self._transform_assay_override = int( + np.where(assay_mappings == transform_assay)[0][0] + ) + else: + self._transform_assay_override = int(transform_assay) + + try: + result = super().get_normalized_expression( + adata=adata, + indices=indices, + transform_batch=transform_batch, + gene_list=gene_list, + library_size=library_size, + n_samples=n_samples, + n_samples_overall=n_samples_overall, + weights=weights, + batch_size=batch_size, + return_mean=return_mean, + return_numpy=return_numpy, + silent=silent, + dataloader=dataloader, + data_loader_kwargs=data_loader_kwargs, + **importance_weighting_kwargs, + ) + finally: + if hasattr(self, "_transform_assay_override"): + del self._transform_assay_override + + return result + + @classmethod + @setup_anndata_dsp.dedent + def setup_anndata( + cls, + adata: AnnData, + layer: str | None = None, + batch_key: str | None = None, + assay_key: str | None = None, + labels_key: str | None = None, + adversarial_group_key: str | None = None, + unlabeled_category: str = "unlabeled", + categorical_covariate_keys: list[str] | None = None, + continuous_covariate_keys: list[str] | None = None, + **kwargs, + ): + """%(summary)s. + + Parameters + ---------- + %(param_adata)s + %(param_layer)s + %(param_batch_key)s + assay_key + Key in ``adata.obs`` that corresponds to the assay of the data. + %(param_labels_key)s + adversarial_group_key + Key in ``adata.obs`` that corresponds to the adversarial group for adversarial + training. If ``None``, performs no conditional adversarial training. + %(param_unlabeled_category)s + %(param_cat_cov_keys)s + %(param_cont_cov_keys)s + """ + setup_method_args = cls._get_setup_method_args(**locals()) + anndata_fields = [ + LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=True), + CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key), + CategoricalObsField(ASSAY_KEY, assay_key), + CategoricalJointObsField(REGISTRY_KEYS.CAT_COVS_KEY, categorical_covariate_keys), + NumericalJointObsField(REGISTRY_KEYS.CONT_COVS_KEY, continuous_covariate_keys), + ] + if labels_key is not None: + anndata_fields.append( + LabelsWithUnlabeledObsField( + REGISTRY_KEYS.LABELS_KEY, labels_key, unlabeled_category + ) + ) + if adversarial_group_key is not None: + anndata_fields.append( + CategoricalObsField(ADVERSARIAL_GROUP_KEY, adversarial_group_key) + ) + # register new fields if the adata is minified + adata_minify_type = _get_adata_minify_type(adata) + if adata_minify_type is not None: + anndata_fields += cls._get_fields_for_adata_minification(adata_minify_type) + adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args) + adata_manager.register_fields(adata, **kwargs) + cls.register_manager(adata_manager) diff --git a/src/scvi/external/scvix/_module.py b/src/scvi/external/scvix/_module.py new file mode 100644 index 0000000000..a0001ee4a9 --- /dev/null +++ b/src/scvi/external/scvix/_module.py @@ -0,0 +1,798 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch +from torch import distributions +from torch.nn.functional import one_hot + +from scvi import REGISTRY_KEYS +from scvi.data._constants import ADATA_MINIFY_TYPE +from scvi.module._classifier import Classifier +from scvi.module._constants import MODULE_KEYS +from scvi.module.base import ( + BaseMinifiedModeModuleClass, + EmbeddingModuleMixin, + GaussianPrior, + LossOutput, + MogPrior, + VampPrior, + auto_move_data, +) +from scvi.nn import Embedding +from scvi.utils import unsupported_if_adata_minified + +from ._components import DecoderSCVIX, EncoderX + +if TYPE_CHECKING: + from collections.abc import Callable + from typing import Literal + + from torch.distributions import Distribution + +logger = logging.getLogger(__name__) + +ASSAY_KEY = "assay" +ASSAY_INDEX_KEY = "assay_index" +ADVERSARIAL_GROUP_KEY = "adversarial_group" +KL_SAMPLE_KEY = "kl_divergence_sample" + + +class VAEX(EmbeddingModuleMixin, BaseMinifiedModeModuleClass): + """Variational auto-encoder :cite:p:`Lopez18`. + + Parameters + ---------- + n_input + Number of input features. + n_batch + Number of batches. If ``0``, no batch correction is performed. + n_labels + Number of labels. + n_hidden + Number of nodes per hidden layer. Passed into :class:`~scvi.nn.Encoder` and + :class:`~scvi.nn.DecoderSCVI`. + n_latent + Dimensionality of the latent space. + n_layers + Number of hidden layers. Passed into :class:`~scvi.nn.Encoder` and + :class:`~scvi.nn.DecoderSCVI`. + n_continuous_cov + Number of continuous covariates. + n_cats_per_cov + A list of integers containing the number of categories for each categorical covariate. + dropout_rate + Dropout rate. Passed into :class:`~scvi.nn.Encoder` but not :class:`~scvi.nn.DecoderSCVI`. + dispersion + Flexibility of the dispersion parameter when ``gene_likelihood`` is either ``"nb"`` or + ``"zinb"``. One of the following: + + * ``"gene"``: parameter is constant per gene across cells. + * ``"gene-batch"``: parameter is constant per gene per batch. + * ``"gene-assay"``: parameter is constant per gene per assay. + * ``"gene-cell"``: parameter is constant per gene per cell. + log_variational + If ``True``, use :func:`~torch.log1p` on input data before encoding for numerical stability + (not normalization). + gene_likelihood + Distribution to use for reconstruction in the generative process. One of the following: + + * ``"nb"``: :class:`~scvi.distributions.NegativeBinomial`. + * ``"zinb"``: :class:`~scvi.distributions.ZeroInflatedNegativeBinomial`. + * ``"poisson"``: :class:`~scvi.distributions.Poisson`. + * ``"normal"``: :class:`~torch.distributions.Normal`. + encode_covariates + If ``True``, covariates are concatenated to gene expression prior to passing through + the encoder(s). Else, only gene expression is used. + deeply_inject_covariates + If ``True`` and ``n_layers > 1``, covariates are concatenated to the outputs of hidden + layers in the encoder(s) (if ``encoder_covariates`` is ``True``) and the decoder prior to + passing through the next layer. + batch_representation + Method for encoding batch information. One of the following: + + * ``"one-hot"``: represent batches with one-hot encodings. + * ``"embedding"``: represent batches with continuously-valued embeddings using + :class:`~scvi.nn.Embedding`. + + Note that batch representations are only passed into the encoder(s) if + ``encode_covariates`` is ``True``. + use_batch_norm + Specifies where to use :class:`~torch.nn.BatchNorm1d` in the model. One of the following: + + * ``"none"``: don't use batch norm in either encoder(s) or decoder. + * ``"encoder"``: use batch norm only in the encoder(s). + * ``"decoder"``: use batch norm only in the decoder. + * ``"both"``: use batch norm in both encoder(s) and decoder. + + Note: if ``use_layer_norm`` is also specified, both will be applied (first + :class:`~torch.nn.BatchNorm1d`, then :class:`~torch.nn.LayerNorm`). + use_layer_norm + Specifies where to use :class:`~torch.nn.LayerNorm` in the model. One of the following: + + * ``"none"``: don't use layer norm in either encoder(s) or decoder. + * ``"encoder"``: use layer norm only in the encoder(s). + * ``"decoder"``: use layer norm only in the decoder. + * ``"both"``: use layer norm in both encoder(s) and decoder. + + Note: if ``use_batch_norm`` is also specified, both will be applied (first + :class:`~torch.nn.BatchNorm1d`, then :class:`~torch.nn.LayerNorm`). + var_activation + Callable used to ensure positivity of the variance of the variational distribution. Passed + into :class:`~scvi.nn.Encoder`. Defaults to :func:`~torch.exp`. + extra_encoder_kwargs + Additional keyword arguments passed into :class:`~scvi.nn.Encoder`. + extra_decoder_kwargs + Additional keyword arguments passed into :class:`~scvi.nn.DecoderSCVI`. + batch_embedding_kwargs + Keyword arguments passed into :class:`~scvi.nn.Embedding` if ``batch_representation`` is + set to ``"embedding"``. + + Notes + ----- + Lifecycle: argument ``batch_representation`` is experimental in v1.2. + """ + + def __init__( + self, + n_input: int, + n_batch: int = 0, + n_assay: int = 0, + n_labels: int = 0, + n_adversarial_group: int = 0, + n_hidden: int = 128, + n_latent: int = 10, + n_layers: int = 1, + n_continuous_cov: int = 0, + n_cats_per_cov: list[int] | None = None, + dropout_rate: float = 0.1, + dispersion: Literal["gene", "gene-batch", "gene-assay", "gene-cell"] = "gene", + log_variational: bool = True, + gene_likelihood: Literal["zinb", "nb", "poisson"] = "nb", + encode_covariates: bool = False, + encode_assay: bool = True, + deeply_inject_covariates: bool = False, + batch_representation: Literal["one-hot", "embedding"] = "one-hot", + use_batch_norm: Literal["encoder", "decoder", "none", "both"] = "none", + use_layer_norm: Literal["encoder", "decoder", "none", "both"] = "both", + var_activation: Callable[[torch.Tensor], torch.Tensor] = None, + extra_encoder_kwargs: dict | None = None, + extra_decoder_kwargs: dict | None = None, + batch_embedding_kwargs: dict | None = None, + conditional_norm: bool = True, + conditional_output: bool = True, + prior: str | None = None, + pseudoinput_data: dict | None = None, + n_prior_components: int | None = 50, + ): + super().__init__() + + self.dispersion = dispersion + self.n_latent = n_latent + self.log_variational = log_variational + self.gene_likelihood = gene_likelihood + self.n_batch = n_batch + self.n_assay = n_assay + self.n_labels = n_labels + self.n_adversarial_group = n_adversarial_group + self.encode_covariates = encode_covariates + self.use_observed_lib_size = True + self.n_hidden = n_hidden + + if self.dispersion == "gene": + self.px_r = torch.nn.Parameter(3.0 * torch.ones(n_input)) + elif self.dispersion == "gene-batch": + self.px_r = torch.nn.Parameter(3.0 * torch.ones(n_input, n_batch)) + elif self.dispersion == "gene-assay": + self.px_r = torch.nn.Parameter(3.0 * torch.ones(n_input, n_assay)) + elif self.dispersion == "gene-cell": + pass + else: + raise ValueError( + "`dispersion` must be one of 'gene', 'gene-batch', 'gene-assay', 'gene-cell'." + ) + + self.batch_representation = batch_representation + self.batch_representation_encoder = False + n_cats_per_cov_ = list([] if n_cats_per_cov is None else n_cats_per_cov) + n_continuous = n_continuous_cov + + if self.batch_representation == "embedding": + self.init_embedding(REGISTRY_KEYS.BATCH_KEY, n_batch, **(batch_embedding_kwargs or {})) + n_continuous += self.get_embedding_dim(REGISTRY_KEYS.BATCH_KEY) + elif self.batch_representation != "one-hot": + raise ValueError("`batch_representation` must be one of 'one-hot' or 'embedding'.") + + use_batch_norm_encoder = use_batch_norm == "encoder" or use_batch_norm == "both" + use_batch_norm_decoder = use_batch_norm == "decoder" or use_batch_norm == "both" + use_layer_norm_encoder = use_layer_norm == "encoder" or use_layer_norm == "both" + use_layer_norm_decoder = use_layer_norm == "decoder" or use_layer_norm == "both" + + if self.batch_representation == "embedding": + cat_list = n_cats_per_cov_ + else: + cat_list = [n_batch] + n_cats_per_cov_ + self.encode_assay = encode_assay + conditional_category = 0 + if encode_assay: + encode_assay_list = [n_assay] + else: + encode_assay_list = [0] + if not encode_covariates: + encoder_cat_list = encode_assay_list + n_cont_encoder = 0 + else: + if not conditional_norm and self.batch_representation == "embedding": + # If we are not using conditional norm, we need to pass the batch information + self.batch_representation_encoder = True + encoder_cat_list = encode_assay_list + n_cats_per_cov_ + n_cont_encoder = n_continuous + else: + self.batch_representation_encoder = False + encoder_cat_list = encode_assay_list + [n_batch] + n_cats_per_cov_ + n_cont_encoder = n_continuous_cov + if conditional_norm and not encode_assay: + conditional_category = 1 + _extra_encoder_kwargs = extra_encoder_kwargs or {} + self.z_encoder = EncoderX( + n_input, + n_latent, + n_continuous=n_cont_encoder, + n_cat_list=encoder_cat_list, + n_layers=n_layers, + n_hidden=n_hidden, + dropout_rate=dropout_rate, + inject_covariates=deeply_inject_covariates, + use_batch_norm=use_batch_norm_encoder, + use_layer_norm=use_layer_norm_encoder, + var_activation=var_activation, + return_dist=True, + conditional_norm=conditional_norm, + conditional_category=conditional_category, + **_extra_encoder_kwargs, + ) + + _extra_decoder_kwargs = extra_decoder_kwargs or {} + self.decoder = DecoderSCVIX( + n_latent, + n_input, + n_cat_list=cat_list, + n_continuous=n_continuous, + n_layers=n_layers, + n_hidden=n_hidden, + n_conditions_output=self.n_assay if conditional_output else 0, + inject_covariates=deeply_inject_covariates, + use_batch_norm=use_batch_norm_decoder, + use_layer_norm=use_layer_norm_decoder, + scale_activation="softmax", + **_extra_decoder_kwargs, + ) + + cls_parameters = { + "n_layers": 1, + "n_hidden": 128, + "dropout_rate": 0.0, + "logits": True, + } + + if self.n_labels > 1: + self.classifier = Classifier( + n_latent, + n_labels=self.n_labels, + use_batch_norm=False, + use_layer_norm=True, + **cls_parameters, + ) + if prior == "gaussian": + self.prior = GaussianPrior() + elif prior == "vamp": + if pseudoinput_data is None: + raise ValueError("`pseudoinput_data` must be specified if using VampPrior.") + pseudoinput_data = self._get_inference_input(pseudoinput_data, full_forward_pass=True) + cat_list = [n_batch] + n_cats_per_cov_ + encode_assay_list + self.prior = VampPrior( + n_components=n_prior_components, + inference=self._regular_inference, + encoder=self.z_encoder, + pseudoinputs=pseudoinput_data, + n_cat_list=cat_list, + trainable_priors=True, + additional_categorical_covariates=["assay_index"], + ) + elif prior == "mog": + self.prior = MogPrior( + n_components=n_prior_components, + n_latent=n_latent, + ) + elif prior == "mog_celltype": + self.prior = MogPrior(n_components=n_labels, n_latent=n_latent, celltype_bias=True) + else: + raise ValueError("`prior` must be one of 'gaussian', 'vamp', 'mog', 'mog_celltype'.") + + @property + def embeddings_dim(self) -> dict[str, int]: + """Dictionary of embedding dimensions before variational doubling.""" + if not hasattr(self, "_embeddings_dim"): + self._embeddings_dim = {} + return self._embeddings_dim + + @property + def embedding_variational(self) -> dict[str, bool]: + """Dictionary of whether each embedding parameterizes a variational distribution.""" + if not hasattr(self, "_embedding_variational"): + self._embedding_variational = {} + return self._embedding_variational + + def init_embedding( + self, + key: str, + num_embeddings: int, + embedding_dim: int = 5, + variational: bool = False, + **kwargs, + ) -> None: + """Initialize a scVI-X embedding without changing the shared embedding mixin API.""" + self.embeddings_dim[key] = embedding_dim + self.embedding_variational[key] = variational + weight_dim = embedding_dim * 2 if variational else embedding_dim + embedding = Embedding(num_embeddings, weight_dim, **kwargs) + torch.nn.init.zeros_(embedding.weight) + self.add_embedding(key, embedding) + + def get_embedding_dim(self, key: str, default_value: int | None = None) -> int: + """Get the non-variational dimension of an embedding.""" + if key not in self.embeddings_dim: + if default_value is not None: + return default_value + raise KeyError(f"Embedding {key} not found.") + return self.embeddings_dim[key] + + def get_embedding_variational(self, key: str, default_value: bool | None = None) -> bool: + """Get whether an embedding is variational.""" + if key not in self.embedding_variational: + if default_value is not None: + return default_value + raise KeyError(f"Embedding {key} not found.") + return self.embedding_variational[key] + + @auto_move_data + def compute_embedding( + self, + key: str, + indices: torch.Tensor, + return_mean: bool = False, + return_dist: bool = False, + ) -> torch.Tensor | distributions.Normal: + """Forward pass for a scVI-X embedding.""" + indices = indices.flatten() if indices.ndim > 1 else indices + embedding = self.get_embedding(key)(indices) + if self.get_embedding_variational(key, default_value=False): + embedding_dim = self.get_embedding_dim(key) + mean = embedding[:, :embedding_dim] + scale = torch.exp(embedding[:, embedding_dim:]) + if return_mean: + return mean + dist = distributions.Normal(mean, scale) + if return_dist: + return dist + return dist.rsample() + return embedding + + def _get_inference_input( + self, + tensors: dict[str, torch.Tensor | None], + full_forward_pass: bool = False, + ) -> dict[str, torch.Tensor | None]: + """Get input tensors for the inference process.""" + if full_forward_pass or self.minified_data_type is None: + loader = "full_data" + elif self.minified_data_type in [ + ADATA_MINIFY_TYPE.LATENT_POSTERIOR, + ADATA_MINIFY_TYPE.LATENT_POSTERIOR_WITH_COUNTS, + ]: + loader = "minified_data" + else: + raise NotImplementedError(f"Unknown minified-data type: {self.minified_data_type}") + + if loader == "full_data": + return { + MODULE_KEYS.X_KEY: tensors[REGISTRY_KEYS.X_KEY], + MODULE_KEYS.BATCH_INDEX_KEY: tensors[REGISTRY_KEYS.BATCH_KEY], + ASSAY_INDEX_KEY: tensors.get(ASSAY_KEY, None), + MODULE_KEYS.CONT_COVS_KEY: tensors.get(REGISTRY_KEYS.CONT_COVS_KEY, None), + MODULE_KEYS.CAT_COVS_KEY: tensors.get(REGISTRY_KEYS.CAT_COVS_KEY, None), + ADVERSARIAL_GROUP_KEY: tensors.get(ADVERSARIAL_GROUP_KEY, None), + } + else: + return { + MODULE_KEYS.QZM_KEY: tensors[REGISTRY_KEYS.LATENT_QZM_KEY], + MODULE_KEYS.QZV_KEY: tensors[REGISTRY_KEYS.LATENT_QZV_KEY], + REGISTRY_KEYS.OBSERVED_LIB_SIZE: tensors[REGISTRY_KEYS.OBSERVED_LIB_SIZE], + } + + def _get_generative_input( + self, + tensors: dict[str, torch.Tensor], + inference_outputs: dict[str, torch.Tensor | Distribution | None], + ) -> dict[str, torch.Tensor | None]: + """Get input tensors for the generative process.""" + return { + MODULE_KEYS.Z_KEY: inference_outputs[MODULE_KEYS.Z_KEY], + MODULE_KEYS.LIBRARY_KEY: inference_outputs[MODULE_KEYS.LIBRARY_KEY], + MODULE_KEYS.BATCH_INDEX_KEY: tensors[REGISTRY_KEYS.BATCH_KEY], + ASSAY_INDEX_KEY: tensors.get(ASSAY_KEY, None), + MODULE_KEYS.CONT_COVS_KEY: tensors.get(REGISTRY_KEYS.CONT_COVS_KEY, None), + MODULE_KEYS.CAT_COVS_KEY: tensors.get(REGISTRY_KEYS.CAT_COVS_KEY, None), + } + + @auto_move_data + def _regular_inference( + self, + x: torch.Tensor, + batch_index: torch.Tensor, + assay_index: torch.Tensor | None = None, + cont_covs: torch.Tensor | None = None, + cat_covs: torch.Tensor | None = None, + adversarial_group: torch.Tensor | None = None, + n_samples: int = 1, + ) -> dict[str, torch.Tensor | Distribution | None]: + """Run the regular inference process.""" + x_ = x + if self.use_observed_lib_size: + library = torch.log(x.sum(1)).unsqueeze(1) + if self.log_variational: + x_ = x_ / x_.mean(1).unsqueeze(1) + x_ = torch.log1p(x_) + + if cat_covs is not None and self.encode_covariates: + categorical_input = torch.split(cat_covs, 1, dim=1) + else: + categorical_input = () + + if self.encode_covariates and self.batch_representation_encoder: + batch_rep = self.compute_embedding(REGISTRY_KEYS.BATCH_KEY, batch_index) + if cont_covs is not None and self.encode_covariates: + cont = torch.cat([cont_covs, batch_rep], dim=-1) + else: + cont = batch_rep + else: + if self.encode_covariates: + cont = cont_covs + else: + cont = None + if not self.encode_assay: + assay_index = None + else: + assay_index = assay_index.long() + qz, z = self.z_encoder(x_, assay_index, batch_index, *categorical_input, cont=cont) + + if n_samples > 1: + untran_z = qz.sample((n_samples,)) + z = self.z_encoder.z_transformation(untran_z) + library = library.unsqueeze(0).expand((n_samples, library.size(0), library.size(1))) + + return { + MODULE_KEYS.Z_KEY: z, + MODULE_KEYS.QZ_KEY: qz, + MODULE_KEYS.LIBRARY_KEY: library, + ADVERSARIAL_GROUP_KEY: adversarial_group, + } + + @auto_move_data + def _cached_inference( + self, + qzm: torch.Tensor, + qzv: torch.Tensor, + observed_lib_size: torch.Tensor, + n_samples: int = 1, + ) -> dict[str, torch.Tensor | None]: + """Run the cached inference process.""" + from torch.distributions import Normal + + qz = Normal(qzm, qzv.sqrt()) + # use dist.sample() rather than rsample because we aren't optimizing the z here + if n_samples == 1: + untran_z = qz.sample() + else: + qz.sample() + untran_z = qz.sample((n_samples,)) + z = self.z_encoder.z_transformation(untran_z) + library = torch.log(observed_lib_size) + if n_samples > 1: + library = library.unsqueeze(0).expand((n_samples, library.size(0), library.size(1))) + + return { + MODULE_KEYS.Z_KEY: z, + MODULE_KEYS.QZ_KEY: qz, + MODULE_KEYS.LIBRARY_KEY: library, + } + + @auto_move_data + def generative( + self, + z: torch.Tensor, + library: torch.Tensor, + batch_index: torch.Tensor, + assay_index: torch.Tensor | None = None, + cont_covs: torch.Tensor | None = None, + cat_covs: torch.Tensor | None = None, + size_factor: torch.Tensor | None = None, # Consistency + y: torch.Tensor | None = None, + transform_batch: torch.Tensor | None = None, + transform_assay: torch.Tensor | None = None, + ) -> dict[str, Distribution | None]: + """Run the generative process.""" + from torch.nn.functional import linear + + from scvi.distributions import ( + NegativeBinomial, + Normal, + Poisson, + ZeroInflatedNegativeBinomial, + ) + + if cat_covs is not None: + categorical_input = torch.split(cat_covs, 1, dim=1) + else: + categorical_input = () + + if transform_batch is not None: + batch_index = torch.ones_like(batch_index) * transform_batch + if transform_assay is not None: + assay_index = torch.ones_like(assay_index) * transform_assay + + if self.batch_representation == "embedding": + batch_rep = self.compute_embedding(REGISTRY_KEYS.BATCH_KEY, batch_index) + if cont_covs is not None: + cont = torch.cat([cont_covs, batch_rep], dim=-1) + else: + cont = batch_rep + else: + cont = cont_covs + + px_scale, px_r, px_rate, px_dropout = self.decoder( + self.dispersion, + z, + library, + batch_index, + *categorical_input, + y, + cont=cont, + output_condition=assay_index.long(), + ) + + if self.dispersion == "gene-assay": + px_r = linear( + one_hot(assay_index.squeeze(-1).long(), self.n_assay).float(), self.px_r + ) # px_r gets transposed - last dimension is nb genes + elif self.dispersion == "gene-batch": + px_r = linear(one_hot(batch_index.squeeze(-1), self.n_batch).float(), self.px_r) + elif self.dispersion == "gene": + px_r = self.px_r + + px_r = torch.exp(px_r) + + if self.gene_likelihood == "zinb": + px = ZeroInflatedNegativeBinomial( + mu=px_rate, + theta=px_r, + zi_logits=px_dropout, + scale=px_scale, + ) + elif self.gene_likelihood == "nb": + px = NegativeBinomial(mu=px_rate, theta=px_r, scale=px_scale) + elif self.gene_likelihood == "poisson": + px = Poisson(rate=px_rate, scale=px_scale) + elif self.gene_likelihood == "normal": + px = Normal(px_rate, px_r, normal_mu=px_scale) + + # Priors + if self.use_observed_lib_size: + pl = None + else: + ( + local_library_log_means, + local_library_log_vars, + ) = self._compute_local_library_params(batch_index) + pl = Normal(local_library_log_means, local_library_log_vars.sqrt()) + pz = Normal(torch.zeros_like(z), torch.ones_like(z)) + + return { + MODULE_KEYS.PX_KEY: px, + MODULE_KEYS.PL_KEY: pl, + MODULE_KEYS.PZ_KEY: pz, + } + + @unsupported_if_adata_minified + def loss( + self, + tensors: dict[str, torch.Tensor], + inference_outputs: dict[str, torch.Tensor | Distribution | None], + generative_outputs: dict[str, Distribution | None], + kl_weight: float = 1.0, + weight_assay_loss: float = 0.0, + weight_kl_sample: float = 1.0, + classification_ratio: float = 500.0, + ) -> LossOutput: + """Compute the loss.""" + from torch.distributions import kl_divergence + + x = tensors[REGISTRY_KEYS.X_KEY] + y = tensors.get(REGISTRY_KEYS.LABELS_KEY, None) + + kl_divergence_z = self.prior.kl( + qz=inference_outputs[MODULE_KEYS.QZ_KEY], + z=inference_outputs[MODULE_KEYS.Z_KEY], + labels=y, + ) + + if self.get_embedding_variational(REGISTRY_KEYS.BATCH_KEY, default_value=False): + pz_sample = self.compute_embedding( + REGISTRY_KEYS.BATCH_KEY, tensors[REGISTRY_KEYS.BATCH_KEY], return_dist=True + ) + qz_sample = distributions.Normal( + torch.zeros_like(pz_sample.loc), torch.ones_like(pz_sample.scale) + ) + kl_divergence_sample = kl_divergence(qz_sample, pz_sample).sum(dim=1) + else: + kl_divergence_sample = torch.zeros_like(kl_divergence_z) + + assay_index = tensors.get(ASSAY_KEY, None) + if weight_assay_loss > 0.0 and assay_index is not None: + assay_loss = self._compute_assay_penalty( + inference_outputs[MODULE_KEYS.QZ_KEY].loc, assay_index + ) + else: + assay_loss = 0.0 + + reconst_loss = -generative_outputs[MODULE_KEYS.PX_KEY].log_prob(x).sum(-1) + + weighted_kl_local = kl_weight * (kl_divergence_z + weight_kl_sample * kl_divergence_sample) + + loss = torch.mean(reconst_loss + weighted_kl_local + weight_assay_loss * assay_loss) + + if self.n_labels > 1: + logits = self.classifier(inference_outputs[MODULE_KEYS.Z_KEY]) + classification_loss_ = torch.nn.functional.cross_entropy( + logits, y.ravel(), reduction="none" + ) + mask = y != self.n_labels + classification_loss = classification_ratio * torch.masked_select( + classification_loss_, mask + ).mean(0) + loss += torch.mean(classification_loss) + + return LossOutput( + loss=loss, + reconstruction_loss=reconst_loss, + kl_local=kl_divergence_z, + classification_loss=classification_loss, + logits=logits, + true_labels=y, + ) + + return LossOutput( + loss=loss, + reconstruction_loss=reconst_loss, + kl_local={ + MODULE_KEYS.KL_Z_KEY: kl_divergence_z, + KL_SAMPLE_KEY: kl_divergence_sample, + }, + ) + + @torch.inference_mode() + def sample( + self, + tensors: dict[str, torch.Tensor], + n_samples: int = 1, + max_poisson_rate: float = 1e8, + ) -> torch.Tensor: + r"""Generate predictive samples from the posterior predictive distribution. + + The posterior predictive distribution is denoted as :math:`p(\hat{x} \mid x)`, where + :math:`x` is the input data and :math:`\hat{x}` is the sampled data. + + We sample from this distribution by first sampling ``n_samples`` times from the posterior + distribution :math:`q(z \mid x)` for a given observation, and then sampling from the + likelihood :math:`p(\hat{x} \mid z)` for each of these. + + Parameters + ---------- + tensors + Dictionary of tensors passed into :meth:`~scvi.module.VAE.forward`. + n_samples + Number of Monte Carlo samples to draw from the distribution for each observation. + max_poisson_rate + The maximum value to which to clip the ``rate`` parameter of + :class:`~scvi.distributions.Poisson`. Avoids numerical sampling issues when the + parameter is very large due to the variance of the distribution. + + Returns + ------- + Tensor on CPU with shape ``(n_obs, n_vars)`` if ``n_samples == 1``, else + ``(n_obs, n_vars,)``. + """ + from scvi.distributions import Poisson + + inference_kwargs = {"n_samples": n_samples} + _, generative_outputs = self.forward( + tensors, inference_kwargs=inference_kwargs, compute_loss=False + ) + + dist = generative_outputs[MODULE_KEYS.PX_KEY] + if self.gene_likelihood == "poisson": + dist = Poisson(torch.clamp(dist.rate, max=max_poisson_rate)) + + # (n_obs, n_vars) if n_samples == 1, else (n_samples, n_obs, n_vars) + samples = dist.sample() + # (n_samples, n_obs, n_vars) -> (n_obs, n_vars, n_samples) + samples = torch.permute(samples, (1, 2, 0)) if n_samples > 1 else samples + + return samples.cpu() + + def _compute_assay_penalty(self, params, assay): + assay = assay.squeeze(-1).long() + unique = torch.unique(assay) + pair_penalty = torch.tensor(0.0, device=assay.device) + if len(unique) > 1: + for i in unique: + pp = self.mmd(params, mask=(assay == i)) + pair_penalty += pp + + return pair_penalty + + def mmd(self, params, mask=None): + if mask is None: + raise ValueError("`mask` is required to compute MMD between two groups.") + mod_1 = params[mask] + mod_2 = params[~mask] + return rbf_kernel(mod_1, mod_2) + + +@auto_move_data +def rbf_kernel(x, y, gammas=None): + """ + Compute the RBF kernel between two tensors. + + Parameters + ---------- + x : torch.Tensor + Input tensor of shape (N, D). + y : torch.Tensor + Input tensor of shape (M, D). + gammas : list of float or None + List of gamma values to compute the kernel with. + + Returns + ------- + kernel_sum : torch.Tensor + Tensor of shape (N, M), the sum of RBF kernels over all gamma values. + """ + if gammas is None: + gammas = [ + 1e-10, + 1e-8, + 1e-6, + 1e-4, + 1e-3, + 1e-2, + 1e-1, + 1, + 2, + 5, + 10, + ] + kxy = torch.cdist(x, y).pow(2) + kxx = torch.cdist(x, x).pow(2) + kyy = torch.cdist(y, y).pow(2) + + kernel_sum_xy = torch.tensor(0.0, device=x.device) + kernel_sum_xx = torch.tensor(0.0, device=x.device) + kernel_sum_yy = torch.tensor(0.0, device=x.device) + for gamma in gammas: + kernel_sum_xy += torch.exp(-gamma * kxy).mean() + kernel_sum_xx += torch.exp(-gamma * kxx).mean() + kernel_sum_yy += torch.exp(-gamma * kyy).mean() + + return (kernel_sum_xx + kernel_sum_yy - 2 * kernel_sum_xy) / len(gammas) diff --git a/src/scvi/external/scvix/_trainingplans.py b/src/scvi/external/scvix/_trainingplans.py new file mode 100644 index 0000000000..590e3df1ce --- /dev/null +++ b/src/scvi/external/scvix/_trainingplans.py @@ -0,0 +1,225 @@ +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +import torch +from torch.optim.lr_scheduler import ReduceLROnPlateau + +from scvi import settings +from scvi.module import Classifier +from scvi.train import AdversarialTrainingPlan + +if TYPE_CHECKING: + from typing import Literal + + from scvi.module.base import BaseModuleClass + from scvi.train._trainingplans import TorchOptimizerCreator + +ASSAY_KEY = "assay" +ADVERSARIAL_GROUP_KEY = "adversarial_group" + + +class SCVIXTrainingPlan(AdversarialTrainingPlan): + """Adversarial training plan with scVI-X-specific classifier inputs.""" + + def __init__( + self, + module: BaseModuleClass, + *, + optimizer: Literal["Adam", "AdamW", "Custom"] = "Adam", + optimizer_creator: TorchOptimizerCreator | None = None, + lr: float = 1e-3, + weight_decay: float = 1e-6, + n_steps_kl_warmup: int = None, + n_epochs_kl_warmup: int = 400, + reduce_lr_on_plateau: bool = False, + lr_factor: float = 0.6, + lr_patience: int = 30, + lr_threshold: float = 0.0, + lr_scheduler_metric: Literal[ + "elbo_validation", "reconstruction_loss_validation", "kl_local_validation" + ] = "elbo_validation", + lr_min: float = 0, + adversarial_classifier: bool | Classifier = False, + adversarial_key: str = ASSAY_KEY, + adversarial_steps: int = 1, + scale_adversarial_loss: float | Literal["auto"] = "auto", + compile: bool = False, + compile_kwargs: dict | None = None, + **loss_kwargs, + ): + super().__init__( + module=module, + optimizer=optimizer, + optimizer_creator=optimizer_creator, + lr=lr, + weight_decay=weight_decay, + n_steps_kl_warmup=n_steps_kl_warmup, + n_epochs_kl_warmup=n_epochs_kl_warmup, + reduce_lr_on_plateau=reduce_lr_on_plateau, + lr_factor=lr_factor, + lr_patience=lr_patience, + lr_threshold=lr_threshold, + lr_scheduler_metric=lr_scheduler_metric, + lr_min=lr_min, + adversarial_classifier=False, + scale_adversarial_loss=scale_adversarial_loss, + compile=compile, + compile_kwargs=compile_kwargs, + **loss_kwargs, + ) + self.adversarial_key = adversarial_key + self.adversarial_steps = adversarial_steps + + if adversarial_classifier is True: + n_adversarial_name = f"n_{adversarial_key}" + if not hasattr(self.module, n_adversarial_name): + raise ValueError( + f"Adversarial key {adversarial_key!r} not found in module setup args." + ) + self.n_output_classifier = getattr(self.module, n_adversarial_name) + if self.n_output_classifier == 1: + warnings.warn( + "Disabling adversarial classifier as there is only one class.", + UserWarning, + stacklevel=settings.warnings_stacklevel, + ) + self.adversarial_classifier = False + else: + self.adversarial_classifier = Classifier( + n_input=self.module.n_latent + getattr(self.module, "n_adversarial_group", 0), + n_hidden=128, + n_labels=self.n_output_classifier, + n_layers=2, + logits=True, + use_batch_norm=False, + use_layer_norm=True, + ) + else: + self.adversarial_classifier = adversarial_classifier + + def loss_adversarial_classifier( + self, + z: torch.Tensor, + adversarial_group: torch.Tensor, + class_index: torch.Tensor, + predict_true_class: bool = True, + ) -> torch.Tensor: + """Loss for the scVI-X adversarial classifier.""" + n_classes = self.n_output_classifier + n_adv_group = getattr(self.module, "n_adversarial_group", 0) + if n_adv_group > 0: + adversarial_group_one_hot = torch.nn.functional.one_hot( + adversarial_group.squeeze(-1).long(), num_classes=n_adv_group + ).float() + z = torch.cat([z, adversarial_group_one_hot], dim=1) + if predict_true_class: + z = z.detach() + + cls_logits = self.adversarial_classifier(z) + if predict_true_class: + return torch.nn.functional.cross_entropy(cls_logits, class_index.squeeze(-1).long()) + + one_hot_class = torch.nn.functional.one_hot(class_index.squeeze(-1), n_classes).float() + cls_target = (1 - one_hot_class) / (n_classes - 1) + return -(cls_target * torch.nn.functional.log_softmax(cls_logits, dim=1)).sum(dim=1).mean() + + def training_step(self, batch, batch_idx): + """Training step for scVI-X adversarial training.""" + if "kl_weight" in self.loss_kwargs: + self.loss_kwargs.update({"kl_weight": self.kl_weight}) + self.log("kl_weight", self.kl_weight, on_step=True, on_epoch=False) + kappa = ( + 1 - self.kl_weight + if self.scale_adversarial_loss == "auto" + else self.scale_adversarial_loss + ) + class_tensor = batch[self.adversarial_key].long() + + opts = self.optimizers() + if not isinstance(opts, list): + opt1 = opts + opt2 = None + else: + opt1, opt2 = opts + + inference_outputs, _, scvi_loss = self.forward(batch, loss_kwargs=self.loss_kwargs) + z = inference_outputs["z"] + adversarial_group = inference_outputs.get(ADVERSARIAL_GROUP_KEY) + if adversarial_group is None: + adversarial_group = torch.zeros(z.size(0), device=z.device, dtype=torch.long) + + loss = scvi_loss.loss + orig_loss = loss + if kappa > 0 and self.adversarial_classifier is not False: + fool_loss = self.loss_adversarial_classifier( + z, + adversarial_group, + class_tensor, + predict_true_class=False, + ) + loss += fool_loss * kappa + + self.log("train_loss", loss, on_step=self.on_step, on_epoch=self.on_epoch, prog_bar=True) + if self.on_step: + self.trainer.logger.log_metrics( + {"train_loss_step": loss}, + step=self.global_step, + ) + self.compute_and_log_metrics(scvi_loss, self.train_metrics, "train") + opt1.zero_grad() + self.manual_backward(loss) + opt1.step() + + if opt2 is not None and kappa > 0: + qz = inference_outputs["qz"] + for _ in range(max(self.adversarial_steps, 0)): + z_classifier = qz.sample() + classifier_loss = kappa * self.loss_adversarial_classifier( + z_classifier, + adversarial_group, + class_tensor, + predict_true_class=True, + ) + opt2.zero_grad() + self.manual_backward(classifier_loss) + opt2.step() + + if scvi_loss.extra_metrics is not None and len(scvi_loss.extra_metrics.keys()) > 0: + self.prepare_scib_autotune(scvi_loss.extra_metrics, "training") + + return orig_loss + + def configure_optimizers(self): + """Configure optimizers for scVI-X adversarial training.""" + params1 = filter(lambda p: p.requires_grad, self.module.parameters()) + optimizer1 = self.get_optimizer_creator()(params1) + config1 = {"optimizer": optimizer1} + if self.reduce_lr_on_plateau: + scheduler1 = ReduceLROnPlateau( + optimizer1, + patience=self.lr_patience, + factor=self.lr_factor, + threshold=self.lr_threshold, + min_lr=self.lr_min, + threshold_mode="abs", + ) + config1.update( + { + "lr_scheduler": { + "scheduler": scheduler1, + "monitor": self.lr_scheduler_metric, + }, + }, + ) + + if self.adversarial_classifier is not False: + params2 = filter(lambda p: p.requires_grad, self.adversarial_classifier.parameters()) + optimizer2 = torch.optim.Adam(params2, lr=3e-4, eps=1e-4, weight_decay=1e-9) + opts = [config1.pop("optimizer"), optimizer2] + if "lr_scheduler" in config1: + return opts, [config1["lr_scheduler"]] + return opts + + return config1 diff --git a/tests/external/scvix/test_scvix.py b/tests/external/scvix/test_scvix.py new file mode 100644 index 0000000000..cb1b775eea --- /dev/null +++ b/tests/external/scvix/test_scvix.py @@ -0,0 +1,246 @@ +import os + +import numpy as np +import pytest +import torch + +import scvi +from scvi.data import synthetic_iid +from scvi.data._constants import ADATA_MINIFY_TYPE +from scvi.data._utils import _is_minified +from scvi.external import SCVIX +from scvi.model.base import BaseMinifiedModeModelClass + + +def assert_approx_equal(a, b): + # Allclose because on GPU, the values are not exactly the same + # as some values are moved to cpu during data minification + np.testing.assert_allclose(a, b, rtol=3e-1, atol=5e-1) + + +@pytest.mark.parametrize("prior", ["gaussian", "mog", "vamp", "mog_celltype"]) +def test_scvix(prior: str): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, prior=prior) + model.train(max_epochs=1) + model.get_latent_representation() + model.get_normalized_expression() + model.get_normalized_expression(transform_batch="batch_1") + model.get_normalized_expression(n_samples=2) + model.get_elbo(indices=model.validation_indices) + model.get_reconstruction_error(indices=model.validation_indices) + model.differential_expression(groupby="labels", group1="label_1") + + +@pytest.mark.parametrize("dispersion", ["gene", "gene-batch", "gene-assay", "gene-cell"]) +def test_scvix_dispersion(dispersion: str): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, dispersion=dispersion) + model.train(max_epochs=1) + model.get_normalized_expression() + + +def test_scvix_transform_assay(): + adata = synthetic_iid(batch_size=100) + # Use batch as both batch and assay (two categories: batch_0, batch_1) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch") + model = SCVIX(adata) + model.train(max_epochs=1) + + # baseline: no transformation + expr_default = model.get_normalized_expression(return_numpy=True) + + # transform by integer index (assay 0) + scvi.settings.seed = 42 + expr_int = model.get_normalized_expression(transform_assay=0, return_numpy=True) + assert expr_int.shape == expr_default.shape + + # transform by string name — must match integer result (same seed, same code) + scvi.settings.seed = 42 + expr_str = model.get_normalized_expression(transform_assay="batch_0", return_numpy=True) + np.testing.assert_array_equal(expr_int, expr_str) + + # transform_assay and transform_batch can be combined + expr_both = model.get_normalized_expression( + transform_assay="batch_0", transform_batch="batch_1", return_numpy=True + ) + assert expr_both.shape == expr_default.shape + + # different assay indices produce different outputs + expr_assay1 = model.get_normalized_expression(transform_assay=1, return_numpy=True) + assert not np.allclose(expr_int, expr_assay1) + + # invalid assay name raises ValueError + with pytest.raises(ValueError, match="not a valid assay category"): + model.get_normalized_expression(transform_assay="nonexistent_assay") + + +def test_scvix_encode_covariates(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, encode_covariates=True) + model.train(max_epochs=1) + model.get_normalized_expression(n_samples=2) + + +def test_scvix_embedding(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, batch_representation="embedding") + model.train(max_epochs=1) + model.get_normalized_expression(n_samples=2) + + +def test_scvix_layernorm(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, conditional_norm=False, use_batch_norm="both", use_layer_norm="none") + model.train(max_epochs=1) + model.get_normalized_expression(n_samples=2) + model = SCVIX(adata, conditional_norm=True, use_batch_norm="both", use_layer_norm="none") + model.train(max_epochs=1) + model.get_normalized_expression(n_samples=2) + + +def test_scvix_batch_representation_encoder_initialized_without_covariates(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch") + model = SCVIX(adata, encode_covariates=False) + + assert model.module.batch_representation_encoder is False + + +def test_scvix_vamp_prior_validates_pseudoinput_shape(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch") + + with pytest.raises(ValueError, match="one-dimensional"): + SCVIX( + adata, + prior="vamp", + pseudoinputs_data_indices=np.array([[0]]), + n_prior_components=1, + ) + + +def test_scvix_mmd_requires_mask(): + adata = synthetic_iid(batch_size=100) + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch") + model = SCVIX(adata) + + with pytest.raises(ValueError, match="mask"): + model.module.mmd(torch.randn(4, 2)) + + +def test_scvix_scarches_one_hot(save_path): + # test transfer_anndata_setup + view + adata1 = synthetic_iid() + SCVIX.setup_anndata(adata1, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata1, batch_representation="one-hot") + model.train(1, train_size=0.5) + dir_path = os.path.join(save_path, "saved_model/") + model.save(dir_path, overwrite=True) + + # adata2 has more genes and a perfect subset of adata1 + adata2 = synthetic_iid(n_genes=110) + adata2.obs["batch"] = adata2.obs.batch.cat.rename_categories(["batch_2", "batch_3"]) + SCVIX.prepare_query_anndata(adata2, dir_path) + SCVIX_query = SCVIX.load_query_data(adata2, dir_path) + SCVIX_query.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + adata3 = SCVIX.prepare_query_anndata(adata2, dir_path, inplace=False) + SCVIX_query2 = SCVIX.load_query_data(adata3, dir_path) + SCVIX_query2.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + # adata4 has more genes and missing 10 genes from adata1 + adata4 = synthetic_iid(n_genes=110) + new_var_names_init = [f"Random {i}" for i in range(10)] + new_var_names = new_var_names_init + adata4.var_names[10:].to_list() + adata4.var_names = new_var_names + + SCVIX.prepare_query_anndata(adata4, dir_path) + assert np.sum(adata4[:, adata4.var_names[:10]].X) == 0 + np.testing.assert_equal(adata4.var_names[:10].to_numpy(), adata1.var_names[:10].to_numpy()) + SCVIX_query3 = SCVIX.load_query_data(adata4, dir_path) + SCVIX_query3.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + adata5 = SCVIX.prepare_query_anndata(adata4, dir_path, inplace=False) + SCVIX_query4 = SCVIX.load_query_data(adata5, dir_path) + SCVIX_query4.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + +def test_scvix_scarches_embedding(save_path): + # test transfer_anndata_setup + view + adata1 = synthetic_iid() + SCVIX.setup_anndata(adata1, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata1, batch_representation="embedding") + model.train(1, train_size=0.5) + dir_path = os.path.join(save_path, "saved_model/") + model.save(dir_path, overwrite=True) + + # adata2 has more genes and a perfect subset of adata1 + adata2 = synthetic_iid(n_genes=110) + adata2.obs["batch"] = adata2.obs.batch.cat.rename_categories(["batch_2", "batch_3"]) + SCVIX.prepare_query_anndata(adata2, dir_path) + SCVIX_query = SCVIX.load_query_data(adata2, dir_path) + SCVIX_query.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + adata3 = SCVIX.prepare_query_anndata(adata2, dir_path, inplace=False) + SCVIX_query2 = SCVIX.load_query_data(adata3, dir_path) + SCVIX_query2.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + # adata4 has more genes and missing 10 genes from adata1 + adata4 = synthetic_iid(n_genes=110) + new_var_names_init = [f"Random {i}" for i in range(10)] + new_var_names = new_var_names_init + adata4.var_names[10:].to_list() + adata4.var_names = new_var_names + + SCVIX.prepare_query_anndata(adata4, dir_path) + assert np.sum(adata4[:, adata4.var_names[:10]].X) == 0 + np.testing.assert_equal(adata4.var_names[:10].to_numpy(), adata1.var_names[:10].to_numpy()) + SCVIX_query3 = SCVIX.load_query_data(adata4, dir_path) + SCVIX_query3.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + adata5 = SCVIX.prepare_query_anndata(adata4, dir_path, inplace=False) + SCVIX_query4 = SCVIX.load_query_data(adata5, dir_path) + SCVIX_query4.train(1, train_size=0.5, plan_kwargs={"weight_decay": 0.0}) + + +def test_scvix_minified(): + adata = synthetic_iid() + SCVIX.setup_anndata(adata, batch_key="batch", assay_key="batch", labels_key="labels") + model = SCVIX(adata, batch_representation="embedding", gene_likelihood="zinb") + model.train(1, train_size=0.5) + + qzm, qzv = model.get_latent_representation(give_mean=False, return_dist=True) + model.adata.obsm["X_latent_qzm"] = qzm + model.adata.obsm["X_latent_qzv"] = qzv + lib_size = np.squeeze(np.asarray(adata.X.sum(axis=-1))) + + scvi.settings.seed = 1 + params_orig = model.get_likelihood_parameters(n_samples=200, give_mean=True) + adata_orig = adata.copy() + + model.minify_adata() + assert model.minified_data_type == ADATA_MINIFY_TYPE.LATENT_POSTERIOR + assert model.adata_manager.registry is model.registry_ + + assert not _is_minified(adata) + assert adata is not model.adata + + orig_obs_df = adata_orig.obs + orig_obs_df[BaseMinifiedModeModelClass._OBSERVED_LIB_SIZE_KEY] = lib_size + assert model.adata.obs.equals(orig_obs_df) + assert model.adata.var_names.equals(adata_orig.var_names) + assert model.adata.var.equals(adata_orig.var) + + scvi.settings.seed = 1 + keys = ["mean", "dispersions", "dropout"] + params_latent = model.get_likelihood_parameters(n_samples=200, give_mean=True) + for k in keys: + assert params_latent[k].shape == params_orig[k].shape + + for k in keys: + assert_approx_equal(params_latent[k], params_orig[k])