WIP: Release 0.1.11

.github/dependabot.yml

@@ -0,0 +1,31 @@
+version: 2
+updates:
+  # Enable version updates for uv (Python dependencies)
+  - package-ecosystem: "uv"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+      time: "09:00"
+    open-pull-requests-limit: 10
+    labels:
+      - "dependencies"
+      - "python"
+    commit-message:
+      prefix: "deps"
+      include: "scope"
+
+  # Enable version updates for GitHub Actions
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+      day: "monday"
+      time: "09:00"
+    open-pull-requests-limit: 5
+    labels:
+      - "dependencies"
+      - "github-actions"
+    commit-message:
+      prefix: "ci"
+      include: "scope"

.pre-commit-config.yaml

@@ -11,15 +11,15 @@ repos:
     hooks:
       - id: prettier
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.2
+    rev: v0.14.0
     hooks:
       - id: ruff
         types_or: [python, pyi, jupyter]
         args: [--fix, --exit-non-zero-on-fix]
       - id: ruff-format
         types_or: [python, pyi, jupyter]
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v5.0.0
+    rev: v6.0.0
     hooks:
       - id: detect-private-key
       - id: check-ast

CHANGELOG.md

@@ -2,6 +2,14 @@
 
 ## [Unreleased]
 
+## [0.1.10] - 2025-11-16
+
+- Add RnaSeqMixin from scvi-tools for RNA-seq specific methods (get_normalized_expression, differential_expression, posterior_predictive_sample, get_likelihood_parameters)
+- Fix bug in decode space handling where library size was not considered (issue #46)
+- Update decode_latent_samples logic (decode in log space by default)
+- Code improvements and bug fixes
+- Add dependabot for dependency update notification (not used now, for next releases)
+
 ## [0.1.9] - 2025-07-01
 
 - Add DRVI-APnoEXP baseline

docs/conf.py

@@ -10,6 +10,9 @@
 from importlib.metadata import metadata
 from pathlib import Path
 
+# Temporary fix
+import drvi  # noqa: F401
+
 HERE = Path(__file__).parent
 sys.path.insert(0, str(HERE / "extensions"))
 

pyproject.toml

@@ -4,7 +4,7 @@ requires = ["hatchling"]
 
 [project]
 name = "drvi-py"
-version = "0.1.9"
+version = "0.1.10"
 description = "Unsupervised Deep Disentangled Representation of Single-Cell Omics"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -60,7 +60,6 @@ optional-dependencies.dev = [
     "twine>=4.0.2",
 ]
 optional-dependencies.doc = [
-    # Disable for now as nvidia servers return 404
     "merlin-dataloader==23.8.0",
     "docutils>=0.8,!=0.18.*,!=0.19.*",
     "sphinx>=4",

src/drvi/nn_modules/embedding.py

@@ -326,6 +326,9 @@ def freeze_top_embs(self, n_freeze_list: list[int]) -> None:
         table, which is useful for transfer learning scenarios.
         """
         for emb, n_freeze in zip(self.emb_list, n_freeze_list, strict=False):
+            # If that specific category has no change in size, skip.
+            if not emb.weight.requires_grad:
+                continue
             emb.freeze(n_freeze, emb.embedding_dim)
 
     @property

src/drvi/nn_modules/noise_model.py

@@ -304,10 +304,14 @@ def dist(self, aux_info, parameters, lib_y):
         Normal
             Normal distribution with specified mean and variance.
         """
+        mean = parameters["mean"]
         var = parameters["var"]
         if self.model_var:
             var = torch.nan_to_num(torch.exp(var), posinf=100, neginf=0) + self.eps
-        return Normal(parameters["mean"], torch.abs(var).sqrt())
+        output_dist = Normal(mean, torch.abs(var).sqrt())
+        output_dist.mu = mean  # Just for scvi RNASeqMixin compatibility
+        output_dist.theta = var  # Just for scvi RNASeqMixin compatibility
+        return output_dist
 
 
 class PoissonNoiseModel(NoiseModel):
@@ -399,14 +403,18 @@ def dist(self, aux_info, parameters, lib_y):
         library_size = lib_y
 
         if self.mean_transformation == "exp":
-            trans_mean = torch.exp(mean)
-            trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False)
+            trans_scale = torch.exp(mean)
+            trans_mean = library_size_correction(trans_scale, library_size, self.library_normalization, log_space=False)
         elif self.mean_transformation == "softmax":
-            trans_mean = torch.softmax(mean, dim=-1)
-            trans_mean = library_size.unsqueeze(-1) * trans_mean
+            trans_scale = torch.softmax(mean, dim=-1)
+            trans_mean = library_size.unsqueeze(-1) * trans_scale
         else:
             raise NotImplementedError()
-        return Poisson(trans_mean)
+        output_dist = Poisson(trans_mean)
+        output_dist.scale = trans_scale  # Just for scvi RNASeqMixin compatibility
+        output_dist.mu = trans_mean  # Just for scvi RNASeqMixin compatibility
+        output_dist.theta = torch.ones_like(trans_mean)  # Just for scvi RNASeqMixin compatibility
+        return output_dist
 
 
 class NegativeBinomialNoiseModel(NoiseModel):
@@ -494,22 +502,22 @@ def dist(self, aux_info, parameters, lib_y):
         library_size = lib_y
 
         if self.mean_transformation == "exp":
-            trans_mean = torch.exp(mean)
-            trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False)
+            px_scale = torch.exp(mean)
+            px_rate = library_size_correction(px_scale, library_size, self.library_normalization, log_space=False)
         elif self.mean_transformation == "softmax":
-            trans_mean = torch.softmax(mean, dim=-1)
-            trans_mean = library_size.unsqueeze(-1) * trans_mean
+            px_scale = torch.softmax(mean, dim=-1)
+            px_rate = library_size.unsqueeze(-1) * px_scale
         # `softplus` and `none` for ablation. Useless in practice.
         elif self.mean_transformation == "softplus":
-            trans_mean = F.softplus(mean)
-            trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False)
+            px_scale = F.softplus(mean)
+            px_rate = library_size_correction(px_scale, library_size, self.library_normalization, log_space=False)
         elif self.mean_transformation == "none":
-            trans_mean = mean
-            trans_mean = library_size_correction(trans_mean, library_size, self.library_normalization, log_space=False)
+            px_scale = mean
+            px_rate = library_size_correction(px_scale, library_size, self.library_normalization, log_space=False)
         else:
             raise NotImplementedError()
         trans_r = torch.exp(r)
-        return NegativeBinomial(mu=trans_mean, theta=trans_r, scale=None)
+        return NegativeBinomial(mu=px_rate, theta=trans_r, scale=px_scale)
 
 
 class LogNegativeBinomial(Distribution):
@@ -536,10 +544,13 @@ class LogNegativeBinomial(Distribution):
     are computed as exp(log_m) and exp(log_r) respectively.
     """
 
-    def __init__(self, log_m, log_r, eps: float = 1e-8, validate_args=False) -> None:
+    def __init__(
+        self, log_m, log_r, log_scale: torch.Tensor | None = None, eps: float = 1e-8, validate_args=False
+    ) -> None:
         self.log_m = log_m
         self.log_r = log_r
         self._eps = eps
+        self.log_scale = log_scale
         super().__init__(validate_args=validate_args)
 
     @property
@@ -633,6 +644,17 @@ def negative_binomial_log_ver(k, m_log, r_log, eps=1e-8):
     def log_prob(self, value: torch.Tensor) -> torch.Tensor:
         return self.negative_binomial_log_ver(value, self.log_m, self.log_r, eps=self._eps)
 
+    # Additional properties for compatibility with scvi RNASeqMixin
+    @property
+    def mu(self) -> torch.Tensor:
+        return torch.exp(self.log_m)
+
+    @property
+    def scale(self) -> torch.Tensor | None:
+        if self.log_scale is None:
+            return None
+        return torch.exp(self.log_scale)
+
 
 class LogNegativeBinomialNoiseModel(NoiseModel):
     """Log-space negative binomial noise model.
@@ -724,11 +746,12 @@ def dist(self, aux_info, parameters, lib_y):
         library_size = lib_y
 
         if self.mean_transformation == "none":
-            trans_mean = library_size_correction(mean, library_size, self.library_normalization, log_space=True)
+            trans_scale = mean
+            trans_mean = library_size_correction(trans_scale, library_size, self.library_normalization, log_space=True)
         elif self.mean_transformation == "softmax":
-            trans_mean = mean - torch.logsumexp(mean, dim=-1, keepdim=True)
-            trans_mean = torch.log(library_size).unsqueeze(-1) + trans_mean
+            trans_scale = mean - torch.logsumexp(mean, dim=-1, keepdim=True)
+            trans_mean = torch.log(library_size).unsqueeze(-1) + trans_scale
         else:
             raise NotImplementedError()
         trans_r = r
-        return LogNegativeBinomial(log_m=trans_mean, log_r=trans_r)
+        return LogNegativeBinomial(log_m=trans_mean, log_r=trans_r, log_scale=trans_scale)

src/drvi/scvi_tools_based/model/_drvi.py

@@ -7,7 +7,7 @@
 from scvi import REGISTRY_KEYS, settings
 from scvi.data import AnnDataManager
 from scvi.data.fields import CategoricalObsField, LayerField, NumericalJointObsField
-from scvi.model.base import BaseModelClass, UnsupervisedTrainingMixin, VAEMixin
+from scvi.model.base import BaseModelClass, RNASeqMixin, UnsupervisedTrainingMixin, VAEMixin
 from scvi.utils import setup_anndata_dsp
 
 import drvi
@@ -28,7 +28,7 @@
 logger = logging.getLogger(__name__)
 
 
-class DRVI(VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass, GenerativeMixin):
+class DRVI(RNASeqMixin, VAEMixin, DRVIArchesMixin, UnsupervisedTrainingMixin, BaseModelClass, GenerativeMixin):
     """DRVI model based on scvi-tools framework for disentangled representation learning.
 
     Parameters
@@ -69,6 +69,7 @@ def __init__(
         decoder_dims: Sequence[int] = (128, 128),
         prior: Literal["normal", "gmm_x", "vamp_x"] = "normal",
         prior_init_obs: np.ndarray | None = None,
+        batch_key: str | None = None,
         categorical_covariates: list[str] = (),
         **model_kwargs,
     ) -> None:
@@ -86,14 +87,34 @@ def __init__(
                 "make sure merlin is installed as a dependency."
             )
 
-        categorical_covariates_info = FeatureInfoList(categorical_covariates, axis="obs", default_dim=10)
+        n_batch = self.summary_stats.n_batch
+        n_cats_per_cov = [n_batch]
         if 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)
+            n_cats_per_cov = n_cats_per_cov + list(cat_cov_stats.n_cats_per_key)
+
+        if batch_key is not None:
+            all_categorical_covariates = [batch_key] + list(categorical_covariates)
+            has_batch = 1
+        else:
+            all_categorical_covariates = list(categorical_covariates)
+            has_batch = 0
+        categorical_covariates_info = FeatureInfoList(all_categorical_covariates, axis="obs", default_dim=10)
+
+        # validations
+        if n_batch > 1:
+            batch_original_key = self.adata_manager.get_state_registry(REGISTRY_KEYS.BATCH_KEY).original_key
+            assert categorical_covariates_info.names[0] == batch_original_key
+            categorical_covariates_dims = categorical_covariates_info.dims
         else:
-            n_cats_per_cov = []
-            assert len(categorical_covariates_info) == 0
+            assert batch_key is None
+            categorical_covariates_dims = [1] + categorical_covariates_info.dims
+        if 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)
+            assert tuple(categorical_covariates_info.names[has_batch:]) == tuple(cat_cov_stats.field_keys)
+        else:
+            assert len(categorical_covariates_info) == has_batch
+
         n_continuous_cov = self.summary_stats.get("n_extra_continuous_covs", 0)
 
         prior_init_dataloader = None
@@ -113,7 +134,7 @@ def __init__(
             n_continuous_cov=n_continuous_cov,
             prior=prior,
             prior_init_dataloader=prior_init_dataloader,
-            categorical_covariate_dims=categorical_covariates_info.dims,
+            categorical_covariate_dims=categorical_covariates_dims,
             **model_kwargs,
         )
 
@@ -132,6 +153,39 @@ def __init__(
 
         logger.info("The model has been initialized")
 
+    @staticmethod
+    def _update_source_registry_for_existing_model(source_registry: dict[str, Any]) -> dict[str, Any]:
+        """Update the source registry for an existing model to the latest version if any updates are needed."""
+        from packaging.version import Version
+
+        source_registry_drvi_version = Version(
+            source_registry.get("drvi_version", "0.1.0")
+        )  # "0.1.0" for legacy code before pypi release
+        logger.info(f"The model is trained with DRVI version {source_registry_drvi_version}.")
+
+        while source_registry_drvi_version < Version(drvi.__version__):
+            if source_registry_drvi_version < Version("0.1.9"):
+                # No braking change up to 0.1.9
+                source_registry_drvi_version = Version("0.1.9")
+            elif source_registry_drvi_version == Version("0.1.9"):
+                # log the transfer
+                logger.info("Modifying model args from 0.1.9 to 0.1.10 (no user action required)")
+                logger.info("Adding empty batch key ...")
+                source_registry["setup_args"]["batch_key"] = None
+                source_registry["field_registries"]["batch"] = {
+                    "data_registry": {"attr_name": "obs", "attr_key": "_scvi_batch"},
+                    "state_registry": {"categorical_mapping": np.array([0]), "original_key": "_scvi_batch"},
+                    "summary_stats": {"n_batch": 1},
+                }
+                source_registry_drvi_version = Version("0.1.10")
+                logger.info("Done updating source registry from 0.1.9 to 0.1.10.")
+            else:
+                # No braking change yet!
+                source_registry_drvi_version = Version(drvi.__version__)
+        logger.info(f"Loading source in DRVI version {drvi.__version__}.")
+
+        return source_registry
+
     @classmethod
     @setup_anndata_dsp.dedent
     def setup_anndata(
@@ -140,6 +194,7 @@ def setup_anndata(
         labels_key: str | None = None,
         layer: str | None = None,
         is_count_data: bool = True,
+        batch_key: str | None = None,
         categorical_covariate_keys: list[str] | None = None,
         continuous_covariate_keys: list[str] | None = None,
         **kwargs,
@@ -152,6 +207,7 @@ def setup_anndata(
         %(param_adata)s
         %(param_labels_key)s
         %(param_layer)s
+        %(param_batch_key)s
         %(param_cat_cov_keys)s
         %(param_cont_cov_keys)s
 
@@ -164,11 +220,19 @@ def setup_anndata(
 
         anndata_fields = [
             LayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=is_count_data),
+            CategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),
             CategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key),
             FixedCategoricalJointObsField(REGISTRY_KEYS.CAT_COVS_KEY, categorical_covariate_keys),
             NumericalJointObsField(REGISTRY_KEYS.CONT_COVS_KEY, continuous_covariate_keys),
         ]
         adata_manager = AnnDataManager(fields=anndata_fields, setup_method_args=setup_method_args)
+
+        # We may need to manupulate in case of version updates (only when loading a model).
+        if "source_registry" in kwargs:
+            source_registry = kwargs["source_registry"]
+            source_registry = cls._update_source_registry_for_existing_model(source_registry)
+            kwargs["source_registry"] = source_registry
+
         adata_manager.register_fields(adata, **kwargs)
         cls.register_manager(adata_manager)
 
@@ -179,6 +243,7 @@ def setup_merlin_data(
         labels_key: str | None = None,
         layer: str = "X",
         is_count_data: bool = True,
+        batch_key: str | None = None,
         categorical_covariate_keys: list[str] | None = None,
         continuous_covariate_keys: list[str] | None = None,
         **kwargs,
@@ -195,6 +260,8 @@ def setup_merlin_data(
             key in `merlin_data` to use as input.
         is_count_data
             Whether the data is count data.
+        batch_key
+            Key in `merlin_data` for batch information.
         categorical_covariate_keys
             List of categorical covariate keys in `merlin_data`.
         continuous_covariate_keys
@@ -212,6 +279,7 @@ def setup_merlin_data(
 
         fields = [
             melin_fields.MerlinLayerField(REGISTRY_KEYS.X_KEY, layer, is_count_data=is_count_data),
+            melin_fields.MerlinCategoricalObsField(REGISTRY_KEYS.BATCH_KEY, batch_key),
             melin_fields.MerlinCategoricalObsField(REGISTRY_KEYS.LABELS_KEY, labels_key),
             melin_fields.MerlinCategoricalJointObsField(REGISTRY_KEYS.CAT_COVS_KEY, categorical_covariate_keys),
             melin_fields.MerlinNumericalJointObsField(REGISTRY_KEYS.CONT_COVS_KEY, continuous_covariate_keys),