Posterior predictive moments and total variance

Your task

Function signature

posterior_predictive_moments(
    mean: torch.Tensor,
    inverse_dispersion: torch.Tensor,
    sample_weights: torch.Tensor,
    sample_mask: torch.Tensor,
    gene_mask: torch.Tensor,
    zero_inflation_logits: torch.Tensor | None = None,
    likelihood: str = 'zinb',
) -> dict[str, torch.Tensor]

Worked example

Analytic two component predictive mixture

import torch

mean = torch.tensor([[[1.0]], [[3.0]]], dtype=torch.float64, requires_grad=True)
inverse_dispersion = torch.tensor([1.0], dtype=torch.float64, requires_grad=True)
zero_inflation_logits = torch.tensor([[[0.0]], [[0.0]]], dtype=torch.float64, requires_grad=True)
sample_weights = torch.tensor([[1.0], [1.0]], dtype=torch.float64, requires_grad=True)
sample_mask = torch.tensor([[True], [True]], dtype=torch.bool)
gene_mask = torch.tensor([[True]], dtype=torch.bool)
likelihood = 'zinb'

result = posterior_predictive_moments(mean=mean, inverse_dispersion=inverse_dispersion, zero_inflation_logits=zero_inflation_logits, sample_weights=sample_weights, sample_mask=sample_mask, gene_mask=gene_mask, likelihood=likelihood)

Expected return

{
    'mean': torch.tensor([[1.0]], dtype=torch.float64),
    'variance': torch.tensor([[5.0]], dtype=torch.float64),
    'likelihood_variance': torch.tensor([[4.75]], dtype=torch.float64),
    'latent_variance': torch.tensor([[0.25]], dtype=torch.float64),
    'zero_probability': torch.tensor([[0.6875]], dtype=torch.float64),
    'effective_sample_size': torch.tensor([2.0], dtype=torch.float64),
    'normalized_weights': torch.tensor([[0.5], [0.5]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Marginalize supplied NB or ZINB decoder distributions across posterior samples, applying the law of total variance. mean and zero_inflation_logits are [S,B,G], with posterior sample, cell, gene axes. inverse_dispersion is [G] or [S,B,G]. sample_weights and bool sample_mask are [S,B]; bool gene_mask is [B,G].
• All dimensions are \ge 1. True masks include entries. \mathrm{Active}=\mathrm{sample\_mask}[:,:,\mathrm{None}] \land \mathrm{gene\_mask}[\mathrm{None},:,:]. Means are nonnegative and \theta strictly positive on active entries. Active sample weights are nonnegative.
• All inputs are finite; finite invalid sentinels in masked positions are ignored. All tested products/ratios are representable. Set likelihood to 'zinb' (default) or 'nb'.
• In ZINB mode \pi=\operatorname{sigmoid}(\mathrm{zero\_inflation\_logits}) and q=1-\pi, computed stably.
• In NB mode q=1, \pi=0 and supplied logits are ignored. Each active conditional distribution has m_{s}=q_{s}\cdot \mu_{s}, v_s=q_s\left(\mu_s+\frac{\mu_s^2}{\theta_s}\right)+q_s\pi_s\mu_s^2, p_{0,s}=\pi_s+q_s\exp\left[-\theta_s\log\left(1+\frac{\mu_s}{\theta_s}\right)\right].
• Here mean inputs are NB component count means, not normalized gene compositions. At \mu=0 the count distribution is exactly zero-valued, so its mean/variance are zero and \mathrm{p0}=1.
• For each cell, w_{s}=\operatorname{where}(\mathrm{sample\_mask},\mathrm{sample\_weights},0), W=\sum_{s} w_{s} and \alpha_{s}=w_{s}/W. Return \mathrm{normalized\_weights}=\alpha. If W=0, \alpha, all predictive moments and effective_sample_size are zero.
• Otherwise \operatorname{mean}=\sum_{s} \alpha_{s}\cdot m_{s}; \mathrm{likelihood\_variance}=\sum_{s} \alpha_{s}\cdot v_{s}; \mathrm{latent\_variance}=\sum_{s} \alpha_{s}\cdot (m_{s}-\operatorname{mean})^{2}; \mathrm{variance}=\mathrm{likelihood\_variance}+\mathrm{latent\_variance}; \mathrm{zero\_probability}=\sum_{s} \alpha_{s}\cdot \mathrm{p0\_s}; \texttt{effective\_sample\_size}=\frac{W^2}{\sum_s w_s^2}.
• The variance is that of the finite weighted predictive mixture, not an unbiased S-1 estimate. A single active sample has zero latent variance and \mathrm{ESS}=1. Masked genes return zero moments/probability; ESS and normalized weights remain per-cell sample diagnostics independent of gene_mask.
• Reject an unknown likelihood, negative active means, nonpositive active \theta, or negative active weights with ValueError. Other stated conditions are preconditions. No sampling, model loading, full encoder/decoder execution, gene-gene covariance, posterior interval estimation, or calibration on experimental data occurs here.
• The supplied decoder/posterior samples already encode the desired library and batch-conditioning choices. Compose sigmoid and variance arithmetic from raw torch; do not use distributions or built-in variance/softmax functions.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def posterior_predictive_moments(
    mean: torch.Tensor,  # (S, B, G)
    inverse_dispersion: torch.Tensor,  # (G,) or (S, B, G)
    sample_weights: torch.Tensor,  # (S, B)
    sample_mask: torch.Tensor,  # (S, B)
    gene_mask: torch.Tensor,  # (B, G)
    zero_inflation_logits: torch.Tensor | None = None,  # (S, B, G)
    likelihood: str = 'zinb',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    mean: (B, G).
    variance: (B, G). (B,G) total predictive variance, not only between-sample variance.
    likelihood_variance: (B, G).
    latent_variance: (B, G).
    zero_probability: (B, G).
    effective_sample_size: (B,).
    normalized_weights: (S, B)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement posterior_predictive_moments")

Sources and further reading

• www.nature.com
• scverse/scvi-tools · _rnamixin.py
• scverse/scvi-tools · _negative_binomial.py