Add DRVI (#3866)

## Add DRVI

This PR adds **DRVI** as `scvi.external.DRVI`.

This is a rewrite of #3699, reusing scvi-tools functionality and
removing extras as much as possible.

### Design:

`DRVI` now subclasses `scvi.model.SCVI` directly and only changes what
is actually different:

- Decoder is replaced by a splitting + additive decoder + pooling
(`DecoderDRVI` extends `DecoderSCVI`, `SplitFCLayers` extends
`FCLayers`). Disentanglement is induced in the decoder — the latent is
reshaped into `n_split_latent` independent groups, each decoded
separately, then aggregated.
- Adds log-space likelihood for negative binomial.
- Adds latent stats and interpretability mixins to interpret latent
dimensions (plotting functions are removed).
- Reuses everything else from SCVI unchanged: encoder, prior,
covariate/batch handling (including `batch_representation="embedding"`),
RNASeqMixin, size factors, minified mode, scArches transfer, training
via a datamodule, and `setup_anndata`.

To enable the decoder swap without forking `FCLayers`,
`scvi.nn.FCLayers` gained a few small overridable seams (`_build_layer`,
`_apply_layer`, `_apply_batch_norm`). These are behavior-preserving for
existing models (SCVI, ...) —`SplitFCLayers` overrides them to use
per-split weights.

### Likelihoods

In addition to scvi's `nb`, `zinb` and `poisson`, DRVI adds log-space
nb:

- `pnb` (default of DRVI) — parametrized NB, same mean as `nb` but
evaluated in log space via the new
`LogNegativeBinomial` distribution (should be numerically more stable
but equivalent to scvi's `NegativeBinomial`).
- `normal` / `normal_unit_var` — Gaussian with the mean modeled directly
(no library/softmax), for
  continuous / log-normalized input.

### Interpretability

Because the decoder is additive over splits, each latent dimension's
contribution to every gene can be measured directly. The included
methods — `set_latent_dimension_stats`,
`calculate_interpretability_scores`, `get_interpretability_scores` —
quantify these contributions and extract top genes per dimension.

### Dependencies

Adds `stacked-linear>=0.2.0` (used for the parallel per-split linear
layers). That package requires no dependencies other than PyTorch. Also,
the same logic can be included as a snippet instead.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: ori-kron-wis <ori.kronfeld@weizmann.ac.il>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

CHANGELOG.md

@@ -9,11 +9,14 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo
 
 #### Added
 
-- Add scvi-tools MCP package that gives any MCP-compatible LLM access to scvi-tools knowledge
+- Add [scvi-tools MCP](https://scvi-tools-mcp.readthedocs.io/en/latest/index.html) package
+    that gives any MCP-compatible LLM access to scvi-tools knowledge.
 - Add {class}`~scvi.dataloaders.AnnbatchDataModule` for out-of-core dataloading via `annbatch`,
     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.DRVI` for unsupervised disentangled representation learning of
+    single-cell omics, {pr}`3866`.
 - 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`.

docs/api/developer.md

@@ -197,6 +197,7 @@ Module classes in the external API with respective generative and inference proc
    external.scviva.NicheLossOutput
    external.sysvi.SysVAE
    external.diagvi.DIAGVAE
+   external.drvi.DRVIModule
    external.JointEmbeddingVAE
 ```
 

docs/api/user.md

@@ -67,6 +67,7 @@ import scvi
    external.SysVI
    external.SCVIVA
    external.DIAGVI
+   external.DRVI
    external.JointEmbeddingSCVI
    external.Tangram
 ```

docs/references.bib

@@ -387,6 +387,16 @@ @article{Weinberger26
   publisher={Nature Publishing Group}
 }
 
+@article{Moinfar2024,
+  title = {Disentangling cellular heterogeneity into interpretable biological factors through structured latent representations},
+  author = {Moinfar, Amir Ali and Theis, Fabian J},
+  doi = {10.1101/2024.11.06.622266},
+  year = {2024},
+  month = nov,
+  journal = {bioRxiv},
+  publisher = {Cold Spring Harbor Laboratory}
+}
+
 @article{Svensson26,
   title={Improving SCVI for low-count cells through self-supervised augmentation},
   author={Valentine Svensson},

docs/tutorials/index_r.md

@@ -19,7 +19,7 @@ Perform basic Python operations in an R environment
 ```
 
 ```{customcard}
-:path: notebooks/quick_start/api_overview_in_R
+:path: notebooks/r/api_overview_in_R
 :tags: Analysis, Differential-comparison, Dimensionality-reduction, Removal-of-variance, Dev, Integration
 
 Go through the typical steps of an scvi-tools workflow in R

docs/tutorials/index_scrna.md

@@ -16,9 +16,10 @@ notebooks/scrna/linear_decoder
 notebooks/scrna/amortized_lda
 notebooks/scrna/AutoZI_tutorial
 notebooks/scrna/sysVI
+notebooks/scrna/JointEmbeddingSCVI_tutorial
+notebooks/scrna/DRVI_pipeline
 notebooks/scrna/decipher_tutorial
 notebooks/scrna/velovi
-notebooks/scrna/JointEmbeddingSCVI_tutorial
 notebooks/scrna/Tahoe100_mrVI
 ```
 
@@ -36,12 +37,6 @@ Perform integration of multiple scRNA-seq datasets both with and without cell ty
 Analyze multi-sample scRNA-seq data with MrVI
 ```
 
-```{customcard}
-:path: notebooks/scrna/MrVI_tutorial_torch
-:tags: Analysis, Differential-comparison, Dimensionality-reduction, Removal-of-variance
-
-Analyze multi-sample scRNA-seq data with MrVI in PyTorch
-```
 
 ```{customcard}
 :path: notebooks/scrna/scanvi_fix
@@ -113,6 +108,13 @@ Use the AutoZI model to enable gene-specific treatment of zero-inflation
 Integrate scRNA-seq datasets with substantial batch effects.
 ```
 
+```{customcard}
+:path: notebooks/scrna/DRVI_pipeline
+:tags: Analysis, Dimensionality-reduction, Interpretability, Factor Analysis
+
+Learn an interpretable, disentangled representation of scRNA-seq data with DRVI and link latent dimensions to genes.
+```
+
 ```{customcard}
 :path: notebooks/scrna/decipher_tutorial
 :tags: Analysis, Integration

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/drvi`
+     - Unsupervised deep disentangled representation learning of single-cell omics
+     - :cite:p:`Moinfar2024`
    * - :doc:`/user_guide/models/jointembeddingscvi`
      - Improving SCVI for low-count cells through self-supervised augmentation
      - :cite:p:`Svensson26`

docs/user_guide/models/drvi.md

@@ -0,0 +1,83 @@
+# DRVI
+
+**DRVI** {cite:p}`Moinfar2024` (Disentangled Representation Variational Inference,
+Python class {class}`~scvi.external.DRVI`)
+is an unsupervised deep generative model that learns an **interpretable, disentangled**
+latent representation of single-cell omics data.
+
+The advantages of DRVI are:
+
+- Disentanglement: DRVI disentangled cell identity, signaling pathways, stress responses, developmental trajectories, and perturbation effects into distinct factors.
+- Interpretability: Each factor (latent dimension) can be linked back to genes via the built-in interpretability analyses.
+- Built on scVI: {class}`~scvi.external.DRVI` subclasses {class}`~scvi.model.SCVI` and only swaps
+  the decoder, so it inherits the full SCVI interface — covariates, batch embedding
+  (`batch_representation="embedding"`), size factors, minified mode, scArches transfer learning,
+  out-of-core training via a datamodule, and the standard likelihoods (plus DRVI's log-space `pnb`).
+
+
+```{topic} Tutorials:
+
+-   {doc}`/tutorials/notebooks/scrna/DRVI_pipeline`
+```
+
+
+## Method background
+
+DRVI is a variational autoencoder (VAE): the integrated, disentangled representation corresponds to
+the latent-space embedding of the cells. It shares the encoder, prior (`N(0, I)`) and reconstruction
+likelihoods with {class}`~scvi.module.VAE`. Disentanglement is induced **in the decoder** rather
+than through the prior.
+
+### Disentanglement via an additive split decoder
+
+The latent vector of dimension `K` is mapped into `n_split_latent` independent groups ("splits").
+Each split is decoded **independently** and the per-split outputs are aggregated to form the final
+reconstruction parameters. Because every split can only influence the output through this shared
+aggregation, the model is encouraged to allocate independent factors of variation to different
+latent dimensions.
+
+Two aggregations over the split dimension are available:
+
+- `logsumexp` (default): aggregation in log space, which makes the decoder additive in the rate.
+- `mean`: average of the per-split parameters.
+
+Two latent-to-split mappings are available (`n_latent` should be divisible by `n_split_latent`):
+
+- `split_map` (default): the latent is reshaped into `n_split_latent` chunks and each chunk is passed through a learned per-split linear projection.
+- `split_mask`: latent is split into chunks and padded with zeros so, each split `i` keeps only its own chunk.
+
+
+### Likelihoods modeled in log space
+
+DRVI models the generative parameters in **log space**, which is what makes the `logsumexp`
+aggregation a genuine additive decoder (summing the per-split contributions of the rate). In
+addition to scvi's `nb`, `zinb` and `poisson`, DRVI adds:
+
+- `pnb` — *parametrized negative binomial*: the same mean as `nb` (``library * softmax``) but
+  parametrized and evaluated in log space via {class}`~scvi.external.drvi.LogNegativeBinomial`,
+  which is numerically stable and composes the per-split log contributions directly.
+- `normal` / `normal_unit_var` — a Gaussian whose mean is modeled **directly** (no library/softmax),
+  with the per-gene variance modeled in log space (fixed to one for `normal_unit_var`); intended for
+  continuous, e.g. log-normalized, input.
+
+### Interpretability
+
+Because the decoder is additive over splits, the contribution of each split to the reconstructed
+expression of every gene can be measured directly. DRVI exposes per-split parameters in an
+inspection mode and provides methods to quantify these contributions both in-distribution (over the
+observed cells) and out-of-distribution (by traversing each latent dimension), and to extract the
+top genes per latent dimension. See {meth}`~scvi.external.DRVI.set_latent_dimension_stats`,
+{meth}`~scvi.external.DRVI.calculate_interpretability_scores` and
+{meth}`~scvi.external.DRVI.get_interpretability_scores`.
+
+## Scope: core model only
+
+Only the **core DRVI model** and its **core interpretability methods** are included in scvi-tools
+(the model, the additive decoder and log-sum-exp pooling, the log-space likelihoods, and the interpretability methods.).
+
+The **utility functions for latent-space visualization and analysis are not** part of scvi-tools.
+For those, install the full [`drvi-py`](https://github.com/theislab/drvi) package:
+
+```bash
+pip install drvi-py
+```

docs/user_guide/models/index.md

@@ -11,6 +11,7 @@ cytovi
 decipher
 destvi
 diagvi
+drvi
 gimvi
 jointembeddingscvi
 linearscvi

pyproject.toml

@@ -47,6 +47,7 @@ dependencies = [
     "scikit-learn",
     "scipy",
     "sparse>=0.14.0",
+    "stacked-linear>=0.2.0",
     "tensorboard",
     "torch",
     "torchmetrics",