scvi.model.mlxSCVI#

class scvi.model.mlxSCVI(adata, n_hidden=128, n_latent=10, dropout_rate=0.1, gene_likelihood='nb', **model_kwargs)[source]#

Single-cell variational inference model using the MLX framework.

This implementation leverages the features of the MLX framework to provide optimized performance on Apple Silicon chips.

Parameters:

adata AnnData: AnnData object registered via mlxSCVI.setup_anndata().
n_hidden int (default: 128): Number of nodes per hidden layer.
n_latent int (default: 10): Dimensionality of the latent space.
dropout_rate float (default: 0.1): Dropout rate for neural networks.
gene_likelihood Literal['nb', 'poisson'] (default: 'nb'): One of: * ‘nb’ - Negative binomial distribution * ‘poisson’ - Poisson distribution

Attributes table#

`adata`	Data attached to model instance.
`adata_manager`	Manager instance associated with self.adata.
`device`	Get the current device.
`get_normalized_function_name`	What the get normalized functions name is
`history`	Returns computed metrics during training.
`is_trained`	Whether the model has been trained.
`registry`	Data attached to model instance.
`run_id`	Returns the run id of the model.
`run_name`	Returns the run name of the model.
`summary_string`	Summary string of the model.
`test_indices`	Observations that are in test set.
`train_indices`	Observations that are in train set.
`validation_indices`	Observations that are in validation set.

Methods table#

`convert_legacy_save`(dir_path, output_dir_path)	Converts a legacy saved model (<v0.15.0) to the updated save format.
`data_registry`(registry_key)	Returns the object in AnnData associated with the key in the data registry.
`deregister_manager`([adata])	Deregisters the `AnnDataManager` instance associated with `adata`.
`differential_abundance`(args, *kwargs)	Not implemented for this model class.
`get_anndata_manager`(adata[, required])	Retrieves the `AnnDataManager` for a given AnnData object.
`get_from_registry`(adata, registry_key)	Returns the object in AnnData associated with the key in the data registry.
`get_latent_representation`([adata, indices, ...])	Get the latent representation for each cell.
`get_normalized_expression`(args, *kwargs)	Not implemented for this model class.
`get_setup_arg`(setup_arg)	Returns the string provided to setup of a specific setup_arg.
`get_state_registry`(registry_key)	Returns the state registry for the AnnDataField registered with this instance.
`get_var_names`([legacy_mudata_format])	Variable names of input data.
`load`(dir_path[, adata, accelerator, device, ...])	Instantiate a model from the saved output.
`load_registry`(dir_path[, prefix])	Return the full registry saved with the model.
`register_manager`(adata_manager)	Registers an `AnnDataManager` instance with this model class.
`save`(dir_path[, prefix, overwrite, ...])	Save the state of the model.
`setup_annbatch`([collection_path, paths, ...])	Set up a model from on-disk AnnData files via the annbatch streaming loader.
`setup_anndata`(adata[, layer, batch_key, ...])	Set up AnnData object for training.
`to_device`(device)	Move the model to a specific device.
`train`([max_epochs, accelerator, devices, ...])	Train the model.
`transfer_fields`(adata, **kwargs)	Transfer fields from a model to an AnnData object.
`update_setup_method_args`(setup_method_args)	Update setup method args.
`view_anndata_setup`([adata, ...])	Print summary of the setup for the initial AnnData or a given AnnData object.
`view_registry`([hide_state_registries])	Prints summary of the registry.
`view_setup_args`(dir_path[, prefix])	Print args used to setup a saved model.
`view_setup_method_args`()	Prints setup kwargs used to produce a given registry.

Attributes#

mlxSCVI.adata#: Data attached to model instance.

mlxSCVI.adata_manager#: Manager instance associated with self.adata.

mlxSCVI.device#

Get the current device.

MLX automatically handles device placement.

Returns:: str Device identifier.

mlxSCVI.get_normalized_function_name#: What the get normalized functions name is

mlxSCVI.history#: Returns computed metrics during training.

mlxSCVI.is_trained#: Whether the model has been trained.

mlxSCVI.registry#: Data attached to model instance.

mlxSCVI.run_id#: Returns the run id of the model. Used in MLFlow

mlxSCVI.run_name#: Returns the run name of the model. Used in MLFlow

mlxSCVI.summary_string#: Summary string of the model.

mlxSCVI.test_indices#: Observations that are in test set.

mlxSCVI.train_indices#: Observations that are in train set.

mlxSCVI.validation_indices#: Observations that are in validation set.

Methods#

classmethod mlxSCVI.convert_legacy_save(dir_path, output_dir_path, overwrite=False, prefix=None, **save_kwargs)[source]#

Converts a legacy saved model (<v0.15.0) to the updated save format.

Parameters:

dir_path str: Path to the directory where the legacy model is saved.
output_dir_path str: Path to save converted save files.
overwrite bool (default: False): Overwrite existing data or not. If False and directory already exists at output_dir_path, an error will be raised.
prefix str | None (default: None): Prefix of saved file names.
**save_kwargs: Keyword arguments passed into save().

Return type:

None

mlxSCVI.data_registry(registry_key)[source]#

Returns the object in AnnData associated with the key in the data registry.

Parameters:

registry_key str: key of an object to get from self.data_registry

Return type:

ndarray | DataFrame

Returns:

The requested data.

mlxSCVI.deregister_manager(adata=None)[source]#

Deregisters the AnnDataManager instance associated with adata.

If adata is None, deregisters all AnnDataManager instances in both the class and instance-specific manager stores, except for the one associated with this model instance.

mlxSCVI.differential_abundance(*args, **kwargs)[source]#

Not implemented for this model class.

Available in models that inherit from VAEMixin.

Raises:: NotImplementedError –

mlxSCVI.get_anndata_manager(adata, required=False)[source]#

Retrieves the AnnDataManager for a given AnnData object.

Requires self.id has been set. Checks for an AnnDataManager specific to this model instance.

Parameters:

adata AnnData | MuData: AnnData object to find a manager instance for.
required bool (default: False): If True, errors on missing manager. Otherwise, returns None when manager is missing.

Return type:

AnnDataManager | None

mlxSCVI.get_from_registry(adata, registry_key)[source]#

Returns the object in AnnData associated with the key in the data registry.

AnnData object should be registered with the model prior to calling this function via the self._validate_anndata method.

Parameters:

registry_key str: key of object to get from the data registry.
adata AnnData | MuData: AnnData to pull data from.

Return type:

ndarray

Returns:

The requested data as a NumPy array.

mlxSCVI.get_latent_representation(adata=None, indices=None, give_mean=True, n_samples=1, batch_size=None)[source]#

Get the latent representation for each cell.

Parameters:

adata anndata.AnnData | None (default: None): AnnData object with the same structure as the initial AnnData object. If None, defaults to the AnnData object used when initializing the model.
indices Sequence[int] | None (default: None): Indices of cells to use from adata. If None, all cells are used.
give_mean bool (default: True): Whether to return the mean of the posterior distribution or a sample.
n_samples int (default: 1): Number of samples to use for computing the latent representation.
batch_size int | None (default: None): Minibatch size for data loading into the model.

Return type:

ndarray

Returns:

-latent_representation (ndarray) Low-dimensional representation for each cell

mlxSCVI.get_normalized_expression(*args, **kwargs)[source]#

Not implemented for this model class.

Available in RNA models that inherit from RNASeqMixin.

Raises:: NotImplementedError –

mlxSCVI.get_setup_arg(setup_arg)[source]#

Returns the string provided to setup of a specific setup_arg.

Return type:: attrdict

mlxSCVI.get_state_registry(registry_key)[source]#

Returns the state registry for the AnnDataField registered with this instance.

Return type:: attrdict

mlxSCVI.get_var_names(legacy_mudata_format=False)[source]#

Variable names of input data.

Return type:: dict

classmethod mlxSCVI.load(dir_path, adata=None, accelerator='auto', device='auto', prefix=None, backup_url=None, datamodule=None, allowed_classes_names_list=None)[source]#

Instantiate a model from the saved output.

Parameters:

dir_path str: Path to saved outputs.
adata AnnData | MuData | None (default: None): AnnData organized in the same way as data used to train model. It is not necessary to run setup_anndata, as AnnData is validated against the saved scvi setup dictionary. If None, will check for and load anndata saved with the model. If False, will load the model without AnnData.
accelerator str (default: 'auto'): Supports passing different accelerator types ("cpu", "gpu", "tpu", "ipu", "hpu", "mps, "auto") as well as custom accelerator instances.
device int | str (default: 'auto'): The device to use. Can be set to a non-negative index (int or str) or "auto" for automatic selection based on the chosen accelerator. If set to "auto" and accelerator is not determined to be "cpu", then device will be set to the first available device.
prefix str | None (default: None): Prefix of saved file names.
backup_url str | None (default: None): URL to retrieve saved outputs from if not present on disk.
datamodule LightningDataModule | None (default: None): EXPERIMENTAL A LightningDataModule instance to use for training in place of the default DataSplitter. Can only be passed in if the model was not initialized with AnnData.
allowed_classes_names_list list[str] | None (default: None): list of allowed classes names to be loaded (besides the original class name)

Returns:

Model with loaded state dictionaries.

Examples

>>> model = ModelClass.load(save_path, adata)
>>> model.get_....

static mlxSCVI.load_registry(dir_path, prefix=None)[source]#

Return the full registry saved with the model.

Parameters:

dir_path str: Path to saved outputs.
prefix str | None (default: None): Prefix of saved file names.

Return type:

dict

Returns:

The full registry saved with the model

classmethod mlxSCVI.register_manager(adata_manager)[source]#

Registers an AnnDataManager instance with this model class.

Stores the AnnDataManager reference in a class-specific manager store. Intended for use in the setup_anndata() class method followed up by retrieval of the AnnDataManager via the _get_most_recent_anndata_manager() method in the model init method.

Notes

Subsequent calls to this method with an AnnDataManager instance referring to the same underlying AnnData object will overwrite the reference to previous AnnDataManager.

mlxSCVI.save(dir_path, prefix=None, overwrite=False, save_anndata=False, save_kwargs=None, legacy_mudata_format=False, datamodule=None, **anndata_write_kwargs)[source]#

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 str: Path to a directory.
prefix str | None (default: None): Prefix to prepend to saved file names.
overwrite bool (default: False): Overwrite existing data or not. If False and directory already exists at dir_path, an error will be raised.
save_anndata bool (default: False): If True, also saves the anndata
save_kwargs dict | None (default: None): Keyword arguments passed into save().
legacy_mudata_format bool (default: False): If True, saves the model var_names in the legacy format if the model was trained with a MuData object. The legacy format is a flat array with variable names across all modalities concatenated, while the new format is a dictionary with keys corresponding to the modality names and values corresponding to the variable names for each modality.
datamodule LightningDataModule | None (default: None): EXPERIMENTAL A LightningDataModule instance to use for training in place of the default DataSplitter. Can only be passed in if the model was not initialized with AnnData.
anndata_write_kwargs: Kwargs for write()

classmethod mlxSCVI.setup_annbatch(collection_path=None, paths=None, batch_key=None, labels_key=None, sample_key=None, unlabeled_category='Unknown', layer=None, categorical_covariate_keys=None, continuous_covariate_keys=None, rebuild=True, batch_size=4096, chunk_size=256, preload_nchunks=32, preload_to_gpu=True, dataset_size='20GB', shuffle=False, var_subset=None, merge=None, adatas=None, use_class_sampler=False, class_sampler_key=None, class_weights=None)[source]#

Set up a model from on-disk AnnData files via the annbatch streaming loader.

Builds (or reuses) a zarr-backed DatasetCollection from the supplied h5ad file paths, then wraps it in a AnnbatchDataModule ready for training.

Parameters:

collection_path str | None (default: None): Directory where the zarr collection is written (or already exists). If None (default), a path is auto-generated as "./{ModelName}_annbatch.zarr" in the current working directory.
paths list[str] | None (default: None): Paths to h5ad files that make up the training dataset. If None, an existing collection at collection_path is opened without rebuilding — useful when the zarr store was created by a previous call. Must be provided when the store does not exist yet.
batch_key str | None (default: None): Column in obs to use as the batch variable.
labels_key str | None (default: None): Column in obs to use as the cell-type / label variable.
sample_key str | None (default: None): Column in obs to use as the sample variable. Used by models like MrVI.
unlabeled_category str (default: 'Unknown'): Value used to mark unlabeled cells in labels_key. Required by semi-supervised models such as SCANVI. Defaults to "Unknown".
layer str | None (default: None): Layer in the h5ad files to use as the count matrix. None uses adata.X (falling back to adata.raw.X when a raw slot exists).
categorical_covariate_keys list[str] | None (default: None): Additional categorical covariate columns in obs.
continuous_covariate_keys list[str] | None (default: None): Additional continuous covariate columns in obs.
rebuild bool (default: True): If True, always rebuild the zarr collection even if it already exists on disk. If False (default), reuse an existing collection and skip the (potentially expensive) add_adatas step.
batch_size int (default: 4096): Number of cells per batch yielded by the Loader.
chunk_size int (default: 256): Number of cells loaded from disk contiguously per read.
preload_nchunks int (default: 32): Number of chunks to preload and shuffle in memory.
preload_to_gpu bool (default: True): Whether the loader should move data to GPU before yielding.
dataset_size int | str (default: '20GB'): Number of observations to load into memory for shuffling / pre-processing when building the collection, or annbatch’s human-readable size strings, e.g. "20GB".
shuffle bool (default: False): Whether to pre-shuffle cells when building the collection.
var_subset list[str] | None (default: None): Optional list of gene names to restrict the collection to (passed as var_subset to add_adatas).
merge Literal['same', 'unique', 'first', 'only'] | None (default: None): How annbatch should merge var metadata across inputs. Passed through to annbatch.DatasetCollection.add_adatas().

Returns:

AnnbatchDataModule A configured datamodule whose registry can be passed directly to the model constructor, e.g. model = SCVI(registry=dm.registry).

Notes

After training, saving and reloading the model requires reconstructing the datamodule manually and passing it to load() via the datamodule argument — the same pattern used by all custom datamodules.

classmethod mlxSCVI.setup_anndata(adata, layer=None, batch_key=None, labels_key=None, categorical_covariate_keys=None, continuous_covariate_keys=None, **kwargs)[source]#

Set up AnnData object for training.

Parameters:

adata AnnData: AnnData object.
layer str | None (default: None): If not None, use this layer instead of X for training.
batch_key str | None (default: None): If not None, use the obs column specified by this key as batch information.
labels_key str | None (default: None): If not None, use the obs column specified by this key as labels.
categorical_covariate_keys list[str] | None (default: None): Keys in adata.obs for additional categorical covariates.
continuous_covariate_keys list[str] | None (default: None): Keys in adata.obs for additional continuous covariates.

mlxSCVI.to_device(device)[source]#

Move the model to a specific device.

MLX automatically handles device placement, so this is a no-op.

Parameters:

device: Target device.

mlxSCVI.train(max_epochs=None, accelerator='auto', devices='auto', train_size=None, validation_size=None, shuffle_set_split=True, batch_size=128, datasplitter_kwargs=None, plan_kwargs=None, **trainer_kwargs)[source]#

Train the model.

Parameters:

max_epochs int | None (default: None): Number of passes through the dataset. If None, defaults to np.min([round((20000 / n_cells) * 400), 400])
accelerator str (default: 'auto'): Accelerator type. MLX automatically selects the best available device.
devices int | list[int] | str (default: 'auto'): Device selection. MLX automatically selects the best available device.
train_size float | None (default: None): Training set size in the range [0.0, 1.0].
validation_size float | None (default: None): Validation set size. If None, defaults to 1 - train_size.
shuffle_set_split bool (default: True): Whether to shuffle indices before splitting.
batch_size int (default: 128): Minibatch size to use during training.
datasplitter_kwargs dict | None (default: None): Additional keyword arguments for DataSplitter.
plan_kwargs dict | None (default: None): Keyword arguments for the training plan.
**trainer_kwargs: Additional keyword arguments for training.

Returns:

self The trained model instance.

mlxSCVI.transfer_fields(adata, **kwargs)[source]#

Transfer fields from a model to an AnnData object.

Return type:: AnnData

mlxSCVI.update_setup_method_args(setup_method_args)[source]#

Update setup method args.

Parameters:

setup_method_args dict: This is a bit of a misnomer, this is a dict representing kwargs of the setup method that will be used to update the existing values in the registry of this instance.

mlxSCVI.view_anndata_setup(adata=None, hide_state_registries=False)[source]#

Print summary of the setup for the initial AnnData or a given AnnData object.

Parameters:

adata AnnData | MuData | None (default: None): AnnData object setup with setup_anndata or transfer_fields().
hide_state_registries bool (default: False): If True, prints a shortened summary without details of each state registry.

Return type:

None

mlxSCVI.view_registry(hide_state_registries=False)[source]#

Prints summary of the registry.

Parameters:

hide_state_registries bool (default: False): If True, prints a shortened summary without details of each state registry.

Return type:

None

static mlxSCVI.view_setup_args(dir_path, prefix=None)[source]#

Print args used to setup a saved model.

Parameters:

dir_path str: Path to saved outputs.
prefix str | None (default: None): Prefix of saved file names.

Return type:

None

mlxSCVI.view_setup_method_args()[source]#

Prints setup kwargs used to produce a given registry.

Parameters:

registry: Registry produced by an AnnDataManager.

Return type:

None

Core

Ecosystem

scvi.model.mlxSCVI#

Attributes table#

Methods table#

Attributes#

Methods#