mixle.task.inverse module

learn_inverse – amortized posteriors q(theta | y) for a simulator, with calibration receipts.

Simulation-based inference done the mixle way: given a forward simulator g: theta -> y (a bare Python callable – no mixle.task.imagine/M2 program required) and a prior p(theta) (any fitted mixle Model), learn_inverse() trains a torch CONDITIONAL density student q(theta | y) on simulated (theta, y) pairs, then ships it wrapped as an Posterior (M0’s type) so downstream condition/do composition and B7 treat a learned inverse exactly like an exactly-conditioned one – except its .receipt carries an explicit amortization warning plus a pointer to InverseReceipts, because a trained student is an APPROXIMATION and the whole point of this module is to ship the numbers that say whether to trust it.

Convention (matches build_mdn/build_conditional_flow’s own log_density(x, y)/ sample_given(x) contract): theta – the quantity being inferred – is the student’s y-ARGUMENT, and the observed data y is its x-argument. So q(theta | y_obs) is module.sample_given(y_obs) -> theta and its density is module.log_density(y_obs, theta) – the inverse of the simulator’s own arrow.

The student is trained through the vendored GradLeaf (a bare torch module IS the model), not the simpler NeuralConditionalDensity adapter – see notes/designs/M3.md for why: sequential refinement (below) re-scores the module against freshly generated round data before the next fit commits, which wants the generic seq_log_density path GradLeaf gives any bare module (warm-started across optimize() calls via the SAME underlying nn.Module object) rather than a second bespoke accumulator. NeuralConditionalDensity remains the simpler documented alternative for callers who don’t need round-conditioned rescoring.

Algorithm (notes/designs/M3.md):

  1. Pair generation (round 1). theta_i ~ p(theta) via the prior’s own sampler; y_i = simulator(theta_i) in a plain Python loop (no batching assumed on simulator).

  2. Student. build_conditional_flow/build_mdn wrapped in GradLeaf, fit via optimize(list(zip(y_pairs, theta_pairs)), leaf, ...) – A4.4’s tuple-default-loss fix is what lets the bare module’s two-arg log_density(x, y) score straight off tuple observations.

  3. Sequential refinement (rounds 2..R). Resolved decision (was open in the design note): eager, via an optional y_obs keyword to THIS function. When y_obs is given, each subsequent round draws theta ~ q(theta | y_obs) from the CURRENT round’s student, re-runs the simulator, and retrains warm-started from the previous round’s module weights (same object, so optimize continues training it in place – the same warm-start pattern GradLeaf’s own M-step uses across EM iterations). rounds > 1 without y_obs has no observation to sharpen toward, so it raises ValueError rather than silently doing nothing – round 1 alone (unconditional pair generation) is valid without y_obs.

  4. Optional exactness stage. reweight=True with true_log_likelihood(theta, y_obs) -> float (a LOG likelihood) treats the final round’s q(theta | y_obs) as a self-normalized- importance-sampling proposal: log w_j = log p(theta_j) + true_log_likelihood(theta_j, y_obs) - log q(theta_j | y_obs), normalized by log-sum-exp (the same construction mixle.inference.condition’s SIR fallback uses), with ESS = 1 / sum(w_norm^2) reported so a low ESS visibly says “don’t trust this reweighted posterior” instead of silently returning a degenerate one.

  5. Calibration receipts (always computed). - SBC. Resolved decision (was open): a chi-square uniformity test on binned ranks (Talts et

    al.), bins = min(20, n_sbc_replications // 5) (clamped to >= 2), one chi-square statistic per theta dimension SUMMED (valid under the standard independence-across-dimensions simplification – a sum of independent chi-square variables is chi-square with summed degrees of freedom), against threshold p-value > 0.01 (no rejection of uniformity).

    • Coverage. Per-dimension, per-replication credible interval containment vs nominal level, averaged; pass when within +/-5% of nominal.

    • Prior-predictive. Empirical mean/std of the round-1 simulated y_i’s, plus (when y_obs is given) its per-dimension z-score against that empirical distribution – a caller-facing warning, not a gate.

class InverseModel(*, module, prior, simulator, family, theta_dim, y_dim, receipts, seed=None)[source]

Bases: object

A fitted amortized posterior q(theta | y) plus its InverseReceipts.

Parameters:

  • module (Any)

  • prior (Any)

  • simulator (Callable[[Any], Any])

  • family (str)

  • theta_dim (int)

  • y_dim (int)

  • receipts (InverseReceipts)

  • seed (int | None)

posterior(y)[source]

Wrap q(theta | y) as an M0 Posterior: sample(n) / log_density(theta) / mean(field) / .receipt – so downstream condition/do composition treats a learned inverse like an exactly-conditioned one, modulo the amortization warning on .receipt and the InverseReceipts pointer at .receipt.inverse_receipts.

Parameters:

y (Any)

Return type:

Posterior

class InverseReceipts(sbc_statistic, sbc_pvalue, sbc_bins, sbc_replications, sbc_pass, coverage, coverage_pass, prior_predictive, rounds_trained, sharpness_by_round=<factory>, ess=None, ess_ratio=None, warnings=<factory>)[source]

Bases: object

The calibration report that ships with every InverseModel – tells the caller whether to trust q(theta | y), not just a point estimate.

Parameters:
learn_inverse(simulator, prior, *, family='flow', n_sims=2000, rounds=1, n_sbc_replications=200, coverage_levels=(0.5, 0.9), reweight=False, true_log_likelihood=None, y_obs=None, seed=None, m_steps=200, lr=5e-3, max_its=1, hidden=32, n_posterior_samples=200, n_reweight_samples=500)[source]

Learn an amortized posterior q(theta | y) for simulator g: theta -> y under prior p(theta). See the module docstring for the full algorithm and the calibration receipts computed unconditionally.

family="flow" (build_conditional_flow) requires theta (the quantity being inferred, the student’s y-argument) to be >= 2-dimensional – build_conditional_flow needs y_dim >= 2 for its coupling layers to be non-trivial (see its own docstring). A 1-D theta (e.g. a scalar-parameter inverse problem) must use family="mdn", which has no such restriction (a mixture of per-component Gaussians is well-defined for scalar theta too, and is the more direct fit for asserting multimodality component-by-component).

rounds > 1 (SNPE-style sequential refinement toward a SPECIFIC observation) requires y_obs: round 1 alone (unconditional pair generation) is the only round that has meaning without an observation to sharpen against.

Parameters:
Return type:

InverseModel