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. Proposal-corrected sequential refinement (rounds 2..R). Each round draws from the frozen current proposal q_r(theta | y_obs), records p(theta) / q_r(theta | y_obs) importance weights, and refits on ALL retained rounds. Round 1 prior draws keep global support; every later block is self-normalized to equal round mass. Thus every round’s weighted conditional-density objective targets the original prior joint p(theta)p(y|theta) – never the proposal joint.

  4. Optional finite-particle target correction. reweight=True with true_log_likelihood(theta, y_obs) turns the final neural posterior into a self-normalized importance proposal. The returned model stores those particles and normalized weights, and posterior(y_obs) samples and computes means from that weighted empirical posterior. It is an importance-resampled approximation, not an exact continuous density; ESS and a warning make that contract explicit.

  5. Calibration receipts (always computed).

    • SBC. Ties are randomized and every discrete rank cell is jittered into a continuous uniform variate. Each parameter dimension gets an equal-probability-bin chi-square test; a Bonferroni-adjusted minimum p-value supplies a valid global test under arbitrary cross-dimension dependence.

    • Coverage. Per-dimension empirical coverage ships with simultaneous 99% Wilson intervals; a level passes only when its nominal probability lies inside every dimension’s interval.

    • 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, reweighted_y=None, reweighted_particles=None, reweighted_weights=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)

  • reweighted_y (np.ndarray | None)

  • reweighted_particles (np.ndarray | None)

  • reweighted_weights (np.ndarray | 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, sbc_pvalues_by_dimension=<factory>, sbc_method='randomized_rank_bonferroni', coverage_by_dimension=<factory>, coverage_intervals=<factory>, coverage_method='simultaneous_99pct_wilson', sharpness_by_round=<factory>, round_training=<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 performs proposal-corrected SNPE toward a specific observation. Round-one prior simulations are retained forever. Every later proposal block is weighted by p(theta)/q_round(theta|y_obs) and the accumulated weighted objective continues to target the declared prior posterior rather than the narrower proposal posterior.

reweight=True requires y_obs and returns an observation-bound, finite-particle importance-resampled posterior. It improves the target measure when its ESS is healthy but does not claim an exact continuous density.

Parameters:
Return type:

InverseModel