diff --git a/CHANGELOG.md b/CHANGELOG.md index 12291316b0..a1a9ebbfe2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,8 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo enabling memory-efficient training on large-scale datasets stored as sharded Zarr collections, with support for batch covariates, {pr}`3620`. - Add support for rapids-singlecell, {pr}`3811`. +- Add {class}`scvi.external.JointEmbeddingSCVI`, a self-supervised SCVI variant using binomial + thinning and a cross-correlation objective (CCO) for robust embeddings, {pr}`3883`. - Add {class}`scvi.external.CYTOVI` KNN imputation backend option to be cuML, {pr}`3821`. #### Fixed diff --git a/docs/api/developer.md b/docs/api/developer.md index 510535beb4..ffb8898fe2 100644 --- a/docs/api/developer.md +++ b/docs/api/developer.md @@ -197,7 +197,7 @@ Module classes in the external API with respective generative and inference proc external.scviva.NicheLossOutput external.sysvi.SysVAE external.diagvi.DIAGVAE - + external.JointEmbeddingVAE ``` ## Module (Base) diff --git a/docs/api/user.md b/docs/api/user.md index fde643922c..d018b411e8 100644 --- a/docs/api/user.md +++ b/docs/api/user.md @@ -67,6 +67,7 @@ import scvi external.SysVI external.SCVIVA external.DIAGVI + external.JointEmbeddingSCVI external.Tangram ``` diff --git a/docs/references.bib b/docs/references.bib index 27ad192046..e0ef7db4a3 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -386,3 +386,12 @@ @article{Weinberger26 month = apr, publisher={Nature Publishing Group} } + +@article{Svensson26, + title={Improving SCVI for low-count cells through self-supervised augmentation}, + author={Valentine Svensson}, + journal={bioRxiv}, + year={2026}, + publisher={Cold Spring Harbor Laboratory}, + doi = {10.64898/2026.02.11.705441} +} diff --git a/docs/tutorials/index_scrna.md b/docs/tutorials/index_scrna.md index 8a375523c9..cf63c245e2 100644 --- a/docs/tutorials/index_scrna.md +++ b/docs/tutorials/index_scrna.md @@ -18,6 +18,7 @@ notebooks/scrna/AutoZI_tutorial notebooks/scrna/sysVI notebooks/scrna/decipher_tutorial notebooks/scrna/velovi +notebooks/scrna/JointEmbeddingSCVI_tutorial notebooks/scrna/Tahoe100_mrVI ``` @@ -126,6 +127,13 @@ Use Decipher to jointly analyze samples from distinct conditions. Use VeloVI to estimate RNA velocity. ``` +```{customcard} +:path: notebooks/scrna/JointEmbeddingSCVI_tutorial +:tags: Analysis + +Use JointEmbeddingSCVI for improving SCVI for low-count cells through self-supervised augmentation +``` + ```{customcard} :path: notebooks/scrna/Tahoe100_mrVI :tags: Analysis, Differential-comparison, Dimensionality-reduction, Removal-of-variance diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index 7f78a34ffb..8917d84c11 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -54,6 +54,9 @@ scvi-tools is composed of models that can perform one or many analysis tasks. In * - :doc:`/user_guide/models/velovi` - Deep generative modeling of transcriptional dynamics for RNA velocity analysis in single cells - :cite:p:`GayosoWeiler23` + * - :doc:`/user_guide/models/jointembeddingscvi` + - Improving SCVI for low-count cells through self-supervised augmentation + - :cite:p:`Svensson26` ``` ## ATAC-seq analysis diff --git a/docs/user_guide/models/index.md b/docs/user_guide/models/index.md index e845022bed..65e706e31a 100644 --- a/docs/user_guide/models/index.md +++ b/docs/user_guide/models/index.md @@ -12,6 +12,7 @@ decipher destvi diagvi gimvi +jointembeddingscvi linearscvi methylanvi methylvi diff --git a/docs/user_guide/models/jointembeddingscvi.md b/docs/user_guide/models/jointembeddingscvi.md new file mode 100644 index 0000000000..8825bbd01d --- /dev/null +++ b/docs/user_guide/models/jointembeddingscvi.md @@ -0,0 +1,20 @@ +# Joint-Embedding SCVI + +:::{note} +This page is under construction. +::: + +**JointEmbeddingSCVI** {cite:p}`Svensson26` (Python class {class}`~scvi.external.JointEmbeddingSCVI`) is a ... + +The advantages of JointEmbeddingSCVI are: + +- ... + +The limitations of JointEmbeddingSCVI include: + +- ... + +```{topic} Tutorials: + +- {doc}`/tutorials/notebooks/scrna/JointEmbeddingSCVI_tutorial` +``` diff --git a/src/scvi/external/__init__.py b/src/scvi/external/__init__.py index 0fe6ddd8ab..5752804bdd 100644 --- a/src/scvi/external/__init__.py +++ b/src/scvi/external/__init__.py @@ -4,6 +4,7 @@ from .decipher import Decipher from .diagvi import DIAGVI from .gimvi import GIMVI +from .joint_embedding_scvi import JointEmbeddingSCVI, JointEmbeddingVAE from .methylvi import METHYLANVI, METHYLVI from .mrvi import MRVI from .poissonvi import POISSONVI @@ -22,6 +23,8 @@ "SCAR", "SOLO", "GIMVI", + "JointEmbeddingSCVI", + "JointEmbeddingVAE", "Decipher", "RNAStereoscope", "SpatialStereoscope", diff --git a/src/scvi/external/joint_embedding_scvi/__init__.py b/src/scvi/external/joint_embedding_scvi/__init__.py new file mode 100644 index 0000000000..3e41250261 --- /dev/null +++ b/src/scvi/external/joint_embedding_scvi/__init__.py @@ -0,0 +1,4 @@ +from ._model import JointEmbeddingSCVI +from ._module import JointEmbeddingVAE + +__all__ = ["JointEmbeddingSCVI", "JointEmbeddingVAE"] diff --git a/src/scvi/external/joint_embedding_scvi/_model.py b/src/scvi/external/joint_embedding_scvi/_model.py new file mode 100644 index 0000000000..58cd994a5b --- /dev/null +++ b/src/scvi/external/joint_embedding_scvi/_model.py @@ -0,0 +1,157 @@ +"""Joint Embedding SCVI model.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from scvi.external.joint_embedding_scvi._module import JointEmbeddingVAE +from scvi.model._scvi import SCVI + +if TYPE_CHECKING: + from typing import Literal + + from anndata import AnnData + + +class JointEmbeddingSCVI(SCVI): + """SCVI with joint embedding loss using binomial thinning and CCO :cite:p:`Svensson26`. + + This model extends the standard SCVI with a cross-correlation objective (CCO) + loss that encourages the embedding of a thinned view to match the embedding + of the original data. This promotes robustness to count dropout/noise. + + Thinning probabilities are dynamically sampled per cell to produce target + library sizes that are log-uniform between min_library_size and the observed + library size. This matches realistic library size variation in single-cell data. + + Parameters + ---------- + adata + AnnData object that has been registered via :meth:`~scvi.model.SCVI.setup_anndata`. + registry + Registry of the datamodule used to train the model, passed through when training + without an :class:`~anndata.AnnData` object (see :class:`~scvi.model.SCVI`). + joint_embedding_weight + Weight for the CCO loss. Default is 100.0. + lambda_off_diag + Off-diagonal penalty in CCO loss. Default is 0.01. + min_library_size + Minimum target library size for thinning. Default is 10. + Thinned library sizes are sampled log-uniformly between this + value and the observed library size. + reconstruction_weight + Weight for reconstruction loss. Default is 1.0. + Set to 0.0 for pure self-supervised training with only CCO loss. + variance_weight + Weight for variance regularization loss (VICReg-style). Default is 0.0. + Set to positive value (e.g., 1.0) to prevent dimension collapse + in self-supervised training. + use_joint_embedding + Whether to use joint embedding loss. Default is True. + Set to False to train as standard SCVI. + 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-label'`` - dispersion can differ between different labels + * ``'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 + use_observed_lib_size + If ``True``, use the observed library size for RNA as the scaling factor in the mean of the + conditional distribution. + latent_distribution + One of: + + * ``'normal'`` - Normal distribution + * ``'ln'`` - Logistic normal distribution (Normal(0, I) transformed by softmax) + **kwargs + Additional keyword arguments for :class:`~scvi.external.JointEmbeddingVAE`. + + Examples + -------- + >>> adata = anndata.read_h5ad(path_to_anndata) + >>> scvi.external.JointEmbeddingSCVI.setup_anndata(adata, batch_key="batch") + >>> model = scvi.external.JointEmbeddingSCVI( + ... adata, + ... joint_embedding_weight=1.0, + ... lambda_off_diag=0.01, + ... ) + >>> model.train() + >>> latent = model.get_latent_representation() + + See Also + -------- + :class:`~scvi.model.SCVI` + :class:`~scvi.external.JointEmbeddingVAE` + """ + + _module_cls = JointEmbeddingVAE + + def __init__( + self, + adata: AnnData | None = None, + registry: dict | None = None, + joint_embedding_weight: float = 100.0, + lambda_off_diag: float = 0.01, + min_library_size: float = 10.0, + reconstruction_weight: float = 1.0, + variance_weight: float = 0.0, + use_joint_embedding: bool = True, + n_hidden: int = 128, + n_latent: int = 10, + n_layers: int = 1, + dropout_rate: float = 0.1, + dispersion: Literal["gene", "gene-batch", "gene-label", "gene-cell"] = "gene", + gene_likelihood: Literal["zinb", "nb", "poisson", "normal"] = "zinb", + use_observed_lib_size: bool = True, + latent_distribution: Literal["normal", "ln"] = "normal", + **kwargs, + ): + # Pass joint embedding params through kwargs to module + super().__init__( + adata=adata, + registry=registry, + n_hidden=n_hidden, + n_latent=n_latent, + n_layers=n_layers, + dropout_rate=dropout_rate, + dispersion=dispersion, + gene_likelihood=gene_likelihood, + use_observed_lib_size=use_observed_lib_size, + latent_distribution=latent_distribution, + joint_embedding_weight=joint_embedding_weight, + lambda_off_diag=lambda_off_diag, + min_library_size=min_library_size, + reconstruction_weight=reconstruction_weight, + variance_weight=variance_weight, + use_joint_embedding=use_joint_embedding, + **kwargs, + ) + + # Update model summary string + self._model_summary_string = ( + "JointEmbeddingSCVI 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}, latent_distribution: {latent_distribution}, " + f"joint_embedding_weight: {joint_embedding_weight}, " + f"lambda_off_diag: {lambda_off_diag}, min_library_size: {min_library_size}, " + f"reconstruction_weight: {reconstruction_weight}, " + f"variance_weight: {variance_weight}, " + f"use_joint_embedding: {use_joint_embedding}." + ) diff --git a/src/scvi/external/joint_embedding_scvi/_module.py b/src/scvi/external/joint_embedding_scvi/_module.py new file mode 100644 index 0000000000..d2cf995a6d --- /dev/null +++ b/src/scvi/external/joint_embedding_scvi/_module.py @@ -0,0 +1,222 @@ +"""Joint Embedding VAE module.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from scvi import REGISTRY_KEYS +from scvi.external.joint_embedding_scvi._utils import ( + binomial_split, + cross_correlation_loss, + sample_thinning_probs, + variance_loss, +) +from scvi.module._constants import MODULE_KEYS +from scvi.module._vae import VAE +from scvi.module.base import LossOutput + +if TYPE_CHECKING: + from torch.distributions import Distribution + + +class JointEmbeddingVAE(VAE): + """VAE with joint embedding loss using binomial thinning and CCO :cite:p:`Svensson26`. + + This module extends the standard VAE with a cross-correlation objective (CCO) + loss that encourages the embedding of a thinned view to match the embedding + of the original data. This promotes robustness to count dropout/noise. + + Thinning probabilities are dynamically sampled per cell to produce target + library sizes that are log-uniform between min_library_size and the observed + library size. This matches realistic library size variation in single-cell data. + + Parameters + ---------- + n_input + Number of input features. + joint_embedding_weight + Weight for the CCO loss. Default is 100.0. + lambda_off_diag + Off-diagonal penalty in CCO loss. Default is 0.01. + min_library_size + Minimum target library size for thinning. Default is 10. + Thinned library sizes are sampled log-uniformly between this + value and the observed library size. + reconstruction_weight + Weight for reconstruction loss. Default is 1.0. + Set to 0.0 for pure self-supervised training with only CCO loss. + variance_weight + Weight for variance regularization loss (VICReg-style). Default is 0.0. + Set to positive value (e.g., 1.0) to prevent dimension collapse + in self-supervised training. + use_joint_embedding + Whether to use joint embedding loss. Default is True. + Set to False to train as standard SCVI. + **kwargs + Additional keyword arguments passed to :class:`~scvi.module.VAE`. + + See Also + -------- + :class:`~scvi.module.VAE` + """ + + def __init__( + self, + n_input: int, + *args, + joint_embedding_weight: float = 100.0, + lambda_off_diag: float = 0.01, + min_library_size: float = 10.0, + reconstruction_weight: float = 1.0, + variance_weight: float = 0.0, + use_joint_embedding: bool = True, + **kwargs, + ): + super().__init__(n_input, *args, **kwargs) + self.joint_embedding_weight = joint_embedding_weight + self.lambda_off_diag = lambda_off_diag + self.min_library_size = min_library_size + self.reconstruction_weight = reconstruction_weight + self.variance_weight = variance_weight + self.use_joint_embedding = use_joint_embedding + + # Binomial thinning requires integer count data. The Gaussian likelihood + # operates on real-valued (e.g. log-normalized) input, which thinning would + # silently corrupt, so the joint embedding objective is incompatible with it. + if self.use_joint_embedding and self.gene_likelihood == "normal": + raise ValueError( + "JointEmbeddingVAE uses binomial thinning, which requires count data, so " + "gene_likelihood='normal' is incompatible with use_joint_embedding=True. " + "Use a count likelihood ('nb', 'zinb', 'poisson') or set " + "use_joint_embedding=False." + ) + + 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, + joint_embedding_weight: float | None = None, + reconstruction_weight: float | None = None, + ) -> LossOutput: + """Compute the loss including optional joint embedding CCO loss. + + Parameters + ---------- + tensors + Dictionary of input tensors. + inference_outputs + Dictionary of inference outputs. + generative_outputs + Dictionary of generative outputs. + kl_weight + Weight for KL divergence term. + joint_embedding_weight + Optional override for joint embedding weight. + reconstruction_weight + Optional override for reconstruction weight. + Set to 0.0 for pure self-supervised training. + + Returns + ------- + LossOutput + Loss output with total loss, reconstruction loss, KL terms, + and CCO loss in extra_metrics. + """ + # Get the reconstruction weight + recon_weight = ( + reconstruction_weight + if reconstruction_weight is not None + else self.reconstruction_weight + ) + + # Get base loss components from parent (we'll reweight them) + base_loss = super().loss(tensors, inference_outputs, generative_outputs, kl_weight) + + # Recompute total loss with reconstruction weight + # base_loss.loss = mean(reconst_loss + kl_weight * kl_z + kl_l) + # We want: mean(recon_weight * reconst_loss + kl_weight * kl_z + kl_l) + # Note: reconstruction_loss is stored as a dict in LossOutput, use dict_sum to get tensor + reconst_loss = base_loss.dict_sum(base_loss.reconstruction_loss) + kl_local = base_loss.kl_local + kl_z = kl_local[MODULE_KEYS.KL_Z_KEY] + kl_l = kl_local[MODULE_KEYS.KL_L_KEY] + + weighted_recon = recon_weight * reconst_loss + weighted_kl = kl_weight * kl_z + kl_l + elbo_loss = torch.mean(weighted_recon + weighted_kl) + + # Joint embedding needs at least two cells to estimate per-dimension + # statistics (correlation / variance); the default splitter can emit a + # singleton final minibatch, so fall back to the reweighted ELBO there. + z_full = inference_outputs[MODULE_KEYS.Z_KEY] + too_small = z_full.shape[0] < 2 + + # If joint embedding is disabled, not training, or the batch is too small, + # return reweighted ELBO + if not self.use_joint_embedding or not self.training or too_small: + return LossOutput( + loss=elbo_loss, + reconstruction_loss=reconst_loss, + kl_local=kl_local, + extra_metrics=base_loss.extra_metrics, + ) + + # Get input data + x = tensors[REGISTRY_KEYS.X_KEY] + + # Sample per-cell thinning probabilities for realistic library size variation + p = sample_thinning_probs(x, min_library_size=self.min_library_size) + + # Create thinned view via binomial thinning with per-cell probabilities + x_thin, _ = binomial_split(x, p=p) + + # Get covariates for inference + batch_index = tensors[REGISTRY_KEYS.BATCH_KEY] + cont_covs = tensors.get(REGISTRY_KEYS.CONT_COVS_KEY) + cat_covs = tensors.get(REGISTRY_KEYS.CAT_COVS_KEY) + + # Encode thinned view + inference_thin = self._regular_inference(x_thin, batch_index, cont_covs, cat_covs) + z_thin = inference_thin[MODULE_KEYS.Z_KEY] + + # Compute CCO loss between thinned and full embeddings (z_full computed above) + cco_components = cross_correlation_loss( + z_thin, z_full, self.lambda_off_diag, return_components=True + ) + cco_loss = cco_components["total"] + + # Compute variance loss for both views (VICReg-style anti-collapse) + var_loss_full = variance_loss(z_full) + var_loss_thin = variance_loss(z_thin) + var_loss = (var_loss_full + var_loss_thin) / 2 + + # Apply weights + je_weight = ( + joint_embedding_weight + if joint_embedding_weight is not None + else self.joint_embedding_weight + ) + var_weight = self.variance_weight + + # Total loss: weighted ELBO + CCO + variance + total_loss = elbo_loss + je_weight * cco_loss + var_weight * var_loss + + # Combine extra metrics - track CCO components separately + extra_metrics = { + **base_loss.extra_metrics, + "cco_loss": cco_loss.detach(), + "cco_invariance": cco_components["invariance"].detach(), # self-supervision + "cco_redundancy": cco_components["redundancy"].detach(), # anti-collapse + "variance_loss": var_loss.detach(), # VICReg-style anti-collapse + } + + return LossOutput( + loss=total_loss, + reconstruction_loss=reconst_loss, + kl_local=kl_local, + extra_metrics=extra_metrics, + ) diff --git a/src/scvi/external/joint_embedding_scvi/_utils.py b/src/scvi/external/joint_embedding_scvi/_utils.py new file mode 100644 index 0000000000..5badf60928 --- /dev/null +++ b/src/scvi/external/joint_embedding_scvi/_utils.py @@ -0,0 +1,159 @@ +"""Utility functions for joint embedding loss.""" + +from __future__ import annotations + +import torch + + +def variance_loss(z: torch.Tensor, target_std: float = 1.0) -> torch.Tensor: + """Compute variance regularization loss (VICReg-style). + + Penalizes dimensions with standard deviation below target_std. + This prevents collapse where all samples map to similar embeddings. + + Parameters + ---------- + z + Embedding tensor of shape (batch_size, n_latent). + target_std + Target standard deviation for each dimension. Default is 1.0. + + Returns + ------- + torch.Tensor + Scalar loss value (mean of hinge losses across dimensions). + """ + # Use the biased estimator so the loss is well-defined (0, not NaN) for + # single-sample minibatches, which the default data splitter can emit. + std = z.std(dim=0, unbiased=False) + return torch.relu(target_std - std).mean() + + +def binomial_split( + x: torch.Tensor, p: float | torch.Tensor = 0.5 +) -> tuple[torch.Tensor, torch.Tensor]: + """Split counts via binomial thinning. + + Parameters + ---------- + x + Input count tensor of shape (batch_size, n_features). + p + Probability for binomial split. Can be a scalar or a tensor of shape + (batch_size, 1) for per-cell probabilities. Default is 0.5. + + Returns + ------- + tuple + Two tensors (x1, x2) where x1 ~ Binomial(x, p) and x2 = x - x1. + Guarantees x1 + x2 = x and both are non-negative. + + Notes + ----- + Binomial thinning requires non-negative integer counts. Inputs are rounded + and clamped to ``>= 0`` so that data stored as integer-valued floats (the + usual single-cell convention) is handled, and non-count inputs do not raise + mid-training. + """ + counts = x.round().clamp(min=0) + x1 = torch.distributions.Binomial(total_count=counts, probs=p).sample() + x2 = counts - x1 + return x1, x2 + + +def sample_thinning_probs(x: torch.Tensor, min_library_size: float = 10.0) -> torch.Tensor: + """Sample per-cell thinning probabilities for realistic library size variation. + + Samples target library sizes from a log-uniform distribution between + min_library_size and the observed library size for each cell. This + matches the variation typically observed in real single-cell data. + + Parameters + ---------- + x + Input count tensor of shape (batch_size, n_features). + min_library_size + Minimum target library size. Default is 10. + + Returns + ------- + torch.Tensor + Per-cell thinning probabilities of shape (batch_size, 1). + """ + # Compute observed library sizes per cell + lib_size = x.sum(dim=1, keepdim=True) # (batch_size, 1) + + # Sample target library size from LogUniform(min_library_size, lib_size) + # log(L_target) ~ Uniform(log(min_lib), log(lib_size)) + log_min = torch.log(torch.tensor(min_library_size, device=x.device, dtype=x.dtype)) + log_lib = torch.log(lib_size.clamp(min=min_library_size + 1e-6)) + + # Sample uniformly in log space + log_target = log_min + torch.rand_like(lib_size) * (log_lib - log_min) + target_lib_size = torch.exp(log_target) + + # Compute thinning probability: p = target_lib_size / lib_size + # Clamp to [0, 1] for safety (handles edge cases where lib_size <= min_library_size) + p = (target_lib_size / lib_size.clamp(min=1e-6)).clamp(0.0, 1.0) + + return p + + +def cross_correlation_loss( + z1: torch.Tensor, + z2: torch.Tensor, + lambda_off_diag: float = 0.01, + return_components: bool = False, +) -> torch.Tensor | dict[str, torch.Tensor]: + """Compute cross-correlation objective (CCO) loss. + + This loss has two components: + - Diagonal (invariance): encourages C_ii = 1, meaning corresponding dimensions + of the two views should be correlated. This is the self-supervision signal. + - Off-diagonal (redundancy reduction): encourages C_ij = 0 for i != j, meaning + different latent dimensions should be decorrelated. This prevents mode collapse. + + Parameters + ---------- + z1 + First embedding tensor of shape (batch_size, n_latent). + z2 + Second embedding tensor of shape (batch_size, n_latent). + lambda_off_diag + Weight for off-diagonal penalty. Default is 0.01. + return_components + If True, return a dict with individual components. Default is False. + + Returns + ------- + torch.Tensor or dict + If return_components is False, returns scalar loss value. + If True, returns dict with keys: 'total', 'invariance', 'redundancy'. + """ + # Normalize (z-score across batch dimension). The biased std keeps this + # finite for single-sample minibatches and makes C a proper correlation matrix. + z1_norm = (z1 - z1.mean(0)) / (z1.std(0, unbiased=False) + 1e-8) + z2_norm = (z2 - z2.mean(0)) / (z2.std(0, unbiased=False) + 1e-8) + + # Cross-correlation matrix: C[i,j] = corr(z1[:,i], z2[:,j]) + batch_size = z1.shape[0] + C = z1_norm.T @ z2_norm / batch_size + + # Diagonal loss (invariance): (1 - C_ii)^2 - self-supervision signal + invariance_loss = ((1 - C.diagonal()) ** 2).sum() + + # Off-diagonal loss (redundancy reduction): sum(C_ij^2 for i != j) - anti-collapse + # Clone C before modifying diagonal to avoid in-place operation issues + C_off_diag = C.clone() + C_off_diag.fill_diagonal_(0) + redundancy_loss = (C_off_diag**2).sum() + + total_loss = invariance_loss + lambda_off_diag * redundancy_loss + + if return_components: + return { + "total": total_loss, + "invariance": invariance_loss, + "redundancy": redundancy_loss, + } + return total_loss diff --git a/tests/external/joint_embedding_scvi/test_joint_embedding_scvi.py b/tests/external/joint_embedding_scvi/test_joint_embedding_scvi.py new file mode 100644 index 0000000000..9725047b3d --- /dev/null +++ b/tests/external/joint_embedding_scvi/test_joint_embedding_scvi.py @@ -0,0 +1,149 @@ +import numpy as np +import pytest +import scipy.sparse as sp + +from scvi.data import synthetic_iid +from scvi.external import JointEmbeddingSCVI + +CCO_METRICS = {"cco_loss", "cco_invariance", "cco_redundancy", "variance_loss"} + +N_LATENT = 5 + + +def _get_adata(sparse: bool = False): + adata = synthetic_iid() + if sparse: + adata.X = sp.csr_matrix(adata.X) + JointEmbeddingSCVI.setup_anndata(adata, batch_key="batch", labels_key="labels") + return adata + + +def test_joint_embedding_scvi_train_and_latent(): + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT) + model.train(1, train_size=0.5) + assert model.is_trained is True + latent = model.get_latent_representation() + assert latent.shape == (adata.n_obs, N_LATENT) + assert np.all(np.isfinite(latent)) + + +def test_joint_embedding_scvi_inference_methods(): + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT) + model.train(1, train_size=0.5) + + assert model.get_elbo().ndim == 0 + assert model.get_marginal_ll(n_mc_samples=3).ndim == 0 + assert model.get_reconstruction_error()["reconstruction_loss"].ndim == 0 + assert model.get_normalized_expression().shape == (adata.n_obs, adata.n_vars) + + +def test_joint_embedding_scvi_sparse(): + adata = _get_adata(sparse=True) + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT) + model.train(1, train_size=0.5) + assert model.is_trained is True + assert model.get_latent_representation().shape == (adata.n_obs, N_LATENT) + + +@pytest.mark.parametrize("gene_likelihood", ["zinb", "nb", "poisson"]) +def test_joint_embedding_scvi_gene_likelihood(gene_likelihood): + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT, gene_likelihood=gene_likelihood) + model.train(1, train_size=0.5) + assert model.is_trained is True + assert model.get_latent_representation().shape == (adata.n_obs, N_LATENT) + + +def test_joint_embedding_scvi_fallback(): + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT, use_joint_embedding=False) + model.train(1, train_size=0.5) + assert model.is_trained is True + assert model.get_latent_representation().shape == (adata.n_obs, N_LATENT) + + +def _one_batch(model, adata): + """Return one correctly-formatted minibatch of registered tensors.""" + dl = model._make_data_loader(adata=adata, indices=np.arange(adata.n_obs)) + return next(iter(dl)) + + +def test_joint_embedding_scvi_cco_branch_active_in_training(): + """The CCO loss branch only runs in training mode and only when enabled.""" + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT) + model.train(1, train_size=0.5) + tensors = _one_batch(model, adata) + module = model.module + + # Training mode + enabled: CCO metrics present and finite + module.train() + _, _, losses = module(tensors) + assert CCO_METRICS.issubset(losses.extra_metrics) + assert all(np.isfinite(losses.extra_metrics[k].item()) for k in CCO_METRICS) + + # Eval mode: CCO branch is skipped (ELBO only) + module.eval() + _, _, losses_eval = module(tensors) + assert not CCO_METRICS.intersection(losses_eval.extra_metrics) + + +def test_joint_embedding_scvi_cco_branch_disabled(): + """With use_joint_embedding=False, the CCO branch never runs, even in training.""" + adata = _get_adata() + model = JointEmbeddingSCVI(adata, n_latent=N_LATENT, use_joint_embedding=False) + model.train(1, train_size=0.5) + tensors = _one_batch(model, adata) + model.module.train() + _, _, losses = model.module(tensors) + assert not CCO_METRICS.intersection(losses.extra_metrics) + + +def test_joint_embedding_scvi_hyperparameters_in_summary(): + adata = _get_adata() + model = JointEmbeddingSCVI( + adata, + n_latent=N_LATENT, + joint_embedding_weight=2.0, + lambda_off_diag=0.05, + min_library_size=20.0, + reconstruction_weight=0.5, + variance_weight=1.0, + ) + summary = model._model_summary_string + assert "joint_embedding_weight: 2.0" in summary + assert "lambda_off_diag: 0.05" in summary + assert "min_library_size: 20.0" in summary + assert "reconstruction_weight: 0.5" in summary + assert "variance_weight: 1.0" in summary + + +def test_joint_embedding_scvi_save_load(tmp_path): + adata = _get_adata() + model = JointEmbeddingSCVI( + adata, + n_latent=N_LATENT, + joint_embedding_weight=2.0, + lambda_off_diag=0.05, + min_library_size=20.0, + reconstruction_weight=0.5, + variance_weight=1.0, + ) + model.train(1, train_size=0.5) + latent_before = model.get_latent_representation() + + save_dir = str(tmp_path) + model.save(save_dir, overwrite=True, save_anndata=True) + + loaded = JointEmbeddingSCVI.load(save_dir) + latent_after = loaded.get_latent_representation() + + assert np.allclose(latent_before, latent_after) + # New init params must round-trip through save/load + assert loaded.module.joint_embedding_weight == 2.0 + assert loaded.module.lambda_off_diag == 0.05 + assert loaded.module.min_library_size == 20.0 + assert loaded.module.reconstruction_weight == 0.5 + assert loaded.module.variance_weight == 1.0 diff --git a/tests/external/joint_embedding_scvi/test_joint_embedding_utils.py b/tests/external/joint_embedding_scvi/test_joint_embedding_utils.py new file mode 100644 index 0000000000..0f02c107af --- /dev/null +++ b/tests/external/joint_embedding_scvi/test_joint_embedding_utils.py @@ -0,0 +1,61 @@ +import torch + +from scvi.external.joint_embedding_scvi._utils import ( + binomial_split, + cross_correlation_loss, + sample_thinning_probs, + variance_loss, +) + + +def test_binomial_split_conserves_counts(): + x = torch.randint(0, 50, (32, 10)).float() + x1, x2 = binomial_split(x, p=0.5) + assert torch.all(x1 >= 0) + assert torch.all(x2 >= 0) + assert torch.allclose(x1 + x2, x) + + +def test_binomial_split_per_cell_probs(): + x = torch.randint(0, 50, (8, 5)).float() + p = torch.rand(8, 1) + x1, x2 = binomial_split(x, p=p) + assert torch.allclose(x1 + x2, x) + + +def test_binomial_split_rounds_non_integer_counts(): + # Non-integer counts must not raise; they are rounded before thinning. + x = torch.full((4, 3), 3.7) + x1, x2 = binomial_split(x, p=0.5) + assert torch.allclose(x1 + x2, torch.full((4, 3), 4.0)) + + +def test_sample_thinning_probs_range_and_shape(): + x = torch.randint(0, 100, (16, 20)).float() + p = sample_thinning_probs(x, min_library_size=10.0) + assert p.shape == (16, 1) + assert torch.all(p >= 0.0) + assert torch.all(p <= 1.0) + + +def test_cross_correlation_loss_identity_components(): + z = torch.randn(64, 8) + comp = cross_correlation_loss(z, z, lambda_off_diag=0.01, return_components=True) + # z vs itself -> near-perfect invariance (diagonal ~1) + assert comp["invariance"].item() < 1e-3 + # total decomposes exactly into invariance + lambda * redundancy + assert torch.allclose(comp["total"], comp["invariance"] + 0.01 * comp["redundancy"]) + + +def test_variance_loss_collapsed_vs_spread(): + collapsed = torch.zeros(32, 5) + spread = torch.randn(2000, 5) * 2.0 + assert variance_loss(collapsed, target_std=1.0).item() > 0.9 + assert variance_loss(spread, target_std=1.0).item() < 0.1 + + +def test_losses_finite_on_single_sample_batch(): + # Single-row minibatches must stay finite (biased std), not NaN. + z = torch.randn(1, 5) + assert torch.isfinite(variance_loss(z)) + assert torch.isfinite(cross_correlation_loss(z, z))