From ddaa220d3eea13357181eb4f32458bcd309a4eb6 Mon Sep 17 00:00:00 2001 From: moinfar Date: Wed, 21 May 2025 23:48:23 +0200 Subject: [PATCH 1/7] Relax dependency versions --- pyproject.toml | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c298e54..fdbbe6e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,23 +35,19 @@ classifiers = [ # Please make an issue if you need wider range of versions dependencies = [ - "torch>=2.1.0,<2.4", - "lightning>=2.0,<2.1", - "scanpy==1.9.5", + "torch>=2.1.0", + "lightning>=2.0", + "scanpy>=1.9.5", "scikit-learn>=1.5.1", - "scipy>=1.11.3,<1.13", - "scvi-tools==1.0.4", + "scipy>=1.11.3", + "scvi-tools>=1.0.4", ## TODO: Remove when scvi-tools and jax become ok again "jax<=0.4.20", "jaxlib<=0.4.20", ## END_TODO - "anndata>=0.10.2,<0.11", - "numpy>=1.16.1", # for np.linspace + "anndata>=0.10.2", + "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 - "urllib3<2", - ## END_TODO # for debug logging (referenced from the issue template) "session-info", ] From 0d2509064019febf04aac21ccee8daf566dee8ea Mon Sep 17 00:00:00 2001 From: moinfar Date: Thu, 22 May 2025 14:51:52 +0200 Subject: [PATCH 2/7] Relax scvi dependency and align the code with the newest changes --- pyproject.toml | 19 +++- src/drvi/nn_modules/prior.py | 4 +- .../model/base/_archesmixin.py | 66 ++++++++----- .../model/base/_generative_mixin.py | 22 +++-- .../scvi_tools_based/module/_constants.py | 48 ++++++++++ src/drvi/scvi_tools_based/module/_drvi.py | 93 ++++++++++--------- 6 files changed, 177 insertions(+), 75 deletions(-) create mode 100644 src/drvi/scvi_tools_based/module/_constants.py diff --git a/pyproject.toml b/pyproject.toml index fdbbe6e..ad7799d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,13 +41,30 @@ dependencies = [ "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", + "scikit-learn>=1.5.1", + "scipy>=1.11.3,<1.13", + "scvi-tools==1.0.4", ## TODO: Remove when scvi-tools and jax become ok again "jax<=0.4.20", "jaxlib<=0.4.20", ## END_TODO - "anndata>=0.10.2", + "anndata>=0.10.2,<0.11", "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 + "urllib3<2", + ## END_TODO # for debug logging (referenced from the issue template) "session-info", ] diff --git a/src/drvi/nn_modules/prior.py b/src/drvi/nn_modules/prior.py index 2d24560..353a3eb 100644 --- a/src/drvi/nn_modules/prior.py +++ b/src/drvi/nn_modules/prior.py @@ -2,6 +2,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 @@ -65,7 +67,7 @@ def get_params(self): 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": x, args, kwargs = self.preparation_function({**self.pi_aux_data, **self.pi_tensor_data}) q_m, q_v, latent = self.encoder(x, *args, **kwargs) diff --git a/src/drvi/scvi_tools_based/model/base/_archesmixin.py b/src/drvi/scvi_tools_based/model/base/_archesmixin.py index 445ae8b..257ba43 100644 --- a/src/drvi/scvi_tools_based/model/base/_archesmixin.py +++ b/src/drvi/scvi_tools_based/model/base/_archesmixin.py @@ -1,9 +1,11 @@ import logging +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 @@ -20,8 +22,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", @@ -34,6 +37,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`. @@ -66,6 +70,11 @@ 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, @@ -73,29 +82,44 @@ def load_query_data( 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: diff --git a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py index 9b54311..8bfb6c3 100644 --- a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py +++ b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py @@ -7,6 +7,8 @@ from anndata import AnnData from torch.nn import functional as F +from drvi.scvi_tools_based.module._constants import MODULE_KEYS + logger = logging.getLogger(__name__) @@ -68,9 +70,9 @@ def iterate_on_decoded_latent_samples( scvi.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) @@ -106,7 +108,7 @@ def decode_latent_samples( return_mean Return the mean of the distribution or the full distribution. """ - step_func = lambda gen_output, store: store.append(gen_output["params"]["mean"].detach().cpu()) + step_func = lambda gen_output, store: store.append(gen_output[MODULE_KEYS.PX_PARAMS_KEY]["mean"].detach().cpu()) aggregation_func = lambda store: torch.cat(store, dim=0).numpy(force=True) return self.iterate_on_decoded_latent_samples( @@ -195,7 +197,9 @@ def get_reconstruction_effect_of_each_split( def calculate_effect(inference_outputs, generative_outputs, losses, store): 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 @@ -203,7 +207,7 @@ def calculate_effect(inference_outputs, generative_outputs, losses, store): 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: @@ -278,13 +282,15 @@ def get_max_effect_of_splits_within_distribution( def calculate_effect(inference_outputs, generative_outputs, losses, store): 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) diff --git a/src/drvi/scvi_tools_based/module/_constants.py b/src/drvi/scvi_tools_based/module/_constants.py new file mode 100644 index 0000000..87d4f0e --- /dev/null +++ b/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() diff --git a/src/drvi/scvi_tools_based/module/_drvi.py b/src/drvi/scvi_tools_based/module/_drvi.py index bbba0ba..5c7d8d9 100644 --- a/src/drvi/scvi_tools_based/module/_drvi.py +++ b/src/drvi/scvi_tools_based/module/_drvi.py @@ -17,6 +17,7 @@ PoissonNoiseModel, ) from drvi.nn_modules.prior import GaussianMixtureModelPrior, StandardPrior, VampPrior +from drvi.scvi_tools_based.module._constants import MODULE_KEYS from drvi.scvi_tools_based.nn import DecoderDRVI, Encoder TensorDict = dict[str, torch.Tensor] @@ -281,7 +282,7 @@ def _construct_prior(self, prior, prior_init_dataloader=None): n_components = int(prior.split("_")[1]) if prior_init_dataloader is not None: inference_output = self.inference(**self._get_inference_input(next(iter(prior_init_dataloader)))) - init_data = inference_output["qz_m"], inference_output["qz_v"] + init_data = inference_output[MODULE_KEYS.QZM_KEY], inference_output[MODULE_KEYS.QZV_KEY] else: init_data = None return GaussianMixtureModelPrior(n_components, self.n_latent, data=init_data) @@ -290,9 +291,9 @@ def _construct_prior(self, prior, prior_init_dataloader=None): if prior_init_dataloader is not None: def preparation_function(prepared_input): - x = prepared_input["encoder_input"] + x = prepared_input[MODULE_KEYS.X_KEY] args = [] - kwargs = {"cat_full_tensor": prepared_input["cat_full_tensor"]} + kwargs = {"cat_full_tensor": prepared_input[MODULE_KEYS.CAT_COVS_KEY]} return x, args, kwargs model_input = self._input_pre_processing(**self._get_inference_input(next(iter(prior_init_dataloader)))) @@ -303,8 +304,8 @@ def preparation_function(prepared_input): self.z_encoder, model_input, input_type="scvi", - trainable_keys=("encoder_input",), - fixed_keys=("cat_full_tensor",), + trainable_keys=(MODULE_KEYS.X_KEY,), + fixed_keys=(MODULE_KEYS.CAT_COVS_KEY,), preparation_function=preparation_function, ) else: @@ -317,7 +318,7 @@ def _get_inference_input(self, tensors): cont_covs = tensors.get(REGISTRY_KEYS.CONT_COVS_KEY) cat_covs = tensors.get(REGISTRY_KEYS.CAT_COVS_KEY) - input_dict = {"x": x, "cont_covs": cont_covs, "cat_covs": cat_covs} + input_dict = {MODULE_KEYS.X_KEY: x, MODULE_KEYS.CONT_COVS_KEY: cont_covs, MODULE_KEYS.CAT_COVS_KEY: cat_covs} return input_dict def _input_pre_processing(self, x, cont_covs=None, cat_covs=None): @@ -330,11 +331,11 @@ def _input_pre_processing(self, x, cont_covs=None, cat_covs=None): encoder_input = x_ return { - "encoder_input": encoder_input, - "cat_full_tensor": cat_covs if self.encode_covariates else None, - "cont_full_tensor": cont_covs if self.encode_covariates else None, - "library": library, - "gene_likelihood_additional_info": gene_likelihood_additional_info, + MODULE_KEYS.X_KEY: encoder_input, + MODULE_KEYS.CAT_COVS_KEY: cat_covs if self.encode_covariates else None, + MODULE_KEYS.CONT_COVS_KEY: cont_covs if self.encode_covariates else None, + MODULE_KEYS.LIBRARY_KEY: library, + MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY: gene_likelihood_additional_info, } @auto_move_data @@ -345,7 +346,7 @@ def inference(self, x, cont_covs=None, cat_covs=None): Runs the inference (encoder) model. """ pre_processed_input = self._input_pre_processing(x, cont_covs, cat_covs).copy() - x_ = pre_processed_input["encoder_input"] + x_ = pre_processed_input[MODULE_KEYS.X_KEY] # Mask if needed if self.fill_in_the_blanks_ratio > 0.0 and self.training: @@ -359,43 +360,45 @@ def inference(self, x, cont_covs=None, cat_covs=None): # Prepare shared emb if self.shared_covariate_emb is not None and self.encode_covariates: - pre_processed_input["cat_full_tensor"] = self.shared_covariate_emb( - pre_processed_input["cat_full_tensor"].int() + pre_processed_input[MODULE_KEYS.CAT_COVS_KEY] = self.shared_covariate_emb( + pre_processed_input[MODULE_KEYS.CAT_COVS_KEY].int() ) # get variational parameters via the encoder networks qz_m, qz_v, z = self.z_encoder( x_, - cat_full_tensor=pre_processed_input["cat_full_tensor"], - cont_full_tensor=pre_processed_input["cont_full_tensor"], + cat_full_tensor=pre_processed_input[MODULE_KEYS.CAT_COVS_KEY], + cont_full_tensor=pre_processed_input[MODULE_KEYS.CONT_COVS_KEY], ) outputs = { - "z": z, - "qz_m": qz_m, - "qz_v": qz_v, - "library": pre_processed_input["library"], - "x_mask": x_mask, - "gene_likelihood_additional_info": pre_processed_input["gene_likelihood_additional_info"], + MODULE_KEYS.Z_KEY: z, + MODULE_KEYS.QZM_KEY: qz_m, + MODULE_KEYS.QZV_KEY: qz_v, + MODULE_KEYS.LIBRARY_KEY: pre_processed_input[MODULE_KEYS.LIBRARY_KEY], + MODULE_KEYS.X_MASK_KEY: x_mask, + MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY: pre_processed_input[ + MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY + ], } return outputs def _get_generative_input(self, tensors, inference_outputs): - z = inference_outputs["z"] + z = inference_outputs[MODULE_KEYS.Z_KEY] if self.fully_deterministic: - z = inference_outputs["qz_m"] - library = inference_outputs["library"] - gene_likelihood_additional_info = inference_outputs["gene_likelihood_additional_info"] + z = inference_outputs[MODULE_KEYS.QZM_KEY] + library = inference_outputs[MODULE_KEYS.LIBRARY_KEY] + gene_likelihood_additional_info = inference_outputs[MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY] cont_covs = tensors.get(REGISTRY_KEYS.CONT_COVS_KEY) cat_covs = tensors.get(REGISTRY_KEYS.CAT_COVS_KEY) input_dict = { - "z": z, - "library": library, - "gene_likelihood_additional_info": gene_likelihood_additional_info, - "cont_covs": cont_covs, - "cat_covs": cat_covs, + MODULE_KEYS.Z_KEY: z, + MODULE_KEYS.LIBRARY_KEY: library, + MODULE_KEYS.LIKELIHOOD_ADDITIONAL_PARAMS_KEY: gene_likelihood_additional_info, + MODULE_KEYS.CONT_COVS_KEY: cont_covs, + MODULE_KEYS.CAT_COVS_KEY: cat_covs, } return input_dict @@ -414,9 +417,9 @@ def generative(self, z, library, gene_likelihood_additional_info, cont_covs=None ) return { - "px": px, - "params": params, - "original_params": original_params, + MODULE_KEYS.PX_KEY: px, + MODULE_KEYS.PX_PARAMS_KEY: params, + MODULE_KEYS.PX_UNAGGREGATED_PARAMS_KEY: original_params, } def loss( @@ -428,10 +431,10 @@ def loss( ): """Loss function.""" x = tensors[REGISTRY_KEYS.X_KEY] - x_mask = inference_outputs["x_mask"] - qz_m = inference_outputs["qz_m"] - qz_v = inference_outputs["qz_v"] - px = generative_outputs["px"] + x_mask = inference_outputs[MODULE_KEYS.X_MASK_KEY] + qz_m = inference_outputs[MODULE_KEYS.QZM_KEY] + qz_v = inference_outputs[MODULE_KEYS.QZV_KEY] + px = generative_outputs[MODULE_KEYS.PX_KEY] kl_divergence_z = self.prior.kl(Normal(qz_m, torch.sqrt(qz_v))).sum(dim=1) if self.fill_in_the_blanks_ratio > 0.0 and self.training: @@ -449,13 +452,15 @@ def loss( loss = torch.mean(reconst_loss + weighted_kl_local) - kl_local = {"kl_divergence_z": kl_divergence_z.sum()} + kl_local = {MODULE_KEYS.KL_Z_KEY: kl_divergence_z.sum()} return LossOutput( loss=loss, reconstruction_loss=reconst_loss, kl_local=kl_local, extra_metrics={ - "mse": torch.nn.functional.mse_loss(x, px.mean, reduction="none").sum(dim=1).mean(dim=0), + MODULE_KEYS.MSE_LOSS_KEY: torch.nn.functional.mse_loss(x, px.mean, reduction="none") + .sum(dim=1) + .mean(dim=0), }, ) @@ -496,7 +501,7 @@ def sample( compute_loss=False, ) - dist = generative_outputs["px"] + dist = generative_outputs[MODULE_KEYS.PX_KEY] if n_samples > 1: exprs = dist.sample().permute([1, 2, 0]) # Shape : (n_cells_batch, n_genes, n_samples) @@ -517,9 +522,9 @@ def marginal_ll(self, tensors: TensorDict, n_mc_samples: int): for i in range(n_mc_samples): # Distribution parameters and sampled variables inference_outputs, _, losses = self.forward(tensors) - qz_m = inference_outputs["qz_m"] - qz_v = inference_outputs["qz_v"] - z = inference_outputs["z"] + qz_m = inference_outputs[MODULE_KEYS.QZM_KEY] + qz_v = inference_outputs[MODULE_KEYS.QZV_KEY] + z = inference_outputs[MODULE_KEYS.Z_KEY] # Reconstruction Loss reconst_loss = losses.dict_sum(losses.reconstruction_loss) From f21a16a6aabc3fc5df54273b9158611cbb401886 Mon Sep 17 00:00:00 2001 From: moinfar Date: Fri, 23 May 2025 16:29:33 +0200 Subject: [PATCH 3/7] Allow passing registry and setting adata=None --- src/drvi/scvi_tools_based/model/_drvi.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/drvi/scvi_tools_based/model/_drvi.py b/src/drvi/scvi_tools_based/model/_drvi.py index 1f232eb..ad51a8e 100644 --- a/src/drvi/scvi_tools_based/model/_drvi.py +++ b/src/drvi/scvi_tools_based/model/_drvi.py @@ -62,7 +62,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), @@ -71,7 +72,7 @@ def __init__( categorical_covariates: list[str] = (), **model_kwargs, ): - super().__init__(adata) + super().__init__(adata, registry) # TODO: Remove later. Currently used to detect autoreload problems sooner. if isinstance(adata, AnnData): @@ -86,10 +87,11 @@ 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: - 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) + if 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", [])) else: n_cats_per_cov = [] assert len(categorical_covariates_info) == 0 From fac03165afb9daea0a69afb0d750c0d006301406 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 9 Jul 2025 16:03:52 +0000 Subject: [PATCH 4/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/drvi/scvi_tools_based/model/base/_generative_mixin.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py index df249c5..0743802 100644 --- a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py +++ b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py @@ -209,6 +209,7 @@ def decode_latent_samples( >>> cat_covs = np.array([0, 1, 0, 1] * 25) # batch labels >>> reconstructed = model.decode_latent_samples(z, cat_values=cat_covs) """ + def step_func(gen_output: dict[str, Any], store: list[Any]) -> None: store.append(gen_output["params"]["mean"].detach().cpu()) From 8695408d3cb873e897626ffc36e4ea3f01472932 Mon Sep 17 00:00:00 2001 From: moinfar Date: Wed, 9 Jul 2025 19:22:37 +0200 Subject: [PATCH 5/7] Fix deprecated one_hot --- src/drvi/scvi_tools_based/model/_drvi.py | 12 ++++++++++-- src/drvi/scvi_tools_based/nn/_base_components.py | 3 +-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/drvi/scvi_tools_based/model/_drvi.py b/src/drvi/scvi_tools_based/model/_drvi.py index 34b21af..8bf3ba4 100644 --- a/src/drvi/scvi_tools_based/model/_drvi.py +++ b/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 @@ -73,7 +74,10 @@ def __init__( categorical_covariates: list[str] = (), **model_kwargs, ) -> None: - super().__init__(adata, registry) + 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): @@ -88,11 +92,15 @@ def __init__( ) categorical_covariates_info = FeatureInfoList(categorical_covariates, axis="obs", default_dim=10) - if REGISTRY_KEYS.CAT_COVS_KEY in self.registry["field_registries"]: + 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) else: n_cats_per_cov = [] assert len(categorical_covariates_info) == 0 diff --git a/src/drvi/scvi_tools_based/nn/_base_components.py b/src/drvi/scvi_tools_based/nn/_base_components.py index 33a8fd0..0e8d59e 100644 --- a/src/drvi/scvi_tools_based/nn/_base_components.py +++ b/src/drvi/scvi_tools_based/nn/_base_components.py @@ -4,7 +4,6 @@ from typing import Any, Literal import torch -from scvi.nn._utils import one_hot from torch import nn from torch.distributions import Normal from torch.nn import functional as F @@ -319,7 +318,7 @@ def forward(self, x: torch.Tensor, cat_full_tensor: torch.Tensor | None) -> torc for n_cat, cat in zip(self.n_cat_list, cat_list, strict=False): if n_cat and cat is None: raise ValueError("cat not provided while n_cat != 0 in init. params.") - concat_list += [one_hot(cat, n_cat)] + concat_list += [F.one_hot(cat.long().squeeze(-1), n_cat).float()] elif self.covariate_vector_modeling == "emb_shared": concat_list = [cat_full_tensor] else: From 114a80670694a5a8121bc04d4e3d5cea2eece3eb Mon Sep 17 00:00:00 2001 From: moinfar Date: Wed, 9 Jul 2025 19:29:46 +0200 Subject: [PATCH 6/7] Bugfix --- src/drvi/scvi_tools_based/model/base/_generative_mixin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py index 0743802..9656b31 100644 --- a/src/drvi/scvi_tools_based/model/base/_generative_mixin.py +++ b/src/drvi/scvi_tools_based/model/base/_generative_mixin.py @@ -95,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() @@ -211,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) From 533aeb7e8c16612757c4fce19fd5d021942ea8b6 Mon Sep 17 00:00:00 2001 From: moinfar Date: Thu, 10 Jul 2025 10:07:30 +0200 Subject: [PATCH 7/7] Change deprecated csr_matrix .A to .toarray() --- src/drvi/utils/tools/interpretability/_latent_traverse.py | 2 +- tests/drvi_model/test_model.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/drvi/utils/tools/interpretability/_latent_traverse.py b/src/drvi/utils/tools/interpretability/_latent_traverse.py index 8e288bb..2960a1d 100644 --- a/src/drvi/utils/tools/interpretability/_latent_traverse.py +++ b/src/drvi/utils/tools/interpretability/_latent_traverse.py @@ -230,7 +230,7 @@ def make_traverse_adata( # Control and effect latent data control_data = noise_vector - effect_data = noise_vector + span_adata.X.A + effect_data = noise_vector + span_adata.X.toarray() print("traversing latent ...") diff --git a/tests/drvi_model/test_model.py b/tests/drvi_model/test_model.py index 045da17..e835a4b 100644 --- a/tests/drvi_model/test_model.py +++ b/tests/drvi_model/test_model.py @@ -51,9 +51,9 @@ def make_test_adata(self, is_sparse=True): adata.layers["lognorm"] = np.log1p(adata.X) if not is_sparse: - adata.X = adata.X.A + adata.X = adata.X.toarray() for l in ["counts", "lognorm"]: - adata.layers[l] = adata.layers[l].A + adata.layers[l] = adata.layers[l].toarray() return adata