Sec. X: End-to-End PE Drivers (jaxpe.drivers)

1. The sampling schedule
2. The args contract
3. API Reference
   1. run_pe
   2. map_laplace
   3. eta_to_q

The bin/ benchmark scripts are not toy examples: they carry the relative-binning likelihoods, the MAP+Laplace initialisation and the adaptive warmup/equilibration/production schedule that produce the numbers in the benchmark ledgers. Two of them — bin/run_bns_ce_pe.py (frequency domain, IMRPhenomD) and bin/run_td_phenomt_pe.py (time domain, IMRPhenomT) — differed only in how they build a likelihood. Everything downstream was duplicated verbatim.

jaxpe.drivers.relative_binning_pe holds that shared core. It lives in the library rather than in bin/ for two reasons: bin/ scripts are not importable without importlib path hacks, and bin/ has no test coverage, so several hundred lines of convergence logic were previously unreachable by the test suite.

The sampling schedule

run_pe implements a three-phase schedule over the unconstrained posterior \(\pi(y \mid d)\), all phases sharing one mass matrix derived from the Laplace covariance at the MAP.

Mass matrix with an eigenvalue floor. The GW posterior piles up against the \(\eta = 1/4\) and \(\chi_i = 0\) prior boundaries, so its soft directions are approximately \(\mathrm{Exponential}\) tails in unconstrained space rather than Gaussian. A mode Hessian systematically reports these too narrow. Writing the Laplace covariance in its eigenbasis \(\Sigma = V \Lambda V^{\top}\), the floor raises only the soft eigenvalues,

\[\tilde\lambda_i = \begin{cases} \max(\lambda_i, 1) & \lambda_i > 10^{-2} \\ \lambda_i & \text{otherwise,} \end{cases} \qquad \tilde\Sigma = V \tilde\Lambda V^{\top},\]

which widens the tail directions together with their correlated chirp-mass and spin compensations, because the eigenbasis is preserved. A diagonal floor does not have this property: it dilutes the correlations and collapses acceptance.

Warmup adapts the step size toward a target acceptance, then re-seeds chains that are stranded on secondary ripples of the oscillatory matched-filter likelihood. Equilibration and production interleave local kernel steps with flow-assisted global moves (see sampler), refitting the normalizing flow on a rolling window of recent samples and reverting any refit whose acceptance collapses.

Convergence is gated on rank-normalized \(\widehat{R}\) rather than the plain statistic — on boundary-piled marginals with heavy unconstrained-space tails, plain \(\widehat{R}\) is noisy and biased high.

The args contract

run_pe and _make_kernel read their configuration off an argparse.Namespace rather than an explicit signature. This is deliberate — each driver exposes roughly forty tuning flags, and threading them individually would be worse — but it is an implicit interface, so it is written down here and in the module docstring.

run_pe(problem, y_map, cov0, args, timings) reads exactly these 27 attributes:

adapt_gain, ensemble_metric, equil_rounds, ess_target, flow_acc_target,
flow_epochs, flow_interval, flow_interval_wide, flow_layers, flow_width,
friction, kernel, max_minutes, max_production_blocks, n_chains, n_global,
n_leapfrog, production_steps, retune_blocks, rhat_target, seed, step_size,
target_acceptance, thin, warmup_blocks, warmup_leapfrog, warmup_steps

Both shipped drivers define all 27. A new caller must too; a missing attribute surfaces as an AttributeError partway through a run, which is the honest cost of this design.

Reproducibility caveat. The production loop breaks when elapsed wall clock exceeds --max-minutes, so a run that hits its budget is not reproducible block-for-block — the block count depends on how fast the machine was that day. When you need a deterministic result, set --max-minutes beyond reach and bound the run with --max-production-blocks instead.

API Reference

run_pe

jaxpe.drivers.relative_binning_pe.run_pe(problem, y_map, cov0, args, timings)

Runs warmup, equilibration and production, returning (phys, log_probs, rhat, ess, converged, kernel). phys has shape (n_kept, n_chains, n_dim) in physical parameters; rhat and ess are per-dimension. timings is mutated in place with warmup, warmup_first_block_incl_compile, flow_fit, equilibration and production.

problem needs .log_posterior and .prior.to_physical.

map_laplace

jaxpe.drivers.relative_binning_pe.map_laplace(problem, y0, n_newton=24, tol=1e-9)

Damped Newton ascent with eigenvalue-clipped curvature and a backtracking line search, returning (y_map, cov, log_posterior_at_map). Every accepted step strictly increases the log-posterior, so the search cannot leave the starting mode’s basin — essential for a likelihood whose chirp-mass direction is an oscillatory needle (\(\sigma_y \sim 10^{-4}\)) flanked by secondary ridges.

The returned covariance is \(\left(-\tfrac{1}{2}(H + H^{\top})\right)^{-1}\) evaluated at the mode, with eigenvalues clipped away from zero, so it is positive definite by construction and usable directly as a dense metric.

eta_to_q

jaxpe.drivers.relative_binning_pe.eta_to_q(eta)

Mass ratio \(q = m_2/m_1 \le 1\) from the symmetric mass ratio \(\eta = q/(1+q)^2\), via \(\delta = \sqrt{1 - 4\eta}\) and \(q = (1-\delta)/(1+\delta)\). Uses the double-where idiom so the gradient stays finite at the equal-mass boundary \(\eta = 1/4\), where \(\partial\delta/\partial\eta\) diverges; \(\eta > 1/4\) is clamped to \(q = 1\) rather than returning NaN.