netmap.model.nbautoencoder.NegativeBinomialAutoencoder

class netmap.model.nbautoencoder.NegativeBinomialAutoencoder(*args: Any, **kwargs: Any)[source]

Bases: Module

Autoencoder with dual mu/theta decoder heads for NB-distributed scRNA-seq data.

The encoder compresses input gene expression into a low-dimensional latent space. Two separate decoder heads predict the NB parameters mu (mean) and theta (dispersion). Both encoder and decoders support configurable hidden layer widths and dropout.

Three mutually exclusive forward-mode flags let the model return a single output tensor, which is required by Captum attribution explainers:

  • forward_mu_only: return only mu.

  • forward_theta_only: return only theta.

  • latent_only: return the latent representation.

When all flags are False (default), forward returns (mu, theta).

Parameters:

  • input_dim (int) – Number of input genes.

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

  • dropout_rate (float) – Dropout probability applied after each hidden layer. Defaults to 0.1.

  • hidden_dims (list of int) – Width of each hidden layer in the encoder; the decoder mirrors this in reverse. Defaults to [64].

  • activation (str) – Activation function name — one of 'relu', 'leaky_relu', 'elu', 'sigmoid'. Defaults to 'relu'.

__init__(input_dim, latent_dim, dropout_rate=0.1, hidden_dims=[64], activation='relu')[source]
compute_loss(x, l1_lambda=0.1)[source]

Compute the NB loss for a batch.

Parameters:

  • x (torch.Tensor) – Input count matrix of shape (batch, input_dim).

  • l1_lambda (float) – Unused regularisation weight (reserved). Defaults to 0.1.

Returns:

Scalar NB negative log-likelihood loss.

Return type:

torch.Tensor

forward(x)[source]

Run a forward pass, returning output controlled by the mode flags.

When all mode flags are False, returns (mu, theta). Set forward_mu_only, forward_theta_only, or latent_only to True before wrapping the model in a Captum explainer so that the model returns a single tensor suitable for scalar-target attribution.

Parameters:

x (torch.Tensor) – Input gene expression tensor of shape (n_cells, input_dim).

Returns:

Depending on the active flag —

mu tensor, theta tensor, latent tensor, or (mu, theta) tuple.

Return type:

torch.Tensor or tuple

latent(x)[source]

Encode input to the latent representation.

Parameters:

x (torch.Tensor) – Input tensor of shape (n_cells, input_dim).

Returns:

Latent embedding of shape (n_cells, latent_dim).

Return type:

torch.Tensor