scxpand.autoencoders.ae_modules

scxpand.autoencoders.ae_modules#

Functions

apply_inverse_transforms_to_mean(raw_mu, ...)

Apply inverse transforms to decoder mean output to match row-normalized target scale.

dropout_activation(x)

Activation function for dropout probability parameter.

theta_activation(x)

Activation function for theta dispersion parameter in negative binomial regression.

Classes

BaseAutoencoder(data_format[, latent_dim, ...])

Base class for autoencoder models.

CategoricalHead([latent_dim, hidden_dims, ...])

ClassificationHead([latent_dim, ...])

Decoder([latent_dim, hidden_dims, ...])

Encoder(in_dim[, latent_dim, hidden_dims, ...])

class scxpand.autoencoders.ae_modules.BaseAutoencoder(data_format, latent_dim=32, encoder_hidden_dims=(64,), decoder_hidden_dims=(64,), classifier_hidden_dims=(16,), dropout_rate=0.3)#

Base class for autoencoder models.

__init__(data_format, latent_dim=32, encoder_hidden_dims=(64,), decoder_hidden_dims=(64,), classifier_hidden_dims=(16,), dropout_rate=0.3)#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

decode(latent_vec)#

Decode latent representation to output parameters.

Parameters:

latent_vec (Tensor) – Latent vector, shape [batch_size, latent_dim]

Returns:

  • mu: Mean parameter after inverse transforms to match row-normalized target scale, shape [batch_size, n_genes], non-negative

  • pi: Zero-inflation parameter, shape [batch_size, n_genes] (None if not needed, in [0, 1])

  • theta: Dispersion parameter, shape [batch_size, n_genes] (None if not needed, positive values)

Return type:

DecoderOutput containing

Note

The decoder handles activation and inverse transforms internally.

encode(x)#

Encode input data to latent representation.

Parameters:

x (Tensor) – Input data tensor, shape [batch_size, n_genes]

Returns:

Latent representation, shape [batch_size, latent_dim]

Return type:

latent_vec

forward(x)#

Forward pass through the autoencoder model.

Return type:

ModelOutput

class scxpand.autoencoders.ae_modules.CategoricalHead(latent_dim=32, hidden_dims=(16,), dropout_rate=0.3, n_classes=2, feature_name=None)#
__init__(latent_dim=32, hidden_dims=(16,), dropout_rate=0.3, n_classes=2, feature_name=None)#

Initialize a classification head for a categorical feature.

Parameters:

  • latent_dim (int (default: 32)) – Input dimension (latent vector size)

  • hidden_dims (tuple[int, ...] (default: (16,))) – Dimensions of hidden layers

  • dropout_rate (float (default: 0.3)) – Dropout probability

  • n_classes (int (default: 2)) – Number of classes for this categorical feature

  • feature_name (str | None (default: None)) – Name of the feature this head predicts (for debugging)

forward(latent_vec)#

Forward pass through the categorical head.

Parameters:

latent_vec (Tensor) – Latent vector from encoder

Return type:

Tensor

Returns:

Categorical logits with shape [batch_size, n_classes]

class scxpand.autoencoders.ae_modules.ClassificationHead(latent_dim=32, hidden_dims=(16,), dropout_rate=0.3)#
__init__(latent_dim=32, hidden_dims=(16,), dropout_rate=0.3)#

Initialize classification head network.

Parameters:

  • latent_dim (int (default: 32)) – Input dimension (latent vector size)

  • hidden_dims (tuple[int, ...] (default: (16,))) – Dimensions of hidden layers

  • dropout_rate (float (default: 0.3)) – Dropout probability

forward(latent_vec)#

Forward pass through the classification head.

Parameters:

latent_vec (Tensor) – Latent vector from encoder

Return type:

Tensor

Returns:

Classification logit

class scxpand.autoencoders.ae_modules.Decoder(latent_dim=32, hidden_dims=(64,), dropout_rate=0.3, n_genes=None, needs_pi=True, needs_theta=True, data_format=None)#
__init__(latent_dim=32, hidden_dims=(64,), dropout_rate=0.3, n_genes=None, needs_pi=True, needs_theta=True, data_format=None)#

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(latent_vec)#

Forward pass through the decoder.

Parameters:

latent_vec (Tensor) – Latent vector from encoder

Returns:

  • mu: Mean parameter after inverse transforms to match row-normalized target scale

  • pi: Dropout probability (None if not needed), in [0, 1]

  • theta: Dispersion parameter (None if not needed)

Return type:

DecoderOutput containing

class scxpand.autoencoders.ae_modules.Encoder(in_dim, latent_dim=32, hidden_dims=(64,), dropout_rate=0.3)#
__init__(in_dim, latent_dim=32, hidden_dims=(64,), dropout_rate=0.3)#

Initialize encoder network.

Parameters:

  • in_dim (int) – Input dimension

  • latent_dim (int (default: 32)) – Dimension of the latent space

  • hidden_dims (tuple[int, ...] (default: (64,))) – Dimensions of hidden layers

  • dropout_rate (float (default: 0.3)) – Dropout probability

forward(x)#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

scxpand.autoencoders.ae_modules.apply_inverse_transforms_to_mean(raw_mu, data_format)#

Apply inverse transforms to decoder mean output to match row-normalized target scale.

Applies inverse transforms in reverse order: 1. Inverse z-score normalization 2. Inverse log transform (if used)

Parameters:

  • raw_mu (Tensor) – Decoder mean output (raw logits from the model)

  • data_format (DataFormat) – Data format containing normalization parameters

Return type:

Tensor

Returns:

Transformed mean that matches row-normalized target scale

scxpand.autoencoders.ae_modules.dropout_activation(x)#

Activation function for dropout probability parameter.

Range: (0, 1) as appropriate for probabilities.

Return type:

Tensor

scxpand.autoencoders.ae_modules.theta_activation(x)#

Activation function for theta dispersion parameter in negative binomial regression.

In negative binomial regression, theta (r) is the dispersion parameter that must be positive.

Return type:

Tensor

Contents