WIP: Update DRVI and align with the latest scvi-tools version

pyproject.toml

@@ -35,6 +35,19 @@ classifiers = [
 
 # Please make an issue if you need wider range of versions
 dependencies = [
+    "torch>=2.1.0",
+    "lightning>=2.0",
+    "scanpy>=1.9.5",
+    "scikit-learn>=1.5.1",
+    "scipy>=1.11.3",
+    "scvi-tools>=1.0.4",
+    "anndata>=0.10.2",
+    "numpy>=1.16.1,<2.0.0",  # for np.linspace
+    "pandas>=1.2.0",
+    # for debug logging (referenced from the issue template)
+    "session-info",
+]
+optional-dependencies.restrict = [
     "torch>=2.1.0,<2.4",
     "lightning>=2.0,<2.1",
     "scanpy==1.9.5",
@@ -46,7 +59,7 @@ dependencies = [
     "jaxlib<=0.4.20",
     ## END_TODO
     "anndata>=0.10.2,<0.11",
-    "numpy>=1.16.1",  # for np.linspace
+    "numpy>=1.16.1,<2.0.0",  # for np.linspace
     "pandas>=1.2.0",
     ## TODO: update this when this is resolved: https://github.com/boto/botocore/issues/2926
     # lightning-cloud depends on boto3 that is currently not compatible with urllib3 so resolution takes forever

src/drvi/nn_modules/prior.py

@@ -5,6 +5,8 @@
 from torch import nn
 from torch.distributions import Normal, kl_divergence
 
+from drvi.scvi_tools_based.module._constants import MODULE_KEYS
+
 # Standard, VaMP, GMM from Karin's CSI repo
 
 
@@ -198,7 +200,7 @@ def get_params(self) -> tuple[torch.Tensor, torch.Tensor]:
         self.encoder.train(False)
         if self.input_type == "scfemb":
             z = self.encoder({**self.pi_aux_data, **self.pi_tensor_data})
-            output = z["qz_mean"], z["qz_var"]
+            output = z[MODULE_KEYS.QZM_KEY], z[MODULE_KEYS.QZV_KEY]
         elif self.input_type == "scvi":
             if self.preparation_function is None:
                 raise ValueError("preparation_function must be provided for scvi input type")

src/drvi/scvi_tools_based/model/_drvi.py

@@ -3,6 +3,7 @@
 from typing import Any, Literal
 
 import numpy as np
+import scvi
 from anndata import AnnData
 from scvi import REGISTRY_KEYS, settings
 from scvi.data import AnnDataManager
@@ -63,7 +64,8 @@ class DRVI(VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass,
 
     def __init__(
         self,
-        adata: AnnData | MerlinData,
+        adata: AnnData | MerlinData | None = None,  # TODO: align with all scvi changes: registry, etc.
+        registry: dict | None = None,  # TODO: align with all scvi changes: registry, etc.
         n_latent: int = 32,
         encoder_dims: Sequence[int] = (128, 128),
         decoder_dims: Sequence[int] = (128, 128),
@@ -72,7 +74,10 @@ def __init__(
         categorical_covariates: list[str] = (),
         **model_kwargs,
     ) -> None:
-        super().__init__(adata)
+        if scvi.__version__ >= "1.3.1":
+            super().__init__(adata, registry)
+        else:
+            super().__init__(adata)
 
         # TODO: Remove later. Currently used to detect autoreload problems sooner.
         if isinstance(adata, AnnData):
@@ -87,7 +92,12 @@ def __init__(
             )
 
         categorical_covariates_info = FeatureInfoList(categorical_covariates, axis="obs", default_dim=10)
-        if REGISTRY_KEYS.CAT_COVS_KEY in self.adata_manager.data_registry:
+        if scvi.__version__ >= "1.3.1" and REGISTRY_KEYS.CAT_COVS_KEY in self.registry["field_registries"]:
+            cat_cov_stats = self.registry["field_registries"][REGISTRY_KEYS.CAT_COVS_KEY]["state_registry"]
+            print(cat_cov_stats)
+            n_cats_per_cov = cat_cov_stats.get("n_cats_per_key", [])
+            assert tuple(categorical_covariates_info.names) == tuple(cat_cov_stats.get("field_keys", []))
+        elif scvi.__version__ < "1.3.1" and REGISTRY_KEYS.CAT_COVS_KEY in self.adata_manager.data_registry:
             cat_cov_stats = self.adata_manager.get_state_registry(REGISTRY_KEYS.CAT_COVS_KEY)
             n_cats_per_cov = cat_cov_stats.n_cats_per_key
             assert tuple(categorical_covariates_info.names) == tuple(cat_cov_stats.field_keys)

src/drvi/scvi_tools_based/model/base/_archesmixin.py

@@ -1,10 +1,12 @@
 import logging
 from collections.abc import Sequence
 
+import scvi
 import torch
 from anndata import AnnData
+from lightning import LightningDataModule
 from scvi import REGISTRY_KEYS
-from scvi.data._constants import _MODEL_NAME_KEY, _SETUP_ARGS_KEY
+from scvi.data._constants import _MODEL_NAME_KEY, _SETUP_ARGS_KEY, _SETUP_METHOD_NAME
 from scvi.model._utils import parse_device_args
 from scvi.model.base import BaseModelClass
 from scvi.model.base._archesmixin import ArchesMixin, _get_loaded_data, _initialize_model, _validate_var_names
@@ -21,8 +23,9 @@ class DRVIArchesMixin(ArchesMixin):
     @classmethod
     def load_query_data(
         cls,
-        adata: AnnData,
-        reference_model: str | BaseModelClass,
+        adata: AnnData = None,
+        reference_model: str | BaseModelClass = None,
+        registry: dict = None,
         inplace_subset_query_vars: bool = False,
         accelerator: str = "auto",
         device: int | str = "auto",
@@ -35,6 +38,7 @@ def load_query_data(
         reset_decoder: bool = False,
         freeze_batchnorm_encoder: bool = True,
         freeze_batchnorm_decoder: bool = False,
+        datamodule: LightningDataModule | None = None,
     ):
         """Online update of a reference model with scArches algorithm :cite:p:`Lotfollahi21`.
 
@@ -67,36 +71,56 @@ def load_query_data(
         freeze_batchnorm_decoder
             Whether to freeze decoder batchnorms' weight and bias during transfer
         """
+        if reference_model is None:
+            raise ValueError("Please provide a reference model as string or loaded model.")
+        if adata is None and registry is None:
+            raise ValueError("Please provide either an AnnData or a registry dictionary.")
+
         _, _, device = parse_device_args(
             accelerator=accelerator,
             devices=device,
             return_device="torch",
             validate_single_device=True,
         )
 
-        attr_dict, var_names, load_state_dict = _get_loaded_data(reference_model, device=device)
+        # We limit to [:3] as from scvi version 1.1.5 additional output (pyro_param_store) is returned
+        attr_dict, var_names, load_state_dict = _get_loaded_data(reference_model, device=device)[:3]
 
-        if inplace_subset_query_vars:
-            logger.debug("Subsetting query vars to reference vars.")
-            adata._inplace_subset_var(var_names)
-        _validate_var_names(adata, var_names)
+        if adata:
+            if inplace_subset_query_vars:
+                logger.debug("Subsetting query vars to reference vars.")
+                adata._inplace_subset_var(var_names)
+            _validate_var_names(adata, var_names)
 
-        registry = attr_dict.pop("registry_")
-        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.")
+            registry = attr_dict.pop("registry_")
+            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.")
+            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,
-            extend_categories=True,
-            allow_missing_labels=True,
-            **registry[_SETUP_ARGS_KEY],
-        )
+            if registry[_SETUP_METHOD_NAME] != "setup_datamodule":
+                setup_method = getattr(cls, registry[_SETUP_METHOD_NAME])
+                setup_method(
+                    adata,
+                    source_registry=registry,
+                    extend_categories=True,
+                    allow_missing_labels=True,
+                    **registry[_SETUP_ARGS_KEY],
+                )
 
-        model = _initialize_model(cls, adata, attr_dict)
+            cls.setup_anndata(
+                adata,
+                source_registry=registry,
+                extend_categories=True,
+                allow_missing_labels=True,
+                **registry[_SETUP_ARGS_KEY],
+            )
+
+        if scvi.__version__ >= "1.3.1":
+            model = _initialize_model(cls, adata, registry, attr_dict, datamodule)
+        else:
+            model = _initialize_model(cls, adata, attr_dict)
         adata_manager = model.get_anndata_manager(adata, required=True)
 
         if REGISTRY_KEYS.CAT_COVS_KEY in adata_manager.data_registry:

src/drvi/scvi_tools_based/model/base/_generative_mixin.py

@@ -9,6 +9,8 @@
 from scvi import REGISTRY_KEYS
 from torch.nn import functional as F
 
+from drvi.scvi_tools_based.module._constants import MODULE_KEYS
+
 logger = logging.getLogger(__name__)
 
 
@@ -93,7 +95,7 @@ def iterate_on_decoded_latent_samples(
         >>> import numpy as np
         >>> # Define custom step function to extract means
         >>> def extract_means(gen_output, store):
-        ...     store.append(gen_output["params"]["mean"].detach().cpu())
+        ...     store.append(gen_output[MODULE_KEYS.PX_PARAMS_KEY]["mean"].detach().cpu())
         >>> # Define aggregation function to concatenate results
         >>> def concatenate_results(store):
         ...     return torch.cat(store, dim=0).numpy()
@@ -138,9 +140,9 @@ def iterate_on_decoded_latent_samples(
                         REGISTRY_KEYS.CAT_COVS_KEY: cat_tensor,
                     },
                     inference_outputs={
-                        "z": z_tensor,
-                        "library": lib_tensor,
-                        "gene_likelihood_additional_info": {},
+                        MODULE_KEYS.Z_KEY: z_tensor,
+                        MODULE_KEYS.LIBRARY_KEY: lib_tensor,
+                        MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY: {},
                     },
                 )
                 gen_output = self.module.generative(**gen_input)
@@ -209,7 +211,7 @@ def decode_latent_samples(
         """
 
         def step_func(gen_output: dict[str, Any], store: list[Any]) -> None:
-            store.append(gen_output["params"]["mean"].detach().cpu())
+            store.append(gen_output[MODULE_KEYS.PX_PARAMS_KEY]["mean"].detach().cpu())
 
         def aggregation_func(store: list[Any]) -> np.ndarray:
             return torch.cat(store, dim=0).numpy(force=True)
@@ -375,15 +377,17 @@ def calculate_effect(
             inference_outputs: dict[str, Any], generative_outputs: dict[str, Any], losses: Any, store: list[Any]
         ) -> None:
             if self.module.split_aggregation == "logsumexp":
-                log_mean_params = generative_outputs["original_params"]["mean"]  # n_samples x n_splits x n_genes
+                log_mean_params = generative_outputs[MODULE_KEYS.PX_UNAGGREGATED_PARAMS_KEY][
+                    "mean"
+                ]  # n_samples x n_splits x n_genes
                 log_mean_params = F.pad(
                     log_mean_params, (0, 0, 0, 1), value=np.log(add_to_counts)
                 )  # n_samples x (n_splits + 1) x n_genes
                 effect_share = -torch.log(1 - F.softmax(log_mean_params, dim=-2)[:, :-1, :]).sum(
                     dim=-1
                 )  # n_samples x n_splits
             elif self.module.split_aggregation == "sum":
-                effect_share = torch.abs(generative_outputs["original_params"]["mean"]).sum(
+                effect_share = torch.abs(generative_outputs[MODULE_KEYS.PX_UNAGGREGATED_PARAMS_KEY]["mean"]).sum(
                     dim=-1
                 )  # n_samples x n_splits
             else:
@@ -474,13 +478,15 @@ def calculate_effect(
             inference_outputs: dict[str, Any], generative_outputs: dict[str, Any], losses: Any, store: list[Any]
         ) -> None:
             if self.module.split_aggregation == "logsumexp":
-                log_mean_params = generative_outputs["original_params"]["mean"]  # n_samples x n_splits x n_genes
+                log_mean_params = generative_outputs[MODULE_KEYS.PX_UNAGGREGATED_PARAMS_KEY][
+                    "mean"
+                ]  # n_samples x n_splits x n_genes
                 log_mean_params = F.pad(
                     log_mean_params, (0, 0, 0, 1), value=np.log(add_to_counts)
                 )  # n_samples x (n_splits + 1) x n_genes
                 effect_share = -torch.log(1 - F.softmax(log_mean_params, dim=-2)[:, :-1, :])
             elif self.module.split_aggregation == "sum":
-                effect_share = torch.abs(generative_outputs["original_params"]["mean"])
+                effect_share = torch.abs(generative_outputs[MODULE_KEYS.PX_UNAGGREGATED_PARAMS_KEY]["mean"])
             else:
                 raise NotImplementedError("Only logsumexp and sum aggregations are supported for now.")
             effect_share = effect_share.amax(dim=0).detach().cpu().numpy(force=True)

src/drvi/scvi_tools_based/module/_constants.py

@@ -0,0 +1,48 @@
+# For backward compatibility
+try:
+    from scvi.module._constants import _MODULE_KEYS as _SCVI_MODULE_KEYS
+except ImportError:
+    from typing import NamedTuple
+
+    class _NEW_SCVI_MODULE_KEYS(NamedTuple):
+        X_KEY: str = "x"
+        # inference
+        Z_KEY: str = "z"
+        QZ_KEY: str = "qz"
+        QZM_KEY: str = "qzm"
+        QZV_KEY: str = "qzv"
+        LIBRARY_KEY: str = "library"
+        QL_KEY: str = "ql"
+        BATCH_INDEX_KEY: str = "batch_index"
+        Y_KEY: str = "y"
+        CONT_COVS_KEY: str = "cont_covs"
+        CAT_COVS_KEY: str = "cat_covs"
+        SIZE_FACTOR_KEY: str = "size_factor"
+        # generative
+        PX_KEY: str = "px"
+        PL_KEY: str = "pl"
+        PZ_KEY: str = "pz"
+        # loss
+        KL_L_KEY: str = "kl_divergence_l"
+        KL_Z_KEY: str = "kl_divergence_z"
+
+    class _SCVI_MODULE_KEYS(_NEW_SCVI_MODULE_KEYS):
+        QZM_KEY: str = "qz_m"
+        QZV_KEY: str = "qz_v"
+
+
+class _DRVI_MODULE_KEYS(_SCVI_MODULE_KEYS):
+    # generative
+    PX_PARAMS_KEY = "px_params"
+    PX_UNAGGREGATED_PARAMS_KEY = "px_unaggregated_params"
+    # Extra
+    LIKELIHOOD_ADDITIONAL_PARAMS_KEY: str = "gene_likelihood_additional_info"
+    X_MASK_KEY: str = "x_mask"
+    # Tensor IO structure
+    CONT_COVS_TENSOR_KEY: str = "cont_full_tensor"
+    CAT_COVS_TENSOR_KEY: str = "cat_full_tensor"
+    # Loss
+    MSE_LOSS_KEY: str = "mse"
+
+
+MODULE_KEYS = _DRVI_MODULE_KEYS()