spatialvi.module.nicheVAE

spatialvi.module.nicheVAE#

class spatialvi.module.nicheVAE(n_input, n_output_niche, n_batch=0, n_labels=0, n_hidden=128, n_latent=10, n_layers=1, n_layers_niche=1, n_layers_compo=1, n_hidden_niche=128, n_hidden_compo=128, n_continuous_cov=0, n_cats_per_cov=None, dropout_rate=0.1, dispersion='gene', log_variational=True, gene_likelihood='poisson', latent_distribution='normal', niche_likelihood='gaussian', cell_rec_weight=1.0, latent_kl_weight=1.0, spatial_weight=10, prior_mixture=False, prior_mixture_k=20, semisupervised=True, linear_classifier=True, inpute_covariates_niche_decoder=True, encode_covariates=False, deeply_inject_covariates=True, batch_representation='one-hot', use_batch_norm='none', use_layer_norm='both', use_size_factor_key=False, use_observed_lib_size=True, library_log_means=None, library_log_vars=None, batch_embedding_kwargs=None, extra_decoder_kwargs=None, extra_encoder_kwargs=None, **vae_kwargs)[source]#

Variational auto-encoder with niche decoders [].

Parameters:

  • n_input (int) – Number of input features.

  • n_batch (int) – Number of batches. If 0, no batch correction is performed.

  • n_labels (int) – Number of labels.

  • n_hidden (int) – Number of nodes per hidden layer. Passed into Encoder and DecoderSCVI.

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

  • n_layers (int) – Number of hidden layers. Passed into Encoder and DecoderSCVI.

  • n_layers_niche (int) – Number of hidden layers in the niche state decoder.

  • n_layers_compo (int) – Number of hidden layers in the composition decoder.

  • n_hidden_niche (int) – Number of nodes per hidden layer in the niche state decoder.

  • n_hidden_compo (int) – Number of nodes per hidden layer in the composition decoder.

  • n_continuous_cov (int) – Number of continuous covariates.

  • n_cats_per_cov (list[int] | None) – A list of integers containing the number of categories for each categorical covariate.

  • dropout_rate (float) – Dropout rate. Passed into Encoder but not DecoderSCVI.

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

    Flexibility of the dispersion parameter when gene_likelihood is either "nb" or "zinb". One of the following:

    • "gene": parameter is constant per gene across cells.

    • "gene-batch": parameter is constant per gene per batch.

    • "gene-label": parameter is constant per gene per label.

    • "gene-cell": parameter is constant per gene per cell.

  • log_variational (bool) – If True, use log1p() on input data before encoding for numerical stability (not normalization).

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

    Distribution to use for reconstruction in the generative process. One of the following:

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

    Distribution to use for the latent space. One of the following:

    • "normal": isotropic normal.

    • "ln": logistic normal with normal params N(0, 1).

  • niche_likelihood (Literal['poisson', 'gaussian']) –

    Distribution to use for the niche state. One of the following:

    • "poisson": Poisson.

    • "gaussian": Normal.

    Default is "gaussian" and Poisson should be used if the niche state is count data.

  • cell_rec_weight (float) – Weight of the cell reconstruction loss.

  • latent_kl_weight (float) – Weight of the latent KL divergence.

  • spatial_weight (float) – Weight of the spatial losses

  • prior_mixture (bool) – If True, use a mixture of Gaussians for the latent space. Else, use unimodal Gaussian.

  • prior_mixture_k (int) – Number of components in the Gaussian mixture.

  • semisupervised (bool) – If True, use a classifier to predict cell type labels from the latent space.

  • linear_classifier (bool) – If True, use a linear classifier. Else, use a neural network.

  • inpute_covariates_niche_decoder (bool) – If True, covariates are concatenated to the input of the niche state decoder.

  • encode_covariates (bool) – If True, covariates are concatenated to gene expression prior to passing through the encoder(s). Else, only gene expression is used.

  • deeply_inject_covariates (bool) – If True and n_layers > 1, covariates are concatenated to the outputs of hidden layers in the encoder(s) (if encoder_covariates is True) and the decoder prior to passing through the next layer.

  • batch_representation (Literal['one-hot', 'embedding']) –

    EXPERIMENTAL Method for encoding batch information. One of the following:

    • "one-hot": represent batches with one-hot encodings.

    • "embedding": represent batches with continuously-valued embeddings using Embedding.

    Note that batch representations are only passed into the encoder(s) if encode_covariates is True.

  • use_batch_norm (Literal['encoder', 'decoder', 'none', 'both']) –

    Specifies where to use BatchNorm1d in the model. One of the following:

    • "none": don’t use batch norm in either encoder(s) or decoder.

    • "encoder": use batch norm only in the encoder(s).

    • "decoder": use batch norm only in the decoder.

    • "both": use batch norm in both encoder(s) and decoder.

    Note: if use_layer_norm is also specified, both will be applied (first BatchNorm1d, then LayerNorm).

  • use_layer_norm (Literal['encoder', 'decoder', 'none', 'both']) –

    Specifies where to use LayerNorm in the model. One of the following:

    • "none": don’t use layer norm in either encoder(s) or decoder.

    • "encoder": use layer norm only in the encoder(s).

    • "decoder": use layer norm only in the decoder.

    • "both": use layer norm in both encoder(s) and decoder.

    Note: if use_batch_norm is also specified, both will be applied (first BatchNorm1d, then LayerNorm).

  • use_size_factor_key (bool) – If True, use the obs column as defined by the size_factor_key parameter in the model’s setup_anndata method as the scaling factor in the mean of the conditional distribution. Takes priority over use_observed_lib_size.

  • use_observed_lib_size (bool) – If True, use the observed library size for RNA as the scaling factor in the mean of the conditional distribution.

  • library_log_means (ndarray | None) – ndarray of shape (1, n_batch) of means of the log library sizes that parameterize the prior on library size if use_size_factor_key is False and use_observed_lib_size is False.

  • library_log_vars (ndarray | None) – ndarray of shape (1, n_batch) of variances of the log library sizes that parameterize the prior on library size if use_size_factor_key is False and use_observed_lib_size is False.

  • extra_decoder_kwargs (dict | None) – Additional keyword arguments passed into DecoderSCVI.

  • batch_embedding_kwargs (dict | None) – Keyword arguments passed into Embedding if batch_representation is set to "embedding".

Notes

Lifecycle: argument batch_representation is experimental in v1.2.

__init__(n_input, n_output_niche, n_batch=0, n_labels=0, n_hidden=128, n_latent=10, n_layers=1, n_layers_niche=1, n_layers_compo=1, n_hidden_niche=128, n_hidden_compo=128, n_continuous_cov=0, n_cats_per_cov=None, dropout_rate=0.1, dispersion='gene', log_variational=True, gene_likelihood='poisson', latent_distribution='normal', niche_likelihood='gaussian', cell_rec_weight=1.0, latent_kl_weight=1.0, spatial_weight=10, prior_mixture=False, prior_mixture_k=20, semisupervised=True, linear_classifier=True, inpute_covariates_niche_decoder=True, encode_covariates=False, deeply_inject_covariates=True, batch_representation='one-hot', use_batch_norm='none', use_layer_norm='both', use_size_factor_key=False, use_observed_lib_size=True, library_log_means=None, library_log_vars=None, batch_embedding_kwargs=None, extra_decoder_kwargs=None, extra_encoder_kwargs=None, **vae_kwargs)[source]#

Methods

__init__(n_input, n_output_niche[, n_batch, ...])

add_embedding(key, embedding[, overwrite])

Add an embedding to the module.

add_module(name, module)

Add a child module to the current module.

apply(fn)

Apply fn recursively to every submodule (as returned by .children()) as well as self.

bfloat16()

Casts all floating point parameters and buffers to bfloat16 datatype.

buffers([recurse])

Return an iterator over module buffers.

children()

Return an iterator over immediate children modules.

compile(*args, **kwargs)

Compile this Module's forward using torch.compile().

compute_embedding(key, indices)

Forward pass for an embedding.

cpu()

Move all model parameters and buffers to the CPU.

cuda([device])

Move all model parameters and buffers to the GPU.

double()

Casts all floating point parameters and buffers to double datatype.

eval()

Set the module in evaluation mode.

extra_repr()

Return the extra representation of the module.

float()

Casts all floating point parameters and buffers to float datatype.

forward(tensors[, ...])

Forward pass through the network.

generative(z, library, batch_index[, ...])

Run the generative process.

get_buffer(target)

Return the buffer given by target if it exists, otherwise throw an error.

get_embedding(key)

Get an embedding from the module.

get_extra_state()

Return any extra state to include in the module's state_dict.

get_parameter(target)

Return the parameter given by target if it exists, otherwise throw an error.

get_submodule(target)

Return the submodule given by target if it exists, otherwise throw an error.

half()

Casts all floating point parameters and buffers to half datatype.

inference(*args, **kwargs)

Main inference call site.

init_embedding(key, num_embeddings[, ...])

Initialize an embedding in the module.

ipu([device])

Move all model parameters and buffers to the IPU.

load_state_dict(state_dict[, strict, assign])

Copy parameters and buffers from state_dict into this module and its descendants.

loss(tensors, inference_outputs, ...[, ...])

Compute the loss.

marginal_ll(tensors, n_mc_samples[, ...])

Compute the marginal log-likelihood of the data under the model.

modules([remove_duplicate])

Return an iterator over all modules in the network.

mtia([device])

Move all model parameters and buffers to the MTIA.

named_buffers([prefix, recurse, ...])

Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.

named_children()

Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.

named_modules([memo, prefix, remove_duplicate])

Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.

named_parameters([prefix, recurse, ...])

Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.

on_load(model, **kwargs)

Callback function run in load().

parameters([recurse])

Return an iterator over module parameters.

register_backward_hook(hook)

Register a backward hook on the module.

register_buffer(name, tensor[, persistent])

Add a buffer to the module.

register_forward_hook(hook, *[, prepend, ...])

Register a forward hook on the module.

register_forward_pre_hook(hook, *[, ...])

Register a forward pre-hook on the module.

register_full_backward_hook(hook[, prepend])

Register a backward hook on the module.

register_full_backward_pre_hook(hook[, prepend])

Register a backward pre-hook on the module.

register_load_state_dict_post_hook(hook)

Register a post-hook to be run after module's load_state_dict() is called.

register_load_state_dict_pre_hook(hook)

Register a pre-hook to be run before module's load_state_dict() is called.

register_module(name, module)

Alias for add_module().

register_parameter(name, param)

Add a parameter to the module.

register_state_dict_post_hook(hook)

Register a post-hook for the state_dict() method.

register_state_dict_pre_hook(hook)

Register a pre-hook for the state_dict() method.

remove_embedding(key)

Remove an embedding from the module.

requires_grad_([requires_grad])

Change if autograd should record operations on parameters in this module.

sample(tensors[, n_samples, ...])

Generate predictive samples from the posterior predictive distribution.

set_extra_state(state)

Set extra state contained in the loaded state_dict.

set_submodule(target, module[, strict])

Set the submodule given by target if it exists, otherwise throw an error.

share_memory()

See torch.Tensor.share_memory_().

state_dict(*args[, destination, prefix, ...])

Return a dictionary containing references to the whole state of the module.

to(*args, **kwargs)

Move and/or cast the parameters and buffers.

to_empty(*, device[, recurse])

Move the parameters and buffers to the specified device without copying storage.

train([mode])

Set the module in training mode.

type(dst_type)

Casts all parameters and buffers to dst_type.

xpu([device])

Move all model parameters and buffers to the XPU.

zero_grad([set_to_none])

Reset gradients of all model parameters.

Attributes

T_destination

call_super_init

device

dump_patches

embeddings_dict

Dictionary of embeddings.

minified_data_type

The type of minified data associated with this module, if applicable.

training