Factorized sampler: a composable Step-based inference spine

[stephengreen]

Summary

This PR replaces the monolithic Sampler/GWSampler/GWSamplerGNPE classes with a factorized inference spine: the posterior is sampled by an explicit chain of steps — factors (networks, pins, tables), reparametrizations, and corrections — composed by a ChainComposer over a per-event GWSamplerContext. Every existing workflow (NPE, multi-iteration GNPE with density recovery, importance sampling with synthetic phase) runs on the spine bit-exactly, and the spine's first new physics capability — DINGO-BNS single-step chirp-mass GNPE (Dax et al., Nature 639, 49 (2025)) — runs end-to-end through dingo_pipe on real GW170817 data.

Architecture

• dingo/core/factors.py — the domain-agnostic spine: Factor / Reparametrization / TargetCorrection step kinds (FlowFactor, DeltaFactor, SampleTableFactor, ProxyOffsetReparam, GibbsBlock), Stage/ChainComposer (fan-out, point-mass and 1:1 multiplicity rules, kind-aware reverse-fold log_prob), and the ComposedSampler base.
• dingo/gw/inference/ — split by responsibility: context.py (GWSamplerContext: event data with its network-input, prior, and likelihood views; importance-sampling representations via immutable derive()), steps.py (GNPE factors, SyntheticPhaseFactor, the RAToEventFrame/RAToTrainingFrame pair, SpinConventionReparam), sampler.py (GWComposedSampler + chain builders), scan.py (the chirp-mass scan).
• Result carries no construction knowledge: prior, domain, and likelihood delegate to the context (reconstructed from metadata for from-file results); synthetic phase runs as a chain rooted in the proposal sample table.
• Results record structured provenance: settings["sampler"] holds the serialized chain description, model paths, and (where applicable) density-recovery and scan recipes.

Binary neutron stars

• Heterodyning transforms (HeterodynePhase, batch-native factor_fiducial_waveform) ported from Add heterodyning transforms for BNS chirp mass GNPE #355; d_inner_h phase decomposition (return_rho2opt) ported from Core GNPE extensions for BNS support #354; matched-filter SNR reporting (return_aux_snr) added for the scan.
• Fixed-proxy chain [DeltaFactor(pins) → RAToTrainingFrame → FlowFactor → ProxyOffsetReparam → RAToEventFrame]: single-pass sampling with preserved density (no Gibbs, no density recovery). The pinned values have one owner (the chain), and the heterodyne consumes the proxy through prepared_data(conditioning=...) — one method with a uniform row-aligned contract serving pins and sweeps alike.
• Chirp-mass scan (--chirp-mass-scan): sweeps the proxy over the training prior on the ordinary chain machinery (grid table root, per-row batched heterodyne, one network pass per block), selects the maximum-likelihood draw, and records the trigger + SNR in provenance. --fixed-context-parameters accepts the published demo-INI form.
• A pinned sky position is rotated into the network's training frame at input and back at output; exactly zero for networks whose reference time is the trigger (all current ones).

Verification

• Legacy-vs-composed bit-exactness through both pipe stages (NPE and GNPE with density recovery), IS evidence identical to recorded digits; the legacy samplers are deleted (~1200 lines) with the last containing commit tagged legacy-samplers-final.
• BNS sampling bit-exact to float32 against the bns_add_dingo_pipe_max research branch on frozen event data; IS side verified live ≡ from-file bit-exactly.
• Real GW170817 (cleaned LOSC frames + GWTC-1 PSDs) through the full pipe: with the demo pins, efficiency ~10–11%, log BF +510, chirp mass 1.19754 ± 0.00007; with only the sky pinned, the scan recovers the trigger (1.1976, SNR 33.0) and the evidence agrees with the pinned run within Monte Carlo error — the prior-conditioning invariance working in production.
• GPU campaign (A100 server): full test suite green on Linux/CUDA; NPE sampling with the reverse-fold log_prob re-evaluated on-device (float32-level agreement); real GW170817 BNS pin and scan runs reproduce the reference evidences within Monte Carlo error (scan trigger 1.19756, snr 33.0, efficiency ~11%); GNPE with density-recovery training and prior-dict-updates exercised end-to-end on CUDA. Two device leaks in the previously CPU-only table-rooted chain path were found and fixed.
• A standing regression battery (full suite + pinned BNS/NPE/scan harnesses) ran after every change in this arc.

Documentation

• New pages: Sampling chains -- the factorized-sampler concepts page (the product-of-conditionals model, step types, the sampler context, multiplicity rules, the reverse density fold, provenance, and chain building, with figures) -- and Binary neutron stars -- the DINGO-BNS workflow (phase heterodyning, prior conditioning, the chirp-mass scan, and the dingo_pipe configuration).
• Rewritten: Inference (the Python-API workflow page) and the injection example, which was executed end to end against a trained GNPE pair as part of the update.
• Updated: GNPE (live inference routes replacing the removed dingo_analyze_event; a new single-step section) and dingo_pipe (the BNS sampler options; the base-domain likelihood default).
• build_model_from_kwargs is re-exported from dingo.core.posterior_models so the documented import works. Build locally with uv run --group docs sphinx-build -b html docs/source docs/build.

Compatibility and deferred work

• Pre-v0.9.0 models need compatibility/remove_domain_window_factor.py (unchanged story).
• Deferred, tracked for follow-ups: Max's decimated/heterodyned fast likelihood (and the JAX layer from the paper), the synthetic-phase compute_likelihood fast path, embed-once on FlowWrapper, calibration-as-a-factor, per-event time scans for pre-merger networks, CUDA/server realism runs.

🤖 Generated with Claude Code

[stephengreen]

@max-dax — this is the factorized-sampler spine we've been coordinating on sampler-revamp, now at a good point to look at: Increments 1a–4 are done, _post_process is fully eliminated, and it's bit-exact vs the old samplers throughout. Opened as a draft for design feedback before any dingo_pipe wiring. Would value your eyes on the Step structure and the two factors.py modules whenever you get a chance — it supersedes the per-factor batching and folds in the GNPE kernel/flow split we worked through.

[stephengreen]

@max-dax For reviewing, I would recommend starting by reading the updated documentation, mainly the sections on "Sampling chains" and "Binary neutron stars".

[stephengreen]

@max-dax flagging this line for a look — it feeds the trigger SNR that the chirp-mass scan logs and stores in provenance.

The deep review caught that this previously reused the phase-marginalized likelihood term ln_i0(|kappa2C|) as the SNR numerator. The intended statistic is the phase-maximized matched-filter SNR: under the (2,2) approximation, <d, h(phi)> = |kappa2C| cos(2 phi - arg kappa2C), so maximizing over phase gives |kappa2C| / sqrt(rho2opt) (with rho2opt phase-invariant). By contrast, ln_i0(|kappa2C|) ~ |kappa2C| - 0.5 ln(2 pi |kappa2C|) is the marginalized (integrated) weight — the maximized value minus the phase-volume Occam penalty — so the old value was biased low (GW170817 scan: 32.89 → 32.99 after the fix). Trigger selection was never affected, since the scan argmaxes the log-likelihood, which was always correct.

A regression test pins the statistic now (test_phase_marginalized_aux_snr_is_phase_maximized). Note return_aux_snr is new in this branch — it did not come from #354; the PR description's earlier attribution was wrong and has been corrected.