spatialvi.model.SpatialVAE

spatialvi.model.SpatialVAE#

class spatialvi.model.SpatialVAE(adata, n_hidden=128, n_latent=10, n_layers=1, dropout_rate=0.1, dispersion='gene', gene_likelihood='zinb', latent_distribution='normal', use_spatial=True, spatial_weight=1.0, **model_kwargs)[source]#

Bases: SpatialTrainingMixin, SpatialMixin, BaseSpatialModel

Spatial Variational Autoencoder for spatial transcriptomics.

This model combines gene expression modeling with spatial context using a VAE framework. It can be used for: - Spatially-aware dimensionality reduction - Batch effect correction in spatial data - Imputation with spatial smoothing

Parameters:

  • adata (AnnData) – AnnData object that has been registered via setup_anndata().

  • n_hidden (int) – Number of nodes per hidden layer.

  • n_latent (int) – Dimensionality of the latent space.

  • n_layers (int) – Number of hidden layers.

  • dropout_rate (float) – Dropout rate for neural networks.

  • dispersion (Literal['gene', 'gene-batch', 'gene-cell']) – Dispersion parameter for negative binomial. One of: - “gene”: single dispersion per gene - “gene-batch”: dispersion per gene and batch - “gene-cell”: dispersion per gene and cell

  • gene_likelihood (Literal['zinb', 'nb', 'poisson']) – Distribution for gene expression. One of: - “zinb”: Zero-inflated negative binomial - “nb”: Negative binomial - “poisson”: Poisson

  • latent_distribution (Literal['normal', 'ln']) – Distribution for latent space. One of: - “normal”: Normal distribution - “ln”: Logistic normal

  • use_spatial (bool) – Whether to use spatial information in encoding.

  • spatial_weight (float) – Weight for spatial regularization.

  • **model_kwargs – Additional keyword arguments for SpatialVAEModule.

  • adata (AnnData)

  • n_hidden (int)

  • n_latent (int)

  • n_layers (int)

  • dropout_rate (float)

  • dispersion (Literal['gene', 'gene-batch', 'gene-cell'])

  • gene_likelihood (Literal['zinb', 'nb', 'poisson'])

  • latent_distribution (Literal['normal', 'ln'])

  • use_spatial (bool)

  • spatial_weight (float)

Examples

>>> import spatialvi
>>> adata = spatialvi.data.synthetic_spatial()
>>> SpatialVAE.setup_anndata(adata, spatial_key="spatial")
>>> model = SpatialVAE(adata)
>>> model.train()
>>> latent = model.get_latent_representation()
__init__(adata, n_hidden=128, n_latent=10, n_layers=1, dropout_rate=0.1, dispersion='gene', gene_likelihood='zinb', latent_distribution='normal', use_spatial=True, spatial_weight=1.0, **model_kwargs)[source]#
Parameters:

  • adata (AnnData)

  • n_hidden (int)

  • n_latent (int)

  • n_layers (int)

  • dropout_rate (float)

  • dispersion (Literal['gene', 'gene-batch', 'gene-cell'])

  • gene_likelihood (Literal['zinb', 'nb', 'poisson'])

  • latent_distribution (Literal['normal', 'ln'])

  • use_spatial (bool)

  • spatial_weight (float)

Methods

__init__(adata[, n_hidden, n_latent, ...])

compute_niche_composition(adata, labels_key)

Compute neighborhood cell type composition.

compute_spatial_neighbors(adata[, ...])

Compute spatial nearest neighbors.

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.

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, ...])

Return the latent representation for each cell.

get_normalized_expression([adata, indices, ...])

Return normalized gene expression.

get_setup_arg(setup_arg)

Returns the string provided to setup of a specific setup_arg.

get_spatial_coordinates([adata, indices, ...])

Get spatial coordinates.

get_spatial_neighbors([adata, indices])

Get spatial neighbor indices and distances.

get_spatial_representation([adata, indices, ...])

Get spatial-aware latent representation.

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_anndata(adata[, layer, batch_key, ...])

Sets up the AnnData object for this model.

to_device(device)

Move the model to the 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

adata

Data attached to model instance.

adata_manager

Manager instance associated with self.adata.

device

The current device that the module's params are on.

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.

n_neighbors

Number of spatial neighbors.

registry

Data attached to model instance.

run_id

Returns the run id of the model.

run_name

Returns the run name of the model.

spatial_key

Key for spatial coordinates in obsm.

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.
get_normalized_expression(adata=None, indices=None, transform_batch=None, gene_list=None, library_size=1.0, n_samples=1, n_samples_overall=None, batch_size=None, return_mean=True, return_numpy=None)[source]#

Return normalized gene expression.

Parameters:

  • adata (AnnData | None) – AnnData object with equivalent structure to initial AnnData.

  • indices (Sequence[int] | None) – Indices of cells to use.

  • transform_batch (str | Sequence[str] | None) – Batch to condition on.

  • gene_list (Sequence[str] | None) – Subset of genes to use.

  • library_size (Union[float, Literal['latent']]) – Library size to use for normalization.

  • n_samples (int) – Number of samples to draw from posterior.

  • n_samples_overall (int | None) – Total number of samples across all cells.

  • batch_size (int | None) – Batch size for data loader.

  • return_mean (bool) – Whether to return the mean of samples.

  • return_numpy (bool | None) – Whether to return numpy array.

  • adata (AnnData | None)

  • indices (Sequence[int] | None)

  • transform_batch (str | Sequence[str] | None)

  • gene_list (Sequence[str] | None)

  • library_size (float | Literal['latent'])

  • n_samples (int)

  • n_samples_overall (int | None)

  • batch_size (int | None)

  • return_mean (bool)

  • return_numpy (bool | None)

Return type:

ndarray[tuple[Any, ...], dtype[TypeVar(_ScalarT, bound= generic)]] | dict[str, ndarray[tuple[Any, ...], dtype[TypeVar(_ScalarT, bound= generic)]]]

Returns:

Normalized expression array.

get_spatial_representation(adata=None, indices=None, batch_size=None)[source]#

Get spatial-aware latent representation.

This method returns a latent representation that incorporates spatial context from neighboring cells.

Parameters:
Return type:

ndarray[tuple[Any, ...], dtype[TypeVar(_ScalarT, bound= generic)]]

Returns:

Spatial latent representation array.

classmethod setup_anndata(adata, layer=None, batch_key=None, labels_key=None, size_factor_key=None, categorical_covariate_keys=None, continuous_covariate_keys=None, spatial_key='spatial', neighbor_index_key=None, neighbor_dist_key=None, **kwargs)[source]#

Sets up the AnnData object for this model.

A mapping will be created between data fields used by this model to their respective locations in adata. None of the data in adata are modified. Only adds fields to adata.

Parameters:

  • adata (AnnData) – AnnData object. Rows represent cells, columns represent features.

  • layer (str | None) – if not None, uses this as the key in adata.layers for raw count data.

  • batch_key (str | None) – key in adata.obs for batch information. Categories will automatically be converted into integer categories and saved to adata.obs['_scvi_batch']. If None, assigns the same batch to all the data.

  • labels_key (str | None) – key in adata.obs for label information. Categories will automatically be converted into integer categories and saved to adata.obs['_scvi_labels']. If None, assigns the same label to all the data.

  • size_factor_key (str | None) – key in adata.obs for size factor information. Instead of using library size as a size factor, the provided size factor column will be used as offset in the mean of the likelihood. Assumed to be on linear scale.

  • categorical_covariate_keys (list[str] | None) – keys in adata.obs that correspond to categorical data. These covariates can be added in addition to the batch covariate and are also treated as nuisance factors (i.e., the model tries to minimize their effects on the latent space). Thus, these should not be used for biologically-relevant factors that you do _not_ want to correct for.

  • continuous_covariate_keys (list[str] | None) – keys in adata.obs that correspond to continuous data. These covariates can be added in addition to the batch covariate and are also treated as nuisance factors (i.e., the model tries to minimize their effects on the latent space). Thus, these should not be used for biologically-relevant factors that you do _not_ want to correct for.

  • spatial_key (str) – Key in adata.obsm for spatial coordinates.

  • neighbor_index_key (str | None) – Key in adata.obsm for neighbor indices. If None, neighbors will not be used during training.

  • neighbor_dist_key (str | None) – Key in adata.obsm for neighbor distances.

  • adata (AnnData)

  • layer (str | None)

  • batch_key (str | None)

  • labels_key (str | None)

  • size_factor_key (str | None)

  • categorical_covariate_keys (list[str] | None)

  • continuous_covariate_keys (list[str] | None)

  • spatial_key (str)

  • neighbor_index_key (str | None)

  • neighbor_dist_key (str | None)

Return type:

None

Contents