feat: Linear SCVI adata minifiction (#3294)

added adata minification for linear scvi

Author: ori-kron-wis

CHANGELOG.md

@@ -16,6 +16,7 @@ to [Semantic Versioning]. Full commit history is available in the
 - Add get normalized function model property for any generative model {pr}`3238` and changed
     get_accessibility_estimates to get_normalized_accessibility, where needed.
 - Add Early stopping KL warmup steps. {pr}`3262`.
+- Add Minification option to {class}`~scvi.model.LinearSCVI` {pr}`3294`.
 
 #### Fixed
 

src/scvi/model/_linear_scvi.py

@@ -7,13 +7,15 @@
 
 from scvi import REGISTRY_KEYS
 from scvi.data import AnnDataManager
+from scvi.data._constants import ADATA_MINIFY_TYPE
+from scvi.data._utils import _get_adata_minify_type
 from scvi.data.fields import CategoricalObsField, LayerField
 from scvi.model._utils import _init_library_size
 from scvi.model.base import UnsupervisedTrainingMixin
 from scvi.module import LDVAE
 from scvi.utils import setup_anndata_dsp
 
-from .base import BaseModelClass, RNASeqMixin, VAEMixin
+from .base import BaseMinifiedModeModelClass, RNASeqMixin, VAEMixin
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -23,7 +25,12 @@
 logger = logging.getLogger(__name__)
 
 
-class LinearSCVI(RNASeqMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClass):
+class LinearSCVI(
+    RNASeqMixin,
+    VAEMixin,
+    UnsupervisedTrainingMixin,
+    BaseMinifiedModeModelClass,
+):
     """Linearly-decoded VAE :cite:p:`Svensson20`.
 
     Parameters
@@ -51,6 +58,9 @@ class LinearSCVI(RNASeqMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClas
         * ``'nb'`` - Negative binomial distribution
         * ``'zinb'`` - Zero-inflated negative binomial distribution
         * ``'poisson'`` - Poisson distribution
+    use_observed_lib_size
+        If ``True``, use the observed library size for RNA as the scaling factor in the mean of the
+        conditional distribution.
     latent_distribution
         One of:
 
@@ -75,6 +85,8 @@ class LinearSCVI(RNASeqMixin, VAEMixin, UnsupervisedTrainingMixin, BaseModelClas
     """
 
     _module_cls = LDVAE
+    _LATENT_QZM_KEY = "ldvae_latent_qzm"
+    _LATENT_QZV_KEY = "ldvae_latent_qzv"
 
     def __init__(
         self,
@@ -85,13 +97,19 @@ def __init__(
         dropout_rate: float = 0.1,
         dispersion: Literal["gene", "gene-batch", "gene-label", "gene-cell"] = "gene",
         gene_likelihood: Literal["zinb", "nb", "poisson"] = "nb",
+        use_observed_lib_size: bool = False,
         latent_distribution: Literal["normal", "ln"] = "normal",
         **model_kwargs,
     ):
         super().__init__(adata)
 
         n_batch = self.summary_stats.n_batch
-        library_log_means, library_log_vars = _init_library_size(self.adata_manager, n_batch)
+        library_log_means, library_log_vars = None, None
+        if (
+            self.minified_data_type != ADATA_MINIFY_TYPE.LATENT_POSTERIOR
+            and not use_observed_lib_size
+        ):
+            library_log_means, library_log_vars = _init_library_size(self.adata_manager, n_batch)
 
         self.module = self._module_cls(
             n_input=self.summary_stats.n_vars,
@@ -105,6 +123,7 @@ def __init__(
             latent_distribution=latent_distribution,
             library_log_means=library_log_means,
             library_log_vars=library_log_vars,
+            use_observed_lib_size=use_observed_lib_size,
             **model_kwargs,
         )
         self._model_summary_string = (
@@ -153,6 +172,10 @@ def setup_anndata(
             CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),
             CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key),
         ]
+        # register new fields if the adata is minified
+        adata_minify_type = _get_adata_minify_type(adata)
+        if adata_minify_type is not None:
+            anndata_fields += cls._get_fields_for_adata_minification(adata_minify_type)
         adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args)
         adata_manager.register_fields(adata, **kwargs)
         cls.register_manager(adata_manager)

tests/model/test_linear_scvi.py

@@ -6,6 +6,8 @@
 import torch
 
 from scvi.data import synthetic_iid
+from scvi.data._constants import ADATA_MINIFY_TYPE
+from scvi.data._utils import _is_minified
 from scvi.model import LinearSCVI
 from scvi.utils import attrdict
 
@@ -131,3 +133,25 @@ def test_linear_scvi_use_observed_lib_size():
     model.get_loadings()
     model.differential_expression(groupby="labels", group1="label_1")
     model.differential_expression(groupby="labels", group1="label_1", group2="label_2")
+
+
+def test_linear_scvi_with_minification(save_path):
+    adata = synthetic_iid()
+    adata = adata[:, :10].copy()
+    LinearSCVI.setup_anndata(adata)
+    model = LinearSCVI(adata, n_latent=10, use_observed_lib_size=True)
+    model.train(1, check_val_every_n_epoch=1, train_size=0.5)
+    assert len(model.history["elbo_train"]) == 1
+    assert len(model.history["elbo_validation"]) == 1
+    qzm, qzv = model.get_latent_representation(give_mean=False, return_dist=True)
+    model.adata.obsm["X_latent_qzm"] = qzm
+    model.adata.obsm["X_latent_qzv"] = qzv
+    model.minify_adata()
+    assert model.minified_data_type == ADATA_MINIFY_TYPE.LATENT_POSTERIOR
+    assert model.adata_manager.registry is model.registry_
+    assert not _is_minified(adata)
+    assert adata is not model.adata
+    assert len(model.adata.X.data) == 0
+    assert model.adata.raw is None
+    minified_model_path = os.path.join(save_path, "linear_scvi_minified")
+    model.save(minified_model_path, save_anndata=True, overwrite=True)