diff --git a/CHANGELOG.md b/CHANGELOG.md index a45a663e83..269a05a1f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,11 @@ to [Semantic Versioning]. The full commit history is available in the [commit lo #### Removed +- Remove {class}`scvi.external.DestVI`, {class}`scvi.external.GimVI`, {class}`scvi.external.ResolVI`, + {class}`scvi.external.SCVIVA`, {class}`scvi.external.DiagVI`, {class}`scvi.external.Stereoscope`, + and {class}`scvi.external.Tangram` spatial models, along with their associated tutorials and + documentation, {pr}`3908`. + ### 1.5.0 (2026-07-08) #### Added diff --git a/docs/api/developer.md b/docs/api/developer.md index d80eb6f72c..557f90f7b5 100644 --- a/docs/api/developer.md +++ b/docs/api/developer.md @@ -154,7 +154,6 @@ Existing module classes with respective generative and inference procedures. module.AutoZIVAE module.Classifier module.LDVAE - module.MRDeconv module.PEAKVAE module.MULTIVAE module.SCANVAE @@ -178,12 +177,9 @@ Module classes in the external API with respective generative and inference proc :template: class_no_inherited.rst :nosignatures: - external.gimvi.JVAE external.cytovi.CytoVAE external.cellassign.CellAssignModule external.contrastivevi.ContrastiveDataSplitter - external.stereoscope.RNADeconv - external.stereoscope.SpatialDeconv external.scbasset.ScBassetModule external.contrastivevi.ContrastiveVAE external.velovi.VELOVAE @@ -191,12 +187,8 @@ Module classes in the external API with respective generative and inference proc external.methylvi.METHYLVAE external.methylvi.METHYLANVAE external.decipher.DecipherPyroModule - external.resolvi.RESOLVAE external.totalanvi.TOTALANVAE - external.scviva.nicheVAE - external.scviva.NicheLossOutput external.sysvi.SysVAE - external.diagvi.DIAGVAE external.drvi.DRVIModule external.JointEmbeddingVAE ``` diff --git a/docs/api/user.md b/docs/api/user.md index c25177f33d..9677a15b9d 100644 --- a/docs/api/user.md +++ b/docs/api/user.md @@ -25,7 +25,6 @@ import scvi model.AUTOZI model.CondSCVI - model.DestVI model.LinearSCVI model.PEAKVI model.SCANVI @@ -49,9 +48,6 @@ import scvi external.CellAssign external.CYTOVI - external.GIMVI - external.RNAStereoscope - external.SpatialStereoscope external.SOLO external.SCAR external.SCBASSET @@ -63,13 +59,9 @@ import scvi external.METHYLANVI external.Decipher external.TOTALANVI - external.RESOLVI external.SysVI - external.SCVIVA - external.DIAGVI external.DRVI external.JointEmbeddingSCVI - external.Tangram ``` ## Data loading diff --git a/docs/installation.md b/docs/installation.md index b956389913..fc3a4f9977 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -90,7 +90,6 @@ It has many optional dependencies which expand its capabilities: - _parallel_ - for parallelization engine - _interpretability_ - for supervised models interpretability - _dataloaders_ - for custom dataloaders use -- _diagvi_ - for DiagVI model (requires torch_geometric and geomloss) - _mlflow_ - for MLflow support - _tests_ - in order to be able to perform tests - _editing_ - for code editing diff --git a/docs/tutorials/index.md b/docs/tutorials/index.md index b99da709ed..bdc5f8fd15 100644 --- a/docs/tutorials/index.md +++ b/docs/tutorials/index.md @@ -19,7 +19,6 @@ index_atac index_cytometry index_scbs index_multimodal -index_spatial index_hub index_use_cases index_custom_dl diff --git a/docs/tutorials/index_multimodal.md b/docs/tutorials/index_multimodal.md index dddf3f84f8..866d5af2a2 100644 --- a/docs/tutorials/index_multimodal.md +++ b/docs/tutorials/index_multimodal.md @@ -9,8 +9,6 @@ notebooks/multimodal/totalVI_reference_mapping notebooks/multimodal/cite_scrna_integration_w_totalVI notebooks/multimodal/MultiVI_tutorial notebooks/multimodal/totalanvi -notebooks/multimodal/DiagVI_spatial_proteomics.ipynb -notebooks/multimodal/DiagVI_spatial_transcriptomics.ipynb ``` ```{customcard} @@ -54,17 +52,3 @@ Go through the MultiVI workflow to perform joint analysis of paired and unpaired Use TotalANVI to perform semi-supervised analysis of CITE-seq data, leveraging partial cell type annotations for label prediction, protein imputation, and differential abundance ``` - -```{customcard} -:path: notebooks/multimodal/DiagVI_spatial_proteomics -:tags: Analysis, Integration, Modality-imputation, Dimensionality-reduction - -Perform integration of spatial proteomics and single-cell transcriptomics data with DiagVI -``` - -```{customcard} -:path: notebooks/multimodal/DiagVI_spatial_transcriptomics -:tags: Analysis, Integration, Modality-imputation, Dimensionality-reduction - -Perform integration of spatial and single-cell transcriptomics data with DiagVI -``` diff --git a/docs/tutorials/index_spatial.md b/docs/tutorials/index_spatial.md deleted file mode 100644 index c67af7294f..0000000000 --- a/docs/tutorials/index_spatial.md +++ /dev/null @@ -1,62 +0,0 @@ -# Spatial transcriptomics - -```{toctree} -:maxdepth: 1 - -notebooks/spatial/resolVI_tutorial -notebooks/spatial/scVIVA_tutorial -notebooks/spatial/DestVI_tutorial -notebooks/spatial/gimvi_tutorial -notebooks/spatial/tangram_scvi_tools -notebooks/spatial/stereoscope_heart_LV_tutorial -notebooks/spatial/cell2location_lymph_node_spatial_tutorial -``` - -```{customcard} -:path: notebooks/spatial/resolVI_tutorial -:tags: Analysis, Integration, Transfer-learning, Dimensionality-reduction - -Use resolVI to correct cellular-resolved spatial transcriptomics data. -``` - -```{customcard} -:path: notebooks/spatial/scVIVA_tutorial -:tags: Analysis, Integration, Dimensionality-reduction, Differential-comparison - -Stratify spatial transcriptomics data into niche-aware cell states with scVIVA -``` - -```{customcard} -:path: notebooks/spatial/DestVI_tutorial -:tags: Deconvolution, Modality-imputation, Differential-comparison - -Perform multi-resolution analysis on spatial transcriptomics data with DestVI -``` - -```{customcard} -:path: notebooks/spatial/gimvi_tutorial -:tags: Modality-imputation, Integration - -Use gimVI to impute missing genes in spatial data -``` - -```{customcard} -:path: notebooks/spatial/tangram_scvi_tools -:tags: Deconvolution, Analysis - -Use Tangram to map spatial transcriptomics data -``` - -```{customcard} -:path: notebooks/spatial/stereoscope_heart_LV_tutorial -:tags: Deconvolution, Integration - -Go through the Stereoscope workflow to map single-cell data -``` - -```{customcard} -:path: notebooks/spatial/cell2location_lymph_node_spatial_tutorial -:tags: Dev, Analysis, Integration - -Spatially map lymph node cell types using Cell2location -``` diff --git a/docs/user_guide/index.md b/docs/user_guide/index.md index a895900f67..76d5d9bac9 100644 --- a/docs/user_guide/index.md +++ b/docs/user_guide/index.md @@ -149,42 +149,9 @@ scvi-tools is composed of models that can perform one or many analysis tasks. In * - :doc:`/user_guide/models/multivi` - Integration of paired/unpaired multiome data, missing modality imputation, normalization of other cell- and sample-level confounding factors - :cite:p:`AshuachGabitto21` - * - :doc:`/user_guide/models/diagvi` - - Diagonal integration of unpaired multiome data, dimensionality reduction, cross-modality imputation, cell label transfer - - [] ``` -## Spatial transcriptomics analysis - -```{eval-rst} -.. list-table:: - :widths: 15 100 25 - :header-rows: 1 - - * - Model - - Tasks - - Reference - * - :doc:`/user_guide/models/destvi` - - Multi-resolution deconvolution, cell-type-specific gene expression imputation, comparative analysis - - :cite:p:`Lopez22` - * - :doc:`/user_guide/models/stereoscope` - - Deconvolution - - :cite:p:`Andersson20` - * - :doc:`/user_guide/models/gimvi` - - Imputation of missing spatial genes - - :cite:p:`Lopez19` - * - :doc:`/user_guide/models/tangram` - - Deconvolution, single cell spatial mapping - - :cite:p:`Biancalani21` - * - :doc:`/user_guide/models/resolvi` - - Generative model of single-cell resolved spatial transcriptomics - - :cite:p:`Ergen25` - * - :doc:`/user_guide/models/scviva` - - Representation of cells and their environments in spatial transcriptomics - - :cite:p:`Levy25` -``` - ## General purpose analysis ```{eval-rst} diff --git a/docs/user_guide/models/destvi.md b/docs/user_guide/models/destvi.md deleted file mode 100644 index 01ce1bcd96..0000000000 --- a/docs/user_guide/models/destvi.md +++ /dev/null @@ -1,258 +0,0 @@ -# DestVI - -**DestVI** {cite:p}`Lopez22` (Deconvolution of Spatial Transcriptomics profiles using Variational Inference; Python class {class}`~scvi.model.DestVI`) -posits a conditional generative model of spatial transcriptomics down to the sub-cell-type variation level which -can be used to explore the spatial organization of a tissue and understanding gene expression variation between tissues and conditions. - -The advantages of DestVI are: - -- Can stratify cells into discrete cell types and model continuous sub-cell-type variation. -- Scalable to very large datasets (>1 million cells). - -The limitations of DestVI include: - -- Effectively requires a GPU for fast inference. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorial: - -- {doc}`/tutorials/notebooks/spatial/DestVI_tutorial` -``` - -## Preliminaries - -DestVI requires training two models, the scLVM (single-cell latent variable model) and the -stLVM (spatial transcriptomic latent variable model). The scLVM takes in as input a scRNA-seq gene -expression matrix of UMI counts $X$ with $N$ cells and $G$ genes, along with -a vector of cell type labels $\vec{c}$. Subsequently, the stLVM takes in the trained scLVM, -along a spatial gene expression matrix $Y$ with $S$ spots and $G$ genes. -Optionally, the user can specify the number of components used for the mixture model underlying the -empirical prior. - -## Generative process - -### scLVM - -For cell $n$, the scLVM assumes observed discrete cell type labels $c_n$ and models -continuous covariates $\gamma_n$ of dimension $d$ to explain variation in gene expression within a cell type. -The scLVM posits that the observed UMI counts for cell $n$ are generated by the following process: - -```{math} -:nowrap: true - -\begin{align} - \gamma_n &\sim \textrm{Normal}(0, I) \tag{1} \\ - x_{ng} &\sim \textrm{NegativeBinomial}(l_nf^g(c_n, \gamma_n), p_g) \tag{2} \\ -\end{align} -``` - -where $l_n$ is the library size, $f$ is a two-layer neural network which outputs a $G$ -dimensional vector, and $p_g$ is the rate parameter of the negative binomial distribution for -a given gene $g$. - -:::{note} -We are using the standard rate-shape parametrization of the negative binomial here, rather than the mean-dispersion -parametrization used in {doc}`/user_guide/models/scvi`. This is to take advantage of the additive property of -negative binomial distributions sharing the same shape parameter. In this case, the rate parameter for the -negative binomial modeling the expression counts for a given gene and spot is equivalent to the sum of the rate -parameters for each contributing cell. -::: - -This generative process is also summarized in the following graphical model: - -:::{figure} figures/scLVM_graphical_model.svg -:align: center -:alt: scLVM graphical model -:class: img-fluid - -scLVM graphical model. -::: - -The latent variables for the scLVM, along with their description, are summarized in the following table: - -```{eval-rst} -.. list-table:: - :widths: 20 90 15 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable (if different) - * - :math:`\gamma_n \in \mathbb{R}^d` - - Low-dimensional representation of sub-cell-type covariates. - - ``z`` - * - :math:`p_g \in (0, \infty)` - - Rate parameter for the negative binomial distribution. - - ``px_r`` -``` - -### stLVM - -For the stLVM, we also model the expression counts with a $\mathrm{NegativeBinomial}$. However, -for spatial data, we assume that each spot $s$ has expression $x_s$ composed of a bulk of cell types, with -cell type abundance, $\beta_{sc}$, for each cell type $c$. We assume that for a given spot $s$ -and gene $g$, the observation is generated as a function of the latent variables $(c, \gamma_s^c)$ by the following process: - -```{math} -:nowrap: true - -\begin{align} - \gamma_x^c &\sim \sum_{k=1}^K m_{kc} q_\Phi(\gamma^c \mid u_{kc}, c) \tag{4} \\ - x_{sg} &\sim \mathrm{NegativeBinomial}(l_s\alpha_g\sum_{c=1}^{C}\beta_{sc}f^g(c, \gamma_s^c), p_g) \tag{5} \\ -\end{align} -``` - -Where $l_s$ is the library size and $\alpha_g$ is a correction term for -difference in experimental assays. Like the scLVM, $f$ is a decoder neural network, and -$p_g$ is the rate parameter for the negative binomial distribution. - -To avoid the latent variable $\gamma_s^c$ from incorporating variation attributed to experimental -assay differences, we assign an empirical prior informed by the scLVM and the corresponding -cells of the same cell type in the scRNA-seq dataset. To compute this function, we cluster the latent space of the -scLVM for each cell type to K cell-type-specific clusters. For each cluster we compute an empirical mean and variance. -Above, $\{u_{kc}\}_{k=1}^K$ designates the set of cell-type-specific subclusters from cell type $c$ in the scRNA-seq dataset, and -$q_\Phi$ designates the empirical normal distribution from the computed cluster mean and variance. -The loss is weighted by the probability of a random cell from this cell type to be in the respective cluster in the -scRNA-seq dataset (mixture probability, $m_{kc}$). -In literature, the prior is referred to as a VampPrior ("variational aggregated mixture of posteriors" prior) [^ref2]. -More can be read on this prior in the DestVI paper. - -Lastly, an additional latent variable, $\eta_g$, is incorporated into the aggregated cell expression profile -as a dummy cell type to represent gene-specific noise. The dummy cell type's expression profile is distributed -as $\epsilon_g := \mathrm{Softplus}(\eta_g)$ where $\eta_g \sim \mathrm{Normal}(0, 1)$. -Like the other cell types, there is an associated cell type abundance parameter $\beta_{sc}$ associated with $\eta$. -We suspect each spot to only contain a fraction of the different cell types. To increase the sparsity of the cell type -proportions, the stLVM supports L1 regularization on the cell types proportions $\beta_{sc}$. By default, this loss is -not used. - -This generative process is also summarized in the following graphical model: - -:::{figure} figures/stLVM_graphical_model.svg -:align: center -:alt: stLVM graphical model -:class: img-fluid - -stLVM graphical model. -::: - -The latent variables for the stLVM, along with their description, are summarized in the following table: - -```{eval-rst} -.. list-table:: - :widths: 20 90 15 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable (if different) - * - :math:`\beta_{sc} \in (0, \infty)` - - Spot-specific cell type abundance. - - ``v_ind`` - * - :math:`\gamma_s^c \in (-\infty, \infty)` - - Low-dimensional representation of sub-cell-type covariates for a given spot and cell type. - - ``gamma`` - * - :math:`\eta_g \in (0, \infty)` - - Gene-specific noise. - - ``eta`` - * - :math:`\alpha_g \in (0, \infty)` - - Correction term for technological differences. - - ``beta`` - * - :math:`p_g \in (0,\infty)` - - Rate parameter for the negative binomial distribution. - - ``px_o`` - -``` - -## Inference - -### scLVM - -DestVI uses variational inference and specifically auto-encoding variational bayes (see {doc}`/user_guide/background/variational_inference`) -to learn both the model parameters (the neural network params, rate params, etc.) and an approximate posterior distribution -for the scLVM. Like {class}`scvi.model.SCVI`, the underlying class used as the encoder for DestVI is {class}`~scvi.nn.Encoder`. - -### stLVM - -For the stLVM, DestVI infers point estimates for latent variables $\gamma^c, \alpha, \beta$ using a penalized -likelihood method. Beyond vanilla MAP inference, to regularize $\alpha$ a variance penalty is applied across all genes. -Additionally, rather than having just $C$ parameters per spot to denote the estimated cell type abundances per spot, the stLVM -has $dC$ parameters per spot as well to account for the latent space learned by the scLVM. - -The loss is defined as: - -```{math} -:nowrap: true - -\begin{align} - L(l, \alpha, \beta, f^g, \gamma, p, \eta) := &-\log p(X \mid l, \alpha, \beta, f^g, \gamma, p, \eta) - \lambda_{\eta} \log p(\eta) \\ - &+ \lambda_{\alpha} \mathrm{Var}(\alpha) - \log p(\gamma \mid \mathrm{VampPrior}) + \lambda_{\beta} \lVert \beta_{sc} \rVert_1 \tag{6} \\ -\end{align} -``` - -Where $\mathrm{Var}(\alpha)$ refers to the empirical variance of the parameters alpha across all genes. We used this as a practical form of regularization (a similar regularizer is used in the ZINB-WaVE model [^ref3]). - -$\lambda_{\beta}$ (`l1_reg` in code), $\lambda_{\eta}$ (`eta_reg` in code) and $\lambda_{\alpha}$ (`beta_reg` in code) are hyperparameters used to scale the loss term. Increasing $\lambda_{\beta}$ leads to increased sparsity of cell type proportions. Increasing $\lambda_{\alpha}$ leads to less model flexibility for technical variation between single cell and spatial sequencing dataset. Increasing $\lambda_{\eta}$ leads to more genes being explained by the dummy cell type (we recommend not changing the default value). -To avoid overfitting, DestVI amortizes inference using a neural network to parametrize the latent variables. -Via the `amortization` parameter of {class}`scvi.module.MRDeconv`, the user can specify which of -$\beta$ and $\gamma^c$ will be parametrized by the neural network. - -## Tasks - -### Cell type deconvolution - -Once the model is trained, one can retrieve the estimated cell type proportions in each spot using the method: - -``` ->>> proportions = st_model.get_proportions() ->>> st_adata.obsm["proportions"] = proportions -``` - -These proportions are computed by normalizing across all learned cell type abundances, $\beta_{sc}$, for a given spot $s$. -I.e., the estimated proportion of cell type $c$ for spot $s$ is $\frac{\beta_{sc}}{\sum_c \beta_{sc}}$. - -Subsequently for a given cell type, users can plot a heatmap of the cell type proportions spatially using scanpy with: - -``` ->>> import scanpy as sc ->>> st_adata.obs['B cells'] = st_adata.obsm['proportions']['B cells'] ->>> sc.pl.spatial(st_adata, color="B cells", spot_size=130) -``` - -### Intra cell type variation - -Users can retrieve the values of $\gamma$, the latent variables corresponding to the -modeled cell-type-specific continuous covariates with: - -``` ->>> gamma = st_model.get_gamma()["B cells"] ->>> st_adata.obsm["B_cells_gamma"] = gamma -``` - -### Cell-type-specific gene expression imputation - -Assuming the user has identified key gene modules that vary within a cell type of interest, they can -impute the spatial pattern of the cell-type-specific gene expression with: - -``` ->>> # Filter spots with low abundance. ->>> indices = np.where(st_adata.obsm["proportions"][ct_name].values > 0.03)[0] ->>> imputed_counts = st_model.get_scale_for_ct("Monocyte", indices=indices)[["Cxcl9", "Cxcl10", "Fcgr1"]] -``` - -### Comparative analysis between samples - -To perform differential expression across samples, one can apply a frequentist test by taking samples -from the parameters of the generative distribution predicted for each spot in question. - -### Utilities function - -To explore the results of the output of the stLVM, we published an utilities function covering functions -for automatic thresholding of cell type proportions, a spatial PCA analysis to find the main axis of variation -in spatial gene expression and the described frequentist test for differential expression. Further information -can be found on [destvi_utils](https://destvi-utils.readthedocs.io/en/latest/installation.html) - -[^ref2]: Jakub Tomczak, Max Welling (2018),_VAE with a VampPrior_, [Proceedings of Machine Learning Research](https://proceedings.mlr.press/v84/tomczak18a.html) -[^ref3]: Davide Risso, Fanny Perraudeau, Svetlana Gribkova, Sandrine Dudoit, Jean-Philippe Vert (2018). _A general and flexible method for signal extraction from single-cell RNA-seq data_, [Nature Communications] (https://www.nature.com/articles/s41467-017-02554-5) diff --git a/docs/user_guide/models/diagvi.md b/docs/user_guide/models/diagvi.md index 554123bf95..e69de29bb2 100644 --- a/docs/user_guide/models/diagvi.md +++ b/docs/user_guide/models/diagvi.md @@ -1,481 +0,0 @@ -# DiagVI - -**DiagVI** (Diagonal multi-modal integration Variational Inference; Python class {class}`~scvi.external.DIAGVI`) is a deep generative model for diagonal integration of unpaired multi-modal single-cell data using prior knowledge about cross-modal feature correspondences. These relationships (such as associative or repressive interactions) are encoded in a guidance graph, where edge weights can represent both positive and negative covariation between features across modalities. - -DiagVI is inspired by the GLUE[^ref1] architecture, which uses modality-specific variational autoencoders (VAEs) to project heterogeneous data types into a shared latent space. In contrast to GLUE’s adversarial alignment strategy, DiagVI aligns modalities using Unbalanced Optimal Transport (UOT) via the Sinkhorn divergence[^ref2], explicitly accounting for differences in cell-type composition across modalities. - -The advantages of DiagVI are: - -- Flexible two-modality integration of various data types (e.g., scRNA-seq, spatial transcriptomics, spatial proteomics). -- Full feature utilization: all features (not only overlapping ones) contribute to model training via the guidance graph and modality-specific VAEs. -- Biologically informed alignment using prior feature correspondences via the guidance graph. -- Robust integration of modality-specific or rare cell populations via UOT. -- Scalable to very large datasets (>1 million cells). - -The limitations of DiagVI include: - -- Currently supports integration of two modalities only. -- Requires prior information on cross-modal feature correspondences (explicitly or implicitly). -- May require tuning of loss weights for optimal performance (for more information, see [Practical guidance](#practical-guidance)). -- Effectively requires a GPU for fast inference. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorials: - -- {doc}`/tutorials/notebooks/multimodal/DiagVI_spatial_transcriptomics` -- {doc}`/tutorials/notebooks/multimodal/DiagVI_spatial_proteomics` -``` - -:::{note} -DiagVI requires additional dependencies (`geomloss`, `torch_geometric`) that are not installed by default. To use DiagVI, install scvi-tools with the `diagvi` extra: - -```bash -pip install scvi-tools[diagvi] -``` -::: - - -## Preliminaries - -### Input - -DiagVI takes as input expression matrices $\mathbf{X}_1 \in \mathbb{R}^{N \times G}$ with $N$ cells and $G$ features and $\mathbf{X}_2 \in \mathbb{R}^{M \times P}$ with $M$ cells and $P$ features from two unpaired modalities. - -For **count data**, such as scRNA-seq, DiagVI expects a raw count expression matrix, where rows correspond to individual cells and columns correspond to features (e.g., genes). - -For **continuous data**, such as antibody-based single-cell or spatial proteomics, DiagVI expects a transformed (and optionally scaled) protein expression matrix, where rows correspond to cells and columns to features (e.g., marker proteins). Because antibody-based measurements are inherently relative, preprocessing is required. Typically, transformations such as arcsinh or log1p are applied to compress dynamic range and stabilize variance, followed by feature-wise scaling (e.g., z-score or min–max scaling) to ensure comparability across markers. This results in approximately normal-like feature distributions, which can be effectively modeled using a normal likelihood. Empirically, we found this approach performs better than modeling raw intensities (see also {class}`~scvi.external.CYTOVI` and the corresponding preprint {cite:p}`Ingelfinger25`). - -For integration with DiagVI, we recommend a simple two-step preprocessing strategy inspired by {class}`~scvi.external.CYTOVI`: - -1. **Arcsinh transformation** to stabilize variance and compress dynamic range. -2. **Feature-wise min–max scaling** to rescale each marker to the [0, 1] range and account for differences in marker brightness. - -Optional: For both count and continuous data, DiagVI can additionally incorporate: -- **Experimental covariates** such as batch annotation or **confounding variables** such as donor sex for one or both modalites. For simplicity, we will describe the case of categorical batch identifiers $s_n, s_m \in \{1,...,S\}$. -- **Cell label annotations** that weakly inform the prior of the latent space and guide a classifier in semi-supervised training for one or both modalities. We assume categorical label identifiers $c_n, c_m \in\{1,...,C\}$. - -Both optional inputs are handled independently for each modality, enabling distinct batch structures or cell type annotations per modality. - -Currently supported modalities include: -- **scRNA-seq** -- **Spatial transcriptomics**: image-based and sequencing-based, e.g., 10x Xenium, Vizgen MERSCOPE, NanoString CosMx, and 10x Visium -- **Spatial proteomics**: antibody-based imaging assays, e.g., PhenoCycler (formerly CODEX), 4i, Imaging Mass Cytometry, and protein add-on panels for Xenium or CosMx -- **Single-cell proteomics**: such as CITE-seq and other antibody-derived tag (ADT)-based assays -- Other count-based or continuous feature measurements - -### Model components - -DiagVI consists of several components which together define the overall training objective (see [Training Objective](#training-objective)). - -#### Modality-specific variational autoencoders - -DiagVI integrates two unpaired modalities by projecting their expression matrices $\mathbf{X}_1 \in \mathbb{R}^{N \times G}$ and $\mathbf{X}_2 \in \mathbb{R}^{M \times P}$ into a shared latent space in which each cell $n$ in modality 1 has a latent representation $\mathbf{z}_n \in \mathbb{R}^d$ and each cell $m$ in modality 2 has a latent representation $\mathbf{z}_m \in \mathbb{R}^d$. To find this shared state, the observed data from each modality is modeled using an independently parametrized VAE. - -#### Guidance graph - -Projecting cells into a common low-dimensional space alone does not result in semantically consistent embeddings in which identical cell types are assigned to the same region across both modalities. To ensure biological consistency in the latent space, a guidance graph establishes logical associations between the two modalities by connecting linked features (see also GLUE[^ref1]). - -Feature correspondences are encoded in the guidance graph $\mathcal{G} = (\mathcal{V}, \mathcal{E})$, where $\mathcal{V} = \mathcal{V}_1 \cup \mathcal{V}_2$ with modality-specific feature sets $\mathcal{V}_1$ and $\mathcal{V}_2$, and $\mathcal{E} \subseteq \mathcal{V} \times \mathcal{V}$ denotes the set of edges. Each edge $(i,j) \in \mathcal{E}$ is associated with - -- a weight $w_{ij} \in (0, 1]$ reflecting the confidence of the link -- a sign $\sigma_{ij} \in \{-1,1\}$ specifying whether the interaction is associative ($\sigma_{ij} = 1$) or repressive ($\sigma_{ij} = -1$) - -The graph loss encourages the inner product between embeddings $\mathbf{v}_i$ and $\mathbf{v}_j$ of linked features to be large and positive for $\sigma_{ij} = 1$ and large and negative for $\sigma_{ij} = -1$, with strength modulated by edge weights. - -#### Unbalanced optimal transport - -To ensure robust alignment of cells from different modalities within the shared latent space, DiagVI leverages unbalanced optimal transport (UOT). Specifically, it minimizes the de-biased Sinkhorn divergence between latent distributions using the [GeomLoss](https://www.kernel-operations.io/geomloss/#) library[^ref4]. UOT relaxes the marginal constraints of classical optimal transport, allowing for unequal total mass between distributions. This enables robust integration in the presence of differing cell type proportions across modalities. - -#### Classifier - -Inspired by {class}`~scvi.external.RESOLVI` {cite:p}`Ergen25`, a simple cell type classifier predicting labels $c_n$ and $c_m$ from cell latent vectors $\mathbf{z}_n$ and $\mathbf{z}_m$, respectively, is integrated into the model in the semi-supervised setting. This classifier is trained jointly with the generative model. - - -## Descriptive model - -DiagVI assumes that observations from each modality are generated from a shared $d$-dimensional latent space. For cell $n$ and feature $g$, the observed data $x_{ng}$ is generated conditionally on the cell latent variable $\mathbf{z}_n \in \mathbb{R}^{d}$, the feature latent variable $\mathbf{v}_g \in \mathbb{R}^{d}$, and its associated batch $s_n$. - -The prior on the feature latent variable is a standard multivariate normal: - -$$ -\mathbf{v}_g \sim \mathcal{N}(\mathbf{0}, \mathbf{I}_d). -$$ - -The prior on the cell latent variable can be either a standard multivariate normal (default) or a Gaussian mixture with $L$ components: - -$$ -\mathbf{z}_n \sim -\begin{cases} -\mathcal{N}(\mathbf{0}, \mathbf{I}_d), \\ -\sum_{l=1}^L \pi_l \, \mathcal{N}(\boldsymbol{\mu}_l, \boldsymbol{\Sigma}_l). -\end{cases} -$$ - -If cell type labels are provided and a Gaussian mixture prior is used, -$L$ is set to the number of unique cell types in the labeled data. - - -## Generative process - -### Generative model - -The generative model follows the formulations used in GLUE[^ref1], scVI {cite:p}`Lopez18`, and CytoVI {cite:p}`Ingelfinger25`, and adapts depending on whether the observed data consist of discrete counts or continuous measurements. In both cases, the specific form of the data likelihood varies with the distribution used to model the observed data in each modality. - -Regardless of the data type and likelihood choice, the inner product between cell and feature -latent embeddings parametrizes the decoder together with batch-specific versions of scaling ($\alpha_{s_n,g}$) and bias ($\beta_{s_n,g}$) parameters: - -$$ -\eta_{ng} = \alpha_{s_n,g} \, \mathbf{z}_n^\top \mathbf{v}_g + \beta_{s_n,g}. -$$ - -For count-based modalities, denoised and normalized expression proportions $\rho_{ng}$ are obtained via a softmax over features: - -$$ -\rho_{ng} = \frac{\exp(\eta_{ng})} {\sum_{g'} \exp(\eta_{ng'})}. -$$ - -The mean of the generative distribution is given by - -$$ -\mu_{ng} = l_n\rho_{ng}, -$$ - -where $l_n$ denotes the observed library size of cell $n$. Counts are then modeled as - -$$ -x_{ng} -\sim -\mathrm{NB}(\mu_{ng}, \theta_{s_n,g}) -$$ - -when using the negative binomial likelihood with mean $\mu_{ng}$ and dispersion $\theta_{s_n,g}$. For other likelihood choices, see [Likelihood Models](#likelihood-models). - -For continuous modalities, no simplex constraint is enforced and no library size normalization is applied. When using the normal likelihood, continuous expression values are modeled as - -$$ -x_{ng} -\sim -\mathcal{N}(\eta_{ng}, \sigma_{s_n,g}^2). -$$ - -In both settings, when a batch covariate is provided, DiagVI learns batch-specific versions of scaling and bias parameters to account for batch effects within each modality. For each modality $k$ and batch $s$, the scaling and bias parameters satisfy - -$$ -\boldsymbol{\alpha}^{(k)}_s \in \mathbb{R}_+^{|\mathcal{V}_k|}, -\qquad -\boldsymbol{\beta}^{(k)}_s \in \mathbb{R}^{|\mathcal{V}_k|}. -$$ - -This generative process is also summarized in the following graphical model: - -:::{figure} figures/diagvi_graphical_model.svg -:align: center -:alt: DiagVI graphical model -:class: img-fluid -:width: 90% - -DiagVI graphical model. Shaded nodes represent observed data, unshaded nodes represent latent variables. -::: - -### Likelihood models - -Depending on the modality and data characteristics, different likelihood functions can be used to reconstruct the input data. The choice of likelihood should reflect whether the data are counts or continuous measurements, and whether they exhibit excess zeros or background signal. - -We generally recommend: - -- **Negative binomial (``nb``)** for count-based single-cell RNA-seq data. -- **Normal (``normal``)** for transformed and scaled continuous measurements. - -More specialized likelihoods such as zero-inflated or mixture models may be beneficial for sparse or background-contaminated measurements. - -DiagVI supports the following likelihood functions: - -```{eval-rst} -.. list-table:: - :widths: 10 20 30 40 - :header-rows: 1 - - * - Likelihood - - Distribution - - Typical modality - - Recommended preprocessing - * - ``nb`` - - Negative Binomial - - Count data (e.g., scRNA-seq) - - Raw counts - * - ``zinb`` - - Zero-Inflated Negative Binomial - - Strongly zero-inflated count data - - Raw counts - * - ``nbmixture`` - - Negative Binomial Mixture - - Protein counts with background signal (e.g., CITE-seq ADTs) - - Raw counts - * - ``normal`` - - Normal - - Continuous data - - Transformed (e.g., arcsinh or log1p) and feature-wise scaled - * - ``log1pnormal`` - - Log1p-Normal - - Non-negative continuous data - - Raw data, optionally scaled - * - ``ziln`` - - Zero-Inflated Log-Normal - - Sparse positive continuous data - - Raw data, optionally scaled - * - ``zig`` - - Zero-Inflated Gamma - - Sparse positive continuous data - - Raw data, optionally scaled -``` - -### Latent variables - -The exact set of variables instantiated during training depends on the chosen likelihood and on whether a Gaussian mixture prior is used for the cell latent space. Below we summarize the principal latent, decoder, and auxiliary variables used in DiagVI. - -```{eval-rst} -.. list-table:: - :widths: 25 55 20 - :header-rows: 1 - - * - Variable - - Description - - Code variable - * - :math:`\mathbf{z}_n \in \mathbb{R}^{d}` - - Low-dimensional latent representation of cell :math:`n`, capturing its underlying biological state. - - ``z`` - * - :math:`\mathbf{v}_g \in \mathbb{R}^{d}` - - Low-dimensional embedding of feature :math:`g`, inferred via the graph encoder from the guidance graph. - - ``v`` - * - :math:`c_n \in \{1,\dots,C\}` - - Observed cell type label of cell :math:`n` (if available), used for semi-supervised training and to inform the Gaussian mixture prior. - - ``y`` - * - :math:`\eta_{ng} \in \mathbb{R}` - - Decoder linear predictor obtained from the bilinear interaction between cell and feature embeddings. - - N/A - * - :math:`\boldsymbol{\rho}_n \in \Delta^{|\mathcal{V}_k|-1}` - - Denoised, normalized expression proportions for cell :math:`n` (count data); constrained to the probability simplex via softmax. - - ``px_scale`` - * - :math:`\mu_{ng} \in \mathbb{R}_+` - - Mean of the generative distribution for feature :math:`g` in cell :math:`n`. - - ``px_rate`` - * - :math:`\theta_{s_n,g} \in (0,\infty)` - - Batch-specific inverse dispersion parameters used in count-based likelihoods. - - ``px_r`` - * - :math:`\sigma_{s_n,g}^2 \in (0,\infty)` - - Batch-specific variance parameters used in continuous likelihoods. - - ``px_r`` - * - :math:`\delta_g \in (0,\infty)` - - Dropout or zero-inflation parameters (used for zero-inflated likelihoods). - - ``px_dropout`` - * - :math:`\pi_l` - - Mixture weights of the Gaussian mixture prior (if enabled). - - ``gmm_logits`` -``` - - -## Inference - -Since the posterior distributions over the latent variables are intractable, DiagVI uses variational inference and specifically auto-encoding variational Bayes (see {doc}`/user_guide/background/variational_inference`) to jointly learn model parameters and approximate posterior distributions. - -For a given modality, the variational distribution factorizes as - -$$ -q_\eta(\mathbf{z}_n, \mathbf{v} \mid x_n, \mathcal{G}) = q_\eta(\mathbf{z}_n \mid x_n) \; q_\eta(\mathbf{v} \mid \mathcal{G}). -$$ - -Here $\eta$ is a set of parameters corresponding to inference neural networks (encoders). - - -## Training objective - -DiagVI is trained by minimizing a weighted sum of loss terms (weighted by `lam_*` parameters) corresponding to the [Model Components](#model-components) introduced above: - -- **Modality-specific VAEs**: The data reconstruction loss (`lam_data`) measures how well each modality-specific decoder reconstructs the observed data, while the KL divergence (`lam_kl`) regularizes the cell latent variables by encouraging adherence to the prior. -- **Guidance graph**: The graph reconstruction loss (`lam_graph`) enforces biological consistency between feature embeddings using the guidance graph. -- **UOT**: The UOT alignment loss (`lam_sinkhorn`) aligns cell distributions across modalities using unbalanced optimal transport. -- **Classifier**: The classification loss (`lam_class`, optional) enables supervised or semi-supervised training via cell-type labels. - -The `lam_*` parameters control the relative importance of within-modality reconstruction (`lam_data`, `lam_kl`), cross-modality alignment (`lam_graph`, `lam_sinkhorn`), and optional supervision (`lam_class`). DiagVI uses sensible defaults for all of these values, but they may require further tuning depending on the type of data being integrated (see [Practical guidance](#practical-guidance)). - - -## Practical guidance - -### Guidance graph creation - -When initializing {class}`~scvi.external.DIAGVI`, three ways to specify the guidance graph are supported: - -1. **Automatic construction**: If neither `guidance_graph` nor `mapping_df` are provided, DiagVI constructs a graph from overlapping feature names across modalities (e.g., shared gene symbols). - -2. **Custom mapping via DataFrame**: Pass a `pandas.DataFrame` to `mapping_df` in which each column corresponds to a modality name (matching the keys in `input_dict`) and each row specifies a feature pair. This is useful when feature naming conventions differ between modalities (e.g., genes vs. proteins). You can also use {func}`~scvi.external.DIAGVI.construct_custom_guidance_graph` to create a graph with custom edge weights and signs. - -3. **Explicit graph specification**: For full control, pass a pre-constructed `torch_geometric.data.Data` object directly to `guidance_graph`. The graph must include node features, edge indices, edge weights, edge signs, and modality-specific feature index tensors.. - -### Loss weights - -DiagVI provides default values for all loss weights that have been tested across multiple datasets and integration tasks. However, depending on the modalities and the integration setting, tuning some of these weights can improve performance. - -In practice, we recommend paying particular attention to: - -- `lam_sinkhorn` (cross-modality alignment strength), and -- `lam_class` (classification strength, if cell-type labels are provided). - -These two weights control the trade-off between alignment across modalities and separation of cell types within the latent space. - -When integrating **very different modalities** (e.g., scRNA-seq + spatial proteomics), stronger alignment is typically required. In such cases: - -- Higher `lam_sinkhorn` -- Lower `lam_class` (if labels are used) - -This encourages the model to prioritize matching global cell distributions across modalities, even if the feature spaces differ substantially. - -When integrating **more similar modalities** (e.g., scRNA-seq + spatial transcriptomics), less alignment pressure is needed. In such cases: - -- Lower `lam_sinkhorn` -- Higher `lam_class` (if labels are used) - -Here, the modalities are already structurally similar, so stronger supervision can help maintain clearer separation of biologically distinct cell types. - -As a general rule, we recommend `lam_sinkhorn` and `lam_class` values between 1 and 100, which work well in most settings. - -The default loss weight values place relatively strong emphasis on cross-modality alignment and are therefore well suited for integrating heterogeneous modalities. For more similar modalities, we recommend reducing `lam_sinkhorn` (e.g., 5-20) and increasing `lam_class` (e.g., 60-80), depending on label availability and desired separation. - -In practice, small grid searches over `lam_sinkhorn` and `lam_class` are often sufficient to find well-performing settings. - -### Sinkhorn parameters - -Three parameters determine the behavior of the Sinkhorn divergence used for cross-modality alignment: - -- `p`: order of the ground cost (Wasserstein-p distance). The default is p = 2, corresponding to a squared Euclidean ground cost. -- `blur`: controls the strength of entropic regularization and therefore the smoothness of the transport plan. Larger values increase regularization, resulting in smoother and more diffuse transport plans. -- `reach`: controls the marginal relaxation parameter in unbalanced optimal transport, penalizing deviations from strict mass conservation. Larger values enforce stricter marginal constraints, while smaller values allow greater mass variation between the matched distributions. - -Although these parameters can be specified manually, we generally **do not recommend modifying them**. - -By default, DiagVI follows the heuristic strategy introduced in [OTT-JAX](https://ott-jax.readthedocs.io/). In this setting: -- `blur` is computed adaptively at each optimization step from the minibatch cost matrix, and -- `reach` is derived as a function of blur. - -This adaptive strategy makes the Sinkhorn loss scale-aware and robust across datasets, and in practice removes the need for manual tuning. Adjusting these parameters is rarely necessary. - - -## Tasks - -Here we provide an overview of some of the tasks that DiagVI can perform. Please see {class}`~scvi.external.DIAGVI` for the full API reference. - -### Dimensionality reduction - -DiagVI provides aligned low-dimensional representations for each modality. By default, the mean of the approximate posterior $q_\eta(\mathbf{z}_n \mid x_n)$ is returned. - -Latent representations can be obtained using: - -``` ->>> latents = model.get_latent_representation() ->>> adata_mod1.obsm["X_diagvi"] = latents["mod1"] ->>> adata_mod2.obsm["X_diagvi"] = latents["mod2"] -``` - -The aligned representations can be used for downstream analyses either within each modality or jointly across modalities. - -For example, within-modality visualization: - -``` ->>> import scanpy as sc ->>> sc.pp.neighbors(adata_mod1, use_rep="X_diagvi") ->>> sc.tl.umap(adata_mod1) -``` - -For joint analysis across modalities: - -``` ->>> adata_combined = sc.concat( -... [adata_mod1, adata_mod2], -... axis=0, -... join="inner", -... label="modality", -... keys=["mod1", "mod2"], -... ) ->>> sc.pp.neighbors(adata_combined, use_rep="X_diagvi") ->>> sc.tl.umap(adata_combined) -``` - -### Cross-modal feature imputation - -DiagVI can impute features from one modality into another using {func}`~scvi.external.DIAGVI.get_imputed_values`. - -For example, to impute protein expression for RNA cells: - -``` ->>> imputed_protein = model.get_imputed_values( -... query_name="rna", -... query_adata=adata_rna, -... ) ->>> adata_rna.obsm["imputed_protein"] = imputed_protein -``` - -It is also possible to perform counterfactual predictions by specifying a reference batch or reference library size (see also {doc}`/user_guide/background/counterfactual_prediction`): - -``` ->>> imputed_rna = model.get_imputed_values( -... query_name="protein", -... reference_batch="batch_1", -... reference_libsize=10000, -... ) ->>> adata_protein.obsm["imputed_rna"] = imputed_rna -``` - -This allows imputation under specified sequencing depth or batch conditions. - -### Cell label transfer - -Aligned latent representations can be used for cross-modal cell label transfer. - -For example, when integrating scRNA-seq reference data with spatial proteomics data, cell-type labels can be transferred from the RNA reference to the spatial data. This can be done either using DiagVI's built-in classifier (when cell-type labels have been provided for the reference modality during training) or with external tools such as [`CellMapper`](https://cellmapper.readthedocs.io/en/latest/), a toolkit for cross-modal cell mapping and evaluation. - -Using DiagVI's trained classifier, cell labels can be transferred as follows: - -``` ->>> preds = model.predict_celltype(labeled_modality="rna") ->>> adata_protein.obs["celltype_pred"] = preds["predictions"] ->>> adata_protein.obs["celltype_conf"] = preds["confidence"] -``` - -Alternatively, using CellMapper: - -``` ->>> latents = model.get_latent_representation() ->>> adata_rna.obsm["X_diagvi"] = latents["rna"] ->>> adata_protein.obsm["X_diagvi"] = latents["protein"] ->>> ->>> import CellMapper ->>> cmap = CellMapper(adata_protein, adata_rna) ->>> cmap.map("cell_type_labels", use_rep="X_diagvi") ->>> cmap.evaluate_label_transfer("cell_type_labels") -``` - -Because DiagVI produces a shared latent space, any downstream method that operates on low-dimensional embeddings (e.g., kNN mapping, clustering, classification) can be applied directly. - -### Data simulation - -Data can be generated from the model using the posterior predictive distribution in {func}`~scvi.external.DIAGVI.posterior_predictive_sample`. This is equivalent to feeding a cell through the model, sampling from the posterior distributions of the latent variables, retrieving the likelihood parameters, and finally, sampling from this distribution. - - -## References - -[^ref1]: - Cao, Zhi-Jie; Gao, Ge (2022), - _Multi-omics single-cell data integration and regulatory inference with graph-linked embedding_, - [Nature Biotechnology](https://www.nature.com/articles/s41587-022-01284-4). - -[^ref2]: - Séjourné, Thibault; Feydy, Jean; Vialard, François-Xavier; Trouvé, Alain; Peyré, Gabriel (2023), - _Sinkhorn Divergences for Unbalanced Optimal Transport_, - [arXiv](https://arxiv.org/abs/1910.12958). - -[^ref4]: - Feydy, Jean; Séjourné, Thibault; Vialard, François‑Xavier; Amari, Shun‑ichi; Trouvé, Alain; Peyré, Gabriel (2019) - _Interpolating between Optimal Transport and MMD using Sinkhorn Divergences_ - [The 22nd International Conference on Artificial Intelligence and Statistics](https://arxiv.org/abs/1810.08278). diff --git a/docs/user_guide/models/figures/diagvi_graphical_model.svg b/docs/user_guide/models/figures/diagvi_graphical_model.svg deleted file mode 100644 index 0fde13968b..0000000000 --- a/docs/user_guide/models/figures/diagvi_graphical_model.svg +++ /dev/null @@ -1,2358 +0,0 @@ - - - -feature latent space - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -batch idlibrary sizelatent spaceobservationcell n feature g - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -feature latent spaceguidance graphbatch idlibrary sizelatent spaceobservationcell m feature p - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -feature latent space - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -batch idlibrary sizeobservationcell n feature g - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -feature latent spaceguidance graphbatch idlibrary sizecell type labelobservationcell m feature platentspacelatentspace - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -cell type label - - - - - - - diff --git a/docs/user_guide/models/gimvi.md b/docs/user_guide/models/gimvi.md deleted file mode 100644 index 776829b36d..0000000000 --- a/docs/user_guide/models/gimvi.md +++ /dev/null @@ -1,204 +0,0 @@ -# gimVI - -**gimVI** {cite:p}`Lopez19` (Generative Imputation Model with Variational Inference; Python class {class}`~scvi.external.GIMVI`) is a deep generative model for the joint analysis of an scRNA-seq dataset and a spatial transcriptomics dataset that measures only a subset of the genes. gimVI learns a single shared latent space for the two modalities and uses it to **impute the genes that are missing from the spatial assay** by transferring information from the (genome-wide) sequencing reference. - -The two datasets are modeled by a single joint variational autoencoder ({class}`~scvi.external.gimvi.JVAE`) with modality-specific encoder and decoder heads but a shared latent representation. Because the spatial panel is a subset of the genes captured by scRNA-seq, the decoder reconstructs into the full sequencing gene space, and each modality's reconstruction loss is masked to only the genes it actually observes. Alignment of the two modalities in the shared latent space is encouraged by an adversarial classifier that is trained to predict the modality of origin from a cell's latent code, while the generative model is trained to fool it. - -The advantages of gimVI are: - -- Imputation of genes that are not measured in the spatial assay, using a paired (genome-wide) scRNA-seq reference. -- A shared latent space that integrates spatial and sequencing cells for joint visualization and clustering. -- Flexible, per-modality generative distributions (e.g. `zinb` for the sequencing data and `nb` for the spatial data) and per-modality choice of whether to model library size. -- Handles the two assays measuring different gene sets and different numbers of cells. - -The limitations of gimVI include: - -- Supports exactly two datasets (one sequencing, one spatial). -- The spatial genes must be a subset of the sequencing genes. -- Imputation quality depends on how well the scRNA-seq reference matches the spatial tissue. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorials: - -- {doc}`/tutorials/notebooks/spatial/gimvi_tutorial` -``` - -:::{note} -gimVI is a spatial transcriptomics model that will be moved to the scvi-tools spatial companion package `scviva-tools` starting in scvi-tools v1.5 and will no longer be supported in scvi-tools; it is scheduled for deprecation in v1.6. -::: - - -## Preliminaries - -gimVI takes as input two raw-count gene expression matrices: - -- a sequencing matrix $X^{(0)} \in \mathbb{N}^{N_0 \times G_0}$ with $N_0$ cells and $G_0$ genes (the full gene set), and -- a spatial matrix $X^{(1)} \in \mathbb{N}^{N_1 \times G_1}$ with $N_1$ cells and $G_1$ genes, - -where the spatial genes are a subset of the sequencing genes ($\mathcal{G}_1 \subseteq \mathcal{G}_0$). Internally the decoder operates on the full sequencing gene space of size $G = G_0$, and a per-modality index mapping $\mathcal{M}_m$ records where each modality's observed genes live in that space. The spatial reconstruction is therefore evaluated only on the columns corresponding to $\mathcal{G}_1$. - -Each dataset can additionally carry a categorical batch covariate $s$ and, optionally, categorical cell-type labels $c$. Both datasets are registered independently with two separate calls to {meth}`~scvi.external.GIMVI.setup_anndata` before the joint model is constructed: - -``` ->>> scvi.external.GIMVI.setup_anndata(adata_seq, batch_key="batch", labels_key="labels") ->>> scvi.external.GIMVI.setup_anndata(adata_spatial, batch_key="batch", labels_key="labels") ->>> model = scvi.external.GIMVI(adata_seq, adata_spatial) ->>> model.train() -``` - -The two datasets share a single batch space: batch indices of the spatial dataset are offset by the number of sequencing batches so that the two modalities never collide on a batch id (see [#2446](https://github.com/scverse/scvi-tools/issues/2446)). - -### Model architecture - -gimVI is built from the following components, which together define the [training objective](#training-objective): - -#### Joint encoder - -A single multi-head encoder ({class}`~scvi.nn.MultiEncoder`) maps both modalities into a shared $d$-dimensional latent space. Each modality has its own input layer(s) (because the two assays have different numbers of genes), followed by shared layers that produce the variational posterior $q(\mathbf{z}_n \mid x_n)$. The `mode` argument selects which input head is used. - -For each modality, an optional, separate library-size encoder produces a latent library size. Whether a modality models library size as a latent variable or instead uses the observed log total counts is controlled by `model_library_size` (default `[True, False]`: the sequencing data models library size, the spatial data uses the observed value). - -#### Joint decoder - -A single, shared multi-decoder ({class}`~scvi.nn.MultiDecoder`) maps the shared latent code back into the full sequencing gene space of size $G$, conditioned on the batch covariate (the decoder weights are **not** modality-specific). Because the decoder reconstructs all $G$ genes for both modalities, the latent representation of a spatial cell can be decoded into genes that were never measured spatially — this is exactly the imputation mechanism (see [Imputation of missing spatial genes](#imputation-of-missing-spatial-genes)). The modality index only selects which gene-index mapping is used to renormalize the decoded expression over each modality's observed genes; it does not route through different decoder networks. - -#### Adversarial modality classifier - -A small classifier ({class}`~scvi.module.Classifier`, 3 layers, 2 output classes) is trained to predict, from a latent code $\mathbf{z}$, which modality it came from. The generative model is trained adversarially to *fool* this classifier, which pushes the sequencing and spatial latent distributions to overlap and yields an integrated, modality-mixed latent space. This is implemented in {class}`~scvi.external.gimvi._task.GIMVITrainingPlan` (a subclass of {class}`~scvi.train.AdversarialTrainingPlan`), and the strength of the adversarial term is set by the `kappa` argument of {meth}`~scvi.external.GIMVI.train`. - - -## Generative process - -For a cell $n$ in modality $m \in \{0, 1\}$ (0 = sequencing, 1 = spatial) with batch $s_n$, gimVI posits the following generative process over the full gene space: - -```{math} -:nowrap: true - -\begin{align} - \mathbf{z}_n &\sim \mathrm{Normal}(0, I) \\ - \ell_n &\sim - \begin{cases} - \mathrm{LogNormal}(\ell_\mu^\top s_n,\ \ell_{\sigma^2}^\top s_n) & \text{if modality } m \text{ models library size} \\ - \delta(\log \textstyle\sum_g x_{ng}) & \text{otherwise (observed library size)} - \end{cases} \\ - \rho_n &= f_w(\mathbf{z}_n, s_n) \\ - x_{ng} &\sim \mathrm{ObservationModel}_m(\ell_n \rho_{ng},\ \theta_{s_n g}),\quad g \in \mathcal{G}_m -\end{align} -``` - -The latent code $\mathbf{z}_n$ is **shared across modalities** — the same prior and the same latent space are used for both datasets. The shared decoder $f_w$ produces normalized expression $\rho_n$ over the full gene set; for modality $m$, $\rho_n$ is renormalized to sum to 1 over only the genes observed in that modality ($\mathcal{M}_m$), and the mean rate is $\mu_{ng} = \ell_n \rho_{ng}$. - -The observation model is chosen per modality via `generative_distributions` (default `["zinb", "nb"]`). gimVI supports: - -- `zinb` — Zero-Inflated Negative Binomial (default for the sequencing data), -- `nb` — Negative Binomial (default for the spatial data), -- `poisson` — Poisson. - -The inverse dispersion $\theta$ is by default gene- and batch-specific (`dispersion="gene-batch"`); `gene` and `gene-label` parameterizations are also available. - -The latent variables are summarized below: - -```{eval-rst} -.. list-table:: - :widths: 25 60 15 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable - * - :math:`\mathbf{z}_n \in \mathbb{R}^{d}` - - Shared low-dimensional representation of a cell, common to both modalities. - - ``z`` - * - :math:`\rho_n \in \Delta^{|\mathcal{M}_m|-1}` - - Denoised, normalized expression for cell :math:`n`, renormalized over the genes observed in its modality. - - ``px_scale`` - * - :math:`\ell_n \in (0, \infty)` - - Library size. Modeled as a latent variable for modalities flagged in ``model_library_size``; otherwise set to the observed log total count. - - ``library`` - * - :math:`\mu_{ng} \in (0, \infty)` - - Mean of the count likelihood, :math:`\ell_n \rho_{ng}`. - - ``px_rate`` - * - :math:`\theta_{s_n g} \in (0, \infty)` - - Inverse dispersion of the (zero-inflated) negative binomial, by default gene- and batch-specific. - - ``px_r`` -``` - - -## Inference - -gimVI uses variational inference (auto-encoding variational Bayes; see {doc}`/user_guide/background/variational_inference`) with the per-modality factorization - -```{math} -:nowrap: true - -\begin{align} - q_\eta(\mathbf{z}_n, \ell_n \mid x_n) := q_\eta(\mathbf{z}_n \mid x_n)\, q_\eta(\ell_n \mid x_n). -\end{align} -``` - -The latent posterior $q_\eta(\mathbf{z}_n \mid x_n)$ is produced by the shared multi-head encoder (using the input head for the cell's modality), and $q_\eta(\ell_n \mid x_n)$ by the modality's optional library encoder. Inputs are log-transformed (`log(1 + x)`) before encoding for numerical stability. - - -## Training objective - -gimVI minimizes a per-modality evidence lower bound — a reconstruction term plus KL terms on the latent variables — summed across both modalities, with an adversarial term added on top: - -- **Reconstruction**: the negative log-likelihood of the observed counts under the chosen observation model, **masked to the genes observed in each modality** (the spatial term ignores genes outside the spatial panel). -- **KL divergence on $\mathbf{z}$**: regularizes the shared latent variable toward the standard normal prior. -- **KL divergence on $\ell$**: applied only for modalities that model library size as a latent variable, with a batch-specific log-normal prior. -- **Adversarial loss**: the modality classifier is trained to predict modality from $\mathbf{z}$, while the generative model is trained to fool it. The relative weight is controlled by `kappa` (the `scale_adversarial_loss` of the training plan); when set to `"auto"` it follows `1 - kl_weight` over the warmup schedule. - -The two objectives (model vs. classifier) are optimized in alternating steps with manual optimization, as in {class}`~scvi.train.AdversarialTrainingPlan`. - -### Cyclic multi-dataset data loading - -The two datasets generally have different numbers of cells, so a custom data loader, {class}`~scvi.external.gimvi._task.CyclicMultiDataLoader`, combines their per-modality dataloaders: it iterates the longest loader once per epoch and **cycles the shorter loader(s)** so that every training step yields one minibatch from each modality. It accepts either a sequence (yielding a tuple of minibatches) or a mapping of modality name → loader (yielding a dict, preserving modality names). - -This loader is shared infrastructure: {class}`~scvi.external.DIAGVI` imports it (`from scvi.external.gimvi._task import CyclicMultiDataLoader as TrainDL`) and uses the mapping form to drive its own multi-modality training, so the same cycling logic backs both gimVI and DiagVI. - - -## Tasks - -Here we provide an overview of some of the tasks that gimVI can perform. Please see {class}`~scvi.external.GIMVI` for the full API reference. - -### Dimensionality reduction - -{meth}`~scvi.external.GIMVI.get_latent_representation` returns the shared latent representation for **each** dataset as a list `[latent_seq, latent_spatial]`. By default the mean of the approximate posterior is returned (`deterministic=True`). - -``` ->>> latent_seq, latent_spatial = model.get_latent_representation() ->>> adata_seq.obsm["X_gimvi"] = latent_seq ->>> adata_spatial.obsm["X_gimvi"] = latent_spatial -``` - -Because the latent space is shared and modality-mixed (via the adversarial classifier), the two representations can be concatenated for a joint neighbor graph and UMAP to assess integration quality: - -``` ->>> import scanpy as sc, anndata as ad, numpy as np ->>> joint = ad.AnnData(np.concatenate([latent_seq, latent_spatial])) ->>> joint.obs["modality"] = ["seq"] * len(latent_seq) + ["spatial"] * len(latent_spatial) ->>> sc.pp.neighbors(joint, use_rep="X") ->>> sc.tl.umap(joint) -``` - -### Imputation of missing spatial genes - -The central task of gimVI is to impute, for the spatial cells, the genes that were only measured in the scRNA-seq data. {meth}`~scvi.external.GIMVI.get_imputed_values` returns the decoded expression for **all** sequencing genes for each dataset: - -``` ->>> imputed_seq, imputed_spatial = model.get_imputed_values(normalized=True) ->>> # imputed_spatial now contains values for every sequencing gene, ->>> # including those absent from the original spatial panel. -``` - -- `normalized=True` (default) returns the normalized expression frequencies $\rho_n$ (decoder `sample_scale`); `normalized=False` returns the unnormalized mean rate $\mu_n = \ell_n \rho_n$ (decoder `sample_rate`). -- `decode_mode` selects which dataset's decoding is applied. The cell is always encoded with its own modality's encoder; gimVI then uses a **single shared decoder**, and `decode_mode` only changes which modality's gene-index mapping is used to renormalize the decoded expression over the probability simplex (`indices_mappings[decode_mode]`). It does not switch to a separate per-modality decoder network. - -A common evaluation (used in the tutorial) holds out a fraction of genes from the spatial panel, imputes them with gimVI, and measures the per-gene Spearman correlation between the held-out true spatial expression and the imputed values. - -### Saving and loading - -gimVI uses a custom {meth}`~scvi.external.GIMVI.save` / {meth}`~scvi.external.GIMVI.load` that persists the two registries and both sets of `var_names` (the model stores two AnnData managers). `save_anndata=True` writes both `adata_seq` and `adata_spatial` alongside the model. Legacy saves (pre-v0.15.0) can be upgraded with {meth}`~scvi.external.GIMVI.convert_legacy_save`. diff --git a/docs/user_guide/models/index.md b/docs/user_guide/models/index.md index fdd90c8d26..238785f4f7 100644 --- a/docs/user_guide/models/index.md +++ b/docs/user_guide/models/index.md @@ -9,10 +9,7 @@ cellassign contrastivevi cytovi decipher -destvi -diagvi drvi -gimvi jointembeddingscvi linearscvi methylanvi @@ -21,16 +18,12 @@ mrvi multivi peakvi poissonvi -resolvi scanvi scar scbasset scvi -scviva solo -stereoscope sysvi -tangram totalanvi totalvi velovi diff --git a/docs/user_guide/models/resolvi.md b/docs/user_guide/models/resolvi.md deleted file mode 100644 index 958a6e937e..0000000000 --- a/docs/user_guide/models/resolvi.md +++ /dev/null @@ -1,195 +0,0 @@ -# ResolVI - -**resolVI** {cite:p}`Ergen25` (Python class {class}`~scvi.external.RESOLVI`) is a generative model of single-cell resolved spatial -transcriptomics that can subsequently be used for many common downstream tasks. - -The advantages of resolVI are: - -- Addresses noise and bias in ST data due to wrong segmentation, unspecific background and limited spatial resolution -- Scalable to very large datasets (>1 million cells). - -The limitations of resolVI include: - -- Effectively requires a GPU for fast inference. -- Latent space is not interpretable, unlike that of a linear method. -- Assumes single cells are observed and do not work with low-resolution ST like Visium or Slide-Seq. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorials: - -- {doc}`/tutorials/notebooks/spatial/resolVI_tutorial` -``` - -## Preliminaries - -ResolVI takes as input spatially resolved RNA_seq count matrices downstream of cellular segmentation and molecule -assignments to cells. These counts can be either derived from sequencing spatially resolved molecules or fluorescent -imaging. ResolVI leverages the gene expression of neighboring cells and reassigns observed gene expression to neighboring -cells as well as an unspecific background. - -ResolVI accepts as input the observed expression of the cell itself, its spatial neighbors and their gene expression -as well as the distance between these cells. Additionally, a vector of categorical covariates $S$, representing -batch, donor, etc., is an optional input to the model. ResolVI provides a semi-supervised mode, adjusting the prior in -the latent space for different cell types and training a classifier to predict cell types from latent embeddings. - -## Generative process - -ResolVI posits that the observed expression of cell $n$ in gene $g$, $x_{ng}$ is generated by the following process: - -```{math} -:nowrap: true - -\begin{align} - z &\sim \mathrm{MixtureOfGaussians}(\mu_1, \dots, \mu_K, \Sigma_1, \dots, \Sigma_K) \\ - \alpha_n &\sim \mathrm{Dirichlet}(C) \\ - r_{ng} &\sim \mathrm{Exponential}(R) \\ - h_{ng} &= - \mathrm{Gamma}(r_{ng}, \frac{r_{ng}}{\alpha_0 f_\theta(z, b) + \alpha_1 \sum\limits_{{N(n)}} \beta_{N(n)} f_\theta(z_{N(n)}, b)}) + \alpha_2 bg\\ - x_{ng} &\sim \mathrm{Poisson}(l_n h_{ng}) -\end{align} -``` - -In particular, $z$ and $z_{N(n)}$ are the latent embeddings of the cell itself as well as its spatial neighbors -both of dimension $L$. ResolVI uses a mixture of Gaussians prior to $z$: - -```{math} -:nowrap: true - -\begin{align} - c_n &\sim \textrm{Categorical}( - \pi_1, \pi_2, \dots, \pi_K - ), \\ - z_n \mid c_n = c &\sim \mathcal{N}(\mu_c, \sigma_c) -\end{align} -``` - -In brief, we assume that observed expression of gene $g$ for cell $n$ can be modelled as a sum over -the components of expression truly expressed by the cell $\alpha_0$, the expression explained by neighboring -cells $\alpha_1$ and wrongly assigned to $n$ and a component due to unspecific background $\alpha_2$. -The expression of neighboring cells $N(n)$ is assigned to each of the neighboring cells $\beta_{N(n)n}$. -Both the expression of cell $n$ and the expression of neighboring cells $N(n)$ are generated using the same -generative network $f_\Theta$ from their respective latent code $z_{N(n)}$ and $z_n$. -This generative process uses a neural network: - -```{math} -:nowrap: true - -\begin{align} - f_{\theta}(z_{n}, s_n) &: \mathbb{R}^{d} \times \{0, 1\}^K \to \Delta^{G-1} -\end{align} -``` - -which estimates the normalized gene expression of cell $n$. We use the observed counts per cell to scale these rates. - -The latent variables, along with their description, are summarized in the following table: - -```{eval-rst} -.. list-table:: - :widths: 20 90 15 45 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable (if different) - - Prior - * - :math:`z_n \in \mathbb{R}^L` - - Low-dimensional representation capturing the state of a cell - - ``latent`` - - Mixture-of-Gaussian - * - :math:`\beta_{N(n)} \in \Delta^{N(n) - 1}` - - Per-neighbor diffusion - - ``per_neighbor_diffusion`` - - Dirichlet - * - :math:`\alpha_{n0 \dots 2} \in \Delta^{2}` - - Per cell true, diffusion and background proportion - - ``mixture_proportions`` - - Dirichlet - * - :math:`bg_{ng} \in \Delta^{G - 1}` - - Per cell estimate of background - - ``background`` - - None - * - :math:`background_{s} \in \mathbb{R}^G` - - Per sample background vector - - ``per_gene_background`` - - Dirichlet - * - :math:`\rho_n \in \Delta^{G - 1}` - - Per cell rate of expression - - ``px_scale`` - - None - * - :math:`\mu_n, \mu_{N(n)} \in \mathbb{R}^G` - - Per cell estimated expression - - ``px_rate and px_rate_n`` - - None -``` - - -## Inference - -ResolVI uses variational inference, specifically auto-encoding variational Bayes -(see {doc}`/user_guide/background/variational_inference`) in Pyro to learn both the model parameters -(the neural network parameters, dispersion parameters, etc.) and an approximate posterior distribution. -We perform amortization using neural network for $z_n$ and $\alpha_n$, while $\beta_{N(n)n}$ is estimated -for each cell. - -## Tasks - -Here we provide an overview of some of the tasks that resolVI can perform. Please see {class}`scvi.external.RESOLVI` -for the full API reference. - -### Dimensionality reduction - -For dimensionality reduction, the mean of the approximate posterior $q_\phi(z_i \mid y_i, n_i)$ is returned by default. -This is achieved using the method: - -``` ->>> adata.obsm["X_resolvi"] = model.get_latent_representation() -``` - -Users may also return samples from this distribution, as opposed to the mean, by passing the argument `give_mean=False`. -The latent representation can be used to create a nearest neighbor graph with scanpy with: - -``` ->>> import scanpy as sc ->>> sc.pp.neighbors(adata, use_rep="X_resolvi") ->>> adata.obsp["distances"] -``` - -### Transfer learning - -A resolVI model can be pre-trained on reference data and updated with query data using {meth}`~scvi.external.RESOLVI.load_query_data`, which then facilitates transfer of metadata like cell type annotations. $\beta_{N(n)n}$ is extended to the new cells and learned on these cells. The encoder by default does not see the batch covariate, and $z_n$ can be predicted without performing query model training. See the {doc}`/user_guide/background/transfer_learning` guide for more information. - -### Estimation of true expression levels - -In {meth}`~scvi.external.RESOLVI.get_normalized_expression` ResolVI returns the expected true expression value of $\rho_n$ under the approximate posterior. For one cell $n$, this can be written as: - -```{math} -:nowrap: true - -\begin{align} - \mathbb{E}_{q_\phi(z_n \mid x_n)}\left[f_{\theta}\left(z_{n}, s_n \right) \right] -\end{align} -``` - -### Differential expression - -Differential expression analysis is achieved with {meth}`~scvi.external.RESOLVI.differential_expression`. -ResolVI tests differences in expression levels $\rho_{n} = f_{\theta}\left(z_n, s_n\right)$. -We allow for importance-based sampling using pyro's built-in function. - -### Cell-type prediction - -Prediction of cell-type labels is performed with {meth}`~scvi.external.RESOLVI.predict`. -A semisupervised model is necessary to perform this analysis as it leverages the cell-type classifier. -ResolVI performs for each cell $n$ $c_{n} = h_{nu}\left(z_n\right)$ and samples from $z_n$ to yield -the cell-type labels. - -### Differential niche abundance - -Differential niche abundance analysis is achieved with {meth}`~scvi.external.RESOLVI.differential_niche_abundance`. -A semisupervised model is necessary to perform this analysis as it leverages the cell-type classifier. -ResolVI tests differences in abundance of various cell-types in the neighborhood of a cell $n$ -$c_{n} = h_{nu}\left(z_n\right)$. Cell-type prediction vectors are averaged weighted by the distance of a specific cell -and differential computation is performed. diff --git a/docs/user_guide/models/scviva.md b/docs/user_guide/models/scviva.md deleted file mode 100644 index 68c7ab993f..0000000000 --- a/docs/user_guide/models/scviva.md +++ /dev/null @@ -1,154 +0,0 @@ -# scVIVA - -**scVIVA** {cite:p}`Levy25` (Python class {class}`~scvi.external.SCVIVA`) is a generative model of single-cell resolved spatial -transcriptomics that can subsequently be used for many common downstream tasks. - -The advantages of scVIVA are: - -- Provides a probabilistic low-dimensional representation of the state of each cell that is corrected for batch effects - and captures its gene expression profile and its environment. -- Enables differential expression analysis across niches while accounting for wrong assignment of molecules to cells. -- Scalable to very large datasets (>1 million cells). - -The limitations of scVIVA include: - -- Effectively requires a GPU for fast inference. -- Latent space is not interpretable, unlike that of a linear method. -- Assumes single cells are observed and do not work with low-resolution ST like Visium or Slide-Seq. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorials: - -- {doc}`/tutorials/notebooks/spatial/scVIVA_tutorial` -``` - -## Preliminaries - -scVIVA takes as input spatially resolved scRNA data. In addition to the gene expression matrix ${X}$ with $N$ cells and $G$ genes, -it requires for each cell $n$: -- the spatial coordinates of the cell $y_n$ -- the cell type assignment (possibly coarse) $c_n \in \{1, ..., T\}$ -- the batch assignment $s_n$. - - -As preprocessing, we take the $K$ nearest neighbors of a cell to define its niche using the Euclidean distance in physical space. -We characterize the niche by its cell-type composition and gene expression. We denote by ${\alpha_n}$ the $T$-dimensional vector of cell type -proportions among the $K$ nearest neighbors of the cell $n$. Its values are in the probability simplex. -The niche gene expression is defined as the average expression of each cell type present in the niche. -In practice, we leverage gene expression embeddings (PCA, scVI or similar) and characterize a cell type expression profile as the local average -embedding of cells of the same type. The average embeddings are stored in the matrix ${\eta_n} \in \mathbb{R}^{T \times D}$, where $D$ is the embedding dimension. -## Descriptive model - -We propose a latent variable model aiming to capture both gene expression heterogeneity and spatial variation resulting from the micro-environment. -We assume these two sources of variability are both captured by a $P$-dimensional latent variable ($P \ll G$): - -```{math} -:nowrap: true -\begin{align} - z_n \sim \mathbf{MixtureOfGaussians}(\mu_1, ..., \mu_M; \Sigma_1, ..,\Sigma_M; \pi_1, ...,\pi_M) -\end{align} -``` - -We assume that the observed counts for cell $n$ and gene $g$, $x_{ng}$, are generated from the following process: - -```{math} -:nowrap: true -\begin{align} - \rho _n &= f_{w}\left( z_n, s_n \right) \\ - x_{ng} &\sim \mathbf{NegativeBinomial}(\ell_n \rho_n, \theta_g), - \end{align} -``` -where $\rho_n$ is the normalized gene expression, $\ell_n$ is the library size of cell $n$ and $\theta_g$ is the dispersion parameter for gene $g$. See {doc}`/user_guide/models/scvi` for more details. -The cell-type proportions of the cell's $K$ nearest neighbors are obtained as - -```{math} -:nowrap: true -\begin{align} - \alpha_n &\sim \mathbf{Dirichlet}\left( f_{\omega}(z_n) \right), -\end{align} -``` - -Last, we assume that the neighboring cells' average expression profiles are obtained as - -```{math} -:nowrap: true -\begin{equation} -\eta_{nt} \sim -\begin{cases} -\mathcal{N} \left(f_{\nu}^{t}(z_n) \right), & \text{if } \alpha_{t} > 0 \\ -0, & \text{otherwise} -\end{cases} -\end{equation} -``` - -where $t=1,...,T$. $w$, $\omega$ and $\nu$ are neural network parameters. - - -## Inference - -We want to maximize the evidence of the data, which can be decomposed as: - -```{math} -:nowrap: true -\begin{align} - \log p \left( \alpha, x, \eta \mid s \right) = \log p \left(x \mid s \right) + \log p \left( \alpha, \eta \mid x, s \right). -\end{align} -``` - -scVIVA uses variational inference, specifically auto-encoding variational Bayes -(see {doc}`/user_guide/background/variational_inference`) to learn both the model parameters -(the neural network parameters, dispersion parameters, etc.) and an approximate posterior distribution. - -## Tasks - -Here we provide an overview of some of the tasks that scVIVA can perform. Please see {class}`scvi.external.SCVIVA` -for the full API reference. - -### Dimensionality reduction - -For dimensionality reduction, the mean of the approximate posterior $q_\phi(z \mid x)$ is returned by default. -This is achieved using the method: - -``` ->>> adata.obsm["X_scVIVA"] = model.get_latent_representation() -``` - -$\phi$ is a set of parameters corresponding to inference neural networks (encoders). -Users may also return samples from this distribution, as opposed to the mean, by passing the argument `give_mean=False`. - -### Estimation of normalized expression - -In {meth}`~scvi.external.SCVIVA.get_normalized_expression` scVIVA returns the expected true expression value of $\rho_n$ under the approximate posterior. For one cell $n$, this can be written as: - -```{math} -:nowrap: true - -\begin{align} - \mathbb{E}_{q_\phi(z_n \mid x_n)}\left[f_{w}\left(z_{n}, s_n \right) \right] -\end{align} -``` - -### Differential Expression (DE) - -Differential expression analysis is achieved with {meth}`~scvi.external.SCVIVA.differential_expression`. \ -We leverage the lvm-DE method (see {doc}`/user_guide/background/differential_expression`) and adapt it to spatial data by taking into account cell neighborhood expression in a bid to discard false positives due to contamination. \ -Considering two groups of cells $\textit{C1}$ and $\textit{C2}$ corresponding to different spatial contexts (for instance, astrocytes in two brain regions), the goal is to determine which genes have different expression levels between the two groups. When setting `niche_mode="true"`, we compute the group spatial neighborhoods $\textit{N1}$ and $\textit{N2}$, which are the spatial nearest neighbors of a different type than the cells in $\textit{C1}$, and $\textit{C2}$ respectively. - - -To determine the upregulated genes of $\textit{C1 vs C2}$, we compute DE between $\{\textit{C1, C2}\}$, $\{\textit{N1, C2}\}$ and $\{\textit{C1, N1}\}$: using lvm-DE, we test differences in expression levels $\rho_{n}$ to compute Log-Fold Changes (LFC). \ -The significantly upregulated genes for $\textit{C1, N1}$ define a set of local cell type markers, denoted $\mathcal{S}_1$. Conversely, if a gene is both higher expressed in $\textit{N1}$ compared to $\textit{C1}$ and $\textit{C1}$ compared to $\textit{C2}$, it is likely that the increased expression in $\textit{C1}$ is spurious. -We argue that the probability of a gene being a $\textit{local marker}$ could be a relevant score to filter spurious genes. To compute this score, we considered the upregulation of a gene in one group relative to the upregulation in its neighborhood: a local marker $g$ should verify - -```{math} -:nowrap: true -\begin{align} - \mathit{LFC^{~g}_{C1~vs~C2}} > \mathit{LFC^{~g}_{N1~vs~C2}}, -\end{align} -``` - -which means that the signal comes from cells in $\textit{C1}$ rather than their neighbors $\textit{N1}$. \ -We select genes for which $\mathit{LFC_{C1~vs~C2}} > 0$ and use the genes $\mathcal{S}_1$ as truly differentially expressed. We also define $\mathcal{N}_1 = \{g|\mathit{LFC^{~g}_{C1~vs~C2}} > 0,~g \notin \mathcal{S}_1 \}$. \ -We train a Gaussian process classifier on $\mathbf{X} = [LFC_{C1~vs~C2}~,~LFC_{N1~vs~C2}]$ to classify between the $\textit{local markers}$ $\mathcal{S}_1$ and the $\textit{neighborhood genes}$ $\mathcal{N}_1$. Once fitted, the classifier returns a local marker probability $p_g=\mathit{p}(g \in \mathcal{S}_1 | \mathbf{X})$ for each gene $g$, that we can compare to a given threshold $\tau$ to filter the neighborhood genes. diff --git a/docs/user_guide/models/stereoscope.md b/docs/user_guide/models/stereoscope.md deleted file mode 100644 index bd59aadb84..0000000000 --- a/docs/user_guide/models/stereoscope.md +++ /dev/null @@ -1,185 +0,0 @@ -# Stereoscope - -**Stereoscope** {cite:p}`Andersson20` (Python classes {class}`~scvi.external.RNAStereoscope` and {class}`~scvi.external.SpatialStereoscope`) posits a probabilistic model of spatial transcriptomics and an associated -method for the deconvoluton of cell type profiles using a single-cell RNA sequencing reference dataset. - -The advantages of Stereoscope are: - -- Can stratify cells into discrete cell types. -- Scalable to very large datasets (>1 million cells). - -The limitations of Stereoscope include: - -- Effectively requires a GPU for fast inference. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorial: - -- {doc}`/tutorials/notebooks/spatial/stereoscope_heart_LV_tutorial` -``` - -## Preliminaries - -Stereoscope requires training two latent variable models (LVMs): one for the single-cell reference -dataset and one for the spatial transcriptomics dataset, which incorporates the learned parameters of the -single-cell reference LVM. The first LVM takes in as input a scRNA-seq gene expression matrix of UMI counts -$Y$ with $N$ cells and $G$ genes, along with a vector of cell type labels $\vec{z}$. -Subsequently, the second LVM takes in the learned parameters of the first LVM, along with a spatial gene -expression matrix $X$ with $S$ spots and $G$ genes. - -## Generative process - -### Single-cell reference LVM - -For cell $c$, the LVM assumes an observed discrete cell type label $z_c$ and models -the UMI count observation for a given gene $g$ as a negative binomial distribution. This LVM posits that the observed -UMI counts for cell $c$ and gene $g$ are generated by the following process: - -```{math} -:nowrap: true - -\begin{align} - y_{gc} &\sim \textrm{NegativeBinomial}(s_{c}r_{gz}, p_{g}) \tag{1} \\ -\end{align} -``` - -where $s_c = \sum_{g\in G} y_{gc}$ is the observed library size of the cell, -$r_{gz}$ is the latent rate parameter for the cell type $z_c$ and gene $g$, -and $p_g$ is the latent variable representing the success probability for gene $g$. - -:::{note} -We are using the standard rate-shape parametrization of the negative binomial here, rather than the mean-dispersion -parametrization used in {doc}`/user_guide/models/scvi`. This is to take advantage of the additive property of -negative binomial distributions sharing the same shape parameter. In this case, the rate parameter for the -negative binomial modeling the expression counts for a given gene and spot is equivalent to the sum of the rate -parameters for each contributing cell. -::: - -This generative process is also summarized in the following graphical model: - -:::{figure} figures/stsc_scLVM_graphical_model.svg -:align: center -:alt: single-cell reference LVM graphical model -:class: img-fluid - -single-cell reference LVM graphical model. -::: - -The latent variables for the single-cell reference LVM, along with their description, are summarized in the following table: - -```{eval-rst} -.. list-table:: - :widths: 20 90 15 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable (if different) - * - :math:`r_{gz} \in (0, \infty)` - - Rate parameter for the negative binomial distribution. - - ``px_scale`` - * - :math:`p_g \in [0, 1]` - - Shape parameter for the negative binomial distribution. - - ``px_o`` :math:`:= \log \left( \frac{p_g}{1 - p_g} \right)` -``` - -### Spatial transcriptomics LVM - -For the second LVM, we also model the expression counts with a $\mathrm{NegativeBinomial}$. However, -for spatial data, we assume that each spot $s$ has expression $x_s$ composed of a bulk of cell types, with -cell type abundance, $v_{sz}$, for each cell type $z$. We assume that for a given spot $s$ and gene $g$, -the observation is generated by the following process: - -```{math} -:nowrap: true - -\begin{align} - x_{sg} &\sim \mathrm{NegativeBinomial}(\beta_g\sum_{z\in Z}v_{sz}r_{gz}, p_g) \tag{2} \\ -\end{align} -``` - -where $\beta_g$ is a gene-specific correction term for technical differences. -The parameters $r_{gz}$ and $p_g$ are the learned parameters from the first LVM. - -An additional latent variable, $\eta_g$, is incorporated into the aggregated cell expression profile -as a dummy cell type to represent gene-specific noise. The dummy cell type's expression profile is distributed -as $\varepsilon_g := \mathrm{Softplus}(\eta_g)$ where $\eta_g \sim \mathrm{Normal}(0, 1)$ to avoid the model -from incorrectly assigning explanatory power to this term. -Like the other cell types, there is an associated cell type abundance parameter $\gamma_s$ associated with $\varepsilon$. - -This generative process is also summarized in the following graphical model: - -:::{figure} figures/stsc_stLVM_graphical_model.svg -:align: center -:alt: spatial transcriptomics LVM graphical model -:class: img-fluid - -spatial transcriptomics LVM graphical model. -::: - -The latent variables for the spatial transcriptomics LVM, along with their description, are summarized in the following table: - -```{eval-rst} -.. list-table:: - :widths: 20 90 15 - :header-rows: 1 - - * - Latent variable - - Description - - Code variable (if different) - * - :math:`v_{sz} \in (0, \infty)` - - Spot-specific cell type abundance. The code variable ``v_ind`` also incorporates the - the abundance term, :math:`\gamma_s` for the dummy noise cell type, :math:`\varepsilon`. - - ``v_ind`` - * - :math:`\eta_g \in (-\infty, \infty)` - - Gene-specific noise. Incorporated into the model as :math:`\varepsilon_g := \mathrm{Softplus}(\eta_g)`. - - ``eta`` - * - :math:`\beta_g \in (0, \infty)` - - Correction term for technological differences. - - ``beta`` - * - :math:`r_{gz} \in (0, \infty)` - - Rate parameter for the negative binomial distribution shared from the single-cell reference LVM. - - ``w`` - * - :math:`p_g \in [0,1]` - - Shape parameter for the negative binomial distribution shared from the single-cell reference LVM. - - ``px_o`` :math:`:= \log \left( \frac{p_g}{1 - p_g} \right)` - -``` - -## Inference - -### Single-cell reference LVM - -Stereoscope uses maximum likelihood estimation to estimate the parameters of the first LVM w.r.t. The negative binomial model of -UMI observations. This is achieved via stochastic gradient ascent on the likelihood function using the Pytorch framework. - -### Spatial transcriptomics LVM - -For the spatial transcriptomics LVM, Stereoscope uses MAP inference to estimate the parameters specific to the model. To be exact, -the only parameter given a non-uniform prior is $\eta_g$ that is posited as a gene-specific random effect distributed by a standard -Normal prior. Note, the $r_{gz}$ and $p_g$ parameters are not inferred in this step, but held fixed as the parameters shared by the -single-cell reference LVM. - -## Tasks - -### Cell type deconvolution - -Once the model is trained, one can retrieve the estimated cell type proportions in each spot using the method: - -``` ->>> proportions = spatial_model.get_proportions() ->>> st_adata.obsm["proportions"] = proportions -``` - -These proportions are computed by normalizing across all learned cell type abundances, $v_{sz}$, for a given spot $s$. -I.e., the estimated proportion of cell type $z$ for spot $s$ is $\frac{v_{sz}}{\sum_{z'} v_{sz'}}$. - -Subsequently, for a given cell type, users can plot a heatmap of the cell type proportions spatially using scanpy with: - -``` ->>> import scanpy as sc ->>> sc.p1.embedding(st_adata, basis="location", color="B cells") -``` diff --git a/docs/user_guide/models/tangram.md b/docs/user_guide/models/tangram.md deleted file mode 100644 index 30a0ce1119..0000000000 --- a/docs/user_guide/models/tangram.md +++ /dev/null @@ -1,122 +0,0 @@ -# Tangram - -:::{note} -This model is deprecated starting v1.5. -::: - -**Tangram** {cite:p}`Biancalani21` (Python class {class}`~scvi.external.Tangram`) maps single-cell RNA-seq data to spatial data, permitting deconvolution of cell types in spatial data like Visium. - -This is a reimplementation of Tangram, which can originally be found [here](https://github.com/broadinstitute/Tangram). - -The advantages of Tangram are: - -- It maps single-cell transcriptomes onto spatial observations with a directly - interpretable mapping matrix. -- It can project cell annotations, such as cell types, from single-cell data to spatial - data. -- It can project gene expression from the single-cell reference into spatial - coordinates. -- The scvi-tools implementation supports Tangram's `"cells"` and `"constrained"` - modes. - -The limitations of Tangram include: - -- The scvi-tools model page is deprecated starting v1.5. -- It requires matched genes between the single-cell and spatial modalities used for - training. -- Training is not mini-batched in the current implementation, so memory use depends on - the number of single-cell observations and spatial observations. -- Tangram is an optimization-based mapping model, not a generative model with posterior - sampling. - -## Overview - -Tangram learns a matrix $M$ with shape ($n_{sc} \times n_{sp}$), in which each row sums to 1. Thus, this matrix can be viewed as a map from single cells to the spatial observations. - -:::{note} -Starting scVI-Tools v1.5 this model is part of scVIVA-Tools, and no longer being maintained here. -::: - -```{topic} Tutorials: - -- {doc}`/tutorials/notebooks/spatial/tangram_scvi_tools` -``` - -## Preliminaries - -Tangram is registered with {meth}`~scvi.external.Tangram.setup_mudata`, not -`setup_anndata`. The input is a MuData object containing a single-cell modality and a -spatial modality. The two modalities used for training must contain the same genes in the -same order. - -The `modalities` argument tells Tangram which MuData modality contains each registered -field. A typical setup registers: - -- `sc_layer`, the single-cell expression matrix used for mapping, -- `sp_layer`, the spatial expression matrix used as the target, and -- `density_prior_key`, an optional spatial observation column containing a density prior. - -If a density prior is supplied, it must sum to 1. The tutorial computes a density prior -from estimated cell counts in spatial observations. - -## Mapping Objective - -The mapping matrix $M$ is parameterized by unconstrained trainable weights and converted -to a row-stochastic matrix with a softmax. The predicted spatial expression matrix is: - -```{math} -:nowrap: true - -\begin{align} - \hat{X}_{sp} = M^\top X_{sc}, -\end{align} -``` - -where $X_{sc}$ is the single-cell expression matrix and $\hat{X}_{sp}$ is the spatial -expression predicted from mapped single cells. - -Tangram optimizes a loss that rewards agreement between measured spatial expression and -the predicted spatial expression. The default expression term uses gene-wise cosine -similarity, and an optional spatial-observation-wise cosine term can also be enabled. If a -density prior is registered, the loss includes a KL-divergence term between the predicted -spatial density and the supplied prior. - -In constrained mode, Tangram also learns a cell filter and requires `target_count`. The -loss then includes a count term that encourages the selected number of cells to match -`target_count`, plus a filter regularizer. - -## Tasks - -Here we provide an overview of common tasks. Please see {class}`~scvi.external.Tangram` -for the full API reference. - -### Mapping Matrix - -After training, {meth}`~scvi.external.Tangram.get_mapper_matrix` returns the mapping -matrix with shape `(n_obs_sc, n_obs_sp)`: - -``` ->>> mapper = model.get_mapper_matrix() -``` - -Each row contains a probability distribution over spatial observations for one -single-cell observation. - -### Projection of Cell Annotations - -{meth}`~scvi.external.Tangram.project_cell_annotations` uses the mapping matrix to -project categorical single-cell labels, such as cell types, onto spatial observations: - -``` ->>> adata_sp.obsm["tangram_ct_pred"] = model.project_cell_annotations( -... adata_sc, adata_sp, mapper, adata_sc.obs["cell_type"] -... ) -``` - -The returned DataFrame has spatial observations as rows and label categories as columns. - -### Projection of Genes - -{meth}`~scvi.external.Tangram.project_genes` multiplies the mapping matrix by the -single-cell expression matrix and returns an AnnData object containing projected gene -expression in spatial coordinates. diff --git a/pyproject.toml b/pyproject.toml index 6cfaaedffa..f1ec1b3c2d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = ["hatchling"] [project] name = "scvi-tools" -version = "1.5.0.post1" +version = "1.5.0" description = "Deep probabilistic analysis of single-cell omics data." readme = "README.md" requires-python = ">=3.12" @@ -78,10 +78,10 @@ docs = [ "myst-nb", "sphinx-autodoc-typehints", ] -docsbuild = ["scvi-tools[docs,autotune,hub,diagvi]","mlx"] +docsbuild = ["scvi-tools[docs,autotune,hub]","mlx"] # scvi.autotune -autotune = ["hyperopt>=0.2", "setuptools>=77.0.3,<82", "ray[tune]", "scib-metrics", "muon"] +autotune = ["hyperopt>=0.2", "ray[tune]", "scib-metrics", "muon"] # scvi.hub dependencies hub = ["huggingface_hub", "dvc[s3]", "boto3"] # scvi.data.add_dna_sequence @@ -94,8 +94,6 @@ parallel = ["dask[array]", "zarr"] interpretability = ["captum", "shap", "decoupler"] # for custom dataloders dataloaders = ["lamindb>=1.12.1", "cellxgene-census", "tiledbsoma", "tiledbsoma_ml", "torchdata","annbatch[zarrs]>=0.2.0"] -# for diagvi -diagvi = ["torch_geometric", "geomloss"] # for mlflow mlflow = ["mlflow","psutil","GPUtil","nvidia-ml-py"] # for rapids @@ -104,12 +102,11 @@ rapids-cuda12 = [ "cugraph-cu12>=24", "cuml-cu12>=24", "cupy-cuda12x", "rapids-s rapids-cuda13 = [ "cugraph-cu13>=24", "cuml-cu13>=24", "cupy-cuda13x", "rapids-singlecell-cu13[rapids]" ] optional = [ - "scvi-tools[autotune,mlflow,hub,file_sharing,regseq,parallel,interpretability,diagvi]", + "scvi-tools[autotune,mlflow,hub,file_sharing,regseq,parallel,interpretability]", "igraph","leidenalg","pynndescent", ] tutorials = [ "biomart", - "cell2location", "jupyter", "seaborn", "matplotlib", diff --git a/src/scvi/external/__init__.py b/src/scvi/external/__init__.py index 931e5ef7f3..83ae5741d4 100644 --- a/src/scvi/external/__init__.py +++ b/src/scvi/external/__init__.py @@ -2,35 +2,25 @@ from .contrastivevi import ContrastiveVI from .cytovi import CYTOVI from .decipher import Decipher -from .diagvi import DIAGVI from .drvi import DRVI -from .gimvi import GIMVI from .joint_embedding_scvi import JointEmbeddingSCVI, JointEmbeddingVAE from .methylvi import METHYLANVI, METHYLVI from .mrvi import MRVI from .poissonvi import POISSONVI -from .resolvi import RESOLVI from .scar import SCAR from .scbasset import SCBASSET -from .scviva import SCVIVA from .solo import SOLO -from .stereoscope import RNAStereoscope, SpatialStereoscope from .sysvi import SysVI -from .tangram import Tangram from .totalanvi import TOTALANVI from .velovi import VELOVI __all__ = [ "SCAR", "SOLO", - "GIMVI", "JointEmbeddingSCVI", "JointEmbeddingVAE", "Decipher", - "RNAStereoscope", - "SpatialStereoscope", "CellAssign", - "Tangram", "TOTALANVI", "SCBASSET", "POISSONVI", @@ -40,9 +30,6 @@ "MRVI", "METHYLVI", "METHYLANVI", - "RESOLVI", - "SCVIVA", "CYTOVI", - "DIAGVI", "DRVI", ] diff --git a/src/scvi/external/diagvi/__init__.py b/src/scvi/external/diagvi/__init__.py deleted file mode 100644 index cb7a89ed06..0000000000 --- a/src/scvi/external/diagvi/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""DIAGVI model for multi-modal integration with guidance graphs.""" - -from ._model import DIAGVI -from ._module import DIAGVAE -from ._task import DiagTrainingPlan - -__all__ = ["DIAGVI", "DIAGVAE", "DiagTrainingPlan"] diff --git a/src/scvi/external/diagvi/_base_components.py b/src/scvi/external/diagvi/_base_components.py deleted file mode 100644 index fd5a71122a..0000000000 --- a/src/scvi/external/diagvi/_base_components.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Base neural network components for DIAGVI model.""" - -from __future__ import annotations - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.distributions import Normal - -from scvi.utils import dependencies - -EPS = 1e-8 - - -class DecoderSinglePathway(nn.Module): - """Single-pathway decoder for unimodal likelihoods. - - Outputs a single set of parameters (scale, r, rate, dropout). - Use for: nb, zinb, normal, log1pnormal, ziln, zig. - - Parameters - ---------- - n_output - Number of output features. - n_batches - Number of batches. - normalize - If True, apply softmax normalization to output scale (for count data). - If False, output raw scale values (for continuous data). - Default is True for backward compatibility. - """ - - def __init__( - self, - n_output: int, - n_batches: int, - normalize: bool = True, - ): - super().__init__() - self.n_output = n_output - self.n_batches = n_batches - self.normalize = normalize - - self.scale_lin = nn.Parameter(torch.zeros(n_batches, n_output)) - self.bias = nn.Parameter(torch.zeros(n_batches, n_output)) - self.log_theta = nn.Parameter(torch.zeros(n_batches, n_output)) - - self.px_dropout_param = nn.Parameter(torch.randn(n_output) * 0.01) - - def forward( - self, - u: torch.Tensor, - l: torch.Tensor, - batch_index: torch.Tensor, - v: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Forward pass through the decoder. - - Parameters - ---------- - u - Cell latent representation, shape (n_cells, n_latent). - l - Library size (log-scale), shape (n_cells, 1). - batch_index - Batch indices for cells, shape (n_cells,) or (n_cells, 1). - v - Feature latent representation, shape (n_features, n_latent). - - Returns - ------- - tuple of 4 :class:`torch.Tensor` - - px_scale: Normalized scale parameters, shape (n_cells, n_features). - - px_r: Overdispersion parameter, shape (n_cells, n_features). - - px_rate: Rate parameters (mean), shape (n_cells, n_features). - - px_dropout: Dropout/zero-inflation probabilities, shape (n_features,). - """ - # Check batch index dimension - if batch_index.dim() > 1: - batch_index = batch_index.squeeze(-1) - - if (batch_index.max() >= self.bias.shape[0]) or (batch_index.min() < 0): - raise IndexError( - f"Batch index out of bounds: valid range is [0, {self.bias.shape[0] - 1}]" - ) - - # Get batch-specific parameters - scale = F.softplus(self.scale_lin[batch_index]) - bias = self.bias[batch_index] - log_theta = self.log_theta[batch_index] - - # Compute decoder output parameters - raw_px_scale = scale * (u @ v.T) + bias - - if self.normalize: - # Count data: softmax produces proportions, library scales to counts - px_scale = torch.softmax(raw_px_scale, dim=-1) - px_rate = torch.exp(l) * px_scale - else: - # Continuous data: no normalization, scale is positive, rate is unbounded - px_scale = F.softplus(raw_px_scale) + EPS - px_rate = raw_px_scale - - px_dropout = F.softplus(self.px_dropout_param) - px_r = log_theta - - return px_scale, px_r, px_rate, px_dropout - - -class DecoderDualPathway(nn.Module): - """Dual-pathway decoder for mixture likelihoods. - - Outputs two sets of parameters for foreground/background mixture models. - Use for: nbmixture. - - Parameters - ---------- - n_output - Number of output features. - n_batches - Number of batches. - normalize - If True, apply softmax normalization to output scales (for count data). - If False, output raw scale values (for continuous data). - Default is True for backward compatibility. - """ - - def __init__( - self, - n_output: int, - n_batches: int, - normalize: bool = True, - ): - super().__init__() - self.n_output = n_output - self.n_batches = n_batches - self.normalize = normalize - - self.scale_lin = nn.Parameter(torch.zeros(n_batches, n_output)) - - self.bias1 = nn.Parameter(torch.zeros(n_batches, n_output)) - self.bias2 = nn.Parameter(torch.zeros(n_batches, n_output)) - - self.log_theta = nn.Parameter(torch.zeros(n_batches, n_output)) - - def forward( - self, - u: torch.Tensor, - l: torch.Tensor, - batch_index: torch.Tensor, - v: torch.Tensor, - ) -> tuple[ - tuple[torch.Tensor, torch.Tensor], - torch.Tensor, - tuple[torch.Tensor, torch.Tensor], - torch.Tensor, - ]: - """Forward pass through the decoder. - - Parameters - ---------- - u - Cell latent representation, shape (n_cells, n_latent). - l - Library size (log-scale), shape (n_cells, 1). - batch_index - Batch indices for cells, shape (n_cells,) or (n_cells, 1). - v - Feature latent representation, shape (n_features, n_latent). - - Returns - ------- - tuple of 4 elements - - scales: Tuple of (foreground, background) scale parameters. - Each has shape (n_cells, n_features). - - px_r: Overdispersion parameter, shape (n_cells, n_features). - - rates: Tuple of (foreground, background) rate parameters (mean). - Each has shape (n_cells, n_features). - - mixture_logits: Logits for mixture component probabilities, - shape (n_cells, n_features). - """ - # Check batch index dimension - if batch_index.dim() > 1: - batch_index = batch_index.squeeze(-1) - - if (batch_index.max() >= self.bias1.shape[0]) or (batch_index.min() < 0): - raise IndexError( - f"Batch index out of bounds: valid range is [0, {self.bias1.shape[0] - 1}]" - ) - - # Get batch-specific parameters - scale = F.softplus(self.scale_lin[batch_index]) - bias1 = self.bias1[batch_index] - bias2 = self.bias2[batch_index] - log_theta = self.log_theta[batch_index] - - # Compute decoder output parameters for both pathways - raw_px_scale_1 = scale * (u @ v.T) + bias1 - raw_px_scale_2 = scale * (u @ v.T) + bias2 - - if self.normalize: - # Count data: softmax produces proportions, library scales to counts - px_scale_1 = torch.softmax(raw_px_scale_1, dim=-1) - px_scale_2 = torch.softmax(raw_px_scale_2, dim=-1) - px_rate_1 = torch.exp(l) * px_scale_1 - px_rate_2 = torch.exp(l) * px_scale_2 - else: - # Continuous data: no normalization, scales are positive, rates are unbounded - px_scale_1 = F.softplus(raw_px_scale_1) + EPS - px_scale_2 = F.softplus(raw_px_scale_2) + EPS - px_rate_1 = raw_px_scale_1 - px_rate_2 = raw_px_scale_2 - - mixture_logits = raw_px_scale_1 - raw_px_scale_2 - px_r = log_theta - - return (px_scale_1, px_scale_2), px_r, (px_rate_1, px_rate_2), mixture_logits - - -class GraphEncoder(nn.Module): - """Graph convolutional encoder for feature embeddings. - - Encodes feature nodes in the guidance graph to learn feature - embeddings that capture cross-modality relationships. - - Parameters - ---------- - vnum - Number of nodes (features) in the graph. - out_features - Dimensionality of the output feature embeddings. - """ - - @dependencies("torch_geometric") - def __init__(self, vnum: int, out_features: int): - import torch_geometric - - super().__init__() - self.vrepr = nn.Parameter(torch.zeros(vnum, out_features)) - self.conv = torch_geometric.nn.GCNConv(out_features, out_features) - self.loc = nn.Linear(out_features, out_features) - self.std_lin = nn.Linear(out_features, out_features) - - def forward(self, edge_index: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Forward pass through the graph encoder. - - Parameters - ---------- - edge_index - Edge indices for the graph, shape (2, n_edges). - - Returns - ------- - tuple of 3 :class:`torch.Tensor` - - z: Sampled latent representation from the normal distribution, - shape (n_nodes, out_features). - - mu: Mean of the latent distribution, shape (n_nodes, out_features). - - logvar: Log variance of the latent distribution, shape (n_nodes, out_features). - """ - h = self.conv(self.vrepr, edge_index) - loc = self.loc(h) - std = F.softplus(self.std_lin(h)) + EPS - - dist = Normal(loc, std) - mu = dist.loc - std = dist.scale - logvar = torch.log(std**2) - z = dist.rsample() - - return z, mu, logvar diff --git a/src/scvi/external/diagvi/_model.py b/src/scvi/external/diagvi/_model.py index ec4552047d..e69de29bb2 100644 --- a/src/scvi/external/diagvi/_model.py +++ b/src/scvi/external/diagvi/_model.py @@ -1,1237 +0,0 @@ -"""DIAGVI model for multi-modal integration with guidance graphs.""" - -from __future__ import annotations - -import logging -import os -import warnings -from typing import TYPE_CHECKING, Literal - -import numpy as np -import scipy.sparse -import torch -from mudata import MuData - -from scvi import REGISTRY_KEYS, settings -from scvi.data import AnnDataManager -from scvi.data._constants import _MODEL_NAME_KEY, _SETUP_ARGS_KEY -from scvi.data.fields import CategoricalObsField, LabelsWithUnlabeledObsField, LayerField -from scvi.dataloaders import DataSplitter -from scvi.external.diagvi._utils import ( - _check_guidance_graph_consistency, - _construct_guidance_graph, - _load_saved_diagvi_files, -) -from scvi.external.gimvi._task import CyclicMultiDataLoader as TrainDL -from scvi.model._utils import get_max_epochs_heuristic, parse_device_args, use_distributed_sampler -from scvi.model.base import BaseModelClass, VAEMixin -from scvi.module._constants import MODULE_KEYS -from scvi.train import Trainer -from scvi.utils import dependencies, setup_anndata_dsp -from scvi.utils._docstrings import devices_dsp - -from ._module import DIAGVAE -from ._task import DiagTrainingPlan - -if TYPE_CHECKING: - from collections.abc import Sequence - from typing import Literal - - import pandas as pd - from anndata import AnnData - from torch_geometric.data import Data - - -logger = logging.getLogger(__name__) - - -class DIAGVI(BaseModelClass, VAEMixin): - """Diagonal Multi-Modal Integration Variational Inference (DIAGVI) model. - - Integrates multi-modal single-cell data using a guidance graph and supports - semi-supervised learning and GMM priors. - - The model architecture is inspired by GLUE (Cao & Gao, 2022). - This scvi-tools implementation is based on GimVI. - Handling of continuous data in decoder is inspired by CytoVI. - - Parameters - ---------- - adatas - Dictionary mapping input names to AnnData objects. - guidance_graph - Precomputed guidance graph. If None, it will be constructed from the data - by using overlapping feature names. - mapping_df - DataFrame specifying feature correspondences between modalities - (used to compute the guidance graph). - n_latent - Dimensionality of the latent space. - n_hidden - Number of nodes per hidden layer. - n_layers - Number of hidden layers used for encoder and decoder NNs. - dropout_rate - Dropout rate for neural networks. - **model_kwargs - Additional keyword arguments for :class:`~scvi.external.diagvi._module.DIAGVAE`. - - Examples - -------- - >>> adatas = {"rna_data": adata_rna, "protein_data": adata_protein} - >>> model = DIAGVI(adatas) - >>> model.train() - """ - - _module_cls = DIAGVAE - _data_splitter_cls = DataSplitter - _training_plan_cls = DiagTrainingPlan - - def __init__( - self, - adatas: dict[str, AnnData] | MuData, - guidance_graph: Data | None = None, - mapping_df: pd.DataFrame | None = None, - n_latent: int = 50, - n_hidden: int = 256, - n_layers: int = 2, - dropout_rate: float = 0.1, - **model_kwargs, - ): - warnings.warn( - "DiagVI is a spatial transcriptomics model that will be moved to the " - "scvi-tools spatial companion package `scviva-tools` starting in scvi-tools v1.5 and " - "will no longer be supported here. It will be deprecated from scvi-tools in v1.6.", - FutureWarning, - stacklevel=settings.warnings_stacklevel, - ) - super().__init__() - # Handle MuData input by extracting modalities as dict - if isinstance(adatas, MuData): - self.adatas = {mod_key: adatas.mod[mod_key] for mod_key in adatas.mod.keys()} - else: - self.adatas = adatas - self.input_names = list(self.adatas.keys()) - self.adata_managers = { - name: self._get_most_recent_anndata_manager(adata, required=True) - for name, adata in self.adatas.items() - } - - self.registries_ = [] - for adm in self.adata_managers.values(): - self._register_manager_for_instance(adm) - self.registries_.append(adm.registry) - - sum_stats = {name: adm.summary_stats for name, adm in self.adata_managers.items()} - self.summary_stats = sum_stats - n_inputs = {name: s["n_vars"] for name, s in sum_stats.items()} - n_batches = {name: s["n_batch"] for name, s in sum_stats.items()} - n_labels = {name: s["n_labels"] for name, s in sum_stats.items()} - - generative_distributions = { - name: adata.uns["diagvi_likelihood"] for name, adata in self.adatas.items() - } - normalize_lib = { - name: adata.uns["diagvi_normalize_lib"] for name, adata in self.adatas.items() - } - gmm_priors = {name: adata.uns["diagvi_gmm_prior"] for name, adata in self.adatas.items()} - semi_supervised = { - name: adata.uns["diagvi_semi_supervised"] for name, adata in self.adatas.items() - } - n_mixture_components = { - name: adata.uns["diagvi_n_mixture_components"] for name, adata in self.adatas.items() - } - - # Compute guidance graph if not provided, do a sanity check - if guidance_graph is not None: - self.guidance_graph = guidance_graph - else: - self.guidance_graph = _construct_guidance_graph(self.adatas, mapping_df) - _check_guidance_graph_consistency(self.guidance_graph, self.adatas) - self.module = self._module_cls( - n_inputs=n_inputs, - n_batches=n_batches, - n_labels=n_labels, - modality_likelihoods=generative_distributions, - normalize_lib=normalize_lib, - guidance_graph=self.guidance_graph, - use_gmm_prior=gmm_priors, - semi_supervised=semi_supervised, - n_mixture_components=n_mixture_components, - n_latent=n_latent, - n_hidden=n_hidden, - n_layers=n_layers, - dropout_rate=dropout_rate, - **model_kwargs, - ) - - self._model_summary_string = ( - f"DiagVI Model with the following params: input names: {self.input_names}, " - f"n_inputs: {n_inputs}, n_batches: {n_batches}, n_labels: {n_labels}, " - f"semi_supervised: {semi_supervised}, gmm_priors: {gmm_priors}, " - f"generative distributions: {generative_distributions}, " - f"n_latent: {n_latent}." - ) - - self.init_params_ = self._get_init_params(locals()) - logger.info(self._model_summary_string) - - @devices_dsp.dedent - def train( - self, - max_epochs: int | None = None, - lr: float = 1e-3, - batch_size: int = 256, - train_size: float = 0.9, - accelerator: str = "auto", - devices: int | list[int] | str = "auto", - shuffle_set_split: bool = True, - datasplitter_kwargs: dict | None = None, - plan_kwargs: dict | None = None, - **kwargs, - ): - """Train the DIAGVI model. - - Parameters - ---------- - max_epochs - Maximum number of training epochs. If None, a heuristic is used. - lr - Learning rate for optimization. - batch_size - Minibatch size for training. - %(param_accelerator)s - %(param_devices)s - train_size - Proportion of data to use for training (rest for validation). - shuffle_set_split - Whether to shuffle data before splitting into train/validation. - datasplitter_kwargs - Additional keyword arguments for the DataSplitter and DataLoaders. - Can be either: - - A dict of shared kwargs applied to all modalities - - A nested dict mapping modality names to their specific kwargs, - e.g., ``{"rna": {"external_indexing": [train, val, test]}, "protein": {...}}`` - plan_kwargs - Additional keyword arguments for the training plan. - **kwargs - Additional keyword arguments for the Trainer. - """ - # Compute max_epochs if not provided - if max_epochs is None: - min_obs = np.min( - [ - self.summary_stats[self.input_names[0]]["n_cells"], - self.summary_stats[self.input_names[1]]["n_cells"], - ] - ) - max_epochs = get_max_epochs_heuristic(min_obs) - logger.info(f"max_epochs was approximated to {max_epochs}") - - accelerator, devices, device = parse_device_args( - accelerator=accelerator, - devices=devices, - return_device="torch", - ) - - # Extract early stopping parameters from kwargs with defaults - early_stopping = kwargs.pop("early_stopping", True) - early_stopping_monitor = kwargs.pop("early_stopping_monitor", "validation_loss") - early_stopping_min_delta = kwargs.pop("early_stopping_min_delta", 1) - early_stopping_patience = kwargs.pop("early_stopping_patience", 10) - early_stopping_mode = kwargs.pop("early_stopping_mode", "min") - - # Initialize Trainer - self.trainer = Trainer( - max_epochs=max_epochs, - enable_progress_bar=True, - progress_bar_refresh_rate=1, - early_stopping=early_stopping, - early_stopping_monitor=early_stopping_monitor, - early_stopping_min_delta=early_stopping_min_delta, - early_stopping_patience=early_stopping_patience, - early_stopping_mode=early_stopping_mode, - accelerator=accelerator, - devices=devices, - **kwargs, - ) - - # Set up DataLoaders with DataSplitter - validation_size = 1 - train_size - self.train_indices_, self.test_indices_, self.validation_indices_ = [], [], [] - train_dls, test_dls, val_dls = {}, {}, {} - datasplitter_kwargs = datasplitter_kwargs if isinstance(datasplitter_kwargs, dict) else {} - - # Support modality-specific datasplitter_kwargs - # If keys match modality names, treat as modality-specific config - is_modality_specific = any(key in self.input_names for key in datasplitter_kwargs.keys()) - - for name, adm in self.adata_managers.items(): - # Get modality-specific kwargs if provided, otherwise use shared kwargs - if is_modality_specific: - modality_kwargs = datasplitter_kwargs.get(name, {}) - else: - modality_kwargs = datasplitter_kwargs - - ds = self._data_splitter_cls( - adm, - train_size=train_size, - validation_size=validation_size, - batch_size=batch_size, - shuffle_set_split=shuffle_set_split, - distributed_sampler=use_distributed_sampler(kwargs.get("strategy", None)), - **modality_kwargs, - ) - ds.setup() - train_dls[name] = ds.train_dataloader() - test_dls[name] = ds.test_dataloader() - val_dls[name] = ds.val_dataloader() - self.train_indices_.append(ds.train_idx) - self.test_indices_.append(ds.test_idx) - self.validation_indices_.append(ds.val_idx) - - data_loader_kwargs = { - "num_workers": settings.dl_num_workers, - "persistent_workers": getattr(settings, "dl_persistent_workers", False), - "pin_memory": getattr(settings, "dl_pin_memory_gpu_training", False), - } - train_dl = TrainDL(train_dls, **data_loader_kwargs) - val_dl = TrainDL(val_dls, **data_loader_kwargs) - - # Initialize and run training plan - plan_kwargs = plan_kwargs if isinstance(plan_kwargs, dict) else {} - update_dict = {"lr": lr} - plan_kwargs.update(update_dict) - self._training_plan = self._training_plan_cls( - self.module, - **plan_kwargs, - ) - if train_size == 1.0: - self.trainer.fit(self._training_plan, train_dl) - else: - self.trainer.fit( - self._training_plan, train_dataloaders=train_dl, val_dataloaders=val_dl - ) - - try: - self.history_ = self.trainer.logger.history - except AttributeError: - self.history_ = None - - self.module.eval() - self.to_device(device) - self.is_trained_ = True - - @classmethod - @setup_anndata_dsp.dedent - def setup_anndata( - cls, - adata: AnnData, - layer: str | None = None, - batch_key: str | None = None, - labels_key: str | None = None, - likelihood: Literal[ - "nb", "zinb", "nbmixture", "normal", "log1pnormal", "ziln", "zig" - ] = "nb", - normalize_lib: bool = True, - gmm_prior: bool = False, - n_mixture_components: int = 10, - unlabeled_category: str = "unknown", - **kwargs, - ): - """%(summary)s. - - Parameters - ---------- - adata - AnnData object to register. - %(param_batch_key)s - %(param_labels_key)s - %(param_layer)s - likelihood - Likelihood model for this modality (default: 'nb'). - One of: - - 'nb' : Negative Binomial - - 'zinb' : Zero-Inflated Negative Binomial - - 'nbmixture' : Negative Binomial Mixture (for protein data with background/foreground) - - 'normal' : Normal distribution - - 'log1pnormal' : Log1p Normal distribution - - 'ziln' : Zero-Inflated Log Normal distribution - - 'zig' : Zero-Inflated Gamma distribution - normalize_lib - Whether to normalize counts with library size in the model. - gmm_prior - Whether to use a GMM prior for this modality. - n_mixture_components - Number of mixture components for the GMM prior. If labels_key is provided, - this parameter is ignored and set to the number of unique labels in labels_key. - unlabeled_category - Category for unlabeled cells in labels_key. - **kwargs - Additional keyword arguments. - """ - # Checks for valid likelihood and data compatibility - continuous_likelihoods = {"ziln", "zig", "log1pnormal", "normal"} - if likelihood in continuous_likelihoods: - # Check for negative values (invalid for most continuous likelihoods) - data = adata.X if layer is None else adata.layers[layer] - if scipy.sparse.issparse(data): - data_min = data.min() - else: - data_min = np.min(data) - - if data_min < 0 and likelihood not in {"normal"}: - raise ValueError( - f"Likelihood '{likelihood}' requires non-negative data, " - f"but found minimum value {data_min}. Consider using 'normal' instead." - ) - - if scipy.sparse.issparse(adata.X) and not isinstance(adata.X, scipy.sparse.csr_matrix): - adata.X = adata.X.tocsr() - if layer is not None: - if layer not in adata.layers: - raise ValueError(f"Layer '{layer}' not found in adata.layers.") - layer_data = adata.layers[layer] - if scipy.sparse.issparse(layer_data) and not isinstance( - layer_data, scipy.sparse.csr_matrix - ): - adata.layers[layer] = layer_data.tocsr() - - adata.uns["diagvi_likelihood"] = likelihood - adata.uns["diagvi_normalize_lib"] = normalize_lib - adata.uns["diagvi_gmm_prior"] = gmm_prior - - # If labels_key is provided, set semi_supervised to True - # and n_mixture_components to number of unique labels - if labels_key is not None: - semi_supervised = True - actual_n_components = adata.obs[labels_key].nunique() - if n_mixture_components != 10 and n_mixture_components != actual_n_components: - warnings.warn( - f"n_mixture_components={n_mixture_components} is ignored; " - f"using {actual_n_components} labels from labels_key='{labels_key}' " - f"as n_mixture_components instead.", - UserWarning, - stacklevel=settings.warnings_stacklevel, - ) - n_mixture_components = actual_n_components - else: - semi_supervised = False - adata.uns["diagvi_semi_supervised"] = semi_supervised - adata.uns["diagvi_n_mixture_components"] = n_mixture_components - - setup_method_args = cls._get_setup_method_args(**locals()) - if labels_key is None: - label_field = CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key) - else: - label_field = LabelsWithUnlabeledObsField( - REGISTRY_KEYS.LABELS_KEY, labels_key, unlabeled_category - ) - - is_count_data = likelihood in {"nb", "zinb", "nbmixture"} - anndata_fields = [ - LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=is_count_data), - CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key), - label_field, - ] - adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args) - adata_manager.register_fields(adata, **kwargs) - cls.register_manager(adata_manager) - - @classmethod - def setup_mudata( - cls, - mdata: MuData, - modalities: list[str], - batch_key: dict[str, str] | str | None = None, - labels_key: dict[str, str] | str | None = None, - layer: dict[str, str | None] | str | None = None, - likelihood: dict[str, Literal["nb", "zinb", "nbmixture", "normal"]] | str = "nb", - normalize_lib: dict[str, bool] | bool = True, - gmm_prior: dict[str, bool] | bool = False, - n_mixture_components: dict[str, int] | int = 10, - unlabeled_category: dict[str, str] | str = "unknown", - **kwargs, - ): - """Register a MuData object for use with DIAGVI. - - Parameters - ---------- - mdata - MuData object containing multiple modalities. - modalities - List of two modality names from mdata.mod to use. - batch_key - Key(s) in adata.obs for batch annotation. Can be a single string - (applied to all modalities) or a dict mapping modality names to keys. - labels_key - Key(s) in adata.obs for cell type labels. - layer - Layer(s) in adata to use as input. - likelihood - Likelihood model(s) for each modality. Options: 'nb', 'zinb', - 'nbmixture', 'normal'. Default is 'nb'. - normalize_lib - Whether to normalize counts with library size in the model for each modality. - gmm_prior - Whether to use GMM prior for each modality. Default is False. - n_mixture_components - Number of mixture components for the GMM prior for each modality. If labels_key - is provided, this parameter is ignored and set to the number of unique labels - in labels_key. - unlabeled_category - Category name for unlabeled cells. Default is 'unknown'. - **kwargs - Additional keyword arguments passed to setup_anndata. - """ - for mod_key in modalities: - adata = mdata.mod[mod_key] - - batch_key_mod = batch_key[mod_key] if isinstance(batch_key, dict) else batch_key - labels_key_mod = labels_key[mod_key] if isinstance(labels_key, dict) else labels_key - layer_mod = layer[mod_key] if isinstance(layer, dict) else layer - likelihood_mod = likelihood[mod_key] if isinstance(likelihood, dict) else likelihood - normalize_lib_mod = ( - normalize_lib[mod_key] if isinstance(normalize_lib, dict) else normalize_lib - ) - gmm_prior_mod = gmm_prior[mod_key] if isinstance(gmm_prior, dict) else gmm_prior - n_mixture_components_mod = ( - n_mixture_components[mod_key] - if isinstance(n_mixture_components, dict) - else n_mixture_components - ) - unlabeled_category_mod = ( - unlabeled_category[mod_key] - if isinstance(unlabeled_category, dict) - else unlabeled_category - ) - - cls.setup_anndata( - adata=adata, - batch_key=batch_key_mod, - labels_key=labels_key_mod, - layer=layer_mod, - likelihood=likelihood_mod, - normalize_lib=normalize_lib_mod, - gmm_prior=gmm_prior_mod, - n_mixture_components=n_mixture_components_mod, - unlabeled_category=unlabeled_category_mod, - **kwargs, - ) - - @staticmethod - @dependencies("torch_geometric") - def construct_custom_guidance_graph( - input_dict: dict[str, AnnData], - mapping_df: pd.DataFrame, - weight: float = 1.0, - sign: float = 1.0, - ) -> Data: - """Constructs a custom guidance graph defined by the mapping_df. - - Parameters - ---------- - input_dict - Dictionary mapping modality names - (e.g., "scRNAseq", "seqFISH") to their respective AnnData objects. - The keys in this dictionary must match the column names in `mapping_df`. - mapping_df - DataFrame specifying feature correspondences between modalities. - Each column should correspond to a modality name in input_dict, - and each row defines a feature pair. - weight - Edge weight assigned to cross-modality edges. - sign - Edge sign assigned to cross-modality edges. - - Returns - ------- - A PyTorch Geometric Data object representing the guidance graph, - including node features, edge indices, edge weights, edge signs, - and modality-specific feature indices. - - Notes - ----- - - Features are renamed to include the modality name as a suffix. - - Self-loops are added for all features with weight=1.0 and sign=1.0. - - Missing features in the mapping are silently skipped. - """ - from torch_geometric.data import Data - - modality_names = list(input_dict.keys()) - adata1, adata2 = list(input_dict.values()) - - ad_1_ft = [f"{f}_{modality_names[0]}" for f in adata1.var_names] - ad_2_ft = [f"{f}_{modality_names[1]}" for f in adata2.var_names] - - all_features = ad_1_ft + ad_2_ft - feature_to_index = {f: i for i, f in enumerate(all_features)} - - edge_index = [] - edge_weight = [] - edge_sign = [] - - adata_adt_vars = [] - adata_gex_vars = [] - - for ft_pair in range(mapping_df.shape[0]): - pair = mapping_df.iloc[ft_pair, :] - diss_ft = f"{pair[modality_names[0]]}_{modality_names[0]}" - sp_ft = f"{pair[modality_names[1]]}_{modality_names[1]}" - - if diss_ft not in feature_to_index or sp_ft not in feature_to_index: - continue - - adata_adt_vars.append(sp_ft) - adata_gex_vars.append(diss_ft) - - i = feature_to_index[diss_ft] - j = feature_to_index[sp_ft] - - edge_index += [[i, j], [j, i]] - edge_weight += [weight, weight] - edge_sign += [sign, sign] - - for feature in all_features: - i = feature_to_index[feature] - edge_index.append([i, i]) - edge_weight.append(1.0) - edge_sign.append(1.0) - - edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous() - edge_weight = torch.tensor(edge_weight, dtype=torch.float) - edge_sign = torch.tensor(edge_sign, dtype=torch.float) - - x = torch.eye(len(all_features)) - - indices_1 = torch.tensor([feature_to_index[f] for f in ad_1_ft], dtype=torch.long) - indices_2 = torch.tensor([feature_to_index[f] for f in ad_2_ft], dtype=torch.long) - - guidance_graph = Data( - x=x, - edge_index=edge_index, - edge_weight=edge_weight, - edge_sign=edge_sign, - **{ - f"{modality_names[0]}_indices": indices_1, - f"{modality_names[1]}_indices": indices_2, - }, - ) - - return guidance_graph - - @torch.inference_mode() - def get_latent_representation( - self, - adatas: dict[str, AnnData] | list[AnnData] | None = None, - indices: dict[str, Sequence[int]] | None = None, - batch_size: int = 1024, - ) -> dict[str, np.ndarray]: - """Return the latent space embedding for each dataset. - - Parameters - ---------- - adatas - Modalities to compute embeddings for. If None, uses all modalities. - indices - Dictionary mapping modality names to indices of cells to use. - If ``None``, all cells are used for each modality. - batch_size - Minibatch size for data loading. - - Returns - ------- - Dictionary mapping modality names to latent embeddings. - """ - if adatas is None: - adatas = self.adatas - if isinstance(adatas, dict): - input_names = list(adatas.keys()) - adata_list = list(adatas.values()) - else: - input_names = self.input_names - adata_list = adatas - - if indices is None: - indices = dict.fromkeys(input_names) - - self.module.eval() - - latents = {} - for input_name, adata in zip(input_names, adata_list, strict=False): - mode_indices = indices.get(input_name) - scdl = self._make_data_loader( - adata, - indices=mode_indices, - batch_size=batch_size, - ) - latent = [] - for tensors in scdl: - latent.append( - self.module.inference( - **self.module._get_inference_input(tensors), mode=input_name - )["z"] - .cpu() - .detach() - ) - latent = torch.cat(latent).numpy() - latents[input_name] = latent - - return latents - - @torch.inference_mode() - def posterior_predictive_sample( - self, - adatas: dict[str, AnnData] | None = None, - indices: dict[str, Sequence[int]] | None = None, - n_samples: int = 1, - batch_size: int | None = None, - ) -> dict[str, np.ndarray]: - """Generate posterior predictive samples for each modality. - - This method generates samples from the posterior predictive distribution - p(x̂|x) for each modality, which can be used for model criticism via - :class:`~scvi.criticism.PosteriorPredictiveCheck`. - - Parameters - ---------- - adatas - Dictionary mapping modality names to AnnData objects. If ``None``, - defaults to the AnnData objects used to initialize the model. - indices - Dictionary mapping modality names to indices of cells to use. - If ``None``, all cells are used for each modality. - n_samples - Number of posterior predictive samples to generate per cell. - batch_size - Minibatch size for data loading. Defaults to ``scvi.settings.batch_size``. - - Returns - ------- - Dictionary mapping modality names to arrays of shape - ``(n_cells, n_features, n_samples)`` if ``n_samples > 1``, else - ``(n_cells, n_features)``. - """ - if adatas is None: - adatas = self.adatas - - if indices is None: - indices = dict.fromkeys(adatas.keys()) - - # Validate that all provided modalities are valid - if not set(adatas.keys()).issubset(set(self.input_names)): - invalid_keys = set(adatas.keys()) - set(self.input_names) - raise ValueError( - f"Invalid modality names: {invalid_keys}. " - f"Must be subset of model input names: {set(self.input_names)}" - ) - - self.module.eval() - - # Create data loaders for each modality - data_loaders = {} - for name, adata in adatas.items(): - mode_indices = indices.get(name) - dl = self._make_data_loader( - adata, - indices=mode_indices, - batch_size=batch_size or settings.batch_size, - ) - data_loaders[name] = dl - - # Collect samples for each modality - samples_per_modality = {name: [] for name in adatas.keys()} - - # Process each modality separately to handle different numbers of cells - for mode in adatas.keys(): - mode_samples = [] - for tensors in data_loaders[mode]: - # Create a dict with only this modality's tensors - mode_tensors_dict = {mode: tensors} - - # Run module.sample which handles inference + generative + sampling - mode_result = self.module.sample(mode_tensors_dict, n_samples=n_samples) - - mode_sample = mode_result[mode] - - if n_samples > 1: - # mode_sample shape: (n_samples, n_cells, n_features) - # Transpose to (n_cells, n_features, n_samples) - mode_sample = mode_sample.permute(1, 2, 0) - - mode_samples.append(mode_sample.numpy()) - - # Concatenate along cell dimension (axis 0) - samples_per_modality[mode] = np.concatenate(mode_samples, axis=0) - - return samples_per_modality - - @torch.inference_mode() - def get_imputed_values( - self, - query_name: str, - query_adata: AnnData | None = None, - indices: Sequence[int] | None = None, - deterministic: bool = True, - batch_size: int = 2048, - reference_batch: int | str | np.ndarray | None = None, - reference_libsize: float | np.ndarray | None = None, - ) -> np.ndarray: - """Return imputed values for a given query modality. - - Parameters - ---------- - query_name - Name of the query modality. - query_adata - AnnData object for the query modality. If None, uses the registered AnnData. - indices - Indices of cells to use. If ``None``, all cells are used. - deterministic - Whether to use deterministic inference (posterior mean). - batch_size - Minibatch size for data loading. - reference_batch - Reference batch index or array for imputation. Can be a scalar, a - length-1 array, or an array with length equal to the number of - queried cells. - reference_libsize - Reference library size(s) for imputation. Can be a scalar, a - length-1 array, or an array with length equal to the number of - queried cells. - - Returns - ------- - Imputed values. - """ - # Choose query adata according to mode - if query_name not in self.adatas: - raise ValueError(f"`query_name` must be one of {list(self.adatas.keys())}!") - if query_adata is None: - query_adata = self.adatas[query_name] - - # Get batch categories for this modality - batch_manager = self.adata_managers[query_name] - batch_categories = batch_manager.get_state_registry(REGISTRY_KEYS.BATCH_KEY).get( - "categorical_mapping", None - ) - - self.module.eval() - dl = self._make_data_loader(query_adata, indices=indices, batch_size=batch_size) - - n_query_cells = query_adata.n_obs if indices is None else len(np.asarray(indices)) - full_reference_batch = None - if reference_batch is not None: - full_reference_batch = np.asarray(reference_batch).squeeze() - if full_reference_batch.ndim == 0: - full_reference_batch = np.repeat(full_reference_batch[np.newaxis], n_query_cells) - elif full_reference_batch.ndim > 1: - raise ValueError("`reference_batch` cannot be >1 dimensional") - elif full_reference_batch.size == 1: - full_reference_batch = np.repeat(full_reference_batch, n_query_cells) - elif full_reference_batch.size != n_query_cells: - raise ValueError( - "`reference_batch` must be scalar, length 1, or have length equal " - "to the number of queried cells." - ) - - if not np.issubdtype(full_reference_batch.dtype, np.integer): - if batch_categories is None: - raise ValueError( - "`reference_batch` contains non-integer values, but no batch " - "category mapping is available for this modality." - ) - full_reference_batch = np.array( - [np.where(batch_categories == lbl)[0][0] for lbl in full_reference_batch] - ) - - full_reference_libsize = None - if reference_libsize is not None: - full_reference_libsize = np.asarray(reference_libsize).squeeze() - if full_reference_libsize.ndim == 0: - full_reference_libsize = np.repeat( - full_reference_libsize[np.newaxis], n_query_cells - ) - elif full_reference_libsize.ndim > 1: - raise ValueError("`reference_libsize` cannot be >1 dimensional") - elif full_reference_libsize.size == 1: - full_reference_libsize = np.repeat(full_reference_libsize, n_query_cells) - elif full_reference_libsize.size != n_query_cells: - raise ValueError( - "`reference_libsize` must be scalar, length 1, or have length equal " - "to the number of queried cells." - ) - - batch_offset = 0 - lib_offset = 0 - reconstructed_values = [] - for tensor in dl: - inference_output = self.module.inference( - **self.module._get_inference_input(tensor), - mode=query_name, - deterministic=deterministic, - ) - # Use the feature embedding of the other modality for reconstruction - inference_output["v"] = inference_output["v_other"] - - generative_input = self.module._get_generative_input( - tensor, - inference_output, - ) - - # Handle reference batch - batch_size_ = generative_input[MODULE_KEYS.LIBRARY_KEY].shape[0] - if full_reference_batch is not None: - b = full_reference_batch[batch_offset : batch_offset + batch_size_] - if b.size != batch_size_: - raise RuntimeError( - "Internal error while slicing `reference_batch` by minibatch." - ) - batch_offset += batch_size_ - # Use batch index zero if no reference batch is provided - else: - b = np.zeros(batch_size_, dtype=int) - b = np.asarray(b, dtype=int).reshape(-1, 1) - generative_input[MODULE_KEYS.BATCH_INDEX_KEY] = torch.tensor( - b, - dtype=torch.long, - device=generative_input[MODULE_KEYS.LIBRARY_KEY].device, - ) - - # Handle reference libsize - if full_reference_libsize is not None: - l = full_reference_libsize[lib_offset : lib_offset + batch_size_] - if l.size != batch_size_: - raise RuntimeError( - "Internal error while slicing `reference_libsize` by minibatch." - ) - lib_offset += batch_size_ - l = l.reshape((-1, 1)) - generative_input[MODULE_KEYS.LIBRARY_KEY] = torch.tensor( - l, - dtype=generative_input[MODULE_KEYS.LIBRARY_KEY].dtype, - device=generative_input[MODULE_KEYS.LIBRARY_KEY].device, - ) - - # Determine reference modality - reference_names = [m for m in self.input_names if m != query_name] - if len(reference_names) != 1: - raise ValueError("There must be exactly two modalities defined.") - reference_name = reference_names[0] - - generative_output = self.module.generative(**generative_input, mode=reference_name) - # Use distribution mean for correct expected value across all likelihoods - px_dist = generative_output[MODULE_KEYS.PX_KEY] - reconstructed_values.append(px_dist.mean.cpu().detach()) - - reconstructed_value = torch.cat(reconstructed_values).numpy() - return reconstructed_value - - def save( - self, - dir_path: str, - prefix: str | None = None, - overwrite: bool = False, - save_anndata: bool = False, - save_kwargs: dict | None = None, - ): - """Save the DIAGVI model and optionally the AnnData objects. - - Parameters - ---------- - dir_path - Directory path to save the model and AnnData objects. - prefix - Prefix for the saved model file name. - overwrite - Whether to overwrite the directory if it already exists. - save_anndata - Whether to save the AnnData objects used to train the model. - save_kwargs - Additional keyword arguments for torch.save. - - Raises - ------ - ValueError - If dir_path exists and overwrite is False. - """ - if not os.path.exists(dir_path) or overwrite: - os.makedirs(dir_path, exist_ok=overwrite) - else: - raise ValueError( - f"{dir_path} already exists. Please provide an unexisting directory for saving." - ) - - file_name_prefix = prefix or "" - save_kwargs = save_kwargs or {} - - adatas = self.adatas - if save_anndata: - for key in self.input_names: - ad = adatas[key] - save_path = os.path.join(dir_path, f"{file_name_prefix}adata_{key}.h5ad") - ad.write(save_path) - - model_state_dict = self.module.state_dict() - - var_names = {} - for key in self.input_names: - var_names[key] = adatas[key].var_names - - # Get all the user attributes - user_attributes = self._get_user_attributes() - - # Only save the public attributes with _ at the very end - user_attributes = {a[0]: a[1] for a in user_attributes if a[0][-1] == "_"} - - model_save_path = os.path.join(dir_path, f"{file_name_prefix}model.pt") - - torch.save( - { - "model_state_dict": model_state_dict, - "names": self.input_names, - **{f"var_names_{mod}": var_names[mod] for mod in self.input_names}, - "attr_dict": user_attributes, - }, - model_save_path, - **save_kwargs, - ) - - @classmethod - @devices_dsp.dedent - def load( - cls, - dir_path: str, - adatas: dict[str, AnnData] | None = None, - accelerator: str = "auto", - device: int | str = "auto", - prefix: str | None = None, - backup_url: str | None = None, - ) -> DIAGVI: - """Load a saved DIAGVI model from disk. - - Parameters - ---------- - dir_path - Directory path where the model was saved. - adatas - Dictionary mapping modality names to AnnData objects. Required when - the model was saved with ``save_anndata=False``. - %(param_accelerator)s - %(param_device)s - prefix - Prefix used when saving the model. - backup_url - URL to download the model from if not found locally. - - Returns - ------- - Loaded DIAGVI model instance. - - Raises - ------ - ValueError - If the saved model is from a different class or missing setup inputs. - """ - _, _, device = parse_device_args( - accelerator=accelerator, - devices=device, - return_device="torch", - ) - - ( - attr_dict, - var_names_dict, - model_state_dict, - loaded_adatas, - ) = _load_saved_diagvi_files( - dir_path, - prefix=prefix, - map_location=device, - backup_url=backup_url, - ) - - if adatas is not None: - for name, adata in adatas.items(): - loaded_adatas[name] = adata - var_names_dict[name] = adata.var_names - - missing = [name for name, adata in loaded_adatas.items() if adata is None] - if missing: - raise ValueError( - f"AnnData objects for modalities {missing} were not saved with the model. " - f"Pass them via the `adatas` parameter: " - f"DIAGVI.load(dir_path, adatas={{...}})." - ) - - adatas = loaded_adatas - - for mod in adatas.keys(): - saved_var_names = var_names_dict[mod] - user_var_names = adatas[mod].var_names.astype(str) - - if not np.array_equal(saved_var_names, user_var_names): - warnings.warn( - "var_names for adata passed in does not match var_names of adata " - "used to train the model. For valid results, the vars need to be the" - "same and in the same order as the adata used to train the model.", - UserWarning, - stacklevel=settings.warnings_stacklevel, - ) - - registries = attr_dict.pop("registries_") - for adata, registry in zip(adatas, registries, strict=True): - if _MODEL_NAME_KEY in registry and registry[_MODEL_NAME_KEY] != cls.__name__: - raise ValueError("It appears you are loading a model from a different class.") - - if _SETUP_ARGS_KEY not in registry: - raise ValueError( - "Saved model does not contain original setup inputs. " - "Cannot load the original setup." - ) - - cls.setup_anndata(adatas[adata], source_registry=registry, **registry[_SETUP_ARGS_KEY]) - - # Get the parameters for the class init signature - init_params = attr_dict.pop("init_params_") - - # New saving and loading, enable backwards compatibility - if "non_kwargs" in init_params.keys(): - # Grab all the parameters except for kwargs (is a dict) - non_kwargs = init_params["non_kwargs"] - kwargs = init_params["kwargs"] - kwargs = {k: v for (i, j) in kwargs.items() for (k, v) in j.items()} - else: - # Grab all the parameters except for kwargs (is a dict) - non_kwargs = {k: v for k, v in init_params.items() if not isinstance(v, dict)} - kwargs = {k: v for k, v in init_params.items() if isinstance(v, dict)} - kwargs = {k: v for (i, j) in kwargs.items() for (k, v) in j.items()} - - non_kwargs.pop("adatas", None) - kwargs.pop("adatas", None) - - model = cls(adatas, **non_kwargs, **kwargs) - - for attr, val in attr_dict.items(): - setattr(model, attr, val) - - model.module.load_state_dict(model_state_dict) - model.module.eval() - model.to_device(device) - return model - - @torch.inference_mode() - def predict_celltype( - self, - labeled_modality: str, - target_modality: str | None = None, - target_adata: AnnData | None = None, - indices: Sequence[int] | None = None, - batch_size: int = 2048, - deterministic: bool = True, - ) -> dict[str, np.ndarray]: - """Predict cell type labels using a trained classifier. - - Parameters - ---------- - labeled_modality - Name of the modality for which the classifier was trained. A classifier - must have been trained for this modality (i.e., labels_key must have been - provided during setup_anndata for this modality). - target_modality - Name of the modality to predict cell types for. If None, defaults to the - other modality (not labeled_modality) for cross-modal label transfer. - target_adata - AnnData object for the target modality. If None, uses the registered AnnData. - indices - Indices of cells to predict. If ``None``, all cells are used. - batch_size - Minibatch size for data loading. - - Returns - ------- - Dictionary with keys: - - "predictions": Array of predicted label categories (strings) - - "probabilities": Array of class probabilities (n_cells, n_classes) - - "confidence": Array of maximum probability for each cell - - Raises - ------ - ValueError - If labeled_modality is not a valid modality name. - If target_modality is not a valid modality name. - If no classifier was trained for the labeled modality. - """ - # Validate labeled modality - if labeled_modality not in self.input_names: - raise ValueError(f"labeled_modality must be one of {self.input_names}!") - - # Default to the other modality for cross-modal transfer - if target_modality is None: - target_modality = ( - self.input_names[1] - if labeled_modality == self.input_names[0] - else self.input_names[0] - ) - - # Validate target modality - if target_modality not in self.input_names: - raise ValueError( - f"Invalid target_modality: '{target_modality}'. Must be one of {self.input_names}." - ) - - # Get the classifier for the labeled modality - if labeled_modality == self.input_names[0]: - classifier = self.module.classifier_0 - else: - classifier = self.module.classifier_1 - - # Check if classifier exists - if classifier is None: - raise ValueError( - f"No classifier was trained for modality '{labeled_modality}'. " - f"To use this function, you must provide a labels_key when calling " - f"setup_anndata for the labeled modality." - ) - - # Get target adata - if target_adata is None: - target_adata = self.adatas[target_modality] - - # Get label categories from the labeled modality's adata manager - label_registry = self.adata_managers[labeled_modality].get_state_registry( - REGISTRY_KEYS.LABELS_KEY - ) - label_categories = label_registry.get("categorical_mapping", None) - - # Get latent representations for target modality - self.module.eval() - - dl = self._make_data_loader( - target_adata, - indices=indices, - batch_size=batch_size, - ) - - all_logits = [] - for tensors in dl: - z = self.module.inference( - **self.module._get_inference_input(tensors), - mode=target_modality, - deterministic=deterministic, - )["z"] - logits = classifier(z) - all_logits.append(logits.cpu()) - - # Concatenate all logits and compute predictions - all_logits = torch.cat(all_logits, dim=0) - probabilities = torch.nn.functional.softmax(all_logits, dim=-1).numpy() - predicted_indices = torch.argmax(all_logits, dim=1).numpy() - confidence = probabilities.max(axis=1) - predictions = label_categories[predicted_indices] - - return { - "predictions": predictions, - "probabilities": probabilities, - "confidence": confidence, - } diff --git a/src/scvi/external/diagvi/_module.py b/src/scvi/external/diagvi/_module.py deleted file mode 100644 index 3c58543980..0000000000 --- a/src/scvi/external/diagvi/_module.py +++ /dev/null @@ -1,597 +0,0 @@ -"""DIAGVAE module for multi-modal variational autoencoder.""" - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.distributions import Categorical, Independent, MixtureSameFamily, kl_divergence - -from scvi import REGISTRY_KEYS -from scvi.distributions import ( - Log1pNormal, - NegativeBinomial, - NegativeBinomialMixture, - Normal, - ZeroInflatedGamma, - ZeroInflatedLogNormal, - ZeroInflatedNegativeBinomial, -) -from scvi.module import Classifier -from scvi.module._constants import MODULE_KEYS -from scvi.module.base import BaseModuleClass, LossOutput, auto_move_data -from scvi.nn import Encoder - -from ._base_components import DecoderDualPathway, DecoderSinglePathway, GraphEncoder - -if TYPE_CHECKING: - from torch_geometric.data import Data - -logger = logging.getLogger(__name__) - -# Decoder selection based on likelihood function -# Single-pathway: outputs (scale, r, rate, dropout) -# Dual-pathway: outputs ((scale1, scale2), r, (rate1, rate2), mixture_logits) -LIKELIHOOD_TO_DECODER = { - # Single-pathway likelihoods (unimodal distributions) - "nb": DecoderSinglePathway, - "zinb": DecoderSinglePathway, - "normal": DecoderSinglePathway, - "log1pnormal": DecoderSinglePathway, - "ziln": DecoderSinglePathway, - "zig": DecoderSinglePathway, - # Dual-pathway likelihoods (mixture distributions) - "nbmixture": DecoderDualPathway, -} - -# Likelihoods that require softmax normalization (count data) -NORMALIZE_LIKELIHOODS = {"nb", "zinb", "nbmixture"} -NON_NORMALIZE_LIKELIHOODS = {"log1pnormal", "ziln", "zig"} -OTHER_LIKELIHOODS = {"normal"} - - -class DIAGVAE(BaseModuleClass): - """Variational autoencoder module for DIAGVI multi-modal integration. - - Supports GMM priors, semi-supervised classification, and flexible - modality-specific decoders. - - Parameters - ---------- - n_inputs - Number of input features for each modality. - n_batches - Number of batches for each modality. - n_labels - Number of labels/classes for each modality. - modality_likelihoods - Likelihood model for each modality. - One of: 'nb', 'zinb', 'normal', 'log1pnormal', 'zig', 'nbmixture'. - normalize_lib - Whether to normalize counts with library size in the model for each modality. - guidance_graph - Graph object encoding feature correspondences. - use_gmm_prior - Whether to use a GMM prior for each modality. - semi_supervised - Whether to use semi-supervised classification for each modality. - n_mixture_components - Number of mixture components for the GMM prior for each modality. If semi_supervised - is True, this parameter is ignored and set to the number of unique labels in labels_key. - n_latent - Dimensionality of the latent space. - n_hidden - Number of nodes per hidden layer. - n_layers - Number of hidden layers. - dropout_rate - Dropout rate for encoders. - """ - - def __init__( - self, - n_inputs: dict[str, int], - n_batches: dict[str, int], - n_labels: dict[str, int], - modality_likelihoods: dict[str, str], - normalize_lib: dict[str, bool], - guidance_graph: Data, - use_gmm_prior: dict[str, bool], - semi_supervised: dict[str, bool], - n_mixture_components: dict[str, int], - n_latent: int, - n_hidden: int, - n_layers: int, - dropout_rate: float, - ): - super().__init__() - - # Learnable parameters - self.gmm_logits = nn.ParameterDict() - self.gmm_means = nn.ParameterDict() - self.gmm_scales = nn.ParameterDict() - - # Model attributes - self.use_gmm_prior = use_gmm_prior - self.semi_supervised = semi_supervised - self.n_mixture_components = n_mixture_components - self.n_labels = n_labels - self.n_input_list = n_inputs - self.n_batches_list = n_batches - self.modality_likelihoods = modality_likelihoods - self.normalize_lib = normalize_lib - self.guidance_graph = guidance_graph - self.n_latent = n_latent - - self.input_names = list(n_inputs.keys()) - - # GMM prior parameters - for name in use_gmm_prior.keys(): - if self.use_gmm_prior[name]: - k = self.n_mixture_components[name] - # Overwrite the number of mixture components if semi_supervised = True - if self.semi_supervised[name]: - k = self.n_labels[name] - self.gmm_logits[name] = nn.Parameter(torch.zeros(k)) - self.gmm_means[name] = nn.Parameter(torch.randn(k, n_latent)) - self.gmm_scales[name] = nn.Parameter(torch.zeros(k, n_latent)) - - # Encoders - self.encoder_0 = Encoder( - n_input=n_inputs[self.input_names[0]], - n_output=n_latent, - n_hidden=n_hidden, - n_layers=n_layers, - dropout_rate=dropout_rate, - return_dist=True, - ) - - self.encoder_1 = Encoder( - n_input=n_inputs[self.input_names[1]], - n_output=n_latent, - n_hidden=n_hidden, - n_layers=n_layers, - dropout_rate=dropout_rate, - return_dist=True, - ) - - # Decoders - selected based on likelihood function - likelihood_0 = modality_likelihoods[self.input_names[0]] - likelihood_1 = modality_likelihoods[self.input_names[1]] - - # Determine decoder class based on likelihood - decoder_class_0 = LIKELIHOOD_TO_DECODER.get(likelihood_0, DecoderSinglePathway) - decoder_class_1 = LIKELIHOOD_TO_DECODER.get(likelihood_1, DecoderSinglePathway) - - # Determine whether to apply softmax normalization based on likelihood - normalize_0 = likelihood_0 in NORMALIZE_LIKELIHOODS - normalize_1 = likelihood_1 in NORMALIZE_LIKELIHOODS - if likelihood_0 in OTHER_LIKELIHOODS: - normalize_0 = normalize_lib[self.input_names[0]] - if likelihood_1 in OTHER_LIKELIHOODS: - normalize_1 = normalize_lib[self.input_names[1]] - - logger.debug( - f"Decoder for '{self.input_names[0]}' (likelihood={likelihood_0}): " - f"{decoder_class_0.__name__}, normalize={normalize_0}" - ) - self.decoder_0 = decoder_class_0( - n_output=n_inputs[self.input_names[0]], - n_batches=n_batches[self.input_names[0]], - normalize=normalize_0, - ) - - logger.debug( - f"Decoder for '{self.input_names[1]}' (likelihood={likelihood_1}): " - f"{decoder_class_1.__name__}, normalize={normalize_1}" - ) - self.decoder_1 = decoder_class_1( - n_output=n_inputs[self.input_names[1]], - n_batches=n_batches[self.input_names[1]], - normalize=normalize_1, - ) - - # Graph encoder - self.graph_encoder = GraphEncoder( - vnum=n_inputs[self.input_names[0]] + n_inputs[self.input_names[1]], - out_features=self.n_latent, - ) - - # Classifiers for semi-supervised learning - if self.semi_supervised[self.input_names[0]]: - cls_parameters = { - "n_layers": 0, - "n_hidden": 128, - "dropout_rate": 0.0, - } - self.classifier_0 = Classifier( - n_latent, - n_labels=self.n_labels[self.input_names[0]], - use_batch_norm=False, - use_layer_norm=True, - **cls_parameters, - ) - else: - self.classifier_0 = None - - if self.semi_supervised[self.input_names[1]]: - cls_parameters = { - "n_layers": 0, - "n_hidden": 128, - "dropout_rate": 0.0, - } - self.classifier_1 = Classifier( - n_latent, - n_labels=self.n_labels[self.input_names[1]], - use_batch_norm=False, - use_layer_norm=True, - **cls_parameters, - ) - else: - self.classifier_1 = None - - def _get_inference_input( - self, tensors: dict[str, torch.Tensor] - ) -> dict[str, torch.Tensor | None]: - """Prepare input dictionary for the inference step.""" - return { - MODULE_KEYS.X_KEY: tensors[REGISTRY_KEYS.X_KEY], - } - - def _get_generative_input( - self, tensors: dict[str, torch.Tensor], inference_outputs: dict[str, torch.Tensor] - ) -> dict[str, torch.Tensor]: - """Prepare input dictionary for the generative step.""" - 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], - MODULE_KEYS.Y_KEY: tensors[REGISTRY_KEYS.LABELS_KEY], - "v": inference_outputs["v"], - } - - def encode_graph( - self, device: torch.device - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Run the graph encoder, moving the guidance graph to device on first call. - - Parameters - ---------- - device - Target device. - - Returns - ------- - Tuple of (v_all, mu_all, logvar_all) graph embeddings. - """ - if self.guidance_graph.edge_index.device != device: - self.guidance_graph = self.guidance_graph.to(device) - return self.graph_encoder(self.guidance_graph.edge_index) - - @auto_move_data - def inference( - self, - x: torch.Tensor, - mode: str | None = None, - v_all: torch.Tensor | None = None, - mu_all: torch.Tensor | None = None, - logvar_all: torch.Tensor | None = None, - deterministic: bool = False, - ) -> dict[str, torch.Tensor]: - """Run the inference (encoder and graph) step for a given modality. - - Parameters - ---------- - x - Input data tensor. - mode - Name of the modality. - v_all - Pre-computed graph node embeddings. If None, computed from graph encoder. - mu_all - Pre-computed graph encoder means. Required when v_all is provided. - logvar_all - Pre-computed graph encoder log-variances. Required when v_all is provided. - deterministic - Whether to use the mean instead of a sample. - - Returns - ------- - Dictionary of inference outputs, including latent variables and - graph embeddings. - """ - library = torch.log(x.sum(1)).unsqueeze(1) - - if v_all is None: - v_all, mu_all, logvar_all = self.encode_graph(x.device) - - if deterministic: - v_all = mu_all - - graph = self.guidance_graph - v = v_all[getattr(graph, f"{mode}_indices")] - other_mode = [m for m in self.input_names if m != mode][0] - v_other_mod = v_all[getattr(graph, f"{other_mode}_indices")] - - # Data inference - if mode == self.input_names[0]: - qz, z = self.encoder_0(x) - else: - qz, z = self.encoder_1(x) - - if deterministic: - z = qz.loc - - return { - MODULE_KEYS.QZ_KEY: qz, - MODULE_KEYS.Z_KEY: z, - MODULE_KEYS.LIBRARY_KEY: library, - "v": v, - "v_other": v_other_mod, - "v_all": v_all, - "mu_all": mu_all, - "logvar_all": logvar_all, - } - - @auto_move_data - def generative( - self, - z: torch.Tensor, - library: torch.Tensor, - batch_index: torch.Tensor | None = None, - y: torch.Tensor | None = None, - v: torch.Tensor | None = None, - mode: str | None = None, - ) -> dict[str, torch.Tensor]: - """Run the generative (decoder) model for a given modality. - - Parameters - ---------- - z - Latent variable tensor. - library - Library size tensor. - batch_index - Batch index tensor. - y - Cell type labels (for semi-supervised GMM prior). - v - Feature embedding from the graph encoder. - mode - Name of the modality. - - Returns - ------- - Dictionary containing generative model outputs and distributions. - """ - EPS = 1e-8 - - # Get decoder outputs (generative distribution parameters) based on modality - if mode == self.input_names[0]: - px_scale, px_r, px_rate, px_dropout = self.decoder_0(z, library, batch_index, v) - elif mode == self.input_names[1]: - px_scale, px_r, px_rate, px_dropout = self.decoder_1(z, library, batch_index, v) - - # Adjust parameters based on likelihood (count vs continuous) - if self.modality_likelihoods[mode] in NORMALIZE_LIKELIHOODS: - px_r = px_r.exp() - elif self.modality_likelihoods[mode] in OTHER_LIKELIHOODS and self.normalize_lib[mode]: - px_r = px_r.exp() - else: - px_r = F.softplus(px_r) + EPS - - # Construct the appropriate generative distribution p(x) based on likelihood - if self.modality_likelihoods[mode] == "nb": - px = NegativeBinomial(mu=px_rate, theta=px_r) - elif self.modality_likelihoods[mode] == "zinb": - px = ZeroInflatedNegativeBinomial( - mu=px_rate, - theta=px_r, - zi_logits=px_dropout, - scale=px_scale, - ) - elif self.modality_likelihoods[mode] == "normal": - px = Normal(px_rate, px_r, normal_mu=px_scale) - elif self.modality_likelihoods[mode] == "nbmixture": - px = NegativeBinomialMixture( - mu1=px_rate[0], - mu2=px_rate[1], - theta1=px_r, - mixture_logits=px_dropout, - ) - elif self.modality_likelihoods[mode] == "log1pnormal": - px = Log1pNormal(mu=px_rate, sigma=px_r, scale=px_scale) - elif self.modality_likelihoods[mode] == "ziln": - px = ZeroInflatedLogNormal( - mu=px_rate, - scale=px_r, - zi_logits=px_dropout, - normal_mu=px_scale, - ) - elif self.modality_likelihoods[mode] == "zig": - # Clamp concentration and rate to avoid lgamma numerical issues - # with negative values - px = ZeroInflatedGamma( - concentration=torch.clamp(px_rate, min=EPS), - rate=torch.clamp(px_r, min=EPS), - zi_logits=px_dropout, - scale=px_scale, - ) - else: - raise ValueError( - f"Unknown likelihood '{self.modality_likelihoods[mode]}' for modality '{mode}'. " - f"Supported likelihoods: {list(LIKELIHOOD_TO_DECODER.keys())}" - ) - - # Construct variational prior p(z) based on GMM settings - if self.use_gmm_prior[mode]: - logits = self.gmm_logits[mode] - means = self.gmm_means[mode] - scales = torch.exp(self.gmm_scales[mode]) + 1e-4 - if self.semi_supervised[mode]: - # Vectorized one-hot encoding - n_labels = self.n_labels[mode] - y_flat = y.ravel().to(z.device) - - # Clamp invalid labels to 0 (will be zeroed out below) - valid_mask = (y_flat >= 0) & (y_flat < n_labels) - y_clamped = y_flat.clamp(0, n_labels - 1).long() - - # One-hot encode on device - logits_input = F.one_hot(y_clamped, num_classes=n_labels).float() - - # Zero out rows for invalid labels (unlabeled cells) - logits_input = logits_input * valid_mask.unsqueeze(-1).float() - logits = logits + 100 * logits_input - means = means.expand(y.shape[0], -1, -1) - scales = scales.expand(y.shape[0], -1, -1) - cats = Categorical(logits=logits) - normal_dists = Independent(Normal(means, scales), reinterpreted_batch_ndims=1) - pz = MixtureSameFamily(cats, normal_dists) - else: - pz = Normal(torch.zeros_like(z), torch.ones_like(z)) - - return { - MODULE_KEYS.PX_KEY: px, - MODULE_KEYS.PZ_KEY: pz, - "px_rate": px_rate, - } - - def loss( - self, - tensors: dict[str, torch.Tensor], - inference_outputs: dict[str, torch.Tensor], - generative_outputs: dict[str, torch.Tensor], - lam_kl: torch.Tensor | float, - lam_data: torch.Tensor | float, - mode: str | None = None, - ) -> LossOutput: - """Compute the loss for a batch. - - Parameters - ---------- - tensors - Input data tensors. - inference_outputs - Outputs from the inference step. - generative_outputs - Outputs from the generative step. - lam_kl - Weight for the KL divergence term. - lam_data - Weight for the reconstruction loss term. - mode - Name of the modality. - - Returns - ------- - Object containing loss components and metrics. - """ - x = tensors[REGISTRY_KEYS.X_KEY] - n_obs = x.shape[0] - n_var = x.shape[1] - - # Data nll calculation (reconstruction loss) - reconst_loss = -generative_outputs[MODULE_KEYS.PX_KEY].log_prob(x).sum(-1) - reconstruction_loss_norm = torch.mean(reconst_loss) - - # KL divergence calculation - if self.use_gmm_prior[mode]: - kl_div = inference_outputs[MODULE_KEYS.QZ_KEY].log_prob(inference_outputs["z"]).sum( - -1 - ) - generative_outputs[MODULE_KEYS.PZ_KEY].log_prob(inference_outputs["z"]) - else: - kl_div = kl_divergence( - inference_outputs[MODULE_KEYS.QZ_KEY], generative_outputs[MODULE_KEYS.PZ_KEY] - ).sum(dim=-1) - kl_local_norm = torch.sum(kl_div) / (n_obs * n_var) - - # Total loss - loss = lam_data * reconstruction_loss_norm + lam_kl * kl_local_norm - - # Graph inference - mu_all = inference_outputs["mu_all"] - logvar_all = inference_outputs["logvar_all"] - v_all = inference_outputs["v_all"] - - # Semi-supervised classification loss (per modality if specified) - classification_loss = 0.0 - if self.classifier_0 is not None and mode == self.input_names[0]: - y = tensors[REGISTRY_KEYS.LABELS_KEY].ravel().long() - z_mean = inference_outputs[MODULE_KEYS.QZ_KEY].loc - y_logits = self.classifier_0(z_mean) - classification_loss += torch.nn.functional.cross_entropy(y_logits, y, reduction="mean") - if self.classifier_1 is not None and mode == self.input_names[1]: - y = tensors[REGISTRY_KEYS.LABELS_KEY].ravel().long() - z_mean = inference_outputs[MODULE_KEYS.QZ_KEY].loc - y_logits = self.classifier_1(z_mean) - classification_loss += torch.nn.functional.cross_entropy(y_logits, y, reduction="mean") - - return LossOutput( - loss=loss, - reconstruction_loss=reconst_loss, - kl_local=kl_local_norm, - extra_metrics={ - "z": inference_outputs[MODULE_KEYS.Z_KEY], - "mu_all": mu_all, - "logvar_all": logvar_all, - "v_all": v_all, - "guidance_graph": self.guidance_graph, - "classification_loss": classification_loss, - }, - ) - - @torch.inference_mode() - def sample( - self, - tensors: dict[str, dict[str, torch.Tensor]], - n_samples: int = 1, - ) -> dict[str, torch.Tensor]: - """Sample from the generative model for each modality. - - Parameters - ---------- - tensors - Dictionary mapping modality names to their respective input tensors. - Can contain one or both modalities. - n_samples - Number of samples to generate per cell. - - Returns - ------- - Dictionary mapping modality names to sampled tensors of shape - ``(n_samples, n_cells, n_features)`` if ``n_samples > 1``, else - ``(n_cells, n_features)``. - """ - samples = {} - - # Iterate only over modalities present in tensors (allows single-modality sampling) - for mode in tensors.keys(): - mode_tensors = tensors[mode] - - # Run inference for this modality - inference_outputs = self.inference( - **self._get_inference_input(mode_tensors), - mode=mode, - ) - - # Run generative model - generative_input = self._get_generative_input(mode_tensors, inference_outputs) - generative_outputs = self.generative(**generative_input, mode=mode) - - # Get the distribution and sample from it - px = generative_outputs[MODULE_KEYS.PX_KEY] - - if n_samples > 1: - # Sample multiple times and stack - sample_list = [px.sample() for _ in range(n_samples)] - mode_sample = torch.stack(sample_list, dim=0) # (n_samples, n_cells, n_features) - else: - mode_sample = px.sample() # (n_cells, n_features) - - samples[mode] = mode_sample.cpu() - - return samples diff --git a/src/scvi/external/diagvi/_task.py b/src/scvi/external/diagvi/_task.py deleted file mode 100644 index 5408a7471f..0000000000 --- a/src/scvi/external/diagvi/_task.py +++ /dev/null @@ -1,525 +0,0 @@ -"""Training plan for DIAGVI model.""" - -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -import torch - -from scvi import REGISTRY_KEYS -from scvi.external.diagvi._utils import ( - compute_graph_loss, - kl_divergence_graph, -) -from scvi.train import TrainingPlan -from scvi.utils import dependencies - -if TYPE_CHECKING: - from typing import Literal - -logger = logging.getLogger(__name__) - - -def squared_distances(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Compute squared Euclidean distances between point sets.""" - if x.dim() == 2: - D_xx = (x * x).sum(-1).unsqueeze(1) - D_xy = torch.matmul(x, y.permute(1, 0)) - D_yy = (y * y).sum(-1).unsqueeze(0) - elif x.dim() == 3: - D_xx = (x * x).sum(-1).unsqueeze(2) - D_xy = torch.matmul(x, y.permute(0, 2, 1)) - D_yy = (y * y).sum(-1).unsqueeze(1) - else: - raise ValueError(f"Incorrect number of dimensions: {x.shape}") - - return D_xx - 2 * D_xy + D_yy - - -def distances(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Compute Euclidean distances between point sets.""" - return torch.sqrt(torch.clamp_min(squared_distances(x, y), 1e-8)) - - -cost_routines = { - 1: lambda x, y: distances(x, y), - 2: lambda x, y: squared_distances(x, y) / 2, -} - - -def _anneal_param( - current_epoch: int, - max_epochs: int, - init_value: float, - target_value: float, -) -> float: - """Linearly anneal a parameter from init_value to target_value over one third of max_epochs. - - Parameters - ---------- - current_epoch - Current epoch number. - max_epochs - Total number of epochs. - init_value - Initial value of the parameter. - target_value - Target value of the parameter. - - Returns - ------- - Annealed parameter value. - """ - anneal_epochs = max_epochs // 3 - if current_epoch >= anneal_epochs: - return target_value - progress = current_epoch / anneal_epochs - return init_value + (target_value - init_value) * progress - - -class DiagTrainingPlan(TrainingPlan): - """Training plan for DIAGVI model. - - Handles multi-modal loss computation including reconstruction, KL divergence, - graph loss, Sinkhorn (optimal transport) loss, and classification loss. - - Parameters - ---------- - module - DIAGVAE module to train. - lam_graph - Weight for the graph reconstruction loss. - lam_kl - Weight for the KL divergence loss. - lam_data - Weight for the data reconstruction loss. - lam_sinkhorn - Weight for the Sinkhorn (optimal transport) loss. - lam_class - Weight for the classification loss. If None, automatically set based - on whether classification is used. - sinkhorn_p - Order of the Wasserstein distance (p-norm). - sinkhorn_blur - Blur for geomloss Sinkhorn (epsilon = blur^p). If None, computed - adaptively from cost matrix like OTT-JAX, by default None. - epsilon_from_cost - Statistic of cost to compute epsilon: "std" or "mean". - epsilon_scale - Scaling factor: `epsilon = epsilon_scale * statistic(C)`. - sinkhorn_reach - Reach parameter for unbalanced OT. If None, calculate adaptive reach from blur. - reach_scale - Scaling factor for adaptive reach: `reach = reach_scale * blur`. - lr - Learning rate. - loss_annealing - Whether to anneal loss parameters during training. - log_train - Whether to log individual train loss components. - log_val - Whether to log individual validation loss components. - *args - Additional positional arguments passed to :class:`~scvi.train.TrainingPlan`. - **kwargs - Additional keyword arguments passed to :class:`~scvi.train.TrainingPlan`. - """ - - def __init__( - self, - module: torch.nn.Module, - lam_graph: float = 15.0, - lam_kl: float = 1.0, - lam_data: float = 0.1, - lam_sinkhorn: float = 85.0, - lam_class: float = 5.0, - sinkhorn_p: int = 2, - sinkhorn_blur: float | None = None, - epsilon_from_cost: Literal["mean", "std"] = "mean", - epsilon_scale: float = 0.5, - sinkhorn_reach: float | None = None, - reach_scale: float = 3.0, - lr: float = 1e-3, - loss_annealing: bool = False, - log_train: bool = True, - log_val: bool = False, - *args, - **kwargs, - ) -> None: - super().__init__(module, *args, **kwargs) - - # Loss weights - self.lam_graph = lam_graph - self.lam_kl = lam_kl - self.lam_data = lam_data - self.lam_sinkhorn = lam_sinkhorn - self.lam_class = lam_class - - # Sinkhorn parameters - self.sinkhorn_p = sinkhorn_p - self.sinkhorn_blur = sinkhorn_blur - self.sinkhorn_reach = sinkhorn_reach - self.epsilon_from_cost = epsilon_from_cost - self.epsilon_scale = epsilon_scale - self.reach_scale = reach_scale - - # Training parameters - self.lr = lr - self.loss_annealing = loss_annealing - self.log_train = log_train - self.log_val = log_val - - # Initial values for annealing (10x larger for smoother optimization start) - self.init_blur = 10 * self.sinkhorn_blur if self.sinkhorn_blur is not None else None - self.init_reach = 10 * self.sinkhorn_reach if self.sinkhorn_reach is not None else None - - # Current sinkhorn parameters - self.current_blur = self.sinkhorn_blur - self.current_reach = self.sinkhorn_reach - - def _compute_modality_losses( - self, batch: dict[str, dict[str, torch.Tensor]], log_prefix: str = "" - ) -> tuple[list[dict], dict, int]: - """Compute per-modality losses (reconstruction + KL). - - Parameters - ---------- - batch - A batch of data containing tensors for each modality. - log_prefix - Prefix for logging keys (e.g., "train_" or "val_"). - - Returns - ------- - loss_outputs - Per-modality loss information containing z, modality_loss, graph_v, classification_loss - last_loss_output - The last loss output object (contains shared info like guidance_graph) - total_batch_size - Total number of samples across all modalities - """ - loss_outputs = [] - - # Compute graph encoding once — identical for all modalities - first_x = next(iter(batch.values()))[REGISTRY_KEYS.X_KEY] - v_all, mu_all, logvar_all = self.module.encode_graph(first_x.device) - - for name, tensors in batch.items(): - batch_size = tensors[REGISTRY_KEYS.X_KEY].shape[0] - - # Update loss kwargs for current modality - self.loss_kwargs.update( - {"lam_kl": self.lam_kl, "lam_data": self.lam_data, "mode": name} - ) - - # Calculate reconstruction, KL, and classification losses - _, _, loss_output = self.forward( - tensors, - loss_kwargs=self.loss_kwargs, - inference_kwargs={ - "mode": name, - "v_all": v_all, - "mu_all": mu_all, - "logvar_all": logvar_all, - }, - generative_kwargs={"mode": name}, - ) - - # Log per-modality metrics - # Validation and training step - if log_prefix: - self.log( - f"{log_prefix}loss_{name}", - loss_output.loss, - batch_size=batch_size, - on_epoch=True, - on_step=True, - ) - - # Validation and training step - if log_prefix in {"train_", "val_"}: - reconstruction_loss = torch.mean( - loss_output.reconstruction_loss["reconstruction_loss"] - ) - self.log( - f"{log_prefix}nll_{name}", - reconstruction_loss, - batch_size=batch_size, - on_epoch=True, - ) - self.log( - f"{log_prefix}kl_{name}", - loss_output.kl_local["kl_local"], - batch_size=batch_size, - on_epoch=True, - ) - - loss_outputs.append( - { - "z": loss_output.extra_metrics["z"], - "modality_loss": loss_output.loss, - "graph_v": loss_output.extra_metrics["v_all"], - "classification_loss": loss_output.extra_metrics["classification_loss"], - } - ) - - total_batch_size = sum(tensors[REGISTRY_KEYS.X_KEY].shape[0] for tensors in batch.values()) - - return loss_outputs, loss_output, total_batch_size - - def _compute_graph_loss( - self, - loss_output: dict, - feature_embeddings: torch.Tensor, - log: bool = False, - batch_size: int = 0, - ) -> torch.Tensor: - """Compute graph reconstruction and KL loss. - - Parameters - ---------- - loss_output - Loss output object containing graph information. - feature_embeddings - Tensor of shape (n_features, embedding_dim) containing feature embeddings. - log - Whether to log graph loss components. - batch_size - Batch size for logging. - - Returns - ------- - Scalar tensor containing the total graph loss. - """ - graph = loss_output.extra_metrics["guidance_graph"] - - # Graph reconstruction loss - graph_nll = compute_graph_loss(graph, feature_embeddings) - - # Graph KL divergence (normalized by number of features) - graph_kl = kl_divergence_graph( - loss_output.extra_metrics["mu_all"], - loss_output.extra_metrics["logvar_all"], - ) - graph_kl_norm = graph_kl / feature_embeddings.shape[0] - - if log: - self.log("nll_graph", graph_nll, batch_size=batch_size, on_epoch=True) - self.log("kl_graph", graph_kl_norm, batch_size=batch_size, on_epoch=True) - - return graph_nll + graph_kl_norm - - @dependencies("geomloss") - def _compute_sinkhorn_loss( - self, z1: torch.Tensor, z2: torch.Tensor, use_annealing: bool = False - ) -> torch.Tensor: - """Compute Sinkhorn (Unbalanced Optimal Transport) loss between latent spaces. - - Sinkhorn parameters (blur and reach) can be adaptively computed from the cost matrix - using OTT-JAX style heuristics or set to fixed values with optional annealing. - - Parameters - ---------- - z1 - Latent representations from modality 1, shape (n_cells, n_latent). - z2 - Latent representations from modality 2, shape (n_cells, n_latent). - use_annealing - Whether to use annealed Sinkhorn parameters. Only applied when both - sinkhorn_blur and sinkhorn_reach are explicitly specified. - - Returns - ------- - Scalar tensor containing the Sinkhorn loss. - """ - import geomloss - - # Check if adaptive computation is needed - needs_adaptive = self.sinkhorn_blur is None or self.sinkhorn_reach is None - - if needs_adaptive: - # TODO: think about how to handle the cost matrix for large batches - # Adaptively compute Sinkhorn parameters via OTT-JAX style heuristics - # Note: annealing is NOT applied when using adaptive computation - - cost_fn = cost_routines[self.sinkhorn_p] - with torch.no_grad(): - # Compute cost matrix C - C_st = cost_fn(z1, z2) - - # Compute adaptive epsilon from cost and set blur - if self.sinkhorn_blur is None: - if self.epsilon_from_cost == "std": - eps = self.epsilon_scale * C_st.std().item() - elif self.epsilon_from_cost == "mean": - eps = self.epsilon_scale * C_st.mean().item() - else: - raise ValueError(f"Unknown epsilon_from_cost: {self.epsilon_from_cost}") - - eps = max(eps, 1e-8) # Ensure positive - - # Convert epsilon to blur for geomloss: epsilon = blur^p → blur = epsilon^(1/p) - self.current_blur = eps ** (1.0 / self.sinkhorn_p) - else: - self.current_blur = self.sinkhorn_blur - - # Compute adaptive reach from blur - if self.sinkhorn_reach is None: - self.current_reach = self.reach_scale * self.current_blur - else: - self.current_reach = self.sinkhorn_reach - - else: - # Both blur and reach are specified - use them with optional annealing - if use_annealing: - max_epochs = self.trainer.max_epochs - self.current_blur = _anneal_param( - self.current_epoch, max_epochs, self.init_blur, self.sinkhorn_blur - ) - self.current_reach = _anneal_param( - self.current_epoch, max_epochs, self.init_reach, self.sinkhorn_reach - ) - else: - self.current_blur = self.sinkhorn_blur - self.current_reach = self.sinkhorn_reach - - sinkhorn = geomloss.SamplesLoss( - loss="sinkhorn", - p=self.sinkhorn_p, - blur=self.current_blur, - reach=self.current_reach, - ) - return sinkhorn(z1, z2) - - def _compute_total_loss( - self, - loss_outputs: list[dict], - loss_output: dict, - total_batch_size: int, - use_annealing: bool = False, - log: bool = False, - ) -> torch.Tensor: - """Compute total loss combining all components. - - Parameters - ---------- - loss_outputs - List of per-modality loss information. - loss_output - The last loss output object (contains shared info like guidance_graph). - total_batch_size - Total number of samples across all modalities. - use_annealing - Whether to use annealed Sinkhorn parameters. - log - Whether to log loss components. - - Returns - ------- - Scalar tensor containing the total loss. - """ - # 1. Data loss (reconstruction + KL for each modality) - data_loss = sum(out["modality_loss"] for out in loss_outputs) - - # 2. Graph loss - feature_embeddings = loss_outputs[0]["graph_v"] - graph_loss = self._compute_graph_loss( - loss_output, feature_embeddings, log=log, batch_size=total_batch_size - ) - - # 3. Classification loss - classification_loss = sum(out["classification_loss"] for out in loss_outputs) - if log: - self.log("class_loss", classification_loss, batch_size=total_batch_size, on_epoch=True) - - # 4. Sinkhorn (UOT) loss - z1, z2 = loss_outputs[0]["z"], loss_outputs[1]["z"] - sinkhorn_loss = self._compute_sinkhorn_loss(z1, z2, use_annealing=use_annealing) - if log: - self.log("uot_loss", sinkhorn_loss, batch_size=total_batch_size, on_epoch=True) - - # Combine all losses - total_loss = ( - data_loss - + self.lam_graph * graph_loss - + self.lam_sinkhorn * sinkhorn_loss - + self.lam_class * classification_loss - ) - - return total_loss - - def training_step(self, batch: dict[str, dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]: - """Training step. - - During training, computes the losses for each modality (NLL, KL, and classification), - the graph loss, the Sinkhorn loss between modalities, and combines them into a total loss. - - Parameters - ---------- - batch - A batch of data containing tensors for each modality. - - Returns - ------- - A dictionary containing the total loss for backpropagation. - """ - loss_outputs, loss_output, total_batch_size = self._compute_modality_losses( - batch, log_prefix="train_" - ) - - total_loss = self._compute_total_loss( - loss_outputs, - loss_output, - total_batch_size, - use_annealing=self.loss_annealing, - log=self.log_train, - ) - - self.log( - "training_loss", - total_loss, - prog_bar=True, - on_step=False, - on_epoch=True, - batch_size=total_batch_size, - ) - - # Return embeddings along with loss for callbacks (detached to avoid graph issues) - embeddings = { - name: out["z"].detach() for name, out in zip(batch.keys(), loss_outputs, strict=False) - } - - return {"loss": total_loss, "embeddings": embeddings} - - def validation_step(self, batch: dict[str, dict[str, torch.Tensor]]): - """Validation step. - - During validation, computes and logs the losses for each modality (NLL, KL, and - classification), the graph loss, and the Sinkhorn loss between modalities, and - combines them into a total loss. - - Parameters - ---------- - batch - A batch of data containing tensors for each modality. - - Returns - ------- - Returns none. Logs validation losses. - """ - loss_outputs, loss_output, total_batch_size = self._compute_modality_losses( - batch, log_prefix="val_" - ) - - total_loss = self._compute_total_loss( - loss_outputs, loss_output, total_batch_size, use_annealing=False, log=self.log_val - ) - - self.log( - "validation_loss", - total_loss, - prog_bar=True, - on_epoch=True, - batch_size=total_batch_size, - ) diff --git a/src/scvi/external/diagvi/_utils.py b/src/scvi/external/diagvi/_utils.py deleted file mode 100644 index 0d3fecea79..0000000000 --- a/src/scvi/external/diagvi/_utils.py +++ /dev/null @@ -1,384 +0,0 @@ -"""Utility functions for the DIAGVI model.""" - -from __future__ import annotations - -import logging -import os -from typing import TYPE_CHECKING - -import anndata as ad -import numpy as np -import scipy.spatial -import torch -import torch.nn.functional as F - -from scvi.data._download import _download -from scvi.utils import dependencies - -if TYPE_CHECKING: - from typing import Literal - - import pandas as pd - from anndata import AnnData - from torch_geometric.data import Data - -logger = logging.getLogger(__name__) - - -@dependencies("torch_geometric") -def _construct_guidance_graph( - adatas: dict[str, AnnData], - mapping_df: pd.DataFrame | None, - weight: float = 1.0, - sign: float = 1.0, -) -> Data: - """Construct a guidance graph linking features across modalities. - - Creates a bipartite graph where nodes represent features from each modality - and edges connect corresponding features based on the mapping DataFrame or - shared feature names. - - Parameters - ---------- - adatas - Dictionary mapping modality names to AnnData objects. - mapping_df - DataFrame with columns matching modality names, containing feature - mappings. If None, uses shared feature names. - weight - Edge weight for cross-modality connections. - sign - Edge sign for cross-modality connections. - - Returns - ------- - PyTorch Geometric Data object with node features, edge indices, - edge weights, edge signs, and modality index tensors. - - Raises - ------ - ValueError - If not exactly two modalities are provided or no overlapping features - exist when mapping_df is None. - """ - from torch_geometric.data import Data - - if len(adatas) != 2: - raise ValueError("Exactly two modalities are required.") - input_names = list(adatas.keys()) - adata1, adata2 = adatas[input_names[0]], adatas[input_names[1]] - - if mapping_df is not None: - features1 = list(adata1.var_names) - features2 = list(adata2.var_names) - else: - shared_features = set(adata1.var_names) & set(adata2.var_names) - if not shared_features: - raise ValueError("No overlapping features between the two modalities.") - - features1 = [f"{f}_{input_names[0]}" for f in adata1.var_names] - features2 = [f"{f}_{input_names[1]}" for f in adata2.var_names] - - all_features = features1 + features2 - feature_to_index = {f: i for i, f in enumerate(all_features)} - - edge_index = [] - edge_weight = [] - edge_sign = [] - - if mapping_df is not None: - for ft_pair in range(mapping_df.shape[0]): - pair = mapping_df.iloc[ft_pair, :] - diss_ft = pair[input_names[0]] - sp_ft = pair[input_names[1]] - - i = feature_to_index[diss_ft] - j = feature_to_index[sp_ft] - - edge_index += [[i, j], [j, i]] - edge_weight += [weight, weight] - edge_sign += [sign, sign] - - else: - for feature in shared_features: - i = feature_to_index[f"{feature}_{input_names[0]}"] - j = feature_to_index[f"{feature}_{input_names[1]}"] - - edge_index += [[i, j], [j, i]] - edge_weight += [weight, weight] - edge_sign += [sign, sign] - - for feature in all_features: - i = feature_to_index[feature] - edge_index.append([i, i]) - edge_weight.append(weight) - edge_sign.append(sign) - - edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous() - edge_weight = torch.tensor(edge_weight, dtype=torch.float) - edge_sign = torch.tensor(edge_sign, dtype=torch.float) - - x = torch.eye(len(all_features)) - - indices1 = torch.tensor([feature_to_index[f] for f in features1], dtype=torch.long) - indices2 = torch.tensor([feature_to_index[f] for f in features2], dtype=torch.long) - - return Data( - x=x, - edge_index=edge_index, - edge_weight=edge_weight, - edge_sign=edge_sign, - **{f"{input_names[0]}_indices": indices1, f"{input_names[1]}_indices": indices2}, - ) - - -def _check_guidance_graph_consistency(graph: Data, adatas: dict[str, AnnData]): - """Validate guidance graph structure and consistency with AnnData objects. - - Performs several consistency checks on the guidance graph: - 1. Node count matches total number of features across modalities - 2. Required edge attributes (edge_weight, edge_sign) are present - 3. Self-loops exist for all nodes - 4. Graph is symmetric (undirected) - - Parameters - ---------- - graph - PyTorch Geometric Data object representing the guidance graph. - adatas - Dictionary mapping modality names to AnnData objects. - - Raises - ------ - ValueError - If any consistency check fails. - """ - n_expected = sum(adata.shape[1] for adata in adatas.values()) - - # 1. Check variable coverage via counts - if graph.num_nodes != n_expected: - raise ValueError( - f"Graph node count {graph.num_nodes} does not match expected {n_expected}." - ) - - # 2. Check edge attributes - for attr in ["edge_weight", "edge_sign"]: - if not hasattr(graph, attr): - raise ValueError(f"Graph missing required edge attribute: {attr}") - if getattr(graph, attr).shape[0] != graph.edge_index.shape[1]: - raise ValueError(f"Edge attribute {attr} does not match number of edges.") - - # 3. Check self-loops - src, tgt = graph.edge_index - self_loops = src == tgt - n_self_loops = self_loops.sum().item() - if n_self_loops < graph.num_nodes: - raise ValueError("Graph is missing self-loops for some nodes.") - - # 4. Check symmetry (for undirected graphs: For every edge (i, j), check that (j, i) exists) - edge_set = {(i.item(), j.item()) for i, j in zip(src, tgt, strict=False)} - for i, j in zip(src, tgt, strict=False): - if (j.item(), i.item()) not in edge_set: - raise ValueError( - f"Graph is not symmetric: edge ({i.item()}, {j.item()}) has no counterpart." - ) - - # If all checks pass - logger.info("Guidance graph consistency checks passed.") - - -def _load_saved_diagvi_files( - dir_path: str, - prefix: str | None = None, - map_location: Literal["cpu", "cuda"] | None = None, - backup_url: str | None = None, -) -> tuple[dict, dict[str, np.ndarray], dict, dict[str, AnnData | None]]: - """Loads saved DiagVI model and AnnData files from a directory. - - Parameters - ---------- - dir_path - Directory path where the model and AnnData files are stored. - prefix - Optional prefix for the file names. - map_location - Device mapping for loading the model. - backup_url - Optional URL to download the model file if not found locally. - - Returns - ------- - A tuple containing: - - attr_dict: Dictionary of model attributes. - - var_names: Dictionary of variable names for each modality. - - model_state_dict: State dictionary of the model. - - adatas: Dictionary of AnnData objects for each modality. - - Raises - ------ - ValueError - If model file cannot be loaded. - """ - file_name_prefix = prefix or "" - - model_file_name = f"{file_name_prefix}model.pt" - model_path = os.path.join(dir_path, model_file_name) - - try: - _download(backup_url, dir_path, model_file_name) - model = torch.load(model_path, map_location=map_location, weights_only=False) - except FileNotFoundError as exc: - raise ValueError(f"Failed to load model file at {model_path}. ") from exc - - names = model["names"] - - adatas = {} - var_names = {} - for name in names: - adata_path = os.path.join(dir_path, f"{file_name_prefix}adata_{name}.h5ad") - if os.path.exists(adata_path): - adatas[name] = ad.read_h5ad(adata_path) - var_names[name] = adatas[name].var_names - else: - adatas[name] = None - key = f"var_names_{name}" - if key in model: - var_names[name] = model[key] - - model_state_dict = model["model_state_dict"] - attr_dict = model["attr_dict"] - - return ( - attr_dict, - var_names, - model_state_dict, - adatas, - ) - - -@dependencies("torch_geometric") -def compute_graph_loss(graph: Data, feature_embeddings: torch.Tensor) -> torch.Tensor: - """Compute graph reconstruction loss using negative sampling. - - Uses structured negative sampling to compute a contrastive loss that - encourages connected nodes to have similar embeddings and unconnected - nodes to have dissimilar embeddings. - - Parameters - ---------- - graph - PyTorch Geometric Data object with edge_index. - feature_embeddings - Tensor of shape (n_features, embedding_dim) containing feature embeddings. - - Returns - ------- - Scalar tensor containing the graph reconstruction loss. - """ - import torch_geometric - - edge_index = graph.edge_index - edge_index_neg = torch_geometric.utils.structured_negative_sampling(edge_index) - - pos_i, pos_j, neg_j = edge_index_neg[0], edge_index_neg[1], edge_index_neg[2] - - vi = feature_embeddings[pos_i] - vj = feature_embeddings[pos_j] - vj_neg = feature_embeddings[neg_j] - - pos_logits = (vi * vj).sum(dim=1) - pos_loss = F.logsigmoid(pos_logits).mean() - - neg_logits = (vi * vj_neg).sum(dim=1) - neg_loss = F.logsigmoid(-neg_logits).mean() - - total_loss = -(pos_loss + neg_loss) / 2 - - return total_loss - - -def kl_divergence_graph(mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor: - """Computes the KL divergence for graph latent variables. - - Parameters - ---------- - mu - Mean tensor of the latent variables. - logvar - Log-variance tensor of the latent variables. - - Returns - ------- - The mean KL divergence as a tensor. - """ - kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1) - kl_mean = kl.mean() - return kl_mean - - -def compute_foscttm( - latents: dict[str, np.ndarray], - indices: list[np.ndarray] | None = None, - downsample: bool = False, - n_obs: int = 10000, -) -> dict[str, float]: - """Compute FOSCTTM (Fraction of Samples Closer Than True Match) metric. - - For paired multi-modal data, measures how well the model aligns corresponding - cells in the latent space. Lower values are better (0 = perfect alignment). - Assumes cells are aligned: mod1 cell i corresponds to mod2 cell i. - - Parameters - ---------- - latents - Dictionary mapping modality names to latent arrays. - indices - List of two arrays: indices for [mod1, mod2]. - Usually obtained from model.validation_indices after training. - downsample - Whether to downsample when the number of observations is large. - n_obs - Number of observations to keep if downsample is True. - """ - # Extract latent representations for the two modalities - latent_mod1 = latents[list(latents.keys())[0]] - latent_mod2 = latents[list(latents.keys())[1]] - - # Subset if indices provided - if indices is not None: - latent_mod1 = latent_mod1[indices[0]] - latent_mod2 = latent_mod2[indices[1]] - - # Validate shapes match - if latent_mod1.shape[0] != latent_mod2.shape[0]: - raise ValueError("Shapes do not match!") - - n_cells = latent_mod1.shape[0] - - # Downsample if requested and number of cells exceeds n_obs - if n_cells > n_obs and downsample: - np.random.seed(0) - sample_indices = np.random.choice(n_cells, size=n_obs, replace=False) - latent_mod1 = latent_mod1[sample_indices] - latent_mod2 = latent_mod2[sample_indices] - n_cells = n_obs - - # Compute pairwise distances - distances = scipy.spatial.distance_matrix(latent_mod1, latent_mod2) - - # Compute FOSCTTM - foscttm_mod1_to_mod2 = (distances < np.expand_dims(np.diag(distances), axis=1)).mean(axis=1) - foscttm_mod2_to_mod1 = (distances < np.expand_dims(np.diag(distances), axis=0)).mean(axis=0) - - # Aggregate metrics - mean_mod1_to_mod2 = foscttm_mod1_to_mod2.mean() - mean_mod2_to_mod1 = foscttm_mod2_to_mod1.mean() - - mod1_name, mod2_name = list(latents.keys()) - foscttm_metrics = { - f"foscttm/{mod1_name}_to_{mod2_name}": float(mean_mod1_to_mod2), - f"foscttm/{mod2_name}_to_{mod1_name}": float(mean_mod2_to_mod1), - "foscttm/mean": float((mean_mod1_to_mod2 + mean_mod2_to_mod1) / 2), - } - - return foscttm_metrics diff --git a/src/scvi/external/gimvi/__init__.py b/src/scvi/external/gimvi/__init__.py deleted file mode 100644 index 93037019a3..0000000000 --- a/src/scvi/external/gimvi/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from ._model import GIMVI -from ._module import JVAE - -__all__ = ["GIMVI", "JVAE"] diff --git a/src/scvi/external/gimvi/_model.py b/src/scvi/external/gimvi/_model.py deleted file mode 100644 index 9a884cca99..0000000000 --- a/src/scvi/external/gimvi/_model.py +++ /dev/null @@ -1,690 +0,0 @@ -from __future__ import annotations - -import logging -import os -import warnings -from typing import TYPE_CHECKING - -import numpy as np -import torch - -from scvi import REGISTRY_KEYS, settings -from scvi.data import AnnDataManager -from scvi.data._compat import registry_from_setup_dict -from scvi.data._constants import _MODEL_NAME_KEY, _SETUP_ARGS_KEY -from scvi.data.fields import CategoricalObsField, LayerField -from scvi.dataloaders import DataSplitter -from scvi.model._utils import _init_library_size, parse_device_args -from scvi.model.base import BaseModelClass, RNASeqMixin, VAEMixin -from scvi.train import Trainer -from scvi.train._config import merge_kwargs -from scvi.utils import setup_anndata_dsp -from scvi.utils._docstrings import devices_dsp - -from ._module import JVAE -from ._task import CyclicMultiDataLoader as TrainDL -from ._task import GIMVITrainingPlan -from ._utils import _load_legacy_saved_gimvi_files, _load_saved_gimvi_files - -if TYPE_CHECKING: - from anndata import AnnData - - from scvi.dataloaders import AnnDataLoader - -logger = logging.getLogger(__name__) - - -def _unpack_tensors(tensors): - x = tensors[REGISTRY_KEYS.X_KEY].squeeze_(0) - batch_index = tensors[REGISTRY_KEYS.BATCH_KEY].squeeze_(0) - y = tensors[REGISTRY_KEYS.LABELS_KEY].squeeze_(0) - return x, batch_index, y - - -class GIMVI(VAEMixin, RNASeqMixin, BaseModelClass): - """Joint VAE for imputing missing genes in spatial data :cite:p:`Lopez19`. - - Parameters - ---------- - adata_seq - AnnData object that has been registered via :meth:`~scvi.external.GIMVI.setup_anndata` - and contains RNA-seq data. - adata_spatial - AnnData object that has been registered via :meth:`~scvi.external.GIMVI.setup_anndata` - and contains spatial data. - n_hidden - Number of nodes per hidden layer. - generative_distributions - List of generative distribution for adata_seq data and adata_spatial data. Defaults to - ['zinb', 'nb']. - model_library_size - List of bool of whether to model library size for adata_seq and adata_spatial. Defaults to - [True, False]. - n_latent - Dimensionality of the latent space. - **model_kwargs - Keyword args for :class:`~scvi.external.gimvi.JVAE` - - Examples - -------- - >>> adata_seq = anndata.read_h5ad(path_to_anndata_seq) - >>> adata_spatial = anndata.read_h5ad(path_to_anndata_spatial) - >>> scvi.external.GIMVI.setup_anndata(adata_seq) - >>> scvi.external.GIMVI.setup_anndata(adata_spatial) - >>> vae = scvi.model.GIMVI(adata_seq, adata_spatial) - >>> vae.train(n_epochs=400) - - Notes - ----- - See further usage examples in the following tutorials: - - 1. :doc:`/tutorials/notebooks/spatial/gimvi_tutorial` - """ - - def __init__( - self, - adata_seq: AnnData, - adata_spatial: AnnData, - generative_distributions: list[str] | None = None, - model_library_size: list[bool] | None = None, - n_latent: int = 10, - **model_kwargs, - ): - warnings.warn( - "GIMVI is a spatial transcriptomics model that will be moved to the " - "scvi-tools spatial companion package `scviva-tools` starting in scvi-tools v1.5 and " - "will no longer be supported here. It will be deprecated from scvi-tools in v1.6.", - FutureWarning, - stacklevel=settings.warnings_stacklevel, - ) - super().__init__() - if adata_seq is adata_spatial: - raise ValueError( - "`adata_seq` and `adata_spatial` cannot point to the same object. " - "If you would really like to do this, make a copy of the object and pass it in as " - "`adata_spatial`." - ) - model_library_size = model_library_size or [True, False] - generative_distributions = generative_distributions or ["zinb", "nb"] - self.adatas = [adata_seq, adata_spatial] - self.adata_managers = { - "seq": self._get_most_recent_anndata_manager(adata_seq, required=True), - "spatial": self._get_most_recent_anndata_manager(adata_spatial, required=True), - } - self.registries_ = [] - for adm in self.adata_managers.values(): - self._register_manager_for_instance(adm) - self.registries_.append(adm.registry) - - seq_var_names = adata_seq.var_names - spatial_var_names = adata_spatial.var_names - - if not set(spatial_var_names) <= set(seq_var_names): - raise ValueError("spatial genes needs to be subset of seq genes") - - spatial_gene_loc = [np.argwhere(seq_var_names == g)[0] for g in spatial_var_names] - spatial_gene_loc = np.concatenate(spatial_gene_loc) - gene_mappings = [slice(None), spatial_gene_loc] - sum_stats = [adm.summary_stats for adm in self.adata_managers.values()] - n_inputs = [s["n_vars"] for s in sum_stats] - - total_genes = n_inputs[0] - - adata_seq_n_batches = sum_stats[0]["n_batch"] - adata_spatial_batch = adata_spatial.obs[ - self.adata_managers["spatial"].data_registry[REGISTRY_KEYS.BATCH_KEY].attr_key - ] - if np.min(adata_spatial_batch) == 0: - # see #2446 - # since we are combining datasets, we need to increment the batch_idx of one of the - # datasets. we only need to do this once so we check if the min is 0 - adata_spatial.obs[ - self.adata_managers["spatial"].data_registry[REGISTRY_KEYS.BATCH_KEY].attr_key - ] += adata_seq_n_batches - - n_batches = sum(s["n_batch"] for s in sum_stats) - - library_log_means = [] - library_log_vars = [] - for adata_manager in self.adata_managers.values(): - adata_library_log_means, adata_library_log_vars = _init_library_size( - adata_manager, n_batches - ) - library_log_means.append(adata_library_log_means) - library_log_vars.append(adata_library_log_vars) - - self.module = JVAE( - n_inputs, - total_genes, - gene_mappings, - generative_distributions, - model_library_size, - library_log_means, - library_log_vars, - n_batch=n_batches, - n_latent=n_latent, - **model_kwargs, - ) - - self._model_summary_string = ( - "GimVI Model with the following params: \nn_latent: {}, n_inputs: {}, n_genes: {}, " - + "n_batch: {}, generative distributions: {}" - ).format(n_latent, n_inputs, total_genes, n_batches, generative_distributions) - self.init_params_ = self._get_init_params(locals()) - - @devices_dsp.dedent - def train( - self, - max_epochs: int = 200, - accelerator: str = "auto", - devices: int | list[int] | str = "auto", - kappa: int = 5, - train_size: float | None = None, - validation_size: float | None = None, - shuffle_set_split: bool = True, - batch_size: int = 128, - datasplitter_kwargs: dict | None = None, - plan_kwargs: dict | None = None, - **kwargs, - ): - """Train the model. - - Parameters - ---------- - max_epochs - Number of passes through the dataset. If `None`, defaults to - `np.min([round((20000 / n_cells) * 400), 400])` - %(param_accelerator)s - %(param_devices)s - kappa - Scaling parameter for the discriminator loss. - 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. - datasplitter_kwargs - Additional keyword arguments passed into :class:`~scvi.dataloaders.DataSplitter`. - plan_kwargs - Keyword args for model-specific Pytorch Lightning task. Keyword arguments passed - to `train()` will overwrite values present in `plan_kwargs`, when appropriate. - **kwargs - Other keyword args for :class:`~scvi.train.Trainer`. - """ - accelerator, devices, device = parse_device_args( - accelerator=accelerator, - devices=devices, - return_device="torch", - ) - datasplitter_kwargs = datasplitter_kwargs or {} - - self.trainer = Trainer( - max_epochs=max_epochs, - accelerator=accelerator, - devices=devices, - **kwargs, - ) - self.train_indices_, self.test_indices_, self.validation_indices_ = [], [], [] - train_dls, test_dls, val_dls = [], [], [] - for i, adm in enumerate(self.adata_managers.values()): - ds = DataSplitter( - adm, - train_size=train_size, - validation_size=validation_size, - batch_size=batch_size, - shuffle_set_split=shuffle_set_split, - **datasplitter_kwargs, - ) - ds.setup() - train_dls.append(ds.train_dataloader()) - test_dls.append(ds.test_dataloader()) - val = ds.val_dataloader() - val_dls.append(val) - val.mode = i - self.train_indices_.append(ds.train_idx) - self.test_indices_.append(ds.test_idx) - self.validation_indices_.append(ds.val_idx) - train_dl = TrainDL(train_dls) - - plan_kwargs = merge_kwargs(None, plan_kwargs, name="plan") - self._training_plan = GIMVITrainingPlan( - self.module, - adversarial_classifier=True, - scale_adversarial_loss=kappa, - **plan_kwargs, - ) - - if train_size == 1.0: - # circumvent the empty data loader problem if all dataset used for training - self.trainer.fit(self._training_plan, train_dl) - else: - # accepts list of val dataloaders - self.trainer.fit(self._training_plan, train_dl, val_dls) - try: - self.history_ = self.trainer.logger.history - except AttributeError: - self.history_ = None - self.module.eval() - - self.to_device(device) - self.is_trained_ = True - - def _make_scvi_dls( - self, adatas: list[AnnData] = None, batch_size: int = 128 - ) -> list[AnnDataLoader]: - if adatas is None: - adatas = self.adatas - post_list = [self._make_data_loader(ad, batch_size=batch_size) for ad in adatas] - for i, dl in enumerate(post_list): - dl.mode = i - - return post_list - - @torch.inference_mode() - def get_latent_representation( - self, - adatas: list[AnnData] = None, - deterministic: bool = True, - batch_size: int = 128, - ) -> list[np.ndarray]: - """Return the latent space embedding for each dataset. - - Parameters - ---------- - adatas - List of adata seq and adata spatial. - deterministic - If true, use the mean of the encoder instead of a Gaussian sample. - batch_size - Minibatch size for data loading into model. - """ - if adatas is None: - adatas = self.adatas - scdls = self._make_scvi_dls(adatas, batch_size=batch_size) - self.module.eval() - latents = [] - for mode, scdl in enumerate(scdls): - latent = [] - for tensors in scdl: - ( - sample_batch, - *_, - ) = _unpack_tensors(tensors) - latent.append( - self.module.sample_from_posterior_z( - sample_batch, mode, deterministic=deterministic - ) - .cpu() - .detach() - ) - - latent = torch.cat(latent).numpy() - latents.append(latent) - - return latents - - @torch.inference_mode() - def get_imputed_values( - self, - adatas: list[AnnData] = None, - deterministic: bool = True, - normalized: bool = True, - decode_mode: int | None = None, - batch_size: int = 128, - ) -> list[np.ndarray]: - """Return imputed values for all genes for each dataset. - - Parameters - ---------- - adatas - List of adata seq and adata spatial - deterministic - If true, use the mean of the encoder instead of a Gaussian sample for the latent - vector. - normalized - Return imputed normalized values or not. - decode_mode - If a `decode_mode` is given, use the encoder specific to each dataset as usual but use - the decoder of the dataset of id `decode_mode` to impute values. - batch_size - Minibatch size for data loading into model. - """ - self.module.eval() - - if adatas is None: - adatas = self.adatas - scdls = self._make_scvi_dls(adatas, batch_size=batch_size) - - imputed_values = [] - for mode, scdl in enumerate(scdls): - imputed_value = [] - for tensors in scdl: - ( - sample_batch, - batch_index, - label, - *_, - ) = _unpack_tensors(tensors) - if normalized: - imputed_value.append( - self.module.sample_scale( - sample_batch, - mode, - batch_index, - label, - deterministic=deterministic, - decode_mode=decode_mode, - ) - .cpu() - .detach() - ) - else: - imputed_value.append( - self.module.sample_rate( - sample_batch, - mode, - batch_index, - label, - deterministic=deterministic, - decode_mode=decode_mode, - ) - .cpu() - .detach() - ) - - imputed_value = torch.cat(imputed_value).numpy() - imputed_values.append(imputed_value) - - return imputed_values - - def save( - self, - dir_path: str, - prefix: str | None = None, - overwrite: bool = False, - save_anndata: bool = False, - save_kwargs: dict | None = None, - **anndata_write_kwargs, - ): - """Save the state of the model. - - Neither the trainer optimizer state nor the trainer history are saved. - Model files are not expected to be reproducibly saved and loaded across versions - until we reach version 1.0. - - Parameters - ---------- - dir_path - Path to a directory. - prefix - Prefix to prepend to saved file names. - overwrite - Overwrite existing data or not. If `False` and directory - already exists at `dir_path`, error will be raised. - save_anndata - If True, also saves the anndata - save_kwargs - Keyword arguments passed into :func:`~torch.save`. - anndata_write_kwargs - Kwargs for anndata write function - """ - if not os.path.exists(dir_path) or overwrite: - os.makedirs(dir_path, exist_ok=overwrite) - else: - raise ValueError( - f"{dir_path} already exists. Please provide an unexisting directory for saving." - ) - - file_name_prefix = prefix or "" - save_kwargs = save_kwargs or {} - - seq_adata = self.adatas[0] - spatial_adata = self.adatas[1] - if save_anndata: - seq_save_path = os.path.join(dir_path, f"{file_name_prefix}adata_seq.h5ad") - seq_adata.write(seq_save_path) - - spatial_save_path = os.path.join(dir_path, f"{file_name_prefix}adata_spatial.h5ad") - spatial_adata.write(spatial_save_path) - - # save the model state dict and the trainer state dict only - model_state_dict = self.module.state_dict() - - seq_var_names = seq_adata.var_names.astype(str).to_numpy() - spatial_var_names = spatial_adata.var_names.astype(str).to_numpy() - - # get all the user attributes - user_attributes = self._get_user_attributes() - # only save the public attributes with _ at the very end - user_attributes = {a[0]: a[1] for a in user_attributes if a[0][-1] == "_"} - - model_save_path = os.path.join(dir_path, f"{file_name_prefix}model.pt") - - torch.save( - { - "model_state_dict": model_state_dict, - "seq_var_names": seq_var_names, - "spatial_var_names": spatial_var_names, - "attr_dict": user_attributes, - }, - model_save_path, - **save_kwargs, - ) - - @classmethod - @devices_dsp.dedent - def load( - cls, - dir_path: str, - adata_seq: AnnData | None = None, - adata_spatial: AnnData | None = None, - accelerator: str = "auto", - device: int | str = "auto", - prefix: str | None = None, - backup_url: str | None = None, - ): - """Instantiate a model from the saved output. - - Parameters - ---------- - dir_path - Path to saved outputs. - adata_seq - AnnData organized in the same way as data used to train model. - It is not necessary to run :meth:`~scvi.external.GIMVI.setup_anndata`, - as AnnData is validated against the saved `scvi` setup dictionary. - AnnData must be registered via :meth:`~scvi.external.GIMVI.setup_anndata`. - adata_spatial - AnnData organized in the same way as data used to train model. - If None, will check for and load anndata saved with the model. - %(param_accelerator)s - %(param_device)s - prefix - Prefix of saved file names. - backup_url - URL to retrieve saved outputs from if not present on disk. - - Returns - ------- - Model with loaded state dictionaries. - - Examples - -------- - >>> vae = GIMVI.load(adata_seq, adata_spatial, save_path) - >>> vae.get_latent_representation() - """ - _, _, device = parse_device_args( - accelerator=accelerator, - devices=device, - return_device="torch", - ) - - ( - attr_dict, - seq_var_names, - spatial_var_names, - model_state_dict, - loaded_adata_seq, - loaded_adata_spatial, - ) = _load_saved_gimvi_files( - dir_path, - adata_seq is None, - adata_spatial is None, - prefix=prefix, - map_location=device, - backup_url=backup_url, - ) - adata_seq = loaded_adata_seq or adata_seq - adata_spatial = loaded_adata_spatial or adata_spatial - adatas = [adata_seq, adata_spatial] - var_names = [seq_var_names, spatial_var_names] - - for i, adata in enumerate(adatas): - saved_var_names = var_names[i] - user_var_names = adata.var_names.astype(str) - if not np.array_equal(saved_var_names, user_var_names): - warnings.warn( - "var_names for adata passed in does not match var_names of adata " - "used to train the model. For valid results, the vars need to be the" - "same and in the same order as the adata used to train the model.", - UserWarning, - stacklevel=settings.warnings_stacklevel, - ) - - registries = attr_dict.pop("registries_") - for adata, registry in zip(adatas, registries, strict=True): - if _MODEL_NAME_KEY in registry and registry[_MODEL_NAME_KEY] != cls.__name__: - raise ValueError("It appears you are loading a model from a different class.") - - if _SETUP_ARGS_KEY not in registry: - raise ValueError( - "Saved model does not contain original setup inputs. " - "Cannot load the original setup." - ) - - cls.setup_anndata(adata, source_registry=registry, **registry[_SETUP_ARGS_KEY]) - - # get the parameters for the class init signature - init_params = attr_dict.pop("init_params_") - - # new saving and loading, enable backwards compatibility - if "non_kwargs" in init_params.keys(): - # grab all the parameters except for kwargs (is a dict) - non_kwargs = init_params["non_kwargs"] - kwargs = init_params["kwargs"] - - # expand out kwargs - kwargs = {k: v for (i, j) in kwargs.items() for (k, v) in j.items()} - else: - # grab all the parameters except for kwargs (is a dict) - non_kwargs = {k: v for k, v in init_params.items() if not isinstance(v, dict)} - kwargs = {k: v for k, v in init_params.items() if isinstance(v, dict)} - kwargs = {k: v for (i, j) in kwargs.items() for (k, v) in j.items()} - model = cls(adata_seq, adata_spatial, **non_kwargs, **kwargs) - - for attr, val in attr_dict.items(): - setattr(model, attr, val) - - model.module.load_state_dict(model_state_dict) - model.module.eval() - model.to_device(device) - return model - - @classmethod - def convert_legacy_save( - cls, - dir_path: str, - output_dir_path: str, - overwrite: bool = False, - prefix: str | None = None, - **save_kwargs, - ) -> None: - """Converts a legacy saved GIMVI model (