Posterior differential-expression evidence and expected FDP

Your task

Function signature

bayesian_differential_expression(
    expression_a: torch.Tensor,
    expression_b: torch.Tensor,
    pair_indices: torch.Tensor,
    pair_weights: torch.Tensor,
    pair_mask: torch.Tensor,
    gene_mask: torch.Tensor,
    prior_log_odds: torch.Tensor,
    mode: str = 'change',
    delta: float = 0.5,
    pseudocount: float = 1e-06,
    probability_epsilon: float = 1e-08,
    fdr_target: float = 0.05,
) -> dict[str, torch.Tensor]

Worked example

Select the longest acceptable prefix

import torch

expression_a = torch.tensor([[4.0, 4.0, 1.0, 1.0], [4.0, 4.0, 1.0, 1.0], [4.0, 1.0, 1.0, 1.0], [4.0, 1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
expression_b = torch.tensor([[1.0, 1.0, 1.0, 1.0]], dtype=torch.float64, requires_grad=True)
pair_indices = torch.tensor([[0, 0], [1, 0], [2, 0], [3, 0]], dtype=torch.int64)
pair_weights = torch.tensor([1.0, 1.0, 1.0, 1.0], dtype=torch.float64, requires_grad=True)
pair_mask = torch.tensor([True, True, True, True], dtype=torch.bool)
gene_mask = torch.tensor([True, True, True, True], dtype=torch.bool)
prior_log_odds = torch.tensor([0.0, 0.0, 0.0, 0.0], dtype=torch.float64, requires_grad=True)
mode = 'change'
delta = 0.5
pseudocount = 1e-06
probability_epsilon = 1e-08
fdr_target = 0.25

result = bayesian_differential_expression(expression_a=expression_a, expression_b=expression_b, pair_indices=pair_indices, pair_weights=pair_weights, pair_mask=pair_mask, gene_mask=gene_mask, prior_log_odds=prior_log_odds, mode=mode, delta=delta, pseudocount=pseudocount, probability_epsilon=probability_epsilon, fdr_target=fdr_target)

Expected return

{
    'pair_log2_fold_change': torch.tensor([[1.9999989179793958, 1.9999989179793958, 0.0, 0.0], [1.9999989179793958, 1.9999989179793958, 0.0, 0.0], [1.9999989179793958, 0.0, 0.0, 0.0], [1.9999989179793958, 0.0, 0.0, 0.0]], dtype=torch.float64),
    'mean_log2_fold_change': torch.tensor([1.9999989179793958, 0.9999994589896979, 0.0, 0.0], dtype=torch.float64),
    'variance_log2_fold_change': torch.tensor([0.0, 0.9999989179796884, 0.0, 0.0], dtype=torch.float64),
    'probability_de': torch.tensor([1.0, 0.5, 0.0, 0.0], dtype=torch.float64),
    'log_bayes_factor': torch.tensor([18.420680753952368, 0.0, -18.420680753952368, -18.420680753952368], dtype=torch.float64),
    'ordered_gene_indices': torch.tensor([0, 1, 2, 3], dtype=torch.int64),
    'prefix_expected_fdp': torch.tensor([0.0, 0.25, 0.5, 0.625], dtype=torch.float64),
    'selected_mask': torch.tensor([True, True, False, False], dtype=torch.bool),
    'normalized_pair_weights': torch.tensor([0.25, 0.25, 0.25, 0.25], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Summarize supplied posterior expression pairs and select genes by posterior expected false-discovery proportion. expression_a is [A,G] and expression_b [B,G], finite nonnegative NORMALIZED expression draws.
• pair_indices is int64 [P,2], addressing one A and one B draw per row; pair_mask is bool [P], pair_weights floating [P], gene_mask bool [G], and prior_log_odds floating [G]. A,B,P,G\ge 1. Supplied pairing may repeat samples and is never generated randomly.
• Active indices must be in range and active weights nonnegative; masked pairs may carry finite negative weights and -1 indices. Masked genes are excluded from all evidence and FDP denominators. Batch matching, population definition and sample generation happen upstream.
• keys: mode='change' (default) or 'vanilla'; \delta\ge 0 (default 0.5); \mathrm{pseudocount}>0 (default 10^{-6}); probability_epsilon in (0,0.5) (default 10^{-8}); fdr_target in [0,1] (default 0.05). All scalar values are finite.
• For each active pair and unmasked gene, d_{pg}=\log_2(a_{\mathrm{pair},g}+\mathrm{pseudocount})-\log_2(b_{\mathrm{pair},g}+\mathrm{pseudocount}), so positive d means higher A expression. Pair effects are zero outside masks.
• In change mode, the DE event is \left|d\right|>\delta, with equality explicitly not selected.
• In vanilla mode the event is a>b (strict); \delta is ignored and the event is directional, not a two-sided difference. Let w_{p}=\operatorname{where}(\mathrm{pair\_mask},\mathrm{pair\_weights},0), W=\sum_{p} w_{p}, \alpha=w/W.
• Return \mathrm{normalized\_pair\_weights}=\alpha, weighted mean of d, weighted population variance \sum(\alpha\cdot (d-\mathrm{mean\_d})^{2}), and \mathrm{probability\_de}=\sum_{p} \alpha_{p}\cdot 1[\mathrm{event\_pg}].
• Use these same pseudocount-adjusted effects for event tests and effect summaries; do not apply a hidden second pseudocount. Return \mathrm{log\_bayes\_factor}=\operatorname{log}(\mathrm{probability\_de}+\varepsilon)-\operatorname{log}(1-\mathrm{probability\_de}+\varepsilon)-\mathrm{prior\_log\_odds}.
• For stability near probability one, compute event mass u=\sum(w\cdot \mathrm{event}), complementary mass v=\sum_p w_p\,\mathbf1[\neg\mathrm{event}_p], and probability u/(u+v); the log evidence is equivalently \operatorname{log}(u+\varepsilon\cdot (u+v))-\operatorname{log}(v+\varepsilon\cdot (u+v))-\mathrm{prior\_log\_odds}.
• This avoids subtracting nearly equal numbers to estimate the tiny complementary probability. Equal prior odds are represented by a supplied all-zero prior_log_odds tensor. This is an EPSILON-STABILIZED approximation to posterior-minus-prior log odds; \varepsilon does not alter probability_de or the FDP calculations.
• Sort eligible genes by descending probability_de, breaking ties by their original gene index. ordered_gene_indices has length G with eligible indices first and -1 padding.
• For each eligible sorted prefix of length k, \mathrm{prefix\_expected\_fdp}[k-1]=\sum_{\mathrm{prefix}}(1-\mathrm{probability\_de})/k. Entries after the final eligible gene are zero. Select the longest eligible prefix whose expected \mathrm{FDP} \le \mathrm{fdr\_target}; selected_mask is returned in original gene order.
• This is a posterior expected-FDP rule, not Benjamini-Hochberg on p-values. With no positive total weight, \alpha and every aggregate diagnostic are zero, order is \mathrm{all} -1 and no gene is selected; pair_log2_fold_change still records any active supplied pairs. There is no evidence to select even when \mathrm{target}=1.
• If every gene is masked, all gene-level outputs are zero or -1 but \alpha still reports the available pair weights. All computation must be raw torch: implement deterministic ordering from comparisons/indexing rather than adding a general sort exception.
• Reject invalid values, negative expression samples, negative active weights and out-of-range active indices with ValueError. Other shape/dtype/finiteness conditions are preconditions.

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 bayesian_differential_expression(
    expression_a: torch.Tensor,  # (A, G)
    expression_b: torch.Tensor,  # (B, G)
    pair_indices: torch.Tensor,  # (P, 2)
    pair_weights: torch.Tensor,  # (P,)
    pair_mask: torch.Tensor,  # (P,)
    gene_mask: torch.Tensor,  # (G,)
    prior_log_odds: torch.Tensor,  # (G,)
    mode: str = 'change',
    delta: float = 0.5,
    pseudocount: float = 1e-06,
    probability_epsilon: float = 1e-08,
    fdr_target: float = 0.05,
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    pair_log2_fold_change: (P, G).
    mean_log2_fold_change: (G,).
    variance_log2_fold_change: (G,).
    probability_de: (G,).
    log_bayes_factor: (G,).
    ordered_gene_indices: (G,).
    prefix_expected_fdp: (G,).
    selected_mask: (G,).
    normalized_pair_weights: (P,)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement bayesian_differential_expression")

Sources and further reading

• www.nature.com
• scverse/scvi-tools · _differential.py
• scverse/scvi-tools · _de_core.py
• docs.scvi-tools.org