Changelog

Versions are derived from git tags by setuptools-scm.

Unreleased

Fixed

  • IMRPhenomT’s coefficient solve evaluated the merger-amplitude boundary at the wrong frequency. _solve_22_amplitude_coefficients used the eta-dependent tcut22 boundary via the inspiral-only omega ansatz, instead of the fixed t=-150M boundary via the full piecewise omega22 dispatch that LALSuite’s own amplitude solve actually uses. Only matters for eta far from the tested grid, where tcut22 can fall outside the inspiral region entirely; zero regression on the existing test suite.

  • IMRPhenomT was a placeholder. _compute_phenom_coefficients returned a fixed dict, ignoring its eta/chi1/chi2 arguments entirely — silently wrong for every call, not just outside some range (see docs/constants.md). Replaced with a real port of LALSuite’s algorithm (Estelles et al. 2020, arXiv:2004.08302): 3.5PN TaylorT3 inspiral plus NR-fitted higher-order corrections solved via collocation, an arcsinh-parametrized merger frequency ansatz, Damour & Nagar 2014’s ringdown ansatz, and phase as the exact analytic integral of each region’s frequency ansatz (no numerical integration). Two LAL convention bugs were found and fixed along the way: the reference phase must be anchored at f_ref (not t=0/merger), and phi_ref enters through the spin-weighted Ylm’s azimuthal argument, not a separate per-mode rotation. Verified against real LALSuite: mismatch 1.4e-6–4.6e-6 across a parameter grid (mass ratio 1:1–5:1, spins to 0.8, varying inclination/phi_ref) — the same precision level as IMRPhenomD.

  • IMRPhenomTHM’s higher modes were a placeholder envelope, not a waveform model — invented t_meco/t_ring constants and ad hoc amplitude/frequency shapes (flagged in docs/constants.md, verified mismatch ~0.8 against LALSuite, i.e. uncorrelated). Replaced with a faithful reimplementation of the 4 higher modes (2,1),(3,3),(4,4), (5,5), generalizing IMRPhenomT’s (2,2) coefficient-solving machinery to arbitrary (l,m): mode-specific PN amplitude coefficients, the same merger/ringdown frequency+phase ansatze with per-mode QNM substitution, and the fixed per-mode phase offsets. Verified against LALSuite’s own IMRPhenomTHM: polarization-level mismatch 1.4e-10–2.5e-6 across a parameter grid, a direct per-mode comparison against SimIMRPhenomTHM_Modes, and an explicit check that odd-m modes vanish exactly at the equal-mass/equal-spin symmetry point. --waveform phenomthm (guarded with a clear error since it was added to the CLI) is now enabled in run-pe, the same pattern as --waveform phenomt.

  • PhaseDistanceMarginalLikelihood.dist_bounds defaulted to (1000, 8000) Mpc regardless of a run’s actual configured distance prior. For a campaign’s injections as close as 200 Mpc, the distance-marginalization quadrature grid silently never covered the true distance, producing a wrong (not merely noisier) marginal likelihood over the intrinsic parameters — caught by luck, not by design. The identical pattern was independently reproduced in ModesNetworkLikelihood.log_marginal_likelihood[_full]’s dist_min/dist_max (100, 5000). Both are now required keyword-only arguments with no library-side default, so an omitted bound fails immediately with TypeError instead of silently answering wrong; MarginalizedIntrinsicLikelihood.__init__ additionally validates its settings dict eagerly. The one place a default may still exist — (100, 8000) Mpc, for a degenerate resolved prior box — now lives in jaxpe/cli.py as a named, documented constant behind _dist_bounds_for_marginalization(). docs/constants.md is the full audit this fix came out of, covering every hardcoded numeric literal in the library functioning as a physical or algorithmic boundary.

  • Every injection in a campaign received the same noise realisation. The CLI passed seeds.noise through verbatim, so generate-injections --noise gaussian --n-injections N followed by run-pe analysed N copies of one realisation. This silently invalidates a PP-plot campaign, which is the main thing such a campaign exists to run. generate-injections now resolves a per-injection seed with derive_noise_seed(seeds.noise, index) and records it as metadata.noise_seed, so the realisation is pinned by the artifact; run-pe reads it back, falling back to deriving from metadata.index for sets already on disk.

  • A detector’s noise depended on its position in detector_names. make_injection advanced a single numpy Generator through the detector loop, so ("H1","L1") and ("L1","H1") gave H1 different data at the same seed — measured at a 130% relative change. Streams are now keyed by the detector’s name.

Added

  • HTCondor DAG generator and job scripts for the PE campaign (bin/condor/generate_campaign_dag.py, run_pe_job.sh, postprocess_job.sh). Emits three per-variant submit-file templates (phenomd_hmc, phenomd_gpry, esigma_gpry), a smoke.dag covering injection 0 of each variant, and one DAG per track — campaign_<variant>.dag — rather than a single bundled campaign.dag, since bundling let the slowest track (phenomd_hmc, the only variant without GPry’s own convergence-based early stopping) block submitting the other two at all. Each track DAG carries its own POSTPROCESS_<variant> node (postprocess_campaign.py --variants <variant>) as a child of that track’s PE jobs, so corner plots and the PP plot are produced automatically as soon as the track finishes. --maxjobs/--track-maxjobs default to 0 (no DAGMAN_MAX_JOBS_SUBMITTED cap) — the condor negotiator already does fair-share scheduling, so a DAGMan-side cap just leaves nodes idle without protecting anything. request_cpus=8 and run_pe_job.sh’s matching taskset -c 0-N wrapper are sized from measurements on the real pool (~25–26 static 32-core/63.7 GB nodes; JAX/XLA sizes its CPU thread pool from sched_getaffinity, not OMP_NUM_THREADS/XLA_FLAGS, so taskset is what actually constrains it).

  • bin/postprocess_campaign.py, a general-purpose post-processing report for PE run suites laid out as <results_dir>/<variant>/<run_id>/. Auto-discovers every variant and completed run, then builds one self-contained, file://-viewable report: per-run corner plots with WebP thumbnails for a gallery grid, a per-variant Beta-order-statistic PP plot with per-parameter KS-test p-values against Uniform(0,1), and a static index.html. Missing/failed runs are silently skipped, so the report stays honest about partial completion of a large campaign.

  • ESIGMA and the marginalized-intrinsic likelihood, wired into run-pe, enabling PhenomD+HMC / PhenomD+GPry / ESIGMA+GPry side by side from the same CLI: --waveform {auto,phenomd,phenomt,phenomthm,esigma}, decoupled from --domain (ESIGMA is always internally TD → FFT’d to FD); --likelihood marginalized_intrinsic, the 4-D GPry-facing likelihood via extrinsic importance sampling, for any waveform exposing mode_dict(); generate-injections --target-snr-range/--target-snr-waveform to redraw only luminosity_distance (via distance_for_target_snr) to hit a target network SNR, holding every other parameter’s draw bit-for-bit identical to an unflagged run under the same seed; run-pe GPry overrides (--gpry-ref-bounds-rel/-abs, --gpry-noise-level, --gpry-svm-threshold, --gpry-trust-region-threshold) for populations spanning a wide SNR range, following the GPry LISA paper’s own high-SNR SMBHB recipe, since GPry’s defaults (tuned for cosmological likelihoods) underfit the sharper log-likelihood curvature of a high-SNR CBC signal; and a persistent JAX compilation cache ($HOME/.jaxpe, NFS-safe) so a campaign’s many single-process CLI invocations reuse compiled XLA across runs sharing static shapes. Backed by a new esigma.* run-configuration section mapping onto ESIGMAInspiral’s constructor kwargs, the esigma conda/pip extra (diffrax, numba, numbalsoda, an editable esigmapy-dev install), and examples/configs/campaign_10to80_aligned.json, the backing config for a chirp-mass-[10,80] M☉ aligned-spin method-comparison campaign.

  • jaxpe.gw.match: a reusable, PSD-weighted match/mismatch between two frequency-domain waveforms, generalizing a test-local helper that required a jaxpe likelihood object into a standalone function taking arrays + PSD + df directly. Pure JAX throughout, so it composes under jax.jit/jax.vmap. This is the shared comparison primitive the IMRPhenomD/IMRPhenomT/IMRPhenomTHM LALSuite-validation tests are built on; IMRPhenomD in particular had never been checked numerically against LAL before (only jaxpe-internal batched-vs-serial self-consistency) — tests/test_phenomd.py now does, over a grid spanning mass ratio, spin and inclination: worst-case mismatch ~4e-9.

  • make_injections builds a whole suite of injections in one vmap, keeping signal, projection and noise on the accelerator. Measured on 4 s @ 1024 Hz H1+L1: the vectorised arithmetic alone is ~4x on CPU and ~6x on GPU at n=256, and end-to-end against a loop over make_injection it is 28x at n=64 on GPU (173.5 s → 6.2 s) — though most of that larger margin is amortising make_injection’s per-call jax.jit compilation rather than parallelism.

    Deliberately scoped: this is for suites that are the product (training sets, banks, systematics studies), not a speedup for validation campaigns, where injection creation is a fraction of a percent of wall time. A batch needs one analysis grid, so it cannot mix trigger times or durations, and oversized requests are refused with an estimate — a 2048 s BNS segment is 67 MB per injection per detector.

    Compare batched against serial by mismatch (~1e-9), not elementwise amplitude: raw amplitudes differ at ~5e-5 because XLA fuses IMRPhenomD differently between jit graphs. That predates this work and is unrelated to batching.

  • jaxpe.drivers.relative_binning_pe, the shared core of the two bin/ relative-binning benchmarks. bin/run_bns_ce_pe.py (FD, IMRPhenomD) and bin/run_td_phenomt_pe.py (TD, IMRPhenomT) had duplicated 670 lines: kernel construction, the MAP+Laplace mode finder, and the whole warmup/equilibration/production schedule. The duplication was not approximate. Compiled and fingerprinted, run_pe, _make_kernel, eta_to_q, _decimate and rank_normalized had byte-identical bytecode in both files; the entire functional difference between the two copies was one leftover print("H HAS NANS:") debug block in map_laplace, now deleted. Two comments naming a driver-specific benchmark ledger were retargeted, since a module both drivers import cannot name one of them.

    Each driver keeps what genuinely differs — build_loglike, validate_rb, main — and imports eta_to_q, map_laplace and run_pe, the only three names either needs. 1406 lines leave bin/; 756 arrive in the library.

    It is a library module rather than a bin/ script so that it is importable and testable: bin/ has zero test coverage, which left ~570 lines of convergence logic unreachable by the suite. tests/test_drivers.py adds the first coverage this code has had.

    The implicit interface is now documented rather than merely present: run_pe reads 27 attributes off an argparse.Namespace, listed in the module docstring and in docs/api/drivers.md. Both drivers define all 27.

    Verified by characterisation baseline, since there was no coverage to rely on: both drivers were run twice before the move and shown bitwise reproducible, then again after, with samples.npz compared key by key. Every array is bitwise identical across the move.

  • network_snr and distance_for_target_snr, replacing the build-measure-rescale recipe open-coded at three sites. h ∝ 1/D exactly, so the target is hit in one measurement with no search.

  • analysis_grid and resolve_f_max, one definition of the segment grid and of the 90%-of-Nyquist convention, replacing five copies.

Changed

  • Noise generation moved to JAX (simulate_noise_fd_jax), unifying the RNG with parameter drawing, which already used PRNGKeys. simulate_noise_fd (numpy) remains exported and unchanged.

    This changes which noise realisation a given seed produces. Nothing in the test suite pins noise sample values, but the cached artifacts under examples/output/ were computed against the old realisation and have been regenerated.

    One reproducibility property is worth knowing: JAX’s underlying random stream is bitwise identical across CPU and GPU, but jax.random.normal is not — measured at up to 3e-15 between backends, because XLA compiles erf_inv differently for each. Realisations are therefore bitwise reproducible on a given platform and equal to ~1e-15 across platforms. Use simulate_noise_fd where bitwise cross-platform equality is required.

  • examples/09_validate_injection_vs_dynesty.py fingerprints the injection data into its bilby label. bilby resumes from a checkpoint keyed only by label, so a checkpoint computed against different data would previously have been resumed silently under the new likelihood.

  • bin/run_bns_ce_pe.py, bin/run_td_phenomt_pe.py and bin/profile_sampler_scaling.py use jaxpe.gw.lalsim_psd("CE", ...) instead of a cosmic_explorer_psd helper duplicated in the first two and loaded by file path in the third. Verified bitwise identical at 8 s @ 2048 Hz and 2048 s @ 4096 Hz.

  • A driver run that hits --max-minutes is not reproducible block-for-block. Not a new behaviour, but newly written down (docs/api/drivers.md): the production loop breaks on elapsed wall clock, so the block count depends on how fast the machine was. For a deterministic run, put --max-minutes beyond reach and bound it with --max-production-blocks. This is why the characterisation baselines guarding the extraction above were captured that way.

  • bin/run_td_phenomt_pe.py: --target-snr no longer changes the analysis window. Its rescale rebuild passed tukey_alpha=0.0 while the first build used the 0.1 default, so requesting a target SNR silently altered the Tukey taper as well as the distance. Both builds now share one hoisted keyword dict.

Documentation

  • docs/gpry_gw_pe_literature.md — a survey of GP-surrogate active-learning samplers in GW PE, written while diagnosing a phenomd_gpry GPAcquisitionError on the 300-run campaign. Answers whether GPry itself has prior GW-PE use (one published application: LISA, not LVK/CE/ET) and whether the high-SNR GP-robustness problem hit here has prior art — RIFT’s publication record documents the same failure shape from a different entry point, with fixes that turned out to be transferable.

  • docs/constants.md gained a full audit of hardcoded numeric literals functioning as physical or algorithmic boundaries, tolerances, or resolutions (the distance-bounds fix above is one entry in it), and its IMRPhenomT/IMRPhenomTHM entries were updated from “placeholder” through “IMRPhenomT fixed, IMRPhenomTHM still open” to fully resolved as each reimplementation landed and was LAL-validated.

Performance

  • Post-processing no longer rebuilds an injection to recover the prior. PostProcessor only uses problem.prior; obtaining it cost a waveform generation, projection, FFT and jit compile per samples file — measured at 1.70 s and 1.83 s. InferenceProblem.log_likelihood now defaults to a stub that raises, so a prior-only problem is expressible without lying about the density.

0.1.0

First tagged release. Cut to mark a significant performance fix in the sampling engine.

Fixed

  • run_chains initialised its chains outside the jit (jaxpe/kernels/base.py). jax.vmap(kernel.init) ran eagerly, dispatching the target’s entire gradient graph one operation at a time. For a gravitational-wave likelihood (~3600 instructions) that cost ~2.2 s per call against ~0.004 s jitted — roughly 550× — and it was a fixed cost, paid on every call regardless of n_steps.

    Initialisation now happens inside _run_chains_jit, which already took logp_fn as a static argument, so there is no additional jit cache and no extra recompilation. Every kernel benefits; the effect is largest for expensive targets and for workflows that make many short run_chains calls.

    Measured on the BNS/Cosmic-Explorer benchmark (bin/run_bns_ce_pe.py), which makes ~33 such calls per run: 6.05 → 3.11 min mean wall clock, and 15.42 → 3.11 min against the original reference. The fixed per-call cost measures 0.028 s after the fix, down from 2.254 s.

    Guarded by two regression tests in tests/test_kernels.py. The primary one asserts the structural cause — that the target is never invoked with concrete, untraced values during run_chains — rather than a wall-clock threshold, so it stays meaningful across hardware.

Added

  • bin/profile_sampler_scaling.py — fits T = fixed + marginal × work for run_chains and _global_block across several sizes, plus gradient cost versus chain count and global-block cost versus flow capacity. This is the tool that found the bug above; point measurements had missed it. Recommended as the first thing to run on new hardware.

Changed

  • Benchmark defaults in bin/run_bns_ce_pe.py, each with its measurement recorded at the argument: --equil-rounds 5 → 3, --production-steps 25 → 12, --flow-layers 8 → 4 (new flag), --max-production-blocks 40 → 80. The first two had previously measured as losses because of the fixed cost above; both win once it is removed.
  • --max-production-blocks raised because 40 turned “needs two more blocks” into a reported non-convergence on a 1.35 + 1.25 M☉ source. --max-minutes is the real budget guard.

Documentation

  • docs/bns_ce_pe_benchmark.md rewritten and pruned (848 → 338 lines), reorganised around the result rather than the chronology of getting there.

Copyright © Prayush Kumar. Distributed under the MIT license.

This site uses Just the Docs, a documentation theme for Jekyll.