diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5102b8b87..d4167d924 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -197,6 +197,14 @@ jobs: python -m pip install coverage pytest --break-system-packages python -m pip install --editable . --break-system-packages - name: Run probe confirm-on-fail accounting and CIP intrinsic-prior tests + # The multi-approximant marginalization suite rides along too (~11 s, and it + # builds its own DAG from synthetic inputs, so it needs no rundir). It guards a + # combination that is silently wrong rather than loud: pooling several waveform + # models flat weights them by REPLICA COUNT instead of by p(m), which shifts lnL + # by a few tenths of a nat and raises nothing. Its DAG cases build with TWO + # approximants deliberately -- with one, every cross-model path is vacuously + # satisfied, which is how that builder kept a severed grid handoff for years. + # # The CIP prior suite rides along here for the same reason: it is pure logic that # runs in seconds (it reads CIP's source with ast instead of importing it, so it # needs only numpy/scipy), and a wrong prior density or a mis-wired @@ -205,8 +213,14 @@ jobs: run: | python -m pytest -q \ MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_probe_confirm.py \ + MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_confirm_regressions.py \ + MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_rift_provenance_guard.py \ + MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_seq_gmm_check.py \ MonteCarloMarginalizeCode/Code/test/test_cip_priors.py \ - MonteCarloMarginalizeCode/Code/test/test_cosmology_single_source.py + MonteCarloMarginalizeCode/Code/test/test_cip_gp_kernel_bounds.py \ + MonteCarloMarginalizeCode/Code/test/test_cosmology_single_source.py \ + MonteCarloMarginalizeCode/Code/test/test_multiapprox_marginalization.py \ + MonteCarloMarginalizeCode/Code/test/test_multiapprox_pseudo_pipe.py q-window-stencil-check: needs: install @@ -231,38 +245,103 @@ jobs: # SILENTLY: picking the wrong sub-sample stencil raises nothing, it just makes Q_lm(t) # less accurate, which surfaces only as a slightly wrong likelihood surface. # - # test_q_window_interp asserts the cubic/sinc crossover in BOTH directions on purpose. - # sinc winning everywhere would mean the Lanczos window had been widened until it was no - # longer a local stencil, so neither direction may be relaxed to make a change pass. - # test_time_interp_choice pins the pipeline thresholds inside the measured ambiguous - # band, and checks the decision uses the sampling rate the run is actually on. + # MEMBERSHIP IS NO LONGER A LIST HERE. It used to be a backslash-continued pytest + # invocation in this block; three PRs appended to that one list in two days and it + # conflicted twice, and taking one side of such a conflict silently unregisters the + # other side's test -- which in this CI simply never runs again, green forever. + # A test file now joins this job by carrying, on its own line, # - # test_calmarg_stencil_gating runs its CPU arms without a GPU (its GPU arm is additive), - # so it belongs here: it is what stops cubic/sinc being routed to the fused calibration - # kernel, which is implemented for 'nearest' only. + # # RIFT-CI-GATE: q-window-stencil # - # test_interpolate_time_cli runs the three scripts as real SUBPROCESSES (~30 s). That - # cost is the point: the unit tests exercise the resolver and the gate predicate, but - # neither can see whether the SCRIPTS are still wired to them. Reverting a parser to - # const=None, or deleting a script's resolver call, leaves every unit test green while - # restoring a bare flag that silently does nothing. All three mutations were checked to - # fail these tests before they landed. + # so a PR adding a test edits only its own new file. .travis/test-q-window-stencil.sh + # discovers those files, and its MANIFEST fails the job if a file matching one of the + # filename patterns this job owns carries no marker and is not explicitly excluded -- + # so forgetting to register a test is RED, not silent. It also pins a per-file + # collection floor of 1 (pytest exits 5, "no tests ran", on a file that collects + # nothing), a total collection floor, and a junit OUTCOME floor on tests that actually + # PASSED. Modelled on .travis/test-slowrot.sh, which solves the same problem for + # slowrot-check; the per-file rationale that used to sit in this block now lives in + # each test file's header, beside its marker. # # The GPU parity files (test_q_window_interp_gpu, test_noloop_gpu_stencils) are - # deliberately NOT here -- there is no GPU on these runners, and they would report as - # skipped. They are run by hand on a GPU node; the numbers are in PR #97. + # deliberately NOT run -- there is no GPU on these runners, and they would report as + # skipped. They are run by hand on a GPU node; the numbers are in PR #97. They are + # in the script's EXCLUDED list, which is what keeps the manifest from demanding them. # # test_slowrot_* files do NOT belong here. They are gated by slowrot-check, whose # manifest requires every test_slowrot_*.py in this directory to be listed or # explicitly excluded; a copy in this job's list is invisible to that manifest and - # would simply run the file twice (issue #169). + # would simply run the file twice (issue #169). This job's scope patterns exclude + # that prefix on purpose. + env: + OMP_NUM_THREADS: 1 + run: bash .travis/test-q-window-stencil.sh + + ci-roster-check: + runs-on: ubuntu-latest + # Census, one level above every other gate in this file. CI membership here is per-file and + # hand-listed -- ci.yml and .gitlab-ci.yml name individual files, .travis/*.sh name individual + # files, and three directories are handed to pytest whole. Two mechanisms already defend a + # SINGLE job's list against silent loss: .travis/test-slowrot.sh (an explicit FILES manifest, + # issue #169) and .travis/test-q-window-stencil.sh (a marker line inside each test file, PR + # #242). Neither can answer "is this file registered with ANY job?", and on 2026-09-03 the + # answer for 86 of 200 test files under Code/ was no. + # + # This job does not run any test. It asserts that every test file is either reachable from + # CI configuration or carries a line in .travis/ci_roster.txt saying why it is not, so a new + # test dropped into a directory no job runs goes RED instead of sitting unrun forever. + # + # No `needs: install`: stdlib only, no numpy and no RIFT import, so it runs in seconds and + # reports independently of whether the install matrix is healthy. + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - name: Census every test file against the CI roster + run: python .travis/test-ci-roster.py + + core-unit-check: + needs: install + runs-on: ubuntu-latest + # The tests the census found unrun and that SHOULD run: ordinary pytest suites, numpy / + # scipy / lal / sklearn only, that collect and pass in ~50 s, spanning calmarg, likelihood + # dispatch, integrator seeding and allocation, CIP evidence, distance export, hyperpipe, and + # the packaging/config contracts. They guard quantities that regress SILENTLY -- an evidence + # normalization, a seeding path, a distance grid, a container manifest -- where a wrong value + # is still a plausible value and nothing raises. + # + # See .travis/test-core-units.sh for why it counts collection PER FILE and asserts on junit + # OUTCOMES rather than invoking pytest on a directory: several files in these same + # directories collect ZERO items, and pytest exits 5, "no tests ran", which reads as a pass. + # + # Membership is an explicit FILES manifest, not the marker line of q-window-stencil-check. + # That is deliberate: this set spans six unrelated subject areas with no shared filename + # pattern, so the marker's SCOPE_GLOBS check -- the half that makes a marker fail-closed -- + # would have nothing to scope over. The census job above is what keeps this set honest. + # + # Python 3.10 to match the sibling numpy jobs. timeout-minutes is a runaway backstop, an + # order of magnitude above the 49 s measured on CIT, not a budget. + timeout-minutes: 20 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.10' + cache: 'pip' + cache-dependency-path: requirements.txt + - name: Enable symlink + run: sudo ln -sf $(which python3) /usr/bin/python + - name: Install dependencies run: | - python -m pytest -q \ - MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_window_interp.py \ - MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_time_interp_choice.py \ - MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_calmarg_stencil_gating.py \ - MonteCarloMarginalizeCode/Code/RIFT/misc/test_psd_bandwidth.py \ - MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_interpolate_time_cli.py + python -m pip install --upgrade pip --break-system-packages + python -m pip install -r requirements.txt --break-system-packages + python -m pip install coverage pytest --break-system-packages + python -m pip install --editable . --break-system-packages + - name: Run core unit gate + env: + OMP_NUM_THREADS: 1 + run: bash .travis/test-core-units.sh slowrot-check: needs: install @@ -329,9 +408,20 @@ jobs: # and means a breaking upstream jax release can redden this job outside any PR's # control. If that becomes noisy, pin here rather than deleting the job. # - # Cost. CURRENT (EXPECTED_TESTS=27 in .travis/test-jax.sh): 27 tests, measured - # 578-834 s of pytest across repeat runs on quiet CPU nodes (ldas-pcdev11/13, jax - # 0.9.2, JAX_PLATFORMS=cpu, JAX_ENABLE_X64=1, OMP_NUM_THREADS=1). + # Cost. CURRENT (EXPECTED_TESTS=144 in .travis/test-jax.sh): 144 tests, measured + # 859 s of pytest on ldas-pcdev11 pinned to 16 cores (jax 0.9.2, + # JAX_PLATFORMS=cpu, OMP_NUM_THREADS=1). The count grew 27 -> 48 -> 64 + # (#180, fair-draw export) -> 95 (this PR, the tempering chooser), and #190 + # Cost. CURRENT (EXPECTED_TESTS=153 in .travis/test-jax.sh): 153 tests, the 139 + # previously measured at 859 s of pytest on ldas-pcdev11 pinned to 16 cores + # (jax 0.9.2, JAX_PLATFORMS=cpu, OMP_NUM_THREADS=1) plus test_angle_marg_exact.py + # (30 tests, ~600 s local on ldas-grid, 16 cores, same stack -- the sizing + # regressions run a deliberately large dense grid once). The count grew + # 27 -> 48 -> 64 + # (#180, fair-draw export) -> 95 (the tempering chooser) -> 139 -> 169 (the + # exact angle-marginalization schemes + their review fixes), and #190 + # added test_interp_choices.py along the way; this note sat at 27 through + # several of those, so re-derive it from the gate rather than trusting it. # test_jax_slowrot.py dominates (the p_max=0/p_max=1 rotation ladders and # freqresponse, each followed by the AD/jit/vmap/hessian checks); it is the first # thing to trim if CI minutes ever bite. timeout-minutes is generous so a slower @@ -342,7 +432,9 @@ jobs: # the actual runner (job 95750869285, python 3.11.15, jax 0.10.2, numpyro 0.21.0) -- # ~10x headroom, with the unpinned install already a jax minor version ahead of the # 0.9.2 measured locally. Do NOT read the 14 as a current count; the file list has - # grown since and the gate asserts 27. + # grown since and the gate asserts 169. That runner-vs-local ratio (286 s runner + # for 964 s local, ~3.4x) is why 859 s local is not a timeout concern at 60 + # minutes. timeout-minutes: 60 steps: - uses: actions/checkout@v4 @@ -370,6 +462,15 @@ jobs: OMP_NUM_THREADS: 1 run: bash .travis/test-jax.sh + # CPU ONLY. GitHub runners have no GPU and no cupy, so this lane cannot + # exercise the cupy branch of RIFT/likelihood/SphericalHarmonics_gpu.py -- the + # branch that used to break every LISA likelihood evaluation on a GPU host + # while passing here. The tests it runs therefore assert on the BACKEND + # ARGUMENT rather than only on the outcome (test_spherical_harmonics_backend.py + # fakes the GPU module; test_lisa_operational_synthetic.py spies on the call + # site), so this lane does catch that class of defect without a device. The + # real-cupy run lives in .travis/test-lisa-gpu.sh, driven by the gpu_integration + # job in .gitlab-ci.yml. lisa-check: needs: install runs-on: ubuntu-latest @@ -470,6 +571,52 @@ jobs: - name: Run Asimov integration test run: bash .travis/test-asimov.sh + rimsky-integration: + needs: install + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - rimsky-series: '0.1.0rc1' + rimsky-spec: 'rimsky==0.1.0rc1' + python-version: '3.12' + - rimsky-series: 'main-2621d15' + rimsky-spec: 'git+https://git.ligo.org/colm.talbot/rimsky.git@2621d15cf9a39ce01145ac5b81c92e1173a1d2d0' + python-version: '3.14' + name: rimsky-integration (${{ matrix.rimsky-series }}) + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + cache: 'pip' + cache-dependency-path: requirements.txt + - name: Enable symlink + run: sudo ln -sf $(which python3) /usr/bin/python + - name: Install system dependencies + run: sudo apt-get update && sudo apt-get install -y libgsl-dev + - name: Install RIFT and Rimsky orchestration stack + run: | + retry_pip_install() { + local attempt + for attempt in 1 2 3; do + if python -m pip install "$@" --break-system-packages; then + return 0 + fi + echo "pip install failed (attempt ${attempt}/3)" + sleep 10 + done + return 1 + } + python -m pip install --upgrade pip --break-system-packages + retry_pip_install '${{ matrix.rimsky-spec }}' + retry_pip_install --upgrade 'asimov>=0.7,<0.8' 'bilby_pipe @ git+https://git.ligo.org/lscsoft/bilby_pipe.git@be6c77021db809690c781a7744f40564e62dd72f' + python -m pip install pytest packaging --break-system-packages + python -m pip install --editable . --no-deps --break-system-packages + - name: Run Rimsky to RIFT end-to-end test + run: python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_rimsky_end_to_end.py + test-run: needs: install runs-on: ubuntu-latest @@ -516,7 +663,10 @@ jobs: # Also run by .travis/test-integrate.sh (integration-check, py3.10). Repeated here # because this job is the only one matrixed over BOTH numpy lanes, and these tests # are the kind that break on numpy API removals (e.g. np.trapz -> np.trapezoid). - run: python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py + run: | + python -m pytest -q \ + MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py \ + MonteCarloMarginalizeCode/Code/test/test_limit_distance.py - name: Run NAL / supplementary-likelihood hook tests # Also run by .travis/test-integrate.sh (integration-check, py3.10). Repeated here for the # same reason as above: nal_io is pure numpy and has a legacy-numpy fallback for diff --git a/.gitignore b/.gitignore index 89deb15d7..c72e06478 100644 --- a/.gitignore +++ b/.gitignore @@ -53,4 +53,4 @@ cover/ # in the .venv directory. It is recommended not to include this directory in version control. .pixi/* !.pixi/config.toml - +asimov.log diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 43c660f20..68c6efa4a 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -91,6 +91,7 @@ import_check: - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_lisa_lalsimutils_compat.py - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_lisa_helper_contract.py + - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py dependency_compat_check: stage: system tests @@ -129,6 +130,20 @@ test_run: # numpy/scipy, seconds). Runs before the heavy scripts so a wrong prior fails # fast rather than after the end-to-end runs. - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_cip_priors.py + # Per-extrinsic-sample time-marginalization offset in the vectorized NoLoop + # likelihood (issue #232). numpy + lal, ~9 s, and it fails SILENTLY in the worst + # way: a batch-wide log-sum-exp offset underflows every extrinsic sample more than + # ~745 nats below the loudest one to lnL = -inf, which above rho ~ 40 is the bulk of + # the prior and collapses mcsamplerAV. An unlisted test never runs in this CI, so + # the wiring ships with the fix. + - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_noloop_time_marg_row_offset.py + # Same defect on the in-loop calibration-marginalization reduction (n_cal>1, + # cal_method='loop', which is the DEFAULT: the fused kernel is opt-in behind + # --calibration-fused-kernel). Its streaming log-sum-exp offset was a SCALAR shared + # by every extrinsic sample, so rows more than ~745 nats below the loudest came back + # lnL = -inf (issue #232). numpy + lal, ~10 s. An unlisted test never runs in this + # CI, so the wiring ships with the fix. + - python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_calmarg_running_max_row_offset.py - . .travis/test-coord.sh - bash .travis/test-integrate.sh - . .travis/test-posterior.sh @@ -160,7 +175,7 @@ gpu_integration: --env GW_SURROGATE="$GW_SURROGATE" --env RIFT_CI_REQUIRE_GPU="$RIFT_CI_REQUIRE_GPU" "$RIFT_CI_APPTAINER_IMAGE" - bash -lc 'cd "$CI_PROJECT_DIR" && export PYTHONPATH="$CI_PROJECT_DIR/MonteCarloMarginalizeCode/Code${PYTHONPATH:+:$PYTHONPATH}" && bash .travis/test-integrate.sh && bash .travis/test-calmarg-gpu.sh' + bash -lc 'cd "$CI_PROJECT_DIR" && export PYTHONPATH="$CI_PROJECT_DIR/MonteCarloMarginalizeCode/Code${PYTHONPATH:+:$PYTHONPATH}" && bash .travis/test-integrate.sh && bash .travis/test-calmarg-gpu.sh && bash .travis/test-lisa-gpu.sh' rules: - if: '$CI_PIPELINE_SOURCE == "web"' when: manual diff --git a/.travis/ci_roster.txt b/.travis/ci_roster.txt new file mode 100644 index 000000000..8456d86e7 --- /dev/null +++ b/.travis/ci_roster.txt @@ -0,0 +1,136 @@ +# CI roster -- every test file under MonteCarloMarginalizeCode/Code that is reachable from NO +# CI job, and the reason it is not. Enforced by .travis/test-ci-roster.py (job ci-roster-check). +# +# Format: +# +# STATUS is one of: +# HANDRUN a hand-run study or demo, not a pytest target. Most of these have a __main__ +# and argparse, assert internally, and collect ZERO items under pytest -- which +# exits 5, "no tests ran", and reads as a pass in a log skim. Gating them as-is +# would buy a green tick over an empty run. +# LEGACY imports pre-package flat module names (factored_likelihood, lalsimutils, ourio, +# ourparams, xmlutils, common_cl, statutils, effectiveFisher, spokes) that have not +# existed since RIFT was packaged into RIFT.*. These cannot be IMPORTED, let alone +# run, in any current environment. +# OPTDEP needs a dependency CI does not install. +# GPU needs a GPU; the runners have none, so it would report as skipped. +# EXPENSIVE opt-in behind an env var by design. +# BROKEN collects but FAILS today. A debt, recorded as one. +# PENDING waiting on a named gate that is not live yet. The reason must say +# `gate:`, and the entry is legal only while that gate is absent -- when +# it lands, the census errors on the entry by name. An earlier version was +# exempt from the staleness check unconditionally, which made it the one status +# nothing could ever force out. No entries currently: the last one, for +# test_noloop_time_interp.py on gate:q-window-stencil, expired when #242 merged. +# +# Everything here was collected and run individually on CIT with the IGWN conda python +# (3.11, numpy 1.26.4, lal 7.7.0) on 2026-09-03; the counts quoted are from that run. +# +# Adding a line here is a DECISION, not a way to silence the check. If a file belongs in CI, +# register it with a job instead -- then delete its line, which the check will demand anyway. + +# --------------------------------------------------------------------------------------- +# LEGACY -- the pre-package driver suite. 28 files, none importable. +# +# These are the original hand-run RIFT test drivers from before the code became a package. +# `pytest --collect-only` on each reports a collection ERROR, not zero tests: e.g. +# test_likelihood.py -> ModuleNotFoundError: No module named 'factored_likelihood'. They are +# not a coverage gap, because there is no coverage to lose; they are a decision nobody has +# made. Recommendation, deliberately NOT taken in the PR that created this roster: delete +# them, or move them under Code/old/ where test_regions.py already sits. That is a judgment +# call about historical value, and it belongs to a maintainer, not to a CI audit. +MonteCarloMarginalizeCode/Code/old/test_regions.py LEGACY imports lalsimutils; also raises IndexError at collection on current numpy +MonteCarloMarginalizeCode/Code/test/factored_likelihood_test.py LEGACY imports factored_likelihood, lalsimutils +MonteCarloMarginalizeCode/Code/test/test_SpokesRefine.py LEGACY imports spokes +MonteCarloMarginalizeCode/Code/test/test_data_vs_template.py LEGACY imports factored_likelihood, lalsimutils, ourio, ourparams +MonteCarloMarginalizeCode/Code/test/test_effectiveFisher.py LEGACY imports effectiveFisher, lalsimutils +MonteCarloMarginalizeCode/Code/test/test_hlm.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_interpolation.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_like_and_samp.py LEGACY imports common_cl, factored_likelihood, lalsimutils, ourio, ourparams, xmlutils +MonteCarloMarginalizeCode/Code/test/test_like_and_samp_margPsi.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_like_and_samp_noisydata.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_like_and_samp_noisydata_margPsi.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_like_and_samp_simplified.py LEGACY imports common_cl, factored_likelihood, lalsimutils, ourio, ourparams, xmlutils +MonteCarloMarginalizeCode/Code/test/test_like_and_samp_singleifo.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_like_singletemplate_emcee.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_like_singletemplate_emcee_pt.py LEGACY imports factored_likelihood, lalsimutils, ourio +MonteCarloMarginalizeCode/Code/test/test_likelihood.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_precompute.py LEGACY imports factored_likelihood, lalsimutils +MonteCarloMarginalizeCode/Code/test/test_precompute_noisydata.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_precompute_singleifo.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_profile_components.py LEGACY imports factored_likelihood, lalsimutils, ourparams +MonteCarloMarginalizeCode/Code/test/test_psd_xml_io.py LEGACY imports factored_likelihood, lalsimutils +MonteCarloMarginalizeCode/Code/test/test_response_functions_Q.py LEGACY imports factored_likelihood, lalsimutils, ourio, ourparams, xmlutils +MonteCarloMarginalizeCode/Code/test/test_rhotimeseries.py LEGACY imports factored_likelihood +MonteCarloMarginalizeCode/Code/test/test_runvar.py LEGACY imports statutils +MonteCarloMarginalizeCode/Code/test/test_sampler_visualization.py LEGACY imports lalsimutils +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp.py LEGACY imports statutils and mpl_toolkits.basemap +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp_pinned.py LEGACY imports statutils +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsampler_gpu.py LEGACY imports ourio + +# --------------------------------------------------------------------------------------- +# HANDRUN -- quantitative studies with real internal gates, run by hand. +# +# These are NOT dead. Each has a __main__, argparse, and its own pass/fail criterion (a +# 4-sigma bias gate, an efficiency ratio, a bias-ordering assertion). They collect ZERO items +# under pytest and exit 5, so wiring them into a pytest job as they stand would report a green +# tick over an empty run -- the exact trap .travis/test-slowrot.sh documents. +# +# The right move for the first six is to convert them the way the shape-recovery suite was +# converted: a pytest wrapper under test/expensive_before_merging/, skipped unless +# RIFT_RUN_EXPENSIVE=1, so the merge gate can invoke them and CI does not pay for them. That +# is a per-suite piece of work with a real cost, and it is not attempted here. +MonteCarloMarginalizeCode/Code/test/integrators/test_AV_bootstrap.py HANDRUN AV warm-start efficiency study with a 4-sigma bias gate; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_AV_warmstart_safety.py HANDRUN anti-bias guard for reusing a proposal across problems; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_adaptive_alloc.py HANDRUN portfolio draw-allocation study vs standalone AV; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_balance_heuristic.py HANDRUN portfolio safety under a decoy member; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_oracle.py HANDRUN needle-target oracle study; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamplerEnsemble.py HANDRUN GMM-vs-mcsampler comparison demo, prints results; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamplerEnsemble_AdaptationDemo.py HANDRUN adaptation demo, plots and prints; 0 collected, exit 5 +MonteCarloMarginalizeCode/Code/demo/rift/export_likelihoods/head_to_head/run_test.py HANDRUN GP-vs-RF figure driver for a demo; --stage picks a stage; 0 collected, exit 5 + +# --------------------------------------------------------------------------------------- +# OPTDEP / GPU -- would run, but not in the environment CI builds. +# +# The two jax_gp files are the strongest candidates for promotion: jax-ile-check already +# installs a CPU jax stack, so adding them there costs only optax and a raised EXPECTED_TESTS. +# Not done here because that job's counts are pinned and this PR does not own them. +MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py OPTDEP needs jax; belongs in jax-ile-check, which already installs a CPU jax stack +MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py OPTDEP needs jax and optax; 10 collected, all 10 error on ModuleNotFoundError optax rather than skipping +# The two cupy parity legs are NOT listed here. PR #242 landed, and its +# .travis/test-q-window-stencil.sh names both in an EXCLUDED array -- with the same reason, +# and with its own fail-closed check that an EXCLUDED path still exists and does not carry +# the marker. That is a better home than this file: the decision sits beside the gate it +# belongs to. Their roster lines were deleted when #242 merged, exactly as the census +# demanded ("listed as GPU but IS now reachable -- delete it"). +MonteCarloMarginalizeCode/Code/test/backends/test_backends_lowlevel.py OPTDEP needs lscsoft-glue and htcondor; 15 collected, 15 pass where both are installed +MonteCarloMarginalizeCode/Code/test/hyperpipe/tests/test_hydra_integration.py OPTDEP needs hydra and omegaconf; skips cleanly without them +MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_mode_sign.py OPTDEP needs the gwsignal TEOB route; 2 collected, 1 passes and 1 skips without EOBRun_module +MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_near_aligned.py OPTDEP needs EOBRun_module; 1 collected, skips without it +MonteCarloMarginalizeCode/Code/test/test_teobresums_compat.py OPTDEP TEOBResumS compat shim; 15 collected and all pass on CIT, unverified on a runner +MonteCarloMarginalizeCode/Code/test/test_rimsky_integration.py OPTDEP companion to test_rimsky_end_to_end.py; belongs in the rimsky-integration job, whose env this PR does not own +MonteCarloMarginalizeCode/Code/test/integrators/test_NF_reuse.py OPTDEP needs nflows for the normalizing-flow store; collection errors without it +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp_vegas.py OPTDEP needs the vegas package, commented out of requirements.txt; NameError at import without it +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsampler_rosenbrock.py HANDRUN Rosenbrock sampler study; its docstring pairs it with plot_posterior_corner.py by hand +MonteCarloMarginalizeCode/Code/test/test_eosmanager_misc.py OPTDEP needs LALSIMULATION_DATADIR set; raises KeyError at import without it +MonteCarloMarginalizeCode/Code/test/test_skysamp.py LEGACY imports lalinference.bayestar.fits, removed upstream; cannot be imported +MonteCarloMarginalizeCode/Code/test/test_mcsampler_foridiots.py BROKEN NameError int_vals at import; a plotting demo that no longer runs at all + +# --------------------------------------------------------------------------------------- +# EXPENSIVE -- correctly gated already, by an env var rather than by CI membership. +MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_shape_recovery.py EXPENSIVE 4 collected, all skip unless RIFT_RUN_EXPENSIVE=1 +MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_escaped_mass_diagnostic.py EXPENSIVE 5 collected, all skip unless RIFT_RUN_EXPENSIVE=1 + +# --------------------------------------------------------------------------------------- +# BROKEN -- collects and FAILS on rift_O4d today. Found only because this audit ran them. +# +# test_replica_pooling.py is the clearest argument for the census. It loads six helpers out of +# bin/integrate_likelihood_extrinsic_batchmode by REGEX and exec()s them into a synthetic +# module. The driver has since been refactored so that _lnZ_of_rvs and _kish_neff_of_rvs +# delegate to a seventh helper, _lw_of, which the regex list does not extract. Inside the +# exec'd module _lw_of is undefined; the driver's own `except Exception: return None` swallows +# the NameError, both helpers return None, and 10 of 15 tests die on `None - float`. Adding +# "_lw_of" to the slice list in the test is the immediate fix. The reimplemented-harness shape +# is the real problem and outlives that fix. +MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py BROKEN 10 of 15 fail; its regex helper-slicer misses _lw_of, added to the driver after the test was written +MonteCarloMarginalizeCode/Code/test/hyperpipe/tests/test_marg_list.py BROKEN 2 of 3 fail; _stage_event_file writes event-N.net into base_dir while the test and assemble_marg_list's own run_dir docstring say run_dir diff --git a/.travis/test-calmarg.sh b/.travis/test-calmarg.sh index 443266e51..87b1ab552 100755 --- a/.travis/test-calmarg.sh +++ b/.travis/test-calmarg.sh @@ -3,7 +3,10 @@ # reduction and the per-realization self-term fix: precompute time-alignment + # identity-cal cross terms, the loop/fused reduction vs a brute-force reference # (incl. n_cal==1), the low-rank SVD self-term basis vs a direct band integral, and -# the backtest of the cal reduction (default + distance-marginalization helpers). +# the backtest of the cal reduction (default + distance-marginalization helpers), +# and the driver's calibration OPTION-COMPATIBILITY gate (which flag combinations +# are refused at startup instead of silently degrading to the zero-calibration +# likelihood, and which must still be accepted). # Any nonzero exit fails the job (set -e). GPU/CUDA paths are exercised separately # on hardware; here every check runs on the numpy backend. set -euo pipefail @@ -13,6 +16,20 @@ command -v "$PY" >/dev/null 2>&1 || PY="$(command -v python3)" CODE="MonteCarloMarginalizeCode/Code" export OMP_NUM_THREADS=1 +# INVARIANT: this gate tests THIS CHECKOUT, never an installed build. Must PREPEND -- +# appending lets a caller's PYTHONPATH win. +# +# Without it the first check below runs as `python /test_precompute_alignment.py`, which +# puts the SCRIPT'S directory on sys.path and not Code/, so bare `import RIFT` resolves to +# whatever is installed. In CI that happens to be the editable install of this checkout, so +# it passes; on a plain checkout with a real RIFT in the environment it silently tests the +# INSTALLED code, and reports its staleness as a failure of this branch -- observed on CIT, +# where the CVMFS IGWN RIFT predates a kwarg the checkout added: +# PrecomputeLikelihoodTerms() got an unexpected keyword argument 'calibration_realizations' +# The `python -m RIFT.calmarg.*` runs below are already safe (cwd is on sys.path under -m); +# this makes the whole script safe the same way its sibling gates are. +export PYTHONPATH="$PWD/$CODE${PYTHONPATH:+:$PYTHONPATH}" + # precompute alignment + identity-cal self-term cross terms == baseline "$PY" "$CODE/RIFT/calmarg/test_precompute_alignment.py" @@ -24,4 +41,19 @@ export OMP_NUM_THREADS=1 && "$PY" -m RIFT.calmarg.backtest_calmarg --backend cpu --n-cal 8 --methods reference,in_loop_B \ && "$PY" -m RIFT.calmarg.backtest_calmarg --backend cpu --n-cal 8 --loglikelihood distmarg --methods reference,in_loop_B ) +# Option-compatibility gate: which calibration option combinations the ILE driver +# REFUSES at startup, and -- the load-bearing half -- which it must still ACCEPT. +# pytest here, not a bare module run: these are parametrized accept/refuse pairs plus +# real subprocess invocations of the driver, and `set -e` turns pytest's exit 5 ("no +# tests ran") into a job failure rather than a silent green. A collection floor is +# asserted for the same reason: a file that stops collecting is a lost gate. +_n=$( cd "$CODE" && "$PY" -m pytest -q --collect-only "RIFT/calmarg/test_option_compat.py" 2>/dev/null | grep -c '::' || true ) +if [ "${_n:-0}" -lt 20 ]; then + echo "test-calmarg.sh: option-compat suite collected ${_n:-0} tests, expected >= 20." >&2 + echo " Tests were renamed, removed, or the file stopped importing. Fix it; as is," >&2 + echo " the gate covers less than it claims." >&2 + exit 1 +fi +( cd "$CODE" && PYTHONPATH="$PWD${PYTHONPATH:+:$PYTHONPATH}" "$PY" -m pytest -q "RIFT/calmarg/test_option_compat.py" ) + echo "calmarg CPU regression gate: PASS" diff --git a/.travis/test-ci-roster.py b/.travis/test-ci-roster.py new file mode 100755 index 000000000..06e8d241f --- /dev/null +++ b/.travis/test-ci-roster.py @@ -0,0 +1,501 @@ +#!/usr/bin/env python3 +"""Repo-wide census: every test file must be reachable from CI, or rostered with a reason. + +WHY THIS EXISTS. CI job membership in this repo is per-file and hand-listed: ci.yml and +.gitlab-ci.yml name individual files, .travis/*.sh name individual files, and three +directories are passed to pytest whole. Two mechanisms already defend a SINGLE job's list +against silent loss -- .travis/test-slowrot.sh (an explicit FILES manifest, issue #169) and +.travis/test-q-window-stencil.sh (a marker line inside each test file, PR #242). Both are +SCOPED: slowrot's manifest covers test_slowrot_*.py, and q-window's SCOPE_GLOBS cover eight +filename patterns. Neither can answer the question one level up -- "is this file registered +with ANY job?" -- and on 2026-09-03 the answer for 86 of 200 test files was no. + +That gap is not the same defect as a conflicted job list, and it does not have the same fix. +Most of those 86 files should NOT be gated: they are hand-run studies, plotting demos, and +scripts that import pre-package flat module names (factored_likelihood, lalsimutils, ourio) +which have not existed since RIFT was packaged, so they cannot even be IMPORTED. A +membership marker has nowhere to record that. What was missing is a place to record a +DECISION, and a check that fails when a file has none. + +So this script does not gate anything. It asserts that every test file under Code/ is either +reachable from CI configuration, or carries a roster entry stating why it is not. A new test +file added to a directory no job runs now fails the build instead of sitting unrun. + +REACHABILITY, AND WHAT EACH SIGNAL CAN AND CANNOT DO. A file counts as covered by one of +three signals, and they are not equally trustworthy. An earlier version of this docstring +claimed all three could only UNDER-report -- miss that a file is unrun, never falsely clear a +file. That was wrong for two of them, and both were live bugs: + + * DIRECTORY TARGET -- the file lives under a directory handed to pytest whole. Exact. + * NAMED IN A CI CONFIG -- the basename stem appears in a non-comment line, so a module + invocation (`python -m RIFT.calmarg.test_selfterm_basis`) counts as well as a path. This + is textual and therefore loose in one direction: it cannot distinguish an invocation from + a mention. Comment-only lines are stripped for exactly that reason (see _cfg_blob) -- + before that, ci.yml's comment explaining why the two cupy parity files are EXCLUDED was + enough to mark them covered. It can still be fooled by a filename appearing in a + non-comment context that does not run it; that residue is accepted and stated. + * RIFT-CI-GATE MARKER -- the file asserts its own membership. This one is the file + claiming coverage for itself, so the claim is checked against KNOWN_GATES rather than + taken on its shape. Unvalidated, `# RIFT-CI-GATE: totally-made-up-job` cleared a file + that no script greps for and nothing runs. + +The roster is where the narrower truth is written down. + +Stdlib only -- no numpy, no RIFT import -- so it can run as its own cheap job. +""" + +import os +import re +import subprocess +import sys + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +CODEDIR = os.path.join("MonteCarloMarginalizeCode", "Code") +ROSTER = os.path.join(".travis", "ci_roster.txt") + +# Directories handed to pytest as a whole, so every file beneath them runs. Kept explicit +# rather than parsed: a wrong guess here would silently EXCUSE files, which is the failure +# mode this script exists to prevent. Each must still be a real directory (checked below). +DIR_TARGETS = ( + CODEDIR + "/RIFT/simulation_manager/tests", # .travis/test-simulation-manager.sh + CODEDIR + "/test/asimov_integration", # .travis/test-asimov.sh + CODEDIR + "/test/jax", # .travis/test-jax.sh +) + +# Membership markers of the PR #242 kind. A file carrying one is registered with that job by +# the job's own script, so the census must not then demand a roster entry for it. +# +# BUT THE NAME MUST BE VALIDATED, or this branch is fail-open. The marker is the only one of +# the three reachability signals that a test file can assert about ITSELF, and the pattern +# below matches shape, not meaning: `# RIFT-CI-GATE: totally-made-up-job` is syntactically +# perfect and no script anywhere greps for it. Honouring it would count the file as covered +# while nothing runs it -- the inert-guard failure this whole census exists to catch, +# reintroduced inside the census. It is not hypothetical: before KNOWN_GATES landed, planting +# that exact line on one of the LEGACY files below -- a file that cannot even be IMPORTED, +# because it wants the pre-package `factored_likelihood` -- moved it into the "reachable from +# CI" count and made its roster entry report as STALE. +# +# So a marker counts only if its gate name appears in KNOWN_GATES below, and only if that +# gate's script really greps for the exact literal. KNOWN_GATES is a shared line, which is +# what PR #242 set out to remove -- but it is one line per JOB, not per test, edited when a +# gate is created rather than when a test is added, so it does not carry the conflict cost +# that motivated the marker in the first place. +MARKER_RE = re.compile(r"^# RIFT-CI-GATE: ([a-z0-9-]+)$", re.M) + +# gate name -> the script that discovers files by that marker. Both directions are checked: +# a name here whose script exists but does not contain the literal is a broken registry, and a +# marker naming anything NOT here is a hard error rather than silent coverage. +# +# A gate whose script does not exist YET is simply not live: markers naming it are not +# honoured, and the files carrying them need roster entries until it lands. That is what keeps +# this independent of PR #242's merge order -- on rift_O4d today no file carries any marker. +KNOWN_GATES = { + "q-window-stencil": ".travis/test-q-window-stencil.sh", # PR #242 +} + + +# PROVING CONSUMPTION, RATHER THAN INFERRING IT. +# +# A gate is honoured only if its script can be ASKED which files it will run. Set this env var +# and the script must print those paths, one per line, and exit 0: +# +# if [ -n "${RIFT_CI_GATE_LIST:-}" ]; then printf '%s\n' "${FILES[@]}"; exit 0; fi +# +# PLACEMENT IS PART OF THE CONTRACT, and getting it wrong is not subtle -- it just fails. The +# discovery and this short-circuit must come BEFORE the script's pytest/numpy/lal probes, which +# means moving the marker/CODEDIR assignment and the discovery line above them. This census job +# deliberately has no `needs: install`: it is stdlib-only and must stay that way, so a list mode +# sitting behind a numpy probe cannot run here. Discovery needs only grep; nothing else in the +# gate's preamble is required to answer "which files". +# +# WHY NOTHING WEAKER WILL DO. Two textual checks were tried here and both were fail-open. +# Requiring the marker literal on a non-comment line passes an ORPHANED `MARKER=` assignment +# whose discovery grep has been deleted. Adding "and a line matching grep.*MARKER" does not fix +# it: test-q-window-stencil.sh contains `grep -qxF -- "${MARKER}" "${e}"` for a completely +# different purpose -- asserting that an EXCLUDED file does NOT carry the marker -- so deleting +# the real discovery line still left the pattern satisfied. Checked; it passed green. No +# regex distinguishes "uses the marker to find files" from "uses the marker"; only running the +# discovery does. +# +# The cost is a small requirement on any gate that wants marker-based membership, and it is +# reported precisely: a marker naming a gate that cannot list is an error against THAT GATE, +# quoting the snippet. It is not raised speculatively -- a registered gate with no marked files +# costs nothing. +LIST_ENV = "RIFT_CI_GATE_LIST" + +LIST_SNIPPET = ('if [ -n "${%s:-}" ]; then printf \'%%s\\n\' "${FILES[@]}"; exit 0; fi' + % LIST_ENV) + + +def _gate_listing(script): + """Ask a gate script which files it would run. (files, error); files is None if unsupported.""" + if LIST_ENV not in _strip_comment_lines(open(script, errors="replace").read()): + return None, None + env = dict(os.environ) + env[LIST_ENV] = "1" + try: + pr = subprocess.run(["bash", script], env=env, timeout=120, + stdout=subprocess.PIPE, stderr=subprocess.PIPE) + except (OSError, subprocess.TimeoutExpired) as exc: + return None, "%s declares %s but could not be run in list mode: %s" % (script, LIST_ENV, exc) + if pr.returncode != 0: + return None, ("%s declares %s but exited %d in list mode.\n" + " A list mode that cannot run is worse than none: it looks authoritative.\n" + " stderr: %s" + % (script, LIST_ENV, pr.returncode, + pr.stderr.decode("utf-8", "replace").strip()[:300])) + found = [l.strip() for l in pr.stdout.decode("utf-8", "replace").splitlines() if l.strip()] + if not found: + return None, ("%s listed no files in list mode. Its discovery returns nothing, so every " + "file carrying its marker is covered by an empty run." % script) + return found, None + + +def _live_gates(live_cfg): + """Gate name -> the set of files that gate's own discovery returns. + + Three conditions, each of which was a hole before it was checked: + + * the gate's script is INVOKED from a workflow entry point. A script on disk that no job + runs is not a gate, and honouring its markers repeats the DIR_TARGETS bug. + * it declares the marker literal outside its comments. (test-q-window-stencil.sh quotes + its own marker in the comment block explaining the mechanism, so a bare + `literal in text` passed with the live MARKER= renamed away.) + * it can LIST what it will run, and that listing is what the census believes. + """ + gates, errs = {}, [] + live_paths = set(live_cfg) + for name, script in sorted(KNOWN_GATES.items()): + literal = "# RIFT-CI-GATE: %s" % name + if not os.path.exists(script): + # Not an error: the gate has not landed (or was retired). Just not live. + continue + if script not in live_paths: + errs.append("KNOWN_GATES maps %r to %s, which exists but is invoked by no CI job.\n" + " Files carrying that marker would be counted as covered by a script " + "nothing runs.\n" + " Restore its job, or drop the registry entry." % (name, script)) + continue + if literal not in _strip_comment_lines(open(script, errors="replace").read()): + errs.append("KNOWN_GATES maps %r to %s, but that script does not contain the " + "literal %r outside its comments.\n" + " Every file carrying that marker is then counted as covered by a " + "gate that never looks for it.\n" + " Fix the script, or drop the registry entry." % (name, script, literal)) + continue + listing, lerr = _gate_listing(script) + if lerr: + errs.append(lerr) + continue + if listing is None: + # Not honoured, and not an error on its own -- only files that actually carry this + # marker are affected, and they are told below, individually. + continue + gates[name] = set(os.path.normpath(x) for x in listing) + return gates, errs + +VALID_STATUS = { + # not gated, and that is the right answer + "HANDRUN": "hand-run study or demo; not a pytest target", + "LEGACY": "imports pre-package flat modules; cannot be imported at all", + "OPTDEP": "needs a dependency CI does not install", + "GPU": "needs a GPU; CI runners have none", + "EXPENSIVE": "opt-in behind an env var by design", + # not gated, and that is NOT the right answer -- these are debts, stated as such + "BROKEN": "collects but fails; needs a fix before it can be gated", + # tolerated in either state while a companion PR is in flight + "PENDING": "unreachable AND waiting on a named gate; expires when either changes", +} + + +SELF = os.path.basename(os.path.abspath(__file__)) + + +def _strip_comment_lines(text): + """Drop lines whose first non-whitespace character is '#'. + + Used twice, and both uses closed a fail-open hole where prose about a name was accepted as + a use of that name. A comment is where a CI file EXPLAINS what it does not run, so it is + precisely the place a name appears without being invoked. + """ + return "\n".join(l for l in text.splitlines() if not l.lstrip().startswith("#")) + + +# Workflow ENTRY POINTS. Everything else counts only if something reachable from one of these +# invokes it. A .travis script that no job runs is not CI, it is a file. +CFG_ROOTS = (".github/workflows", ".gitlab-ci.yml", ".travis.yml") + + +def _live_configs(): + """CI files reachable from a workflow entry point, by transitive invocation. + + WHY NOT just "every file in .travis/". Coverage has to depend on a job actually running + something, not on a file existing. Deleting the asimov-integration job from ci.yml and + .gitlab-ci.yml left .travis/test-asimov.sh on disk, still naming test/asimov_integration/ -- + and the census went on reporting those three files as covered, green, with nothing running + them. Checked before this closure landed; that is the P2 this answers. + + A script joins the live set when its basename appears on a non-comment line of something + already live, so `bash .travis/test-asimov.sh` pulls it in and a deleted job drops it out. + Iterated to a fixed point, because scripts can invoke scripts. + """ + live = {} + for root in CFG_ROOTS: + if os.path.isdir(root): + for f in sorted(os.listdir(root)): + if f.endswith((".yml", ".yaml")): + live[os.path.join(root, f)] = open(os.path.join(root, f), errors="replace").read() + elif os.path.exists(root): + live[root] = open(root, errors="replace").read() + + candidates = {} + if os.path.isdir(".travis"): + for f in sorted(os.listdir(".travis")): + if f.endswith((".sh", ".py")) and f != SELF: + candidates[os.path.join(".travis", f)] = f + + changed = True + while changed: + changed = False + blob = "\n".join(_strip_comment_lines(t) for t in live.values()) + for path, base in sorted(candidates.items()): + if path in live: + continue + # NOTE the lookbehind allows '/': every invocation is a PATH, `bash + # .travis/test-asimov.sh`, so excluding '/' matched nothing and emptied the live set. + if re.search(r"(? (status, reason). Duplicate paths are an error: two verdicts, one file.""" + entries, errs = {}, [] + if not os.path.exists(ROSTER): + return entries, ["%s does not exist" % ROSTER] + for n, raw in enumerate(open(ROSTER, errors="replace"), 1): + line = raw.split("#", 1)[0].strip() if raw.lstrip().startswith("#") else raw.rstrip("\n") + if not line.strip() or raw.lstrip().startswith("#"): + continue + bits = line.split(None, 2) + if len(bits) < 3: + errs.append("%s:%d: need ' ', got %r" % (ROSTER, n, raw.strip())) + continue + path, status, reason = bits[0], bits[1], bits[2].strip() + if status not in VALID_STATUS: + errs.append("%s:%d: unknown status %r (valid: %s)" + % (ROSTER, n, status, ", ".join(sorted(VALID_STATUS)))) + if len(reason) < 12: + errs.append("%s:%d: reason for %s is too short to be a reason: %r" + % (ROSTER, n, path, reason)) + if path in entries: + errs.append("%s:%d: %s is listed twice" % (ROSTER, n, path)) + entries[path] = (status, reason) + return entries, errs + + +def main(): + os.chdir(REPO) + errs = [] + + live_cfg = _live_configs() + blob = "\n".join(_strip_comment_lines(t) for t in live_cfg.values()) + + # A directory target excuses every file beneath it, so BOTH halves must hold: the directory + # exists, AND some live CI file still hands it to pytest. Checking only the first meant the + # asimov-integration job could be deleted outright while its three files kept reporting as + # covered -- existence of a directory is not evidence that anything runs it. + live_dirs = set() + for d in DIR_TARGETS: + if not os.path.isdir(d): + errs.append("DIR_TARGETS names %s, which is not a directory. It was renamed or " + "removed; left as is it silently EXCUSES files from this census." % d) + continue + if not re.search(r"(?`, and the entry is legal only while that + # gate is NOT live. The moment the gate lands, the entry is an error naming itself. + # + # The cost is honest and stated in the PR: merging the companion needs a one-line deletion + # here. That is a forcing function, not a failure. + for f, (status, reason) in sorted(roster.items()): + if f not in reachable: + errs.append("%s: %s no longer exists. A roster entry for a deleted file is a " + "silent no-op; drop the line." % (ROSTER, f)) + continue + # PENDING carries an EXTRA condition, not a weaker one. It must still go stale the + # moment the file is covered -- by ANY job, not only by the gate it names. An earlier + # version checked the gate and then `continue`d unconditionally, so a file that became + # reachable through some other job while its named gate stayed dormant kept a PENDING + # entry for ever: the one escape left in this file, and the same "never expires" defect + # that removing the blanket exemption was meant to close. So fall through to the + # staleness check below rather than returning here. + if status == "PENDING": + m = re.search(r"gate:([a-z0-9-]+)", reason) + if not m: + errs.append("%s: %s is PENDING but its reason names no gate. Write `gate:` " + "in the reason so the entry has a condition that can expire, or use " + "a status that does not need one." % (ROSTER, f)) + elif m.group(1) not in KNOWN_GATES: + errs.append("%s: %s is PENDING on gate %r, which is not in KNOWN_GATES. A " + "condition that can never be met never expires." + % (ROSTER, f, m.group(1))) + elif m.group(1) in gates: + errs.append("%s: %s is PENDING on gate %r, and that gate is now LIVE.\n" + " The wait is over: either the gate registers this file (delete " + "this line) or it does not (give the file a real status)." + % (ROSTER, f, m.group(1))) + if reachable[f] is not None: + extra = ("\n PENDING is not an exemption from this: it waits on a named gate, but " + "the file is covered NOW, by this job." if status == "PENDING" else "") + errs.append("%s: %s is listed as %s but IS now reachable (%s). The entry is stale " + "-- delete it.%s" % (ROSTER, f, status, reachable[f], extra)) + + n_reach = sum(1 for v in reachable.values() if v is not None) + print("test-ci-roster: %d test files under %s" % (len(files), CODEDIR)) + print(" reachable from CI : %d" % n_reach) + print(" rostered : %d" % len(roster)) + counts = {} + for status, _ in roster.values(): + counts[status] = counts.get(status, 0) + 1 + for s in sorted(counts): + # An unknown status is already an error above; do not crash the summary on it, or the + # specific message gets buried under a traceback. + print(" %-10s %3d (%s)" % (s, counts[s], VALID_STATUS.get(s, "UNKNOWN STATUS"))) + + if errs: + print("\ntest-ci-roster: FAIL", file=sys.stderr) + for e in errs: + print(" " + e, file=sys.stderr) + return 1 + print("test-ci-roster: PASS -- every test file is gated or has a stated reason it is not.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.travis/test-core-units.sh b/.travis/test-core-units.sh new file mode 100755 index 000000000..c82fbb8d9 --- /dev/null +++ b/.travis/test-core-units.sh @@ -0,0 +1,174 @@ +#!/usr/bin/env bash +# CPU gate for unit suites that were reachable from NO CI job until this script landed. +# +# WHY. On 2026-09-03 an audit of every test_*.py under Code/ found 86 of 200 named by no +# workflow and no .travis script (see .travis/ci_roster.txt and .travis/test-ci-roster.py for +# the census that now keeps that number honest). Most of those 86 should stay out -- they are +# hand-run studies, plotting demos, or scripts importing pre-package flat modules that have not +# existed since RIFT was packaged. The files below are the ones that should NOT: they are +# ordinary pytest suites, numpy/scipy/lal/sklearn only, that collect and PASS in seconds, and +# they guard things that regress SILENTLY -- an evidence accounting, a seeding path, a +# distance grid, a container manifest, a parameter port. A wrong number there is still a +# plausible number. +# +# Every file listed here was run individually on CIT (IGWN conda python 3.11, numpy 1.26.4, +# lal 7.7.0) before it was added; the measured collection counts are the floors below. +# +# SHAPE. Modelled on .travis/test-slowrot.sh, and it keeps that script's defences, because +# the trap it documents is live in this very set: several files elsewhere in these directories +# collect ZERO items and pytest exits 5, "no tests ran", which reads as a pass. Membership +# here is an explicit FILES manifest rather than the marker line used by +# .travis/test-q-window-stencil.sh: this set spans six unrelated subject areas with no shared +# filename pattern, so a marker's scope check would have nothing to scope over. The two +# mechanisms answer different questions -- see the PR body. +set -uo pipefail +# NOTE: deliberately no -e, matching test-slowrot.sh. Every command's rc is handled below. + +cd "$(dirname "$0")/.." || { echo "test-core-units.sh: cannot cd to repo root" >&2; exit 1; } + +# INVARIANT: test THIS CHECKOUT, never an installed build. Must PREPEND -- appending lets a +# caller's PYTHONPATH win. +export PYTHONPATH="$PWD/MonteCarloMarginalizeCode/Code${PYTHONPATH:+:$PYTHONPATH}" + +PYTHON_BIN="${RIFT_COREUNIT_PYTHON:-${PYTHON:-python}}" +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + PYTHON_BIN="$(command -v python3)" +fi + +# Guard every probe whose pass condition is empty output: a missing interpreter plus a +# redirected stderr is indistinguishable from a clean result. +"${PYTHON_BIN}" -c 'import pytest' || { echo "test-core-units.sh: pytest unavailable" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import numpy, scipy; print("numpy", numpy.__version__)' \ + || { echo "test-core-units.sh: numpy/scipy unavailable" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import lal; print("lal", lal.__version__)' \ + || { echo "test-core-units.sh: lal unavailable" >&2; exit 1; } + +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" +export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}" +export MPLBACKEND="${MPLBACKEND:-Agg}" + +C="MonteCarloMarginalizeCode/Code" + +FILES=( + # -- calibration marginalization (module-level suites the calmarg gate never picked up) + "$C/RIFT/calmarg/test_cal_mc_error.py" + "$C/RIFT/calmarg/test_seed_fallback.py" + "$C/test/test_calmarg_calibration.py" + # -- likelihood dispatch + "$C/RIFT/likelihood/test_td_dispatch_epoch.py" + "$C/test/test_ile_scalar_edge_cases.py" + "$C/test/test_srate_resample_time_marginalization.py" + # -- integrators: seeding, allocation, weight derivation + "$C/test/integrators/test_convergence_sample_order.py" + "$C/test/integrators/test_gmm_adaptive.py" + "$C/test/integrators/test_portfolio_gmm_member_trains.py" + "$C/test/integrators/test_portfolio_restrict_and_warm.py" + "$C/test/integrators/test_rvs_weight_derivation.py" + "$C/test/integrators/test_seeding_public_paths.py" + "$C/test/integrators/test_seeding_reproducibility.py" + "$C/test/test_mc_error.py" + # -- CIP / evidence / distance export + "$C/test/test_cip_evidence_consolidation.py" + "$C/test/test_cip_pipeline.py" + "$C/test/test_distance_grid.py" + "$C/test/test_distance_tail.py" + "$C/test/test_dslice_device_native.py" + # -- hyperpipe (paper4 area; the hydra leg is rostered OPTDEP, not here) + "$C/test/hyperpipe/tests/test_config.py" + "$C/test/hyperpipe/tests/test_coords.py" + "$C/test/hyperpipe/tests/test_drivers.py" + "$C/test/test_hyperpipeline_io.py" + # -- packaging / config contracts / waveform conventions + "$C/test/test_advanced_parameter_ports.py" + "$C/test/test_container_manifest.py" + "$C/test/test_lisa_ini_contract.py" + "$C/test/test_tracer_placement_gp.py" + "$C/test/waveforms/test_uv_symmetry.py" +) + +# A manifest entry that stops existing is a SILENT no-op: the gate keeps passing while +# covering less. Same defence as test-slowrot.sh's DESELECT-still-resolves check. +missing=0 +for f in "${FILES[@]}"; do + [ -f "$f" ] || { echo "test-core-units.sh: manifest names $f, which does not exist." >&2; missing=1; } +done +[ "$missing" -eq 0 ] || { echo " Fix the manifest or restore the file; left as is it covers nothing." >&2; exit 1; } + +# PER-FILE collection floor of 1. A file that collects nothing is the exit-5 trap arriving +# through the front door: inside a multi-file run pytest's exit 5 never appears at all, so it +# has to be checked per file. +echo "== per-file collection floor ==" +floor_rc=0 +for f in "${FILES[@]}"; do + n=$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "$f" 2>/dev/null | grep -c '::') + printf ' %-72s %3d\n' "$f" "$n" + if [ "$n" -lt 1 ]; then + echo "test-core-units.sh: $f collects 0 tests." >&2 + echo " pytest exits 5 on that ('no tests ran'), which reads as a pass. Either it is not" >&2 + echo " a pytest target and belongs in .travis/ci_roster.txt, or its entry points broke." >&2 + floor_rc=1 + fi +done +[ "$floor_rc" -eq 0 ] || exit 1 + +# Pinned TOTAL floor, so a renamed file or a dropped test_* entry point goes red rather than +# green-on-fewer-tests. MEASURED 2026-09-03 on CIT with the IGWN conda python (3.11, numpy +# 1.26.4, scipy 1.14.1, lal 7.7.0), whole manifest in one run: 278 collected, 266 passed, +# 12 skipped (11 pytest.skip + 1 xfail), 49 s. +EXPECTED_TESTS=278 +# Outcomes, not just exit status: a collection floor cannot see a test that collects, runs and +# asserts nothing, and a pytest.skip can quietly absorb a lost gate. The 12 skips are +# environment legs -- cupy in test_seeding_reproducibility, device legs in +# test_dslice_device_native, and the xfail in test_uv_symmetry -- and a GitHub runner has no +# GPU either, so they skip there too. +# +# CONFIRMED ON A RUNNER. The first CI run of this job (PR #243, ubuntu-latest, python 3.10 + +# editable install) reported the same 278 / 266 / 12, in 24.7 s. So these floors are exact on +# both stacks, not merely the CIT numbers copied across, and a future divergence is a real +# change rather than an environment difference to be explained away. +EXPECTED_PASSED=266 +MAX_SKIPPED=12 + +junit="$(mktemp -t core-units-junit-XXXXXX.xml)" +echo "== running ==" +"${PYTHON_BIN}" -m pytest -q -p no:cacheprovider --durations=10 --junit-xml="${junit}" "${FILES[@]}" +rc=$? +if [ "$rc" -ne 0 ]; then + echo "test-core-units.sh: pytest exited ${rc}" >&2 + [ "$rc" -eq 5 ] && echo " Exit 5 is 'no tests ran'. That is a FAILURE here, not a pass." >&2 + rm -f "${junit}" + exit 1 +fi + +read -r TOT FAIL ERR SKIP < <("${PYTHON_BIN}" - "${junit}" <<'PY' +import sys, xml.etree.ElementTree as ET +r = ET.parse(sys.argv[1]).getroot() +s = r if r.tag == 'testsuite' else r.find('testsuite') +g = lambda k: int(s.get(k, 0)) +print(g('tests'), g('failures'), g('errors'), g('skipped')) +PY +) +rm -f "${junit}" +PASSED=$(( TOT - FAIL - ERR - SKIP )) +echo "== outcomes: ${TOT} collected, ${PASSED} passed, ${SKIP} skipped, ${FAIL} failed, ${ERR} errored ==" + +out_rc=0 +if [ "${TOT}" -lt "${EXPECTED_TESTS}" ]; then + echo "test-core-units.sh: collected ${TOT} tests, expected at least ${EXPECTED_TESTS}." >&2 + echo " A file was renamed, or a test_* entry point was dropped. Restore it, or lower this" >&2 + echo " floor DELIBERATELY in the same commit that removes the tests." >&2 + out_rc=1 +fi +if [ "${PASSED}" -lt "${EXPECTED_PASSED}" ]; then + echo "test-core-units.sh: only ${PASSED} PASSED, expected at least ${EXPECTED_PASSED}." >&2 + out_rc=1 +fi +if [ "${SKIP}" -gt "${MAX_SKIPPED}" ]; then + echo "test-core-units.sh: ${SKIP} SKIPPED, at most ${MAX_SKIPPED} expected." >&2 + echo " A pytest.skip can absorb a lost gate without failing anything. If the new skip is" >&2 + echo " legitimate, raise MAX_SKIPPED here and say which test and why." >&2 + out_rc=1 +fi +[ "$out_rc" -eq 0 ] || exit 1 + +echo "core unit gate: PASS" diff --git a/.travis/test-integrate.sh b/.travis/test-integrate.sh index 51cdd1b26..9f4a3cd5a 100755 --- a/.travis/test-integrate.sh +++ b/.travis/test-integrate.sh @@ -37,6 +37,11 @@ fi # coordinate-transform + prior-mass identities, so they belong with the integrator gate # rather than with the end-to-end run tests. python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py +# --limit-distance: the SAMPLING-only distance box. Listed separately from the cosine +# file because the failure it guards is a different one -- not "the flag is ignored" but +# "the flag silently renormalized the prior", which looks like success in the obvious +# check (lnZ comes back unchanged) and is only separable with a constant likelihood. +python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_distance.py python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_mcsampler_ensemble_log_contract.py # Supplementary-likelihood plugin hook: the NAL reader/evaluator (pure numpy, no data) and the @@ -45,6 +50,105 @@ python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_mcsampler_ensemble_ python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_nal_io.py \ MonteCarloMarginalizeCode/Code/test/test_supplementary_likelihood_hook.py +# Time-marginalization quadrature. The historical rule integrates exp(lnL(t)) with Simpson at +# the FIXED spacing deltaT=1/srate, while the integrand's width sigma_t = 1/(2 pi rho sigma_f) is +# set by the SIGNAL and shrinks as 1/rho -- so production under-resolves its own integrand, worse +# at higher SNR (measured: the reported lnL moves 1.649 nats when the grid phase is scanned over +# 2*deltaT at srate 4096, rho=40). This gate covers the opt-in band-limited quadrature against an +# ANALYTIC continuous reference, plus its finite-window reconstruction and resolution guards and +# -- the part that matters most +# here -- that the option actually reaches the shipped likelihood rather than being inert. +# The PIPELINE file is listed alongside it deliberately: the quadrature is inert unless it +# survives util_RIFT_pseudo_pipe.py -> helper_LDG_Events.py -> args_ile.txt -> +# create_event_parameter_pipeline_BasicIteration -> ILE*.sub, and the last link is an +# INHERITANCE (ile_args_extr = ile_args + ...), not an explicit forward. An unlisted test never +# runs in this CI, so wiring the file in is part of shipping the wiring. +_TMARG_TESTS=( + MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py + MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature_pipeline.py + MonteCarloMarginalizeCode/Code/test/test_continuous_time_posterior_export.py +) +# Count guard, matching .travis/test-slowrot.sh and test-jax.sh. `set -e` already +# catches a total collection failure (pytest exits 5), but a silent shrink from 60 +# tests to 3 -- a rename, a stale -k, a decorator that stops matching -- reads as +# green. Raise EXPECTED by RUNNING collection, never by arithmetic. +_TMARG_EXPECTED=161 +_TMARG_FOUND=$(python -m pytest -q --collect-only "${_TMARG_TESTS[@]}" 2>/dev/null | grep -c '::' || true) +if [ "$_TMARG_FOUND" -ne "$_TMARG_EXPECTED" ]; then + echo "time-marginalization gate: collected $_TMARG_FOUND tests, expected $_TMARG_EXPECTED" >&2 + exit 1 +fi +# SKIP guard. `pytest -q` exits 0 with skips, so a test that quietly stops +# running reads as green, and the count guard above catches DESELECTION, not +# SKIPPING. +# +# This IDENTIFIES rather than COUNTS, which an earlier version of it did not. +# Counting is wrong twice over: a compensating pair (one importorskip starts +# firing while the GPU test stops) keeps the total unchanged, and the expected +# total is a property of the RUNNER, not of the code -- on a GPU-equipped runner +# that does not set RIFT_CI_REQUIRE_GPU=1 the cupy test legitimately stops +# skipping, and a count guard then fails a perfectly good run. So: allow skips +# whose REASON names cupy/GPU, and fail on any other skip whatever the total. +_TMARG_OUT=$(python -m pytest -q -rs "${_TMARG_TESTS[@]}" 2>&1) || { echo "$_TMARG_OUT"; exit 1; } +echo "$_TMARG_OUT" | tail -20 +_TMARG_BAD=$(echo "$_TMARG_OUT" | grep -E '^SKIPPED' | grep -vciE 'cupy|gpu|cuda' || true) +_TMARG_BAD=${_TMARG_BAD:-0} +if [ "$_TMARG_BAD" -ne 0 ]; then + echo "time-marginalization gate: $_TMARG_BAD test(s) skipped for a reason other than an absent GPU:" >&2 + echo "$_TMARG_OUT" | grep -E '^SKIPPED' | grep -viE 'cupy|gpu|cuda' >&2 + exit 1 +fi + +# Peak-local time-marginalization quadrature. Same defect, same derived-resolution +# discipline; what changes is WHERE the refined grid is placed -- around the enumerated +# peaks of a band-limited kappa rather than over the whole window, so the cost stops +# growing with SNR. What this gate has to protect, beyond accuracy: that the intervals +# are MERGED (the un-merged variant double-counts the overlap, +1.6 nats at rho~6), that +# the omitted mass is BOUNDED rather than assumed (a deliberately sabotaged enumeration +# must be caught and sent to the dense path), that the local evaluator reconstructs the +# same interpolant the dense FFT does at every production npts including the odd ones, +# and that the option reaches the shipped likelihood instead of being inert. +_TMARG_PL_TESTS=MonteCarloMarginalizeCode/Code/test/test_time_marginalization_peak_local.py +# Raise EXPECTED by RUNNING collection, never by arithmetic. +_TMARG_PL_EXPECTED=121 +_TMARG_PL_FOUND=$(python -m pytest -q --collect-only "$_TMARG_PL_TESTS" 2>/dev/null | grep -c '::' || true) +if [ "$_TMARG_PL_FOUND" -ne "$_TMARG_PL_EXPECTED" ]; then + echo "peak-local gate: collected $_TMARG_PL_FOUND tests, expected $_TMARG_PL_EXPECTED" >&2 + exit 1 +fi +# SKIP guard, IDENTIFYING rather than counting, for the reasons the band-limited gate +# above gives: a compensating pair leaves the total unchanged, and the expected total is +# a property of the RUNNER (a GPU-equipped runner that does not set RIFT_CI_REQUIRE_GPU=1 +# legitimately stops skipping, and a count guard then fails a good run). Allow skips +# whose REASON names cupy/GPU, and fail on any other skip whatever the total. +_TMARG_PL_OUT=$(python -m pytest -q -rs "$_TMARG_PL_TESTS" 2>&1) || { echo "$_TMARG_PL_OUT"; exit 1; } +echo "$_TMARG_PL_OUT" | tail -20 +_TMARG_PL_BAD=$(echo "$_TMARG_PL_OUT" | grep -E '^SKIPPED' | grep -vciE 'cupy|gpu|cuda' || true) +_TMARG_PL_BAD=${_TMARG_PL_BAD:-0} +if [ "$_TMARG_PL_BAD" -ne 0 ]; then + echo "peak-local gate: $_TMARG_PL_BAD test(s) skipped for a reason other than an absent GPU:" >&2 + echo "$_TMARG_PL_OUT" | grep -E '^SKIPPED' | grep -viE 'cupy|gpu|cuda' >&2 + exit 1 +fi + +# Joint (phi,psi) peak-local angle marginalization, numpy reference kernel. Gated here +# because test/ has no manifest check of its own: an unlisted test file is simply never +# run, which is the failure test-jax.sh exists to prevent one level up. What this has to +# protect: that the outside supremum is CERTIFIED (a straddling cell must count as +# outside -- classifying grid centres once returned "nothing uncovered" and accepted +# unconditionally), that a distance node is only dropped when the drop is provable +# against the computed value, and that an undersized region is DECLINED rather than +# returned. +_JOINT_PL_TESTS=MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py +# Raise EXPECTED by RUNNING collection, never by arithmetic. +_JOINT_PL_EXPECTED=26 +_JOINT_PL_FOUND=$(python -m pytest -q --collect-only "$_JOINT_PL_TESTS" 2>/dev/null | grep -c '::' || true) +if [ "$_JOINT_PL_FOUND" -ne "$_JOINT_PL_EXPECTED" ]; then + echo "joint peak-local gate: collected $_JOINT_PL_FOUND tests, expected $_JOINT_PL_EXPECTED" >&2 + exit 1 +fi +python -m pytest -q "$_JOINT_PL_TESTS" + python MonteCarloMarginalizeCode/Code/test/test_mcsamplerEnsemble_extended.py --as-test --n-max 100000 python MonteCarloMarginalizeCode/Code/test/test_mcsamplerEnsemble_extended.py --as-test --n-max 100000 --use-lnL diff --git a/.travis/test-jax.sh b/.travis/test-jax.sh index c84ea04c4..96b4efe5a 100755 --- a/.travis/test-jax.sh +++ b/.travis/test-jax.sh @@ -65,6 +65,78 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # test_network_coords.py 1 network-frame sky fold on a real injection # test_nuts_phimarg.py 1 fisher_nuts_sample_phimarg vs an analytic 4-D # target (needs numpyro; no lal) +# test_jax_fairdraw_export.py 34 the --save-samples export contract of +# bin/integrate_likelihood_extrinsic_jax: +# that it is a FAIR DRAW (reweighted against +# the sampler's own importance weights, then +# multinomial-resampled as ILE does), that +# ILE's 1.5*ESS cap binds, that the count +# options act exactly where the driver reports +# them implemented and nowhere else, that the +# export RNG is never the science generator, +# that a fair draw which CANNOT be performed +# exports nothing at all (and clears a stale +# file at that path) instead of shipping the +# raw cloud, that a refused export leaves no +# `_.dat` result row for the event either, that +# an SMC ladder which stops short of inv_T=1 +# publishes neither artifact (and that the +# sampler reports the exponent it reached), and +# that the provenance header +# describes the file it sits on. Needs no lal or GPU: the +# driver is imported by path and driven on an +# analytic 4-D target with known moments. +# Several of these are AST guards on the +# DRIVER SOURCE (the F1 post_weight gate, the +# write_samples call site, the export-before- +# result write order) because the defects +# they pin live at call sites, where a +# helper-level assertion cannot see them. +# test_jax_tempering_chooser.py 45 the --adapt-weight-exponent chooser and the +# tempering-cost law +# ESS/N = [beta(2-beta)]^(dim/2) it rests on. +# Pins the law against the EXACT sweep measured +# on the real BNS likelihood (both directions: +# a law that under-predicts the cost would +# silently under-budget a run), that it takes +# no SNR argument -- the non-JAX helper's rule +# keys on SNR and does not transfer -- that the +# chooser RETURNS the exponent rather than +# writing it back to opts (writing it made +# event 1 of a batch read event 0's choice), +# and that the degenerate-export branch WARNS +# (it used to raise; the calibrated estimate +# cannot support a hard floor -- see +# DESIGN_jax_tempering.md 4a). Includes a +# RETIRED-claim +# guard asserting the SNR rule has not crept +# back into the driver. Needs no lal, no GPU +# and no flowMC. +# test_jax_stencil_parity.py 23 #193: the JAX 'sinc' gatherer is the SAME +# stencil as the numpy/cupy/CUDA paths. Those +# three share one weight array and cannot drift; +# JAX re-expresses the formula independently +# (its weights depend on the traced sub-sample +# offset), so this is what converts that +# duplication from "trust the reviewer" into +# "CI fails". Landed WITHOUT a manifest entry, +# which made rift_O4d fail its own manifest +# check; added here. 24 tests, of which the +# cupy leg is deselected on this CPU runner -- +# see DESELECTED_TESTS -- so 23 are gated. +# test_flow_reuse_default.py 7 flow re-use is OFF by default, and --flow-reuse +# still reaches the old behaviour. A store_true +# flag cannot express its own negation, so simply +# flipping default=True would have made +# --no-flow-reuse inert AND deleted the capability; +# both directions and last-one-wins are pinned, as +# is the batch loop still reading the flag. +# test_interp_choices.py 3 #190: --interp cubic is reachable from the +# CLI and selects _gather_cubic. Merged in +# from rift_O4d, which added it to FILES +# without a count in this ledger; recorded +# here so the block stays a complete +# accounting of EXPECTED_TESTS. # test_tvals_grid_convention.py 13 issue #146: the time-marginalization window # grid the JAX wrapper and # bin/integrate_likelihood_extrinsic_batchmode @@ -75,9 +147,103 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # wrapper against the production driver, and # because 16384 is the rate test_jax_endtoend # (4096) structurally cannot cover. +# test_angle_marg_smoke.py 8 CHEAP mutation-bearing floor for the whole +# angle-marg feature: scheme selection (a +# previous head could never return 'exact'), +# both dense-sizing levers, required +# amp_sizing, the host failsafe record and +# its cond-guard, the driver AST guard on the +# VALUE node (hardcoding angle_marg="grid" +# passes a weaker guard), and that BOTH +# artifacts are labelled and never imply +# verification. Seconds, not minutes. +# test_angle_marg_compile_cost.py 6 the laplace path's COMPILE- and RUN-cost +# structure (2026-08-28: an unrolled kernel +# x 64 distance blocks put a production +# SNR-40 run >88 min / 22 GiB into XLA +# compilation; the fix then exposed a +# 36.41 GiB RESOURCE_EXHAUSTED at the +# default eval chunk). Trace-only where +# possible: the traced graph must not grow +# with the distance grid, the kernel must +# stay rolled (equation-count ceiling), the +# distance tail padding must be exactly- +# zero-weight, and the anglemarg eval-chunk +# cap must stay WIRED in samplers and the +# driver. Each fails under a verified +# mutation (see the PR). Seconds. +# test_angle_marg_block_dispatch.py 4 the laplace path's EXECUTION-cost +# structure (2026-08-28: with compilation +# fixed, the kernel executed ~2,950x the +# grid scheme because BOTH blend branches +# ran at every lattice point, the 320-pt +# quadrature everywhere included the 99.5% +# of points needing N ~ 32-96). Pins the +# lax.switch block dispatch: the N ladder +# keeps the shipped aliasing exponent, the +# dispatcher matches the undispatched +# kernel in every branch, and the fused +# driver actually CALLS the dispatcher +# (wiring). Each fails under a verified +# mutation (see the PR). Seconds. +# test_distance_grid_loguniform.py 33 the OPT-IN log-uniform ("peak-resolving") +# distance quadrature for the dense +# angle-marg schemes. Pins the spacing +# contract (Delta ln d <= c/rho_max), the +# TWO-SIDED calibration of c against the +# Gaussian trapezoid error law it is +# derived from (a one-sided check is +# satisfied by c -> 0, which is accurate +# and arbitrarily expensive), that +# --distance-grid-scheme still DEFAULTS to +# the historical uniform grid node for +# node, and -- the safety property -- that +# the dense angle lattice is sized from the +# amplitude on the FULL prior support, so +# no distance grid can shrink it. Includes +# driver AST guards on the option VALUE +# node, on the forwarded (not hardcoded) +# keyword, and on the fail-closed refusal +# when the flag is set on a mode that does +# not implement it. One numerical +# execution test against a 1024-node +# uniform reference; the rest are numpy or +# AST. ~19 s. Each fails under a verified +# mutation (matrix in the PR). +# test_angle_marg_sizing_rule.py 1 the m_max-aware dense phi sizing rule. # -# DELIBERATELY EXCLUDED (measured on ldas-pcdev11, JAX_PLATFORMS=cpu, OMP_NUM_THREADS=1): -# +# test_angle_marg_gh_laplace.py 15 the psi-marginal distance-node placement +# that lets 'laplace' honour +# JAX_ILE_DISTMARG_GH. GATED despite +# costing ~3.5 min: it is a NEW numerical +# path, and every constant in it is +# pinned by MUTATION (collapse the +# half-span, drop the node floor, force +# the sigma cap) rather than by a +# pass-through assertion. The two +# agreement legs (converged uniform grid, +# exact scheme under the same quadrature) +# are the expensive ones; they are also +# the only ones that would catch a wiring +# error, so they stay. +# Pure numpy, milliseconds, closed-form I0 +# reference. FAILS under the old m_max-blind +# rule (0.498 nats vs 1.17e-10), which every +# low-scale brute-force test passes -- so this +# is the only gated check that distinguishes +# the corrected sizing. The rest of the +# angle-marg suite is EXCLUDED; see below. +# test_limit_distance_jax.py 21 --limit-distance on this arm: the distance +# QUADRATURE narrows while the prior keeps its +# [d_min,d_max] normalization. Includes the +# bitwise no-op of the default call (both the +# uniform and the adaptive grid), the ACCEPTANCE +# comparison (narrowed vs full-range lnZ at equal +# n_grid: 0.0 nats, measured 2.8e-14), and its +# power check -- the pre-change call signature on +# the same box moves lnZ by +4.16 nats. ~110 s, +# CPU, one synthetic precompute. + # test_nuts_phimarg_injection.py Not a pytest file at all: it runs the whole study at # module scope and calls sys.exit() there. WITHOUT numpyro # that surfaces as a fast COLLECTION ERROR; WITH numpyro -- @@ -101,7 +267,63 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # demo_*.py, debug_*.py, Demos, debugging scripts and a figure generator, not # benchmark_snr_sequence.py, assertions. None defines a test_* function and none # make_3g_figdata.py is intended as a gate. +# test_jax_time_quadrature.py 12 band-limited time marginalization. The +# stock path integrates exp(lnL_t) with fixed +# Simpson weights at the DATA spacing while the +# integrand width sigma_t = 1/(2 pi rho sigma_f) +# SHRINKS with SNR -- 61.2 us against grid +# spacings of 244/122/61 us at srate +# 4096/8192/16384 on a 35+30 HLV injection at +# rho=40. Simpson is not a safeguard: it is +# (4 T_h - T_2h)/3, so it carries the coarser +# T_2h alias and is WORSE than trapezoid when +# under-resolved. Pins that upsampling is EXACT +# (sampling theorem, not an approximation), that +# the Nyquist bin is split rather than dumped, +# that the band-limited result is grid-phase +# INDEPENDENT where stock Simpson swings 4.26 +# nats, convergence in the free upsample factor, +# and that an unknown time_quad RAISES instead of +# silently giving the old behaviour. Also pins +# the two ways the reconstruction can silently +# change the ANSWER rather than the resolution: +# that it integrates the ORIGINAL (n-1)*deltaT +# window and not the periodic FFT continuation +# past the last sample (a constant integrand +# makes that a pure normalization shift), and +# that it REFUSES data whose depends on +# arrival time (the slow-rotation post-phase), +# where holding the norm at one bin would be a +# different likelihood. And the GUARD +# SAMPLES: the window is a CROP, so its +# ends do not join, and the FFT's +# periodic seam rings into the inserted +# samples while every retained sample +# stays exact -- invisible to an +# exactness test built from whole-period +# modes. Pins the defect on a +# non-periodic cropped tone, its removal +# by guard samples, that the guard is +# support and never enters the integral, +# that it has no default, and that the +# band-limited path actually widens the +# accumulation window (while Simpson +# does not). Pure numpy +# and jax, no lal, no GPU. +# test_jax_terminal_time_marginalization.py +# 18 adaptive primitive-field integration: +# odd/even reflection, exact normalization, +# Event-B high-SNR convergence, AD, bounded +# batch-independent dispatch, a near-Nyquist +# phase-marginalization counterexample, explicit +# nonlinear-endpoint refusal, driver wiring, and +# honest phase-marginalized sky/psi export, +# K=14/K=88 independent guarded references, +# and executable baseline/banded support refusal. + FILES=( + "${JAXDIR}/test_jax_time_quadrature.py" + "${JAXDIR}/test_jax_terminal_time_marginalization.py" "${JAXDIR}/test_jax_likelihood.py" "${JAXDIR}/test_jax_endtoend.py" "${JAXDIR}/test_jax_slowrot_coeffs.py" @@ -110,18 +332,77 @@ FILES=( "${JAXDIR}/test_jax_slowrot_cauchy_schwarz.py" "${JAXDIR}/test_network_coords.py" "${JAXDIR}/test_nuts_phimarg.py" + "${JAXDIR}/test_jax_fairdraw_export.py" + "${JAXDIR}/test_jax_tempering_chooser.py" "${JAXDIR}/test_tvals_grid_convention.py" + "${JAXDIR}/test_interp_choices.py" + "${JAXDIR}/test_jax_stencil_parity.py" + "${JAXDIR}/test_flow_reuse_default.py" + "${JAXDIR}/test_angle_marg_sizing_rule.py" + "${JAXDIR}/test_angle_marg_smoke.py" + "${JAXDIR}/test_angle_marg_compile_cost.py" + "${JAXDIR}/test_angle_marg_block_dispatch.py" + "${JAXDIR}/test_distance_grid_loguniform.py" + "${JAXDIR}/test_angle_marg_gh_laplace.py" + "${JAXDIR}/test_angle_marg_default.py" + "${JAXDIR}/test_angle_marg_gh_selection.py" + "${JAXDIR}/test_joint_anglemarg_peaklocal.py" + "${JAXDIR}/test_angle_marg_peaklocal_wiring.py" + "${JAXDIR}/test_limit_distance_jax.py" ) # EXCLUDED: files in JAXDIR matching test_*.py that are deliberately NOT gated. The # manifest check below fails if a file is in neither FILES nor EXCLUDED, so adding a new # test_*.py to test/jax/ forces a decision instead of being silently unrun -- which is # this gate's own failure mode, one level up. +DESELECTED_TESTS=( + "${JAXDIR}/test_jax_stencil_parity.py::test_gpu_gather_parity_against_numpy_window" +) EXCLUDED=( + # test_angle_marg_exact.py -- the angle-marginalization VALIDATION suite. + # + # NOT gated per-PR, and this is a deliberate, measured decision rather than a + # convenience. It is a development check in the same sense that full RIFT + # analysis runs are: it establishes the schemes' ERROR LAW at production + # amplitude, a property of the mathematics that does not change commit to + # commit. Three separate CI failures forced the split, each a different + # symptom of the same cost: the 169-test gate was CANCELLED at the job's + # 60-minute cap; a later head OOM-killed the runner at 19 min; and the run + # after that reached 83% and then died with "the runner has received a + # shutdown signal" (exit 143). The 139-test baseline ran in 13m53s. + # + # What remains GATED is the coverage that actually bites: + # test_angle_marg_sizing_rule.py pins the m_max-aware dense sizing with a + # pure-numpy, millisecond test against a closed-form I0 reference, and FAILS + # under the old m_max-blind rule (0.498 nats vs 1.17e-10). The low-scale + # brute-force comparisons in the excluded file prove exactness but do NOT + # distinguish the sizing rule -- the broken rule passes them all -- which is + # why extracting that one test was necessary before excluding the rest. + # + # RUN IT BY HAND when touching anglemarg.py, on a quiet host with >=16 cores: + # PYTHONPATH=/MonteCarloMarginalizeCode/Code JAX_PLATFORMS=cpu \ + # JAX_ENABLE_X64=1 OMP_NUM_THREADS=1 JAX_COMPILATION_CACHE_DIR="" \ + # taskset -c 0-15 python -m pytest -q \ + # /MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py + # and record the numbers in the PR, per records-protocol. + "${JAXDIR}/test_angle_marg_exact.py" "${JAXDIR}/test_nuts_phimarg_injection.py" "${JAXDIR}/test_flow_reuse.py" ) +# DESELECT: individual tests inside a GATED file that cannot run on this CPU runner. +# File-level EXCLUDED is too blunt for these -- dropping test_jax_stencil_parity.py to +# silence its one GPU leg would also drop the 23 CPU tests that are the whole point of +# #193. Deselecting is not the same as tolerating a skip: a skip leaves the gate green +# while asserting nothing, whereas a deselected test is accounted for HERE, in writing. +# +# test_jax_stencil_parity.py::test_gpu_gather_parity_against_numpy_window +# The cupy leg of the sinc-stencil parity check. It needs a real CUDA device; +# this job has none, so it self-skips. It is a genuine gate on a GPU host -- +# run it by hand there when touching Q_inner_product_sinc_cupy. +DESELECT=() +for t in "${DESELECTED_TESTS[@]}"; do DESELECT+=( --deselect "$t" ); done + echo "== manifest check (every test_*.py is gated or explicitly excluded) ==" manifest_rc=0 for f in "${JAXDIR}"/test_*.py; do @@ -139,12 +420,67 @@ if [ "${manifest_rc}" -ne 0 ]; then exit 1 fi -# Sum of the per-file counts above. Pinned deliberately: a bare `pytest test/jax/` +# Sum of the per-file counts above. +# Pinned deliberately: a bare `pytest test/jax/` # that collected 0 would exit 5, and a partial loss (say 14 -> 3) would still exit 0. -EXPECTED_TESTS=27 +# NOTE: this environment collects ONE MORE test than the CI runner does (local +# 147 vs CI 146; the delta was 3 earlier in this branch's life). So "recount by +# collection" must mean collection IN THE GATE'S ENVIRONMENT -- a local count has +# tripped this floor twice. When in doubt, take the number from a CI log line +# ("collected N tests from M files") rather than from your shell. +# Raised 153 -> 155 by the two new test_jax_time_quadrature.py pins (original +# integration window; refusal of arrival-time-dependent norms), then 155 -> 160 by +# the five guard-sample pins in the same file (periodic-seam defect on a +# non-periodic crop; its removal by guard samples; guard is support, not window; +# no default guard; the band-limited path widens the accumulation window). +# PR #209 then adds six test_angle_marg_compile_cost.py pins, raising 160 -> 166, +# and PR #210 adds five test_angle_marg_block_dispatch.py pins, raising 166 -> 171. +# PR #216 adds eighteen adaptive primitive-time pins, raising 171 -> 189. +# The psi-marginal GH placement (#225) adds 36, raising 189 -> 225: 15 in +# test_angle_marg_gh_laplace.py, 5 in test_angle_marg_default.py, 8 in +# test_angle_marg_gh_selection.py, plus 4 answering external review on the +# identity gate (imaginary-A0 coefficient, B1 in the conjugate slice, the gate +# applying to an explicit laplace, the kernel guard staying trace-safe). +# The log-uniform distance quadrature adds 37 on top of the base, raising +# it by that amount wherever the base then sat: 30 for the scheme itself, 3 from external re-review (the +# zero-clipped-amplitude extreme of the F1 detector, the DRIVER half of the F2 +# refusal, and a guard on the sky-doubling path -- each because a mutation +# SURVIVED the 33-mutation matrix without it), and 3 covering the truncated- +# endpoint precondition added by the automated review pass, which shipped with +# none (the estimator against an independent numpy measurement, the +# interior-but-too-close refusal, and the non-positive-clearance window that +# built at 6.7x tol). +# Raising the floor +# by exactly the number of tests ADDED is safe whatever the environment delta above, +# since it preserves the margin the previous floor already had. +# Raised 189 -> 217 by the 28 tests added in the angle-marg GH branch (#225); then +# 217 -> 225 by the tests answering external review on its identity gate; then +# 225 -> 234 while this branch was open, by the peak-local framework (#224), the +# joint (phi,psi) peak-local kernel (#230) and the AV batch-max change (#234); +# then 234 -> 272 by the adaptive distance quadrature branch's own 37 (#221). +# Raised 272 -> 293 on merging --limit-distance, by that branch's 21 +# test_limit_distance_jax.py pins: 15 shipped; 4 added after an adversarial +# mutation sweep found two INERT guards (the adaptive grid's d_prior_range branch +# could be deleted with the suite green, worth +5.5 nats silently, and +# sample_prior could draw over the full range while the box correction stayed); +# and 2 more for this merge's own interactions with the #221 distance quadrature +# -- the box must not shrink the angle lattice, and the log-uniform grid must +# refuse a box rather than renormalize the prior onto it. +# THREE branches have now raised this constant, so it is the single place this +# merge is most likely to go quietly wrong; the FILES array above is the other. +# Taken from a collection RUN, never by adding the three accountings. +# +# The peak-local ILE WIRING branch adds 13: 11 in +# test_angle_marg_peaklocal_wiring.py (the scheme reaches the likelihood, +# matches exact, is absent from 'auto', joins the amp failsafe / batch-memory +# cap / artifact label, and the CLI rejects a misspelling) and 2 in +# test_joint_anglemarg_peaklocal.py (twice differentiable, and the gradient stays +# finite as the quartic leading coefficient vanishes). 293 + 13 = 306, re-derived +# by RUNNING the gate's own collection after rebasing over #221/#238/#223. +EXPECTED_TESTS=306 echo "== collection floor check (expect >= ${EXPECTED_TESTS} tests) ==" -collect_out="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${FILES[@]}" 2>&1)" +collect_out="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${DESELECT[@]}" "${FILES[@]}" 2>&1)" collect_rc=$? if [ "${collect_rc}" -ne 0 ]; then printf '%s\n' "${collect_out}" @@ -165,11 +501,32 @@ if [ "${n_collected}" -lt "${EXPECTED_TESTS}" ]; then exit 1 fi +# A --deselect whose nodeid does not resolve is SILENTLY IGNORED by pytest: rename the +# test, or fat-finger the path, and the deselect quietly stops applying while this script +# still claims the test is accounted for. The skip would then come back and the count +# would be off by one, which is exactly the confusion the deselect was added to end. So +# verify both halves: the test still EXISTS under that name, and it is actually GONE from +# the collection. +for t in "${DESELECTED_TESTS[@]}"; do + f="${t%%::*}"; nm="${t##*::}" + if [ ! -f "${f}" ]; then + echo "test-jax.sh: DESELECTED_TESTS names ${f}, which does not exist." >&2; exit 1 + fi + if ! grep -qE "^[[:space:]]*def ${nm}\\(" "${f}"; then + echo "test-jax.sh: DESELECTED_TESTS names ${nm}, which ${f} no longer defines." >&2 + echo " It was probably renamed. Update the nodeid, or drop it from DESELECTED_TESTS." >&2 + exit 1 + fi + if printf '%s\n' "${collect_out}" | grep -qE "^${f}::${nm}(\\[|$)"; then + echo "test-jax.sh: --deselect did not take effect for ${t}." >&2; exit 1 + fi +done + junit="$(mktemp -t jaxci-junit-XXXXXX.xml)" trap 'rm -f "${junit}"' EXIT echo "== running ==" -"${PYTHON_BIN}" -m pytest -q -p no:cacheprovider --durations=0 --junit-xml="${junit}" "${FILES[@]}" +"${PYTHON_BIN}" -m pytest -q -p no:cacheprovider --durations=0 --junit-xml="${junit}" "${DESELECT[@]}" "${FILES[@]}" rc=$? if [ "${rc}" -ne 0 ]; then # rc 5 == "no tests ran"; it is a FAILURE here, not a pass. @@ -196,7 +553,8 @@ if tests < expected: bad.append("ran %d tests, expected at least %d" % (tests, expected)) if skipped: bad.append("%d SKIPPED -- a skip silently disables a gate here; if a skip is " - "legitimate, exclude the file in FILES and say why" % skipped) + "legitimate, exclude the file in FILES -- or, for a single test, add it to " + "DESELECTED_TESTS -- and say why" % skipped) if failures or errors: bad.append("%d failures, %d errors" % (failures, errors)) if bad: diff --git a/.travis/test-lisa-gpu.sh b/.travis/test-lisa-gpu.sh new file mode 100755 index 000000000..a7832df9f --- /dev/null +++ b/.travis/test-lisa-gpu.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# LISA gate on a cupy-capable host. +# +# The LISA likelihood is host numpy, but it calls helpers in +# RIFT/likelihood/SphericalHarmonics_gpu.py whose backend used to default to cupy +# on any host where cupy merely IMPORTS -- no GPU work intended or wanted. That +# combination raised `TypeError: Unsupported type ` on +# every GPU host and passed on every cupy-free one, which is why neither the +# GitHub `lisa-check` job nor the GitLab `import_check` job (both cupy-free) ever +# saw it. +# +# So this script REFUSES to run without cupy rather than degrading to the CPU +# backend: a green LISA suite that silently only exercised numpy is the exact +# hole being closed. Run it on the GitLab `gpu` runner (see gpu_integration in +# .gitlab-ci.yml), or by hand on a GPU node. +set -euo pipefail +# Same interpreter resolution as .travis/test-lisa.sh, so the two LISA gates +# cannot silently run under different pythons. +PY="${RIFT_LISA_PYTHON:-${PYTHON:-python}}" +command -v "$PY" >/dev/null 2>&1 || PY="$(command -v python3)" +export OMP_NUM_THREADS=1 + +"$PY" - <<'PY' +try: + import cupy +except Exception as exc: + raise SystemExit(f"test-lisa-gpu.sh requires cupy, which did not import: {exc}") from exc + +try: + n_devices = cupy.cuda.runtime.getDeviceCount() +except Exception as exc: + raise SystemExit(f"test-lisa-gpu.sh requires a CUDA device; cupy could not query one: {exc}") from exc +if n_devices < 1: + raise SystemExit("test-lisa-gpu.sh requires a CUDA device; cupy reported zero") + +from RIFT.likelihood import SphericalHarmonics_gpu as sh + +# Assert this host is genuinely in the configuration that used to break. Without +# this the suite would report PASS on a cupy-free runner and prove nothing. +if not sh.cupy_here: + raise SystemExit( + "cupy imports, but RIFT.likelihood.SphericalHarmonics_gpu did not enable it " + "(cupy_here=False) -- the GPU dispatch path is not under test" + ) +if sh.xpy_default is not cupy: + raise SystemExit( + f"expected SphericalHarmonics_gpu.xpy_default to be cupy on this host, got {sh.xpy_default!r}" + ) + +print(f"LISA GPU preflight OK: cupy={cupy.__version__}, cuda_devices={n_devices}") +PY + +"$PY" -m pytest -q \ + MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py \ + MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py + +echo "LISA GPU-host gate: PASS" diff --git a/.travis/test-lisa.sh b/.travis/test-lisa.sh index f4f03adda..7028e40c9 100644 --- a/.travis/test-lisa.sh +++ b/.travis/test-lisa.sh @@ -23,4 +23,6 @@ fi MonteCarloMarginalizeCode/Code/test/test_lisa_use_lnL_branches.py \ MonteCarloMarginalizeCode/Code/test/test_lisa_portfolio_method_integrity.py \ MonteCarloMarginalizeCode/Code/test/test_lisa_mc_error_replicas.py \ - MonteCarloMarginalizeCode/Code/test/test_lisa_driver_drift.py + MonteCarloMarginalizeCode/Code/test/test_lisa_driver_drift.py \ + MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py \ + MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py diff --git a/.travis/test-q-window-stencil.sh b/.travis/test-q-window-stencil.sh new file mode 100755 index 000000000..8aad3b979 --- /dev/null +++ b/.travis/test-q-window-stencil.sh @@ -0,0 +1,308 @@ +#!/usr/bin/env bash +# CPU gate for the Q_lm sub-sample stencil / time-interpolation / log-sum-exp-offset +# tests, driven from ci.yml's q-window-stencil-check job. +# +# WHY THIS SCRIPT EXISTS. Until it landed, this job's membership was a hand-maintained +# backslash-continued file list inside ci.yml. Three PRs appended to that one list in +# two days and it conflicted twice; resolving such a conflict by taking one side rather +# than both silently UNREGISTERS the other side's test, and in this CI an unlisted test +# never runs and the job reports green forever. That is coverage that looks like a guard +# and is not. +# +# It is modelled on .travis/test-slowrot.sh and keeps that script's defences, with ONE +# deliberate difference. slowrot's membership is an explicit FILES array; here membership +# is declared by a MARKER LINE INSIDE EACH TEST FILE, so a PR adding a test to this area +# edits only its own new file and no shared line exists to conflict on. Everything else +# -- the fail-closed manifest, the pinned floors, the junit OUTCOME assertion, the +# EXCLUDED list with a reason per entry -- follows slowrot deliberately. +# +# Five defences. Do not simplify any of them into a bare `pytest `: +# +# 1. MARKER-BASED membership, not a directory glob. A glob would sweep up files that +# collect ZERO items, and pytest exits 5 on those -- "no tests ran" reads as a pass +# in a log skim. (test_slowrot_*.py in one of these very directories is five such +# files; they belong to slowrot-check, not here.) +# 2. A fail-closed MANIFEST over the filename patterns this job owns: a file matching +# one of them that carries no marker and is not explicitly EXCLUDED fails the job. +# This is what makes "added a test and forgot to register it" RED instead of silent. +# 3. A PER-FILE collection floor of 1. A registered file that collects nothing is the +# exit-5 trap arriving through the front door; it must be visible, not green. +# 4. A pinned TOTAL collection floor, so a renamed file or a dropped test_* entry point +# goes red instead of green-on-fewer-tests. +# 5. A hard fail on ANY nonzero pytest exit (5 included) plus a junit OUTCOME assertion +# on tests that actually PASSED. The collection floors count COLLECTION, which +# cannot see a test that collects, runs, and asserts nothing. +# +# Needs numpy + lal only: no GPU, no jax, no numpyro. +set -uo pipefail +# NOTE: deliberately no -e, matching test-slowrot.sh. Every command below has its rc +# handled explicitly so the failure messages stay specific; if you add one, guard it. + +# The paths below are repo-relative, so anchor cwd rather than trusting the caller. +cd "$(dirname "$0")/.." || { echo "test-q-window-stencil.sh: cannot cd to repo root" >&2; exit 1; } + +# INVARIANT: this gate always tests THIS CHECKOUT, never an installed build. Must +# PREPEND -- appending lets a caller's PYTHONPATH win. (test_interpolate_time_cli.py +# launches the RIFT scripts as real subprocesses, and they inherit this.) +export PYTHONPATH="$PWD/MonteCarloMarginalizeCode/Code${PYTHONPATH:+:$PYTHONPATH}" + +# --------------------------------------------------------------------------------- +# DISCOVERY RUNS FIRST, before the dependency probes below. +# +# Nothing here needs pytest, numpy or lal: finding the registered files is a grep. Doing +# it first is what lets this script answer RIFT_CI_GATE_LIST=1 (see below) for the +# repo-wide census in .travis/test-ci-roster.py, which runs as its own job with NO +# `needs: install` and so has none of those libraries. Keep this above the probes. +CODEDIR="MonteCarloMarginalizeCode/Code" + +# --------------------------------------------------------------------------------- +# MEMBERSHIP. A test file joins this gate by carrying this line, on its own, verbatim: +# +# # RIFT-CI-GATE: q-window-stencil +# +# The match is whole-line and fixed-string (grep -x -F), so prose mentioning the tag -- +# including the explanatory line the registered files put directly underneath it -- does +# NOT register a file. Search is limited to test_*.py under Code/, so a doc or a script +# quoting the tag cannot enrol itself either. +MARKER="# RIFT-CI-GATE: q-window-stencil" + +mapfile -t FILES < <(grep -rlxF --include='test_*.py' -- "${MARKER}" "${CODEDIR}" 2>/dev/null | LC_ALL=C sort) + +# LIST MODE. Print the files this gate would run, one per line, and stop. +# +# ci-roster-check honours marker-based membership ONLY for a gate whose own discovery it +# can execute. That is not fussiness: no text pattern distinguishes "uses the marker to +# find files" from "uses the marker", and this script contains both -- the mapfile above, +# and the `grep -qxF -- "${MARKER}"` further down that asserts an EXCLUDED file does NOT +# carry it. A census that pattern-matched would keep passing with the mapfile deleted. +# So it asks, and this is the answer. +if [ -n "${RIFT_CI_GATE_LIST:-}" ]; then printf '%s\n' "${FILES[@]}"; exit 0; fi + +PYTHON_BIN="${RIFT_QWINDOW_PYTHON:-${PYTHON:-python}}" +if ! command -v "${PYTHON_BIN}" >/dev/null 2>&1; then + PYTHON_BIN="$(command -v python3)" +fi + +# Guard the tool checks: a missing interpreter plus a redirected stderr is +# indistinguishable from a clean result. +"${PYTHON_BIN}" -c 'import pytest' || { echo "test-q-window-stencil.sh: pytest unavailable" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import numpy; print("numpy", numpy.__version__)' \ + || { echo "test-q-window-stencil.sh: numpy unavailable" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import lal, lalsimulation; print("lal", lal.__version__)' \ + || { echo "test-q-window-stencil.sh: lal/lalsimulation unavailable" >&2; exit 1; } + +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" +export MKL_NUM_THREADS="${MKL_NUM_THREADS:-1}" + + + +if [ "${#FILES[@]}" -eq 0 ]; then + echo "test-q-window-stencil.sh: no test file carries the marker line" >&2 + echo " Expected at least one file under ${CODEDIR} containing exactly:" >&2 + echo " ${MARKER}" >&2 + echo " Either the marker text was edited here without updating the files, or the" >&2 + echo " registrations were lost. This is a hard failure, not an empty run." >&2 + exit 1 +fi + +# --------------------------------------------------------------------------------- +# MANIFEST SCOPE. The filename patterns this job owns. Every file matching one of them +# must carry the marker or appear in EXCLUDED below, so a new test in this area forces a +# decision instead of being silently unrun -- this gate's own failure mode, one level up. +# +# The patterns are deliberately NOT "every test_*.py in these directories": Code/test/ +# holds ~90 files belonging to other jobs, and Code/RIFT/likelihood/ holds the +# test_slowrot_*.py suite, which is slowrot-check's business (issue #169). The cost of +# that narrowness is real and stated: a new stencil test filed under some OTHER prefix +# escapes the manifest. Name it to match one of these, or add a pattern here. +# +# Registration itself needs no edit to this file -- only the marker in the new test. +SCOPE_GLOBS=( + "${CODEDIR}/RIFT/likelihood/test_q_window_*.py" + "${CODEDIR}/RIFT/likelihood/test_time_interp_*.py" + "${CODEDIR}/RIFT/likelihood/test_interpolate_time_*.py" + "${CODEDIR}/RIFT/likelihood/test_calmarg_stencil_*.py" + "${CODEDIR}/RIFT/likelihood/test_batchmode_stencil_*.py" + "${CODEDIR}/RIFT/likelihood/test_noloop_*.py" + "${CODEDIR}/RIFT/misc/test_psd_bandwidth*.py" + "${CODEDIR}/test/test_noloop_*.py" + "${CODEDIR}/test/test_calmarg_running_max_*.py" +) + +# EXCLUDED, with the reason each is out. +# +# test_q_window_interp_gpu.py Need a GPU. On a CPU runner they report as SKIPPED +# test_noloop_gpu_stencils.py with exit 0, and the junit check below treats extra +# skips as a failure. Run by hand on a GPU node; the +# numbers are in PR #97. Same treatment as the GPU files +# in slowrot-check. +EXCLUDED=( + "${CODEDIR}/RIFT/likelihood/test_q_window_interp_gpu.py" + "${CODEDIR}/RIFT/likelihood/test_noloop_gpu_stencils.py" +) + +echo "== registered files (marker: ${MARKER}) ==" +printf ' %s\n' "${FILES[@]}" + +# A scope pattern that stops matching anything is a SILENT no-op: the manifest keeps +# passing while covering less. Assert each still matches at least one file, exactly as +# test-slowrot.sh asserts each DESELECT nodeid still resolves. +echo "== scope pattern check (every pattern still matches something) ==" +scope_rc=0 +for g in "${SCOPE_GLOBS[@]}"; do + # shellcheck disable=SC2206 + matches=( ${g} ) + if [ ! -e "${matches[0]}" ]; then + echo "test-q-window-stencil.sh: scope pattern ${g} matches no file." >&2 + echo " Those tests were renamed or removed. Fix the pattern or drop it; left as" >&2 + echo " is, it silently covers nothing." >&2 + scope_rc=1 + fi +done +[ "${scope_rc}" -eq 0 ] || exit 1 + +# An EXCLUDED entry for a file that no longer exists is the same silent no-op, and an +# EXCLUDED file that ALSO carries the marker is a contradiction that would run it anyway. +echo "== exclusion check ==" +excl_rc=0 +for e in "${EXCLUDED[@]}"; do + if [ ! -f "${e}" ]; then + echo "test-q-window-stencil.sh: EXCLUDED entry ${e} does not exist." >&2 + echo " It was renamed or removed; drop it from EXCLUDED or fix the path." >&2 + excl_rc=1 + continue + fi + if grep -qxF -- "${MARKER}" "${e}"; then + echo "test-q-window-stencil.sh: ${e} is EXCLUDED but carries the marker." >&2 + echo " Remove the marker, or remove the file from EXCLUDED. As it stands the" >&2 + echo " exclusion's stated reason does not describe what this gate does." >&2 + excl_rc=1 + fi +done +[ "${excl_rc}" -eq 0 ] || exit 1 + +echo "== manifest check (every in-scope test file is registered or explicitly excluded) ==" +manifest_rc=0 +for g in "${SCOPE_GLOBS[@]}"; do + for f in ${g}; do + [ -f "${f}" ] || continue + known=0 + for k in "${FILES[@]}" "${EXCLUDED[@]}"; do + [ "${f}" = "${k}" ] && { known=1; break; } + done + if [ "${known}" -eq 0 ]; then + echo "test-q-window-stencil.sh: ${f} is neither registered nor explicitly excluded." >&2 + manifest_rc=1 + fi + done +done +if [ "${manifest_rc}" -ne 0 ]; then + echo " Register it by adding this line, on its own, near the top of that file:" >&2 + echo " ${MARKER}" >&2 + echo " and raise EXPECTED_TESTS / EXPECTED_PASSED below. If it must NOT run here" >&2 + echo " (needs a GPU, is a print-only study, belongs to another job), add it to" >&2 + echo " EXCLUDED in this script WITH A REASON." >&2 + exit 1 +fi + +# --------------------------------------------------------------------------------- +# The pinned floors, as of this commit. Re-derive both after adding or removing a test: +# EXPECTED_TESTS `pytest --collect-only -q` over the registered files. +# EXPECTED_PASSED the "N passed" from a full run (tests minus skips). +# Never lower either without saying why in the commit message. +EXPECTED_TESTS=69 +EXPECTED_PASSED=67 + +# The only legitimate skips here are the two cupy legs -- one in +# test_noloop_time_marg_row_offset.py, one in test_calmarg_running_max_row_offset.py -- +# which pytest.importorskip's away on these GPU-less runners. A THIRD skip means a gate +# was disabled, which is the exact shape this script exists to prevent, so cap it rather +# than letting skips absorb losses silently. +MAX_SKIPS=2 + +# PER-FILE collection floor. A registered file that collects nothing contributes zero +# gates while looking like membership; on its own pytest would exit 5 on it, and inside a +# multi-file run that exit code never appears at all. +echo "== per-file collection check (each registered file must collect >= 1 test) ==" +perfile_rc=0 +n_collected=0 +for f in "${FILES[@]}"; do + out="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${f}" 2>&1)" + rc=$? + if [ "${rc}" -ne 0 ] && [ "${rc}" -ne 5 ]; then + printf '%s\n' "${out}" + echo "test-q-window-stencil.sh: collection of ${f} failed (exit ${rc})" >&2 + exit 1 + fi + # Anchor to '.py::' at line start. An unanchored grep -c '::' also counts merged + # stderr and warning text, and because these are >= tests, OVER-counting is the + # dangerous direction: one stray line masks exactly one lost test. + n="$(printf '%s\n' "${out}" | grep -cE '^[^[:space:]]+\.py::')" + printf ' %3d %s\n' "${n}" "${f}" + if [ "${n}" -eq 0 ]; then + echo "test-q-window-stencil.sh: ${f} carries the marker but collects 0 tests." >&2 + perfile_rc=1 + fi + n_collected=$(( n_collected + n )) +done +if [ "${perfile_rc}" -ne 0 ]; then + echo " pytest exits 5 (\"no tests ran\") on such a file when run alone, and inside a" >&2 + echo " multi-file run it contributes nothing while looking registered. Give it" >&2 + echo " test_* functions, or remove the marker and add it to EXCLUDED with a reason." >&2 + exit 1 +fi + +echo "== collection floor check (expect >= ${EXPECTED_TESTS} tests) ==" +echo "collected ${n_collected} tests from ${#FILES[@]} files" +if [ "${n_collected}" -lt "${EXPECTED_TESTS}" ]; then + echo "test-q-window-stencil.sh: collected ${n_collected} tests, expected at least ${EXPECTED_TESTS}." >&2 + echo " A file was renamed/moved, or a test_* entry point was dropped and pytest is" >&2 + echo " now passing on fewer tests than this gate promises. Fix the file, or update" >&2 + echo " EXPECTED_TESTS in this script and say why." >&2 + exit 1 +fi + +junit="$(mktemp -t qwindowci-junit-XXXXXX.xml)" || { echo "test-q-window-stencil.sh: mktemp failed" >&2; exit 1; } +trap 'rm -f "${junit}"' EXIT + +echo "== pytest ==" +"${PYTHON_BIN}" -m pytest -q -p no:cacheprovider --durations=10 --junit-xml="${junit}" "${FILES[@]}" +rc=$? +if [ "${rc}" -ne 0 ]; then + # rc 5 == "no tests ran"; it is a FAILURE here, not a pass. + echo "test-q-window-stencil.sh: pytest exited ${rc}" >&2 + exit "${rc}" +fi + +# OUTCOME check. The floors above count COLLECTION, which cannot see a test that +# collects, runs, and asserts nothing: one pytest.skip() or importorskip() disables a gate +# while both the collected count and the pytest exit status stay green. So assert what +# the RUN did, in PASSED tests, not in collected ones. +"${PYTHON_BIN}" - "${junit}" "${EXPECTED_PASSED}" "${MAX_SKIPS}" <<'PYCHECK' +import sys, xml.etree.ElementTree as ET +path, expected, max_skips = sys.argv[1], int(sys.argv[2]), int(sys.argv[3]) +root = ET.parse(path).getroot() +ts = root if root.tag == "testsuite" else root.find("testsuite") +if ts is None: + sys.stderr.write("test-q-window-stencil.sh: no in the junit report\n"); sys.exit(1) +g = lambda k: int(ts.get(k, 0) or 0) +tests, skipped, failures, errors = g("tests"), g("skipped"), g("failures"), g("errors") +passed = tests - skipped - failures - errors +print("junit: tests=%d passed=%d skipped=%d failures=%d errors=%d" + % (tests, passed, skipped, failures, errors)) +bad = [] +if failures or errors: + bad.append("%d failures, %d errors" % (failures, errors)) +if passed < expected: + bad.append("only %d tests PASSED, expected at least %d -- tests were lost, not just " + "reported differently" % (passed, expected)) +if skipped > max_skips: + bad.append("%d SKIPPED, at most %d expected (the cupy legs) -- a skip silently " + "disables a gate here; if a new skip is legitimate, raise MAX_SKIPS and " + "say which test and why" % (skipped, max_skips)) +if bad: + sys.stderr.write("test-q-window-stencil.sh: " + "; ".join(bad) + "\n"); sys.exit(1) +PYCHECK +if [ $? -ne 0 ]; then exit 1; fi + +echo "q-window stencil gate: PASS (${#FILES[@]} registered files, ${n_collected} tests collected)" diff --git a/CHANGES.rst b/CHANGES.rst index 9b8aeea1d..1745cd2c7 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -2,6 +2,47 @@ 0.0.18.0 ------------ development tree is rift_O4d. + +** BEHAVIOUR CHANGE, jax ILE: flow re-use across ``--n-events-to-analyze`` is now + OFF by default. ``--flow-reuse`` restores the old behaviour; ``--no-flow-reuse`` + is kept and now restates the default, so existing command lines keep working. + Measured over an 8-event batch at two seeds: re-using the trained flow contracts + the extrinsic posterior monotonically in slot index -- psi to ~40% of its + no-re-use width by slot 7, on BOTH seeds, with slot 0 (where no re-use has yet + happened) at ~1.0 as a control -- while costing no measurable wall time (1589 s + mean with re-use, 1567 s without; the seed-to-seed spread is larger than the + difference and its sign flips). Anyone relying on re-use to amortize a batch + must now pass ``--flow-reuse`` explicitly, and should not do so for any run whose + extrinsic SAMPLES are used. Evidence is less affected than samples. + +** jax ILE ``--save-samples`` is now a FAIR DRAW. Previously the driver wrote + whatever cloud the sampler produced, with no weight column, so every consumer + read a Gaussian-proposal cloud (``--mode laplace-is``, the default) or raw + PRIOR draws (``--mode prior-mc``) as if it were a posterior. Each estimator + now returns its per-sample importance weight and the export is + multinomial-resampled against it, matching production ILE's convention + (``RIFT/integrators/mcsampler.py::integrate``). The data columns and their + header line are unchanged for a given ``--mode``; a second header line now + records the mode and the export ESS. ``--fairdraw-extrinsic-output``, + ``--fairdraw-extrinsic-output-n-max`` and ``--n-fairdraw-extrinsic-samples`` + are implemented (gated per mode, and honoured as a COUNT contract even when + the weights are uniform). The requested count is clamped ONLY by the rows + available; those already carry ILE's ``1.5*ESS`` cap wherever the export + weights were non-uniform, applied against the EXPORT weights. There is + deliberately no second clamp by the evidence estimator's ``neff``: on the + flowMC modes and on ``laplace-is`` that number describes a separate cloud + (the moment-matched Gaussian evidence proposal, or the annealing ladder's + minimum rung ESS), and clamping by it truncated valid equal-weight chains. NOTE ``--fairdraw-extrinsic-output-n-max`` + defaults to 5, as in ILE, so passing ``--fairdraw-extrinsic-output`` without + an explicit maximum now yields 5 rows where it previously yielded the whole + cloud. If the weights admit no fair draw at all (degenerate or + unnormalizable), the driver writes NO ``*_samples.dat`` -- and deletes a stale + one at that path -- and fails the event (``--soft-fail-event-range`` still + skips to the next one) rather than exporting an unreweighted cloud under the + name that means "posterior draws". That check runs BEFORE the + ``__.dat`` result row is written, so such an event leaves no + normal ILE result behind either (a stale row is likewise removed); otherwise a + soft-failed, collapsed integration was still collectable as a success. - (rc0) O4d base refresh, from rift_O4c to rift_O4d: Python/numpy CI modernization (py3.10-py3.13, numpy 2.x checks), Asimov/RIFT smoke tests, docs deployment, pluggable workflow backends and simulation-manager prototypes, distance-grid/distance-slice likelihood export, container-family and pixi/SWIG @@ -214,10 +255,6 @@ development tree is rift_O4d. certify correctness; k-hat does not catch confidently-wrong runs from support mismatch; the L0 'doubles landed fraction' claim and the cap24 lnZ-bias claim are retracted). -0.0.17.13 ---------- -MR https://git.ligo.org/rapidpe-rift/rift/-/merge_requests/55 , for ln(e) parameter access in pipeline - 0.0.17.12 --------- MR https://git.ligo.org/rapidpe-rift/rift/-/merge_requests/54 , for eccentricity prior (log-uniform) diff --git a/INSTALL_OPTIONAL_DEPENDENCIES.md b/INSTALL_OPTIONAL_DEPENDENCIES.md index 46d53882c..b3f5d262b 100644 --- a/INSTALL_OPTIONAL_DEPENDENCIES.md +++ b/INSTALL_OPTIONAL_DEPENDENCIES.md @@ -3,11 +3,27 @@ ## TEOBResumS -TEOBResumS is available from [git source](https://bitbucket.org/eob_ihes/teobresums/src) and via pypi. We recommend you install via +TEOBResumS is available from [git source](https://bitbucket.org/eob_ihes/teobresums/src) and via pypi. For the standard non-hyperbolic interface, install via ``` pip install teobresums ``` -which provides the ``EOBRun_module``. However, we have noticed some incompatibilties with numpy can arise if this is done. If needed, please instead install from source as described above +which provides the ``EOBRun_module``. However, we have noticed some incompatibilties with numpy can arise if this is done. If needed, please instead install from source as described above. + +The PyPI release is not sufficient for every advanced-physics workflow. In +particular, hyperbolic parameters (``H_hyp``, ``j_hyp``, and ``r_hyp``) require +a DALI-capable source build. This compatibility layer has been exercised with +DALI commits ``9c4482d95c51b9b4db634d3e432da23a3c0543ed`` and +``5504fdb736d2c63df49632f0afc90ddf553b2693``; record the exact source commit +and compiled ``EOBRun_module`` hash with each production analysis. + +RIFT selects a TEOBResumS parameter profile automatically. The default is +``RIFT_TEOBRESUMS_PROFILE=auto``; validated deployments may pin ``dali``, +``default``, or ``legacy``. An unknown explicit profile fails rather than +guessing. The first call for each parameter schema runs in a child process so +an incompatible native extension cannot crash the parent RIFT process. For +unusually long validation calls, tune ``RIFT_TEOBRESUMS_PROBE_TIMEOUT`` (in +seconds). Use ``RIFT_TEOBRESUMS_SKIP_PROBE=1`` only with a separately +validated and pinned binary. ## gwsurrogate diff --git a/MonteCarloMarginalizeCode/Code/RIFT/asimov/README.md b/MonteCarloMarginalizeCode/Code/RIFT/asimov/README.md index 9290ca086..e1e795d6c 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/asimov/README.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/asimov/README.md @@ -20,3 +20,12 @@ finished and does not submit a duplicate postprocessing job. contract for separate postprocessing adapters: samples (always a list), the RIFT configuration, PSDs, calibration envelopes, likelihood products, and basic event/analysis provenance. Consumers should tolerate additional keys. + +Rimsky integration +------------------ + +The ``rift-rimsky-analysis`` command generates a RIFT follow-up document for +Rimsky's ``sample_sink.asimov_configuration`` hook. It bootstraps from the +PESummary metafile produced by Rimsky's online Bilby analysis and normalizes +Rimsky's underscore-separated prior names for the RIFT template. See +``RIFT/rimsky/README.md`` for configuration and operational details. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini index d34121f81..ccc3a225a 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini +++ b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini @@ -62,12 +62,17 @@ accounting_group_user={{ config['condor']['user'] }} [datafind] url-type=file +{% if data contains 'frame types' %} types = { {% for ifo in ifos %}"{{ifo}}":"{{data['frame types'][ifo]}}",{% endfor %} } +{% endif %} [data] channels = { {% for ifo in ifos %}"{{ifo}}":"{{data['channels'][ifo]}}",{% endfor %} } [lalinference] +{% if data contains 'frame cache' %} +fake-cache = { {% for ifo in ifos %}"{{ifo}}":"{{data['frame cache'][ifo]}}",{% endfor %} } +{% endif %} {% if likelihood contains 'minimum frequency' %} flow = { {% for ifo in ifos %}"{{ifo}}":{{likelihood['minimum frequency'][ifo]}},{% endfor %} } {% else %} @@ -293,11 +298,12 @@ l-max={{ waveform['maximum mode'] | default: 4 }} # * distance prior if this argument is *not* set is dL^2 {%- if priors.keys() contains "luminosity distance" %} {%- assign p = priors['luminosity distance'] %} -{% if p['type'] contains 'PowerLaw' %} +{%- assign distance_prior_type = p['type'] | default: '' %} +{% if distance_prior_type contains 'PowerLaw' %} # Default distance prior no text here, assume alpha=2 -{% elsif p['type'] contains 'UniformSourceFrame' %} +{% elsif distance_prior_type contains 'UniformSourceFrame' %} ile-distance-prior='cosmo_sourceframe' -{% elsif p['type'] contains 'UniformComovingVolume' %} +{% elsif distance_prior_type contains 'UniformComovingVolume' %} ile-distance-prior='cosmo' {% else %} ile-distance-prior="pseudo_cosmo" diff --git a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.py b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.py index 8b70232db..e339a5d05 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.py @@ -5,6 +5,7 @@ import os import re import subprocess +from pathlib import Path from ligo.gracedb.rest import HTTPError @@ -72,6 +73,12 @@ def __init__(self, production, category=None): def _create_ledger_entries(self): """Create entries in the ledger which might be required in the templating.""" + # Rimsky writes Bilby-style prior names into the shared Asimov event. + # Add RIFT's legacy aliases without removing the keys Bilby consumes. + from RIFT.rimsky import normalize_event_metadata + normalized = normalize_event_metadata(self.production.meta) + self.production.meta.clear() + self.production.meta.update(normalized) if "sampler" not in self.production.meta: self.production.meta["sampler"] = {} required_args = { @@ -85,7 +92,7 @@ def _create_ledger_entries(self): section_data[section_arg] = {} def _get_psds(self, format="ascii"): - """Return PSD assets across the ASIMOV 0.5 and 0.7 APIs.""" + """Return PSD assets across the ASIMOV 0.5, 0.6, and 0.7 APIs.""" legacy_getter = getattr(self.production, "get_psds", None) if callable(legacy_getter): assets = legacy_getter(format) @@ -95,6 +102,120 @@ def _get_psds(self, format="ascii"): if format == "xml" and isinstance(assets, dict): return list(assets.values()) return assets + + def _detector_for_psd(self, psdfile): + """Identify a PSD's detector without assuming a filename ordering.""" + filename = Path(psdfile).name.upper() + matches = [ + ifo.upper() + for ifo in self.production.meta.get("interferometers", []) + if ifo.upper() in filename + ] + if len(matches) != 1: + raise PipelineException( + "RIFT cannot identify a unique detector for PSD {}".format(psdfile), + production=self.production.name, + ) + return matches[0] + + def _convert_psd(self, ascii_format, ifo, dryrun=False): + """Convert one on-disk ASCII PSD into RIFT's XML representation.""" + ascii_format = os.path.abspath(os.path.expanduser(ascii_format)) + if not os.path.isfile(ascii_format): + raise PipelineException( + "RIFT PSD for {} does not exist: {}".format(ifo, ascii_format), + production=self.production.name, + ) + + command = [ + "convert_psd_ascii2xml", + "--fname-psd-ascii", + ascii_format, + "--ifo", + ifo.upper(), + "--conventional-postfix", + ] + if dryrun: + print(" ".join(command)) + return command + + try: + converted = subprocess.run( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + check=False, + ) + except FileNotFoundError as exc: + raise PipelineException( + "RIFT PSD conversion executable is unavailable: {}".format( + command[0] + ), + production=self.production.name, + ) from exc + if converted.returncode != 0: + output = converted.stdout.decode(errors="replace") + raise PipelineException( + "RIFT could not convert the {} PSD {}:\n{}".format( + ifo, ascii_format, output + ), + production=self.production.name, + ) + + xml_path = os.path.abspath("{}-psd.xml.gz".format(ifo.upper())) + if not os.path.isfile(xml_path): + raise PipelineException( + "RIFT PSD conversion did not create {}".format(xml_path), + production=self.production.name, + ) + return xml_path + + def _prepare_frame_caches(self): + """Create LAL cache files for local frames supplied by Rimsky.""" + if self.production.meta.get("orchestrator") != "rimsky": + return {} + data = self.production.meta.get("data", {}) + data_files = data.get("data files", {}) + if not isinstance(data_files, dict) or not data_files: + return {} + + cache_dir = Path(self.production.event.work_dir) + cache_dir.mkdir(parents=True, exist_ok=True) + caches = {} + for detector, files in data_files.items(): + if isinstance(files, (str, os.PathLike)): + files = [files] + if not isinstance(files, (list, tuple)): + raise PipelineException( + "RIFT Rimsky frame list for {} is malformed".format(detector), + production=self.production.name, + ) + + entries = [] + for filename in files: + frame = Path(filename).expanduser().resolve() + match = re.search(r"-(\d+)-(\d+)\.gwf$", frame.name) + if not frame.is_file() or match is None: + raise PipelineException( + "RIFT Rimsky frame is missing or has no GPS/duration suffix: {}".format( + frame + ), + production=self.production.name, + ) + start, duration = match.groups() + entries.append( + "{} RIMSKY {} {} {}".format( + detector[0].upper(), start, duration, frame.as_uri() + ) + ) + + cache = cache_dir / "{}-rimsky.cache".format(detector) + cache.write_text("\n".join(entries) + "\n", encoding="utf-8") + caches[detector] = str(cache) + + data["frame cache"] = caches + return caches + # Top-level groups a PESummary metafile carries that are not analysis labels _PESUMMARY_RESERVED = ('version', 'history') @@ -298,29 +419,42 @@ def before_config(self, dryrun=False): """ event = self.production.event category = config.get("general", "calibration_directory") + self._prepare_frame_caches() # XML PSDs self.logger.info("Checking for XML format PSDs") if len(self._get_psds("xml")) == 0 and "psds" in self.production.meta: self.logger.info("Did not find XML format PSDs") + project_dir = Path.cwd() + repository_dir = Path(event.repository.directory) + if not repository_dir.is_absolute(): + repository_dir = (project_dir / repository_dir).resolve() for ifo in self.production.meta["interferometers"]: with set_directory(f"{event.work_dir}"): sample = self.production.meta["likelihood"]["sample rate"] self.logger.info(f"Converting {ifo} {sample}-Hz PSD to XML") - self._convert_psd( + asset = self._convert_psd( self.production.meta["psds"][sample][ifo], ifo, dryrun=dryrun ) - asset = f"{ifo.upper()}-psd.xml.gz" - self.logger.info(f"Conversion complete as {asset}") - git_location = os.path.join(category, "psds") - saveloc = os.path.join( - git_location, str(sample), f"psd_{ifo}.xml.gz" - ) - self.production.event.repository.add_file( + if dryrun: + continue + self.logger.info(f"Conversion complete as {asset}") + git_location = os.path.join(category, "psds") + saveloc = os.path.join( + git_location, str(sample), f"psd_{ifo}.xml.gz" + ) + # EventRepo paths may be relative to the Asimov project. Add + # the converted file after leaving the event work directory so + # it cannot be nested beneath that directory accidentally. + with set_directory(project_dir): + event.repository.add_file( asset, saveloc, commit_message=f"Added the xml format PSD for {ifo}.", ) - self.logger.info(f"Saved at {saveloc}") + xml_psds = getattr(self.production, "xml_psds", None) + if isinstance(xml_psds, dict): + xml_psds[ifo] = str(repository_dir / saveloc) + self.logger.info(f"Saved at {saveloc}") # calmarg: find bilby ini file if needed self.logger.info(" About to check for calmarg ") if 'likelihood' in self.production.meta['sampler']: @@ -673,7 +807,7 @@ def build_dag(self, user=None, dryrun=False): if self.production.event.repository: # with set_directory(os.path.abspath(self.production.rundir)): for psdfile in self._get_psds("xml"): - ifo = psdfile.split("/")[-1].split("-")[1].split(".")[0] + ifo = self._detector_for_psd(psdfile) os.system(f"cp {psdfile} {ifo}-psd.xml.gz") # os.system("cat *_local.cache > local.cache") @@ -716,7 +850,7 @@ def submit_dag(self, dryrun=False): """ self.before_submit() for psdfile in self._get_psds("xml"): - ifo = psdfile.split("/")[-1].split("-")[1].split(".")[0] + ifo = self._detector_for_psd(psdfile) os.system(f"cp {psdfile} {ifo}-psd.xml.gz") command = [ diff --git a/MonteCarloMarginalizeCode/Code/RIFT/calmarg/DESIGN_calmarg_in_loop.md b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/DESIGN_calmarg_in_loop.md index 3f531a04a..9f3e18b8f 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/calmarg/DESIGN_calmarg_in_loop.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/DESIGN_calmarg_in_loop.md @@ -193,6 +193,89 @@ calibration marginalization is active, else 1) and threads it into the three production `DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` call sites (plain / distance-marg / distance+phase-marg). +## Option-compatibility gate (which configurations are refused, and why) + +`RIFT/calmarg/option_compat.py`, called once at driver startup after `opts.gpu` is +resolved and before any precompute. It only **refuses**; it changes no accepted +configuration's arithmetic. + +The problem it solves is structural. In-loop calibration marginalization is switched +on by exactly one option, `--calibration-envelope-directory`, but the `n_cal>1` +reduction exists at exactly one family of call sites: the +`DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` calls inside the +time-marginalized, `--vectorized`, xpy (`--gpu` / `--force-xpy`) branch. Every other +dispatch in the driver reaches a likelihood that takes no `n_cal` argument at all. On +those paths the realizations are drawn, precomputed and paid for, and then never enter +the likelihood: the run costs more, carries `--calibration-envelope-directory` in its +banner and its `.sub` file, and reports the **zero-calibration** answer. Nothing +raises. This is the same failure the sub-sample stencil and +`--time-marginalization-quadrature` gates already refuse in the same driver. + +Classification, traced from the source rather than from flag names: + +| Combination | Category | Source evidence | +|---|---|---| +| calmarg + `--rotation-slow` | **not implemented** | `if opts.rotation_slow:` precedes the `else:` that carries `n_cal`/`cal_method` at every production call site; `PrecomputeLikelihoodTermsWithRotation` takes no `calibration_realizations` and returns no cal cross terms | +| calmarg + `--freqresponse` | **not implemented** | same dispatch structure via `elif opts.freqresponse:`; `PrecomputeLikelihoodTermsFreqResponse` likewise | +| calmarg without `--vectorized` | **cannot take effect** | the packed rholm/cross-term arrays the reduction indexes are built only under `if opts.vectorized:`; the likelihood called is the scalar `FactoredLogLikelihoodTimeMarginalized` | +| calmarg without `--time-marginalization` | **cannot take effect** | the `if not opts.time_marginalization:` branch calls `FactoredLogLikelihood`, which has no `n_cal` | +| calmarg without `--gpu`/`--force-xpy` | **cannot take effect** | plain `--vectorized` calls `DiscreteFactoredLogLikelihoodViaArrayVector` (not `...NoLoop`), which has no `n_cal`. `--gpu` is silently downgraded when cupy is absent, so `--gpu` alone on a CPU node lands here | +| any calibration opt-in without `--calibration-envelope-directory` | **cannot take effect** | `calibration_marginalization` is set by that option and by nothing else; every opt-in is read only under it. `--calibration-fused-kernel` in particular gates on `use_fused_calmarg = calibration_marginalization and opts.calibration_fused_kernel` — with no envelope there is nothing to fuse | +| `--calibration-burn-in-nmax` without `--calibration-burn-in-neff` | **cannot take effect** | the cap is read only inside `if opts.calibration_burn_in_neff ... :`, the branch that runs the zero-cal burn-in. With no burn-in target there is no burn-in to cap, so the explicit cap is silently ignored — the envelope directory alone does *not* make this opt-in live, which is what distinguishes it from the row above | +| `--calibration-dump-responsibilities` (the cal pilot) | **legitimate** | a diagnostic pilot that deliberately uses `cal_method='loop'` and `return 0.0`s from inside the precompute block. Its prerequisites are genuinely different — it needs `--vectorized` and nothing else — and `util_CalPilotStage.py` depends on that. **Exempted** from the `--time-marginalization` and xpy rules, **not** from `--vectorized`, and **not** from the two 3G refusals (it evaluates the baseline packed arrays, so under `--rotation-slow` it would report responsibilities for a likelihood nobody asked for) | +| `--calibration-dump-responsibilities` + `--calibration-fused-kernel` | **legitimate** (notice, not refusal) | `util_CalPilotStage.py` inherits the wide `args_ile.txt` verbatim, so a `--calmarg-fused-kernel` campaign hands its pilot this flag. The pilot uses the loop reduction; the flag is inert here and the run says so | +| `--calibration-export-posterior` on the wide stage | **legitimate** | `util_RIFT_pseudo_pipe.py` emits it there deliberately, documented as harmless (it fires only at the fairdraw stage) | + +### The two guards this replaces were inert + +`--rotation-slow` and `--freqresponse` each carried + +```python +if opts.gpu and getattr(opts, 'calibration_marginalization', False): + raise ValueError("... does not support calibration/glitch marginalization (n_cal>1)") +``` + +There is no `--calibration-marginalization` option and nothing ever sets that attribute +on `opts` — `calibration_marginalization` is a module-level variable assigned ~700 lines +later — so the `getattr` default made both guards **always false**. Verified by running +the driver with `--rotation-slow --calibration-envelope-directory`: it proceeded. They +were also GPU-only and ran before `opts.gpu` is resolved, while the incompatibility is +neither. Both are removed in favour of the central gate. + +### Why "untested" is an empty category here — and why the boundary is still written down + +R. O'Shaughnessy's decision (2026-09) is that it is **fine for these paths not to be +implemented, because calibration marginalization is not tested against the +third-generation machinery at all** — so a combination that silently degrades produces +plausible numbers from a path nothing has ever checked. Failing loudly at startup is +strictly better. + +The gate distinguishes *cannot take effect* from *not implemented*, and a +`not implemented` refusal additionally carries `enable_requires`: what would have to +exist first. For both 3G refusals that is (a) the 3G precompute returning +per-realization calibration cross terms as `PrecomputeLikelihoodTerms` does; (b) an +`n_cal>1` reduction in the corresponding NoLoop, with the same per-realization +self-term ``; and (c) a brute-force agreement check of that reduction, of +the kind `test_selfterm_reduction.py` applies to the baseline likelihood, wired into +`calmarg-check`. + +A third category — *runnable, accepted, but unvalidated* — has **no members** in this +driver, and that is a finding, not an omission: every calibration x 3G combination fails +the stronger test first, because the dispatch does not exist. There is therefore +nothing to leave untested and no `UNTESTED` code path in the gate (an unused kind would +itself be an inert guard). If a 3G calibration reduction is ever implemented, that +category becomes real and belongs here. + +### Tests + +`RIFT/calmarg/test_option_compat.py`, wired into the `calmarg-check` CI gate via +`.travis/test-calmarg.sh`. Every refusal is paired with the **nearest configuration +that must still be accepted**, because these are all over-broad-condition risks and only +the accepting edge can catch one. All fourteen guards were mutation-checked: deleting +each rule, and four over-broad mutations (removing the pilot exemption, making the 3G +refusals pilot-exempt, applying the opt-in rule with an envelope present, and adding a +numeric-default option to `CAL_OPT_IN_FLAGS`). + ## Bug fixed In `ComputeModeIPTimeSeries`'s calibration branch the inner product was being taken diff --git a/MonteCarloMarginalizeCode/Code/RIFT/calmarg/option_compat.py b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/option_compat.py new file mode 100644 index 000000000..b6b328788 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/option_compat.py @@ -0,0 +1,394 @@ +"""Option-compatibility gate for in-loop calibration marginalization (ILE batchmode). + +WHY THIS EXISTS. ``integrate_likelihood_extrinsic_batchmode`` turns calibration +marginalization on from ONE option -- ``--calibration-envelope-directory`` -- but the +``n_cal>1`` reduction it enables lives at exactly one family of call sites: the +``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`` calls inside the time-marginalized, +vectorized, xpy (``--gpu`` / ``--force-xpy``) branch. Every other dispatch in that driver +reaches a likelihood function that takes no ``n_cal`` argument at all. The calibration +realizations are still drawn, still precomputed, and then never enter the likelihood -- +so the run costs more, prints ``calibration`` in its banner and its submit file, and +reports the ZERO-CALIBRATION answer. Nothing raises. + +That is the failure mode this module refuses. It is the same discipline the sub-sample +stencil gate and the ``--time-marginalization-quadrature`` gate already apply in the same +driver, and for the same reason: an option that silently does nothing is worse than one +that is unavailable, because a campaign can be run against it and believed. + +TWO KINDS OF REFUSAL, deliberately distinguished (they mean different things to whoever +reads the message): + +``KIND_INERT`` + The requested option cannot take effect in this configuration. The request is + self-contradictory: what was asked for is unreachable from the flags given. The + remedy is always local -- add the missing flag, or drop the one that does nothing. + +``KIND_UNIMPLEMENTED`` + The combination is not implemented, and running it would silently evaluate a + DIFFERENT likelihood than the one requested. The remedy is to drop one side. These + refusals additionally carry ``enable_requires``: what would have to be built and + validated before the combination could be allowed. R. O'Shaughnessy's decision + (2026-09) is that it is fine for these paths not to be implemented, BECAUSE + calibration marginalization is not tested against the third-generation machinery at + all -- so a silent degradation there produces plausible numbers from a path nothing + has ever checked. That is a scope boundary, not an oversight; ``enable_requires`` + is where the boundary is written down. + +A THIRD KIND -- runnable, accepted, but unvalidated -- IS NOT REPRESENTED HERE, and that +is a finding rather than an omission. Tracing every calibration x 3G combination in the +driver, each one fails the STRONGER test first: the dispatch does not exist, so there is +nothing to leave untested. See DESIGN_calmarg_in_loop.md ("Option-compatibility gate"). + +This module only REFUSES. It changes no accepted configuration's arithmetic. +""" +from __future__ import print_function + +import collections + +__all__ = ["KIND_INERT", "KIND_UNIMPLEMENTED", "CAL_OPT_IN_FLAGS", "Refusal", + "calibration_refusals", "calibration_notices", "refusals_from_opts", + "notices_from_opts", "refuse_incompatible_calibration_options"] + +KIND_INERT = "cannot take effect" +KIND_UNIMPLEMENTED = "not implemented" + +Refusal = collections.namedtuple("Refusal", "kind options message enable_requires") + + +# Calibration opt-ins whose ONLY effect is inside in-loop calibration marginalization, +# i.e. inside the `if opts.calibration_envelope_directory:` block or gated on the module +# flag it sets. Restricted on purpose to options whose default is False/None: an option +# with a numeric default (--calibration-n-realizations, --calibration-spline-count, +# --calibration-pilot-extrinsic, --calibration-mc-error-extrinsic, +# --calibration-neff-cal-target, --calibration-n-realizations-max) is ALWAYS "set", so a +# rule over those would refuse every run that merely left the defaults in place. +# +# --calibration-burn-in-nmax is listed here like the rest (it too does nothing without the +# envelope), but it needs a SECOND rule below: the envelope alone does not make it live, +# because the driver reads it only under --calibration-burn-in-neff. +# +# (cli_flag, opts_attribute). +CAL_OPT_IN_FLAGS = ( + ("--calibration-fused-kernel", "calibration_fused_kernel"), + ("--calibration-conjugate-phase", "calibration_conjugate_phase"), + ("--calibration-global-norm", "calibration_global_norm"), + ("--calibration-proposal-breadcrumb", "calibration_proposal_breadcrumb"), + ("--calibration-dump-responsibilities", "calibration_dump_responsibilities"), + ("--calibration-export-posterior", "calibration_export_posterior"), + ("--calibration-burn-in-neff", "calibration_burn_in_neff"), + ("--calibration-burn-in-nmax", "calibration_burn_in_nmax"), +) + +_ENVELOPE = "--calibration-envelope-directory" + +# What would have to exist before calibration marginalization could be allowed on a +# third-generation likelihood path. ONE definition, interpolated into both 3G refusals, +# so the two messages cannot drift apart. +_3G_ENABLE_REQUIRES = ( + "calibration marginalization has NO test coverage against the third-generation " + "machinery -- not a weak test, none -- so this is a scope boundary and not an " + "oversight. Enabling it needs, at minimum: (a) {precompute} to return the " + "per-realization calibration cross terms that PrecomputeLikelihoodTerms returns " + "for the baseline likelihood; (b) an n_cal>1 reduction in {noloop}, with the same " + "per-realization self-term the baseline reduction uses; and (c) a " + "brute-force agreement check of that reduction, of the kind " + "RIFT/calmarg/test_selfterm_reduction.py already applies to the baseline " + "likelihood, wired into the calmarg-check CI gate." +) + + +def _inert(options, message): + return Refusal(KIND_INERT, tuple(options), message, None) + + +def _unimplemented(options, message, enable_requires): + return Refusal(KIND_UNIMPLEMENTED, tuple(options), message, enable_requires) + + +def calibration_refusals(calibration_envelope_directory=None, + opt_in_flags=(), + time_marginalization=False, + vectorized=False, + xpy_evaluator=False, + rotation_slow=False, + freqresponse=False, + dump_responsibilities=False, + burn_in_neff=None, + burn_in_nmax=None, + n_realizations=None, + fused_kernel=False): + """Return the list of Refusals for one resolved ILE configuration (possibly empty). + + Pure: booleans in, Refusals out. No option namespace, no I/O, no raising. + + Parameters + ---------- + calibration_envelope_directory : str or None + The value of --calibration-envelope-directory. This ALONE decides whether + in-loop calibration marginalization is active in the driver. + opt_in_flags : sequence of str + CLI spellings of the calibration opt-ins the user actually set (see + CAL_OPT_IN_FLAGS). Order is preserved in the output. + time_marginalization, vectorized, rotation_slow, freqresponse : bool + The corresponding CLI booleans. + xpy_evaluator : bool + ``opts.gpu`` AFTER the driver has resolved it, i.e. a real CUDA device OR + --force-xpy. It must be the resolved value: `--gpu` is silently downgraded to + False when cupy is unavailable, and the downgraded configuration is precisely + one of the ones that drops calibration on the floor. + dump_responsibilities : bool + Whether --calibration-dump-responsibilities is set. The cal PILOT is a + legitimate diagnostic mode with DIFFERENT prerequisites: it runs inside the + `if opts.vectorized:` precompute block, uses cal_method='loop' on whatever xpy + is available, and returns before the driver ever selects a production + likelihood_function -- so it needs neither --time-marginalization nor an xpy + evaluator. Refusing it on those grounds would break the shipped adaptive + pipeline (util_CalPilotStage.py). + burn_in_neff, burn_in_nmax : float/int or None + --calibration-burn-in-neff and --calibration-burn-in-nmax. The cap is a + DEPENDENT option: the driver reads it only inside `if + opts.calibration_burn_in_neff:`, so on its own it is silently ignored. + n_realizations : int or None + --calibration-n-realizations. NOT an opt-in (it has a non-None default, so a rule + over its mere presence would refuse every run), but it is a PREREQUISITE for the + two options whose consumers sit behind `n_cal > 1`. None means "not supplied"; + the caller passes the resolved value. + """ + out = [] + + if not calibration_envelope_directory: + # Calibration marginalization is OFF. Every calibration opt-in is then a silent + # no-op: `calibration_marginalization` is set by the envelope directory and by + # nothing else, and each of these flags is read only under it. Nothing else in + # this function applies -- --rotation-slow and --freqresponse are perfectly fine + # on their own, and must stay fine. + for flag in opt_in_flags: + out.append(_inert( + (flag, _ENVELOPE), + "%s was requested without %s, so in-loop calibration marginalization is " + "never switched on and %s is read by nothing: it is silently ignored. " + "Add %s (the per-IFO .txt envelope directory is what activates " + "calibration marginalization), or drop %s." + % (flag, _ENVELOPE, flag, _ENVELOPE, flag))) + return out + + if rotation_slow: + out.append(_unimplemented( + (_ENVELOPE, "--rotation-slow"), + "%s (in-loop calibration marginalization) with --rotation-slow is not " + "implemented. --rotation-slow dispatches to " + "DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation, which takes " + "no n_cal/cal_method argument, and PrecomputeLikelihoodTermsWithRotation " + "builds no calibration cross terms. The realizations would be drawn and " + "paid for, then never entered: the run would advertise a " + "calibration-marginalized analysis and report the zero-calibration one. " + "Drop --rotation-slow to marginalize over calibration on the baseline " + "likelihood, or drop %s to run the slow-rotation likelihood." + % (_ENVELOPE, _ENVELOPE), + _3G_ENABLE_REQUIRES.format( + precompute="PrecomputeLikelihoodTermsWithRotation", + noloop="DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation"))) + + if freqresponse: + out.append(_unimplemented( + (_ENVELOPE, "--freqresponse"), + "%s (in-loop calibration marginalization) with --freqresponse is not " + "implemented. --freqresponse dispatches to " + "DiscreteFactoredLogLikelihoodFreqResponseNoLoop, which takes no " + "n_cal/cal_method argument, and PrecomputeLikelihoodTermsFreqResponse " + "builds no calibration cross terms. The realizations would be drawn and " + "paid for, then never entered: the run would advertise a " + "calibration-marginalized analysis and report the zero-calibration one. " + "Drop --freqresponse to marginalize over calibration on the baseline " + "likelihood, or drop %s to run the finite-size-response likelihood." + % (_ENVELOPE, _ENVELOPE), + _3G_ENABLE_REQUIRES.format( + precompute="PrecomputeLikelihoodTermsFreqResponse", + noloop="DiscreteFactoredLogLikelihoodFreqResponseNoLoop"))) + + if burn_in_nmax and not burn_in_neff: + # A DEPENDENT opt-in: unlike the others, the envelope directory is not enough to + # make it live. The driver reads opts.calibration_burn_in_nmax only inside + # `if opts.calibration_burn_in_neff ...`, which is the option that switches the + # zero-cal burn-in on; with no burn-in there is nothing for the cap to cap. + out.append(_inert( + ("--calibration-burn-in-nmax", "--calibration-burn-in-neff"), + "--calibration-burn-in-nmax caps the zero-cal burn-in, but the burn-in is " + "switched on by --calibration-burn-in-neff and by nothing else -- the driver " + "reads the cap only inside that option's branch. Set on its own the cap is " + "silently ignored, and the run goes straight to the full cal-marginalized " + "integration at the production sample budget, which is exactly the number of " + "samples the cap was meant to hold down. Add --calibration-burn-in-neff " + "(the burn-in target), or drop --calibration-burn-in-nmax.")) + + # ---- the n_cal > 1 prerequisites ------------------------------------------------- + # Two opt-ins have consumers guarded by n_cal > 1, so ONE realization makes them inert + # even on an otherwise perfect configuration. Traced to the source, not inferred from + # the flag names: + # --calibration-fused-kernel factored_likelihood's `if n_cal == 1:` branch RETURNS + # before `cal_method == 'fused'` is ever read, so the + # fused reduction is unreachable and the run silently + # uses the ordinary one. + # --calibration-burn-in-neff the driver's burn-in block is + # `if opts.calibration_burn_in_neff and + # calibration_marginalization and + # n_cal_for_likelihood > 1:` -- with one realization the + # zero-cal burn-in never runs. + # NOT listed, deliberately: --calibration-conjugate-phase and --calibration-global-norm + # reach the likelihood through extra_kwargs on the PRECOMPUTE and are honoured at + # n_cal == 1 too (that branch uses rho_sq_cal[0]), so refusing them here would be a + # false positive. + if n_realizations is not None and int(n_realizations) <= 1: + if fused_kernel: + out.append(_inert( + ("--calibration-fused-kernel", "--calibration-n-realizations"), + "--calibration-fused-kernel with --calibration-n-realizations %d: the " + "fused reduction is selected by cal_method='fused', which " + "factored_likelihood reads only AFTER its `n_cal == 1` branch has already " + "returned. With one realization the flag is therefore silently ignored " + "and the ordinary reduction runs -- the answer is right, but the run " + "advertises a kernel it did not use. Raise " + "--calibration-n-realizations above 1, or drop the flag." + % (int(n_realizations),))) + if burn_in_neff: + out.append(_inert( + ("--calibration-burn-in-neff", "--calibration-n-realizations"), + "--calibration-burn-in-neff with --calibration-n-realizations %d: the " + "zero-cal burn-in runs only under `n_cal_for_likelihood > 1`, so with one " + "realization there is no burn-in and the target is silently ignored. " + "Raise --calibration-n-realizations above 1, or drop the flag." + % (int(n_realizations),))) + + if not vectorized: + out.append(_inert( + (_ENVELOPE, "--vectorized"), + "%s requires --vectorized. Without it the driver never builds the packed " + "rholm/cross-term arrays the calibration reduction indexes, and the " + "likelihood it calls is the scalar FactoredLogLikelihoodTimeMarginalized " + "(or FactoredLogLikelihood without --time-marginalization), neither of " + "which takes n_cal -- so the calibration realizations are drawn and then " + "ignored. Add --vectorized, or drop %s." % (_ENVELOPE, _ENVELOPE))) + + if not dump_responsibilities: + # The cal PILOT (--calibration-dump-responsibilities) is exempt from both of + # these: it returns 0.0 from inside the precompute block, before the driver + # picks a production likelihood_function at all. Its own prerequisite + # (--vectorized) is checked above and is NOT exempted. + if not time_marginalization: + out.append(_inert( + (_ENVELOPE, "--time-marginalization"), + "%s requires --time-marginalization. Without it the driver takes the " + "`if not opts.time_marginalization` branch and calls " + "FactoredLogLikelihood, which takes no n_cal, so the calibration " + "realizations are drawn and then ignored. Add --time-marginalization, " + "or drop %s. (The one configuration that legitimately runs " + "calibration without it is the pilot, " + "--calibration-dump-responsibilities, which returns before this branch " + "is reached.)" % (_ENVELOPE, _ENVELOPE))) + if not xpy_evaluator: + out.append(_inert( + (_ENVELOPE, "--gpu"), + "%s requires the maintained NoLoop evaluator, selected by --gpu. " + "Plain --vectorized without it calls " + "DiscreteFactoredLogLikelihoodViaArrayVector, which takes no n_cal, so " + "the calibration realizations are drawn and then ignored. NOTE that " + "--gpu is SILENTLY DOWNGRADED when cupy is unavailable, so `--gpu` on a " + "host with no CUDA device lands here: add --force-xpy, which keeps the " + "identical NoLoop code path on numpy. Otherwise drop %s. (The pilot, " + "--calibration-dump-responsibilities, is exempt: it evaluates on " + "whatever xpy is present and returns before a production likelihood is " + "selected.)" % (_ENVELOPE, _ENVELOPE))) + + return out + + +def calibration_notices(calibration_envelope_directory=None, + fused_kernel=False, + dump_responsibilities=False): + """Non-fatal notices: accepted configurations where an option has no effect BY DESIGN. + + These are deliberately NOT refusals. Each is emitted by the shipped pipeline on a + stage where it is a documented no-op, so refusing them would break production: + + * util_CalPilotStage.py inherits the WIDE args_ile.txt (which may carry + --calibration-fused-kernel) and appends --calibration-dump-responsibilities. + The pilot deliberately uses cal_method='loop'. + * util_RIFT_pseudo_pipe.py emits --calibration-export-posterior on the wide stage + too, where it is documented as harmless (it only fires at the fairdraw stage). + + A notice says so out loud instead of leaving the reader to infer it from a banner. + """ + out = [] + if calibration_envelope_directory and dump_responsibilities and fused_kernel: + out.append( + "[calmarg] --calibration-dump-responsibilities: this is the cal PILOT. It " + "uses the loop reduction (cal_method='loop') and returns before production " + "integration, so --calibration-fused-kernel has no effect on this run. " + "Accepted, not refused: util_CalPilotStage.py inherits the wide ILE " + "arguments verbatim, so the pilot legitimately carries this flag.") + return out + + +def _opt_ins_set_on(opts): + return tuple(flag for flag, attr in CAL_OPT_IN_FLAGS if getattr(opts, attr, None)) + + +def refusals_from_opts(opts): + """Adapter: read a driver option namespace and return calibration_refusals(...). + + THE SEAM. The predicate above can be perfectly correct and still be wired to the + wrong attribute; that is exactly how the two `getattr(opts, + 'calibration_marginalization', False)` guards this gate replaces came to be inert + (there is no such option, so both always evaluated False). So this adapter is + covered by subprocess tests through the real CLI, not only by unit calls. + + ``opts.gpu`` MUST already be resolved (see calibration_refusals). + """ + return calibration_refusals( + calibration_envelope_directory=getattr(opts, "calibration_envelope_directory", None), + opt_in_flags=_opt_ins_set_on(opts), + time_marginalization=bool(getattr(opts, "time_marginalization", False)), + vectorized=bool(getattr(opts, "vectorized", False)), + xpy_evaluator=bool(getattr(opts, "gpu", False)), + rotation_slow=bool(getattr(opts, "rotation_slow", False)), + freqresponse=bool(getattr(opts, "freqresponse", False)), + dump_responsibilities=bool(getattr(opts, "calibration_dump_responsibilities", None)), + burn_in_neff=getattr(opts, "calibration_burn_in_neff", None), + burn_in_nmax=getattr(opts, "calibration_burn_in_nmax", None), + n_realizations=getattr(opts, "calibration_n_realizations", None), + fused_kernel=bool(getattr(opts, "calibration_fused_kernel", False)), + ) + + +def notices_from_opts(opts): + return calibration_notices( + calibration_envelope_directory=getattr(opts, "calibration_envelope_directory", None), + fused_kernel=bool(getattr(opts, "calibration_fused_kernel", False)), + dump_responsibilities=bool(getattr(opts, "calibration_dump_responsibilities", None)), + ) + + +def format_refusals(refusals): + """One message for the whole configuration, so a user fixes it in one pass.""" + lines = ["Refusing this calibration-marginalization configuration rather than " + "silently degrading it to the zero-calibration likelihood:"] + for i, r in enumerate(refusals, 1): + lines.append(" (%d) [%s] %s" % (i, r.kind, r.message)) + if r.enable_requires: + lines.append(" To enable it: %s" % r.enable_requires) + return "\n".join(lines) + + +def refuse_incompatible_calibration_options(opts, printer=print): + """Raise ValueError if this configuration cannot honour its calibration options. + + Emits the non-fatal notices first, so an accepted-but-inert-by-design flag is still + visible in the log. Returns the notice list (for tests). + """ + notices = notices_from_opts(opts) + for n in notices: + printer(n) + refusals = refusals_from_opts(opts) + if refusals: + raise ValueError(format_refusals(refusals)) + return notices diff --git a/MonteCarloMarginalizeCode/Code/RIFT/calmarg/test_option_compat.py b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/test_option_compat.py new file mode 100644 index 000000000..1a733e5a9 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/test_option_compat.py @@ -0,0 +1,464 @@ +#!/usr/bin/env python3 +"""Calibration option-compatibility gate: the refusals AND their nearest legal neighbours. + +Every guard in RIFT/calmarg/option_compat.py is an OVER-BROAD-CONDITION risk: a rule that +refuses too much is invisible to a test that only checks the refusal fires. So each +refusal here is paired with the closest configuration that must still be ACCEPTED -- the +one flag away. Those accepting cases are the load-bearing half. + +Three of them are not hypothetical. They are configurations the SHIPPED pipeline emits: + + * --calibration-dump-responsibilities without --time-marginalization / without an xpy + evaluator. The cal pilot returns from inside the precompute block; refusing it on + the production likelihood's prerequisites would break util_CalPilotStage.py. + * --calibration-fused-kernel together with --calibration-dump-responsibilities. + util_CalPilotStage.py inherits the WIDE args_ile.txt verbatim, so a run configured + with --calmarg-fused-kernel gives its pilot that flag. The pilot uses the loop + reduction; the flag is inert, not wrong. + * --calibration-export-posterior on the wide stage. util_RIFT_pseudo_pipe.py emits it + there deliberately, documented as harmless (it only fires at the fairdraw stage). + + python3 test_option_compat.py # or: pytest test_option_compat.py +""" +from __future__ import print_function + +import os +import re +import subprocess +import sys + +import RIFT.calmarg.option_compat as oc + +_HERE = os.path.dirname(os.path.abspath(__file__)) +CODE_ROOT = os.path.normpath(os.path.join(_HERE, '..', '..')) +DRIVER = os.path.join(CODE_ROOT, 'bin', 'integrate_likelihood_extrinsic_batchmode') + +ENV_DIR = '/tmp/cal_env_that_need_not_exist' + + +class _Opts(object): + """A stand-in for the driver's optparse namespace, with the driver's own defaults.""" + + def __init__(self, **kw): + self.calibration_envelope_directory = None + self.time_marginalization = False + self.vectorized = False + self.gpu = False + self.rotation_slow = False + self.freqresponse = False + for _flag, attr in oc.CAL_OPT_IN_FLAGS: + setattr(self, attr, None) + self.calibration_fused_kernel = False + self.calibration_conjugate_phase = False + self.calibration_global_norm = False + self.calibration_export_posterior = False + # the driver's own default (optp.add_option("--calibration-n-realizations", + # default=100)). NOT an opt-in; it is the prerequisite for the two options + # whose consumers sit behind n_cal > 1. + self.calibration_n_realizations = 100 + for k, v in kw.items(): + if not hasattr(self, k): + raise AttributeError("no such driver option: %s" % k) + setattr(self, k, v) + + +def _honoured(**kw): + """The configuration in-loop calmarg actually runs on (the demo's own).""" + base = dict(calibration_envelope_directory=ENV_DIR, time_marginalization=True, + vectorized=True, gpu=True) + base.update(kw) + return _Opts(**base) + + +def _kinds(refusals): + return sorted(set(r.kind for r in refusals)) + + +def _flags(refusals): + return set(f for r in refusals for f in r.options) + + +# ------------------------------------------------------------------ ACCEPTING EDGE + +def test_the_honoured_configuration_is_accepted(): + """The configuration every calmarg run and the calmarg demo use. If this ever + starts failing, the gate has been tightened into breaking production.""" + assert oc.refusals_from_opts(_honoured()) == [] + + +def test_every_calibration_opt_in_is_accepted_on_the_honoured_configuration(): + """Each opt-in, one at a time, on the good configuration. A rule keyed on the flag + rather than on the missing prerequisite would fail here. + + --calibration-burn-in-nmax is the one opt-in the envelope does not suffice for: it is + a cap on a burn-in that only --calibration-burn-in-neff switches on, so it is carried + with its dependency here and refused on its own below.""" + for flag, attr in oc.CAL_OPT_IN_FLAGS: + value = 'x' if attr.endswith(('breadcrumb', 'responsibilities')) else True + if attr.startswith('calibration_burn_in'): + value = 10 + extra = {'calibration_burn_in_neff': 100.} if attr.endswith('burn_in_nmax') else {} + opts = _honoured(**dict(extra, **{attr: value})) + assert oc.refusals_from_opts(opts) == [], (flag, oc.refusals_from_opts(opts)) + + +def test_calibration_off_never_refuses_anything(): + """No envelope directory and no opt-ins: every other flag combination must pass, + including the two 3G paths, which are perfectly legal on their own.""" + for kw in (dict(), dict(rotation_slow=True, vectorized=True), + dict(freqresponse=True, vectorized=True), + dict(time_marginalization=True, vectorized=True, gpu=True)): + assert oc.refusals_from_opts(_Opts(**kw)) == [], kw + + +def test_export_posterior_on_the_wide_stage_is_accepted(): + """util_RIFT_pseudo_pipe.py emits --calibration-export-posterior on the WIDE stage, + where it is a documented no-op (it fires only at the fairdraw stage). Refusing it + would break every calmarg campaign built with --calmarg-export-posterior.""" + assert oc.refusals_from_opts(_honoured(calibration_export_posterior=True)) == [] + + +# --------------------------------------------------------- R1/R2: the 3G likelihoods + +def test_calmarg_with_rotation_slow_is_refused(): + r = oc.refusals_from_opts(_honoured(rotation_slow=True)) + assert len(r) == 1 and r[0].kind == oc.KIND_UNIMPLEMENTED + assert r[0].options == ('--calibration-envelope-directory', '--rotation-slow') + assert 'PrecomputeLikelihoodTermsWithRotation' in r[0].message + assert 'NO test coverage against the third-generation machinery' in r[0].enable_requires + + +def test_calmarg_with_freqresponse_is_refused(): + r = oc.refusals_from_opts(_honoured(freqresponse=True)) + assert len(r) == 1 and r[0].kind == oc.KIND_UNIMPLEMENTED + assert r[0].options == ('--calibration-envelope-directory', '--freqresponse') + assert 'PrecomputeLikelihoodTermsFreqResponse' in r[0].message + assert 'NO test coverage against the third-generation machinery' in r[0].enable_requires + + +def test_the_3g_refusals_do_not_depend_on_the_backend(): + """The refusals must fire on the CPU-vectorized path too, not only under --gpu. + + This is the regression the guards this gate replaced actually had. They read + `opts.gpu and getattr(opts, 'calibration_marginalization', False)` -- two defects in + one line: the attribute does not exist (so they never fired at all), and even had it + existed the `opts.gpu and` would have left the CPU-vectorized replacement likelihood + unprotected. Both --rotation-slow and --freqresponse are wired into the CPU + branch as well as the xpy one, and neither PrecomputeLikelihoodTermsWithRotation nor + PrecomputeLikelihoodTermsFreqResponse builds calibration cross terms on EITHER + backend, so the backend has nothing to do with it. + + Asserted by MEMBERSHIP, not by `len(r) == 1`: at gpu=False the configuration also + earns the separate --gpu refusal (in-loop calmarg needs the xpy evaluator), so a + length assertion here would pass for the wrong reason. Re-adding an `opts.gpu and` + condition to either rule must fail this test. + """ + for kw, flag in ((dict(rotation_slow=True), '--rotation-slow'), + (dict(freqresponse=True), '--freqresponse')): + r = oc.refusals_from_opts(_honoured(gpu=False, **kw)) + match = [x for x in r + if x.options == ('--calibration-envelope-directory', flag)] + assert len(match) == 1, (flag, r) + assert match[0].kind == oc.KIND_UNIMPLEMENTED, (flag, match) + + +def test_the_3g_refusals_survive_the_pilot_exemption(): + """--calibration-dump-responsibilities exempts a configuration from the PRODUCTION + likelihood's prerequisites. It must NOT exempt it from the 3G refusals: the pilot + evaluates the BASELINE packed arrays, so under --rotation-slow it would report + calibration responsibilities for a likelihood the user did not ask for.""" + for kw in (dict(rotation_slow=True), dict(freqresponse=True)): + r = oc.refusals_from_opts(_honoured(calibration_dump_responsibilities='d.npz', **kw)) + assert [x.kind for x in r] == [oc.KIND_UNIMPLEMENTED], (kw, r) + + +def test_3g_without_calibration_is_accepted(): + """The nearest legal neighbour of R1/R2: the same 3G run with no calibration.""" + assert oc.refusals_from_opts(_Opts(rotation_slow=True, vectorized=True)) == [] + assert oc.refusals_from_opts(_Opts(freqresponse=True, vectorized=True)) == [] + + +# ------------------------------------------------- R3/R4/R5: the honoured-path prereqs + +def test_calmarg_without_vectorized_is_refused(): + r = oc.refusals_from_opts(_honoured(vectorized=False)) + assert any(x.options == ('--calibration-envelope-directory', '--vectorized') for x in r), r + assert _kinds(r) == [oc.KIND_INERT] + + +def test_calmarg_without_time_marginalization_is_refused(): + r = oc.refusals_from_opts(_honoured(time_marginalization=False)) + assert any(x.options == ('--calibration-envelope-directory', '--time-marginalization') + for x in r), r + + +def test_calmarg_without_an_xpy_evaluator_is_refused(): + r = oc.refusals_from_opts(_honoured(gpu=False)) + assert any(x.options == ('--calibration-envelope-directory', '--gpu') for x in r), r + assert '--force-xpy' in [x for x in r if x.options[-1] == '--gpu'][0].message + + +def test_all_missing_prerequisites_are_reported_at_once(): + """A gate that reports one missing flag per run costs a queue cycle per flag.""" + r = oc.refusals_from_opts(_Opts(calibration_envelope_directory=ENV_DIR)) + assert _flags(r) >= {'--vectorized', '--time-marginalization', '--gpu'}, r + + +# ---------------------------------------------------------- the pilot's own exemptions + +def test_pilot_is_exempt_from_time_marginalization_and_the_xpy_evaluator(): + """The cal PILOT returns 0.0 from inside the `if opts.vectorized:` precompute block, + before the driver selects a production likelihood_function at all. Its prerequisites + are genuinely different, and util_CalPilotStage.py depends on that.""" + opts = _Opts(calibration_envelope_directory=ENV_DIR, vectorized=True, + calibration_dump_responsibilities='resp.npz') + assert oc.refusals_from_opts(opts) == [], oc.refusals_from_opts(opts) + + +def test_pilot_is_NOT_exempt_from_vectorized(): + """The pilot block lives inside `if opts.vectorized:`; without it the pilot never + runs and never writes its dump, and util_CalPilotFit.py then fails on a missing + file. The exemption must be scoped to the two prerequisites it actually earns.""" + opts = _Opts(calibration_envelope_directory=ENV_DIR, + calibration_dump_responsibilities='resp.npz') + r = oc.refusals_from_opts(opts) + assert any(x.options == ('--calibration-envelope-directory', '--vectorized') for x in r), r + + +def test_pilot_with_the_fused_kernel_is_accepted_with_a_notice(): + """util_CalPilotStage.py inherits the wide args verbatim, so a --calmarg-fused-kernel + campaign hands its pilot --calibration-fused-kernel. The pilot uses cal_method='loop' + and returns early: the flag is inert here, which is a NOTICE, not a refusal.""" + opts = _honoured(calibration_dump_responsibilities='resp.npz', + calibration_fused_kernel=True) + assert oc.refusals_from_opts(opts) == [] + notices = oc.notices_from_opts(opts) + assert len(notices) == 1 and '--calibration-fused-kernel' in notices[0] + # ... and no notice when the pilot is not involved + assert oc.notices_from_opts(_honoured(calibration_fused_kernel=True)) == [] + + +# --------------------------------------------- R7: the burn-in cap needs the burn-in + +def test_burn_in_nmax_without_neff_is_refused(): + """The envelope is NOT enough to make --calibration-burn-in-nmax live. The driver + reads opts.calibration_burn_in_nmax only inside the `if opts.calibration_burn_in_neff` + branch, so on the otherwise-honoured configuration the cap is silently ignored and the + burn-in it was sizing never happens: the exact failure mode this gate exists for.""" + r = oc.refusals_from_opts(_honoured(calibration_burn_in_nmax=4000)) + assert len(r) == 1 and r[0].kind == oc.KIND_INERT, r + assert r[0].options == ('--calibration-burn-in-nmax', '--calibration-burn-in-neff'), r + + +def test_burn_in_nmax_with_neff_is_accepted(): + """The nearest legal neighbours: the cap together with the option that reads it, and + the burn-in target on its own (no cap needed -- it falls back to the run's --n-max).""" + assert oc.refusals_from_opts( + _honoured(calibration_burn_in_neff=100., calibration_burn_in_nmax=4000)) == [] + assert oc.refusals_from_opts(_honoured(calibration_burn_in_neff=100.)) == [] + + +# ------------------------------------------------------ R6: opt-ins without the envelope + +def test_each_opt_in_without_the_envelope_directory_is_refused(): + for flag, attr in oc.CAL_OPT_IN_FLAGS: + value = 'x' if attr.endswith(('breadcrumb', 'responsibilities')) else True + if attr.startswith('calibration_burn_in'): + value = 10 + opts = _Opts(time_marginalization=True, vectorized=True, gpu=True, **{attr: value}) + r = oc.refusals_from_opts(opts) + assert len(r) == 1 and r[0].kind == oc.KIND_INERT, (flag, r) + assert r[0].options == (flag, '--calibration-envelope-directory'), (flag, r) + + +# ------------------------------------- R8: the n_cal > 1 prerequisites + +def test_fused_kernel_with_one_realization_is_refused(): + """--calibration-fused-kernel is inert at n_cal == 1: factored_likelihood's + `if n_cal == 1:` branch returns before cal_method is read, so the fused reduction is + unreachable and the run advertises a kernel it did not use.""" + r = oc.refusals_from_opts(_honoured(calibration_fused_kernel=True, + calibration_n_realizations=1)) + match = [x for x in r if x.options == ('--calibration-fused-kernel', + '--calibration-n-realizations')] + assert len(match) == 1, r + assert match[0].kind == oc.KIND_INERT, match + + +def test_burn_in_neff_with_one_realization_is_refused(): + """The zero-cal burn-in block is guarded by `n_cal_for_likelihood > 1`, so with one + realization the burn-in never runs and the target is silently ignored.""" + r = oc.refusals_from_opts(_honoured(calibration_burn_in_neff=100., + calibration_n_realizations=1)) + match = [x for x in r if x.options == ('--calibration-burn-in-neff', + '--calibration-n-realizations')] + assert len(match) == 1, r + assert match[0].kind == oc.KIND_INERT, match + + +def test_the_n_cal_rules_do_not_fire_at_the_shipped_default(): + """--calibration-n-realizations defaults to 100, so these rules must be silent on a + command line that simply leaves it alone. A rule that fires on the default would + refuse every calmarg run there is.""" + for kw in (dict(calibration_fused_kernel=True), dict(calibration_burn_in_neff=100.)): + assert oc.refusals_from_opts(_honoured(**kw)) == [], kw + + +def test_conjugate_and_global_norm_are_NOT_refused_at_one_realization(): + """The deliberate non-members of the rule above. Both reach the likelihood through + extra_kwargs on the PRECOMPUTE and are honoured at n_cal == 1 as well (that branch + uses rho_sq_cal[0]), so refusing them would be a false positive -- the failure mode + this whole gate is supposed to avoid.""" + for kw in (dict(calibration_conjugate_phase=True), dict(calibration_global_norm=True)): + r = oc.refusals_from_opts(_honoured(calibration_n_realizations=1, **kw)) + assert r == [], (kw, r) + + +def test_numeric_default_options_are_not_treated_as_opt_ins(): + """--calibration-n-realizations and friends have non-None DEFAULTS, so a rule over + them would refuse every run that merely left the defaults alone. They are + deliberately absent from CAL_OPT_IN_FLAGS; this pins that.""" + attrs = set(a for _f, a in oc.CAL_OPT_IN_FLAGS) + for forbidden in ('calibration_n_realizations', 'calibration_spline_count', + 'calibration_pilot_extrinsic', 'calibration_mc_error_extrinsic', + 'calibration_neff_cal_target', 'calibration_n_realizations_max'): + assert forbidden not in attrs, forbidden + + +def test_every_opt_in_flag_exists_in_the_driver(): + """CAL_OPT_IN_FLAGS is a hand-maintained list of CLI spellings. A renamed or removed + option would leave a rule that can never fire -- an inert guard.""" + src = open(DRIVER).read() + for flag, attr in oc.CAL_OPT_IN_FLAGS: + assert '"%s"' % flag in src, "%s is not an option of the driver" % flag + assert re.search(r'\bopts\.%s\b' % attr, src), \ + "%s is never read by the driver" % attr + + +# ------------------------------------------------------------------- the wiring itself + +def _driver_code_lines(): + """The driver with comment-only lines removed. + + Needed because the replaced guard is QUOTED in the comment that explains why it went, + and a naive substring test over the whole file would match that comment -- i.e. it + would pass only while nobody documented the change, and fail the moment somebody did. + """ + out = [] + for line in open(DRIVER): + if line.lstrip().startswith('#'): + continue + out.append(line) + return ''.join(out) + + +def test_the_inert_getattr_guards_are_gone(): + """The two guards this gate replaces read + `getattr(opts, 'calibration_marginalization', False)`. There is no such option and + nothing sets that attribute, so both were ALWAYS FALSE. Re-adding one would restore + a guard that looks like coverage and is not.""" + code = _driver_code_lines() + assert "getattr(opts, 'calibration_marginalization'" not in code + assert 'getattr(opts, "calibration_marginalization"' not in code + # ... and the pointer to the replacement must survive, or the next reader re-adds the + # guard where it used to be rather than extending the gate. + src = open(DRIVER).read() + assert src.count('RIFT/calmarg/option_compat.py') >= 2, \ + 'the removed guards no longer point at the gate that replaced them' + + +def test_the_gate_runs_after_the_gpu_downgrade(): + """opts.gpu is silently downgraded to False when cupy is absent. The gate reads it, + so it must run AFTER that -- a call placed above the downgrade would accept exactly + the CPU-node configuration that drops calibration on the floor.""" + src = open(DRIVER).read() + downgrade = src.index('Override --gpu (not available)') + call = src.index('refuse_incompatible_calibration_options(opts)') + precompute = src.index('factored_likelihood.PrecomputeLikelihoodTerms(') + assert downgrade < call < precompute, (downgrade, call, precompute) + + +# ------------------------------------------------------- the CLI seam, in subprocesses + +def _run(args, timeout=300): + env = dict(os.environ) + env['PYTHONPATH'] = CODE_ROOT + os.pathsep + env.get('PYTHONPATH', '') + env['OMP_NUM_THREADS'] = '1' + env.setdefault('CUDA_VISIBLE_DEVICES', '') + proc = subprocess.Popen([sys.executable, DRIVER] + args, env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + out, _ = proc.communicate() + if not isinstance(out, str): + out = out.decode('utf-8', 'replace') + return re.sub(r'\s+', ' ', out) + + +_HONOURED_CLI = ['--vectorized', '--gpu', '--force-xpy', '--time-marginalization'] +_REFUSED = 'Refusing this calibration-marginalization configuration' + + +def test_driver_refuses_through_the_real_cli(): + """The predicate can be right and the wiring wrong; that is how the guards this + replaces became inert. Each case names the source evidence for its category.""" + cases = [ + (['--calibration-envelope-directory', ENV_DIR, '--rotation-slow'] + _HONOURED_CLI, + '--rotation-slow', 'rotation dispatch takes no n_cal'), + (['--calibration-envelope-directory', ENV_DIR, '--freqresponse'] + _HONOURED_CLI, + '--freqresponse', 'freqresponse dispatch takes no n_cal'), + (['--calibration-envelope-directory', ENV_DIR, '--vectorized', '--gpu', '--force-xpy'], + '--time-marginalization', 'FactoredLogLikelihood takes no n_cal'), + (['--calibration-envelope-directory', ENV_DIR, '--vectorized', '--time-marginalization'], + '--force-xpy', 'ViaArrayVector (no NoLoop) takes no n_cal'), + (['--calibration-fused-kernel'] + _HONOURED_CLI, + '--calibration-envelope-directory', 'nothing to fuse'), + (['--calibration-envelope-directory', ENV_DIR, + '--calibration-burn-in-nmax', '4000'] + _HONOURED_CLI, + '--calibration-burn-in-neff', 'no burn-in for the cap to cap'), + ] + for args, expect, why in cases: + out = _run(args) + assert _REFUSED in out and expect in out, \ + 'driver accepted %s (%s); output tail: %s' % (args, why, out[-600:]) + print('driver refuses %-34s : OK' % why) + + +def test_driver_accepts_the_honoured_calmarg_configuration(): + """THE regression test for over-broadness at the CLI. This is the calmarg demo's own + BACKEND=cpu command line; it must get past the gate (and fail later, on the data it + was not given).""" + out = _run(['--calibration-envelope-directory', ENV_DIR] + _HONOURED_CLI) + assert _REFUSED not in out, out[-600:] + print('driver accepts the honoured calmarg configuration : OK') + + +def test_driver_accepts_the_burn_in_cap_with_its_target(): + """The over-broadness half of R7 at the CLI: the cap IS legal alongside the burn-in + target, and a rule wired to the wrong attribute would refuse it here.""" + out = _run(['--calibration-envelope-directory', ENV_DIR, + '--calibration-burn-in-neff', '100', + '--calibration-burn-in-nmax', '4000'] + _HONOURED_CLI) + assert _REFUSED not in out, out[-600:] + print('driver accepts the burn-in cap with its target : OK') + + +def test_driver_accepts_the_pilot_without_time_marginalization(): + """The pilot's exemption, through the CLI. util_CalPilotStage.py depends on it.""" + out = _run(['--calibration-envelope-directory', ENV_DIR, '--vectorized', + '--calibration-dump-responsibilities', '/tmp/resp_unused.npz']) + assert _REFUSED not in out, out[-600:] + print('driver accepts the cal pilot on --vectorized alone : OK') + + +if __name__ == '__main__': + fails = 0 + for name, fn in sorted(globals().items()): + if name.startswith('test_') and callable(fn): + try: + fn() + except Exception as e: # noqa: BLE001 + fails += 1 + print('FAIL %s: %s' % (name, e)) + else: + print('ok %s' % name) + print('\n%s' % ('FAILED' if fails else 'PASS')) + sys.exit(1 if fails else 0) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py index ab4a447ed..c9ccea34b 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py @@ -1044,6 +1044,85 @@ def cosine_sampler_limits(lo, hi, kind): return z_lo, z_hi +### +### Distance: a SAMPLING-only restriction (--limit-distance) +### +# The angular zoom boxes above narrow the PRIOR SUPPORT: the angle priors are never +# renormalized to the box, so restricting one costs exactly the prior mass it should +# and lnZ moves by an analytic amount. --limit-distance is a DIFFERENT animal, and +# deliberately so: it narrows only the range the sampler draws from, while the +# distance prior keeps the normalization it has over the FULL physical [d_min,d_max]. +# The reported lnZ is then the same number the full-range run reports -- no +# correction, directly comparable across runs and across samplers -- to the extent +# the likelihood is negligible outside the box. That is the whole point: the +# distance posterior narrows as 1/rho, so at high amplitude a box tracking it costs +# nothing in evidence and buys back the resolution the quadrature was wasting. +# +# THE TRAP this pair of helpers exists to close: the natural implementation narrows +# param_limits["distance"], which the Euclidean prior lambda also reads for its own +# normalization -- so the density silently renormalizes to the box, lnZ comes back +# UNCHANGED at exactly the value that says the prior moved, and nothing looks wrong. +# Hence the two ranges are separate arguments here and cannot be conflated by +# accident: `sampling_range` is what the sampler draws from, `prior_range` is what +# the prior integrates to one over. + +def distance_limit_range(value, d_min, d_max): + """Parse and validate a --limit-distance 'LO,HI' request (Mpc). + + Returns (lo, hi) as floats with d_min <= lo < hi <= d_max. Raises ValueError + on an unparseable, non-finite, empty, inverted, or out-of-prior-range request: + a box outside [d_min,d_max] would ask the sampler for distances the prior gives + zero mass, which is a configuration error, not something to clip silently. + """ + try: + lo, hi = [float(_x) for _x in str(value).split(',')] + except ValueError: + raise ValueError("distance_limit_range: expected 'LO,HI' in Mpc, got '{}'".format(value)) + if not numpy.isfinite(lo) or not numpy.isfinite(hi): + raise ValueError("distance_limit_range: non-finite distance range [{}, {}]".format(lo, hi)) + if not (hi > lo): + raise ValueError("distance_limit_range: empty or inverted distance range [{}, {}] (need LO < HI, in Mpc)".format(lo, hi)) + d_min = float(d_min) + d_max = float(d_max) + if lo < d_min or hi > d_max: + raise ValueError( + "distance_limit_range: requested sampling range [{}, {}] Mpc is not inside the prior range [{}, {}] Mpc " + "(--d-min/--d-max). --limit-distance narrows the SAMPLING only; it cannot extend the prior.".format( + lo, hi, d_min, d_max)) + return lo, hi + + +def distance_sampler_kwargs(sampler_module, sampling_range, prior_range, + d_prior='Euclidean', xpy=numpy, adaptive_sampling=False): + """Build the add_parameter() kwargs for the distance coordinate. + + `sampler_module` supplies the backend-appropriate uniform sampler helpers + (mcsampler / mcsamplerGPU / mcsamplerAdaptiveVolume all export them, and the + GPU ones return device arrays). `sampling_range` sets the sampler's proposal + and its (llim, rlim); `prior_range` sets the normalization of `prior_pdf` and + is the ONLY thing that determines the reported evidence scale. Pass them equal + for the historical behaviour. + """ + lo, hi = float(sampling_range[0]), float(sampling_range[1]) + p_lo, p_hi = float(prior_range[0]), float(prior_range[1]) + if d_prior == 'pseudo_cosmo': + import RIFT.likelihood.priors_utils as _priors_utils + nm = _priors_utils.dist_prior_pseudo_cosmo_eval_norm(p_lo, p_hi) + prior_pdf = functools.partial(_priors_utils.dist_prior_pseudo_cosmo, nm=nm, xpy=xpy) + elif d_prior == 'Euclidean': + prior_pdf = lambda x: x**2/(p_hi**3/3. - p_lo**3/3.) + else: + raise ValueError("distance_sampler_kwargs: unknown distance prior '{}'".format(d_prior)) + return dict( + pdf=sampler_module.ret_uniform_samp_vector_alt(lo, hi), + cdf_inv=functools.partial(sampler_module.uniform_samp_cdf_inv_vector, lo, hi), + left_limit=lo, + right_limit=hi, + prior_pdf=prior_pdf, + adaptive_sampling=adaptive_sampling, + ) + + def ret_dec_samp_vector(dec_lo, dec_hi): """Sampling pdf in DECLINATION for a uniform-in-sin(dec) draw truncated to [dec_lo, dec_hi]. Normalized to unity over that box (the samplers that use diff --git a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/distance_tail.py b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/distance_tail.py new file mode 100644 index 000000000..0e75d6224 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/distance_tail.py @@ -0,0 +1,279 @@ +#!/usr/bin/env python +"""distance_tail_fit.py -- make the CIP fit fall off in distance beyond the exported support. + +THE DEFECT. An ExtraTrees/RF ensemble is piecewise constant outside its training envelope. The +.dslice export trains it on ~50 discrete distances per intrinsic point, so beyond that point's own +outermost slice the ensemble returns the EDGE value of lnL forever instead of letting it decay. +Measured directly (tools/test_distance_extrapolation.py): the held-out lnL bias grows with how far +past the training edge you ask, +0.00 / +0.14 / +0.84 nats at reach <5% / 5-15% / >15%. + +WHY IT MATTERS SO MUCH. The recovered distance posterior is exp(lnL) * pi(d), and pi(d) is +volumetric -- it GROWS like d^2. Hold lnL flat and the integrand grows without bound out to +whatever --d-max the CIP range allows, so probability that should have died survives, and the +posterior comes out ~18% too wide in every quantile span including the core, median untouched. +Confirmed at Spearman rho=+0.952 (p=4e-5) between an event's off-support sample fraction and its +width excess. Confining d to the exported support removes 18 of those 18 points. + +THIS IS NOT A COORDINATE PROBLEM. Fitting in 1/d instead of d was tried and made things WORSE +(off-support mass 0.098 -> 0.194). A tree extrapolates flat in whatever coordinate you give it; +changing coordinates only moves where "flat" lands. What is missing is a BOUNDARY CONDITION. + +-------------------------------------------------------------------------------------------------- +WHAT WAS TRIED FIRST AND REJECTED, because it looks right on paper and fails on the data. + +For fixed intrinsic parameters AND fixed extrinsic angles, RIFT's log likelihood ratio is exactly +quadratic in x = 1/d through the origin, lnL = a x - b x^2/2, since the signal enters as h/d. That +suggests fitting each slice with a global through-the-origin polynomial: exact at d -> infinity by +construction, which is precisely where the tree fails. + +It does not work, and the reason is instructive. The export marginalises over the extrinsic angles, +and the distance-inclination degeneracy makes the MARGINALISED lnL nearly FLAT across the whole +exported support -- on S240615ea, 14.3 down to 13.1 while d runs 1085 -> 4909 Mpc. A low-order +polynomial through the origin cannot be both flat over the support and zero at x=0: the per-slice +residual is 1.7 nats at order 2 and 0.55 at order 3, against a per-slice noise of ~0.13. Forcing it +anyway drives the conditional width 58% LOW. `polyfit_slice` below is retained as the diagnostic +that measures this, not as the fix. + +The flatness is also the reason the defect is severe rather than subtle: the fit reaches the edge of +its support with lnL still near its peak, so "hold the last value" is not a small error. + +-------------------------------------------------------------------------------------------------- +THE FIX. Keep the base fit -- it is right ON the support, which the `nnexp` rung established +independently -- and give it a tail. Write x = 1/d and u = x/x_edge, and continue each slice past +its outermost exported distance with the exact single-angle form THROUGH THE ORIGIN, matched to the +value and the slope the data actually has at that edge: + + lnL(u) / lnL_edge = (2 - s) u + (s - 1) u^2, s = x_edge * dlnL/dx |_edge / lnL_edge + +s is the slice's dimensionless log-slope at its edge, and it is the whole content of the model: + + s = 0 (slice still FLAT at its edge, the common case) -> 2u - u^2, a gentle roll-off + s = 1 (slice already decaying linearly in 1/d) -> u, the pure asymptotic form + +Clamped to s in [0, 2), which is exactly the range over which the continuation is monotone in u, +so walking out in distance always walks lnL down, never up. + +WHY VALUE-AND-SLOPE AND NOT A FITTED POLYNOMIAL. Fitting a through-the-origin quadratic to the +outer part of a slice directly was tried first and over-corrects: over a narrow span of x the +origin constraint sits far outside the data, so the fit has an enormous extrapolation lever and +comes back much steeper than the slice really is. On the catalog it recovered only ~50% of the +available width error and made the two events that had NO defect measurably worse. Matching the +value and a locally regressed slope has no lever: a flat slice provably gets the gentle 2u - u^2 +roll-off, because that is the unique through-the-origin quadratic that is flat at the edge. + +Three properties, and all three matter: + + * CONTINUOUS. The ratio is 1 at u=1, so this changes nothing on the support and introduces no + step that the sampler would read as structure. + * CORRECT LIMIT. The form is through the origin, so lnL -> 0 as d -> infinity. RIFT's lnL is a + likelihood RATIO, so an infinitely distant source is exactly the noise hypothesis: this limit + is an identity, not a modelling choice. + * LOCALLY MATCHED. A flat-edged slice rolls off gently rather than being guillotined, and a + slice that is already decaying is continued at its own rate. + +The decay still comfortably beats the d^2 prior, which is all that is needed to remove the spurious +mass: at twice the edge distance (u=1/2) a flat-edged slice retains 0.75 of its edge lnL, so for a +typical lnL_edge ~ 13 that is a 3.3-nat drop against a 1.4-nat prior gain. The integrand falls +instead of growing. + +WHAT THIS DOES NOT CLAIM. Beyond the support the data does not constrain the turnover, so the tail +is a physically-motivated continuation, not a measurement. Its job is to be DECAYING and correct in +the d -> infinity limit, not to be quantitatively exact where nothing was exported. The honest test +is whether the recovered posterior matches the reference; that is what the validation check scores. +The near-d end is left to the base fit: it is bounded by the CIP distance range and suppressed by +the d^2 prior, and it carries no measured defect. +""" +import numpy as np + +__all__ = ["wrap_distance_tail", "polyfit_slice", "slice_fit_report"] + +MIN_DISTINCT_D = 6 # below this a slice cannot support the tail fit +OUTER_FRAC = 0.5 # fraction of each slice, from the FAR end, used to fit the tail + + +def polyfit_slice(d, y, y_err=None, order=2, sel=None): + """WLS fit of y(x) = sum_{k=1..order} c_k x^k, x = 1/d, THROUGH THE ORIGIN. + + Returns (c, c_err, rms_resid, n) or None. The intercept is not free: lnL(x=0)=0 is a physical + identity for a likelihood ratio, and letting it float discards the only information the tree is + missing. `sel` restricts the fit to a subset of the points (used to fit only the outer, small-x + end, where the local expansion is good and where extrapolation will actually happen). + """ + d = np.asarray(d, dtype=float) + y = np.asarray(y, dtype=float) + ok = np.isfinite(d) & np.isfinite(y) & (d > 0) + if sel is not None: + ok &= sel + if ok.sum() < order + 2 or len(np.unique(d[ok])) < MIN_DISTINCT_D: + return None + d, y = d[ok], y[ok] + x = 1.0 / d + M = np.stack([x ** k for k in range(1, order + 1)], axis=1) + if y_err is None: + w = np.ones(len(y)) + else: + w = 1.0 / np.maximum(np.asarray(y_err, dtype=float)[ok], 1e-3) ** 2 + Mw = M * w[:, None] + try: + cov = np.linalg.inv(M.T @ Mw) + except np.linalg.LinAlgError: + return None + c = cov @ (Mw.T @ y) + r = y - M @ c + dof = max(len(y) - order, 1) + chi2_red = float(np.sum(w * r ** 2) / dof) + c_err = np.sqrt(np.maximum(np.diag(cov), 0.0) * max(chi2_red, 1e-12)) + return c, c_err, float(np.sqrt(np.mean(r ** 2))), len(y) + + +def slice_fit_report(d, y, y_err=None, orders=(1, 2, 3)): + """{order: rms residual} for a through-the-origin polynomial on one slice. Compare against the + slice's own lnL noise: this is what showed the GLOBAL parametric route to be untenable.""" + return {o: (None if (f := polyfit_slice(d, y, y_err, order=o)) is None else f[2]) + for o in orders} + + +def _slice_index(X_int): + """Group rows by identical intrinsic coordinates. The export holds the intrinsic point fixed + and varies only distance, so rows within a slice match exactly and np.unique on the raw rows is + correct and cheap. Rounding would be wrong: two genuinely distinct grid points can be + arbitrarily close in a dense grid.""" + _, inv = np.unique(X_int, axis=0, return_inverse=True) + return inv + + +def wrap_distance_tail(base_fit, X, Y, coord_names, y_errors=None, dist_name="dist", + outer_frac=OUTER_FRAC, lnL_offset=0.0, law="chord", power=None, report=None): + """Wrap a fitted CIP callable so it decays beyond each intrinsic point's exported distance + support instead of holding flat. + + base_fit the already-fitted callable from fit_rf / fit_gp / ... : f(X) -> lnL + X, Y the SAME training arrays that base_fit was fitted on, in the CIP fit basis + coord_names names of X's columns; must contain `dist_name` + y_errors per-row lnL errors, used to weight the per-slice tail fit + lnL_offset CIP fits Y = lnL_physical - lnL_shift, so pass lnL_shift here. The lnL(d->inf)=0 + identity holds for the PHYSICAL likelihood ratio, so the tail must be built and + applied on lnL + lnL_offset and the offset removed again on return. Leaving this + at 0 when CIP has applied a shift anchors the decay to the wrong asymptote -- + silently, and in the direction that reintroduces the bug. + report optional dict, filled with diagnostics + + Returns a callable with the same signature as base_fit. On the support it returns base_fit + unchanged, so this is a strict addition: nothing that currently works is altered. + """ + from scipy.spatial import cKDTree + + names = list(coord_names) + if dist_name not in names: + raise ValueError("distance tail fix needs a %r column in the fit basis; got %r" + % (dist_name, names)) + idist = names.index(dist_name) + # 'inv_dist' is a redundant reparametrisation of the same axis; it must not enter the intrinsic + # key or two rows of one slice would look like two different intrinsic points. + keep = [i for i, n in enumerate(names) if n not in (dist_name, "inv_dist")] + + X = np.asarray(X, dtype=float) + Y = np.asarray(Y, dtype=float) + d = X[:, idist] + sig = None if y_errors is None else np.asarray(y_errors, dtype=float) + Xint = X[:, keep] + inv = _slice_index(Xint) + nsl = int(inv.max()) + 1 if len(inv) else 0 + + cent, dmax, slope = [], [], [] + n_flat = 0 + for si in range(nsl): + m = np.where(inv == si)[0] + ds = d[m] + if len(m) < MIN_DISTINCT_D or len(np.unique(ds)) < MIN_DISTINCT_D: + continue + # Outer end in DISTANCE == small x, which is the side we extrapolate to. The slope is fit + # with a FREE intercept: we want the local gradient, and forcing this local regression + # through the origin is exactly the lever that made the earlier version over-correct. + thr = np.quantile(ds, 1.0 - outer_frac) + sel = ds >= thr + if sel.sum() < 3 or len(np.unique(ds[sel])) < 3: + sel = np.ones(len(ds), bool) + xs = 1.0 / ds[sel] + ys = Y[m][sel] + lnL_offset + w = (np.ones(sel.sum()) if sig is None + else 1.0 / np.maximum(sig[m][sel], 1e-3) ** 2) + # weighted linear regression ys ~ p + q xs + sw = w.sum() + mx = float((w * xs).sum() / sw) + my = float((w * ys).sum() / sw) + vxx = float((w * (xs - mx) ** 2).sum()) + if vxx <= 0: + continue + q = float((w * (xs - mx) * (ys - my)).sum() / vxx) + if q <= 0: + n_flat += 1 + cent.append(Xint[m][0]); dmax.append(ds.max()); slope.append(q) + + if len(cent) < 8: + raise ValueError("distance tail fix: only %d usable slices (need >=8). Is this a dslice " + "export, and did the intrinsic grouping work?" % len(cent)) + + cent = np.asarray(cent); dmax = np.asarray(dmax); slope = np.asarray(slope) + # Nearest slice in intrinsic space, standardised so no coordinate dominates the metric by units + # alone (mc is O(30), spins O(1)). + scale = cent.std(axis=0) + scale[scale <= 0] = 1.0 + tree = cKDTree(cent / scale) + + if report is not None: + report.update(n_slices_total=int(nsl), n_slices_used=int(len(cent)), law=law, + n_flat_or_rising_edge=int(n_flat), outer_frac=float(outer_frac), + dmax_median=float(np.median(dmax))) + + def fn_return(x_in): + x_in = np.asarray(x_in, dtype=float) + out = np.asarray(base_fit(x_in), dtype=float).copy() + dq = x_in[:, idist] + ok = np.isfinite(dq) & (dq > 0) & np.all(np.isfinite(x_in), axis=-1) + if not ok.any(): + return out + idx = np.where(ok)[0] + _, j = tree.query(x_in[idx][:, keep] / scale, k=1) + de = dmax[j] + beyond = dq[idx] > de + if not beyond.any(): + return out + sel = idx[beyond] + j = j[beyond]; de = de[beyond] + # value at the edge, from the base fit itself (denoised), same intrinsic coordinates + Xe = x_in[sel].copy() + Xe[:, idist] = de + if "inv_dist" in names: + Xe[:, names.index("inv_dist")] = 1.0 / de + Le = np.asarray(base_fit(Xe), dtype=float) + lnL_offset + xq = 1.0 / dq[sel] + xe = 1.0 / de + u = xq / xe + # THE DECAY LAW. Measured, not assumed (see the module docstring): + # "chord" ratio = u. Parameter-free. lnL is convex in x=1/d wherever Var(a) > for + # the marginalised likelihood, and a convex function through the origin lies + # BELOW its chord, so u is an upper bound on lnL beyond the edge as well as the + # exact asymptotic form. This is the default and the only law used in production. + # "slope" the through-the-origin quadratic matched to the slice's local value AND slope. + # Faithful to the data at the edge and provably gentle -- and that is exactly why + # it fails: the marginalised lnL is a plateau across the whole exported support, + # so the local slope cannot see the turnover, and this law barely decays at all. + # Kept because measuring that is what ruled it out. + # power=p DIAGNOSTIC ONLY, not a production setting. Scanning p says whether the best + # decay sits at the principled p=1 or only at a tuned value; a fix that needs + # tuning to the reference is not a fix. Never set this in a production run. + if power is not None: + ratio = u ** float(power) + elif law == "chord": + ratio = u + else: + sdl = np.clip(slope[j] * xe / np.where(np.abs(Le) < 1e-12, 1e-12, Le), 0.0, 1.999) + ratio = (2.0 - sdl) * u + (sdl - 1.0) * u ** 2 + # Guard the continuation rather than trusting it: it may only shrink the edge value, never + # grow it or flip its sign. Without this a pathological slice could turn the tail fix into + # a second, worse extrapolation bug. + out[sel] = Le * np.clip(ratio, 0.0, 1.0) - lnL_offset + return out + + return fn_return diff --git a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py index ead18a87d..d90fc5a0c 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py @@ -21,6 +21,7 @@ import sys import copy import types +from RIFT.physics import teobresums_compat has_external_teobresum=False import os info_use_ext = True @@ -47,6 +48,7 @@ from numpy import sin, cos from scipy import interpolate from scipy import signal +from scipy.optimize import fsolve import scipy # for decimate try: import precession @@ -86,6 +88,32 @@ def safe_int(mystr): rosDebugMessagesContainer = [False] + + +def _hyperbolic_endpoint_outcome(dynamics, normalized_amplitude): + """Distinguish an outgoing scatter from an ingoing plunge. + + Prefer the final radial momentum from TEOBResumS dynamics. The fallback + uses only the endpoint-to-peak amplitude ratio, so the result cannot + depend on the arbitrary luminosity distance assigned to the parameters. + """ + if isinstance(dynamics, dict) and "Prstar" in dynamics: + radial_momentum = np.asarray(dynamics["Prstar"]) + if radial_momentum.size: + final_momentum = radial_momentum.flat[-1] + if np.isfinite(final_momentum) and final_momentum != 0: + return "scatter" if final_momentum > 0 else "plunge" + + normalized_amplitude = np.asarray(normalized_amplitude) + if not normalized_amplitude.size or not np.isfinite(normalized_amplitude[-1]): + return "plunge" + return "scatter" if normalized_amplitude[-1] > 1e-4 else "plunge" + + +def _zero_mode_data(modes): + """Zero every time-series mode in place.""" + for series in modes.values(): + series.data.data *= 0.0 rosDebugMessagesLongContainer = [False] if log_loud: print( "[Loading lalsimutils.py : MonteCarloMarginalization version]",file=sys.stderr) @@ -333,6 +361,31 @@ def check_FD_pending(code): def modes_to_k(modes): return [int(x[0]*(x[0]-1)/2 + x[1]-2) for x in modes] +def mean_anomaly_from_true(true_anomaly, eccentricity): + return ( + np.arctan2( + -np.sqrt( + 1 - eccentricity**2) + * np.sin(true_anomaly), + - eccentricity - np.cos(true_anomaly) + ) + np.pi + - eccentricity + * np.sqrt(1 - eccentricity**2) + * np.sin(true_anomaly) + / ( + 1 + eccentricity + * np.cos(true_anomaly) + ) + ) +def eccentric_anomaly_from_mean(mean_anomaly, eccentricity): + func = lambda E : E - eccentricity*np.sin(E) - mean_anomaly + return fsolve(func, x0=mean_anomaly) +def true_anomaly_from_eccentric(eccentric_anomaly, eccentricity): + # formula from https://ui.adsabs.harvard.edu/abs/1973CeMec...7..388B/abstract + # avoids numerical issues + + beta = eccentricity / (1+np.sqrt(1 - eccentricity**2)) + return eccentric_anomaly + 2 * np.arctan2(beta * np.sin(eccentric_anomaly), 1 - beta * np.cos(eccentric_anomaly)) # https://www.lsc-group.phys.uwm.edu/daswg/projects/lal/nightly/docs/html/_l_a_l_sim_inspiral_8c_source.html#l02910 def lsu_StringFromPNOrder(order): @@ -361,7 +414,7 @@ def lsu_StringFromPNOrder(order): # Class to hold arguments of ChooseWaveform functions # -valid_params = ['m1', 'm2', 's1x', 's1y', 's1z', 's2x', 's2y', 's2z', 'chi1_perp', 'chi2_perp', 'chi1_perp_bar', 'chi2_perp_bar','chi1_perp_u', 'chi2_perp_u', 's1z_bar', 's2z_bar', 'lambda1', 'lambda2', 'theta','phi', 'phiref', 'psi', 'incl', 'tref', 'dist', 'mc', 'mc_ecc', 'eta', 'delta_mc', 'chi1', 'chi2', 'thetaJN', 'phiJL', 'theta1', 'theta2', 'cos_theta1', 'cos_theta2', 'theta1_Jfix', 'theta2_Jfix', 'psiJ', 'beta', 'cos_beta', 'sin_phiJL', 'cos_phiJL', 'phi12', 'phi1', 'phi2', 'LambdaTilde', 'DeltaLambdaTilde', 'lambda_plus', 'lambda_minus', 'q', 'mtot','xi','chiz_plus', 'chiz_minus', 'chieff_aligned','fmin','fref', "SOverM2_perp", "SOverM2_L", "DeltaOverM2_perp", "DeltaOverM2_L", "shu","ampO", "phaseO",'eccentricity','eccentricity_squared','eccentricity_ln', 'chi_pavg','mu1','mu2','eos_table_index','meanPerAno'] +valid_params = ['m1', 'm2', 's1x', 's1y', 's1z', 's2x', 's2y', 's2z', 'chi1_perp', 'chi2_perp', 'chi1_perp_bar', 'chi2_perp_bar','chi1_perp_u', 'chi2_perp_u', 's1z_bar', 's2z_bar', 'lambda1', 'lambda2', 'theta','phi', 'phiref', 'psi', 'incl', 'tref', 'dist', 'mc', 'mc_ecc', 'eta', 'delta_mc', 'chi1', 'chi2', 'thetaJN', 'phiJL', 'theta1', 'theta2', 'cos_theta1', 'cos_theta2', 'theta1_Jfix', 'theta2_Jfix', 'psiJ', 'beta', 'cos_beta', 'sin_phiJL', 'cos_phiJL', 'phi12', 'phi1', 'phi2', 'LambdaTilde', 'DeltaLambdaTilde', 'lambda_plus', 'lambda_minus', 'q', 'mtot','xi','chiz_plus', 'chiz_minus', 'chieff_aligned','fmin','fref', "SOverM2_perp", "SOverM2_L", "DeltaOverM2_perp", "DeltaOverM2_L", "shu","ampO", "phaseO",'eccentricity','eccentricity_squared','eccentricity_ln', 'chi_pavg','mu1','mu2','eos_table_index','meanPerAno','a6c','E0','p_phi0','hypclass'] # so far, used for puffball, to prevent insanity (infinite growth) and/or death to downselect # - note we also provide for extrinsic: RA (phi), phiref, psi, just in case we need it in the future @@ -407,6 +460,9 @@ def lsu_StringFromPNOrder(order): "s2y": r"$\chi_{2,y}$", "eccentricity":"$e$", "meanPerAno":"$l_{gw}$", + "E0" : "$E_0 / M$", + "p_phi0" : r"$p_{\phi}^0$", + "a6c":"$a^c_6$", # tex labels for inherited LI names "a1z": r'$\chi_{1,z}$', "a2z": r'$\chi_{2,z}$', @@ -455,8 +511,11 @@ def __init__(self, phiref=0., deltaT=1./4096., m1=10.*lsu_MSUN, deltaF=None, fmax=0., # for use w/ FD approximants taper=lsu_TAPER_NONE, # for use w/TD approximants eccentricity=0., # make eccentricity a parameter - meanPerAno=0. # make meanPerAno a parameter - ): + meanPerAno=0., # make meanPerAno a parameter + E0=0., # make E0/M a parameter + p_phi0=0., # make j_hyp a parameter + a6c=10000. # EOB Parameter + ): self.phiref = phiref self.deltaT = deltaT self.m1 = m1 @@ -481,10 +540,12 @@ def __init__(self, phiref=0., deltaT=1./4096., m1=10.*lsu_MSUN, self.theta = theta # DEC. DEC =0 on the equator; the south pole has DEC = - pi/2 self.phi = phi # RA. self.psi = psi - self.meanPerAno = 0.0 # port + self.a6c=a6c self.longAscNodes = self.psi # port to master self.eccentricity=eccentricity self.meanPerAno=meanPerAno + self.E0 = E0 + self.p_phi0 = p_phi0 self.tref = tref self.radec = radec self.detector = "H1" @@ -922,6 +983,119 @@ def extract_param(self,p): return (self.m2+self.m1) if p == 'q': return self.m2/self.m1 + if p == 'hypclass': + # Checks type of hyperbolic waveform: scatter, plunge, zoomwhirl, or meaningless + # check if valid + if self.E0 == 0.0: + print('Invalid use of hypclass: non-hyperbolic configuration') + return None + + # Generate waveform + pars = { + 'M' : (self.m1+self.m2)/lal.MSUN_SI, + 'q' : self.m1/self.m2, + 'H_hyp' : self.E0, # energy at initial separation + 'j_hyp' : self.p_phi0, # angular momentum at initial separation + 'r_hyp' : 6000, + 'LambdaAl2' : self.lambda1, + 'LambdaBl2' : self.lambda2, + 'chi1' : self.s1z, + 'chi2' : self.s2z, + 'dt' : self.deltaT, + 'domain' : 0, # 0 sets time domain + 'arg_out' : 'yes', # Request multiples and dynamics as output + 'nqc' : 'no', + 'nqc_coefs_hlm' : 'none', + 'nqc_coefs_flx' : 'none', + 'use_mode_lm' : [1], # 22 mode + 'output_lm' : [1], + 'srate_interp' : 1./self.deltaT, + 'use_geometric_units': 'no', + 'interp_uniform_grid': 'yes', + 'initial_frequency' : self.fmin, + 'ode_tmax' : 3e4, + 'distance' : self.dist/(lal.PC_SI*1e6), + 'inclination' : self.incl, + 'output_hpc' : 'no' + } + + t, hptmp, hctmp, hlmtmp, dym = teobresums_compat.run( + EOBRun_module, pars, purpose="hyperbolic_classification" + ) + + # peak finding to classify + + # wf amplitude - 22 mode only + tmp_22 = np.array(hlmtmp['1']) + distance_s = self.dist/lal.C_SI + m_total_s = MsunInSec*(self.m1+self.m2)/lal.MSUN_SI + M1=self.m1/lal.MSUN_SI + M2=self.m2/lal.MSUN_SI + nu=M1*M2/((M1+M2)**2) + tmp_22[0] *= (m_total_s/distance_s)*nu + amp = np.abs(tmp_22[0] * np.exp(-1j*(2*(np.pi/2.)+tmp_22[1]))) + amp_norm = amp / np.amax(amp) # normalize amplitude for peak finding + + # Check for cases where the amplitude is max at the start or end + if np.argmax(amp_norm) == 0 or np.argmax(amp_norm) == len(amp_norm) - 1: + reclassify = True + else: + reclassify = False + # peak finding to determine system type + height_thresh = 0.25 + prom_thresh = 0.1 + peaks, props = signal.find_peaks(amp_norm, height = height_thresh, prominence = prom_thresh) + peak_heights = props['peak_heights'] + # filtering out peaks so we only keep the local maxima + indices_to_keep = set() + sorted_indices = np.argsort(peak_heights)[::-1] + tol = int(pars['srate_interp'] / 13.65) # 300 samples at srate of 4096 - minimum distance between peaks. + for i in sorted_indices: + peak = peaks[i] + keep = True + for kept_index in indices_to_keep: + if abs(peaks[kept_index] - peak) <= tol: + keep = False + break + if keep: + indices_to_keep.add(i) + # sorted(): indices_to_keep is a set, and downstream epoch choices + # need the retained peaks in chronological order + filtered_peaks = peaks[sorted(indices_to_keep)] + + # parsing number of peaks after filtering against distance tolerance + if len(filtered_peaks) == 1: + return _hyperbolic_endpoint_outcome(dym, amp_norm) + elif len(filtered_peaks) == 0: + # meaningless waveform + reclassify = True + else: + # zoom whirl waveform + return 'zoomwhirl' + + if reclassify: + print('Running re-classification') + # run minimal threshold peak finder + all_peaks, all_props = signal.find_peaks(amp_norm, height=0.0001, prominence=0.0001) + + # checking for minima if no peaks found + if len(all_peaks) == 0: + print("No peaks found, checking for minima instead...") + all_peaks, all_props = signal.find_peaks((-1*amp_norm + 1.0), height=0.0001, prominence=0.0001) + + if len(all_props['prominences']) > 3: + print('MANY peaks detected on reclassification, evaluating...') + # these can be scatter or plunge + outcome = _hyperbolic_endpoint_outcome(dym, amp_norm) + print('Reclassifying to {}'.format(outcome.capitalize())) + return outcome + elif len(all_props['prominences']) == 3 or len(all_props['prominences']) == 2 or len(all_props['prominences']) == 1: + # these are always scatters + print('Reclassifying to Scatter') + return 'scatter' + else: + print('No peaks detected after reclassifcation') + return 'meaningless' if p == 'delta' or p=='delta_mc': # Same access routine return (self.m1-self.m2)/(self.m1+self.m2) if p == 'mc': @@ -1577,6 +1751,7 @@ def print_params(self,show_system_frame=False): Lhat = np.array( [np.sin(self.incl),0,np.cos(self.incl)]) # does NOT correct for psi polar anogle! Uses OLD convention for spins! print( " : hat(L). s1 x s2 = ", vecDot( Lhat, vecCross([self.s1x,self.s1y,self.s1z],[self.s2x,self.s2y,self.s2z]))) print( " : hat(L).(S1(1+q)+S2(1+1/q)) = ", vecDot( Lhat, S1vec*(1+qval) + S2vec*(1+1./qval) )/(self.m1+self.m2)/(self.m1+self.m2)) + print( " chi_p = ", self.extract_param('chi_p')) if show_system_frame: thePrefix = "" thetaJN, phiJL, theta1, theta2, phi12, chi1, chi2, psiJ = self.extract_system_frame() @@ -1591,12 +1766,16 @@ def print_params(self,show_system_frame=False): print( thePrefix, " :+ beta = ", self.extract_param('beta')) print( "lambda1 =", self.lambda1) print( "lambda2 =", self.lambda2) + print("EOB Parameters:") + print("a6c = ",self.a6c) print( "inclination =", self.incl) print( "distance =", self.dist / 1.e+6 / lsu_PC, "(Mpc)") print( "reference orbital phase =", self.phiref) print( "polarization angle =", self.psi) print( "eccentricity = ", self.eccentricity) print( "meanPerAno = ", self.meanPerAno) + print("E0 / M = ", self.E0) + print("j_hyp = ", self.p_phi0) print( "time of coalescence =", float(self.tref), " [GPS sec: ", int(self.tref), ", GPS ns ", (self.tref - int(self.tref))*1e9, "]") print( "detector is:", self.detector) if self.radec==False: @@ -1785,6 +1964,9 @@ def copy_sim_inspiral(self, row): # FAKED COLUMNS (nonstandard) self.lambda1 = row.alpha5 self.lambda2 = row.alpha6 + self.E0 = row.psi3 # hyperbolic parameter + self.p_phi0 = row.beta # hyperbolic parameter + self.a6c = row.psi0 self.eccentricity=row.alpha4 self.meanPerAno=row.alpha self.snr = row.alpha3 # lnL info @@ -1850,6 +2032,9 @@ def create_sim_inspiral(self): # NONSTANDARD row.alpha5 = self.lambda1 row.alpha6 = self.lambda2 + row.psi0 = self.a6c + row.psi3 = self.E0 + row.beta = self.p_phi0 row.alpha4 = self.eccentricity row.alpha = self.meanPerAno if self.eos_table_index and not self.eccentricity: @@ -1910,14 +2095,15 @@ def copy_lsctables_sim_inspiral(self, row): # Call the function to read lalmetaio.SimInspiral format self.copy_sim_inspiral(swigrow) - def scale_to_snr(self,new_SNR,psd, ifo_list,analyticPSD_Q=True): + def scale_to_snr(self,new_SNR,psd, ifo_list,analyticPSD_Q=True, **kwargs): """ scale_to_snr - evaluates network SNR in the ifo list provided (assuming *constant* psd for all..may change) - uses network SNR to rescale the distance of the source, so the SNR is now new_SNR - returns current_SNR, for sanity """ - deltaF=findDeltaF(self) + Lmax = kwargs.get('Lmax', 4) # Default to 4 if not specified in kwargs + deltaF=findDeltaF(self, Lmax=Lmax) det_orig = self.detector IP = Overlap(fLow=self.fmin, fNyq=1./self.deltaT/2., deltaF=deltaF, psd=psd, full_output=True,analyticPSD_Q=analyticPSD_Q) @@ -1926,7 +2112,7 @@ def scale_to_snr(self,new_SNR,psd, ifo_list,analyticPSD_Q=True): for det in ifo_list: self.detector = det self.radec = True - h=hoff(self) + h=hoff(self, Lmax=Lmax) rho_ifo[det] = IP.norm(h) current_SNR_squared +=rho_ifo[det]*rho_ifo[det] current_SNR = np.sqrt(current_SNR_squared) @@ -2305,6 +2491,7 @@ def ip(self, h1, h2,include_epoch_differences=False): Compute inner product between two COMPLEX16Frequency Series Accounts for time shfit """ +# print(h1.data.length,h2.data.length,self.len2side,self.fNyq) assert h1.data.length==h2.data.length==self.len2side assert abs(h1.deltaF-h2.deltaF) <= TOL_DF\ and abs(h1.deltaF-self.deltaF) <= TOL_DF @@ -2832,7 +3019,7 @@ def nextPow2(length): """ return int(2**np.ceil(np.log2(length))) -def findDeltaF(P): +def findDeltaF(P,**kwargs): """ Given ChooseWaveformParams P, generate the TD waveform, round the length to the next power of 2, @@ -2840,7 +3027,8 @@ def findDeltaF(P): This is useful b/c deltaF is needed to define an inner product which is needed for norm_hoft and norm_hoff functions """ - h = hoft(P) + Lmax = kwargs.get('Lmax', 4) # Default to 4 if not specified in kwargs + h = hoft(P,Lmax=Lmax) return 1./(nextPow2(h.data.length) * P.deltaT) def estimateWaveformDuration(P,LmaxEff=2): @@ -2922,73 +3110,173 @@ def hoft(P, Fp=None, Fc=None,**kwargs): extra_waveform_args.update(kwargs['extra_waveform_args']) extra_params = P.to_lal_dict_extended(extra_args_dict=extra_waveform_args) if P.approx==lalsim.TEOBResumS and has_external_teobresum and info_use_ext: - Lmax=8 + Lmax = kwargs.get('Lmax', 4) # Default to 4 if not specified in kwargs modes_used = [] distance_s = P.dist/lal.C_SI m_total_s = MsunInSec*(P.m1+P.m2)/lal.MSUN_SI + k_coprecessing_frame_check = [1, 0, 4, 8] for l in np.arange(2,Lmax+1,1): for m in np.arange(0,l+1,1): if m !=0: modes_used.append((l,m)) + print(modes_used) k = modes_to_k(modes_used) + k_coprecessing_frame = [] + for count,value in enumerate(k): + if value in k_coprecessing_frame_check: + k_coprecessing_frame.append(value) + print(" k inertial modes: ", k, "k coprecessing frames: ", k_coprecessing_frame) + if kwargs.get('force_22_mode', False): + # Restrict EVERY requested list. The precessing branch below asks the + # backend for inertial modes through `k`, so narrowing only the + # coprecessing list still returns all modes through Lmax. + modes_used = [(2,2)] + k = modes_to_k(modes_used) + k_coprecessing_frame = list(k) + print("Forcing ONLY the 22 modes, so k inertial modes: ", k, " k coprecessing frames: ", k_coprecessing_frame) M1=P.m1/lal.MSUN_SI M2=P.m2/lal.MSUN_SI nu=M1*M2/((M1+M2)**2) - if (P.eccentricity == 0.0): - print("Using ResumS master; not eccentric") + hyp_wav = False + ecc_ano=eccentric_anomaly_from_mean(P.meanPerAno,P.eccentricity) + true_ano=true_anomaly_from_eccentric(ecc_ano,P.eccentricity) + if P.E0 == 0.0: + print("Using TEOBResumSDALI/GIOTTO standard call") pars = { 'M' : M1+M2, 'q' : M1/M2, - 'LambdaAl2' : P.lambda1, - 'LambdaBl2' : P.lambda2, + 'LambdaAl2' : P.lambda1, + 'LambdaBl2' : P.lambda2, 'chi1x' : P.s1x, 'chi1y' : P.s1y, 'chi1z' : P.s1z, 'chi2x' : P.s2x, 'chi2y' : P.s2y, 'chi2z' : P.s2z, - 'domain' : 0, - 'arg_out' : "yes", - 'use_mode_lm' : k, + 'ecc' : P.eccentricity, + 'inclination' : P.incl, + 'coalescence_angle' : np.pi / 2 - P.phiref, 'srate_interp' : 1./P.deltaT, - 'use_geometric_units': "no", 'initial_frequency' : P.fmin, - 'interp_uniform_grid': "yes", 'distance' : P.dist/(lal.PC_SI*1e6), - 'inclination' : P.incl, - "coalescence_angle": np.pi / 2 - P.phiref, + 'anomaly' : true_ano, +# 'a6c' : P.a6c, + 'domain' : 0, + 'arg_out' : "yes", + 'interp_uniform_grid': "yes", + 'use_geometric_units': "no", + # spin_flx can be "EOB" or "PN" prescription. Currently suggested to use EOB as standard (use PN for GIOTTO like prescription) + 'spin_flx' : "EOB", + 'spin_interp_domain' : 0, + # This can be "QNMs" or "constant". Currently it looks like without QNMs enabled gives basically only prior for precessing events (even for low mass events) + 'ringdown_eulerangles': "QNMs", + 'use_mode_lm' : k_coprecessing_frame, + # 'output_lm' : k, +# 'df' : P.deltaF, 'output_hpc' : "no" } + # Match the backend's own aligned/precessing decision: requesting + # inertial modes for a point DALI evolves as aligned is unsafe. + if teobresums_compat.is_precessing_for_resums(P.s1x, P.s1y, P.s2x, P.s2y): + pars.update({'use_mode_lm_inertial': k}) + if P.a6c < 1000 and P.a6c != 0.0: + pars.update({'a6c' : P.a6c}) else: - print("Using eccentric call") + print("Using TEOBResumSDALI hyperbolic call") + hyp_wav = True # convenient way to know if the waveform is hyperbolic + pars = { 'M' : M1+M2, 'q' : M1/M2, - 'Lambda2' : P.lambda1, - 'Lambda2' : P.lambda2, - 'chi1' : P.s1z, - 'chi2' : P.s2z, - 'domain' : 0, - 'arg_out' : 1, - 'use_mode_lm' : k, - 'output_lm' : k, + 'H_hyp' : P.E0, # energy at initial separation + 'j_hyp' : P.p_phi0, # angular momentum at initial separation + 'r_hyp' : 6000.0, #hardcoded separation: may need to increase with certain systems; may want to make a option + 'nqc' : "no", + 'nqc_coefs_hlm' : "none", + 'nqc_coefs_flx' : "none", + 'ode_tmax' : 3e4, # controls how long end of waveform is; may need to change for some systems + 'LambdaAl2' : P.lambda1, + 'LambdaBl2' : P.lambda2, + 'chi1x' : P.s1x, + 'chi1y' : P.s1y, + 'chi1z' : P.s1z, + 'chi2x' : P.s2x, + 'chi2y' : P.s2y, + 'chi2z' : P.s2z, + 'inclination' : P.incl, + 'coalescence_angle' : np.pi / 2 - P.phiref, 'srate_interp' : 1./P.deltaT, - 'use_geometric_units': 0, 'initial_frequency' : P.fmin, - 'df' : P.deltaF, - 'interp_uniform_grid': 1, 'distance' : P.dist/(lal.PC_SI*1e6), - 'inclination' : P.incl, - "coalescence_angle": np.pi / 2 - P.phiref, - 'output_hpc' : 0, - 'ecc' : P.eccentricity, - 'ecc_freq' : 1 #Use periastron (0), average (1) or apastron (2) frequency for initial condition computation. Default = 1 + 'domain' : 0, + 'arg_out' : "yes", + 'interp_uniform_grid': "yes", + 'use_geometric_units': "no", + # spin_flx can be "EOB" or "PN" prescription. Currently suggested to use EOB as standard (use PN for GIOTTO like prescription) + 'spin_flx' : "EOB", + 'spin_interp_domain' : 0, + # This can be "QNMs" or "constant". Currently it looks like without QNMs enabled gives basically only prior for precessing events (even for low mass events) + 'ringdown_eulerangles': "QNMs", + 'use_mode_lm' : k_coprecessing_frame, + # 'output_lm' : k, +# 'df' : P.deltaF, + 'output_hpc' : "no" } + # Match the backend's own aligned/precessing decision: requesting + # inertial modes for a point DALI evolves as aligned is unsafe. + if teobresums_compat.is_precessing_for_resums(P.s1x, P.s1y, P.s2x, P.s2y): + pars.update({'use_mode_lm_inertial': k}) + print("Starting EOBRun_module") print(pars) - t, hptmp, hctmp, hlmtmp, dyn = EOBRun_module.EOBRunPy(pars) + t, hptmp, hctmp, hlmtmp, dyn = teobresums_compat.run( + EOBRun_module, + pars, + purpose="hyperbolic_waveform" if hyp_wav else "eccentric_waveform", + ) print("EOBRun_module done") - hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + if not(hyp_wav): + ## Set the epoch for non-hyperbolic cases ## + hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + else: + ## custom epoch for the hyperbolic case ## + # wf amplitude + amp = np.sqrt(np.abs(hptmp)**2+np.abs(hctmp)**2) + amp_norm = amp / np.amax(amp) # normalize amplitude for peak finding + # peak finding to determine system type + height_thresh = 0.25*np.abs(amp_norm) + prom_thresh = 0.1*np.abs(amp_norm) + peaks, props = signal.find_peaks(amp_norm, height = height_thresh, prominence = prom_thresh) + peak_heights = props['peak_heights'] + # filtering out peaks so we only keep the local maxima + indices_to_keep = set() + sorted_indices = np.argsort(peak_heights)[::-1] + tol = int(pars['srate_interp'] / 13.65) # 300 samples at srate of 4096 - minimum distance between peaks. + for i in sorted_indices: + peak = peaks[i] + keep = True + for kept_index in indices_to_keep: + if abs(peaks[kept_index] - peak) <= tol: + keep = False + break + if keep: + indices_to_keep.add(i) + # sorted(): indices_to_keep is a set, and downstream epoch choices + # need the retained peaks in chronological order + filtered_peaks = peaks[sorted(indices_to_keep)] + # parsing number of peaks after filtering against distance tolerance + if len(filtered_peaks) == 1: + # scatter case OR plunge case, we can set the epoch normally + hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + elif len(filtered_peaks) == 0: + # meaningless waveform, essentially a non-interacting case. Set the epoch normally + # These points should return very low likelihood and not interfere with the analysis + print('WARNING: no peak detected; non-interacting hyperbolic case') + hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + else: + # capture case, we need to force the epoch to be the last peak + hpepoch = -P.deltaT*filtered_peaks[-1] + hplen = len(hptmp) hp = {} hc = {} @@ -3039,15 +3327,94 @@ def hoft(P, Fp=None, Fc=None,**kwargs): ht = lalsim.SimDetectorStrainREAL8TimeSeries(hp, hc, P.phi, P.theta, P.psi, lalsim.DetectorPrefixToLALDetector(str(P.detector))) - if P.taper != lsu_TAPER_NONE: # Taper if requested + if P.taper != lsu_TAPER_NONE and P.approx != lalsim.TEOBResumS: # Taper if requested lalsim.SimInspiralREAL8WaveTaper(ht.data, P.taper) + if P.approx == lalsim.TEOBResumS: + # Create a taper similar to NRWaveformCatalogManager for TEOBResumS waveforms + if (P.E0 == 0.0): + # Tapering for eccentric/master branch + t_samp_1 = ht.data.length*P.deltaT*0.05 + t_samp_2 = 2/P.fmin + t_samp = np.max([t_samp_1,t_samp_2]) + n_samp = int(t_samp/P.deltaT) + vectaper= 0.5 + 0.5*np.cos(np.pi* (1-np.arange(n_samp)/(1.*n_samp))) + ht.data.data[0:n_samp] *= vectaper + else: + if P.deltaF is not None: + TDlen = int(1./P.deltaF * 1./P.deltaT) + for count,value in enumerate(ht.data.data): + if count == 0: + continue + if np.abs(value-ht.data.data[count-1]) < 0.01 *np.abs(ht.data.data[0]): + ht = lal.ResizeREAL8TimeSeries(ht, count, ht.data.length) + break + else: + continue + for count,value in enumerate(ht.data.data): + if count == 0: + continue + if np.abs(value-ht.data.data[0]) > 0.01 * np.abs(ht.data.data[0]): + n_samp=int(count/2) + break + + # peak finding to determine system type + ht_norm = ht.data.data/np.amax(ht.data.data) + height_thresh = 0.25 + prom_thresh = 0.1 + + peaks, props = signal.find_peaks(ht_norm, height = height_thresh, prominence = prom_thresh) + peak_heights = props['peak_heights'] + indices_to_keep = set() + sorted_indices = np.argsort(peak_heights)[::-1] + tol = int(pars['srate_interp'] / 13.65) # 300 samples at srate of 4096 - minimum distance between peaks. + for i in sorted_indices: + peak = peaks[i] + keep = True + for kept_index in indices_to_keep: + if abs(peaks[kept_index] - peak) <= tol: + keep = False + break + if keep: + indices_to_keep.add(i) + # sorted(): indices_to_keep is a set, and downstream epoch choices + # need the retained peaks in chronological order + filtered_peaks = peaks[sorted(indices_to_keep)] + + vectaper= 0.5 + 0.5*np.cos(np.pi* (1-np.arange(n_samp)/(1.*n_samp))) # this tapers the start + + nmax = np.argmax(ht.data.data) + ht.data.data[0:n_samp] *= vectaper + + if len(filtered_peaks) == 1: + # check if a scatter or a plunge + if np.abs(ht.data.length-nmax) > 3e3: + #scatter + print('Tapering for scatter waveform') + n_samp2=n_samp + vectaper2 = 0.5 + 0.5 * np.cos(np.pi * np.arange(n_samp2 + 1) / (1. * n_samp2)) # end taper + ht.data.data[-(n_samp2+1):] *= vectaper2 + else: + # plunge, only need to taper the start + print('Plunge waveform, only start taper') + + elif len(filtered_peaks) == 0: + print('Non-interactive hyperbolic waveform, tapering both ends') + n_samp2=n_samp + vectaper2 = 0.5 + 0.5 * np.cos(np.pi * np.arange(n_samp2 + 1) / (1. * n_samp2)) + ht.data.data[-(n_samp2+1):] *= vectaper2 + + else: + # zoom-whirl case, taper just the start + print('Zoom-whirl waveform, only start taper') + + if P.deltaF is not None: TDlen = int(1./P.deltaF * 1./P.deltaT) assert TDlen >= ht.data.length ht = lal.ResizeREAL8TimeSeries(ht, 0, TDlen) return ht -def hoff(P, Fp=None, Fc=None, fwdplan=None): +def hoff(P, Fp=None, Fc=None, fwdplan=None, **kwargs): """ Generate a FD waveform from ChooseWaveformParams P. Will return a COMPLEX16FrequencySeries object. @@ -3063,6 +3430,7 @@ def hoff(P, Fp=None, Fc=None, fwdplan=None): If P.deltaF == None, the TD waveform will be zero-padded to the next power of 2. """ + Lmax = kwargs.get('Lmax', 4) # Default to 4 if not specified in kwargs # For FD approximants, use the ChooseFDWaveform path = hoff_FD if lalsim.SimInspiralImplementedFDApproximants(P.approx)==1: # Raise exception if unused arguments were specified @@ -3072,11 +3440,11 @@ def hoff(P, Fp=None, Fc=None, fwdplan=None): # For TD approximants, do ChooseTDWaveform + FFT path = hoff_TD else: - hf = hoff_TD(P, Fp, Fc, fwdplan) + hf = hoff_TD(P, Fp, Fc, fwdplan, Lmax=Lmax) return hf -def hoff_TD(P, Fp=None, Fc=None, fwdplan=None): +def hoff_TD(P, Fp=None, Fc=None, fwdplan=None, **kwargs): """ Generate a FD waveform from ChooseWaveformParams P by creating a TD waveform, zero-padding and @@ -3096,6 +3464,7 @@ def hoff_TD(P, Fp=None, Fc=None, fwdplan=None): Returns a COMPLEX16FrequencySeries object """ + Lmax = kwargs.get('Lmax', 4) # Default to 4 if not specified in kwargs ht = hoft(P, Fp, Fc) if P.deltaF == None: # h(t) was not zero-padded, so do it now @@ -3395,16 +3764,15 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil # see discussion in https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c if is_precessing and fd_L_frame and P.fref: alpha,beta,gamma =extract_JL_angles(P,return_inverse=True) - # alpha0, beta==theta_JN already identified above. Missing polarization rotation factor as well - # zeta_pol will not be used since it is inclination-dependent and therefore depends on extrinsic choices! Also normally calling with P.incl ==0 - _, _, _, theta_JN, alpha0, misc_phi, zeta_pol = lalsim.SimIMRPhenomXPCalculateModelParametersFromSourceFrame(P.m1,P.m2, P.fref, P.phiref, P.incl, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, extra_params) - # Get remaining psiJ quantity (should add to extract_JL_angles) - thetaJN, phiJL, theta1, theta2, phi12, chi1, chi2, psiJ = P.extract_system_frame() - # theta_JN is not useful, for rotations will be close to P.incl in most cases - alpha+= np.pi # empirically validated sign for XPHM, comparing precessing radiation to SEOBv4PHM -# print(alpha, beta, gamma, zeta_pol) -# print(alpha0, thetaJN, np.pi - phiJL,psiJ) - hlmsT_alt = rotate_hlm_static(hlmsT, -gamma - np.pi/2, -beta,-alpha ,extra_polarization=psiJ) + # The modes must remain intrinsic: hoft_from_hlm applies P.psi when + # constructing the strain. A further alpha += pi changes odd-m + # modes relative to even-m modes, rather than supplying a global + # polarization convention. The fixed pi/2 polarization below only + # cancels the historical global minus sign applied above. + hlmsT_alt = rotate_hlm_static( + hlmsT, -gamma - np.pi/2, -beta, -alpha, + extra_polarization=np.pi/2, + ) hlmsT = hlmsT_alt # phase shift ChooseFDModes # for mode in hlmsT: @@ -3500,22 +3868,40 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil P.fmin, P.fref, P.dist, extra_params, \ Lmax, approx_here) elif P.approx ==lalsim.TEOBResumS and has_external_teobresum and not(info_use_resum_polarizations): # don't call external if fallback to polarizations - print("Using TEOBResumS hlms") + print("Using TEOBResumS hlms (only using (2,2); (2,1); (3,3); (4,4) in coprecessing frame)") modes_used = [] + modes_used_check = [] + hypclass = 0.0 + hyp_wav = False distance_s = P.dist/lal.C_SI m_total_s = MsunInSec*(P.m1+P.m2)/lal.MSUN_SI + k_coprecessing_frame_check = [1, 0, 4, 8] for l in np.arange(2,Lmax+1,1): for m in np.arange(0,l+1,1): if m !=0: modes_used.append((l,m)) print(modes_used) k = modes_to_k(modes_used) - print(k) + k_coprecessing_frame = [] + for count,value in enumerate(k): + if value in k_coprecessing_frame_check: + k_coprecessing_frame.append(value) + print(" k inertial modes: ", k, "k coprecessing frames: ", k_coprecessing_frame) + if kwargs.get('force_22_mode', False): + # Restrict EVERY requested list. The precessing branch below asks the + # backend for inertial modes through `k`, so narrowing only the + # coprecessing list still returns all modes through Lmax. + modes_used = [(2,2)] + k = modes_to_k(modes_used) + k_coprecessing_frame = list(k) + print("Forcing ONLY the 22 modes, so k inertial modes: ", k, " k coprecessing frames: ", k_coprecessing_frame) M1=P.m1/lal.MSUN_SI M2=P.m2/lal.MSUN_SI nu=M1*M2/((M1+M2)**2) - if P.eccentricity == 0.0: - print("Using ResumS master; not eccentric") + ecc_ano=eccentric_anomaly_from_mean(P.meanPerAno,P.eccentricity) + true_ano=true_anomaly_from_eccentric(ecc_ano,P.eccentricity) + if P.E0 == 0.0: + print("Using TEOBResumSDALI/GIOTTO standard call") pars = { 'M' : M1+M2, 'q' : M1/M2, @@ -3527,50 +3913,187 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil 'chi2x' : P.s2x, 'chi2y' : P.s2y, 'chi2z' : P.s2z, - 'domain' : 0, - 'arg_out' : "yes", - 'use_mode_lm' : k, -# 'output_lm' : k, + 'ecc' : P.eccentricity, + 'inclination' : P.incl, + 'coalescence_angle' : np.pi / 2 - P.phiref, 'srate_interp' : 1./P.deltaT, -# 'df' : P.deltaF, - 'use_geometric_units': "no", 'initial_frequency' : P.fmin, - 'interp_uniform_grid': "yes", 'distance' : P.dist/(lal.PC_SI*1e6), - 'inclination' : P.incl, + 'anomaly' : true_ano, +# 'a6c' : P.a6c, + 'domain' : 0, + 'arg_out' : "yes", + 'interp_uniform_grid': "yes", + 'use_geometric_units': "no", + # spin_flx can be "EOB" or "PN" prescription. Currently suggested to use EOB as standard (use PN for GIOTTO like prescription) + 'spin_flx' : "EOB", + 'spin_interp_domain' : 0, + # This can be "QNMs" or "constant". Currently it looks like without QNMs enabled gives basically only prior for precessing events (even for low mass events) + 'ringdown_eulerangles': "QNMs", + 'use_mode_lm' : k_coprecessing_frame, + # 'output_lm' : k, +# 'df' : P.deltaF, 'output_hpc' : "no" } + # Match the backend's own aligned/precessing decision: requesting + # inertial modes for a point DALI evolves as aligned is unsafe. + if teobresums_compat.is_precessing_for_resums(P.s1x, P.s1y, P.s2x, P.s2y): + pars.update({'use_mode_lm_inertial': k}) + if P.a6c < 1000 and P.a6c != 0.0: + pars.update({'a6c' : P.a6c}) else: - print("Using eccentric call") + print("Using hyperbolic call") + hyp_wav = True # convenient way to know if the waveform is hyperbolic + pars = { 'M' : M1+M2, 'q' : M1/M2, - 'LambdaAl2' : P.lambda1, - 'LambdaBl2' : P.lambda2, - 'chi1' : P.s1z, - 'chi2' : P.s2z, - 'domain' : 0, - 'arg_out' : 1, - 'use_mode_lm' : k, - 'output_lm' : k, + 'H_hyp' : P.E0, # energy at initial separation + 'j_hyp' : P.p_phi0, # angular momentum at initial separation + 'r_hyp' : 6000.0, #hardcoded separation: may need to increase with certain systems; may want to make a option + 'nqc' : "no", + 'nqc_coefs_hlm' : "none", + 'nqc_coefs_flx' : "none", + 'ode_tmax' : 3e4, # controls how long end of waveform is; may need to change for some systems + 'LambdaAl2' : P.lambda1, + 'LambdaBl2' : P.lambda2, + 'chi1x' : P.s1x, + 'chi1y' : P.s1y, + 'chi1z' : P.s1z, + 'chi2x' : P.s2x, + 'chi2y' : P.s2y, + 'chi2z' : P.s2z, + 'inclination' : P.incl, + 'coalescence_angle' : np.pi / 2 - P.phiref, 'srate_interp' : 1./P.deltaT, - 'use_geometric_units': 0, 'initial_frequency' : P.fmin, - 'df' : P.deltaF, - 'interp_uniform_grid': 1, 'distance' : P.dist/(lal.PC_SI*1e6), - 'inclination' : P.incl, - 'output_hpc' : 0, - 'ecc' : P.eccentricity, - 'ecc_freq' : 1 #Use periastron (0), average (1) or apastron (2) frequency for initial condition computation. Default = 1 + 'domain' : 0, + 'arg_out' : "yes", + 'interp_uniform_grid': "yes", + 'use_geometric_units': "no", + # spin_flx can be "EOB" or "PN" prescription. Currently suggested to use EOB as standard (use PN for GIOTTO like prescription) + 'spin_flx' : "EOB", + 'spin_interp_domain' : 0, + # This can be "QNMs" or "constant". Currently it looks like without QNMs enabled gives basically only prior for precessing events (even for low mass events) + 'ringdown_eulerangles': "QNMs", + 'use_mode_lm' : k_coprecessing_frame, + # 'output_lm' : k, +# 'df' : P.deltaF, + 'output_hpc' : "no" } + # Match the backend's own aligned/precessing decision: requesting + # inertial modes for a point DALI evolves as aligned is unsafe. + if teobresums_compat.is_precessing_for_resums(P.s1x, P.s1y, P.s2x, P.s2y): + pars.update({'use_mode_lm_inertial': k}) + # Run the WF generator print("Starting EOBRun_module") print(pars) - t, hptmp, hctmp, hlmtmp, dym = EOBRun_module.EOBRunPy(pars) - print("EOBRun_module done") + t, hptmp, hctmp, hlmtmp, dym = teobresums_compat.run( + EOBRun_module, + pars, + purpose="hyperbolic_modes" if hyp_wav else "eccentric_modes", + ) + k_list_orig = hlmtmp.keys() - hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + if not(hyp_wav): + rho_net = np.zeros(len(t)) + for mode in hlmtmp: + rho_net += np.abs(hlmtmp[mode][0])**2 + hlmepoch = -P.deltaT*np.argmax(rho_net) +# hpepoch = -P.deltaT*np.argmax(np.abs(hptmp)**2+np.abs(hctmp)**2) + + else: + ## custom epoch for the hyperbolic case ## + + # wf amplitude - 22 mode only + tmp_22 = np.array(hlmtmp['1']) + tmp_22[0] *= (m_total_s/distance_s)*nu + amp = np.abs(tmp_22[0] * np.exp(-1j*(2*(np.pi/2.)+tmp_22[1]))) + amp_norm = amp / np.amax(amp) # normalize amplitude for peak finding + + # Check for cases where the amplitude is max at the start or end + if np.argmax(amp_norm) == 0 or np.argmax(amp_norm) == len(amp_norm) - 1: + reclassify = True + else: + reclassify = False + + # initial peak finding to determine system type + height_thresh = 0.25 + prom_thresh = 0.1 + peaks, props = signal.find_peaks(amp_norm, height = height_thresh, prominence = prom_thresh) + peak_heights = props['peak_heights'] + # filtering out peaks so we only keep the local maxima + indices_to_keep = set() + sorted_indices = np.argsort(peak_heights)[::-1] + tol = int(pars['srate_interp'] / 13.65) # 300 samples at srate of 4096 - minimum distance between peaks. + for i in sorted_indices: + peak = peaks[i] + keep = True + for kept_index in indices_to_keep: + if abs(peaks[kept_index] - peak) <= tol: + keep = False + break + if keep: + indices_to_keep.add(i) + # sorted(): indices_to_keep is a set, and downstream epoch choices + # need the retained peaks in chronological order + filtered_peaks = peaks[sorted(indices_to_keep)] + # parsing number of peaks after filtering against distance tolerance + if len(filtered_peaks) == 1: + # scatter case OR plunge case, we can set the epoch normally + hpepoch = -P.deltaT*np.argmax(amp) + hypclass = _hyperbolic_endpoint_outcome(dym, amp_norm) + + + elif len(filtered_peaks) == 0: + hypclass = 'meaningless' + hpepoch = -P.deltaT*np.argmax(amp) + reclassify = True + + else: + # capture case, we need to force the epoch to be the last peak + hypclass = 'zoomwhirl' + hpepoch = -P.deltaT*filtered_peaks[-1] + + if reclassify: + print('Running re-classification') + # run minimal threshold peak finder + all_peaks, all_props = signal.find_peaks(amp_norm, height=0.000001, prominence=0.000001) + + # checking for minima if no peaks found + + if len(all_peaks) == 0: + print("No peaks found, checking for minima instead...") + all_peaks, all_props = signal.find_peaks((-1*amp_norm + 1.0), height=0.000001, prominence=0.000001) + + if len(all_props['prominences']) > 3: + print('MANY peaks detected on reclassification, evaluating...') + hypclass = _hyperbolic_endpoint_outcome(dym, amp_norm) + if hypclass == 'plunge': + print('Reclassifying to Plunge') + hpepoch = -P.deltaT*np.argmax(amp) + else: + print('Reclassifying to Scatter') + max_peak_index = all_peaks[np.argmax(all_props['peak_heights'])] + print(f"Largest peak is at index {max_peak_index} with height {amp_norm[max_peak_index]}") + # setting the epoch to the largest amplitude peak + hpepoch = -P.deltaT*max_peak_index + elif len(all_props['prominences']) == 3 or len(all_props['prominences']) == 2 or len(all_props['prominences']) == 1: + print('Reclassifying to Scatter') + hypclass = 'scatter' + + max_peak_index = all_peaks[np.argmax(all_props['peak_heights'])] + print(f"Largest peak is at index {max_peak_index} with height {amp_norm[max_peak_index]}") + + # setting the epoch to the largest amplitude peak + hpepoch = -P.deltaT*max_peak_index + else: + print('No peaks detected after reclassifcation') + hypclass='meaningless' + if hyp_wav: + hlmepoch=hpepoch hlmlen = len(hptmp) hlm = {} hlmtmp2 = {} @@ -3641,7 +4164,7 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil modes_used_new.append((4,0)) hlmtmp2[(4,0)]=np.array(hlmtmp[k]) modes_used=modes_used_new - print(modes_used,hlmtmp,hlmtmp2) +# print(modes_used,hlmtmp,hlmtmp2) # for count,mode in enumerate(modes_used): # hlmtmp2[mode]=np.array(hlmtmp[str(count)]) check_if_only_positive_m = False @@ -3650,21 +4173,27 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil check_if_only_positive_m = not( (mode_keys < 0).any()) for mode in modes_used: hlmtmp2[mode][0]*=(m_total_s/distance_s)*nu - hlm[mode] = lal.CreateCOMPLEX16TimeSeries("Complex hlm(t)", hpepoch, 0, + hlm[mode] = lal.CreateCOMPLEX16TimeSeries("Complex hlm(t)", hlmepoch, 0, P.deltaT, lsu_DimensionlessUnit, hlmlen) hlm[mode].data.data = (hlmtmp2[mode][0] * np.exp(-1j*(mode[1]*(np.pi/2.)+hlmtmp2[mode][1]))) - if not (P.deltaF is None): + if not ((P.deltaF is None)): # or (hyp_wav)): TDlen = int(1./P.deltaF * 1./P.deltaT) - print("TDlen: ", TDlen, "data length: ", hlm[mode].data.length) +# print("TDlen: ", TDlen, "data length: ", hlm[mode].data.length) if TDlen < hlm[mode].data.length: -# print("TDlen < hlm[mode].data.length: need to increase segment length; Instead Truncating from left!") -# sys.exit() - hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode],hlm[mode].data.length-TDlen,TDlen) + print("TDlen < hlm[mode].data.length: need to increase segment length; Instead Truncating from left!") + if hypclass == 'scatter': + j_peak_dyn = np.argmin(dym['r']) + j_peak = int((dym['t'][j_peak_dyn] - dym['t'][0])*(M1+M2)*lal.MTSUN_SI/P.deltaT) + hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode], max(int(j_peak - 0.5*TDlen), 0), TDlen) + else: + hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode],hlm[mode].data.length-TDlen,TDlen) elif TDlen >= hlm[mode].data.length: hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode],0,TDlen) - if check_if_only_positive_m or (np.abs(P.s1x) < 1e-4 and P.s2x == 0.0 and P.s1y == 0.0 and P.s2y == 0.0): - print("Conjugating modes") + # Complement of the inertial-mode request above: the aligned path + # returns only positive m, so those modes need conjugates here. + if check_if_only_positive_m or not teobresums_compat.is_precessing_for_resums(P.s1x, P.s1y, P.s2x, P.s2y): mode_conj = (mode[0],-mode[1]) + print("Conjugating mode: ",mode_conj) if not mode_conj in hlm: hC = hlm[mode] hC2 = lal.CreateCOMPLEX16TimeSeries("Complex h(t)", hC.epoch, hC.f0, @@ -3672,16 +4201,95 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil hC2.data.data = (-1.)**mode[0] * np.conj(hC.data.data) # h(l,-m) = (-1)^ell hlm^* for reflection symmetry hlm[mode_conj] = hC2 - # Create a taper, matching exactly what is used in hoft - hp = lal.CreateREAL8TimeSeries('junk', - lal.LIGOTimeGPS(0.), 1., P.deltaT, - lsu_DimensionlessUnit, len(hlm[(2,2)].data.data ) ) - hp.data.data = np.ones(len(hp.data.data)) - lalsim.SimInspiralREAL8WaveTaper(hp.data, P.taper) - # apply taper to all modes - for mode in hlm: - hlm[mode].data.data*= hp.data.data - + if not(hyp_wav): + # Create a taper, matching exactly what is used in hoft + hp = lal.CreateREAL8TimeSeries('junk', + lal.LIGOTimeGPS(0.), 1., P.deltaT, + lsu_DimensionlessUnit, len(hlm[(2,2)].data.data ) ) + hp.data.data = np.ones(len(hp.data.data)) + lalsim.SimInspiralREAL8WaveTaper(hp.data, P.taper) + # apply taper to all modes + for mode in hlm: + hlm[mode].data.data*= hp.data.data + else: + # Taper-boundary detection must be TOTAL. A non-interacting (flat) + # mode never crosses the 1% threshold below, and a fully zeroed mode + # has no nonzero sample at all. Those are exactly the meaningless + # waveforms the hypclass branch further down zeroes out, so they have + # to reach it: start from safe taper bounds and only narrow them if a + # crossing is actually found. + data_22 = hlm[(2,2)].data.data + + # determine location of start taper + if not 'n_samp' in locals(): + n_samp = 1 + for count,value in enumerate(data_22): + if count ==0: + continue + if np.abs(np.real(value)-np.real(data_22[0])) > 0.01 * np.abs(np.real(data_22[0])): + n_samp=int(count/2) + break + + # determine location of end taper + j_nonzero = np.nonzero(data_22 != 0.)[0] + if len(j_nonzero) == 0: + print("Identically zero (2,2) mode; no signal endpoint to locate.") + j_signal_end = hlm[(2,2)].data.length + elif data_22[-1] == 0.0: + print("Signal shorter than seglen; probably can use smaller value.") + j_signal_end = j_nonzero[-1] + 1 + else: + j_signal_end = hlm[(2,2)].data.length + if not 'n_samp2' in locals(): + n_samp2 = 1 + for count, value in enumerate(reversed(data_22[:j_signal_end])): # Scan backwards + if count == 0: + continue + if np.abs(np.real(value) - np.real(data_22[j_signal_end - 1])) > 0.01 * np.abs(np.real(data_22[j_signal_end - 1])): + n_samp2 = int(count / 2) + break + # A zero-length end taper (first preceding sample already crosses the + # threshold) would divide by zero below and write nan into the endpoint + n_samp2 = max(int(n_samp2), 1) + j_taper_end = range(j_signal_end - (n_samp2 + 1), j_signal_end) + + + # Always taper the start + + vectaper= 0.5 + 0.5*np.cos(np.pi* (1-np.arange(n_samp)/(1.*n_samp))) + nmax = np.argmax(hlm[(2,2)].data.data) + for mode in hlm: + #pass + hlm[mode].data.data[0:n_samp] *= vectaper + + if hypclass == 'scatter': + # Taper for scatter + print('Scatter waveform, taper start and end') + vectaper2= 0.5 + 0.5 * np.cos(np.pi * np.arange(n_samp2 + 1) / (1. * n_samp2)) + for mode in hlm: +# hlm[mode].data.data[-(n_samp2+1):] *= vectaper2 + hlm[mode].data.data[j_taper_end] *= vectaper2 + elif hypclass == 'plunge': + # taper for plunge + print('Plunge waveform, only start taper') + elif hypclass == 'zoomwhirl': + # taper for ZW + print('Zoom-whirl waveform, only start taper') + elif hypclass =='meaningless': + # zero out meaningless + _zero_mode_data(hlm) + +# for mode in hlm: +# print(mode) + # For hyp waveforms, resize after tapering. Non-hyp waveforms should have done this earlier (before tapering). +# if not (P.deltaF is None): +# TDlen = int(1./P.deltaF * 1./P.deltaT) +# print("TDlen for your Hyperbolic waveform: ", TDlen," hlm[{}].data.length: ".format(mode), hlm[mode].data.length) +# if TDlen < hlm[mode].data.length: +# hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode],hlm[mode].data.length-TDlen,TDlen) +# print("TDlen < hlm[mode].data.length: need to increase segment length; Instead Truncating from left!") +# elif TDlen >= hlm[mode].data.length: +# hlm[mode] = lal.ResizeCOMPLEX16TimeSeries(hlm[mode],0,TDlen) return hlm else: # (P.approx == lalSEOBv4 or P.approx == lalsim.SEOBNRv2 or P.approx == lalsim.SEOBNRv1 or P.approx == lalsim.EOBNRv2 extra_params = P.to_lal_dict_extended(extra_args_dict=extra_waveform_args) @@ -4094,7 +4702,7 @@ def hoft_from_hlm(hlms,P, return_complex=False, extra_phase_shift=0): # Create complex strain object hT = lal.CreateCOMPLEX16TimeSeries("hT", h22.epoch, h22.f0, h22.deltaT, h22.sampleUnits, h22.data.length) - hT.data.data*=0 # fill with zeros + hT.data.data[:] = 0 # fill with zeros # create for loop over elements of the series to add it for mode in hlms: @@ -5909,13 +6517,13 @@ def rotate_hlm_static(hlm,alphaA,betaA,gammaA,extra_polarization=None): hCatOut[(L,M)] = lal.CreateCOMPLEX16TimeSeries("Template h(t)", h0.epoch, h0.f0, h0.deltaT, lsu_DimensionlessUnit, h0.data.length) - hCatOut[(L,M)].data.data *=0 # initialize + hCatOut[(L,M)].data.data[:] = 0 # initialize lal_type=True elif isinstance(hlm[(L,M)] , lal.Complex16FrequencySeries): hCatOut[(L,M)] = lal.CreateCOMPLEX16FrequencySeries("Template h(t)", h0.epoch, h0.f0, h0.deltaF, lsu_HertzUnit , h0.data.length) - hCatOut[(L,M)].data.data *=0 # initialize + hCatOut[(L,M)].data.data[:] = 0 # initialize lal_type=True elif isinstance(hlm[(L,M)], np.ndarray): hCatOut[(L,M)] = np.zeros(h0.shape) @@ -5954,29 +6562,27 @@ def extract_JL_angles(P,return_inverse=True,phase_factor=1): theta_JL, phi_JL = polar_angles_in_frame(VectorToFrame(Lhat), Jhat) # To make review-stable, be VERY explicit : try to use lalsuite function calls, not above my_cos = np.dot(Lhat, Jhat) - theta_JL = P.incl # this is pretty close. Good enough if we are nearly aligned - if my_cos < 1-1e-5: # but if we are even slightly misaligned, use the acos of above - theta_JL = np.arccos( my_cos ) #P.extract_param('beta') # this is just Jhat.Jhat + theta_JL = np.arccos(np.clip(my_cos, -1., 1.)) # Remaining angle: nhat_obs = np.array([ np.sin(P.incl)*np.cos(np.pi/2*phase_factor-P.phiref), np.sin(P.incl)*np.sin(np.pi/2*phase_factor-P.phiref), np.cos(P.incl)]) # base vector - vecZ = np.array([0,0,1]) - vecY = np.array([0,1,0]) - - # thetaJL == beta. However, we're noit going to use this + # thetaJL == beta. # theta_jn, _, theta_1, theta_2, phi_12, a_1, a_2 =lalsim.SimInspiralTransformPrecessingWvf2PE( # P.incl, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, P.m1, P.m2, P.fref, P.phiref) - # rotation from J to point along L for example - rot = np.matmul(rotation_matrix( vecY, -theta_JL), rotation_matrix( vecZ, -phi_JL)) - - # Rotation around L needed to rotate viewing direction - nhat_rot = np.matmul(rot, nhat_obs) - last_angle = np.angle( nhat_rot[0]+1j*nhat_rot[1]) + # Rotate N by Rz(-phi_JL), then Ry(-theta_JL), following Appendix C + # explicitly. rotation_matrix uses a different active/passive convention, + # so using it here changes the final Euler angle substantially. + cos_phi = np.cos(phi_JL) + sin_phi = np.sin(phi_JL) + nx_rot = nhat_obs[0]*cos_phi + nhat_obs[1]*sin_phi + ny_rot = -nhat_obs[0]*sin_phi + nhat_obs[1]*cos_phi + nx_jframe = nx_rot*np.cos(theta_JL) - nhat_obs[2]*np.sin(theta_JL) + last_angle = np.arctan2(ny_rot, nx_jframe) # Phase conventions as in XPHM paper https://journals.aps.org/prd/abstract/10.1103/PhysRevD.103.104056 if return_inverse: - # the two np.pi factors end up producing a mode sign (-1)^m (-1)^m' to deal with thetaJL < 0 , so we don't have a negative angle there. + # The two pi factors produce (-1)^m (-1)^m' while keeping beta positive. return np.pi - last_angle, theta_JL, np.pi-phi_JL else: # diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md new file mode 100644 index 000000000..b90d8b8f1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md @@ -0,0 +1,759 @@ +# Peak-local marginalization as a framework: the axis contract + +Companion to no single module yet — this is the planning note for consolidating a +method that now exists three times. Read +`DESIGN_time_marginalization_quadrature.md` and then +`DESIGN_time_marginalization_peak_local.md` first: this note assumes the time +instance's vocabulary (enumeration, localisation, merged intervals, the tail bound, +fail-closed delegation) and does not re-derive it. + +Everything measured here was measured on `citlogin6` (CIT), CVMFS IGWN python 3.11, +`OMP_NUM_THREADS=1`, on `rift_O4d` at `6e5bd4b1` (PR #205 merged). Where a claim is +NOT measured it says so. This note is a design record, not a shipped-behaviour +record: nothing in it has been implemented. + +Harnesses (host-local, `ldas-*` NFS home): `~/pl_framework_harness/`. +`psi_warrant.py` the psi extrema count, the `M_k` bound, and the shipped-function +check; `effective_bandwidth.py` the `exp(A cos phi)` coefficient envelope; +`distance_unimodality.py` the distance stationary-point structure and the node-centre +offset; `psi_bracket_adversarial.py` the near-annihilation extremum separation that +falsified this note's first bracketing claim; `psi_backstop.py` the trapezoid aliasing +table and the calibrated-vs-certified sizing cost. Each runs standalone with `PYTHONPATH` set to `MonteCarloMarginalizeCode/Code`, +so a reviewer can re-run rather than take the tables on assertion. + +## The method has been written three times, independently + +| | **time**
`time_marginalization_peak_local` | **angle**
`jax_ile/anglemarg.py` | **distance**
`jax_ile/core.py` | +|---|---|---|---| +| sharpness measure | `peak_width_from_lnL` → σ_t | `estimate_angle_amplitude` → A | `R`, analytic | +| resolution derived from it | `required_upsample_factors`, h ≤ σ/2 | `_dense_grid_sizes`, N = K√A | node scale 1/√R | +| rule selection | `_classify_rows` | `choose_angle_marg_scheme` | — | +| local placement | interval around each enumerated crest | dense (φ,u) grid, or Laplace | nodes centred on x\*=K/R | +| runtime check | dense remeasure + tail-bound certificate | `_runtime_amp_failsafe` | — | +| policy on failure | **fail-closed** → dense rule | **fail-open** → warn + label | — | + +`_distmarg_gh_logL` (`jax_ile/core.py:1405`) states the shared motivation in its own +words: *"the (1/SNR)-narrow high-SNR distance peak that a fixed uniform-in-d grid +under-resolves (biasing the average ~1% low) is resolved at EVERY SNR."* That is the +time module's opening argument with one noun changed. Three authors, three axes, one +skeleton — measure the sharpness, derive the resolution, place work where the mass is, +check afterwards that you were entitled to. + +## The organizing principle: the EXPONENT is band-limited, never the integrand + +The single most costly misreading available here — I made it once already this +session — is to ask whether `exp(lnL)` is band-limited. It never is. On every axis +`exp(lnL)` is a needle, and resolving the needle is precisely the expensive thing +peak-local exists to avoid. + +What the method actually requires is that the **smooth exponent** be band-limited, or +otherwise structurally constrained. For time that object is `kappa(t)`; for +polarization it is `g(ψ)`; for distance it is `A u − B u²` plus the prior term. Every +question about whether the method transfers to a new axis is a question about that +object, and about nothing else. + +Read that way, the band limit plays **four** distinct roles for time, and they +generalize separately: + +1. **Enumeration RESOLUTION — and note what it does and does not buy.** `kappa` is + band-limited below Nyquist, so its narrowest possible LOBE is `deltaT`, and a fixed, + SNR-independent `PEAK_ENUM_FACTOR` places 8 points across it whatever the SNR. That + is what makes "enumerate, do not search" affordable. It is NOT a guarantee that every + extremum is separately bracketed: a band-limited function can carry a max and a saddle + approaching annihilation, arbitrarily close together (measured for the ψ case below). + The shipped module is already careful about exactly this — *"Completeness of the + enumeration buys SPEED; the bound buys CORRECTNESS"* — and a generalized core must + inherit that division rather than the looser reading. +2. **The derivative bound `M_k`.** `sum_j |X_j| |w_j|^k` is a TRUE bound because the + interpolant is a finite trig sum. Three consumers: the crest pre-filter (`M2`), the + certificate remainder (`M4`), and — the round-6 lesson — it is the only place an + inequality rather than a targeting estimate is available at all. +3. **The reconstruction / evaluator.** `bandlimited_spectrum`, `eval_bandlimited_*`, + `enum_grid_derivatives`, and the even-reflection contract: ~400 lines, and the source + of a whole defect family (Nyquist-bin splitting, odd/even `n`, periodic-vs-reflected + drift). +4. **It DEFINES the ground truth.** For time the continuous integrand does not exist + independently of the samples; it *is* the band-limited interpolant. "Exact" is + therefore a claim about reconstruction. + +Roles 1 and 2 are what the method needs. Role 3 is a tax time pays because it only has +samples. Role 4 is the deep one: on every other axis the integrand exists on its own +and every evaluation of it is truth, so role 3 disappears and role 4 inverts — but +roles 1 and 2 must then be re-sourced, because "the samples determine the function" was +doing the certifying. + +## The completeness warrant + +Whether an axis can be certified at all, and therefore what its fail policy may be, is +decided by ONE property: what guarantees that no mass was missed. Three kinds occur in +this codebase, and they are not interchangeable. + +**`exact-band-limit` — time.** `kappa`'s spectrum is identically zero above `fmax`. +Certificate available; fail-closed correct. + +**`exact-trig-degree` — polarization ψ, and the strongest case of the three.** The +exponent at `factored_likelihood.py:943` is + + g(ψ) = term2a + Re(term1 · e^{+2iψ}) + Re(term2b · e^{−4iψ}) + +Harmonics 2 and 4 only, so period π and an exact degree-4 trig polynomial. Measured +over 4000 random coefficient draws spanning six decades of amplitude: + +| | measured | bound | +|---|---|---| +| extrema of `g` on [0,π) | **4**, amplitude-independent | 4, from degree 4 | +| violations of `M_k = 2^k\|term1\| + 4^k\|term2b\|` | **0** | — | + +So ψ gets all of roles 1–3 exactly and in closed form, and more cheaply than time: +`M_k` is two terms rather than a spectral sum, evaluation at arbitrary ψ is O(1), and the +domain is genuinely periodic so none of the reflection machinery applies. ψ is not the +hard axis. It is the easiest one, and it is the right second instance. + +**But an earlier draft of this note claimed a fixed grid "brackets every extremum, +forever", and that is FALSE.** It is recorded here rather than quietly fixed, because it +is the project's characteristic error — a bound on the extremum COUNT silently promoted +to a guarantee about RESOLVING them. A degree-2 trig polynomial can carry a maximum and +a saddle that approach annihilation, so adjacent extrema come arbitrarily close. Measured +on a plain 3000-draw random family: **minimum adjacent-extremum separation 0.068 rad, +against a 0.196 rad cell on a 32-point grid** — two extrema in one cell, with no tuning +required. Driven deliberately toward the bifurcation: + +| `c1/c2` | 3.9 | 3.99 | 3.999 | 3.99999 | +|---|---|---|---|---| +| min separation (rad) | 0.224 | 0.0707 | 0.0224 | 0.0022 | +| value spread of the merging pair (nats) | 1.3e-3 | 1.3e-5 | 1.3e-7 | **1.2e-11** | + +Two things follow, and both matter for the contract. + +*What is actually true* is the weaker pair of statements the method needs: the extremum +COUNT is bounded by 4 and is amplitude-independent (which is what sizes `MAX_INTERVALS`), +and the pair that can hide inside one cell is a max/saddle whose value difference vanishes +as roughly the fourth power of their separation — 1.2e-11 nats at 0.0022 rad — so missing +it cannot lose mass. + +*And the certificate is not optional for ψ.* `segment_sup_bound` bounds `max q` over a +cell from its endpoint values, slopes and `M4` remainder, **regardless of how many extrema +are inside it**. So an unenumerated maximum in a covered cell is bounded anyway, and one +in an uncovered cell is carried by the tail bound. An adapter with an `exact-trig-degree` +warrant might look like it could skip the certificate on the strength of its extremum +count; it cannot, and the contract must not let it. + +**`bounded-stationary-set` — distance.** In `u = Dref/D` the plain-callback exponent is +`A u − B u²`, plus `−4 ln u` for the volumetric prior `p(d) ∝ d²` — which is the exponent +the shipped `_distmarg_gh_logL` sums over its nodes (`jax_ile/core.py:1454`: +`K u − ½R u² − 4 ln u`, so `A = K` and `B = R/2`). It is **not** concave, and — an +earlier draft of this note said otherwise — it is **not** unimodal either, nor does it +always have an interior maximum. What is provable is weaker and support-dependent: + + u·e′(u) = −(2B u² − A u + 4) = −(R u² − K u + 4) + +so the stationary points are the roots of a quadratic, which are **real only when +`A² ≥ 32B`, i.e. `K² ≥ 16R`**; when they exist they are +`u_± = (K ± √(K² − 16R)) / (2R)`, the smaller a minimum and the larger a maximum. Two +regimes follow, and the second is the one an implementation would get wrong: + +* **`K² < 16R`.** The quadratic has no real root and is positive throughout, so + `e′ < 0` on all of (0,∞): the exponent is strictly DECREASING and there is no interior + maximum at all. On the physical support the maximum sits at the lower endpoint + `u_min = Dref/D_max`. +* **`K² ≥ 16R`.** `u_+` is an interior maximum only if it lies inside the support, and + even then it must be compared against the endpoints — on (0,∞) the exponent runs to + `+∞` as `u → 0⁺` (the `−4 ln u` prior term), so there is no global maximum on (0,∞) and + the integral is finite only because `u_min > 0`. + +So the completeness that is actually available is over the **support-aware** candidate set +`{u_min, u_max} ∪ ({u_+} ∩ [u_min, u_max])` — at most three points, at most two of them +local maxima, all in closed form. That is still cheap and still certificate-free, but it +is *not* "nothing to miss": an adapter that enumerated interior stationary points alone +would silently drop the boundary-dominated mass. The shipped code already anticipates +that regime (how often it fires is NOT measured here) — it clips the node centre into +`[x_min, x_max]` exactly so +that bins whose Gaussian peak falls outside the support "still get their +boundary-dominated integral resolved" (`jax_ile/core.py:1437-1442`). A peak-local +distance adapter inherits that obligation. + +Note also the shipped node centre `K/R` is the peak of the Gaussian factor only — it +ignores both the `−4 ln u` prior term and the endpoints — so it is offset from the true +maximum: measured **−0.08σ at K=50,R=1** but **−1.6σ at K=4.1,R=1**, and that second case +sits barely above the `K² ≥ 16R` threshold (16.8 against 16), which is precisely where the +interior maximum is weakest and the endpoints take over. Benign at the SNR this exists +for, but a heuristic, not a bound, and undocumented at the call site. + +**`effective-bandwidth` — the existing angle grid, and the one that CANNOT be +certified.** `_dense_grid_sizes` sizes N = K√A. That constant is not arbitrary: +`exp(A cos φ)` has a Gaussian-in-k coefficient envelope `exp(−k²/2A)`, and the measured +cutoff index over `A = 25 … 10000` is + +| A | 25 | 100 | 450 | 2000 | 10000 | +|---|---|---|---|---|---| +| k(10⁻⁸)/√A | 6.40 | 6.20 | 6.13 | 6.08 | 6.08 | + +Stable to 5% over three decades — the rule is well-founded. But the coefficients never +reach zero. The cutoff is a TOLERANCE, not a band limit, so no true `M_k` exists and no +certificate is possible. Hence the 2× margin and the runtime failsafe, which is the +correct design for that warrant. + +**This is why the fail policies differ, and why unifying them would be a bug.** Time +is fail-closed because it can prove what it dropped. The angle grid is deliberately +fail-OPEN — warn and label, do not poison — because NaN is silently filtered by flowMC, +the SMC path, and `write_samples`, so failing closed there would EXCISE the hot sky +region and publish a clean-looking posterior over what remains. Both are right. A core +that imposes one on the other is a regression in whichever axis it is imposed on. + +## What the existing angle instance is, and is not + +`anglemarg.py` sizes a dense grid for `exp(lnL)` and adds a Laplace branch above a +measured crossover. In the time module's vocabulary that is the **dense rule** — the +analogue of `time_marginalize_bandlimited` — plus a peak-approximation shortcut. It is +*not* a peak-local rule: nothing enumerates, nothing localises, nothing bounds omitted +mass, and there is no fallback ladder. So the first concrete payoff of this framework +is not deduplication. It is giving the angle axis a peak-local branch it does not have, +with a certificate that ψ's structure makes exact. + +## The seam is "supply a spectrum", not "supply an evaluator" + +Both smooth exponents in scope are **finite exponential sums**. Time: +`q(t) = Re sum_j Xw_j exp(w_j t)`, an `npts`-term spectrum. Polarization, in the +`u = 2ψ` chart, is `g(u) = a + Re(c1 e^{iu}) + Re(c2 e^{2iu})` — a **two-term spectrum +with frequencies {1, 2}** (`jax_ile/anglemarg.py:961` states exactly this). + +That is not an analogy. MEASURED: calling the shipped time function +`spectral_derivative_bound(Xw=(c1,c2), fk=(1,2), period=2π, order=k)` — unmodified, +no new code — returns ψ's bound `|c1| + 2^k|c2|` over 300 draws at orders 1, 2, 4, +with **zero violations and tightness exactly 1.000**, i.e. the bound is *attained* +(with two terms the triangle inequality is achievable). + +So the primitives the time rule built are already axis-free; they are merely misfiled +as time code. `spectral_derivative_bound`, the Newton in `localise_peaks`, +`eval_bandlimited_points`, and the whole certified-supremum stack (`parabolic_sup`, +`segment_sup_bound`, the sub-cell cuts) consume only *(coefficients, frequencies, +period)* and cell geometry. What an axis must supply beyond a spectrum is the short +list a spectrum cannot carry: domain topology, the lnL map, the width source, the +backstop, and the fail policy. + +The reflection machinery is not on that list. It exists because time's spectrum is +*implicit in coarse samples of a truncated window* — reflection is how the time adapter +**manufactures** its spectrum. ψ's spectrum is explicit input. Reflection is adapter +internals, invisible to the core. + +## The axis contract + +**The warrant is a closed, typed union owned by the core**, with four kinds — the three +above plus `effective-bandwidth-with-margin`. The fourth is in the type *so the core can +refuse it by name*: + + if axis.warrant.kind is EFFECTIVE_BANDWIDTH_WITH_MARGIN: + raise NotImplementedError( + "peak-local requires a warrant that certifies enumeration completeness and " + "true derivative bounds; effective-bandwidth-with-margin certifies neither " + "(it is the sizing rule for a DENSE grid with a runtime failsafe). Use the " + "axis's dense rule; peak-local cannot be built on it.") + +Same texture as the existing `phase_marginalization` refusal: name the reason, name the +alternative, refuse rather than degrade. + +**A fitted `M_k` is made unrepresentable, not merely discouraged.** `derivative_bound` +is NOT an adapter-overridable method. The adapter supplies `spectrum(rows)` and the +**core** computes `M_k` by the triangle inequality — the one construction that cannot be +a fit. An axis that cannot express its exponent as an exponential sum but holds a +`bounded-stationary-set` warrant gets no `M_k` and no Hermite certificate; its tail control +is the bracket around the closed-form candidate set — interior stationary maximum AND both +support endpoints, per the distance case above. There is no third path. This is the structural answer to the +defect that reopened four times. + +**The core never applies a fail policy.** It returns `(values, ok, decline_ledger)`. +The axis-owning wrapper decides: time hands `~ok` rows to `time_marginalize_bandlimited`; +a jitted consumer labels and continues. That is how "the fail policy differs by design" +survives consolidation — it becomes a property of the call site, expressed on the ledger, +rather than a branch inside the core. + +Protocol members, each named for the site that consumes it: `domain` (span, periodic), +`warrant`, `spectrum`, `enum_grid`, `exponent_on_enum_grid`, `exponent_deriv_on_enum_grid`, +`eval_points`, `eval_uniform`, `lnL` (must be monotone), `peak_scale`, `viable_rows`, +`backstop`, `backstop_cost`, and the per-axis constants **together with the discharge +inequality each must register** for the suite to assert. + +## ψ worked through, and where it falsifies the contract + +| member | ψ supplies | +|---|---| +| domain | span 2π in `u`, **periodic** — no reflection, no pinning, no edge guard | +| warrant | `exact-trig-degree(harmonics=(1,2), period=2π)` | +| spectrum | `(c1, c2)`, `(1, 2)` → core `M_k = \|c1\| + 2^k\|c2\|`, exact and attained | +| enum_grid | 32 points on [0,2π): 8 per period of the top harmonic, the same discipline as `PEAK_ENUM_FACTOR`. Sizes the enumeration, and does NOT by itself guarantee separation — see the bracketing measurement above | +| eval | closed form, O(1)/point; `q'`, `q''` analytic — every evaluation is truth | +| lnL | identity | +| backstop | trapezoid on the full circle at **N derived from amplitude**, `N ≈ 6.2·√A` for a 1e-8 relative tolerance. NOT certified at that N — see below | + +Two places ψ pushes back on the contract, and both should be treated as findings rather +than smoothed over: + +* **`peak_scale` must admit truth-grade curvature.** ψ has exact `σ = 1/√(−g''(u*))`. + If the protocol forces the shared stencil estimator, ψ is pushed through an + approximation it does not need. (`σ` may be fitted in general — it *targets*, it never + *bounds* — but the contract must not forbid exactness.) +* **`max_intervals` means different things on different axes.** For time it is a cost + guard. For ψ the warrant *proves* ≤2 maxima, so exceeding it is a broken-adapter + assertion. The core may only count and decline; the meaning belongs to the adapter. + +**Fail policy for ψ: fail-closed, but to an UNCERTIFIED backstop — and that distinction +has to be carried, not glossed.** Declining is affordable per row in numpy, so +fail-closed is the right mechanism. What an earlier draft got wrong is the *status* of +the thing it declines to: at the calibrated `N ≈ 6.2·√A` the backstop's accuracy rests on +a sizing rule with a margin, exactly like `anglemarg`'s dense scheme, and it is therefore +an `effective-bandwidth-with-margin` object. A certified backstop is available (below) +but costs ~24× the nodes at `A = 10⁴`. So the honest statement is: peak-local-ψ's own +answer is certified by the sub-cell Hermite bound on `g`; the rule it falls back to is +not, unless the certified sizing is paid for. A future jitted port inherits +`anglemarg`'s constraint (static shapes, no per-row backstop) and must go +fail-open-with-label. + +### RETRACTED: "the backstop carries an aliasing bound" + +An earlier draft called the ψ backstop a *fixed*-N trapezoid, "super-exponentially +convergent for a periodic band-limited exponent, so the backstop itself carries an +aliasing bound". **That is false**, and it is the same error this note diagnoses +elsewhere: super-exponential convergence is in `N` at FIXED coefficients, not uniform +over amplitude. `g(u) = A cos u` is degree one, but `exp(g)` has width `∝ A^{-1/2}`, so +any fixed grid eventually under-resolves it. The note already contained the disproof — +the `exp(A cos φ)` coefficient envelope measured two sections above is the same object. + +For the `N`-point trapezoid on a `2π`-periodic `f` the error is *exactly* the aliasing +sum `T_N − I = 2π Σ_{m≠0} ĉ_{mN}`, and for `f = exp(A cos u)` the coefficients are +`I_k(A)`, so the relative error is `2 Σ_{m≥1} I_{mN}(A) / I_0(A)` — computable, not +estimated. Measured: + +| relative error | N=16 | N=32 | N=64 | N=128 | N=256 | +|---|---|---|---|---|---| +| A = 25 | 1.3e-2 | 1.4e-8 | 4.6e-28 | 7.6e-85 | 4.8e-236 | +| A = 450 | 2.3e0 | 6.6e-1 | 2.1e-2 | 2.7e-8 | 2.7e-31 | +| A = 10000 | 1.2e1 | 6.8e0 | 2.9e0 | 9.6e-1 | 7.6e-2 | + +Read the `N=32` column: the error grows without limit in `A`. The required `N` follows +the same `√A` law as everything else here — `N(10⁻⁸)/√A` = 6.80, 6.40, 6.22, 6.22, 6.20, +6.19 over `A = 25 … 5×10⁴`. + +**Can it be certified?** Yes, and the price is measured. For the two-term exponent the +coefficients of `exp(g)` are a Bessel convolution, so +`|ĉ_k| ≤ Σ_j I_j(|c2|) I_{k−2j}(|c1|)` is a true bound. Converting it to a *relative* +bound needs a rigorous lower bound on `ĉ_0`; the only cheap one is Jensen — the harmonics +have zero mean, so `ĉ_0 ≥ e^{⟨g⟩} = 1` — which is valid but conservative by roughly +`e^A`: + +| A | 25 | 450 | 2000 | 10000 | 50000 | +|---|---|---|---|---|---| +| `N` calibrated (relative, 1e-8) | 34 | 132 | 278 | 620 | 1384 | +| `N` certified (Jensen) | 50 | 692 | 3030 | **15100** | **75456** | +| ratio | 1.5× | 5.2× | 10.9× | 24.4× | 54.5× | + +So the certified sizing grows like `1.5·A` against the calibrated `6.2·√A`, and the +penalty grows as roughly `√A/4.4`. **This is an open design choice with a measured +price, not a settled fact**, and the contract must record which one an adapter takes — +because only one of them may be described as certified. + +(Computed with a uniform-asymptotic `log I_k(A)`; a first attempt used `scipy.special.ive` +directly and silently underflowed to zero, which read as "tolerance met" and produced +certified-`N` values too small by up to 9×. Validated against `scipy` wherever `scipy` +is still finite: max disagreement 5e-3 nats.) + +**Reconciliation with the existing `anglemarg`:** its *exact/dense* scheme is the ψ +analogue of `time_marginalize_bandlimited`; peak-local does not compete with it and the +core refuses its warrant. It stays. Its *Laplace branch* is what peak-local-ψ is a +certified replacement for — that branch already enumerates all maxima, then applies an +O(1/A) width model with documented worst-phase error and a blend band that absorbed three +review rounds. But not in one step: the core does host-side ragged bookkeeping while +`anglemarg`'s kernel runs under `lax.scan` with static shapes. The numpy ψ instance +exists to prove the *contract*; porting it into the Laplace slot is a separate PR with its +own gates. + +## Migration, and the gate at each step + +The house rule is `_classify_rows`: no extraction before the second instance exists in +draft. Gate numbers below are measured on `rift_O4d` at `6e5bd4b1`. + +**Step 0 — pin the baseline.** Hash `time_marginalize_peak_local` outputs over a mixed +block. The bar for every later step is **exact identity, not ULP-close**: pure code +motion cannot reassociate floats. Plus, with the test files UNTOUCHED: + +| gate | files | count | +|---|---|---| +| peak-local | `test_time_marginalization_peak_local.py` | **121** collected | +| band-limited | `..._quadrature.py` (81) + `..._quadrature_pipeline.py` (57) + `test_continuous_time_posterior_export.py` (23) | **161** collected, 160 passed, 1 skipped (cupy absent) | + +**Step 1 — write ψ as a draft that imports from the time module and duplicates +deliberately.** Importing across modules is the safe direction; it is *copying policy* +that rotted. Gate: agreement with a converged `quad` reference over an amplitude ladder +and adversarial phases; enumeration-completeness and `M_k` property tests; time hash +unchanged (trivially — nothing has moved). + +**Step 2 — diff the two spines, and extract only what the diff proves shared.** Move the +axis-free helpers to the core verbatim, leaving re-export shims so `__all__` keeps +resolving. Gate: time hash identical, both suites green *unmodified*, plus a new test +asserting a helper is the **same object** in both consumers — the `is`-identity anti-drift +analogue of the existing classification-itself test, since a parity test cannot see a +change both sides read. + +**Step 3 — formalize the contract types**, deliberately *after* step 2: the protocol is +transcribed from the seam the diff exposed, not designed ahead of it. Gate: refusal tests +executable, time hash still identical, public signature unchanged. + +**Falsifiable prediction, recorded now so step 2 can settle it:** the plan/bucket skeleton +of `_peak_local_chunk` will **not** extract. It exists because time's interval and point +counts are ragged; ψ has ≤2 maxima and a fixed per-interval point count, so its natural +shape is fixed-slot arrays with no plan, no buckets, no host round trip. What should +extract from that spine is the two-stage keep discipline, the accounting-reconciliation +ledger, and the accept predicate. + +## The primitive is PER-AXIS {localize, dense}, not "1-D vs 2-D marginalizer" + +The natural-looking decomposition — a 1-D marginalizer primitive and a 2-D one — is the +wrong cut. We do not always dual-localize, and the reason is physical, not a limitation +to be engineered away. The right primitive is a per-axis CHOICE, and the existing schemes +are already members of one family rather than a ladder: + +| | φ dense | φ localized | +|---|---|---| +| **ψ dense** | `..._exact`, cost ~A | (no use case found) | +| **ψ localized** | `..._laplace`, cost ~√A — SHIPPED | the joint high-SNR target | + +Today's `laplace` is not a stepping stone to be replaced. It is the correct member +whenever ψ is localizable and φ is not, and a core that only offers "all axes localized" +would delete a working configuration. + +**The joint exponent is a 2-D trig polynomial, so the warrant extends verbatim.** This is +already in the shipped code, not an assumption: `angle_coefficient_tables` builds *"Exact +2-D Fourier coefficient tables"*, `_reconstruct_field` evaluates *"the real trig polynomial +... at (phi, u)"*, and `angle_sample_grid_sizes` states the bidegree as derived — +φ-harmonics up to `2*m_max`, and u-harmonics **at most 2 for ANY mode set** (spin-2), so +the ψ degree never grows. Measured on random tables of bidegree (4,2), 150 draws × 7 +derivative orders: `M_(a,b) = sum_kq |C_kq| k^a |q|^b` had **0 violations**, tightness 0.89. +`spectral_derivative_bound` generalizes to a multi-index unchanged in construction. + +## "Localized" ALWAYS means multi-mode. Single-centre is not in the family + +Stated first because it is the easy thing to get wrong when extending to a second axis. +The shipped ψ branch is already multi-mode — `_psi_lnI_lap_branch` is *"the +enumerated-maxima Laplace branch"*, resolving up to four roots — so "Laplace" in this +module has never meant "expand about one point". The joint (φ,ψ) kernel must inherit that +structure on BOTH axes. The measurement below is the quantitative reason, not a criticism +of anything shipped: it sizes what a naive single-centre joint extension would cost. + +Measured on the shipped coefficient tables (`make_synth` fixture, real (φ,ψ) structure, +not random coefficients), at fixed `x`: + +| kappa boost | 1 | 10 | 100 | 1000 | +|---|---|---|---|---| +| exponent amplitude | 3.4 | 32.5 | 325 | 3250 | +| co-dominant maxima | 8 | 12 | 8 | 8 | +| **single-centre Laplace error (nats)** | **−1.66** | **−1.48** | **−1.40** | **−1.40** | + +Two readings, and the second is the load-bearing one. + +*The obstruction is MULTIPLICITY, not conditioning.* `cond(H)` at the dominant mode +measured 11–23 — unremarkable. What breaks a single-centre Laplace is that with dominant +(2,±2) content the (φ,ψ) surface carries ~8–12 maxima whose values are equal to machine +precision (measured second-best gap 0 to 4.6e-13 nats). Adding (3,±3) leaves 11 maxima +within 0.031 nats. These are exact structural degeneracies, and the count is +amplitude-independent — consistent with the fixed bidegree. + +*And the single-centre error DOES NOT DECAY WITH AMPLITUDE.* It sits at −1.4 nats from +amplitude 3 to 3250. That is the opposite of the usual "Laplace improves at high SNR" +intuition, and the reason is that the deficit is combinatorial rather than curvature: one +centre represents one of k equal modes however sharp each becomes. High SNR does not +rescue a single centre — it is precisely where enumerating the modes matters. Which is +why the joint kernel is peak-local (enumerate, then integrate near each) rather than a +Laplace refinement, and why the enumeration budget, not the curvature model, is the thing +the φ-axis selector has to size. + +HONEST LIMIT OF THIS HARNESS: the multi-mode sum measured +0.42 to +1.01 nats and did not +show clean `O(1/A)` convergence. The per-mode Hessians here are grid finite differences on +a 512² torus, so that residual is plausibly the harness rather than the method — it is NOT +evidence that a multi-mode estimator converges, and a real localizer plus certificate is +needed before any such claim. Recorded so the number is not quoted as a validation. + +## Selection must not become an option matrix + +The tree already shows the failure mode: `..._laplace` REFUSES `JAX_ILE_DISTMARG_GH` +(`jax_ile/anglemarg.py:1119`) because its node placement is defined per fixed-ψ exponent. +That is a pairwise option corner case, and a family of per-axis choices multiplies them if +each pair is hand-checked. Two rules keep the surface from growing: + +* **Selection stays keyed on measured quantities, per axis.** `choose_angle_marg_scheme` + already picks from the data-derived amplitude rather than a user flag; a new member of + the family is a new regime on that same measured axis, not a new option. The degeneracy + measurement above is what the φ-axis selector must key on — multiplicity and mode gap, + not amplitude alone. +* **An incompatibility is a property of a scheme's warrant, declared once**, never an `if` + per pair. The GH refusal is really "this scheme's node-placement warrant is defined per + fixed-ψ exponent"; stated that way it is checked generically. NOTE: distance + marginalization is expected to arrive on the other path, so this particular refusal must + be expressed as a warrant property and not frozen as a permanent pairwise rule. + +## ENUMERATION AS A PRIMITIVE + +The reusable building block, worked out abstractly because it is where the long-term +compute is saved: every axis needs it, and re-deriving it per axis is how the certificate +gets weakened by accident. + +### The output is a certified COVER, not a list of peaks + +Three tempting contracts all fail at the operating point: + +* *"all critical points"* — the count is DISCONTINUOUS in the coefficients at a max/saddle + annihilation, so it cannot be honoured continuously; and annihilation is normal here, not + a corner (measured, above). +* *"all local maxima"* — with 8–12 maxima equal to machine precision, which of them are + "the" maxima is a rounding accident. +* *"all modes above a mass threshold"* — the right question, wrong contract: mode mass is + not knowable before integrating, so a threshold on estimated mass is exactly an estimate + promoted to a bound. + +What survives — and what the time module already implements, distributed across +`enumerate_peak_indices`, `merge_intervals_by_row` and `segment_sup_bound` — is: + +> **ModeCover**: a finite set of disjoint regions `{R_i}`, each with a representative point +> and targeting data, plus a CERTIFIED bound `B_out >= sup{ g : outside the union }`, plus a +> ledger. + +Three promises, in decreasing strength. **(1) The bound** is the only correctness-bearing +one: it converts "did I miss a mode?" into an inequality the caller discharges, omitted mass +`<= |domain \ union| * exp(B_out)`. **(2) Targeting** — each representative is a stationary +point to a stated residual — buys speed only, and is allowed to fail; failure surfaces as +`B_out` too large and the row declines. **(3) There is NO count promise.** The number of +regions is an output. A merged max/saddle pair is one region; twelve machine-equal maxima +may be twelve regions or fewer if their covers overlap. + +Mass therefore SELECTS (drop a region whose certified sup is below tolerance — a bound-based +rejection, safe) but never CERTIFIES PRESENCE. + +### Warrants, by what finite object exhausts the stationary set + +| class | axis | certificate | cost | amplitude? | +|---|---|---|---|---| +| **exact trig poly, 1-D** | ψ | fundamental theorem of algebra: `z = e^{iu}` gives degree `2n`; those `2n` roots are all of them | one `2n × 2n` companion eigenproblem | **independent** | +| **exact trig poly, multi-D** | joint (φ,ψ) | BKK / mixed volume | one algebraic solve of that size | **independent** | +| **band-limited, large spectrum** | time | algebraically possible (`2·npts` companion) but `O(npts³)` unaffordable | grid seeds + `segment_sup_bound` | independent | +| **closed-form stationary set** | distance | the algebra of that functional form; ≤3 candidates incl. support endpoints | `O(1)` | independent | +| **effective bandwidth** | the dense angle grid | NONE | — | refused by name | + +The economic argument, stated plainly: under the first two, everything that sizes the +enumeration — degree, mixed volume, `M_k` — is invariant under `C -> λC`. Amplitude enters +only downstream, in how narrow each mode's LOCAL grid must be. Dense pays `√A` per axis +(so `~A` for the 2-D product); enumeration converts an amplitude-scaling cost into a +physics-scaling one. That conversion is the whole justification. + +**Time is the instructive case**: it is class-1 in principle but not in practice, so it +enumerates on a grid — which is NOT a certificate — and restores correctness at +CERTIFICATION time via the cover bound. Classes 1 and 2 certify at ENUMERATION time. Both +are certified; only where the cost is paid differs. + +### The algebraic core, and the tolerance that must not exist + +1-D: `P(z) = c2 z⁴ + (c1/2) z³ − (c̄1/2) z − c̄2` for ψ. Roots on `|z|=1` are the critical +points. Measured: residual `|g'|` at the roots is **3.4e-15** relative to `M_1` over 3000 +draws spanning six decades; never fewer than a 2e5-point grid; **~20 µs/call, flat from +amplitude 1 to 1e6** — the amplitude-independence, demonstrated. + +Grid seeding fails where it matters. Driving `c1/c2 -> 4`: + +| `c1/c2` | separation | grid-32 | grid-256 | grid-4096 | algebraic | +|---|---|---|---|---|---| +| 3.99 | 0.0707 | 1 | 3 | 3 | **4** | +| 3.99999 | 0.0022 | 1 | 1 | 3 | **4** | + +(Verified this is genuine resolution loss, not a seam-wrap artefact: a circular counter gives +identical numbers.) + +**AND A TRAP THAT WAS WALKED PAST ONCE, recorded so it is not walked past twice.** The +obvious implementation filters roots by `||z| − 1| < tol`. That tolerance is itself an +estimate promoted to a bound. At EXACT multiplicity `m` the computed roots smear off the +circle by `ε_machine^(1/m)` — measured **4.6e-6 for a triple root** — so: + +| on-circle tol | 1e-9 | 1e-7 | 1e-6 | 1e-3 | +|---|---|---|---|---| +| roots found at exact degeneracy (true: 4) | 1 | 1 | **1** | 4 | + +A `1e-6` filter — a perfectly reasonable-looking choice, and the one first written here — +silently returns ONE mode where there are four, in precisely the machine-degenerate +configuration that is the production regime. A conjugate-reciprocal pairing test does not +escape it either: measured, it returns counts identical to the naive filter at every +tolerance, because the partner is off-circle by the same amount. + +**The fix is to have no tolerance at all**, and it falls out of promise 3. Seed regions +from `arg(z)` of ALL `2n` roots and filter nothing. Over-covering is free — redundant +regions merge — while under-covering is the only real danger. Measured over 4000 draws +spanning six decades, every true extremum lies within **3.1e-4 rad** of a seed, which is the +reference grid's own resolution, i.e. exact. And at exact degeneracy, where sign-change +enumeration finds **zero** extrema (`g'` touches zero without crossing), the algebraic seeds +still cover the point exactly. + +### Degeneracy + +Cluster, and treat a cluster as one region. The codebase already contains this as +`merge_intervals_by_row`, whose docstring makes the deeper point: merging is not an +optimization, it is what prevents double-counting and what makes the method degrade +CONTINUOUSLY into the dense grid with no threshold anywhere. Under ModeCover, clustering is +not even a special case — regions are built around every seed, and overlapping ones merge. + +The invariant that expresses "do not split an annihilating pair" is upper-semicontinuity IN +THE MASS SENSE: as coefficients vary, the integral over the union plus the certified outside +bound must vary continuously, while the region COUNT is free to jump. A primitive promising +only the union and `B_out` cannot be broken by the bifurcation, because nothing it promises +changes discontinuously there. That is a testable property: perturb through a measured +bifurcation and assert the accepted mass moves continuously. + +SPECULATION, not verified: at exact degeneracy the co-dominant modes look like orbit-mates +of a discrete symmetry, and quotienting would divide the algebraic degree by `|G|` and +replace `k` numerically-rediscovered copies by an exact `+ln k`. Attractive, but the +measured gaps are 0 *to 4.6e-13* nats and the tables are built by sampled accumulation, so +the symmetry can be broken at roundoff. A symmetry assumed exact when it is 1e-13-broken is +the same defect in a new costume. Correct layering: numerical clustering stays load-bearing; +a declared symmetry may SEED clustering and tighten the budget, and the certificate verifies. + +### The 2-D enumerator COMPOSES the 1-D one — the pencil may not be needed at all + +The obvious route to joint (φ,ψ) is a full 2-D algebraic solve: two Laurent equations, BKK +mixed volume `8mn = 64` as the certificate, hidden-variable pencil to solve it. The flagged +blocker was that pencil's conditioning on the machine-degenerate production tables — the 2-D +analogue of the on-circle-tolerance trap. + +**That blocker is dissolved rather than solved, by composition.** The u-degree is pinned at +2 for ANY mode set, so at every fixed φ the u-critical points are the unit-circle roots of +the SAME degree-4 polynomial the ψ primitive already solves. The variety `{∂_u g = 0}` is +therefore obtained EXACTLY, with no grid in u and no tolerance. The 2-D critical points lie +on that curve, so the remaining search is **one-dimensional in φ along a curve known +exactly** — no resultant, no pencil, no BKK machinery. + +Measured on the shipped tables (`make_synth`, bidegree (4,2) — note `A` and `B` have +DIFFERENT bidegrees, `A` linear in the waveform (φ≤m_max, u≤1) and `B` quadratic +(φ≤2m_max, u≤2), which is why `c2` carries no `A` contribution exactly as +`_laplace_psi_lnI` states): + +| κ boost | 1 | 10 | 100 | 1000 | +|---|---|---|---|---| +| mass-carrying maxima (brute force) | 16 | 12 | 12 | 12 | +| **recovered, at 64 φ-seeds** | **16** | **12** | **12** | **12** | +| worst candidate-to-maximum gap (rad) | 0.067 | 0.026 | 0.070 | 0.069 | + +Every mass-carrying maximum is recovered at every amplitude, and the gap shrinks as φ is +refined (0.070 → 0.039 at 128 seeds). Candidate count is `4 × N_φ` — **amplitude-independent**. + +Against the SHIPPED `_dense_grid_sizes` product grid: + +| amplitude | 325 | 3 250 | 3.25e4 | 3.25e5 | +|---|---|---|---|---| +| dense (φ,u) points | 48 640 | 430 592 | 4 216 576 | 41 806 336 | +| composed (4 × 64) | 256 | 256 | 256 | 256 | +| **ratio** | 190× | 1 682× | 16 471× | 163 306× | + +The ratio grows linearly in `A`, which is the amplitude-independence argument made concrete. + +**Be precise about what is and is not certified here.** This is a HYBRID: the u axis is +certified at enumeration time (exact quartic, all roots, no filtering), while the φ axis is +GRID-SEEDED and therefore is not — it carries exactly the same "a grid is a resolution, not +a certificate" caveat as the time axis. Correctness on φ must come from the cover bound, as +it does for time. What composition buys is not a φ certificate; it is the removal of the +entire 2-D algebraic apparatus and its conditioning risk, at a cost that does not grow with +amplitude. A full 2-D solve remains the route to an enumeration-time certificate on BOTH +axes if one is ever needed; this measurement says it is not needed to get the cost win. + +### Where enumeration loses — the exclusion region is part of the design + +* **Low amplitude**: modes are wide, regions merge toward the whole domain, and the method + IS the dense grid plus overhead. Selection should not enter peak-local below crossover. +* **High degree**: cost is the algebraic solve regardless of how much mass matters. +* **Dimension ≥ 3**: mixed volume and solve cost grow multiplicatively; do not extend + without a new measurement campaign. Composing over peak SETS is the designed alternative. +* **No warrant**: refuse by name. + +## Anti-goals + +* **Do not unify the fail policy.** It follows from the warrant. See above. +* **Do not let the core own per-axis constants.** `W_SIGMA`, `PEAK_KEEP_NATS`, + `TAIL_LOG_TOL`, `MAX_INTERVALS` are inequalities over *time's* dynamic range. Each + axis re-derives its own and asserts its own; the core demands that the assertions + exist, and supplies no values. +* **Do not accept a fitted `M_k`.** The recurring defect in this module's history is a + targeting model promoted to a bound (Door 4: "off by 122 nats"; the crest estimate was + one octave too optimistic four times running). An adapter interface that accepts an + estimated bound institutionalizes that defect. Either the adapter supplies a proof- + carrying bound or a completeness certificate, or the axis is REFUSED — in the style of + the existing `phase_marginalization` refusal, not approximated. +* **Do not let sub-cell certificate geometry leak into adapters.** Whole-cell bounding + does not work: on sharp rows the merged interval is narrower than one enumeration cell + (half-width 0.05 of a cell at derived factor 4096), a cell-granular covered mask marks + nothing, the crest's own cell counts as outside, the bound then bounds the CREST, and + every row is rejected — the option goes inert. Measured. If this geometry lives in + the adapter, every new axis rediscovers it. +* **Do not own row classification or viability.** `_classify_rows` stays where it is + for time; another axis brings its own dispatcher. The core never sees `factors`, + ceilings, or edge guards. +* **Do not call the likelihood on a full axis grid, and do not call the backstop.** All + exponent evaluation routes through the adapter, and the backstop is invoked only by the + policy-owning wrapper — otherwise the core smuggles a fail policy in through the back + door. +* **Do not carry cross-call state.** Batch-local only; any persistent scale makes results + batch-order-dependent. +* **Do not silently widen.** Every decline goes on the ledger under a named reason, with + the reconcile invariant that the sub-counts sum to the declined rows. A change that adds + an unledgered decline path must fail a reconcile test. +* **Do not extract before the second instance exists in draft.** `_classify_rows` is + the house precedent and it was extracted only AFTER the duplicated policy had drifted + three times. Of ~1600 lines here, perhaps 300 are the generic skeleton; the rest is + the time evaluator plus eight rounds of measured adversarial fixes that a speculative + core would inherit without having earned. + +## Evidence that the certificate arithmetic belongs in ONE place + +`parabolic_sup` shipped, then needed a second numerical-robustness pass after merge +(`87a7a98b`): normalized discriminant to stop `inf − inf = NaN` erasing genuine +stationary points from a purported upper bound, cancellation-safe quadratic roots +(`q = −½(B + copysign(√disc, B))`, other root from `C/q`) because the direct form loses +the in-range root of a nearly-quadratic Hermite cell — which happens naturally for a +symmetric band-limited crest, where `a` should vanish and the endpoint arithmetic leaves +a few ulps. Two rounds of hardening on 30 lines of pure arithmetic. Copied per axis, +that is a defect per axis. + +## Open, and not established here + +* **The composition/nesting contract**, and it is the sharpest open question. Nesting + means the outer axis's exponent is the inner MARGINAL, `G(t) = log int exp(g(t,psi)) dpsi`, + and `G` inherits neither the warrant nor the derivative bounds of `g`. Concretely, under + the tilted inner measure `mu`, + + d2G = E_mu[d2 g] + Var_mu(d_t g) + + and the variance term scales with inner amplitude — so **the enumeration factor for `G` + is not SNR-independent**, which is the pillar the whole time construction stands on. + The existing refusal of phase marginalization is the special case: the outer exponent + becomes `log I0(|kappa|/D)`-shaped, its width carries `(I1/I0)(|kappa|/D)`, and `|kappa|` + peaks elsewhere than `Re kappa`, so both the width derivation and the monotone-argmax + separation die at once. + + Two classes, and only one is open: + + - **Monotone-reduction inner — shipped, works, and should be NAMED.** If the inner + marginalization is presented as a callback monotone in the outer exponent at fixed row, + argmaxes are preserved, the outer warrant is untouched, and the width is measured + through the callback. Distance-inside-time is this class and is the production + configuration today. It is already the `lnL` + monotone member of the contract. + - **Genuine inner peak-local — DEFERRED, with named preconditions.** The core composes + nothing. Before the contract can be written, two things must exist: (i) a cumulant + bound giving `sup|d_t^k G|` from `M_k(g)` and inner-measure moment bounds, with the + `I1/I0` case as its acceptance test — SPECULATION: such bounds plausibly exist but grow + with inner amplitude, and whether they stay tight enough to beat the dense alternative + is unmeasured; and (ii) a measured completeness study for enumerating extrema of `G`, + since a fixed factor is provably insufficient. Until both exist, every nested pairing + refuses per pairing, naming the missing lemma. + - **The practical bridge that needs neither:** `return_peaks=True` already exports + `(t_star, sigma)` per row, callback-independent. A time-first reordering composes over + PEAK SETS rather than over nested exponents — which is why the peaks are an output and + not a temporary. +* Whether ψ-marginalization is actually moving into the vectorized likelihood. + `NetworkLogLikelihoodPolarizationMarginalized` is on the old non-vectorized API and + production samples ψ by Monte Carlo. NOT verified. +* The harmonic degree of the φ_orb exponent (asserted ≤ 2·Lmax via the Ylm crossterms). + NOT verified — check the crossterm conventions before relying on it. +* Whether cosmological distance priors preserve the stationary-point structure derived + above for the volumetric prior — the quadratic, its discriminant, and hence the size of + the candidate set all come from the `−4 ln u` term specifically. NOT checked. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md index 2c7587039..5e8654894 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md @@ -65,7 +65,8 @@ explicit stencil name and the retired "choose for me" spelling raises. ## 3. Mass ladder (fmin 30) -SEOBNRv4, an IMR model. Against an exact FFT-zero-padded reference; paired, K=2000, 3 seeds; each +SEOBNRv4, an IMR model. Against an FFT-zero-padded reference (**not** exact -- see the note in +§9.6.1); paired, K=2000, 3 seeds; each mass normalised to SNR_lik = 100. srate 4096, fmax 1700, fmin 30, Lmax 2. max|ΔlnL| in nats: | M/M☉ | nearest | cubic | sinc | winner | @@ -178,7 +179,562 @@ winner. The entire fmin sweep is at srate 4096. --- -## 9. Provenance +## 9. Backend coverage, and the one place the stencil is deliberately absent + +**Added 2026-08-26.** The stencil now has **four** implementations across **three** likelihood +variants, but that is fewer moving parts than it sounds, because the variants all funnel through +two primitives. + +| backend | entry point | weights from | +|---|---|---| +| numpy | `_sinc_Q_window_numpy` | `_sinc_lanczos_weight_matrix` | +| cupy | `Q_inner_product_sinc_cupy` | `_sinc_lanczos_weight_matrix` (built on device) | +| CUDA | `Q_inner_sinc` kernel | passed in from the above -- **not** re-derived in C | +| JAX | `jax_ile.core._make_gather_sinc` | `_sinc_lanczos_weights_jax` (see below) | + +`factored_likelihood_freqresponse` (finite-size) and `factored_likelihood_with_rotation` +(slow-rotation) do **not** carry their own stencils: both call +`FL._q_window_numpy_interp` / `FL._q_inner_product_gpu`, so they inherit whatever the two +primitives do. Three variants, three implementations of the weights, not nine. + +**The fused calibration-marginalization kernels (`cuda_Q_fused_calmarg.cu`, +`cuda_Q_fused_calmarg_distmarg.cu`) implement `nearest` only, and that is deliberate.** They do +an integer gather at `ifirst + c*N_window + i_time` inside an already register-heavy fused +reduction. `time_interp != 'nearest'` therefore raises `NotImplementedError` at the library level +and falls back to the `loop` path at the driver level; `test_calmarg_stencil_gating.py` pins both, +and pins them against `== 'nearest'` rather than a hard-coded list of the stencils that existed +when the guard was written. + +### 9.1 JAX is the one backend that cannot share the weight array + +The other three consume one `(n_extrinsic, 2a)` array from `_sinc_lanczos_weight_matrix`, so they +cannot drift. JAX cannot: the weights depend on the sub-sample offset, which is a *traced* +function of sky location, and `jax.grad` has to see through them. `_sinc_lanczos_weights_jax` is +therefore a second, independent expression of the same formula -- the one real divergence risk in +this feature. + +What keeps it honest is `test/jax/test_jax_stencil_parity.py`, which compares the two generators +directly rather than trusting review. Measured agreement, ldas-grid + ldas-pcdev13, igwn python +3.11 / jax 0.7.1 / cupy 12.0.0, 2026-08-26: + +| comparison | measured | gate | +|---|---|---| +| JAX vs numpy weights | 4.4e-16 | `1e-14` | +| JAX vs numpy assembled window (interior / left edge / right edge) | 1.1e-13 / 7.1e-15 / 1.8e-13 | `1e-12` | +| JAX vs CUDA `Q_inner_sinc` (RTX 2080 Ti, sm_75) | passes | `1e-12` | + +Two conventions are easy to get right in one backend and wrong in another, so both are pinned +explicitly: the `|x| >= a` hard zero, and the fact that the unit-sum renormalisation is applied +over the **full** stencil and is **not** redone after out-of-buffer taps are dropped. A backend +that renormalised after masking agrees on weights and disagrees at the buffer edge. + +**The parity suite was mutation-tested, not merely written.** Four source-level mutants of the +JAX gatherer: pointing `_GATHERERS['sinc']` at the cubic stencil (kills 6 of 15 tests), dropping +the normalisation (5), shifting the tap offsets by one (7), and replacing zero-extension with an +edge clamp (5). A fifth -- deleting the `|x| >= a` clause -- initially **survived**, and is an +*equivalent* mutation on the wired path: for `u` in `[0,1)` the only guarded tap sits at `x = -a` +where `sinc(-1) = 3.9e-17`, so the clause is worth `1.5e-33`. Off the wired path it is worth +`2.2e-3`, so it is pinned by `test_weight_parity_outside_the_unit_interval` rather than deleted. + +### 9.2 Half-width is a library parameter, NOT reachable from any driver + +`SINC_HALFWIDTH_DEFAULT = 8` (16 taps). `_sinc_Q_window_numpy(..., a=)` and +`Q_inner_product_sinc_cupy(..., halfwidth=)` accept a value, but **no CLI, pipeline flag or +`time_interp` spelling exposes it**, so every production path is locked to `a = 8`. That is worth +stating because the sky-offset diagnosis measured `a = 16` as visibly better on the 3G +finite-size demo (§9.3), and reaching it requires a code change, not a flag. Adding a name for it +would have to flow through `TIME_INTERP_CHOICES` in two modules, the CLI help pinning in +`test_interpolate_time_cli.py`, and the calmarg gating truth table -- and §2 is explicit that +stencil naming is deliberately locked down -- so it is left as a decision to be taken, not a +silent addition. + +### 9.3 The 3G finite-size demo, where the JAX gap was found + +Archived configuration of the 3-site 3G demo (CE+ET+K, m1 1.6 / m2 1.4, fmin 50, fmax 1024, +srate 2048, self-consistent Qmax 4, SNR 600, zero noise). Deterministic hierarchical-zoom peak +scan -- no sampler, no seed -- so the stencil is the only variable. Offset of the lnL peak from +the injected sky position, and `dlnL` above lnL at the truth (which for a zero-noise +self-consistent injection should be ~0): + +| stencil | peak offset | `dlnL` above truth | lnL at truth | shortfall vs rho^2/2 = 180000 | +|---|---|---|---|---| +| `cubic` | 9.22' | 64.53 | 178214.14 | 1785.9 | +| **`sinc` (shipped, a = 8)** | **0.74'** | **0.515** | **179824.16** | **175.8** | +| `lanczos8` (prototype, a = 8, unnormalised) | 0.33' | 0.117 | 179915.91 | 84.1 | +| `lanczos16` (prototype, a = 16, unnormalised) | 0.13' | 0.018 | 179944.54 | 55.5 | + +Against the marginalisation overhead of 14-23 nats (measured separately; roughly constant in +nats, NOT a constant fraction, so an absolute threshold is right at every SNR), shipping `sinc` +takes the shortfall from 1786 nats to 176 -- a 10x reduction, and the peak from 9.2 arcmin to +0.74. That is the improvement this stencil buys on this configuration. + +**But note rows 2 and 3.** The shipped `sinc` is 2.3x worse in offset, and 92 nats worse in lnL +at the injection, than the `lanczos8` prototype **at the same half-width**. The cause is +identified exactly, not inferred: deleting the unit-sum renormalisation from the shipped stencil +reproduces the prototype **bit-for-bit** (max|diff| = 0.000e+00 over 5000 positions including +out-of-buffer). The renormalisation -- which moves the weights by at most 3.1e-4, and which the +docstring justifies as making interpolation exact for constants -- is the entire difference. + +**The mechanism is NOT established, and two plausible ones were tested and died.** (i) "the +normalisation fixes DC gain and thereby tilts the passband": a bandwidth sweep from f/fNyq 0.05 +to 0.98 on a synthetic band-limited signal shows no trend, ratios scattering over 0.79-1.70 with +1.00 at the top of the band. (ii) "it is a systematic gain deficit worth rho^2*(1-g)": the +best-fit complex gain differs between the two by ~1.5e-4, implying ~0.2 nats, three orders of +magnitude short of 92. Whatever produces the 3G effect is not reproduced by a single-detector +synthetic and has not been isolated -- it may involve the 3-site network, the finite-size +response, or the distance marginalisation. **Do not quote a mechanism for this row.** + +**A second configuration does not replicate the ordering, so do not generalise the row above.** +Same demo, same everything, at fmax 2048 / srate 4096: + +| stencil | peak offset | lnL at truth | shortfall | +|---|---|---|---| +| `cubic` | 0.571' | 179799.15 | 200.9 | +| `sinc` (shipped, a = 8) | **0.030'** | 179960.42 | 39.6 | +| `lanczos8` (a = 8, unnormalised) | 0.145' | **179999.42** | **0.6** | +| `lanczos16` (a = 16, unnormalised) | 0.031' | 179982.00 | 18.0 | + +Two things break here that hold at fmax 1024. **(i) The two metrics disagree.** In lnL at truth +the renormalisation still costs (39 nats, same sign as the 92 at fmax 1024), but in sky offset it +has *flipped*: shipped `sinc` at 0.030' now beats the unnormalised `lanczos8` at 0.145'. Judge +this knob on both moments or it will tell you whatever you asked. **(ii) Half-width is +non-monotone**: `lanczos16` claims *less* lnL at truth than `lanczos8` (179982 against 179999), +which a purely stencil-limited error cannot do. Something other than the stencil is binding at +this sampling -- consistent with the `slowrot_fs_lib` at-Nyquist defect being tracked separately. + +That is also an independent cross-check on that separate defect, arrived at without using any of +the sampling evidence from the sky-offset diagnosis: at fmax 1024 even `lanczos16` leaves 55.5 +nats against a 14-23 nat marginalisation floor, so **33-41 nats there are not stencil error**. + +So "should the shipped stencil renormalise?" is an open question with a measured cost in one +metric at two configurations and a measured *reversal* in the other, **not** a settled defect. It +is deliberately NOT changed here: the renormalisation is shared by all four backends, so flipping +it changes results for every existing `sinc` user, and the §3-§4 crossover tables were all +measured with it on. Anyone revisiting it needs more than two configurations and should report +offset and lnL shortfall side by side. + + +`cubic` is not merely less accurate here: it moves the likelihood peak 9.2 arcmin off a +zero-noise self-consistent injection and buys 64.5 nats for doing so. This configuration is far +below the crossover of §1 (M = 3 Msun, fmin 50), which is exactly the regime §1 assigns to +`sinc` -- the demo is a confirmation of that table, not a counter-example to it. + +### 9.4 The JAX default moved from `linear` to `sinc` (2026-08-26) + +`RIFT.likelihood.jax_ile.core.JAX_INTERP_DEFAULT` is the single definition, consumed by every +entry point in the package and by `--interp` in `bin/integrate_likelihood_extrinsic_jax`. +**This changes results for any caller that did not pass `interp=`/`--interp`.** Pass +`interp="linear"` to reproduce a pre-2026-08-26 run. + +**Why it moved at all.** `linear` was never chosen on merit -- it was the first differentiable +stencil the JAX path had, and the default was inherited. It is the *worst* option in the registry +at high SNR, worse than `nearest`, and this driver is used **exclusively** at high SNR, so the +one regime the default was tuned for is the one regime it is wrong in. On the 3G demo `linear` +sits between `nearest` and `cubic`; §3's ladder puts `nearest` at 200-443 nats, crossing 1 nat of +error by SNR 2-6. + +**Why `sinc` rather than `cubic`,** given §1 says neither is universally better: a default is +chosen for its *worst* case, not its average, because it is what people get without thinking. +§5 measured `sinc`'s error as FLAT -- 3.1-7.9 nats across the fmin-30 mass ladder, 2.3-5.6 across +the 20-point fmin sweep -- while all the variation belongs to `cubic`, whose error ranges 0.143 +to 69.3 nats over the same points. Error also grows as SNR^2 (measured exponent 1.999-2.006), so +at 3G sensitivities the tail is what matters. `sinc` costs ~1.6-3.0x `cubic` in the Q product on +GPU (§7), which is the price of that bound. + +**THE CONCERN, recorded because it is not resolved by the above.** `sinc` is the bounded choice, +not the better one, and there is a real population for which this default is *worse* than the old +behaviour would have been had anyone set `cubic`: §4 measures `cubic` winning by 2.1-4.4x over +9-55 Msun at fmin <= 50, and by 9.1x at 80 Msun and 55x at 120 Msun on the fmin-30 ladder. A +high-mass, low-fmin BBH -- much of the O4 catalogue -- is squarely in `cubic`'s regime, and will +now silently get `sinc` unless the caller says otherwise. The mitigation is that `sinc`'s loss in +that regime is bounded (its error never leaves 2.3-7.9 nats) whereas `cubic`'s loss in the other +regime is not, so the asymmetry favours `sinc` for a default -- but "bounded loss" is still loss, +and anyone running a high-mass campaign should pass `--interp cubic` explicitly. + +Two further limits on that reasoning, both from §8: every crossover in this document is at +**srate 4096**, which has *not* been swept and is the numerator of the ratio §6 says sets the +answer; and the flatness claim for `sinc` rests on the same measurements. If srate turns out to +move the crossover as strongly as fmin did, this default should be revisited -- it would be the +third time a rule here was overturned by an axis that had not been swept. + +### 9.5 The tap-axis memory trap, and why `u` is passed separately + +Found by adversarial review of §9.4. The first version of this section was **wrong**, and how it +was wrong is the point: it quoted 2.7x, measured on the **CPU** backend. On CPU nothing fuses +well, so the sinc overhead hid inside an already-large baseline. On GPU -- the backend this +driver actually runs on -- XLA fuses `nearest`, `linear` and `cubic` *completely*, and only sinc +materialised its `(S, npts, 2a)` weight array. The real figure was **65x**, not 2.7x. + +Whole-likelihood XLA temp, `compile().memory_analysis()`, 3 detectors, npts 614: + +| S | backend | `nearest` | `linear` | `cubic` | `sinc` before | `sinc` after | +|---|---|---|---|---|---|---| +| 20000 | CPU | 1285 MB | 2464 MB | 4821 MB | 6586 MB | -- | +| 20000 | **GPU** | 101 MB | 101 MB | 101 MB | **6583 MB** | **1279 MB** | +| 2000 | GPU | 10.1 MB | 10.1 MB | 10.3 MB | 658 MB | 128 MB | + +**The fix: pass the fractional offset separately.** Both accumulators build +`pos = p0[:, None] + arange(npts)` with INTEGER offsets, so `frac(pos)` does not vary along the +time axis -- only `S` distinct values exist, not `S * npts`. `_separable_u(p0)` returns that +`(S, 1)` array and every gatherer takes it as an optional third argument, so the weight array +becomes `(S, 1, 2a)` and is small enough that the surrounding product and reduction fuse, exactly +as cubic's inline weights already do. Isolated gather at S=20000/npts=614: +**1719.2 MB -> 2.7 MB, a 637x reduction.** + +**It is a strict win, not a trade** -- measured on GPU (RTX 3080, container jaxlib): + +| axis | before | after | +|---|---|---| +| whole-likelihood temp, S=20000 | 6583 MB | **1279 MB** | +| runtime, S=2000 | 0.00540 s/call | **0.00269 s/call** | +| accuracy | -- | unchanged, see below | + +Runtime *halves* because the old form recomputed 16 sinc pairs per (sample, time-bin) when only +`S` distinct weight rows exist -- 614-fold redundant arithmetic. Post-fix sinc costs 1.33x cubic +(0.00269 against 0.00203) and 1.44x linear, which is the honest price of the stencil. + +**And it is slightly MORE accurate.** `p0` is a sample index of order 1e5-1e6, so `p0 + t` can +cross a binade and drop a low mantissa bit; `frac(p0 + t)` then differs from `frac(p0)` by up to +an ulp of the position (~1.5e-11 at 65536, measured). The numpy, cupy and CUDA backends all +compute one fractional offset per sample *from the sample position* -- i.e. the separable form -- +so this brings JAX into line with them rather than away. + +Residual: sinc is still 12.7x cubic at the whole-likelihood level (1279 vs 101 MB), because the +per-(detector, mode) `(S, npts, 2a)` *gathered-value* array is only partly fused. At the +production default `--n-chunk 8000` that is ~512 MB against ~40 MB -- comfortable on any card +this runs on. Squeezing the last factor would mean fusing the lm contraction into the gather, as +the CUDA kernels do; not attempted. + +**`u` is passed only to the weight-building stencils.** Feeding it to `nearest`, which ignores +it, is not free: an unused extra input in the banded slow-rotation trace -- which is compile-bound +rather than arithmetic-bound -- cost **>60% wall** (`test_rotation_path_a`, 69.8 s -> >113 s, +same tree, same load, one variable). So the call sites read +`u_sep = None if interp == "nearest" else _separable_u(p0)`. The memory win is unaffected: it +belongs to `cubic`/`sinc`, which still receive it. This is the same lesson as the unrolled-tap +compile blow-up in `_make_gather_sinc` -- in a large trace, graph shape can cost more than +arithmetic. + +Getting `u` wrong is **silent** -- the gather returns the right shape evaluated at the wrong +offsets -- so it is pinned two ways: `test_separable_u_matches_the_general_path` compares the two +paths for all four stencils at production magnitudes, and `test_accumulators_pass_separable_u` +parses `core.py` to assert both call sites actually pass it. Mutation-tested: perturbing `u` by +0.05 fails the first for `linear`/`cubic`/`sinc` (and correctly not for `nearest`, which ignores +it); dropping the argument at either call site fails the second. + +A `lax.scan` variant was also measured and rejected -- it cut the temp only to 1427 MB and cost +2.9x runtime, against the separable form's 1279 MB at 0.5x runtime. + +### 9.6 The batchmode default moved from `nearest` to `sinc` (2026-09-02) + +`bin/integrate_likelihood_extrinsic_batchmode`'s `--interpolate-time` defaulted to `False`, i.e. +`nearest` — the stencil §3 measures at 200–443 nats and which reaches 1 nat of error by SNR 2–6. +It now defaults to `RIFT.likelihood.time_interp_choice.TIME_INTERP_DEFAULT`, which is `'sinc'`. +**This changes results for any ILE run that did not pass `--interpolate-time`,** which is +essentially every pipeline-driven run: `helper_LDG_Events.py` emits the flag only when +`--internal-ile-interpolate-time` is given. **Pass `--interpolate-time nearest` to reproduce a +pre-2026-09-02 run.** Filed as issue #233. + +`JAX_INTERP_DEFAULT` is now an *alias* of the same constant rather than a second literal `"sinc"`. +That is the structural half of #233: the two drivers previously shipped opposite defaults for the +same physical choice, so any cross-implementation comparison run at defaults was measuring a flag, +and — because stencil error grows as SNR² — the disagreement presented as an amplitude-dependent +bug in one of the codes rather than as a configuration difference. + +**Why `sinc` and not `cubic`.** Same reasoning as §9.4, re-measured independently on a +production-shaped grid (below): a default is chosen for its *worst* case. Over 15 (mass, fmin) +points spanning 10–120 M☉ and fmin 20/30/100, `sinc`'s error stays inside **1.59–4.58 nats** while +`cubic`'s ranges **0.078–27.0 nats**. `cubic` is the better choice over much of that grid and is +one flag away; it is not the safer *default*. + +#### 9.6.1 Accuracy, re-measured (2026-09-02) + +`study_stencil_lnL_sensitivity.py --mode mass-ladder`, SEOBNRv4, H1L1V1 zero noise, aLIGO ZDHP, +srate 4096, fmax 1700, Lmax 2, every mass normalised to SNR_lik = 100, K = 400 × 2 seeds, against +the same FFT-zero-padded reference §3 uses. max|ΔlnL| in nats; **winner** in bold. + +**THE REFERENCE IS NOT "EXACT", AND THIS DOCUMENT SAID IT WAS.** Both §3 and the paragraph above +described it as an *exact* FFT-zero-padded reference. `study_stencil_lnL_sensitivity.eval_reference` +zero-pads by M = 32 -- that step *is* exact for a periodic band-limited signal -- and then does a +**`cubic` lookup on the fine grid** (`REF_STENCIL = 'cubic'`). The reference therefore shares a +method with one of the three things it is judging, which is the shape of error a reader is entitled +to be told about. The harness's own docstring is honest about it; the summaries were not, and are +corrected here. + +It does **not** move any conclusion. Audited (2026-09-03) against an independent construction whose +reference is an exact frequency-domain sub-sample shift, `N·ifft(W·exp(2πi f u Δt))`, sharing no +machinery with any stencil: this reference's own error is **7.5e-6 nats at M = 20 and 7.7e-7 nats at +M = 120** (fmin 20, peak lnL normalised to 5000), i.e. six orders of magnitude below the 1.6-4.8 nats +it is used to resolve, and consistent with the O((1/M)^4) ~ 1e-6 bound `eval_reference` claims and +with the per-stencil reference floors recorded in §10. What has **not** been audited independently is +the periodic-wrap / Gibbs component from zero-padding a *cut* of rho(t); that still rests on the +harness's own 6.5e-3 nat internal bound. + +| fmin | M/M☉ | nearest | cubic | sinc | winner, margin | +|---|---|---|---|---|---| +| 20 | 10 | 145.9 | 3.507 | **1.625** | sinc 2.16× | +| 20 | 20 | 403.2 | 4.810 | **1.739** | sinc 2.77× | +| 20 | 35 | 176.4 | **1.974** | 2.460 | cubic 1.25× | +| 20 | 65 | 339.8 | **0.396** | 2.174 | cubic 5.49× | +| 20 | 120 | 211.8 | **0.078** | 2.703 | cubic 34.5× | +| 30 | 10 | 189.3 | 5.520 | **2.306** | sinc 2.39× | +| 30 | 20 | 142.1 | 4.306 | **1.938** | sinc 2.22× | +| 30 | 35 | 169.9 | **1.616** | 2.175 | cubic 1.35× | +| 30 | 65 | 190.4 | **0.510** | 1.586 | cubic 3.11× | +| 30 | 120 | 209.9 | **0.095** | 2.924 | cubic 30.7× | +| 100 | 10 | 652.8 | 17.13 | **4.582** | sinc 3.74× | +| 100 | 20 | 515.1 | 27.04 | **2.387** | sinc 11.3× | +| 100 | 35 | 620.4 | 5.940 | **2.104** | sinc 2.82× | +| 100 | 65 | 411.1 | **1.688** | 2.726 | cubic 1.61× | +| 100 | 120 | 203.2 | **0.183** | 2.701 | cubic 14.8× | + +**Does this reproduce §3–§4?** The *winner* agrees at every one of the ten points these two +measurements share, including the two cells §4 flags as flipping with fmin: M = 35 goes cubic at +fmin 30 and **sinc at fmin 100**, exactly as §4 records. Ratios agree closely where §4 quotes one +(M = 20, fmin 30: §4 2.2×, here 2.22×; M = 35, fmin 100: §4 2.5×, here 2.82×; M = 20, fmin 100: +§4 8.6×, here 11.3×). + +**The ABSOLUTE numbers here are systematically smaller than §3's and that is expected, not a +disagreement.** max|ΔlnL| is an extreme-value statistic over the drawn extrinsic points, and this +table uses K = 400 × 2 seeds against §3's K = 2000 × 3. Compare the ratios, or re-run at §3's K +before comparing the nats. + +**fmin 20 is new here** — §3–§4 start at fmin 20 only in the sweep, and §3's ladder is at fmin 30. +It matters because fmin 20 is the O4 production value, and it is the fmin at which `cubic`'s +advantage at high mass is largest (34.5× at 120 M☉). The default is still `sinc` because the +comparison that decides a default is between the two WORST cases, and at fmin 20 those are +`sinc` 2.70 nats against `cubic` 4.81 nats. + +#### 9.6.2 Cost, measured (2026-09-02) + +Per `DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` call — **that is the denominator**: an ILE +job's wall time also carries precompute and sampler overhead, so its end-to-end ratio is bounded +above by these and was not measured. 3 detectors, Lmax 2, SEOBNRv4, srate 4096, fmax 1700. Median +over repeats with the first call discarded (numba/cupy compile, first-touch allocation); on GPU +every timed call is device-synchronized. + +**CPU** (`citlogin6`, uncapped and quiet — load average 3.6 before, 4.6 after; IGWN CVMFS python +3.11, `OMP_NUM_THREADS=1`). K = 2000 extrinsic points, npts = 411, median of 5: + +| M/M☉ | fmin | nearest (s) | cubic (s) | sinc (s) | cubic/nearest | sinc/nearest | sinc/cubic | +|---|---|---|---|---|---|---|---| +| 20 | 20 | 0.108 | 0.520 | 1.897 | 4.83 | 17.6 | 3.65 | +| 20 | 100 | 0.115 | 0.668 | 1.879 | 5.79 | 16.3 | 2.81 | +| 35 | 20 | 0.107 | 0.504 | 1.864 | 4.70 | 17.4 | 3.70 | +| 35 | 100 | 0.108 | 0.501 | 1.858 | 4.66 | 17.3 | 3.71 | +| 65 | 20 | 0.112 | 0.508 | 1.829 | 4.54 | 16.4 | 3.60 | +| 65 | 100 | 0.105 | 0.496 | 1.824 | 4.74 | 17.4 | 3.67 | + +**GPU** (`ldas-pcdev13` device 0, RTX PRO 4000 Blackwell, **idle — one job on the device**; +`rift_o4d_cc90-120_cuda128_20260717.sif`, cupy 14.1.1). Time-marginalized (the production +reduction), median of 7: + +| K (≈ `--n-chunk`) | npts | nearest (s) | cubic (s) | sinc (s) | cubic/nearest | sinc/nearest | +|---|---|---|---|---|---|---| +| 2000 | 411 | 0.0137 | 0.0140 | 0.0157 | 1.02 | 1.15 | +| 10000 | 411 | 0.0165–0.0176 | 0.0175–0.0200 | 0.0205–0.0214 | 1.01–1.21 | 1.18–1.30 | +| 40000 | 411 | 0.0386 | 0.0432 | 0.0543 | 1.12 | 1.41 | +| 10000 | 1229 | 0.0263 | 0.0289 | 0.0372 | 1.10 | 1.42 | + +**Read the small-K rows as overhead-diluted, not as the cost of the stencil.** At K = 2000 the +call is launch-bound and every stencil looks free; the ratio grows toward ~1.4× as K and npts rise +into and past the production chunk. The K = 10000 row is quoted as a range because it was measured +in two separate sweeps whose results differ by ~8% — run-to-run spread on a shared node, not a +mass or fmin dependence (the six-point mass × fmin grid at K = 10000 is flat to 1.18–1.21×, which +is what a stencil cost should do). + +**Provenance of these timings: they were taken at `6565d68c`, before PR #234's +per-extrinsic-sample log-sum-exp offset landed in the same function.** That change replaces a +scalar `max(lnL_t)` with a `keepdims` row-wise max in the time-marginalized reduction, so the +absolute seconds in the GPU table may shift slightly. It is stencil-independent by construction -- +the same reduction runs for `nearest`, `cubic` and `sinc` -- so the RATIOS, which are what this +section is for, are unaffected. Not re-timed. + +**So the cost objection is a CPU objection.** Production ILE runs `--gpu`, where `sinc` costs +1.15–1.42× `nearest` and `cubic` 1.01–1.21×. On CPU `sinc` costs 16.3–17.6× `nearest`, which +reproduces the ~16× end-to-end figure reported in issue #233 and identifies that measurement as +CPU-bound. + +**Two comparisons against §7 and the shipped help text, both with different denominators — do not +read either as a contradiction.** §7's "~4.2–4.5× on CPU, ~1.6–3.0× on GPU" is `sinc` vs `cubic` +**in the Q product alone**; measured here at the whole-likelihood level it is 2.8–3.7× on CPU +(consistent) and 1.07–1.29× on GPU (much smaller, because the surrounding likelihood dominates). +§7's end-to-end CPU figures (nearest 9.3 s, cubic 25.1 s, sinc 85.3 s → 2.7× and 9.2×) are at a +fixed n_max on a different configuration; the per-call ratios here are larger. Quote whichever +denominator matches what you are describing, and say which. + +#### 9.6.3 What a DEFAULT had to be prevented from doing + +The one-line change is not the whole change. `--interpolate-time` had three behaviours keyed on +`!= 'nearest'` that were written when `nearest` was the default, and each would have fired on +runs that pass no flag at all. The driver therefore distinguishes an **explicit request** from an +**inherited default** (`opts._interp_time_from_default`, decided on `is None` before any string +coercion, because `str(None) == 'none'` is itself a legal explicit spelling meaning `nearest`): + +1. **The honoured-path gate would have become a startup crash.** A stencil is only honoured under + `--time-marginalization --vectorized` plus one of `--gpu`/`--rotation-slow`/`--freqresponse`, + and anything else is *refused*. As a default that refusal turns every other configuration — + including a bare invocation — from working into `ValueError`, with no command line changed + anywhere. An explicit request is still refused, unchanged; a default falls back to `nearest` + and prints why. +2. **The time-posterior export mode would have flipped.** `resolve_time_posterior_export_mode` + maps `auto` to `continuous` for any non-`nearest` stencil, so the same edit would have changed + the fair-draw time export of every `--resample-time-marginalization` run: a re-evaluation of + the whole likelihood on a ≥4× denser grid, a different draw algorithm, two extra output + columns, and `validate_time_posterior_working_set`'s `MemoryError` newly reachable. The export + now keys on an explicit stencil only. Asking for a stencil still opts in; + `--time-posterior-export continuous` still works on its own. +3. **The fused calibration kernel would have been silently abandoned.** The fused calmarg kernels + implement `nearest` only (§9), and the driver's three NoLoop call sites fall back to + `cal_method='loop'` — and drop the `cal_distmarg` table — for any other stencil. A default must + not spend someone else's `--calibration-fused-kernel` that way, so it stays `nearest` there. + With an *explicit* stencil the behaviour is unchanged but is no longer silent: the driver now + prints that the fused kernel is not in use. + + **Corrected 2026-09-03 (P2 review, PR #237): this guard was BROADER THAN THE THING IT + PROTECTS.** It keyed on `--calibration-fused-kernel` alone, but `use_fused_calmarg` (`:3261`) + is `calibration_marginalization and opts.calibration_fused_kernel`, and + `calibration_marginalization` (`:1317`) is exactly `bool(opts.calibration_envelope_directory)` + — so with no envelope configured **no fused kernel can run under any stencil**, and downgrading + there protected nothing while silently costing the accuracy the default exists to provide. That + failure mode is invisible by construction: a needless downgrade looks exactly like the + historical behaviour. Both the downgrade and the `NOT USED` notice now key on + `fused_calmarg_in_use`. Pinned by + `test_a_bare_fused_kernel_flag_no_longer_downgrades_the_default`, whose command line is the + fused-kernel test's minus the envelope, with the opposite required outcome. + + **THE FINDING IS THE DUPLICATION, NOT ANY ONE MISSING CONJUNCT.** That correction was made + twice more the same day — `flag` → `flag + envelope` → `flag + envelope + path` — each round + adding a term the dispatch-time expression already had in effect, each round found by a + reviewer rather than by a test. It could not be otherwise: **an over-broad predicate downgrades + the inherited stencil to `nearest`, which is exactly the historical behaviour**, so nothing + fails, nothing logs, and no value test distinguishes it. The condition is now a single function, + `fused_calmarg_in_use`, called by *both* the startup stencil guard and `use_fused_calmarg`, and + `test_the_fused_predicate_has_ONE_definition_shared_by_BOTH_call_sites` pins the **structure** + (one definition, ≥2 call sites, no re-inlined copy) rather than a value, because the next + forgotten conjunct will produce a fourth expression that agrees with every case anyone thought + to write down. + + A fused kernel runs only where all of these hold, and each is a **reachability fact** verified + against the call sites rather than a guess: calibration marginalization configured; + `--calibration-fused-kernel`; `--calibration-n-realizations > 1` (`factored_likelihood` returns + from its `n_cal == 1` branch before `cal_method` is read); not `--rotation-slow` and not + `--freqresponse` (both *replace* the likelihood — the non-distmarg fused call site is in the + `else` of their dispatch, and the two distmarg call sites are unreachable because both options + refuse distance marginalization at startup); and not `--calibration-dump-responsibilities` + (that pilot evaluates with an explicit `cal_method='loop'` and `return_cal_components=True`, + then returns before the production integration is built). The five negative cases are pinned + one per reason in `test_no_fused_configuration_that_cannot_fuse_downgrades_the_default`, + because they fail for different reasons and one expression has already been wrong about three + of them. + + **What the early call cannot see, stated rather than left to be rediscovered.** At startup + `calibration_marginalization` does not exist yet; it is set several hundred lines later by + `if opts.calibration_envelope_directory:` from a `False` initial value, so the early call + substitutes that option — a restatement of the same condition, not an approximation. It is + also *checked*: the dispatch-time call passes the real variable and the driver **refuses** if + the two disagree, so a change to how `calibration_marginalization` is derived fails loudly + instead of silently re-opening the drift. Nothing else in the predicate is unavailable at + startup. + + **Related, found while surveying the other consumers, NOT fixed here.** + `bin/integrate_likelihood_extrinsic_batchmode:535` and `:548` guard + `--rotation-slow` / `--freqresponse` on GPU against calibration marginalization with + `getattr(opts, 'calibration_marginalization', False)`. **There is no + `--calibration-marginalization` option and nothing ever assigns that attribute**, so both + guards are dead and have never fired; the live state is the module-level + `calibration_marginalization`, which is not set until much later, so the check would have to + read `opts.calibration_envelope_directory` — the same substitution `fused_calmarg_in_use` + documents. Reported on PR #237 for the separate early-rejection change, which covers exactly + these combinations. + + **And the downgrade notice promised behaviour the driver does not have.** One shared sentence + served both downgrades and told everyone that naming the stencil explicitly would get them + *refused*. That is true of the prerequisite downgrade and **false** of this one: + `--calibration-fused-kernel --interpolate-time sinc` is **accepted**, the fused kernel is + dropped, and the driver says so. Contradictory guidance is worse than none in a + result-changing startup notice, so each downgrade now carries its own remedy and they are + printed one line per reason. `test_each_downgrade_states_the_remedy_that_actually_applies` + asserts each remedy against the driver's actual behaviour on the command line it recommends, + not just against the string. + +4. **The legacy scalar path would have started interpolating.** `opts._legacy_interpolate_time` + is derived from the *provisional* default a hundred lines before the downgrade runs, and it is + not a stencil: it is the plain boolean handed to + `FactoredLogLikelihoodTimeMarginalized(..., interpolate=...)` at `batchmode:3645`/`:4002`, which + is the likelihood that actually runs whenever `--vectorized` is absent. A non-`nearest` + provisional default makes it `True`, so an omitted flag would have switched every non-vectorized + run onto that path's unrelated cubic interpolation — on a path that has no sub-sample stencil at + all. The downgrade resets it. **This one was missed on the first pass**: the code was correct but + nothing pinned it, and adversarial review (2026-09-03) found it by mutation — deleting the reset + left the whole gate green. Now pinned by + `test_the_downgrade_also_takes_the_legacy_scalar_path_down_with_it`. + +A fifth, in the pipeline: `resolve_interpolate_time_request` collapses "flag absent" and an +explicit off-request (`--internal-ile-interpolate-time False`) to the same `None`, and both used +to emit nothing. While the driver default was `nearest` those were the same answer; they are now +opposites, so `helper_LDG_Events.py` re-expresses an off-request as an explicit +`--interpolate-time nearest`. Without that, **"off" would have meant "on"**. + +#### 9.6.4 The concerns, recorded because they are not resolved by the above + +- **High-mass, low-fmin BBH — much of the O4 catalogue — is the population this default is worse + for.** §9.6.1 measures `cubic` winning by 5.5× at 65 M☉ and 34.5× at 120 M☉ at fmin 20. The + mitigation is the same asymmetry §9.4 relied on and this table re-measures: `sinc`'s loss there + is bounded (2.17 and 2.70 nats at SNR 100), `cubic`'s loss in the other regime is not (27.0 nats + at M = 20, fmin 100). Anyone running a high-mass campaign should pass `--interpolate-time cubic`. +- **"Bounded" means FLAT IN MASS at fixed SNR, not small — and at O4 SNRs this choice barely + matters.** Scaling this table's own SNR² exponent off its SNR = 100 normalisation, a loud O4 + event at network ρ = 30 sees, at fmin 20: `nearest` 13–36 nats (this is what actually justifies + the change), `sinc` 0.15–0.24 nats, `cubic` 0.007–0.43 nats. **The `sinc`-vs-`cubic` tie-break is + sub-nat for the current catalogue, high-mass included.** The real case for this PR is therefore + "anything but `nearest`", plus agreement with the jax driver (#233) — not an accuracy win over + `cubic`. At 3G amplitudes *both* are ruinous at high mass (`sinc`'s 2.70 nats at ρ = 100 is + 270 nats at ρ = 1000, where `cubic` would be 7.8), so the earlier phrasing "`cubic`'s loss is not + bounded" should be read as **"not flat across the grid"**: both are finite. The minimax ordering + is unchanged and still selects `sinc` (worst case 4.58 against 27.0, or 458 against 6900 scaled + to ρ = 1000), so the decision stands — but it stands on predictability, not on magnitude. +- **`--time-marginalization-quadrature bandlimited` was measured against `nearest`, and its + prerequisites are a strict SUBSET of the stencil's honoured set.** `bandlimited` requires + `--time-marginalization --vectorized --gpu`, which is exactly what makes the stencil honoured — + so **every** `bandlimited` user who does not pass `--interpolate-time` now lands in the regime + below, and the driver prints nothing about it (verified 2026-09-03). That + module's own docstring records +0.0002 nats for `nearest` against an analytic truth where + Simpson is −521, but −2.29 for `sinc` where Simpson is +1.28, "and over a scan of seeds and + grid phases Simpson wins about half the cases". The stencil default change moves that opt-in + quadrature into the regime where its advantage is not established. **This pairing has not been + re-measured and is an open item**, not a settled result. + + **The SILENCE is closed (2026-09-03, P2 review, PR #237); the MEASUREMENT is not.** The driver + now prints `--time-marginalization-quadrature : ADVANTAGE NOT ESTABLISHED with the + DEFAULT/explicitly requested Q_lm stencil ` whenever a band-limited quadrature runs under + a non-`nearest` stencil, carrying both numbers and the `--interpolate-time nearest` reproduce + instruction. Not a downgrade — the quadrature is not wrong, its *advantage* is unestablished — + and not a refusal, because either stencil is a legitimate choice. `peak-local` is included + because its accuracy is *defined* against `bandlimited` (max 1.9e-11 nats), so a stencil bias in + the reference is a stencil bias in it. Pinned at both edges: the notice must fire for + `bandlimited` and `peak-local` under the default, and must stay silent for `simpson` and for an + explicit `nearest` — a notice that always prints carries no information, and one that fired on + `nearest` would contradict its own advice. + **Still NEEDS AN OWNER: whether `bandlimited`'s advantage survives against `sinc` is unmeasured.** +- **The measurements above are for the BASELINE likelihood only.** `--rotation-slow` and + `--freqresponse` also satisfy the honoured set, so they inherit the new default too. Neither goes + silently inert — `factored_likelihood_with_rotation` and `factored_likelihood_freqresponse` both + implement all three stencils and call `validate_time_interp` — but §9.6.1 did not measure them. +- **One line in the downgrade block is belt-and-braces, not a guard.** The + `resolve_time_posterior_export_mode(..., 'nearest', ...)` recompute inside it can never change + the stored value, because the block only runs when the stencil came from the default and + `_ti_for_export` was already `'nearest'`. No test can distinguish it (verified by mutation, + 2026-09-03). Recorded so it is not mistaken for coverage; annotated in the source. +- **srate is still unswept.** Every crossover in this document is at srate 4096, which §8 names as + the numerator of the ratio §6 says sets the answer. If srate moves the crossover as strongly as + fmin did, this default should be revisited. +- **The LISA twin was deliberately not changed.** `integrate_likelihood_extrinsic_batchmode_lisa` + keeps `--interpolate-time default=False` and its own `legacy_time_interpolation_enabled` + parsing. The two executables therefore now differ in default. That is a scope cut, not an + oversight: the LISA driver has a separate drift-ledger gate and its own export contract. + + +## 10. Provenance The fmin sweep was measured against a pinned `git archive` of the #97 merge commit `c1a2e2df`, not a shared checkout, so a branch switch could not move code mid-run. Its fmin-30 column diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_peak_local.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_peak_local.md new file mode 100644 index 000000000..94cb621e8 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_peak_local.md @@ -0,0 +1,1244 @@ +# Peak-local time marginalization: measured record + +Companion to `time_marginalization_peak_local.py`, and a follow-up to +`DESIGN_time_marginalization_quadrature.md` — read that one first. The module +docstring carries the argument; this file carries the numbers behind it and the +harnesses that produced them. + +**Generalizing this method to other axes:** see `DESIGN_peak_local_framework.md`. +The same skeleton has since been written twice more, independently, for angles and +for distance; that note works out which parts of what follows are about PEAK-LOCAL +and which are about TIME, and it is the place to look before copying anything here +onto a new parameter. + +Everything here was measured on `ldas-pcdev` class CPU (CIT), CVMFS IGWN python 3.11, +`OMP_NUM_THREADS=1`, on branch `rift_O4d_tmarg_peaklocal` (based on +`rift_O4d_tmarg_bandlimited`, PR #203). Harnesses: `~/tmarg_harness/` for the +prototypes, `~/pl_work/` for the measurements below. + +## THE BUG THIS SHIPPED WITH, and the requirement that was missing + +An adversarial review of PR #205 found a critical correctness bug. It is recorded +first because the design statement below was *wrong*, not merely incomplete, and the +shape of the error is the reusable lesson. + +The design named TWO resolution requirements. **There are three.** Enumeration +returns a grid INDEX, and an index is not a location: the true crest can lie +`h_enum/2` from the sample that reports it. Localising it to a fraction of `sigma_t` +is a third requirement, it is SNR-DEPENDENT, and it belonged to neither of the two +named mechanisms — so it was assigned to neither and simply did not happen. The +interval was built around `cols_np * h_enum`, the grid sample. Whenever + + W_SIGMA * sigma_t < h_enum / 2 i.e. sigma_t/deltaT < 0.0052 + +— any row with a derived factor of 512 or more — **the peak lay entirely outside its +own interval**. Measured on the synthetic fixture at `sigma_t/deltaT = 0.0024`, error +against the dense path as the crest is walked off the enumeration grid: + +| crest offset from the sample | 0 | h_enum/4 | h_enum/2 | +|---|---|---|---| +| peak-local − reference | +0.000000 | **−6.52** | **−164.93** | + +Always negative: it deletes mass, never adds it. At production scale (64 rows, +arrival times uniform w.r.t. the sample grid, `sigma_t/deltaT = 0.0023`): **median +−1.15 nats, worst −131 nats, 56% of rows wrong by >0.01 nats, and every one of them +ACCEPTED.** A bias, not noise, silently deleting extrinsic samples. + +**Why the suite was blind, which is the part worth internalising.** Every sharp +fixture used `peak_sample = NPTS//2 + 0.25`, and the phase sweep in +`test_exact_on_a_periodic_window` was `[0.0, 0.25, 0.5]`. All are exact multiples of +`1/PEAK_ENUM_FACTOR = 1/8`, so the crest landed EXACTLY on an enumeration sample in +every single test. A bug that only exists *between* samples was invisible to all of +them. That test's own docstring already said the phase is swept because "a fixture +pinned to one phase can be exactly wrong and look exactly right" — it was pinned, in +the only sense this module cares about. **One added phase of 0.3125 catches all of +it.** Sweeping a parameter is not the same as sweeping it over the quantisation that +matters. + +Two further holes the same review found, both fixed: + +* the tail bound could not catch this, because `q_out_max` is a maximum over + *enumeration-grid samples* — the same grid that failed. In the worst case it + reported `tail_bound_worst = -275` (claiming `e^-275` of omitted mass) while dropping + 164.9 nats. `T_outside` was also counted in grid indices, so a gap narrower than + `h_enum` contributed to neither the length nor the maximum; +* `enumerate_peak_indices` excluded endpoints on the justification that the edge guard + had already routed such rows away. That is false: `exposed` keys on the row's + GLOBAL coarse argmax, so a row with a mid-window dominant peak and a secondary peak + hard against an edge was refined with the secondary peak never enumerated. + +### The fix + +1. **`localise_peaks`** — Newton on the spectral interpolant, seeded at the enumerated + sample and confined to the bracket `[t_i - h_enum, t_i + h_enum]` that the + enumeration already established. It places peaks; it cannot find or lose one, which + is what distinguishes it from the seed-and-hope that sank RIFT PR #201. Quadratic + convergence, so the SNR-dependence of this step is logarithmic: 2–4 iterations + across the whole production range. Convergence to `LOCALISE_SAFETY * sigma_t` is + ASSERTED, the interval is widened by that residual, and a peak that misses it sends + its row to the dense path. +2. **An a-posteriori containment check** — the local integration grids must ATTAIN the + localised crest's `lnL` to within `CONTAINMENT_SLACK_NATS`. Free (that maximum is + already computed for the log-sum-exp offset) and it is what actually catches a + mis-placed interval. It compares against the LOCALISED crest, not the enumeration + sample: against the sample it would pass precisely in the case it must catch, since + an off-grid crest leaves the sample tens of nats low. +3. **Exact `T_outside`** from interval geometry rather than a grid-index count. +4. **Endpoints enumerated**, and the false justification deleted. +5. **The ceiling checked before the cost gate** (see below). + +After the fix, the same sweeps: **0.000000 at every crest offset**, and at production +scale median +0.000000, worst +0.000000, 0/64 rows outside 1e-3. + +`W_SIGMA`'s `erfc(W/sqrt2)` argument is a statement about a Gaussian truncated about +its CREST, so centring on the sample carried an unstated precondition +`W_SIGMA * sigma_t >= h_enum/2`, coupling `W_SIGMA` to `PEAK_ENUM_FACTOR` and nowhere +asserted — note that raising `PEAK_ENUM_FACTOR`, the intuitively safer move, would have +made it *worse*. Localisation discharges the precondition rather than asserting it. + +### F3: the ceiling was bypassed for exactly the sharpest rows + +`viable = c_lo < c_dn` compares against the dense cost, so the sharper the row the more +certainly peak-local kept it — and a row past `UPSAMPLE_FACTOR_MAX` is the sharpest kind +there is. At a derived factor of 8192 the dense path RAISED, as designed, while +peak-local returned −24451 nats and reported `tail_bound_worst = -2721`. The ceiling is +now checked before the cost gate. The old test only passed because it set +`UPSAMPLE_FACTOR_MAX = 2`, broad enough that the *cost* gate declined the row first; the +ceiling was never what routed it. + +## G-class: the defect was a CLASS, and fixing one site was not enough + +A third review, after the localisation fix, found the same quantisation error at three +more consumers of the enumeration index. The audit that finds them is one grep: +`cols_p` / `q_up[`. Every quantity computed from `q_up[rows_p, cols_p]` inherits it; +the index is legitimate only as a Newton SEED and bracket centre, never as a value. + +**G1 (critical).** `PEAK_KEEP_NATS` was applied to the SAMPLE value, before +localisation. A crest `d` from its sample reads `(d/sigma)^2/2` nats low, so a peak +between samples was compared against a peak on a sample and dropped. On the two-peak +fixture at rho ~ 700, shipped code, no monkeypatching: + +| B offset (h_enum) | peak-local − truth | n_peaks_total | +|---|---|---| +| 0.00 | +0.000000 | 2 | +| 0.25 | **−0.693147** | 1 | +| 0.50 | **−0.693147** | 1 | + +`−log 2` exactly: one of two equal peaks deleted. The secondary crest was 1.003 nats +below the dominant crest while its SAMPLE was 70.99 nats below. Fixed by comparing +crests: the vertex height of the parabola through the three stencil points is the crest +value exactly for a Gaussian peak, at any spacing — the same property the width +estimator uses, applied to the zeroth moment instead of the second. It costs nothing. +After: 2 peaks kept and +0.000000 at every offset. + +**G2.** The claim that a dropped peak is "safe because it enters the tail bound" was +vacuous. Since `q_out_max` is at least the dropped peak's own sample, the bound accepts +the row unless `sigma_t < 2.6e-18 s`, while `UPSAMPLE_FACTOR_MAX` bounds the sharpest +legal row at `3.0e-08 s` — ten orders of magnitude away. For EVERY legal row a +keep-filter drop is automatically accepted. The claim is removed; the filter now rests +on a direct magnitude argument tying `PEAK_KEEP_NATS` to `UPSAMPLE_FACTOR_MAX`, and that +inequality is asserted. + +**G3.** The containment check's scope is narrower than it looked: `row_star` is a per-row +maximum over KEPT peaks, so it verifies the dominant crest only and a dropped peak can +never enter it. Documented, not widened. + +**G4.** The "conservative superset" was not conservative: `localise_peaks` brackets at +`+/- h_enum` and accepts anything strictly inside, so `|t* - t_grid|` is bounded by +`h_enum`, not half of it — and the bound is approached (0.959 `h_enum` observed over +14,182 peaks). Widened to `+ h_enum`. Narrowing the bracket instead would be wrong: an +asymmetric peak's crest genuinely can sit more than half a cell from its sample. + +**G5.** The curvature stencil CENTRE was clipped inward by `maxd = 8`, so a peak within 8 +enumeration samples of an end had its width measured elsewhere, returning `sigma = inf` +at index 0 and dropping the peak before localisation could help. Now the centre stays +put and out-of-range half-widths are masked. + +**G6 — why the suite missed G1.** Every multi-peak fixture ran at `amp = 200` +(`sigma_t/deltaT = 0.0072`, just above the 0.0057 threshold); every sharp fixture was +single-peak; and the updated F1 tests moved BOTH peaks off-grid by similar amounts, so +the deficits cancelled. **The defect needs asymmetry** — one peak near a sample, one +between — and the suite never crossed "more than one peak" with "sharp". That cell now +exists. + +## What this changes + +The dense band-limited rule refines the WHOLE window to a peak whose width shrinks as +`1/rho`, so its cost grows exactly where the peak occupies least of the domain. This +splits the two resolution requirements the dense rule conflates: + +* enumerating the extrema of `kappa` — a small, **SNR-independent** factor, because + `kappa` is band-limited at Nyquist by construction; +* integrating `exp(lnL)` — the rho-dependent part, needed only within a few `sigma_t` + of each enumerated peak. + +Intervals of half-width `W_SIGMA * sigma_i` are built around every enumerated maximum, +**merged into disjoint intervals**, and each is integrated at its own derived spacing. + +## THE COST CLAIM: what the prototype's number was, and what this one is + +`DESIGN_time_marginalization_quadrature.md` tabulates the prototype +(`~/tmarg_harness/peaklocal2.py`) at a flat ~97 evaluation points against the dense +grid's 19,648 → 628,736, i.e. **"200x at rho ~ 15, 6,482x at rho ~ 692"**. + +**That measurement had the analytic `kappa` in hand**, so evaluating `kappa(t)` at an +arbitrary time was one closed-form call. In the shipped code only the coarse samples +exist and the band-limited interpolant must be evaluated at the local grid points, +which costs `O(npts)` per point. A point count is therefore not a cost, and **the +prototype's speedups are not this module's speedups.** Nothing below is inherited +from it. + +## Cost, measured END-TO-END through the shipped likelihood + +`DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`, `n_extrinsic = 4000`, 3 IFOs, +srate 4096, `npts = 614`, CPU `process_time`, same call and same inputs for all three +rules (`~/pl_work/cost_pl.py`, an extension of `~/tmarg_harness/cost_e2e.py`): + +| `sigma_t/deltaT*` | simpson | bandlimited | **peak-local** | bandlimited / peak-local | rows peak-local handled | +|---|---|---|---|---|---| +| 1.735 | 0.519 s | 0.559 s | 0.549 s | 1.02x | 0 / 20 | +| 0.549 | 0.536 s | 1.456 s | 1.535 s | 0.95x | 0 / 2345 | +| 0.174 | 0.569 s | 5.583 s | 7.465 s | 0.75x | 0 / 3486 | +| 0.055 | 0.550 s | 19.493 s | 9.328 s | **2.09x** | 1866 / 3834 | +| 0.017 | 0.569 s | 46.011 s | 6.005 s | **7.66x** | 3335 / 3950 | + +All five rows are on `ldas-pcdev13`, re-measured AFTER the localisation fix. An earlier +table in this file was taken on a different, faster host (Simpson baseline 0.18–0.30 s +rather than 0.55 s) and has been replaced rather than merged: seconds are not comparable +across hosts and mixing them would invent a trend. Read the RATIOS. + +What the fix cost: at the sharp end essentially nothing (7.78x → 7.66x, 2.03x → 2.09x, +measured before/after on their respective hosts), and about 19% at `sigma_t/deltaT = +0.17` (0.93x → 0.75x), which is a regime where this rule has nothing to offer anyway and +delegates every row. + +Read against Simpson instead, the same table says peak-local costs 0.9x / 2.5x / 9.6x / +12.8x / 8.8x the historical rule, where the dense rule costs 1.0x / 2.4x / 8.8x / +26.0x / 68.5x — i.e. **peak-local's cost stops growing with rho and the dense rule's +does not**, which is the structural claim, and is visible in the last two rows. + +A first attempt at the fix ran the localiser BEFORE the gate that discards the row, +which is the same error as the one recorded below and cost far more: **0.14x** at +`sigma_t/deltaT = 0.17`, i.e. 7x slower than the path it delegates to, on a block where +every single row fell back. The gate now runs first, on intervals built about the +sample and widened by `h_enum/2` — a CONSERVATIVE SUPERSET of whatever localisation will +produce, so a gate on it can only decline rows, never wrongly keep one. + +`sigma_t/deltaT*` is the sharpest row in the block. The Simpson baseline is +rho-independent by construction; a run where it moves with rho is contaminated. +Host-sensitive: the O4c effort measured the same quantity moving up to 2x between +hosts, so read the RATIOS, not the seconds. + +Two honest readings of that table: + +* The win grows with rho and is real at the sharp end, but it is **an order of + magnitude, not three**. The gap between this and the prototype's figure is + entirely the cost of evaluating the interpolant, which the prototype did not pay. +* At low rho this method LOSES to the dense rule, which is why the per-row cost gates + exist. They are cost decisions only — both branches satisfy the same derived + resolution criterion, so no accuracy is traded. + +### The gates, and a mistake worth recording + +A first version applied the cost comparison only AFTER enumerating. Rows that then +fell back had paid for an enumeration FFT they did not use, and the method came out +**slower than the dense path it delegated to** (0.43x and 0.74x at +`sigma_t/deltaT` = 0.55 and 0.17) even though every single row fell back. Two changes +fixed it, and the worst case is now 0.93x: + +1. before any work: `MIN_LOCAL_POINTS = 2*W_SIGMA*UPSAMPLE_SAFETY + 1 = 49` is the + fewest local points any row can need, since one interval at spacing + `sigma/UPSAMPLE_SAFETY` is `2*W_SIGMA*UPSAMPLE_SAFETY` sub-intervals whatever + `sigma` is, and merging only adds points; +2. the merge itself was vectorised across rows. A broad integrand has many enumerated + peaks, and a Python loop over them cost more than the enumeration did. The running + maximum is restarted at each row boundary by offsetting row `r` by `r * big`, so one + global `maximum.accumulate` gives the per-row running maximum with no segmentation. + +The residual 0.93x at `sigma_t/deltaT = 0.17` is the enumeration FFT for rows that then +fall back — the price of a gate that cannot be decided without looking. It is a 7% +tax on a regime where this rule has nothing to offer, and it is reported here rather +than tuned away. + +## Accuracy + +Against the **analytic** truth (`kappa` a sum of exponentials below Nyquist, so the +continuous integral is closed-form), srate 4096, npts 614, peak at sample 307.25: + +| `sigma_t/deltaT` | rho ~ | bandlimited − truth | **peak-local − truth** | rule used | local points | intervals | +|---|---|---|---|---|---|---| +| 0.743 | 2 | +2.3e-09 | +2.3e-09 | dense (cost) | — | — | +| 0.470 | 3 | +7.5e-10 | +7.5e-10 | dense (cost) | — | — | +| 0.255 | 6 | 0 | 0 | dense (cost) | — | — | +| 0.149 | 11 | +5.7e-14 | +5.7e-14 | dense (cost) | — | — | +| 0.105 | 15 | +5.7e-14 | +5.7e-14 | peak-local | 64 | 1 | +| 0.047 | 35 | +6.0e-13 | +8.1e-13 | peak-local | 64 | 1 | +| 0.017 | 98 | +6.4e-12 | +7.3e-12 | peak-local | 64 | 1 | +| 0.0074 | 219 | 0 | −3.6e-12 | peak-local | 64 | 1 | +| 0.0023 | 692 | +1.5e-10 | +1.5e-10 | peak-local | 64 | 1 | + +Errors at the 1e-14…1e-10 level are the REFERENCE's own resolution (a 2048x uniform +refinement of the closed-form `kappa`), not the method's; the two rules agree with each +other far more closely than either is being measured to here. The flat **64 local +points** across four decades of `rho` is the structural point: the local point count is +set by `W_SIGMA` and `UPSAMPLE_SAFETY`, not by the peak width, so it does not grow with +SNR while the dense grid's does (`npts * factor` = 19,648 → 628,736 over the same rows). +What that flat point count is NOT is a flat cost — see above. + +And against the dense band-limited path row by row over the 4000-row blocks above: +max |peak-local − band-limited| = **1.9e-11 nats**, median 6.8e-13. + +The local evaluator reconstructs the same interpolant the dense zero-padded FFT does, +to 2e-14 … 4e-13 relative, at every production `npts` — 153, 307, 613, 614, 1228, 2457 +— and at 8, 9, 3. Odd `npts` is the common case (three of the five production sample +rates), and the failure it invites is exact AT the samples and wrong between them, so +it is parametrised rather than spot-checked. + +### Memory chunking moves the last bits, and the docstring that said otherwise was wrong + +The extrinsic axis is chunked so one dense temporary stays inside a working-set budget. +The per-row plan — which rule the row gets, how many intervals, and the point count, +which is bucketed to a power of two — depends only on that row, so chunking cannot +change any of it. It DOES change the leading dimension of the FFT and of the reduction +inside `eval_bandlimited_uniform`, and both numpy's FFT and its pairwise summation +reassociate with batch shape. + +Measured, one row per chunk versus all rows at once, on a three-row block spanning +rho ~ 100 to 700: **0, 0 and 2 ULPs** (2.4e-16 relative). The first version of this +module carried the inherited claim that chunking "cannot change the answer"; the test +that asserted bit-identity failed, which is how the overclaim was found. The test now +pins a few ULPs — still sharp enough that a dropped or mis-ordered chunk, which moves a +row by nats, cannot hide behind it. + +## Why the truncation is rigorous and not hopeful + +RIFT PR #201 was caught seeding a Newton solve at guessed points, missing genuine +maxima and returning `-inf` for a finite integral. There is no seeded root-finder +here. Two separate mechanisms: + +* **Enumeration.** Every local maximum of the band-limited interpolant is found on a + grid that resolves `kappa` itself. Checked against a factor-64 enumeration: every + peak carrying representable mass at factor 64 is also found at the shipped factor 8, + to within one coarse sample. +* **A computed bound.** `log(T_outside) + max_{outside} lnL` upper-bounds the omitted + integral and is compared per row against the value computed; a row that cannot meet + `TAIL_LOG_TOL` goes to the dense path. **This does not depend on the enumeration + being complete** — a missed peak lands outside the intervals, so the sampled maximum + outside sits on it and the bound fails. Enumeration buys speed; the bound buys + correctness. + +That is asserted, not asserted-about: `test_a_sabotaged_enumeration_is_caught_by_the_tail_bound` +monkeypatches the enumeration down to a single maximum on a two-peak integrand — which +discards half the mass, `log 2 = 0.69` nats — and requires the module to detect the +shortfall and return the dense value. + +Evaluating the outside maximum is free because every shipped `loglikelihood` callback +is monotone increasing in `Re kappa` (plain, distance-marginalized) or `|kappa|` +(phase-marginalized), and `rho_sq` is time-independent on this path, so +`argmax_t lnL = argmax_t term(kappa)` whatever the distance, the distance prior or the +callback. Verified over 0.05 … 40 in `1/D` and across three callback shapes: identical +peak sets. This is what keeps the callback — a table interpolation in production — +off the full time axis entirely. + +## Merging is correctness, not tidiness + +Two overlapping windows integrated separately both contain the shared region, so the +log-sum-exp of the parts counts it twice. Prototype (`~/tmarg_harness/peaklocal.py`): +**+1.6 nats at rho ~ 6**. `test_unmerged_intervals_double_count` rebuilds the +un-merged variant against this module's own evaluator and requires it to be wrong, +because on a sharply peaked row the intervals do not overlap at all and every other +accuracy test stays green when the merge is deleted. + +Merging is also what makes this ONE algorithm rather than a regime switch: isolated +peaks give a tiny union, crowded peaks grow the union to the whole window and the +method degenerates continuously into the dense grid. No threshold anywhere. + +## Round 4: the class was still open, at two more doors + +### Door 1 — the parabolic vertex obeys the SAME `(delta/sigma)^2` law + +`_parabolic_vertex_height` was introduced to fix G1 and its docstring claimed the vertex +is the crest "EXACTLY for a Gaussian peak, at any spacing". That is true only if `lnL` +is quadratic across the whole `+/-h_enum` stencil. It is band-limited `q(t)`; the +quartic term survives. Under-read of the crest at half-cell phase, in nats: + +| `sigma_t/deltaT` | factor | sample | **vertex** | +|---|---|---|---| +| 0.004934 | 512 | −108.75 | −1.17 | +| 0.002467 | 1024 | −434.98 | −4.66 | +| 0.001233 | 2048 | −1739.93 | −18.66 | +| 0.000617 | 4096 | −6959.72 | **−74.63** | + +`UPSAMPLE_FACTOR_MAX = 4096` permits down to `sigma_t/deltaT = 0.000488`, so the last row +is a LEGAL configuration, and end-to-end against the exact interpolant at 16384x it +reproduced **−0.693147** — `−log 2`, one of two equal peaks deleted — with both defences +silent. Roughly an 8x improvement in reach over comparing samples, not a fix. + +**The lesson, and the reason this was a third round on one class: every approximation +substituted for the crest fails the same way one octave further out.** So the keep +decision is now taken in two stages, and neither treats an estimate as the answer: + +1. a CONSERVATIVE PRE-FILTER, whose only job is to bound how many peaks reach + localisation. It compares an UPPER bound on each crest — the sample plus the + worst-case correction `(h_enum/2)^2/(2 sigma^2)` — against a LOWER bound on the + highest crest — the largest sample, which cannot exceed its own crest. It can only + ever keep too many; +2. after localisation, the EXACT filter on `lnL_star`, which is the crest by + construction rather than to second order. + +Verified on the reviewer's own probe at every legal factor: **+0.000000 with +`n_peaks_total = 2` at factor 1024, 2048 and 4096, at offsets 0, 0.25 and 0.5.** + +### Door 2 — both estimators degraded to the raw sample at the array ends + +At `cols_p == 0` the whole left half of the curvature stencil is out of range at EVERY +half-width, so `d2` was NaN throughout and `sigma = inf`; the vertex height fell back to +`y0`, i.e. pre-G1 behaviour. Revision 2's "mask instead of clip" change looked more +principled than the `maxd`-clipping it replaced and was worse exactly at the two indices +its own justification named. + +Fixed by shifting the stencil centre inward by the MINIMUM needed for a three-point +stencil to exist — one sample — which is a genuine one-sided fit at an endpoint, not a +compromise. The localiser's bracket is also clamped to `[0, t_last]` and a peak pinned +at a window boundary counts as converged. + +**Partially closed, and measured as such.** Against the reviewer's `f2_edge_sigma`, +comparing `761cafb3` with the fix, same fixtures: + +| case | before | after | +|---|---|---| +| `sigma@edge = inf`, dH 2.0 / 1.0 / 0.5 / 0.2 | −0.126 / −0.312 / −0.473 / −0.598 | **+0.000 / −0.000 / +0.000 / +0.000** | +| `sigma@edge` finite, `sig2/sig1 = 0.58` | −6.001 | −6.001 (unchanged) | +| `sigma@edge` finite, `sig2/sig1 = 0.41` | −0.335 | −0.335 (unchanged) | + +So the family Door 2 describes — no finite width obtainable at an endpoint — is closed +with no regression anywhere. **A residual family remains and is NOT that mechanism**: +those rows already had a finite edge sigma at `761cafb3`, up to **−6.0 nats**, accepted. + +**Re-measured after Round 5's rebase, this family changed character but did NOT close.** +Against a converged reflected reference (32768x; successive refinements differ by +−1.8e-4), on the sharpest case (`sig2/sig1 = 0.58`, secondary at `0.4 h_enum` from t=0): + +| | value | vs reference | +|---|---|---| +| converged reference | 39844.73839 | — | +| `bandlimited`, derived factor 256 | 39845.34836 | **+0.610** | +| `peak-local` | 39845.87353 | **+1.135** | + +So the disagreement with the rule it delegates to is **+0.53 nats, ACCEPTED** — +`tail_bound_worst = −250.71` and the containment check passes, so neither defence fires. +It was −6.0 nats and is now +0.53, a 10x improvement and a sign flip, but it is still an +accepted error three orders above the 1e-3 nat bar the rest of this rule meets. + +**Located, not yet fixed.** The even reflection is value-continuous but its DERIVATIVE +flips sign at the join, so a peak sitting at t=0 is a two-half-peak cusp rather than a +locally Gaussian peak. Both rules derive a width — and hence a factor and an interval +half-width — from a Gaussian curvature model, so both under-resolve it; `bandlimited` is +wrong here too (+0.610 at its own derived factor), which places the root cause in the +shared reconstruction rather than in the local placement. peak-local roughly doubles the +error because it also sizes its interval as `W_SIGMA * sigma` from that same model. +Diagnostics: the row enumerates 295 maxima, the keep filter retains 1, and that one is at +enumeration index 0 with `q` decreasing monotonically away from it. + +**The quantisation class should not be called closed on my say-so, and this is not the +quantisation class**: the enumeration-index defect is closed and re-verified above, while +this is a curvature-model defect at the reflection join, inherited from the dense path. +It is the top open item. + +### The `W_SIGMA` coupling is now asserted + +The sampled `q_out_max` survived a determined attempt to break it (24 accepted rows, +honest supremum on a 4096x grid, worst honest margin −63.42 against `TAIL_LOG_TOL = +-23`). But the reason is structural slack, not adequate sampling: the outside supremum +sits at an interval edge, already `(W_SIGMA + LOCALISE_SAFETY)**2/2 = 75.03` nats below the crest. Dropping +`W_SIGMA` below ~8–9 would silently invalidate the bound. The inequality + + W_SIGMA**2 / 2 > |TAIL_LOG_TOL| + log(T_out / (sqrt(2 pi) sigma_min)) + +is now asserted, tying `W_SIGMA` to `TAIL_LOG_TOL` and `UPSAMPLE_FACTOR_MAX` so none can +move alone. + +## Round 5: the rebase onto `rift_O4d` changed the reconstruction underneath this module + +Retargeting this PR from the merged `rift_O4d_tmarg_bandlimited` onto `rift_O4d` moves it +across `e4ed25c7` ("Avoid Gibbs ringing in time marginalization"), which + +* replaced the raw periodic zero-padded FFT with an EVEN REFLECTION — periodize + `[kappa forward, kappa backward]`, keep the forward interval — because a zero-padded FFT + of the gathered slice alone identifies its unlike endpoints and rings globally + (+140.9 nats measured on an adversarial row); and +* demoted `EDGE_GUARD_FRACTION` to a diagnostic, on the grounds that crossing an arbitrary + threshold must not silently move an under-resolved row back to Simpson. + +**The diff did not change; its meaning did.** This module was written against the older +contract, and nothing in the seven-file delta says so. Three clauses had drifted, and each +one made peak-local return a WORSE value than `time_marginalize_bandlimited` for the same +row — the one thing this rule promises never to do. + +| drift | measured | +|---|---| +| enumeration, localisation and local evaluation still on the PERIODIC interpolant, while fallback rows got the reflected one — two different continuous functions inside one call | **−3.79 nats** on a row with peaks at both window ends; **9.0e-6 nat** median bias on the uniform-arrival block | +| `EDGE_GUARD_FRACTION` still routing a near-edge row to SIMPSON | row never entered the rule; every fallback counter read zero | +| `boundary_unresolved` missing: an endpoint maximum whose inward-clipped stencil reads positive curvature is mislabelled "flat" and silently keeps Simpson | **+4.60 nats** against the reflected reference, where the dense path was +0.81 | + +Fixed by taking the spectrum of `concatenate((kappa, flip(kappa)))` at period +`2*npts*deltaT`. The local evaluator then reproduces `reflected_bandlimited_upsample` to +**2.3e-13 relative** at every production `npts` (153 / 307 / 614 / 1228 / 2457, odd and +even). + +### The lesson is structural, and the fix is too + +Fixing the copy leaves the copy. Row classification is now `_classify_rows` in +`time_marginalization_quadrature.py`, the SINGLE definition, called by both rules; +peak-local keeps only `refined = has_peak & (factors > 1)`. Verified behaviour-preserving +rather than asserted: `time_marginalize_bandlimited` is **bit-identical** over a 32-row +battery spanning flat, edge, near-edge, broad, sharp and both-ends-peaked rows, with every +`last_report()` counter identical. + +`test_row_classification_matches_the_dense_path_exactly` compares the CLASSIFICATION, not +the values, because a value check passes whenever the two rules happen to agree — which on +most rows they are designed to, which is exactly why this drifted unnoticed three times. + +### Door 1 re-verified at the sharpest legal row, against the right reference + +The round-3 reopen (`g1e_reopen2`) reproduced at derived factor 4096, off = half a cell: +**−0.693147** = −log 2, one of two equal crests deleted. Re-run against a REFLECTED 16384x +reference — the object the module now integrates — all nine rows are exact and both peaks +survive: + +| sigma_t/deltaT | factor | offset | peak_local − truth | n_peaks | +|---|---|---|---|---| +| 0.000615 | 4096 | 0.00 | +0.000000 | 2 | +| 0.000615 | 4096 | 0.25 | +0.000000 | 2 | +| 0.000615 | 4096 | 0.50 | **+0.000000** | 2 | + +### Door 2's mechanism is closed, and measured at the source + +The estimator is now genuinely one-sided at the array ends rather than degrading to the raw +sample. On a parabola of known width `3 h_enum` sampled on the enumeration grid, the +recovered sigma at peak index 0, 1, 7, 8, 100, n−9, n−2 and n−1 is **finite and exact +(3.0000 h) at every one**, including both endpoints. The dead-code condition — 22 endpoint +maxima enumerated, zero able to obtain a finite width — is gone. + +End to end, a row whose dominant maximum sits at the very edge is either handled correctly +or delegated, never silently approximated: at tau = 0.5 and 612.5 samples peak-local is +exact; at 0.2 and 612.8 it is **better** than the dense path (+1.44 against +1.71); at 0.0 +and 613.0 it declines and returns the dense value **bit-identically**. + +### The contract, checked reference-free + +`|peak_local − bandlimited|` over 46 rows spanning sharp phase-scan, broad, near-edge, +two-peak, both-ends and flat/null families: **worst 2.6e-10 nats**, no family above 1e-3. +All ten near-edge rows are now handled by the rule; before this round they were routed to +Simpson. + +## Round 6 — DOOR 4: the pre-filter's "upper bound" was not one + +The fourth independent re-attack reopened the class at the one site the previous three had +left alone, and it is the site whose own comment said it could not fail: the +**conservative pre-filter**, described as comparing "an UPPER bound on each crest against a +LOWER bound on the highest crest, so it can only ever keep too many." + + crest_upper = lnL_sample + (h_enum/2)**2 / (2 sigma**2) + +is not an upper bound, for two independent reasons, and the second is the one that matters. + +**The Laplace model is not what failed.** `lnL` is of course not a parabola. The whole design +uses the Gaussian/curvature picture to **target** — where the mass is, how wide an interval +must be, how fine a spacing it needs — and then either the Laplace quadrature is reliable +enough or the row falls back to the grid. Used that way the model owes nothing: the answer +comes from the quadrature actually performed and from the checks that verify it. What failed +was using the same model where an **inequality** was required. A targeting estimate that is +off by 122 nats costs a slightly wrong interval, which the verification catches; the same +estimate asked to justify DELETING a peak costs the peak. + +* the localiser's bracket is `+/- h_enum` and displacements of `0.959*h_enum` have been + observed, so the correction covers less than half the distance it must; and +* **the ANHARMONIC part of the crest deficit carries the same `1/sigma**2` amplification** + as the quadratic part, so the error does not shrink where it matters. At + derived factor 1024 the pure quantisation excess is **4.4 nats** and the true shortfall + is **122.30**. + +Being short, it DELETED peaks — before localisation, so the exact filter never saw them. + +| derived factor | shortfall of `crest_upper` | `pl - bl` | accepted? | +|---|---|---|---| +| 1024 | +122.30 nats | **−0.358385** | yes, both defences silent | +| 2048 | +489.19 | **−0.358427** | yes | +| 4096 | +1956.74 | **−0.358383** | yes | + +and the magnitude is not bounded by `log 2`. Raising the deleted peak above the survivor +by `Delta`: + +| Delta | `pl - ref` | accepted | +|---|---|---| +| 0 | −0.358 | yes | +| 800 | −799.159 | yes | +| **1850** | **−1849.159** | **yes** | +| 1950 | +0.000 | yes (the peak survives) | + +The cutoff is `shortfall - PEAK_KEEP_NATS`, so **a peak may sit ~1900 nats ABOVE the one +that survives and still be deleted.** The tail bound cannot backstop it: `q_out_max` reads +the deleted peak at its SAMPLE, the very quantity the defect corrupts. Recomputed as an +honest supremum on a 4096x refinement, those rows score **+11.5 / +12.1 / +7.2** against +`TAIL_LOG_TOL = -23` — every one would be REJECTED. + +### The fix is an inequality, not a better fit + +Widening the constant would have been the fifth version of the same mistake. Expanding +about the CREST, where `q'` vanishes by definition, + + q(t_s) = q(t*) + q''(xi) (t_s - t*)^2 / 2, |t_s - t*| <= h_enum + +so `q(t*) <= q(t_s) + max|q''| * h_enum**2 / 2`, and `max|q''| <= sum_j |Xw_j| |w_j|**2` +by the triangle inequality on the spectral sum. **Nothing is fitted**, so there is no model +error left to amplify, and it uses `h_enum` rather than half of it. `loglikelihood` is +monotone in `q`, so a bound on `q` is a bound on `lnL`. + +Verified as a PROPERTY rather than on the fixture that motivated it: over **9311 enumerated +peaks** across npts 153/307/614/1228/2457 and amplitudes 2e4–2.5e7, **zero violations**, +where the old bound failed on **0.3–3 %** of peaks by up to **19456 nats**. + +`spectral_curvature_bound` and `crest_upper_bound` are module functions precisely so the +suite tests the SHIPPED formula and not a copy of it. + +### Deleting the pre-filter outright is not the answer, and was tried + +Without it every row enumerates its whole oscillation — 295 maxima on one fixture — the +structure gate sees more than `MAX_INTERVALS` and declines every row, and the option goes +**inert**: `n_peak_local_rows = 0` on all six fixture families, which is the W1 hazard the +suite already guards against. With the rigorous bound, coverage is exactly what it was +(24/24 sharp, 10/10 near-edge, 3/3 two-peak) and `|peak_local - bandlimited|` is +**2.6e-10 nats** over 46 rows. + +### The same defect at a third site, fixed with it + +The pre-filter's LOWER bound read `lnL_st[:, maxd]` — the callback at the stencil CENTRE, +which is clipped inward by one at the array ends. At enumeration index 0 or `n-1` that is a +full cell from the peak it describes: **132 nats** low at rho ~ 40, **8449** at rho ~ 700, +growing with SNR. A stencil centre is clipped so the CURVATURE can be measured; the peak's +own value must be read where the peak is. + +### What makes the regression test worth anything + +`test_a_codominant_crest_is_not_deleted_by_the_prefilter` **FAILS on the parent commit** and +passes here. Its fixture solves for the second amplitude by bisection rather than hard-coding +a ratio, because the crest gap is linear in amplitude — a ratio tuned at rho ~ 40 leaves the +second peak thousands of nats down at rho ~ 700 and the fixture silently stops testing +anything, which is exactly how round 2's suite missed round 3's defect. + +The helper that finds the two crests LOCALISES them. The first version of it did not, and +was wrong in the same way as everything else in this file: peak A sits on an enumeration +sample so its sample is its crest, while peak B is deliberately off-grid and its sample +understates its crest by ~3125 nats in the sharp regime. Equalising the SAMPLES leaves B +thousands of nats above A, A is correctly dropped, and the test passes while measuring +nothing. + +### Still open after round 6 + +* **The ceiling is a limit on the DENSE GRID, and this rule is not bound by it.** RO's + ruling, recorded rather than re-litigated. `over_ceiling` here catches the saturation + sentinel; it deliberately does not mirror the dense path's refined re-measurement. See the + round-8 entry for the measured size of the difference and why forcing agreement costs more + than it buys. +* The tail bound is still a SAMPLED maximum, and its safety still comes from the + `(W_SIGMA + LOCALISE_SAFETY)**2/2 = 75.03` nat structural slack rather than from the sampling being adequate. + +## Round 7 — DOOR 5: a crest pinned at a window end is not a peak + +The round-6 bound held. An independent re-attack could not violate it — **0 violations**, +measured slack **15x to 84.5x** — so the class did not reopen where it had been patched. It +reopened one site further along, at the **resolution scale** rather than the crest location. + +`localise_peaks` counted a peak pinned against a window end as CONVERGED, reasoning that +the maximum over the integration domain really is the boundary. It is. **But it is not a +STATIONARY point, and everything downstream assumes one.** At an interior crest `q'` +vanishes, so `exp(lnL)` is locally Gaussian and a spacing derived from the curvature +resolves it. At a boundary the maximum is a CORNER: `q'` is non-zero there, so the local +integrand is an **exponential decay of rate `|q'|`**, and a spacing derived from `sigma` +does not resolve that scale at all. The trapezoid's half-weight endpoint then over-counts by + + log( lam * (0.5 + 1/(exp(lam) - 1)) ) -> log(lam/2), lam = |q'| * h_loc + +MEASURED on a single band-limited bump centred at `t = 0`, against the same interpolant +integrated on **the very interval the module chose** — so this is the rule's own quadrature +error on its own domain, not a reference artifact: + +| amplitude | coarse factor | `lam` | over-count | accepted | +|---|---|---|---|---| +| 4e4 | 512 | 7.12 | **+1.27410** | yes | +| 1.6e5 | 1024 | 14.24 | **+1.96381** | yes | +| 6.4e5 | 2048 | 28.49 | **+2.65648** | yes | +| 2.56e6 | 4096 | 56.98 | **+3.34949** | yes | +| 1e7 | 4096 | 112.6 | **+4.03071** | yes | +| 2e7 | 4096 (ceiling) | — | **+4.37722** | yes | + +`log(112.6/2) = 4.031` against a measured `+4.03071` — the closed form is exact. **The error +grows by +log 2 per factor 4 in amplitude and is unbounded inside the legal range.** + +### Neither defence could ever have caught it + +This is structural, not bad luck, and it is the reason the fix is a refusal rather than +another check: + +* **Containment** compares the local grid's attained maximum against `row_star`. The grid's + FIRST POINT is the pinned crest, so `attained == row_star` identically. It cannot fire on + this family, at any parameters. +* **The tail bound** is a statement about mass OUTSIDE the intervals. This error is entirely + inside. `tail_bound_worst` reads −250 to −32990 on exactly these rows. +* `localise_peaks` declared the peak converged, so `n_dense_fallback_localise` was 0. +* `_classify_rows` flags these rows `exposed` and, by design since `e4ed25c7`, `exposed` + selects nothing. + +### The fix, and what it does not fix + +A pinned peak is reported UNCONVERGED and its row goes to the dense path. Fail-closed, and +it restores the contract this rule actually makes — **never a worse value than the +backstop**. Coverage is unchanged (24/24 sharp, 10/10 near-edge, 3/3 two-peak; worst +`|peak_local - bandlimited|` **2.6e-10** over 46 rows), because only rows whose crest is +literally at the boundary now decline. + +It also closed a **regression round 6 had introduced and the author had not found**: the +looser pre-filter kept a second peak on a near-edge row that round 5 correctly REJECTED, so +the row became accepted carrying this defect — `pl - bl` went `0.000000` (rejected) to +`-0.778178` (accepted). It is back to `+0.000000`, declined. The edge/cusp family that this +note carried at −6.0 nats, then +0.53 and growing with SNR, is now `pl - bl = 0.000000` +throughout: every row either matches the dense path exactly or declines. + +**DISCLOSED, NOT FIXED.** `time_marginalize_bandlimited` carries the same defect in milder +form — its refined grid also begins at the boundary with a half weight. Measured against a +converged spectral reference: **+0.768 / +1.429 / +2.121 / +2.813 / +3.479 nats** at +amplitude 4e4 / 1.6e5 / 6.4e5 / 2.56e6 / 1e7, i.e. derived factor 256 → 4096. It obeys the +same `+log 2 per factor 4` law, and finer references move the top row DOWN, so **+3.48 nats +at the legal ceiling is a lower bound.** (An earlier measurement here stopped one octave +short and quoted a maximum of +2.74; that was wrong and is corrected.) + +peak-local now matches it **exactly** — `pl - bl = +0.000000` at every amplitude with +`n_peak_local_rows = 0` — so it inherits the defect rather than adding to it, which is what +makes the residual cleanly attributable: it belongs to the shared reconstruction and to +#203's rule, not to this one. **It should be raised upstream, not patched here.** + +### A correction to a stated precondition + +The module claimed every shipped `loglikelihood` callback is monotone increasing in +`Re kappa`. **That is false**: the distance-marginalized callback returns `-inf` ABOVE its +table as well as below. The direction is safe — an upper bound landing in the hole evaluates +to `-inf`, the peak is dropped, and the row loses peaks until it falls back — so it costs +coverage and never accuracy, and with the shipped table the boundary sits at +`D_eff < d_min/10`, unreachable at default settings. But the claim as written was wrong, and +the pre-filter's `q`-bound-to-`lnL`-bound step rests on it, so it is now stated precisely. + +## Round 8 — the first independent YES, and what is left + +A sixth reviewer, handed the round-6 and round-7 fixes and told to break them, returned +**YES: the quantisation/resolution class is closed.** That is the first affirmative verdict +in six adversarial passes, and what it rests on is worth recording rather than just the word. + +### Why `~pinned` is the right predicate, and not a patch over one fixture + +The obvious next door is a crest that is NOT pinned but sits close enough to a boundary that +its interval is CLIPPED, leaving a non-stationary edge. That band was scanned directly — +auto-bracketing the pinned-to-interior transition, then bisecting the bump offset to land +`t*/sigma` on 13 targets in (0.25, 11.9), with the module's own local grid captured and +re-integrated at 512x/1024x on the same interval (convergence ≤ 9.3e-9): + +| derived factor | worst `pl − exact(own interval)` | worst `pl − bl` | worst `lam` | +|---|---|---|---| +| 256 | −0.00238 | −0.00193 | 0.172 | +| 1024 | −0.00272 | −0.00224 | 0.80 | +| 4096 | −0.00268 | −0.00222 | 1.24 | + +**Flat in amplitude** — 0.0024 to 0.0027 across 250x in amplitude and 16x in resolution — +and negative. The opposite signature to the pinned defect, which grew +1.27 → +4.03 nats at ++log 2 per factor 4. + +The mechanism is why, and it is the argument that actually closes the class. At an INTERIOR +crest the edge deficit and the decay rate are TIED: for a peak of width `sigma` integrated at +`h <= sigma/2`, an edge at distance `d` has `lam = d*h/sigma^2 <= d/(2 sigma)` and depth +`Delta = d^2/(2 sigma^2) ~ 2 lam^2`. So `lam >> 1` forces `Delta >> 1` and the over-counted +mass is exponentially suppressed, while `Delta ~ 0` forces `lam ~ 0` and the half weight is +correct. **The pinned case escaped only because a boundary CORNER has `q' != 0` at +`Delta = 0`** — there is no such relation there. Round 7 removes exactly the configurations +that break the tie. + +Confirmed at the extreme: an accepted row whose clipped edge is **0.005 nats** below its +interval's own maximum — a 72-nat violation of what `W_SIGMA` assumes — still has +`lam_lo = 0.031` and an error of −6.1e-4 nats. The other end of that same interval has +`lam_hi = 5.67` at depth 72.85. The two conditions never co-occur. + +**Merged edges** cannot be the door either, and structurally so: the union's endpoints are +`min_i(t_i − W sigma_i)` and `max_i(t_i + W sigma_i)`, so merging can only push an edge +FURTHER from every crest. Over 250 random rows (1–3 bumps, log-uniform amplitude +10^3.5–10^7.5, 70% of positions biased hard against the window ends) the worst +`pl − bandlimited` was **+0.000000 nats**. + +### What is still open, in severity order + +* **The ceiling: a DECISION, not a defect.** `over_ceiling` is taken on the COARSE derived + factor, while `time_marginalize_bandlimited` re-measures on the refined grid and doubles + until the criterion holds, raising if the chain leaves the legal range. So the dense path + can refuse a row whose coarse factor was legal, and this rule returns it. Measured over 220 + random multi-bump rows: the disagreement is confined to `factor_coarse == + UPSAMPLE_FACTOR_MAX` exactly — 27 rows landed there, 8 of them refused by the dense path and + returned here. Below 4096 it never arose (it would need the coarse estimate to be optimistic + by three doublings); at the 8192 sentinel both refuse, 44 of 44. + + Those 8 rows are not approximated — scored exact against a converged spectral reference. + **`UPSAMPLE_FACTOR_MAX` bounds the dense grid**, the point past which a zero-padded FFT over + the whole window stops being affordable or believable, and this rule never builds that grid: + its resolution comes from a per-peak width measured on the enumeration grid, already 8x + finer than the coarse grid the ceiling factor is derived from, with the containment check + and tail bound verifying the outcome. So a row the dense path declines on grid size is one + this rule may legitimately still resolve. **RO's ruling: it may.** Forcing agreement would + mean `>=`, which sends all 27 rows at the legal ceiling to the dense path — including the 19 + it handles perfectly well — and the sharpest legal rows are precisely the regime this rule + exists to serve. Recorded as a documented difference between the two paths. +* **The dense path's own boundary defect** — merged code, up to **+3.48 nats** at the legal + ceiling, inherited rather than caused, and peak-local now matches it exactly. +* **A cost regression from round 7.** A row carrying an interior dominant crest AND a + boundary crest within `PEAK_KEEP_NATS` is now declined whole: `exact_keep` keeps the pinned + peak and `bad_loc` then condemns the row. Fail-closed and correct, but a real coverage loss + on edge-peaked rows. Measured: a uniform-arrival block is 60/60 accepted and a + near-far-end block 6/7, so it bites on edge-peaked rows specifically, not broadly. +* **A latent accounting fragility.** `n_dense_fallback_localise` is incremented + unconditionally while `keep_row` later ANDs `~bad_loc` with `~too_much`/`~too_slow`, which + increment their own counters, so a row that is both would be double counted and the + sub-counts would exceed `n_dense_fallback_rows`. Three early returns also skip the + `n_dense_fallback_nopeak` accounting. **Not reachable on any fixture tried** — a 47-row + mixed block reconciles exactly — because the `bad_loc & too_much` path needs peaks + separated by between `24.5 sigma` and `24.5 sigma + 2 h_enum`, a window the band limit + closes at high SNR. Recorded as fragility, not as a demonstrated bug. Relatedly `too_slow` + looks like dead code: the provisional point count is an over-estimate, so `p_slow` fires + first. + +## Real data, through the shipped driver + +Everything above is synthetic band-limited rows. This is not. + +**Case:** the paper-1 3G study, `RIFT_roboto_paper/analyses/slowrot_finite-size/3g`, +`run_ET_snr100` — a zero-spin BNS (1.6 + 1.4 Msun) injected into the ET triangle with the +finite-size response, network SNR 100, analysed with the **real ILE** +(`integrate_likelihood_extrinsic_batchmode`) in its `static` configuration: `--srate 2048 +--fmin-template 50 --fmax 1024 --approximant IMRPhenomD --l-max 2 --vectorized --force-xpy +--gpu --internal-use-lnL --time-marginalization --sampler-method AV --n-eff 300 --n-max +800000`, intrinsic grid point 4, `--seed 1234`. Only `--time-marginalization-quadrature` +varied between runs. + +| quadrature | lnL | sigma/L | n_eff | samples | wall | +|---|---|---|---|---|---| +| `simpson` (default) | −296.1915649136 | 0.747 | 1.36 | 802476 | **64.6 s** | +| `bandlimited` | −294.8959884975 | 0.413 | 3.06 | 809950 | **749.0 s** | +| `peak-local` | −294.8959884975 | 0.413 | 3.06 | 809950 | **915.9 s** | + +**`peak-local` and `bandlimited` agree to 0.000e+00 nats — the two `.dat` files are +BYTE-IDENTICAL (same md5), and `n_eff` matches to all 16 digits.** Re-run after the outside +supremum was certified on sub-cell geometry: **the same md5 again**, so certifying the bound +left the real-data answer bit-for-bit unchanged. (Wall times moved between the two runs -- +`bandlimited`, which was not touched, went 749 s to 596 s -- so the run-to-run ratio is machine +load and must not be read as a cost of the certificate.) That is the rule's +central contract, verified end to end on real data through the shipped driver rather than on +a fixture, and it is a stronger statement than the synthetic `2.6e-10` because the two runs +followed the same adaptive sample path and still landed on the same bits. + +**Cost, on this real case: 11.6x `simpson` for `bandlimited` and 14.2x for `peak-local`**, on +comparable sample counts. So `peak-local` is **1.22x SLOWER than the dense rule here**, which +is the disclosed behaviour — SNR 100 is below the ~180 where the local placement starts to +pay, and this measurement is a confirmation of that claim on real data, not a refutation of +it. It is also the first end-to-end cost number for this rule that is not a microbenchmark. + +**The `simpson` difference is NOT a clean bias measurement, and must not be quoted as one.** +It reads −1.296 nats against the other two, but its adaptive sample path diverged (`n_eff` +1.36 vs 3.06, different draw count), so the number confounds the quadrature with the +sampling. What it does show is that the choice is not inert at this amplitude. A clean bias +number needs the same samples under all three rules, which this run does not provide. + +**Caveat on the absolute values:** `n_eff` is 1.4-3.1, so the lnL values themselves are not +converged — this grid point is in the high-SNR `n_eff` lottery the case's own README +describes. That does not weaken the `peak-local == bandlimited` result, which is exact and +independent of convergence, but it does mean no accuracy claim should be read off the lnL +column here. + +### Two things the real run surfaced that the fixtures could not + +* **`--force-xpy` alone does not satisfy the guard, though its label says it does.** The + prerequisite is printed as `--gpu (accepts --force-xpy)`, but the predicate is + `bool(opts.gpu)`, and `--force-xpy` only re-enables `opts.gpu` when `--gpu` was ALSO passed + and no device was found. The paper case ships `--vectorized --force-xpy` with no `--gpu`, so + a CPU run of it is refused with a message that reads as though the flag it already has + should have been enough. The refusal is CORRECT -- without `--gpu` the run takes + `DiscreteFactoredLogLikelihoodViaArrayVector`, not the maintained NoLoop path -- but the + label is misleading in exactly the configuration a CPU user hits. +* **`last_report()` is never called by the shipped driver.** The counters exist and the suite + asserts them, but nothing surfaces them in a production run, so an operator gets no + confirmation of how many rows `peak-local` actually handled, how many fell back, or what the + tail-bound margin was. Several comments in this module appeal to "an operator reading the + columns"; there are no columns to read outside the tests. + +## P1 from an outside reviewer: the outside supremum is now EVALUATED, not sampled + +**The complaint.** `q_out_max` is the maximum on the factor-8 enumeration samples, so it is a +LOWER bound on the continuous outside supremum. As amplitude grows a peak left outside can read +arbitrarily lower at its sample than it is, so `margin` passes and the returned integral is +missing mass. The containment check cannot catch it — it examines only retained intervals. + +**Measured, and the premise was right.** Honest supremum recomputed on a +64x-per-enumeration-cell grid over the module's own uncovered set, on rows it ACCEPTED: + +| amplitude | module reported | honest | error | +|---|---|---|---| +| 2e4 | −79.57 | −65.53 | **14.0 optimistic** | +| 2e5 | −289.51 / −220.32 / −141.08 | −71.14 / −71.83 / −66.55 | **218 / 148 / 75** | +| 2e6 | −2998.61 / −2307.03 / −1514.75 | −75.82 / −90.20 / −70.17 | **2923 / 2217 / 1445** | + +The error grows without bound. `tail_bound_worst` got *more* flattering the sharper the row, +which is the wrong direction for a safety number. + +### Where the gap came from — not a peak the enumeration missed + +The supremum over a union of closed intervals is attained at an interior stationary point or at +an END. The candidates are therefore exactly the crests of uncovered enumerated maxima and the +ends of the covered intervals, **and the ends dominate**. An interval end sits `W_SIGMA * sigma` +from its crest, i.e. `(W_SIGMA + LOCALISE_SAFETY)**2/2 = 75.03` nats below it at any amplitude — but the nearest SAMPLE +outside that end is a further `W_SIGMA * h_enum / sigma` down, and that term diverges as the +peak sharpens. It is the whole measured error. + +### The fix: evaluate the candidates on the double-copy Fourier model + +`q` is band-limited and the even reflection reconstructs it exactly between samples, so the +interval ends are evaluated directly on that model (`eval_bandlimited_at`), and the localised +crests are already in hand from Newton — **including the peaks the exact filter dropped**, which +is precisely the set the sampled version read at their samples. The membership test that decides +which crests are outside is vectorised by the row-offset trick `merge_intervals_by_row` uses; a +loop over intervals there is `O(groups x peaks)` and is the shape that once made this rule slower +than the path it delegates to. The exclusion is not optional — feeding a COVERED crest to the +outside maximum makes the bound `log(T_out) + crest - result`, which rejects every row. + +After, on the same rows: + +| amplitude | module reported | honest | error | +|---|---|---|---| +| 2e4 | −65.13 / −65.01 / −64.89 | −65.14 / −65.42 / −65.53 | **−0.01 / −0.41 / −0.64** | +| 2e5 | −64.13 / −64.00 / −63.88 | −71.14 / −71.83 / −66.55 | −7.0 / −7.8 / −2.7 | +| 2e6 | −63.00 / −62.86 / −62.75 | −75.82 / −90.20 / −70.17 | −12.8 / −27.3 / −7.4 | + +The amplitude-divergent term is gone and the sign has flipped: the reported margin is now +CONSERVATIVE relative to the honest one (the module evaluates the end exactly, where the +reference grid still under-resolves it). Accuracy, coverage and the `peak-local == bandlimited` +contract are unchanged — worst `|pl - bl|` over the six fixture families is still 2.6e-10, with +24/24 sharp, 10/10 near-edge and 3/3 two-peak rows kept. + +### What this does and does not claim + +It does NOT certify. A peak the enumeration never found would still be missed, and the reviewer's +proposed certification — adding `0.5 * q_ddot_max * h_enum**2` — remains non-viable here: that +bound's slack is 15x to 84.5x (494 nats at amplitude 2e4, 5.7e5 at 2.5e7), harmless against a +60-nat KEEP window and fatal against a 23-nat REJECT window, where it makes every high-amplitude +row fall back. Per RO: off-grid peaks are not a realistic failure mode; between-sample accuracy +is the real issue, and that is what an exact interpolant supplies. Measured in support: across +npts 153/307/614 and spectral widths 30/120/300, **zero material continuous maxima** (within 100 +nats of the row top) were missed by enumeration, which in fact returns one or two MORE than the +interior continuous count. + +Coverage-boundary cells were already included in the sampled version — the mask uses `ceil`/ +`floor`, so only samples strictly inside an interval are covered. That was never the gap; the gap +was that the nearest sample is not the end. + +### A defence that stopped being tested, and was restored + +The off-grid bound now catches the mis-placed-interval fixture BEFORE the containment check does, +which is an improvement — but it silently cost the containment check its coverage. Disabling +containment outright (`contained = all True`) left the whole suite green. The two are NOT +redundant: the tail bound excludes covered points by construction, so it cannot see a crest that +is INSIDE its interval and that the local grid failed to attain, which is the only thing +containment can catch. `test_containment_still_catches_a_mis_placed_interval_with_the_tail_bound +_disabled` isolates it by making the tail bound unconditionally pass, and it kills that mutation +again. + +## What the tail bound is actually measuring, and it is mostly not the data + +The reported margin is suspiciously flat: −65.1 / −65.0 / −64.9 at amplitude 2e4, −64.1 / +−64.0 / −63.9 at 2e5, −63.0 / −62.9 / −62.8 at 2e6. Two decades of amplitude move it by two +nats. That is not a coincidence and it is not the data speaking. + +For a single-peak row the interval end sits `(W_SIGMA + LOCALISE_SAFETY) * sigma` from the +crest, and the Laplace value of the integral is `lnL_crest + log(sqrt(2 pi) sigma)`, so + + margin = log(T_out) + lnL(end) - result + = log(T_out / (sqrt(2 pi) sigma)) - (W_SIGMA + LOCALISE_SAFETY)**2 / 2 + +**The second term is 75.031 nats and is a pure design constant** — it says where we chose to +cut, nothing else. The only data-dependence is the logarithm, worth `log 10 = 2.3` nats per +decade of amplitude. Measured against that closed form: + +| amplitude | sigma/deltaT | `T_out/(sqrt(2 pi) sigma)` | predicted | reported | diff | +|---|---|---|---|---|---| +| 2e4 | 0.00691 | 3.54e4 | −64.56 | −65.13 | −0.58 | +| 2e5 | 0.00218 | 1.12e5 | −63.41 | −64.13 | −0.73 | +| 2e6 | 0.00069 | 3.54e5 | −62.25 | −63.00 | −0.74 | + +Agreement to under a nat across the range; the residual is the Laplace approximation's own +error (trapezoid, non-Gaussian tails), and it is systematic and small. + +### What follows, and it is not comfortable + +**On a clean single-peak row the tail bound is very nearly vacuous.** It passes iff +`(W_SIGMA + LOCALISE_SAFETY)**2/2 - log(T_out/(sqrt(2 pi) sigma)) > |TAIL_LOG_TOL|`, which is a +statement about three constants and a logarithm — not about the row. It cannot fail on such a +row for any data this module will see. So a `tail_bound_worst` of −64 is NOT evidence that the +truncation was safe here; it is the arithmetic of `W_SIGMA` restated, and reading it as a +per-row safety margin over-reads it. + +This also explains, exactly rather than by observation, the "~40 nats of structural slack" two +independent reviewers measured and could not eat into. The slack is +`(W_SIGMA + LOCALISE_SAFETY)**2/2 - log(T_out/(sqrt(2 pi) sigma))`, and no amount of probing the +SAMPLING could touch it, because sampling was never what set it. + +**Where the check does earn its place is structure OUTSIDE the intervals** — a peak the keep +filter dropped, a comb-like row, an enumeration that has gone wrong. There the outside supremum +is set by that structure and not by the interval end, the closed form above does not apply, and +the bound fires: a deliberately sabotaged enumeration is CAUGHT at all four amplitudes tested, +and the mis-placed-interval fixture now trips it. **The tail bound is a structure detector, not +an accuracy margin**, and it should be read that way. + +(The edge deficit is `(W_SIGMA + LOCALISE_SAFETY)**2/2 = 75.03`, not `W_SIGMA**2/2 = 72` as +earlier drafts of this note said: the interval is widened by the localisation residual. The +suite's coupling assertion uses 72, which remains a valid CONSERVATIVE lower bound on the real +slack, so it is left as it is.) + +## Two more reviewer P1s: one fixed, one answered with a measurement + +### Time resampling with `peak-local` is now REFUSED + +`--resample-time-marginalization` under `peak-local` had two available answers and both were +silently wrong. `--time-posterior-export grid` draws `t_ref` from the original coarse `lnLt` +bins — so the integral would be sub-sample accurate and the exported TIME quantised to the very +grid the option exists to escape, computed and then discarded. With `--interpolate-time nearest` +the `auto` mode resolves to `grid`, so that was the DEFAULT outcome, not a corner. And +`continuous` was never available: it needs `return_time_draw`, which needs a validated dense +reconstruction over the whole window, and this rule by construction never forms one. + +Refused at startup, with the message naming `bandlimited` — which mandates the continuous +export and is one flag away. This is the same disclosed limitation as "`resample_samples()` +unserved", now enforced instead of documented. + +### The certified outside supremum, on SUB-CELL geometry + +The reviewer asked for the uncovered segments to be certified with the spectral derivative +bounds. They now are, and the shipped tail bound is a genuine upper bound on `q` over the +uncovered set — it assumes nothing about where extrema are, so the "a peak between samples was +never enumerated" objection is closed outright rather than argued to be unrealistic. + +**The machinery that arrived for it was wrong.** `enum_grid_derivatives` came in uncalled and +untested and was **52% off on `q'` and 46% on `q''`**: for the reflected row `n` is EVEN, so +`bandlimited_spectrum` splits the Nyquist bin and returns `n+1` coefficients, and that version +used the count as the transform LENGTH — wrong grid, and the two split bins collided on one +index. It now differentiates spectrally and reuses the validated upsampler, agreeing with +pointwise evaluation to **1e-14 relative**. `parabolic_sup` and `segment_sup_bound`, named in +`__all__` but never defined, are implemented. + +**Fetching the slopes is what makes certification affordable.** Slack per cell, measured: + +| bound | amp 2e4 | amp 2e6 | amp 2e7 | +|---|---|---|---| +| one sample + `M2 h^2/2` | 407 | 4.1e4 | 4.1e5 | +| endpoint values + `M2 h^2/8` | 102 | 1.0e4 | 1.0e5 | +| **endpoints + SLOPES + `M4 h^4/384`** | **0.12** | **12.2** | **122** | + +Three orders of magnitude. An earlier revision of this note called certification non-viable; it +had costed only the crude form and was wrong. + +**Sub-cell geometry is the whole point, and bounding whole cells does not work.** This rule +earns its keep on sharp rows, and there **the merged interval is narrower than one enumeration +cell** — measured, the half-width falls to 0.05 of a cell at derived factor 4096. A +cell-granular notion of "covered" then marks nothing, the crest's own cell counts as outside, a +bound over that whole cell bounds the CREST, and every row is rejected. That was measured, and +it is the inert-option outcome. So each cell is bounded on the part that is genuinely outside: +the interval cuts the cell at `lo` and `hi`, leaving `[0, lo]` and `[hi, 1]` in cell +coordinates, and the Hermite is maximised over those. Empty pieces return `-inf` and drop out. +A cell containing TWO merged intervals would leave an uncovered gap between them that a +two-piece decomposition misses; it is detected and the ROW is declined, fail-closed. + +**Result.** Every row still accepted, margins `-65.0 / -63.1 / -55.0` at amplitude 2e4 / 2e5 / +2e6 against `TAIL_LOG_TOL = -23`, conservative relative to the honest supremum throughout. +Contract, accuracy and coverage unchanged: worst `|pl - bl|` over six fixture families still +**2.6e-10**, with 24/24 sharp, 10/10 near-edge and 3/3 two-peak rows kept, and Door 1 at derived +factor 4096 still exact with both peaks retained — now reporting a margin of -42 where the +sampled version reported -6904. + +**A reference trap, for the third time in this file.** The test that pins this first failed +because its own reference was a 512x uniform grid, and on a sub-cell interval (`sigma = 5.1e-7` +s) the grid spacing is comparable to `sigma`, so the nearest grid point outside an interval end +under-reads that end by `W_SIGMA/sigma * spacing ~ 11` nats — and the supremum sits exactly at +an end. A correct bound looked 11 nats loose. The reference now uses the grid for the bulk and +EXACT evaluation at the ends, against which the bound is 0.46 nats high, inside its own 0.64 nat +remainder. Sampled references have now been wrong here three times: periodic instead of +reflected, sample instead of crest, and grid instead of exact-at-the-ends. + + +## Mutation sweep + +**34 mutations against the current code** (`e03dde95`), baseline **109 collected / 108 +passed / 1 skipped** (the skip is the GPU parity test on a CPU runner). Restores from +`git show HEAD:` — a pristine source, never a reverse-edit, never a snapshot taken while a +mutation was live — with every anchor required to match EXACTLY ONCE so a stale anchor +reports a HARNESS FAILURE rather than a false survivor. **0 harness failures**, and the +anchor check earned its keep: it caught one anchor that had gone stale under the round-7 +edit before the sweep could report a phantom survivor. + +**22 killed, 12 survived.** + +### The harness had a defect of its own, and it inverted every verdict + +The first run reported `L1_skip_localisation` as **SURVIVED** beside the words +"11 failed, 92 passed". `pytest ... | tail -4` returns **TAIL's** exit status, which is +always 0, so every killed mutation read as a survivor. This is the same family as any check +whose pass condition is empty output. The runner now captures the full run and takes the +last lines in Python, so the return code is pytest's own. Quoted here because a sweep that +silently inverts its verdicts is worse than no sweep. + +### Killed (22), grouped by the defence that did it + +| defence | mutations killed | +|---|---| +| localisation | `L1` skip it, `L2` one Newton step, `L3` drop the convergence test, `L5` drop the tol widening | +| **the round-7 pinned refusal** | **`P1` count a pinned crest as converged** | +| **the round-6 crest bound** | **`D4c` revert to the parabolic bound**, `D4e` remove the pre-filter entirely (the option goes inert) | +| curvature stencil | `G5` re-clip the stencil centre, `E1` re-exclude endpoints | +| containment / tail bound | `C1` containment always passes, `C3` `T_outside` from indices, `C4` drop `log T_outside`, `C6` ceiling after the cost gate, `C10` disable the tail tolerance | +| constants' coupling | `C9` break the `LOCALISE_SAFETY` relation, **`W2` drop `W_SIGMA` below its slack**, `W3` lower `PEAK_KEEP_NATS` | +| quadrature | `E7` drop the trapezoid half-weights | +| **the round-5 reconstruction** | **`R1` enumerate on the periodic interpolant, `R2` evaluate on it, `R3` route the edge guard back to Simpson** | +| the shipped wiring | **`W1` make the branch delegate to `bandlimited`** — same signature, same numbers, whole cost benefit gone | + +Three of these matter more than the rest. **`W1`** is the inert-option hazard, invisible to +any value comparison because agreeing with `bandlimited` *is* the design; only the +`last_report()` assertion catches it. **`W2`** confirms the `W_SIGMA` coupling inequality is +load-bearing rather than decorative. **`E1` was a SURVIVOR before round 4** and was reported +then as "a no-op over dead code" — it is now killed, which is the cleanest evidence that the +endpoint enumeration Door 2 was meant to enable is actually live. + +### The 12 survivors, and which of them are evidence + +A mutation that cannot be shown to change behaviour is a **harness artifact, not a coverage +gap**, and reporting it as one would be the sweep lying to itself. Each survivor was +re-applied and run over six fixture families chosen to hit the shape it targets, comparing +values AND every `last_report()` counter against pristine. + +**Ten of the twelve are no-ops** — identical values *and* identical counters on all six +families: `L4` (widen the Newton bracket), `G1` (disable the exact keep filter — it only ever +keeps MORE, so it cannot change an answer), `D4a` (shrink the bound's constant 4x), `D4b` +(halve the curvature bound), `D4d` (read the sample at the clipped stencil centre), `C7`, +`C8`, `E5`, `E6`, `G4`. + +**Two changed behaviour, and only these are coverage gaps:** + +* **`R4` drop the `boundary_unresolved` seed — max |Δ| = 6.06 nats** on an endpoint-max row, + with every counter unchanged. A genuine gap, and instructive about why it hid: it survived + `test_row_classification_matches_the_dense_path_exactly` because both rules read the SAME + `_classify_rows`, so removing the clause moves them together and parity still holds. **A + parity test cannot see a change made to the definition both sides read.** + `test_boundary_unresolved_rows_are_refined_and_not_called_flat` now asserts the + classification itself and kills it. **Closed.** +* **`C5` one extra covered sample per end** — values identical to the last bit; only + `tail_bound_worst` moves (−109.7 → −462.9, −57.7 → −67.0, −366.2 → −1591.7). Diagnostics + only. **Open gap, and it is a reporting number, not an answer.** + +Two of the no-op verdicts deserve their reasons stated rather than assumed, because a right +verdict on a false premise is how the next defect hides: + +* **`D4a` / `D4b`.** These weaken the round-6 crest bound by 4x and 2x, and nothing notices. + The reason is not that the bound is untested — `D4c`, which reverts it to the parabolic + FORM, is killed. It is that the spectral bound is **loose by construction**: measured slack + is 15x to 84.5x, so scaling its constant by 2–4x leaves it a valid bound on every fixture + available. The suite pins the bound's FORM but not its CONSTANT, and closing that would + need a fixture where the true crest deficit comes within 4x of the bound. Given the bound + is deliberately conservative, such a fixture may not exist. **Stated as a limit of the + sweep, not as a passed test.** +* **`D4d`.** Reading the peak's value at the inward-clipped stencil centre instead of the + enumerated index — the Round-6 third-site defect — is a no-op on these fixtures because the + round-6 bound's slack absorbs the 132-to-8449 nat mis-read. `test_the_peak_sample_is_read_at + _the_enumerated_index` asserts the property that makes the mis-read wrong, but it does not + observe the module's choice, so it cannot kill this mutation. **Open gap, precisely located.** + + +## A costing lever that is measured, DANGEROUS, and deliberately not in this PR + +Refinement currently runs on every row. On the real 3G case the extrinsic integral is carried +by `n_eff ~ 3` out of 810k samples, which invites the obvious question: could rows far below +the batch's best be left on the coarse rule? + +**Measured**, per-batch coarse peak `lnL` on the real `run_ET_snr100` case (6 batches, 120493 +rows), fraction of rows below their OWN batch maximum by more than: + +| batch | rows | max | median | >60 | >500 | >1000 | >2000 | +|---|---|---|---|---|---|---|---| +| 0 | 20000 | −372.50 | −3183.05 | 1.000 | 0.979 | 0.896 | 0.651 | +| 3 | 20181 | −276.76 | −789.85 | 0.998 | 0.524 | 0.000 | 0.000 | +| 5 | 20196 | −275.08 | −650.01 | 0.992 | 0.202 | 0.000 | 0.000 | +| **all** | **120493** | | | **99.72%** | **60.82%** | **25.31%** | **10.95%** | + +So the row-count opportunity is large, and the spread NARROWS as AV contracts -- batch 0 spans +~2800 nats, batch 5 ~375. + +**It is not taken, and should not be taken lightly, for a reason that is about the physics and +not about the gate.** RO's ruling: *the lnL range in a posterior is wide, and accurate lnL is +needed WAY OFF PEAK to probe the tails.* The rows this would skip are not spare capacity -- +they carry the posterior's tails, and the adaptive sampler ADAPTS on them, so degrading their +`lnL` changes where the next batch looks and therefore the whole run. "Contributes negligibly +to the evidence integral" is not the same property as "may be computed badly", and only the +first is measured above. + +Two further constraints, recorded so the work starts from them: + +* **Batch-local only.** A maximum carried ACROSS calls would be a hidden, order-dependent + variable: each call is independent by contract, and a persistent scale would make results + depend on batch order and defeat reproducibility. +* **The ranking cannot use the coarse value.** Simpson's error is precisely what this whole + line of work exists to remove, so a gate keyed on the coarse `lnL` is unsafe by exactly that + amount. A sound gate needs a certified UPPER bound on each row's refined value compared + against an ACHIEVED lower bound from the same batch -- refine the top rows first, then bound + the rest. The machinery for the upper bound now exists (`spectral_derivative_bound`, + `parabolic_sup`, `segment_sup_bound`) and is cheap on the COARSE grid: the Hermite remainder + there is ~492 nats at amplitude 2e4, against ~6500 for a curvature-only bound. + +**Disposition (RO): OPT-IN AND EXPLICITLY EXPERIMENTAL if it is ever built, expected to need +WEEKS of calibration to establish safety, and it belongs in a NEW PR -- not this one.** Recorded +here only so the measurement and the hazard are not lost. + +## Not done in this draft + +* **The evaluator is a direct spectral sum**, `O(npts)` per output point. A chirp-z + (Bluestein) evaluation would make it `O((npts + M) log(npts + M))` and is the single + largest remaining cost lever. Not attempted here. +* **No GPU measurement.** The path is `xpy`-generic and there is a cupy parity test, + but the cost table above is CPU only. +* **Real data: ONE run, one grid point.** See "Real data, through the shipped driver" -- + `peak-local == bandlimited` byte-identical end to end, cost 14.2x `simpson`. That is a + single ET-SNR100 grid point at low `n_eff`; there is no campaign, no SNR ladder through this + path, and no converged accuracy comparison against `simpson`. Formerly Accuracy is against analytic truth and against the dense path; + the dense path's own real-injection comparison has not been repeated for this rule. +* **`MAX_INTERVALS`, `PEAK_KEEP_NATS`** are fail-closed guards with an argument behind + them but no sweep behind the specific values. +* ~~**The tail bound is still a sampled maximum, not a supremum.**~~ SUPERSEDED, and + the entry is kept rather than deleted because it records what changed. As of + `6b4467ec` `q_out_max` is no longer a sampled maximum: it is + `max` over cells of `segment_sup_bound`, a cubic Hermite through each cell's endpoint + values AND slopes plus the classical remainder `M4 h^4/384`, with `M4` from + `spectral_derivative_bound` at order 4. That is the between-samples bound this bullet + called "not attempted", and it is certified on SUB-CELL geometry because bounding whole + cells makes the option inert (see `DESIGN_peak_local_framework.md`). +* **No re-measure-and-double loop.** The dense path ENFORCES its resolution criterion; + this path derives the spacing and then verifies the outcome two other ways. Both are + checked; they are not the same criterion, and this file no longer claims they are. +* **Phase marginalization is REFUSED**, at the library and at driver startup — a + deliberate scope cut, not an omission. Production marginalizes over distance, not + phase, and under phase marginalization the time peak's Laplace width picks up an + `(I1/I0)(|kappa|/D)` factor that does not reduce, so the local spacing would stop + being derivable from `rho_sq` and the curvature alone. `bandlimited` still supports + it and is unchanged. +* **`resample_samples()` is not served.** The extrinsic time-export path needs a full + `lnL(t)` array on the original grid, which this rule by construction does not + produce. `return_lnLt=True` therefore still runs the coarse path, exactly as it does + under `bandlimited`. `t_star` and the local widths ARE exposed + (`return_peaks=True`), which is the piece a future export or a time-first + reordering of the marginalizations would build on. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md new file mode 100644 index 000000000..c337106c1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md @@ -0,0 +1,342 @@ +# Time-marginalization quadrature: measured record + +Companion to `time_marginalization_quadrature.py`. The module docstring carries the +argument; this file carries the numbers behind it and the harnesses that produced them, +so a reviewer can re-run rather than take them on assertion. + +Downstream: `DESIGN_time_marginalization_peak_local.md` is the follow-up rule that +delegates to this one, and `DESIGN_peak_local_framework.md` is the planning note for +generalizing that rule to other parameters. This module is the BACKSTOP in that +picture — the thing peak-local falls back to — so changes here move both rules. + +Harnesses (host-local, `ldas-*` NFS home): `~/tmarg_harness/`. +`probe.py` periodic-window accuracy, `wrap.py` non-periodic window, `adv.py` edge sweep and +mixed blocks, `detrend.py` the rejected endpoint-detrend, `cost.py` quadrature-only cost, +`cost_e2e.py` end-to-end through the shipped likelihood, `peaklocal2.py` the peak-local +prototype below, `real_path.py` / `simps_iso.py` the GPU correctness runs, `cost_gpu.py` the +GPU cost table. + +## The defect + +Simpson at the fixed spacing `deltaT = 1/srate` against an integrand of width +`sigma_t = 1/(2 pi rho sigma_f)`. Grid-phase span of the reported lnL over `2*deltaT`, +35+30 Msun SEOBNRv4 H1L1V1, rho=40 — **measured on the JAX mirror**, quoted as the physical +scale of the defect: + +| srate | 4096 | 8192 | 16384 | +|---|---|---|---| +| span | 1.649 nats | 0.385 | 0.0095 | + +## Accuracy, against an analytic truth + +Synthetic band-limited kappa, srate 4096, npts 614, error in nats at three grid phases: + +| sigma_t/deltaT | Simpson | band-limited | factor | +|---|---|---|---| +| 2.27–2.61 | +5e-6 … +0 | 0 | 1 | +| 0.72–0.83 | +2.5e-3 … -1.1e-4 | 0 | 4 | +| 0.25–0.28 | +0.844 / +0.242 / -1.101 | 0 | 16 | +| 0.10–0.12 | +1.742 / -1.833 / -11.68 | 0 | 32 | +| 0.046–0.052 | +2.548 / -15.33 / -66.18 | 0 | 64 | +| 0.016–0.019 | +3.589 / -139.4 / -549.0 | 0 | 256 | +| 0.007–0.008 | +4.393 / -710.6 / -2760.4 | 0 | 512 | + +Non-periodic window (segment of a longer band-limited signal, peak centred): the literal +2N forward/backward reflection measured errors from 2e-8 to 2.5e-6 nats over amplitudes +0.05–5, where Simpson can be wrong by hundreds of nats. + +## Finite-window reconstruction: why decay is not an eligibility test + +`kappa_rows` contains only the gathered integration-window slice of a longer inverse-FFT +series. Zero-padding its FFT directly treats that slice as one period; centring the coarse +argmax does not bound the artificial wrap or its ringing. More subtly, decay of the +integrand does not bound it either: `exp(lnL)` may be negligible at the edges while the +quantity being FFT-interpolated, complex `kappa`, has a large endpoint mismatch. + +The adversarial review pinned that distinction with a centred analytic row: outer-eighth +log-likelihood drops 51.4 and 49.6 nats, endpoint mismatch 1720, derived refinement factor +64. The raw periodic-slice FFT was **+140.88 nats** wrong. Thus a 30-nat tail guard would +have certified precisely the catastrophic case it was intended to exclude. + +The shipped mitigation forms the literal length-2N sequence +`[kappa[0], ..., kappa[-1], kappa[-1], ..., kappa[0]]`, FFT-interpolates that periodic, +value-continuous sequence, and retains only the forward interval. On the counterexample the +error is **+2.4e-4 nats**. The superficially standard 2(N-1) reflection was tested too and +was worse (+0.078 nats) because it places the turn on the endpoint sample rather than between +the duplicated endpoints. On the ordinary longer-period fixtures the literal 2N form was +also better: 2e-8–2.5e-6 nats versus 2e-6–2.5e-4 for 2(N-1). + +An internal randomized review added 400 coherent sinusoids (random amplitude, mode and +phase) to the longer-period analytic fixture. Of 387 rows that passed the former centring +and 30-nat tail guards, raw periodization still reached 130 nats error. Reflection reduced +the worst error to 0.0137 nats; 364/387 were below 1e-3 and only one exceeded 1e-2. The +worst residual used a deliberately coherent near-Nyquist mode of amplitude 1894, comparable +to the signal peak. Production remeasurement doubled its quadrature factor from 128 to 256 +and left the value unchanged, establishing that 0.0137 is reconstruction, not integration, +error. Local Lanczos reconstruction was also tested on that worst row and was inferior: +errors 7.24, -0.779, 0.164, 0.0646 and 0.0273 nats for half-widths 8, 16, 32, 64 and 128. +This records the measured limitation instead of implying that reflection recovers the +unavailable full-period series exactly. + +This is a numerical boundary condition, not a claim that the physical correlation reverses +outside the window. The integration domain remains exactly `[t0, t_{N-1}]`; the backward +half contributes no probability mass. + +## Boundary peaks: diagnostic, never a rule switch + +Peak swept toward the window edge, `sigma_t/deltaT = 0.042`: + +| distance from edge (samples) | 307 | 100 | 30 | 8 | 2 | 0 | +|---|---|---|---|---|---|---| +| band-limited | 5e-6 | 4.6e-3 | 5.2e-2 | 5.6e-2 | -3.3 | **+88.8** | +| Simpson | -29.2 | -29.9 | -29.3 | -29.4 | -29.7 | -29.9 | + +`lnL` is linear in `kappa`, so the error at a fixed distance scales with amplitude. Just +outside the guard: -8.0e-4 / -8.1e-3 / -8.1e-2 / **-0.846** nats at peak lnL 5.3e2 / 5.3e3 / +5.3e4 / 5.3e5 (rho ~ 33 / 103 / 326 / 1031). The O4c effort measured the same linear scaling +on a different fixture and implementation and reached the same magnitude at the same +amplitude — corroboration across lines, not a single-fixture artefact. + +The table records the rejected raw-slice periodic reconstruction. An endpoint-ramp detrend +was also rejected: it halves the interior error but is worse at 8 and 2 samples from the edge +(`detrend.py`). + +The former outer-eighth guard is now diagnostic-only. A sweep immediately across its +boundary (peak samples 75.3, 76.3, 77.3 for N=614) gave reflected errors of order 1e-6 nats; +switching to Simpson there would create a discontinuous, arbitrary loss of accuracy. Peaks +at the actual integration endpoint remain a distinct physical truncation problem. The +reflected result is best-effort on the unchanged domain and the boundary count remains in +`last_report()` so truncation is visible, but it neither returns a lower-resolution rule nor +raises: in the production pipeline either behavior can silently bias selection, because an +exception may be interpreted as waveform failure and excise that configuration. + +This distinction is load-bearing: + +* a **numerical reconstruction boundary** is mitigated by forward/backward reflection; +* a **physical integration boundary** is reported, not silently reclassified; +* Simpson fallback is retained only where no peak width can be measured or no refinement is + needed, not because a row crossed a location or tail-height threshold. + +One endpoint corner needs explicit classification. At an argmax on the first or last sample, +the nominal centred curvature stencil is clipped inward; on a severely under-resolved row it +can measure positive curvature away from the peak and call the row flat. A nonconstant row in +that state now receives a seed factor of 4, after which dense-grid remeasurement derives the +needed resolution (factor 16 in the pinned sharp-endpoint fixture). A truly constant antenna +null remains flat and unrefined. This prevents an exact-boundary row from silently reaching +Simpson through a different classification path. + +## Odd npts + +`marginalization_time_grid(0.075, 1/srate)` gives npts = 153 / 307 / 614 / 1228 / 2457 at +srate 1024 / 2048 / 4096 / 8192 / 16384 — **odd at three of five**. A spectrum split at +`h = n//2` files the highest positive frequency under a negative frequency for odd `n`: +max error 1.4e-12 at n=614 but 4.1e-1 at n=613, 5.4e-1 at n=307, 6.0e-2 at n=2457. Exact at +the samples, wrong between them — so a "reproduces its input" test cannot see it, and a +fixture with an empty top bin cannot either. + +## CPU/GPU: a PRE-EXISTING divergence, not introduced here + +`factored_likelihood` integrates with scipy's `simpson` on CPU and the vendored +`optimized_gpu_tools.simps` on GPU. The latter is an old scipy with `even='avg'`; modern +scipy uses the Cartwright correction. **Odd N agree exactly; even N do not** — and +production npts is even at srate 4096/8192: + + n=613 random scipy=318.0505029932 gpu=318.0505029932 reldiff 0 + n=614 random scipy=309.1334604435 gpu=305.7046546187 reldiff 1.1e-2 + n=615 random scipy=298.5401667745 gpu=298.5401667745 reldiff 0 + +Through the shipped likelihood on `ldas-pcdev13`, numpy vs cupy differ by up to **0.405 nats** +under the historical `simpson` quadrature. Worth its own investigation; not repaired here. + +Measured *after* this change, same inputs: band-limited numpy-vs-cupy max 1.1e-6, median +4.9e-14, versus Simpson's max 0.405, median 5.4e-3 — refined rows integrate with trapezoid on +the dense grid, which has no even/odd ambiguity, so this path removes the divergence where it +applies. + +## Cost, and why the strategy should change + +The table below predates the finite-window fix and measures the rejected raw-slice FFT. +Forward/backward reconstruction doubles the FFT period and the chunking budget accounts for +that larger temporary. Treat these numbers as the lower-bound historical record, not as a +current performance claim; correctness is the gate for this opt-in quadrature and the +peak-local follow-up remains the intended cost reduction. + +End-to-end through the shipped likelihood, n_extrinsic 4000, 3 IFOs, CPU time: + +| sigma_t/deltaT | Simpson | band-limited | ratio | +|---|---|---|---| +| 1.74 | 0.212 s | 0.241 s | 1.1x | +| 0.55 | 0.302 s | 0.610 s | 2.0x | +| 0.17 | 0.180 s | 1.741 s | 9.7x | +| 0.055 | 0.250 s | 6.652 s | 26.6x | + +Host-sensitive: O4c measured the same quantity moving up to 2x between hosts. The Simpson +baseline is rho-independent by construction, so a run where it moves with rho is contaminated. + +### On GPU, at production settings -- which is what the table above is not + +The CPU table above is `n_extrinsic = 4000` with the module's DEFAULT affine callback. +Production ILE runs `--vectorized --gpu`, at `--n-chunk 40000` by default, and passes +`distmarg_loglikelihood` at every call site. All three matter, and all three make it worse. + +Method. `~/tmarg_harness/cost_gpu.py` and, for everything below, +`~/adv_tmarg_gpu_audit/adv_cost2.py` -- an independently written harness whose numbers agree +with the first to a few percent at the shared operating point. Same shipped +`DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` call, same synthetic band-limited fixture, +3 IFOs. `ldas-pcdev13` `CUDA_VISIBLE_DEVICES=3` (GeForce RTX 2080 Ti, cc75), cupy 12.0.0 from +the IGWN CVMFS python, `OMP_NUM_THREADS=1`, GPU otherwise idle. Arms INTERLEAVED within a +replicate with the order balanced across replicates; 4-6 replicates; quoted spread is min-max +of the per-replicate ratios. Timing is WALL CLOCK with an explicit device synchronize, on +both arms; CUDA events agree to <1%, and removing the sync moves the Simpson arm, so the sync +is load-bearing and correctly placed. Two significant figures throughout: the device-to-device +spread below is 2.6x, which is what actually bounds these numbers. + +**Production configuration** -- `distmarg_loglikelihood`, `rho_sq = 100` (the affine fixture +ships `rho_sq = 0`, which makes `x0 = kappa_sq/rho_sq` NaN for every row and refines nothing, +so a distmarg run at `rho_sq = 0` measures pure overhead and must be discarded). srate 4096, +npts 614: + +| n_extrinsic | sigma_t/deltaT | Simpson | band-limited | ratio | +|---|---|---|---|---| +| 4,000 | 0.044 | 0.019 s | 0.90 s | **49x** | +| 16,000 | 0.035 | 0.038 s | 3.34 s | **89x** | +| **40,000 (the default `--n-chunk`)** | 0.031 | 0.084 s | **15.7 s** | **190x** | + +srate 16384, npts 2457 (**odd** -- three of five production rates). `n_extrinsic = 40000` at +this rate exceeds the 11 GB card, which is itself worth knowing: + +| n_extrinsic | sigma_t/deltaT | Simpson | band-limited | ratio | +|---|---|---|---|---| +| 4,000 | 0.130 | 0.036 s | 0.59 s | **16x** | +| 16,000 | 0.107 | 0.121 s | 3.93 s | **32x** | +| 40,000 | -- | out of memory on 11 GB | | | + +**The ratio triples between the measured 4,000 and the production 40,000, and it does so for a +reason worth reading.** It is not only that the GPU baseline is nearly free. The refinement +factor is derived ONCE PER GROUP of rows and re-doubled until the criterion holds for the +group MINIMUM (`_integrate_group`: `sigma_dense_min = min(...)` over the chunk). Ten times as +many rows reach ten times deeper into the tail of that minimum, so the whole group pays an +extra octave: the factor histogram at the worst rung moves from mostly 32 at n=4,000 +(`{16: 233, 32: 3126, 64: 235}`) to mostly 64 at n=40,000 (`{32: 610, 64: 35236, 128: 188}`). +The cost per row therefore GROWS with the chunk size rather than staying flat. Anyone reading +the earlier "the baseline is nearly free, so any added work reads as a large multiple" +explanation would expect the factor to shrink once the baseline does real work; it does the +opposite. + +**The affine, n=4,000 table, kept because it is what the CPU table compares against.** Same +device, `--callback affine`, `rho_sq = 0`: + +| sigma_t/deltaT | GPU 4096 (npts 614) | GPU 16384 (npts 2457) | CPU 4096 | CPU 16384 | +|---|---|---|---|---| +| ~1.7 | 1.9x | 1.7x | 1.0x | 1.1x | +| ~0.53 | 4.5x | 4.8x | 2.5x | 2.2x | +| ~0.17 | 12x | 16x | 9.1x | 6.3x | +| ~0.055 | 35x | 56x | 31x | 20x | + +The CPU columns are the same host and script as the GPU ones, so the GPU/CPU difference is a +backend effect and not a host difference; the srate-4096 CPU column reproduces the published +CPU table (1.1 / 2.0 / 9.7 / 26.6). + +#### What this costs in seconds, stated correctly + +An earlier draft of this section claimed GPU band-limited was "~20-40x cheaper in seconds than +CPU Simpson". That was wrong by about 20x: the two numbers it divided (17.0 s and 36.3 s) are +the CPU BAND-LIMITED times, not the CPU Simpson baseline, so the quotient was "this GPU is 40x +faster than this CPU at the same task" -- a statement about two devices, not about the option. +Measured, same host, worst rung, affine n=4,000: + +| | srate 4096 | srate 16384 | +|---|---|---| +| CPU Simpson (the historical cost) | 0.53 s | 1.80 s | +| CPU band-limited | 16.9 s | 36.1 s | +| GPU band-limited | 0.43 s | 0.92 s | + +So GPU band-limited is **1.2x and 2.0x** cheaper than simply running Simpson on CPU -- not 20-40x. +At production settings (distmarg, n=40,000) the GPU band-limited call is 15.7 s against a GPU +Simpson baseline of 0.084 s, i.e. it is far more expensive than any CPU-Simpson comparison. + +#### Three things this table does not control for + +* **Device: 2.6x spread, larger than anything else here.** Identical operating point and + identical refinement histograms, worst rung, affine n=4,000: + RTX 2080 Ti (cc75) **35x**, RTX 3080 (cc86) **58x**, A100-PCIE-40GB (cc80) **22x**. + The ratio divides an overhead-dominated quantity by a work-dominated one, so it substantially + measures the device's launch overhead. Quote it to two figures and expect a factor of ~2.5 + either way on unseen hardware. +* **`_DENSE_CHUNK_BYTES = 128 MB` is a tunable constant worth 7-17% of the cost.** Each chunk + forces a host sync. At 1 GB, worst rung 35x -> 32x and the third rung 12x -> 9.6x. It is not + part of the derivation and was never tuned. +* **The rung-to-rung spread of the Simpson arm is a DRIFT MONITOR, not a control.** The Simpson + arm does identical arithmetic at every rung, and the ratios are formed from paired interleaved + calls, so common-mode drift has already cancelled before that check runs. It has no power + against n_extrinsic, callback or device, which are the confounds that actually govern this + table. For scale: the published CPU table above would FAIL it (baselines 0.212 / 0.302 / + 0.180 / 0.250 s, max/min = 1.68), so "reproduces the published CPU table" means agreement with + a table carrying about +-40% internal baseline noise. + +## Selecting it from a campaign + +`--internal-ile-time-marginalization-quadrature {simpson|bandlimited}` on +`util_RIFT_pseudo_pipe.py`, default `None` meaning "emit nothing" so the default workflow is +byte-identical. Four hops: pseudo_pipe forwards it to `helper_LDG_Events.py`, which validates +it and appends `--time-marginalization-quadrature` to `helper_ile_args`; that becomes +`args_ile.txt`; and `create_event_parameter_pipeline_BasicIteration` inherits the whole +argument string into `ILE.sub`, `ILE_extr.sub`, `ILE_puff.sub` and `ILE_fetch.sub`. + +* **The exclusion list lives in `time_marginalization_quadrature.py`** + (`_PIPELINE_REQUIRED_ILE_FLAGS` / `_PIPELINE_EXCLUDING_ILE_FLAGS`), mirroring the `_tq_prereqs` + block in `bin/integrate_likelihood_extrinsic_batchmode`. Both pipeline layers import it; it is + never re-typed. Matching handles optparse's equals form and unique-prefix abbreviations, + including short legal spellings such as `--g` for `--gpu`; there is no invented minimum + abbreviation length. Exact-option precedence keeps `--time-marginalization` distinct from + an abbreviated `--time-marginalization-quadrature`. +* **The guard checks the BYTES, not the parsed options.** `refuse_unless_time_quadrature_emitted` + requires the flag to be present exactly once with the requested value in the argument string + about to be written. A guard keyed on the options approves an `args_ile.txt` that never + received the flag. The helper is invoked by name, so pseudo-pipe also removes the generated + `helper_ile_args.txt` before invocation and refuses a nonzero helper status; otherwise a + same-value stale file can satisfy even the byte guard. +* **The extrinsic stage is only half covered.** The flag reaches `ILE_extr.sub`, but + `--resample-time-marginalization` calls the likelihood with `return_lnLt=True`, which returns + `lnL(t)` on the original grid and never reaches the quadrature. The marginalized `lnL` is + refined; the exported `t_ref` is still quantised at `deltaT = 1/srate`. The helper prints this + at build time. +* **Never set it in an ini.** The RIFT ini parser overrides the command line for non-boolean + options, so an ini value silently wins over a Makefile's. + +### The follow-up: enumerate peaks, integrate locally (RO'S, 2026-08-27) + +The dense strategy refines the WHOLE window to a peak whose width shrinks as 1/rho, so it +works hardest exactly where the peak occupies least of the domain. It conflates two +requirements that should be separated: + +* resolving `kappa(t)` enough to **enumerate its extrema** — a small factor, and + **SNR-independent**, because kappa is band-limited at Nyquist by construction; +* resolving `exp(lnL)` to integrate it — the rho-dependent part, needed only over a few + `sigma_t` around each enumerated peak. + +Enumeration is also what makes the truncation *rigorous* rather than hopeful, which is the +brief's warning about PR #201's seed-and-hope: every maximum of the band-limited interpolant +is found, so the mass outside the local windows is bounded rather than assumed. The O4c +effort sharpened the argument usefully — every shipped callback is monotone in `Re kappa` / +`|kappa|`, so the maxima of `lnL` ARE the maxima of `kappa`, and enumeration on kappa alone +suffices whatever the callback. + +Prototype (`peaklocal2.py`), same analytic truth, windows merged into disjoint intervals: + +| sigma_t/deltaT | rho~ | dense err | peak-local err | dense pts | local pts | speedup | +|---|---|---|---|---|---|---| +| 0.74 | 2 | 0.000000 | 0.000000 | 2,456 | 3,710 | 0.7x | +| 0.15 | 11 | 0.000000 | 0.000000 | 9,824 | 430 | 23x | +| 0.105 | 15 | 0.000000 | -0.000000 | 19,648 | 98 | 200x | +| 0.017 | 98 | 0.000000 | 0.000000 | 78,592 | 98 | 802x | +| 0.0023 | 692 | 0.000000 | 0.000000 | 628,736 | 97 | **6,482x** | + +Exact at every SNR, flat at ~97 points above rho ~ 15. **Merging the windows is what makes +this one algorithm rather than a regime switch**: isolated peaks give a tiny union; +overlapping peaks grow the union to the whole window and the method degenerates continuously +into the dense grid. Without merging it double-counts and gives +1.6 nats at rho ~ 6. + +Sequencing (RO'S): land the dense implementation first as the reviewed reference, then this +as a separate PR that can be A/B'd against it. The default stays `simpson` regardless of how +cheap it turns out. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SphericalHarmonics_gpu.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SphericalHarmonics_gpu.py index 1f62d77bd..fcbba5338 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SphericalHarmonics_gpu.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SphericalHarmonics_gpu.py @@ -162,10 +162,23 @@ +def _infer_xpy(theta): + """Pick the array backend from the data when the caller did not name one. + + The module-level ``xpy_default`` is cupy on any host where cupy merely + imports, so an unadorned call from host-numpy code used to dispatch a GPU + kernel onto host arrays and raise. Explicit ``xpy=`` always wins; this only + covers callers that leave it unset. + """ + if cupy_here and isinstance(theta, cupy.ndarray): + return cupy + return np + + def SphericalHarmonicsVectorized( lm, theta, phi, - xpy=xpy_default, dtype=np.complex128, + xpy=None, dtype=np.complex128, l_max=8, ): """ @@ -180,8 +193,10 @@ def SphericalHarmonicsVectorized( Array of polar angles to evaluate harmonics at. phi : array_like, shape = (n_params,) Array of azimuth angles to evaluate harmonics at. - xpy : numpy or cupy (default is cupy if loaded, else numpy) - Numpy implementation to use. + xpy : numpy or cupy (default: inferred from `theta`) + Numpy implementation to use. If left as None, cupy is used when `theta` + is a cupy array and numpy otherwise, so host arrays never reach a GPU + kernel. dtype : numpy.dtype (default is numpy.complex128) Datatype to use for output. Must be complex. l_max : int (default is 8) @@ -196,6 +211,9 @@ def SphericalHarmonicsVectorized( Array of spherical harmonics. First axis varies `theta, phi`, and second axis varies `l, m`. """ + if xpy is None: + xpy = _infer_xpy(theta) + l, m = lm.T n_indices = l.size diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py index 1ab3f86c9..8334227ef 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py @@ -63,6 +63,18 @@ from itertools import product, combinations import math +from . import time_marginalization_quadrature as time_quadrature_module +from . import time_marginalization_peak_local as time_peak_local_module +from .time_marginalization_quadrature import TIME_QUADRATURE_CHOICES + +#: Time-marginalization quadrature used when a caller does not pass +#: ``time_quadrature`` explicitly. 'simpson' is the historical fixed-deltaT +#: Simpson rule and remains the DEFAULT; 'bandlimited' refines the time grid to +#: the integrand actually present, using only the samples already computed (see +#: RIFT.likelihood.time_marginalization_quadrature). Drivers set this once from +#: their CLI so every call site inherits it; tests pass the kwarg directly. +TIME_QUADRATURE_DEFAULT = 'simpson' + from .vectorized_lal_tools import ComputeDetAMResponse,TimeDelayFromEarthCenter import os @@ -205,7 +217,7 @@ def internal_hlm_generator(P, extra_waveform_kwargs={}, use_gwsignal=False, use_gwsignal_approx=None, - use_external_EOB=False,nr_lookup=False,nr_lookup_valid_groups=None,no_memory=True,perturbative_extraction=False,perturbative_extraction_full=False,hybrid_use=False,hybrid_method='taper_add',use_provided_strain=False,ROM_group=None,ROM_param=None,ROM_use_basis=False,ROM_limit_basis_size=None,skip_interpolation=False,**kwargs): + use_external_EOB=False,nr_lookup=False,nr_lookup_valid_groups=None,no_memory=True,perturbative_extraction=False,perturbative_extraction_full=False,hybrid_use=False,hybrid_method='taper_add',use_provided_strain=False,ROM_group=None,ROM_param=None,ROM_use_basis=False,ROM_limit_basis_size=None,skip_interpolation=False,force_22_mode=False,**kwargs): """ internal_hlm_generator: top-level front end to all waveform generators used. Needs to be restructured so it works on a 'hook' basis, so we are not constantly changing the source code @@ -253,7 +265,7 @@ def internal_hlm_generator(P, elif use_gwsignal and (has_GWS): # this MUST be called first, so the P.approx is never tested if not quiet: print( " FACTORED LIKELIHOOD WITH hlmoff (GWsignal) " ) - hlms, hlms_conj = rgws.std_and_conj_hlmoff(P,Lmax,approx_string=use_gwsignal_approx,**extra_waveform_kwargs) + hlms, hlms_conj = rgws.std_and_conj_hlmoff(P,Lmax,approx_string=use_gwsignal_approx,force_22_mode=force_22_mode,**extra_waveform_kwargs) elif (not nr_lookup) and (not NR_group) and ( P.approx ==lalsim.SEOBNRv2 or P.approx == lalsim.SEOBNRv1 or P.approx==lalsim.SEOBNRv3 or P.approx == lsu.lalSEOBv4 or P.approx ==lsu.lalSEOBNRv4HM or P.approx == lalsim.EOBNRv2 or P.approx == lsu.lalTEOBv2 or P.approx==lsu.lalTEOBv4 ): # note: alternative to this branch is to call hlmoff, which will actually *work* if ChooseTDModes is propertly implemented for that model @@ -307,7 +319,7 @@ def internal_hlm_generator(P, # hlms_conj = hlms_conj_list if not('fd_standoff_factor' in extra_waveform_kwargs): extra_waveform_kwargs['fd_standoff_factor'] = 0.9 # IMPORTANT to match SimInspiralTD. But allow user to override - hlms, hlms_conj = lsu.std_and_conj_hlmoff(P,Lmax,**extra_waveform_kwargs) + hlms, hlms_conj = lsu.std_and_conj_hlmoff(P,Lmax,force_22_mode=force_22_mode,**extra_waveform_kwargs) elif (nr_lookup or NR_group) and useNR: # look up simulation # use nrwf to get hlmf @@ -433,7 +445,7 @@ def PrecomputeLikelihoodTerms(event_time_geo, t_window, P, data_dict, extra_waveform_kwargs={}, use_gwsignal=False, use_gwsignal_approx=None, - use_external_EOB=False,nr_lookup=False,nr_lookup_valid_groups=None,no_memory=True,perturbative_extraction=False,perturbative_extraction_full=False,hybrid_use=False,hybrid_method='taper_add',use_provided_strain=False,ROM_group=None,ROM_param=None,ROM_use_basis=False,ROM_limit_basis_size=None,skip_interpolation=False, calibration_realizations=None, calibration_conjugate=False, return_calibration_crossterms=False, calibration_self_term=True): + use_external_EOB=False,nr_lookup=False,nr_lookup_valid_groups=None,no_memory=True,perturbative_extraction=False,perturbative_extraction_full=False,hybrid_use=False,hybrid_method='taper_add',use_provided_strain=False,ROM_group=None,ROM_param=None,ROM_use_basis=False,ROM_limit_basis_size=None,skip_interpolation=False,force_22_mode=False, calibration_realizations=None, calibration_conjugate=False, return_calibration_crossterms=False, calibration_self_term=True): """ Compute < h_lm(t) | d > and < h_lm | h_l'm' > @@ -493,7 +505,7 @@ def PrecomputeLikelihoodTerms(event_time_geo, t_window, P, data_dict, hybrid_use=hybrid_use,hybrid_method=hybrid_method,use_provided_strain=use_provided_strain, ROM_group=ROM_group,ROM_param=ROM_param,ROM_use_basis=ROM_use_basis,ROM_limit_basis_size=ROM_limit_basis_size, extra_waveform_kwargs=extra_waveform_kwargs,use_gwsignal=use_gwsignal,use_gwsignal_approx=use_gwsignal_approx, - skip_interpolation=skip_interpolation) + skip_interpolation=skip_interpolation,force_22_mode=force_22_mode) if not(ignore_threshold is None) and (not ROM_use_basis): @@ -777,7 +789,16 @@ def FactoredLogLikelihoodTimeMarginalized(tvals, extr_params, rholms_intp, rholm # Said another way, the m^th harmonic of the waveform should transform as # e^{- i m phiref}, but the Ylms go as e^{+ i m phiref}, so we must give # - phiref as an argument so Y_lm h_lm has the proper phiref dependence - Ylms = ComputeYlms(Lmax, incl, -phiref, selected_modes=rholms_intp[list(rholms.keys())[0]].keys()) + # InterpolateRholms deliberately drops identically-zero data/mode overlaps. + # Keep the scalar nearest-neighbour path on that same active-mode set: using + # every raw rholm below while building Ylms from only the retained modes + # otherwise raises KeyError for symmetry-suppressed modes (for example the + # (2,1) mode of an equal-mass source) and for exactly zero data. Detectors + # can retain different modes, so build harmonics for their union. + active_modes = set() + for det in detectors: + active_modes.update(rholms_intp[det].keys()) + Ylms = ComputeYlms(Lmax, incl, -phiref, selected_modes=active_modes) # lnL = 0. lnL = np.zeros(len(tvals),dtype=RiftFloat) @@ -795,7 +816,8 @@ def FactoredLogLikelihoodTimeMarginalized(tvals, extr_params, rholms_intp, rholm det_rholms[key] = func(float(t_det)+tvals) else: # do not interpolate, just use nearest neighbors. - for key, rhoTS in rholms[det].items(): + for key in rholms_intp[det]: + rhoTS = rholms[det][key] tfirst = float(t_det)+tvals[0] ifirst = int(np.round(( float(tfirst) - float(rhoTS.epoch)) / rhoTS.deltaT) + 0.5) ilast = ifirst + len(tvals) @@ -2208,8 +2230,11 @@ def _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts): return Qlms -SINC_HALFWIDTH_DEFAULT = 8 # taps per side for time_interp='sinc' (stencil 2a); see - # _sinc_Q_window_numpy for the accuracy-vs-oversampling crossover +# taps per side for time_interp='sinc' (stencil 2a); see _sinc_Q_window_numpy for the +# accuracy-vs-oversampling crossover. DEFINED IN time_interp_choice, not here, so the JAX +# gatherer can read it without importing this module (numba/lal); re-exported so every existing +# `factored_likelihood.SINC_HALFWIDTH_DEFAULT` reference keeps resolving. +from .time_interp_choice import SINC_HALFWIDTH_DEFAULT def _sinc_lanczos_weight_matrix(u, a=SINC_HALFWIDTH_DEFAULT, xpy=np): @@ -2295,10 +2320,13 @@ def _sinc_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts, RIFT/likelihood/DESIGN_q_window_stencil.md. Automatic selection was removed as measurably unreliable. - NO stencil is applied by default -- time_interp defaults to 'nearest', as does - --interpolate-time when omitted, so a caller who asks for nothing gets the nearest-bin gather - and neither interpolating stencil; 'cubic' is only the legacy truthy --interpolate-time - mapping. + THE LIBRARY DEFAULT AND THE DRIVER DEFAULT ARE DIFFERENT, deliberately. This function's + own ``time_interp`` argument still defaults to 'nearest', so a library caller who asks for + nothing gets the nearest-bin gather and no interpolating stencil -- every existing caller is + unaffected. bin/integrate_likelihood_extrinsic_batchmode's --interpolate-time, by contrast, + defaults to ``time_interp_choice.TIME_INTERP_DEFAULT`` as of 2026-09-02 (issue #233), which is + what an ILE run now gets when the flag is omitted; 'cubic' remains only the legacy truthy + --interpolate-time mapping. COST, measured (not estimated from the tap count): CPU ~4.2-4.5x cubic -- 2a=16 taps against 4, and this path IS tap-count bound. @@ -2392,6 +2420,55 @@ def _q_inner_product_gpu(Q, A, start_indices, fractional_offsets, npts, time_int % (time_interp, TIME_INTERP_CHOICES)) +_Q_EXPLICIT_TEMP_MAX_BYTES = 64 * 1024 * 1024 + + +def _q_inner_product_explicit_times(Q, A, start_indices, fractional_offsets, + time_interp, xpy=np): + """Evaluate a Q-window stencil at every explicitly supplied time. + + ``start_indices`` has shape ``(n_extrinsic, n_time)``. Work is chunked on + BOTH axes so neither one very finely refined row nor many fair-draw rows can + allocate the full ``n_extrinsic*n_time*n_modes`` gather at once. Each + flattened entry asks the existing, tested stencil for one sample; unlike + the historical window gather, no implicit ``+j*deltaT`` is introduced. + """ + if start_indices.ndim != 2: + raise ValueError("explicit start_indices must have shape (n_extrinsic, n_time)") + n_ext, n_time = start_indices.shape + n_modes = A.shape[-1] + # Q gather + repeated antenna/mode row + interpolation workspace, kept + # below ~64 MiB even when n_time alone is enormous. The previous row-only + # chunk had a minimum of one row, so a 2.5M-time, 21-mode row still made + # >800 MiB A_rows and Q_one temporaries apiece. + max_cells = max(1, _Q_EXPLICIT_TEMP_MAX_BYTES // + max(1, n_modes * 16 * 3)) + time_chunk = max(1, min(n_time, max_cells)) + ext_chunk = max(1, min(n_ext, max_cells // time_chunk)) + out = xpy.empty((n_ext, n_time), dtype=np.complex128) + for ext_start in range(0, n_ext, ext_chunk): + ext_stop = min(n_ext, ext_start + ext_chunk) + for time_start in range(0, n_time, time_chunk): + time_stop = min(n_time, time_start + time_chunk) + starts = start_indices[ext_start:ext_stop, + time_start:time_stop].reshape(-1) + fracs = (None if fractional_offsets is None else + fractional_offsets[ext_start:ext_stop, + time_start:time_stop].reshape(-1)) + width = time_stop - time_start + A_rows = xpy.repeat(A[ext_start:ext_stop], width, axis=0) + if xpy is np: + Q_one = _q_window_numpy_interp( + Q, starts, fracs, 1, time_interp, xpy=xpy)[:, 0, :] + values = np.einsum("ej,ej->e", A_rows, Q_one) + else: + values = _q_inner_product_gpu( + Q, A_rows, starts, fracs, 1, time_interp)[:, 0] + out[ext_start:ext_stop, time_start:time_stop] = values.reshape( + (ext_stop - ext_start, width)) + return out + + def _nearest_Q_window_numpy(Q_block, start_indices, npts, xpy=np): """Return nearest-grid Q windows with zero extension.""" npts_extrinsic = len(start_indices) @@ -2406,7 +2483,7 @@ def _nearest_Q_window_numpy(Q_block, start_indices, npts, xpy=np): return Qlms -def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDict, rholmsArrayDict, ctUArrayDict,ctVArrayDict,epochDict,Lmax=2,array_output=False,xpy=np, loglikelihood=_factored_lnL_helper,return_lnLt=False,phase_marginalization=False,n_cal=1,cal_method='loop',cal_distmarg=None,cal_log_weights=None,return_cal_components=False,time_interp='nearest',ctUArrayDict_cal=None,ctVArrayDict_cal=None): +def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDict, rholmsArrayDict, ctUArrayDict,ctVArrayDict,epochDict,Lmax=2,array_output=False,xpy=np, loglikelihood=_factored_lnL_helper,return_lnLt=False,phase_marginalization=False,n_cal=1,cal_method='loop',cal_distmarg=None,cal_log_weights=None,return_cal_components=False,time_interp='nearest',ctUArrayDict_cal=None,ctVArrayDict_cal=None,time_quadrature=None,explicit_time_values=False,return_time_draw=False,time_draw_uniforms=None): """ DiscreteFactoredLogLikelihoodViaArray uses the array-ized data structures to compute the log likelihood, either as an array vs time *or* marginalized in time. @@ -2480,6 +2557,38 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic --interpolate-time value maps to. Ask for a stencil explicitly if you want one. All three stencils have CPU and GPU implementations. See _sinc_Q_window_numpy and RIFT/likelihood/DESIGN_q_window_stencil.md for the measured tables. + + time_quadrature : {'simpson', 'bandlimited', 'peak-local'} or None + Rule used for the time integral. None (the default) defers to the + module-level ``TIME_QUADRATURE_DEFAULT``, which is 'simpson'. + + 'simpson' is the historical rule: Simpson's rule at the FIXED spacing + deltaT=1/srate. The integrand's width, however, is set by the signal -- + sigma_t = 1/(2 pi rho sigma_f) -- not by the data's sample rate, so this + under-resolves its own integrand at high SNR, and Simpson's (4T_h-T_2h)/3 + form makes an under-resolved peak worse than trapezoid rather than better. + + 'bandlimited' refines the grid to the integrand using ONLY the samples + already computed: kappa(t) is band-limited below Nyquist by construction + and rho_sq is time-independent on this path, so the continuous lnL(t) is + recovered exactly by a zero-padded FFT per row. The refinement factor is + DERIVED from the measured peak width and re-asserted on the refined grid; + it is not a settable accuracy knob. Restricted to n_cal==1 and to the + integrated outputs and continuous posterior draws (not return_lnLt / + return_cal_components); anything else raises rather than quietly falling + back. A time draw returns ``(time_offset, lnL_at_draw)`` and uses the + same validated dense representation as the integral. Rationale, measured + before/after and the exclusions: RIFT.likelihood.time_marginalization_quadrature. + + 'peak-local' is the same argument with the refined grid placed only where + the integrand has support: kappa's extrema are ENUMERATED on a small, + SNR-independent upsample, an interval of a few sigma_t is built around each, + overlapping intervals are MERGED into disjoint ones, and each is integrated + at its own derived spacing. The mass left outside is bounded per row and + checked, not assumed; a row whose bound is not small enough, or whose local + grid would cost more than the dense one, is given the 'bandlimited' value. + Same exclusions as 'bandlimited', PLUS phase marginalization, which it + refuses. RIFT.likelihood.time_marginalization_peak_local. """ global distMpcRef @@ -2487,6 +2596,60 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic if time_interp != 'nearest' and cal_method == 'fused': raise NotImplementedError("time_interp='{}' is not implemented for cal_method='fused'".format(time_interp)) + _time_quadrature_explicit = time_quadrature is not None + if time_quadrature is None: + time_quadrature = TIME_QUADRATURE_DEFAULT + time_quadrature_module.validate_time_quadrature(time_quadrature) + if return_time_draw and return_lnLt: + raise ValueError("return_time_draw and return_lnLt are mutually exclusive") + if return_time_draw and return_cal_components: + raise ValueError("return_time_draw and return_cal_components are mutually exclusive") + if return_time_draw and time_quadrature != 'bandlimited': + raise ValueError( + "return_time_draw requires time_quadrature='bandlimited'; the continuous " + "draw must share the quadrature's validated reconstruction") + if time_quadrature != 'simpson': + # Refuse loudly wherever the band-limited argument does not hold, rather + # than falling back to Simpson: a silently inert accuracy option is worse + # than an unavailable one. 'peak-local' rests on exactly the same argument + # as 'bandlimited' -- it only moves WHERE the refined grid is placed -- so it + # inherits the same exclusions rather than getting a parallel set. + if n_cal != 1: + raise NotImplementedError( + "time_quadrature=%r is not implemented for calibration " % time_quadrature + + "marginalization (n_cal=%d). The cal reduction sums exp() over " + "realizations, so each realization's kappa row must be refined and the " + "derived factor reconciled across them; that is untested." % n_cal) + if return_cal_components: + raise NotImplementedError( + "time_quadrature=%r is not implemented for return_cal_components, " + "which takes a per-realization time integral." % time_quadrature) + if time_quadrature == 'peak-local' and phase_marginalization: + # Refused BEFORE the integration runs, for the same reason the + # return_lnLt guard above is scoped the way it is: raising late means the + # whole run happens and then dies. Production marginalizes over + # distance, not phase; under phase marginalization the peak's Laplace + # width picks up an (I1/I0)(|kappa|/D) factor that does not reduce. + # 'bandlimited' supports it and is unchanged. + raise NotImplementedError( + "time_quadrature='peak-local' does not support phase marginalization " + "(the local width is no longer derivable from rho_sq and the curvature " + "alone). Use time_quadrature='bandlimited', which does.") + if return_lnLt and _time_quadrature_explicit: + # Explicitly ASKING for a quadrature on a call that takes no integral is + # a caller error and is refused. Merely INHERITING the module default is + # not: return_lnLt hands back lnL(t) on the original grid and never + # integrates, so the quadrature is inapplicable rather than ignored. + # Raising on the inherited default instead broke the group's standard + # extrinsic stage -- --add-extrinsic --add-extrinsic-time-resampling maps + # to --resample-time-marginalization, whose resample_samples() calls this + # function with return_lnLt=True and no explicit quadrature -- so enabling + # the option ran the whole integration and then died at the export step. + raise NotImplementedError( + "time_quadrature=%r was requested explicitly on a return_lnLt call, " + "which returns lnL(t) on the original grid and takes no time integral. " + "Drop the argument." % time_quadrature) + detectors = rholmsArrayDict.keys() npts = len(tvals) npts_extrinsic = len(P_vec.phi) @@ -2598,15 +2761,24 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic float(greenwich_mean_sidereal_time_tref), xpy=xpy ) - tfirst = t_det + tvals[0] - - sample_first = tfirst / deltaT - if time_interp == 'nearest': - ifirst = (xpy.rint(sample_first) + 0.5).astype(np.int32) # C uses 32 bit integers : be careful - frac_first = None + if explicit_time_values: + sample_at_times = ((t_det[:, None] + + xpy.asarray(tvals)[None, :]) / deltaT) + if time_interp == 'nearest': + ifirst = (xpy.rint(sample_at_times) + 0.5).astype(np.int32) + frac_first = None + else: + ifirst = xpy.floor(sample_at_times).astype(np.int32) + frac_first = (sample_at_times - xpy.floor(sample_at_times)).astype(np.float64) else: - ifirst = xpy.floor(sample_first).astype(np.int32) - frac_first = (sample_first - xpy.floor(sample_first)).astype(np.float64) + tfirst = t_det + tvals[0] + sample_first = tfirst / deltaT + if time_interp == 'nearest': + ifirst = (xpy.rint(sample_first) + 0.5).astype(np.int32) # C uses 32 bit integers : be careful + frac_first = None + else: + ifirst = xpy.floor(sample_first).astype(np.int32) + frac_first = (sample_first - xpy.floor(sample_first)).astype(np.float64) # ilast = ifirst + npts @@ -2707,25 +2879,39 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # Shape Q = (npts_time_full, nlms) # Shape A=FY_conj = (npts_extrinsic, nlms) # shape result = (npts_extrinsic, npts_time_*window* = npts) - Q_prod_result = _q_inner_product_gpu( - Q, FY_conj, ifirst, frac_first, npts, time_interp) + if explicit_time_values: + Q_prod_result = _q_inner_product_explicit_times( + Q, FY_conj, ifirst, frac_first, time_interp, xpy=xpy) + else: + Q_prod_result = _q_inner_product_gpu( + Q, FY_conj, ifirst, frac_first, npts, time_interp) else: # Use old code completely unchanged ... very wasteful on memory management! Q_block = rholmsArrayDict[det].T - Qlms = _q_window_numpy_interp(Q_block, ifirst, frac_first, npts, time_interp, - xpy=xpy) + if explicit_time_values: + Q_prod_result = _q_inner_product_explicit_times( + Q_block, np.conj(F_vec_dummy_lm * Ylms_vec), ifirst, + frac_first, time_interp, xpy=xpy) + Qlms = None + else: + Qlms = _q_window_numpy_interp(Q_block, ifirst, frac_first, npts, time_interp, + xpy=xpy) if phase_marginalization: + if explicit_time_values: + raise NotImplementedError( + "explicit time values with CPU phase marginalization are untested") Qlms[:, :, 1] = xpy.conj(Qlms[:, :, 1]) - FY_dummy_t = np.broadcast_to( - (F_vec_dummy_lm * Ylms_vec)[:, np.newaxis], - Qlms.shape, - ) + if not explicit_time_values: + FY_dummy_t = np.broadcast_to( + (F_vec_dummy_lm * Ylms_vec)[:, np.newaxis], + Qlms.shape, + ) - Q_prod_result = np.einsum( - "...i,...i", - np.conj(FY_dummy_t), Qlms, - ) + Q_prod_result = np.einsum( + "...i,...i", + np.conj(FY_dummy_t), Qlms, + ) kappa_sq += Q_prod_result * (distMpcRef/distMpc)[..., np.newaxis] else: @@ -2767,15 +2953,65 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic lnL_t = loglikelihood(kappa_sq.real, rho_sq_here) # Take exponential of the log likelihood in-place. - lnLmax = xpy.max(lnL_t) + # PER-ROW offset, not the batch max: lnL_t is (npts_extrinsic, npts_time), and a + # single scalar max shifts every extrinsic sample by the LOUDEST sample's peak. + # Any row more than ~745 nats below it then underflows exp() to 0 across the whole + # time axis -> L=0 -> lnL=-inf where the likelihood is finite. With lnL~rho^2/2 at + # the peak and ~0 for a typical prior draw that fires above rho~40 and takes out the + # BULK of the prior, collapsing mcsamplerAV. keepdims=True is load-bearing: with a + # bare axis=-1 the (n,) result broadcasts along the TIME axis instead, silently. + # See oshaughnessy-junior/research-projects-RIT#232. + lnLmax = xpy.max(lnL_t, axis=-1, keepdims=True) if return_lnLt: return lnL_t #- lnLmax # we want the verbatim lnL_t values, no shift + + if time_quadrature == 'bandlimited': + # Same integrand and the same closed domain [tvals[0], tvals[-1]] -- + # ONLY the resolution and the rule change, so a before/after difference + # is attributable to the quadrature and to nothing else. (The internal + # log-sum-exp offset does differ: it must be taken on the refined grid, + # whose maximum can exceed the coarse one by hundreds of nats. That is + # a numerical detail of an offset-invariant expression, not a second + # change of estimator.) + # Hand over THIS path's Simpson rule, not a private copy. On GPU + # that is optimized_gpu_tools.simps and on CPU it is scipy's, and the + # two are NOT interchangeable: the vendored GPU copy is an old scipy + # with even='avg' while modern scipy uses the Cartwright correction, + # so for EVEN npts (production is 614 at srate 4096) they disagree -- + # by 0.405 nats on an under-resolved peak. Rows that fall back must + # reproduce what the run they are in would have returned. Omitting + # this also made the module default to scipy, which RAISES on a cupy + # array: every --vectorized --gpu run of this option crashed. + _time_result = time_quadrature_module.time_marginalize_bandlimited( + kappa_sq, rho_sq_here, float(deltaT), loglikelihood, + phase_marginalization=phase_marginalization, simps=simps, + lnL_coarse=lnL_t, return_time_draw=return_time_draw, + draw_uniforms=time_draw_uniforms, t0=float(tvals[0]), xpy=xpy) + if return_time_draw: + _, _drawn_t, _drawn_lnL = _time_result + return _drawn_t, _drawn_lnL + return _time_result + + if time_quadrature == 'peak-local': + # Same integrand, same closed domain, same derived resolution criterion + # and the same fallback-to-Simpson row classification as 'bandlimited'; + # what changes is that the refined grid is placed only around the + # enumerated peaks instead of over the whole window. Rows this rule + # declines -- on its cost estimate, or because it could not bound the + # mass it left out -- are given the 'bandlimited' value, so the reviewed + # dense implementation is the backstop rather than Simpson. + return time_peak_local_module.time_marginalize_peak_local( + kappa_sq, rho_sq_here, float(deltaT), loglikelihood, + phase_marginalization=phase_marginalization, simps=simps, + lnL_coarse=lnL_t, xpy=xpy) + L_t = xpy.exp(lnL_t - lnLmax, out=lnL_t) L = simps(L_t, dx=deltaT, axis=-1) - # Compute log likelihood in-place. - lnL = lnLmax + xpy.log(L, out=L) + # Compute log likelihood in-place. lnLmax carries the kept trailing axis; drop it + # so the add-back lines up with L, which simps has already reduced over that axis. + lnL = lnLmax[..., 0] + xpy.log(L, out=L) return lnL @@ -2865,13 +3101,23 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic Q_block = Q_det[c*N_window_block:(c+1)*N_window_block] # (N_window, n_lms) ifirst_within = ifirst_det.astype(np.int32) if not (xpy is np): - Q_prod_result = _q_inner_product_gpu( - Q_block, FY_conj_det, ifirst_within, frac_first_det, npts, time_interp) + if explicit_time_values: + Q_prod_result = _q_inner_product_explicit_times( + Q_block, FY_conj_det, ifirst_within, frac_first_det, + time_interp, xpy=xpy) + else: + Q_prod_result = _q_inner_product_gpu( + Q_block, FY_conj_det, ifirst_within, frac_first_det, npts, time_interp) else: - Qlms = _q_window_numpy_interp(Q_block, ifirst_within, frac_first_det, npts, - time_interp, xpy=xpy) - # Q_det and FY_conj_det already encode any phase-marg conjugation - Q_prod_result = np.einsum("ej,etj->et", FY_conj_det, Qlms) + if explicit_time_values: + Q_prod_result = _q_inner_product_explicit_times( + Q_block, FY_conj_det, ifirst_within, frac_first_det, + time_interp, xpy=xpy) + else: + Qlms = _q_window_numpy_interp(Q_block, ifirst_within, frac_first_det, npts, + time_interp, xpy=xpy) + # Q_det and FY_conj_det already encode any phase-marg conjugation + Q_prod_result = np.einsum("ej,etj->et", FY_conj_det, Qlms) kappa_sq_c += Q_prod_result * invDistMpc[..., np.newaxis] # Fused-calmarg self-term fix: use this realization's rho_sq_c = @@ -2895,12 +3141,34 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # fold in this realization's importance log-weight lnL_t_c = lnL_t_c + cal_log_w[c] - m_c = xpy.max(lnL_t_c) + # PER-EXTRINSIC-SAMPLE offset, not the batch max. lnL_t_c is + # (npts_extrinsic, npts_time) and `running_max` is the streaming log-sum-exp + # offset for S. A SCALAR running_max shifts every extrinsic sample by the + # LOUDEST sample's peak, so any row more than ~745 nats below it underflows + # exp() to 0 across its whole time axis and every realization -> S row 0 -> + # lnL = -inf where the likelihood is finite. With lnL~rho^2/2 at the peak and + # ~0 for a typical prior draw that fires above rho~40 and takes out the BULK of + # the prior, collapsing mcsamplerAV. keepdims=True is load-bearing: a bare + # axis=-1 gives (n,), which broadcasts along the TIME axis instead -- silently, + # when npts_extrinsic == npts_time. Same defect and same fix as the n_cal==1 + # offset above, and as the return_cal_components branch a few lines up, which + # was already per-row. See oshaughnessy-junior/research-projects-RIT#232. + m_c = xpy.max(lnL_t_c, axis=-1, keepdims=True) # (npts_extrinsic, 1) + # A row that is -inf at every time (e.g. a distance-marginalization callback + # that rejects the whole row) has no finite offset of its own. Offset it by 0 + # instead: exp(-inf - 0) = 0 keeps its S row at 0 and its lnL at -inf, whereas + # exp(-inf - -inf) = nan would poison the running sum for that row for good. + # The scalar offset was shielded from this by any other finite row in the batch; + # a per-row offset is not, so the guard ships with the per-row offset. + m_c = xpy.where(xpy.isfinite(m_c), m_c, 0.0) if running_max is None: running_max = m_c - elif m_c > running_max: - S *= xpy.exp(running_max - m_c) - running_max = m_c + else: + # Elementwise: each row rescales S by its OWN change of offset (a no-op + # multiply by 1 for rows whose running max did not move). + new_max = xpy.maximum(running_max, m_c) + S *= xpy.exp(running_max - new_max) + running_max = new_max S += xpy.exp(lnL_t_c - running_max) if return_cal_components: @@ -2917,7 +3185,10 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic L = simps(S, dx=deltaT, axis=-1) # lnL = max + log( sum_c exp(log_w[c]) \int dt exp(lnL_t - max) ) - log(n_cal) - lnL = running_max + xpy.log(L) - cal_log_w_norm + # running_max carries the kept trailing axis; drop it so the add-back lines up with + # L, which simps has already reduced over that axis. (The return_lnLt branch above + # keeps the axis on purpose -- there S is still (npts_extrinsic, npts_time).) + lnL = running_max[..., 0] + xpy.log(L) - cal_log_w_norm return lnL diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_LISA.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_LISA.py index 57e3a880b..7b14cf3d7 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_LISA.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_LISA.py @@ -328,8 +328,12 @@ def FactoredLogLikelihoodAlignedSpinLISA(Q_lm, U_lm_pq, beta, lam, psi, inclinat factor = np.ones(modes.shape) factor[:,1] = -factor[:,1] negative_m_modes = modes * factor - spherical_harmonics = SphericalHarmonicsVectorized(modes, inclination, -phi_ref) - negative_m_harmonics = SphericalHarmonicsVectorized(negative_m_modes, inclination, -phi_ref) + # xpy=np is mandatory: this whole routine works in host numpy, while + # SphericalHarmonicsVectorized would otherwise pick cupy on any host where + # cupy imports. Same reason factored_likelihood_freqresponse.py passes + # xpy=np to TimeDelayFromEarthCenter. + spherical_harmonics = SphericalHarmonicsVectorized(modes, inclination, -phi_ref, xpy=np) + negative_m_harmonics = SphericalHarmonicsVectorized(negative_m_modes, inclination, -phi_ref, xpy=np) term_lm_conj_conjterm_lm__ = {} conjterm_lm_term_lm__conj = {} diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py index 4d12cc8f2..f9378ee45 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py @@ -338,6 +338,21 @@ def DiscreteFactoredLogLikelihoodFreqResponseNoLoop( array_output=True returns lnL_t of shape (npts_ex, npts); else time-marginalized. """ + # The band-limited time quadrature is REFUSED here, not silently ignored. Its + # correctness rests on lnL(t) being a pointwise function of a band-limited + # kappa(t) and a CONSTANT self-term; on this path rho_sq is time-dependent + # (and the antenna response carries sidereal post-phases), so the samples on + # the deltaT grid do NOT determine the continuous integrand and refining it + # by FFT would produce a confident wrong number. A caller that set the + # module-level default globally must be told, not quietly given Simpson. + from . import factored_likelihood as _fl + if getattr(_fl, 'TIME_QUADRATURE_DEFAULT', 'simpson') != 'simpson': + raise NotImplementedError( + "factored_likelihood.TIME_QUADRATURE_DEFAULT=%r, but the finite-size response likelihood has a " + "time-DEPENDENT rho_sq, so the band-limited argument does not apply here. " + "Audit this path separately rather than enabling the option globally." + % (_fl.TIME_QUADRATURE_DEFAULT,)) + import lal from . import factored_likelihood as FL on_gpu = not (xpy is np) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py index 981659438..c932e5b4e 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py @@ -786,6 +786,21 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( array_output=True returns lnL_t of shape (npts_ex, npts) (before time marginalization); array_output=False returns the time-marginalized lnL of shape (npts_ex,). """ + # The band-limited time quadrature is REFUSED here, not silently ignored. Its + # correctness rests on lnL(t) being a pointwise function of a band-limited + # kappa(t) and a CONSTANT self-term; on this path rho_sq is time-dependent + # (and the antenna response carries sidereal post-phases), so the samples on + # the deltaT grid do NOT determine the continuous integrand and refining it + # by FFT would produce a confident wrong number. A caller that set the + # module-level default globally must be told, not quietly given Simpson. + from . import factored_likelihood as _fl + if getattr(_fl, 'TIME_QUADRATURE_DEFAULT', 'simpson') != 'simpson': + raise NotImplementedError( + "factored_likelihood.TIME_QUADRATURE_DEFAULT=%r, but the slow-rotation likelihood has a " + "time-DEPENDENT rho_sq, so the band-limited argument does not apply here. " + "Audit this path separately rather than enabling the option globally." + % (_fl.TIME_QUADRATURE_DEFAULT,)) + require_post_phase_bank( meta, 'DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation') diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_gh_laplace.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_gh_laplace.md new file mode 100644 index 000000000..da36a3474 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_gh_laplace.md @@ -0,0 +1,310 @@ +# psi-marginal distance-node placement for `laplace` + `JAX_ILE_DISTMARG_GH` + +Branch `ghlaplace_o4d`, off `524331989ba7c135f2226ef4f33916684792c5a4`. +All numbers here are measured; the code carries the constants and a pointer. + +## The problem + +`core._distmarg_gh_logL` places frozen distance nodes at +`clip(K/R, x_min, x_max) +- 7/sqrt(R)` for a FIXED psi. On the `laplace` path +psi has already been integrated out analytically, so `K` and `R` do not exist +pointwise: the nodes must bracket the psi-MARGINAL distance integrand + + I(x) = (1/pi) int dpsi exp(x A(u) - x^2/2 B(u)), u = 2 psi + A(u) = A0 + Re(A1 e^{iu}), B(u) = B0 + Re(B1 e^{iu}) + Re(B2 e^{2iu}) + +which is a MIXTURE over u of Gaussians of centre `x*(u) = A(u)/B(u)` and width +`1/sqrt(B(u))`. The function used to `raise` rather than guess. + +## The rule that ships + + w = B0*A1 - conj(A1)*B2 + e^{iu*}= +- conj(w)/|w| (sign chosen so A(u*) > 0) + centre = stop_gradient(clip(A(u*)/B(u*), x_min, x_max)) + sigma = stop_gradient(min(1/sqrt(max(R_lo, 1e-30)), (x_max-x_min)/44)) + R_lo = B0 - |B1| - |B2| (== min_u B; see the identity below) + nodes = clip(centre + sigma * z, x_min, x_max), z = linspace(-22, +22, n) + +with `n = max(49, 1 + ceil((N-1) * 22/7))` for `JAX_ILE_DISTMARG_GH = N`, so the +node DENSITY the caller asked for at +-7 sigma is preserved, not diluted. +Trapezoid weights `0.5*(x[k+1]-x[k-1])` with the index clamped at both ends -- +algebraically identical to `_distmarg_gh_logL`'s `diff`/`concatenate` form, but +computable one block at a time so the distance axis stays scanned. +Gated on `m_max <= 2`; richer mode content still raises. + +`u*` is the EXACT maximiser of `A(u)^2/(2 B(u))`, not of `A(u)`. With +`A0 = B1 = 0` the stationary condition `2 A' B = A B'` reduces in `z = e^{iu}` +to `z^2 w = conj(w)`, so the maximiser is closed form and angle-free, and +reduces to `conj(A1)/|A1|` when `B2 = 0`. Verified against a 400,001-point +brute-force argmax on 20,000 random `(A1, B0, B2)` with `|B2|/B0` up to +0.99999: brute force never beats it by more than 6.5e-16 relative. + +## Task 1: why the closed form is enough (m_max = 2 only) + +Ladder-2 injection (35+30 Msun, H1/L1/V1, SEOBNRv4, `--l-max 2`), sky points +drawn from the measured rho-40.77 whole-sky AV posterior (the peer session's +draw, reproduced verbatim), 24 dense-phi x 16 sky x 614 time = 235,776 bins per +rung, psi ranked by the CLIPPED exponent `x_c A - x_c^2 B/2`, +`x_c = clip(A/B, x_min, x_max)`, u grid 721 points (du = 8.71e-3 rad). + +| quantity | rho 40.77 | rho 163.08 | pre-registered cut | branch taken | +|---|---|---|---|---| +| `W = sqrt(min_u B / R_lo)` | med 1.0000 max 1.0000 | med 1.0000 max 1.0000 | `<= 1.25` -> closed form as-is | closed form as-is | +| `C = |x*(u_cf)-x*(u_exact)|/sigma` | med 0.0014 p99 0.6361 max 0.6886 | med 0.0090 p99 0.0968 max 0.1117 | `<= 1` -> `u* = arg(A1)` adequate | closed-form centring | +| `S` (psi span, sigma) | med 2.607 p99 3.899 max 4.085 | med 0.641 p99 0.863 max 0.887 | half-width `(7+ceil(S_p99))` | 11 sigma, budget 12 | +| `R_lo <= 0` | 0.0000% of ALL 235,776 bins | 0.0000% | HARD REJECT if any | not triggered | + +(W, C, S over the bins within 100 nats of the best bin's clipped exponent; the +`R_lo <= 0` fraction over every bin, as the cut demands. Per-sample and global +weight masks give the same numbers to 3 dp.) + +Directly measured operational quantity -- the one-sided reach from the +closed-form centre to the furthest weight-carrying component centre, which is +what the half-width must actually cover: + +| rung | reach p99 | reach max | needed half-width `7 + reach` | +|---|---|---|---| +| 40.77 | 3.280 | 3.461 | 10.46 sigma | +| 163.08 | 0.762 | 0.772 | 7.77 sigma | +| 652.31 (spot check, 8 sky x 32 phi, nu 8192) | 4.510 | 4.510 | 11.51 sigma | + +so 12 sigma covers all three rungs with margin, and the pre-registered +`(7 + ceil(S_p99)) = 11` plus `C_p99 = 0.64` gives the same 12. + +**Control reproduction.** The peer session's committed record +(`analyses/va_rebuild_20260902/records/angle_coeff_structure.json`, paper repo +branch `claude/elated-merkle-c4dda4`) reports, from an INDEPENDENT harness that +goes through `anglemarg._reconstruct_field`, W inflation median 1.0000157 max +1.0000610 at rho 40.77 and `frac_R_lo_nonpositive` 0.0 at both rungs. Its +`W - 1` is pure u-grid discretization: its `bound_tightness_rel_median` is +8.53e-05, exactly `(1/2)|B''|(du/2)^2 / min B` at its du = 2pi/361. This +harness at du = 2pi/721 gives W - 1 = 4e-6 and at du = 2pi/16384 gives 1e-8 -- +the same identity seen at three resolutions. C and S are NOT in that record +(the file the coordinator named, `scripts/bracket_stats.py`, does not exist on +that branch or any other in the paper repo); the C/S values relayed +(med 0.00151 / p99 0.63555; med 2.607 / p99 3.899 / max 4.085) are reproduced +here to 3-4 significant figures, so whatever produced them agrees with this. + +## Why it is m_max = 2 ONLY + +`W == 1` is an IDENTITY, not a lucky bound. The spin-2 antenna response makes +A0 and B1 vanish identically for (2,+-2) content -- measured +`|A0|max/|A1|max = 6.7e-17`, `|B1|max/|B0|max = 5.6e-16` at rho 40.77 +(5.99e-17 / 5.53e-16 at rho 163.08), and independently 4.4e-17 / 4.5e-16 on the +synthetic fixture with random U/V, so it is algebra, not a property of this +injection. Then `B(u) = B0 + Re(B2 e^{2iu})` and `R_lo = B0 - |B2|` IS +`min_u B`, and `A(u)` is a pure first harmonic whose maximiser is available in +closed form. Neither statement survives odd-m or l >= 3 content, which is why +the ship gate keys on `m_max`, following `angle_sample_grid_sizes`'s precedent. +The higher-mode verdict is owned by another session; the gate stays until it +lands, whatever it says. + +## Beyond the fixture: an exhaustive scan of the reachable family + +Because `A0 == 0` and `B1 == 0` hold for EVERY mode set, after scaling +`B0 -> 1`, `sigma0 = 1/sqrt(B0)` and rotating `arg(A1) -> 0` the bracket problem +depends on exactly three numbers: + + rho = |A1|/sqrt(B0) r = |B2|/B0 in [0,1) delta = relative phase + +and `rho` enters only as an overall factor of the exponent, so the whole family +can be SCANNED rather than sampled. 222,950 points (r to 0.9999, rho 0.2 to +3000, 91 phases, v grid 262,144, 100-nat weight threshold), of which 214,849 +have >= 32 samples inside the weight-carrying window; the rest are grid-limited +and excluded. Clipping into `[x_min, x_max]` is 1-Lipschitz, so these unclipped +reaches are UPPER bounds on the clipped ones. + +| centring | reach p50 | reach p99 | reach MAX | +|---|---|---|---| +| `argmax A(u)` (the naive closed form) | 6.02 | 3136.7 | 10417.2 | +| `argmax A(u)^2/(2B(u))` (shipped) | 2.79 | 12.80 | **14.134** | + +The naive centring is catastrophic at large `|B2|/B0` -- and `|B2|/B0` is a +property of the network response, not of the mode content, so no `m_max` gate +would have caught it. The ladder-2 fixture never reaches that corner +(`|B2|/B0` median 0.011, p99 0.023 at rho 40.77 and 163.08), which is precisely +why a fixture-only validation would have shipped the wrong rule: this is the +"exact on a fixture, loose in production" defect class in its usual form. + +With the shipped centring the reach is bounded across every peak-exponent +decade: + +| peak exponent | rows | reach p99 | reach MAX | +|---|---|---|---| +| [0, 30) | 74,092 | 7.24 | 7.74 | +| [30, 100) | 13,463 | 13.86 | 14.11 | +| [100, 300) | 12,503 | 13.80 | 14.13 | +| [300, 1e3) | 13,727 | 12.05 | 13.82 | +| [1e3, 1e4) | 26,402 | 11.20 | 13.50 | +| [1e4, 1e5) | 26,282 | 10.62 | 13.06 | +| >= 1e5 | 48,380 | 10.11 | 12.65 | + +The maximum saturates at `sqrt(2T) = sqrt(200) = 14.142` -- attained in the +weak-signal corner, where the 100-nat window is the whole circle. Hence the +shipped half-span `7 + 14.14 -> 22 sigma`, and the node floor +`2*22/0.92 + 1 = 49`. The pre-registered `(7 + ceil(S_p99)) = 11` plus +`C_p99 = 0.64` gives 12 sigma, which the ladder-2 operating point needs; 22 is +the family-wide requirement and is what ships. Widening is the conservative +direction, and it costs nothing at the operating point: the answer at rho 40.77 +is bit-identical between 27 nodes at 12 sigma and 49 nodes at 22 sigma to the +six decimals printed. + +## Task 3: validation, in nats + +Ladder-2, `--data-integration-window-half 0.005` (npts 40), 4 sky points from +the same posterior draw, dense phi grid from `estimate_angle_amplitude` (the +production route, amp 1092.55 -> nphi_d 544, nu_d 272), max over sky points, +with the SHIPPED constants (exact-argmax centring, 22 sigma, 49-node floor): + +| comparison | rho 40.77 | rho 163.08 | +|---|---|---| +| laplace+GH16 vs exact+GH16 | 9.196e-05 | 5.128e-06 | +| laplace+GH33 vs exact+GH33 | 9.196e-05 | 5.128e-06 | +| laplace+GH65 vs exact+GH65 | 9.196e-05 | 5.128e-06 | +| laplace+GH129 vs exact+GH129 | 9.196e-05 | 5.128e-06 | +| laplace+GH65 vs laplace+uniform-4096 | 3.662e-04 | 1.481e-03 | +| laplace+GH16 vs laplace+GH129 (self-convergence) | 4.20e-09 | 1.273e-10 | +| laplace+GH33/65 vs laplace+GH129 | 0.0 | 0.0 | + +Every cell is measured with the shipped constants; nothing is carried over. +The rho 163.08 `laplace+uniform-4096` leg finished last and reproduced +1.4810625671e-03, the value the first-cut run gave, as it must -- that leg does +not use the adaptive nodes at all. The `laplace+GH` values at BOTH rungs are +identical to six decimals between the first cut (argmax A, 12 sigma, 27 nodes) +and what ships (exact argmax, 22 sigma, 49 nodes). + +The laplace-vs-exact residual is FLAT in node count, so it is the psi-Laplace +error alone, not the distance quadrature. The uniform-4096 residual is that +grid's own discretization: at rho 163 a 4096-point uniform grid over +[1, 10000] Mpc puts well under one point across the distance peak, so it is NOT +a converged reference there -- reported because it was asked for, not as truth. +The placement is converged at the 49-node floor (differences of 0.0 to the +printed precision from 102 nodes up), and the rho 40.77 answers are identical +to six decimals between the first cut (argmax A, 12 sigma, 27 nodes) and what +ships (exact argmax, 22 sigma, 49 nodes). + +### Gates + +* `test_angle_marg_gh_laplace.py`: 15 tests, all passing, wired into + `.travis/test-jax.sh` (`EXPECTED_TESTS` 189 -> 204). +* Full jax gate: **206 passed, 1 deselected, 11m34s**, collected 206 from 21 + files (204 floor). Interpreter `~/.conda/envs/rift_jax/bin/python` + (numpy 2.4.6, jax 0.9.2, numpyro 0.21.0), resolved by `test-jax.sh` through + `RIFT_JAX_PYTHON`. **That env has no pytest of its own**; pytest 9.1.1 was + supplied from a private `pip install --no-deps --target + pip --target` directory placed on `PYTHONPATH` (nothing + installed into the shared env). The gate CANNOT run on the CVMFS IGWN python + at all: `test-jax.sh` aborts at its `import numpyro` precheck, which CVMFS + lacks. So on this host the gate needs rift_jax for jax+numpyro AND an + external pytest -- that combination is the blocker a peer hit. +* `test_angle_marg_exact.py` (excluded from the gate for cost; the gate's own + comment says to run it by hand when touching `anglemarg.py`). **Every run + below names its interpreter, because the result depends on the numpy + version:** + + **That step is avoidable, and the gate has been run without it.** + `~/.conda/envs/gwkokab_stable` carries pytest 9.0.3 (a real `bin/pytest`), + jax/jaxlib 0.9.2, numpyro 0.21.0, numpy 2.4.6, flowMC, lal and lalsimulation + -- everything `test-jax.sh` prechecks. The full gate on the integration + branch ran there with NO `PYTHONPATH` additions: **219 passed, 1 deselected, + 0 failed, 0 errors** in 11m01s, collection floor 217, 23 files. + + A first attempt on that interpreter was OOM-killed 41 tests in (rc 137). The + 25 GiB cgroup is `user.slice/user-.slice` -- PER-UID, shared by every + session on the host -- so "run serially" bounds only your own contribution and + a neighbouring session can kill you. Exit 137 with no message is + indistinguishable from a code failure; `dmesg -T | grep -E 'oom-kill|Killed + process'` separates them. Recorded in the infra-atlas (lvk-cit) rather than + here, since it is a host fact and not a property of this code. + + | numpy / jax | interpreter | tree | result | + |---|---|---|---| + | 2.4.6 / 0.9.2 | `~/.conda/envs/rift_jax/bin/python` + private pytest 9.1.1 | base `52433198` | 2 failed, 33 passed | + | 2.4.6 / 0.9.2 | same | this branch | 2 failed, 33 passed | + | 1.26.4 / 0.7.1 | `/cvmfs/software.igwn.org/conda/envs/igwn/bin/python`, pytest 8.3.5 | base, before the shim | **6 failed**, 29 passed | + | 1.26.4 / 0.7.1 | same | this branch, before the shim | 6 failed, 29 passed | + | 1.26.4 / 0.7.1 | same | base kernel + shimmed reference | 3 failed, 32 passed* | + | 1.26.4 / 0.7.1 | same | this branch, shimmed | 2 failed, 33 passed | + + \* the third is this branch's own `test_laplace_refuses_gh_env_above_the_ + covered_mode_content`, run against base-commit code that still refuses GH + unconditionally -- an accidental positive control that the test discriminates. + + The four extra failures on numpy 1.26 were + `test_laplace_kernel_{error_law,first_harmonic_cancellation,randomized_sweep, + branch_window}`, all raising + `AttributeError: module 'numpy' has no attribute 'trapezoid'` from inside the + test file's OWN brute-force reference integrator (numpy renamed `trapz` -> + `trapezoid` in 2.0). They died before comparing anything, so they said + nothing about the laplace kernel; with the one-line shim RIFT's library + already uses (`misc/distance_grid.py`, `misc/distance_slices.py`) all four + PASS on numpy 1.26 against the unmodified base kernel. **There is no kernel + failure in either environment.** + + The two genuine pre-existing failures, present in every environment and on + the pristine base commit, are: + - `test_laplace_high_amplitude_accuracy_and_trend` -- its `errs[1] < errs[0]` + trend assertion compares 4.17e-10 against 5.68e-14, i.e. two numbers both + at machine precision; the docstring's reference values (0.055 / 0.028) are + stale by ~12 orders of magnitude, so the test now compares noise. Wants a + separate fix; NOT touched here. + - `test_driver_labels_a_suspect_angle_grid_in_provenance`. + +### Mutation results (every new guard) + +| guard | mutation | fired? | +|---|---|---| +| `_GH_PSI_HALF_SIGMA = 22` sufficient | raise to 44 -> answer must not move | yes, < 1e-4 nats | +| `_GH_PSI_HALF_SIGMA` is live at all | collapse to 0.05 -> answer must move | yes, > 1e-2 nats | +| `_GH_PSI_MIN_NODES = 49` | drop the floor to 3, request 2 nodes | yes, > 1e-3 nats | +| sigma cap is live | force it to bind (half-span 1e-9) | yes, > 1e-2 nats | +| sigma cap is inactive in production | compare against production support | yes (>100x margin) | +| `m_max` gate | m_max = 3 data raises; positive control: same data with GH off returns finite, and m_max = 2 data with GH on returns finite | yes | +| `JAX_ILE_DISTMARG_GH` is not inert on `laplace` | same call with GH on vs off must differ | yes, > 1e-6 nats | +| A0 == B1 == 0 identity | no mode set can break it, so plant the harmonics instead | yes | +| closed-form psi argmax | naive `argmax A` must be measurably worse on > 50% of random triples | yes | + +## What makes the placement admissible, and what does not + +The closed-form psi maximiser is DERIVED from `A0 == 0` and `B1 == 0`. Two +conditions gate it, and they do different jobs. + +**The guarantee is the response model, and it is angle-independent.** The static +path builds `F = F+ + i Fx` through `compute_detamresponse` (LAL's +`ComputeDetAMResponse`), where polarization enters as an exact rotation: + + F+(psi) + i Fx(psi) = (F+(0) + i Fx(0)) e^{-2 i psi} + +a SINGLE u-harmonic (u = 2 psi). `kappa` is linear in `F` and `rho^2` quadratic, +so `A` carries only u-harmonics +-1 and `B` only {0, +-2} -- at EVERY +(ra, dec, incl). That is why one probe angle is not being asked to establish a +global property: the property comes from the detector, not from the sample. + +`_GH_PSI_STATIC_FEATURES = (None,)`. **The only admitted value is the static +response, which is the ABSENCE of a feature tag.** Every NAMED feature is +refused -- `"rotation"` and `"freqresponse"` build their coefficients from the +arm vectors and a time-varying orientation and have no such factorization -- and +so is any value added later, which must be added to the tuple deliberately. +(Stated explicitly because the inverse reading is easy to fall into and was made +in correspondence: "unnamed" is the thing that PASSES, not the thing refused.) + +The precondition is enforced in the PUBLIC kernel, not only in the wrapper: +`feature` is a plain Python attribute, so it is static and trace-safe, and a +direct caller with a banded response would otherwise execute the unsupported +placement while the wrapper correctly refused it. + +**The numerical A0/B1 check is an implementation assertion, not the guarantee.** +It confirms the code matches the structure, pointwise per bin (a ratio of +unrelated global maxima hides a locally invalid bin behind a large denominator +elsewhere). It lives in the wrapper because it needs CONCRETE coefficient +tables: measuring it inside the kernel raises `TracerArrayConversionError` under +`jax.grad`, since the tables there are tracers. + +**Scope, so this is not over-read.** The identity and this gate concern the PSI +axis. Under precession the PHI content of the tables IS materially +redistributed -- measured on SEOBNRv5PHM, the `A` phi-slot-0 weight moves from +1.1e-16 aligned to 2.9e-2 precessing while staying band-limited to ~6e-15 -- and +that is a property of the SOURCE. It is not an identity failure and the gate is +right to admit it, but "the identity holds under precession" is a statement +about psi only. + diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_distance_quadrature.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_distance_quadrature.md new file mode 100644 index 000000000..79561fd53 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_distance_quadrature.md @@ -0,0 +1,796 @@ +# Distance quadrature for the dense angle-marginalization schemes + +**This is an accuracy change, not a speedup.** At the shipped tolerance, ON THE +REFERENCE CONFIGURATION (injected SNR 40), it uses ~46% MORE distance nodes than +the current default and is ~2000x more accurate. (It was ~3% until the spacing +was sized from the fail-safe's TRIP threshold rather than from `amp_sizing` +itself -- see section 1's `rho_max` note. The cost went up; the accuracy claim +did not change.) + +**That ~46% does not travel.** The count is derived from the run's own +amplitude, so it GROWS with it -- which is the scheme working as designed, not a +defect. At `rho_sampled = 163.08` the same contract derives n = 1558, i.e. ~6x +the default rather than 1.46x. Carrying "+46%" to a loud event is wrong by more +than a factor of four; see "quote a node-count error only with its amplitude +attached" in section 2a. +Matched to the current default's own accuracy it is 2.0-2.6x cheaper on the +distance axis, which is one axis of one kernel and does not rescue any campaign +(section 3). Read the cost section before citing a factor from this document. + +Scope: the `(x_grid, log_w_grid)` that +`fused_log_likelihood_distphipsimarg_{exact,laplace}` integrate distance over, +reached from `JAXDistPhiPsiMargLikelihood` / `--mode flowmc-phipsimarg`. + +Everything below was measured on `origin/rift_O4d` @ `52433198` plus this +branch, on the ladder-2 reference configuration: SEOBNRv4, 35+30 Msun, +H1L1V1, `--srate 4096` (npts = 614), `--l-max 2`, injected SNR 40 at +d = 633.92 Mpc, prior range `[--d-min 1, --d-max 10000]` Mpc, `interp="sinc"`. +Environment: `~/.conda/envs/rift_jax` (python 3.13, jax 0.9.2), `JAX_ENABLE_X64=1`, +`OMP_NUM_THREADS=1`, `taskset -c 0-15`, `JAX_PLATFORMS=cpu`, `RIFT.__file__` +asserted into the tree under test by every script. + +--- + +## 1. The contract, before and after + +**Before (still the default).** `make_distance_grid` lays `--distance-grid-points` +nodes **uniformly in d** across the whole prior range and gives each node the +same interval `dd = (d_max - d_min)/(n-1)`. What that guarantees is a *fixed* +discretization, independent of the data: the same grid for every event, and a +quadrature error that nobody has to think about because it never changes. What +it does **not** guarantee is that the error is small. The distance integrand's +width is set by the data (see below), and the shipped 256-node grid's error at +this operating point is 0.17-0.22 nats, not a design target anyone chose. + +**After (opt-in, `--distance-grid-scheme loguniform`).** Nodes are laid +**uniformly in ln d** with a count derived from the run's own data: + + Delta(ln d) <= c(tol) / rho_max, c(tol) = pi * sqrt(2 / ln(2/tol)) + n = ceil(rho_max * ln(d_max/d_min) / c) + 1 + +The guarantee changes from "a fixed grid" to "**a stated fractional error on +the distance integral, uniformly over the prior range and over every angle +sample -- PROVIDED the maximizing distance is interior to that range**". The +error is no longer data-independent; it is *bounded* by a number the caller +states, conditional on a sky-sweep estimator (section 1's `rho_max` +discussion) and on that interiority precondition, which is checked at build +time and REFUSED when violated (section 1a). Both conditions are stated here +because neither is an inequality anyone has proved. + +### Why relative spacing is the right invariant + +Per angle sample and time bin the distance integrand is + + exp(K x - 0.5 R x^2), x = distMpcRef / d, + K = Re, R = at the reference distance + +a Gaussian in `x` peaked at `x* = K/R` with standard deviation `1/sqrt(R)`. +Its **relative** width is + + sigma / x* = sqrt(R)/K = 1 / rho, rho = K / sqrt(R) = that sample's matched SNR + +which does not depend on where the peak sits. So one relative spacing resolves +every peak everywhere, with no peak location entering anywhere. Uniform-in-d +spacing has the opposite property: it over-resolves large d (where the +likelihood is flat) and starves small d (where the peaks are narrowest in +absolute terms). + +### Why `c(tol)` is derived rather than tuned + +The trapezoid rule on a Gaussian converges super-algebraically. By Poisson +summation, for `f = exp(-(u-mu)^2/2s^2)` sampled at spacing `h`, + + (h * sum_k f(u_k)) / integral(f) = 1 + 2 exp(-2 pi^2 s^2 / h^2) cos(2 pi mu / h) + ... + +so the worst-case fractional error over the peak phase `mu` is +`2 exp(-2 pi^2 s^2 / h^2)`. Setting that to `tol` and writing `h = c * s` gives +`c = pi * sqrt(2/ln(2/tol))`. `test_tolerance_constant_matches_the_gaussian_trapezoid_error_law` +pins this **two-sided**: a one-sided "error <= tol" check is satisfied by any +smaller `c`, including `c -> 0`, which is perfectly accurate and arbitrarily +expensive. + +### Where `rho_max` comes from -- and why it needs no new estimator + +`anglemarg.estimate_angle_amplitude` returns, by its own construction, +`ANGLE_AMP_MARGIN` times the maximum over a SAMPLED set of angles, and over the +distance support, of `x A - 0.5 x^2 B`, whose closed-form maximum in `x` is +`A^2/(2B) = rho^2/2`. So + + rho_max = sqrt(2 * A) = sqrt(ANGLE_AMP_MARGIN) * rho_sampled_max + = sqrt(2) * rho_sampled_max (measured: exactly 1.4142 at rho 25/40/80) + +**This is not an identity and `rho_max` is not a proven bound.** It is +`sqrt(2)` times the largest matched SNR a finite sky sweep found, and that +estimator's own docstring says so: *"THIS IS AN ESTIMATOR, NOT A PROVEN +BOUND"*. Two consequences worth stating plainly: + +* the shipped node count is **41% above what the stated tolerance requires** -- + the reference configuration's 373 nodes meet the `tol = 1e-2` contract at + ~187. That margin is deliberate, and it is earning its keep: on this + configuration the sweep's own empirical maximum was 758.17 against the + injection's 800, i.e. it DID under-read the true maximum, and the margin + absorbed it; +* the contract is therefore conditional on the sweep, not on an inequality. It + is stated that way here and should be read that way. + +Three consequences, all of them the point of this design: + +* no peak is located, so there is no peak estimate to be wrong; +* the number that sizes the distance grid is the **same** number that sizes the + dense angle lattice, so the two cannot disagree; +* the fused kernels' existing runtime fail-safe (`_runtime_amp_failsafe`) + recomputes `A` from the coefficient tables on **every** likelihood call and + warns if the shipped sizing was exceeded. That fail-safe therefore now + covers the distance grid too, at no extra cost. The residual failure mode is + named in §5. + +`A` here is `amp_sizing` -- the amplitude **floored at +`ANGLE_MARG_CROSSOVER_AMPLITUDE = 450`** -- and not the unfloored `amp_data`. +That is deliberate and it is what makes the fail-safe coverage in section 5 an +identity rather than an approximation: `_runtime_amp_failsafe` compares the +per-call amplitude against `amp_sizing`, so sizing the distance spacing from +anything else leaves a gap. Concretely, on a quiet target (`amp_data < 450`) a +runtime amplitude anywhere between `amp_data` and 450 would under-resolve the +distance peak and trip nothing. The floor costs a minimum of ~144 nodes on an +event whose run is cheap anyway. Scheme *selection* still uses the unfloored +`amp_data`, unchanged, so quiet targets still take the `exact` branch. + +### 1a. Where the contract does NOT hold, and what happens instead + +The derivation above assumes the integrand is a Gaussian **peak inside the +support**. When the maximizing distance `x* = A/B` is EXTERIOR to +`[x_min, x_max]` -- the distance posterior rails against a prior edge -- the +integrand is monotone on the support instead: a boundary layer at `d_max` (or +`d_min`). The log-uniform grid is the wrong instrument for that, twice over: + +* its ABSOLUTE spacing is coarsest exactly at `d_max`, where the layer sits; +* refining it adds nodes proportionally *everywhere*, so the layer never + resolves. Measured on the reference configuration with the peak pushed to + `d* = 20000` Mpc against a `[1, 10000]` Mpc prior and a 2^21-node reference, + tightening `tol` from 0.5 to 1e-9 moves the error only **5.23 -> 3.92 nats**, + while uniform 256 -> 4096 moves **2.52 -> 0.36**. + +And the node count moves the WRONG WAY, because `estimate_angle_amplitude` +maximizes at `clip(A/B, x_min, x_max)`: an exterior `A/B` makes the clipped +value UNDER-read `rho^2/2`, and in the extreme it returns exactly 0, at which +point the crossover floor pins `rho_max = 30` and the grid collapses to 145 +nodes over `[1, 10^4]` Mpc. Measured: + +| rho | d* (Mpc) | support | n at tol=1e-2 | log-uniform err | uniform-256 err | +|---|---|---|---|---|---| +| 40 | 12000 | [1, 10000] | 271 | **+1.89 nats** | +0.67 | +| 40 | 20000 | [1, 10000] | **145** (amp -> 0) | **+4.60** | +2.53 | +| 40 | 3000 | [100, 2000] | 89 | **+3.01** | +1.51 | +| 40 | 634 (interior control) | [1, 10000] | 271 | +5.6e-05 | +0.049 | + +So at the UPPER edge the scheme is **1.3-3x WORSE in nats than the default it +replaces**. It is REFUSED at build time, not silently mis-sized and not +silently fallen back to uniform. The detector is a single scalar from the same +sky sweep: the amplitude recomputed WITHOUT the clip (`A^2/(2B)`, the true +stationary value) against the clipped one. `clip_excess > 1 + 1e-3` means the +maximizer is exterior. Verified to fire on truncated supports and to stay at +exactly 1.0 on interior ones. + +**The refusal is symmetric in the two prior edges; the harm is NOT, and the +table above measures only one of them.** Every row of it puts the maximizer +ABOVE `d_max`, which is where the mechanism argument applies -- the log grid's +absolute spacing is coarsest at `d_max`, so a layer there is the worst case for +it. At the LOWER edge the same grid is *finest* exactly where the layer sits, +and the sign reverses. Measured (external re-review, 2026-09-02) on a +`rho_max = 30` target with the maximizer at 86 Mpc, sweeping `d_min` past it +with `d_max = 10^4` Mpc. **The reference is two references**: a uniform +65536-node grid and a log-uniform grid at `tol = 1e-10`, and their mutual +disagreement is reported as a column rather than assumed away -- an earlier +draft of this table used a single uniform-8192 reference and argued its +convergence from the fact that the log grid agreed with it, which is a +common-mode argument and was ~25% low. + +| `d_min` | `clip_excess` | verdict | n (log) | refs disagree | log-uniform err | uniform-256 err | +|---|---|---|---|---|---|---| +| 86 | 1.0 | accepted | 75 | 3.68e-5 | 1.85e-4 | 1.41e-3 | +| 92 | 1.0044 | REFUSED | 74 | 3.67e-5 | 1.85e-4 | 1.41e-3 | +| 120 | 1.0879 | REFUSED | 70 | 3.69e-5 | 1.84e-4 | 1.41e-3 | +| 400 | 2.5579 | REFUSED | 52 | 3.79e-5 | 1.81e-4 | 1.36e-3 | + +**Read this to two significant figures and no further.** The two references +disagree by 3.7e-5, which is ~20% of the log-uniform error being measured, so +that column is good to about that and no better. What survives the uncertainty +is the ratio and its stability: log-uniform is **~7.6x more accurate than the +uniform-256 default** here, and both errors are flat to a few percent while the +support shrinks 4.6x and the node count falls 75 -> 52. An independent +reconstruction of this measurement by a second reviewer reproduced the +DIRECTION (log-uniform better at the lower edge, by 2.3x-540x on their +`(A, B)`) but not the flatness; their fixture is not this one, and the +disagreement is unresolved. Treat the direction as established and the +magnitude as fixture-specific. + +We refuse both edges anyway, and that is a deliberate choice rather than an +oversight: the CONTRACT is what fails once the maximizer leaves the support -- +`c(tol)/rho_max` is derived from a Gaussian peak's relative width, and there is +no peak on the support to have a width -- so the stated fractional error is not +being delivered even where the realised error happens to be small. Refusing on +the condition we can actually detect (`clip_excess`) rather than on a realised +error we cannot compute at build time keeps the option's promise honest in both +directions. The cost is stated here rather than hidden: at the lower edge the +refusal sends the caller back to a grid that is measurably worse. Recourse is +the same one the message names -- narrow `--d-min` so the posterior is interior. + +**What `tol` actually delivers near an edge, measured.** The endpoint guard +budgets the whole of `tol` to the endpoint term while the alias law has already +spent it: `c(tol)` is derived from `2 exp(-2 pi^2 / c^2) = tol`, so at the +sizing peak the alias error IS `tol` by construction. The two add. Measured on +the scheme's own quadrature in pure numpy (uniform in `ln d`, half-width end +intervals, against a 2e6-node reference): + +| clearance (peak widths) | measured error | guard | +|---|---|---| +| 1.0 | 0.113 | refuses | +| 2.93 (the reported minimum) | 0.0132 | accepts | +| 30 (deep interior) | 0.0093 | accepts | + +So deep in the interior the contract is met (0.0093 against 0.01), and **at the +acceptance boundary for a peak at the sizing SNR the delivered error is ~1.3x +`tol`, not `tol`**. That is stated rather than fixed: making the comparison +`alias + endpoint <= tol` would refuse EVERY loud peak, because the alias term +alone already equals `tol` at `rho = rho_max`. The honest bound is that the +option delivers `tol` deep in the interior and up to ~2x `tol` (the guard's own +`ENDPOINT_ERROR_MARGIN`) in the last few widths before a refusal. Quieter peaks +sit on a proportionally finer grid and are covered with room to spare -- measured +0.0053 against a 0.0081 estimate at `rho_pk = 10.3` on a grid sized for 42.4. + +**The non-positive-clearance window, and why it is a separate refusal.** +`_endpoint_bell` clamps `k <= 0` to zero, on the correct observation that an +endpoint sitting ON the peak is a stationary point -- measured error 1.2e-9 at +exactly `k = 0`. But the clamp cannot distinguish that from `k < 0`, where the +peak has left the support and the error climbs steeply, so the term scored its +worst case as its safest. Between that and the `clip_excess` trip lay a window +neither guard saw: measured, clearance -0.656 built with a true error of +**0.0668, 6.7x `tol`**, while `clip_excess` still read 1.0. A non-positive +clearance is now refused outright. Measured transition, monotone with no +building window: + +| clearance | true error | verdict | +|---|---|---| +| 6.97 | -- | builds | +| 2.07 | 0.0074 | refused (endpoint) | +| -0.66 | 0.0668 | refused (clearance) | +| -2.40 | 0.485 | refused (exterior) | + +**Why refuse rather than fall back to uniform.** Three reasons. A fallback +would make `--distance-grid-scheme loguniform` silently produce the *other* +scheme's grid -- the silently-inert-flag class this module keeps being bitten +by, and the same defect as the `JAX_ILE_DISTMARG_GH` combination refused +alongside it. Both grids are bad here anyway: uniform 256 is itself 2.5 nats +out, so a fallback substitutes one wrong answer for another rather than fixing +anything. And the regime is a *physics* signal -- the posterior is railing +against `--d-max` -- which the caller should see rather than have papered over. +The refusal names the recourse: widen `--d-max` (or narrow `--d-min`) so the +posterior is interior, or stay on the uniform default and raise +`--distance-grid-points`. + +Residual limitation, stated: the detector reads the same finite sky sweep as +the amplitude. If that sweep misses an exterior-peak angle configuration +entirely, the build proceeds. Nothing here bounds that. + +### The decoupling property (the most important safety property here) + +`estimate_angle_amplitude` reads only `x_grid.min()` and `x_grid.max()`: the +per-angle distance maximum is closed form at `clip(A/B, x_min, x_max)`, and +`A/B` is interior to any window that contains it. + +**CORRECTION, and it weakens this section: no in-tree scheme can currently +trigger the coupling.** `make_distance_grid_adaptive` always concatenates a +full-range `linspace` backbone (`coarse = np.linspace(d_min, d_max, n_coarse)`) +before dedup, so it spans the whole support and returns an identical amplitude; +and both grids this PR ships span the full support by construction. The +`[0.8 d, 1.25 d]` window in the table below is a hand-built diagnostic, not a +grid any code path produces. So the decoupling change is **prospective +insurance**, not an active protection, and the mutation that reverses it +survives every value-level assertion for exactly that reason (see the PR's +mutation matrix, M12). It is still worth having -- it makes the invariant +structural rather than incidental -- but it should not be sold as fixing a live +bug. Measured on this configuration (`amp_vs_distgrid.py`, same directory): + +| distance grid | n | amp | dense (n_phi, n_u) | +|---|---|---|---| +| shipped `[1, 10000]` | 256 | 1516.33 | 624, 320 | +| `[1, 10000]` | 64 | 1516.33 | 624, 320 | +| `[1, 10000]` | 24 | 1516.33 | 624, 320 | +| `[0.5d, 2d]` | 24 | 1516.33 | 624, 320 | +| `[0.8d, 1.25d]` | 24 | **1325.78** | **592, 304** | + +Identical to every printed digit until the window stops containing `A/B`, and +then the **angle lattice silently shrinks**. This branch removes that coupling +by construction rather than bounding it: the amplitude is computed on a +full-support uniform grid built for that purpose, never on the grid the +likelihood integrates over. For `--distance-grid-scheme uniform` the two are +the same object, so the default path is unchanged node for node. + +--- + +## 2. Accuracy, in nats + +### 2a. The distance quadrature in isolation + +`dist_quad_error2.py` (measurement scripts: RIFT_roboto_paper `analyses/jax_anglemarg_exec_cost/`). The dense kernels integrate +distance as a plain log-sum-exp over `(x_grid, log_w_grid)` of +`exp(x A - 0.5 x^2 B)`; with `A, B` replaced by the `(K, R)` that +`core._accumulate_unit` returns, the identical quadrature can be evaluated on +the real precompute at negligible cost, for 256 angle samples x 614 time bins, +against a **65536-node uniform reference** (self-converged: the 32768 -> 65536 +step moves the result by 2.3e-5 nats). The reported quantity is the per-sample +value after the time reduction, `L_s = logsumexp_t lnZ_d(s,t)` -- the number the +sampler consumes -- and `dlnZ`, the cloud evidence proxy. + +Prior-draw cloud (S = 256): + +| grid | n | max abs dL_s | dlnZ (nats) | +|---|---|---|---| +| uniform 24 ("the 10.7x") | 24 | **27.4** | **+1.57** | +| uniform 64 | 64 | 4.50 | +0.238 | +| uniform 128 | 128 | 0.640 | -0.104 | +| **uniform 256 (shipped default)** | 256 | **0.170** | **-0.170** | +| uniform 512 | 512 | 0.0029 | -0.0029 | +| log-uniform 96 | 96 | 0.208 | +0.184 | +| log-uniform 128 | 128 | 0.0415 | -0.026 | +| log-uniform 160 | 160 | 0.0061 | -0.0061 | +| log-uniform 192 | 192 | 0.0030 | +8.9e-6 | +| log-uniform 256 | 256 | 0.0017 | +2.4e-5 | +| `make_distance_grid_adaptive` (in tree) | 144 | **9.44** | -0.312 | + +**Quote a node-count error only with its AMPLITUDE attached.** Every row in +these two tables is at the header's operating point (injected SNR 40). The +distance peak narrows as `1/rho` against a grid fixed by the PRIOR range, so a +256-node uniform grid is a different instrument at a different amplitude: the +same family measures 0.170-0.216 nats here and 43.16 nats at `rho = 163.08` +(paper-1 ladder, `snr160_laplace_gh0_dg256`). Same rule, opposite sides of the +threshold -- not a discrepancy. + +It is NOT a conversion, and two sessions have now been tempted to use it as one. +Matching the two series on points-per-peak-width does not collapse them: on +`N/rho` the louder series is 9.2-17.0x worse and roughly flat; on `N/rho^2` (the +better-motivated scaling, since a grid uniform in `d` over a fixed range +resolves a peak of width `d*/rho`, and `d* ~ 1/rho` when amplitude is set by +injected distance) it is 0.67x, 0.52x, 0.06x -- overshooting and not flat. The +amplitude threshold is the dominant effect and explains the sign and most of the +size; the residual is real, not constant, and carries the rest of the +configuration (seglen, deltaF, probe points against a cloud max). So compare +grids WITHIN one operating point and re-measure across them. + +For scale at the loud end: `rho_sampled = 163.08` derives +`rho_max = sqrt(ANGLE_AMP_MARGIN * AMP_FAILSAFE_TRIP_FACTOR) * rho = 326.16` and +**n = 1558** at `tol = 1e-2` over `[1, 10^4]` Mpc -- against 4096 linear nodes +for that rung's converged reference, so under a factor of three, not the ninth a +transfer of the SNR-40 count (373) would suggest. + +Cloud concentrated near the injection (S = 256): + +| grid | n | max abs dL_s | dlnZ (nats) | +|---|---|---|---| +| uniform 24 | 24 | **73.0** | **-34.0** | +| uniform 64 | 64 | 9.04 | -2.70 | +| uniform 128 | 128 | 2.30 | -0.594 | +| **uniform 256 (shipped default)** | 256 | **0.216** | -0.022 | +| uniform 512 | 512 | 0.0032 | -0.0029 | +| log-uniform 128 | 128 | 0.184 | +0.158 | +| log-uniform 160 | 160 | 0.040 | -0.033 | +| log-uniform 192 | 192 | 0.0066 | -6.3e-4 | +| log-uniform 256 | 256 | 1.0e-4 | -1.0e-4 | +| `make_distance_grid_adaptive` (in tree) | 144 | **22.6** | -1.73 | + +Independently confirmed by the coordinator with a different reference (uniform +`n_grid = 8192`), a different cloud (the exported 4800-row ladder-2 cloud) and +the **`grid`** angle scheme, so the distance path is isolated from the dense +machinery entirely. That sweep gives, against its own reference: `n_grid = 24` +mean +34.64 / max 56.05 / lnZ +50.86 nats; `n_grid = 128` mean +0.32 / max 3.05; +`n_grid = 256` (shipped) mean +0.010 / max 0.326 / lnZ +0.013; `n_grid = 512` +mean +0.0028. Two references, two clouds and two angle schemes sharing no +estimator agree on the shipped default's max error (0.326 against 0.216 here) +and on the catastrophe at 24 nodes. + +Read off: + +* **The shipped default's error is 0.17-0.22 nats**, not a chosen tolerance. +* **Equal accuracy to the shipped default is reached at n ~ 100-128 log-uniform + nodes: a 2.0-2.6x reduction on the distance axis.** +* At the shipped tolerance default `tol = 1e-2`, at the reference + configuration's amplitude, the derived count is **n = 373** -- ~46% *more* + nodes than the default -- for **1.0e-4 nats instead of 0.216**, i.e. ~2000x + more accurate at ~1.46x the distance-axis cost. Both the count and the ratio + are amplitude-dependent: n = 1558 (~6x the default) at `rho_sampled = 163.08`. +* **There is no 10.7x at fixed accuracy.** `n_grid = 24` costs 27-73 nats of + lnL and 1.6-34 nats of evidence. See §4. + +### 2a'. What these numbers do and do not establish + +* **One operating point, one event, zero noise.** `load_injection` builds + `data_dict[det] = non_herm_hoff(P)` -- pure signal, no noise realization -- + with `SimNoisePSDaLIGOZeroDetHighPower` for all three detectors including V1. + So these are quadrature-error measurements on a clean signal, not a + population statement. Nothing here is averaged over noise realizations, + masses, or sky positions. +* **The "equal accuracy at n ~ 100-128" figure is empirical, not the + contract.** It is where the measured error happens to match the shipped + grid's on THIS configuration. The shipped default `tol = 1e-2` deliberately + sizes above it (n = 373), because the option's selling point is a bound that + holds without this measurement, not the smallest number that passed it. +* **The error metric is the per-sample lnL after the time reduction, plus a + cloud evidence proxy.** It is not a posterior-level statement: a bias that is + constant across the cloud cancels in the posterior and not in the evidence, + and this metric reports both. +* The reference is a uniform grid, i.e. the SAME quadrature family refined -- + a common-mode assumption. It is defensible here because the integrand is a + smooth 1-D Gaussian-times-polynomial with no singular structure, and because + the two grid families (uniform and log-uniform) agree with each other to + 1e-4 nats at their converged ends, which a shared systematic would not + produce. + +### 2b. The same thing through the real `laplace` kernel + +`e2e_laplace.py` drives +`fused_log_likelihood_distphipsimarg_laplace` itself with `amp_sizing` held +fixed across arms, so the dense angle lattice is bit-identical and the distance +grid is the only variable. Numbers in §2c of the PR body. + +--- + +## 3. Cost + +Cost on the distance axis is linear in the node count: the laplace kernel scans +`ceil(n/dist_block)` blocks of a fixed kernel (#209), and the exact kernel calls +`_logsumexp_grid_blocked` over the same nodes. Timing method: interleaved arms +inside each replicate, order flipped between replicates, explicit +`block_until_ready()` on every arm, minimum and full spread reported -- a +sequential A/B on a quiet host has overstated a speedup by ~50% on this code +before. Numbers in the PR body. + +**There is no speedup to quote at the recommended operating point**, and the +matched-accuracy one is small. At `tol = 1e-2` and the reference amplitude the +derived count is 373 against the shipped 256: ~46% more distance work -- and ~6x +more at `rho_sampled = 163.08`, since the count follows the amplitude. Matched to the shipped grid's own +accuracy the count is ~100-128, i.e. 2.0-2.6x less work **on the distance axis +only** — the dense `(phi, u)` lattice is unchanged by construction. Combined +with the other measured lever (per-distance-block phi sizing: 1.45x standalone, +1.11x after this one) the total reaches ~2.9-4x against a campaign requirement +of order 140x, so this does not rescue that campaign and must not be cited as +if it might. + +### Compile and execute, separated, at production `npts` + +`compile_vs_execute.py` times the two phases apart using jax's explicit +lowering API (`jax.jit(f).lower(...)` then `.compile()`), so the compile is not +hidden inside a first call. `npts = 614`, S = 64 sky samples, `amp_sizing` +pinned so the dense angle lattice is identical across arms. RTX PRO 4000 +Blackwell. + +| arm | n | trace | **compile** | **first execute** | +|---|---|---|---|---| +| uniform 256 (shipped) | 256 | 1.76 s | **29.16 s** | **421.4 s** | +| log-uniform `tol = 0.5` | 136 | 2.31 s | **28.00 s** | **243.4 s** | + +Interleaved execute timings, order flipped between replicates, explicit +`block_until_ready()` on every arm, 5 replicates: + +| arm | n | compile | exec min | exec med | exec max | +|---|---|---|---|---|---| +| uniform 256 (shipped) | 256 | 29.16 s | 419.993 s | 420.100 s | 420.165 s | +| log-uniform `tol = 0.5` | 136 | 28.00 s | 243.132 s | 243.138 s | 243.245 s | + +**Matched-accuracy execute speedup: 1.727x** for a 1.88x node reduction. The +run-to-run spread is 0.17 s on a 420 s arm (0.04%), and the sequential +first-execute pair measured before the replicate loop gave 1.732x -- so the +sequential-vs-interleaved concern that has bitten this code before did NOT bite +here, and that is a measured statement rather than an assumption. Fitting the +two arms plus the second GPU's n = 1024 point gives +`t_exec = 41.5 s + 1.484 s/node`: linear in the node count with a small fixed +term, which is why the wall-clock ratio (1.73x) sits just under the node ratio +(1.88x). The fit predicts 1561 s at n = 1024 against 1685 s measured on the +other card (6%). + +The two arms differ by **0.597 nats** at their largest over the 64-sample prior +cloud. That is an arm-to-arm DIFFERENCE, not either arm's error: section 2a +puts the shipped uniform 256 at 0.17-0.33 nats from a converged reference and a +136-node log-uniform grid at a comparable distance from it, so a difference of +this size is what two independently-wrong approximations should show. It is +also why `tol = 0.5` is the loose end of the option and not what this document +recommends -- at the shipped `tol = 1e-2` the count is 373 and the error is +~1e-4 nats, and there is no speedup at all. + +The practical number: **~6.6 s per sky sample at `npts = 614`, n = 256**. That +predicts ~94 minutes for an 853-sample chunk, which is consistent with a +separate `laplace` job observed spending over 78 minutes of GPU at 99% without +completing one such chunk — and it is ~2.9x slower than the rate #210 published. +**#209 and #210 were both benchmarked at `npts = 64`**, about a tenth of the +production value on the axis that drives this cost. That discrepancy is a +finding about the published cost tables, not about this change, and it deserves +its own issue. + +Two scaling consequences that follow from the node count +`n = ceil(rho_max * ln(d_max/d_min) / c) + 1`, and that a reader should have in +front of them before choosing this scheme: + +* **Cost is linear in `rho_max`, i.e. in SNR.** At the reference SNR 40, + `rho_max = 77.88` (= sqrt(TRIP) x 55.07) and `n = 373` at `tol = 1e-2`. At + the top of the cost bake-off ladder (SNR 640) `rho_max ~ 905` and `n ~ 4300`, + ~17x the shipped grid. That is not a regression introduced here: a uniform + 256-node grid at that SNR is wrong by far more than the numbers in section 2 + (its spacing is ~80 sigma). It does mean **the log-uniform grid is a good + deal at moderate SNR and an expensive one at extreme SNR**, where the + per-sample quadrature of section 4(d) is the right answer instead. The + derived node count is printed by the driver, so the cost is visible before the + run rather than after it. +* **Cost is linear in `ln(d_max/d_min)`.** The driver defaults are + `[1, 10000]` Mpc, `ln = 9.21`, and roughly half of that -- everything below + 100 Mpc -- carries essentially no posterior for a 35+30 source. Narrowing an + unphysically wide distance prior now buys proportional speed; under the + uniform grid it buys nothing, because the spacing is set by `d_max` alone. + `--d-min 50` on this configuration takes `n` from 373 to 215. + +--- + +## 4. Rejected alternatives, with the measurement that rejected them + +**First, a distinction this document must not blur.** "The adaptive distance +machinery" names TWO different mechanisms and only one of them is being +rejected here. + +* `estimate_distance_peak` + `make_distance_grid_adaptive` build ONE static + window from an EXTERNAL peak estimate. That estimate is a 300-step gradient + ascent on `0.5*K^2/R`, measured 15.5-19.8 sigma from truth and varying + 224.8-1231.5 Mpc with the random seed alone. This is what is rejected + below, and the rejection is about the ESTIMATOR, not about adaptivity. +* `core._distmarg_gh_logL` is a per-sample adaptive quadrature: centre + `stop_gradient(clip(K/R, x_min, x_max))`, width `1/sqrt(R)`, both derived + from the data AT THAT (sample, time-bin), with no external estimator + anywhere. Nothing here rejects it, and the published `JAX_ILE_DISTMARG_GH` + rows are not tarred by the paragraph above. It is REFUSED in combination + with this option only because it consumes just the SUPPORT of `x_grid`, so + the option would be bit-identically inert beside it -- a flag-inertness + refusal, not a quality judgement. Its centre-clipping is designed for + exactly the boundary-layer regime section 1a refuses, so the two may well be + complementary rather than competing; that is unmeasured here and is an open + lead, not a claim. + +**(a) `make_distance_grid_adaptive` + `estimate_distance_peak` (the in-tree +machinery behind `JAX_ILE_DISTGRID_ADAPTIVE`).** Not shipped. Three defects, +all measured: + +1. `estimate_distance_peak` returned `d_peak = 1037.58 Mpc, sigma_d = 25.96 Mpc` + for a `d_inj = 633.92 Mpc` injection. Its `+-12 sigma` fine window is + `[726.1, 1349.1]` Mpc, which **does not contain the injected distance**. +2. Its 300-step fixed-schedule gradient ascent is not converged: the `rho` it + implies (`d_peak/sigma_d = 39.97`) is far below the amplitude bound's 55.07. +3. The returned peak does not respond to `guess_snr` at all: it is + byte-identical (`d_peak = 1037.5768181433305`, + `sigma_d = 25.95911760361334`) for `guess_snr` of `None`, 1.0, 17.38, 25.36, + 40.0 and 400.0 — a 400x span, one distinct pair. Recorded as an + OBSERVATION. The only thing added here is a citation, not a mechanism: the + function's own docstring says `guess_snr` "is accepted only as a fallback if + the sweep finds no K>0 sample", so on any path where the sky sweep succeeds + the argument is inert by design. `d_peak / d_inj = 1.636763` is **not** + explained in this document and should not be guessed at. +4. Its trapezoid gives the first and last nodes a **full** rather than half + interval. On a coarse backbone that misplaces percent-level volumetric + prior mass onto `d_max`; measured as a **0.018-0.026 nat error floor that no + refinement removes** (see the "FULL-endpoint" rows: 0.052 nats at n = 256, + still 0.026 at n = 512). + +Net: **9.4 nats (prior cloud) / 22.6 nats (near-injection cloud)** of lnL error +at 144 nodes. The helper's own docstring already carries a `LIMITATION` +paragraph saying a single static window is insufficient here; this measures it. +The env-var branch is left reachable so nothing that sets it changes behaviour, +but it now prints a deprecation warning and is mutually exclusive with the new +option. **Recommend removing it outright in a follow-up.** + +One behaviour change does reach that branch, and it is a fix rather than a +regression: the dense angle lattice used to be sized from the *adaptive* grid +and is now sized from the full prior support. Because +`estimate_angle_amplitude` clips to the grid's own `[x_min, x_max]`, the old +ordering could only ever make the lattice SMALLER than the full-support answer +(measured: 12.6% smaller amplitude, `(624, 320) -> (592, 304)`, on a +`[0.8d, 1.25d]` window). So the new ordering can only make it the same or +larger. `--distance-grid-scheme uniform` without the environment variable -- +i.e. every existing command line -- is unaffected: there the two grids are the +same object. + +**(b) A data-derived *window* (fine zone + coarse backbone) instead of a +full-range log-uniform grid.** Rejected, and this is the more interesting +result. `support_window.py` reads the amplitude +estimator's own sky sweep (reproduced bit-exactly: the amplitude it recomputes +matches `estimate_angle_amplitude` to 0.00e+00 relative) and asks how wide the +set of `x* = clip(A/B)` is over angle points within `T` nats of the maximum: + +| T (nats) | window (Mpc, padded) | ln-width | fine nodes at spacing 1/rho_max | +|---|---|---|---| +| 30 | [932, 1260] | 0.30 | 17 | +| 50 | [895, 1331] | 0.40 | 22 | +| 100 | [830, 1489] | 0.58 | 33 | +| 200 | [669, 1798] | 0.99 | 55 | +| 400 | [193, 2609] | 2.60 | 144 | + +That looks like a large win -- 22 fine nodes instead of 373 -- and it is a trap. +The window is centred on **1073 Mpc**, and `d_inj = 633.92 Mpc` is **outside it +until T = 400**. The reason is structural: the sweep is 64 random +sky/inclination draws plus deterministic extremes, and its empirical maximum +(758.17) is *below* the injection's own value (800). `ANGLE_AMP_MARGIN = 2.0` +bounds the sweep's error in the amplitude **value**, which is all the angle +lattice needs -- but nothing bounds its error in the **location** `x*`, which is +what a window needs. A window built from that sweep can therefore exclude the +true peak while the amplitude bound it shares is perfectly sound. So: no +window. The full-range log-uniform grid is location-free by construction, and +that is precisely why it is the one shipped. + +**(c) Per-distance-block dense phi sizing.** Measured separately and rejected: +1.45x standalone (not the 2-3x claimed), and 1.11x once an adaptive distance +grid has run, because the adaptive grid deletes exactly the low-amplitude +far-distance nodes that per-block sizing feeds on. Sizing goes as `sqrt(A)`, so +halving `n_phi` needs a 4x smaller amplitude. Documented dead end; do not +re-propose. + +**(d) The per-sample adaptive quadrature (`core._distmarg_gh_logL`, env +`JAX_ILE_DISTMARG_GH`).** This is the numerically superior answer -- nodes at +`x* +- 7 sigma` per sample, ~32-64 of them, gradient-stable via +`stop_gradient` -- and it is already implemented and already wired into the +`exact` scheme. It is **not** touched here because `laplace` +(1.9x faster than `exact` since #210, and the scheme the SNR-40 selector picks) +explicitly refuses it: its node placement is defined per fixed-`psi` exponent, +and the Laplace path has already integrated `psi` out analytically at each node. +Extending it would need a psi-marginal node-placement rule and its own +validation. The log-uniform grid works on **both** dense schemes today. +Promoting `JAX_ILE_DISTMARG_GH` from an environment variable to a first-class +option for `exact` is a good, separate PR. + +--- + +## 5. What can still go wrong, and what detects it + +Three distinct failure modes. They have different detectors and one of them is +NOT covered by the runtime fail-safe, which an earlier draft of this document +wrongly claimed it was. + +**(a) `rho_max` underestimates the interior peak's sharpness.** Then the +spacing is too coarse and the distance marginal is biased. `rho_max` derives +from `amp_sizing`, the same number the dense lattice is sized from, and +`anglemarg._runtime_amp_failsafe` recomputes the amplitude from the coefficient +tables inside every jitted likelihood call and warns when it exceeds +`amp_sizing`. That coverage is inherited, and it only holds because the +spacing is sized from `amp_sizing` (floored) rather than the unfloored +`amp_data` -- see section 1. + +**Limits, and they are more severe than an earlier draft of this section +said.** (i) The fail-safe is a `jax.debug.callback`, which XLA may drop (the +driver already labels artifacts `BEST-EFFORT` for this reason, and silence is +not verification). (ii) It compares the amplitude, not the spacing, so the +claim lapses if a future change sources `rho_max` from anywhere else. (iii) +**Under `vmap` the label is uninformative in BOTH directions**, so on the +`--mode flowmc-*` paths -- which is what production runs -- inheriting its +coverage buys nothing. `_runtime_amp_failsafe` guards its callback with +`jax.lax.cond`; a BATCHED predicate lowers to a select, so both branches +execute, and the callback branch passes a literal `True` rather than the +predicate, so the recorded state cannot distinguish tripped from not-tripped. +Measured on a sound run with the predicate false by four orders of magnitude +(`amp_sizing` forced to 1e12 against a worst reported amplitude of 10.49): + +| transformation | `tripped` | `n_calls` | +|---|---|---| +| plain call | False | 0 | +| under `jit` | False | 0 | +| under `vmap` | **True** | 2 | + +That is a defect in `_runtime_amp_failsafe`, NOT in this PR, and it is +deliberately not fixed here -- it is pre-existing, it needs a host-side +predicate evaluation and its own validation, and folding it in would widen a +PR under review. It is recorded here because it bounds what the sentence +above may be used for: for the interior-undersizing mode (a), **treat the +runtime fail-safe as absent on any vmapped path** and rely on the build-time +sizing. (Independently reported by the chip-05 session; confirmed here by the +measurement above.) + +**(b) The maximizing distance is EXTERIOR to the prior support** (section 1a). +**The runtime fail-safe is BLIND to this one.** +`_runtime_amp_failsafe` applies the identical `jnp.clip(M_A/B0, x_min, x_max)`, +so it under-reads by exactly the same mechanism the build-time amplitude does: +`amp_call <= amp_sizing`, and its `amp_call > 2 * amp_sizing` trigger never +fires. Nothing at runtime detects this regime. It is handled by REFUSING at +build time instead, using the unclipped-amplitude diagnostic, which is why that +refusal is not optional and must not be softened into a fallback. + +**(c) The detector's own accumulation, on the sky re-draw path.** +`estimate_angle_amplitude` re-draws the sky when its split-half check says the +maximum is still growing. Both maxima -- clipped and unclipped -- must be +updated there; updating only the clipped one leaves `clip_excess` reading the +first batch against a denominator that grew, so it falls BELOW 1 and the +refusal in (b) stops firing on exactly the events whose sky sampling was too +coarse to trust. External re-review deleted that update and all 30 tests of +the gate stayed green, so the two are now accumulated in the same idiom on +adjacent lines and guarded at the source +(`test_sky_doubling_updates_the_unclipped_maximum_too`). A behavioural test is +not available, and the reason is worth stating precisely because the obvious +one is wrong: the re-draw branch IS reachable (19 of 120 searched +`(data seed, n_sky, seed)` combinations enter it, one of them on an exterior +support). What is not reachable is a DIFFERENCE. The deterministic +face-on/face-off extremes are appended to the FIRST batch and are what attains +the unclipped maximum, so the second batch's unclipped contribution was a no-op +in every configuration measured -- deleting the update leaves `clip_excess` +bit-identical at 2.42571586419 on the one exterior doubling case available. +The corruption is therefore real but silent, and nothing bounds a dataset whose +unclipped maximum comes from a second-batch draw. A test that cannot be made +to fail is the kind this file has already deleted once, so the guard is on the +source instead. + +The contract in section 1 is therefore stated as holding uniformly over the +prior range and over every angle sample **given that the maximizing distance is +interior** -- a precondition that is checked, and refused when violated, rather +than assumed. + +`--distance-grid-scheme uniform` (the default) is untouched by all of this. + +### Refused combinations + +Two of these are refused in the CONSTRUCTOR only; the rest also fail at +option-validation time, before any precompute. + +* the exterior maximizing distance, and it cannot be otherwise: it is a + property of the DATA, not of the option set, so nothing before the precompute + can see it; +* `JAX_ILE_DISTGRID_ADAPTIVE=1` together with a non-uniform scheme. This one + COULD be caught at parse time -- it is visible from the option set plus the + environment, exactly like the `JAX_ILE_DISTMARG_GH` row -- and is not, which + is an inconsistency rather than a necessity. It is left alone here because + that variable is deprecated and its branch additionally requires + `guess_snr`, so the parse-time check could not reproduce the constructor's + condition without duplicating it. Recorded so the asymmetry is deliberate + rather than overlooked. (Found by external review of the fix round.) + +| combination | why | +|---|---| +| exterior maximizing distance | section 1a: 1.9-4.6 nats, worse than the default | +| `JAX_ILE_DISTMARG_GH` set | `core._distmarg_gh_logL` places its own per-sample nodes and reads only the SUPPORT of `x_grid`, so the option would be bit-identically inert while still reported as active. Reachable without the user naming a dense scheme: under GH `choose_angle_marg_scheme` resolves to one regardless, and this refusal does not depend on which | +| `--angle-marg-scheme grid` | the sizing amplitude is not computed on that path | +| a mode other than `flowmc-phipsimarg` | not validated there | +| `--distance-grid-points` also given | two options setting the same node count | +| `--distance-grid-tol` with the uniform scheme | would be inert | +| an unrecognised scheme value | optparse `choices` | + +## 5a. Verified through the driver, not only through the library + +Run on the reference configuration with `--srate 4096`, `--mode +flowmc-phipsimarg --angle-marg-scheme exact --distance-grid-scheme loguniform`, +the driver reports: + + Distance + phi_ref + psi marginalization: ON (grid=256, nphi=8, npsi=8, d in [1,10000] Mpc) + angle-marg scheme: exact (requested exact): amp_sizing=1109.1705986675768; amplitude=1109.1705986675768; ... + distance grid: dlnd=0.04093484609767193; mode=loguniform; n=226; n_uniform_requested=256; rho_max=47.09926960511334; tol=0.01 + +(`amp` differs from section 2's 1516.33 because the driver computes its own +fiducial epoch rather than the fixed one the study scripts pass; the derived +node count tracks it, which is the behaviour under test.) Note that the +`grid=256` on the first line is `--distance-grid-points`, which the log-uniform +scheme does not use; the authoritative line is the one below it, and it is +printed unconditionally. + +All four fail-closed refusals were exercised the same way and each raised: + +| what was passed | result | +|---|---| +| `--distance-grid-scheme loguniform --mode flowmc-phimarg` | `... apply only to --mode flowmc-phipsimarg ...` | +| `--distance-grid-scheme loguniform --distance-grid-points 256` | `... both set the distance node count ...` | +| `--distance-grid-tol 0.1` with the uniform scheme | `... it would be silently inert here.` | +| `--distance-grid-scheme adaptive` | optparse: `invalid choice: 'adaptive'` | + +--- + +## 6. Reproducing + + export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 + export JAX_ENABLE_X64=1 JAX_PLATFORMS=cpu + export ADAPT_TREE=$PWD PYTHONPATH=$PWD/MonteCarloMarginalizeCode/Code + taskset -c 0-15 python -u dist_quad_error2.py 40 256 # section 2a + taskset -c 0-15 python -u support_window.py 40 # section 4b + taskset -c 0-15 python -u e2e_laplace.py 40 64 4 # sections 2b, 3 + taskset -c 0-15 python -u compile_vs_execute.py 64 5 # section 3, compile vs execute + taskset -c 0-7 python -u verify_peak_snr.py # section 4a, guess_snr + +Gate: `MonteCarloMarginalizeCode/Code/test/jax/test_distance_grid_loguniform.py`, +run by `.travis/test-jax.sh`. Every test in it was verified to fail under a +named mutation; the matrix is in the PR body. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_tempering.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_tempering.md new file mode 100644 index 000000000..c7789a705 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_tempering.md @@ -0,0 +1,355 @@ +# Choosing `--adapt-weight-exponent` on the JAX extrinsic path + +**Status: a record, not a specification.** It is expected to be superseded. Where +it disagrees with the code, the code wins; where it disagrees with a later +measurement, the measurement wins. Dated 2026-08-23, RIFT branch +`rift_O4d_jax_adapt_weight_chooser` (stacked on `claude/jax-fairdraw-extrinsic`, +PR #180). + +The live decision is `samplers.export_ess_fraction` / +`samplers.beta_for_export_ess` plus the driver's `resolve_tempering_exponent`, +pinned by `test/jax/test_jax_tempering_chooser.py`. + +## The question + +RO'S, 2026-08-23: *"of course it's the default, we use intelligent choices to set +that weight-exponent with the helper historically. We just don't have that +intelligence yet for the JAX protocol (only some of the non-JAX samplers use it) +.. but we should have something."* + +## 1. The historical rule, and what it actually controls + +`bin/helper_LDG_Events.py:1472-1481`: + +```python +prefactor = 0.1 +snr_fac = max(1.0, event_dict["SNR"]/15.) # line 708 +if snr_fac > 1.5: # i.e. SNR > 22.5 + beta = prefactor/np.power(snr_fac/1.5, 2) +else: + beta = prefactor +``` + +Equivalently, with `lnLmax = SNR^2/2`: + +``` +beta = min(0.1, 25.31 / lnLmax) +``` + +— the rule holds **`beta * lnLmax` constant at ~25.3 nats** above SNR 22.5 and +caps `beta` at 0.1 below. It is a *fixed tempered dynamic range* rule. +`util_RIFT_pseudo_pipe.py` does not set the exponent itself; the helper is the +only place the choice is made. + +### Who actually consumes it + +`bin/integrate_likelihood_extrinsic_batchmode:1806` passes +`tempering_exp = opts.adapt_weight_exponent` (0.0 under `--no-adapt`). + +| sampler | consumes it? | how | +|---|---|---| +| `mcsampler` | yes | 1-D adapted histogram weights | +| `mcsamplerGPU` | yes | `log_weights = tempering_exp*lnL + ln p - ln p_s` | +| `mcsamplerEnsemble` | yes | `update_sampling_prior`, `ln_weights *= tempering_exp` | +| `mcsamplerNFlow` | yes | `update_sampling_prior` | +| `mcsamplerPortfolio` | yes | via the member samplers' `update_sampling_prior` | +| `MonteCarloEnsemble` | yes | `_solve_tempering_exp` (adaptive beta) | +| **`mcsamplerAdaptiveVolume`** | **NO** | reads it (line 1575) and never acts on it | + +**AV is not a consumer.** It has no `update_sampling_prior` (only +`update_sampling_prior_selfish`, which ignores `tempering_exp`), and the single +line that would apply it is commented out at line 1735: +`# log_weights = tempering_exp*lnL + log_joint_p_prior`. Its *only* effect in AV +is the side effect at line 1595, forcing `save_intg = True`. Anything that says +AV honours the exponent is wrong. + +### The structural point + +In every real consumer the exponent shapes **only the adaptive sampling prior**. +The estimator stays `log_integrand = lnL + ln p - ln p_s` +(`mcsamplerGPU.py:753`), so the integral is **unbiased at any beta** and the +exported sample count is untouched. `beta` there is a pure proposal-breadth knob: +free to use, and smaller = more defensive. + +## 2. Why the rule does not transfer + +On the JAX path (`samplers.flowmc_sample_phimarg`, driver line ~1178) +`beta = inv_T` is the exponent of the target the MCMC **samples**: + +```python +inv_T = 1.0/temper +def logpdf(theta, data): return inv_T*like._scalar(theta) + log_prior(theta) +``` + +The draws **are** the deliverable, so the export must be reweighted by +`post_weight ∝ L^(1-beta)` (PR #180) — and that reweight costs effective samples, +which the non-JAX path never pays. + +For a locally Gaussian peak `ln Z_g = g*lnLmax - (d/2) ln g + const`, so + +> **ESS/N = Z_1² / (Z_beta · Z_{2-beta}) = [beta (2 - beta)]^(d/2)** + +with `d` the **sampled dimension**. **No `lnLmax`. No SNR.** The two pictures make +opposite predictions, and the JAX one is the one that holds here. + +*(Recorded because it was the first guess and it is a plausible one: the +lognormal-weight estimate `exp(-(1-beta)² Var_beta[lnL])` is **wrong**, by ~70 +orders of magnitude at beta=0.1. The reweight weights are heavy-tailed, not +lognormal. `Var_beta[lnL] = d/(2 beta²)` is right — it is the ESS step that fails.)* + +## 3. Measurements + +### 3a. Exact offline sweep — real BNS likelihood, SNR 23.8, 4-D + +`~/rift_bns_jax_run_scratch/beta_ess_offline.py`, event 0 of +`~/rift_bns_jax_run/rundir/final_post.xml.gz`. Defensive importance sample, +300k draws, prior-pilot centres over a 5-scale bandwidth ladder + 12% prior +floor. **Reference ESS(g=1) = 5235**, and every `g` in [0,2] the estimator needs +is resolved (worst-case IS ESS 508), so the whole sweep is converged. + +| beta | measured ESS/N | law `[b(2-b)]²` | measured/law | export rows (of 4800) | +|---|---|---|---|---| +| 0.05 | 7.55e-3 | 9.51e-3 | 0.79 | 55 | +| **0.0951** (historical) | **2.76e-2** | 3.28e-2 | 0.84 | **199** | +| 0.20 | 1.04e-1 | 1.30e-1 | 0.80 | 751 | +| 0.50 | 5.00e-1 | 5.63e-1 | 0.89 | 3600 | +| 0.70 | 7.92e-1 | 8.28e-1 | 0.96 | 4800 | +| 0.90 | 9.75e-1 | 9.80e-1 | 0.995 | 4800 | +| 1.00 | 1.000 | 1.000 | 1.00 | 4800 | + +The law holds to 0.79–1.00 across the range; it is therefore slightly +**optimistic** and is safe as a design rule, not as a promise. + +`Var_beta[lnL]` independently tracks `d/(2 beta²)`: 213 vs 200 (beta=0.1), +8.08 vs 8.00 (0.5), 2.50 vs 2.00 (1.0). + +### 3b. End-to-end on the real sampler — no reference involved + +Full `integrate_likelihood_extrinsic_jax` runs, same event, 156 s each. The +driver's own `fairdraw: N weighted samples (ESS=...)` line is a direct +measurement, so this shares no machinery with 3a. + +| arm | lnZ | sigma_lnL | export ESS | rows written | law predicted ESS | +|---|---|---|---|---|---| +| beta = 1.0 (default) | 253.686 | 0.0118 | n/a (uniform) | 4800 | 4800 | +| **beta = 0.0951 (historical)** | 253.908 | **0.0779** | **184.9** | **278** | 157 | +| beta = 0.7735 | see §4 | | 4202.8 | 4800 | 4320 | + +Predicted vs measured: 157 vs 185 (ratio 1.17) and 4320 vs 4203 (0.97). The law +is validated within ~20% on the real sampler by a route that never touches the +offline reference. + +**Porting the historical exponent literally costs 17x in exported rows +(4800 -> 278) and 6.6x in evidence precision**, and the driver itself prints +*"NOT a usable posterior sample however it is drawn."* + +### 3c. Is the cost SNR-set or dimension-set? (reference-free) + +Same source, SNR set by injected distance (network SNR is exactly propto 1/d), +beta held fixed, reading the driver's own reported export ESS. No reference cloud +is involved, so this cannot inherit a reference's convergence problems. + +| injected d | lnZ | implied SNR | beta=0.5 ESS/N | beta=0.1 ESS/N | +|---|---|---|---|---| +| 1200 Mpc | 113.1 | ~15.0 | 0.620 | 0.0472 | +| 600 Mpc | 533.1 | ~32.7 | 0.517 | 0.0371 | +| 300 Mpc | 2230.7 | ~66.8 | 0.560 | 0.00823 | +| 150 Mpc | — | ~134 | 0.126 | 0.00763 | +| **law** | | | **0.5625** | **0.0361** | + +**At beta = 0.5 the cost is flat to +-10% from SNR 15 to 67** — a factor 20 in +lnLmax, over which the historical rule would have demanded beta fall by the same +factor 20. That is the discriminator, and it settles it: the reweight cost is not +SNR-set. + +**Two honest caveats, both visible in that table.** + +1. **The law is optimistic at small beta, and increasingly so at high SNR.** At + beta=0.1 the measured cost falls from 0.047 (SNR 15) to 0.0082 (SNR 67) while + the law says 0.036 throughout. The Gaussian-peak approximation degrades where + the target is both very sharp and very heavily tempered. This matters for the + guard: near its threshold the guard can *under*-refuse, because it trusts a + law that is too kind in exactly that corner. It never over-refuses. +2. **beta = 0.5 breaks down by SNR ~134** (0.126 against 0.5625). At the highest + rung the flow itself is struggling, not just the reweight. Static tempering is + therefore NOT the tool for the 3G regime, which is the regime it is usually + proposed for. + +An earlier offline sweep over the same injections (`beta_ess_vs_snr.py`) is +**not** quoted: its reference collapsed above SNR 20 (ESS(g=1) = 4.1 at SNR 40 — +the Hessian at the truth is near-flat in inclination and the eigenvalue floor +mis-scales the proposal). Those rows are not evidence. + +### 3d. Does beta < 1 buy accuracy? (two seeds) + +Scored against the independent defensive-IS reference, floor at the arm's own +row count. + +| arm | rows | JS psi s0/s1 | JS incl s0/s1 | sd psi s0/s1 | +|---|---|---|---|---| +| beta = 1.0 | 4800 | 0.0735 / 0.0358 | 0.0617 / 0.0325 | 0.940 / 0.990 | +| beta = 0.7735 | 4800 | 0.0407 / 0.0289 | 0.0445 / 0.0296 | 0.983 / 1.037 | +| beta = 0.0951 | 278 / 203 | 0.1354 / 0.1438 | 0.1310 / 0.0954 | 1.081 / 0.872 | +| `--adapt-adapt` | 4800 | 0.0462 / **0.5690** | 0.0774 / **0.3825** | 1.043 / **0.031** | + +**Not resolved: whether beta=0.7735 beats beta=1.** Lower psi and incl JS in both +seeds, but beta=1's own psi JS varies 2x across seeds — a spread comparable to +the gap. Two seeds cannot settle it, and **no part of this change rests on it**. +What is solid is that the auto exponent costs nothing measurable: same 4800 rows, +export ESS 4203 / 4266. + +**Resolved: `--adapt-adapt` collapsed on one seed of two.** Seed 1 returned psi +spanning only [1.599, 1.769] rad of the full [0, pi] — 3% of the reference width +— against [0.026, 3.131] for seed 0 and both beta=1 arms. Verified in the raw +export, not only in the score. Seed 0 was also 21% narrow in inclination. At +SNR 23.8, not an extreme case. + +## 4. What was built, and what was deliberately NOT + +**Built.** + +- `samplers.export_ess_fraction(beta, n_dim)` / `beta_for_export_ess(target, n_dim)` + — the law and its inverse. The signature carries no SNR argument, on purpose; + a test pins that. +- `--auto-adapt-weight-exponent` + `--target-export-ess-frac` (default 0.9): + picks the smallest beta meeting the export budget, keyed on the **sampled + dimension** (at a 0.9 target: d=3 -> 0.786, d=4 -> 0.807, d=5 -> 0.824). + These solve the CALIBRATED lower bound, not the bare law -- see §4a. +- A guard: any beta whose predicted export ESS falls below the driver's own + usability floor of 200 is **refused** (exit 1, no file written), with a message + naming the historical rule as the trap. `--allow-degenerate-tempering` overrides. + +**Not built: a port of the SNR rule to the static exponent.** It keys on a +variable the cost does not depend on, and at this event it selects a value the +driver already declares unusable. + +**Not built: `--adapt-adapt` on by default.** It is a different mechanism — an +annealing *schedule* that ladders `inv_T` up and always terminates at +`inv_T = 1` (the loop breaks on `inv_T >= 1`, and `post_weight` is then uniform), +so it costs nothing at export. But it is **not** a free robustness win: it +collapsed on one seed of two at SNR 23.8 (§3d), and cost **>24 min against 156 s** +for a static run. On this evidence it must not be a default. PR #183 needs this: +the anneal cannot be assumed safe. + +### The closer structural analogue, for whoever picks this up + +Old-RIFT's beta broadens a *proposal* while the estimator stays exact. The JAX +constructs that do that are **not** `inv_T` — they are the hardcoded defensive +inflations: `cov = cov * 2.0` (`samplers.py:816`, `:1462`, the moment-matched IS +evidence proposal) and `fisher_is_inflate=1.3` (`:1134`, the high-SNR Fisher-IS +fallback). Those are where "intelligence" could go without paying any export-ESS +cost. Not touched here — out of scope, and unmeasured. + +### 4a. The law is optimistic — and the calibration is NOT a bound either + +Raised in review of this change, and correct. `export_ess_fraction` is the +Gaussian-peak law, which the §3a sweep shows is optimistic by up to 21% +(measured/law 0.79 at beta=0.05, rising to 1.00 at beta=1). Inverting it to pick +beta, and guarding on it, both hand back something already known to fall short of +what was asked for: at d=4 a 0.9 target returned **beta=0.77347, whose measured +lower bound is 0.866**. + +`export_ess_estimate(beta, n_dim) = cal(beta) * law(beta, n_dim)` is what the +chooser solves and what the export check reports. `cal` is a piecewise-linear +envelope over the measured ratios, every knot at or below every measured point: + +| beta | 0.05 | 0.20 | 0.40 | 0.60 | 0.80 | 1.00 | +|---|---|---|---|---|---|---| +| cal | 0.79 | 0.79 | 0.83 | 0.91 | 0.97 | 1.00 | + +A single flat factor would have been the wrong shape twice over: the shortfall is +strongly beta-dependent, and a flat 0.79 would make any target above 0.79 +unreachable, since it never rises to 1. The inverse has no closed form once the +envelope is included and is solved by bisection (both factors are monotone in +beta, so the product is). + +**IT IS NOT A LOWER BOUND, AND A SECOND REVIEW ROUND CAUGHT IT BEING USED AS +ONE.** `cal` is fitted at SNR ~= 23.8, and the shortfall grows with SNR. §3c's +own ladder already said so, and it supplies the counterexample: + +| | beta=0.1, d=4 | +|---|---| +| estimate (`cal x law`) | 0.0285 | +| **measured, SNR ~= 67** | **0.00823** | + +a factor 3.5 the wrong way. On a 10 000-row cloud that is an estimated ESS of +285 against a measured ~82: comfortably over a 200-row floor while being well +under it. The first version of this section inverted the estimate to pick beta +AND refused runs on it, i.e. built a hard guarantee out of a single-SNR fit while +the caveat disproving it sat two sections above. + +A real bound needs a calibration in **(beta, SNR)**, and the driver has no +trustworthy SNR where the choice is made — `guess_snr` is an explicit +guesstimate (10.32 against a true network 23.78 on the study event). So: + +* the export check **warns, it does not refuse** — a refusal is a guarantee; +* `--auto-adapt-weight-exponent` is documented **EXPERIMENTAL and not + sufficiently validated**, and no paper result uses it; +* **measured at d=4 only** — the ratio at other dimensions is assumed, not + measured, and is the first thing to check if a d=3 or d=5 budget comes up short. + +## 5. Limitations — axes swept, and axes presumed load-bearing + +**Swept:** beta over [0.05, 1]; SNR ~15 to ~134 at fixed beta; two seeds on the +accuracy arms; sampled dimension via the closed form (3/4/5, verified +analytically, only d=4 measured); two independent estimators (offline defensive +IS, and the driver's own reported ESS). + +**NOT swept, presumed load-bearing:** + +- **Accuracy above SNR 24.** The accuracy arms (§3d) are one event at SNR 23.8. + The ESS ladder reaches SNR ~134, but only measures export ESS there, not + whether the resulting posterior is right. +- **The guard's threshold in the corner where the law is optimistic** (§3c + caveat 1): near ESS ~200 at small beta and high SNR the guard trusts a law that + over-predicts. It errs toward passing, not refusing. Not characterised. +- **The calibration is single-SNR** (§4a). Fitted at SNR ~= 23.8; the ladder + shows it optimistic by 3.5x at SNR ~= 67 and beta=0.1. Nothing may refuse a + run on it, and the chooser's target is met only on that calibration. +- **The calibration at dimensions other than 4** (§4a): assumed from the d=4 + sweep, not measured. +- **Only two seeds.** Enough to show the `--adapt-adapt` collapse (it is a 30x + effect) and to leave the beta=0.7735-vs-1 question open. Not enough for either + to be a width claim. +- **Two paths where the exponent is inert**, both found by reading rather than + running. `--smc-puffball` routes to `smc_puffball_sample`, which swallows + `temper` in `**_ignore` and exports uniform weights: the guard is skipped there + and the no-op reported instead. `--fisher-is-samples` is conditional — a + *successful* Fisher-IS pass replaces the cloud with an already-fair-drawn + uniform-weight set so the reweight cost never materialises, but it falls back + to the tempered draws when it fails. The guard still refuses there + (fail-closed), with `--allow-degenerate-tempering` as the documented escape. + Neither path was measured. +- **Non-Gaussian / strongly multimodal targets.** The law is a Gaussian-peak + result; the measured 0.79 shortfall at small beta is that approximation + failing. A target with well-separated equal-mass modes may do worse. +- **Modes other than `flowmc-phimarg`.** `flowmc` (5-D), `flowmc-phipsimarg` + (3-D) and `flowmc-dpsimarg` (4-D) take the same code path and the same `dim`, + but were not run. + +## 5b. A defect class found by READING, not by running + +The chooser originally wrote its answer back to `opts`. `opts` is per-RUN; +`analyze_one` is per-EVENT. On event 1 of a batch the chooser read its own +event-0 output as a user-supplied exponent and aborted with `SystemExit` — and +`ILE_extr.sub` runs batches, so every real multi-event `--auto` run would have +died. **No single-event test can see this**, and every measurement in §3 is +single-event. The repo's own `RIFT/integrators/REVIEW_CHECKLIST.md` names the +shape ("Is per-point state cleared on ENTRY"). The chooser is now pure and +returns the exponent; `analyze_one` keeps it in a local. + +## 6. Reproduce + +``` +# offline exact sweep (needs the BNS run products, read-only) +python beta_ess_offline.py --code /MonteCarloMarginalizeCode/Code \ + --event 0 --n-pilot 100000 --n-ref 300000 --tag REF0 +# one end-to-end arm +./run_arm.sh ARM_b095_s0 0 --adapt-weight-exponent 0.09508 +# score an arm against the independent reference +python score_arm.py --ref REF0_ev0_ref.npz --files ARM_b095_s0_0_samples.dat +``` + +All three live in `~/rift_bns_jax_run_scratch/` on CIT (NFS home — **not** on +condor execute nodes). Run on `ldas-pcdev11`/`13` with `OMP_NUM_THREADS=1` and +`taskset`; JAX sizes its XLA pool from the visible CPU count. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md index ff24f0538..c6a98d4d8 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md @@ -1,5 +1,20 @@ # `jax_ile` — an AD-compatible JAX reimplementation of the ILE extrinsic likelihood +> **`--save-samples` output is a FAIR DRAW**, not the raw sampler cloud: +> equal-weight rows, no weight column, and the same columns as before *for a +> given `--mode`* (different modes export different column sets — see the Driver +> section). A second header line records the mode and the export ESS, e.g. +> `# mode=laplace-is fairdraw: ESS=5.5 n_in=300000 n_out=9`. **Check that ESS +> before trusting a file**: a low-ESS export is not a usable posterior sample +> however it is drawn, and the driver warns on stderr when it is below 200. +> When the weights admit no fair draw at all (degenerate/unnormalizable), the +> event fails and **no samples file is written** (any stale one at that path is +> removed) — there is no mode in which this product holds unreweighted rows. +> That refusal is checked *before* the `__.dat` result row is +> written, so a failed event leaves **no result row either** (a stale one is +> removed too): with `--soft-fail-event-range` the batch goes on, and a row left +> behind would be collected as a successful integration. + A `jax.numpy`, automatic-differentiation-compatible reimplementation of RIFT's ILE extrinsic likelihood, mirroring the production `factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` @@ -27,6 +42,69 @@ response, geometric time delay, spin-(-2) spherical harmonics, the `kappa`/`rho^2` assembly, continuous time-shift interpolation, time marginalization, and **analytic distance marginalization**. +### Time quadrature + +All JAX likelihood wrappers accept the conventional ILE keyword +`time_quadrature={"simpson","bandlimited"}`. Simpson remains the default. +The opt-in `bandlimited` path is currently supported by +`JAXExtrinsicLikelihood`, including analytic phase marginalization. It forms +the endpoint-nonduplicating even extension +`[kappa[0], ..., kappa[-1], kappa[-2], ..., kappa[1]]`, FFT-interpolates it, +applies the phase reduction on the +fine grid, and integrates the original closed interval with a stable trapezoid +rule. The per-row power-of-two factor is derived from fine-grid peak curvature, +remeasured after interpolation, and doubled until the integral agrees within +1e-3 nat. Row-local `lax.map` execution bounds scratch memory independently of +the sampler batch. There is deliberately no public factor knob; a row that +cannot meet the criterion fails closed. + +The supported signal regime assumes spectral headroom below the sampled +Nyquist frequency and negligible likelihood mass at both ends of the short +integration window. The latter is checked on the refined grid: either endpoint +must be at least 15 natural-log units below the peak, otherwise `bandlimited` +fails closed rather than trusting a boundary extension that can affect the +answer. Increase the physical time window or use Simpson when this diagnostic +fires. + +The primitive gather includes support outside that window. Its initial guard +is the established half-window default rounded up to a power of two; one guard +doubling is gathered at the same time. A raised-cosine pad acts only across +the support samples, reaching exactly one at the integration crop and zero +with zero slope at the remote even-reflection turns. The value is accepted +only when both guard widths agree within 1e-3 nat, independently of the fine +quadrature-factor doubling check. Thus short-window truncation and fine-grid +resolution have separate certificates. + +The JAX driver derives this support requirement before waveform precompute and +widens `--internal-data-storage-window-half` when necessary. It includes the +full certified guard, a conservative 50 ms detector-delay allowance (larger +than the Earth-diameter light time), and the +largest shipped interpolation stencil. The accumulator also validates every +guarded gather index per row; missing support produces a fail-closed likelihood +instead of inheriting the ordinary gatherer's out-of-buffer zero fill. The +baseline and banded finite-size/frequency-response accumulators enforce the +same check; rotation remains refused because its norm depends on arrival time. +The +curvature-derived starting fine factor is capped at 1024 and certified once at +2048; a sharper row is refused with guidance to increase the input/rholm sample +rate rather than allocating multi-gigabyte FFT branches. + +Distance, phi, psi, exact-angle, and Laplace-marginalized wrappers currently +refuse `bandlimited`. Those nonlinear reductions generate time harmonics, so +interpolating their already-reduced `lnL(t)` can converge to the wrong function; +they require endpoint-specific primitive refinement before they can safely opt +in. They continue to use the unchanged Simpson default. +The driver exposes the same public spelling as conventional ILE: +`--time-marginalization-quadrature`. `--interpolate-time` is an alias for the +JAX-native `--interp` with conflict detection. Conditional nuisance recovery +is outside this implementation. For drop-in CLI compatibility, +`--resample-time-marginalization`, `--srate-resample-time-marginalization`, and +`--time-posterior-export` are accepted and reported as ignored: JAX ILE's +sample export keeps time terminally marginalized rather than reconstructing one +conditional time per exported row. This intentionally differs from +conventional ILE's XML export semantics, but a high-level DAG can swap +executables without dying during option parsing. + ## Modules - `detector.py` — `compute_detamresponse`, `time_delay_from_earth_center` @@ -136,13 +214,27 @@ Modes (`--mode`): sample the 5-D angular posterior). Implemented in `samplers.py`. **Efficiency / robustness options:** -- **Flow re-use across a batch** (`--mode flowmc` + `--n-events-to-analyze`): - the trained normalizing flow is bootstrapped from one intrinsic template to - the next (its NF weights warm-start the next event and its posterior draws - initialize the chains). For nearby templates the posterior changes slowly, so - this is the partial-flow-reuse win — most visible at scale and at high SNR. - `--no-flow-reuse` disables it (re-train each event). Validated: a re-used run - recovers the truth sky with neff ≥ the fresh run (`test/jax/test_flow_reuse.py`). +- **Flow re-use across a batch** (`--mode flowmc` + `--n-events-to-analyze`) — + **OFF by default; opt in with `--flow-reuse`.** When enabled, the trained + normalizing flow is bootstrapped from one intrinsic template to the next (its + NF weights warm-start the next event and its posterior draws initialize the + chains). `--no-flow-reuse` is accepted and now restates the default. + + **Do not enable it for any run whose extrinsic SAMPLES are used.** Re-use + contracts the posterior in later slots: across an 8-event batch at two seeds, + psi fell to ~40% of its no-re-use width by slot 7 on both seeds, with slot 0 + (no re-use yet) at ~1.0 as a control, and inclination to 0.49/0.61. Confirmed + against independent per-slot references, not just arm-vs-arm. It also + reproduces an earlier, independent measurement (mean incl 0.5795 → 0.3465, + sd(psi) 0.9122 → 0.3738) that caused an amortization claim to be retracted from + the companion paper. + + *What the earlier validation actually showed.* `test/jax/test_flow_reuse.py` + checks that a re-used run **recovers the truth sky with neff ≥ the fresh run**. + That remains true and is not contradicted here — it tests sky location and an + evidence-side neff, neither of which is posterior *width*. The contraction is + in the width of the orientation parameters, an observable that check does not + look at. - **Network sky coordinates** (`--sky-coordinates network`, `multistart-nuts` only): sample the sky in the two-detector baseline frame `(cosθ_n, φ_n)` to fold the time-delay ring (the prior stays uniform there). Falls back to @@ -158,8 +250,14 @@ neff / wall time — the data for the skymap-vs-SNR figure and the high-SNR efficiency comparison vs the adaptive (AV) integrator. Preliminary small-budget run (H1/L1/V1): the flow **recovers the truth sky at every SNR through 640**, the 90% sky credible area shrinks with SNR (≈0.05 deg² at 40 → 3.6e-5 deg² at 80 → -sub-sample-resolution above), and **flow re-use cuts the per-event wall time ~2×** -(first event ≈114 s, warm-started events ≈60 s). The simple moment-matched +sub-sample-resolution above), and in **that small-budget configuration** flow +re-use cut the per-event wall time ~2× (first event ≈114 s, warm-started events +≈60 s). **That saving does not carry to production settings:** on an 8-event BNS +batch at full settings it measured 1589 s with re-use against 1567 s without +(1549/1629 vs 1644/1489 across two seeds) — a difference smaller than the +seed-to-seed spread, with its sign flipping. Treat the ~2× as specific to the +small-budget benchmark, not as a general amortization argument, and see the +accuracy warning above before enabling re-use at all. The simple moment-matched Gaussian importance evidence is reliable only at moderate SNR (it is flagged `nan` once `neff` collapses, since `logZ ≤ lnL_max` is violated by an ill-conditioned proposal at sub-resolution peaks) — a robust narrow-peak evidence @@ -178,9 +276,14 @@ Self-test (no frames needed): ``` PYTHONPATH=<...>/Code python bin/integrate_likelihood_extrinsic_jax \ --inj-mode --mass1 35 --mass2 30 --spin1z 0.1 --spin2z -0.2 \ - --mode nuts --d-max 5000 --save-samples --output-file out + --mode nuts --distance-marginalization --d-max 5000 \ + --save-samples --output-file out ``` +(`--mode nuts` requires `--distance-marginalization`: `run_nuts` raises +`SystemExit` without it, because the bare 5-D angular+distance likelihood is +degenerate. The command above previously omitted the flag and could not run.) + Output: `out_0_.dat` (`event_id m1 m2 s1x..s2z lnL sigma_lnL ntotal neff`) and, with `--save-samples`, `out_0_samples.dat`. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py new file mode 100644 index 000000000..c29a4bea3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py @@ -0,0 +1,2021 @@ +""" +Exact (phi_ref, psi) angle marginalization for the JAX factored likelihood. + +WHY THIS MODULE EXISTS +---------------------- +:func:`core.fused_log_likelihood_distphipsimarg` marginalizes (phi_ref, psi) +by averaging exp(lnL) over the SAME small grid it evaluates the likelihood on +(nphi x npsi, production 8x8). That is an 8-node quadrature of a function +whose peak width is ~1/SNR, so the error grows without bound with SNR +(measured against a converged reference: ~1e2 nats at SNR 40, ~7e3 nats at +SNR 320 for the psi marginal at npsi=8). The 8-point phi grid additionally +puts the n=4 phi harmonic (present for any m_max=2 signal through rho^2) +exactly AT Nyquist, where it aliases onto n=-4. + +THE STRUCTURE THE FIX EXPLOITS (analytic, not tunable) +------------------------------------------------------ +At fixed time and distance the factored lnL is a bivariate trigonometric +polynomial of KNOWN low order in (phi_ref, psi): + + * the antenna pattern enters as F(psi) = F(0) e^{-2 i psi} -- linearly in + kappa (u-harmonics +-1, with u = 2 psi) and quadratically in rho^2 + (u-harmonics {0, +-2}); + * the harmonics Y_lm carry e^{i m phi_ref} -- kappa has phi-harmonics up to + m_max, rho^2 up to 2*m_max. + +So the two unit-distance fields the likelihood is built from, + + A(phi, psi; t) = Re kappa_unit (phi order <= m_max, u order <= 1) + B(phi, psi; t) = rho^2_unit (phi order <= 2*m_max, u order <= 2) + +are EXACTLY determined by their values on a small Nyquist-sized sample grid. +The expensive :func:`core._accumulate_unit` evaluations are needed only to +pin those Fourier coefficients; every subsequent evaluation of +lnL_t(phi, psi | x) = x*A - 0.5*x^2*B (x = distMpcRef/d) is pure arithmetic. +The number of expensive evaluations is fixed by MODE CONTENT, never by SNR, +and is asserted -- there is no accuracy-vs-cost knob to set too small. + +TWO MARGINALIZATION SCHEMES (plus a selector, see the wrapper/driver) +--------------------------------------------------------------------- +exact : reconstruct lnL_t on a dense (phi, u) product grid from the + coefficient tables and average exp(.) over it. The dense grid is + free (no likelihood calls); its size is derived from a DATA-DERIVED + amplitude bound (:func:`estimate_angle_amplitude`, computed from + the coefficient tables themselves -- never from a caller's SNR + estimate) via :func:`_dense_grid_sizes`, floored at the crossover + amplitude. Best at low/moderate amplitude. +laplace : marginalize psi ANALYTICALLY by Laplace's method at every + (phi, distance-node, time) point -- at fixed (phi, x, t) the + u-exponent is exactly a + b cos(u-beta) + d cos(2u-delta), whose + stationary points Newton finds from u0 = beta, beta+pi in a few + elementary iterations, and whose curvature is closed-form. This + removes the psi axis entirely (cost ~SNR instead of ~SNR^2) and its + O(1/amplitude) error SHRINKS as SNR grows. Best at high amplitude. + +Both schemes marginalize distance with the same quadrature machinery as the +grid path (:func:`core._logsumexp_grid_blocked`, or the adaptive +:func:`core._distmarg_gh_logL` when JAX_ILE_DISTMARG_GH is set; the laplace +scheme cannot call that function, whose nodes are placed per FIXED psi, and +uses the psi-MARGINAL placement documented above +:func:`_gh_psi_node_offsets` instead -- for m_max <= 2 only, raising above +it), and use the same normalization convention (mean over uniform +angle grids, i.e. the uniform priors dphi/2pi, dpsi/pi), so they are +drop-in replacements for the grid function and agree with it wherever the +grid is converged (pinned in test/jax/test_angle_marg_exact.py). + +Everything here is jit/vmap/grad-compatible: no numpy in the hot path, no +scipy special functions, lax.scan (checkpointed) bounds memory by CHUNK, not +by grid size. +""" + +import numpy as np +import jax +import jax.numpy as jnp + +from . import core as _core +from .core import (JAX_INTERP_DEFAULT, TIME_QUAD_DEFAULT, _accumulate_unit, + _time_marginalize_terminal, + _logsumexp_grid_blocked, _distmarg_gh_logL, + make_distance_gh) + +__all__ = [ + "ANGLE_MARG_DEFAULT", + "ANGLE_MARG_LEGACY", + "ANGLE_MARG_CHOICES", + "angle_sample_grid_sizes", + "angle_coefficient_tables", + "estimate_angle_amplitude", + "fused_log_likelihood_distphipsimarg_exact", + "fused_log_likelihood_distphipsimarg_laplace", + "choose_angle_marg_scheme", + "fused_log_likelihood_distphipsimarg_peaklocal", + "gh_laplace_supported", + "ANGLE_MARG_CROSSOVER_AMPLITUDE", +] + + +# --------------------------------------------------------------------------- +# Selector crossover and dense-grid sizing constants. +# +# Calibrated 2026-08-27 on the SEOBNRv4 35+30 Msun HLV injection harness +# (the configuration behind the measured production-grid errors quoted in the +# module docstring), on the full distance+phi+psi-marginalized lnL, against a +# brute-force dense product-grid reference (which agrees with the exact +# scheme to ~2e-12 wherever it is affordable). Measured errors in nats: +# +# SNR A=rho^2/2 exact(self-conv) laplace-exact grid 32x8 grid 8x8 +# 10 50 7e-15 -1.1e-03 -5.9e-02 -9.5e-02 +# 20 200 0 -1.8e-04 -8.8e-01 -1.7e+00 +# 40 800 2.4e-09 -1.6e-05 -5.6e+00 -1.1e+01 +# 80 3200 4.6e-13 -2.8e-06 -2.7e+01 -5.1e+01 +# +# The Laplace error falls FASTER than 1/A here; the isolated-kernel error law +# is ~0.1/b nats (pinned in test_angle_marg_exact.py). The crossover sits +# where BOTH schemes are deep in their accurate regimes (laplace ~1e-4, +# exact ~machine), so the switch is insensitive to the O(1) slack in the +# measured amplitude bound that drives it, and tests evaluate both schemes in +# the overlap region and assert agreement -- the crossover is a validated +# constant, not a tuning knob. +# --------------------------------------------------------------------------- +# THE (phi_ref, psi) SCHEME DEFAULT. One definition; every entry point imports +# it, so the driver flag and the wrapper argument cannot drift. This is the +# same discipline JAX_INTERP_DEFAULT already uses, and for the same reason: the +# last default move on this path (interp linear -> sinc) had the value re-typed +# in many places. +# +# CHANGED 2026-09-02: 'grid' -> 'exact'. THIS CHANGES RESULTS for any caller +# that does not pass the scheme explicitly. Pass --angle-marg-scheme grid (or +# angle_marg=ANGLE_MARG_LEGACY) to reproduce a pre-2026-09-02 run. +# +# Why 'exact' and not 'auto': 'auto' selects 'laplace' above +# ANGLE_MARG_CROSSOVER_AMPLITUDE (rho ~21-30), which is an ACCURACY crossover. +# But 'laplace' cannot use the per-sample adaptive distance quadrature and the +# log-uniform distance grid is opt-in, so on the default uniform grid 'laplace' +# was measured 43.2 nats from 'exact'+GH16 at rho 163 (mean; 16.3 median) -- an +# error on the DISTANCE axis, not the angular one, which is ~1e-6 nats there. +# A default that is correct and slow beats one that is fast and tens of nats +# wrong. 'auto' becomes the right default once laplace has a sound distance +# quadrature, and ANGLE_MARG_CROSSOVER_AMPLITUDE should then be re-derived from +# COST as well as accuracy -- the measured cost crossover is rho ~200-326, an +# order of magnitude above the accuracy one. +# +# Why not 'grid': its quadrature error grows without bound with SNR (it averages +# exp(lnL), whose peak width is ~1/SNR, on n_phi x n_psi nodes). Measured on +# the paper-1 ladder-2 injection at rho 652: the best of the 4 distinct +# n_phi=8 nodes is 37,419 nats below the true phi_ref profile peak, and the +# recovered sky position is displaced 0.53 deg -- the grid scheme ranks that +# artifact ABOVE the injection and the correct peak BELOW it, by 900.6 nats. +# Evidence: RIFT_roboto_paper analyses/sky_offset_diagnosis/ +# RESULTS_phigrid_2026-09-02.md (commit 3f1f66f). +ANGLE_MARG_DEFAULT = "exact" +ANGLE_MARG_LEGACY = "grid" # the spelling that reproduces pre-2026-09-02 runs +ANGLE_MARG_CHOICES = ("grid", "exact", "laplace", "peak-local", "auto") + +#: 'peak-local' is deliberately NOT reachable from 'auto' yet. It agrees with 'exact' +#: to 1e-13 nats on the tables measured so far and is device-independent (the same answer +#: on CPU and on an NVIDIA Blackwell GPU), but nothing has yet compared the two head to +#: head on a production campaign, and a scheme that changes the likelihood must not +#: become reachable by default on the strength of unit tests. Explicit-only is what lets +#: a pilot run both and decide; promoting it into `choose_angle_marg_scheme` is a +#: separate change with its own evidence. + +# --------------------------------------------------------------------------- +ANGLE_MARG_CROSSOVER_AMPLITUDE = 450.0 # A = rho^2/2; rho = 30. NOTE the +# auto selector compares the MARGINED data-derived bound (~2x the true +# amplitude) to this, so laplace engages from true A ~ 225 (SNR ~ 21). That +# early engagement is safe by measurement: laplace is at -1.8e-4 nats by +# A = 200 on the injection ladder and improves upward, while exact remains +# valid (crossover-floored sizing) below. +# Dense-size rule N = ceil(K * sqrt(A)) points, from the trapezoid aliasing +# error of exp(trig poly): relative error ~ exp(-c N^2 / A). The constants +# carry a >= 2x margin in N over the empirically adequate values (error +# floor reached in the calibration ladder). +_DENSE_K_U = 8.0 # u = 2 psi axis +_DENSE_K_PHI = 16.0 # phi axis: harmonics up to 2*m_max, dominant n=2 +_DENSE_FLOOR_U = 64 +_DENSE_FLOOR_PHI = 128 + + +def angle_sample_grid_sizes(m_max): + """Nyquist-derived (nphi_s, npsi_s) sample-grid sizes for mode content m_max. + + lnL_t carries phi-harmonics up to 2*m_max and u-harmonics up to 2 + (u = 2 psi), so unaliased sampling needs nphi_s > 2*(2*m_max) and + npsi_s > 2*2. These are DERIVED and asserted, not options: the historical + defect was precisely a settable sample size (nphi=8 aliases the n=4 + harmonic; npsi=8 under-resolves nothing at the SAMPLING stage but the + grid was also used as the quadrature). + """ + m_max = int(m_max) + if m_max < 1: + raise ValueError("m_max must be >= 1, got %r" % m_max) + nphi_s = 4 * m_max + 8 # >= 2*(2 m_max)+2, margin >= 6 harmonics + npsi_s = 8 # u content is <= 2 for ANY mode set (spin-2) + assert nphi_s >= 2 * (2 * m_max) + 2 + assert npsi_s >= 2 * 2 + 2 + return nphi_s, npsi_s + + +def _data_m_max(data): + lms = np.asarray(data.lms) + return int(np.max(np.abs(lms[:, 1]))) + + +def angle_coefficient_tables(data, ra, dec, incl, interp=JAX_INTERP_DEFAULT, + sample_chunk=None): + """Exact 2-D Fourier coefficient tables of A = Re kappa_unit, B = rho^2_unit. + + Samples :func:`core._accumulate_unit` on the Nyquist-sized + (nphi_s x npsi_s) grid from :func:`angle_sample_grid_sizes` and + accumulates the discrete Fourier coefficients + + C[kp, ks] = (1/Ns) sum_j f(phi_j, psi_j) e^{-i kp phi_j} e^{-i ks u_j} + + for the harmonics the physics allows: A keeps kp in 0..m_max, + ks in -1..1; B keeps kp in 0..2*m_max, ks in -2..2 (u = 2 psi; + negative-kp coefficients follow from Hermitian symmetry of the real + fields and are not stored). Reconstruction weights are w_kp = 1 for + kp = 0 and 2 for kp > 0 (the kp = 0 row stores both ks signs, whose + conjugate pairing is already real). + + Memory: the tables are (m_max+1, 3, S, npts) and (2*m_max+1, 5, S, npts) + complex -- independent of every grid size. The sample scan runs in + chunks of ``sample_chunk`` grid points (default npsi_s, i.e. one phi row + per step), checkpointed so reverse-mode AD does not store per-step + intermediates. + + Returns ``(C_A, C_B, meta)`` with ``meta = dict(m_max, nphi_s, npsi_s)``. + """ + m_max = _data_m_max(data) + nphi_s, npsi_s = angle_sample_grid_sizes(m_max) + if sample_chunk is None: + sample_chunk = npsi_s + KPA, KSA = m_max + 1, 1 # ks in -KSA..KSA + KPB, KSB = 2 * m_max + 1, 2 + + phi_s = np.linspace(0.0, 2.0 * np.pi, nphi_s, endpoint=False) + psi_s = np.linspace(0.0, np.pi, npsi_s, endpoint=False) + PH, PS = np.meshgrid(phi_s, psi_s, indexing="ij") + pairs = np.stack([PH.ravel(), PS.ravel()], axis=-1) # (Ns, 2) + Ns = pairs.shape[0] + if Ns % sample_chunk: + raise ValueError("sample_chunk must divide nphi_s*npsi_s") + + def _phase_table(kp_max, ks_max): + kp = np.arange(kp_max + 1) + ks = np.arange(-ks_max, ks_max + 1) + return np.exp(-1j * (pairs[:, 0, None, None] * kp[None, :, None] + + 2.0 * pairs[:, 1, None, None] * ks[None, None, :]) + ) / Ns # (Ns, KP, KS) + + phase_A = _phase_table(m_max, KSA) + phase_B = _phase_table(2 * m_max, KSB) + + ra = jnp.asarray(ra, dtype=jnp.float64) + dec = jnp.asarray(dec, dtype=jnp.float64) + incl = jnp.asarray(incl, dtype=jnp.float64) + S = ra.shape[0] + npts = data.npts + c = int(sample_chunk) + nsteps = Ns // c + + xs = (jnp.asarray(pairs.reshape(nsteps, c, 2)), + jnp.asarray(phase_A.reshape(nsteps, c, KPA, 2 * KSA + 1)), + jnp.asarray(phase_B.reshape(nsteps, c, KPB, 2 * KSB + 1))) + + def _step(carry, x): + CA, CB = carry + prs, pA, pB = x # (c,2),(c,KPA,3),(c,KPB,5) + # batch the c grid points against the S parameter rows: (c*S,) + ra_b = jnp.broadcast_to(ra[None, :], (c, S)).reshape(-1) + dec_b = jnp.broadcast_to(dec[None, :], (c, S)).reshape(-1) + incl_b = jnp.broadcast_to(incl[None, :], (c, S)).reshape(-1) + phi_b = jnp.broadcast_to(prs[:, 0][:, None], (c, S)).reshape(-1) + psi_b = jnp.broadcast_to(prs[:, 1][:, None], (c, S)).reshape(-1) + ku, rs = _accumulate_unit(data, ra_b, dec_b, psi_b, incl_b, phi_b, + interp, False) + A = ku.real.reshape(c, S, npts) + B = rs.reshape(c, S, npts) + CA = CA + jnp.einsum("ckq,cst->kqst", pA, A) + CB = CB + jnp.einsum("ckq,cst->kqst", pB, B) + return (CA, CB), None + + CA0 = jnp.zeros((KPA, 2 * KSA + 1, S, npts), dtype=jnp.complex128) + CB0 = jnp.zeros((KPB, 2 * KSB + 1, S, npts), dtype=jnp.complex128) + (C_A, C_B), _ = jax.lax.scan(jax.checkpoint(_step), (CA0, CB0), xs) + meta = dict(m_max=m_max, nphi_s=nphi_s, npsi_s=npsi_s) + return C_A, C_B, meta + + +def _kp_weights(KP): + w = np.ones(KP) + w[1:] = 2.0 + return jnp.asarray(w) + + +def _reconstruct_field(C, phi, u): + """Evaluate the real trig polynomial with coefficient table C at (phi, u). + + C: (KP, 2*KS+1, S, npts) complex from :func:`angle_coefficient_tables`; + phi, u: (c,) points. Returns (c, S, npts) float64. + """ + KP = C.shape[0] + KS = (C.shape[1] - 1) // 2 + kp = jnp.arange(KP, dtype=jnp.float64) + ks = jnp.arange(-KS, KS + 1, dtype=jnp.float64) + E = jnp.exp(1j * (phi[:, None, None] * kp[None, :, None] + + u[:, None, None] * ks[None, None, :])) # (c,KP,KS) + E = E * _kp_weights(KP)[None, :, None] + return jnp.einsum("ckq,kqst->cst", E, C).real + + +def _dense_grid_sizes(amp, m_max=2): + """(nphi_d, nu_d) dense reconstruction sizes adequate for amplitude ``amp`` + and mode content ``m_max``. + + Derived from the trapezoid aliasing error of exp(trig poly of amplitude + A): N = K*sqrt(A) with the calibrated constants above (>= 2x margin, + calibrated at m_max = 2) and hard floors. The phi axis ADDITIONALLY + scales linearly with m_max: the exponent's phi content extends to + harmonic ~(2*m_max)*sqrt(A)-ish, so amplitude alone under-resolves + higher modes (external review 3, P1: a pure order-8 term at A = 450 was + phase-dependently wrong by ~0.037 nats, order 16 by up to 1.2 nats, + under the amplitude-only rule). The u axis never scales with m_max -- + psi enters at spin-2 for every mode. This is NOT a settable knob; + callers pass the DATA-DERIVED amplitude bound from + :func:`estimate_angle_amplitude` (floored at the crossover) and the + data's m_max. + """ + amp = max(float(amp), 25.0) + m_scale = max(1.0, float(m_max) / 2.0) # calibration point is m_max=2 + n_u = max(_DENSE_FLOOR_U, int(np.ceil(_DENSE_K_U * np.sqrt(amp)))) + n_phi = max(int(np.ceil(_DENSE_FLOOR_PHI * m_scale)), + int(np.ceil(_DENSE_K_PHI * m_scale * np.sqrt(amp)))) + # round up to multiples of 16 so chunking stays regular + n_u = ((n_u + 15) // 16) * 16 + n_phi = ((n_phi + 15) // 16) * 16 + return n_phi, n_u + + +ANGLE_AMP_SKY_POINTS = 64 # sky/inclination draws for the amplitude bound +ANGLE_AMP_MARGIN = 2.0 # covers the finite sky sample: the amplitude is + # a smooth O(1)-varying function of sky position, + # and the sizing error enters only through + # sqrt(amp), so a 2x margin in amplitude is a + # 1.4x margin in grid size +ENDPOINT_GUARD_BAND = 15.0 # nats below the amplitude maximum within which a + # sky point's dominant peak still counts for the + # near-boundary diagnostic below. Applied per sky + # point against that point's own maximum AND again + # against the global one, so what survives is + # within [BAND, 2*BAND] nats of the global maximum + # -- a bound worth stating, since the two-stage + # form is what avoids a second pass over the + # per-entry arrays the sky loop exists to keep + # transient. + + +def _endpoint_bell(k): + """``k exp(-k^2/2)``, the shape of a Gaussian's DERIVATIVE at a point ``k`` + widths from its peak, clamped at 0 for exterior (non-positive) clearances. + + This is the only ``k`` dependence in the truncated-endpoint error term of + :func:`core.loguniform_endpoint_error`; it is maximal (0.6065) at exactly + ONE width of clearance, and vanishes both far out in the tail AND at the + peak itself -- an endpoint sitting ON the peak is a stationary point, where + the Euler-Maclaurin endpoint term has nothing to correct. So "closer to the + edge" is NOT monotonically worse, and a guard shaped as a bare minimum + clearance would refuse the k -> 0 case that the alias law already covers. + """ + k = np.maximum(np.asarray(k, dtype=float), 0.0) + return k * np.exp(-0.5 * np.square(k)) + + +def _peak_clearance(A, B, x_min, x_max): + """``(rho, clearance from the d_min edge, clearance from the d_max edge)``. + + The distance integrand ``exp(x A - x^2 B/2)`` is a Gaussian in ``x`` peaked + at ``x* = A/B``, whose width in ``ln d`` is ``1/rho`` with + ``rho = A/sqrt(B)`` -- scale free, which is the whole basis of the + log-uniform grid (see :func:`core.make_distance_grid_loguniform`). The + clearances are that peak's distance from each prior edge IN THOSE UNITS, + which is what decides whether the untruncated alias law applies. + + ``x = distMpcRef/d`` inverts the edges: the ``d_min`` edge is ``x_max``. + Non-positive clearances mean the maximizer is EXTERIOR -- the boundary-layer + regime section 1a refuses outright -- not a resolved peak, and callers must + treat them as such rather than as "very close to the edge". + """ + A = np.asarray(A, dtype=float) + B = np.asarray(B, dtype=float) + Bs = np.maximum(B, 1e-300) + rho = np.where((A > 0.0) & (B > 0.0), A / np.sqrt(Bs), 0.0) + x_star = np.where(rho > 0.0, A / Bs, 0.0) + ok = x_star > 0.0 + xs = np.where(ok, x_star, 1.0) + k_lo = np.where(ok, rho * np.log(float(x_max) / xs), 0.0) + k_hi = np.where(ok, rho * np.log(xs / float(x_min)), 0.0) + return rho, k_lo, k_hi + + +def estimate_angle_amplitude(data, x_grid, interp=JAX_INTERP_DEFAULT, + n_sky=ANGLE_AMP_SKY_POINTS, seed=0, + margin=ANGLE_AMP_MARGIN, + _n_phi_e=None, _n_u_e=24, + return_diagnostics=False): + """DATA-DERIVED bound on the (phi, psi)-exponent amplitude A. + + This is the number that sizes the dense reconstruction grids, and it is + computed from the very function being integrated -- NOT from a caller's + SNR estimate. (External review, correctly: sizing from ``guess_snr`` + meant a missing or underestimated SNR silently under-resolved the dense + quadrature, quietly reintroducing exactly the failure mode this module + exists to remove -- the n_psi=8 defect again, one level up.) + + Method: evaluate the coefficient tables EAGERLY (concrete numpy inputs, + build time -- grid sizes must be static under jit, so this cannot run + inside the traced likelihood) at ``n_sky`` random sky/inclination + points. The PRIMARY estimate is the EMPIRICAL maximum of the exponent + over a dense angular reconstruction: A and B are trig polynomials of + known order (<= (2*m_max, 2)), so a 96 x 24 grid reconstructs them + exactly up to interpolation and the grid max understates the continuum + max by < 1% (peak offset <= half a cell, curvature <= (k_max)^2 A) -- + absorbed in ``margin``. Per angle point the distance max is closed + form: B >= 0 makes x*A - x^2/2*B concave in x, so the max over the + ACTUAL x support is at clip(A/B, x_min, x_max). + + THIS IS AN ESTIMATOR, NOT A PROVEN BOUND (review 3, P2): the sky + maximum is located by sampling, and no finite sample proves a maximum + was not missed. Three mechanisms stand behind it instead of a + bound-shaped claim: + + 1. the sky sample is random draws PLUS deterministic extremes (the + inclination poles and a coarse uniform sky grid, which cover the + known response maximizers); + 2. a split-half convergence check: if the second half of the sample + moves the running maximum by more than 20%, the sample is doubled + (up to twice) and the growth is printed -- an under-sampled sky + variation is detected empirically rather than assumed away; + 3. a RUNTIME fail-safe in the fused functions: every jitted likelihood + call recomputes the amplitude its own coefficient tables reach and + prints a loud warning if it exceeds the amp_sizing the dense grids + were built for -- so an underestimate is DETECTED at the point of + use, with the recourse (rebuild with the reported amplitude) named + in the message. The failure mode is never silent. + + A second, analytic expression max_x (x*M_A - x^2/2*B0)+ (M_A = sum + w|C_A| pointwise-bounds |A|; B0 = angular mean of B) is kept as a + build-time cross-check: it pairs the max of A with the MEAN of B, which + review 2 already noted is heuristic (B can dip below its mean where A + peaks); empirically it over-reads by 1.5-1.9x, and if it ever reads + BELOW the empirical max the disagreement is printed. + + Returns ``margin`` times the empirical max, UNfloored: the auto selector + compares it to the crossover (a floor here would push every quiet target + into the laplace branch); the WRAPPER floors the SIZING amplitude at the + crossover separately, so grids are never sized below the calibration + point. + """ + x = np.asarray(x_grid) + x_min, x_max = float(x.min()), float(x.max()) + + def _per_sky_amps(ra, dec, incl): + """Per-sky-point empirical amplitude maxima from exact reconstruction.""" + C_A, C_B, meta = angle_coefficient_tables(data, ra, dec, incl, + interp=interp) + C_A = np.asarray(C_A) + C_B = np.asarray(C_B) + # dense angular reconstruction matrices (numpy mirror of + # _reconstruct_field). The grid is DERIVED from the mode content + # and ASSERTED, exactly like the sample grid (fail-closed): sampled + # at >= 8x per highest harmonic, a band-limited trig polynomial's + # grid max under-reads its continuum max by < 1% (absorbed in + # `margin`), while a coarser grid can miss an adversarially-phased + # n = 2*m_max harmonic entirely -- so a too-small grid is refused, + # not trusted. _n_phi_e/_n_u_e are test hooks; production callers + # take the derived defaults. + m_max_e = meta["m_max"] + n_phi_e = int(_n_phi_e) if _n_phi_e is not None \ + else max(96, 16 * (2 * m_max_e)) + n_u_e_ = int(_n_u_e) + assert n_phi_e >= 8 * (2 * m_max_e), \ + "estimator phi grid %d under-samples 2*m_max=%d content" \ + % (n_phi_e, 2 * m_max_e) + assert n_u_e_ >= 8 * 2, \ + "estimator u grid %d under-samples order-2 content" % (n_u_e_,) + PH, UU = np.meshgrid( + np.linspace(0, 2 * np.pi, n_phi_e, endpoint=False), + np.linspace(0, 2 * np.pi, n_u_e_, endpoint=False), + indexing="ij") + phis, us = PH.ravel(), UU.ravel() + + def _recon_matrix(KP, KS): + kp = np.arange(KP) + ks = np.arange(-KS, KS + 1) + E = np.exp(1j * (phis[:, None, None] * kp[None, :, None] + + us[:, None, None] * ks[None, None, :])) + w = np.ones(KP) + w[1:] = 2.0 + return (E * w[None, :, None]).reshape(len(phis), -1) + + E_A = _recon_matrix(C_A.shape[0], (C_A.shape[1] - 1) // 2) + E_B = _recon_matrix(C_B.shape[0], (C_B.shape[1] - 1) // 2) + amps = [] + amps_unclipped = [] + pk_A = [] + pk_B = [] + pk_ep = [] + pk_ext = [] + for j in range(C_A.shape[2]): # per-sky loop bounds the transient + A_g = (E_A @ C_A[:, :, j].reshape(-1, C_A.shape[-1])).real + B_g = np.maximum( + (E_B @ C_B[:, :, j].reshape(-1, C_B.shape[-1])).real, 0.0) + # B >= 0 makes x*A - x^2/2*B concave in x: the max over the + # actual support is at the clipped stationary point + x_hat = np.clip(A_g / np.maximum(B_g, 1e-300), x_min, x_max) + val = x_hat * A_g - 0.5 * np.square(x_hat) * B_g + amps.append(max(float(val.max()), 0.0)) + # UNCLIPPED companion: the stationary value A^2/(2B) itself, which + # is what the exponent would reach if the maximizing distance were + # inside the prior support. Only used to DETECT that it is not -- + # the returned amplitude is unchanged. A <= 0 puts the stationary + # point at negative x, where the max over x >= 0 is 0. + # The A_g > 0 arm is inert FOR THIS CONSUMER and load-bearing for + # a different one, which is the only useful thing to say about it. + # A^2 is even in A, so the arm changes nothing whenever the + # selection is an ARGMAX: the lattice maximum of A^2/(2B) sits at + # A > 0 anyway, and removing the arm leaves clip_excess + # bit-identical on every mode set tried ((2,+-2), +(2,+-1), + # +(3,+-3), interior and exterior supports). Nothing here can fail + # if it is deleted. It becomes load-bearing the moment a consumer + # selects a BAND rather than a maximum -- over a threshold the + # A < 0 branch is exactly degenerate with its mirror and survives + # the cut, which is a documented trap (a weight cut ranked on the + # unconstrained A^2/(2B) came out a factor ~300 wrong, and the tell + # was that the answer did not move between a 10-nat and a 100-nat + # threshold). So: keep it, and if you add a banded selection here, + # that is the point at which it needs a test. + val_u = np.where(A_g > 0.0, + np.square(A_g) / (2.0 * np.maximum(B_g, 1e-300)), + 0.0) + amps_unclipped.append(max(float(val_u.max()), 0.0)) + # (A, B) of THIS sky point's DOMINANT configuration -- the entry + # that attains its clipped maximum, i.e. the one whose distance + # integral this sky direction's marginal is made of. Kept (two + # scalars per sky point, not the arrays) so the near-boundary + # diagnostic below can be formed after the global maximum is known, + # without a second pass over the per-entry arrays. + i_hat = int(np.argmax(val)) + pk_A.append(float(A_g.ravel()[i_hat])) + pk_B.append(float(B_g.ravel()[i_hat])) + # ...and the endpoint term over EVERY entry of this sky point, not + # just its dominant one. A sky direction can have a far-interior + # maximum and a near-equal secondary (phi, psi, time) entry whose + # distance peak sits a width from an edge; that entry contributes + # materially to the marginal and carries the full endpoint error, + # while the dominant one scores ~0. Reducing per sky to a scalar + # keeps the memory bounded (the arrays never leave this loop) and + # is CONSERVATIVE: it is a max over more entries than the band + # would keep, and the rho^2 factor self-limits the quiet ones. + # Measured under-read of the argmax-only form on a quiet + # synthetic: up to 7562x (external review, P1). + rho_e, klo_e, khi_e = _peak_clearance(A_g.ravel(), B_g.ravel(), + x_min, x_max) + # The interior mask differs from no mask only when ONE edge is + # exterior and the OTHER is within a few widths -- i.e. on a + # support only a few 1/rho wide. No prior this code is run with is + # anywhere near that: [1,1000], [1,10000] and [1,3000] Mpc are + # 69-947 widths across at rho 10-103, and even a narrow [100,2000] + # box is 30-308, so the far edge's bell underflows to exactly 0 and + # the two forms are numerically identical. Dropping the mask + # therefore survives the gate, and that is an equivalent mutant + # rather than a coverage gap -- do not add a test for it. + ep = np.where((klo_e > 0.0) & (khi_e > 0.0), + np.square(rho_e) + * (_endpoint_bell(klo_e) + _endpoint_bell(khi_e)), + 0.0) + pk_ep.append(float(ep.max()) if ep.size else 0.0) + # ...and the loudest EXTERIOR entry of this sky point. The + # endpoint term above cannot represent one: its bell is clamped at + # k <= 0, and the Euler-Maclaurin expansion is the wrong instrument + # for a peak outside the support anyway (that is a boundary layer, + # section 1a). The exterior guards do not see it either -- + # clip_excess is a ratio of GLOBAL maxima and peak_clearance + # describes the global argmax -- so an interior dominant entry with + # a near-equal exterior secondary reads completely clean. Carried + # as its own scalar so the wrapper can band it. + _ext = ~((klo_e > 0.0) & (khi_e > 0.0)) + pk_ext.append(float(val.ravel()[_ext].max()) if _ext.any() + else -np.inf) + return (np.array(amps), np.array(amps_unclipped), + np.array(pk_A), np.array(pk_B), np.array(pk_ep), + np.array(pk_ext), C_A, C_B) + + def _draw(n, rng): + ra = rng.uniform(0.0, 2.0 * np.pi, n) + dec = np.arcsin(rng.uniform(-1.0, 1.0, n)) + incl = np.arccos(rng.uniform(-1.0, 1.0, n)) + return ra, dec, incl + + rng = np.random.default_rng(seed) + ra, dec, incl = _draw(n_sky, rng) + # deterministic extremes: face-on/face-off inclinations over a coarse + # uniform sky grid (the known maximizers of the response amplitude) + g_ra, g_dec = np.meshgrid(np.linspace(0, 2 * np.pi, 6, endpoint=False), + np.array([-1.0, -0.35, 0.35, 1.0]), + indexing="ij") + for i0_ in (0.0, np.pi): + ra = np.concatenate([ra, g_ra.ravel()]) + dec = np.concatenate([dec, g_dec.ravel()]) + incl = np.concatenate([incl, np.full(g_ra.size, i0_)]) + + (amps, amps_u, pk_A, pk_B, pk_ep, pk_ext, + C_A, C_B) = _per_sky_amps(ra, dec, incl) + # Concatenated across sky BATCHES, in the same idiom the two maxima use: + # the near-boundary diagnostic is formed after the loop, so it must see the + # re-drawn batches too or it reads a first batch that a later one displaced. + amps_cat, pk_A_cat, pk_B_cat = amps, pk_A, pk_B + pk_ep_cat, pk_ext_cat = pk_ep, pk_ext + # split-half convergence check (mechanism 2 of the docstring): compare + # the max WITHOUT the second half of the random draws against the max + # with them; growth > 20% means the sky variation is under-sampled, so + # draw again (at most twice) and say so. + half = np.concatenate([amps[: n_sky // 2], amps[n_sky:]]) # + extremes + amp_emp = float(amps.max()) + # Accumulated in the SAME idiom as amp_emp, and adjacently: dropping the + # unclipped update in the re-draw loop below takes clip_excess below 1 and + # disarms the F1 exteriority refusal, with no observable difference on any + # fixture. DESIGN_jax_distance_quadrature.md section 5(c); + # test_sky_doubling_updates_the_unclipped_maximum_too pins it at the source. + amp_u_emp = float(amps_u.max()) if len(amps_u) else 0.0 + amp_ref = float(half.max()) + grows = amp_emp > 1.2 * amp_ref + 1e-12 + n_extra = 0 + while grows and n_extra < 2 * n_sky: + print("estimate_angle_amplitude: sky maximum still growing " + "(%.4g -> %.4g); doubling the sample." % (amp_ref, amp_emp)) + ra2, dec2, incl2 = _draw(n_sky, rng) + (amps2, amps_u2, pk_A2, pk_B2, pk_ep2, pk_ext2, + _, _) = _per_sky_amps(ra2, dec2, incl2) + amp_ref = amp_emp + amp_emp = max(amp_emp, float(amps2.max())) + amp_u_emp = max(amp_u_emp, float(amps_u2.max())) + grows = amp_emp > 1.2 * amp_ref + 1e-12 + n_extra += n_sky + amps_cat = np.concatenate([amps_cat, amps2]) + pk_A_cat = np.concatenate([pk_A_cat, pk_A2]) + pk_B_cat = np.concatenate([pk_B_cat, pk_B2]) + pk_ep_cat = np.concatenate([pk_ep_cat, pk_ep2]) + pk_ext_cat = np.concatenate([pk_ext_cat, pk_ext2]) + + # analytic cross-check (mechanism documented above; heuristic direction) + w = np.ones(C_A.shape[0]) + w[1:] = 2.0 + M_A = np.einsum("k,kqst->st", w, np.abs(C_A)) + ks0 = (C_B.shape[1] - 1) // 2 + B0 = np.maximum(C_B[0, ks0].real, 0.0) + expo = (x[:, None, None] * M_A[None] + - 0.5 * np.square(x)[:, None, None] * B0[None]) + amp_analytic = float(np.clip(expo, 0.0, None).max()) + if amp_analytic < amp_emp * (1.0 - 1e-9): + print("estimate_angle_amplitude: analytic cross-check %.6g fell " + "BELOW the empirical max %.6g (the review-flagged heuristic " + "direction); the empirical value governs." % (amp_analytic, + amp_emp)) + if return_diagnostics: + amp_unclipped = amp_u_emp + # NEAR-BOUNDARY diagnostic, for the log-uniform distance grid's OTHER + # precondition: clip_excess sees a maximizer that has left the support, + # but the spacing law is a statement about an effectively UNTRUNCATED + # Gaussian, and a peak that is interior yet only ~1 width inside an edge + # breaks it while clip_excess reads exactly 1. Reduced here to one + # scalar the wrapper turns into an error with its own spacing: + # max over the loud sky points of rho^2 * (bell(k_lo) + bell(k_hi)), + # which is the Euler-Maclaurin endpoint term stripped of the grid factor + # (core.loguniform_endpoint_error puts it back). BOTH edges, summed: + # they are separate corrections and a narrow support has both. + # + # WHAT IS AND IS NOT COVERED. EVERY reconstructed (phi, psi, time) + # entry of every sky point within ENDPOINT_GUARD_BAND of the maximum -- + # not just each point's dominant configuration. An earlier revision + # kept only the per-sky argmax, so a sky direction with a far-interior + # maximum and a near-equal secondary entry one width from an edge + # scored ~0 while the secondary carried the full endpoint error; + # measured under-read up to 7562x (external review, P1). The per-entry + # max is formed inside _per_sky_amps and reduced to one scalar per sky + # there, so the per-entry arrays never leave that loop. The band is + # still on the sky point, which makes this conservative: it covers + # entries a per-entry band would drop, and the rho^2 factor self-limits + # them. The band is a threshold on the CLIPPED value, never on + # A^2/(2B): the A < 0 mirror is exactly degenerate under an + # unconstrained ranking and survives such a cut (trap in + # _per_sky_amps). + # The band is a threshold on the CLIPPED value, never on A^2/(2B): the + # A < 0 mirror is exactly degenerate under an unconstrained ranking and + # survives such a cut, which is the trap recorded in _per_sky_amps. + keep = amps_cat >= amp_emp - ENDPOINT_GUARD_BAND + rho_pk, k_lo, k_hi = _peak_clearance(pk_A_cat, pk_B_cat, x_min, x_max) + endpoint_scale = (float(np.where(keep, pk_ep_cat, 0.0).max()) + if pk_ep_cat.size else 0.0) + _ext_best = (float(np.max(pk_ext_cat)) if pk_ext_cat.size + and np.isfinite(pk_ext_cat).any() else -np.inf) + i_dom = int(np.argmax(amps_cat)) if amps_cat.size else 0 + return margin * amp_emp, dict( + amp_clipped=float(amp_emp), + amp_unclipped=amp_unclipped, + # the Euler-Maclaurin endpoint term, grid-independent half + endpoint_scale=endpoint_scale, + # nats by which the loudest EXTERIOR entry sits BELOW the global + # maximum. inf when there is none. Small means a boundary-layer + # configuration is contributing materially while every other + # diagnostic here reads clean -- see the wrapper's refusal. + exterior_gap=(float(amp_emp - _ext_best) + if np.isfinite(_ext_best) else float("inf")), + # the globally dominant peak itself, reported so a refusal can name + # a number the caller can act on (and so a test can BUILD a support + # with a chosen clearance rather than hunt for one) + peak_x=(float(pk_A_cat[i_dom] / max(pk_B_cat[i_dom], 1e-300)) + if amps_cat.size else 0.0), + peak_rho=float(rho_pk[i_dom]) if rho_pk.size else 0.0, + peak_clearance=float(min(k_lo[i_dom], k_hi[i_dom])) + if rho_pk.size else 0.0, + # > 1 means the exponent's maximizing distance x* = A/B lies + # OUTSIDE [x_min, x_max] for the dominant angles, i.e. the + # distance posterior rails against a prior edge. See + # DESIGN_jax_distance_quadrature.md section 1a. + clip_excess=(amp_unclipped / amp_emp) if amp_emp > 0.0 + else (float("inf") if amp_unclipped > 0.0 else 1.0), + x_min=x_min, x_max=x_max) + return margin * amp_emp + + +AMP_FAILSAFE_TRIP_FACTOR = 2.0 +# The multiple of amp_sizing at which _runtime_amp_failsafe speaks up, and +# therefore the largest amplitude a run may reach while remaining UNLABELLED. +# It is a named constant because a second consumer now sizes from it: the +# log-uniform distance grid resolves peaks up to +# rho = sqrt(2 * TRIP * amp_sizing), so that every call the guard admits in +# silence is one the distance spacing already covers. Sizing that grid from +# amp_sizing itself (as an earlier draft did) leaves a factor sqrt(TRIP) in SNR +# -- and hence, through the Gaussian alias law, a tol -> sqrt(2*tol) hole -- +# open BELOW the trigger, where nothing is printed and nothing is recorded. +# Raise this and the distance grid follows automatically; that coupling is the +# point of the constant. + + +_AMP_FAILSAFE = {"tripped": False, "n_calls": 0, "worst_amp": 0.0, + "amp_sizing": None, "scheme": None} + + +def reset_amp_failsafe(): + """Clear the undersizing record (call once per event, before sampling). + + Barriers first: an in-flight callback from the PREVIOUS event must not land + after the reset and mislabel this one. + """ + try: + jax.effects_barrier() + except Exception: + pass + _AMP_FAILSAFE.update(tripped=False, n_calls=0, worst_amp=0.0, + amp_sizing=None, scheme=None) + + +def amp_failsafe_state(barrier=True): + """Host-side record of whether the dense grids were ever undersized. + + ``barrier=True`` calls :func:`jax.effects_barrier` first, so queued debug + callbacks have landed before the record is read. Without it a caller can + read CLEAN while a tripped callback is still in flight, or reset for the + next event before the previous event's callback arrives. + + Returns a dict; ``tripped`` is the load-bearing field. Consumers should + LABEL their output rather than discard it -- see the note in + :func:`_runtime_amp_failsafe` about why this is not fatal and not a NaN. + """ + if barrier: + try: + jax.effects_barrier() + except Exception: + pass + return dict(_AMP_FAILSAFE) + + +def _record_amp_failsafe(tripped, amp_call, amp_sizing, scheme_name): + """Host callback. Runs outside the traced graph; never alters a value.""" + _AMP_FAILSAFE["n_calls"] += 1 + if bool(tripped): + _AMP_FAILSAFE["tripped"] = True + _AMP_FAILSAFE["worst_amp"] = max(_AMP_FAILSAFE["worst_amp"], float(amp_call)) + _AMP_FAILSAFE["amp_sizing"] = float(amp_sizing) + _AMP_FAILSAFE["scheme"] = scheme_name + + +def _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, scheme_name): + """Mechanism 3 of :func:`estimate_angle_amplitude`'s contract: DETECT, at + the point of use, a call whose coefficient tables reach amplitudes the + dense grids were not sized for (i.e. the build-time estimator missed a + hotter sky region -- it is an estimator, not a proven bound). + + Uses the cheap analytic expression max over (S, t) of the concave-in-x + maximum of x*M_A - x^2/2*B0 (closed form, no distance-grid axis). That + expression OVER-reads the true amplitude by a measured 1.5-1.9x, so the + trigger threshold is 2*amp_sizing: it fires when the true local + amplitude exceeds ~1.3-2x the sizing bound -- comfortably BEFORE the + dense grids actually degrade (their calibrated constants carry a 2x + margin in N, i.e. 4x in amplitude). The warning prints from inside jit + via jax.debug.print (no value is altered; the recourse is named in the + message). Everything under stop_gradient: the check must not appear in + the AD graph. + """ + w = _kp_weights(C_A.shape[0]) + M_A = jnp.einsum("k,kqst->st", jnp.asarray(w), jnp.abs(C_A)) + ks0 = (C_B.shape[1] - 1) // 2 + B0 = jnp.maximum(C_B[0, ks0].real, 0.0) + x_min = jnp.min(x_grid) + x_max = jnp.max(x_grid) + x_hat = jnp.clip(M_A / jnp.maximum(B0, 1e-300), x_min, x_max) + amp_call = jnp.max(jnp.clip( + x_hat * M_A - 0.5 * jnp.square(x_hat) * B0, 0.0, None)) + amp_call = jax.lax.stop_gradient(amp_call) + # The hazard this check answers: a warning printed from inside jit does not stop + # anything, so a production run could finish and publish biased likelihoods, samples + # and evidence while the "fail-safe" scrolled past in a log. The recourse chosen is + # a HOST-RECORDED LABEL, not a poisoned value -- see the block below, which gives the + # reasoning and the two rejected alternatives. This function returns None; it alters + # no value. Everything is under stop_gradient so the check never enters the AD graph. + jax.lax.cond( + amp_call > AMP_FAILSAFE_TRIP_FACTOR * amp_sizing, + lambda a_: jax.debug.print( + "WARNING anglemarg/" + scheme_name + ": this call's coefficient " + "tables reach an amplitude scale ~{a:.4g} (analytic over-reading " + "expression), above " + + "%gx the amp_sizing=%.4g" % (AMP_FAILSAFE_TRIP_FACTOR, amp_sizing) + + " the dense (phi,psi) grids were built for. " + "estimate_angle_amplitude underestimated the sky maximum; the " + "marginal may be under-resolved at such points. Rebuild the " + "likelihood with amp_sizing >= the reported amplitude.", a=a_), + lambda a_: None, + amp_call) + # DELIBERATELY NOT FATAL, AND DELIBERATELY NOT A NaN. + # + # An earlier version returned NaN to "fail closed". That was worse than the + # warning it replaced: every consumer of this likelihood FILTERS non-finite + # values -- flowMC/MALA reject such proposals as invalid, the SMC path drops + # non-finite lnL, and write_samples discards non-finite rows. So a NaN over + # a hot sky region the estimator missed does not stop anything; it silently + # EXCISES exactly that region and publishes a clean-looking posterior over + # what remains. Invisible mutilation beats visible failure only from the + # code's point of view, never from the operator's. + # + # Aborting is also wrong here: this is a configuration estimate, and hard + # failure would destroy a multi-hour run over a recoverable condition. + # + # So: the value is untouched, the run completes, and the condition is + # recorded on the HOST so the driver can LABEL the result as suspect in its + # provenance. A labelled result an operator can judge beats both a vanished + # region and a dead run. + # The callback sits INSIDE lax.cond so the ORDINARY path has no host + # callback at all. An unconditional callback fires once per likelihood + # evaluation -- once per MALA/flowMC proposal, per chain -- transferring to + # the host and destroying accelerator throughput even when undersizing never + # happens. Only the rare tripped branch pays. + # + # Reliability caveat, stated because it bounds what this record can be used + # for: jax.debug.callback effects may be dropped, duplicated or reordered + # under transformation, and may land AFTER the result is ready. So this is + # a best-effort DIAGNOSTIC LABEL, not a correctness gate -- consumers must + # call jax.effects_barrier() before reading or resetting the state, and must + # not treat a clean read as proof of adequacy. + jax.lax.cond( + amp_call > AMP_FAILSAFE_TRIP_FACTOR * amp_sizing, + lambda a_: jax.debug.callback( + _record_amp_failsafe, True, a_, + jnp.asarray(amp_sizing, dtype=jnp.float64), scheme_name), + lambda a_: None, + amp_call) + + +def _require_amp_sizing(amp_sizing): + if amp_sizing is None: + raise ValueError( + "amp_sizing is required: pass a sound UPPER bound on the " + "(phi,psi)-exponent amplitude A ~ rho^2/2, e.g. " + "estimate_angle_amplitude(data, x_grid). There is deliberately " + "no default: a too-small value silently under-resolves the dense " + "quadrature, which is the defect this module exists to fix.") + return float(amp_sizing) + + +def _lse_update(m, s, e, axis=0): + """Running log-sum-exp: fold block ``e`` (reduced over ``axis``) into (m, s). + + Robust to all--inf blocks and an all--inf carry (both yield exp(-inf)=0 + rather than the exp(-inf - -inf) = nan of the naive update): padded chunk + tails and rejected Laplace bins produce -inf entries by design. + """ + m_blk = jnp.max(e, axis=axis) + m_safe = jnp.where(jnp.isfinite(m_blk), m_blk, 0.0) + s_blk = jnp.sum(jnp.exp(e - jnp.expand_dims(m_safe, axis)), axis=axis) + s_blk = jnp.where(jnp.isfinite(m_blk), s_blk, 0.0) + m_new = jnp.maximum(m, m_blk) + m_new_safe = jnp.where(jnp.isfinite(m_new), m_new, 0.0) + s_new = (s * jnp.exp(jnp.where(jnp.isfinite(m), m - m_new_safe, -jnp.inf)) + + s_blk * jnp.exp(jnp.where(jnp.isfinite(m_blk), + m_blk - m_new_safe, -jnp.inf))) + return m_new, s_new + + +def _pad_chunks(values, chunk): + """Split (N,) point arrays into (nsteps, chunk) with -inf log-pad weights.""" + N = values[0].shape[0] + nsteps = (N + chunk - 1) // chunk + pad = nsteps * chunk - N + lw = np.zeros(N) + out = [] + for v in values: + out.append(np.pad(v, (0, pad), mode="edge").reshape(nsteps, chunk)) + lw = np.pad(lw, (0, pad), constant_values=-np.inf).reshape(nsteps, chunk) + return [jnp.asarray(o) for o in out] + [jnp.asarray(lw)] + + +def fused_log_likelihood_distphipsimarg_exact( + data, ra, dec, incl, x_grid, log_w_grid, + interp=JAX_INTERP_DEFAULT, amp_sizing=None, + dense_chunk=8, grid_block=32, + time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False): + """Distance-, phi_ref- AND psi-marginalized lnL: exact-coefficient scheme. + + Drop-in replacement for :func:`core.fused_log_likelihood_distphipsimarg` + (same signature contract minus the two grid arguments, same normalization + convention: uniform priors dphi/2pi, dpsi/pi). The expensive likelihood + is sampled ONLY on the Nyquist grid fixed by mode content; the (phi, psi) + quadrature runs on a dense reconstruction whose size follows + :func:`_dense_grid_sizes` for ``amp_sizing`` -- a REQUIRED upper bound on + the exponent amplitude A ~ rho^2/2, obtained from + :func:`estimate_angle_amplitude` (the wrapper does this automatically). + There is no default: a silently-undersized grid is the defect this + module exists to fix. Honors JAX_ILE_DISTMARG_GH exactly as the grid + path does. + + Memory is bounded by ``dense_chunk`` (points per scan step), never by the + dense grid size: the largest transient is the inner distance-quadrature + slab (dense_chunk * S, npts, grid_block), ~0.8 GB f64 at the defaults for + a batched S=64, npts=614 call -- these two are COST/MEMORY knobs only, + with no effect on the result. + """ + x_grid = jnp.asarray(x_grid, dtype=jnp.float64) + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64) + C_A, C_B, meta = angle_coefficient_tables(data, ra, dec, incl, interp) + S = ra.shape[0] + npts = data.npts + + amp_sizing = _require_amp_sizing(amp_sizing) + _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, "exact") + nphi_d, nu_d = _dense_grid_sizes(amp_sizing, m_max=meta["m_max"]) + phi_d = np.linspace(0.0, 2.0 * np.pi, nphi_d, endpoint=False) + u_d = np.linspace(0.0, 2.0 * np.pi, nu_d, endpoint=False) # u = 2 psi + PH, UU = np.meshgrid(phi_d, u_d, indexing="ij") + c = int(dense_chunk) + phi_x, u_x, lw_x = _pad_chunks([PH.ravel(), UU.ravel()], c) + n_dense = nphi_d * nu_d + + a_g = x_grid + b_g = -0.5 * jnp.square(x_grid) + _use_gh = _core._DISTMARG_GH_N > 0 + if _use_gh: + gh_xi, gh_logw = make_distance_gh(_core._DISTMARG_GH_N) + x_min = jnp.min(x_grid) + x_max = jnp.max(x_grid) + + def _step(carry, x): + m, s = carry + phw, uw, lww = x + A = _reconstruct_field(C_A, phw, uw) # (c,S,npts) + B = _reconstruct_field(C_B, phw, uw) + K2 = A.reshape(c * S, npts) + R2 = B.reshape(c * S, npts) + if _use_gh: + lnL = _distmarg_gh_logL(K2, R2, gh_xi, gh_logw, x_min, x_max) + else: + lnL = _logsumexp_grid_blocked(K2, R2, a_g, b_g, log_w_grid, + grid_block) + lnL = lnL.reshape(c, S, npts) + lww[:, None, None] + m_new, s_new = _lse_update(m, s, lnL, axis=0) + return (m_new, s_new), None + + m0 = jnp.full((S, npts), -jnp.inf, dtype=jnp.float64) + s0 = jnp.zeros((S, npts), dtype=jnp.float64) + (m, s), _ = jax.lax.scan(jax.checkpoint(_step), (m0, s0), + (phi_x, u_x, lw_x)) + lnL_t = m + jnp.log(s) - jnp.log(float(n_dense)) + if return_lnLt: + return lnL_t + return _time_marginalize_terminal(lnL_t, data, time_quadrature) + + +# --------------------------------------------------------------------------- +# Analytic psi Laplace +# --------------------------------------------------------------------------- + +# Quadrature/Laplace handover (external reviews, items 2 and P2-local). +# +# HISTORY, because two revisions got this boundary wrong in two different +# ways: a hard series/Laplace switch at b + 2d = 0.5 left a wide window of +# 0.27-0.53 nats value error with a SIGN-INVERTED |c2| gradient just above +# the cut (review 2, item 2); a truncated-Bessel-series fix moved the +# handover to b + 2d ~ 16, but a third review showed the subdominance +# argument used to tolerate the remaining O(0.25)-nat worst-phase Laplace +# error near the handover is GLOBAL while the kernel is evaluated at every +# proposed sky position -- at a low-response proposal such bins are locally +# dominant (counterexample: b = 11.887, d = 3.516, beta = -1.890, +# delta = -0.650 -> +0.251 nats). +# +# The resolution: below the handover the u-integral is done by FIXED-N +# trapezoid quadrature of exp(f) -- for a periodic band-limited exponent the +# trapezoid rule converges super-exponentially, so with N = 320 the branch +# is machine-accurate for every t = b + 2d <= BLEND_HI = 300 (aliasing +# ~ I_{N/2}(t)/I_0(t) ~ e^-40 at the band edge; the counterexample bin is +# now exact). N is fixed by the HANDOVER amplitude, not by the data, so +# the laplace scheme keeps its ~sqrt(A) cost scaling. Above the band the +# enumerated-maxima Laplace applies, where its worst-phase error +# (~3.9/(b+d), adversarial constant from review 3) is <= ~0.026 and typical +# error is ~1e-3, falling as 1/A. The two are blended C^1-smoothly over +# [BLEND_LO, BLEND_HI]: the blend gradient carries dw * (quad - laplace), +# i.e. the branch disagreement, which the band placement bounds. +_LAPLACE_BLEND_LO = 220.0 # pure quadrature below this in t = b + 2d +_LAPLACE_BLEND_HI = 300.0 # pure Laplace above this +_LAPLACE_QUAD_N = 320 # u-quadrature points; content of exp(f) extends + # to ~2*5.25*sqrt(t/2) harmonics (d-dominated + # worst case) = 128 at t = 300, Nyquist 160 +_LAPLACE_BRACKET_CELLS = 24 # sign-scan cells for the stationary points of + # f(u): f' is a degree-2 trig polynomial, so it + # has AT MOST 4 transversal zeros on the circle + # (its resultant quartic 2d e^{-i delta} z^4 + + # b e^{-i beta} z^3 - b e^{i beta} z - + # 2d e^{i delta} = 0, z = e^{iu}); 24 cells + # (width 0.26) place adjacent zeros in distinct + # cells except for a MERGING max-min pair + # (f' = f'' = 0), which cannot contain the + # global maximum and carries negligible weight. +_LAPLACE_MAX_ROOTS = 4 + + +def _psi_lnI_amplitudes(c1, c2): + """(b, d, t_amp) for the kernel and the block dispatcher: harmonic + magnitudes and the blend variable t = b + 2 d. Factored out so the + dispatcher can bound t without tracing either branch.""" + mag1 = jnp.square(c1.real) + jnp.square(c1.imag) + mag2 = jnp.square(c2.real) + jnp.square(c2.imag) + b = jnp.sqrt(mag1 + 1e-300) + d = jnp.sqrt(mag2 + 1e-300) + return b, d, b + 2.0 * d + + +def _psi_lnI_lap_branch(a, c1, c2, b, t_amp): + """The enumerated-maxima Laplace branch of :func:`_laplace_psi_lnI` + (verbatim code motion; see that docstring for the contract). Returns + ln_laplace; -inf is impossible for a finite integral by the tolerant + acceptance, and hypothetical nonfinite values are guarded by callers.""" + lap_dummy = t_amp < _LAPLACE_BLEND_LO # blend weight is exactly 1 here + # jnp.where's VJP sends a ZERO cotangent through the unselected branch, + # and 0 * inf = nan: the Laplace branch must have BOUNDED derivatives + # (including second, for .fisher()) even on the pure-quadrature bins. + # Feed it safe dummy amplitudes there (their contribution is weighted 0 + # by the blend) and floor the curvature RELATIVE to the amplitude scale. + c1l = jnp.where(lap_dummy, 5.0 + 0.0j, c1) + c2l = jnp.where(lap_dummy, 0.25 + 0.0j, c2) + bl = jnp.sqrt(jnp.square(c1l.real) + jnp.square(c1l.imag) + 1e-300) + dl = jnp.sqrt(jnp.square(c2l.real) + jnp.square(c2l.imag) + 1e-300) + h_floor = 1e-6 * (bl + 4.0 * dl) + + def fval(u): + eiu = jnp.exp(1j * u) + return (c1l * eiu).real + (c2l * eiu * eiu).real + + def fp(u): + eiu = jnp.exp(1j * u) + return -(c1l * eiu).imag - 2.0 * (c2l * eiu * eiu).imag + + def fpp(u): + eiu = jnp.exp(1j * u) + return -(c1l * eiu).real - 4.0 * (c2l * eiu * eiu).real + + def fpppp(u): + eiu = jnp.exp(1j * u) + return (c1l * eiu).real + 16.0 * (c2l * eiu * eiu).real + + def _guard(H): + # sign-preserving denominator floor + return jnp.where(jnp.abs(H) >= h_floor, H, + jnp.where(H >= 0, h_floor, -h_floor)) + + # ---- bracket every transversal zero of f' (at most 4; see the constant) + # Signs are taken one grid node at a time so the transient stays one + # X-sized array; roots are assigned to at most _LAPLACE_MAX_ROOTS slot + # registers in encounter order. Interval-based bracketing cannot yield + # duplicate roots: the sign sequence flips exactly once per transversal + # crossing, including a crossing that sits exactly on a grid node. + # + # STRUCTURE, not mathematics: the cell walk is a lax.scan and the slot + # registers are one stacked (_LAPLACE_MAX_ROOTS, X) array rather than a + # Python loop over cells and slots. The elementwise updates are the + # same; the Python version unrolled ~24 x 4 select chains into the traced + # graph, and this kernel is instantiated once per distance block, which + # multiplied that unroll into an XLA graph that took over an hour (and + # >20 GiB) to compile at production sizes. Same fix pattern below for + # the bisection and the fixed-N quadrature. + N = _LAPLACE_BRACKET_CELLS + ug = np.linspace(0.0, 2.0 * np.pi, N + 1) + cell = ug[1] - ug[0] + # a, c1 and c2 may have different but broadcast-compatible shapes. The + # scan carry must start at their combined shape: lax.scan forbids the + # Python-loop behaviour of expanding a scalar carry on the first step. + zero_f = jnp.zeros_like(t_amp) + false_x = jnp.zeros_like(t_amp, dtype=bool) + nR = _LAPLACE_MAX_ROOTS + slot_ids = jnp.arange(nR, dtype=zero_f.dtype).reshape( + (nR,) + (1,) * zero_f.ndim) + s_prev0 = fp(jnp.asarray(ug[0])) >= 0 + los0 = jnp.stack([zero_f] * nR) + s_los0 = jnp.stack([false_x] * nR) + filled0 = jnp.stack([false_x] * nR) + + def _bracket_step(carry, edges): + s_prev, count, los, s_los, filled = carry + u_left, u_right = edges + s_next = fp(u_right) >= 0 + flip = s_prev != s_next + take = flip[None] & (count[None] == slot_ids) + los = jnp.where(take, u_left, los) + s_los = jnp.where(take, s_prev[None], s_los) + filled = filled | take + count = count + flip.astype(count.dtype) + return (s_next, count, los, s_los, filled), None + + (_, _, los, s_los, filled), _ = jax.lax.scan( + _bracket_step, (s_prev0, zero_f, los0, s_los0, filled0), + (jnp.asarray(ug[:-1]), jnp.asarray(ug[1:]))) + + # ---- per-slot bisection (value-only) + one differentiable polish step + # (batched over the slot axis; the 20 halvings are a fori_loop) + slo = s_los + def _bisect_step(_, lohi): # cell/2^20 ~ 2.5e-7, then Newton + lo, hi = lohi + mid = 0.5 * (lo + hi) + go_right = (fp(mid) >= 0) == slo + return (jnp.where(go_right, mid, lo), jnp.where(go_right, hi, mid)) + lo, hi = jax.lax.fori_loop(0, 20, _bisect_step, (los, los + cell)) + u0 = jax.lax.stop_gradient(0.5 * (lo + hi)) + u = u0 - fp(u0) / _guard(fpp(u0)) + H = fpp(u) + # tolerant acceptance: a maximum with H in [-h_floor, +h_floor) is a + # (near-)degenerate flat top; drop it and a finite integral could + # come back -inf, so keep it with the floored curvature instead + # (its Laplace weight is then merely inaccurate, never absent). + ok = filled & (H < h_floor) + Hm = jnp.minimum(H, -h_floor) + # Peak width: the Gaussian factor sqrt(2 pi/|H|) OVERESTIMATES a + # near-degenerate (quartic-flat) maximum by nats -- at the exactly + # aligned b = 4d configuration the floored-curvature form was ~5 + # nats high (review 3's local-error standard). The quartic width + # int exp(f4 u^4/24) du = Gamma(1/4)/2 * (24/|f4|)^(1/4) is closed + # form (f'''' is elementary for a trig polynomial). The widths are + # combined as W = W_g (1 + rho)^-1/2 with rho = (W_g/W_q)^2 -- exact + # at both ends and at most 0.083 nats off on the scale-free + # Gaussian x quartic family (vs 0.26 for min() and ~5 for the + # floored Gaussian alone) -- but GATED on rho: for a REGULAR + # maximum rho ~ 1/sqrt(b) and the ungated correction would inject + # an O(1/sqrt(A)) systematic where plain Laplace errs only O(1/A) + # (measured: sweep worst at t = 1000 rose 3.7e-3 -> 4.6e-2 + # ungated). The gate turns the correction on smoothly over + # rho in [0.2, 0.8], i.e. only where the peak is genuinely + # quartic-contaminated. + F4 = fpppp(u) + f4_floor = 1e-6 * (bl + 16.0 * dl) + F4m = jnp.minimum(F4, -f4_floor) + lnW_gauss = 0.5 * jnp.log(2.0 * jnp.pi / (-Hm)) + lnW_quart = 0.5949217316 + 0.25 * jnp.log(24.0 / (-F4m)) + rho = jnp.exp(jnp.clip(2.0 * (lnW_gauss - lnW_quart), -50.0, 50.0)) + g8 = jnp.clip((rho - 0.2) / 0.6, 0.0, 1.0) + g8 = g8 * g8 * (3.0 - 2.0 * g8) + lnW = lnW_gauss - 0.5 * jnp.log1p(rho * g8) + terms = jnp.where(ok, + a + fval(u) + lnW + - jnp.log(2.0 * jnp.pi), # (1/2 du/dpsi) * (1/pi) + -jnp.inf) # (nR,) + X + + # guarded log-add-exp over the root slots: an all--inf slot set has a NaN + # backward pass under the naive form, and the NaN leaks through jnp.where. + # Explicit left fold (not jnp.max/sum) preserves the pre-restructure + # accumulation order bit for bit. + mt = terms[0] + for j in range(1, nR): + mt = jnp.maximum(mt, terms[j]) + mts = jnp.where(jnp.isfinite(mt), mt, 0.0) + ssum = zero_f + for j in range(nR): + ssum = ssum + jnp.exp(terms[j] - mts) + ln_laplace = jnp.where(ssum > 0, + mts + jnp.log(jnp.maximum(ssum, 1e-300)), + -jnp.inf) + return ln_laplace + + +def _psi_lnI_quad_branch(a, c1, c2, b, n_quad): + """The fixed-N trapezoid u-quadrature branch of :func:`_laplace_psi_lnI` + (verbatim code motion, N parameterized; n_quad must be a multiple of the + 16-point scan chunk). Bit-identical to the pre-split code at + n_quad = _LAPLACE_QUAD_N. + """ + # ---- fixed-N u-quadrature branch: mean of exp(f) over a uniform u grid + # equals (1/pi) int dpsi. Uses the TRUE c1, c2 (no dummies needed: no + # divisions, and the running-max log-sum-exp keeps exp() in range even + # for the huge-t bins whose blend weight is 0). Chunked (as a lax.scan + # over precomputed host phase tables, one traced body instead of + # n_quad unrolled evaluations) so the transient stays a few + # X-sized arrays. + uq = np.linspace(0.0, 2.0 * np.pi, n_quad, endpoint=False) + QCH = 16 + e1 = np.exp(1j * uq) # host phases, as before + e2 = e1 * e1 # == eiu * eiu elementwise + p1 = jnp.asarray(e1.reshape(-1, QCH)) # (nq, QCH) + p2 = jnp.asarray(e2.reshape(-1, QCH)) + bshape = jnp.broadcast_shapes(jnp.shape(c1), jnp.shape(c2)) + pshape = (QCH,) + (1,) * len(bshape) + + def _quad_step(carry, phases): + mq, sq = carry + p1k, p2k = phases # (QCH,) complex + blk = ((c1[None] * p1k.reshape(pshape)).real + + (c2[None] * p2k.reshape(pshape)).real) + return _lse_update(mq, sq, blk, axis=0), None + + mq0 = jnp.full(bshape, -jnp.inf, dtype=b.dtype) + sq0 = jnp.zeros(bshape, dtype=b.dtype) + (mq, sq), _ = jax.lax.scan(_quad_step, (mq0, sq0), (p1, p2)) + ln_quad = (a + mq + jnp.log(jnp.maximum(sq, 1e-300)) + - jnp.log(float(n_quad))) + return ln_quad + + +def _laplace_psi_lnI(a, c1, c2): + """log[(1/pi) int_0^pi exp(a + Re(c1 e^{iu}) + Re(c2 e^{2iu})) dpsi], u = 2 psi. + + Two regimes, C^1-blended on t = b + 2d (b = |c1|, d = |c2|); see the + constants block above for the placement rationale and review history. + + t < BLEND_HI: fixed-N trapezoid quadrature of exp(f) over u -- machine- + accurate for a periodic band-limited exponent up to the handover, which + is what makes the kernel's LOCAL error small at every reachable bin (a + global-amplitude subdominance argument is not available: the kernel runs + at every proposed sky position, review 3). + + t > BLEND_LO: Laplace's method with ALL maxima enumerated. An early + revision seeded Newton only at the extrema of the FIRST harmonic, which + fails outright when that harmonic cancels (c1 = 0, c2 = -d: both seeds + are minima; -inf was returned for a finite integral -- review 1). Every + transversal zero of f' is bracketed by a sign scan (interval-based, so + coincident roots cannot be double-counted), bisected under + stop_gradient, polished by one differentiable Newton step (a contraction + step from the converged point carries the implicit derivative without a + deep 1/H^2 gradient chain); near-degenerate maxima are kept with floored + curvature rather than dropped, so -inf is impossible for a finite + integral. Angle-free throughout: f, f', f'' are evaluated directly from + c1, c2, so arg(0) never appears and b = 0 is a regular point. + + Elementary functions only (no scipy, no eigensolvers); differentiable; + any input shape (elementwise over broadcast a, c1, c2). + """ + # Materialize the documented elementwise broadcast before introducing + # the leading root-slot axis. Otherwise a data axis contributed only by + # ``a`` can collide with the four-root axis (silently when its length is + # four, or as a shape error for any other length). + a, c1, c2 = jnp.broadcast_arrays(a, c1, c2) + b, d, t_amp = _psi_lnI_amplitudes(c1, c2) + ln_laplace = _psi_lnI_lap_branch(a, c1, c2, b, t_amp) + ln_quad = _psi_lnI_quad_branch(a, c1, c2, b, _LAPLACE_QUAD_N) + # ---- C^1 blend: pure quadrature below LO, pure Laplace above HI + r = jnp.clip((_LAPLACE_BLEND_HI - t_amp) + / (_LAPLACE_BLEND_HI - _LAPLACE_BLEND_LO), 0.0, 1.0) + wgt = r * r * (3.0 - 2.0 * r) # smoothstep + # ln_laplace cannot be -inf for t_amp >= LO (a periodic f' has >= 2 sign + # flips and the tolerant acceptance keeps the global maximum), but guard + # the 0-weight product against a hypothetical -inf anyway. + ln_lap = jnp.where(jnp.isfinite(ln_laplace), ln_laplace, ln_quad) + return wgt * ln_quad + (1.0 - wgt) * ln_lap + + +# --------------------------------------------------------------------------- +# Block-dispatched execution of the kernel (2026-08-28 execution-cost fix). +# +# The C^1 blend evaluates BOTH branches at every lattice point, and the +# quadrature is sized for the handover amplitude (N = 320 at t = 300) +# regardless of the local t. Measured on the production-shaped lattice +# (amp_sizing ~ 1109, SNR-40 scale), 99.5% of (distance x dense-phi x sample +# x time) points sit at t < BLEND_LO -- 89% at t < 20 -- while every point +# that carries posterior weight sits at t > 900: each branch does needed +# work on a small, DISJOINT part of the lattice, yet both were paid +# everywhere (quad 55% / root-finding 39% of execution, additively). +# +# Per-point branching cannot save work under SIMD (select evaluates both +# sides), but the fused driver already evaluates the kernel in blocks +# (dist_block x phi_chunk x S x npts), and a block-level scalar bound on +# t = b + 2d makes the choice discrete via lax.switch (one branch executes): +# - every point has t >= BLEND_HI -> Laplace branch only; +# - every point has t < BLEND_LO -> quadrature only (weight is exactly +# 1 and the blend gradient exactly 0 there, so values AND derivatives +# equal the shipped kernel's), with N from the ladder below; +# - otherwise -> the full blended kernel, unchanged. +# +# The N ladder applies the SHIPPED sizing rule locally: the aliasing error +# of the N-point trapezoid rule on exp(f) is ~ I_{N/2}(t)/I_0(t), and the +# shipped pair (N = 320, t = 300) fixes the accepted exponent +# E(N, t) = sqrt(nu^2 + t^2) - nu asinh(nu/t) - t = -41.73 (nu = N/2) +# (the constants block above quotes e^-40 for the same pair). Each rung's +# threshold t_ok is rounded DOWN from the exact E = -41.73 contour, so every +# rung is at least as accurate as the shipped band edge: relative aliasing +# <= e^-41.7 ~ 8e-19, i.e. below f64 roundoff of the result. Rungs are +# multiples of the 16-point scan chunk. Exact contour values: +# N=32: 0.918, 48: 3.583, 64: 8.101, 96: 22.38, 128: 43.27, 160: 70.54, +# 224: 143.8, 320: 300 (test_angle_marg_block_dispatch.py recomputes these). +_QUAD_LADDER_N = (32, 48, 64, 96, 128, 160, 224, 320) +_QUAD_LADDER_TOK = (0.9, 3.5, 8.0, 22.0, 43.0, 70.0, 143.0, + _LAPLACE_BLEND_HI) + + +def _laplace_psi_lnI_block(a, c1, c2): + """Same value and derivatives as :func:`_laplace_psi_lnI` (see the + dispatch comment above for the exact equivalence statement), evaluated + with one lax.switch branch chosen by scalar bounds of t over the WHOLE + input block. Intended for the fused driver's per-(distance-block, + phi-chunk) kernel calls; for pointwise use, call _laplace_psi_lnI. + + Two deliberate differences from the shipped kernel, both confined to + cases the review history establishes as unreachable or sub-roundoff: + (1) in the pure-Laplace branch a hypothetical nonfinite ln_laplace + falls back to ``a`` instead of ln_quad (ln_quad is not computed there; + the fallback cannot fire for t >= BLEND_LO, see the blend comment); + (2) pure-quadrature blocks use the ladder N instead of N = 320, with + relative aliasing <= e^-41.7 at every rung edge (vs e^-41.7 at the + shipped band edge itself). + """ + # Keep the data axes distinct from the dispatcher's root-slot axis just + # as _laplace_psi_lnI does. The block dispatcher is the production call + # path, so it must preserve the same three-input broadcast contract. + a, c1, c2 = jnp.broadcast_arrays(a, c1, c2) + b, d, t_amp = _psi_lnI_amplitudes(c1, c2) + tmin = jnp.min(t_amp) + tmax = jnp.max(t_amp) + + def _pure_lap(_): + ln_l = _psi_lnI_lap_branch(a, c1, c2, b, t_amp) + return jnp.where(jnp.isfinite(ln_l), ln_l, a) + + def _quad_rung(nq): + def _q(_): + return _psi_lnI_quad_branch(a, c1, c2, b, nq) + return _q + + def _full(_): + return _laplace_psi_lnI(a, c1, c2) + + branches = ([_pure_lap] + [_quad_rung(n) for n in _QUAD_LADDER_N] + + [_full]) + n_rungs = len(_QUAD_LADDER_N) + rung = jnp.zeros((), dtype=jnp.int32) + for tok in _QUAD_LADDER_TOK[:-1]: + rung = rung + (tmax > tok).astype(jnp.int32) + idx = jnp.where(tmin >= _LAPLACE_BLEND_HI, jnp.int32(0), + jnp.where(tmax < _LAPLACE_BLEND_LO, + jnp.int32(1) + rung, jnp.int32(1 + n_rungs))) + return jax.lax.switch(idx, branches, None) + + +# --------------------------------------------------------------------------- +# psi-marginal node placement for the adaptive distance quadrature +# (JAX_ILE_DISTMARG_GH on the laplace path) +# +# core._distmarg_gh_logL places its frozen nodes at clip(K/R) +- 7/sqrt(R) for a +# FIXED psi. On this path psi has already been integrated out analytically, so +# the nodes have to bracket the psi-MARGINAL distance integrand +# I(x) = (1/pi) int dpsi exp(x A(u) - x^2/2 B(u)), u = 2 psi, +# a mixture over u of Gaussians of centre x*(u) = A(u)/B(u) and width +# 1/sqrt(B(u)). Two facts make a closed-form bracket sufficient: +# +# * every component is NARROWER than sigma = 1/sqrt(R_lo) with +# R_lo = B0 - |B1| - |B2| <= min_u B, so 7 sigma past the extreme +# component centre covers the mixture to the same 1e-11 the fixed-psi +# rule covers its single Gaussian; +# * the component centres that carry weight span a BOUNDED number of sigma. +# +# Both were measured on the ladder-2 injection (35+30 Msun, H1/L1/V1, SEOBNRv4, +# l_max = 2 -> lms {(2,-2),(2,2)}), at the sky points the campaign's own +# posterior occupies, over 235,776 (dense-phi, sample, time) bins per rung: +# +# rho W = sqrt(min_u B/R_lo) C = |x*(u_cf)-x*(u_exact)|/sigma S span/sigma +# 40.77 median 1.0000 max 1.0000 median 0.0014 p99 0.636 max 0.689 p99 3.899 max 4.085 +# 163.1 median 1.0000 max 1.0000 median 0.0090 p99 0.097 max 0.112 p99 0.863 max 0.887 +# +# with R_lo <= 0 at 0.0000% of ALL bins at both rungs. W == 1 is an IDENTITY, +# not a lucky bound: the spin-2 response F(psi) ~ e^{-2 i psi} puts the kappa +# term at exactly one u-harmonic and the rho^2 term at exactly harmonics 0 and +# 2, so A0 and B1 vanish for EVERY mode set -- measured |A0|/|A1| ~ 7e-17 and +# |B1|/|B0| ~ 6e-16 on real IMRPhenomXHM data at m_max = 2, 3 AND 4, and on +# synthetic data with random U/V. Hence B(u) = B0 + Re(B2 e^{2iu}) and +# R_lo = B0 - |B2| IS min_u B. The m_max gate below is therefore CONSERVATIVE +# rather than load-bearing for the width; it stands because the half-span +# constant was measured on (2,+-2) fixtures and the higher-mode verdict is +# owned by another session. +# +# CENTRING. A0 == 0 makes A(u) a pure first harmonic, so the u that maximises +# A is closed form; but the u that maximises the DISTANCE-maximum exponent +# A(u)^2/(2 B(u)) is what the bracket must sit on, and the two part company as +# |B2|/B0 grows. With A0 = B1 = 0 the stationary condition 2 A' B = A B' +# reduces, in z = e^{iu}, to +# z^2 (B0 A1 - conj(A1) B2) = conj(B0 A1 - conj(A1) B2) +# so with w = B0*A1 - conj(A1)*B2 the maximiser is EXACTLY +# e^{i u*} = +- conj(w)/|w| (sign chosen so A(u*) > 0) +# -- closed form, angle-free, and equal to conj(A1)/|A1| when B2 = 0. (B has +# only even u-harmonics, so E(u) = E(u+pi) and the two roots of z^2 are the +# same maximum; the other two stationary points are the A = 0 minima, divided +# out.) Checked against a 400,001-point brute-force argmax on 20,000 random +# (A1, B0, B2) with |B2|/B0 up to 0.99999: the brute force never beats it by +# more than 6.5e-16 relative. Using argmax A instead costs nothing on the +# ladder (the two centres differ by 0.64 sigma at rho 40.77, 0.10 at 163.08) +# but is catastrophic elsewhere in the family -- see below. +# +# HALF-SPAN. Because A0 = B1 = 0 hold for ANY mode set, the whole problem +# reduces after scaling to three numbers -- rho = |A1|/sqrt(B0), r = |B2|/B0, +# and the relative phase -- so the bracket can be scanned EXHAUSTIVELY instead +# of sampled on a fixture. Over 57,082 well-resolved points of that family +# (r up to 0.999, rho 1..1500, 61 phases, v grid 262144, weight threshold 100 +# nats), the one-sided reach that the half-span must cover is +# centre = argmax A : p50 8.1 p99 2701 MAX 7172 sigma +# centre = argmax A^2/(2B) : p50 3.5 p99 13.1 MAX 14.14 sigma +# and the 14.14 = sqrt(2 * 100 nats) bound is attained in the weak-signal +# corner where the 100-nat window is the whole circle. Hence 7 + 14.14 -> 22. +# On the ladder-2 injection itself the requirement is far smaller (7 + 3.46 = +# 10.5 sigma at rho 40.77, 7.8 at 163.08, 11.5 at 652), so the shipped span is +# ~2x what the operating point needs. ("Weight-carrying" = within 100 nats of +# the bin's own maximum over u of the clipped exponent; psi below that +# contribute < e^-100, and an under-reaching bracket can only UNDER-estimate +# them, never inflate them, the trapezoid being exponentially accurate at this +# spacing.) +_GH_PSI_HALF_SIGMA = 22.0 # node half-span, in units of sigma = 1/sqrt(R_lo) +_GH_PSI_MIN_NODES = 49 # floor: 44 sigma / 48 gaps = 0.92 sigma spacing, + # trapezoid aliasing on a Gaussian ~ 2e^-2pi^2/h^2 + # = 2e-10 -- below the f64 noise of the result +_GH_PSI_M_MAX = 2 # mode content the path is SHIPPED for. The + # A0 == B1 == 0 identity itself is structural and + # measured through m_max = 4, but the higher-mode + # verdict is owned elsewhere; keyed on mode + # content the way angle_sample_grid_sizes is. + + +def _gh_psi_node_offsets(n_nodes): + """Node offsets ``(z, z_prev, z_next, n)`` for the psi-marginal bracket. + + ``z`` spans +-``_GH_PSI_HALF_SIGMA`` instead of the fixed-psi rule's +-7, + and the count is raised in proportion so the NODE DENSITY the caller asked + for through JAX_ILE_DISTMARG_GH is preserved rather than diluted by the + wider bracket (floored at ``_GH_PSI_MIN_NODES``). + + ``z_prev``/``z_next`` are ``z`` with the INDEX clamped at the ends. The + composite-trapezoid weight of node k is then 0.5*(x[k+1] - x[k-1]) with the + same end convention as :func:`core._distmarg_gh_logL`'s + ``diff``-and-concatenate form -- identical weights, but computable one + block at a time, so the distance axis stays scanned and memory stays + bounded by ``dist_block``. + """ + n = max(int(_GH_PSI_MIN_NODES), + 1 + int(np.ceil((int(n_nodes) - 1) * _GH_PSI_HALF_SIGMA / 7.0))) + z = np.linspace(-_GH_PSI_HALF_SIGMA, _GH_PSI_HALF_SIGMA, n) + idx = np.arange(n) + return (z, z[np.maximum(idx - 1, 0)], z[np.minimum(idx + 1, n - 1)], n) + + +def fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, x_grid, log_w_grid, + interp=JAX_INTERP_DEFAULT, amp_sizing=None, + phi_chunk=16, dist_block=4, + time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False): + """Distance-, phi_ref- AND psi-marginalized lnL: analytic psi-Laplace scheme. + + Same contract and normalization as + :func:`fused_log_likelihood_distphipsimarg_exact`, but the psi axis is + removed analytically (see :func:`_laplace_psi_lnI`): at every + (dense-phi, distance-node, time) point the u-exponent coefficients follow + directly from the SAME coefficient tables, + + a = x A0(phi) - x^2/2 B0(phi) + c1 = x A1(phi) - x^2/2 B1(phi) (order e^{iu}) + c2 = - x^2/2 B2(phi) (order e^{2iu}) + + so no additional likelihood evaluations are needed. Cost scales ~sqrt(A) + (the dense phi axis) instead of ~A; the Laplace error is O(1/A) and + SHRINKS with SNR. + + The adaptive distance quadrature (JAX_ILE_DISTMARG_GH) is honoured for + ``m_max <= _GH_PSI_M_MAX`` via the psi-marginal node placement documented + above ``_gh_psi_node_offsets``; ``x_grid``/``log_w_grid`` then only supply + the support [x_min, x_max] and the prior normalization, exactly as on the + exact path. Richer mode content still RAISES rather than being silently + accepted: the placement rests on an A0 == B1 == 0 identity that is + established for (2,+-2) only. + + Two DIFFERENT axes, and conflating them has already misled a reader. The + paragraph above is about the PER-SAMPLE adaptive quadrature. The STATIC + distance grid is separate and is not restricted here at all: + ``--distance-grid-scheme loguniform`` is supported and gated on this path, + and needs no node-placement rule because it locates no peak -- one relative + spacing resolves every per-sample peak wherever it sits. See + DESIGN_jax_distance_quadrature.md. The two cannot be combined: with + JAX_ILE_DISTMARG_GH set the per-sample quadrature consumes only the SUPPORT + of ``x_grid``, so the log-uniform option would be bit-identically inert and + is refused rather than silently ignored. + + Memory is bounded by ``phi_chunk`` x ``dist_block``, never by grid sizes. + """ + # RESPONSE-MODEL PRECONDITION, before anything is built. This function is + # public (__all__) and is called directly by the wrapper and by several test + # modules, so a wrapper-only gate leaves a live bypass: a direct call with a + # banded response and m_max <= 2 would execute the unsupported placement + # while the wrapper correctly refused it. `feature` is a plain Python + # attribute -- static and trace-safe -- so unlike the numerical A0/B1 + # measurement (which needs concrete tables and therefore stays in the + # wrapper) it costs nothing, and checking it here also avoids paying for a + # coefficient-table build that is about to be rejected. + if _core._DISTMARG_GH_N > 0: + _feature = getattr(data, "feature", None) + if _feature not in _GH_PSI_STATIC_FEATURES: + raise ValueError( + "JAX_ILE_DISTMARG_GH is set, but the 'laplace' angle-marg " + "scheme's psi-marginal distance-node placement requires the " + "static detector response: it is DERIVED from A0 == 0 and " + "B1 == 0, which follow from F+(psi) + i Fx(psi) = " + "(F+(0) + i Fx(0)) e^{-2i psi}. This data has feature=%r, " + "which does not have that factorization. Use " + "--angle-marg-scheme exact, or unset JAX_ILE_DISTMARG_GH." + % (_feature,)) + x_grid = jnp.asarray(x_grid, dtype=jnp.float64) + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64) + C_A, C_B, meta = angle_coefficient_tables(data, ra, dec, incl, interp) + m_max = meta["m_max"] + S = ra.shape[0] + npts = data.npts + + _use_gh = _core._DISTMARG_GH_N > 0 + # This runs under jit/grad, where C_A and C_B are TRACERS, so the identity + # cannot be measured here -- it is a property of the DATA and is enforced + # once, on concrete tables, by JAXDistPhiPsiMargLikelihood (which gates + # EVERY scheme, not just 'auto'). A caller invoking this kernel directly + # under GH is responsible for calling gh_laplace_supported itself; the + # m_max test below is the only check available at trace time. + if _use_gh and int(m_max) > _GH_PSI_M_MAX: + raise ValueError( + "JAX_ILE_DISTMARG_GH is set and the 'laplace' angle-marg scheme's " + "psi-marginal distance-node placement is validated for mode " + "content m_max <= %d only (it rests on the A0 == 0 / B1 == 0 " + "identity); this data has m_max = %d. Use --angle-marg-scheme " + "exact, or unset JAX_ILE_DISTMARG_GH." + % (_GH_PSI_M_MAX, int(m_max))) + + amp_sizing = _require_amp_sizing(amp_sizing) + # x_grid is still the right argument under GH: the adaptive nodes are + # CLIPPED into [min x_grid, max x_grid], so the amplitude bound the + # failsafe computes over x_grid bounds the nodes actually used. + _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, "laplace") + nphi_d, _ = _dense_grid_sizes(amp_sizing, m_max=m_max) + phi_d = np.linspace(0.0, 2.0 * np.pi, nphi_d, endpoint=False) + c = int(phi_chunk) + phi_x, lw_x = _pad_chunks([phi_d], c) + + wA = _kp_weights(m_max + 1) + wB = _kp_weights(2 * m_max + 1) + kpA = jnp.arange(m_max + 1, dtype=jnp.float64) + kpB = jnp.arange(2 * m_max + 1, dtype=jnp.float64) + G = x_grid.shape[0] + blk = int(dist_block) + # Distance nodes packed into (n_dblk, blk) for the lax.scan below; the + # tail block (if G % blk) is edge-padded with -inf log-weights, exactly + # the _pad_chunks convention, so padded nodes contribute exactly 0 to the + # running log-sum-exp. A Python loop here instantiated the FULL + # _laplace_psi_lnI kernel once per distance block inside the (already + # checkpointed) phi scan body -- at the production n_grid=256, blk=4 that + # is 64 copies of an already-large kernel, and XLA compile time/memory on + # the resulting graph was the >1 h, >20 GiB wall the 2026-08-28 bake-off + # hit. The scan traces the kernel ONCE. Numerics are unchanged. + n_dblk = (G + blk - 1) // blk + pad_d = n_dblk * blk - G + if pad_d: + x_pad = jnp.concatenate( + [x_grid, jnp.broadcast_to(x_grid[-1], (pad_d,))]) + lw_pad = jnp.concatenate( + [log_w_grid, jnp.full((pad_d,), -jnp.inf, dtype=jnp.float64)]) + else: + x_pad, lw_pad = x_grid, log_w_grid + xg_blk = x_pad.reshape(n_dblk, blk) + lwg_blk = lw_pad.reshape(n_dblk, blk) + + if _use_gh: + # Adaptive nodes replace the static grid entirely; x_grid survives only + # as the physical support and the (dref-independent) prior norm C0, the + # same two roles it plays inside core._distmarg_gh_logL. + x_min = jnp.min(x_grid) + x_max = jnp.max(x_grid) + gh_C0 = jnp.log(3.0) - jnp.log(jnp.min(x_grid) ** (-3.0) + - jnp.max(x_grid) ** (-3.0)) + z_np, zp_np, zn_np, n_gh = _gh_psi_node_offsets(_core._DISTMARG_GH_N) + n_zblk = (n_gh + blk - 1) // blk + pad_z = n_zblk * blk - n_gh + pad_lw = np.zeros(n_gh) + if pad_z: + z_np = np.pad(z_np, (0, pad_z), mode="edge") + zp_np = np.pad(zp_np, (0, pad_z), mode="edge") + zn_np = np.pad(zn_np, (0, pad_z), mode="edge") + pad_lw = np.pad(pad_lw, (0, pad_z), constant_values=-np.inf) + zg_blk = jnp.asarray(z_np.reshape(n_zblk, blk), jnp.float64) + zpg_blk = jnp.asarray(zp_np.reshape(n_zblk, blk), jnp.float64) + zng_blk = jnp.asarray(zn_np.reshape(n_zblk, blk), jnp.float64) + zpad_blk = jnp.asarray(pad_lw.reshape(n_zblk, blk), jnp.float64) + # Never let the bracket exceed the physical support: as R_lo -> 0 (a + # bin with no response at all, where the exponent is flat in x) sigma + # would otherwise blow up and every node would clip onto one of the two + # rails. Capped, such a bin degrades to a uniform-in-x tiling of the + # support instead of a 2-point one. Inactive by ~3 orders of magnitude + # wherever the data carry signal (test_angle_marg_gh_laplace.py pins it). + gh_sigma_cap = (x_max - x_min) / (2.0 * _GH_PSI_HALF_SIGMA) + + def _step(carry, x): + m, s = carry + phw, lww = x # (c,) + # psi-Fourier coefficient FIELDS at the dense phi points (c,S,npts). + # Shared with gh_laplace_supported so the identity that predicate + # measures IS the one this placement depends on. + A0, A1, B0, B1, B2 = psi_harmonics_at_phi(C_A, C_B, phw, m_max) + + # distance quadrature: blocked, vectorized over the block (AD-fast), + # running log-sum-exp across blocks (a lax.scan; see the packing note + # above -- one traced kernel instead of G/blk unrolled copies) + def _dist_step(carry, xw): + mx, sx = carry + xgb, lwgb = xw # (blk,) + xg = xgb[:, None, None, None] # (g,1,1,1) + lwg = lwgb[:, None, None, None] + av = xg * A0[None] - 0.5 * jnp.square(xg) * B0[None] + c1 = xg * A1[None] - 0.5 * jnp.square(xg) * B1[None] + c2 = -0.5 * jnp.square(xg) * B2[None] + e = _laplace_psi_lnI_block(av, c1, c2) + lwg # (g,c,S,npts) + return _lse_update(mx, sx, e, axis=0), None + + if _use_gh: + # ---- psi-marginal adaptive node placement, all FROZEN ---------- + # Centre on the psi that maximises the (unclipped) distance-maximum + # exponent A(u)^2/(2 B(u)) -- available in CLOSED FORM here, see + # the derivation above _gh_psi_node_offsets: + # e^{i u*} = +- conj(w)/|w|, w = B0*A1 - conj(A1)*B2 + # with the sign picking the branch where A(u*) > 0 (x must be + # positive). Angle-free, so arg(0) never appears and w = 0 is a + # regular point; reduces to conj(A1)/|A1| -- the maximiser of A + # itself -- when B2 = 0. + w_st = B0 * A1 - jnp.conj(A1) * B2 + ph1 = jnp.conj(w_st) / jnp.maximum(jnp.abs(w_st), 1e-300) + sgn = jnp.where((A1 * ph1).real >= 0, 1.0, -1.0) + ph1 = ph1 * sgn # e^{i u*} + A_st = A0 + (A1 * ph1).real # A(u*) + B_st = B0 + (B1 * ph1).real + (B2 * ph1 * ph1).real + R_lo = B0 - jnp.abs(B1) - jnp.abs(B2) # <= min_u B + gh_center = jax.lax.stop_gradient( + jnp.clip(A_st / jnp.maximum(B_st, 1e-30), x_min, x_max)) + gh_sigma = jax.lax.stop_gradient( + jnp.minimum(1.0 / jnp.sqrt(jnp.maximum(R_lo, 1e-30)), + gh_sigma_cap)) + + def _gh_dist_step(carry, zw): + mx, sx = carry + zb, zpb, znb, zpadb = zw # (blk,) + + def _node(zz): + return jnp.clip( + gh_center[None] + gh_sigma[None] * zz[:, None, None, None], + x_min, x_max) + + xg = _node(zb) # (g,c,S,npts) + # composite-trapezoid weight, index-clamped at both ends: + # identical to core._distmarg_gh_logL's diff/concatenate form. + w = 0.5 * (_node(znb) - _node(zpb)) + pos = w > 0 # live (unclipped) + lwg = jnp.where(pos, jnp.log(jnp.where(pos, w, 1.0)) + - 4.0 * jnp.log(xg), -jnp.inf) + lwg = lwg + zpadb[:, None, None, None] # -inf on pad slots + av = xg * A0[None] - 0.5 * jnp.square(xg) * B0[None] + c1 = xg * A1[None] - 0.5 * jnp.square(xg) * B1[None] + c2 = -0.5 * jnp.square(xg) * B2[None] + e = _laplace_psi_lnI_block(av, c1, c2) + lwg + return _lse_update(mx, sx, e, axis=0), None + + mx0 = jnp.full((c, S, npts), -jnp.inf, dtype=jnp.float64) + sx0 = jnp.zeros((c, S, npts), dtype=jnp.float64) + (mx, sx), _ = jax.lax.scan( + _gh_dist_step, (mx0, sx0), + (zg_blk, zpg_blk, zng_blk, zpad_blk)) + lnI = (mx + jnp.where(sx > 0, jnp.log(jnp.maximum(sx, 1e-300)), + -jnp.inf) + + gh_C0 + lww[:, None, None]) # (c,S,npts) + m_new, s_new = _lse_update(m, s, lnI, axis=0) + return (m_new, s_new), None + + mx0 = jnp.full((c, S, npts), -jnp.inf, dtype=jnp.float64) + sx0 = jnp.zeros((c, S, npts), dtype=jnp.float64) + (mx, sx), _ = jax.lax.scan(_dist_step, (mx0, sx0), (xg_blk, lwg_blk)) + lnI = (mx + jnp.where(sx > 0, jnp.log(jnp.maximum(sx, 1e-300)), -jnp.inf) + + lww[:, None, None]) # (c,S,npts) + m_new, s_new = _lse_update(m, s, lnI, axis=0) + return (m_new, s_new), None + + m0 = jnp.full((S, npts), -jnp.inf, dtype=jnp.float64) + s0 = jnp.zeros((S, npts), dtype=jnp.float64) + (m, s), _ = jax.lax.scan(jax.checkpoint(_step), (m0, s0), (phi_x, lw_x)) + lnL_t = m + jnp.log(s) - jnp.log(float(nphi_d)) + if return_lnLt: + return lnL_t + return _time_marginalize_terminal(lnL_t, data, time_quadrature) + + +# Relative size at which A0 / B1 count as nonzero. The identity the psi-marginal +# node placement rests on (A0 == 0, B1 == 0, so R_lo = B0 - |B2| IS min_u B) is a +# property of the SPIN-2 detector response, not of the source, and is measured at +# ~1e-16 relative on every non-precessing mode set tried through m_max = 4. It is +# NOT measured under precession. +# +# WHY 1e-8, and what is NOT claimed for it. The identity is a STRUCTURAL +# precondition, not a numerical one: the closed-form psi maximiser below is +# DERIVED from A0 == 0 and B1 == 0 (that is what reduces stationarity to +# z^2 w = conj(w)). So the tolerance's job is to separate "numerically zero" +# from "structurally nonzero", not to bound an error. Observed values are +# ~1e-16 relative on every mode set tried, and the mutation sweep in +# test_angle_marg_gh_selection.py shows planted harmonics at 1e-3 are caught, +# so 1e-8 sits ~8 orders above the noise and ~5 below the smallest breach the +# tests exercise. +# +# An earlier revision of this comment also claimed 1e-8 was "~8 orders below a +# value that would move the bracket". That was never measured and is removed +# rather than left standing: an attempt to measure it produced a centre error +# FLAT at ~2 sigma across six decades of planted A0/B1, including where the +# identity holds -- a hand-rolled reimplementation of the maximiser failing its +# own flatness check, not a property of the code. If the upper end is ever +# wanted, measure it through the shipped kernel, not a re-derivation. +# +# Values are not bit-portable: they come through BLAS-heavy reconstruction in +# angle_coefficient_tables, so anything pinned off them needs a RELATIVE +# tolerance. This comparison is already relative and one-sided. +GH_PSI_IDENTITY_TOL = 1e-8 + +# The response models for which A0 == 0 / B1 == 0 hold at EVERY extrinsic point. +# This is the angle-independent half of the gate and it is the actual guarantee: +# the static path builds F = F+ + i Fx through compute_detamresponse (LAL's +# ComputeDetAMResponse), where the polarization enters as an exact rotation, +# F+(psi) + i Fx(psi) = (F+(0) + i Fx(0)) e^{-2 i psi}, +# a SINGLE u-harmonic (u = 2 psi). kappa is linear in F and rho^2 quadratic, so +# A carries only u-harmonics +-1 and B only {0, +-2}, for every (ra, dec, incl). +# The banded features do not use that response -- "freqresponse" and "rotation" +# build their coefficients from the arm vectors and a time-varying orientation -- +# so the factorization, and with it the identity, is not guaranteed there. +# +# READ THE ALLOWLIST POSITIVELY, because the negative phrasing inverts: the ONLY +# admitted value is the static response, which is the ABSENCE of a feature tag +# (None). Every named feature -- "rotation", "freqresponse", and anything added +# later -- is refused. Fail-closed by construction: a new response model must be +# added to this tuple deliberately rather than inherit a placement whose premise +# nobody checked. +# +# SCOPE, so this is not over-read: the identity and this gate are about the PSI +# axis. Under precession the PHI content of the coefficient tables IS +# materially redistributed (measured 2026-09-02 on SEOBNRv5PHM: the A phi-slot-0 +# weight moves from 1.1e-16 aligned to 2.9e-2 precessing, while staying +# band-limited to ~6e-15), and that is a property of the SOURCE, not the +# detector. It is not an identity failure and this gate is right to admit it -- +# the psi harmonics are unchanged -- but "the identity holds under precession" +# must not be read as a statement about the phi axis. +_GH_PSI_STATIC_FEATURES = (None,) + +# Bin denominators are floored at this fraction of their own global maximum, so +# that response-free bins (numerator and denominator both ~0) are not scored as +# 0/0 violations. Small enough that a bin carrying any real response is judged +# on its own scale. +_GH_PSI_BIN_FLOOR = 1e-6 + + +def psi_harmonics_at_phi(C_A, C_B, phi, m_max): + """The psi-Fourier FIELDS (A0, A1, B0, B1, B2) at the given phi points. + + A(u) = A0 + Re(A1 e^{iu}); B(u) = B0 + Re(B1 e^{iu}) + Re(B2 e^{2iu}), u = 2 psi. + + THE ONLY DEFINITION. Both the laplace kernel and :func:`gh_laplace_supported` + call this, so the identity the predicate measures is by construction the one + the kernel's node placement depends on. They were separate once, and the + predicate silently measured a DIFFERENT quantity: it read the real part of the + coefficient SLICE ``C_A[:, 1]`` rather than of the phi-reconstruction + ``MA(1)``, so a purely imaginary coefficient gave a nonzero A0(phi) that the + check reported as zero; and it read only ``C_B[:, 3]`` while the field also + carries ``conj(C_B[:, 1])``. Either could pass a dataset whose identity does + not hold. Do not re-derive these five lines anywhere. + """ + phi = jnp.asarray(phi, dtype=jnp.float64) + wA = _kp_weights(m_max + 1) + wB = _kp_weights(2 * m_max + 1) + kpA = jnp.arange(m_max + 1, dtype=jnp.float64) + kpB = jnp.arange(2 * m_max + 1, dtype=jnp.float64) + EA = jnp.exp(1j * phi[:, None] * kpA[None, :]) * wA[None, :] + EB = jnp.exp(1j * phi[:, None] * kpB[None, :]) * wB[None, :] + MA = lambda k: jnp.einsum("ck,kst->cst", EA, C_A[:, k]) + MB = lambda k: jnp.einsum("ck,kst->cst", EB, C_B[:, k]) + return (MA(1).real, # A0 (ks index 1 == ks 0) + MA(2) + jnp.conj(MA(0)), # A1 (ks +1 plus conj(ks -1)) + MB(2).real, # B0 + MB(3) + jnp.conj(MB(1)), # B1 + MB(4) + jnp.conj(MB(0))) # B2 + + +def gh_laplace_supported(C_A, C_B, m_max, feature=None): + """May 'laplace' use the per-sample adaptive distance quadrature on THIS data? + + Returns ``(ok, info)``. Two conditions, both necessary: + + 1. ``m_max <= _GH_PSI_M_MAX`` -- what the path is shipped and validated for. + 2. The A0 == 0 / B1 == 0 identity actually HOLDS on these coefficient + tables, MEASURED rather than assumed. + + (2) exists because (1) does not imply it. m_max is the largest |m| in the + mode list, so a PRECESSING l=2 system has m_max = 2 and passes (1) while + breaking the aligned-spin symmetry h_{l,-m} = (-1)^l conj(h_lm) that the + identity has only ever been tested under. Every measurement of the identity + to date -- IMRPhenomXHM through m_max = 4, and a zero-spin SEOBNRv5PHM run -- + is non-precessing. The analytic argument (F+ + i Fx ~ e^{-2i psi} is one + psi-harmonic, so A is linear in it and B quadratic, making this a property of + the DETECTOR) says it should extend; what has been measured is the CODE's + tables, and two non-precessing tests cannot separate those. So the code + CHECKS instead of trusting the argument: cost is O(size of the coefficient + tables), once, at build time. + """ + import numpy as _np + # Measure the RECONSTRUCTED fields the kernel uses, on a phi grid dense + # enough to resolve their phi content (harmonics to 2*m_max), NOT the + # coefficient slices -- see psi_harmonics_at_phi's docstring for the two + # ways reading slices gave the wrong answer. + ok_modes = int(m_max) <= _GH_PSI_M_MAX + if not ok_modes: + # Return before reconstructing: the tables are SIZED by m_max, so a + # mismatched m_max is a shape error rather than a measurement. + return False, dict(gh_laplace_ok=False, m_max=int(m_max), + identity_A0_over_A1=None, identity_B1_over_B0=None, + feature=feature, + gh_laplace_reason="mode content m_max=%d above the " + "validated %d" + % (int(m_max), _GH_PSI_M_MAX)) + # ANGLE-INDEPENDENT CONDITION, and the one that actually generalises. A + # numerical check can only ever speak for the angles it was evaluated at, + # and the placement runs at arbitrary sampled angles; the response model is + # a property of the packed data and holds for all of them. + if feature not in _GH_PSI_STATIC_FEATURES: + return False, dict(gh_laplace_ok=False, m_max=int(m_max), + identity_A0_over_A1=None, identity_B1_over_B0=None, + feature=feature, + gh_laplace_reason="response model %r does not give " + "the exact e^{-2i psi} polarization " + "factorization the A0 == 0 / B1 == 0 " + "identity rests on" % (feature,)) + n_phi_probe = max(8 * int(m_max) + 8, 16) + phi_probe = _np.linspace(0.0, 2.0 * _np.pi, n_phi_probe, endpoint=False) + A0f, A1f, B0f, B1f, B2f = psi_harmonics_at_phi(C_A, C_B, phi_probe, m_max) + A0f = _np.abs(_np.asarray(A0f)); A1f = _np.abs(_np.asarray(A1f)) + B0f = _np.abs(_np.asarray(B0f)); B1f = _np.abs(_np.asarray(B1f)) + # POINTWISE, not a ratio of global maxima. The placement uses the centre + # and width computed at EACH (phi, sample, time) bin independently, so a + # single locally invalid bin is enough to invalidate it there -- and + # max|A0| / max|A1| hides exactly that, because a large A1 somewhere else + # shrinks the ratio (bins (1e-3, 1) and (0, 1e6) give a passing 1e-9 while + # the first violates by 1e-3). Denominators are floored at a small fraction + # of their own global maximum so that bins with no response at all -- where + # numerator and denominator are both ~0 and nothing is at stake -- do not + # register as 0/0 violations. + a_floor = _GH_PSI_BIN_FLOOR * A1f.max() if A1f.size else 0.0 + b_floor = _GH_PSI_BIN_FLOOR * B0f.max() if B0f.size else 0.0 + r_A0 = float((A0f / _np.maximum(A1f, a_floor)).max()) if a_floor > 0 else _np.inf + r_B1 = float((B1f / _np.maximum(B0f, b_floor)).max()) if b_floor > 0 else _np.inf + ok_ident = (r_A0 <= GH_PSI_IDENTITY_TOL) and (r_B1 <= GH_PSI_IDENTITY_TOL) + if not ok_ident: + reason = ("the A0==0/B1==0 identity does NOT hold pointwise on this data " + "(worst-bin |A0|/|A1|=%.3g, |B1|/B0=%.3g, tol %.0e)" + % (r_A0, r_B1, GH_PSI_IDENTITY_TOL)) + else: + reason = ("m_max=%d, response %r, and the A0==0/B1==0 identity holds at " + "every probed bin" % (int(m_max), feature)) + return ok_ident, dict(gh_laplace_ok=bool(ok_ident), feature=feature, + gh_laplace_reason=reason, + identity_A0_over_A1=r_A0, + identity_B1_over_B0=r_B1, + m_max=int(m_max)) + + +def fused_log_likelihood_distphipsimarg_peaklocal( + data, ra, dec, incl, x_grid, log_w_grid, + interp=JAX_INTERP_DEFAULT, amp_sizing=None, + time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False, + phi_chunk=None): + """Distance-, phi_ref- AND psi-marginalized lnL: PEAK-LOCAL scheme. + + Same contract and normalization as + :func:`fused_log_likelihood_distphipsimarg_exact`. What changes is the psi axis: + rather than a dense grid sized ``~sqrt(A)``, the u-stationary points are obtained + EXACTLY -- they are the unit-circle roots of a quartic, the u-degree being pinned at + 2 for any mode set -- the sorted points partition the circle, and each cell is + integrated on a window set by its own curvature. The node count on that axis is + therefore INDEPENDENT of amplitude: 4 cells x 48 nodes, against the dense rule's 896 + at amplitude 1.25e4. + + THE PHI AXIS IS STILL DENSE HERE and is sized by + :func:`~RIFT.likelihood.jax_ile.joint_anglemarg_peaklocal.required_n_phi` from the + same ``amp_sizing`` the other schemes use. Localizing phi as well exists as a numpy + reference; it is not in this jitted path. See DESIGN_peak_local_framework.md. + + Measured against ``..._exact``: -3.6e-05, -7.1e-15 and -9.1e-13 nats at kappa boost + 1, 10 and 100, and the same figure on a CUDA device as on CPU. + + The adaptive distance quadrature (``JAX_ILE_DISTMARG_GH``) is REFUSED rather than + silently ignored, for the reason the laplace branch refuses it: this kernel sums the + caller's distance grid directly and implements no psi-marginal node placement. + """ + if _core._DISTMARG_GH_N > 0: + raise ValueError( + "JAX_ILE_DISTMARG_GH is set, but the 'peak-local' angle-marg scheme does " + "not implement the adaptive distance quadrature (it sums the caller's " + "distance grid directly). Use --angle-marg-scheme exact, or unset " + "JAX_ILE_DISTMARG_GH.") + _require_amp_sizing(amp_sizing) + from . import joint_anglemarg_peaklocal as _jp + + C_A, C_B, _meta = angle_coefficient_tables(data, ra, dec, incl, interp=interp) + + # THE RUNTIME AMPLITUDE FAILSAFE APPLIES HERE TOO, and omitting it was a review + # finding rather than a judgement call. The u axis is localized and needs no + # sizing, but THE PHI AXIS IS STILL DENSE and is sized from `amp_sizing`, which + # `estimate_angle_amplitude` is explicit about being an estimator and NOT a proven + # bound -- so a hotter sampled sky location can under-resolve phi exactly as it can + # for the exact and laplace schemes. Skipping the check would publish that + # silently, and would also leave the artifact without the standing best-effort + # label, which is worse than the undersizing itself. + _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, "peak-local") + + n_phi = _jp.required_n_phi(amp_sizing, m_max=_data_m_max(data)) + kw = {} if phi_chunk is None else {"phi_chunk": int(phi_chunk)} + + # tables are (KP, 2KS+1, S, npts); move the batch axes to the front so one nested + # vmap covers both and the kernel sees a plain 2-D table per (sample, time). + A = jnp.moveaxis(jnp.asarray(C_A), (2, 3), (0, 1)) + B = jnp.moveaxis(jnp.asarray(C_B), (2, 3), (0, 1)) + + def _one(a, b): + return _jp.joint_lnL_phi_dense(a, b, x_grid, log_w_grid, n_phi=n_phi, **kw) + + lnL_t = jax.vmap(jax.vmap(_one))(A, B) # (S, npts) + if return_lnLt: + return lnL_t + return _time_marginalize_terminal(lnL_t, data, time_quadrature) + + +def choose_angle_marg_scheme(amplitude, gh_enabled=None, + gh_laplace_ok=None): + """Select 'exact' or 'laplace' from a measured amplitude bound. + + ``amplitude`` is the DATA-DERIVED bound from + :func:`estimate_angle_amplitude` (A ~ rho^2/2 scale) -- deliberately not + an SNR guess: selection and dense-grid sizing both key on the measured + coefficient tables, so a missing or wrong external SNR estimate can + affect neither (external-review defect 2). + + The crossover ANGLE_MARG_CROSSOVER_AMPLITUDE sits where BOTH schemes are + deep in their accurate regimes (see the constant's derivation note): + laplace error ~1e-4 nats and falling, exact at machine precision with the + crossover-sized dense grid. The switch therefore tolerates the O(1) + slack in the amplitude bound, and tests evaluate both schemes in the + overlap region and assert agreement -- a validated constant, not a + tuning knob. + + Returns ``(scheme, info)``; ``info`` is a provenance dict the caller MUST + surface in the run log (this pipeline has a documented history of + silently-inert flags). + """ + if gh_enabled is None: + gh_enabled = _core._DISTMARG_GH_N > 0 + if amplitude is None: + return "exact", dict(reason="no amplitude bound available; exact " + "scheme is the conservative branch", + amplitude=None, + crossover=ANGLE_MARG_CROSSOVER_AMPLITUDE) + amp = float(amplitude) + if gh_enabled and not gh_laplace_ok: + # 'laplace' CAN use the adaptive distance quadrature now, but only where + # gh_laplace_supported() says so. Selecting it anywhere else would route + # 'auto' into the raise inside the laplace kernel, so this branch is + # deliberately more conservative than that kernel's own gate: when the + # caller does not supply the predicate at all (gh_laplace_ok=None) we + # take the safe branch rather than guess. + return "exact", dict(reason="JAX_ILE_DISTMARG_GH set and the laplace " + "psi-marginal node placement is not " + "available for this data", + amplitude=amp, + crossover=ANGLE_MARG_CROSSOVER_AMPLITUDE) + # NOTE, and it is a live limitation rather than a subtlety: + # ANGLE_MARG_CROSSOVER_AMPLITUDE is an ACCURACY crossover (A=450, rho~30) -- + # the point above which BOTH schemes are accurate. The measured COST + # crossover is rho ~200-326 (A ~2e4-5e4), an order of magnitude higher, so + # between them 'auto' picks the accurate-but-slower scheme. Re-deriving the + # constant from cost as well as accuracy is a separate, measured change; it + # is deliberately NOT folded in here. + scheme = "laplace" if amp >= ANGLE_MARG_CROSSOVER_AMPLITUDE else "exact" + return scheme, dict(reason="measured amplitude bound %s crossover" + % ("above" if scheme == "laplace" else "below"), + amplitude=amp, + crossover=ANGLE_MARG_CROSSOVER_AMPLITUDE) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py index 36586c20a..59cd4162f 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py @@ -19,15 +19,31 @@ ``PrecomputeLikelihoodTerms`` / ``PackLikelihoodDataStructuresAsArrays`` and are passed in here as plain arrays. -Two time-interpolation modes are provided: +Time-interpolation modes (``interp=``, ``--interp``). ``nearest`` and ``linear`` are +strictly the crudest; between ``cubic`` and ``sinc`` there is no universal ordering -- +see the note under ``sinc``: * ``interp="nearest"`` -- reproduces the production discrete-shift behaviour (round the per-detector arrival to the nearest sample) bit-for-bit, used to *validate* the JAX path against the numpy reference. -* ``interp="linear"`` (default) -- evaluates the rholm timeseries at the +* ``interp="linear"`` -- evaluates the rholm timeseries at the *continuous* arrival time, so the likelihood is differentiable with respect to sky location (through the geometric time delay) and the other extrinsic parameters. This is the AD-friendly path used for gradient-based exploration. + It WAS the default until 2026-08-26 and is no longer, because at high SNR it is + the *worst* option here, worse than ``nearest``: it undershoots the sharp rholm + peak and so biases the recovered arrival time and hence the sky location. +* ``interp="cubic"`` -- the 4-point cubic-Lagrange stencil the numpy/cupy/CUDA + paths spell ``time_interp='cubic'``. +* ``interp="sinc"`` (**default** since 2026-08-26) -- the 2a-tap Lanczos windowed + sinc (a = ``SINC_HALFWIDTH_DEFAULT``), matching ``time_interp='sinc'`` on those + paths. Chosen as the default because its error is BOUNDED across the measured + sweep rather than lowest on average; see ``JAX_INTERP_DEFAULT``. + Which of ``cubic`` and ``sinc`` is more accurate depends on how oversampled Q + is -- on fmin and srate as well as on mass -- and there is no automatic rule; + see ``RIFT/likelihood/DESIGN_q_window_stencil.md`` and + ``RIFT.likelihood.time_interp_choice.CROSSOVER_GUIDANCE``. Both are + differentiable in ``pos``. Time marginalization uses a precomputed Simpson quadrature weight vector (``scipy.integrate.simpson`` applied to the identity, exactly as the production @@ -54,6 +70,11 @@ import jax.numpy as jnp from scipy import integrate as _scipy_integrate +# The 'sinc' stencil half-width, shared with the numpy/cupy/CUDA backends. Imported from the +# leaf module rather than from factored_likelihood so this stays free of numba and lal. +from RIFT.likelihood.time_interp_choice import (SINC_HALFWIDTH_DEFAULT, + TIME_INTERP_DEFAULT) + # Adaptive (per-sample) distance marginalization. The distance integrand is # exp(K x - 0.5 R x^2) with x = d_ref/d -- a Gaussian in x (peak x*=K/R, width # 1/sqrt(R)) times the d^2 prior (∝ x^-4). A uniform-in-d grid under-resolves @@ -169,9 +190,13 @@ def build_likelihood_data(packed_per_detector, deltaT, tref, tvals, return JAXLikelihoodData(detectors, deltaT, gmst, tvals, tref, distMpcRef) -def _gather_nearest(Q_col, pos): +def _gather_nearest(Q_col, pos, u=None): """Q_col[(round(pos))] with the reference's (rint(.)+0.5)->int32 rounding. + ``u`` is accepted and ignored: every gatherer takes the same signature so the call sites + can pass the separable fractional offset unconditionally (see :func:`_separable_u`), and a + discrete gather has no use for it. + ``pos`` has shape (S, npts); ``Q_col`` shape (npts_full,). Positions that fall outside the rholm buffer contribute ZERO (the rholm timeseries is zero beyond its computed support). This matches the production "over-running @@ -187,7 +212,7 @@ def _gather_nearest(Q_col, pos): return jnp.where(valid, Q_col[idx], 0.0 + 0.0j) -def _gather_linear(Q_col, pos): +def _gather_linear(Q_col, pos, u=None): """Linear interpolation of Q_col at continuous positions ``pos``. Differentiable with respect to ``pos`` (the sub-sample arrival time). @@ -198,7 +223,7 @@ def _gather_linear(Q_col, pos): """ n = Q_col.shape[0] i0f = jnp.floor(pos) - frac = pos - i0f + frac = (pos - i0f) if u is None else u i0 = jnp.clip(i0f.astype(jnp.int32), 0, n - 1) i1 = jnp.clip(i0 + 1, 0, n - 1) val = Q_col[i0] * (1.0 - frac) + Q_col[i1] * frac @@ -206,7 +231,7 @@ def _gather_linear(Q_col, pos): return jnp.where(valid, val, 0.0 + 0.0j) -def _gather_cubic(Q_col, pos): +def _gather_cubic(Q_col, pos, u=None): """Four-point cubic-Lagrange interpolation of Q_col at continuous ``pos``. Mirrors the production ``factored_likelihood._cubic_Q_window_numpy`` / @@ -222,8 +247,9 @@ def _gather_cubic(Q_col, pos): the reference), so an over-running window falls off to zero. """ n = Q_col.shape[0] - i0 = jnp.floor(pos).astype(jnp.int32) - u = pos - jnp.floor(pos) + fl = jnp.floor(pos) + i0 = fl.astype(jnp.int32) + u = (pos - fl) if u is None else u w = (-u * (u - 1.0) * (u - 2.0) / 6.0, (u + 1.0) * (u - 1.0) * (u - 2.0) / 2.0, -(u + 1.0) * u * (u - 2.0) / 2.0, @@ -236,12 +262,162 @@ def _gather_cubic(Q_col, pos): return out +def _sinc_lanczos_weights_jax(u, a): + """Lanczos tap weights, mirroring ``factored_likelihood._sinc_lanczos_weight_matrix``. + + The numpy/cupy/CUDA paths all consume ONE weight array built by that function, so they + cannot drift. JAX cannot: the weights depend on ``u``, which is a traced function of the + sky location, so they must be built inside the trace for ``jax.grad`` to see through them. + This is therefore a deliberate SECOND definition of the same formula, and the thing that + keeps it honest is ``test/jax/test_jax_stencil_parity.py``, which compares this against the + numpy generator element-by-element -- including the two details that are easy to get wrong: + + * the ``|x| >= a`` hard zero. On the wired path (u in [0,1)) it reaches only u == 0, + where tap k = a sits exactly at x = -a and is worth 1.5e-33 -- but these are library + helpers, and for u outside [0,1) the clause is worth 2.2e-3, so it is pinned there; and + * the renormalisation to unit sum, which is applied over the FULL stencil and is NOT + redone after out-of-buffer taps are dropped. The CUDA kernel does the same, so the + three backends agree in the zero-extension region as well as the interior. + + ``u`` has the shape of ``pos``; the return has that shape plus a trailing (2a,) tap axis. + """ + k = jnp.arange(-a + 1, a + 1) + x = u[..., None] - k + w = jnp.sinc(x) * jnp.sinc(x / float(a)) + w = jnp.where(jnp.abs(x) >= a, 0.0, w) + total = jnp.sum(w, axis=-1) + total = jnp.where(total == 0.0, 1.0, total) + return k, w / total[..., None] + + +def _make_gather_sinc(a): + """Build the 2a-tap Lanczos gatherer used by ``interp="sinc"``. + + VECTORISED OVER TAPS, and that is load-bearing rather than tidiness. Written as a Python + loop over taps -- the shape the 4-tap :func:`_gather_cubic` above can afford -- a 16-tap + stencil unrolls into 16 separate gathers, and inside a numpyro NUTS trace the resulting + graph is large enough that XLA compilation dominates: measured >1 h of compile at 0-1% GPU + against seconds for cubic. Building the tap axis as an array gives one gather and one + reduction, so graph size is independent of ``a`` (0.11 s against 0.84 s for the unrolled + form at a=8). The two forms agree to a few ulp -- not bit-for-bit, since ``jnp.sum`` over + the tap axis and a sequential accumulation are free to associate differently -- and + test_jax_stencil_parity.test_vectorised_matches_unrolled asserts that, to stop a later + "simplification" back into a loop. + + Accuracy against ``cubic`` is NOT universal: it depends on how oversampled Q is, hence on + fmin and srate as well as mass. See RIFT/likelihood/DESIGN_q_window_stencil.md and + RIFT.likelihood.time_interp_choice.CROSSOVER_GUIDANCE for the measured crossover. + + MEMORY. XLA does not fuse the tap axis away on its own -- it materialises the + ``(..., 2a)`` weight array -- so **pass ``u``**; see :func:`_separable_u`, which is what makes + this stencil affordable. Without it, GPU whole-likelihood scratch at S=20000/npts=614 is + 6583 MB against 101 MB for ``cubic``; with it, 1279 MB, and the gather itself drops + 1719.2 -> 2.7 MB. Runtime halves as well, because the general form recomputes 16 sinc pairs + per (sample, time-bin) when only ``S`` distinct weight rows exist. Measured figures and the + rejected ``lax.scan`` alternative are in DESIGN_q_window_stencil.md §9.5. + """ + def _gather(Q_col, pos, u=None): + n = Q_col.shape[0] + fl = jnp.floor(pos) + i0 = fl.astype(jnp.int32) + k, w = _sinc_lanczos_weights_jax((pos - fl) if u is None else u, a) + idx = i0[..., None] + k + valid = (idx >= 0) & (idx < n) + vals = Q_col[jnp.clip(idx, 0, n - 1)] + return jnp.sum(w * jnp.where(valid, vals, 0.0 + 0.0j), axis=-1) + return _gather + + +def _make_gather_sinc_unrolled(a): + """Tap-by-tap form of :func:`_make_gather_sinc`. Reference for the equivalence test ONLY. + + Do not wire this into ``_GATHERERS``; see the compile-time note there. + """ + def _gather(Q_col, pos, u=None): + n = Q_col.shape[0] + fl = jnp.floor(pos) + i0 = fl.astype(jnp.int32) + k, w = _sinc_lanczos_weights_jax((pos - fl) if u is None else u, a) + out = jnp.zeros(pos.shape, dtype=jnp.complex128) + for j in range(2 * a): + idx = i0 + int(k[j]) + valid = (idx >= 0) & (idx < n) + out = out + w[..., j] * jnp.where(valid, Q_col[jnp.clip(idx, 0, n - 1)], + 0.0 + 0.0j) + return out + return _gather + + +def _separable_u(p0): + """Fractional sample offset for a window built as ``pos = p0[:, None] + arange(npts)``. + + THIS IS A MEMORY FIX, and a large one. Both accumulators build their window that way, with + INTEGER time offsets, so ``frac(pos)`` does not vary along the time axis -- only ``S`` + distinct values exist, not ``S * npts``. Letting a gatherer rediscover ``u`` from the full + ``pos`` makes it build an ``(S, npts, 2a)`` weight array that XLA then has to materialise: + at S = 20000, npts = 614 that is 1719 MB of scratch on GPU, against 0 MB for the 4-tap cubic, + whose weights are cheap enough to stay fused. Passing ``u`` with shape ``(S, 1)`` instead + makes the weight array ``(S, 1, 2a)`` -- **2.7 MB, measured, a 637x reduction** -- and small + enough that the surrounding product and reduction fuse, exactly as cubic's already do. + + It is also MORE accurate, not a trade. ``p0`` is a sample index of order 1e5-1e6, so + ``p0 + t`` can cross a binade and drop a low mantissa bit; ``frac(p0 + t)`` then differs from + ``frac(p0)`` by up to an ulp of the position (~1.5e-11 at 65536, measured). The numpy, cupy + and CUDA backends all compute one fractional offset per sample from the sample position -- + i.e. this form -- so using it here IMPROVES cross-backend agreement rather than costing it. + + Callers that do not have a separable window simply omit ``u`` and every gatherer falls back + to ``pos - floor(pos)``. Getting it wrong is silent, so + test_separable_u_matches_the_general_path pins the two against each other at production + magnitudes, and test_accumulators_pass_separable_u pins that the call sites actually pass it. + """ + return (p0 - jnp.floor(p0))[:, None] + + _GATHERERS = {"nearest": _gather_nearest, "linear": _gather_linear, - "cubic": _gather_cubic} + "cubic": _gather_cubic, + "sinc": _make_gather_sinc(SINC_HALFWIDTH_DEFAULT)} + +# The default stencil for every entry point in this package AND for the --interp flag of +# bin/integrate_likelihood_extrinsic_jax, which imports it from here so the two cannot drift. +# +# CHANGED 2026-08-26: 'linear' -> 'sinc'. This CHANGES RESULTS for any caller that did not pass +# interp= explicitly; pass interp="linear" to reproduce a pre-2026-08-26 run. Rationale, and the +# concern that goes with it, are recorded in DESIGN_q_window_stencil.md §9.4 -- in one line: +# linear is the worst stencil here at high SNR (worse than 'nearest'), this path is used +# exclusively at high SNR, and 'sinc' is the option whose error is BOUNDED (measured flat at +# 2.3-7.9 nats across the whole mass/fmin sweep) rather than the one with the best best-case. +# ALIAS, not a second literal (2026-09-02). It was a re-typed "sinc", which is exactly how +# the two drivers came to ship opposite defaults in the first place (issue #233); the value +# now lives once, in time_interp_choice.TIME_INTERP_DEFAULT. The NAME is kept because every +# entry point in this package and bin/integrate_likelihood_extrinsic_jax import it. +JAX_INTERP_DEFAULT = TIME_INTERP_DEFAULT + + +def _guarded_window(data, guard): + """``(npts, t_offsets)`` for the accumulation window widened by ``guard`` samples. + + ``guard == 0`` is the production window -- ``npts == data.npts`` and + ``t_offsets == arange(data.npts)``, unchanged. A positive ``guard`` adds that + many samples at EACH end, so the window runs from ``-guard`` to + ``data.npts + guard - 1`` and the caller gets ``data.npts + 2*guard`` columns. + + The extra columns are RECONSTRUCTION SUPPORT, not window: only + :func:`_time_marginalize_bandlimited` asks for them, and it integrates the + original ``data.npts`` columns after using the guard samples to keep the + FFT's periodic seam out of the integrated region. One definition here so + the two accumulators cannot drift on the offset convention -- an off-by-one + between them would misplace the arrival time of every guarded evaluation. + """ + guard = int(guard) + if guard < 0: + raise ValueError("guard must be >= 0 samples, got %r" % (guard,)) + return (data.npts + 2 * guard, + jnp.arange(-guard, data.npts + guard, dtype=jnp.float64)) def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, - phase_marginalization): + phase_marginalization, guard=0): """Network kappa and rho^2 at the *fiducial* distance (invDist == 1). Returns ``(kappa_unit, rho_sq_unit)`` each shape (S, npts). ``kappa_unit`` @@ -254,10 +430,16 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, the multi-band accumulator is used instead. It returns the *identical* ``(kappa_unit, rho_sq_unit)`` contract, so every downstream marginalization variant (distance, phi_ref, psi, ...) inherits the feature for free. + + ``guard`` widens the evaluated window by that many samples at each end (see + :func:`_guarded_window`), giving ``(S, data.npts + 2*guard)``. It defaults + to 0, i.e. every existing caller gets the production window unchanged; the + band-limited time quadrature is the one caller that asks for more, and it + pays for them in gathers (cost is linear in the number of columns). """ if getattr(data, "feature", None) is not None: return _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, - phase_marginalization) + phase_marginalization, guard=guard) ra = jnp.asarray(ra, dtype=jnp.float64) dec = jnp.asarray(dec, dtype=jnp.float64) psi = jnp.asarray(psi, dtype=jnp.float64) @@ -268,11 +450,11 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, gmst = data.gmst inv_deltaT = 1.0 / data.deltaT S = ra.shape[0] - npts = data.npts - t_offsets = jnp.arange(npts, dtype=jnp.float64) + npts, t_offsets = _guarded_window(data, guard) kappa_unit = jnp.zeros((S, npts), dtype=jnp.complex128) rho_sq_unit = jnp.zeros((S, npts), dtype=jnp.float64) + support_valid = jnp.ones((S,), dtype=bool) for det in data.detector_names: dd = data.detectors[det] @@ -306,10 +488,21 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, + time_delay_from_earth_center(dd["location"], ra, dec, gmst)) p0 = (t_det + data.tval0) * inv_deltaT pos = p0[:, None] + t_offsets[None, :] + if guard: + stencil_margin = {"nearest": 1, "linear": 2, "cubic": 3, + "sinc": SINC_HALFWIDTH_DEFAULT + 1}[interp] + support_valid = support_valid & jnp.all( + (pos >= stencil_margin) + & (pos <= Q.shape[0] - 1 - stencil_margin), axis=-1) + # None for 'nearest': it ignores u, and feeding an unused value into this trace + # is NOT free -- it cost >60% wall on the banded slow-rotation path (measured: + # test_rotation_path_a 69.8 s -> >113 s), which is compile-bound, not arithmetic- + # bound. Only the weight-building stencils get it. See _separable_u. + u_sep = None if interp == "nearest" else _separable_u(p0) kappa_det = jnp.zeros((S, npts), dtype=jnp.complex128) for k in range(K): - Qi = gather(Q[:, k], pos) + Qi = gather(Q[:, k], pos, u_sep) kappa_det = kappa_det + FY_conj[:, k][:, None] * Qi kappa_unit = kappa_unit + kappa_det # NOT a gap, and worth saying so because an earlier revision wrongly marked it as one: @@ -321,9 +514,29 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, # slow-rotation model does have that dependence -- see _accumulate_unit_banded. rho_sq_unit = rho_sq_unit + rho_sq_det[:, None] + if guard: + kappa_unit = jnp.where(support_valid[:, None], kappa_unit, + jnp.nan + 0.0j) + rho_sq_unit = jnp.where(support_valid[:, None], rho_sq_unit, jnp.nan) return kappa_unit, rho_sq_unit +def _norm_is_arrival_time_dependent(data): + """True when the model norm ```` depends on the template's arrival time. + + Only the slow-rotation bank has that dependence: its post-phase + ``C~_a(t) = C_a exp(i n_a Omega (t - tref))`` multiplies the data term AND the + norm, so :func:`_accumulate_unit_banded` returns a genuinely ``(S, npts)`` + ``rho_sq`` there. The baseline accumulator (static ``F``) and the finite-size + ``freqresponse`` bank (no sidereal modulation) both return a norm that is + constant along the time axis, broadcast into the ``(S, npts)`` contract. + + One definition on purpose: the quadratures that hold the norm fixed refuse + exactly the data this predicate flags, so the two must not drift apart. + """ + return getattr(data, "feature", None) == "rotation" + + def _banded_coefficients(data, det, ra, dec, psi): """Per-sample response coefficients ``C`` of shape (A, S) for detector ``det``. @@ -346,7 +559,7 @@ def _banded_coefficients(data, det, ra, dec, psi): def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, - phase_marginalization): + phase_marginalization, guard=0): """Multi-band (slow-rotation / finite-size) network kappa and rho^2. Generalizes :func:`_accumulate_unit` by an extra summed "band" index @@ -401,6 +614,11 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, ``freqresponse`` (Path D) has NO post-phase -- its basis is not a sidereal modulation -- and keeps the arrival-time-independent ``rho_sq``. + ``guard`` widens the window as in :func:`_accumulate_unit` / :func:`_guarded_window`. + The post-phase follows it: ``jgrid`` is built from the same (now partly negative) + integer offsets, so a guarded sample is phased at the arrival time it was gathered + from, exactly as an unguarded one is. + ``phase_marginalization`` is not supported for banded features. """ if phase_marginalization: @@ -418,14 +636,13 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, gmst = data.gmst inv_deltaT = 1.0 / data.deltaT S = ra.shape[0] - npts = data.npts - t_offsets = jnp.arange(npts, dtype=jnp.float64) + npts, t_offsets = _guarded_window(data, guard) refl_idx = data.band["refl_idx"] # (A,) int, static # Arrival-time post-phase: rotation only (see the docstring). Honour the bank # convention flag rather than assuming it, so a future change fails loudly. band = data.band - post_phase = (data.feature == "rotation") + post_phase = _norm_is_arrival_time_dependent(data) if post_phase: if not bool(band.get("post_phase_required", False)): raise ValueError( @@ -445,6 +662,7 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, kappa_unit = jnp.zeros((S, npts), dtype=jnp.complex128) rho_sq_unit = jnp.zeros((S, npts), dtype=jnp.float64) + support_valid = jnp.ones((S,), dtype=bool) for det in data.detector_names: dd = data.detectors[det] @@ -465,6 +683,17 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, + time_delay_from_earth_center(dd["location"], ra, dec, gmst)) p0 = (t_det + data.tval0) * inv_deltaT pos = p0[:, None] + t_offsets[None, :] # (S, npts) + if guard: + stencil_margin = {"nearest": 1, "linear": 2, "cubic": 3, + "sinc": SINC_HALFWIDTH_DEFAULT + 1}[interp] + support_valid = support_valid & jnp.all( + (pos >= stencil_margin) + & (pos <= Q_bank.shape[1] - 1 - stencil_margin), axis=-1) + # None for 'nearest': it ignores u, and feeding an unused value into this trace + # is NOT free -- it cost >60% wall on the banded slow-rotation path (measured: + # test_rotation_path_a 69.8 s -> >113 s), which is compile-bound, not arithmetic- + # bound. Only the weight-building stencils get it. See _separable_u. + u_sep = None if interp == "nearest" else _separable_u(p0) if post_phase: # delta_ij = (arrival time of output bin j for sample i) - tref, in seconds. @@ -497,7 +726,7 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, inner_a = jnp.zeros((S, npts), dtype=jnp.complex128) Qa = Q_bank[a] # (npts_full, K) for k in range(K): - inner_a = inner_a + conjY[:, k][:, None] * gather(Qa[:, k], pos) + inner_a = inner_a + conjY[:, k][:, None] * gather(Qa[:, k], pos, u_sep) if post_phase: i1 = int(pp_t1[a]) kappa_det = kappa_det + ((jnp.conj(C[a]) * pe[i1])[:, None] @@ -527,9 +756,189 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, rho_sq_unit = rho_sq_unit + (rho_sq_det if post_phase else rho_sq_det[:, None]) + if guard: + kappa_unit = jnp.where(support_valid[:, None], kappa_unit, + jnp.nan + 0.0j) + rho_sq_unit = jnp.where(support_valid[:, None], rho_sq_unit, jnp.nan) return kappa_unit, rho_sq_unit +TIME_QUAD_DEFAULT = "simpson" # unchanged behaviour; "bandlimited" is opt-in +_TIME_QUAD_CHOICES = ("simpson", "bandlimited") +_TIME_UPSAMPLE_DEFAULT = 8 +_TIME_ADAPTIVE_FACTOR_MAX = 1024 +_TIME_ADAPTIVE_SAFETY = 2.0 +_TIME_ADAPTIVE_RTOL = 1e-3 +_TIME_ENDPOINT_LOG_GAP_MIN = 15.0 + + +def default_time_guard(npts): + """Guard width (samples per end) used by ``time_quad="bandlimited"`` by default. + + Half the window, floored at 32 samples: the reconstruction's seam error falls + off only like 1/distance (see :func:`_time_marginalize_bandlimited`), so the + guard has to be a FRACTION of the window rather than a small fixed pad, and + the cost is linear -- half a window at each end doubles the gathers of an + opt-in quadrature that exists to buy accuracy. It is a convergence knob, not + a constant of nature: pass ``time_guard=`` and raise it until the answer stops + moving, exactly as with ``time_upsample``. + """ + return max(32, int(npts) // 2) + + +def _upsample_bandlimited(x, factor, axis=-1): + """Band-limited resampling of ``x`` by an integer ``factor``, ASSUMING PERIODICITY. + + Zero-pads the spectrum, which is interpolation only in the sense that the + sampling theorem is: for a signal whose Fourier content the grid already + resolves AND WHOSE PERIOD IS THE ARRAY, the padded inverse transform + reproduces the underlying continuous function at the finer spacing, not an + approximation of it. + + THE PERIODICITY CAVEAT IS NOT DECORATIVE, and it is not something the caller + can check on the output. The array is a period, so this treats ``x[-1]`` and + ``x[0]`` as adjacent; when they are not -- as for any window CROPPED out of a + longer timeseries, which is what the accumulators produce -- the seam's jump + is spectral content as far as the FFT is concerned, and the reconstruction + rings. The retained samples stay exact (the interpolant passes through every + input sample), so a test that only checks ``fine[::factor] == x`` cannot see + it; the INSERTED samples carry the error, largest at the ends and falling off + only like 1/(distance from the seam). Callers reconstructing a crop must + supply guard samples and discard the contaminated region -- + :func:`_time_marginalize_bandlimited` does, and explains the arithmetic. + + Nyquist handling: for even ``n`` the +n/2 bin is split evenly between the + +n/2 and -n/2 positions. Dumping it entirely into one of them biases the + result by a term that oscillates at Nyquist -- small, but exactly the kind + of grid-phase-dependent error this whole change exists to remove. + """ + x = jnp.asarray(x) + n = x.shape[axis] + if factor == 1: + return x + X = jnp.fft.fft(x, axis=axis) + X = jnp.moveaxis(X, axis, -1) + n_out = n * factor + half = n // 2 + pad_shape = X.shape[:-1] + (n_out - n,) + if n % 2 == 0: + lo = X[..., :half] + hi = X[..., half + 1:] + nyq = X[..., half:half + 1] * 0.5 + Y = jnp.concatenate( + [lo, nyq, jnp.zeros(X.shape[:-1] + (n_out - n - 1,), X.dtype), + nyq, hi], axis=-1) + else: + Y = jnp.concatenate( + [X[..., :half + 1], jnp.zeros(pad_shape, X.dtype), + X[..., half + 1:]], axis=-1) + y = jnp.fft.ifft(Y, axis=-1) * factor + return jnp.moveaxis(y, -1, axis) + + +def _time_marginalize_bandlimited(kappa_t, rho_sq, deltaT, factor, guard, + phase_marginalization=False): + """Time marginal evaluated on a band-limited RECONSTRUCTION of kappa(t). + + Why this exists. ``_time_marginalize`` integrates exp(lnL_t) with fixed + Simpson weights at the DATA sample spacing, but the integrand's width is + sigma_t = 1/(2 pi rho sigma_f) -- it SHRINKS as the signal gets louder while + the grid does not. Measured on a 35+30 Msun HLV injection at rho=40: + sigma_t = 61.2 us against grid spacings of 244/122/61 us at srate + 4096/8192/16384, i.e. under-resolved at the sample rates people actually + use, and worse at higher SNR. The required condition is + srate >~ 2 pi sigma_f rho (~16 kHz at rho=40 here, growing linearly with + SNR). Simpson is not a safeguard here: Simpson = (4 T_h - T_2h)/3 carries + the coarser T_2h alias, so when under-resolved it is WORSE than trapezoid -- + the measured lnL span over a rigid grid-phase scan was 1.649 / 0.385 / + 0.0095 nats at those three sample rates, dominated by the period-2h term. + + The fix costs no new likelihood evaluations. kappa(t) is band-limited (it + is a cross-correlation of band-limited data with a band-limited template), + and rho_sq is time-independent (the PRECONDITION below), so the samples ALREADY + COMPUTED determine the continuous integrand exactly. One zero-padded FFT + per row recovers it; the quadrature then runs on the reconstruction. + + Measured against a converged window-shift reference: -0.007 nats, versus + +0.745 nats for stock Simpson at the same grid phase. (That comparison was + taken with ``guard == 0``, i.e. before the guard argument below existed, so + it is quoted for the SIMPSON contrast and not as a bound on the seam error.) + + GUARD SAMPLES, and why this argument has no default. ``kappa_t`` is a window + CROPPED out of a longer correlation buffer, so its two ends are not + neighbours -- while :func:`_upsample_bandlimited`, being an FFT, necessarily + treats them as though they were. The fictitious jump across that seam is + spectral content as far as the transform is concerned, and it rings into the + INSERTED samples, changing the integrand before any quadrature touches it. + The rings are invisible at the input samples (the interpolant reproduces + those exactly), which is why this has to be handled here rather than caught + downstream. + + The cure is support, not smoothing: ``kappa_t`` carries ``guard`` extra + samples at EACH end (:func:`_accumulate_unit` gathers them), the seam moves + out to those ends, and only the middle + ``npts = kappa_t.shape[-1] - 2*guard`` samples -- the window the stock + quadrature integrates -- are kept. The residual is then set by the seam's + distance: writing the reconstruction as a Whittaker sum over the + PERIODICALLY REPEATED window, the error at a point ``d`` samples inside the + kept region is the tail + + sum_{j outside} (xtilde_j - x_j) sinc(t - j) ~ |seam jump| / (pi d), + + which falls off like 1/d and NOT exponentially. So ``guard`` is a + convergence knob of the same standing as ``factor``: raise it until the + answer stops moving (:func:`fused_log_likelihood` starts from + :func:`default_time_guard`). Passing zero on real cropped data is the defect + itself, so there is no default to fall into; ``guard=0`` stays legitimate + only for a window whose ends genuinely join -- a constant, or an integrand + that has decayed to its pedestal at both ends. + + PRECONDITION on ``rho_sq``: the model norm must NOT depend on arrival time. + Only ``rho_sq[..., :1]`` is used here, because the reconstruction is of kappa + alone. That is the complete norm for the baseline and finite-size + accumulators, whose ``rho_sq`` is a constant column broadcast along time, but + the slow-rotation post-phase makes ``rho_sq`` a genuine function of the + arrival bin (:func:`_accumulate_unit_banded`), and taking its first bin there + would evaluate a DIFFERENT likelihood. Callers must screen such data with + :func:`_norm_is_arrival_time_dependent`; :func:`fused_log_likelihood` does. + """ + guard = int(guard) + if guard < 0: + raise ValueError("guard must be >= 0 samples, got %r" % (guard,)) + npts = kappa_t.shape[-1] - 2 * guard + if npts < 2: + raise ValueError( + "guard=%d leaves %d sample(s) of a %d-sample array to integrate; the " + "guard samples are reconstruction support, not window" + % (guard, npts, kappa_t.shape[-1])) + kappa_f = _upsample_bandlimited(kappa_t, factor, axis=-1) + # Integrate the ORIGINAL interval (npts-1)*deltaT, and nothing else. Two + # distinct pieces of the fine grid are NOT part of it: + # * the leading and trailing `guard*factor` samples, whose coarse samples + # were gathered only to hold the periodic seam away from the window; and + # * the last factor-1 samples of the array, which lie PAST the final coarse + # sample and are the periodic continuation wrapping back toward sample 0. + # Integrating either rescales the result -- for a constant integrand by + # exactly log(kept length / ((npts-1)*deltaT)) -- against every other + # quadrature in the module. Starting at guard*factor and taking + # (npts-1)*factor+1 points lands on the original two endpoints exactly. + start = guard * factor + kappa_f = kappa_f[..., start:start + (npts - 1) * factor + 1] + # The norm is taken at the first bin OF THE KEPT WINDOW, not of the guarded + # array: identical under the precondition (the column is constant), and the + # one that stays right if a future caller ever relaxes it. + rho_sq_0 = rho_sq[..., guard:guard + 1] + if phase_marginalization: + lnL_f = jnp.abs(kappa_f) - 0.5 * rho_sq_0 + else: + lnL_f = kappa_f.real - 0.5 * rho_sq_0 + n_f = lnL_f.shape[-1] + w_f = jnp.asarray(_simpson_weights(n_f, deltaT / factor)) + m = jnp.max(lnL_f, axis=-1, keepdims=True) + L = jnp.sum(w_f[None, :] * jnp.exp(lnL_f - m), axis=-1) + return m[:, 0] + jnp.log(L) + + def _time_marginalize(lnL_t, w_t): """log integral_t exp(lnL_t) dt via constant Simpson weights, log-sum-exp stable.""" m = jnp.max(lnL_t, axis=-1, keepdims=True) @@ -537,8 +946,261 @@ def _time_marginalize(lnL_t, w_t): return m[:, 0] + jnp.log(L) +def _reflected_fft_upsample(x, factor): + """FFT-interpolate a finite row with the standard even extension. + + Periodize ``[x[0], ..., x[-1], x[-2], ..., x[1]]``. Omitting duplicate + turning samples is mathematically essential at Nyquist: duplicating them + inserts an artificial flat pair, so ``(-1)**j`` no longer reconstructs + ``cos(pi*t)`` and phase marginalization can converge to the wrong integral. + Only the original closed forward interval is returned. + """ + x = jnp.asarray(x) + factor = int(factor) + if factor == 1: + return x + n = x.shape[-1] + reflected = jnp.concatenate((x, jnp.flip(x[..., 1:-1], axis=-1)), axis=-1) + dense = _upsample_bandlimited(reflected, factor, axis=-1) + return dense[..., :(n - 1) * factor + 1] + + +def _peak_width_from_lnL_jax(lnL_t, dx): + """Measure per-row Gaussian peak width from finite centred differences.""" + n = lnL_t.shape[-1] + finite = jnp.isfinite(lnL_t) + safe = jnp.where(finite, lnL_t, -jnp.inf) + jmax = jnp.argmax(safe, axis=-1) + sigma = jnp.full(jmax.shape, jnp.inf, dtype=jnp.float64) + measurable = jnp.zeros(jmax.shape, dtype=bool) + + def take(j): + return jnp.take_along_axis(lnL_t, j[..., None], axis=-1)[..., 0] + + for d in (1, 2, 4, 8): + if 2 * d >= n: + break + jc = jnp.clip(jmax, d, n - 1 - d) + d2 = (take(jc - d) - 2.0 * take(jc) + take(jc + d)) / (d * dx) ** 2 + fresh = jnp.isfinite(d2) & (~measurable) + neg = fresh & (d2 < 0) + sigma = jnp.where(neg, 1.0 / jnp.sqrt(jnp.where(neg, -d2, 1.0)), sigma) + measurable = measurable | fresh + return sigma, measurable + + +def _log_trapezoid(lnL_t, dx): + """Stable row-wise log trapezoid over exactly the represented interval.""" + m = jnp.max(lnL_t, axis=-1, keepdims=True) + m_safe = jnp.where(jnp.isfinite(m), m, 0.0) + y = jnp.exp(lnL_t - m_safe) + total = dx * (0.5 * y[..., 0] + jnp.sum(y[..., 1:-1], axis=-1) + + 0.5 * y[..., -1]) + return m_safe[..., 0] + jnp.log(total) + + +def _terminal_reflected_fft_at_factor(lnL_t, deltaT, factor): + dense = _reflected_fft_upsample(lnL_t, factor).real + value = _log_trapezoid(dense, deltaT / float(factor)) + sigma, measurable = _peak_width_from_lnL_jax(dense, deltaT / float(factor)) + resolved = (~measurable) | (~jnp.isfinite(sigma)) | ( + deltaT / float(factor) <= sigma / _TIME_ADAPTIVE_SAFETY) + return value, resolved + + +def _time_marginalize_reflected_fft(lnL_t, deltaT, w_t): + """Adaptive reflected-FFT terminal time marginalization. + + A per-row power-of-two factor is selected from the coarse-row curvature. + ``lax.map`` keeps the switched FFT scratch row-local, so one sharp sample + neither changes its batchmates' result nor materializes its fine grid for + the whole sampler batch. The selected grid is remeasured and doubled until + both the width criterion and a 1e-3-nat convergence check pass. Rows with + any non-finite coarse bin retain the historical Simpson value; their + sanitized values still enter traced FFT branches because JAX evaluates both + sides of ``where``. + """ + lnL_t = jnp.asarray(lnL_t, dtype=jnp.float64) + simpson = _time_marginalize(lnL_t, w_t) + finite_rows = jnp.all(jnp.isfinite(lnL_t), axis=-1) + clean = jnp.where(finite_rows[:, None], lnL_t, 0.0) + + sigma, measurable = _peak_width_from_lnL_jax(clean, deltaT) + need = jnp.where(measurable & jnp.isfinite(sigma) & (sigma > 0), + _TIME_ADAPTIVE_SAFETY * deltaT / sigma, 1.0) + need = jnp.maximum(need, 1.0) + factor_float = jnp.exp2(jnp.ceil(jnp.log2(need))) + factor_float = jnp.where(factor_float < need, factor_float * 2.0, factor_float) + too_sharp = (~jnp.isfinite(factor_float)) | ( + factor_float > _TIME_ADAPTIVE_FACTOR_MAX) + # Clamp before the integer cast: inf or an out-of-range float can wrap to a + # negative integer and otherwise masquerade as factor 1. + factor = jnp.minimum(factor_float, float(_TIME_ADAPTIVE_FACTOR_MAX)).astype( + jnp.int32) + + powers = tuple(1 << k for k in range(11)) # q0 <= 1024; certificate <= 2048 + + def make_branch(base): + def branch(x): + v0, r0 = _terminal_reflected_fft_at_factor(x, deltaT, base) + v1, r1 = _terminal_reflected_fft_at_factor(x, deltaT, 2 * base) + c1 = r1 & (jnp.abs(v1 - v0) <= _TIME_ADAPTIVE_RTOL) + return jnp.where(c1, v1, jnp.nan) + return branch + + def refine_one(args): + row, row_factor = args + index = jnp.clip( + jnp.ceil(jnp.log2(row_factor.astype(jnp.float64))).astype(jnp.int32), + 0, len(powers) - 1) + return jax.lax.switch(index, tuple(make_branch(f) for f in powers), row) + + refined = jax.lax.map(jax.checkpoint(refine_one), (clean, factor)) + refined = jnp.where(too_sharp, jnp.nan, refined) + return jnp.where(finite_rows, refined, simpson) + + +def _time_marginalize_reflected_primitive(kappa_t, rho_sq, deltaT, + phase_marginalization=False, + guard=0): + """Adaptive integral after refining the band-limited complex primitive. + + This is required for phase marginalization: interpolating ``abs(kappa)`` + cannot recover intersample structure lost to that nonlinear operation. + Arrival-time-dependent norms remain unsupported by the bandlimited mode. + A refined row whose endpoint is within 15 nats of its peak fails closed: + the even-extension boundary condition is not trustworthy when the finite + window carries appreciable posterior mass at either turn. + """ + guard = int(guard) + if guard < 0 or 2 * guard >= kappa_t.shape[-1] - 1: + raise ValueError("guard must leave at least two integration samples") + npts = kappa_t.shape[-1] - 2 * guard + inner_guard = guard // 2 + if guard and inner_guard < 1: + raise ValueError("guard convergence requires at least two samples per end") + kappa_t = jnp.asarray(kappa_t, dtype=jnp.complex128) + rho_sq = jnp.asarray(rho_sq, dtype=jnp.float64) + coarse_full = ((jnp.abs(kappa_t) if phase_marginalization else kappa_t.real) + - 0.5 * rho_sq) + coarse = coarse_full[..., guard:guard + npts] + finite_rows = jnp.all(jnp.isfinite(coarse), axis=-1) + clean_kappa = jnp.where(finite_rows[:, None], kappa_t, 0.0) + clean_rho = jnp.where(finite_rows[:, None], rho_sq, 0.0) + + def taper_support(x, support_guard): + if not support_guard: + return x + u = jnp.arange(support_guard + 1, dtype=jnp.float64) / support_guard + ramp = 0.5 * (1.0 - jnp.cos(jnp.pi * u)) + taper = jnp.concatenate((ramp[:-1], jnp.ones((npts,)), + jnp.flip(ramp[:-1]))) + return x * taper + + # Probe the primitive at half a sample before deriving curvature. A + # near-Nyquist real kappa can alternate +/-A, making coarse ``abs(kappa)`` + # exactly constant even though the continuous phase-marginalized field has + # a zero between every pair of samples. No statistic of the coarse + # nonlinear field can detect that alias. + probe_kappa = _reflected_fft_upsample(taper_support(clean_kappa, guard), 2) + probe_rho = jnp.broadcast_to(clean_rho[:, :1], probe_kappa.shape) + probe = ((jnp.abs(probe_kappa) if phase_marginalization + else probe_kappa.real) - 0.5 * probe_rho) + probe = probe[..., 2 * guard:2 * guard + (npts - 1) * 2 + 1] + sigma, measurable = _peak_width_from_lnL_jax(probe, deltaT / 2.0) + need = jnp.where(measurable & jnp.isfinite(sigma) & (sigma > 0), + _TIME_ADAPTIVE_SAFETY * deltaT / sigma, 1.0) + need = jnp.maximum(need, 1.0) + factor_float = jnp.exp2(jnp.ceil(jnp.log2(need))) + factor_float = jnp.where(factor_float < need, factor_float * 2.0, factor_float) + too_sharp = (~jnp.isfinite(factor_float)) | ( + factor_float > _TIME_ADAPTIVE_FACTOR_MAX) + factor = jnp.minimum(factor_float, float(_TIME_ADAPTIVE_FACTOR_MAX)).astype( + jnp.int32) + powers = tuple(1 << k for k in range(11)) + + def make_branch(base): + def at_factor(kappa, rho, f, support_guard): + if support_guard < guard: + trim = guard - support_guard + kappa = kappa[trim:-trim] + rho = rho[trim:-trim] + if support_guard: + # Smoothly pad the primitive to zero only in the support + # samples. The raised-cosine value and slope both vanish at + # the remote reflection turns and reach exactly one at the + # integration crop. This removes the derivative cusp whose + # global FFT ringing survives even when endpoint likelihood + # mass is negligible. + kappa = taper_support(kappa, support_guard) + dense_kappa = _reflected_fft_upsample(kappa, f) + # Conventional baseline data have a time-independent model norm. + # Keeping the first value avoids inventing high-frequency structure + # in a constant primitive through roundoff. + dense_rho = jnp.broadcast_to(rho[0], dense_kappa.shape) + dense = ((jnp.abs(dense_kappa) if phase_marginalization + else dense_kappa.real) - 0.5 * dense_rho) + start = support_guard * f + dense = dense[start:start + (npts - 1) * f + 1] + value = _log_trapezoid(dense, deltaT / float(f)) + width, measured = _peak_width_from_lnL_jax(dense, deltaT / float(f)) + resolved = ((~measured) | (~jnp.isfinite(width)) + | (deltaT / float(f) <= width / _TIME_ADAPTIVE_SAFETY)) + peak = jnp.max(dense) + endpoint = jnp.maximum(dense[0], dense[-1]) + boundary_ok = endpoint <= peak - _TIME_ENDPOINT_LOG_GAP_MIN + return value, resolved, boundary_ok + + def branch(args): + kappa, rho = args + v0, r0, b0 = at_factor(kappa, rho, base, guard) + v1, r1, b1 = at_factor(kappa, rho, 2 * base, guard) + if guard: + vg1, _, _ = at_factor(kappa, rho, 2 * base, inner_guard) + g1 = jnp.abs(v1 - vg1) <= _TIME_ADAPTIVE_RTOL + else: + g1 = True + c1 = r1 & b1 & g1 & (jnp.abs(v1 - v0) <= _TIME_ADAPTIVE_RTOL) + return jnp.where(c1, v1, jnp.nan) + return branch + + branches = tuple(make_branch(f) for f in powers) + + def refine_one(args): + kappa, rho, row_factor = args + index = jnp.clip( + jnp.ceil(jnp.log2(row_factor.astype(jnp.float64))).astype(jnp.int32), + 0, len(powers) - 1) + return jax.lax.switch(index, branches, (kappa, rho)) + + # Rematerialize a row's selected branch during reverse mode instead of + # retaining every dense abs/exp/FFT residual across the sampler batch. + refined = jax.lax.map( + jax.checkpoint(refine_one), (clean_kappa, clean_rho, factor)) + refined = jnp.where(too_sharp, jnp.nan, refined) + return jnp.where(finite_rows, refined, jnp.nan) + + +def _time_marginalize_terminal(lnL_t, data, time_quadrature=TIME_QUAD_DEFAULT, + bandlimited_safe=False): + """Common terminal selector used by every JAX time-marginalized endpoint.""" + if time_quadrature not in _TIME_QUAD_CHOICES: + raise ValueError("time_quadrature must be one of %r, got %r" + % (_TIME_QUAD_CHOICES, time_quadrature)) + if time_quadrature == "simpson": + return _time_marginalize(lnL_t, data.w_t) + if not bandlimited_safe: + raise ValueError( + "bandlimited terminal interpolation is invalid after nonlinear " + "distance/phase/polarization marginalization; use 'simpson'") + return _time_marginalize_reflected_fft(lnL_t, data.deltaT, data.w_t) + + def fused_log_likelihood(data, ra, dec, psi, incl, phiref, distMpc, - interp="linear", phase_marginalization=False): + interp=JAX_INTERP_DEFAULT, phase_marginalization=False, + time_quad=TIME_QUAD_DEFAULT, + time_upsample=None, time_guard=None, + time_quadrature=None, return_lnLt=False): """Time-marginalized factored log-likelihood at a fixed distance, lnL(theta). Parameters @@ -550,28 +1212,104 @@ def fused_log_likelihood(data, ra, dec, psi, incl, phiref, distMpc, Time-interpolation of the rholm timeseries (see module docstring). phase_marginalization : bool Marginalize the coalescence phase via ``|kappa|``. + time_quad : {"simpson", "bandlimited"} + Time quadrature. ``"simpson"`` (default) is the unchanged behaviour. + ``"bandlimited"`` integrates a band-limited reconstruction of kappa(t) + over the SAME interval; it holds the model norm fixed in time and is + therefore refused for slow-rotation data, whose ```` depends on the + template arrival time (see :func:`_time_marginalize_bandlimited`). + time_upsample : int + Reconstruction factor for ``time_quad="bandlimited"``; costs no + likelihood evaluations, so raise it until the answer stops moving. + time_guard : int or None + Guard samples per end for ``time_quad="bandlimited"``. The window is a + crop, the reconstruction is periodic, and the seam between the two ends + rings into the integrand; these samples move that seam outside the + integrated window (see :func:`_time_marginalize_bandlimited`). Unlike + ``time_upsample`` they DO cost gathers -- the accumulation runs on + ``npts + 2*time_guard`` bins -- and the seam error falls off only like + 1/distance, so this is the second knob to raise when checking that the + answer has stopped moving. ``None`` takes :func:`default_time_guard`. + Ignored (and forced to 0) by ``time_quad="simpson"``, which integrates + the sampled window itself and has no seam. Returns ------- lnL : array_like, shape (S,) """ + canonical_time_api = time_quadrature is not None + if canonical_time_api: + if time_quad != TIME_QUAD_DEFAULT and time_quad != time_quadrature: + raise ValueError("time_quad and time_quadrature disagree") + time_quad = time_quadrature + if time_quad not in _TIME_QUAD_CHOICES: + # Fail on an unrecognised value rather than silently falling through to + # the default: a typo'd quadrature name that quietly gives you the OLD + # behaviour is exactly the silent-no-op pattern this module keeps + # getting bitten by. + raise ValueError("time_quad must be one of %r, got %r" + % (_TIME_QUAD_CHOICES, time_quad)) + # ``time_quad`` / ``time_upsample`` is the PR-208 low-level API. Preserve + # its primitive-kappa behavior for direct callers; wrappers and drivers use + # the conventional ILE ``time_quadrature`` spelling and the adaptive + # terminal implementation. + legacy_primitive_refinement = (time_quad == "bandlimited" + and not canonical_time_api) + if time_quad == "bandlimited" and _norm_is_arrival_time_dependent(data): + # Same reason, other direction: the band-limited quadrature reconstructs + # kappa(t) and holds the model norm at one time bin, so on data whose + # depends on the arrival time it would quietly return a DIFFERENT + # likelihood rather than a better-integrated one. Refuse it here instead. + raise ValueError( + "time_quad='bandlimited' holds the model norm fixed in time, but this " + "likelihood data carries the slow-rotation post-phase, whose " + "depends on the template arrival time; use time_quad='simpson' for " + "rotation data.") + # Only the band-limited path widens the window; "simpson" integrates the + # sampled window itself, so it must keep gathering exactly data.npts bins. + if legacy_primitive_refinement: + guard = (default_time_guard(data.npts) if time_guard is None + else int(time_guard)) + elif canonical_time_api and time_quad == "bandlimited": + # Start at the established half-window guard, rounded upward to a power + # of two, then gather one doubling as an independent certificate. + g_default = default_time_guard(data.npts) + g_initial = 1 << int(np.ceil(np.log2(g_default))) + guard = 2 * g_initial + else: + guard = 0 distMpc = jnp.asarray(distMpc, dtype=jnp.float64) invDist = data.distMpcRef / distMpc kappa_unit, rho_sq_unit = _accumulate_unit( - data, ra, dec, psi, incl, phiref, interp, phase_marginalization) + data, ra, dec, psi, incl, phiref, interp, phase_marginalization, + guard=guard) kappa_sq = kappa_unit * invDist[:, None] rho_sq = rho_sq_unit * jnp.square(invDist)[:, None] + if legacy_primitive_refinement: + return _time_marginalize_bandlimited( + kappa_sq, rho_sq, data.deltaT, + int(_TIME_UPSAMPLE_DEFAULT if time_upsample is None else time_upsample), guard, + phase_marginalization=phase_marginalization) if phase_marginalization: lnL_t = jnp.abs(kappa_sq) - 0.5 * rho_sq else: lnL_t = kappa_sq.real - 0.5 * rho_sq - return _time_marginalize(lnL_t, data.w_t) + if return_lnLt: + return lnL_t + if canonical_time_api and time_quad == "bandlimited": + return _time_marginalize_reflected_primitive( + kappa_sq, rho_sq, data.deltaT, + phase_marginalization=phase_marginalization, guard=guard) + return _time_marginalize_terminal( + lnL_t, data, time_quad, bandlimited_safe=not phase_marginalization) def fused_log_likelihood_distmarg(data, ra, dec, psi, incl, phiref, x_grid, log_w_grid, - interp="linear", phase_marginalization=False, - grid_block=64): + interp=JAX_INTERP_DEFAULT, phase_marginalization=False, + grid_block=64, + time_quadrature=TIME_QUAD_DEFAULT, + return_lnLt=False): """Distance- AND time-marginalized factored log-likelihood, lnL(angles). Marginalizes the luminosity distance analytically (numerical quadrature over @@ -616,7 +1354,9 @@ def fused_log_likelihood_distmarg(data, ra, dec, psi, incl, phiref, a = x_grid # (G,) b = -0.5 * jnp.square(x_grid) # (G,) lnL_t = _logsumexp_grid_blocked(K, R, a, b, log_w_grid, grid_block) - return _time_marginalize(lnL_t, data.w_t) + if return_lnLt: + return lnL_t + return _time_marginalize_terminal(lnL_t, data, time_quadrature) def _logsumexp_grid_blocked(K, R, a, b, log_w, block): @@ -745,7 +1485,9 @@ def phi_ref_grid(nphi: int) -> np.ndarray: def fused_log_likelihood_phimarg(data, ra, dec, psi, incl, distMpc, - phi_grid, interp="linear"): + phi_grid, interp=JAX_INTERP_DEFAULT, + time_quadrature=TIME_QUAD_DEFAULT, + return_lnLt=False): """Time-marginalized factored lnL with φ_ref marginalized via uniform grid sum. Evaluates the standard factored lnL at each φ_ref in ``phi_grid`` and @@ -782,13 +1524,18 @@ def _phi_step(carry, phi_val): (m, s), _ = jax.lax.scan(_phi_step, (m0, s0), phi_grid_jax) lnL_t_marg = m + jnp.log(s) - jnp.log(nphi) - return _time_marginalize(lnL_t_marg, data.w_t) + if return_lnLt: + return lnL_t_marg + return _time_marginalize_terminal(lnL_t_marg, data, time_quadrature) def fused_log_likelihood_distphimarg(data, ra, dec, psi, incl, x_grid, log_w_grid, - phi_grid, interp="linear", - grid_block=64): + phi_grid, interp=JAX_INTERP_DEFAULT, + grid_block=64, + time_quadrature=TIME_QUAD_DEFAULT, + return_lnLt=False, + return_phi_lnLt=False): """Distance- AND φ_ref-marginalized factored lnL over (ra, dec, psi, incl). Marginalises over both luminosity distance (via quadrature grid, as in @@ -838,14 +1585,18 @@ def _phi_step(carry, phi_val): kappa_unit.real, rho_sq_unit, a, b, log_w_grid, grid_block) m_new = jnp.maximum(m, lnL_t) s_new = s * jnp.exp(m - m_new) + jnp.exp(lnL_t - m_new) - return (m_new, s_new), None + return (m_new, s_new), (lnL_t if return_phi_lnLt else None) m0 = jnp.full((S, data.npts), -jnp.inf, dtype=jnp.float64) s0 = jnp.zeros((S, data.npts), dtype=jnp.float64) - (m, s), _ = jax.lax.scan(_phi_step, (m0, s0), phi_grid_jax) + (m, s), lnL_phi_t = jax.lax.scan(_phi_step, (m0, s0), phi_grid_jax) lnL_t_marg = m + jnp.log(s) - jnp.log(nphi) - return _time_marginalize(lnL_t_marg, data.w_t) + if return_phi_lnLt: + return lnL_phi_t + if return_lnLt: + return lnL_t_marg + return _time_marginalize_terminal(lnL_t_marg, data, time_quadrature) def psi_grid(npsi: int) -> np.ndarray: @@ -860,7 +1611,9 @@ def psi_grid(npsi: int) -> np.ndarray: def fused_log_likelihood_distphipsimarg(data, ra, dec, incl, x_grid, log_w_grid, phi_grid, psi_grid_, - interp="linear", grid_block=64): + interp=JAX_INTERP_DEFAULT, grid_block=64, + time_quadrature=TIME_QUAD_DEFAULT, + return_lnLt=False): """Distance-, phi_ref- AND psi-marginalized factored lnL over (ra, dec, incl). Marginalizes luminosity distance (quadrature grid), orbital phase phi_ref and @@ -909,12 +1662,16 @@ def _step(carry, pair): # backward pass -> memory O(1) in the grid size. (m, s), _ = jax.lax.scan(jax.checkpoint(_step), (m0, s0), pairs) lnL_t_marg = m + jnp.log(s) - jnp.log(npair) - return _time_marginalize(lnL_t_marg, data.w_t) + if return_lnLt: + return lnL_t_marg + return _time_marginalize_terminal(lnL_t_marg, data, time_quadrature) def fused_log_likelihood_distpsimarg(data, ra, dec, phiref, incl, x_grid, log_w_grid, psi_grid_, - interp="linear", grid_block=64): + interp=JAX_INTERP_DEFAULT, grid_block=64, + time_quadrature=TIME_QUAD_DEFAULT, + return_lnLt=False): """Distance- AND psi-marginalized factored lnL over (ra, dec, phi_ref, incl). Marginalizes luminosity distance (quadrature grid) and polarization psi @@ -956,11 +1713,14 @@ def _psi_step(carry, psi_val): # remat: cheap insurance (psi grid is small, but keeps gradient memory O(1)). (m, s), _ = jax.lax.scan(jax.checkpoint(_psi_step), (m0, s0), psi_g) lnL_t_marg = m + jnp.log(s) - jnp.log(npsi) - return _time_marginalize(lnL_t_marg, data.w_t) + if return_lnLt: + return lnL_t_marg + return _time_marginalize_terminal(lnL_t_marg, data, time_quadrature) def phi_ref_conditional_lnL(data, ra, dec, psi, incl, distMpc, - phi_grid, interp="linear"): + phi_grid, interp=JAX_INTERP_DEFAULT, + time_quadrature=TIME_QUAD_DEFAULT): """Log-likelihood vs φ_ref given the other extrinsic parameters. Returns a ``(nphi, S)`` array of time-marginalized lnL values, one per @@ -980,37 +1740,79 @@ def _phi_step(_, phi_val): kappa = kappa_unit * invDist[:, None] rho_sq = rho_sq_unit * jnp.square(invDist)[:, None] lnL_t = kappa.real - 0.5 * rho_sq - return None, _time_marginalize(lnL_t, data.w_t) # carry=None, out=(S,) + return None, _time_marginalize_terminal(lnL_t, data, time_quadrature) _, lnL_per_phi = jax.lax.scan(_phi_step, None, phi_grid_jax) return lnL_per_phi # (nphi, S) +def _distance_prior_density(d, d_prior): + """Unnormalized distance prior density on the nodes ``d``.""" + if d_prior in ("euclidean", "volumetric"): + return d ** 2 + if d_prior == "uniform": + return np.ones_like(d) + raise NotImplementedError("d_prior=%r" % d_prior) + + +def _adaptive_distance_nodes(d_min, d_max, d_peak, sigma_d, n_fine_max, n_coarse, + n_sigma, oversample): + """Node positions for the adaptive grid, or None if the request is degenerate.""" + half = n_sigma * sigma_d # additive: peak is well-located + d_lo = max(float(d_min), d_peak - half) + d_hi = min(float(d_max), d_peak + half) + if not (d_hi > d_lo) or not (sigma_d > 0): + return None + n_fine = int(np.clip((d_hi - d_lo) / (sigma_d / float(oversample)), + 32, int(n_fine_max))) + fine = np.linspace(d_lo, d_hi, n_fine) + coarse = np.linspace(float(d_min), float(d_max), int(n_coarse)) + return np.unique(np.concatenate([coarse, fine])) # sorted, deduped + + +def _trapezoidal_spacing(d): + dd = np.empty_like(d) # trapezoidal spacing + dd[1:-1] = 0.5 * (d[2:] - d[:-2]) + dd[0] = d[1] - d[0] + dd[-1] = d[-1] - d[-2] + return dd + + def make_distance_grid(d_min, d_max, n_grid=256, d_prior="euclidean", - distMpcRef=DIST_MPC_REF): + distMpcRef=DIST_MPC_REF, d_prior_range=None): """Build (x_grid, log_w_grid) for distance marginalization. Uniform grid in distance ``d``; ``x = distMpcRef/d``. Returns the log quadrature weights ``log( p(d) * Delta_d )`` for the requested prior, normalized so ``sum_g exp(log_w_g) == 1`` (a proper distance average). ``d_prior='euclidean'`` is the volumetric ``p(d) ∝ d^2`` prior. + + ``d_prior_range`` (the ILE ``--limit-distance`` hook) SPLITS the two roles + ``[d_min,d_max]`` otherwise plays at once. Left at None the grid range is + also the normalization range -- the historical behaviour, and the reason a + narrowed grid used to renormalize the prior onto itself: the marginal comes + back looking "unchanged" while the evidence scale has silently moved. Given + a ``(lo,hi)``, the NODES span ``[d_min,d_max]`` (the range actually + integrated) while the weights are divided by the prior mass over ``(lo,hi)`` + (the physical range), computed with the SAME discrete rule -- so passing + ``d_prior_range == (d_min,d_max)`` reproduces the None branch bitwise. """ d = np.linspace(d_min, d_max, n_grid) dd = d[1] - d[0] - if d_prior in ("euclidean", "volumetric"): - pd = d ** 2 - elif d_prior == "uniform": - pd = np.ones_like(d) - else: - raise NotImplementedError("d_prior=%r" % d_prior) + pd = _distance_prior_density(d, d_prior) w = pd * dd - w = w / np.sum(w) # normalize the distance average + if d_prior_range is None: + norm = np.sum(w) # normalize the distance average + else: + D = np.linspace(d_prior_range[0], d_prior_range[1], n_grid) + norm = np.sum(_distance_prior_density(D, d_prior) * (D[1] - D[0])) + w = w / norm x = distMpcRef / d log_w = np.log(w) return jnp.asarray(x), jnp.asarray(log_w) -def estimate_distance_peak(data, guess_snr=None, n_sky=4000, seed=0, interp="linear"): +def estimate_distance_peak(data, guess_snr=None, n_sky=4000, seed=0, interp=JAX_INTERP_DEFAULT): """Characteristic distance peak/width directly from the precompute. The distance integrand per (sky, time-bin) is exp(K x - 0.5 R x^2) with @@ -1087,7 +1889,7 @@ def _peak(th5): def make_distance_grid_adaptive(d_min, d_max, d_peak, sigma_d, d_prior="euclidean", distMpcRef=DIST_MPC_REF, n_fine_max=160, n_coarse=48, - n_sigma=12.0, oversample=4.0): + n_sigma=12.0, oversample=4.0, d_prior_range=None): """Non-uniform distance grid: fine near the (SNR-set) peak, coarse on the tail. Concentrates resolution where the distance posterior lives while staying @@ -1112,26 +1914,229 @@ def make_distance_grid_adaptive(d_min, d_max, d_peak, sigma_d, d_prior="euclidea nodes (8x LESS memory than a 256 static grid) and the 1/R only enters through stop_gradient -> gradient-stable. That is a kernel change (TODO). """ - half = n_sigma * sigma_d # additive: peak is well-located - d_lo = max(float(d_min), d_peak - half) - d_hi = min(float(d_max), d_peak + half) - if not (d_hi > d_lo) or not (sigma_d > 0): # degenerate -> uniform fallback - return make_distance_grid(d_min, d_max, n_fine_max + n_coarse, d_prior, distMpcRef) - n_fine = int(np.clip((d_hi - d_lo) / (sigma_d / float(oversample)), - 32, int(n_fine_max))) - fine = np.linspace(d_lo, d_hi, n_fine) - coarse = np.linspace(float(d_min), float(d_max), int(n_coarse)) - d = np.unique(np.concatenate([coarse, fine])) # sorted, deduped + d = _adaptive_distance_nodes(d_min, d_max, d_peak, sigma_d, n_fine_max, + n_coarse, n_sigma, oversample) + if d is None: # degenerate -> uniform fallback + return make_distance_grid(d_min, d_max, n_fine_max + n_coarse, d_prior, + distMpcRef, d_prior_range=d_prior_range) + w = _distance_prior_density(d, d_prior) * _trapezoidal_spacing(d) + if d_prior_range is None: + norm = np.sum(w) + else: + # --limit-distance: nodes span the narrowed range, but the weights carry the + # prior mass over the PHYSICAL range, evaluated with this same construction so + # d_prior_range == (d_min,d_max) reproduces the branch above bitwise. + D = _adaptive_distance_nodes(d_prior_range[0], d_prior_range[1], d_peak, + sigma_d, n_fine_max, n_coarse, n_sigma, oversample) + if D is None: + D = np.linspace(d_prior_range[0], d_prior_range[1], n_fine_max + n_coarse) + norm = np.sum(_distance_prior_density(D, d_prior) * (D[1] - D[0])) + else: + norm = np.sum(_distance_prior_density(D, d_prior) * _trapezoidal_spacing(D)) + w = w / norm + return jnp.asarray(distMpcRef / d), jnp.asarray(np.log(w)) + + +# Log-uniform ("peak-resolving") distance quadrature. Contract and evidence: +# DESIGN_jax_distance_quadrature.md, beside this file. +DIST_GRID_TOL_DEFAULT = 1e-2 +DIST_GRID_SCHEMES = ("uniform", "loguniform") + + +def loguniform_spacing_for_tolerance(tol): + """Relative node spacing ``c`` (in ln d) that a Gaussian peak of unit + relative width tolerates at fractional quadrature error ``tol``. + + The trapezoid rule on a Gaussian converges super-algebraically: by Poisson + summation the fractional error of ``sum_k h f(u_k)`` against ``int f du`` + for ``f = exp(-(u-mu)^2 / 2 s^2)`` is ``2 exp(-2 pi^2 s^2 / h^2)`` (the + k = +-1 aliases; higher ones are negligible), independent of ``mu`` up to + its sign. Setting that equal to ``tol`` and writing ``h = c * s``: + + c = pi * sqrt(2 / ln(2 / tol)) + + So the spacing is DERIVED from a stated tolerance, not tuned. ``tol`` is a + FRACTIONAL error on the distance integral, i.e. ~``tol`` nats on lnL. + """ + tol = float(tol) + if not (0.0 < tol < 2.0): + raise ValueError("dist_grid_tol must be in (0, 2); got %r" % (tol,)) + return float(np.pi * np.sqrt(2.0 / np.log(2.0 / tol))) + + +ENDPOINT_ERROR_MARGIN = 2.0 +# The endpoint model below is the LEADING Euler-Maclaurin term, and c(tol) puts +# the node spacing at ~2 sigma, which is not an asymptotic regime: against a +# directly evaluated truncated trapezoid the leading term under-reads by ~1.5x +# at one width of clearance (the h^4 term and the truncated normalization both +# push the same way). So the guard carries a stated 2x margin, in the same +# spirit as anglemarg.ANGLE_AMP_MARGIN, rather than pretending the series is +# converged. test_endpoint_error_model_tracks_the_measured_truncated_trapezoid +# pins the model against the measurement, which is what bounds this factor. + + +def loguniform_endpoint_error(dlnd, endpoint_scale): + """Fractional error the PRIOR ENDPOINTS add to the alias law of + :func:`loguniform_spacing_for_tolerance`. + + That law is Poisson summation on an UNTRUNCATED Gaussian: it bounds the + aliasing of an infinite trapezoid sum, and is independent of where the peak + sits. A distance prior is finite, so the sum stops, and Euler-Maclaurin + contributes a term the alias law knows nothing about -- proportional to the + integrand's DERIVATIVE at each retained endpoint: + + eps = h^2 / 12 * |g'(edge)| / integral(g) + = h^2 / 12 * rho^2 k exp(-k^2/2) / sqrt(2 pi) + + for a peak of width ``1/rho`` in ``ln d`` sitting ``k`` widths inside the + edge. ``endpoint_scale`` is the ``rho^2 (k exp(-k^2/2))`` half of that, + summed over the two edges and maximized over the loud angle configurations + by ``anglemarg.estimate_angle_amplitude``; ``dlnd`` is the grid's spacing. + + The two errors do NOT cancel and are not the same effect: a peak one width + inside an edge is ~11% wrong (measured) at the shipped ``tol = 1e-2`` + spacing while the alias term is at its promised 1%. A guard on ``tol`` + alone therefore does not deliver ``tol``, which is what this exists to say. + """ + # array-friendly: loguniform_min_clearance sweeps k through it + return (np.asarray(dlnd, dtype=float) ** 2 + * np.asarray(endpoint_scale, dtype=float) + / (12.0 * np.sqrt(2.0 * np.pi))) + + +def loguniform_min_clearance(tol=DIST_GRID_TOL_DEFAULT): + """Clearance, in peak widths, that an edge needs at the WORST spacing. + + Solves ``ENDPOINT_ERROR_MARGIN * loguniform_endpoint_error(c(tol), k) = + tol`` for the peak that sizes the grid (``rho = rho_max``, so ``h = c``), + taking the LARGE root: ``k exp(-k^2/2)`` rises to 0.6065 at one width and + falls away on both sides, so the small root is the "endpoint on the peak" + case the alias law already covers and must not be reported as a requirement. + Reported in refusals; the guard itself uses the error, not this number, + because a quieter peak (``rho < rho_max``) is on a proportionally finer grid + and needs less. + """ + c = loguniform_spacing_for_tolerance(tol) + k = np.linspace(0.0, 12.0, 24001) + over = (ENDPOINT_ERROR_MARGIN + * loguniform_endpoint_error(c, k * np.exp(-0.5 * np.square(k))) + > float(tol)) + return float(k[over].max()) if np.any(over) else 0.0 + + +def loguniform_grid_size(d_min, d_max, rho_max, tol=DIST_GRID_TOL_DEFAULT): + """Node count for :func:`make_distance_grid_loguniform` (pure, testable).""" + rho_max = float(rho_max) + if not np.isfinite(rho_max) or rho_max <= 0.0: + raise ValueError( + "rho_max must be a finite positive matched-SNR bound; got %r. " + "It is sqrt(2*A) with A the data-derived amplitude from " + "anglemarg.estimate_angle_amplitude on the FULL prior support; " + "there is deliberately no fallback -- a missing bound must not " + "silently produce an under-resolved grid." % (rho_max,)) + if not (0.0 < float(d_min) < float(d_max)): + raise ValueError("need 0 < d_min < d_max; got (%r, %r)" % (d_min, d_max)) + L = np.log(float(d_max) / float(d_min)) + c = loguniform_spacing_for_tolerance(tol) + return int(np.ceil(rho_max * L / c)) + 1 + + +def make_distance_grid_loguniform(d_min, d_max, rho_max, d_prior="euclidean", + distMpcRef=DIST_MPC_REF, + tol=DIST_GRID_TOL_DEFAULT, n_max=8192): + """Distance grid whose RELATIVE spacing resolves every per-sample peak. + + WHAT IS BEING INTEGRATED. Per angle sample and time bin the distance + integrand is ``exp(K x - 0.5 R x^2)`` with ``x = distMpcRef / d``, + ``K = Re`` and ``R = `` at the reference distance. That is a + Gaussian in ``x`` peaked at ``x* = K/R`` with standard deviation + ``1/sqrt(R) = x* / rho``, where ``rho = K / sqrt(R)`` is that sample's + matched SNR. Its RELATIVE width ``sigma/x* = 1/rho`` is therefore SCALE + FREE: it does not depend on where the peak sits. + + CONSEQUENCE, and the whole content of this function. A grid that is + uniform in ``ln d`` has constant relative spacing, so ONE spacing resolves + every peak anywhere in ``[d_min, d_max]`` as soon as + + Delta(ln d) <= c / rho_max, c = loguniform_spacing_for_tolerance(tol) + + with ``rho_max`` an estimate of the largest ``rho`` over the angles. No + peak has to be located. Contrast :func:`make_distance_grid_adaptive`, which + centres a window on an ESTIMATED peak and is wrong by ~13 nats when that + estimate is wrong (measured; DESIGN_jax_distance_quadrature.md). + + PRECONDITION -- the contract above holds only where the integrand is a + Gaussian PEAK INSIDE ``[d_min, d_max]``. If the maximizing distance + ``x* = A/B`` is EXTERIOR the integrand is a boundary layer at a prior edge + instead, and this grid is the wrong instrument for it: its absolute spacing + is coarsest exactly at ``d_max``, and refining it adds nodes proportionally + everywhere so the layer never resolves (measured 1.9-4.6 nats, WORSE than + the uniform default, and tightening ``tol`` from 0.5 to 1e-9 recovers only + 5.23 -> 3.92). Callers must detect and refuse that regime; the wrapper + does, via ``estimate_angle_amplitude(..., return_diagnostics=True)`` and + its ``clip_excess``. Design note section 1a. + + PRECONDITION, SECOND HALF -- interior is not enough: the peak must be + interior BY A MARGIN. Poisson summation bounds an UNTRUNCATED trapezoid + sum; a support that cuts the peak's tail adds an Euler-Maclaurin endpoint + term the alias law does not see, and at the shipped ``tol = 1e-2`` spacing + (``c = 1.93``) a peak ONE width inside an edge is ~11% wrong while + ``clip_excess`` still reads exactly 1. The requirement is on the error, not + on a bare distance: see :func:`loguniform_endpoint_error` and + :func:`loguniform_min_clearance`, which the wrapper evaluates on the loud + angle configurations and refuses on. ``rho_max`` alone cannot express it -- + two supports with the same ``rho_max`` differ entirely in this respect -- + which is why this function still takes only ``rho_max`` and the check lives + with the estimator that knows where the peak is. + + WHERE ``rho_max`` COMES FROM. ``A = anglemarg.estimate_angle_amplitude`` + is ``ANGLE_AMP_MARGIN`` times the ``max`` over a SAMPLED sky, and over the + distance support, of ``x A_ang - 0.5 x^2 B_ang``, whose closed-form maximum + in ``x`` is ``A_ang^2 / (2 B_ang) = rho^2 / 2``. So + ``rho_max = sqrt(2 A) = sqrt(ANGLE_AMP_MARGIN) * rho_sampled_max``. This is + deliberately NOT called an identity and NOT a proven bound -- that + estimator's own docstring says it is an estimator -- but it introduces no + NEW estimator: it is the same number that sizes the dense angle lattice, so + the two cannot disagree, and the kernels' runtime fail-safe + (``anglemarg._runtime_amp_failsafe``) rechecks it on every call. That + fail-safe covers an underestimated INTERIOR peak; it is blind to the + exterior regime above, because it applies the identical clip. ``A`` must be + computed on the FULL prior support, never on this grid. + + WEIGHTS. Trapezoidal ``p(d) * Delta d`` with HALF-WIDTH end intervals, + normalized to ``sum exp(log_w) == 1`` (the same "proper distance average" + convention as :func:`make_distance_grid`, whose own constant-``Delta d`` + weights are a right-open rectangle rule). The half-width endpoints matter + here and are not cosmetic: on a log grid the last interval is ~1% of + ``d_max``, and giving the last node a full interval (the convention + :func:`make_distance_grid_adaptive` uses) misplaces several percent of the + volumetric prior mass onto ``d_max`` -- measured as a ~0.018 nat error + floor that no refinement removes. + + Returns ``(x_grid, log_w_grid)`` -- drop-in for every fused kernel. + """ + n = loguniform_grid_size(d_min, d_max, rho_max, tol) + if n > int(n_max): + raise ValueError( + "loguniform distance grid needs %d nodes for rho_max=%.4g over " + "[%g, %g] Mpc at tol=%g, above n_max=%d. Raise n_max (and accept " + "the cost, which is linear in the node count), loosen tol, or " + "narrow [d_min, d_max]. Clamping is deliberately NOT done: a " + "silently clamped grid violates the spacing contract this " + "function exists to provide." + % (n, float(rho_max), float(d_min), float(d_max), float(tol), + int(n_max))) + d = np.geomspace(float(d_min), float(d_max), n) if d_prior in ("euclidean", "volumetric"): pd = d ** 2 elif d_prior == "uniform": pd = np.ones_like(d) else: raise NotImplementedError("d_prior=%r" % d_prior) - dd = np.empty_like(d) # trapezoidal spacing + dd = np.empty_like(d) dd[1:-1] = 0.5 * (d[2:] - d[:-2]) - dd[0] = d[1] - d[0] - dd[-1] = d[-1] - d[-2] + dd[0] = 0.5 * (d[1] - d[0]) + dd[-1] = 0.5 * (d[-1] - d[-2]) w = pd * dd w = w / np.sum(w) return jnp.asarray(distMpcRef / d), jnp.asarray(np.log(w)) @@ -1146,14 +2151,16 @@ def _tref_minus_epoch(self, det): JAXLikelihoodData.tref_minus_epoch = _tref_minus_epoch -def make_log_likelihood(data, interp="linear", phase_marginalization=False, +def make_log_likelihood(data, interp=JAX_INTERP_DEFAULT, phase_marginalization=False, jit=True): """Return a closure ``f(ra, dec, psi, incl, phiref, distMpc) -> lnL``. The returned function closes over ``data`` (treated as constant) and is, by default, ``jax.jit``-compiled. It is differentiable with respect to all six - extrinsic arguments when ``interp="linear"``; combine with ``jax.grad`` / - ``jax.value_and_grad`` / ``jax.vmap`` as needed. + extrinsic arguments for any INTERPOLATING stencil -- ``linear``, ``cubic`` or + ``sinc`` -- but NOT for ``nearest``, whose gather is piecewise constant in the + arrival time and therefore has zero gradient through the sky. Combine with + ``jax.grad`` / ``jax.value_and_grad`` / ``jax.vmap`` as needed. """ def f(ra, dec, psi, incl, phiref, distMpc): return fused_log_likelihood( diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py new file mode 100644 index 000000000..41ceda4aa --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py @@ -0,0 +1,287 @@ +"""Joint (phi, psi) peak-local angle marginalization, JAX kernel. + +The numpy reference is ``RIFT.likelihood.joint_angle_peak_local``; this is the jittable +form of the same rule. It is NOT a transcription -- the reference builds 2-D regions +and merges overlapping ones, which is data-dependent control flow and does not jit. The +formulation here removes the need to merge at all. + +THE PARTITION THAT REPLACES MERGING. At fixed ``phi`` the exponent is +``a + Re(c1 e^{iu}) + Re(c2 e^{2iu})``, whose u-stationary points are the roots of a +quartic. On the circle maxima and minima ALTERNATE, so the sorted stationary points +already tile the domain: the cell of a maximum is the arc between its two neighbouring +minima. Those cells are disjoint by construction and cover the circle, so there is +nothing to merge and nothing to double-count -- the failure the reference spends +``_merge_boxes`` on cannot arise. Everything is then static: 4 roots, 4 candidate +cells, a fixed number of quadrature nodes in each. + +WHY THE ROOTS ARE TAKEN WITHOUT A ``|z| = 1`` FILTER. At exact multiplicity the +computed roots smear off the unit circle by ``eps^(1/m)`` -- measured 4.6e-6 for a +triple root -- so a fixed tolerance drops real modes in precisely the degenerate regime +that is normal here. Every root contributes its angle; a spurious one produces a +zero-length or redundant cell, which is harmless, whereas a dropped one loses mass. + +WHAT SCALES WITH AMPLITUDE AND WHAT DOES NOT. The stationary points of ``g`` do not +move when the data amplitude grows -- ``g -> lambda g`` leaves them fixed -- so the +CELLS are amplitude-independent, while the peak inside each cell narrows as +``A^-1/2``. The local window is therefore sized from the local curvature and clipped +to the cell, which keeps the node count fixed. This is the u axis's whole economy: the +shipped dense scheme spends ``~sqrt(A)`` points on this axis, and this spends a +constant. + +SCOPE OF THIS KERNEL. The u axis is localized; the phi axis is a dense grid, scanned +in chunks. That is deliberately the same cost shape as the shipped ``laplace`` scheme +(``~sqrt(A)`` on phi) and a strict improvement on its u treatment, which uses a blended +O(1/A) width model rather than the exact stationary points. Localizing phi as well -- +the (phi localized, psi localized) cell of the family -- needs the profile ``F(phi)`` +and its envelope derivative, and is not attempted here. + +MEMORY. Bounded by ``phi_chunk`` through ``lax.scan``, never by the grid: the largest +transient is ``(phi_chunk, n_x, 4, n_u)``. It is a cost knob and cannot change the +result beyond floating-point reassociation. +""" + +import numpy as np +import jax +import jax.numpy as jnp +from jax import lax + +__all__ = [ + "required_n_phi", + "U_WINDOW_SIGMA", + "U_NODES_PER_CELL", + "PHI_CHUNK_DEFAULT", + "u_stationary_roots", + "log_inner_u_integral", + "joint_lnL_phi_dense", +] + +#: Local u-window half-width in units of the local sigma, CLIPPED to the cell. The cell +#: boundaries are minima of the exponent, so clipping loses nothing that the cell itself +#: does not already exclude; this only decides where the window stops being the binding +#: constraint. +U_WINDOW_SIGMA = 12.0 + +#: Trapezoid nodes per cell. Fixed, because the window is sized from the LOCAL +#: curvature: at 48 nodes over +-12 sigma the spacing is sigma/2, and the trapezoid's +#: Poisson-summation error on a Gaussian is 2 exp(-2 pi^2 * 4) = 5e-35. Derived from +#: that bound, not tuned -- the same argument as UPSAMPLE_SAFETY in the band-limited +#: time quadrature. This is the u axis's entire cost: 4 cells x 48 nodes = 192 points +#: per phi, INDEPENDENT of amplitude, against the shipped dense rule's ~6.2 sqrt(A) +#: (896 at amplitude 1.25e4). +U_NODES_PER_CELL = 48 + +#: phi points per scan step. +PHI_CHUNK_DEFAULT = 16 + + +def required_n_phi(amplitude, m_max=2): + """phi-grid size for a given exponent amplitude. + + THE PHI AXIS IS NOT LOCALIZED IN THIS KERNEL, so it inherits the dense scaling and + must be sized, not guessed: ``exp(g)`` has phi-width ``~A^-1/2`` and its harmonic + content reaches ``~6.2 sqrt(A)`` (measured), scaled by mode content. Hard-coding a + value instead cost 191 nats at amplitude 1.25e4 during development -- recorded + because a fixed grid looks harmless right up to the point where it is not. + + Mirrors the phi half of ``anglemarg._dense_grid_sizes`` so the two rules are sized + by one argument rather than two that must agree. + """ + from . import anglemarg as _am + return int(_am._dense_grid_sizes(float(amplitude), m_max=int(m_max))[0]) + + +def _a_c1_c2(C, phi): + """The u-independent term and the ``e^{iu}``, ``e^{2iu}`` coefficients at ``phi``. + + The u-independent term is NOT the single ``(k=0, q=0)`` coefficient: the whole + ``q = 0`` column is u-independent and every one of its ``k`` harmonics depends on + phi. Taking only ``C[0, KS]`` drops that phi structure entirely and biases the + result low by an amount that grows with amplitude -- measured -2.0e-3, -2.6 and + -191 nats at exponent amplitude 6.6, 624 and 1.25e4 before this was fixed. + """ + KP = C.shape[0] + KS = (C.shape[1] - 1) // 2 + k = jnp.arange(KP) + w = jnp.where(k > 0, 2.0, 1.0) + ph = jnp.exp(1j * phi[..., None] * k) * w + D = lambda q: (ph * C[:, KS + q]).sum(-1) + # BOTH SIGNS OF q ARE STORED, so the e^{iu} coefficient is not the +1 column alone: + # Re[D_{-q} e^{-iqu}] = Re[conj(D_{-q}) e^{+iqu}], hence c_q = D_{+q} + conj(D_{-q}). + # Using only the +q column drops half the u-dependence. It is survivable where the + # roots are mere SEEDS -- the numpy reference refines them with 2-D Newton and was + # unaffected -- but here the roots define the cell PARTITION, which is load-bearing, + # and the error reached 17 nats at a single phi. + a = D(0).real + return a, D(1) + jnp.conj(D(-1)), D(2) + jnp.conj(D(-2)) + + +def u_stationary_roots(c1, c2): + """The four u-stationary angles, as a companion eigenproblem. Static shape (4,). + + ``P(z) = c2 z^4 + (c1/2) z^3 - (conj(c1)/2) z - conj(c2)``; the roots' arguments are + the stationary points. No ``|z| = 1`` filtering -- see the module docstring. + """ + a4 = c2 + lead = jnp.where(jnp.abs(a4) > 0, a4, 1.0 + 0j) + co = jnp.stack([c1 / 2.0, jnp.zeros_like(c1), -jnp.conj(c1) / 2.0, + -jnp.conj(c2)]) / lead + comp = jnp.zeros((4, 4), dtype=jnp.complex128) + comp = comp.at[0, :].set(-co) + comp = comp.at[1:, :-1].set(jnp.eye(3, dtype=jnp.complex128)) + # the stop_gradient goes on the INPUT: placing it on the output still leaves JAX + # needing eigvals' JVP rule to build the trace, and that is the rule that does not + # exist. Cutting the tangent before the eigensolve means it is never asked for. + z = jnp.linalg.eigvals(jax.lax.stop_gradient(comp)) + # a vanishing quartic leading coefficient degenerates to a cubic; the extra root is + # spurious but produces only a redundant cell, never a lost one. + # + # STOP_GRADIENT, and it is a correctness statement rather than a convenience. + # (i) It is REQUIRED: jnp.linalg.eigvals has no second derivative in JAX ("the + # derivatives of eigenvectors are not implemented"), so without it any Hessian + # through this kernel raises -- and the caller that matters, _fisher_whitening, + # swallows that in an `except Exception` and silently returns None, so + # --fisher-precondition would degrade to raw coordinates with the flag still + # recorded as supplied. It also removes a NaN: as c2 -> 0 the companion matrix + # acquires ~1/c2 entries and the eig JVP degenerates (measured grad 0.567 at + # c2=1, -1.2e14 at 1e-20, nan at 1e-30). + # (ii) It is CORRECT: these angles are cell BOUNDARIES of an exact partition of the + # circle, so a boundary shift adds to one cell exactly what it removes from its + # neighbour and the contribution cancels identically. Where a window stops short + # of its cell edge the integrand there is ~exp(-W^2/2) of the peak, so that + # residue is far below the truncation already accepted. The same argument, and + # the same device, is used for the distance nodes in core._distmarg_gh_logL. + return jnp.mod(jnp.angle(z), 2.0 * jnp.pi) + + +def _g_u(a, c1, c2, u, order=0): + """``d^order/du^order`` of ``a + Re(c1 e^{iu}) + Re(c2 e^{2iu})``.""" + t1 = ((1j) ** order) * c1 * jnp.exp(1j * u) + t2 = ((2j) ** order) * c2 * jnp.exp(2j * u) + base = (t1 + t2).real + return base + (a if order == 0 else 0.0) + + +def log_inner_u_integral(a, c1, c2, n_nodes=U_NODES_PER_CELL, + window_sigma=U_WINDOW_SIGMA): + """``log int_0^{2pi} du exp(a + Re(c1 e^{iu}) + Re(c2 e^{2iu}))``. + + Exact partition, static shapes. The sorted stationary points alternate max/min, so + the cell of the maximum at ``u_(i)`` is ``[u_(i-1), u_(i+1)]`` and the maxima's cells + tile the circle. Non-maxima contribute a masked ``-inf`` and drop out of the + log-sum-exp, so no filtering or compaction is needed. + """ + u = jnp.sort(u_stationary_roots(c1, c2)) # (4,) + + # MIDPOINT cells, not root-bounded cells. Bounding a maximum's cell by its + # neighbouring roots is only a partition if maxima and minima strictly ALTERNATE, + # and they do not: when two roots leave the unit circle as a conjugate-reciprocal + # pair (measured: |z| = 1.268 and 0.789, product 1.0003) their shared angle is not a + # stationary point at all, and using it as a boundary leaves an arc belonging to no + # cell. That silently dropped 0.23 nats on a low-amplitude draw. Midpoints of the + # sorted angles tile the circle for ANY four angles, so no arc can be orphaned + # however the roots behave -- which is what lets the roots stay unfiltered. + mid = 0.5 * (u + jnp.roll(u, -1) + jnp.where(jnp.arange(4) == 3, 2 * jnp.pi, 0.0)) + lo_c = jnp.roll(mid, 1) - jnp.where(jnp.arange(4) == 0, 2 * jnp.pi, 0.0) + hi_c = mid + + # locate the maximum INSIDE each cell: the cell's own root is only a seed, and for a + # spurious root it is not even a stationary point. + def _newton(uc, _): + g1 = _g_u(a, c1, c2, uc, 1) + g2 = _g_u(a, c1, c2, uc, 2) + step = jnp.where(jnp.abs(g2) > 0, -g1 / jnp.where(jnp.abs(g2) > 0, g2, 1.0), 0.0) + step = jnp.clip(step, -0.5, 0.5) + return jnp.clip(uc + step, lo_c, hi_c), None + + ustar, _ = lax.scan(_newton, u, None, length=8) + + # A CLIPPED NEWTON POINT IS NOT A PEAK, however negative the curvature. The + # iteration is clamped to [lo_c, mid], so it can come to rest ON a boundary with a + # large stationary residual; curvature alone then centres a +-W sigma window on a + # non-stationary point and sizes sigma from the wrong curvature. Measured in the + # numpy twin: 18% of cells that g'' < 0 accepted fail this gate, the worst at + # |g_u|/M_1 = 0.33. A cell failing it is integrated WHOLE, which can only add nodes. + g1s = _g_u(a, c1, c2, ustar, 1) + g2s = _g_u(a, c1, c2, ustar, 2) + m1u = jnp.abs(c1) + 2.0 * jnp.abs(c2) # exact bound on |d g / du| + edge = 1e-9 * jnp.max(mid - lo_c) + peaked = ((g2s < 0.0) + & (jnp.abs(g1s) <= 1e-8 * jnp.maximum(m1u, 1e-300)) + & (ustar > lo_c + edge) & (ustar < mid - edge)) + sigma = jnp.where(peaked, 1.0 / jnp.sqrt(jnp.where(peaked, -g2s, 1.0)), jnp.inf) + # a cell with no interior maximum is integrated whole; a peaked one is integrated on + # +-window_sigma, which is self-limiting -- when the integrand is flat sigma is large + # and the window IS the cell. + lo = jnp.where(peaked, jnp.maximum(ustar - window_sigma * sigma, lo_c), lo_c) + hi = jnp.where(peaked, jnp.minimum(ustar + window_sigma * sigma, hi_c), hi_c) + width = jnp.maximum(hi - lo, 0.0) + + s = jnp.linspace(0.0, 1.0, n_nodes) # (n,) + uu = lo[:, None] + width[:, None] * s[None, :] # (4, n) + gg = _g_u(a, c1, c2, uu, 0) + wq = jnp.full(n_nodes, 1.0 / (n_nodes - 1)) + wq = wq.at[0].mul(0.5).at[-1].mul(0.5) + logw = jnp.log(wq)[None, :] + jnp.log(jnp.where(width > 0, width, 1.0))[:, None] + cell = jax.scipy.special.logsumexp(gg + logw, axis=-1) # (4,) + cell = jnp.where(width > 0, cell, -jnp.inf) + return jax.scipy.special.logsumexp(cell) + + +def _joint_table(C_A, C_B, x): + """``x A - x^2/2 B`` with A zero-padded into B's (larger) bidegree.""" + ksa = (C_A.shape[1] - 1) // 2 + ksb = (C_B.shape[1] - 1) // 2 + out = (-0.5 * x * x) * C_B + return out.at[:C_A.shape[0], ksb - ksa:ksb + ksa + 1].add(x * C_A) + + +def joint_lnL_phi_dense(C_A, C_B, x_grid, log_w_grid, n_phi=256, + phi_chunk=PHI_CHUNK_DEFAULT, + n_nodes=U_NODES_PER_CELL): + """Distance-, phi- and psi-marginalized value at one ``(sample, time)``. + + Same normalization as ``anglemarg.fused_log_likelihood_distphipsimarg_*``: uniform + priors ``dphi/2pi`` and ``dpsi/pi``, which in ``u = 2 psi`` is ``(2 pi)^-2`` times + the torus integral. + + ``phi`` is a dense grid scanned in chunks; ``u`` is exact per the cell partition. + """ + C_A = jnp.asarray(C_A, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + x_grid = jnp.asarray(x_grid, dtype=jnp.float64).ravel() + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64).ravel() + KS = (C_B.shape[1] - 1) // 2 + + tables = jax.vmap(lambda x: _joint_table(C_A, C_B, x))(x_grid) # (nx, KP, 2KS+1) + + phis = jnp.linspace(0.0, 2.0 * jnp.pi, n_phi, endpoint=False) + n_chunk = int(np.ceil(n_phi / phi_chunk)) + pad = n_chunk * phi_chunk - n_phi + phis_p = jnp.concatenate([phis, jnp.zeros(pad)]) + live = jnp.concatenate([jnp.ones(n_phi, bool), jnp.zeros(pad, bool)]) + + def one_phi(phi): + a, c1, c2 = jax.vmap(lambda T: _a_c1_c2(T, jnp.atleast_1d(phi)))(tables) + return jax.vmap(log_inner_u_integral, in_axes=(0, 0, 0, None))( + a[:, 0], c1[:, 0], c2[:, 0], n_nodes) # (nx,) + + def step(carry, args): + ph, lv = args + vals = jax.vmap(one_phi)(ph) # (chunk, nx) + vals = jnp.where(lv[:, None], vals, -jnp.inf) + return carry, vals + + # jax.checkpoint on the scan body, as the shipped exact scheme does. Without it a + # REVERSE-mode pass keeps every chunk's intermediates: the wrapper's Hessian tried to + # allocate 135 GB and died RESOURCE_EXHAUSTED, so --fisher-precondition would have + # OOMed rather than run. Forward evaluation was never affected, which is exactly why + # this was invisible until a second derivative was taken. + _, out = lax.scan(jax.checkpoint(step), None, + (phis_p.reshape(n_chunk, phi_chunk), + live.reshape(n_chunk, phi_chunk))) + vals = out.reshape(n_chunk * phi_chunk, -1)[:n_phi] # (n_phi, nx) + + # phi is a periodic trapezoid == plain mean; then the distance sum; then (2pi)^-2 + per_x = jax.scipy.special.logsumexp(vals, axis=0) - jnp.log(n_phi) \ + + jnp.log(2.0 * jnp.pi) + return jax.scipy.special.logsumexp(per_x + log_w_grid) - 2.0 * jnp.log(2.0 * jnp.pi) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py index 90a365f01..b53dc39c6 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py @@ -54,6 +54,129 @@ _PI = float(np.pi) +# --------------------------------------------------------------------------- +# Likelihood tempering: what --adapt-weight-exponent costs on THIS path +# --------------------------------------------------------------------------- +# Decision + provenance: DESIGN_jax_tempering.md, beside this module (2026-08-23). +# The short version, because it is the thing that gets ported wrong: +# +# non-JAX RIFT beta shapes only the adaptive sampling PRIOR +# (mcsamplerGPU: log_weights = beta*lnL + ln p - ln p_s, while the +# estimator stays log_integrand = lnL + ln p - ln p_s). Unbiased +# at any beta; beta costs nothing in exported samples. +# JAX flowMC beta = inv_T is the exponent of the target the MCMC SAMPLES. +# The draws are the deliverable, so the export must be reweighted +# by L^(1-beta) (post_weight) -- and that reweight has an ESS cost +# the non-JAX path never pays. +# +# helper_LDG_Events.py keys its beta on SNR. That is right there and wrong here: +# the cost below is set by the SAMPLED DIMENSION and is independent of lnLmax. +# Measured ratio (measured ESS/N) / (Gaussian law) at dim=4, as a CONSERVATIVE +# piecewise-linear envelope: every knot sits at or below every measured point of +# the sweep in DESIGN_jax_tempering.md. The law is optimistic and increasingly so +# at small beta, so a single flat factor is the wrong shape -- a flat 0.79 would +# also make any target above 0.79 unachievable, since it never reaches 1. +# +# MEASURED AT dim=4 ONLY. Applying the same ratio at other dimensions is an +# assumption, not a measurement; it is the conservative direction (the law is +# optimistic in dim too, since the exponent grows), but it is not verified. +_ESS_CAL_BETA = (0.05, 0.20, 0.40, 0.60, 0.80, 1.00) +_ESS_CAL_RATIO = (0.79, 0.79, 0.83, 0.91, 0.97, 1.00) +_TEMPER_ESS_LAW_CAL = _ESS_CAL_RATIO[0] # worst case, retained for reference + + +def export_ess_fraction(beta, n_dim): + """Fraction of a beta-tempered cloud that survives the post_weight reweight. + + For a locally Gaussian peak ``ln Z_g = g lnLmax - (n_dim/2) ln g + const``, + so the self-normalised reweight ``L^(1-beta)`` from the tempered target + ``L^beta pi`` back to the posterior has + + ESS/N = Z_1^2 / (Z_beta Z_{2-beta}) = [beta (2 - beta)]^(n_dim/2) + + **It depends on n_dim and NOT on lnLmax** -- i.e. not on SNR. That is the + whole reason the non-JAX helper's SNR-keyed exponent must not be carried + over to this path. + + Measured against the real phi-marginalised BNS likelihood (SNR 23.8, 4-D), + the law holds to a ratio 0.79-1.00 over beta in [0.05, 1]; it is therefore a + slightly OPTIMISTIC closed form. Callers sizing a budget should apply + ``_TEMPER_ESS_LAW_CAL``. Provenance and the full sweep: DESIGN_jax_tempering.md. + """ + beta = float(beta) + if not (0.0 < beta <= 1.0): + raise ValueError("beta must be in (0, 1]; got %r" % (beta,)) + return float((beta * (2.0 - beta)) ** (0.5 * int(n_dim))) + + +def _ess_law_calibration(beta): + """Conservative lower-bound ratio (measured/law) at this beta, from the sweep.""" + beta = float(beta) + if beta <= _ESS_CAL_BETA[0]: + return _ESS_CAL_RATIO[0] + if beta >= _ESS_CAL_BETA[-1]: + return _ESS_CAL_RATIO[-1] + return float(np.interp(beta, _ESS_CAL_BETA, _ESS_CAL_RATIO)) + + +def export_ess_estimate(beta, n_dim): + """Calibrated ESTIMATE of the surviving export fraction. NOT a bound. + + ``export_ess_fraction`` is the Gaussian-peak law; the SNR-23.8 sweep shows it + optimistic by up to 21% (ratio 0.79 at beta=0.05, rising to 1.00 at beta=1), + and this applies that ratio. + + IT IS NOT A GUARANTEE, AND THIS FUNCTION WAS ONCE NAMED AS IF IT WERE. The + calibration is fitted at SNR ~= 23.8 only, and the shortfall grows with SNR: + the SNR ladder in DESIGN_jax_tempering.md measures ESS/N = 0.00823 at + beta=0.1, d=4, SNR ~= 67, against 0.0285 from this estimate -- a factor 3.5 + the wrong way. A real bound needs a calibration in (beta, SNR), and this + driver has no trustworthy SNR at the point the choice is made (`guess_snr` is + an explicit guesstimate: 10.32 against a true network 23.78 on the study + event). So callers must treat the result as advisory and must not refuse a + run on it. Reported by review on #186. + """ + return _ess_law_calibration(beta) * export_ess_fraction(beta, n_dim) + + +def beta_for_export_ess(target_frac, n_dim): + """Inverse of :func:`export_ess_fraction`: the SMALLEST beta meeting a target. + + Solves ``[beta(2-beta)]^(n_dim/2) = target_frac`` on ``beta in (0, 1]``: + + beta = 1 - sqrt(1 - target_frac^(2/n_dim)) + + Smallest is the useful root: beta is a breadth knob, so among exponents that + meet the export budget the broadest target is the one that explores most. + + Solves against :func:`export_ess_estimate`, so the returned beta meets + ``target_frac`` **on the SNR ~= 23.8 calibration**. That is not a guarantee + at other SNRs -- the shortfall grows with SNR (see that function) -- which is + why ``--auto-adapt-weight-exponent`` is documented as experimental and why + nothing refuses a run on this number. + """ + t = float(target_frac) + if not (0.0 < t <= 1.0): + raise ValueError("target_frac must be in (0, 1]; got %r" % (t,)) + # Solve on the CALIBRATED lower bound, not the bare law. Both factors are + # non-decreasing in beta, so the product is monotone and bisection is safe. + # There is no closed form once the piecewise-linear calibration is included, + # and the previous closed-form inverse of the optimistic law returned betas + # that retained less than the caller asked for. + if export_ess_estimate(1.0, n_dim) < t: + raise ValueError( + "target_frac %g is unreachable in %d-D even at beta=1 (lower bound " + "%.4f)" % (t, int(n_dim), export_ess_estimate(1.0, n_dim))) + lo, hi = 1e-6, 1.0 + for _ in range(200): + mid = 0.5 * (lo + hi) + if export_ess_estimate(mid, n_dim) >= t: + hi = mid + else: + lo = mid + return float(hi) + + # --------------------------------------------------------------------------- # Prior (physical, uniform sky + orientation), numpy + JAX flavors # --------------------------------------------------------------------------- @@ -108,6 +231,66 @@ def _log_prior_jax(theta5): # --------------------------------------------------------------------------- # Batched lnL evaluation (chunked to bound memory) # --------------------------------------------------------------------------- +# Largest single XLA buffer of the anglemarg laplace path, per sample per +# time point: the (quad_chunk=16, dist_block=4, phi_chunk=16) stacked +# quadrature block, 16*4*16*8 = 8192 bytes. Measured 2026-08-28: at the +# default chunk 4000 with npts=1193 XLA requested exactly 36.41 GiB for that +# buffer and the SNR-40 acceptance run died RESOURCE_EXHAUSTED on a 25 GiB +# cgroup -- the pre-fix code never got past COMPILATION at production size, +# so this execution-side wall was previously unreachable. The exact scheme's +# dense reconstruction has the same batch-multiplied structure (smaller +# constant); the laplace constant is used for both as the worst case. +_ANGLE_MARG_BYTES_PER_SAMPLE_PT = 8192 +_ANGLE_MARG_BUFFER_TARGET = 4 << 30 # ~4 GiB largest single buffer + + +def angle_marg_eval_chunk(like, chunk): + """Cap the batched-eval chunk when ``like`` runs an anglemarg scheme. + + Slices of the batched eval are INDEPENDENT (lnL is elementwise in the + sample axis), so this changes peak memory and nothing else -- same + pattern as the _GH_NODES shrink above. Grid-scheme and 4/5-param + likelihoods pass through unchanged. + """ + # NOT the scheme default. "grid" here is a SENTINEL meaning "this object + # runs no dense angle scheme" -- it is what a JAXDistanceMarginalized/ + # JAXExtrinsic likelihood, which has no angle_marg_scheme at all, must fall + # back to. Do NOT sync it to ANGLE_MARG_DEFAULT: that would shrink the eval + # chunk for every likelihood that does not need it. Two independent things + # that happened to be the same string; the last default move on this path + # (interp linear -> sinc) was bitten by exactly that. + # 'peak-local' is capped WITH the dense schemes, not exempted from them. Its u + # axis is localized, but it still nests sample/time vmaps over the distance grid, + # phi chunks, four cells and 48 u nodes, so the batch multiplies the same way the + # dense schemes do; the laplace bytes-per-sample-point constant is used for it as + # the worst case, exactly as it already is for exact. Leaving it out kept an + # uncapped 8000-sample batch and reopened the 36.4 GiB failure documented above. + if getattr(like, "angle_marg_scheme", "grid") not in ("exact", "laplace", + "peak-local"): + return chunk + npts = int(getattr(getattr(like, "data", None), "npts", 0) or 0) + if npts <= 0: + return chunk + bytes_per = _ANGLE_MARG_BYTES_PER_SAMPLE_PT + if getattr(like, "angle_marg_scheme", None) == "peak-local": + # ITS COST MODEL IS NOT THE DENSE ONE, and enrolling it in the cap without + # saying so was a review finding. peak-local carries the WHOLE distance grid + # inside every phi chunk, so its live slab is + # phi_chunk * n_x * (4 cells) * (u nodes) * 8 bytes + # per (sample, time-point) -- about 6.3 MB at phi_chunk=16 and n_x=256, roughly + # 770x the 8192-byte dense model, before intermediates. Using the dense + # constant would have applied a cap that looks protective and is not. + from . import joint_anglemarg_peaklocal as _jp + n_x = int(np.size(getattr(like, "x_grid", ())) or 1) + bytes_per = max( + bytes_per, + _jp.PHI_CHUNK_DEFAULT * n_x * 4 * _jp.U_NODES_PER_CELL * 8) + cap = max(1, _ANGLE_MARG_BUFFER_TARGET // (bytes_per * npts)) + return min(chunk, cap) + # A floor larger than one defeats the memory bound for long, valid time + # windows (for example npts=65537 made a floor of 64 request ~32 GiB). + + def eval_lnL(like, theta, chunk=_EVAL_CHUNK): """Evaluate the distance-marginalized lnL on an ``(N, 5)`` array in chunks. @@ -115,6 +298,7 @@ def eval_lnL(like, theta, chunk=_EVAL_CHUNK): dimension inside the likelihood). """ theta = np.atleast_2d(theta) + chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] out = np.empty(N) for i in range(0, N, chunk): @@ -655,10 +839,19 @@ def flowmc_sample(like, d_min, d_max, n_chains=20, n_local_steps=20, For a *batch* of nearby intrinsic templates (``--n-events-to-analyze``) the posterior geometry changes only slowly, so the re-used flow is a strong - initialization -- the partial-flow-reuse efficiency gain (most visible at - scale, and at high SNR where peaks are narrow). Re-use degrades gracefully: - a model-shape/version mismatch falls back to a fresh flow, mismatched - ``positions`` fall back to a high-lnL prior draw. + initialization. Re-use degrades gracefully: a model-shape/version mismatch + falls back to a fresh flow, mismatched ``positions`` fall back to a high-lnL + prior draw. + + **RE-USE IS OFF BY DEFAULT AND SHOULD STAY OFF FOR SAMPLE-PRODUCING RUNS.** + It CONTRACTS the extrinsic posterior in later slots -- psi to ~40% of its + no-re-use width by slot 7 of an 8-event batch, on both of two seeds, with + slot 0 as a control at ~1.0 -- and the efficiency argument that used to sit + here does not survive production settings: 1589 s mean wall with re-use + against 1567 s without, a difference smaller than the seed-to-seed spread and + of flipping sign. The ~2x speed-up reported in this module's README is + specific to the small-budget SNR-sequence benchmark. Enable it only where + the EVIDENCE, not the samples, is the product. Returns ------- @@ -825,6 +1018,7 @@ def _log_prior_4_jax(theta4): def eval_lnL_4(like, theta, chunk=_EVAL_CHUNK, desc="lnL"): """Evaluate the 4-param (phi-marginalised) lnL on an ``(N, 4)`` array.""" theta = np.atleast_2d(theta) + chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] out = np.empty(N) try: @@ -921,6 +1115,7 @@ def _log_prior_3_jax(theta3): def eval_lnL_3(like, theta, chunk=_EVAL_CHUNK, desc="lnL"): """Evaluate the 3-param (phi+psi-marginalised) lnL on an ``(N, 3)`` array.""" theta = np.atleast_2d(theta) + chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] out = np.empty(N) try: @@ -1489,8 +1684,12 @@ def smc_puffball_sample(like, d_min, d_max, n_walkers=2000, seed=0, collapse). This is the SMC analogue of RIFT-AV's "sample -> puffball -> sample" and of nested sampling's hill-climb; robust at LISA-loud SNR. - Returns the same dict shape as :func:`flowmc_sample_phimarg`. Evidence is the - standard SMC normalizing-constant estimator logZ = sum_t logmeanexp(dbeta_t lnL). + Returns the same dict shape as :func:`flowmc_sample_phimarg`, plus ``inv_T``: + the tempering exponent the ladder ACTUALLY reached (< 1 when it stopped at + ``max_stages``), with ``post_weight`` the matching ``L**(1-inv_T)`` correction + to the posterior. Evidence is the standard SMC normalizing-constant estimator + logZ = sum_t logmeanexp(dbeta_t lnL) -- which is log Z(inv_T), not log Z, on a + ladder that stopped short. """ _param_order = getattr(like, "ANGULAR_PARAM_ORDER", ("ra", "dec", "psi", "incl")) n_dim = len(_param_order) @@ -1535,7 +1734,11 @@ def _ess(db): return (s * s) / np.sum(w * w) if s > 0 else 0.0 target = float(ess_frac) * W - hi_db = min(1.0 - inv_T, float(max_dbeta)) + # Largest rung allowed here: never past inv_T == 1, and never past the + # per-stage cap (max_dbeta <= 0 disables that cap, as on the flowMC path). + hi_db = 1.0 - inv_T + if float(max_dbeta) > 0: + hi_db = min(hi_db, float(max_dbeta)) if _ess(hi_db) >= target: db = hi_db else: @@ -1546,7 +1749,13 @@ def _ess(db): a = mid else: b = mid - db = max(a, 1e-4) + # The floor keeps a stalled bisection moving, but it must never + # carry the rung PAST the rung cap: db > hi_db advances inv_T beyond + # 1 (or beyond max_dbeta), and the resample/Metropolis moves below + # then target L**inv_T with inv_T > 1 -- an OVER-tempered cloud that + # the final min(inv_T, 1) would report as temper=1 with uniform + # post_weight, i.e. exactly the mislabelling the tail guards against. + db = min(max(a, 1e-4), hi_db) # SMC evidence increment: logZ += logmeanexp(db * lnL) z = db * lnL z = z[np.isfinite(z)] @@ -1632,9 +1841,37 @@ def _ess(db): if verbose: print(" [smc-IS-Z] failed (%r); keeping SMC logZ" % e) + # THE LADDER CAN STOP SHORT OF THE POSTERIOR. The loop above also exits on + # ``stage == max_stages`` (and on a cloud with fewer than two finite lnL), and + # the cloud then still targets ``L**inv_T * prior`` with ``inv_T < 1``. + # Reporting ``temper=1.0`` with uniform ``post_weight`` in that case handed the + # caller a TEMPERED cloud labelled as a posterior draw. Report the exponent + # actually reached, plus the correction weight ``L**(1-inv_T)`` that carries + # the cloud to the posterior -- the same contract flowmc_sample_phimarg uses. + # The weight is identically uniform once inv_T == 1, so the converged path is + # unchanged. Each rung is capped at the distance left to 1, so the clip + # below only absorbs the rounding of the accumulated sum -- it must never be + # covering for a ladder that genuinely ran past 1 (see the ``db`` cap above). + lnL = np.asarray(lnL, dtype=float) + inv_T_final = float(min(inv_T, 1.0)) + if len(lnL) and inv_T_final < 1.0: + lw = (1.0 - inv_T_final) * lnL + lw = np.where(np.isfinite(lw), lw, -np.inf) + mx = np.max(lw) + # All -inf stays all-zero: a cloud whose correction cannot be normalised + # must be REFUSED by the caller, not quietly restored to uniform weights, + # which is the mislabelling this block exists to prevent. + post_weight = np.exp(lw - mx) if np.isfinite(mx) else np.zeros(len(lnL)) + s = post_weight.sum() + if s > 0: + post_weight = post_weight / s + else: + post_weight = np.ones(W) / W return dict(theta=cloud, lnL=lnL, logZ=float(logZ), sigma_over_Z=float(sigma_over_Z), neff=float(neff), - flow_state=None, post_weight=np.ones(W) / W, temper=1.0, + flow_state=None, post_weight=post_weight, inv_T=inv_T_final, + temper=(float(1.0 / inv_T_final) if inv_T_final > 0 + else float("inf")), logZ_laplace=float(logZ_smc), lnL_map=float(np.max(lnL)) if len(lnL) else np.nan) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py index bb3281fc8..43d02c9cd 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py @@ -22,6 +22,8 @@ import RIFT.likelihood.factored_likelihood as factored_likelihood +from . import core as _core + import os from .core import (build_likelihood_data, fused_log_likelihood, @@ -29,11 +31,49 @@ fused_log_likelihood_distphipsimarg, fused_log_likelihood_distpsimarg, make_distance_grid, make_distance_grid_adaptive, + make_distance_grid_loguniform, loguniform_grid_size, + DIST_GRID_TOL_DEFAULT, DIST_GRID_SCHEMES, estimate_distance_peak, phi_ref_grid, psi_grid, - phi_ref_conditional_lnL, DIST_MPC_REF) + phi_ref_conditional_lnL, DIST_MPC_REF, JAX_INTERP_DEFAULT, + TIME_QUAD_DEFAULT, _TIME_QUAD_CHOICES, default_time_guard) +# Generic probe direction for the build-time identity check. The A0==0/B1==0 +# identity is a property of the spin-2 detector response, so it does not depend +# on where we probe; a single generic (ra, dec, incl) away from any pole or +# face-on/edge-on special case is enough, and keeps the check O(1). +_ANGLE_MARG_PROBE_RA = [1.0] +_ANGLE_MARG_PROBE_DEC = [0.3] +_ANGLE_MARG_PROBE_INCL = [1.0] +from . import core as _core +from .anglemarg import (ANGLE_MARG_DEFAULT, ANGLE_MARG_LEGACY, # noqa: F401 + ANGLE_MARG_CHOICES) # Parameter order used throughout the wrapper's vectorized interface. EXTRINSIC_PARAM_ORDER = ("ra", "dec", "psi", "incl", "phiref", "distMpc") +_TIME_SUPPORT_DELAY_MARGIN = 0.05 + + +def bandlimited_storage_requirement(deltaT, integration_window_half): + """Return ``(storage_half, g0, g_certificate)`` for adaptive time support.""" + tvals = factored_likelihood.marginalization_time_grid( + integration_window_half, deltaT, xpy=np) + g_default = default_time_guard(len(tvals)) + g0 = 1 << int(np.ceil(np.log2(g_default))) + g_certificate = 2 * g0 + # Fifty milliseconds exceeds the Earth-diameter light time (~42.6 ms), so + # this support guarantee does not encode an HLV-only network assumption. + storage_half = (float(integration_window_half) + g_certificate * float(deltaT) + + _TIME_SUPPORT_DELAY_MARGIN + 16 * float(deltaT)) + return storage_half, g0, g_certificate + + +def _validate_nonlinear_time_quadrature(time_quadrature, endpoint): + if time_quadrature not in _TIME_QUAD_CHOICES: + raise ValueError("time_quadrature must be one of %r" % (_TIME_QUAD_CHOICES,)) + if time_quadrature == "bandlimited": + raise ValueError( + "time_quadrature='bandlimited' is not valid for %s: the primitive " + "time fields must be refined before its nonlinear marginalization; " + "use 'simpson'" % endpoint) def build_rotation_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, @@ -221,15 +261,24 @@ class JAXExtrinsicLikelihood: arrays of shape (S,). """ - def __init__(self, data, interp="linear", phase_marginalization=False): + def __init__(self, data, interp=JAX_INTERP_DEFAULT, phase_marginalization=False, + *, time_quadrature=TIME_QUAD_DEFAULT): self.data = data self.interp = interp self.phase_marginalization = phase_marginalization + if time_quadrature not in _TIME_QUAD_CHOICES: + raise ValueError("time_quadrature must be one of %r" % (_TIME_QUAD_CHOICES,)) + self.time_quadrature = time_quadrature + if time_quadrature == "bandlimited": + g_default = default_time_guard(data.npts) + self.time_guard_initial = 1 << int(np.ceil(np.log2(g_default))) + self.time_guard_certified = 2 * self.time_guard_initial def _batched(ra, dec, psi, incl, phiref, distMpc): return fused_log_likelihood( data, ra, dec, psi, incl, phiref, distMpc, - interp=interp, phase_marginalization=phase_marginalization) + interp=interp, phase_marginalization=phase_marginalization, + time_quadrature=time_quadrature) self._batched = jax.jit(_batched) @@ -239,7 +288,8 @@ def _scalar(theta6): data, theta6[0:1], theta6[1:2], theta6[2:3], theta6[3:4], theta6[4:5], theta6[5:6], - interp=interp, phase_marginalization=phase_marginalization) + interp=interp, phase_marginalization=phase_marginalization, + time_quadrature=time_quadrature) return v[0] self._scalar = _scalar @@ -281,16 +331,24 @@ class JAXDistanceMarginalizedLikelihood: ANGULAR_PARAM_ORDER = ("ra", "dec", "psi", "incl", "phiref") def __init__(self, data, d_min, d_max, n_grid=256, d_prior="euclidean", - interp="linear", phase_marginalization=False): + interp=JAX_INTERP_DEFAULT, phase_marginalization=False, + *, time_quadrature=TIME_QUAD_DEFAULT, d_prior_range=None): self.data = data + self.interp = interp # the instance's stencil; sample_phi_ref defaults to it + self.phase_marginalization = phase_marginalization + _validate_nonlinear_time_quadrature( + time_quadrature, "distance marginalization") + self.time_quadrature = time_quadrature self.x_grid, self.log_w_grid = make_distance_grid( - d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) def _batched(ra, dec, psi, incl, phiref): return fused_log_likelihood_distmarg( data, ra, dec, psi, incl, phiref, self.x_grid, self.log_w_grid, - interp=interp, phase_marginalization=phase_marginalization) + interp=interp, phase_marginalization=phase_marginalization, + time_quadrature=time_quadrature) self._batched = jax.jit(_batched) @@ -298,7 +356,8 @@ def _scalar(theta5): v = fused_log_likelihood_distmarg( data, theta5[0:1], theta5[1:2], theta5[2:3], theta5[3:4], theta5[4:5], self.x_grid, self.log_w_grid, - interp=interp, phase_marginalization=phase_marginalization) + interp=interp, phase_marginalization=phase_marginalization, + time_quadrature=time_quadrature) return v[0] self._scalar = _scalar @@ -343,8 +402,13 @@ class JAXDistPhiMargLikelihood: ANGULAR_PARAM_ORDER = ("ra", "dec", "psi", "incl") def __init__(self, data, d_min, d_max, nphi=32, n_grid=256, - d_prior="euclidean", interp="linear", guess_snr=None): + d_prior="euclidean", interp=JAX_INTERP_DEFAULT, guess_snr=None, + *, time_quadrature=TIME_QUAD_DEFAULT, d_prior_range=None): self.data = data + self.interp = interp # the instance's stencil; sample_phi_ref defaults to it + _validate_nonlinear_time_quadrature( + time_quadrature, "distance/phase marginalization") + self.time_quadrature = time_quadrature self.nphi = int(nphi) self._phi_grid = phi_ref_grid(self.nphi) # Adaptive distance quadrature: concentrate grid resolution on the @@ -354,29 +418,36 @@ def __init__(self, data, d_min, d_max, nphi=32, n_grid=256, # same stable logsumexp kernel and is gradient-stable. Enable with env # JAX_ILE_DISTGRID_ADAPTIVE=1; falls back to uniform otherwise. if int(os.environ.get("JAX_ILE_DISTGRID_ADAPTIVE", "0")) and guess_snr: - d_peak, sigma_d = estimate_distance_peak(data, guess_snr) + # interp= must be forwarded: this sizes the distance grid the likelihood then + # integrates on, so leaving it at the module default silently mixes stencils -- + # and would break the documented 'pass interp="linear" to reproduce a + # pre-2026-08-26 run' recipe, which is the whole mitigation for that default move. + d_peak, sigma_d = estimate_distance_peak(data, guess_snr, interp=interp) self.x_grid, self.log_w_grid = make_distance_grid_adaptive( - d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="adaptive", d_peak=float(d_peak), sigma_d=float(sigma_d), n=int(self.x_grid.shape[0])) else: self.x_grid, self.log_w_grid = make_distance_grid( - d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="uniform", n=int(self.x_grid.shape[0])) xg, lwg, pg = self.x_grid, self.log_w_grid, self._phi_grid def _batched(ra, dec, psi, incl): return fused_log_likelihood_distphimarg( - data, ra, dec, psi, incl, xg, lwg, pg, interp=interp) + data, ra, dec, psi, incl, xg, lwg, pg, interp=interp, + time_quadrature=time_quadrature) self._batched = jax.jit(_batched) def _scalar(theta4): v = fused_log_likelihood_distphimarg( data, theta4[0:1], theta4[1:2], theta4[2:3], theta4[3:4], - xg, lwg, pg, interp=interp) + xg, lwg, pg, interp=interp, time_quadrature=time_quadrature) return v[0] self._scalar = _scalar @@ -402,7 +473,7 @@ def fisher(self, theta4): return -H def sample_phi_ref(self, ra, dec, psi, incl, distMpc, rng=None, - n_samples=1, interp="linear"): + n_samples=1, interp=None): """Draw φ_ref from its conditional posterior given the other params. Evaluates ``phi_ref_conditional_lnL`` on the grid, normalises, draws @@ -414,11 +485,23 @@ def sample_phi_ref(self, ra, dec, psi, incl, distMpc, rng=None, ra, dec, psi, incl, distMpc : float scalars or (S,) arrays rng : numpy.random.Generator (optional) n_samples : int — draws per input sample + interp : str or None — stencil to evaluate the conditional with. None (the default) + means **this instance's** stencil, not the module default. + + Notes + ----- + This argument carried its own module-level default until 2026-08-26, which was harmless + only while that string happened to equal the constructor's: an instance built with any + other stencil drew its phases from a DIFFERENT likelihood than the one it reports lnL and + evidence from, and nothing raised. Moving the module default to 'sinc' made it bite the + documented backward-compatibility recipe -- constructing with interp="linear" gave a + linear evidence and sinc phase draws. Pass interp= only to override deliberately. Returns ------- phi_ref : (S,) float array (or (S, n_samples) when n_samples > 1) """ + interp = self.interp if interp is None else interp rng = rng or np.random.default_rng() ra_ = np.atleast_1d(np.asarray(ra, float)) dec_ = np.atleast_1d(np.asarray(dec, float)) @@ -430,7 +513,8 @@ def sample_phi_ref(self, ra, dec, psi, incl, distMpc, rng=None, self.data, jnp.asarray(ra_), jnp.asarray(dec_), jnp.asarray(psi_), jnp.asarray(incl_), - jnp.asarray(dist_), self._phi_grid, interp=interp)) # (nphi, S) + jnp.asarray(dist_), self._phi_grid, interp=interp, + time_quadrature=self.time_quadrature)) # (nphi, S) phi_vals = np.asarray(self._phi_grid) dphi = float(phi_vals[1] - phi_vals[0]) @@ -460,36 +544,443 @@ class JAXDistPhiPsiMargLikelihood: ANGULAR_PARAM_ORDER = ("ra", "dec", "incl") def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, - d_prior="euclidean", interp="linear", guess_snr=None): + d_prior="euclidean", interp=JAX_INTERP_DEFAULT, guess_snr=None, + angle_marg=ANGLE_MARG_DEFAULT, *, + time_quadrature=TIME_QUAD_DEFAULT, d_prior_range=None, + dist_grid="uniform", dist_grid_tol=DIST_GRID_TOL_DEFAULT): self.data = data + self.interp = interp # the instance's stencil; sample_phi_ref defaults to it + _validate_nonlinear_time_quadrature( + time_quadrature, "distance/phase/polarization marginalization") + self.time_quadrature = time_quadrature self.nphi = int(nphi) self.npsi = int(npsi) self._phi_grid = phi_ref_grid(self.nphi) self._psi_grid = psi_grid(self.npsi) + # (phi_ref, psi) marginalization scheme. "grid" is the historical + # nphi x npsi quadrature, kept as the DEFAULT so existing command + # lines reproduce existing runs; "exact" / "laplace" are the + # exact-coefficient schemes of RIFT.likelihood.jax_ile.anglemarg + # (which fix the grid path's SNR-unbounded quadrature error and its + # nphi=8 Nyquist aliasing); "auto" selects between them. Both the + # selection and the dense-grid sizing key on a DATA-DERIVED amplitude + # bound (estimate_angle_amplitude, computed below on the FULL prior + # distance support) -- never on guess_snr: an absent or underestimated SNR + # must not be able to silently under-resolve the quadrature + # (external-review defect 2). self.angle_marg_info records what + # actually ran -- callers must surface it in the run log. + if angle_marg not in ANGLE_MARG_CHOICES: + raise ValueError("angle_marg must be one of grid/exact/laplace/" + "peak-local/auto, got %r" % (angle_marg,)) + if dist_grid not in DIST_GRID_SCHEMES: + # An unrecognised value must NEVER fall through to the default: a + # typo that silently returns the old answer is precisely the + # silently-inert-flag failure this module keeps being bitten by. + raise ValueError("dist_grid must be one of %r, got %r" + % (DIST_GRID_SCHEMES, dist_grid)) + if dist_grid != "uniform" and _core._DISTMARG_GH_N > 0: + # core._distmarg_gh_logL places its own per-sample nodes and reads + # ONLY min(x_grid)/max(x_grid); the node positions and the whole + # log_w_grid are unused. Both schemes span the same support, so the + # arms would be bit-identical while dist_grid_info still reported + # mode='loguniform'. That is the silently-inert-flag class the + # other refusals here exist to prevent, and it is reachable + # without the user naming a dense scheme: under GH + # choose_angle_marg_scheme resolves to one regardless. Which one + # is irrelevant to this refusal -- the per-sample quadrature reads + # only the support on every dense path -- so do not re-tie this + # comment to a particular selector outcome. + raise ValueError( + "dist_grid=%r cannot be combined with JAX_ILE_DISTMARG_GH=%d: " + "the per-sample Gauss-Hermite distance quadrature places its " + "own nodes and uses only the SUPPORT of x_grid, so this option " + "would be bit-identically inert while still being reported as " + "active. Unset JAX_ILE_DISTMARG_GH, or use " + "dist_grid='uniform'." % (dist_grid, _core._DISTMARG_GH_N)) + if dist_grid != "uniform" and d_prior_range is not None and ( + float(d_prior_range[0]) != float(d_min) + or float(d_prior_range[1]) != float(d_max)): + raise ValueError( + "dist_grid=%r cannot be combined with a narrowed distance range " + "(--limit-distance): make_distance_grid_loguniform takes no " + "d_prior_range, so it would normalize the distance prior onto the " + "BOX -- the exact renormalization --limit-distance exists to " + "prevent, and worth several nats of evidence silently. The two " + "features address the same problem from opposite sides (narrow the " + "range vs resolve the peak anywhere in it); composing them needs a " + "prior-range split in the log-uniform builder and its own " + "validation. Use one or the other." % (dist_grid,)) + from . import anglemarg as _anglemarg + # THE FULL-SUPPORT distance grid. Two distinct roles are deliberately + # separated here (see DESIGN_jax_distance_quadrature.md, "decoupling"): + # + # x_grid_full sizes the ANGLE lattice. It always spans the whole + # prior range [d_min, d_max], whatever grid the + # likelihood ends up integrating on. + # self.x_grid is what the fused kernel integrates over. + # + # estimate_angle_amplitude reads only min/max of the grid it is given + # (the per-angle distance maximum is closed form at + # clip(A/B, x_min, x_max)), so a narrowed distance grid that still + # contains A/B leaves the amplitude untouched -- but one that does NOT + # contain it silently SHRINKS the angle lattice (measured: a + # [0.8 d, 1.25 d] window drops the amplitude 12.6% and the lattice from + # (624, 320) to (592, 304)). Sizing from the full support costs one + # build-time scalar and removes that coupling by construction instead + # of bounding it. For dist_grid="uniform" this IS self.x_grid, so the + # default path is unchanged, node for node. + # --limit-distance is EXACTLY the "narrowed distance grid" the note above + # measures (12.6% amplitude, (624,320) -> (592,304) lattice). d_min/d_max + # are then the SAMPLED box, so the full-support grid is built over + # d_prior_range instead -- which is what "the whole prior range" means once + # the two roles are split. With no box the two are equal and this is the + # same call, node for node. + _full_lo, _full_hi = ((d_min, d_max) if d_prior_range is None + else (d_prior_range[0], d_prior_range[1])) + x_grid_full, log_w_full = make_distance_grid( + _full_lo, _full_hi, n_grid, d_prior, distMpcRef=data.distMpcRef) if int(os.environ.get("JAX_ILE_DISTGRID_ADAPTIVE", "0")) and guess_snr: - d_peak, sigma_d = estimate_distance_peak(data, guess_snr) + if dist_grid != "uniform": + raise ValueError( + "JAX_ILE_DISTGRID_ADAPTIVE=1 and dist_grid=%r both ask to " + "replace the distance grid. Unset the environment " + "variable (it is deprecated; see " + "DESIGN_jax_distance_quadrature.md)." % (dist_grid,)) + # DEPRECATED. Kept reachable so nothing that sets this variable + # today changes behaviour, but it is measurably unsafe: its window + # is centred on estimate_distance_peak, a 300-step gradient ascent + # that is NOT converged (measured rho 39.97 against the amplitude + # bound's 55.07), and its trapezoid gives the last node a full + # rather than half interval, misplacing ~3% of the volumetric prior + # mass onto d_max. Measured 9.4 nats of lnL error at SNR 40. + print("WARNING: JAX_ILE_DISTGRID_ADAPTIVE is DEPRECATED and " + "measurably unsafe (9.4 nats at SNR 40 on the reference " + "configuration). Use dist_grid='loguniform' " + "(--distance-grid-scheme loguniform); see " + "DESIGN_jax_distance_quadrature.md.") + # interp= must be forwarded: this sizes the distance grid the likelihood then + # integrates on, so leaving it at the module default silently mixes stencils -- + # and would break the documented 'pass interp="linear" to reproduce a + # pre-2026-08-26 run' recipe, which is the whole mitigation for that default move. + d_peak, sigma_d = estimate_distance_peak(data, guess_snr, interp=interp) self.x_grid, self.log_w_grid = make_distance_grid_adaptive( - d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="adaptive", d_peak=float(d_peak), sigma_d=float(sigma_d), n=int(self.x_grid.shape[0])) else: - self.x_grid, self.log_w_grid = make_distance_grid( - d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef) + if d_prior_range is None: + # Default path, unchanged node for node: the integration grid IS + # the full-support grid. + self.x_grid, self.log_w_grid = x_grid_full, log_w_full + else: + # --limit-distance: integrate on the box, normalize on the prior. + self.x_grid, self.log_w_grid = make_distance_grid( + d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="uniform", n=int(self.x_grid.shape[0])) + if angle_marg == "grid": + if dist_grid != "uniform": + # Fail closed. The log-uniform grid is sized from the + # data-derived angle amplitude, which the grid scheme neither + # computes nor rechecks at runtime; applying it there would be + # an unvalidated path, and silently ignoring the request would + # be a silent no-op. + raise ValueError( + "dist_grid=%r requires angle_marg in " + "('exact', 'laplace', 'auto'): the log-uniform grid is " + "sized from the data-derived angle amplitude, which the " + "'grid' scheme does not compute." % (dist_grid,)) + scheme, sel_info = "grid", dict(reason="default grid quadrature") + amp_sizing = None + else: + # Eager, build-time (grid sizes must be static under jit): bound + # the exponent amplitude from the coefficient tables themselves, + # over a sky sample and the FULL prior distance support. + amp_data, amp_diag = _anglemarg.estimate_angle_amplitude( + data, x_grid_full, interp=interp, return_diagnostics=True) + # sizing is FLOORED at the crossover (never below the calibration + # point); the SELECTION below uses the UNfloored bound, so quiet + # targets stay on the exact branch. Computed here, before the + # distance grid, because the distance grid is sized from this same + # floored number -- see the rho_max note directly below. + amp_sizing = max(amp_data, + _anglemarg.ANGLE_MARG_CROSSOVER_AMPLITUDE) + if dist_grid == "loguniform": + # REFUSE the truncated regime. The spacing contract assumes the + # integrand is a Gaussian PEAK inside the support, whose relative + # width 1/rho is what c(tol)/rho_max resolves. When the + # maximizing distance x* = A/B lies OUTSIDE [x_min, x_max] the + # integrand is monotone on the support instead -- a boundary + # layer at one prior edge -- and a log-uniform grid is the wrong + # instrument for it twice over: its ABSOLUTE spacing is coarsest + # exactly at d_max where the layer sits, and refining it adds + # nodes proportionally everywhere so the layer never resolves + # (measured: tol 0.5 -> 1e-9 moves the error only 5.23 -> 3.92 + # nats, while uniform 256 -> 4096 moves 2.52 -> 0.36). Worse, + # the clip makes the amplitude UNDER-read, so the derived node + # count moves the wrong way -- in the extreme it reads 0, the + # crossover floor pins rho_max = 30, and the grid collapses to + # 145 nodes. We refuse rather than fall back to uniform: a + # fallback would make this flag silently produce the other + # scheme's grid, and this regime is a physics signal (the + # posterior rails against a prior edge) that the caller should + # see rather than have papered over. Neither grid is good here + # -- uniform 256 is itself 2.5 nats out. + if amp_diag["clip_excess"] > 1.0 + 1e-3: + raise ValueError( + "dist_grid='loguniform' refuses this event: the " + "likelihood's maximizing distance lies OUTSIDE " + "[d_min, d_max] = [%g, %g] Mpc, so the distance " + "integrand is a boundary layer at a prior edge rather " + "than an interior peak, and the log-uniform spacing " + "contract does not apply (measured 1.9-4.6 nats of " + "error there, worse than the uniform default). " + "Diagnostic: unclipped amplitude %.6g against clipped " + "%.6g (excess %.4g). Recourse: widen --d-max (or " + "narrow --d-min) so the posterior is interior, or stay " + "on --distance-grid-scheme uniform and raise " + "--distance-grid-points. See " + "DESIGN_jax_distance_quadrature.md section 1a." + % (float(d_min), float(d_max), + amp_diag["amp_unclipped"], amp_diag["amp_clipped"], + amp_diag["clip_excess"])) + # rho_max = sqrt(2 A): A is the max over angles of the + # closed-form distance maximum A_ang^2/(2 B_ang) = rho^2/2. NOT + # an identity and NOT a proven bound -- A carries + # ANGLE_AMP_MARGIN and the max is over a SAMPLED sky, so this is + # sqrt(margin) * rho_sampled_max. No NEW estimator is + # introduced and no peak is located, which is the point. + # + # A is amp_SIZING, not amp_data, and that choice is what makes + # the runtime fail-safe cover this grid. _runtime_amp_failsafe + # compares the per-call amplitude against amp_sizing; sizing the + # distance spacing from the unfloored amp_data instead would + # leave a silent gap for quiet targets (amp_data < crossover), + # where a runtime amplitude between amp_data and amp_sizing + # under-resolves the distance peak WITHOUT tripping anything. + # Flooring costs a minimum of ~144 nodes on a quiet event, whose + # run is cheap anyway. + # + # ...and the sizing amplitude is the one the guard ADMITS, not + # the one it was built from. _runtime_amp_failsafe stays silent + # until amp_call > AMP_FAILSAFE_TRIP_FACTOR * amp_sizing, so a + # call just under that threshold carries an interior peak at + # rho = sqrt(TRIP) * sqrt(2*amp_sizing) -- sqrt(2) above the + # spacing's design point at the shipped factor -- and is neither + # printed nor recorded. Through the Gaussian alias law + # (2 exp(-2 pi^2/c^2), c fixed by tol) that turns the advertised + # tol into ~sqrt(2*tol): 0.01 becomes 0.14, unlabelled. Sizing + # from TRIP*amp_sizing closes the window by construction and + # costs sqrt(TRIP) = 1.41x the nodes; the alternative -- a + # second, distance-specific runtime guard at a tighter threshold + # -- adds a mechanism where a constant will do, and the two + # would then have to be kept consistent by hand. + rho_max = float(np.sqrt( + 2.0 * _anglemarg.AMP_FAILSAFE_TRIP_FACTOR + * max(float(amp_sizing), 0.0))) + # SECOND precondition (the first is clip_excess, above): a peak + # that is interior but sits ~1 width inside an edge breaks the + # spacing law while clip_excess reads exactly 1. The alias law + # is Poisson summation on an UNTRUNCATED Gaussian; a truncated + # support adds an Euler-Maclaurin endpoint term proportional to + # the integrand's derivative there, worth ~11% at the shipped + # tol (c = 1.93) against a promised 1%. Refuse rather than + # widen silently: the option's whole claim is the stated + # fractional error, and both alternatives -- shipping the error + # or moving the user's distance prior for them -- are worse than + # saying so. Evaluated at the CONTRACT spacing c/rho_max, which + # is the coarsest the built grid can be (the node count ceils). + dlnd_contract = (_core.loguniform_spacing_for_tolerance( + dist_grid_tol) / rho_max) + eps_end = float(_core.ENDPOINT_ERROR_MARGIN + * _core.loguniform_endpoint_error( + dlnd_contract, + amp_diag["endpoint_scale"])) + # ...and a NON-POSITIVE clearance is not a small correction. + # _endpoint_bell clamps k <= 0 to zero on the argument that an + # endpoint sitting ON the peak is a stationary point with + # nothing to correct -- true, and MEASURED true at exactly + # k = 0 (1.2e-9). But the clamp cannot tell k = 0 from k < 0, + # where the peak has left the support entirely and the error + # climbs steeply, so the term scores zero exactly where it is + # worst. Between that and the clip_excess trip there is a + # window neither guard sees: measured on the reference synthetic + # at rho_pk 51.4, clearance -0.66 builds with a TRUE error of + # 0.0668 -- 6.7x the promised tol -- while clip_excess is still + # ~1. Refuse it here. This also refuses the benign k = 0 + # point, which costs nothing: the error is 2e-2 by half a width + # either side of it. (External review of the endpoint guard.) + # A loud EXTERIOR entry that no other diagnostic can see is a + # REAL structural hole and is deliberately NOT refused here. + # External review (P2) is right that clip_excess is a ratio of + # GLOBAL maxima, that peak_clearance describes the global + # argmax, and that the endpoint term cannot represent an + # exterior entry (its bell is clamped at k <= 0, and + # Euler-Maclaurin is the wrong expansion for a peak outside the + # support). So an interior dominant entry with a near-equal + # exterior secondary reads clean on all three. + # + # It is not refused because no threshold on the available + # quantity survives contact with the fixtures. The obvious + # one -- refuse when the exterior entry's weight exp(-gap) + # exceeds tol, i.e. gap < ln(1/tol) = 4.6 nats -- refuses every + # QUIET event: measured, _synth() has amp_clipped 3.30 nats in + # TOTAL and an exterior gap of 3.29, and a quieter one 0.0083 + # and 0.0082. Their whole exponent range is smaller than the + # threshold, and nothing is wrong with them: at low amplitude + # the distance integrand is smooth and the grid over-resolves + # it, so the exterior entry carries weight but not error. The + # honest condition is weight TIMES that entry's own resolution + # error, which needs a model of the latter that cannot be + # validated against any configuration reachable here. + # + # What IS established: on every loud fixture and every prior + # this code is run with, the loudest exterior entry sits 32-5273 + # nats below the maximum, against a 15-nat contribution band -- + # so the hole is not reachable there. That premise is pinned by + # test_exterior_entries_stay_far_below_the_dominant_one, which + # fails if it ever stops holding, and exterior_gap is reported + # so the condition is visible rather than silent. + if amp_diag["peak_clearance"] <= 0.0: + raise ValueError( + "dist_grid='loguniform' refuses this event: the " + "dominant peak sits AT or OUTSIDE a prior edge of " + "[d_min, d_max] = [%g, %g] Mpc (clearance %.4g peak " + "widths, rho = %.4g), so the spacing contract's " + "precondition -- a Gaussian peak INSIDE the support -- " + "does not hold. This is the same physics signal as " + "the exterior refusal, caught earlier: the truncated " + "endpoint term is what dominates here, and it is not " + "bounded by the alias law the node count is derived " + "from. Recourse: widen the distance prior so the peak " + "is interior with clearance >= %.2f widths, or stay on " + "--distance-grid-scheme uniform. See " + "DESIGN_jax_distance_quadrature.md section 1a." + % (float(d_min), float(d_max), + amp_diag["peak_clearance"], amp_diag["peak_rho"], + _core.loguniform_min_clearance(dist_grid_tol))) + if eps_end > float(dist_grid_tol): + raise ValueError( + "dist_grid='loguniform' refuses this event: the " + "likelihood's maximizing distance is INTERIOR to " + "[d_min, d_max] = [%g, %g] Mpc but too close to an " + "edge for the spacing contract, which assumes an " + "effectively untruncated Gaussian peak. The dominant " + "peak sits %.3g peak widths (1/rho, rho = %.4g) from " + "the nearer edge; at this grid's spacing dlnd = %.4g " + "the truncated-endpoint term alone is ~%.3g of the " + "distance integral, against the requested tol=%g. A " + "peak at the sizing SNR needs >= %.2f widths. " + "Recourse: widen the distance prior so the peak has " + "clearance (this is the same physics signal as the " + "exterior refusal -- the posterior is close to a prior " + "edge), or stay on --distance-grid-scheme uniform and " + "raise --distance-grid-points. Note that TIGHTENING " + "--distance-grid-tol does not help: the endpoint term " + "falls as c(tol)^2 while the budget falls faster. See " + "DESIGN_jax_distance_quadrature.md section 1a." + % (float(d_min), float(d_max), + amp_diag["peak_clearance"], amp_diag["peak_rho"], + dlnd_contract, eps_end, float(dist_grid_tol), + _core.loguniform_min_clearance(dist_grid_tol))) + self.x_grid, self.log_w_grid = make_distance_grid_loguniform( + d_min, d_max, rho_max, d_prior, + distMpcRef=data.distMpcRef, tol=dist_grid_tol) + self.dist_grid_info = dict( + mode="loguniform", n=int(self.x_grid.shape[0]), + tol=float(dist_grid_tol), rho_max=rho_max, + dlnd=float(np.log(float(d_max) / float(d_min)) + / (int(self.x_grid.shape[0]) - 1)), + n_uniform_requested=int(n_grid)) + # The A0==0/B1==0 identity that the GH psi-marginal node placement + # is DERIVED from is measured ONCE here, on concrete tables, and + # gates EVERY route to that placement -- not just 'auto'. An + # earlier revision checked it only in the 'auto' branch, so an + # explicit --angle-marg-scheme laplace walked past it and an + # m_max == 2 dataset whose identity fails was evaluated with a + # placement whose premise was absent (external review). It cannot + # be checked inside the kernel: that runs under jit/grad, where the + # coefficient tables are tracers. + gh_ok, gh_info = None, {} + if _core._DISTMARG_GH_N > 0 and angle_marg in ("auto", "laplace"): + gh_ok, gh_info = _anglemarg.gh_laplace_supported( + *_anglemarg.angle_coefficient_tables( + data, + jnp.asarray(_ANGLE_MARG_PROBE_RA), + jnp.asarray(_ANGLE_MARG_PROBE_DEC), + jnp.asarray(_ANGLE_MARG_PROBE_INCL), + interp)[:2], + _anglemarg._data_m_max(data), + feature=getattr(data, "feature", None)) + if angle_marg == "auto": + scheme, sel_info = _anglemarg.choose_angle_marg_scheme( + amp_data, gh_laplace_ok=gh_ok) + sel_info.update(gh_info) + else: + if angle_marg == "laplace" and gh_ok is False: + raise ValueError( + "--angle-marg-scheme laplace was requested with " + "JAX_ILE_DISTMARG_GH set, but its psi-marginal " + "distance-node placement is not valid for this data: " + "%s. The placement is DERIVED from A0 == 0 and " + "B1 == 0 (that is what reduces stationarity to " + "z^2 w = conj(w)), so it must not be used where they " + "do not hold. Use --angle-marg-scheme exact, or unset " + "JAX_ILE_DISTMARG_GH." + % gh_info.get("gh_laplace_reason", "identity absent")) + scheme, sel_info = angle_marg, dict( + reason="forced by caller", amplitude=amp_data, + crossover=_anglemarg.ANGLE_MARG_CROSSOVER_AMPLITUDE) + sel_info.update(gh_info) + self.angle_marg_scheme = scheme + self.angle_marg_info = dict(sel_info, requested=angle_marg, + scheme=scheme) + # Bound AFTER the distance grid is final: dist_grid="loguniform" + # replaces it inside the block above. xg, lwg, pg, sg = (self.x_grid, self.log_w_grid, self._phi_grid, self._psi_grid) + if scheme in ("exact", "laplace", "peak-local"): + self.angle_marg_info["amp_sizing"] = amp_sizing + self.angle_marg_info["sample_grid"] = tuple( + _anglemarg.angle_sample_grid_sizes( + _anglemarg._data_m_max(data))) + + if scheme == "grid": + def _fused(data_, ra, dec, incl, return_lnLt=False): + return fused_log_likelihood_distphipsimarg( + data_, ra, dec, incl, xg, lwg, pg, sg, interp=interp, + time_quadrature=time_quadrature, return_lnLt=return_lnLt) + elif scheme == "exact": + def _fused(data_, ra, dec, incl, return_lnLt=False): + return _anglemarg.fused_log_likelihood_distphipsimarg_exact( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, time_quadrature=time_quadrature, + return_lnLt=return_lnLt) + elif scheme == "peak-local": + # psi localized on the exact cell partition, phi still dense. Reachable + # only when asked for by name -- see the note on ANGLE_MARG_CHOICES for why + # it is not in 'auto' until a head-to-head pilot has run. + def _fused(data_, ra, dec, incl, return_lnLt=False): + return _anglemarg.fused_log_likelihood_distphipsimarg_peaklocal( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, time_quadrature=time_quadrature, + return_lnLt=return_lnLt) + else: # laplace + def _fused(data_, ra, dec, incl, return_lnLt=False): + return _anglemarg.fused_log_likelihood_distphipsimarg_laplace( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, time_quadrature=time_quadrature, + return_lnLt=return_lnLt) def _batched(ra, dec, incl): - return fused_log_likelihood_distphipsimarg( - data, ra, dec, incl, xg, lwg, pg, sg, interp=interp) + return _fused(data, ra, dec, incl) self._batched = jax.jit(_batched) def _scalar(theta3): - v = fused_log_likelihood_distphipsimarg( - data, theta3[0:1], theta3[1:2], theta3[2:3], - xg, lwg, pg, sg, interp=interp) + v = _fused(data, theta3[0:1], theta3[1:2], theta3[2:3]) return v[0] self._scalar = _scalar self._value_and_grad = jax.jit(jax.value_and_grad(_scalar)) @@ -529,33 +1020,45 @@ class JAXDistPsiMargLikelihood: ANGULAR_PARAM_ORDER = ("ra", "dec", "phiref", "incl") def __init__(self, data, d_min, d_max, npsi=8, n_grid=256, - d_prior="euclidean", interp="linear", guess_snr=None): + d_prior="euclidean", interp=JAX_INTERP_DEFAULT, guess_snr=None, + *, time_quadrature=TIME_QUAD_DEFAULT, d_prior_range=None): self.data = data + self.interp = interp # the instance's stencil; sample_phi_ref defaults to it + _validate_nonlinear_time_quadrature( + time_quadrature, "distance/polarization marginalization") + self.time_quadrature = time_quadrature self.npsi = int(npsi) self._psi_grid = psi_grid(self.npsi) if int(os.environ.get("JAX_ILE_DISTGRID_ADAPTIVE", "0")) and guess_snr: - d_peak, sigma_d = estimate_distance_peak(data, guess_snr) + # interp= must be forwarded: this sizes the distance grid the likelihood then + # integrates on, so leaving it at the module default silently mixes stencils -- + # and would break the documented 'pass interp="linear" to reproduce a + # pre-2026-08-26 run' recipe, which is the whole mitigation for that default move. + d_peak, sigma_d = estimate_distance_peak(data, guess_snr, interp=interp) self.x_grid, self.log_w_grid = make_distance_grid_adaptive( - d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, d_peak, sigma_d, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="adaptive", d_peak=float(d_peak), sigma_d=float(sigma_d), n=int(self.x_grid.shape[0])) else: self.x_grid, self.log_w_grid = make_distance_grid( - d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef) + d_min, d_max, n_grid, d_prior, distMpcRef=data.distMpcRef, + d_prior_range=d_prior_range) self.dist_grid_info = dict(mode="uniform", n=int(self.x_grid.shape[0])) xg, lwg, sg = self.x_grid, self.log_w_grid, self._psi_grid def _batched(ra, dec, phiref, incl): return fused_log_likelihood_distpsimarg( - data, ra, dec, phiref, incl, xg, lwg, sg, interp=interp) + data, ra, dec, phiref, incl, xg, lwg, sg, interp=interp, + time_quadrature=time_quadrature) self._batched = jax.jit(_batched) def _scalar(theta4): v = fused_log_likelihood_distpsimarg( data, theta4[0:1], theta4[1:2], theta4[2:3], theta4[3:4], - xg, lwg, sg, interp=interp) + xg, lwg, sg, interp=interp, time_quadrature=time_quadrature) return v[0] self._scalar = _scalar self._value_and_grad = jax.jit(jax.value_and_grad(_scalar)) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py new file mode 100644 index 000000000..fb20f348e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py @@ -0,0 +1,867 @@ +"""Joint (phi, psi) peak-local marginalization: numpy reference kernel. + +WHAT THIS COMPUTES + + L = (1/(2 pi)) int dphi (1/pi) int dpsi exp(g(phi, psi)) + +with the same normalization as `jax_ile.anglemarg` (uniform priors dphi/2pi, +dpsi/pi). In the chart u = 2 psi this is a plain double average over the torus, +``(2 pi)^-2 int int dphi du exp(g)``, which is the form everything below uses. + +WHY NOT A DENSE GRID. ``exp(g)`` is a needle whose width falls as ``A^-1/2``, so a +product grid costs ``~A``: with the shipped ``_dense_grid_sizes`` constants that is +4.3e5 points at amplitude 3250 and 4.2e7 at 3.25e5, per row. The exponent ``g``, +however, is an exact 2-D trig polynomial of bidegree ``(2 m_max, 2)`` -- its +structure does not change with amplitude at all. Enumerating its modes and +integrating only near them converts an amplitude-scaling cost into a +physics-scaling one. + +WHY NOT A SINGLE-CENTRE LAPLACE. Measured on the shipped coefficient tables, the +(phi, psi) surface carries 8-16 maxima whose values agree to machine precision, and +a one-centre Laplace is 1.4 nats low with an error that does NOT decay with +amplitude, because the deficit is combinatorial rather than curvature. Localisation +here must be multi-mode; that is the whole point. + +HOW THE MODES ARE FOUND, and why this is not a 2-D root solve. The u-degree of the +exponent is pinned at 2 for ANY mode set (spin-2), so at fixed ``phi`` the +u-stationary points are the unit-circle roots of a degree-4 polynomial -- the same +object ``anglemarg._laplace_psi_lnI`` already solves. The curve ``{d_u g = 0}`` is +therefore available EXACTLY, with no grid in u, and the 2-D stationary points lie on +it. What remains is a one-dimensional search in ``phi`` along that curve. No +resultant, no hidden-variable pencil, no BKK machinery -- and no exposure to the +conditioning of a 2-D solve at the machine-degenerate configurations that are the +normal operating point here. + +NO ON-CIRCLE TOLERANCE, deliberately. The obvious filter ``| |z| - 1 | < tol`` is an +estimate promoted to a bound: at exact multiplicity ``m`` the computed roots smear off +the unit circle by ``eps_machine^(1/m)`` (measured 4.6e-6 for a triple root), so a +1e-6 filter returns ONE mode where there are four, in precisely the degenerate regime +that is production. Every root is therefore kept and used only as a SEED; the region +machinery below is what decides what is real. Over-covering is free because regions +merge; under-covering is the only failure that matters. + +WHAT IS CERTIFIED, AND WHAT IS NOT. Read this before quoting the accuracy. + + * The u axis is certified at enumeration time (all roots of an exact quartic). + * The phi axis is GRID-SEEDED and is therefore NOT certified at enumeration time. + It carries exactly the caveat the time module carries: a grid is a resolution, + not a certificate. + * Correctness is restored the way the time module restores it -- by a bound on the + part of the domain the regions do not cover. ``outside_bound`` below is a TRUE + upper bound on ``g`` outside the covered set: a grid maximum plus the Lipschitz + remainder ``M_1 * h / 2``, with ``M_1 = sum |C_kq| |k| (or |q|)`` by the triangle + inequality over the exact coefficient table. Nothing there is fitted. + + A row whose omitted-mass bound is not small enough is NOT returned with a caveat: + it is declined, and the caller falls back to the dense rule. +""" + +import numpy as np + +__all__ = [ + "W_SIGMA", + "MERGE_MAX_PASSES", + "OUTSIDE_TOL_NATS", + "joint_table", + "eval_g", + "u_stationary_at_phi", + "enumerate_modes", + "derivative_bound", + "outside_bound", + "joint_marginalize_peak_local", + "joint_marginalize_over_distance", + "u_profile", + "phi_local_marginalize", +] + +#: Local integration half-width, in units of the mode's MARGINAL Gaussian sigma, per +#: axis. It is a CERTIFICATE-COVERAGE constant, not an accuracy knob, and the +#: distinction is measured: over W = 8, 14, 20, 30 the returned value does not move at +#: all (2.27e-13 nats from a converged reference at every W), while the omitted-mass +#: margin goes -18.4, -70.8, -160.2, -308.9 nats. Widening buys provability, not +#: accuracy. +#: +#: WHY 8 IS NOT ENOUGH, although exp(-8^2/2) = 1.3e-14 suggests it should be. That +#: estimate assumes the mode is locally Gaussian out to 8 sigma. These modes sit on +#: RIDGES -- the Hessian condition number at the dominant mode measures 11-23, and the +#: co-dominant modes come in pairs only ~0.009 rad apart -- so an axis-aligned box of +#: marginal sigmas under-covers the ridge, and the outside supremum then lands on a +#: shoulder only 18 nats below the peak rather than on a genuine subdominant maximum +#: 432 nats down. At W = 8 that fails a -23 nat tolerance; at 14 it passes with -71. +#: 16 is chosen with margin over the measured 14, and is re-derived rather than +#: inherited if the mode structure changes. +W_SIGMA = 16.0 + +#: Region merging is iterated to a fixed point; this only bounds pathological input. +MERGE_MAX_PASSES = 12 + +#: Accept a row when log(omitted area) + sup_outside - log(integral) is below this. +#: exp(-23) ~ 1e-10 of the mass. +OUTSIDE_TOL_NATS = -23.0 + + +def joint_table(C_A, C_B, x=1.0): + """Coefficient table of ``g = x*A - x**2/2 * B`` from the anglemarg tables. + + ``C_A`` and ``C_B`` have DIFFERENT bidegrees -- ``A`` is linear in the waveform + (phi <= m_max, u <= 1), ``B`` is quadratic (phi <= 2 m_max, u <= 2) -- which is + why the ``e^{2iu}`` coefficient carries no ``A`` contribution at all, exactly as + ``anglemarg._laplace_psi_lnI`` documents. ``A`` is zero-padded into ``B``'s shape. + """ + a = np.asarray(C_A) + b = np.asarray(C_B) + out = (-0.5 * x * x) * b.astype(complex) + kp = a.shape[0] + ksa = (a.shape[1] - 1) // 2 + ksb = (b.shape[1] - 1) // 2 + out[:kp, ksb - ksa:ksb + ksa + 1] += x * a + return out + + +def _kq(C): + KP = C.shape[0] + KS = (C.shape[1] - 1) // 2 + k = np.arange(KP)[:, None] + q = np.arange(-KS, KS + 1)[None, :] + w = np.ones((KP, 1)) + w[1:] = 2.0 # k > 0 stored once, counted twice (real field) + return k, q, w, KS + + +#: Points per chunk in :func:`eval_g`. The temporary is ``(chunk, KP, 2KS+1)`` +#: complex, so an unchunked call on a fine reference grid allocates tens of GB -- a +#: 8192^2 torus grid would ask for ~27 GB. Internal memory parameter only: it cannot +#: change the answer beyond floating-point reassociation, and does not even do that +#: here because each point is summed independently. +_POINT_CHUNK = 200_000 + + +def eval_g(C, phi, u, order=(0, 0)): + """``d^a_phi d^b_u g`` at points ``(phi, u)``; ``order=(a, b)``. + + Chunked over points: see :data:`_POINT_CHUNK`. + """ + k, q, w, _ = _kq(C) + a, b = order + phi = np.atleast_1d(np.asarray(phi, dtype=float)) + u = np.atleast_1d(np.asarray(u, dtype=float)) + fac = ((1j * k) ** a * (1j * q) ** b)[None] + wC = (w * C)[None] + n = phi.shape[0] + out = np.empty(n, dtype=np.float64) + for i in range(0, n, _POINT_CHUNK): + j = min(i + _POINT_CHUNK, n) + E = np.exp(1j * (phi[i:j, None, None] * k[None] + + u[i:j, None, None] * q[None])) + out[i:j] = (E * fac * wC).sum(axis=(1, 2)).real + return out + + +def derivative_bound(C, order=(0, 0)): + """TRUE bound on ``|d^a_phi d^b_u g|`` by the triangle inequality on the table. + + Not overridable and not fitted -- the one construction that cannot be a fit. This + is the 2-D multi-index form of the time module's ``spectral_derivative_bound``. + """ + k, q, w, _ = _kq(C) + a, b = order + return float((w * np.abs(C) * (np.abs(k) ** a) * (np.abs(q) ** b)).sum()) + + +def u_stationary_at_phi(C, phi): + """EXACT u-stationary points at fixed ``phi``, as angles in [0, 2 pi). + + At fixed ``phi`` the exponent is ``a + Re(c1 e^{iu}) + Re(c2 e^{2iu})``; with + ``z = e^{iu}`` its u-derivative vanishes on the roots of a quartic. ALL roots are + returned as seeds -- see the module docstring on why there is no ``|z| = 1`` filter. + """ + k, q, w, KS = _kq(C) + ph = (np.exp(1j * phi * k) * w).ravel() + # both signs of q are stored, and Re[D_{-q} e^{-iqu}] = Re[conj(D_{-q}) e^{+iqu}], + # so the effective coefficient is D_{+q} + conj(D_{-q}). Only the +q column was + # used here originally; the error was invisible because these roots are SEEDS that + # 2-D Newton then corrects -- the jax kernel, where the roots define the integration + # partition, is where it showed up. + _D = lambda qq: complex((ph * C[:, KS + qq]).sum()) + c1 = _D(1) + np.conj(_D(-1)) + c2 = (_D(2) + np.conj(_D(-2))) if KS >= 2 else 0.0 + 0.0j + P = np.array([c2, c1 / 2.0, 0.0, -np.conj(c1) / 2.0, -np.conj(c2)]) + nz = np.nonzero(np.abs(P) > 0.0)[0] + if nz.size < 2: + return np.zeros(0) + return np.mod(np.angle(np.roots(P[nz[0]:])), 2.0 * np.pi) + + +def _wrap(d): + """Signed periodic difference in (-pi, pi].""" + return (np.asarray(d) + np.pi) % (2.0 * np.pi) - np.pi + + +def enumerate_modes(C, n_phi=64, newton_iters=12): + """Local maxima of ``g`` on the torus, as ``(points, hessians)``. + + Seeds are ``phi`` grid x EXACT u-roots (see :func:`u_stationary_at_phi`), refined + by 2-D Newton. Seeds are targeting only: a seed that converges nowhere useful is + dropped, and a mode found twice is deduplicated. Neither costs correctness -- + what the regions miss is carried by :func:`outside_bound`. + """ + phis = np.linspace(0.0, 2.0 * np.pi, int(n_phi), endpoint=False) + seeds = [(p, u) for p in phis for u in u_stationary_at_phi(C, p)] + if not seeds: + return np.zeros((0, 2)), np.zeros((0, 2, 2)) + P = np.array(seeds, dtype=float) + + for _ in range(int(newton_iters)): + gp = eval_g(C, P[:, 0], P[:, 1], (1, 0)) + gu = eval_g(C, P[:, 0], P[:, 1], (0, 1)) + gpp = eval_g(C, P[:, 0], P[:, 1], (2, 0)) + guu = eval_g(C, P[:, 0], P[:, 1], (0, 2)) + gpu = eval_g(C, P[:, 0], P[:, 1], (1, 1)) + det = gpp * guu - gpu * gpu + ok = np.abs(det) > 1e-300 + dp = np.where(ok, -(guu * gp - gpu * gu) / np.where(ok, det, 1.0), 0.0) + du = np.where(ok, -(-gpu * gp + gpp * gu) / np.where(ok, det, 1.0), 0.0) + step = np.hypot(dp, du) + # Trust region: an unbounded Newton step means the seed is on a saddle ridge, + # not that the mode is far away. + scale = np.where(step > 0.5, 0.5 / np.maximum(step, 1e-300), 1.0) + P[:, 0] = np.mod(P[:, 0] + dp * scale, 2.0 * np.pi) + P[:, 1] = np.mod(P[:, 1] + du * scale, 2.0 * np.pi) + + gpp = eval_g(C, P[:, 0], P[:, 1], (2, 0)) + guu = eval_g(C, P[:, 0], P[:, 1], (0, 2)) + gpu = eval_g(C, P[:, 0], P[:, 1], (1, 1)) + res = np.hypot(eval_g(C, P[:, 0], P[:, 1], (1, 0)), + eval_g(C, P[:, 0], P[:, 1], (0, 1))) + m1 = derivative_bound(C, (1, 0)) + derivative_bound(C, (0, 1)) + is_max = (gpp < 0) & (gpp * guu - gpu * gpu > 0) & (res <= 1e-6 * max(m1, 1e-300)) + P = P[is_max] + H = np.stack([np.stack([gpp[is_max], gpu[is_max]], -1), + np.stack([gpu[is_max], guu[is_max]], -1)], -2) + if P.shape[0] == 0: + return P, H + + # deduplicate: modes closer than 1e-6 rad are the same mode found twice + keep = [] + for i in range(P.shape[0]): + d = np.hypot(_wrap(P[i, 0] - P[keep, 0]), _wrap(P[i, 1] - P[keep, 1])) \ + if keep else np.array([np.inf]) + if d.min() > 1e-6: + keep.append(i) + return P[keep], H[keep] + + +def _merge_boxes(cen, half): + """Merge overlapping axis-aligned boxes on the torus, to a fixed point. + + Merging is not tidiness: overlapping regions would double-count the mass between + them. It is also what makes the rule degrade CONTINUOUSLY into the dense grid -- + as amplitude falls the regions widen, merge, and the union grows to the whole + torus with no threshold anywhere. + """ + cen = cen.copy() + half = half.copy() + for _ in range(MERGE_MAX_PASSES): + n = cen.shape[0] + if n < 2: + break + merged = False + out_c, out_h, used = [], [], np.zeros(n, dtype=bool) + for i in range(n): + if used[i]: + continue + c, h = cen[i].copy(), half[i].copy() + for j in range(i + 1, n): + if used[j]: + continue + d = np.abs(_wrap(cen[j] - c)) + if np.all(d < h + half[j]): + lo = np.minimum(-h, _wrap(cen[j] - c) - half[j]) + hi = np.maximum(h, _wrap(cen[j] - c) + half[j]) + c = np.mod(c + 0.5 * (lo + hi), 2.0 * np.pi) + h = np.minimum(0.5 * (hi - lo), np.pi) + used[j] = True + merged = True + used[i] = True + out_c.append(c) + out_h.append(h) + cen = np.array(out_c) + half = np.array(out_h) + if not merged: + break + # VERIFY, do not assume. Exiting on the pass limit with boxes still overlapping + # would double-count the mass between them -- silently, and in the direction that + # inflates the answer. The caller declines on `converged=False`. + converged = True + for i in range(cen.shape[0]): + for j in range(i + 1, cen.shape[0]): + if np.all(np.abs(_wrap(cen[j] - cen[i])) < half[i] + half[j]): + converged = False + break + if not converged: + break + return cen, half, converged + + +def outside_bound(C, cen, half, n_grid=256): + """TRUE upper bound on ``g`` outside the covered boxes, and the uncovered area. + + A grid maximum alone is a LOWER bound on a supremum and the gap grows with + amplitude, so it is corrected by the Lipschitz remainder ``(M_phi + M_u) * h / 2`` + with the ``M`` from :func:`derivative_bound` -- a true bound from the exact + coefficient table, nothing fitted. + """ + step = 2.0 * np.pi / int(n_grid) + t = np.linspace(0.0, 2.0 * np.pi, int(n_grid), endpoint=False) + PHI, U = np.meshgrid(t, t, indexing='ij') + + # A CELL IS COVERED ONLY IF THE WHOLE CELL IS INSIDE A BOX, not merely its centre. + # Classifying centres leaves cells that straddle a box edge unexamined: their + # uncovered part contributes to neither the supremum search nor the area. The + # failure is not hypothetical or small -- a box positioned half a step off-axis can + # cover every grid CENTRE while leaving several rad^2 genuinely uncovered, and this + # function would then return (-inf, 0.0), i.e. UNCONDITIONAL ACCEPTANCE. Shrinking + # each box by half a cell before testing makes a straddling cell count as outside; + # the uncovered area is then an OVER-estimate, which is the safe direction. + # A box that already spans the full circle on an axis covers it whatever the + # shrink does; without this the low-amplitude case -- where the regions have merged + # to the whole torus, which is the rule degenerating into the dense grid exactly as + # intended -- would report an uncovered band and decline every such row. + shrink = 0.5 * step + inside = np.zeros(PHI.shape, dtype=bool) + for c, h in zip(cen, half): + eff = np.where(h >= np.pi, np.pi + 1.0, h - shrink) + inside |= ((np.abs(_wrap(PHI - c[0])) <= eff[0]) + & (np.abs(_wrap(U - c[1])) <= eff[1])) + area_out = float((~inside).sum()) * step * step + if not np.any(~inside): + return -np.inf, 0.0 + + # THE SLOPES ARE WHAT MAKE THIS AFFORDABLE, exactly as in the time module's + # certificate. A zeroth-order bound `max_grid + M1 * r` carries the GLOBAL first + # derivative bound, and M1 grows linearly with amplitude, so the remainder swamps + # the tolerance on any grid a peak-local rule can afford (measured: +1225 nats of + # remainder at amplitude 3.2e4 on a 256^2 grid, a "bound" above the integral + # itself). Using each cell's own gradient and paying M2 only on the quadratic term + # makes the remainder local: it is small wherever the surface is flat, which is + # precisely where the outside supremum lives. + r = 0.5 * np.sqrt(2.0) * step # half-diagonal of a cell + m = ~inside + ph = PHI[m].ravel() + uu = U[m].ravel() + g0 = eval_g(C, ph, uu) + gp = eval_g(C, ph, uu, (1, 0)) + gu = eval_g(C, ph, uu, (0, 1)) + # Why this M2 is a valid remainder, written out because it is not obvious and was + # doubted on review. For |d| <= r, |0.5 d^T H d| <= 0.5 r^2 [max(M20,M02) + M11] + # (maximise M20 cos^2 + M11|sin 2t| + M02 sin^2). The value used here, + # M20 + 2 M11 + M02, dominates that for any non-negative M, so the bound holds -- + # conservatively. Checked as well as argued: 0 violations over 1800 (point, radius) + # samples on 300 random tables. + m2 = (derivative_bound(C, (2, 0)) + 2.0 * derivative_bound(C, (1, 1)) + + derivative_bound(C, (0, 2))) + local = g0 + np.hypot(gp, gu) * r + 0.5 * m2 * r * r + return float(local.max()), area_out + + +#: Trapezoid points per local sigma. Derived, not tuned: the trapezoidal rule on a +#: Gaussian of width sigma at spacing h has relative error 2 exp(-2 pi^2 sigma^2/h^2) +#: by Poisson summation, so sigma/h = 3 gives 2e-77 -- the same argument, and the same +#: kind of margin, as UPSAMPLE_SAFETY in the band-limited time quadrature. +#: +#: Measured when dropping 6 -> 3: the value is UNCHANGED (0.0 nats) at amplitudes 325, +#: 3249 and 3.2e4, and the local point count falls from 147k to 82k. At amplitude 3.4 +#: it moves by 1.1e-8 nats -- because there the single region has merged to the whole +#: torus and the integrand is not a Gaussian bump at all, so the Poisson argument above +#: does not apply to it. That is the LOW-AMPLITUDE end of this rule's exclusion region, +#: where the selector should be routing to the dense rule anyway; the residual is +#: recorded rather than hidden because "bit-identical" would have been the wrong claim. +_PTS_PER_SIGMA = 3 + + +def _log_box_integral(C, c, h, pts_per_sigma=_PTS_PER_SIGMA, max_pts=256): + """``log int_box exp(g)`` by a tensor trapezoid sized from the LOCAL curvature.""" + n = [] + for ax in (0, 1): + order = (2, 0) if ax == 0 else (0, 2) + curv = abs(float(eval_g(C, c[0], c[1], order)[0])) + sig = 1.0 / np.sqrt(curv) if curv > 0 else h[ax] + want = int(np.ceil(2.0 * h[ax] / max(sig, 1e-12) * pts_per_sigma)) + 1 + n.append(int(np.clip(want, 9, max_pts))) + a = c[0] + np.linspace(-h[0], h[0], n[0]) + b = c[1] + np.linspace(-h[1], h[1], n[1]) + A, B = np.meshgrid(a, b, indexing='ij') + g = eval_g(C, A.ravel(), B.ravel()).reshape(A.shape) + wa = np.full(n[0], 2.0 * h[0] / (n[0] - 1)); wa[0] *= 0.5; wa[-1] *= 0.5 + wb = np.full(n[1], 2.0 * h[1] / (n[1] - 1)); wb[0] *= 0.5; wb[-1] *= 0.5 + W = np.log(wa)[:, None] + np.log(wb)[None, :] + m = g.max() + return m + np.log(np.sum(np.exp(g - m + W))), n[0] * n[1] + + +def joint_marginalize_peak_local(C, n_phi=64, n_bound_grid=256, + tol_nats=OUTSIDE_TOL_NATS): + """``log[(2 pi)^-2 int int dphi du exp(g)]``, refining only near the modes. + + Returns ``(value, ok, report)``. ``ok`` is False when the omitted-mass bound could + not be made small enough; the caller must then use the dense rule. The value is + returned either way for diagnosis, but a value with ``ok=False`` is NOT to be used. + """ + C = np.asarray(C) + rep = {'n_modes': 0, 'n_regions': 0, 'n_local_points': 0, + 'margin': np.inf, 'area_outside': np.nan, 'sup_outside': np.nan, + 'decline': None} + + P, H = enumerate_modes(C, n_phi=n_phi) + rep['n_modes'] = int(P.shape[0]) + if P.shape[0] == 0: + rep['decline'] = 'no modes enumerated' + return -np.inf, False, rep + + # marginal sigmas of the local Gaussian: sqrt of the diagonal of (-H)^-1 + half = np.empty_like(P) + for i in range(P.shape[0]): + Ci = np.linalg.inv(-H[i]) + half[i, 0] = W_SIGMA * np.sqrt(max(Ci[0, 0], 1e-300)) + half[i, 1] = W_SIGMA * np.sqrt(max(Ci[1, 1], 1e-300)) + half = np.minimum(half, np.pi) + + cen, half, merged_ok = _merge_boxes(P, half) + rep['n_regions'] = int(cen.shape[0]) + if not merged_ok: + rep['decline'] = 'regions still overlap after MERGE_MAX_PASSES' + return -np.inf, False, rep + + parts, npts = [], 0 + for c, h in zip(cen, half): + v, k = _log_box_integral(C, c, h) + parts.append(v) + npts += k + rep['n_local_points'] = int(npts) + parts = np.array(parts) + m = parts.max() + log_inside = m + np.log(np.exp(parts - m).sum()) + + sup_out, area_out = outside_bound(C, cen, half, n_grid=n_bound_grid) + rep['sup_outside'] = sup_out + rep['area_outside'] = area_out + if area_out <= 0.0: + rep['margin'] = -np.inf + else: + rep['margin'] = float(np.log(area_out) + sup_out - log_inside) + + ok = rep['margin'] < tol_nats + if not ok: + rep['decline'] = 'omitted-mass bound too large' + return float(log_inside - 2.0 * np.log(2.0 * np.pi)), bool(ok), rep + + +def joint_marginalize_over_distance(C_A_st, C_B_st, x_grid, log_w_grid, + n_phi=64, n_bound_grid=256, + tol_nats=OUTSIDE_TOL_NATS, keep_nats=None, + _retry=False): + """Distance-, phi- and psi-marginalized value at ONE ``(sample, time)`` point. + + ``log sum_x exp(log_w_x) * (2 pi)^-2 int int exp(x A - x^2/2 B)``, i.e. the same + quantity ``anglemarg.fused_log_likelihood_distphipsimarg_exact`` produces before + time marginalization, with the same normalization. + + THE DISTANCE AXIS IS NOT FREE, and this is where the joint rule's cost actually + lives. The mode locations depend on ``x``, so the dense scheme's trick -- form one + ``(phi, u)`` grid and reuse it for every distance node -- is not available: the + enumeration is per node. What rescues it is that the number of nodes CARRYING MASS + falls as the signal sharpens (measured on the synthetic fixture: 16 of 64 nodes + within 20 nats at exponent amplitude 32, and 1 of 64 at amplitude 3249), so a cheap + pre-pass on the nodes bounds the work before any enumeration happens. + + READ THIS BEFORE QUOTING THE HIGH-SNR NUMBERS. That same collapse means the FIXED + distance grid is itself under-resolving the distance peak there -- the peak is + ``~1/SNR`` narrow, which is precisely the defect ``core._distmarg_gh_logL`` exists + to fix with adaptive nodes. So "1 node carries the mass" is simultaneously a cost + win for this rule and a warning about the grid it was handed. This function + inherits the caller's distance quadrature and does not repair it. + + ``keep_nats`` selects the nodes to work on, using a CHEAP upper bound on each node's + contribution rather than the node's actual value: ``log_w_x + max_(phi,u) g_x``, + where the maximum is taken over the coarse bound grid and lifted by the same local + slope/curvature remainder the outside bound uses. Dropping a node therefore drops + something provably below the kept mass, not something estimated to be -- and the + drop is CHECKED against ``tol_nats``, not merely reported. + + ``keep_nats`` is DERIVED from the tolerance by default, not an independent constant. + It was one (25.0), and the two never had to agree: the dropped set's certified + contribution came out only 15.7 nats below the kept value against a 23 nat + tolerance, so the filter was quietly discarding more than the rule was allowed to + lose. The requirement is ``log(n_dropped) + max(ub_dropped) - value < tol_nats``; + keeping everything within ``|tol_nats| + log(n_nodes)`` of the best bound satisfies + it with room to spare, and adapts automatically if either is changed. + """ + if keep_nats is None: + keep_nats = abs(float(tol_nats)) + np.log(max(len(x_grid), 1)) + 5.0 + keep_nats = float(keep_nats) + x_grid = np.asarray(x_grid, dtype=float).ravel() + log_w_grid = np.asarray(log_w_grid, dtype=float).ravel() + + # --- cheap pre-pass: a true upper bound on each node's contribution + t = np.linspace(0.0, 2.0 * np.pi, 96, endpoint=False) + PHI, U = np.meshgrid(t, t, indexing='ij') + r = 0.5 * np.sqrt(2.0) * (2.0 * np.pi / 96) + ub = np.empty(x_grid.size) + for i, x in enumerate(x_grid): + C = joint_table(C_A_st, C_B_st, x=float(x)) + g0 = eval_g(C, PHI.ravel(), U.ravel()) + gp = eval_g(C, PHI.ravel(), U.ravel(), (1, 0)) + gu = eval_g(C, PHI.ravel(), U.ravel(), (0, 1)) + m2 = (derivative_bound(C, (2, 0)) + 2.0 * derivative_bound(C, (1, 1)) + + derivative_bound(C, (0, 2))) + ub[i] = log_w_grid[i] + float((g0 + np.hypot(gp, gu) * r + + 0.5 * m2 * r * r).max()) + # The filter is a COST optimization and must never change the answer. A threshold + # alone cannot guarantee that: `ub` is certified but LOOSE -- measured about 13 nats + # above the value it bounds -- so a cut that looks safe against `ub.max()` can still + # leave the dropped set above tolerance relative to the ACTUAL value. That is why + # the check below is made against the computed value, and why failing it retries + # with every node rather than widening by a guess: adding nodes only raises the + # value and empties the dropped set, so the retry always succeeds. + live = np.nonzero(ub > ub.max() - keep_nats)[0] + + parts, ok_all, rep = [], True, {'n_nodes': int(x_grid.size), + 'n_nodes_live': int(live.size), + 'worst_margin': -np.inf, 'declines': []} + for i in live: + C = joint_table(C_A_st, C_B_st, x=float(x_grid[i])) + v, ok, r_i = joint_marginalize_peak_local( + C, n_phi=n_phi, n_bound_grid=n_bound_grid, tol_nats=tol_nats) + if not ok: + ok_all = False + rep['declines'].append((int(i), r_i['decline'])) + rep['worst_margin'] = max(rep['worst_margin'], r_i['margin']) + parts.append(log_w_grid[i] + v) + + if not parts: + return -np.inf, False, rep + parts = np.array(parts) + m = parts.max() + value = m + np.log(np.exp(parts - m).sum()) + + # THE DROPPED NODES MUST GATE THE RESULT, not merely be recorded. An earlier + # revision computed `ub` for them, stored the maximum in the report, and then + # returned `ok` without ever comparing it to anything -- the docstring promised a + # provably-negligible drop while the code performed an unchecked one. Each dropped + # node contributes at most `exp(ub_i)`, so the whole dropped set contributes at most + # `log(n_dropped) + max(ub)`, and that must sit below the tolerance relative to the + # kept value on the same scale. + dropped = np.setdiff1d(np.arange(x_grid.size), live) + if dropped.size: + # logsumexp over the dropped set, not log(n) + max: both are valid upper + # bounds on the omitted contribution, but the aggregate one is tighter and so + # declines fewer rows for the same guarantee. + _dm = ub[dropped].max() + drop_bound = float(_dm + np.log(np.exp(ub[dropped] - _dm).sum())) + rep['dropped_bound'] = drop_bound + rep['dropped_margin'] = drop_bound - value + if rep['dropped_margin'] >= tol_nats and not _retry: + # the cut was not justified against the real value: redo with every node. + rep2 = dict(rep) + v2, ok2, r2 = joint_marginalize_over_distance( + C_A_st, C_B_st, x_grid, log_w_grid, n_phi=n_phi, + n_bound_grid=n_bound_grid, tol_nats=tol_nats, + keep_nats=np.inf, _retry=True) + r2['prefilter_retried'] = True + r2['prefilter_first_margin'] = rep['dropped_margin'] + return v2, ok2, r2 + if rep['dropped_margin'] >= tol_nats: + ok_all = False + rep['declines'].append(('dropped-nodes', 'pre-filter bound too large')) + return float(value), bool(ok_all), rep + + +# ------------------------------------------------------------------ phi-local + +def _g_uu_at(C, phi, u): + """``d^2 g / du^2`` at one ``phi`` and several ``u``.""" + return eval_g(C, np.full(np.size(u), float(phi)), np.asarray(u, dtype=float), (0, 2)) + + +def u_profile(C, phi, n_nodes=64, window_sigma=12.0): + """``F(phi) = log int du exp(g)``, and its first two EXACT derivatives. + + The u integral is done on the cell partition (the sorted u-stationary points tile the + circle), so ``F`` is exact rather than a Laplace model. Its derivatives come from + the same nodes at no extra evaluation cost, by differentiating under the integral: + + F' = E[d_phi g] + F'' = E[d^2_phi g] + Var(d_phi g) + + with the expectation under the normalized ``exp(g) du`` on the same axis. That + variance term is why the phi axis cannot inherit the u axis's economy: it grows with + amplitude, so ``F`` sharpens as the signal does even though ``g`` itself does not. + + Returns ``(F, dF, ddF)``, each shaped like ``phi``. + """ + phi = np.atleast_1d(np.asarray(phi, dtype=float)) + k, q, w, KS = _kq(C) + out = np.empty((3, phi.size)) + for i, p in enumerate(phi): + ph = (np.exp(1j * p * k) * w).ravel() + _D = lambda qq: complex((ph * C[:, KS + qq]).sum()) + c1 = _D(1) + np.conj(_D(-1)) + c2 = (_D(2) + np.conj(_D(-2))) if KS >= 2 else 0.0 + 0.0j + P = np.array([c2, c1 / 2.0, 0.0, -np.conj(c1) / 2.0, -np.conj(c2)]) + nz = np.nonzero(np.abs(P) > 0.0)[0] + roots = (np.mod(np.angle(np.roots(P[nz[0]:])), 2 * np.pi) + if nz.size >= 2 else np.linspace(0, 2 * np.pi, 4, endpoint=False)) + u = np.sort(np.concatenate([roots, np.zeros(max(0, 4 - roots.size))]))[:4] + mid = 0.5 * (u + np.roll(u, -1) + np.where(np.arange(4) == 3, 2 * np.pi, 0.0)) + lo_c = np.roll(mid, 1) - np.where(np.arange(4) == 0, 2 * np.pi, 0.0) + + # WINDOW EACH CELL BY ITS OWN sigma, do not spread a fixed node count over the + # whole cell. The cells do NOT shrink with amplitude -- the stationary points of + # g are invariant under g -> lambda g -- while the peak inside them does, so a + # fixed uniform rule over the full cell silently under-resolves as the signal + # sharpens. Measured before this change, against a converged reference: error + # 8.8e-07 at exponent amplitude 1265 and 5.4e-04 at 4217, both of which fall to + # EXACTLY 0.0 when the node count is raised -- i.e. the entire residual was this + # rule, not the method. A resolution that is a fixed number whose default is + # assumed ample is the defect this whole line of work exists to remove. + # REFINE THE CENTRE INSIDE EACH CELL FIRST. The quartic roots are SEEDS, not + # located maxima: this module says elsewhere that they may leave the unit circle + # (a conjugate-reciprocal pair does exactly that), and a spurious root's angle is + # not a stationary point at all. Windowing +-W sigma around the raw root then + # centres the window on the wrong place and takes sigma from the wrong curvature, + # which under-resolves the peak that IS in the cell -- an inside-cell quadrature + # error the phi omitted-mass certificate cannot see. Measured on a constructed + # table before this change: -7.2e-04 nats returned with ok=True, converging to + # the reference only as n_nodes was raised. The jax kernel already refines; + # this path did not, and the two must not differ on something load-bearing. + ustar = u.copy() + pv = np.full(ustar.size, float(p)) + for _ in range(8): + g1 = eval_g(C, pv, ustar, (0, 1)) + g2 = eval_g(C, pv, ustar, (0, 2)) + step = np.where(np.abs(g2) > 0.0, + -g1 / np.where(np.abs(g2) > 0.0, g2, 1.0), 0.0) + ustar = np.clip(ustar + np.clip(step, -0.5, 0.5), lo_c, mid) + + # A CLIPPED NEWTON POINT IS NOT A PEAK, however negative the curvature there. + # The iteration is clamped to [lo_c, mid], so it can come to rest ON a cell + # boundary with a large stationary residual; classifying that as a maximum + # centres a +-W sigma window on a non-stationary point and sizes sigma from the + # wrong curvature. Require, as well as g'' < 0, that the residual is small + # relative to the axis's own derivative bound AND that the point is interior. + # A cell failing either is integrated WHOLE rather than windowed, which is the + # conservative branch: it can only add nodes, never move the centre. + g1c = eval_g(C, pv, ustar, (0, 1)) + g2c = _g_uu_at(C, p, ustar) + _m1u = max(derivative_bound(C, (0, 1)), 1e-300) + _edge = 1e-9 * max(float(np.max(mid - lo_c)), 1e-300) + peaked = ((g2c < 0.0) + & (np.abs(g1c) <= 1e-8 * _m1u) + & (ustar > lo_c + _edge) & (ustar < mid - _edge)) + sig_c = np.where(peaked, 1.0 / np.sqrt(np.where(peaked, -g2c, 1.0)), np.inf) + lo = np.where(peaked, np.maximum(ustar - window_sigma * sig_c, lo_c), lo_c) + hi = np.where(peaked, np.minimum(ustar + window_sigma * sig_c, mid), mid) + # DERIVE THE NODE COUNT; the fallback cell is where a fixed one fails. A + # windowed cell spans +-W sigma so a fixed count resolves it, but a cell that + # FELL BACK spans the whole cell with the same nodes -- and an earlier comment + # here claimed that branch "can only add nodes", which was simply false: it adds + # none and spreads them wider, so rejecting a peak made the resolution WORSE. + # Measured on a searched counterexample: 1.7e-03 nats at 64 nodes, converging + # only by n = 1024. + # + # The requirement is derived, not tuned: |d^2 g/du^2| <= M2u = |c1| + 4|c2| + # EXACTLY, so no feature of exp(g) on this axis is narrower than + # sigma_min = 1/sqrt(M2u), and a spacing of sigma_min/_PTS_PER_SIGMA resolves the + # sharpest thing the coefficients admit. + # SCALE, AND WHY THIS ONE. |d2g/du2| <= M2u = |c1| + 4|c2| exactly, so nothing + # on this axis is narrower than sigma_min = 1/sqrt(M2u) and a spacing of + # sigma_min/_PTS_PER_SIGMA resolves the sharpest feature the coefficients admit. + # Measured on a searched counterexample, this takes the 64-vs-1024-node gap from + # 1.7e-03 to 2.2e-04 nats at essentially no cost. + # + # A BOUNDARY-PEAKED FALLBACK CELL IS NOT GAUSSIAN THERE -- exp(g) falls off like + # exp(-|g'| du), so full convergence would need a spacing set by 1/M1u, and that + # was measured to cost a 25x slowdown for the remaining 2.2e-04 nats. Not taken: + # 2e-04 nats is orders below anything this rule is asked to decide, and the + # residual is reported (`n_fallback_cells`) rather than hidden. If a future + # caller needs it, raise _PTS_PER_SIGMA or pass n_nodes -- both are honest knobs + # and both cost what they cost. + m2u = abs(c1) + 4.0 * abs(c2) + width_c = np.maximum(hi - lo, 0.0) + need = int(np.ceil(float(width_c.max()) * np.sqrt(max(m2u, 1e-300)) + * _PTS_PER_SIGMA)) + 1 + n_use = int(np.clip(max(n_nodes, need), n_nodes, 8192)) + s = np.linspace(0.0, 1.0, n_use) + uu = lo[:, None] + width_c[:, None] * s[None, :] + pp = np.full(uu.size, p) + g = eval_g(C, pp, uu.ravel()) + gp = eval_g(C, pp, uu.ravel(), (1, 0)) + gpp = eval_g(C, pp, uu.ravel(), (2, 0)) + wq = np.full(n_use, 1.0 / (n_use - 1)); wq[0] *= 0.5; wq[-1] *= 0.5 + lw = (np.log(np.maximum(hi - lo, 1e-300))[:, None] + np.log(wq)[None, :]).ravel() + m = g.max() + wgt = np.exp(g - m + lw) + Z = wgt.sum() + e1 = float((wgt * gp).sum() / Z) + out[0, i] = m + np.log(Z) + out[1, i] = e1 + out[2, i] = float((wgt * (gpp + gp * gp)).sum() / Z) - e1 * e1 + return out[0], out[1], out[2] + + +def phi_local_marginalize(C, n_seed=64, w_sigma=12.0, n_nodes=64, + n_bound_grid=512, tol_nats=OUTSIDE_TOL_NATS): + """``log[(2 pi)^-2 int int dphi du exp(g)]`` with BOTH axes localized. + + u is exact on the cell partition; phi is localized around the maxima of the profile + ``F`` using its exact derivatives. The phi axis has no algebraic completeness + warrant -- ``F`` is a log-integral, not a trig polynomial -- so it is the framework's + grid-seeded class and its correctness rests on the cover bound, exactly as the time + axis does. + + Returns ``(value, ok, report)``; ``ok=False`` means the omitted-mass bound on phi + could not be made small enough and the caller must fall back. + """ + rep = {'n_phi_modes': 0, 'n_phi_regions': 0, 'margin': np.inf, 'decline': None} + seeds = np.linspace(0.0, 2.0 * np.pi, int(n_seed), endpoint=False) + p = seeds.copy() + for _ in range(24): + _, d1, d2 = u_profile(C, p, n_nodes=n_nodes) + step = np.where(np.abs(d2) > 0, -d1 / np.where(np.abs(d2) > 0, d2, 1.0), 0.0) + p = np.mod(p + np.clip(step, -0.3, 0.3), 2.0 * np.pi) + F, d1, d2 = u_profile(C, p, n_nodes=n_nodes) + keep = (d2 < 0) & (np.abs(d1) < 1e-6 * max(derivative_bound(C, (1, 0)), 1e-300)) + p, F, d2 = p[keep], F[keep], d2[keep] + if p.size == 0: + rep['decline'] = 'no phi modes' + return -np.inf, False, rep + order = np.argsort(p); p, F, d2 = p[order], F[order], d2[order] + uniq = np.concatenate([[True], np.diff(p) > 1e-6]) + p, F, d2 = p[uniq], F[uniq], d2[uniq] + rep['n_phi_modes'] = int(p.size) + + sig = 1.0 / np.sqrt(-d2) + lo = p - w_sigma * sig + hi = p + w_sigma * sig + # 1-D merge: sort by lo and absorb overlaps. Same argument as the time module -- + # merging is what stops the mass between two windows being counted twice. + # MERGE ON THE CIRCLE, NOT ON THE LINE. A linear sweep over raw + # [p - W sigma, p + W sigma] never joins a window near 0 to one near 2 pi -- but + # each region is afterwards integrated at mod(., 2 pi), so BOTH regions cover BOTH + # peaks and that mass is counted twice. Measured: +log 2 = +0.693 nats returned + # with ok=True and a margin of -437, because the error is INSIDE the regions and an + # omitted-mass certificate cannot see it. Same family as the wrapped-circuit bug, + # one step over. + # + # Reduce every interval to the circle, SPLIT any that crosses the seam, merge the + # pieces linearly, then close the circle by joining a piece touching 2 pi to one + # touching 0. + pieces = [] + for a, b in zip(lo, hi): + wdt = min(float(b - a), 2.0 * np.pi) + a = float(np.mod(a, 2.0 * np.pi)) + if a + wdt <= 2.0 * np.pi: + pieces.append((a, a + wdt)) + else: + pieces.append((a, 2.0 * np.pi)) + pieces.append((0.0, a + wdt - 2.0 * np.pi)) + pieces.sort() + ml, mh = [pieces[0][0]], [pieces[0][1]] + for a, b in pieces[1:]: + if a <= mh[-1]: + mh[-1] = max(mh[-1], b) + else: + ml.append(a) + mh.append(b) + if len(ml) > 1 and ml[0] <= 1e-12 and mh[-1] >= 2.0 * np.pi - 1e-12: + ml[0] = ml[-1] - 2.0 * np.pi # the two seam pieces are one region + ml.pop() + mh.pop() + ml, mh = np.array(ml, dtype=float), np.array(mh, dtype=float) + + # CLAMP TO ONE CIRCUIT. At low amplitude F is nearly flat, so sigma is huge and + # [p - W sigma, p + W sigma] spans far MORE than 2 pi; integrating that range + # literally wraps the circle several times and counts the same mass repeatedly. + # Found on real coefficient tables, not synthetic ones: exponent amplitude 1.09, + # one merged region, +1.84 nats too high -- a factor of e^1.84 = 6.3, i.e. six + # circuits -- and ACCEPTED, because a region covering everything leaves nothing + # outside for the omitted-mass certificate to object to. The certificate bounds + # what is OUTSIDE the regions; it cannot see an error made INSIDE one. + if float((mh - ml).sum()) >= 2.0 * np.pi: + ml, mh = np.array([0.0]), np.array([2.0 * np.pi]) + else: + span = np.minimum(mh - ml, 2.0 * np.pi) + mh = ml + span + covered = float(np.minimum(mh - ml, 2 * np.pi).sum()) + rep['n_phi_regions'] = int(ml.size) + # exposed so the disjointness invariant can be ASSERTED rather than inferred from a + # value comparison: overlapping regions double-count, and that error lives inside + # the regions where the omitted-mass certificate is blind to it. + rep['phi_regions'] = list(zip(ml.tolist(), mh.tolist())) + + parts = [] + for a, b in zip(ml, mh): + # spacing <= sigma/4 for the SHARPEST mode inside this region, and never fewer + # than 64 points across a full circuit -- a nearly-flat F has a huge sigma, so a + # sigma-derived count alone would leave a wrapped region with a handful of nodes. + inside_r = (p >= a - 1e-12) & (p <= b + 1e-12) + sloc = sig[inside_r].min() if np.any(inside_r) else sig.min() + n = int(np.ceil((b - a) / max(sloc, 1e-12) * 4)) + 1 + n = max(n, int(np.ceil(64 * (b - a) / (2 * np.pi))) + 1) + n = max(16, min(2048, n)) + gp = np.linspace(a, b, n) + Fv, _, _ = u_profile(C, np.mod(gp, 2 * np.pi), n_nodes=n_nodes) + wq = np.full(n, (b - a) / (n - 1)); wq[0] *= 0.5; wq[-1] *= 0.5 + m = Fv.max() + parts.append(m + np.log(np.sum(wq * np.exp(Fv - m)))) + parts = np.array(parts); m = parts.max() + value = m + np.log(np.exp(parts - m).sum()) - 2.0 * np.log(2.0 * np.pi) + + # cover bound on phi: same slope-plus-curvature form as the 2-D outside bound, with + # F'' bounded by M_(2,0) + M_(1,0)^2 (the variance term cannot exceed the square of + # the first-derivative bound). + t = np.linspace(0.0, 2.0 * np.pi, int(n_bound_grid), endpoint=False) + inside = np.zeros(t.size, dtype=bool) + step_t = 2.0 * np.pi / n_bound_grid + for a, b in zip(ml, mh): + d = _wrap(t - 0.5 * (a + b)) + inside |= np.abs(d) <= 0.5 * (b - a) - 0.5 * step_t + T_out = float((~inside).sum()) * step_t + if T_out <= 0.0 or covered >= 2 * np.pi: + rep['margin'] = -np.inf + return float(value), True, rep + # BOUND F THROUGH g, NOT THROUGH F''. The obvious route -- Taylor on F with + # F'' <= M_(2,0) + M_(1,0)^2 -- is useless: that variance bound grows as the SQUARE + # of the amplitude, so the remainder swamped everything (measured margins of +51 and + # +1196 nats at amplitude 1265 and 4217, i.e. no bound at all). Instead use + # F(phi) = log int du exp(g) <= log(2 pi) + sup_u g(phi, u), + # and bound that supremum with the SAME slope-plus-M2 form already validated for the + # 2-D outside bound, whose remainder grows only linearly in amplitude. + phi_out = t[~inside] + # The u axis SETS the covering radius here: at n_bound_grid = 512 the phi half-step + # is 0.0061 while a 128-point u grid gives pi/128 = 0.0245, so u dominates r and the + # remainder is four times larger than it need be. Tie the u resolution to the same + # knob so refining the bound refines BOTH axes; the cost is paid only on the + # UNCOVERED phi, which is the small set by construction. + n_ug = int(n_bound_grid) + ug = np.linspace(0.0, 2.0 * np.pi, n_ug, endpoint=False) + PH2, UU2 = np.meshgrid(phi_out, ug, indexing='ij') + # HALF-DIAGONAL OF THE CELL, with no extra factor. The grid spacings are step_t in + # phi and 2*pi/128 in u, so the farthest a point can sit from its grid point is + # hypot(step_t/2, pi/128). An earlier revision multiplied that by a further 0.5, so + # the Taylor remainder covered only half the cell and the "bound" was not one -- it + # happened not to be violated in the trials run, which is exactly why an unjustified + # constant is dangerous rather than harmless. + r = np.sqrt((0.5 * step_t) ** 2 + (np.pi / n_ug) ** 2) + g0 = eval_g(C, PH2.ravel(), UU2.ravel()) + gpv = eval_g(C, PH2.ravel(), UU2.ravel(), (1, 0)) + guv = eval_g(C, PH2.ravel(), UU2.ravel(), (0, 1)) + m2 = (derivative_bound(C, (2, 0)) + 2.0 * derivative_bound(C, (1, 1)) + + derivative_bound(C, (0, 2))) + sup_g = float((g0 + np.hypot(gpv, guv) * r + 0.5 * m2 * r * r).max()) + sup_out = np.log(2.0 * np.pi) + sup_g + rep['margin'] = float(np.log(T_out) + sup_out - 2.0 * np.log(2 * np.pi) - value) + ok = rep['margin'] < tol_nats + if not ok: + rep['decline'] = 'phi omitted-mass bound too large' + return float(value), bool(ok), rep diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_batchmode_stencil_default.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_batchmode_stencil_default.py new file mode 100644 index 000000000..f0c11768f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_batchmode_stencil_default.py @@ -0,0 +1,605 @@ +#!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Added with the test rather than to a shared list in +# ci.yml: that list is what this PR conflicted with when #242 replaced it. +"""test_batchmode_stencil_default -- what an ILE run gets when it says NOTHING. + +WHY THIS FILE EXISTS. Every other stencil test passes an EXPLICIT --interpolate-time, so the +whole suite was blind to the value the flag takes when it is absent -- which is the value +essentially every production run uses, because neither pipeline entry point emits the flag unless +asked (helper_LDG_Events.py only appends '--interpolate-time ' inside +`if time_interp_choice is not None`). That blindness is exactly how the two ILE drivers came to +ship OPPOSITE defaults for the same physical choice, batchmode 'nearest' against jax 'sinc' +(issue #233), with the disagreement growing as SNR^2 so it reads as an amplitude-dependent bug in +one of the codes rather than as a configuration difference. + +The default moved to time_interp_choice.TIME_INTERP_DEFAULT on 2026-09-02. A default change is +the most reachable kind of result change there is, so it is pinned three ways here: + + 1. the CONSTANT -- one definition, shared with the jax driver, so they cannot drift again; + 2. the WIRING -- the driver's add_option really reads that constant, checked with `ast` rather + than by trusting a re-typed literal (this is the check test_jax_stencil_parity already has + for --interp and the batchmode driver did not have for --interpolate-time); + 3. the BEHAVIOUR -- real subprocesses, because the interesting part of this change is not the + new value but the FOUR places where a DEFAULT must behave differently from a REQUEST. + +ON (3), STATED PLAINLY, because it is the part a reviewer should attack. The driver refuses an +explicit --interpolate-time it cannot honour. As a DEFAULT that same refusal would convert every +configuration lacking --time-marginalization/--vectorized/--gpu from working to a startup +ValueError, so the default is downgraded to 'nearest' instead, and the same distinction keeps the +default out of the time-posterior export mode, off the fused calibration kernel, and off the +LEGACY scalar path's `interpolate` boolean (the fourth, found by adversarial review on 2026-09-03: +correct in the code, pinned by nothing). Each of those four is a separate test below; deleting the +distinction makes at least one of them fail. + +TWO FURTHER THINGS THE SAME REVIEW FOUND, and the reason they belong in this file. A guard can be +wrong by being TOO BROAD as easily as by being absent, and a NOTICE can be wrong by describing +behaviour the code does not have -- neither shows up as a failure anywhere, because a needless +downgrade looks exactly like the historical behaviour and a wrong remedy string still prints. So +this file now also pins the OTHER edge of each: that a bare --calibration-fused-kernel with no +calibration envelope does NOT downgrade (nothing to protect), that each downgrade's stated remedy +matches what the driver actually does on that command line, and that the band-limited quadrature +notice fires under the new default and stays silent for 'simpson' and for an explicit 'nearest'. + +Subprocess cases cost a few seconds of lal/numba import each, so the list is kept to the ones +that DISTINGUISH behaviours. + + python3 test_batchmode_stencil_default.py # or: pytest test_batchmode_stencil_default.py +""" +from __future__ import print_function + +import ast +import os +import re +import shutil +import subprocess +import sys +import tempfile + +from RIFT.likelihood.time_interp_choice import (TIME_INTERP_CHOICES, + TIME_INTERP_DEFAULT) + +_HERE = os.path.dirname(os.path.abspath(__file__)) +CODE_ROOT = os.path.normpath(os.path.join(_HERE, '..', '..')) +BIN = os.path.join(CODE_ROOT, 'bin') +DRIVER = os.path.join(BIN, 'integrate_likelihood_extrinsic_batchmode') +HELPER = os.path.join(BIN, 'helper_LDG_Events.py') +PSEUDO = os.path.join(BIN, 'util_RIFT_pseudo_pipe.py') + +# The smallest command line the driver accepts that satisfies all three stencil prerequisites. +# --force-xpy keeps the identical NoLoop code path on numpy, so this runs on a CI box with no GPU. +HONOURED = ['--time-marginalization', '--vectorized', '--gpu', '--force-xpy'] + +# --calibration-fused-kernel can only LOSE a fused kernel if one would have run, and that needs a +# calibration envelope: use_fused_calmarg (batchmode:3261) is +# `calibration_marginalization and opts.calibration_fused_kernel`, and calibration_marginalization +# (:1317) is exactly bool(opts.calibration_envelope_directory). The path is never opened this +# early -- it is first read around :1300, long after the banner these tests parse -- so a +# non-existent directory is the cheapest way to put the driver in the calmarg configuration. Using +# the flag WITHOUT this was the P2 review finding these tests missed: the guard fired on runs where +# no kernel could run at all. +CALMARG = ['--calibration-envelope-directory', '/nonexistent-calibration-envelope-for-tests'] +FUSED = ['--calibration-fused-kernel'] + CALMARG +# NOTE the default --calibration-n-realizations is 100, i.e. > 1, so FUSED alone is a genuine +# fused configuration. test_no_fused_configuration_that_cannot_fuse_downgrades_the_default turns +# that knob off explicitly as one of its cases. + + +def _run(script, args, timeout=300, in_tmpdir=False): + """Run a script and return its combined output. Never raises on non-zero exit. + + Same idiom as test_interpolate_time_cli._run, deliberately: sys.executable rather than a + hard-coded interpreter, CODE_ROOT prepended to PYTHONPATH, and CUDA hidden so the cases are + CPU-only and deterministic. The driver exits non-zero on every case here (there are no data + files) -- what is under test is the text it prints BEFORE it gets that far. + """ + env = dict(os.environ) + env['PYTHONPATH'] = CODE_ROOT + os.pathsep + env.get('PYTHONPATH', '') + env['OMP_NUM_THREADS'] = '1' + env.setdefault('CUDA_VISIBLE_DEVICES', '') + # in_tmpdir: helper_LDG_Events writes helper_*_args.txt and local.cache into its CWD, so + # running it from the source tree leaves untracked files behind. Give it a scratch cwd. + tmp = tempfile.mkdtemp(prefix='stencil_default_') if in_tmpdir else None + try: + proc = subprocess.Popen([sys.executable, script] + args, env=env, cwd=tmp, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + out, _ = proc.communicate() + finally: + if tmp is not None: + shutil.rmtree(tmp, ignore_errors=True) + if not isinstance(out, str): + out = out.decode('utf-8', 'replace') + return out + + +def _squash(text): + return re.sub(r'\s+', ' ', text) + + +def _legacy_scalar_flag(out): + """The `legacy scalar path interpolate=` field of the same banner. + + Separate from _stencil_banner because it pins a DIFFERENT variable with a different + consumer: opts._legacy_interpolate_time, which is what + FactoredLogLikelihoodTimeMarginalized(..., interpolate=...) is handed on the NON-VECTORIZED + path. The stencil name and this boolean can disagree, and one mutation makes them. + """ + m = re.search(r'legacy scalar path interpolate=(\w+)', _squash(out)) + assert m, "driver printed no legacy-scalar field; output was: %s" % out[-1500:] + return m.group(1) + + +def _stencil_banner(out): + """The resolved stencil, read off the driver's own startup line. + + Reading the banner rather than re-deriving the value is the point: the banner is what a + configuration audit reads off a completed run's log, so if it and the code disagree the test + should fail. + """ + # Anchored on '(from --interpolate-time' so it cannot match the DOWNGRADE line, which is a + # different statement about the same subject. An earlier version of this helper was not + # anchored and read the downgrade notice as the banner, reporting the wrong stencil. + m = re.search(r'Q_lm sub-sample time stencil: (\S+) \(from --interpolate-time', _squash(out)) + assert m, "driver printed no stencil banner; output was: %s" % out[-1500:] + return m.group(1) + + +# --------------------------------------------------------------------------- +# 1. the constant +# --------------------------------------------------------------------------- +def test_default_is_a_real_stencil_and_is_sinc(): + assert TIME_INTERP_DEFAULT in TIME_INTERP_CHOICES, ( + "TIME_INTERP_DEFAULT %r is not a stencil this tree implements (%r)" + % (TIME_INTERP_DEFAULT, TIME_INTERP_CHOICES)) + assert TIME_INTERP_DEFAULT == 'sinc', ( + "default stencil changed to %r -- intentional? It changes results for every ILE run that " + "does not pass --interpolate-time, and the error grows as SNR^2, so update " + "DESIGN_q_window_stencil.md 9.6 and the --interpolate-time help text in the same commit." + % (TIME_INTERP_DEFAULT,)) + + +def test_the_two_ile_drivers_ship_the_same_default(): + """Issue #233 in one assertion, checked WITHOUT importing jax. + + This used to `from RIFT.likelihood.jax_ile.core import JAX_INTERP_DEFAULT` and skip when the + import failed. That skip fired in the very job this file belongs to -- q-window-stencil-check + runs on a numpy+lal image with no jaxlib by design -- so the one assertion tying the two + drivers together was the one assertion CI never evaluated. A gate that is skipped exactly + where it is needed is not a gate. + + What actually forbids drift is that core.py BINDS the name rather than copying the value, so + this reads the binding out of the source with ast and never imports the module. That is a + stronger check than the old equality as well as a runnable one: two literals that happen to + read 'sinc' today would have satisfied the import version and would fail here. + """ + core = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'jax_ile', 'core.py') + assert os.path.exists(core), core + with open(core) as handle: + tree = ast.parse(handle.read()) + bound = None + for node in ast.walk(tree): + if isinstance(node, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == 'JAX_INTERP_DEFAULT' for t in node.targets): + bound = node.value + assert bound is not None, "no module-level JAX_INTERP_DEFAULT assignment in %s" % core + assert isinstance(bound, ast.Name) and bound.id == 'TIME_INTERP_DEFAULT', ( + "jax_ile.core must ALIAS the shared default (JAX_INTERP_DEFAULT = TIME_INTERP_DEFAULT), " + "not re-type it: found %r. Two independently written literals are how the drivers came to " + "ship opposite defaults in the first place, and a comparison run at defaults then measures " + "a flag, with the difference growing as SNR^2. See issue #233." + % (ast.dump(bound),)) + + +# --------------------------------------------------------------------------- +# 2. the wiring: the driver reads the constant, it does not re-type the value +# --------------------------------------------------------------------------- +def _driver_option_default_expr(option): + with open(DRIVER) as handle: + tree = ast.parse(handle.read()) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + if not (isinstance(node.func, ast.Attribute) and node.func.attr == 'add_option'): + continue + if not (node.args and isinstance(node.args[0], ast.Constant) + and node.args[0].value == option): + continue + for kw in node.keywords: + if kw.arg == 'default': + return kw.value + return None + raise AssertionError("no add_option(%r) found in %s" % (option, DRIVER)) + + +def test_driver_default_is_the_absent_sentinel_not_a_stencil_literal(): + """`default=None` is load-bearing, not a stylistic choice. + + Every guard downstream keys on "was this asked for, or inherited?", and that question is + decided by `opts.interpolate_time is None`. Baking the stencil name straight into + `default=` would make an inherited default indistinguishable from an explicit request, which + silently re-arms the refusal, the export flip and the fused-kernel downgrade for runs that + passed no flag at all. + """ + default = _driver_option_default_expr('--interpolate-time') + assert isinstance(default, ast.Constant) and default.value is None, ( + "--interpolate-time default is %s; it must be None so the driver can tell an omitted " + "flag from an explicit one (see _interp_time_from_default)." + % ast.dump(default) if default is not None else "--interpolate-time has no default=") + + +def test_driver_resolves_the_absent_flag_through_the_shared_constant(): + """The name TIME_INTERP_DEFAULT must actually appear in the resolution, not just be imported. + + Mutation that this kills: replacing `opts._noloop_time_interp = TIME_INTERP_DEFAULT` with a + re-typed `'sinc'`. That passes every value-level test in this file and re-creates the exact + condition of issue #233 -- two literals in two files that agree today. + """ + with open(DRIVER) as handle: + source = handle.read() + assert 'from RIFT.likelihood.time_interp_choice import TIME_INTERP_DEFAULT' in source + assert 'opts._noloop_time_interp = TIME_INTERP_DEFAULT' in source, ( + "the driver no longer resolves an absent --interpolate-time through the shared constant") + + +# --------------------------------------------------------------------------- +# 3. the behaviour: a DEFAULT is downgraded exactly where a REQUEST is refused +# --------------------------------------------------------------------------- +def test_honoured_configuration_gets_the_new_default(): + out = _run(DRIVER, HONOURED) + assert _stencil_banner(out) == TIME_INTERP_DEFAULT, ( + "a configuration that CAN honour a stencil did not get the default %r: %s" + % (TIME_INTERP_DEFAULT, out[-1500:])) + + +def test_unhonourable_configuration_downgrades_the_default_instead_of_refusing(): + """The regression this change most plausibly introduces, and the reason for the whole design. + + '--vectorized' alone runs a likelihood with no time_interp argument at all. Under the old + 'nearest' default that configuration started normally; a naive default flip turns it into a + startup ValueError with no command line changed anywhere. + """ + out = _run(DRIVER, ['--vectorized']) + squashed = _squash(out) + assert 'cannot honour it' not in squashed, ( + "the honoured-path gate REFUSED a run that passed no --interpolate-time at all. A " + "default must never do that -- it breaks working configurations with no flag changed. " + "Output: %s" % out[-1500:]) + assert _stencil_banner(out) == 'nearest', ( + "an unhonourable configuration must fall back to 'nearest' (the pre-2026-09-02 default), " + "so the run is unchanged: %s" % out[-1500:]) + assert 'Q_lm stencil DEFAULT' in squashed and 'NOT APPLIED' in squashed, ( + "the fallback must be ANNOUNCED. A stencil that is not running is the one thing the log " + "has to say: %s" % out[-1500:]) + + +def test_the_downgrade_also_takes_the_legacy_scalar_path_down_with_it(): + """The FIFTH thing a default had to be prevented from doing, found by adversarial review. + + opts._legacy_interpolate_time is derived from the PROVISIONAL default a hundred lines before + the downgrade runs, and it is not the stencil: it is the plain boolean handed to + FactoredLogLikelihoodTimeMarginalized(..., interpolate=...) at batchmode:3645 and :4002, which + is the likelihood that ACTUALLY RUNS whenever --vectorized is absent. A non-'nearest' + provisional default makes it True, i.e. an omitted --interpolate-time would silently switch + every non-vectorized run onto the legacy path's unrelated cubic interpolation -- a result + change of the same class as the three in 9.6.3, on a path where no sub-sample stencil exists + at all. The downgrade block resets it; nothing pinned that until this test. + + '--time-marginalization' alone is deliberate: it is unhonourable (no --vectorized, no + --gpu/--rotation-slow/--freqresponse), so the downgrade fires, AND it is the branch where the + legacy scalar call site is reachable. Base behaviour is 'False' and must stay 'False'. + """ + out = _run(DRIVER, ['--time-marginalization']) + assert _stencil_banner(out) == 'nearest', out[-1500:] + assert _legacy_scalar_flag(out) == 'False', ( + "the downgrade left opts._legacy_interpolate_time True. On this configuration the " + "likelihood that runs is the LEGACY scalar path, whose `interpolate` argument this flag " + "is -- so the default would silently turn on the legacy cubic interpolation for every " + "non-vectorized run: %s" % out[-1500:]) + + +def test_an_explicit_request_is_still_refused_on_the_same_configuration(): + """The downgrade must not disarm the refusal. Same command line as the test above, one flag + added, opposite required outcome -- which is why they are separate tests and not one.""" + out = _squash(_run(DRIVER, ['--interpolate-time', TIME_INTERP_DEFAULT, '--vectorized'])) + assert 'cannot honour it' in out, ( + "an EXPLICIT --interpolate-time %r was not refused on a configuration that cannot honour " + "it. That refusal is what stops a comparison campaign being run against an inert flag." + % TIME_INTERP_DEFAULT) + + +def test_default_does_not_change_the_time_posterior_export_mode(): + """resolve_time_posterior_export_mode maps `auto` to 'continuous' for any non-'nearest' + stencil, so the same one-line default change would otherwise have flipped the fair-draw time + export of every --resample-time-marginalization run: a denser re-evaluation of the whole + likelihood, a different draw algorithm, two extra output columns, and a new reachable + MemoryError. The export must key on an EXPLICIT stencil only.""" + out = _squash(_run(DRIVER, HONOURED + ['--resample-time-marginalization', '--fairdraw-extrinsic-output'])) + assert 'Time-posterior export: grid' in out, ( + "an inherited default changed the time-posterior export mode; it must stay 'grid' unless " + "the stencil or the export was asked for. Output: %s" % out[-1500:]) + + +def test_an_explicit_stencil_still_opts_into_the_continuous_export(): + """The other half of the previous test: asking for a stencil is still an opt-in to the better + export, so this change narrows the trigger rather than removing the feature.""" + out = _squash(_run(DRIVER, HONOURED + ['--resample-time-marginalization', + '--fairdraw-extrinsic-output', + '--interpolate-time', 'sinc'])) + assert 'Time-posterior export: continuous' in out, ( + "an EXPLICIT --interpolate-time sinc no longer resolves `auto` to the continuous export: " + "%s" % out[-1500:]) + + +def test_default_stays_off_the_fused_calibration_kernel(): + """The fused calmarg kernels implement 'nearest' only, and the driver's three call sites fall + back to cal_method='loop' (and drop cal_distmarg) for any other stencil, silently. A default + must not spend someone else's --calibration-fused-kernel that way.""" + out = _run(DRIVER, HONOURED + FUSED) + assert _stencil_banner(out) == 'nearest', ( + "the default stencil was applied on top of --calibration-fused-kernel, which silently " + "moves the run off the fused kernel it explicitly asked for: %s" % out[-1500:]) + + +def test_an_explicit_stencil_with_the_fused_kernel_says_so(): + """Unchanged behaviour (the user named both flags), but it used to be silent at all three + call sites, which contradicts this option's own 'REFUSED, not ignored' promise.""" + out = _squash(_run(DRIVER, HONOURED + FUSED + ['--interpolate-time', 'sinc'])) + assert '--calibration-fused-kernel: NOT USED' in out, ( + "losing the fused kernel to an explicit stencil is still silent: %s" % out[-1500:]) + + +def test_the_fused_predicate_has_ONE_definition_shared_by_BOTH_call_sites(): + """The fourth P2 finding on PR #237, and the reason it is the last one of its kind. + + The condition "will a fused calibration kernel actually run?" is needed twice: at startup, to + decide whether the inherited stencil is downgraded, and at dispatch, to pick cal_method. + Written as two expressions it drifted THREE TIMES IN ONE DAY, always in the same direction -- + flag -> flag+envelope -> flag+envelope+path -- because an over-broad predicate downgrades to + 'nearest', which is indistinguishable from the historical behaviour, so nothing ever fails. + + A value test cannot catch that; the next conjunct someone forgets will be a fourth expression + that agrees with these on every case anyone thought to write down. So this pins the STRUCTURE: + one function, called at both sites, with no second expression left behind. + """ + with open(DRIVER) as handle: + source = handle.read() + tree = ast.parse(source) + funcs = [n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == 'fused_calmarg_in_use'] + assert len(funcs) == 1, ( + "the shared fused-calibration-kernel predicate is gone; if it was inlined again, the " + "startup guard and use_fused_calmarg are two expressions once more and will drift") + calls = [n for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == 'fused_calmarg_in_use'] + assert len(calls) >= 2, ( + "fused_calmarg_in_use is called %d time(s); BOTH the startup stencil guard and " + "use_fused_calmarg must go through it, or one of them is a second copy again" + % len(calls)) + # and no site re-derives it: the old shapes, in either order, must not reappear + for shape in ("bool(calibration_marginalization and opts.calibration_fused_kernel)", + "bool(opts.calibration_fused_kernel) and bool(\n opts.calibration_envelope_directory)"): + assert shape not in source, ( + "a hand-written copy of the fused-kernel predicate is back in the driver: %r" % shape) + # THE TRIPWIRE, pinned STRUCTURALLY and deliberately not claimed as behaviour coverage. + # + # The one term the early call cannot read is `calibration_marginalization`, for which it + # substitutes bool(opts.calibration_envelope_directory) -- that IS the condition the later + # assignment uses, so the two agree for every command line that exists, and removing the + # comparison changes nothing testable (verified by mutation, 2026-09-03: deleting it leaves + # all 23 tests green). Its whole purpose is to fail if someone later derives + # calibration_marginalization differently, and no test can reach that without editing the + # driver. So this asserts the tripwire is PRESENT rather than that it fires: an unexercisable + # guard that is silently deleted is worse than one that is honestly labelled. + compares = [n for n in ast.walk(tree) if isinstance(n, ast.Compare) + and isinstance(n.left, ast.Name) and n.left.id == 'use_fused_calmarg' + and any(isinstance(c, ast.Name) and c.id == '_fused_calmarg_would_run' + for c in n.comparators)] + assert compares, ( + "the startup/dispatch agreement check on the fused-kernel predicate is gone. It is the " + "only thing that makes a future change to how calibration_marginalization is derived " + "fail loudly instead of silently re-opening the drift this function exists to end.") + + +def test_no_fused_configuration_that_cannot_fuse_downgrades_the_default(): + """The OTHER edge, enumerated -- a guard must not be broader than the thing it protects. + + Each of these passes --calibration-fused-kernel, and in each the fused kernel CANNOT run, so + downgrading the inherited stencil to 'nearest' buys nothing and silently costs the accuracy + the new default exists to provide. They are listed one per reason rather than folded into a + single case because they fail for DIFFERENT reasons and a single expression has already been + wrong about three of them. + """ + cases = [ + ("no calibration envelope: nothing to marginalize over", + HONOURED + ['--calibration-fused-kernel']), + ("--calibration-n-realizations 1: the library returns from its n_cal==1 branch " + "before cal_method is read", + HONOURED + FUSED + ['--calibration-n-realizations', '1']), + ("--rotation-slow REPLACES the likelihood; the fused call site is in the else branch", + ['--time-marginalization', '--vectorized', '--rotation-slow', + '--calibration-fused-kernel'] + CALMARG), + ("--freqresponse REPLACES the likelihood, same dispatch", + ['--time-marginalization', '--vectorized', '--freqresponse', + '--calibration-fused-kernel'] + CALMARG), + ("--calibration-dump-responsibilities is a pilot: it evaluates with cal_method='loop' " + "and returns before the production integration exists", + HONOURED + FUSED + ['--calibration-dump-responsibilities', 'pilot_resp.npz']), + ] + for why, argv in cases: + out = _run(DRIVER, argv) + assert _stencil_banner(out) == TIME_INTERP_DEFAULT, ( + "the default stencil was downgraded to protect a fused kernel that cannot run (%s). " + "A needless downgrade is INVISIBLE -- it looks exactly like the historical behaviour " + "-- which is why this edge is pinned case by case: %s" % (why, out[-1500:])) + assert 'NOT APPLIED' not in _squash(out), ( + "the driver announced a downgrade it did not need to make (%s): %s" % (why, out[-1200:])) + + +def test_a_bare_fused_kernel_flag_no_longer_downgrades_the_default(): + """A guard must not be BROADER than the thing it protects. (P2 review finding, PR #237.) + + --calibration-fused-kernel with no --calibration-envelope-directory cannot run a fused kernel + under ANY stencil: use_fused_calmarg is `calibration_marginalization and the flag`. Keying the + downgrade on the flag alone therefore protected nothing on this command line and silently cost + the accuracy the new default exists to provide -- the failure mode is invisible, because a + needless downgrade looks exactly like the historical behaviour. + + Deliberately the SAME command line as test_default_stays_off_the_fused_calibration_kernel minus + the envelope, with the opposite required outcome, so the pair pins both edges of the condition. + """ + out = _run(DRIVER, HONOURED + ['--calibration-fused-kernel']) + assert _stencil_banner(out) == TIME_INTERP_DEFAULT, ( + "a --calibration-fused-kernel flag with no calibration envelope still downgraded the " + "default stencil, though no fused kernel can run: %s" % out[-1500:]) + squashed = _squash(out) + assert 'NOT APPLIED' not in squashed, ( + "the driver announced a downgrade it did not need to make: %s" % out[-1500:]) + assert 'NOT USED' not in squashed, ( + "the driver reported losing a fused kernel that was never going to run: %s" % out[-1500:]) + + +def test_each_downgrade_states_the_remedy_that_actually_applies(): + """A startup notice must not promise behaviour the driver does not have. (P2, PR #237.) + + The two downgrades have OPPOSITE remedies and one shared sentence lied about one of them. + Naming the stencil explicitly on a prerequisite downgrade gets you REFUSED; naming it on the + fused-kernel downgrade gets you ACCEPTED, with the kernel dropped and a notice. The single + message said "pass ... explicitly to be refused instead" in both cases. + + Asserted against the driver's real behaviour on the same two command lines, not just against + the strings: the refusal claim is checked by actually adding the flag and seeing a refusal, and + the acceptance claim by adding it and seeing the run continue. + """ + prereq = _squash(_run(DRIVER, ['--vectorized'])) + assert 'REFUSED rather than downgraded' in prereq, ( + "the prerequisite downgrade no longer states its remedy: %s" % prereq[-1200:]) + # ... and that claim is true: + assert 'cannot honour it' in _squash( + _run(DRIVER, ['--interpolate-time', TIME_INTERP_DEFAULT, '--vectorized'])), \ + "the prerequisite message promises a refusal the driver does not perform" + + fused = _squash(_run(DRIVER, HONOURED + FUSED)) + assert 'ACCEPTED, not refused' in fused, ( + "the fused-kernel downgrade still borrows the prerequisite downgrade's remedy: %s" + % fused[-1200:]) + assert 'REFUSED rather than downgraded' not in fused, ( + "the fused-kernel downgrade tells the user they will be refused; they will not: %s" + % fused[-1200:]) + # ... and THAT claim is true: same command line plus the explicit stencil is accepted. + accepted = _squash(_run(DRIVER, HONOURED + FUSED + ['--interpolate-time', TIME_INTERP_DEFAULT])) + assert 'cannot honour it' not in accepted, ( + "the fused-kernel message says an explicit stencil is accepted, but it was refused: %s" + % accepted[-1200:]) + assert '--calibration-fused-kernel: NOT USED' in accepted, ( + "the fused-kernel message says the driver will say so; it did not: %s" % accepted[-1200:]) + + +# --------------------------------------------------------------------------- +# 4. the band-limited quadratures were measured against a stencil that is no +# longer the default, and the pairing is now universal rather than rare +# --------------------------------------------------------------------------- +def test_bandlimited_says_its_advantage_is_unestablished_under_the_new_default(): + """The third P2 review finding on PR #237, and the reason it is not a rare corner. + + --time-marginalization-quadrature's prerequisites (--time-marginalization --vectorized --gpu) + are a strict SUBSET of the stencil's honoured set, so EVERY run that opts into a band-limited + quadrature without naming a stencil now gets the default one -- the regime where that module's + own docstring measures -2.29 nats against Simpson's +1.28. Startup said nothing. + """ + for quadrature in ('bandlimited', 'peak-local'): + out = _squash(_run(DRIVER, HONOURED + ['--time-marginalization-quadrature', quadrature])) + assert 'ADVANTAGE NOT ESTABLISHED' in out, ( + "--time-marginalization-quadrature %s ran under the DEFAULT stencil with no notice " + "that its measured advantage is for 'nearest': %s" % (quadrature, out[-1500:])) + assert 'DEFAULT Q_lm stencil' in out, ( + "the notice does not say the stencil was inherited rather than chosen: %s" + % out[-1500:]) + + +def test_the_quadrature_notice_does_not_fire_where_the_numbers_still_hold(): + """The other edge, without which the notice is unfalsifiable. + + A notice that always prints carries no information. It must be silent for 'simpson' (the + quadrature default, which these measurements are not about) and silent for an explicit + 'nearest' (the regime the numbers WERE measured in, and the reproduce instruction the notice + itself gives -- so if it still fired there the advice would be self-contradicting). + """ + assert 'ADVANTAGE NOT ESTABLISHED' not in _squash(_run(DRIVER, HONOURED)), \ + "the quadrature notice fired for the historical simpson quadrature" + out = _squash(_run(DRIVER, HONOURED + ['--time-marginalization-quadrature', 'bandlimited', + '--interpolate-time', 'nearest'])) + assert 'ADVANTAGE NOT ESTABLISHED' not in out, ( + "the notice fired for the very configuration it tells the user to switch to: %s" + % out[-1500:]) + + +# --------------------------------------------------------------------------- +# 5. spellings that must keep meaning what they meant +# --------------------------------------------------------------------------- +def test_explicit_nearest_and_explicit_off_still_mean_nearest(): + for value in ('nearest', 'none', 'False'): + out = _run(DRIVER, HONOURED + ['--interpolate-time', value]) + assert _stencil_banner(out) == 'nearest', ( + "--interpolate-time %r no longer resolves to 'nearest'; that is the only way to " + "reproduce a pre-2026-09-02 run: %s" % (value, out[-1500:])) + + +def test_legacy_truthy_still_means_cubic(): + out = _run(DRIVER, HONOURED + ['--interpolate-time', 'True']) + assert _stencil_banner(out) == 'cubic', ( + "the legacy truthy spelling must still map to 'cubic', not to the new default: %s" + % out[-1500:]) + + +def test_a_typo_is_still_loud(): + out = _squash(_run(DRIVER, HONOURED + ['--interpolate-time', 'lanczos'])) + assert 'unrecognised value' in out, ( + "a misspelled stencil was absorbed. Before this check it fell through to 'nearest'; " + "after the default change it would fall through to the new default, which is a different " + "wrong answer but still a silent one: %s" % out[-1500:]) + + +def test_pipeline_off_request_still_means_off(): + """"off" must still mean off, now that emitting nothing means the new default. + + resolve_interpolate_time_request collapses "flag absent" and an explicit off-request to the + same None. While the driver default was 'nearest' those were the same answer; they are now + opposites, so helper_LDG_Events has to re-express an off-request as an explicit + '--interpolate-time nearest' rather than emitting nothing. + """ + # --fmin is needed only to get the helper PAST its PSD/parameter setup and as far as the + # stencil log line; it then dies on missing frames, which _run tolerates. + out = _squash(_run(HELPER, ['--internal-ile-interpolate-time', 'False', + '--event-time', '1000000000', '--fmin', '20'], + in_tmpdir=True)) + assert '--interpolate-time nearest' in out or "stencil 'nearest'" in out, ( + "helper_LDG_Events did not turn an explicit off-request into an explicit " + "'--interpolate-time nearest'; omitting the flag now means the NEW default, so 'off' " + "would silently mean 'on': %s" % out[-2000:]) + + +def test_pseudo_pipe_forwards_an_off_request_instead_of_swallowing_it(): + """The other half of the off-request repair, and the half a subprocess test cannot reach. + + util_RIFT_pseudo_pipe forwards --internal-ile-interpolate-time to the helper only when + resolve_interpolate_time_request returns non-None -- which it does NOT for an off-request. So + the helper's repair above is unreachable through the pipeline entry point unless this + condition also stops swallowing it. Pinned in the source because driving the pseudo pipe far + enough to emit a helper command line needs a whole event configuration; the condition is one + line and this is the cheap check that it is still there. + """ + with open(PSEUDO) as handle: + source = handle.read() + assert "or opts.internal_ile_interpolate_time is not None" in source, ( + "util_RIFT_pseudo_pipe no longer forwards an explicit off-request to the helper. Since " + "the ILE default stopped being 'nearest', dropping the flag means the NEW default, so " + "'--internal-ile-interpolate-time False' would silently turn interpolation ON.") + + +if __name__ == "__main__": + for name, fn in sorted(list(globals().items())): + if name.startswith('test_') and callable(fn): + fn() + print("PASS %s" % name) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_calmarg_stencil_gating.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_calmarg_stencil_gating.py index ca2ff3e49..60c1d02ab 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_calmarg_stencil_gating.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_calmarg_stencil_gating.py @@ -1,3 +1,16 @@ +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_calmarg_stencil_gating runs its CPU arms without a GPU (its GPU arm is additive), +# so it belongs here: it is what stops cubic/sinc being routed to the fused calibration +# kernel, which is implemented for 'nearest' only. +# --------------------------------------------------------------------------------- """ test_calmarg_stencil_gating : the fused calibration-marginalization kernel is implemented ONLY for time_interp='nearest', and everything else must fall back to the 'loop' path. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_interpolate_time_cli.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_interpolate_time_cli.py index aa55a0535..a2f72b5ca 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_interpolate_time_cli.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_interpolate_time_cli.py @@ -1,4 +1,20 @@ #!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_interpolate_time_cli runs the three scripts as real SUBPROCESSES (~30 s). That +# cost is the point: the unit tests exercise the resolver and the gate predicate, but +# neither can see whether the SCRIPTS are still wired to them. Reverting a parser to +# const=None, or deleting a script's resolver call, leaves every unit test green while +# restoring a bare flag that silently does nothing. All three mutations were checked to +# fail these tests before they landed. +# --------------------------------------------------------------------------------- """test_interpolate_time_cli -- the stencil flag AT THE COMMAND LINE, in real subprocesses. WHY SUBPROCESSES AND NOT UNIT CALLS. test_time_interp_choice exercises @@ -196,16 +212,26 @@ def test_driver_refuses_configurations_that_cannot_honour_the_stencil(): print("driver rejects, missing %-22s : OK" % expect_missing) -def test_driver_does_not_gate_the_default_stencil(): +def test_driver_does_not_gate_an_explicit_nearest(): """'nearest' is the historical behaviour and must never be refused. - Without this, the gate could be tightened into breaking every run that does not ask for - interpolation at all -- a far worse regression than the one it prevents. + Without this, the gate could be tightened into breaking every run that asks for the + nearest-bin gather -- a far worse regression than the one it prevents. + + THIS TEST PASSES 'nearest' EXPLICITLY, so it says nothing about what an OMITTED + --interpolate-time does. While 'nearest' was also the default the two were the same case and + the name of this test did not lie; since 2026-09-02 the default is + time_interp_choice.TIME_INTERP_DEFAULT and they are different cases. The omitted-flag case -- + the one essentially every production run takes, and the one that would have been converted + into a startup ValueError by the same gate -- is covered by + test_batchmode_stencil_default.test_unhonourable_configuration_downgrades_the_default_instead_of_refusing. + Renamed rather than left alone precisely so this file cannot go on reporting green under a + name that claims coverage it does not have. """ out = _squash(_run(DRIVER, ['--interpolate-time', 'nearest', '--vectorized'])) assert 'cannot honour it' not in out, \ - "the gate must not fire for the default 'nearest' stencil: %s" % out[-400:] - print("driver does not gate 'nearest': OK") + "the gate must not fire for an explicit 'nearest' stencil: %s" % out[-400:] + print("driver does not gate an explicit 'nearest': OK") if __name__ == "__main__": @@ -215,5 +241,5 @@ def test_driver_does_not_gate_the_default_stencil(): test_help_text_carries_the_same_crossover_guidance_in_both_entry_points() test_error_messages_carry_the_canonical_guidance_too() test_driver_refuses_configurations_that_cannot_honour_the_stencil() - test_driver_does_not_gate_the_default_stencil() + test_driver_does_not_gate_an_explicit_nearest() print("\nPASS") diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_window_interp.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_window_interp.py index 626e950d4..de71e3fe8 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_window_interp.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_window_interp.py @@ -1,4 +1,17 @@ #!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_q_window_interp asserts the cubic/sinc crossover in BOTH directions on purpose. +# sinc winning everywhere would mean the Lanczos window had been widened until it was no +# longer a local stencil, so neither direction may be relaxed to make a change pass. +# --------------------------------------------------------------------------------- """test_q_window_interp.py -- accuracy of the Q(t) sub-sample interpolation stencils. Q^a_lm(t) is the inverse transform of something supported on [fmin, fmax], so it is BAND-LIMITED, diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_time_interp_choice.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_time_interp_choice.py index 4cdb58cc1..bae43f3d9 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_time_interp_choice.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_time_interp_choice.py @@ -1,4 +1,16 @@ #!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_time_interp_choice pins the pipeline thresholds inside the measured ambiguous +# band, and checks the decision uses the sampling rate the run is actually on. +# --------------------------------------------------------------------------------- """test_time_interp_choice -- the pipeline's Q_lm stencil handling. Automatic selection was REMOVED after measurement (see time_interp_choice's docstring for the diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_interp_choice.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_interp_choice.py index 9f5a3c32d..f9ac31cf1 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_interp_choice.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_interp_choice.py @@ -27,6 +27,30 @@ # factored_likelihood.TIME_INTERP_CHOICES must agree and test_time_interp_choice asserts it does. TIME_INTERP_CHOICES = ('nearest', 'cubic', 'sinc') +# Taps per SIDE for the 'sinc' stencil (the full stencil is 2a wide). It lives in this leaf +# module, rather than beside the weight builder in factored_likelihood, because three backends +# now need it -- the CPU window builder, the cupy kernel wrapper, and the JAX gatherer -- and the +# JAX path must not pay factored_likelihood's numba/lal import cost to learn one integer. +# factored_likelihood re-exports it, so `FL.SINC_HALFWIDTH_DEFAULT` keeps working. +SINC_HALFWIDTH_DEFAULT = 8 + +# THE DEFAULT SUB-SAMPLE STENCIL, for every driver. ONE definition, because the two ILE drivers +# previously shipped OPPOSITE defaults for the same physical choice -- batchmode 'nearest' and +# jax 'sinc' (issue #233) -- so a cross-implementation comparison run at defaults was measuring a +# flag, and the discrepancy grew as SNR^2, which reads as an amplitude-dependent bug in one of the +# codes rather than as a configuration difference. jax_ile.core.JAX_INTERP_DEFAULT is now an +# alias of this constant and test_batchmode_stencil_default pins that they are equal, so the two +# cannot drift apart again. +# +# WHY 'sinc' AND NOT 'cubic'. A default is chosen for its WORST case, not its average, because +# it is what people get without thinking. sinc's error is measured FLAT (3.1-7.9 nats over the +# fmin-30 mass ladder, 2.3-5.6 over the 20-point fmin sweep) while all the variation belongs to +# cubic (0.143-69.3 nats over the same points); error grows as SNR^2. This is the same reasoning, +# on the same measurements, that DESIGN_q_window_stencil.md 9.4 recorded for the JAX default. +# High-mass low-fmin campaigns are the population where cubic is measurably better and should +# pass '--interpolate-time cubic' explicitly -- see 9.6 of that document. +TIME_INTERP_DEFAULT = 'sinc' + # Values of --internal-ile-interpolate-time that mean "don't interpolate at all". These matter # because the flag takes a VALUE: '--internal-ile-interpolate-time False' passes the STRING # 'False', which is truthy in Python, so without this it would sail past an `if opts...:` guard diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_peak_local.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_peak_local.py new file mode 100644 index 000000000..04d7ef4f6 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_peak_local.py @@ -0,0 +1,1644 @@ +"""Peak-local time marginalization: enumerate the peaks, integrate only near them. + +READ ``time_marginalization_quadrature.py`` FIRST. That module states the defect +(Simpson at the fixed spacing ``deltaT = 1/srate`` against an integrand of width +``sigma_t = 1/(2 pi rho sigma_f)``), the reason the existing samples already +determine the continuous integrand (``kappa(t)`` is band-limited below Nyquist and +``rho_sq`` is time-independent on this path), and the derived-not-configured +discipline for resolution. All of that is inherited here unchanged. This module +changes ONE thing: WHERE the refined grid is placed. + +WHAT IS WRONG WITH REFINING THE WHOLE WINDOW +-------------------------------------------- +The dense strategy refines the entire window to the peak's width, so its cost +grows as ``rho`` while the peak it is resolving gets NARROWER as ``1/rho``: the +work grows exactly where it is least needed. It conflates two resolution +requirements that are not the same requirement. There are in fact THREE, and the +first version of this module named only two -- which is exactly where its one +correctness bug lived: + +* Resolving ``kappa(t)`` enough to **enumerate its extrema**. ``kappa`` is + band-limited at Nyquist by construction, so the narrowest feature it can have is + a half-cycle of width ``deltaT``. Enumerating its extrema therefore needs a + small FIXED factor and is **SNR-INDEPENDENT**. +* **Localising** each enumerated extremum to a fraction of ``sigma_t``. Enumeration + returns a grid INDEX, and an index is not a location: the crest can lie + ``h_enum/2`` from the sample that reports it. This requirement is + **SNR-DEPENDENT**, and it belongs to neither of the other two -- which is why it + went missing. Building the interval around the sample instead drops the peak + entirely once ``W_SIGMA * sigma_t < h_enum/2`` and cost up to **165 nats**; see + :data:`LOCALISE_SAFETY`. It is met by Newton on the spectral interpolant, which + converges quadratically, so its cost is logarithmic in rho rather than linear. +* Resolving ``exp(lnL(t))`` well enough to integrate it. That is the rho-dependent + part, and it is only needed over a few ``sigma_t`` around each enumerated peak. + +So: upsample ``kappa`` by ``PEAK_ENUM_FACTOR``, enumerate every local maximum of +the band-limited interpolant, put an interval of half-width ``W_SIGMA * sigma_i`` +around each, **merge overlapping intervals into disjoint ones**, integrate each +merged interval at its own derived spacing, and sum in log space. + +TWO PROPERTIES THAT MAKE THIS AN ALGORITHM RATHER THAN A HACK +------------------------------------------------------------- +**1. Merging makes it one algorithm, not a regime switch.** Isolated peaks give a +tiny union of intervals; overlapping peaks grow the union until it is the whole +window, and the method degenerates *continuously* into the dense grid. There is no +threshold to tune and no regime to detect. The un-merged version is not a slightly +worse variant of this -- it double-counts the shared region and is measurably wrong: +``~/tmarg_harness/peaklocal.py`` errs **+1.6 nats at rho ~ 6** where the merged +version is exact, and ``test_unmerged_intervals_double_count`` reproduces that here +so the merge cannot be quietly removed. + +**2. Enumeration plus a computed tail bound makes the truncation rigorous.** The +danger in any "integrate near the peak" scheme is the mass you did not look at. +RIFT PR #201 was caught by exactly that: a Newton solve seeded at two guessed points +missed genuine maxima and returned ``-inf`` for a finite integral. There is +deliberately no seeded root-finder here. Instead: + +* every local maximum of the band-limited interpolant is enumerated on a grid that + resolves ``kappa`` itself, and +* the mass outside the merged intervals is **bounded and checked, per row**, not + assumed. ``log(T_outside) + max_{t outside} lnL(t)`` is an upper bound on the + omitted integral; it is compared against the value actually computed, and a row + whose bound is not below ``TAIL_LOG_TOL`` relative is handed to the dense + band-limited path instead of being reported. + +That second bullet is the load-bearing one, and it is worth being precise about why: +**the bound does not depend on the enumeration being complete.** If a peak were +missed entirely, its neighbourhood would be outside the intervals, the sampled +maximum outside would sit on it, and the bound would fail -- sending the row to the +dense path. Completeness of the enumeration buys SPEED; the bound buys CORRECTNESS. +Those are deliberately not the same mechanism, because a missed-peak argument that +rests on "the grid is fine enough" is exactly the kind of claim that is true until +it is not. + +Evaluating the maximum outside is free, and that is not an accident. Every shipped +``loglikelihood`` callback in scope here -- plain and distance-marginalized -- is +monotone increasing in ``Re kappa`` at fixed ``rho_sq``, and ``rho_sq`` is +time-independent on this path, so a monotone map cannot move a maximum: + + argmax_t lnL(t) == argmax_t Re kappa(t) + +whatever the distance, the distance prior, or the callback. (Verified over an 800x +range in ``1/D`` and across three callback shapes; +``test_peak_positions_do_not_depend_on_distance_or_callback`` pins it.) So both the +enumeration AND the outside-maximum are computed on ``Re kappa`` alone, with no +callback evaluation on the full time axis -- which matters, because with the +distance-marginalized callback that evaluation is a table interpolation over +``n_extrinsic * npts * factor`` points and is the dominant cost of the dense path. +The callback is evaluated only at the enumerated peaks, at their curvature stencils, +at one point per row for the bound, and on the local grids. + +THE COST CLAIM, AND WHAT THE PROTOTYPE'S NUMBER DOES NOT MEAN +------------------------------------------------------------- +The prototype (``~/tmarg_harness/peaklocal2.py``, tabulated in +``DESIGN_time_marginalization_quadrature.md``) reports a flat ~97 evaluation points +against the dense grid's 19,648 -> 628,736, i.e. "200x at rho ~ 15, 6,482x at +rho ~ 692". **Those numbers were measured with the analytic ``kappa`` in hand**, so +evaluating ``kappa(t)`` at an arbitrary time cost one closed-form call. In the +shipped code it does not: only the coarse samples exist, and the band-limited +interpolant has to be evaluated at the local grid points. A point count is +therefore not a cost, and the prototype's speedups are NOT this module's speedups. +See ``DESIGN_time_marginalization_peak_local.md`` for what was measured here, through +this code path, including the interpolation. + +The evaluator used is a direct spectral sum, ``kappa(t) = sum_k Xw_k exp(2 pi i f_k +t / T)``, advanced along a uniform local grid by a phase recurrence so that no +transcendental is evaluated per output point. Its cost is ``O(npts)`` per output +point and, crucially, **independent of rho**: the number of local points is set by +``W_SIGMA`` and ``UPSAMPLE_SAFETY``, not by the peak's width. The dense path costs +``O(npts * factor * log(npts * factor))`` with ``factor`` proportional to rho, plus a +callback and an ``exp`` over every one of those points. So the two cross over, this +one wins by more as rho grows, and it LOSES at low rho -- which is why a row whose +estimated local cost exceeds its estimated dense cost is given the dense path. That +switch is a cost decision only -- it never substitutes an approximation for a value, it +chooses which of two paths computes it. + +Be precise about the sense in which the two branches agree, because the looser statement +this module used to make was false. The dense path ENFORCES its criterion: it remeasures +the width on the grid it actually integrated and doubles until the criterion holds there. +This path has no such loop. It DERIVES the local spacing from the coarse width and then +verifies the outcome two other ways -- the localisation must converge to +``LOCALISE_SAFETY * sigma_t``, and the local grids must ATTAIN the localised crest's +``lnL`` to within ``CONTAINMENT_SLACK_NATS``. Those are a-posteriori checks rather than +an enforcement loop, and a row failing either goes to the dense path. So both branches +are checked and neither returns a number it cannot defend, but they are NOT the same +criterion and this module should not say they are. + +A chirp-z (Bluestein) evaluation would reduce the local evaluation from +``O(npts * M)`` to ``O((npts + M) log(npts + M))`` and is the obvious next step; it +is deliberately not in this draft. + +SCOPE +----- +The band-limited path's scope, MINUS phase marginalization: baseline (non-rotating) +likelihood, ``n_cal == 1``, time-independent ``rho_sq`` (checked, not assumed). The +rotating-response and finite-size-response likelihoods refuse it for the same reason +they refuse ``bandlimited``. + +``phase_marginalization=True`` is REFUSED here, and that is a deliberate scope cut +rather than an oversight. Production marginalizes over distance, not phase, so it +would be a corner nobody runs -- and it is the one corner that genuinely complicates +this design. Under phase marginalization the Laplace width of the time peak picks up +a factor ``(I1/I0)(|kappa|/D)``, which depends on distance in a way that is not a +power law, so the width does NOT reduce the way it does for the plain and +distance-marginalized callbacks and the derived local spacing would no longer be +derivable from ``rho_sq`` and the curvature alone. ``bandlimited`` still supports it +and is unchanged; a caller who needs it should use that. Refusing rather than +silently falling back is the same discipline the rest of this option follows: an +accuracy option that quietly does something else is worse than one that is +unavailable. +""" + +import numpy as np + +from . import time_marginalization_quadrature as _tmq +from .time_marginalization_quadrature import ( + UPSAMPLE_SAFETY, + CURVATURE_STENCIL_HALFWIDTHS, + bandlimited_upsample, + reflected_bandlimited_upsample, + time_marginalize_bandlimited, + _log_simps_rows, + _safe_offset, + _require_time_independent_rho_sq, + _default_simps, +) + +__all__ = [ + "PEAK_ENUM_FACTOR", + "W_SIGMA", + "PEAK_KEEP_NATS", + "TAIL_LOG_TOL", + "MAX_INTERVALS", + "MIN_LOCAL_POINTS", + "LOCALISE_SAFETY", + "CONTAINMENT_SLACK_NATS", + "localise_peaks", + "bandlimited_spectrum", + "spectral_curvature_bound", + "spectral_derivative_bound", + "parabolic_sup", + "segment_sup_bound", + "crest_upper_bound", + "enum_grid_derivatives", + "eval_bandlimited_uniform", + "eval_bandlimited_points", + "enumerate_peak_indices", + "merge_intervals_by_row", + "time_marginalize_peak_local", + "last_report", +] + + +#: Upsampling factor used ONLY to enumerate the extrema of ``kappa`` and to sample +#: the maximum outside the intervals. SNR-INDEPENDENT, and that independence is the +#: whole point of the split: ``kappa`` is band-limited below Nyquist, so its fastest +#: possible oscillation is a half-cycle of width ``deltaT``, and this places 8 grid +#: points across that narrowest possible lobe regardless of how sharp ``exp(lnL)`` +#: has become. It is NOT an accuracy knob for the integral: the integral's accuracy +#: is set by the local spacing (``UPSAMPLE_SAFETY``) and by the tail bound, and a +#: value too small here shows up as a FAILED tail bound and a dense-path fallback, +#: not as a wrong number. Measured on the synthetic band-limited fixture, the +#: enumerated peak set at this factor is identical to the set found at factor 64 on +#: every row of the accuracy sweep (see DESIGN_time_marginalization_peak_local.md). +PEAK_ENUM_FACTOR = 8 + +#: Localise each enumerated crest to this fraction of its own ``sigma_t``, and widen +#: its interval by the same amount. +#: +#: THIS IS THE THIRD RESOLUTION REQUIREMENT, and the first version of this module did +#: not have it. The design names two -- resolve ``kappa`` to ENUMERATE, resolve +#: ``exp(lnL)`` to INTEGRATE -- but enumeration returns a grid INDEX, and an index is +#: not a location. The true crest can be up to ``h_enum/2`` from the sample that +#: reports it, while the interval half-width is ``W_SIGMA * sigma_t``, so whenever +#: +#: W_SIGMA * sigma_t < h_enum / 2 +#: +#: the peak falls entirely OUTSIDE its own interval. At ``PEAK_ENUM_FACTOR = 8`` and +#: ``W_SIGMA = 12`` that is ``sigma_t/deltaT < 0.0052``, i.e. any row whose derived +#: factor reaches 512. MEASURED on the synthetic fixture at ``sigma_t/deltaT = +#: 0.0024``, error against the dense path as the crest is walked off the enumeration +#: grid: **0.00 nats at offset 0, -6.52 at h_enum/4, -164.93 at h_enum/2**. Always +#: negative -- it silently deletes mass -- and the old tail bound reported -275 nats +#: while dropping 165. +#: +#: Unlike the other two requirements this one is SNR-DEPENDENT: the crest must be found +#: to a fraction of ``sigma_t``, and ``sigma_t`` shrinks as ``1/rho``. It is met by +#: Newton on the spectral interpolant (see :func:`localise_peaks`), which converges +#: quadratically, so its cost grows only logarithmically. +#: +#: Value: 0.25 puts the residual at ``sigma/4``, so widening by it costs 2% of the +#: interval, while the trapezoid error it can induce is ``exp(-(12)^2/2)``-scale -- +#: nothing. It is not an accuracy knob that can be set too small: convergence to this +#: tolerance is ASSERTED, and a row that misses it goes to the dense path. +LOCALISE_SAFETY = 0.25 + +#: Newton iterations allowed per peak. From a parabolic seed the initial error is a few +#: percent of ``h_enum`` and convergence is quadratic, so 2-4 is typical; the cap exists +#: so a pathological row FAILS (and is handed to the dense path) rather than spinning. +LOCALISE_MAX_ITER = 16 + +#: A row is accepted only if its local integration grids actually ATTAIN the localised +#: crest value, to within this many nats. The a-posteriori half of the F1 fix, and it +#: costs nothing -- the maximum over each local grid is already computed for the +#: log-sum-exp offset. +#: +#: Why the slack is safe and why it is small: the local spacing is at most +#: ``sigma/UPSAMPLE_SAFETY = sigma/2``, so the grid's nearest point to the crest is +#: within ``sigma/4``, i.e. its ``lnL`` is within ``1/32`` of the crest's. 0.5 nats is +#: a 16x margin on that, and is still 13x smaller than the smallest miss F1 produced +#: (-6.52 nats). +#: +#: NOTE this compares against ``lnL`` at the LOCALISED crest, not at the enumeration +#: sample. Comparing against the sample cannot work and it is worth saying why: when +#: the crest sits off-grid the sample is already tens of nats below it, so a check +#: against the sample passes precisely in the case it is meant to catch. +#: +#: SCOPE, stated because it is narrower than it looks: ``row_star`` is a per-row MAXIMUM +#: over the peaks that survived the keep filter. So this verifies that the DOMINANT +#: crest was reached; it does not independently verify each secondary crest, and a peak +#: dropped by the keep filter can never enter it at all. That is tolerable only because +#: the keep filter now compares crests (G1) and its own magnitude argument stands on its +#: own -- not because this check covers it. +CONTAINMENT_SLACK_NATS = 0.5 + +#: Local interval half-width, in units of the peak's own ``sigma_t``. A Gaussian +#: peak truncated at ``W_SIGMA`` sigma omits ``erfc(W/sqrt2) ~ exp(-W^2/2)`` of its +#: mass: ``exp(-72) = 5.4e-32`` here, which is below double precision against the +#: largest window-to-sigma dynamic range this path can see (``UPSAMPLE_FACTOR_MAX`` +#: bounds it at ~1e4). +#: +#: ``erfc(W/sqrt2)`` is a statement about a Gaussian truncated symmetrically ABOUT ITS +#: CREST, so it is only a truncation bound if the interval is actually centred there. +#: The first version of this module centred on the enumeration SAMPLE, which left an +#: unstated precondition ``W_SIGMA * sigma_t >= h_enum/2`` -- coupling ``W_SIGMA`` to +#: ``PEAK_ENUM_FACTOR``, violated by every sharp row, and nowhere asserted. It is +#: discharged, not asserted: :func:`localise_peaks` finds the crest to +#: ``LOCALISE_SAFETY * sigma_t`` and the interval is widened by that residual, so the +#: two constants are decoupled and raising ``PEAK_ENUM_FACTOR`` is no longer the +#: dangerous move it used to be. +W_SIGMA = 12.0 + +#: Enumerated peaks more than this far below a row's highest CREST -- not its highest +#: sample; see the G1 note in ``_peak_local_chunk`` -- are dropped before intervals are +#: built. +#: +#: The justification is DIRECT and does not go through the tail bound. An earlier +#: version of this docstring said a dropped peak was safe "because it enters the tail +#: bound", and that claim is vacuous: since ``q_out_max`` is at least the dropped peak's +#: own sample, the bound accepts the row unless ``log(T_out / 2.5 sigma) >= 37``, i.e. +#: ``sigma_t < 2.6e-18 s``, while ``UPSAMPLE_FACTOR_MAX`` bounds the sharpest legal row +#: at ``sigma_t = 3.0e-08 s`` -- ten orders of magnitude away. For EVERY row this +#: module can legally handle, a keep-filter drop is automatically accepted by the bound. +#: The bound cannot backstop this filter and must not be cited as if it could. +#: +#: What justifies it instead is magnitude, crest to crest: a peak ``PEAK_KEEP_NATS`` +#: below the highest crest contributes ``exp(-60) = 8.8e-27`` of its mass per unit +#: width, and the widest window-to-sigma ratio this module can reach is ``T/sigma_min +#: = 2.5e6`` (again from ``UPSAMPLE_FACTOR_MAX``), so the omitted relative mass is below +#: ``2e-20`` -- under double precision. That inequality is asserted in the suite, +#: because it ties this constant to ``UPSAMPLE_FACTOR_MAX`` and neither is free. +PEAK_KEEP_NATS = 60.0 + +#: A row is accepted only if ``log(T_outside) + max_{outside} lnL - result`` is below +#: this, i.e. the bounded omitted mass is under ``e^-23 = 1e-10`` of the value +#: reported. A relative error ``eps`` in the integral is an ABSOLUTE error ``eps`` +#: in the returned log, so this is 1e-10 nats -- seven orders below the ~1e-3 nats +#: scale at which a difference in this quantity means anything. Rows that fail are +#: handed to the dense band-limited path, not reported with a caveat. +TAIL_LOG_TOL = -23.0 + +#: Ceiling on the number of DISJOINT intervals a row may have after merging. This +#: is a fail-closed cost guard, not an accuracy one: a row with more structure than +#: this is not approximated, it is sent to the dense path. (In the degenerate +#: many-overlapping-peaks case merging collapses the count, so this bites only on a +#: genuinely comb-like integrand.) +MAX_INTERVALS = 32 + +#: Re-anchor the phase recurrence in ``eval_bandlimited_uniform`` every this many +#: steps. Purely numerical: ``exp(i x)`` is not exactly unit modulus in floating +#: point, so ``z**m`` by repeated multiplication drifts as ``m * eps``. Re-anchoring +#: bounds that at ``64 * eps ~ 1e-14`` regardless of how long the local grid is. It +#: cannot change the answer beyond rounding and is not a tunable. +_RECURRENCE_REANCHOR = 64 + +#: Working-set budget for one dense temporary, in bytes. Internal memory chunking over +#: the extrinsic axis. Rows are independent and the per-row plan -- interval count and +#: point count, which is bucketed to a power of two -- depends only on that row, so this +#: cannot change WHICH rule a row gets or how finely it is integrated. +#: +#: It does move the last bits, and saying otherwise would be an overclaim: chunking +#: changes the leading dimension of the FFT and of the reduction inside +#: :func:`eval_bandlimited_uniform`, and both numpy's FFT and its pairwise summation +#: reassociate with batch shape. MEASURED at one row per chunk versus all rows at once, +#: on a three-row block spanning rho ~ 100 to 700: 0, 0 and 2 ULPs. +#: ``test_the_memory_chunking_path_assembles_its_result`` pins that, at a bar sharp +#: enough that a dropped or mis-ordered chunk -- which moves a row by nats -- cannot hide +#: behind it. +_CHUNK_BYTES = 128 * 1024 * 1024 + +_LAST_REPORT = {} + + +def last_report(): + """Diagnostics from the most recent :func:`time_marginalize_peak_local` call. + + Shares the row-classification keys of + :func:`time_marginalization_quadrature.last_report` -- ``n_rows``, + ``n_wrap_exposed_rows``, ``n_unmeasurable_rows``, ``n_flat_rows``, + ``n_refined_rows`` -- which mean exactly what they mean there, and adds: + + ``n_peak_local_rows`` rows actually integrated by this module's rule. + ``n_dense_fallback_rows`` refined rows handed to the dense band-limited path. + ``n_dense_fallback_cost_pregate`` / ``n_dense_fallback_cost`` cost declines, split + by WHICH gate fired: before enumeration (from a point-count floor) and after it + (from the merged intervals). Kept apart because only the first can prevent work + being done, and a single shared counter made the pre-gate removable without any + test noticing. + ``n_dense_fallback_nopeak`` rows that passed the gates but enumerated no usable + peak. Exists so the sub-counts RECONCILE with ``n_dense_fallback_rows``. + ``n_dense_fallback_tail`` of those, how many went because the omitted-mass + bound was not small enough. **This is the count to watch**: it is the + method admitting it could not justify its own truncation, and a run where it + is not ~0 is a run where the enumeration is not doing its job. + ``n_dense_fallback_structure`` rows exceeding ``MAX_INTERVALS`` or whose sub-cell + certificate geometry contains more than one disjoint interval in a cell. + ``n_intervals_total`` / ``n_local_points_total`` the work actually done. + ``n_peaks_total`` enumerated maxima kept, over the peak-local rows. + ``tail_bound_worst`` the worst (largest) ``bound - result`` among ACCEPTED + rows, in nats. A number near ``TAIL_LOG_TOL`` means the truncation is only + just being justified. + """ + return dict(_LAST_REPORT) + + +# --------------------------------------------------------------- the evaluator + +def bandlimited_spectrum(x, xpy=np): + """Spectral representation of the band-limited interpolant through rows of ``x``. + + Returns ``(Xw, fk)`` with + + x(t) = sum_j Xw[..., j] * exp(2 pi i * fk[j] * t / (n * deltaT)) + + exact at ``t = i * deltaT`` for every integer ``i``, and equal to the unique + band-limited interpolant that :func:`bandlimited_upsample` evaluates on a + uniform refinement -- so the two agree wherever both are defined, which + ``test_local_evaluator_matches_the_dense_upsample`` asserts on every production + ``npts``. + + ODD ``n`` IS THE COMMON CASE and is the trap this function shares with + :func:`bandlimited_upsample`: ``marginalization_time_grid(0.075, 1/srate)`` + returns 153 / 307 / 614 / 1228 / 2457 at srate 1024 / 2048 / 4096 / 8192 / 16384, + odd at three of five. The positive-frequency block is ``0 .. (n-1)//2``; putting + the top positive bin at a negative frequency (which a split at ``n//2`` does for + odd ``n``) leaves the reconstruction EXACT AT THE SAMPLES and wrong everywhere + between them, so a "reproduces its input" test cannot see it. + + For even ``n`` the Nyquist bin is genuinely ambiguous -- ``+fNyq`` and ``-fNyq`` + are the same sequence on the samples -- so it is split evenly between the two, + matching :func:`bandlimited_upsample`. No interpolant can recover which it was; + for rholm data the bin is empty anyway. + """ + n = x.shape[-1] + X = xpy.fft.fft(x, axis=-1) / float(n) + n_pos = (n - 1) // 2 + pos = np.arange(0, n_pos + 1) + if n % 2 == 0: + neg = np.arange(n_pos + 2, n) - n + fk = np.concatenate([pos, [n // 2], [-(n // 2)], neg]).astype(np.float64) + half = 0.5 * X[..., n_pos + 1:n_pos + 2] + Xw = xpy.concatenate([X[..., :n_pos + 1], half, half, X[..., n_pos + 2:]], + axis=-1) + else: + neg = np.arange(n_pos + 1, n) - n + fk = np.concatenate([pos, neg]).astype(np.float64) + Xw = X + return Xw, xpy.asarray(fk) + + +def eval_bandlimited_uniform(Xw, fk, t0, dt_local, n_local, period, xpy=np): + """Evaluate the interpolant on a per-row uniform grid ``t0[r] + m * dt_local[r]``. + + ``Xw, fk`` come from :func:`bandlimited_spectrum`; ``t0`` and ``dt_local`` are + per-row (shape ``(n_rows,)``); ``period`` is ``npts * deltaT``. Returns + ``(n_rows, n_local)`` complex. + + THIS IS THE COST THE PROTOTYPE DID NOT PAY. With the analytic ``kappa`` in hand + a local time costs one closed-form call; here it costs a sum over the spectrum. + The loop below is ``n_local`` steps of one complex multiply and one reduction + over ``(n_rows, n_freq)``, i.e. ``O(n_rows * npts * n_local)`` and -- the property + that makes the method worth having -- INDEPENDENT OF RHO, because ``n_local`` is + fixed by ``W_SIGMA`` and ``UPSAMPLE_SAFETY`` while the dense path's point count + grows linearly with it. + + No transcendental is evaluated per output point: the grid is uniform, so the + phase advances by a constant factor. It is re-anchored every + ``_RECURRENCE_REANCHOR`` steps because ``exp(i x)`` is not exactly unit modulus. + """ + n_rows = Xw.shape[0] + two_pi_i = 2j * np.pi + scale = fk[None, :] / float(period) + z = xpy.exp(two_pi_i * scale * dt_local[:, None]) + out = xpy.empty((n_rows, n_local), dtype=Xw.dtype) + acc = None + for m in range(n_local): + if acc is None or (m % _RECURRENCE_REANCHOR) == 0: + acc = Xw * xpy.exp(two_pi_i * scale * (t0 + m * dt_local)[:, None]) + out[:, m] = xpy.sum(acc, axis=-1) + acc = acc * z + return out + + +def eval_bandlimited_points(Xw, fk, rows, t, period, xpy=np, point_chunk=1024): + """``q``, ``q'`` and ``q''`` at arbitrary ``(row, time)`` pairs. + + The uniform-grid evaluator above cannot reach these points: the ones that matter are + OFF THE ENUMERATION GRID -- the ends of the merged intervals, wherever localisation + put them. Cost is one exponential array per point and ``O(npts)`` per point. + + NOT ON THE SHIPPED PATH, and kept deliberately. It WAS the production evaluator for + the omitted-mass bound, until :func:`segment_sup_bound` replaced pointwise evaluation + with a sub-cell Hermite certificate that needs values only on the enumeration grid. + What it still does is keep two claims from being circular, which is why it is exported + rather than deleted: + + * it is the INDEPENDENT route :func:`enum_grid_derivatives` is checked against. That + helper arrived 52% wrong on ``q'``, and a check built from the same upsampler it + now reuses could not have caught it. + * it evaluates the merged-interval ends EXACTLY, which the tail-bound test needs. A + uniform reference grid will not do there: on a sharp row its spacing is comparable + to ``sigma`` and it under-reads an end by ~11 nats, which makes a correct bound look + loose. + + Anything that gives it a production caller again should say so here, because "exported, + tested, and called only by tests" is otherwise indistinguishable from dead code. + + Same three sums as :func:`localise_peaks` -- the derivatives are the spectral sum + with ``w_j`` and ``w_j**2`` folded in -- and chunked over points for the same + reason: the temporary is ``(n_points, n_freq)``. + """ + w = (2j * np.pi / float(period)) * fk + n_pt = int(t.shape[0]) + q0 = xpy.zeros(n_pt, dtype=np.float64) + q1 = xpy.zeros(n_pt, dtype=np.float64) + q2 = xpy.zeros(n_pt, dtype=np.float64) + for a in range(0, n_pt, point_chunk): + b = min(a + point_chunk, n_pt) + E = Xw[rows[a:b]] * xpy.exp(w[None, :] * t[a:b][:, None]) + q0[a:b] = xpy.sum(E, axis=-1).real + q1[a:b] = xpy.sum(E * w[None, :], axis=-1).real + q2[a:b] = xpy.sum(E * (w * w)[None, :], axis=-1).real + return q0, q1, q2 + + +def enum_grid_derivatives(x_reflected, factor, n_keep, deltaT, orders=(1, 2), + xpy=np): + """``q'`` and ``q''`` on the ENUMERATION grid, from the same reflected row. + + Differentiation is a spectral multiply, so the derivative of the band-limited + interpolant is the interpolant OF the differentiated row -- which means the validated + upsampler can be reused rather than its zero-padding convention reimplemented. That + matters: the first version of this reimplemented it and got 52% relative error on ``q'`` + and 46% on ``q''``, because for the reflected row ``n`` is EVEN, so + :func:`bandlimited_spectrum` splits the Nyquist bin and returns ``n+1`` coefficients -- + a length the padding logic then used as the transform size, mapping two bins onto one + index and losing the rest. + + Differencing ``q`` on the grid would not do either: a difference of an under-resolved + sample sequence is an estimate, and everything downstream of these arrays is an + inequality. + """ + n = int(x_reflected.shape[-1]) + k = xpy.fft.fftfreq(n, d=1.0) * n # signed bin index, matching numpy's fft + w = (2j * np.pi / (n * float(deltaT))) * k + X = xpy.fft.fft(x_reflected, axis=-1) + out = [] + for order in orders: + row = xpy.fft.ifft(X * (w ** order)[None, :], axis=-1) + out.append(bandlimited_upsample(row, factor, xpy=xpy)[..., :n_keep].real) + return tuple(out) + + +def parabolic_sup(y0, y1, d0, d1, s_lo=0.0, s_hi=1.0, xpy=np): + """``max`` of the cubic Hermite through ``(0, y0, d0)`` and ``(1, y1, d1)``, over the + SUB-RANGE ``[s_lo, s_hi]`` of the cell (default: the whole cell). + + ``d0``/``d1`` are the slopes ALREADY SCALED BY THE CELL WIDTH, i.e. ``h * q'``. The maximum + of a cubic on a closed range is at an end or at a stationary point inside it, so this is + exact -- no search and no iteration. + + The sub-range is what lets a cell that a merged interval CUTS be bounded on the part that + is actually outside. A range with ``s_hi <= s_lo`` is empty and returns ``-inf``, so empty + pieces fall out of a maximum without needing to be filtered. + """ + a = 2.0 * y0 + d0 - 2.0 * y1 + d1 + b = -3.0 * y0 - 2.0 * d0 + 3.0 * y1 - d1 + c = d0 + empty = s_hi <= s_lo + + def _H(t): + return y0 + c * t + b * t * t + a * t ** 3 + + best = xpy.maximum(_H(s_lo), _H(s_hi)) + + def _try(root, live): + # Evaluate only inside the requested range. This avoids overflow from the + # irrelevant enormous root of a nearly-linear quadratic on eager array backends. + root_eval = xpy.clip(root, s_lo, s_hi) + val = _H(root_eval) + return xpy.where(live & (root > s_lo) & (root < s_hi), + xpy.maximum(best, val), best) + + # H'(s) = 3a s^2 + 2b s + c. The CUBIC term vanishes whenever the cell's slopes and its + # secant conspire -- a symmetric bump does it exactly -- so the degenerate branch is not an + # edge case to skip: missing it returns the endpoint maximum and under-bounds precisely the + # cells that contain a peak. + A = 3.0 * a + B = 2.0 * b + # The roots are invariant under a common coefficient scale. Normalize before forming + # the discriminant so finite coefficients cannot produce ``inf - inf = NaN`` and make + # genuine stationary points disappear from a purported upper bound. + coef_scale = xpy.maximum(xpy.maximum(xpy.abs(A), xpy.abs(B)), xpy.abs(c)) + live_poly = coef_scale > 0.0 + safe_scale = xpy.where(live_poly, coef_scale, 1.0) + An, Bn, Cn = A / safe_scale, B / safe_scale, c / safe_scale + cubic = xpy.abs(An) > 0.0 + disc = Bn * Bn - 4.0 * An * Cn + sq = xpy.sqrt(xpy.where(cubic & (disc > 0), disc, 0.0)) + # Cancellation-safe quadratic roots. The direct ``(-b +/- sq)/(3a)`` loses the + # in-range root of a nearly quadratic Hermite cell when ``b`` and ``sq`` agree. This + # occurs naturally for a symmetric band-limited crest: ``a`` should vanish, but the + # endpoint/derivative arithmetic leaves a few ulps behind. Form the large root from the + # non-cancelling sign and the other from the product of the roots, ``c/A``. + q = -0.5 * (Bn + xpy.copysign(sq, Bn)) + q_live = cubic & (disc > 0) & (xpy.abs(q) > 0.0) + root_large = q / xpy.where(cubic, An, 1.0) + root_small = Cn / xpy.where(q_live, q, 1.0) + best = _try(root_large, cubic & (disc > 0)) + best = _try(root_small, q_live) + lin = (~cubic) & (xpy.abs(Bn) > 0.0) + best = _try(-Cn / xpy.where(lin, Bn, 1.0), lin) + return xpy.where(empty, -np.inf, best) + + +def _certificate_acceptance_masks(margin, contained, cert_bad, planned): + """Partition planned rows into one accepted class and three disjoint fallback reasons.""" + margin_ok = margin[planned] < TAIL_LOG_TOL + structure_fail = cert_bad[planned] + tail_fail = (~structure_fail) & (~margin_ok) + containment_fail = (~structure_fail) & margin_ok & (~contained[planned]) + good = (~structure_fail) & margin_ok & contained[planned] + return good, structure_fail, tail_fail, containment_fail + + +def segment_sup_bound(q0, q1, dq0, dq1, h, m4, s_lo=0.0, s_hi=1.0, xpy=np): + """CERTIFIED upper bound on ``max q`` over one enumeration cell. + + Cubic Hermite through the cell's two endpoint values and slopes, plus the classical + Hermite remainder ``M4 * h**4 / 384`` where ``M4`` bounds the FOURTH derivative. ``m4`` + comes from :func:`spectral_derivative_bound` at order 4, so it is a true bound and nothing + is fitted. + + WHY THE SLOPES ARE WORTH FETCHING. A bound from endpoint VALUES alone carries + ``M2 * h**2 / 8``, and from a single sample ``M2 * h**2 / 2``. Measured on a realistic row + at amplitude 2e4 / 2e6 / 2e7 those are 102 / 1.0e4 / 1.0e5 nats and 407 / 4.1e4 / 4.1e5 -- + useless against a 23 nat tolerance, which is why an earlier revision of this note called + certification non-viable. With the slopes the remainder is **0.12 / 12.2 / 122 nats**, + three thousand times tighter, and usable across the range that matters. Above it the + certificate simply stops being small enough, the margin fails, and the row falls back -- + the intended fail-closed behaviour, not a special case. + """ + return (parabolic_sup(q0, q1, dq0, dq1, s_lo=s_lo, s_hi=s_hi, xpy=xpy) + + m4 * (h ** 4) / 384.0) + + +# -------------------------------------------------------------- enumeration + +def spectral_derivative_bound(Xw, fk, period, order, xpy=np): + """``max_t |q^(order)(t)|`` for the interpolant, bounded rather than estimated. + + ``q(t) = Re sum_j Xw_j exp(w_j t)`` with ``w_j = 2 pi i fk_j / period``, so + ``|q^(k)| <= sum_j |Xw_j| |w_j|^k`` everywhere, by the triangle inequality. One + reduction over the spectrum per row; nothing is fitted and no shape is assumed. + """ + w = xpy.abs(2.0 * np.pi * xpy.asarray(fk) / float(period)) ** int(order) + return xpy.sum(xpy.abs(Xw) * w[None, :], axis=-1) + + +def spectral_curvature_bound(Xw, fk, period, xpy=np): + """``max_t |q''(t)|`` for the interpolant, bounded rather than estimated. + + The ``order = 2`` case of :func:`spectral_derivative_bound`, kept under its own + name because the crest pre-filter is the one caller that must not be read as + depending on anything else. + """ + return spectral_derivative_bound(Xw, fk, period, 2, xpy=xpy) + + +def crest_upper_bound(q_at_peak, q_ddot_max, h_enum): + """Upper bound on a crest, from its enumeration SAMPLE and a bound on ``|q''|``. + + Expanding about the crest ``t*``, where ``q'`` vanishes by definition, + + q(t_s) = q(t*) + q''(xi) (t_s - t*)^2 / 2, |t_s - t*| <= h_enum + + so ``q(t*) <= q(t_s) + q_ddot_max * h_enum^2 / 2``. A Taylor remainder with a TRUE + bound on the second derivative -- not a parabolic fit through three samples, which is + what the previous version of this was and which is why it was not a bound: ``lnL`` is + not quadratic across an enumeration cell, and the anharmonic part of the deficit + carries the same ``1/sigma^2`` amplification as the quadratic part. Measured at + derived factor 1024, the pure quantisation excess is 4.4 nats and the true shortfall + was 122.3. + + ``h_enum``, not ``h_enum/2``: the localiser's bracket is ``+/- h_enum`` and + displacements up to 0.959 of it have been observed. + """ + return q_at_peak + 0.5 * q_ddot_max * h_enum ** 2 + + +def enumerate_peak_indices(q, xpy=np): + """Boolean mask of INTERIOR local maxima of each row of ``q``. + + ``q`` is ``Re kappa`` on the enumeration grid, NOT ``lnL``. Every callback in + this path's scope is monotone increasing in it at fixed ``rho_sq``, and + ``rho_sq`` is time-independent here, so a monotone map cannot move a maximum and + the extrema of ``lnL`` are exactly the extrema of ``Re kappa``. Enumerating here + rather than on ``lnL`` keeps the callback -- a table interpolation, for the + distance-marginalized case -- off the full time axis entirely. (Phase + marginalization would make the relevant quantity ``|kappa|``, which peaks + elsewhere; it is refused by this path, see the module docstring.) + + The two comparisons are deliberately asymmetric (``>=`` left, ``>`` right): a + plateau then yields exactly one index, its last, instead of none or all of them. + + ENDPOINTS ARE INCLUDED, and the reason is worth recording because an earlier version + excluded them on a justification that was simply false. That version said a maximum + at the window edge means the window is mis-centred and the inherited edge guard has + already routed such rows away. It has not: ``exposed`` keys on the row's GLOBAL + coarse argmax, so a row whose dominant peak sits comfortably mid-window is refined + even when it also carries a SECONDARY maximum hard against an edge. That secondary + peak was then never enumerated, never covered, and -- because the outside maximum is + sampled on this same grid -- under-represented in the tail bound too. + """ + interior = (q[..., 1:-1] >= q[..., :-2]) & (q[..., 1:-1] > q[..., 2:]) + left = (q[..., :1] > q[..., 1:2]) + right = (q[..., -1:] > q[..., -2:-1]) + return _cat_last(left, interior, right, xpy=xpy) + + +def _cat_last(*parts, **kw): + xpy = kw.get('xpy', np) + return xpy.concatenate(parts, axis=-1) + + +def localise_peaks(Xw, fk, rows, t_grid, h_enum, tol, period, xpy=np, + peak_chunk=4096, t_last=None): + """Newton on the band-limited interpolant: turn a grid INDEX into a LOCATION. + + ``t_grid`` are the enumeration-grid times of the enumerated maxima and ``rows`` says + which row each belongs to. Returns ``(t_star, q_star, converged)``. + + An enumerated extremum is a grid index, and the crest it stands for can be anywhere + within ``+/- h_enum/2`` of it. Building the integration interval around the index + instead of the crest is what made this module drop up to 165 nats -- see + :data:`LOCALISE_SAFETY` for the measured table. This is the missing step. + + It is NOT a seeded root-finder in the sense that sank RIFT PR #201. That failure + was Newton seeded at GUESSED points, used to FIND extrema, so the ones it did not + guess were never found. Here every extremum has already been enumerated, and Newton + only refines a location inside the bracket ``[t_grid - h_enum, t_grid + h_enum]`` + that the enumeration already established. It cannot discover or lose a peak; it can + only place one. An iterate that leaves the bracket, or a peak that does not reach + ``tol``, is reported as NOT converged and its row goes to the dense path. + + ``q``, ``q'`` and ``q''`` come from the spectral representation directly -- + ``q(t) = Re sum_j Xw_j exp(w_j t)`` with ``w_j = 2 pi i f_j / T``, so the + derivatives are the same sum with ``w_j`` and ``w_j^2`` folded in and cost one + exponential array between them. Convergence is quadratic, so the SNR-dependence of + this step is logarithmic: 2-4 iterations over the whole production range. + """ + n_pk = int(t_grid.shape[0]) + w = (2j * np.pi / float(period)) * fk + t_out = xpy.zeros(n_pk, dtype=np.float64) + q_out = xpy.zeros(n_pk, dtype=np.float64) + ok_out = xpy.zeros(n_pk, dtype=bool) + + for a in range(0, n_pk, peak_chunk): + b = min(a + peak_chunk, n_pk) + Xr = Xw[rows[a:b]] # (P, nf) + tg = t_grid[a:b] + tol_c = tol[a:b] + t = tg.copy() + step = xpy.zeros(t.shape, dtype=np.float64) + for _ in range(LOCALISE_MAX_ITER): + E = Xr * xpy.exp(w[None, :] * t[:, None]) + q1 = xpy.sum(E * w[None, :], axis=-1).real + q2 = xpy.sum(E * (w * w)[None, :], axis=-1).real + # Only a strictly concave point is a maximum to walk towards. A + # non-concave iterate is left where it is and reported unconverged. + safe = q2 < 0 + step = xpy.where(safe, -q1 / xpy.where(safe, q2, -1.0), 0.0) + step = xpy.clip(step, -h_enum, h_enum) + t_new = xpy.clip(t + step, tg - h_enum, tg + h_enum) + if t_last is not None: + # The crest of a peak enumerated at an endpoint can lie OUTSIDE the + # window; within it the maximum is then the boundary itself. Clamping + # here keeps the reported crest inside the integration domain, which is + # what the interval is built around. + t_new = xpy.clip(t_new, 0.0, t_last) + step = t_new - t + t = t_new + if bool(xpy.all(xpy.abs(step) <= tol_c)): + break + E = Xr * xpy.exp(w[None, :] * t[:, None]) + q0 = xpy.sum(E, axis=-1).real + q2 = xpy.sum(E * (w * w)[None, :], axis=-1).real + inside = xpy.abs(t - tg) <= h_enum + if t_last is not None: + pinned = (t <= 0.0) | (t >= t_last) + else: + pinned = xpy.zeros(t.shape, dtype=bool) + t_out[a:b] = t + q_out[a:b] = q0 + # A PINNED PEAK IS NOT CONVERGED, and calling it converged cost up to +4.38 nats. + # + # The crest of a peak enumerated at an endpoint can lie outside the window; the + # maximum over the domain is then the boundary itself, and it was tempting to + # call that "found". It is found -- but it is not a PEAK, and everything + # downstream assumes it is. At an interior crest `q'` vanishes, so `exp(lnL)` + # is locally Gaussian and a spacing derived from the curvature resolves it. At a + # boundary the maximum is a CORNER: `q'` is non-zero there, so the local integrand + # is an exponential decay of rate `|q'|`, and the spacing derived from `sigma` + # does not resolve that scale at all. The trapezoid's half-weight endpoint then + # over-counts by + # + # log( lam * (0.5 + 1/(exp(lam) - 1)) ) -> log(lam/2), lam = |q'| * h_loc + # + # which GROWS WITHOUT BOUND in amplitude -- +log 2 per factor 4. MEASURED on a + # single band-limited bump centred at t = 0, against the same interpolant + # integrated on the very interval the module chose: **+1.27 / +1.96 / +2.66 / + # +3.35 / +4.03 nats** at rho-scale 4e4 / 1.6e5 / 6.4e5 / 2.56e6 / 1e7, and +4.38 + # at the ceiling. All ACCEPTED, and both a-posteriori defences are silent BY + # CONSTRUCTION: the containment check compares the grid's attained maximum against + # `row_star`, and the grid's FIRST POINT is the pinned crest, so `attained == + # row_star` identically and it can never fire; the tail bound is a statement about + # mass OUTSIDE the intervals and the error is entirely inside. + # + # So the row is declined and handed to the dense path. That is the fail-closed + # direction and it restores the contract this rule actually makes -- never a worse + # value than the backstop. (The dense path carries a milder form of the same + # defect, +0.77 to +2.74 nats on these rows, because its grid also begins at the + # boundary with a half weight; that is a defect of the shared reconstruction and + # is recorded in the design note, not papered over here.) + ok_out[a:b] = inside & (xpy.abs(step) <= tol_c) & (q2 < 0) & (~pinned) + return t_out, q_out, ok_out + + +def merge_intervals_by_row(rows, lo, hi, span): + """Merge per-row interval lists into disjoint intervals, ALL ROWS AT ONCE. + + ``rows`` must be ascending; ``lo``/``hi`` are the interval ends and ``span`` an + upper bound on ``hi``. Returns ``(order, gid, g_row, g_lo, g_hi)``: the sort + permutation, the merged-interval index each sorted input fell into, and the row, + start and stop of each merged interval. ``gid`` is what lets per-peak quantities + -- the widths that set the local spacing -- be reduced over the merge. + + MERGING IS NOT AN OPTIMISATION. Two overlapping windows integrated separately + both contain the shared region, so the log-sum-exp of the parts double-counts it. + On a broad integrand at rho ~ 6 that is +1.6 nats -- measured, and reproduced by + ``test_unmerged_intervals_double_count``. It is also what makes this ONE + algorithm rather than a regime switch: as peaks crowd together the union grows + continuously to the whole window and the method becomes the dense grid, with no + threshold anywhere. + + Vectorised across rows deliberately, not for elegance: a broad integrand has MANY + enumerated peaks per row, and a Python loop over them made this rule slower than + the dense path it delegates to (measured 0.69x at sigma_t/deltaT = 0.17, on a + block where every row fell back anyway). The running maximum is restarted at + each row boundary by offsetting row ``r`` by ``r * big`` -- every value in row + ``r-1`` is then below every value in row ``r``, so one global + ``maximum.accumulate`` gives the per-row running maximum with no segmentation. + """ + order = np.lexsort((lo, rows)) + r_s, lo_s, hi_s = rows[order], lo[order], hi[order] + big = 2.0 * (float(span) + 1.0) + cummax = np.maximum.accumulate(hi_s + r_s * big) - r_s * big + new = np.ones(r_s.size, dtype=bool) + new[1:] = (r_s[1:] != r_s[:-1]) | (lo_s[1:] > cummax[:-1]) + gid = np.cumsum(new) - 1 + heads = np.nonzero(new)[0] + tails = np.append(heads[1:] - 1, r_s.size - 1) + return order, gid, r_s[heads], lo_s[heads], cummax[tails] + + +# ------------------------------------------------------------------ the rule + +def _peak_curvature_sigma(lnL_stencil, h, xpy=np): + """``sigma`` per peak from a widening centred second difference of ``lnL``. + + ``lnL_stencil`` is ``(n_peaks, 2*maxd+1)``, the callback evaluated on the + enumeration grid at offsets ``-maxd .. +maxd`` about each peak. Same estimator + and same widening rationale as + :func:`time_marginalization_quadrature.peak_width_from_lnL`: the second + difference of a parabola is its second derivative at ANY spacing, so an + under-resolved peak still reports its own width honestly, and stepping the + stencil out over a ``-inf`` hole (the distance-marginalized callback returns + ``-inf`` outside its table) costs nothing. A peak with no finite negative + curvature at any half-width gets ``inf`` and is not given an interval -- which + routes it into the tail bound rather than into a guess. + """ + maxd = (lnL_stencil.shape[-1] - 1) // 2 + centre = lnL_stencil[:, maxd] + sigma = xpy.full(lnL_stencil.shape[0], np.inf, dtype=np.float64) + done = xpy.zeros(lnL_stencil.shape[0], dtype=bool) + for d in CURVATURE_STENCIL_HALFWIDTHS: + if d > maxd: + break + with np.errstate(invalid='ignore'): + d2 = (lnL_stencil[:, maxd - d] - 2.0 * centre + + lnL_stencil[:, maxd + d]) / float(d * h) ** 2 + fresh = xpy.isfinite(d2) & (~done) + neg = fresh & (d2 < 0) + sigma = xpy.where(neg, 1.0 / xpy.sqrt(xpy.where(neg, -d2, 1.0)), sigma) + done = done | fresh + if bool(xpy.all(done)): + break + return sigma + + +#: Fewest local points any row can ever need: one interval of half-width +#: ``W_SIGMA * sigma`` at spacing ``sigma / UPSAMPLE_SAFETY`` is +#: ``2 * W_SIGMA * UPSAMPLE_SAFETY`` sub-intervals however large ``sigma`` is. Used to +#: reject a row BEFORE any work is done for it. +#: +#: It is NOT a strict lower bound, and calling it one was wrong: an interval clipped by a +#: window end carries fewer points than this. Such a row is wrap-exposed or nearly so and +#: has almost always been routed away by the edge guard already -- but "almost always" is +#: not "always", so this is a COST heuristic and nothing more. It cannot affect a +#: returned value: a row it wrongly declines is computed by the dense path instead. +MIN_LOCAL_POINTS = int(2 * W_SIGMA * UPSAMPLE_SAFETY) + 1 + + +def _estimated_costs(n_local_total, npts, factor): + """(local, dense) cost estimates for one row, in complex-multiply units. + + The local evaluator is a spectral sum advanced by a recurrence: one complex + multiply and one add per (frequency, output point), hence ``npts * + n_local_total``. The dense path is a zero-padded FFT of length ``npts*factor``, + hence ``npts*factor*log2(npts*factor)``, and it additionally pays the likelihood + callback and an ``exp`` on every one of those points while the local path pays + them on ``n_local_total`` -- so this estimate is CONSERVATIVE against the local + path (it counts the arithmetic the two share and ignores the per-point cost the + local path avoids). + + It decides only WHICH of two paths runs, and both satisfy the same derived + resolution criterion, so it cannot trade accuracy for cost. + """ + dense_n = npts * np.asarray(factor, dtype=np.float64) + return (npts * np.asarray(n_local_total, dtype=np.float64), + dense_n * np.maximum(np.log2(dense_n), 1.0)) + + +def _log_trapz_local(lnL_loc, h, xpy=np): + """``log \\int exp(lnL) dt`` by trapezoid on one uniform local grid, per row. + + Trapezoid, not Simpson, for the reason the dense path gives: on a peak that has + decayed to nothing inside the interval every Euler-Maclaurin boundary term + vanishes and the trapezoidal rule is spectrally accurate, while Simpson's + ``(4 T_h - T_2h)/3`` reintroduces the ``2h`` alias that is the original defect. + The offset is per row and taken on this grid; the result is offset-invariant. + """ + w = xpy.full(lnL_loc.shape[-1], 1.0, dtype=np.float64) + w[0] = 0.5 + w[-1] = 0.5 + off = _safe_offset(xpy.max(lnL_loc, axis=-1, keepdims=True), xpy=xpy) + return off[..., 0] + xpy.log(xpy.sum(xpy.exp(lnL_loc - off) * w, axis=-1)) \ + + xpy.log(h) + + +def _logaddexp_reduce(parts, xpy=np): + """``log sum_j exp(parts[:, j])``, NaN-safe for an all ``-inf`` row. + + Rows with no intervals at all are all ``-inf`` here by construction -- they are + the ones being handed to the dense path -- so ``log(0)`` is the CORRECT answer for + them and its warning is noise, not a signal. The offset guard is what keeps that + case at ``-inf`` instead of ``NaN``. + """ + off = _safe_offset(xpy.max(parts, axis=-1, keepdims=True), xpy=xpy) + with np.errstate(divide='ignore'): + return off[..., 0] + xpy.log(xpy.sum(xpy.exp(parts - off), axis=-1)) + + +def time_marginalize_peak_local(kappa, rho_sq, deltaT, loglikelihood, + phase_marginalization=False, simps=None, + lnL_coarse=None, xpy=np, return_peaks=False): + """``log \\int dt exp(lnL(t))``, refining only around the enumerated peaks. + + Signature, preconditions, row classification and fallback semantics are those of + :func:`time_marginalization_quadrature.time_marginalize_bandlimited`, which see; + the parameters mean the same things and the same rows fall back to the caller's + Simpson rule for the same reasons. What changes is the rule applied to the rows + that ARE refined. + + ``return_peaks=True`` additionally returns a list, one entry per input row, of + ``(t_star, sigma)`` arrays for the peaks that were enumerated and kept -- + ``None`` for a row this rule did not handle. These are the same ``t_star`` a + time-first reordering of the marginalizations would need, and they are + distance- and callback-independent (see the module docstring), so they are + exposed as an output rather than kept as a temporary. + """ + if phase_marginalization: + # Deliberate scope cut, not an oversight -- see the module docstring. Refuse + # rather than silently running something else. + raise NotImplementedError( + "time_marginalize_peak_local does not support phase marginalization: the " + "Laplace width of the time peak then carries an (I1/I0)(|kappa|/D) factor " + "that does not reduce, so the local spacing is no longer derivable from " + "rho_sq and the curvature alone. Production marginalizes over distance, " + "not phase. Use time_quadrature='bandlimited', which does support it.") + + simps = _default_simps(simps, xpy) + + kappa = xpy.asarray(kappa) + rho_sq = xpy.asarray(rho_sq) + npts = kappa.shape[-1] + n_rows = kappa.shape[0] + deltaT = float(deltaT) + period = npts * deltaT + t_last = (npts - 1) * deltaT + + _require_time_independent_rho_sq(rho_sq, xpy=xpy, rule='peak-local') + rho_col = rho_sq[..., :1] + + _term = lambda k: k.real # phase marginalization is refused above + if lnL_coarse is None: + lnL_coarse = loglikelihood(_term(kappa), rho_sq) + + # ROW CLASSIFICATION IS THE DENSE PATH'S -- literally, by calling it. This module's + # contract is that it changes WHERE the refined grid is placed and nothing about + # WHICH rows get one, so any row `time_marginalize_bandlimited` refines must be a row + # peak-local refines. That was previously kept true by copying, and the copy went + # stale across the rebase onto rift_O4d in three separate clauses; see _classify_rows. + (sigma, jmax, measurable, has_peak, flat, exposed, unmeasurable, + factors) = _tmq._classify_rows(lnL_coarse, deltaT, npts, xpy=xpy) + refined = has_peak & (factors > 1) + + out = _log_simps_rows(lnL_coarse, deltaT, simps, xpy=xpy) + peaks_out = [None] * n_rows if return_peaks else None + + stats = dict(n_peak_local_rows=0, n_dense_fallback_cost=0, + n_dense_fallback_cost_pregate=0, n_dense_fallback_nopeak=0, + n_dense_fallback_tail=0, n_dense_fallback_structure=0, + n_dense_fallback_ceiling=0, n_dense_fallback_localise=0, + n_dense_fallback_containment=0, + n_intervals_total=0, n_local_points_total=0, n_peaks_total=0, + tail_bound_worst=-np.inf) + idx_all = np.asarray(xpy.where(refined)[0] if xpy is np + else xpy.where(refined)[0].get()) + n_dense = 0 + if idx_all.size: + per_row = npts * PEAK_ENUM_FACTOR * 16 * 4 + chunk = max(1, min(int(idx_all.size), int(_CHUNK_BYTES // max(per_row, 1)))) + dense_rows = [] + for start in range(0, int(idx_all.size), chunk): + sel = idx_all[start:start + chunk] + sel_x = xpy.asarray(sel) + vals, ok, peaks = _peak_local_chunk( + kappa[sel_x], rho_col[sel_x], factors[sel_x], npts, deltaT, period, + t_last, loglikelihood, _term, stats, xpy=xpy, want_peaks=return_peaks) + if ok.any(): + got = xpy.asarray(np.where(ok)[0]) + out[xpy.asarray(sel[ok])] = vals[got] + if return_peaks: + for j in np.where(ok)[0]: + peaks_out[int(sel[j])] = peaks[int(j)] + if (~ok).any(): + dense_rows.append(sel[~ok]) + if dense_rows: + di = np.concatenate(dense_rows) + di_x = xpy.asarray(di) + # The dense band-limited path is the BACKSTOP, deliberately: a row this + # rule declines is given a value from the reviewed reference + # implementation, not an approximation with a caveat attached. + out[di_x] = time_marginalize_bandlimited( + kappa[di_x], rho_sq[di_x], deltaT, loglikelihood, + phase_marginalization=phase_marginalization, simps=simps, + lnL_coarse=lnL_coarse[di_x], xpy=xpy) + n_dense = int(di.size) + + _LAST_REPORT.clear() + _LAST_REPORT.update( + n_rows=n_rows, + n_wrap_exposed_rows=int(xpy.sum(exposed)), + n_unmeasurable_rows=int(xpy.sum(unmeasurable)), + n_flat_rows=int(xpy.sum(flat)), + n_refined_rows=int(xpy.sum(refined)), + n_dense_fallback_rows=n_dense, + **stats) + if return_peaks: + return out, peaks_out + return out + + +def _host(a, xpy=np): + """Copy a device array to host numpy. The RAGGED bookkeeping -- which peaks + belong to which row, how their intervals merge -- is scalar work on a handful of + values per row, and doing it on the host keeps one implementation instead of two. + Everything whose size scales (the enumeration grid, the spectrum, the local + evaluation and the callback) stays on the device.""" + return np.asarray(a) if xpy is np else np.asarray(a.get()) + + +def _peak_local_chunk(kappa_rows, rho_col_rows, factors, npts, deltaT, period, + t_last, loglikelihood, _term, stats, xpy=np, want_peaks=False): + """Peak-local integration of one chunk of refined rows. + + Returns ``(values, ok, peaks)``; ``ok[r]`` is False for a row that must be given + to the dense path -- because it had no usable enumerated peak, because it had + more disjoint structure than ``MAX_INTERVALS``, because its estimated cost was + worse here than there, or because the omitted-mass bound could not be met. + """ + n_rows = kappa_rows.shape[0] + h_enum = deltaT / PEAK_ENUM_FACTOR + last = (npts - 1) * PEAK_ENUM_FACTOR + + values = xpy.full(n_rows, -np.inf, dtype=np.float64) + ok = np.zeros(n_rows, dtype=bool) + peaks = [None] * n_rows + factors_np = _host(factors, xpy) + + # ---- gate 1, before any work is done for the row. MIN_LOCAL_POINTS is a lower + # bound on this rule's point count, so a row that already loses on it can never + # win, and enumerating it would be pure waste -- measured: without this gate a + # block whose rows all fall back still paid 0.66 s of enumeration on top of the + # dense path it ended up using anyway. A gate applied only AFTER enumeration + # makes the method slower than the path it delegates to, which is the opposite + # of the point. + c_lo, c_dn = _estimated_costs(MIN_LOCAL_POINTS, npts, factors_np) + # THE CEILING FIRST. `required_upsample_factors` saturates a row it cannot justify + # at 2*UPSAMPLE_FACTOR_MAX precisely so the dense path will RAISE on it -- that is + # the fail-closed behaviour of the whole option. The cost gate below compares + # against the dense cost, so the sharper the row the more certainly this rule keeps + # it, and a row past the ceiling is the sharpest kind there is: it would be handled + # here and silently returned instead of raising. Measured before this check: at a + # derived factor of 8192 the dense path raised (as designed) while peak-local + # returned -24451 nats and reported tail_bound_worst = -2721. + # + # WHAT THIS CHECK IS, AND DELIBERATELY IS NOT. It catches the SATURATION SENTINEL -- + # `required_upsample_factors` returning 2*UPSAMPLE_FACTOR_MAX because it could not + # justify a factor at all. It does NOT mirror the dense path's ceiling test, and that + # is a decision rather than an oversight. + # + # `time_marginalize_bandlimited` re-measures the width on the grid it just refined and + # doubles until the criterion holds, raising if the chain leaves the legal range, so it + # can refuse a row whose COARSE factor was legal. This module tests the coarse factor + # once. The two therefore disagree on exactly one band, `factor_coarse == + # UPSAMPLE_FACTOR_MAX`: measured over 220 random multi-bump rows, 27 landed there and 8 + # of them are rows the dense path refuses and this one returns. (Below 4096 the + # disagreement never arose -- it would need the coarse estimate to be optimistic by + # three doublings -- and at the 8192 sentinel both refuse, 44 of 44.) + # + # Those 8 rows are NOT approximated: their values were scored exact against a converged + # spectral reference. UPSAMPLE_FACTOR_MAX bounds the DENSE GRID -- the point past which + # a zero-padded FFT over the whole window stops being affordable or believable -- and + # this rule never builds that grid. Its resolution comes from a per-peak width measured + # on the enumeration grid, which is already 8x finer than the coarse grid the ceiling + # factor is derived from, and the containment check and tail bound then verify the + # outcome. So a row the dense path declines on grid size is one this rule may legitimately + # still resolve, and RO's ruling is that it may: the ceiling is a limit on the dense + # grid, not on the physics. The cost of forcing agreement would be `>=` here, which + # sends every row at the legal ceiling to the dense path -- 19 of those 27 rows, and the + # sharpest legal rows are precisely the regime this rule exists to serve. + over_ceiling = factors_np > _tmq.UPSAMPLE_FACTOR_MAX + viable = (c_lo < c_dn) & (~over_ceiling) + stats['n_dense_fallback_ceiling'] += int(np.sum(over_ceiling)) + # Counted SEPARATELY from the post-enumeration gate. Both are cost decisions, but + # only this one runs before any work is done for the row, and it is the one credited + # with removing the regression where this rule came out slower than the path it + # delegates to. A single shared counter made it removable with the suite green: + # nothing could tell which gate had fired. + stats['n_dense_fallback_cost_pregate'] += int(np.sum((~viable) & (~over_ceiling))) + if not viable.any(): + return values, ok, peaks + + # ---- enumeration. One FFT upsample of kappa at a FIXED, SNR-independent + # factor. The callback is NOT evaluated on this grid: only term(kappa) is + # needed, because a monotone callback cannot move an extremum. + k_up = reflected_bandlimited_upsample( + kappa_rows, PEAK_ENUM_FACTOR, xpy=xpy)[..., :last + 1] + q_up = _term(k_up) + del k_up + n_enum = q_up.shape[-1] + + # The reflected spectrum is needed for the rigorous crest bound below and again for + # localisation; it is one FFT, built once here. + kappa_reflected = xpy.concatenate( + (kappa_rows, xpy.flip(kappa_rows, axis=-1)), axis=-1) + Xw, fk = bandlimited_spectrum(kappa_reflected, xpy=xpy) + period_ref = 2.0 * period + # RIGOROUS bound on |q''| for this row, from the spectrum rather than from a model: + # q(t) = Re sum_j Xw_j exp(w_j t), so |q''| <= sum_j |Xw_j| |w_j|^2 everywhere, by the + # triangle inequality. No parabola is assumed and nothing is fitted, which is the + # whole point -- see the keep note below. + q_ddot_max = spectral_curvature_bound(Xw, fk, period_ref, xpy=xpy) + + mask = enumerate_peak_indices(q_up, xpy=xpy) + mask = mask & xpy.asarray(viable)[:, None] + rows_p, cols_p = xpy.where(mask) # full-width mask: index is the sample + if int(rows_p.shape[0]) == 0: + return values, ok, peaks + + # ---- widths. These need lnL, so the callback runs HERE -- on a short stencil + # about each peak, never on the full time axis. The stencil CENTRE is clipped + # inward rather than the offsets, so it stays centred; the second difference of a + # parabola is its second derivative at any spacing and about any centre, which is + # the same property peak_width_from_lnL relies on. + maxd = max(CURVATURE_STENCIL_HALFWIDTHS) + if 2 * maxd >= n_enum: + return values, ok, peaks + # The stencil centre is shifted inward by the MINIMUM needed for a three-point + # stencil to exist -- one sample -- and out-of-range half-widths are masked off. + # + # Two wrong versions preceded this. The first clipped the centre by `maxd = 8`, + # which measured a point up to 8 enumeration samples away. The second stopped + # clipping entirely and masked instead, which looks more principled and is worse at + # the two indices that matter: at `cols_p == 0` the whole left half of the stencil + # is out of range at EVERY half-width, so `d2` is NaN throughout, `sigma = inf`, and + # the peak is dropped for "no resolvable curvature" before anything else runs. The + # vertex height degrades the same way, falling back to the raw sample. Measured: + # indices 0 and n_enum-1 dropped, 1 / 7 / 8 / n-2 kept -- so revision 2's endpoint + # enumeration was DEAD CODE (22 endpoint maxima enumerated, zero able to obtain a + # finite width), and the near-edge defect it was meant to fix re-ran byte-identical. + # + # Shifting by one gives a genuine ONE-SIDED fit at an endpoint -- the parabola + # through (0,1,2) -- which is the correct estimator there, not a compromise. + offs = xpy.arange(-maxd, maxd + 1) + c_st = xpy.clip(cols_p, 1, n_enum - 2) + take_raw = c_st[:, None] + offs[None, :] + st_valid = (take_raw >= 0) & (take_raw < n_enum) + q_st = q_up[rows_p[:, None], xpy.clip(take_raw, 0, n_enum - 1)] + lnL_st = loglikelihood(q_st, xpy.broadcast_to(rho_col_rows[rows_p], q_st.shape)) + lnL_st = xpy.where(st_valid, lnL_st, np.nan) + sigma_pk = _peak_curvature_sigma(lnL_st, h_enum, xpy=xpy) + + # G1/Door 1: the KEEP filter must compare CRESTS, not samples -- and not estimates + # of crests either. See the two-stage note below. + # + # This is the same defect as the interval-centring bug, at a different site, and it + # is the reason to grep for every consumer of the enumeration index rather than fix + # the one that was reported. `q_up[rows_p, cols_p]` is the SAMPLE value; a crest a + # distance d from its sample reads (d/sigma)^2/2 nats low, which for a sharp row is + # tens of nats. Comparing a between-samples peak against an on-sample peak then + # drops the former. MEASURED on the two-peak fixture at rho ~ 700: the secondary + # crest is 1.003 nats below the dominant crest, but its SAMPLE is 70.99 nats below, + # past PEAK_KEEP_NATS -- so one of two equal peaks was deleted and the answer came + # back exactly -log(2) = -0.693147 low. + # + # THE PARABOLIC VERTEX IS NOT THE CREST, and believing it was is how this class of + # defect survived a second round. The vertex is exact only if `lnL` is quadratic + # across the whole stencil; it is band-limited `q(t)`, so the quartic term survives + # and the vertex obeys the SAME `(delta/sigma)^2` law as the sample, just smaller. + # MEASURED at half-cell phase -- under-read of the crest, in nats: + # + # sigma_t/deltaT factor sample vertex + # 0.004934 512 -108.75 -1.17 + # 0.002467 1024 -434.98 -4.66 + # 0.001233 2048 -1739.93 -18.66 + # 0.000617 4096 -6959.72 -74.63 <-- past PEAK_KEEP_NATS + # + # `UPSAMPLE_FACTOR_MAX = 4096` permits sigma_t/deltaT down to 0.000488, so that last + # row is a LEGAL configuration, and end-to-end it reproduced -log(2) exactly with + # both defences silent. Every approximation substituted for the crest fails the + # same way one octave further out. + # + # ROUND 6: A TARGETING MODEL WAS PROMOTED TO A BOUND, AND THAT IS THE FOURTH DOOR. + # + # THE LAPLACE MODEL IS NOT THE DEFECT. `lnL` is of course not a parabola; the whole + # design uses the Gaussian/curvature picture to TARGET -- where the mass is, how wide + # an interval must be, how fine a spacing it needs -- and then either the Laplace + # quadrature is reliable enough or the row falls back to the grid. Used that way the + # model owes nothing, because the answer comes from the quadrature actually performed + # and from the checks that verify it. + # + # The defect is using the same model where an INEQUALITY is required. The pre-filter + # compared `crest_upper = lnL_sample + (h_enum/2)^2 / (2 sigma^2)` against the largest + # sample in the row, and was described as an upper bound that "can only ever keep too + # many". It is a targeting estimate wearing the word "bound", and an independent + # re-attack broke it: + # + # * the displacement is bounded by `h_enum`, not `h_enum/2` -- the localiser's own + # bracket says so and 0.959*h_enum has been observed -- so the correction is + # taken at less than half the distance it must cover; and, much worse, + # * the ANHARMONIC part of the deficit carries the same 1/sigma^2 amplification as + # the quadratic part, so the error does not shrink where it matters. MEASURED at + # derived factor 1024 on a skewed peak: the pure quantisation excess is 4.4 nats + # while the actual shortfall is 122.30. + # + # A targeting model that is off by 122 nats costs a slightly wrong interval, which the + # verification catches. The same model asked to justify DELETING a peak costs the peak. + # + # So `crest_upper` fell short of the true crest by 122 / 489 / 1957 nats at derived + # factor 1024 / 2048 / 4096 -- and being short, it DROPPED co-dominant peaks. + # End-to-end, shipped code, both peaks well inside PEAK_KEEP_NATS: **-0.358 nats** at + # factor 1024 and, when the deleted peak is raised above the survivor, **-1849 nats**, + # ACCEPTED, with the tail bound and the containment check both silent. A peak may sit + # ~1900 nats ABOVE the one that survives and still be deleted. The tail bound cannot + # backstop it because `q_out_max` reads the dropped peak at its SAMPLE -- the very + # quantity the defect corrupts; recomputed as an honest supremum, those rows' margins + # are +11.5 / +12.1 / +7.2 and every one would be REJECTED. + # + # This is the fourth time this class has reopened, and the fourth time the estimate of + # the crest was one octave too optimistic. Widening the constant would be the fifth. + # NO QUANTITY DERIVED FROM THE ENUMERATION INDEX MAY DROP A PEAK. The index survives + # only as a Newton seed and bracket centre, which is the one thing it is entitled to + # be. The keep decision is now taken ONCE, after localisation, on `lnL_star`, which + # is the crest rather than an estimate of it. + # + # What used to justify the pre-filter was COST -- bounding how many peaks reach + # localisation. That job is already done, and done safely, by the gate below: it runs + # BEFORE localisation, it is built from the enumeration samples alone, and it declines + # the whole ROW to the dense path rather than selecting which peaks to believe. A + # gate that declines a row is safe in a way that a filter which deletes a peak is not. + # THE PRE-FILTER, REBUILT ON AN INEQUALITY INSTEAD OF A FIT. + # + # Cost still has to be bounded -- without any pre-filter every row enumerates its + # whole oscillation (295 maxima on one fixture here), the gate below sees more + # intervals than `MAX_INTERVALS` and declines EVERY row, and the option becomes inert. + # Measured: `n_peak_local_rows = 0` on all six fixture families. + # + # So a peak may still be dropped, but only against a bound that holds unconditionally. + # Let `t*` be the crest and `t_s` its enumeration sample. Expanding about the CREST, + # where `q'` vanishes by definition, + # + # q(t_s) = q(t*) + q''(xi) (t_s - t*)^2 / 2, |t_s - t*| <= h_enum + # + # so q(t*) <= q(t_s) + q_ddot_max * h_enum^2 / 2 with `q_ddot_max` the spectral + # bound computed above. This is a Taylor remainder with a TRUE bound on the second + # derivative, not a parabolic fit, so it is immune to the anharmonicity that broke the + # previous version: there is no assumption that `lnL` is quadratic across a cell, and + # no `1/sigma^2` amplification of a modelling error. It uses `h_enum`, the localiser's + # actual bracket, not `h_enum/2`. + # + # `loglikelihood` is monotone in its first argument over its DOMAIN, so bounding `q` + # bounds `lnL`. Stated precisely because the obvious stronger claim is false: the + # shipped distance-marginalized callback returns `-inf` ABOVE its table as well as + # below, so it is not monotone increasing everywhere. That direction is safe here -- + # an upper bound that evaluates into the hole becomes `-inf`, the peak is dropped, and + # the row loses peaks until it falls back to the dense path -- so it costs coverage, + # never accuracy. With the shipped table the boundary sits at `D_eff < d_min/10` and + # is not reachable at default settings. The + # comparison is then a genuine upper bound against a genuine lower bound (the largest + # SAMPLE in the row, which cannot exceed the crest above it), and a peak is discarded + # only when it cannot be within `PEAK_KEEP_NATS` however the quantisation falls. + # + # NOTE the sample is read AT THE ENUMERATED INDEX. The previous version read + # `lnL_st[:, maxd]`, which sits at the stencil centre -- clipped inward by one at the + # array ends -- so at `cols_p` 0 or `n_enum-1` it was a full enumeration cell away from + # the peak, measured 132 nats low at rho ~ 40 and 8449 nats low at rho ~ 700, growing + # with SNR. That is the same defect at a third site. + q_at_peak = q_up[rows_p, cols_p] + q_crest_upper = crest_upper_bound(q_at_peak, q_ddot_max[rows_p], h_enum) + rho_at_peak = rho_col_rows[rows_p, 0] + lnL_upper = loglikelihood(q_crest_upper, rho_at_peak) + lnL_lower = loglikelihood(q_at_peak, rho_at_peak) + + rows_np = _host(rows_p, xpy) + cols_np = _host(cols_p, xpy) + sig_np = _host(sigma_pk, xpy) + up_np = _host(lnL_upper, xpy) # rigorous UPPER bound on this crest + low_np = _host(lnL_lower, xpy) # LOWER bound (the sample cannot exceed it) + + row_best = np.full(n_rows, -np.inf) + np.maximum.at(row_best, rows_np, low_np) + # `isfinite(sig_np)` is not a magnitude decision: a peak with no finite negative + # curvature at any stencil half-width has no width, so no interval can be built for + # it. What is dropped here lies outside the intervals and enters the tail bound. + keep = np.isfinite(sig_np) & (up_np > row_best[rows_np] - PEAK_KEEP_NATS) + rows_np, cols_np, sig_np = rows_np[keep], cols_np[keep], sig_np[keep] + if rows_np.size == 0: + return values, ok, peaks + + # ---- rows that survived every gate but enumerated no usable peak. Counted, so the + # sub-counts of last_report() RECONCILE against n_dense_fallback_rows: an operator + # reading the columns should not find an unexplained residual. + _has_peak = np.zeros(n_rows, dtype=bool) + _has_peak[rows_np] = True + stats['n_dense_fallback_nopeak'] += int(np.sum(viable & (~_has_peak))) + + # ---- gate 2, on a CONSERVATIVE SUPERSET of the intervals, BEFORE localising. + # Localisation is the expensive step (Newton over the spectrum, per peak), and a + # broad row has many peaks, so running it before the gate that discards the row is + # pure waste -- the same mistake that once made this rule slower than the path it + # delegates to, and it cost 7x here before this gate was added (0.14x vs the dense + # path at sigma_t/deltaT = 0.17, on a block where every row fell back anyway). + # + # The crest is somewhere within h_enum/2 of its sample, so an interval built about + # the sample and widened by h_enum/2 CONTAINS the interval that localisation will + # produce, whatever the answer turns out to be. Its point count is therefore an + # over-estimate, and a gate on an over-estimate can only decline rows -- never keep + # one it should have declined. Where this rule actually wins the conservatism is + # irrelevant: at sigma_t/deltaT ~ 0.002 the superset costs 66k against a dense cost + # of 12M, so the row is kept with four orders of magnitude to spare. + tol_np = LOCALISE_SAFETY * sig_np + t_grid_np = cols_np * h_enum + # h_enum, NOT h_enum/2. The bracket in `localise_peaks` is +/- h_enum and accepts + # anything strictly inside it, so |t_star - t_grid| is bounded by h_enum and not by + # half of it -- and that bound is reached: over 14,182 localised peaks the largest + # observed displacement was 0.959 * h_enum, and 2.62% of intervals were NOT + # contained by a half-cell margin. Narrowing the bracket instead would be wrong: + # an asymmetric peak's crest genuinely can sit more than half a cell from its + # sample, which is what that 0.959 measures. + prov_half = W_SIGMA * sig_np + tol_np + h_enum + p_order, p_gid, pg_row, pg_lo, pg_hi = merge_intervals_by_row( + rows_np, np.maximum(t_grid_np - prov_half, 0.0), + np.minimum(t_grid_np + prov_half, t_last), t_last) + p_smin = np.full(pg_row.size, np.inf) + np.minimum.at(p_smin, p_gid, sig_np[p_order]) + p_nloc = np.maximum(3, np.ceil( + (pg_hi - pg_lo) / np.minimum(p_smin / UPSAMPLE_SAFETY, h_enum) + ).astype(np.int64) + 1) + p_niv_row = np.bincount(pg_row, minlength=n_rows) + p_cl, p_cd = _estimated_costs(np.bincount(pg_row, weights=p_nloc, + minlength=n_rows), npts, factors_np) + p_much = p_niv_row > MAX_INTERVALS + p_slow = (~p_much) & (p_niv_row > 0) & (p_cl >= p_cd) + stats['n_dense_fallback_structure'] += int(np.sum(p_much)) + stats['n_dense_fallback_cost'] += int(np.sum(p_slow)) + prov_keep = (p_niv_row > 0) & (~p_much) & (~p_slow) + sel_pk = prov_keep[rows_np] + if not sel_pk.any(): + return values, ok, peaks + rows_np, cols_np = rows_np[sel_pk], cols_np[sel_pk] + sig_np, tol_np = sig_np[sel_pk], tol_np[sel_pk] + t_grid_np = t_grid_np[sel_pk] + + # ---- LOCALISE. An enumerated extremum is a grid INDEX; the interval has to be + # built around the CREST. Centring on the index instead cost up to 165 nats, always + # negative -- see LOCALISE_SAFETY. Newton is confined to the bracket the + # enumeration already established, so it places peaks, it cannot find or lose them. + # THE RECONSTRUCTION MUST BE THE ONE THE DENSE PATH USES, and it is no longer the + # raw periodic interpolant. `time_marginalize_bandlimited` periodizes the EVEN + # REFLECTION `[kappa forward, kappa backward]` (rift_O4d e4ed25c7, "Avoid Gibbs + # ringing"), because a zero-padded FFT of the gathered slice alone identifies its + # unlike endpoints and rings globally -- measured +140.9 nats on an adversarial row. + # + # This module was written against the older periodic contract and the rebase onto + # rift_O4d changed it underneath. Leaving it periodic makes peak-local integrate a + # DIFFERENT continuous function from the one its own fallback rows get, inside a + # single call: measured -3.79 nats on a row with peaks near both window ends + # (`test_intervals_are_clipped_to_the_integration_domain`) and a residual 9.0e-6 nat + # median bias on the uniform-arrival block. So enumeration, localisation and local + # evaluation all run on the reflected row. + # + # The reflected row has length `2*npts` -- always EVEN -- so `bandlimited_spectrum` + # takes its even branch and splits the Nyquist bin exactly as `bandlimited_upsample` + # does inside `reflected_bandlimited_upsample`. The two therefore agree on the + # forward interval, which is the only part this module ever evaluates. + t_star, q_star, loc_ok = localise_peaks( + Xw, fk, xpy.asarray(rows_np), xpy.asarray(t_grid_np), h_enum, + xpy.asarray(tol_np), period_ref, xpy=xpy, t_last=t_last) + t_np = _host(t_star, xpy) + lnL_star = _host(loglikelihood(q_star, rho_col_rows[xpy.asarray(rows_np), 0]), xpy) + loc_ok_np = _host(loc_ok, xpy).astype(bool) + + # A row with ANY peak the localiser could not place is not approximated -- it goes + # to the dense path. Fail closed: an unplaced crest is exactly the condition that + # produced the bias. + bad_loc = np.zeros(n_rows, dtype=bool) + + # ---- KEEP, stage 2: EXACT, on the localised crest values. + # + # This is the decision that matters, and it is taken here rather than before + # localisation for one reason: `lnL_star` is the crest, not an estimate of it. The + # pre-filter above only bounded how many peaks reached this point, and it errs + # toward keeping, so nothing that matters has been discarded on an approximation. + row_top = np.full(n_rows, -np.inf) + np.maximum.at(row_top, rows_np, lnL_star) + exact_keep = lnL_star > row_top[rows_np] - PEAK_KEEP_NATS + if not exact_keep.all(): + rows_np, cols_np = rows_np[exact_keep], cols_np[exact_keep] + sig_np, tol_np = sig_np[exact_keep], tol_np[exact_keep] + t_np, lnL_star = t_np[exact_keep], lnL_star[exact_keep] + loc_ok_np = loc_ok_np[exact_keep] + if rows_np.size == 0: + return values, ok, peaks + + # The interval is centred on the crest and widened by the localisation residual, so + # containment does not depend on the crest happening to sit near a grid sample. + half_np = W_SIGMA * sig_np + tol_np + lo_np = np.maximum(t_np - half_np, 0.0) + hi_np = np.minimum(t_np + half_np, t_last) + + # Localisation failure is judged on the peaks that SURVIVED the exact filter: a peak + # about to be discarded as immaterial should not condemn its row. + bad_loc[rows_np[~loc_ok_np]] = True + stats['n_dense_fallback_localise'] += int(np.sum(bad_loc)) + + # Per-row crest value, for the a-posteriori containment check after integration. + row_star = np.full(n_rows, -np.inf) + np.maximum.at(row_star, rows_np, lnL_star) + bounds = np.searchsorted(rows_np, np.arange(n_rows + 1)) + + # ---- merge, and derive each merged interval's own spacing. All rows at once: + # the per-peak Python loop this replaces was the reason the rule could come out + # slower than the dense path it delegates to. + order, gid, g_row, g_lo, g_hi = merge_intervals_by_row(rows_np, lo_np, hi_np, + t_last) + n_groups = g_row.size + # Spacing per merged interval: set by the SHARPEST peak inside it, and never + # coarser than the grid on which the structure was established. + s_min = np.full(n_groups, np.inf) + np.minimum.at(s_min, gid, sig_np[order]) + h_want = np.minimum(s_min / UPSAMPLE_SAFETY, h_enum) + n_loc = np.maximum(3, np.ceil((g_hi - g_lo) / h_want).astype(np.int64) + 1) + + n_iv_row = np.bincount(g_row, minlength=n_rows) + n_loc_row = np.bincount(g_row, weights=n_loc, minlength=n_rows) + c_local, c_dense = _estimated_costs(n_loc_row, npts, factors_np) + + # These can only catch what the provisional pass could not. The localised intervals + # are contained in the provisional ones, so their point count -- and hence the cost + # test -- can only have improved; the INTERVAL COUNT can go up, though, because + # narrower intervals merge less readily, so the structure test is not redundant. + too_much = n_iv_row > MAX_INTERVALS + too_slow = (~too_much) & (n_iv_row > 0) & (c_local >= c_dense) + stats['n_dense_fallback_structure'] += int(np.sum(too_much)) + stats['n_dense_fallback_cost'] += int(np.sum(too_slow)) + keep_row = (n_iv_row > 0) & (~too_much) & (~too_slow) & (~bad_loc) + + gbounds = np.searchsorted(g_row, np.arange(n_rows + 1)) + plan = [] + for r in np.nonzero(keep_row)[0]: + ga, gb = int(gbounds[r]), int(gbounds[r + 1]) + plan.append((int(r), g_lo[ga:gb], g_hi[ga:gb], int(n_loc[ga:gb].max()), + int(bounds[r]), int(bounds[r + 1]))) + + if not plan: + return values, ok, peaks + + # ---- batched evaluation. Rows are grouped by (interval count, point-count + # bucket) so padding to a common shape can waste at most a factor of two, and + # every interval slot of a group is one batched call. + covered = np.zeros((n_rows, n_enum), dtype=bool) + covered_len = np.zeros(n_rows) + attained = np.full(n_rows, -np.inf) + parts = xpy.full((n_rows, MAX_INTERVALS), -np.inf, dtype=np.float64) + buckets = {} + for entry in plan: + key = (entry[1].size, int(2 ** np.ceil(np.log2(max(entry[3], 2))))) + buckets.setdefault(key, []).append(entry) + + for (n_iv, m_pad), members in buckets.items(): + rr = np.array([m[0] for m in members]) + rr_x = xpy.asarray(rr) + for j in range(n_iv): + a_h = np.array([m[1][j] for m in members], dtype=np.float64) + b_h = np.array([m[2][j] for m in members], dtype=np.float64) + h_h = (b_h - a_h) / float(m_pad - 1) + # A zero-length merged interval (a peak pinned against a window end) + # would give h=0 and a degenerate grid; give it the enumeration spacing so + # the trapezoid has a domain. Note which way the risk runs: this makes the + # integration domain EXCEED the merged interval, so the hazard is + # double-counting against a neighbour, not missing mass. Unreachable as + # written -- a zero-length interval needs a peak exactly at a window end, + # which the edge guard has already routed away -- but the thing to check is + # over-coverage. + h_h = np.where(h_h > 0, h_h, h_enum) + k_loc = eval_bandlimited_uniform(Xw[rr_x], fk, xpy.asarray(a_h), + xpy.asarray(h_h), m_pad, period_ref, + xpy=xpy) + lnL_loc = loglikelihood( + _term(k_loc), xpy.broadcast_to(rho_col_rows[rr_x], k_loc.shape)) + parts[rr_x, j] = _log_trapz_local(lnL_loc, xpy.asarray(h_h), xpy=xpy) + # Highest lnL the integration grid actually REACHED, for the containment + # check below. Free: the same maximum is already taken for the offset. + attained[rr] = np.maximum(attained[rr], + _host(xpy.max(lnL_loc, axis=-1), xpy)) + stats['n_local_points_total'] += int(rr.size) * m_pad + for i_m, m in enumerate(members): + lo_i = int(np.ceil(a_h[i_m] / h_enum)) + hi_i = int(np.floor(b_h[i_m] / h_enum)) + if hi_i >= lo_i: + covered[m[0], max(lo_i, 0):hi_i + 1] = True + # EXACT covered length, not a count of grid indices. A merged interval + # narrower than h_enum, or a gap between two of them that happens to + # contain no integer index, contributes nothing to the index count -- so + # a T_outside built from that count omits real, uncovered time. Measured + # on a comb-like row before this change: 4.8% of the true mass sat in + # such gaps, worth -0.050 nats, while the bound reported -43. + covered_len[m[0]] += float(b_h[i_m] - a_h[i_m]) + stats['n_intervals_total'] += int(rr.size) * n_iv + + result = _logaddexp_reduce(parts, xpy=xpy) + + # ---- the tail bound. log(T_outside) + max_{outside} lnL is an upper bound on + # the omitted integral. It is evaluated on term(kappa) -- one callback value per + # row -- because the callback is monotone in it, so no evaluation on the full + # time axis is needed. A row whose bound is not small enough is NOT reported + # with a caveat: it goes to the dense path. + # ---- the outside supremum, CERTIFIED on sub-cell geometry. + # + # `max(q_up over uncovered SAMPLES)` is a LOWER bound on the continuous supremum and the + # gap GROWS WITH AMPLITUDE: measured, the reported margin was 14 nats optimistic at + # amplitude 2e4, 75-218 at 2e5 and 1445-2923 at 2e6. `tail_bound_worst` got more + # flattering the sharper the row, the wrong direction for a safety number. + # + # The fix is a true bound on every uncovered piece. Per enumeration cell, the cubic + # Hermite through the cell's endpoint VALUES AND SLOPES plus the remainder `M4 h^4/384` is + # an upper bound on `q` there -- it assumes nothing about where extrema are, so it closes + # the "a peak between samples was never enumerated" objection outright rather than arguing + # the case is unrealistic. The slopes are what make it affordable: from endpoint values + # alone the remainder is `M2 h^2/8` and from one sample `M2 h^2/2`, which measure 102 / + # 1.0e4 / 1.0e5 and 407 / 4.1e4 / 4.1e5 nats at amplitude 2e4 / 2e6 / 2e7 -- useless + # against a 23 nat tolerance. With the slopes: 0.12 / 12.2 / 122. + # + # SUB-CELL GEOMETRY IS THE WHOLE POINT, and bounding whole cells does not work. This rule + # earns its keep on sharp rows, and there the merged interval is NARROWER THAN ONE + # ENUMERATION CELL -- measured, the half-width falls to 0.05 of a cell at derived factor + # 4096. A cell-granular notion of "covered" then marks nothing, the crest's own cell counts + # as outside, a bound over that whole cell bounds the CREST, and every row is rejected. + # (Measured: exactly that, the option goes inert.) So each cell is bounded on the part + # that is genuinely outside: the interval cuts the cell at `lo` and `hi`, leaving `[0, lo]` + # and `[hi, 1]` in cell coordinates, and the Hermite is maximised over those. Empty pieces + # return -inf and drop out of the maximum. + # + # A cell containing TWO merged intervals would leave a gap between them that this + # two-piece decomposition does not cover. It is detected and the ROW is declined -- + # fail-closed, and rare: it needs two crests inside one enumeration cell. + n_cells = n_enum - 1 + m4 = _host(spectral_derivative_bound(Xw, fk, period_ref, 4, xpy=xpy), xpy) + (dq,) = enum_grid_derivatives(kappa_reflected, PEAK_ENUM_FACTOR, n_enum, deltaT, + orders=(1,), xpy=xpy) + q_np = _host(q_up, xpy) + dq_np = _host(dq, xpy) + del dq + + big = 2.0 * (float(t_last) + 1.0) + t_l = (np.arange(n_cells) * h_enum)[None, :] + row_off = (np.arange(n_rows) * big)[:, None] + if g_row.size: + key_iv = g_lo + g_row * big + j = np.searchsorted(key_iv, (t_l + h_enum + row_off).ravel(), + side='right').reshape(n_rows, n_cells) - 1 + jc = np.maximum(j, 0) + hit = (j >= 0) & (g_row[jc] == np.arange(n_rows)[:, None]) & (g_hi[jc] > t_l) + j2 = np.maximum(jc - 1, 0) + two = (hit & (jc >= 1) & (g_row[j2] == np.arange(n_rows)[:, None]) + & (g_hi[j2] > t_l)) + cert_bad = two.any(axis=1) + lo_n = np.clip(np.where(hit, (g_lo[jc] - t_l) / h_enum, 1.0), 0.0, 1.0) + hi_n = np.clip(np.where(hit, (g_hi[jc] - t_l) / h_enum, 1.0), 0.0, 1.0) + else: + cert_bad = np.zeros(n_rows, dtype=bool) + lo_n = np.ones((n_rows, n_cells)) + hi_n = np.ones((n_rows, n_cells)) + + q0, q1 = q_np[:, :-1], q_np[:, 1:] + d0, d1 = h_enum * dq_np[:, :-1], h_enum * dq_np[:, 1:] + m4c = m4[:, None] + cell_sup = np.maximum( + segment_sup_bound(q0, q1, d0, d1, h_enum, m4c, 0.0, lo_n, xpy=np), + segment_sup_bound(q0, q1, d0, d1, h_enum, m4c, hi_n, 1.0, xpy=np)) + q_out_max = xpy.asarray(np.max(cell_sup, axis=-1)) + del cell_sup, q0, q1, d0, d1, lo_n, hi_n, q_np, dq_np + T_out = np.maximum(t_last - covered_len, 0.0) + lnL_out = loglikelihood(q_out_max, rho_col_rows[:, 0]) + with np.errstate(divide='ignore', invalid='ignore'): + bound = np.where(T_out > 0, np.log(np.maximum(T_out, 1e-300)) + + _host(lnL_out, xpy), -np.inf) + margin = bound - _host(result, xpy) + + planned = np.array([m[0] for m in plan]) + # TWO conditions, and they fail differently on purpose. The tail bound is a + # statement about mass OUTSIDE the intervals, computed from a sampled maximum, so it + # is only as good as the grid it samples -- which is the grid that produced F1. The + # containment check is a statement about the crest being INSIDE, verified from the + # integration grid's own values, and it is what actually catches a mis-placed + # interval. Neither subsumes the other and a row must satisfy both. + contained = attained >= row_star - CONTAINMENT_SLACK_NATS + (good_mask, structure_fail, tail_fail, + containment_fail) = _certificate_acceptance_masks( + margin, contained, cert_bad, planned) + stats['n_dense_fallback_structure'] += int(np.sum(structure_fail)) + accepted = planned[good_mask] + rejected = planned[~good_mask] + stats['n_dense_fallback_tail'] += int(np.sum(tail_fail)) + stats['n_dense_fallback_containment'] += int(np.sum(containment_fail)) + if accepted.size: + acc_x = xpy.asarray(accepted) + values[acc_x] = result[acc_x] + ok[accepted] = True + stats['tail_bound_worst'] = max(stats['tail_bound_worst'], + float(margin[accepted].max())) + stats['n_peak_local_rows'] += int(accepted.size) + stats['n_peaks_total'] += int(np.isin(rows_np, accepted).sum()) + + if want_peaks: + for r, starts, stops, m_max, a, b in plan: + if ok[r]: + # t_np is the LOCALISED crest, not the enumeration sample -- which is the + # whole point of exposing it for a time-first reordering. + peaks[r] = (t_np[a:b].copy(), sig_np[a:b].copy()) + + return values, ok, peaks diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py new file mode 100644 index 000000000..f9e437f93 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py @@ -0,0 +1,1093 @@ +"""Band-limited time marginalization for the factored likelihood. + +WHAT IS WRONG WITH THE HISTORICAL PATH +-------------------------------------- +``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`` forms ``lnL(t)`` on the +grid built by ``factored_likelihood.marginalization_time_grid`` -- spacing +``deltaT = 1/srate``, fixed -- and integrates ``exp(lnL(t))`` over it with +Simpson's rule at that fixed spacing. The grid spacing is a property of the +DATA; the integrand's width is a property of the SIGNAL, and the two are +unrelated. After marginalizing the angles, ``exp(lnL(t))`` is a near-Gaussian +peak of width + + sigma_t = 1 / (2 pi rho sigma_f) + +with ``sigma_f`` the noise-weighted template frequency spread. Resolving it +needs ``deltaT <~ sigma_t``, i.e. ``srate >~ 2 pi sigma_f rho`` -- a requirement +that grows LINEARLY WITH SNR and that production, which runs at the data sample +rate, does not meet. Simpson makes the under-resolved case worse rather than +better: ``simpson = (4 T_h - T_2h)/3`` carries the coarser trapezoid ``T_2h``, +so it inherits an alias with period ``2h``. + +Measured on a 35+30 Msun SEOBNRv4 H1L1V1 injection at rho=40 (sigma_t = 61.2 us, +a property of the signal and so srate-independent), rigidly scanning the grid +phase over 2*deltaT moves the reported lnL by + + srate 4096: 1.649 nats 8192: 0.385 nats 16384: 0.0095 nats + +i.e. the answer depends on where the sample grid happens to fall relative to the +peak, by over a nat at the production sample rate. (Those three numbers were +taken on the JAX mirror of this quadrature, which integrates the SAME grid with +the SAME fixed Simpson weights; they are quoted as the physical scale of the +defect. The numbers for THIS path are the synthetic ones below, which are +against an analytic truth rather than against another estimate.) + +WHY THE EXISTING SAMPLES ALREADY CONTAIN THE ANSWER +--------------------------------------------------- +The data term ``kappa(t) = sum_det (t)`` is built from the +precomputed rholm cross-correlation timeseries, which are inverse FFTs of a +frequency-domain product band-limited to ``[fmin, fmax]`` with +``fmax <= fNyq = 1/(2 deltaT)``. So ``kappa(t)`` is band-limited below Nyquist. +The template self-term ``rho_sq`` is time-independent on this path. Therefore +``lnL(t) = f(kappa(t), rho_sq)`` is recoverable on an arbitrarily fine grid from +samples covering a full underlying period. + +The samples in hand are only a gathered INTEGRATION-WINDOW SLICE, however. A +zero-padded FFT of that slice alone identifies its generally unlike endpoints. +Decay of ``exp(lnL)`` does NOT make that safe: the FFT acts on ``kappa``, whose +endpoint mismatch can remain large where the likelihood is negligible, and the +resulting Gibbs overshoot is global. A centred adversarial row with both outer +eighths more than 49 nats below its peak still moved by +140.9 nats when the raw +slice was periodized. + +The implementation instead doubles the finite row by literal even reflection, +``[kappa forward, kappa backward]``, before FFT interpolation, then discards the +backward half. Both joins are value-continuous, the supplied samples are +unchanged, and no unavailable samples outside the deliberately narrow +integration domain are invented. The same adversarial row is accurate to +2.4e-4 nats. Boundary proximity and coarse tail height are reported or tested, +but neither selects a lower-accuracy Simpson fallback: such a discontinuous +switch silently changes likelihood quality, while an exception can be mapped by +the calling pipeline to waveform failure and silently excise configurations. + +Two independent checks of that claim. On the real injection above, against a +converged dense reference built by re-gathering the rholms at shifted window +offsets -- the expensive, genuinely independent construction -- at srate 4096, +rho=40: + + band-limited upsampling: -0.007 nats + Simpson at deltaT : +0.745 nats + +And for this path specifically, against an ANALYTIC truth rather than another +numerical estimate (see test_time_marginalization_quadrature.py: kappa is built +as a sum of complex exponentials below Nyquist, so the continuous function is +known in closed form). srate 4096, npts 614, error in nats at three grid phases: + + sigma_t/deltaT Simpson band-limited factor + 2.27-2.61 +5e-6 .. +0e0 0 1 + 0.72-0.83 +2.5e-3 .. -1.1e-4 0 4 + 0.25-0.28 +0.844 / +0.242 / -1.101 0 16 + 0.10-0.12 +1.742 / -1.833 / -11.68 0 32 + 0.046-0.052 +2.548 / -15.33 / -66.18 0 64 + 0.016-0.019 +3.589 / -139.4 / -549.0 0 256 + 0.007-0.008 +4.393 / -710.6 / -2760.4 0 512 + +Two things to read off that table. The first row is the reassuring one: where +the integrand is already resolved the historical rule is fine, the derivation +returns a factor of 1 and nothing is paid. The row at sigma_t/deltaT = 0.25 +spans 1.95 nats across grid phase, which is the same scale as the 1.649 nats +measured on the real injection at the same ratio -- the synthetic reproduces the +defect's magnitude rather than a caricature of it. On windows cut from a longer +band-limited signal, the reflected reconstruction measured 2e-8 to 2.5e-6 nats +error over the tested sharpness range; the raw periodic-slice reconstruction +could ring at the artificial wrap. + +RESOLUTION IS DERIVED, NOT CONFIGURED +------------------------------------- +There is deliberately no "upsampling factor" option. This defect class exists +because a resolution was once a settable number whose docstring claimed it was +ample. The factor here is derived from the integrand actually in hand: the +three-point second difference of ``lnL`` about its peak is an EXACT estimator of +``-1/sigma_t^2`` for a Gaussian peak at any grid spacing and any peak-vs-grid +phase, so the width can be measured on the coarse grid even when the peak is +badly under-resolved by it. The factor is then raised to a power of two +satisfying ``h_dense <= sigma_t / SAFETY``, and -- this is the part that makes it +an assertion rather than a guess -- the width is REMEASURED on the dense grid and +the factor doubled until the criterion holds there too. A flat integrand +measures an infinite width, derives a factor of 1 and costs nothing. + +``SAFETY = 2`` is not tunable and is not a compromise. The trapezoidal rule on +a Gaussian of width ``sigma`` at spacing ``h`` has relative error +``2 exp(-2 pi^2 sigma^2 / h^2)`` (Poisson summation); at ``h = sigma/2`` that is +``1.0e-34``. Even ``h = sigma`` would give ``5.3e-9``. + +That bound is the DESIGN CRITERION FOR THE REFINED GRID, where ``sigma/h >= 2`` +puts us deep in the exponential regime. It is NOT a model of the defect's +magnitude, and checking it against the measured Simpson errors above will not +work: at ``sigma/h ~ 0.1-0.4`` the peak is barely sampled at all, the asymptotic +alias picture has not engaged, and the observed spans scale roughly as ``h^2``. +Both statements are correct about different regimes. The quadrature on the dense +grid is TRAPEZOID, not Simpson: for a peak that decays to nothing inside the +window every Euler-Maclaurin boundary term vanishes, so the trapezoidal rule is +spectrally accurate there, while Simpson would reintroduce the ``2h`` alias that +is the original defect. + +WHAT "EXACTLY" IS EXACT ABOUT (read before quoting the accuracy numbers) +----------------------------------------------------------------------- +The reconstruction is exact for the integrand THE CODE ACTUALLY FORMS, which is +the true ``kappa(t)`` only when ``time_interp='nearest'`` -- which since 2026-09-02 +is no longer the ILE driver's default (issue #233; the driver now defaults to +``time_interp_choice.TIME_INTERP_DEFAULT``), so the paragraph below is now the +ORDINARY case rather than the exceptional one -- where +the gathered values are exact samples of ``Q`` (on a grid offset by up to +deltaT/2, which is a pre-existing property of that stencil). + +With ``time_interp='cubic'`` or ``'sinc'`` the gathered values are a fixed FIR +filter applied to ``Q``. That is still a band-limited sequence, so the sampling +argument survives and the refinement is still exact -- but exact for the FILTERED +function. It then converges precisely to the integral of a stencil-biased +integrand, and the stencil bias, not the quadrature, is the larger term. +Measured at srate 4096, peak lnL ~5300, peak centred: with 'nearest' this path is ++0.0002 nats against an analytic truth where Simpson is -521; with 'sinc' it is +-2.29 where Simpson is +1.28, and over a scan of seeds and grid phases Simpson +wins about half the cases. Neither number says the quadrature is wrong -- they +say that once a stencil is in use its own error dominates, and fixing the +quadrature exposes it rather than adding to it. The advantages quoted above are +for ``time_interp='nearest'``, which is NOT the driver default any more: pass +``--interpolate-time nearest`` alongside ``--time-marginalization-quadrature +bandlimited`` to reproduce them. Re-measuring this pairing under the new default +is an OPEN item, not a settled result. + +SCOPE +----- +Applies to the baseline (non-rotating) likelihood with ``n_cal == 1``. The +banded / slow-rotation path has a TIME-DEPENDENT ``rho_sq`` and sidereal +post-phases, so ``lnL(t)`` there is not a pointwise function of a band-limited +``kappa`` and a constant, and the argument above does not carry; that path +refuses this quadrature rather than silently mis-applying it. Calibration +marginalization (``n_cal > 1``) is excluded for now for the same +refuse-rather-than-guess reason: the reduction sums ``exp`` over realizations, so +each realization's kappa row would have to be upsampled and the derived factor +reconciled across realizations, which is untested here. +""" + +import numpy as np + +__all__ = [ + "TIME_QUADRATURE_CHOICES", + "UPSAMPLE_SAFETY", + "UPSAMPLE_FACTOR_MAX", + "EDGE_GUARD_FRACTION", + "bandlimited_upsample", + "reflected_bandlimited_upsample", + "peak_width_from_lnL", + "required_upsample_factors", + "validate_time_quadrature", + "time_quadrature_pipeline_prereqs", + "ILE_TIME_QUADRATURE_FLAG", + "refuse_unless_time_quadrature_emitted", + "refuse_unhonourable_time_quadrature", + "find_time_quadrature_in_ile_args", + "draw_piecewise_linear_log_posterior", + "time_marginalize_bandlimited", + "last_report", +] + +TIME_QUADRATURE_CHOICES = ("simpson", "bandlimited", "peak-local") + +#: 'peak-local' lives in RIFT.likelihood.time_marginalization_peak_local and +#: reuses this module's helpers wholesale (width estimator, derived factor, row +#: classification, edge guard, Simpson hand-over). It is named here rather than +#: there so that validate_time_quadrature stays the single place a quadrature name +#: is checked; the import runs the other way, so there is no cycle. + +#: ``h_dense <= sigma_t / UPSAMPLE_SAFETY``. See the module docstring: at this +#: value the trapezoidal rule's Poisson-summation error on a Gaussian peak is +#: 2e-34, so this is a hard-coded constant and not an accuracy/cost trade. +UPSAMPLE_SAFETY = 2.0 + +#: Fail-closed ceiling. The band limit bounds the useful factor: with +#: ``sigma_t >= deltaT / (pi rho)`` the derivation cannot legitimately ask for +#: more than ``~2 rho``. Exceeding this raises rather than silently truncating +#: the resolution. +UPSAMPLE_FACTOR_MAX = 4096 + +#: Fraction of the window at EACH end used only to REPORT a peak close to a +#: truncated integration boundary. It must not select a different quadrature: +#: crossing an arbitrary threshold cannot silently move an under-resolved row +#: back to Simpson, and raising on such a row can be interpreted upstream as a +#: waveform failure and silently excise that configuration. The reflected +#: reconstruction below has no value discontinuity at either boundary, so the +#: old periodic-wrap exclusion is no longer needed for numerical safety. +#: +#: Measured on a window cut from a longer band-limited signal (peaked kernel plus +#: a 12%-amplitude coloured background, so the two ends genuinely disagree), +#: sigma_t/deltaT = 0.042, error against the analytic continuous truth, in nats: +#: +#: peak distance from edge 307 100 30 8 2 0 +#: band-limited 5e-6 4.6e-3 5.2e-2 5.6e-2 -3.3 +88.8 +#: Simpson (for scale) -29.2 -29.9 -29.3 -29.4 -29.7 -29.9 +#: +#: That table is the REJECTED raw-slice periodic reconstruction and records why +#: boundary proximity must remain visible: +88.8 is in the dangerous direction, +#: where a spurious row can dominate importance weights. It does not justify an +#: algorithm switch. With reflection, peaks immediately on either side of the +#: old one-eighth line (samples 75.3, 76.3, 77.3 at npts=614) agree with analytic +#: truth at order 1e-6 nats. Switching them to under-resolved Simpson because +#: they crossed that line would itself create the quality regression. +#: +#: In a well-posed run nothing comes close: the grid is centred on the trigger's +#: geocentre time, so the peak sits within the trigger timing uncertainty (a few +#: ms, tens of samples) of the CENTRE, not of an edge. Rows that do violate it +#: are still reconstructed and integrated over the caller's unchanged, possibly +#: truncated domain; the count in ``last_report()`` makes that separate physical +#: window-containment issue auditable without changing the likelihood rule. +EDGE_GUARD_FRACTION = 0.125 + +#: Half-widths, in coarse samples, tried in turn for the curvature stencil. A +#: centred three-point stencil at the peak is the natural choice, but ``lnL_t`` +#: genuinely contains ``-inf`` in production -- the distance-marginalization +#: callback returns ``-inf`` outside its interpolation table -- and +#: ``(-inf) - 2*(-inf) + (-inf)`` is NaN, while ``NaN < 0`` is False. A row whose +#: stencil straddles the table edge would therefore report "no resolvable peak", +#: derive a factor of 1, and be SILENTLY UNDER-RESOLVED: no raise, no warning, +#: and the exact failure this whole change exists to remove. Widening the +#: stencil steps over the hole, and costs nothing in accuracy because the second +#: difference of a parabola is its second derivative at ANY spacing. A row where +#: no half-width yields a finite curvature is not guessed at: it is counted and +#: given the historical value. +CURVATURE_STENCIL_HALFWIDTHS = (1, 2, 4, 8) + +#: Working-set budget for one dense temporary, in bytes. An internal +#: memory-chunking parameter: it changes how many extrinsic rows are processed at +#: a time. Rows are independent, so it cannot change the answer BEYOND +#: FLOATING-POINT REASSOCIATION -- the batch shape reaches numpy's FFT and its +#: pairwise summation, so a differently-chunked run can differ in the last +#: bit or two. Measured on the companion peak-local implementation, which +#: inherited this same wording and then failed a bit-identity test that was +#: written to it: 0, 0 and 2 ULPs. "Cannot change the answer" was too strong; +#: the honest statement is that it cannot change the answer at any scale that +#: is not floating-point noise. +_DENSE_CHUNK_BYTES = 128 * 1024 * 1024 + +_LAST_REPORT = {} + + +def last_report(): + """Diagnostics from the most recent :func:`time_marginalize_bandlimited` call. + + Keys: ``upsample_factor`` (the largest used), ``factor_histogram`` + (factor -> row count, over the rows that were refined), ``n_refinements``, + ``sigma_t_min``, ``n_rows``, ``n_wrap_exposed_rows``, ``n_unmeasurable_rows``, + ``n_flat_rows``, ``n_refined_rows``. + + The diagnostic row counts are deliberately kept apart because they mean + different things: + + ``n_wrap_exposed_rows`` -- compatibility name for a resolvable peak sitting + inside ``EDGE_GUARD_FRACTION`` of a window edge. It is diagnostic only: the + row is reconstructed by the same reflected rule as every other measurable + under-resolved row. The name records the old implementation; there is no + periodic wrap in the current reconstruction. + + ``n_unmeasurable_rows`` -- ``lnL(t)`` non-finite around its maximum at every + stencil half-width, so no width can be justified. Given the historical value. + + ``n_flat_rows`` -- finite ``lnL(t)`` with no resolvable curvature: an + extrinsic sample with no signal in it. Nothing is wrong and nothing is paid. + + ``n_refined_rows`` is the count that matters for auditing a change: the + QUADRATURE RULE changes for these rows and for no others. Every other row -- + unmeasurable, flat, or already resolved -- is integrated by the caller's own + Simpson rule over the same domain. + + Read that precisely: it is a statement about the RULE, not about the returned + VALUE. The log-sum-exp offset also changes, from the historical single + global maximum over the whole block to a per-row maximum, and that applies to + every row including the unrefined ones. It is unavoidable on the refined + path (see :func:`_log_trapz_over_window`) and it has a visible consequence: + a row far enough below the block maximum that ``exp(lnL - global_max)`` + underflowed -- which happens once a batch spans more than ~745 nats, routine + at rho >~ 40 -- returned ``-inf`` historically and now returns a finite + value, refined or not. On rows that did not underflow the two agree to + floating-point rounding, not bit-for-bit. + """ + return dict(_LAST_REPORT) + + +def validate_time_quadrature(time_quadrature): + if time_quadrature not in TIME_QUADRATURE_CHOICES: + raise ValueError( + "time_quadrature must be one of {}, got {!r}".format( + TIME_QUADRATURE_CHOICES, time_quadrature)) + return time_quadrature + + +#: Flags in an assembled ILE argument string that the band-limited path REQUIRES, +#: and flags whose presence EXCLUDES it. This is the pipeline-side mirror of the +#: ``_tq_prereqs`` block in ``bin/integrate_likelihood_extrinsic_batchmode``: the +#: driver refuses at first-job time, which costs a queue-slot cycle, so the +#: workflow builder refuses at DAG-BUILD time on the same conditions. Kept here +#: rather than re-typed in the pipeline scripts for the reason the whole option is +#: validated through this module -- one list, one place to change it. +_PIPELINE_REQUIRED_ILE_FLAGS = ( + ('--time-marginalization', + 'without it ILE takes the non-time-marginalized branch, which has no time quadrature at all'), + ('--vectorized', + 'without it ILE calls the SCALAR time-marginalized likelihood, which has no quadrature argument'), + ('--gpu', + 'the band-limited path lives in the maintained NoLoop likelihood; --force-xpy runs it on numpy'), +) +_PIPELINE_EXCLUDING_ILE_FLAGS = ( + ('--rotation-slow', + 'time-DEPENDENT rho_sq: the band-limited argument does not hold'), + ('--freqresponse', + 'separate likelihood, not audited for this'), + ('--calibration-envelope-directory', + 'calibration marginalization: the reduction sums exp over realizations, untested here'), +) + + +ILE_TIME_QUADRATURE_FLAG = '--time-marginalization-quadrature' + +def _ile_tokens(ile_args): + """Tokenise an ILE argument string the way optparse will see it. + + Splits ``--flag=value`` (optparse accepts it, and a naive split does not) and + strips the quotes an ini file leaves behind. Returns ``(flags, pairs)`` where + ``pairs`` is the token list with values still attached in order. + """ + raw = str(ile_args).split() + toks = [] + for t in raw: + t = t.strip().strip('"').strip("'") + if not t: + continue + if t.startswith('--') and '=' in t: + k, v = t.split('=', 1) + toks.append(k) + toks.append(v) + else: + toks.append(t) + return toks + + +def _matches(flag, token): + """True if ``token`` is ``flag`` or a possible optparse abbreviation. + + optparse has no fixed minimum abbreviation length: in the current ILE parser + ``--g`` uniquely selects ``--gpu`` and ``--vec`` selects ``--vectorized``. + Ambiguous prefixes are rejected by ILE itself; treating them as a match here + can only move that refusal to DAG-build time, while imposing an invented + length floor falsely rejects legal configurations. + """ + if token == flag: + return True + return (flag.startswith(token) and token.startswith('--') + and len(token) > 2) + + +def find_time_quadrature_in_ile_args(ile_args): + """Every value given to ``--time-marginalization-quadrature`` in ``ile_args``. + + Returns a list, in order, so the caller can tell "absent" from "present once" + from "given twice with different values" -- optparse takes the LAST + occurrence, so a duplicate silently decides the quadrature. + """ + toks = _ile_tokens(ile_args) + out = [] + for n, t in enumerate(toks): + # optparse accepts unique long-option prefixes. The exact + # ``--time-marginalization`` flag wins as an exact match, but anything + # through the following '-' is a unique prefix of the quadrature flag. + # Treat those spellings exactly as ILE does or a hand-passed abbreviated + # bandlimited request can be invisible to the prerequisite guard. + is_quadrature = ( + t == ILE_TIME_QUADRATURE_FLAG + or (t.startswith('--time-marginalization-') + and ILE_TIME_QUADRATURE_FLAG.startswith(t)) + ) + if is_quadrature: + out.append(toks[n + 1] if n + 1 < len(toks) else None) + return out + + +def time_quadrature_pipeline_prereqs(time_quadrature, ile_args): + """Missing/violated prerequisites for ``time_quadrature`` in an ILE argument string. + + ``ile_args`` is the assembled ILE command line the workflow is about to write + (``args_ile.txt``). Returns a list of human-readable reasons; empty means the + configuration can honour the request. ``'simpson'`` -- the default -- always + returns an empty list, since it is what ILE does anyway. + + Refusing rather than ignoring is the point: a silently-inert accuracy option is + worse than an unavailable one, because a comparison campaign can be run against + it and believed. + """ + validate_time_quadrature(time_quadrature) + if time_quadrature == 'simpson': + return [] + toks = _ile_tokens(ile_args) + missing = [] + for flag, why in _PIPELINE_REQUIRED_ILE_FLAGS: + # Direction matters: a token satisfies a required flag when the FLAG starts + # with the TOKEN (the token is an abbreviation). The reverse test would let + # '--time-marginalization-quadrature' satisfy '--time-marginalization'. + if not any(_matches(flag, t) for t in toks): + missing.append("missing {} ({})".format(flag, why)) + for flag, why in _PIPELINE_EXCLUDING_ILE_FLAGS: + if any(_matches(flag, t) for t in toks): + missing.append("incompatible {} ({})".format(flag, why)) + return missing + + +def refuse_unhonourable_time_quadrature(time_quadrature, ile_args, where): + """Raise unless ``ile_args`` can honour ``time_quadrature``. + + The raise lives HERE, not at the call sites, so that it is executable in a unit + test: both pipeline scripts are top-level scripts that need real data before + they reach their guard, and a guard whose only coverage is "an ast walk found a + call by this name" survives being turned into a print. + """ + missing = time_quadrature_pipeline_prereqs(time_quadrature, ile_args) + if missing: + raise ValueError( + "time-marginalization quadrature {!r} was requested, but {} cannot honour it: " + "{}. Refusing rather than running the historical Simpson quadrature while " + "reporting that you asked for something else.".format( + time_quadrature, where, "; ".join(missing))) + + +def refuse_unless_time_quadrature_emitted(time_quadrature, ile_args, where): + """Raise unless the REQUESTED quadrature is the one the bytes actually carry. + + The prerequisite check above reads the prerequisites in ``ile_args`` but takes + the INTENT from the caller's parsed options, so it approves an argument string + that never received the flag at all. Three ways that happens in practice, all + ending in a silent fall back to Simpson while the pipeline logs the opposite: + + * a helper that predates the option argparse-errors while a re-run directory + still holds a STALE ``helper_ile_args.txt`` (the caller now also removes the + generated file first and checks the helper's exit status); + * ``--manual-extra-ile-args`` appends a second ``--time-marginalization-quadrature`` + after the helper's, and optparse takes the LAST occurrence; + * any future refactor that drops the emission. + + So this checks the bytes for the flag itself. ``time_quadrature`` of ``None`` + means nothing was requested, in which case the flag must be ABSENT unless the + user put it there by hand -- and if they did, it is validated and prerequisite + checked like any other request. + """ + found = find_time_quadrature_in_ile_args(ile_args) + if len(found) > 1: + raise ValueError( + "{} carries {} occurrences of {} ({!r}). optparse takes the LAST, so the " + "quadrature actually used would not be the one this workflow reports -- and " + "the .sub file would read as though it were. Refusing.".format( + where, len(found), ILE_TIME_QUADRATURE_FLAG, found)) + if time_quadrature is None: + if found: + # Set by hand (--manual-extra-ile-args or an ini). Not our flag, but it is + # about to run, so hold it to the same standard rather than none at all. + validate_time_quadrature(found[0]) + refuse_unhonourable_time_quadrature(found[0], ile_args, where) + return + if not found: + raise ValueError( + "time-marginalization quadrature {!r} was requested, but {} contains no {} at " + "all. The request was lost between the pipeline and the ILE arguments -- a " + "stale or version-skewed helper path can do exactly this. Refusing rather " + "than submitting a " + "campaign that would silently run Simpson.".format( + time_quadrature, where, ILE_TIME_QUADRATURE_FLAG)) + if found[0] != time_quadrature: + raise ValueError( + "time-marginalization quadrature {!r} was requested but {} carries {!r}. " + "Refusing.".format(time_quadrature, where, found[0])) + refuse_unhonourable_time_quadrature(time_quadrature, ile_args, where) + + +def bandlimited_upsample(x, factor, xpy=np): + """Zero-padded-FFT upsample of complex rows ``x`` (..., n) by ``factor``. + + Exact for a sequence of samples of a function band-limited below Nyquist and + periodic on the window, which is what the rholm timeseries are by + construction (they are inverse FFTs of a band-limited product). The output + has ``n*factor`` columns and reproduces the input exactly at every + ``factor``-th column. + + ODD ``n`` is not a special case to be waved through: an earlier version split + the spectrum at ``h = n//2`` unconditionally, which for odd ``n`` puts the + HIGHEST POSITIVE frequency at a negative frequency in the padded array. The + reconstruction then stays exact at the original samples -- so "it reproduces + the input" still passes -- while being wrong everywhere in between. Measured + against the analytic band-limited truth at ``factor=4``: max error 1.4e-12 at + ``n=614`` but 4.1e-1 at ``n=613``, 5.4e-1 at ``n=307`` and 6.0e-2 at + ``n=2457``. That matters because ``marginalization_time_grid`` produces ODD + ``npts`` at three of the five production sample rates -- 153 at srate 1024, + 307 at 2048 and 2457 at 16384 -- so the broken case was the common one. + + The Nyquist bin exists only for even ``n``, and is split evenly between + ``+fNyq`` and ``-fNyq``. For the rholm data it is empty anyway (the + two-sided weight construction never populates it), so that half is a + formality; the positive/negative boundary is not. + """ + factor = int(factor) + if factor == 1: + return x + n = x.shape[-1] + lead = x.shape[:-1] + X = xpy.fft.fft(x, axis=-1) + Xup = xpy.zeros(lead + (n * factor,), dtype=xpy.asarray(X).dtype) + n_pos = (n - 1) // 2 # DC plus n_pos strictly-positive bins + Xup[..., :n_pos + 1] = X[..., :n_pos + 1] + if n % 2 == 0: + nyq = X[..., n // 2] + Xup[..., n // 2] = 0.5 * nyq + Xup[..., -(n // 2)] = 0.5 * nyq + Xup[..., -n_pos:] = X[..., n // 2 + 1:] + else: + Xup[..., -n_pos:] = X[..., n_pos + 1:] + return xpy.fft.ifft(Xup, axis=-1) * factor + + +def reflected_bandlimited_upsample(x, factor, xpy=np): + """Upsample a finite row after a literal ``2*n`` even reflection. + + Upsampling the gathered row directly identifies its unlike endpoints and + can create global Gibbs ringing. Instead periodize + ``[x[0], ..., x[-1], x[-1], ..., x[0]]`` and return only the original + forward interval. Both periodic joins are value-continuous, every input + sample is reproduced exactly, and no unavailable samples outside the + integration domain are invented. The duplicated turning samples are + deliberate: the ``2*(n-1)`` reflection was less accurate in measured + realistic and adversarial cases because it locates the turn differently. + + This is a numerical boundary condition, not a claim that the physical + correlation reverses outside the caller's finite integration window. + """ + x = xpy.asarray(x) + factor = int(factor) + if factor == 1: + return x + n = x.shape[-1] + reflected = xpy.concatenate((x, xpy.flip(x, axis=-1)), axis=-1) + dense = bandlimited_upsample(reflected, factor, xpy=xpy) + return dense[..., :(n - 1) * factor + 1] + + +def peak_width_from_lnL(lnL_t, dx, xpy=np): + """Per-row Gaussian width ``sigma_t`` of ``exp(lnL_t)``, from its peak curvature. + + Uses a centred second difference of ``lnL`` (not of ``exp lnL``) about the + peak sample. For a Gaussian ``lnL`` this returns ``sigma`` EXACTLY at any + spacing, any peak-vs-grid phase, and any stencil half-width, because the + second difference of a parabola is its second derivative; that is what lets an + under-resolved peak still report its own width honestly, and it is why + stepping the stencil out over a ``-inf`` hole costs nothing. + + Returns ``(sigma_t, jmax, measurable)``. ``sigma_t`` is ``inf`` for a row + with non-negative curvature -- flat or monotone, where no refinement is + warranted. ``measurable`` distinguishes that legitimate case from a row whose + curvature could not be evaluated at all; the caller must not treat the two + alike, since "flat" means no refinement is NEEDED while "unmeasurable" means + none can be JUSTIFIED. + """ + n = lnL_t.shape[-1] + if n < 3: + raise ValueError("need at least 3 time samples to measure a peak width") + jmax = xpy.argmax(xpy.where(xpy.isfinite(lnL_t), lnL_t, -np.inf), axis=-1) + # numpy.take_along_axis was introduced in 1.15, while RIFT still declares a + # NumPy >=1.14 floor. Flatten the leading axes and use ordinary advanced + # indexing, which has the same semantics on NumPy and CuPy at that floor. + lead_shape = jmax.shape + flat_lnL = lnL_t.reshape((-1, n)) + row_index = xpy.arange(flat_lnL.shape[0]) + take = lambda j: flat_lnL[row_index, j.reshape(-1)].reshape(lead_shape) + + sigma = xpy.full(jmax.shape, np.inf, dtype=np.float64) + measurable = xpy.zeros(jmax.shape, dtype=bool) + for d in CURVATURE_STENCIL_HALFWIDTHS: + if 2 * d >= n: + break + jc = xpy.clip(jmax, d, n - 1 - d) + with np.errstate(invalid='ignore'): + # inf - inf is exactly the case being handled; the NaN it produces is + # the signal that this half-width straddles a hole, not an anomaly. + d2 = (take(jc - d) - 2.0 * take(jc) + take(jc + d)) / float(d * dx) ** 2 + fresh = xpy.isfinite(d2) & (~measurable) + if not bool(xpy.any(fresh)): + continue + neg = fresh & (d2 < 0) + sigma = xpy.where(neg, 1.0 / xpy.sqrt(xpy.where(neg, -d2, 1.0)), sigma) + measurable = measurable | fresh + if bool(xpy.all(measurable)): + break + return sigma, jmax, measurable + + +def required_upsample_factors(sigma, dx, xpy=np): + """Per-row power-of-two factor with ``dx/factor <= sigma/UPSAMPLE_SAFETY``. + + PER ROW, deliberately. A single block-wide factor is correct but ruinous: + the handful of rows near the source impose their resolution on every other + row in the batch, and the cost is dominated by the likelihood callback and + ``exp`` over ``n_rows * npts * factor`` points. Measured on the companion + O4c line at ``--n-chunk 10000``, srate 4096, one block-wide factor cost 18x + the Simpson likelihood call at rho=40 and 39x at rho=80 -- in the ILE inner + loop. Grouping by the derived factor leaves every row meeting its own + criterion while the broad majority stop paying for the sharpest few. + """ + need = UPSAMPLE_SAFETY * float(dx) / xpy.where(xpy.isfinite(sigma) & (sigma > 0), + sigma, np.inf) + need = xpy.where(need > 1.0, need, 1.0) + factor = xpy.exp2(xpy.ceil(xpy.log2(need))) + # 2**ceil(log2(need)) can land one power of two SHORT when log2 rounds down + # for a `need` a hair above a power of two -- erring LOW, i.e. silently + # under-resolving, which is the failure mode this whole module exists to + # remove. The criterion is `factor >= need`, so test exactly that and bump. + # (The margin can only ever be one ulp, so a single bump closes it.) + factor = xpy.where(factor < need, 2.0 * factor, factor) + # Clamp BEFORE the cast. A float factor above 2**63 wraps to a large + # NEGATIVE int64, which downstream `maximum(factors, 1)` turns into 1 -- so an + # unresolvably sharp row would be classified "nothing to refine" and silently + # given the coarse value, which is the failure this module exists to remove. + # Saturating at the ceiling instead sends it into the loop, which raises. + factor = xpy.where(factor > UPSAMPLE_FACTOR_MAX, + float(2 * UPSAMPLE_FACTOR_MAX), factor) + return factor.astype(np.int64) + + +def _default_simps(simps, xpy): + """The caller's Simpson rule, defaulting to scipy's ONLY on the numpy backend. + + scipy's ``simpson`` raises ``TypeError: Implicit conversion to a NumPy array is + not allowed`` on a cupy array, and that default is exactly how every + ``--vectorized --gpu`` run of this option crashed. Refuse rather than leave the + trap armed -- and note the two rules are not interchangeable even where both run + (the vendored GPU copy is an old scipy with ``even='avg'``), so a fallback row + must be integrated by the rule the caller's own likelihood uses. + """ + if simps is not None: + return simps + if xpy is not np: + raise ValueError( + "time marginalization: `simps` must be supplied for a non-numpy backend " + "-- scipy's Simpson rule cannot consume a device array, and the fallback " + "rows must use the rule the caller's own likelihood uses (on GPU, " + "optimized_gpu_tools.simps).") + from scipy import integrate + return getattr(integrate, 'simpson', None) or integrate.simps + + +def _require_time_independent_rho_sq(rho_sq, xpy=np, rule='band-limited'): + """Verify the load-bearing precondition rather than trusting the caller. + + A time-dependent self-term (the banded / slow-rotation response) would make the + refined ``lnL`` wrong in a way no downstream check would catch. + + Compare only where both sides are finite. A NaN self-term is NORMAL: the + defensive proposal component deliberately draws physically-extreme points where + the likelihood is NaN, and the historical path just returns NaN for that row and + lets the sampler move on. A bare ``==`` makes ``nan != nan`` trip this tripwire + and abort the whole ILE process, blaming a rotating-response path that is not + even in use. + """ + rho_col = rho_sq[..., :1] + _cmp = xpy.isfinite(rho_sq) & xpy.isfinite(xpy.broadcast_to(rho_col, rho_sq.shape)) + if not bool(xpy.all(xpy.where(_cmp, rho_sq == rho_col, True))): + raise NotImplementedError( + "%s time marginalization requires a time-independent rho_sq; the supplied " + "self-term varies with time (banded / rotating-response path)" % rule) + return rho_col + + +def _classify_rows(lnL_coarse, deltaT, npts, xpy=np): + """Which rule each row gets. THE SINGLE DEFINITION, shared with 'peak-local'. + + Returns ``(sigma, jmax, measurable, has_peak, flat, exposed, unmeasurable, + factors)``. A row is REFINED by the caller iff ``has_peak & (factors > 1)``. + + It lives here, and is called rather than copied, because + :mod:`RIFT.likelihood.time_marginalization_peak_local` promises to change WHERE the + refined grid is placed and nothing about WHICH rows get one. That promise was made + good by duplication and it did not survive: three clauses below -- reflection's + demotion of the edge guard, ``boundary_unresolved``, and the guard's Simpson routing + -- reached this module and not that one, and each showed up as peak-local silently + returning a lower-accuracy value than this function for the same row. Duplicated + policy that MUST agree is policy that will eventually not. + + The boundary diagnostic applies only to rows that HAVE a resolvable peak: a row whose + lnL(t) is constant -- an extrinsic sample in an antenna null, where kappa is + numerically zero -- has an argmax of 0 by convention and would otherwise be reported + as boundary-exposed. That is harmless numerically (Simpson is exact on a constant) + but it makes the diagnostic lie: measured on a random-sky batch of 4000, it reported + 810 "wrap-exposed" rows, which in a production log reads as a mis-centred window + rather than as 810 rows with no signal in them. + + ``boundary_unresolved``: at the first/last sample the centred stencil is clipped + inward, so for a severely under-resolved endpoint peak it can see positive curvature + away from the maximum and label a strongly varying row "flat". That would silently + retain Simpson for exactly the truncated-boundary case we intend to report and + reconstruct. Such rows get a small seed factor; dense-grid remeasurement takes over + as soon as the reflected peak is measurable. + + ``exposed`` REPORTS possible physical truncation and selects nothing. Reflection + removes the endpoint value jump, so neither boundary proximity nor a tail threshold + may select a Simpson fallback: crossing an arbitrary threshold cannot silently change + likelihood quality, and raising on such a row can be read upstream as a waveform + failure and silently excise that configuration. + """ + sigma, jmax, measurable = peak_width_from_lnL(lnL_coarse, deltaT, xpy=xpy) + guard = max(1, int(npts * EDGE_GUARD_FRACTION)) + finite_lnL = xpy.isfinite(lnL_coarse) + row_max = xpy.max(xpy.where(finite_lnL, lnL_coarse, -np.inf), axis=-1) + row_min = xpy.min(xpy.where(finite_lnL, lnL_coarse, np.inf), axis=-1) + varies = xpy.isfinite(row_max) & xpy.isfinite(row_min) & (row_max > row_min) + boundary_unresolved = (measurable & (~xpy.isfinite(sigma)) & varies + & ((jmax == 0) | (jmax == npts - 1))) + has_peak = measurable & (xpy.isfinite(sigma) | boundary_unresolved) + flat = measurable & (~xpy.isfinite(sigma)) & (~boundary_unresolved) + exposed = has_peak & ((jmax < guard) | (jmax > npts - 1 - guard)) + # Counted unconditionally, NOT `& ~exposed`: an all -inf row also has an argmax of 0, + # so a conditional counter would hide it behind the edge guard. + unmeasurable = ~measurable + factors = xpy.maximum(required_upsample_factors(sigma, deltaT, xpy=xpy), 1) + factors = xpy.where(boundary_unresolved, xpy.maximum(factors, 4), factors) + return (sigma, jmax, measurable, has_peak, flat, exposed, unmeasurable, factors) + + +def _safe_offset(off, xpy=np): + """Log-sum-exp offset, guarded for a row that is ``-inf`` everywhere. + + Such a row has zero likelihood over the whole window, and the historical path + returns ``-inf`` for it (its GLOBAL offset is finite, so every term underflows + to zero and ``log(0)`` follows). A per-row offset would instead compute + ``-inf - (-inf) = NaN`` and hand back a NaN that propagates into the sampler + weights. Substituting a finite offset reproduces ``-inf`` exactly. + """ + return xpy.where(xpy.isfinite(off), off, 0.0) + + +def _log_simps_rows(lnL_t, dx, simps, xpy=np): + """``log \\int exp(lnL) dt`` by the caller's Simpson rule, per row. + + The caller's rule, not a private copy: on GPU the likelihood integrates with + ``optimized_gpu_tools.simps``, so a scipy copy here would agree on CPU and + quietly disagree on the device -- and the whole point of this path is to + reproduce what the historical code would have returned for these rows. + """ + off = _safe_offset(xpy.max(lnL_t, axis=-1, keepdims=True), xpy=xpy) + return off[..., 0] + xpy.log(simps(xpy.exp(lnL_t - off), dx=float(dx), axis=-1)) + + +def _log_trapz_over_window(lnL_dense, dx_dense, npts_coarse, factor, xpy=np): + """``log \\int exp(lnL) dt`` by trapezoid over the ORIGINAL window span. + + The dense grid returned by the FFT upsample is periodic on + ``[t_0, t_0 + npts*deltaT)``, i.e. it carries ``factor-1`` samples PAST the + last coarse sample. Those lie across the periodic wrap and are dropped, so + the integration domain is exactly ``[t_0, t_{npts-1}]`` -- identical to the + domain Simpson used. Changing the domain would have been a second, + confounded change. + + The log-sum-exp offset is PER ROW and taken on the DENSE grid, not the single + global coarse maximum the Simpson path uses. It has to be: the whole point of + refining the grid is that the true peak sits between coarse samples, so the + dense maximum can exceed the coarse one -- by thousands of nats for a sharp + peak -- and offsetting by the coarse maximum overflows ``exp()`` precisely in + the regime this quadrature exists to serve. (A per-row offset also avoids the + underflow-to-``log(0) = -inf`` that a shared global offset gives rows far below + the block maximum.) The result is offset-invariant, so this is a numerical + choice and not a change of estimator. + """ + last = (npts_coarse - 1) * factor + v = lnL_dense[..., :last + 1] + w = xpy.full(v.shape[-1], dx_dense, dtype=np.float64) + w[0] *= 0.5 + w[-1] *= 0.5 + off = _safe_offset(xpy.max(v, axis=-1, keepdims=True), xpy=xpy) + return off[..., 0] + xpy.log(xpy.sum(xpy.exp(v - off) * w, axis=-1)) + + +def draw_piecewise_linear_log_posterior(lnL_t, dx, t0=0.0, + uniforms=None, xpy=np): + """Draw one continuous time per row from a nodal log density. + + Between adjacent nodes the *density* ``exp(lnL)`` is linear. Its interval + mass is therefore exactly the trapezoid used by the refined quadrature. We + first choose an interval by those masses, then invert the linear-density CDF + analytically inside it. The result has no output lattice: ``dx`` describes + the representation's knots, not the support of the returned variate. + + ``uniforms`` may be supplied as shape ``(n_rows, 2)``. The first variate + selects the interval and the second selects the position inside it. This is + both the reproducibility seam and the way CPU/GPU tests ask the two backends + exactly the same question. If omitted, numpy's global generator is used, + preserving the driver's existing ``--seed`` contract. + + Returns ``(times, lnL_at_times)`` in the input backend. ``-inf`` nodes are + supported and carry zero density. NaN/+inf nodes and rows with no positive + finite mass are rejected rather than assigned an invented timestamp. + """ + values = xpy.asarray(lnL_t) + if values.ndim == 1: + values = values[xpy.newaxis, :] + if values.ndim != 2 or values.shape[-1] < 2: + raise ValueError("lnL_t must have shape (n_rows, n_time>=2)") + if not bool(xpy.all((~xpy.isnan(values)) & (~xpy.isposinf(values)))): + raise ValueError("continuous time posterior contains NaN or +inf") + + n_rows, n_time = values.shape + if uniforms is None: + uniforms = np.random.random((n_rows, 2)) + uniforms = xpy.asarray(uniforms, dtype=np.float64) + if uniforms.shape != (n_rows, 2): + raise ValueError("uniforms must have shape (n_rows, 2)") + if not bool(xpy.all((uniforms >= 0.0) & (uniforms < 1.0))): + raise ValueError("uniforms must lie in [0, 1)") + + finite = xpy.isfinite(values) + off = xpy.max(xpy.where(finite, values, -np.inf), axis=-1) + if not bool(xpy.all(xpy.isfinite(off))): + raise ValueError("time posterior has no finite positive mass") + density = xpy.where(finite, xpy.exp(values - off[:, xpy.newaxis]), 0.0) + interval_mass = 0.5 * float(dx) * (density[:, :-1] + density[:, 1:]) + total = xpy.sum(interval_mass, axis=-1) + if not bool(xpy.all(xpy.isfinite(total) & (total > 0.0))): + raise ValueError("time posterior has no finite positive mass") + + cdf = xpy.cumsum(interval_mass, axis=-1) / total[:, xpy.newaxis] + # `uniforms < 1` guarantees an interval, but clip defensively against a + # backend whose final cumsum rounds a hair below one. + # `<=` skips a leading/embedded zero-mass plateau even when the supplied + # variate is exactly zero or exactly on a cumulative boundary. `<` would + # select a zero-density interval at u=0 and return lnL=-inf for a posterior + # that has positive mass later in the window. + interval = xpy.sum(cdf <= uniforms[:, :1], axis=-1).astype(np.int64) + interval = xpy.minimum(interval, n_time - 2) + row = xpy.arange(n_rows) + a = density[row, interval] + b = density[row, interval + 1] + delta = b - a + # A pseudo-random float can (very rarely) be exactly zero. On an interval + # whose left endpoint has zero density, the literal inverse-CDF endpoint + # would then return lnL=-inf and could be silently excised downstream. Use + # the centre of the lowest float64 RNG bin for that one endpoint, matching + # the open-interval variate required by a continuous posterior draw. + r = xpy.maximum(uniforms[:, 1], 0.5 * np.finfo(float).eps) + + # For density p(u)=a+(b-a)u on u in [0,1], inverse-CDF sampling gives + # p(u)^2 = a^2 + r*(b^2-a^2). Use the uniform limit when the interval is + # numerically flat to avoid cancellation in (p-a)/(b-a). + scale = xpy.maximum(xpy.maximum(xpy.abs(a), xpy.abs(b)), 1.0) + flat = xpy.abs(delta) <= 16.0 * np.finfo(float).eps * scale + # Scale locally before squaring. A selected far-tail interval can have + # representable endpoint densities whose squares underflow; the CDF inverse + # must not turn that positive interval into zero density. + local_scale = xpy.maximum(a, b) + a_scaled = a / local_scale + b_scaled = b / local_scale + endpoint_density = local_scale * xpy.sqrt(xpy.maximum( + 0.0, a_scaled * a_scaled + + r * (b_scaled * b_scaled - a_scaled * a_scaled))) + frac = xpy.where(flat, r, (endpoint_density - a) / + xpy.where(flat, 1.0, delta)) + frac = xpy.clip(frac, 0.0, 1.0) + drawn_density = a + delta * frac + times = float(t0) + (interval.astype(np.float64) + frac) * float(dx) + lnL_draw = off + xpy.log(drawn_density) + return times, lnL_draw + + +def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, + phase_marginalization=False, simps=None, + lnL_coarse=None, return_time_draw=False, + draw_uniforms=None, t0=0.0, xpy=np): + """``log \\int dt exp(lnL(t))`` with the time grid refined to the integrand. + + Parameters + ---------- + kappa : (n_extrinsic, npts) complex + The accumulated data term on the coarse grid, exactly as the caller + already builds it. Band-limited below Nyquist by construction. + rho_sq : (n_extrinsic, npts) float + The template self-term. MUST be constant along the time axis on this + path; that is what makes ``lnL(t)`` a pointwise function of a + band-limited quantity. Checked, not assumed. + loglikelihood : callable + ``f(kappa_term, rho_sq) -> lnL``, the same callback the caller passes to + the coarse path (default helper, phase- or distance-marginalized). + simps : callable, optional + The caller's Simpson rule, ``simps(y, dx=..., axis=-1)``, used for rows + that fall back to the historical path. Defaults to scipy's -- which + RAISES on a cupy array, so the GPU caller must supply its own. + lnL_coarse : array, optional + ``loglikelihood`` already evaluated on the coarse grid. The caller + normally has it; passing it avoids re-evaluating the callback over + ``n_extrinsic * npts`` points, which for the distance-marginalized + callback is a table interpolation over millions of points and is the + difference between "no extra likelihood evaluations" being true and + being nearly true. + return_time_draw : bool, optional + Also return one continuous conditional-posterior draw per row and its + instantaneous log likelihood. Refined rows use the exact same validated + dense representation as the trapezoid integral. Unrefined rows are + already resolved and are drawn continuously between their coarse knots. + draw_uniforms : array, optional + Shape ``(n_extrinsic, 2)`` uniforms for deterministic draws. Omit to use + numpy's global RNG, matching the batch driver's ``--seed`` behavior. + t0 : float, optional + Time of the first coarse knot; returned draws are in this coordinate. + + Returns + ------- + lnL : (n_extrinsic,) float + With ``return_time_draw=True``, returns + ``(lnL, time_draw, lnL_at_draw)``. + """ + simps = _default_simps(simps, xpy) + + kappa = xpy.asarray(kappa) + rho_sq = xpy.asarray(rho_sq) + npts = kappa.shape[-1] + n_rows = kappa.shape[0] + deltaT = float(deltaT) + + _require_time_independent_rho_sq(rho_sq, xpy=xpy, rule='band-limited') + rho_col = rho_sq[..., :1] + + _term = (lambda k: xpy.abs(k)) if phase_marginalization else (lambda k: k.real) + if lnL_coarse is None: + lnL_coarse = loglikelihood(_term(kappa), rho_sq) + + # THE SINGLE DEFINITION, shared with 'peak-local' -- see _classify_rows. + (sigma, jmax, measurable, has_peak, flat, exposed, unmeasurable, + factors) = _classify_rows(lnL_coarse, deltaT, npts, xpy=xpy) + + # A row is REFINED only if it has a trustworthy peak AND the derivation + # actually asks for a finer grid. Reflection removes the endpoint value + # jump, so neither boundary proximity nor a tail threshold selects a Simpson + # fallback. ``exposed`` reports possible physical truncation only. + # Everything else -- unmeasurable, flat, or already resolved -- gets the + # historical Simpson value. (The log-sum-exp offset changes for every row -- + # see last_report() -- so this is a statement about the rule, not a promise + # that unrefined rows come back bit-identical.) + # The alternative, letting an unrefined row fall through to + # a coarse TRAPEZOID, is numerically a non-event but silently changes the + # rule for rows this option was never meant to touch, and costs the property + # a reviewer can actually check. (Trapezoid is in fact slightly the better + # rule on a resolved integrand -- 5e-6 nats against an analytic truth, versus + # Simpson's 5e-6 the other way -- so this trades nothing measurable for an + # auditable claim.) + refined = has_peak & (factors > 1) + + out = _log_simps_rows(lnL_coarse, deltaT, simps, xpy=xpy) + time_draw = None + lnL_at_draw = None + if return_time_draw: + if draw_uniforms is None: + draw_uniforms = np.random.random((n_rows, 2)) + draw_uniforms = xpy.asarray(draw_uniforms, dtype=np.float64) + if draw_uniforms.shape != (n_rows, 2): + raise ValueError("draw_uniforms must have shape (n_rows, 2)") + # Seed every row from the already-resolved coarse representation. Rows + # refined below are overwritten with draws from their final validated + # dense representation; flat/already-resolved rows remain continuous + # rather than being snapped back to a coarse knot. + time_draw, lnL_at_draw = draw_piecewise_linear_log_posterior( + lnL_coarse, deltaT, t0=t0, uniforms=draw_uniforms, xpy=xpy) + + hist = {} + n_refine_total = 0 + sigma_seen = np.inf + for f in xpy.unique(xpy.where(refined, factors, 1)): + f = int(f) + if f == 1: + continue + sel = refined & (factors == f) + n_sel = int(xpy.sum(sel)) + if not n_sel: + continue + idx = xpy.where(sel)[0] + vals, f_used, n_ref, s_min, drawn_t, drawn_lnL = _integrate_group( + kappa[idx], rho_col[idx], npts, deltaT, f, loglikelihood, _term, + draw_uniforms_rows=(draw_uniforms[idx] if return_time_draw else None), + t0=t0, xpy=xpy) + out[idx] = vals + if return_time_draw: + time_draw[idx] = drawn_t + lnL_at_draw[idx] = drawn_lnL + hist[int(f_used)] = hist.get(int(f_used), 0) + n_sel + n_refine_total += n_ref + sigma_seen = min(sigma_seen, s_min) + + _LAST_REPORT.clear() + _LAST_REPORT.update( + upsample_factor=max(hist) if hist else 1, + factor_histogram=dict(hist), + n_refinements=n_refine_total, + sigma_t_min=sigma_seen, + n_rows=n_rows, + n_wrap_exposed_rows=int(xpy.sum(exposed)), + n_unmeasurable_rows=int(xpy.sum(unmeasurable)), + n_flat_rows=int(xpy.sum(flat)), + n_refined_rows=int(xpy.sum(refined)), + ) + if return_time_draw: + return out, time_draw, lnL_at_draw + return out + + +def _integrate_group(kappa_rows, rho_col_rows, npts, deltaT, factor, + loglikelihood, _term, draw_uniforms_rows=None, t0=0.0, + xpy=np): + """Refine and integrate one group of rows that share a derived factor. + + Returns ``(values, factor_used, n_refinements, sigma_dense_min, + time_draws, lnL_at_draws)``. The final two entries are ``None`` unless + ``draw_uniforms_rows`` is supplied. + """ + n_rows = kappa_rows.shape[0] + n_refine = 0 + while True: + if factor > UPSAMPLE_FACTOR_MAX: + raise RuntimeError( + "band-limited time marginalization needs an upsampling factor above " + "the ceiling UPSAMPLE_FACTOR_MAX=%d (deltaT=%.3e s). This is far " + "beyond what the band limit can justify: suspect a pathological " + "lnL(t), not an under-resolved one." % (UPSAMPLE_FACTOR_MAX, deltaT)) + + dx_dense = deltaT / factor + # Chunk the extrinsic axis so one dense temporary stays inside the + # working-set budget. Rows are independent; this cannot change results. + # The FFT period is 2*n after reflection; budget for it and the forward + # kappa/rho/lnL temporaries. + per_row = npts * factor * 16 * 8 + chunk = max(1, min(n_rows, int(_DENSE_CHUNK_BYTES // max(per_row, 1)))) + + pieces = [] + draw_time_pieces = [] + draw_lnL_pieces = [] + sigma_dense_min = np.inf + for start in range(0, n_rows, chunk): + k_up = reflected_bandlimited_upsample( + kappa_rows[start:start + chunk], factor, xpy=xpy) + rho_up = xpy.broadcast_to(rho_col_rows[start:start + chunk], k_up.shape) + lnL_up = loglikelihood(_term(k_up), rho_up) + s_d, _, meas = peak_width_from_lnL(lnL_up, dx_dense, xpy=xpy) + s_d = xpy.where(meas, s_d, np.inf) + sigma_dense_min = min(sigma_dense_min, float(xpy.min(s_d))) + pieces.append(_log_trapz_over_window(lnL_up, dx_dense, npts, factor, xpy=xpy)) + if draw_uniforms_rows is not None: + drawn_t, drawn_lnL = draw_piecewise_linear_log_posterior( + lnL_up, dx_dense, t0=t0, + uniforms=draw_uniforms_rows[start:start + chunk], xpy=xpy) + draw_time_pieces.append(drawn_t) + draw_lnL_pieces.append(drawn_lnL) + + # The assertion that turns the derivation into a guarantee: the width + # remeasured on the grid we actually integrated on must still satisfy the + # criterion. A coarse-grid estimate can be optimistic when the peak is + # strongly non-Gaussian; this catches that and pays for another doubling + # instead of reporting a number it cannot defend. + if (not np.isfinite(sigma_dense_min)) or dx_dense <= sigma_dense_min / UPSAMPLE_SAFETY: + values = xpy.concatenate(pieces) if len(pieces) > 1 else pieces[0] + drawn_t = (xpy.concatenate(draw_time_pieces) if len(draw_time_pieces) > 1 + else (draw_time_pieces[0] if draw_time_pieces else None)) + drawn_lnL = (xpy.concatenate(draw_lnL_pieces) if len(draw_lnL_pieces) > 1 + else (draw_lnL_pieces[0] if draw_lnL_pieces else None)) + return values, factor, n_refine, sigma_dense_min, drawn_t, drawn_lnL + + factor *= 2 + n_refine += 1 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_posterior.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_posterior.py new file mode 100644 index 000000000..5a254d968 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_posterior.py @@ -0,0 +1,195 @@ +"""Continuous draws from an interpolated time posterior. + +The ILE time-marginalization likelihood is evaluated on a regular FFT grid. +When sub-sample interpolation is enabled, exporting a time by choosing one of +those grid points throws that resolution away. This module draws directly +from ``exp(CubicSpline(t, lnL))`` instead. + +The sampler uses a piecewise-constant rejection envelope. On every spline +interval the envelope is the exact maximum of the cubic (endpoints plus all +stationary points), so accepted draws are continuous and have the requested +interpolated density without introducing another export lattice. +""" + +import numpy as np +from scipy.interpolate import CubicSpline, PchipInterpolator + + +TIME_POSTERIOR_EXPORT_MODES = ("auto", "continuous", "grid") +TIME_POSTERIOR_WORKING_SET_MAX = 512 * 1024 * 1024 +TIME_POSTERIOR_BYTES_PER_CELL = 128 + + +def validate_time_posterior_working_set(n_rows, n_times, + limit=TIME_POSTERIOR_WORKING_SET_MAX): + """Refuse a dense export before its dominant full matrices can exhaust GPU RAM.""" + if not isinstance(n_rows, (int, np.integer)) or n_rows < 0: + raise ValueError("n_rows must be a non-negative integer") + if not isinstance(n_times, (int, np.integer)) or n_times < 2: + raise ValueError("n_times must be an integer >= 2") + estimated = int(n_rows) * int(n_times) * TIME_POSTERIOR_BYTES_PER_CELL + if estimated > int(limit): + raise MemoryError( + "continuous time export would require an unsafe dense working set " + "(estimated at least {:.3g} MiB for {} rows x {} times; limit {:.3g} " + "MiB). Reduce --fairdraw-extrinsic-output-n-max or use " + "--time-posterior-export grid.".format( + estimated / 2.0**20, n_rows, n_times, int(limit) / 2.0**20)) + return estimated + + +def legacy_time_interpolation_enabled(value): + """Parse the LISA driver's historical boolean/string interpolation flag.""" + normalized = str(value).strip().lower() + if normalized in ("true", "t", "yes", "y", "1", "on", "cubic", "sinc"): + return True + if normalized in ("false", "f", "no", "n", "0", "off", "none", "nearest"): + return False + raise ValueError( + "--interpolate-time: unrecognised LISA value {!r}; use a boolean, " + "nearest, cubic, or sinc".format(value)) + + +def resolve_time_posterior_export_mode(requested, time_interpolation, + continuous_available=True): + """Resolve ``auto`` against interpolation and driver capabilities. + + ``continuous_available=False`` affects only ``auto``: it preserves the + driver's historical grid export. An explicit ``continuous`` request is + returned unchanged so the caller can reject it with a capability-specific + error instead of silently downgrading it. + """ + if requested not in TIME_POSTERIOR_EXPORT_MODES: + raise ValueError("unknown time-posterior export mode %r" % (requested,)) + if requested == "auto": + return ("continuous" if continuous_available and + time_interpolation != "nearest" else "grid") + return requested + + +def _interval_log_envelopes(spline, knots, log_values): + """Return the exact maximum of a cubic spline on every knot interval.""" + maxima = np.maximum(log_values[:-1], log_values[1:]).astype(float, copy=True) + # SciPy 1.0's PPoly.roots() has no ``extrapolate`` keyword. Calling it + # without the keyword is cross-version-safe because the in-domain test + # below already rejects every extrapolated root. + roots = np.asarray(spline.derivative().roots(), dtype=float) + roots = roots[np.isfinite(roots)] + if roots.size: + interval = np.searchsorted(knots, roots, side="right") - 1 + interval = np.clip(interval, 0, len(knots) - 2) + inside = (roots >= knots[interval]) & (roots <= knots[interval + 1]) + for i, value in zip(interval[inside], np.asarray(spline(roots[inside]))): + maxima[i] = max(maxima[i], float(value)) + return maxima + + +def draw_continuous_time_posterior(tvals, lnlt, rng=None, max_attempts=100000): + """Draw one continuous time per row from an interpolated ``lnL(t)``. + + Parameters + ---------- + tvals : array-like, shape (n_time,) + Strictly increasing time knots. + lnlt : array-like, shape (n_rows, n_time) or (n_time,) + Log likelihood at the knots. + rng : numpy RNG-like object, optional + Must provide ``choice`` and ``uniform``. The default is + ``numpy.random`` so the driver's existing ``--seed`` contract remains + unchanged. + max_attempts : int, optional + Maximum rejection proposals per row. Exhaustion raises rather than + allowing a pathological interpolant or RNG to hang an ILE job. + + Returns + ------- + times, log_likelihoods : ndarray + One accepted continuous draw and its interpolated log likelihood for + each input row. + """ + tvals = np.asarray(tvals, dtype=float) + values = np.asarray(lnlt, dtype=float) + if values.ndim not in (1, 2): + raise ValueError("lnlt must be a 1-D or 2-D array") + one_row = values.ndim == 1 + values = np.atleast_2d(values) + if tvals.ndim != 1 or tvals.size < 2 or not np.all(np.diff(tvals) > 0): + raise ValueError("tvals must be a strictly increasing 1-D grid") + if values.shape[1] != tvals.size: + raise ValueError("lnlt's final axis must match tvals") + if np.any(np.isnan(values)) or np.any(np.isposinf(values)): + raise ValueError("continuous time export does not accept NaN or +inf lnL(t)") + if not isinstance(max_attempts, (int, np.integer)) or max_attempts <= 0: + raise ValueError("max_attempts must be a positive integer") + if rng is None: + rng = np.random + + widths = np.diff(tvals) + times = np.empty(values.shape[0], dtype=float) + log_likelihoods = np.empty(values.shape[0], dtype=float) + for row, log_values in enumerate(values): + finite = np.isfinite(log_values) + if not np.any(finite): + raise ValueError("time posterior has no finite positive mass") + + if np.all(finite): + # Preserve the requested cubic interpolation of lnL for the normal + # path. Work relative to the envelope maximum to avoid overflow. + spline = CubicSpline(tvals, log_values) + maxima = _interval_log_envelopes(spline, tvals, log_values) + shift = float(np.max(maxima)) + envelope_mass = widths * np.exp(maxima - shift) + + def evaluate(candidate): + log_candidate = float(spline(candidate)) + interval_maximum = maxima[interval] + ratio = np.exp(log_candidate - interval_maximum) + return log_candidate, min(1.0, float(ratio)) + else: + # -inf is a valid zero-posterior-mass value. A log spline cannot + # represent it. Interpolate the shifted density with PCHIP, which + # preserves non-negativity and the zero knots without manufacturing + # the huge ringing that replacing -inf by an arbitrary log floor can + # cause. This branch is only used for rows containing -inf. + shift = float(np.max(log_values[finite])) + density_values = np.zeros_like(log_values) + density_values[finite] = np.exp(log_values[finite] - shift) + spline = PchipInterpolator(tvals, density_values) + maxima = _interval_log_envelopes( + spline, tvals, density_values) + maxima = np.maximum(maxima, 0.0) + envelope_mass = widths * maxima + + def evaluate(candidate): + density = max(0.0, float(spline(candidate))) + interval_maximum = maxima[interval] + ratio = density / interval_maximum if interval_maximum > 0 else 0.0 + log_candidate = shift + np.log(density) if density > 0 else -np.inf + return log_candidate, min(1.0, ratio) + + total = float(np.sum(envelope_mass)) + if not np.isfinite(total) or total <= 0: + raise ValueError("time posterior has no finite positive mass") + probabilities = envelope_mass / total + + for _attempt in range(max_attempts): + interval = int(rng.choice(len(widths), p=probabilities)) + candidate = float(rng.uniform(tvals[interval], tvals[interval + 1])) + log_candidate, accept_probability = evaluate(candidate) + # A float RNG can return exactly zero. With ``<=``, a proposal at + # a zero-density endpoint then satisfies 0 <= 0 and escapes with + # lnL=-inf, which downstream fair-draw filtering can silently + # excise. Zero posterior density must never be accepted. + if (accept_probability > 0.0 and + float(rng.uniform()) < accept_probability): + times[row] = candidate + log_likelihoods[row] = log_candidate + break + else: + raise RuntimeError( + "continuous time-posterior rejection sampler exhausted " + "{} proposals for row {}".format(max_attempts, row)) + + if one_row: + return times[0], log_likelihoods[0] + return times, log_likelihoods diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/DESIGN_multiapprox_marginalization.md b/MonteCarloMarginalizeCode/Code/RIFT/misc/DESIGN_multiapprox_marginalization.md new file mode 100644 index 000000000..311ec4f0e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/DESIGN_multiapprox_marginalization.md @@ -0,0 +1,125 @@ +# Cross-model marginalization on a shared grid + +A record of what `create_event_parameter_pipeline_BasicMultiApproxIteration` +does and why, as of 2026-08-27. It is expected to be superseded; where this +note and the code disagree, the code and +`test/test_multiapprox_marginalization.py` are live truth. + +## What the workflow computes + +Running several waveform models is not the same as running several pipelines. +This builder evaluates every model on **one shared intrinsic grid** and +marginalizes over models at each grid point: + + L_marg(lambda) = sum_m p(m) L_m(lambda) + +The iteration loop fits `L_marg`, so the adaptive grid refines on the +model-marginalized posterior — points get placed where the *mixture* has +support, including where one model is good and another is not. This is the +difference from running N single-model pipelines and mixing posteriors +afterwards: there, each run's grid refines on its own model's posterior, and no +amount of post-hoc mixing recovers support that neither run sampled. + +The combination is **linear in L**, not in lnL. Averaging lnL would give the +geometric mean — a logarithmic opinion pool — which is not Bayesian model +marginalization and is dominated differently when models disagree. + +## Where each piece lives + +| stage | model-tagged? | why | +|---|---|---| +| ILE input grid `overlap-grid-N.xml.gz` | no | the shared grid; all models evaluate the same lambda | +| ILE output → `approx__consolidated_N.composite` | yes | per-model evaluations, kept apart | +| `unify.sh` → `all.net` | no | pools every model, **marginalizes** | +| in-loop CIP → `overlap-grid-N+1` | no | one fit over the marginalized net | +| `unify_model.sh` → `approx__all.net` | yes | one model only, for the fork | +| terminal CIP → `approx__overlap-grid-N` + `+annotation.dat` | yes | per-model posterior **and** ln Z | +| extrinsic ILE / convert / resample / cat | yes | per model | +| `util_CombineApproximantPosteriors.py` | — | mixture weighted by p(m) Z_m | + +**The loop merges; the terminal stage forks.** The marginalized grid is what +should steer exploration, but the deliverables of a systematics study are the +per-model posteriors and their evidences — so those are produced separately, +left on disk, and combined only at the very end. + +## The marginalization is in util_CleanILE.py + +It keys on the intrinsic parameters (rounded to 5 decimals), so the same lambda +under several models collapses to one entry. `--model-group-regex` takes the +model label from each input filename and switches on a **two-level** combine: + +* **within** a model, replicas estimate one number → ntot-weighted linear mean; +* **across** models, the quantity itself differs → marginalize with p(m). + +Flat pooling of everything — the behaviour without the flag — silently uses +*replica counts* as model weights. On a two-model point where one model has two +replicas and the other one, with lnL 100 and 104, flat pooling gives 102.937 +where uniform-prior marginalization gives 103.325: **0.39 nats**, growing with +model disagreement and replica imbalance. Correct for replicas of one model, +wrong across models, which is why it is opt-in and why one model reduces to it +exactly (verified byte-identical). + +`p(m)` is a *prior over waveform models*, not sampling effort. `ntot` is +sampling effort. Conflating them is the defect above. + +## Partial coverage changes the estimator, point by point + +A lambda evaluated under only some models is marginalized over that subset, so +the estimator differs across the grid — a model- and lambda-dependent change in +the effective prior. Causes: the `sigmaOverL > 0.9` resolution cut firing for +one model and not another, or one model's ILE failing. + +This is reported (never silent) and `--require-all-approx` drops such points +instead. Dropping is not free either — it biases toward regions where the +*worst-resolved* model converged — so it is a choice, not a default. + +## Prior work + +Done before by hand (Yelikar; Jan), not efficiently in production. The claim +here is the production workflow — one DAG, one shared adaptive grid, evidence- +weighted recombination — not the idea of waveform marginalization. + +## Known limits + +* The final mixture resamples with replacement, so a model carrying most of the + weight can be drawn more times than it has samples. Reported as a warning; + it degrades effective sample size without changing the row count. +* `p(m)` is uniform unless `--approx-prior` is given, and applies at both the + in-loop marginalization and the final mixture. They are the same p(m) and are + passed to both from one option. +* A large ln Z gap between waveform models usually means the models disagree by + far more than the statistical error, not that one is right. The combiner + warns above 0.99 mixture weight. + +## Also fixed here + +This builder is not reachable from `util_RIFT_pseudo_pipe.py` (`--pipeline-builder` +offers only BasicIteration, AlternateIteration and Hyperpipe) and asimov drives +pseudo_pipe, so it is a hand-run `bin/` script with no caller — which is how the +following survived. All were read off an emitted two-approximant DAG. + +* The terminal extrinsic stage read `overlap-grid-` while the + final CIP wrote `overlap-grid-`: a guard, + `if not ('it' in globals())`, preserved an `it` the iteration loop had already + left one short. At one iteration it read the raw seed grid. +* `join_grids.sh` interpolated `$(macroapprox)` **inside a bash script** — + command substitution, not a condor macro, since the `.sub` passes only + `$1 $2`. It matched nothing. Gone now that the loop grid is not model-tagged. +* `CIP.sub`'s `initialdir` named `iteration_N_cip` while the mkdir loop created + `approx__iteration_N_cip`; the consolidate/unify log directories were the + mirror image. Either holds the job on the execute node, and no DAG-shape + check sees it — hence `test_every_job_directory_exists`. +* `parent_fit_node` was a single variable spanning `for it: for approx:`, so + model B's iteration-0 ILE waited on model A's iteration-1 convert. The models + were serialized into one chain and cross-coupled; they now run in parallel. +* The terminal convert wrote one untagged `posterior_samples-N.dat` from every + model, while the convergence test read a per-model name nothing wrote. + +## Testing status + +`test/test_multiapprox_marginalization.py` covers the combination arithmetic +against hand-computed values, and the emitted DAG's shape. **Neither validates +the inference on data.** The prototype for that is a rerun of one of the +Jan/Yelikar cases — cheap, and with large between-model differences — framed as +a reproduction claim first and an efficiency claim second. Nothing from this +workflow belongs in a paper until that runs. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/dag_utils_generic.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/dag_utils_generic.py index f3f516ce0..832167dcb 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/misc/dag_utils_generic.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/dag_utils_generic.py @@ -3320,7 +3320,7 @@ def write_extrconsolidate_sub(tag='extrconsolidate', exe=None, log_dir=None, uni return job, sub_name -def write_unify_sub_simple(tag='unify', exe=None, base=None,target=None,universe="vanilla",arg_str=None,log_dir=None, use_eos=False,ncopies=1,no_grid=False, max_runtime_minutes=60,extra_text='',**kwargs): +def write_unify_sub_simple(tag='unify', exe=None, base=None,target=None,universe="vanilla",arg_str=None,log_dir=None, use_eos=False,ncopies=1,no_grid=False, max_runtime_minutes=60,extra_text='',script_name=None,glob_pattern='*.composite',fail_on_error=False,**kwargs): """ Write a submit file for launching a consolidation job util_ILEdagPostprocess.sh # suitable for ILE consolidation. @@ -3342,34 +3342,53 @@ def write_unify_sub_simple(tag='unify', exe=None, base=None,target=None,universe # Write unify.sh # - problem of globbing inside condor commands # - problem that *.composite files from intermediate results will generally NOT be present - cmdname ='unify.sh' + # script_name: a workflow with more than one unify job (e.g. a pooled + # cross-model net AND a per-model net) needs a distinct script per job, + # or the second write clobbers the first. + cmdname = script_name if script_name else 'unify.sh' base_str = '' if not (base is None): base_str = ' ' + base +"/" + # A condor macro CANNOT be interpolated into this shell script: inside bash, + # $(macroapprox) is COMMAND SUBSTITUTION, so it runs a nonexistent command, + # expands to the empty string, and the glob silently matches nothing. When + # the pattern carries a macro, pass it as an ARGUMENT ($1) and let condor + # expand it in the submit file -- which is how join_grids.sh already works. + pattern_is_macro = "$(" in glob_pattern + glob_str = base_str + ("$1" if pattern_is_macro else glob_pattern) with open(cmdname,'w') as f: f.write("#! /usr/bin/env bash\n") if len(extra_text) > 0: f.write(extra_text+"\n") - f.write( "ls " + base_str+"*.composite 1>&2 \n") # write filenames being concatenated to stderr + f.write( "ls " + glob_str+" 1>&2 \n") # write filenames being concatenated to stderr # Sometimes we need to pass --eccentricity or --tabular-eos-file etc to util_CleanILE.py extra_args = '' if arg_str: extra_args = arg_str - f.write( exe + extra_args+ base_str+ "*.composite \n") - # Backstop code for untify.sh + f.write( exe + extra_args+ glob_str+ " \n") + # Historical single-model workflows fall back to raw concatenation if + # CleanILE fails. A model-aware workflow must fail closed: otherwise + # a validation error silently changes model marginalization into flat + # replica pooling while the DAG reports success. f.write("""ret_value=$? if [ $ret_value -eq 0 ]; then exit 0 -else +""") + if fail_on_error: + f.write("else\n exit $ret_value\n") + else: + f.write("""else cat {} -fi -""".format(base_str+"*.composite")) +""".format(glob_str)) + f.write("fi\n") st = os.stat(cmdname) import stat os.chmod(cmdname, st.st_mode | stat.S_IEXEC) ile_job = CondorDAGJob(universe=universe, executable=base_str+cmdname) # force full prefix + if pattern_is_macro: + ile_job.add_arg(glob_pattern) # condor expands the macro here, not bash requirements=[] if universe=='local': requirements.append("IS_GLIDEIN=?=undefined") @@ -3443,7 +3462,10 @@ def write_convert_sub(tag='convert', exe=None, file_input=None,file_output=None, arg_str = arg_str.lstrip('-') ile_job.add_opt(arg_str,'') # because we must be idiotic in how we pass arguments, I strip off the first two elements of the line # ile_job.add_opt(arg_str[2:],'') # because we must be idiotic in how we pass arguments, I strip off the first two elements of the line - ile_job.add_arg(file_input) + if file_input is not None: + # a tool whose inputs are all named options has no trailing positional; + # adding one unguarded appends the literal string "None" to the command + ile_job.add_arg(file_input) # # Logging options @@ -4099,7 +4121,11 @@ def write_resample_sub(tag='resample', exe=None, file_input=None,file_output=Non -def write_cat_sub(tag='cat', exe=None, file_prefix=None,file_postfix=None,file_output=None,universe="vanilla",arg_str='',log_dir=None, use_eos=False,ncopies=1, no_grid=False,**kwargs): +def write_cat_sub(tag='cat', exe=None, file_prefix=None,file_postfix=None, + file_output=None, universe="vanilla", + arg_str='',log_dir=None, use_eos=False,ncopies=1, + no_grid=False, search_root='.', expected_batches=None, + events_per_batch=None, **kwargs): """ Write a submit file for launching a 'resample' job util_ResampleILEOutputWithExtrinsic.py @@ -4110,10 +4136,26 @@ def write_cat_sub(tag='cat', exe=None, file_prefix=None,file_postfix=None,file_o exe_switch = which("switcheroo") # tool for patterend search-replace, to fix first line of output file cmdname = 'catjob.sh' + # Condor expands macros only in submit-file fields, never inside the shell + # script. Pass BOTH the search root and output name as arguments. Besides + # avoiding bash command substitution for $(macroapprox), this lets callers + # scope the input tree per model instead of running `find .` over a shared + # top-level directory and mixing every model's posterior samples. with open(cmdname,'w') as f: f.write("#! /bin/bash\n") - f.write(exe+" . -name '"+file_prefix+"*"+file_postfix+r"' -exec cat {} \; | sort -r | uniq > "+file_output+";\n") - f.write(exe_switch + " 'm1 ' '# m1 ' "+file_output) # add standard prefix + if expected_batches is not None and events_per_batch is not None: + f.write("set -e -o pipefail\n") + f.write("{ for ((batch=0; batch<$3; batch++)); do " + "event=$((batch*$4)); " + "for ((idx=0; idx<$4; idx++)); do " + "file=\"$1/EXTR_out-${event}.xml_${idx}_.dat\"; " + "if [ ! -s \"$file\" ]; then " + "echo \"catjob: missing expected input $file\" >&2; exit 1; fi; " + "cat \"$file\"; done; done; } | sort -r | uniq > \"$2\";\n") + else: + f.write(exe+" \"$1\" -name '"+file_prefix+"*"+file_postfix+ + "' -exec cat {} \\; | sort -r | uniq > \"$2\";\n") + f.write(exe_switch + " 'm1 ' '# m1 ' \"$2\"") # add standard prefix os.system("chmod a+x "+cmdname) ile_job = CondorDAGJob(universe=universe, executable='catjob.sh') @@ -4135,6 +4177,11 @@ def write_cat_sub(tag='cat', exe=None, file_prefix=None,file_postfix=None,file_o ile_sub_name = tag + '.sub' ile_job.set_sub_file(ile_sub_name) + ile_job.add_arg(search_root or '.') + ile_job.add_arg(file_output) + if expected_batches is not None and events_per_batch is not None: + ile_job.add_arg(str(expected_batches)) + ile_job.add_arg(str(events_per_batch)) # ile_job.add_arg(" . -name '" + file_prefix + "*" +file_postfix+"' -exec cat {} \; ") @@ -4449,10 +4496,33 @@ def write_calibration_uncertainty_reweighting_sub(tag='Calib_reweight', exe=None singularity_image_used = "{}".format(singularity_image) # make copy extra_files = [] - if singularity_image: - if 'osdf:' in singularity_image: - singularity_image_used = "./{}".format(singularity_image.split('/')[-1]) - extra_files += [singularity_image] + # Calibration reweighting is CPU-only. A container-family capability + # expression cannot be evaluated on a CPU slot, so collapse the family to + # its explicitly configured fallback image (the same policy as CIP). + # Never pass the manifest YAML itself to the container runtime. + singularity_is_family = False + singularity_container_universe = False + singularity_container_image = None + singularity_fallback_runtime = None + if singularity_image and is_container_manifest(singularity_image): + singularity_is_family = True + _manifest = load_container_manifest(singularity_image) + singularity_container_universe = bool( + use_singularity and os.environ.get('RIFT_CONTAINER_UNIVERSE') + ) + if singularity_container_universe: + singularity_container_image = build_container_image_select( + _manifest, request_gpu=False + ) + else: + singularity_fallback_runtime, _fb_transfer = build_fallback_single_image( + _manifest + ) + if _fb_transfer: + extra_files += [_fb_transfer] + elif singularity_image and 'osdf:' in singularity_image: + singularity_image_used = "./{}".format(singularity_image.split('/')[-1]) + extra_files += [singularity_image] @@ -4471,7 +4541,10 @@ def write_calibration_uncertainty_reweighting_sub(tag='Calib_reweight', exe=None singularity_base_exe_path = "/usr/bin/" # should not hardcode this ...! exe=singularity_base_exe_path + exe_base - ile_job = CondorDAGJob(universe="vanilla", executable=exe) + ile_job = CondorDAGJob( + universe=("container" if singularity_container_universe else "vanilla"), + executable=exe, + ) # This is a hack since CondorDAGJob hides the queue property ile_job._CondorJob__queue = ncopies @@ -4487,8 +4560,18 @@ def write_calibration_uncertainty_reweighting_sub(tag='Calib_reweight', exe=None # Compare to https://github.com/lscsoft/lalsuite/blob/master/lalinference/python/lalinference/lalinference_pipe_utils.py ile_job.add_condor_cmd('request_CPUs', str(1)) ile_job.add_condor_cmd('transfer_executable', 'False') - ile_job.add_condor_cmd("MY.SingularityBindCVMFS", 'True') - ile_job.add_condor_cmd("MY.SingularityImage", '"' + singularity_image_used + '"') + if singularity_container_universe: + ile_job.add_condor_cmd("container_image", singularity_container_image) + else: + ile_job.add_condor_cmd("MY.SingularityBindCVMFS", 'True') + if singularity_is_family: + ile_job.add_condor_cmd( + "MY.SingularityImage", '"' + singularity_fallback_runtime + '"' + ) + else: + ile_job.add_condor_cmd( + "MY.SingularityImage", '"' + singularity_image_used + '"' + ) ile_job.add_condor_cmd("transfer_output_files", "weight_files") requirements.append("HAS_SINGULARITY=?=TRUE") print(" WARNING: cal reweighting requires bilby. Directories are moved to cal_evelopes") diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/hyperpipeline_io.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/hyperpipeline_io.py index 3530b0387..88d3a831b 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/misc/hyperpipeline_io.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/hyperpipeline_io.py @@ -81,6 +81,9 @@ "lambda1", "lambda2", "eos_table_index", + "a6c", + "E0", + "p_phi0", "distance", "ecliptic_longitude", "ecliptic_latitude", @@ -103,6 +106,7 @@ def is_active(env=None): def build_column_list(use_eccentricity=False, use_meanPerAno=False, use_tides=False, use_eos_index=False, + use_eob_parameters=False, use_hyperbolic=False, use_distance=False, use_sky=False): """Compose the column-name tuple for a given physics configuration. @@ -110,7 +114,8 @@ def build_column_list(use_eccentricity=False, use_meanPerAno=False, appends the requested optional groups in a fixed canonical order:: (eccentricity, [meanPerAno,] [lambda1, lambda2,] - [eos_table_index,] [distance,] [ecliptic_longitude, ecliptic_latitude]) + [eos_table_index,] [a6c,] [E0, p_phi0,] [distance,] + [ecliptic_longitude, ecliptic_latitude]) """ cols = list(DEFAULT_BASE_COLUMNS) if use_eccentricity: @@ -121,6 +126,10 @@ def build_column_list(use_eccentricity=False, use_meanPerAno=False, cols.extend(["lambda1", "lambda2"]) if use_eos_index: cols.append("eos_table_index") + if use_eob_parameters: + cols.append("a6c") + if use_hyperbolic: + cols.extend(["E0", "p_phi0"]) if use_distance: cols.append("distance") if use_sky: @@ -288,6 +297,7 @@ def read_table(fname): def to_legacy_dat(arr, use_eccentricity=False, use_meanPerAno=False, use_tides=False, use_eos_index=False, use_distance=False, + use_eob_parameters=False, use_hyperbolic=False, use_sky=False): """Reshape a hyperpipeline structured array into the legacy CIP layout. @@ -295,7 +305,8 @@ def to_legacy_dat(arr, use_eccentricity=False, use_meanPerAno=False, columns are:: [event_id, m1, m2, s1x, s1y, s1z, s2x, s2y, s2z, - (distance?), (lambda1, lambda2, (eos_index)?)?, + (distance?), (lambda1, lambda2, (eos_index)?)?, (a6c?)?, + (E0, p_phi0)?, (eccentricity, (meanPerAno)?)?, (ecliptic_longitude, ecliptic_latitude)?, lnL, sigma_lnL] @@ -315,6 +326,10 @@ def to_legacy_dat(arr, use_eccentricity=False, use_meanPerAno=False, cols.extend(["lambda1", "lambda2"]) if use_eos_index: cols.append("eos_table_index") + if use_eob_parameters: + cols.append("a6c") + if use_hyperbolic: + cols.extend(["E0", "p_phi0"]) if use_eccentricity: cols.append("eccentricity") if use_meanPerAno: @@ -640,7 +655,8 @@ def read_grid_to_P_list(fname, P_factory, lal_module=None, def legacy_column_indices(use_eccentricity=False, use_meanPerAno=False, use_tides=False, use_eos_index=False, - use_distance=False, use_sky=False): + use_distance=False, use_eob_parameters=False, + use_hyperbolic=False, use_sky=False): """Return the positional column indices the legacy CIP loop expects. Mirrors the layout produced by :func:`to_legacy_dat` so callers can @@ -649,14 +665,16 @@ def legacy_column_indices(use_eccentricity=False, use_meanPerAno=False, CIP indexing logic without having to recompute them. Returns a dict keyed by ``'lnL'``, ``'sigma_lnL'``, ``'distance'``, - ``'lambda1'``, ``'eccentricity'``, ``'meanPerAno'``, + ``'lambda1'``, ``'a6c'``, ``'E0'``, ``'p_phi0'``, + ``'eccentricity'``, ``'meanPerAno'``, ``'ecliptic_longitude'``, ``'ecliptic_latitude'``. Any column not present in the configuration maps to ``None``. """ # event_id m1 m2 a1x a1y a1z a2x a2y a2z = 9 leading columns. idx = 9 out = {"lnL": None, "sigma_lnL": None, "distance": None, - "lambda1": None, "eccentricity": None, "meanPerAno": None, + "lambda1": None, "a6c": None, "E0": None, "p_phi0": None, + "eccentricity": None, "meanPerAno": None, "ecliptic_longitude": None, "ecliptic_latitude": None} if use_distance: out["distance"] = idx @@ -666,6 +684,13 @@ def legacy_column_indices(use_eccentricity=False, use_meanPerAno=False, idx += 2 # lambda1, lambda2 if use_eos_index: idx += 1 # eos_table_index (no positional alias used by CIP) + if use_eob_parameters: + out["a6c"] = idx + idx += 1 + if use_hyperbolic: + out["E0"] = idx + out["p_phi0"] = idx + 1 + idx += 2 if use_eccentricity: out["eccentricity"] = idx idx += 1 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/mc_error.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/mc_error.py new file mode 100644 index 000000000..cc46b175b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/mc_error.py @@ -0,0 +1,32 @@ +"""Helpers for normalizing Monte Carlo integration error conventions.""" + +import math + + +def relative_mc_error(result, variance, *, log_space=False): + """Return the relative standard error for linear- or log-space results. + + Linear-space integrators return ``(Z, Var[Z])``. Log-space integrators + return ``(ln Z, ln Var[Z])``; taking ``sqrt`` of the latter is invalid and + was the source of ``nan`` CIP evidence annotations for AV sampling. + """ + result = float(result) + variance = float(variance) + if not math.isfinite(result) or math.isnan(variance): + return math.nan + if log_space: + if variance == -math.inf: + return 0.0 + if variance == math.inf: + return math.inf + try: + return math.exp(0.5 * variance - result) + except OverflowError: + return math.inf + if variance == math.inf: + return math.inf + if variance < 0: + raise ValueError("linear-space integration variance must be non-negative") + if result == 0: + return 0.0 if variance == 0 else math.inf + return math.sqrt(variance) / abs(result) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/samples_utils.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/samples_utils.py index 68d9ed559..81760a668 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/misc/samples_utils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/samples_utils.py @@ -16,6 +16,7 @@ "theta1":"tilt1", "theta2":"tilt2", "xi":"chi_eff", + "shu":"shu", "chiMinus":"chi_minus", "delta":"delta", "delta_mc":"delta", @@ -60,6 +61,13 @@ def extract_combination_from_LI(samples_LI, p): a1z = samples_LI['a1z'] a2z = samples_LI['a2z'] return (m1 * a1z + m2*a2z)/(m1+m2) + if (p == 'shu') and 'a1z' in samples_LI.dtype.names: + m1 = samples_LI['m1'] + m2 = samples_LI['m2'] + a1z = samples_LI['a1z'] + a2z = samples_LI['a2z'] + xi = (m1 * a1z + m2*a2z)/(m1+m2) + return xi - 0.5*(a1z+a2z) * ((m1*m2)/ (m1+m2)**2) # Return cartesian components of spin1, spin2. NOTE: I may already populate these quantities in 'Add important quantities' if p == 'chiz_plus': print(" Transforming ") diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/test_psd_bandwidth.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/test_psd_bandwidth.py index a71743f28..6780f0171 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/misc/test_psd_bandwidth.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/test_psd_bandwidth.py @@ -1,4 +1,16 @@ #!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_psd_bandwidth guards the representative-detector choice and the fallback contract +# behind the stencil decision; see this file's docstring. numpy-only, seconds. +# --------------------------------------------------------------------------------- """test_psd_bandwidth -- representative-detector choice, and the fallback contract. Two things are guarded, both of which are about behaviour under imperfect input rather than diff --git a/MonteCarloMarginalizeCode/Code/RIFT/misc/xmlutils.py b/MonteCarloMarginalizeCode/Code/RIFT/misc/xmlutils.py index 8a4f3476c..bf7fc4c8a 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/misc/xmlutils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/misc/xmlutils.py @@ -17,6 +17,24 @@ def assign_time(row, t): setattr(row, "geocent_end_time",( int(t))) setattr(row, "geocent_end_time_ns",int( (t-int(t))*1e9 ) ) + +def gps_add_seconds_exact(gps_seconds, gps_nanoseconds, offset_seconds): + """Add float relative offsets to an exact GPS pair using integer nanoseconds. + + The offset is rounded once to the SimInspiral schema's nanosecond resolution. + Absolute GPS arithmetic is never performed in float64, whose spacing at + current epochs is coarser than some sub-sample time-posterior draws. + Returns arrays ``(seconds, nanoseconds)`` with nanoseconds normalized to + ``[0, 1e9)``; negative offsets and second-boundary carries are supported. + """ + billion = numpy.int64(1000000000) + offset_ns = numpy.rint(numpy.asarray(offset_seconds, dtype=float) * 1.e9).astype(numpy.int64) + total_ns = (numpy.int64(gps_seconds) * billion + + numpy.int64(gps_nanoseconds) + offset_ns) + return (numpy.floor_divide(total_ns, billion).astype(numpy.int64), + numpy.mod(total_ns, billion).astype(numpy.int64)) + + CMAP = { "right_ascension": "longitude", "longitude":"longitude", "latitude":"latitude", @@ -24,6 +42,13 @@ def assign_time(row, t): "inclination": "inclination", "polarization": "polarization", "t_ref": assign_time,#r.set_time_geocent(LIGOTimeGPS(float(t))), + # Exact fields, when present, deliberately follow t_ref and overwrite the + # legacy float decomposition. float64 absolute GPS times are only spaced by + # ~0.12--0.24 microseconds at current epochs, while SimInspiral stores an + # integer nanosecond field and band-limited time export can resolve below the + # float spacing. + "t_ref_gps_seconds": "geocent_end_time", + "t_ref_gps_nanoseconds": "geocent_end_time_ns", "coa_phase": "coa_phase", "distance": "distance", "mass1": "mass1", @@ -37,11 +62,17 @@ def assign_time(row, t): "alpha4":"alpha4", "alpha5":"alpha5", "alpha6":"alpha6", + "psi3":"psi3", + "beta":"beta", + "psi0":"psi0", "loglikelihood": "alpha1", "joint_prior": "alpha2", "joint_s_prior": "alpha3", "eccentricity":"alpha4", "meanPerAno":"alpha", + "E0":"psi3", + "p_phi0":"beta", + "a6c":"psi0", "lambda1":"alpha5", "lambda2":"alpha6", "spin1x":"spin1x", @@ -54,7 +85,7 @@ def assign_time(row, t): # FIXME: Find way to intersect given cols with valid cols when making table. # Otherwise, we'll have to add them manually and ensure they all exist -sim_valid_cols = ["simulation_id", "inclination", "longitude", "latitude", "polarization", "geocent_end_time", "geocent_end_time_ns", "coa_phase", "distance", "mass1", "mass2", "alpha", "alpha1", "alpha2", "alpha3", "alpha4", "alpha5", "alpha6", "spin1x", "spin1y", "spin1z", "spin2x", "spin2y", "spin2z"] +sim_valid_cols = ["simulation_id", "inclination", "longitude", "latitude", "polarization", "geocent_end_time", "geocent_end_time_ns", "coa_phase", "distance", "mass1", "mass2", "alpha", "alpha1", "alpha2", "alpha3", "alpha4", "alpha5", "alpha6","psi3","beta", "spin1x", "spin1y", "spin1z", "spin2x", "spin2y", "spin2z","psi0"] sngl_valid_cols = [ "event_id", "snr", "tau0", "tau3"] multi_valid_cols = ["process_id", "event_id", "snr"] diff --git a/MonteCarloMarginalizeCode/Code/RIFT/physics/GWSignal.py b/MonteCarloMarginalizeCode/Code/RIFT/physics/GWSignal.py index 9ada39b95..231fed719 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/physics/GWSignal.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/physics/GWSignal.py @@ -12,6 +12,7 @@ import lal import lalsimulation as lalsim import RIFT.lalsimutils as lalsimutils +from RIFT.physics import teobresums_compat import numpy as np import astropy.units as u from astropy.time import Time @@ -86,7 +87,8 @@ def hlmoft(P, Lmax=2,approx_string=None,no_trust_align_method=None,internal_phas approx_string (str): Approximant string. If None, P.approx is used. no_trust_align_method (str): If 'peak', shifts epoch to the peak of the total signal power. internal_phase_shift (float): Phase shift applied to the modes. Default is pi/2. - **kwargs: Additional arguments (e.g., 'lmax_nyquist'). + **kwargs: Additional arguments (e.g., 'lmax_nyquist', 'force_22_mode'). + force_22_mode (bool): If True, return only the (2,+-2) modes. Returns: dict: A dictionary mapping (l, m) to LAL COMPLEX16TimeSeries objects. @@ -94,6 +96,8 @@ def hlmoft(P, Lmax=2,approx_string=None,no_trust_align_method=None,internal_phas assert Lmax >= 2 + force_22_mode = kwargs.get('force_22_mode', False) + # Check that masses are not nan! assert (not np.isnan(P.m1)) and (not np.isnan(P.m2)), " masses are NaN " taper=0 @@ -127,6 +131,16 @@ def hlmoft(P, Lmax=2,approx_string=None,no_trust_align_method=None,internal_phas if not(approx_string): approx_string_here = lalsim.GetStringFromApproximant(P.approx) + # DO NOT remove this as cosmetic spin rounding. TEOBResumS-DALI's C code + # classifies sum(chi_perp) <= 1e-4 as aligned, while its GWSignal wrapper + # requests inertial modes for any exactly nonzero transverse component. + # That disagreement segfaults EOBRunPy in production DALI builds. Zeroing + # only the backend's own aligned interval makes both layers take the same + # path; genuinely precessing spins above the boundary remain untouched. + python_dict = teobresums_compat.guard_gwsignal_transverse_spins( + python_dict, approx_string_here + ) + # Fork on calling different generators gen = gws.models.gwsignal_get_waveform_generator(approx_string_here) # if "NRSur7dq4_gwsurr" == approx_string_here: @@ -148,6 +162,13 @@ def hlmoft(P, Lmax=2,approx_string=None,no_trust_align_method=None,internal_phas continue if mode[0] > Lmax: # skip modes with L > Lmax continue + # force_22_mode must actually produce a 22-only waveform here too, not + # just on the lalsimutils path. The restriction is applied to the + # returned modes rather than to the generator arguments, because the + # mode-restriction keyword is not uniformly supported by the generators + # reachable through gwsignal_get_waveform_generator. + if force_22_mode and not(mode[0] == 2 and abs(mode[1]) == 2): + continue # h = lal.CreateCOMPLEX16TimeSeries("hlm", lal.LIGOTimeGPS(0.), 0., P.deltaT, lal.DimensionlessUnit, @@ -277,6 +298,13 @@ def hoft(P, Fp=None, Fc=None,approx_string=None, **kwargs): if not(approx_string): approx_string_here = lalsim.GetStringFromApproximant(P.approx) + # Apply the same native-backend safety boundary as hlmoft. Keep this on + # the polarization path too: callers may reach TEOBResumS through either + # GWSignal entry point. + python_dict = teobresums_compat.guard_gwsignal_transverse_spins( + python_dict, approx_string_here + ) + # Fork on calling different generators gen = gws.models.gwsignal_get_waveform_generator(approx_string_here) @@ -389,6 +417,12 @@ def complex_hoft(P, Fp=None, Fc=None,approx_string=None,sgn=-1, **kwargs): if not(approx_string): approx_string_here = lalsim.GetStringFromApproximant(P.approx) + # complex_hoft reaches the same GWSignal polarization generator as hoft; + # keep its ResumS native call behind the same near-aligned safety boundary. + python_dict = teobresums_compat.guard_gwsignal_transverse_spins( + python_dict, approx_string_here + ) + # Fork on calling different generators gen = gws.models.gwsignal_get_waveform_generator(approx_string_here) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/physics/teobresums_compat.py b/MonteCarloMarginalizeCode/Code/RIFT/physics/teobresums_compat.py new file mode 100644 index 000000000..65e189549 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/physics/teobresums_compat.py @@ -0,0 +1,305 @@ +"""Compatibility boundary for the optional TEOBResumS Python extension. + +``EOBRun_module`` is a native extension with several incompatible interfaces in +active use. Invalid enum values can terminate Python instead of raising an +exception, so callers must normalize parameters and probe a new interface in a +child process before making the first in-process call. +""" + +import hashlib +import json +import math +import os +import subprocess +import sys + +try: + from importlib import metadata as importlib_metadata +except ImportError: # Python 3.7 compatibility + try: + import importlib_metadata + except ImportError: # fingerprinting is optional on minimal Python 3.6 installs + importlib_metadata = None + + +class TEOBResumSCompatibilityError(RuntimeError): + pass + + +_DALI_MARKERS = { + "eob_dyn_j0_py", + "eob_ham_s_py", + "eob_metric_A5PNlog_py", +} + +# Profiles intentionally contain semantic values, not C enum ordinals. The +# default is the common string API verified against legacy-hyperbolic and DALI +# builds. Future incompatible profiles should be added here rather than +# branching throughout lalsimutils. +_PROFILE_VALUES = { + "default": { + "arg_out": "yes", + "nqc": "no", + "nqc_coefs_hlm": "none", + "nqc_coefs_flx": "none", + "use_geometric_units": "no", + "interp_uniform_grid": "yes", + "output_hpc": "no", + }, + "dali": {}, + "legacy": {}, +} + +_PROBED_SCHEMAS = set() + + +# TEOBResumS-DALI decides whether a system is precessing from +# hypot(chi1x, chi1y) + hypot(chi2x, chi2y) > 1e-4. Keep the value here in +# sync with TEOBResumSPars.c. It is a native-backend safety boundary, not a +# generic numerical-zero tolerance. +DALI_TRANSVERSE_SPIN_THRESHOLD = 1e-4 +DALI_INITIAL_TRANSVERSE_SPIN_RANGE = (1e-3, 3e-3) +LEGACY_INITIAL_TRANSVERSE_SPIN_RANGE = (1e-5, 3e-5) + + +def is_teobresums_approximant(approximant): + """Return whether an approximant name selects the TEOBResumS family.""" + return str(approximant or "").lower().startswith("teobresums") + + +def initial_transverse_spin_range(approximant): + """Return a non-aligned initial-grid seed appropriate to ``approximant``. + + Other precessing models retain RIFT's long-standing tiny seed. ResumS + needs a seed safely above its native 1e-4 aligned/precessing boundary. + """ + if is_teobresums_approximant(approximant): + return DALI_INITIAL_TRANSVERSE_SPIN_RANGE + return LEGACY_INITIAL_TRANSVERSE_SPIN_RANGE + + +def _dimensionless_float(value): + return float(value.value if hasattr(value, "value") else value) + + +def total_transverse_spin(s1x, s1y, s2x, s2y): + """Return the transverse spin magnitude TEOBResumS-DALI classifies with.""" + return math.hypot(_dimensionless_float(s1x), _dimensionless_float(s1y)) + math.hypot( + _dimensionless_float(s2x), _dimensionless_float(s2y) + ) + + +def is_precessing_for_resums(s1x, s1y, s2x, s2y): + """Return whether TEOBResumS-DALI evolves these spins on its precessing path. + + Callers must ask the same question the backend does before requesting + inertial-frame modes: a component that is merely nonzero can still leave the + summed transverse magnitude inside the native aligned interval, and asking + for inertial modes there disagrees with the dynamics DALI actually runs. + """ + return total_transverse_spin(s1x, s1y, s2x, s2y) > DALI_TRANSVERSE_SPIN_THRESHOLD + + +def guard_gwsignal_transverse_spins(parameters, approximant): + """Return GWSignal parameters safe at the ResumS alignment boundary. + + TEOBResumS-DALI treats total transverse spin at or below 1e-4 as aligned, + but its GWSignal wrapper requests inertial modes for *any* exactly nonzero + transverse component. Some native builds segfault when those two choices + disagree. Match the backend's own classification by zeroing only that + interval; do not mutate the caller's dictionary. + """ + if not is_teobresums_approximant(approximant): + return parameters + + keys = ("spin1x", "spin1y", "spin2x", "spin2y") + transverse_spin = total_transverse_spin(*(parameters[key] for key in keys)) + if not 0.0 < transverse_spin <= DALI_TRANSVERSE_SPIN_THRESHOLD: + return parameters + + safe_parameters = dict(parameters) + for key in keys: + value = parameters[key] + safe_parameters[key] = 0.0 * value.unit if hasattr(value, "unit") else 0.0 + return safe_parameters + + +def _json_compatible(value): + if isinstance(value, dict): + return {key: _json_compatible(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_json_compatible(item) for item in value] + if hasattr(value, "item"): + try: + return value.item() + except (TypeError, ValueError): + pass + return value + + +def detect_profile(module, requested=None): + """Return the configured or best-effort TEOBResumS interface profile. + + ``auto`` is the default. Unknown extensions deliberately fall back to the + conservative ``default`` profile; an explicit unknown profile is rejected + so misspellings cannot silently change waveform settings. + """ + requested = requested or os.environ.get("RIFT_TEOBRESUMS_PROFILE", "auto") + requested = requested.lower() + if requested != "auto": + if requested not in _PROFILE_VALUES: + raise TEOBResumSCompatibilityError( + "Unknown RIFT_TEOBRESUMS_PROFILE={!r}; expected auto, default, dali, or legacy".format( + requested + ) + ) + return requested + if _DALI_MARKERS.issubset(set(dir(module))): + return "dali" + return "default" + + +def _profile_values(profile): + values = dict(_PROFILE_VALUES["default"]) + values.update(_PROFILE_VALUES.get(profile, {})) + return values + + +def normalize_parameters(parameters, profile="default"): + """Return a copy using the selected profile's stable enum representation.""" + if profile not in _PROFILE_VALUES: + raise TEOBResumSCompatibilityError( + "Unknown TEOBResumS profile {!r}".format(profile) + ) + normalized = dict(parameters) + values = _profile_values(profile) + + # These integer spellings occur in older RIFT branches. Their apparent C + # enum ordinals changed meaning in DALI, so translate their RIFT semantics + # before applying the profile representation. + legacy_semantics = { + "arg_out": {0: "no", 1: "yes"}, + "nqc": {2: "no"}, + "nqc_coefs_hlm": {0: "none"}, + "nqc_coefs_flx": {0: "none"}, + "use_geometric_units": {0: "no", 1: "yes"}, + "interp_uniform_grid": {0: "no", 1: "yes"}, + "output_hpc": {0: "no", 1: "yes"}, + } + for key, translations in legacy_semantics.items(): + if key not in normalized: + continue + value = normalized[key] + if isinstance(value, bool): + value = int(value) + if isinstance(value, int) and value not in translations: + raise TEOBResumSCompatibilityError( + "Unsupported numeric TEOBResumS value {}={!r}; use a semantic string".format( + key, value + ) + ) + semantic_value = translations.get(value, value) + if semantic_value == values.get(key): + normalized[key] = values[key] + elif isinstance(semantic_value, str): + normalized[key] = semantic_value + return normalized + + +def runtime_fingerprint(module, profile=None): + """Return reproducibility metadata without assuming package version is unique.""" + module_path = os.path.realpath(getattr(module, "__file__", "")) + package_version = None + if importlib_metadata is not None: + try: + package_version = importlib_metadata.version("teobresums") + except importlib_metadata.PackageNotFoundError: + pass + digest = None + if module_path and os.path.isfile(module_path): + hasher = hashlib.sha256() + with open(module_path, "rb") as module_file: + for block in iter(lambda: module_file.read(1024 * 1024), b""): + hasher.update(block) + digest = hasher.hexdigest() + return { + "profile": profile or detect_profile(module), + "module_path": module_path or None, + "module_sha256": digest, + "package_version": package_version, + "exported_symbols": sorted( + name for name in _DALI_MARKERS if hasattr(module, name) + ), + } + + +def _probe_schema(module, parameters, profile, purpose, timeout): + module_path = os.path.realpath(getattr(module, "__file__", "")) + schema = (profile, purpose, module_path, tuple(sorted(parameters))) + if schema in _PROBED_SCHEMAS: + return + if os.environ.get("RIFT_TEOBRESUMS_SKIP_PROBE", "").lower() in { + "1", + "true", + "yes", + }: + return + + probe_code = r""" +import json +import importlib.util +import os +import sys + +expected_path = os.path.realpath(sys.argv[1]) +if expected_path: + spec = importlib.util.spec_from_file_location("EOBRun_module", expected_path) + EOBRun_module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(EOBRun_module) +else: + import EOBRun_module +parameters = json.loads(sys.argv[2]) +result = EOBRun_module.EOBRunPy(parameters) +if not isinstance(result, tuple) or len(result) < 4: + raise RuntimeError("EOBRunPy returned an unexpected result") +""" + probe_parameters = _json_compatible(parameters) + try: + completed = subprocess.run( + [sys.executable, "-c", probe_code, module_path, json.dumps(probe_parameters)], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + universal_newlines=True, + timeout=timeout, + ) + except subprocess.TimeoutExpired as exc: + raise TEOBResumSCompatibilityError( + "TEOBResumS {} compatibility probe exceeded {} seconds".format( + purpose, timeout + ) + ) from exc + if completed.returncode != 0: + detail = (completed.stderr or completed.stdout or "").strip()[-1000:] + raise TEOBResumSCompatibilityError( + "TEOBResumS {} compatibility probe failed with return code {}. {}".format( + purpose, completed.returncode, detail + ) + ) + _PROBED_SCHEMAS.add(schema) + + +def run(module, parameters, purpose="waveform", profile=None, probe=True, timeout=None): + """Normalize, safely preflight, and call ``EOBRunPy``. + + The first call for each parameter schema is duplicated in a child process. + That cost is intentional: an incompatible native extension may segfault. + Set ``RIFT_TEOBRESUMS_SKIP_PROBE=1`` only for a separately validated, + pinned runtime. + """ + selected_profile = profile or detect_profile(module) + normalized = normalize_parameters(parameters, selected_profile) + if probe: + if timeout is None: + timeout = float(os.environ.get("RIFT_TEOBRESUMS_PROBE_TIMEOUT", "60")) + _probe_schema(module, normalized, selected_profile, purpose, timeout) + return module.EOBRunPy(normalized) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/rimsky/README.md b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/README.md new file mode 100644 index 000000000..b29a97f64 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/README.md @@ -0,0 +1,60 @@ +# Rimsky integration + +Rimsky performs online Bilby parameter estimation and can launch follow-up +analyses through its Asimov hook. RIFT supplies a bridge for that hook: + +1. `rift-rimsky-analysis rimsky.yaml rift-followup.yaml` reads the Rimsky + configuration and writes both a RIFT Asimov analysis document and a runnable + `rimsky-rift.yaml`. +2. Initialize an Asimov 0.7 project named by `asimovdir` once, install RIFT in its + environment, and run `rimsky rimsky-rift.yaml`. + +The generated Rimsky configuration defaults `event_sink.bilby_pipe_format` to +`full-submit`, points `sample_sink.asimov_configuration` at the generated RIFT +analysis, and makes relative output paths absolute. Thus the first online +Bilby result is written as a PESummary metafile and Rimsky immediately adds the +ready RIFT follow-up to Asimov. A running Asimov manager then builds and submits +that production. Existing explicit Bilby run modes and Asimov project paths +are preserved. Use `--configured-rimsky PATH` to choose a different filename. + +For example: + +```yaml +output_dir: ./output +asimovdir: ./asimov +detectors: [H1, L1, V1] + +sample_sink: + # Written automatically in rimsky-rift.yaml. + asimov_configuration: /absolute/path/to/rift-followup.yaml + +# Optional. Rimsky ignores this extra section; the RIFT generator consumes it. +rift: + name: rift-online + waveform: + approximant: IMRPhenomXPHM + scheduler: + accounting group: ligo.dev.o4.cbc.pe.rift + osg: false +``` + +Rimsky writes a PESummary metafile before applying the follow-up file. The +generated analysis uses an absolute `output_dir/*/*/{event}/...` glob to find +that event's metafile, sets its dataset to `bilby-online`, and bootstraps RIFT +and its coincidence XML from the online posterior. Exactly one metafile must match; RIFT fails closed +if the path is missing or ambiguous. + +Rimsky 0.1 event documents use Bilby-style prior names (`chirp_mass`, +`mass_ratio`, `a_1`, and so on). The RIFT pipeline retains those keys and adds +the space-separated aliases expected by its Asimov template. This makes the +same event usable by both Bilby and RIFT analyses. + +The bridge itself consumes plain YAML mappings and does not import Rimsky. Its +unit tests remain isolated from streaming, GraceDB, and HTCondor. Dedicated +end-to-end lanes install Rimsky `0.1.0rc1` on Python 3.12 and pinned current main +on Python 3.14, both forced onto Asimov 0.7 and the merged bilby_pipe 0.7 adapter. +They load the generated configuration through Rimsky, invoke its real post-PE +Asimov hook, discover the RIFT pipeline, and resolve the first metafile as the +bootstrap input. External scheduler submission is the only mocked boundary. +The current-main pin is commit `2621d15` (2026-09-01); the bilby_pipe adapter pin +is `be6c770` pending its next release. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/rimsky/__init__.py b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/__init__.py new file mode 100644 index 000000000..6020ad1de --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/__init__.py @@ -0,0 +1,27 @@ +"""Rimsky-to-RIFT orchestration helpers. + +The public API is intentionally independent of Rimsky's Python internals. Rimsky +configuration and event documents are plain mappings, which keeps this bridge +usable across Rimsky release candidates without importing its large online-PE +runtime stack. +""" + +from .integration import ( + RimskyIntegrationError, + build_analysis, + configure_rimsky, + load_rimsky_config, + normalize_event_metadata, + write_analysis, + write_rimsky_config, +) + +__all__ = [ + "RimskyIntegrationError", + "build_analysis", + "configure_rimsky", + "load_rimsky_config", + "normalize_event_metadata", + "write_analysis", + "write_rimsky_config", +] diff --git a/MonteCarloMarginalizeCode/Code/RIFT/rimsky/integration.py b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/integration.py new file mode 100644 index 000000000..c45a41765 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/rimsky/integration.py @@ -0,0 +1,311 @@ +"""Generate a RIFT follow-up for the Rimsky online-PE orchestrator.""" + +import argparse +import copy +import json +from pathlib import Path + + +class RimskyIntegrationError(ValueError): + """Raised when a Rimsky configuration cannot define a RIFT follow-up.""" + + +def _deep_update(target, updates): + """Recursively apply mapping ``updates`` without sharing mutable values.""" + for key, value in updates.items(): + if isinstance(value, dict) and isinstance(target.get(key), dict): + _deep_update(target[key], value) + else: + target[key] = copy.deepcopy(value) + return target + + +def load_rimsky_config(path): + """Load a Rimsky YAML configuration and return it as a mapping.""" + path = Path(path) + try: + import yaml + except ImportError as exc: # Rimsky itself depends on PyYAML. + raise RimskyIntegrationError( + "PyYAML is required to read a Rimsky configuration" + ) from exc + + with path.open("r", encoding="utf-8") as stream: + config = yaml.safe_load(stream) or {} + if not isinstance(config, dict): + raise RimskyIntegrationError("Rimsky configuration must be a mapping") + return config + + +def _detectors(config): + detectors = config.get("detectors", ["H1", "L1"]) + if isinstance(detectors, str): + detectors = [item.strip() for item in detectors.split(",") if item.strip()] + if not isinstance(detectors, (list, tuple)) or not detectors: + raise RimskyIntegrationError("Rimsky 'detectors' must be a non-empty list") + if not all(isinstance(detector, str) and detector for detector in detectors): + raise RimskyIntegrationError("Each Rimsky detector must be a non-empty string") + return list(detectors) + + +def _frequency_dict(value, detectors, field): + if isinstance(value, (int, float)): + return {detector: value for detector in detectors} + if not isinstance(value, dict): + raise RimskyIntegrationError( + "{} must be a number or detector mapping".format(field) + ) + missing = [detector for detector in detectors if detector not in value] + if missing: + raise RimskyIntegrationError( + "{} is missing detector(s): {}".format(field, ", ".join(missing)) + ) + selected = {detector: value[detector] for detector in detectors} + if not all(isinstance(item, (int, float)) for item in selected.values()): + raise RimskyIntegrationError("{} values must be numeric".format(field)) + return selected + + +def _resolve_output_dir(config, config_path=None): + output_dir = Path(config.get("output_dir", "output")).expanduser() + if not output_dir.is_absolute(): + base = Path(config_path).resolve().parent if config_path else Path.cwd() + output_dir = base / output_dir + return output_dir.resolve() + + +def _resolve_config_path(value, *, config_path=None): + """Resolve a Rimsky path using the source configuration as its anchor.""" + path = Path(value).expanduser() + if not path.is_absolute(): + base = Path(config_path).resolve().parent if config_path else Path.cwd() + path = base / path + return path.resolve() + + +def build_analysis(config, *, config_path=None, overrides=None): + """Build the Asimov analysis document consumed by Rimsky's follow-up hook. + + Parameters + ---------- + config : mapping + Parsed Rimsky configuration. An optional top-level ``rift`` mapping is + ignored by Rimsky and accepted here as user overrides. + config_path : path-like, optional + Location of the Rimsky YAML file. Relative ``output_dir`` values are + resolved relative to this file, matching Rimsky's current behaviour. + overrides : mapping, optional + Programmatic overrides applied after the top-level ``rift`` mapping. + + Returns + ------- + dict + One Asimov analysis document suitable for + ``sample_sink.asimov_configuration``. + """ + if not isinstance(config, dict): + raise RimskyIntegrationError("Rimsky configuration must be a mapping") + + detectors = _detectors(config) + event_sink = config.get("event_sink") or {} + bilby = event_sink.get("bilby_pipe_defaults") or {} + minimum = _frequency_dict( + bilby.get("minimum_frequency", 20), detectors, "minimum_frequency" + ) + maximum = _frequency_dict( + bilby.get("maximum_frequency", 1024), detectors, "maximum_frequency" + ) + output_dir = _resolve_output_dir(config, config_path=config_path) + + # Rimsky stores each event under output_dir/YYMM/DD/SID and writes this + # metafile immediately before invoking the configured Asimov follow-ups. + bootstrap = output_dir / "*" / "*" / "{event}" / "results_page" / "metafile.hdf5" + + analysis = { + "kind": "analysis", + "name": "rift-online", + "status": "Ready", + "pipeline": "RIFT", + "orchestrator": "rimsky", + "comment": "RIFT follow-up launched by Rimsky after online Bilby PE", + "dataset": "bilby-online", + "likelihood": { + "start frequency": min(minimum.values()), + "minimum frequency": minimum, + "assume": {"precessing": True}, + "marginalization": {"distance": True}, + }, + "quality": { + "minimum frequency": minimum, + "maximum frequency": maximum, + }, + "waveform": { + "approximant": "IMRPhenomXPHM", + "pn amplitude order": 5, + "maximum mode": 4, + }, + "priors": { + "mass 1": {"minimum": 1, "maximum": 1000}, + }, + "sampler": {"cip": {}, "ile": {}}, + "scheduler": { + "accounting group": "ligo.dev.o4.cbc.pe.rift", + "bootstrap coinc": True, + "bootstrap file": str(bootstrap), + "osg": False, + }, + } + + configured = config.get("rift") or {} + if not isinstance(configured, dict): + raise RimskyIntegrationError( + "Optional Rimsky 'rift' settings must be a mapping" + ) + _deep_update(analysis, configured) + if overrides is not None: + if not isinstance(overrides, dict): + raise RimskyIntegrationError("RIFT overrides must be a mapping") + _deep_update(analysis, overrides) + return analysis + + +def configure_rimsky(config, analysis_path, *, config_path=None): + """Return a runnable Rimsky configuration wired to the RIFT follow-up. + + Rimsky interprets paths relative to its launch directory, rather than the + YAML file. Emit absolute paths so the online output observed by Rimsky is + the same output searched by RIFT's bootstrap glob. Missing orchestration + settings default to a local Asimov project and a submitted Bilby run; an + operator's explicit values are retained. + """ + if not isinstance(config, dict): + raise RimskyIntegrationError("Rimsky configuration must be a mapping") + + configured = copy.deepcopy(config) + configured["output_dir"] = str( + _resolve_output_dir(configured, config_path=config_path) + ) + configured["asimovdir"] = str( + _resolve_config_path( + configured.get("asimovdir", "asimov"), config_path=config_path + ) + ) + event_sink = configured.setdefault("event_sink", {}) + if not isinstance(event_sink, dict): + raise RimskyIntegrationError("Rimsky 'event_sink' must be a mapping") + event_sink.setdefault("bilby_pipe_format", "full-submit") + + sample_sink = configured.setdefault("sample_sink", {}) + if not isinstance(sample_sink, dict): + raise RimskyIntegrationError("Rimsky 'sample_sink' must be a mapping") + sample_sink["asimov_configuration"] = str( + _resolve_config_path(analysis_path) + ) + return configured + + +def normalize_event_metadata(metadata): + """Return RIFT-compatible metadata from a Rimsky-created event mapping. + + Rimsky 0.1 emits Bilby parameter names with underscores. RIFT's Asimov + template predates that convention and uses names containing spaces. Keep + both spellings so other analyses in the same ledger are unaffected. + """ + normalized = copy.deepcopy(metadata) + priors = normalized.setdefault("priors", {}) + aliases = { + "chirp_mass": "chirp mass", + "mass_ratio": "mass ratio", + "luminosity_distance": "luminosity distance", + "mass_1": "mass 1", + } + for source, destination in aliases.items(): + if destination not in priors and source in priors: + priors[destination] = copy.deepcopy(priors[source]) + + for source, destination in (("a_1", "spin 1"), ("a_2", "spin 2")): + if destination not in priors and source in priors: + prior = priors[source] + if isinstance(prior, dict) and "maximum" in prior: + priors[destination] = {"maximum": prior["maximum"]} + for source, destination in (("chi_1", "spin 1"), ("chi_2", "spin 2")): + if destination not in priors and source in priors: + priors[destination] = {"maximum": 0.99} + + # Asimov's RIFT PSD convention is sample-rate -> detector -> path, while + # Rimsky records detector -> path. This copy belongs only to the RIFT + # production, so the event document retained by Rimsky remains unchanged. + psds = normalized.get("psds") + sample_rate = normalized.get("likelihood", {}).get("sample rate") + detectors = normalized.get("interferometers", []) + if ( + isinstance(psds, dict) + and sample_rate is not None + and detectors + and all(detector in psds for detector in detectors) + ): + normalized["psds"] = { + sample_rate: { + detector: copy.deepcopy(psds[detector]) for detector in detectors + } + } + return normalized + + +def write_analysis(analysis, path): + """Write one analysis document as YAML (or JSON, which is valid YAML).""" + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + try: + import yaml + except ImportError: + with path.open("w", encoding="utf-8") as stream: + json.dump(analysis, stream, indent=2) + stream.write("\n") + else: + with path.open("w", encoding="utf-8") as stream: + yaml.safe_dump(analysis, stream, sort_keys=False) + return path + + +def write_rimsky_config(config, path): + """Write a Rimsky configuration containing the automatic RIFT hook.""" + return write_analysis(config, path) + + +def main(argv=None): + """Create the RIFT analysis and a Rimsky config that invokes it.""" + parser = argparse.ArgumentParser( + description="Generate a RIFT follow-up analysis for a Rimsky configuration" + ) + parser.add_argument("rimsky_config", help="Rimsky YAML configuration") + parser.add_argument("output", help="Destination analysis YAML") + parser.add_argument( + "--configured-rimsky", + help=( + "Destination for the runnable Rimsky YAML " + "(default: -rift.yaml)" + ), + ) + args = parser.parse_args(argv) + + config = load_rimsky_config(args.rimsky_config) + analysis_path = write_analysis( + build_analysis(config, config_path=args.rimsky_config), args.output + ) + source = Path(args.rimsky_config) + configured_path = Path(args.configured_rimsky) if args.configured_rimsky else ( + source.with_name(source.stem + "-rift" + source.suffix) + ) + write_rimsky_config( + configure_rimsky( + config, analysis_path.resolve(), config_path=args.rimsky_config + ), + configured_path, + ) + print(analysis_path) + print(configured_path) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/DESIGN.md b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/DESIGN.md index 5c4c42ec8..66fab023a 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/DESIGN.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/DESIGN.md @@ -597,6 +597,76 @@ append-only alternative: | `transfer_output_files` | `extra_transfer_output_files` (appended, `{level}`/`{sim_name}` substituted) | | `periodic_release` | `extra_periodic_release` (OR'd in) | | `request_memory` | the `request_memory` argument, or `Archive.set_resources` per sim | +| `container_image` | the `container_image` argument | +| `when_to_transfer_output` | the `when_to_transfer_output` argument | + +`universe` is **not** refused. Setting `container_image` selects the +container universe for you, but declaring `universe` yourself is legal +and sometimes necessary (`local`, `scheduler`, `grid`). + +### Containers + +`container_image` selects HTCondor's container universe and supplies the +image in one argument. Not because the two can disagree — measured +against condor 25.13.1, `vanilla` + `container_image`, `container` + +`container_image`, and declaring no universe at all produce byte-identical +job ads, and the JDL says `universe` "can either be optionally set to +`container` or not declared at all". It is one argument because that is +one fewer thing to state, not because stating it twice is dangerous: + +```python +DualCondorRunQueue( + container_image="osdf:///ospool/ap41/data//supernu-v2.sif", + when_to_transfer_output="ON_SUCCESS", +) +``` + +Both are keyword-only, and last in the signature: they arrived after the +constructor's positional sequence was already in use, and inserting them +in the middle would have rebound every positional argument after them. + +The reference is not resolved or fetched. An `osdf://` or `docker://` URL +is not readable from the submit host, so requiring that would refuse the +ordinary OSG case; only a value that could corrupt the submit file is +rejected. + +`use_singularity` / `singularity_image` remain, for sites that honour only +the legacy `+SingularityImage` form. Setting **both** is refused: emitted +together, which one takes effect is decided by the site. + +`when_to_transfer_output` takes HTCondor's three legal values — +`ON_EXIT` (default), `ON_EXIT_OR_EVICT`, `ON_SUCCESS`. `NEVER` is **not** +one of them: condor accepts it with rc=0 and silently materialises +`ON_EXIT`, so a caller setting it to suppress transfer gets transfer. + +`ON_EXIT_OR_EVICT` is **refused by this queue**, and the refusal explains +why. HTCondor: *"If a file listed in transfer_output_files does not exist +at eviction time, the job will go on hold."* This queue always emits an +explicit `transfer_output_files` led by `level_.json`, which the +bootstrap writes only after the generator returns — so every mid-run +eviction would hold the job rather than reschedule it, which is worse +than the `ON_EXIT` behaviour the setting gets reached for. It also spools +the evicted sandbox to the access point's `SPOOL`, a shared volume that +is easy to fill. To preserve partial work, use HTCondor +self-checkpointing (`checkpoint_exit_code` / `transfer_checkpoint_files`), +which does not require the final outputs to exist. + +The trap is specific to **listing an output that only exists on success**, +not to `ON_EXIT_OR_EVICT` itself. A submitter whose `transfer_output_files` +names only paths created early — a working directory, say — can use it +safely. This queue cannot, because it always leads that list with the +level marker. If that ever changes, this refusal should lift with it. + +`container_image`, a non-default `when_to_transfer_output`, and the +`extra_transfer_*` lists are all emitted by `build_worker` — which +`submit()` bypasses when `subdag_factory` is set, because the sub-DAG +writes its own submit descriptions. Combining any of them with a factory +is therefore **refused**, in the constructor and again in `submit()` +(both are plain attributes, so a constructor-only check is walked past by +assigning either one afterwards). Put the setting in the DAG the factory +generates. For `ON_EXIT_OR_EVICT` the refusal does double duty: the +sub-DAG path never reaches the check in `build_worker`, so accepting it +would route the one value this queue rejects around its own rejection. `extra_periodic_release` takes a single-line ClassAd expression for sites whose pool holds jobs for reasons the queue does not model — an diff --git a/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/database.py b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/database.py index a069899df..dfcb18d24 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/database.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/database.py @@ -286,6 +286,69 @@ def _validate_subcode_exclusions(value: Any, *, what: str return out +#: HTCondor's `when_to_transfer_output` vocabulary, as the JDL defines +#: it. Three values, not four: NEVER is NOT one of them, and condor does +#: not say so -- measured against condor 25.13.1, `NEVER` submits with +#: rc=0 and materialises as `WhenToTransferOutput="ON_EXIT"`, silently. +#: A caller setting it to suppress transfer gets transfer. `BANANA` is +#: rejected loudly, so listing NEVER blessed the one value condor +#: discards while catching nothing condor would not have caught. +WHEN_TO_TRANSFER_OUTPUT = ("ON_EXIT", "ON_EXIT_OR_EVICT", "ON_SUCCESS") +DEFAULT_WHEN_TO_TRANSFER_OUTPUT = "ON_EXIT" + +def _validate_when_to_transfer(value: Any) -> str: + """One of HTCondor's four, normalised to upper case.""" + if value is None: + return DEFAULT_WHEN_TO_TRANSFER_OUTPUT + if not isinstance(value, str): + raise TypeError( + "when_to_transfer_output must be a string, got {0!r}".format( + type(value).__name__)) + text = value.strip().upper() + if text not in WHEN_TO_TRANSFER_OUTPUT: + raise ValueError( + "when_to_transfer_output must be one of {0}, got {1!r}".format( + ", ".join(WHEN_TO_TRANSFER_OUTPUT), value)) + return text + + +def _validate_container_image(value: Any) -> str: + """A container image reference for `universe = container`. + + Not resolved or fetched. The reference is routinely an osdf:// or + docker:// URL the submit host cannot read, so a checker demanding + local readability would refuse the ordinary OSG case. Only the shape + that could corrupt the submit file is refused. + """ + if value is None: + return "" + if not isinstance(value, str): + raise TypeError( + "container_image must be a string, got {0!r}".format( + type(value).__name__)) + text = value.strip() + if not text: + return "" + if "\n" in text or "\r" in text: + raise ValueError( + "container_image must be a single line: a newline would end " + "the submit command and let the rest be read as further " + "commands") + if text.endswith("\\"): + # A trailing backslash is a submit-file line continuation, so the + # NEXT command is swallowed into this value. Measured: it eats + # `arguments`, condor_submit returns 0, and every worker then + # runs the bootstrap with no --sim-name/--level, exiting 2 with + # nothing in the submit file to explain it. + raise ValueError( + "container_image must not end in a backslash: the submit " + "parser reads it as a line continuation and swallows the " + "next command") + if "\x00" in text: + raise ValueError("container_image must not contain a NUL byte") + return text + + def _validate_release_expression(value: Any, *, what: str) -> str: """Check a ClassAd expression destined for a submit command. @@ -408,6 +471,21 @@ def _validate_transfer_entries(entries: Any, *, what: str, # than losing the release arm. Per-sim sizes go through # Archive.set_resources, which composes rather than substitutes. "request_memory", + # Composed by the queue since container support landed, and each has + # an argument now, so extra_condor_cmds would silently replace a line + # the queue had already written. + # + # `universe` is deliberately NOT here. An earlier draft protected it + # on the claim that a container_image under a vanilla universe is + # ignored by condor. Measured against condor 25.13.1 that is false: + # vanilla+container_image, container+container_image, and declaring + # no universe at all produce byte-identical job ads, and the JDL says + # universe "can either be optionally set to container or not declared + # at all". Protecting it would be a breaking change with no defect + # behind it, and would leave no route to `local`, `scheduler` or + # `grid` at all. + "container_image", + "when_to_transfer_output", }) #: What to use instead of each refused key. Kept beside the set so a new @@ -422,6 +500,8 @@ def _validate_transfer_entries(entries: Any, *, what: str, "expression instead of replacing it", "request_memory": "the request_memory argument, or " "Archive.set_resources for a per-sim override", + "container_image": "the container_image argument", + "when_to_transfer_output": "the when_to_transfer_output argument", } #: Basenames the archive itself stages into the worker sandbox. Condor @@ -1680,6 +1760,26 @@ class DualCondorRunQueue(RunQueue): them); the allowlist is the OSG-blessed alternative. Pass getenv='True' explicitly only on sites that allow it. + container_image : str -- KEYWORD-ONLY. Image reference for + HTCondor's container universe + (`osdf://`, `docker://`, a local + .sif). Setting it switches the job + from vanilla to `universe = + container`. Mutually exclusive with + use_singularity, which is the legacy + +SingularityImage form. + when_to_transfer_output: str -- KEYWORD-ONLY. One of ON_EXIT + (default), + ON_EXIT_OR_EVICT, ON_SUCCESS, NEVER. + ON_EXIT discards the sandbox on + eviction, so a preemptable pool loses + whatever the job had already written. + Like container_image and the transfer + extras, this is emitted by build_worker + and so is refused together with + subdag_factory: a sub-DAG writes its own + submit descriptions, and its nodes would + transfer at the default instead. use_singularity : bool singularity_image: str -- required if use_singularity=True oom_hold_codes : seq -- hold codes this site reports when a @@ -1782,21 +1882,52 @@ def __init__(self, oom_memory_factor: float = 1.5, subdag_factory: Optional[Callable[[Any, str, int], str]] = None, submit_mode: str = "submit", + # Keyword-only, and last: these arrived after the + # positional sequence above was already in use. Splicing + # them in next to use_singularity, where they belong by + # topic, would have rebound every positional argument + # from that point on — a caller's positional True for + # use_singularity would land in container_image and + # raise TypeError from its validator. + *, + container_image: Optional[str] = None, + when_to_transfer_output: Optional[str] = None, **submit_kwargs: Any): self.run_pool = run_pool self.run_collector = run_collector self.extra_transfer_input_files = extra_transfer_input_files self.extra_transfer_output_files = extra_transfer_output_files - if (self.extra_transfer_input_files or self.extra_transfer_output_files) \ - and subdag_factory is not None: + self.container_image = container_image + # Assigned before the guards below, which read it: the setter + # normalises None to the ON_EXIT default, and what the guard asks + # is whether a NON-default policy was requested. + self.when_to_transfer_output = when_to_transfer_output + if (self.extra_transfer_input_files or self.extra_transfer_output_files + or self.container_image) and subdag_factory is not None: # Fail early for the common case. submit() re-checks, because # both of these are plain attributes and assigning either after # construction reaches the same silently-ignoring path. raise ValueError( - "extra_transfer_{input,output}_files are applied by " - "build_worker, which is bypassed when subdag_factory is set: " - "the sub-DAG owns its own submit descriptions. Put the extra " - "entries in the sub-DAG the factory generates instead.") + "extra_transfer_{input,output}_files and container_image are " + "applied by build_worker, which is bypassed when " + "subdag_factory is set: the sub-DAG owns its own submit " + "descriptions. Put them in the sub-DAG the factory generates " + "instead.") + if (self.when_to_transfer_output != DEFAULT_WHEN_TO_TRANSFER_OUTPUT + and subdag_factory is not None): + # Same bypass, and worse than merely dropped: build_worker is + # also where ON_EXIT_OR_EVICT is REFUSED, so accepting it here + # would route the one value this queue rejects around its own + # rejection while the sub-DAG's nodes ran under the default. + raise ValueError( + "when_to_transfer_output={0} is applied by build_worker, " + "which is bypassed when subdag_factory is set: the sub-DAG " + "owns its own submit descriptions, so its nodes would " + "transfer at the default {1} with nothing to show the " + "request was dropped. Set the timing in the sub-DAG the " + "factory generates instead.".format( + self.when_to_transfer_output, + DEFAULT_WHEN_TO_TRANSFER_OUTPUT)) self.request_memory = int(request_memory) self.request_disk = request_disk self.accounting_group = accounting_group or os.environ.get("LIGO_ACCOUNTING") @@ -1900,6 +2031,22 @@ def extra_periodic_release(self, value: Any) -> None: self._extra_periodic_release = _validate_release_expression( value, what="extra_periodic_release") + @property + def container_image(self) -> str: + return self._container_image + + @container_image.setter + def container_image(self, value: Any) -> None: + self._container_image = _validate_container_image(value) + + @property + def when_to_transfer_output(self) -> str: + return self._when_to_transfer_output + + @when_to_transfer_output.setter + def when_to_transfer_output(self, value: Any) -> None: + self._when_to_transfer_output = _validate_when_to_transfer(value) + @property def oom_hold_codes(self) -> Tuple[int, ...]: return self._oom_hold_codes @@ -2085,9 +2232,22 @@ def build_worker(self, archive: Archive, sim_name: str, lines: List[str] = [ "# Auto-generated by RIFT.simulation_manager.database." "DualCondorRunQueue", - "universe = vanilla", + "universe = {}".format( + "container" if self.container_image else "vanilla"), "executable = {}".format(bootstrap), ] + if self.container_image: + if self.use_singularity: + # Two ways to say "run this in a container", emitted + # together, resolved by the site. Refuse rather than let + # the job run under whichever the pool happens to honour. + raise ValueError( + "container_image and use_singularity are alternative " + "ways to request a container: the modern container " + "universe and the legacy +SingularityImage form. Set " + "one. container_image is the one OSG documents now.") + lines.append("container_image = {}".format( + self.container_image)) args_tail = ["--sim-name", sim_name, "--level", str(int(level))] if prev_basenames: args_tail.append("--prev-levels") @@ -2097,7 +2257,34 @@ def build_worker(self, archive: Archive, sim_name: str, if transfer_in: lines.append("transfer_input_files = {}".format(",".join(transfer_in))) lines.append("should_transfer_files = YES") - lines.append("when_to_transfer_output = ON_EXIT") + if self.when_to_transfer_output == "ON_EXIT_OR_EVICT": + # Refused, not emitted. HTCondor: "If a file listed in + # transfer_output_files does not exist at eviction time, the + # job will go on hold." This queue ALWAYS writes an explicit + # transfer_output_files led by level_.json, and the + # bootstrap writes that marker only after run() returns -- so + # every mid-run eviction would hold the job instead of + # rescheduling it. That is strictly worse than the ON_EXIT + # behaviour the setting gets reached for. + # + # It also spools the evicted sandbox to the AP. On the access + # point this was developed against, /var/lib/condor/spool is + # 100% full with 28 MB free and shared by every user; filling + # it has caused a campaign-wide outage before. + # + # What does preserve partial work is HTCondor + # self-checkpointing, which does not require the final + # outputs to exist. That is an archive-design change, not a + # knob, so this raises rather than pretending to offer it. + raise ValueError( + "when_to_transfer_output=ON_EXIT_OR_EVICT cannot be used " + "with this queue: it emits transfer_output_files led by " + "level_.json, which does not exist until the job " + "succeeds, and HTCondor holds a job whose listed output " + "is missing at eviction. Use checkpoint_exit_code / " + "transfer_checkpoint_files to preserve partial work.") + lines.append("when_to_transfer_output = {}".format( + self.when_to_transfer_output)) # Backend-supplied products, beyond the level_.json marker. # transfer_output_files is explicit, so HTCondor returns ONLY what # is named here: anything else the worker wrote is destroyed with @@ -2296,6 +2483,35 @@ def submit(self, archive: Archive, sim_names: Iterable[str] "sub-DAG owns its own submit descriptions. Put the " "entries in the DAG the factory generates, or clear " "subdag_factory.") + # Same reasoning, same path. Without this a queue + # configured with a container image submits a sub-DAG + # whose nodes run the science OUTSIDE the container, + # with no error and nothing in the submit files to + # show it -- the exact silent substitution the + # container argument exists to remove. + if self.container_image: + raise ValueError( + "container_image is applied by build_worker, which " + "this sub-DAG path bypasses, so the sub-DAG's nodes " + "would run outside the container. Set the container " + "in the DAG the factory generates, or clear " + "subdag_factory.") + # And the transfer timing, which build_worker both + # emits and polices. Silently ignoring it here loses a + # deliberate ON_SUCCESS, and lets ON_EXIT_OR_EVICT -- + # the one value build_worker refuses outright -- reach + # a submit with no complaint from either end. + if (self.when_to_transfer_output + != DEFAULT_WHEN_TO_TRANSFER_OUTPUT): + raise ValueError( + "when_to_transfer_output={0} is applied by " + "build_worker, which this sub-DAG path bypasses, " + "so the sub-DAG's nodes would transfer at the " + "default {1} instead. Set the timing in the DAG " + "the factory generates, or clear " + "subdag_factory.".format( + self.when_to_transfer_output, + DEFAULT_WHEN_TO_TRANSFER_OUTPUT)) work_path = self.subdag_factory(archive, sim, lvl) nodes.append((sim, lvl, work_path, True)) else: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/tests/test_condor_container_universe.py b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/tests/test_condor_container_universe.py new file mode 100644 index 000000000..3778e1988 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/simulation_manager/tests/test_condor_container_universe.py @@ -0,0 +1,403 @@ +"""DualCondorRunQueue's container universe and output-transfer timing. + +Two things a backend previously had to reach through `extra_condor_cmds`, +which is emitted last and therefore *replaces* the queue's own lines +rather than extending them. + +`universe = container` + `container_image` is how OSG documents running +in a container today. The queue only knew the legacy `+SingularityImage` +form, so every backend targeting OSG hand-rolled the modern one — and in +doing so silently replaced the queue's `universe` line, with the winner +decided by which line condor read last. + +`when_to_transfer_output` was hardcoded `ON_EXIT`. On a preemptable pool +that discards the sandbox when a job is evicted, throwing away whatever +the job had already written. `ON_EXIT_OR_EVICT` keeps it, and a backend +whose science *is* output files needs that. + +Run with the RIFT-importable interpreter, e.g.: + + PYTHONPATH=<...>/MonteCarloMarginalizeCode/Code \ + python -m pytest -q .../tests/test_condor_container_universe.py +""" + +from __future__ import annotations + +import inspect +import shutil +import subprocess + +import pytest + +from RIFT.simulation_manager.database import ( + Archive, DualCondorRunQueue, Manifest, WHEN_TO_TRANSFER_OUTPUT, +) + +IMAGE = "osdf:///ospool/ap41/data/example/supernu-v2.sif" + + +def _generator_src(): + return ( + "import json, os\n" + "def run(params, sim_dir, level, prev_levels):\n" + " p = os.path.join(sim_dir, 'level_%d.json' % level)\n" + " with open(p, 'w') as f:\n" + " json.dump({'level': level}, f)\n" + " return p\n" + ) + + +@pytest.fixture +def archive(tmp_path): + code = tmp_path / "src" + code.mkdir() + (code / "generator.py").write_text(_generator_src()) + manifest = Manifest.new(name="container", request_queue_kind="condor", + run_queue_kind="condor") + return Archive( + base_location=tmp_path / "arch", manifest=manifest, + generator_spec={"module_path": str(code / "generator.py"), + "entrypoint": "generator:run"}, + ) + + +def _build(archive, queue, level=1): + name = archive.register({"x": 1}, target_level=level) + return open(queue.build_worker(archive, name, level)).read() + + +def _command(sub_text, key): + hits = [l for l in sub_text.splitlines() + if l.split("=")[0].strip().lower() == key] + assert len(hits) == 1, hits + return hits[0].split("=", 1)[1].strip() + + +# -------------------------------------------------------------------- +# container universe +# -------------------------------------------------------------------- + +def test_no_image_is_a_vanilla_job(archive): + """The default must not move: every existing deployment submits + vanilla and none of them asked for this.""" + sub = _build(archive, DualCondorRunQueue()) + assert _command(sub, "universe") == "vanilla" + assert "container_image" not in sub + + +def test_an_image_selects_the_container_universe(archive): + """One argument, not two. A backend that has to remember to set the + universe as well will eventually forget, and a container_image under + a vanilla universe is silently ignored by condor.""" + sub = _build(archive, DualCondorRunQueue(container_image=IMAGE)) + assert _command(sub, "universe") == "container" + assert _command(sub, "container_image") == IMAGE + + +def test_the_image_is_not_resolved_or_fetched(archive): + """An osdf:// or docker:// reference is not readable from the submit + host, and demanding that it be would refuse the ordinary OSG case.""" + for ref in ("docker://library/python:3.11", + "osdf:///ospool/ap41/data/nobody/nothing.sif", + "/no/such/path/local.sif"): + sub = _build(archive, DualCondorRunQueue(container_image=ref)) + assert _command(sub, "container_image") == ref + + +def test_the_two_container_mechanisms_are_refused_together(archive): + """Emitting both leaves the outcome to whichever the site honours.""" + q = DualCondorRunQueue(container_image=IMAGE, use_singularity=True, + singularity_image="/cvmfs/x.sif") + with pytest.raises(ValueError, match="container_image"): + _build(archive, q) + + +def test_the_legacy_form_still_works_alone(archive): + """Not a deprecation. Sites that only honour +SingularityImage exist.""" + sub = _build(archive, DualCondorRunQueue( + use_singularity=True, singularity_image="/cvmfs/x.sif")) + assert _command(sub, "universe") == "vanilla" + assert "SingularityImage" in sub + + +@pytest.mark.parametrize("bad", [ + "osdf:///x.sif\ngetenv = True", + "osdf:///x.sif\r\ntransfer_output_files = nothing", +]) +def test_a_newline_in_the_image_cannot_add_a_submit_command(bad): + with pytest.raises(ValueError): + DualCondorRunQueue(container_image=bad) + + +@pytest.mark.parametrize("bad", [17, ["a"], {"x": 1}, object()]) +def test_a_non_string_image_is_refused(bad): + with pytest.raises(TypeError): + DualCondorRunQueue(container_image=bad) + + +# -------------------------------------------------------------------- +# when_to_transfer_output +# -------------------------------------------------------------------- + +def test_the_default_transfer_timing_is_unchanged(archive): + assert _command(_build(archive, DualCondorRunQueue()), + "when_to_transfer_output") == "ON_EXIT" + + +@pytest.mark.parametrize("value", [v for v in WHEN_TO_TRANSFER_OUTPUT + if v != "ON_EXIT_OR_EVICT"]) +def test_every_usable_value_is_emitted(archive, value): + sub = _build(archive, DualCondorRunQueue(when_to_transfer_output=value)) + assert _command(sub, "when_to_transfer_output") == value + + +def test_the_vocabulary_is_condors(archive): + """Guards the set itself. It drifted once already -- NEVER was in it, + and no test noticed because every test read the emitted text rather + than asking condor.""" + assert set(WHEN_TO_TRANSFER_OUTPUT) == { + "ON_EXIT", "ON_EXIT_OR_EVICT", "ON_SUCCESS"} + + +def test_it_is_normalised_not_passed_through(archive): + sub = _build(archive, + DualCondorRunQueue(when_to_transfer_output=" on_success ")) + assert _command(sub, "when_to_transfer_output") == "ON_SUCCESS" + + +def test_never_is_refused_because_condor_discards_it(archive): + """NEVER is not in the JDL. Measured against condor 25.13.1 it + submits rc=0 and materialises as ON_EXIT -- so a caller setting it to + suppress transfer gets transfer, silently. An illegal value like + BANANA condor rejects loudly on its own, so accepting NEVER was the + only thing this validator actually changed, in the wrong direction.""" + with pytest.raises(ValueError, match="when_to_transfer_output"): + DualCondorRunQueue(when_to_transfer_output="NEVER") + + +def test_on_exit_or_evict_is_refused_by_this_queue(archive): + """HTCondor holds a job whose listed output is missing at eviction, + and this queue always lists level_.json, which exists only after + the job succeeds. Every mid-run eviction would hold rather than + reschedule -- worse than the ON_EXIT it is reached for. Constructing + is allowed; building the submit description is where it raises, so + the message lands with the job that would have been broken.""" + q = DualCondorRunQueue(when_to_transfer_output="ON_EXIT_OR_EVICT") + with pytest.raises(ValueError, match="checkpoint_exit_code"): + _build(archive, q) + + +@pytest.mark.parametrize("bad", ["ON_EVICT", "always", "", "ON_EXIT_OR_EVIC"]) +def test_a_value_condor_does_not_know_is_refused(bad): + """Caught here rather than by the schedd: by the time condor_submit + refuses it the archive has already recorded the sim as dispatched, so + it presents as a stuck simulation, not a configuration error.""" + with pytest.raises(ValueError, match="when_to_transfer_output"): + DualCondorRunQueue(when_to_transfer_output=bad) + + +@pytest.mark.parametrize("bad", [17, ["ON_EXIT"], object()]) +def test_a_non_string_timing_is_refused(bad): + with pytest.raises(TypeError): + DualCondorRunQueue(when_to_transfer_output=bad) + + +# -------------------------------------------------------------------- +# the old way in is closed +# -------------------------------------------------------------------- + +@pytest.mark.parametrize("key,expected", [ + ("container_image", "container_image"), + ("when_to_transfer_output", "when_to_transfer_output"), +]) +def test_the_old_route_is_refused_and_names_the_argument(archive, key, + expected): + """These were reachable only through extra_condor_cmds before there + were arguments for them. Leaving that open beside the argument means + the next backend author finds it first, and a universe set twice is + decided by whichever line condor reads last.""" + q = DualCondorRunQueue(extra_condor_cmds={key: "whatever"}) + with pytest.raises(ValueError, match=expected): + _build(archive, q) + + +def test_assignment_after_construction_is_validated(archive): + q = DualCondorRunQueue() + with pytest.raises(ValueError): + q.container_image = "osdf:///x.sif\ngetenv = True" + with pytest.raises(ValueError): + q.when_to_transfer_output = "ON_EVICT" + assert "getenv = True" not in _build(archive, q) + + +def test_the_policy_survives_the_manifest(tmp_path): + from RIFT.simulation_manager.database import make_queues_from_manifest + + code = tmp_path / "src" + code.mkdir() + (code / "generator.py").write_text(_generator_src()) + manifest = Manifest.new( + name="container_manifest", request_queue_kind="condor", + run_queue_kind="condor", + run_queue_extra={"container_image": IMAGE, + "when_to_transfer_output": "ON_SUCCESS"}) + Archive(base_location=tmp_path / "arch", manifest=manifest, + generator_spec={"module_path": str(code / "generator.py"), + "entrypoint": "generator:run"}) + reopened = Archive(base_location=tmp_path / "arch") + _, run_queue = make_queues_from_manifest(reopened) + assert run_queue.container_image == IMAGE + sub = _build(reopened, run_queue) + assert _command(sub, "universe") == "container" + assert _command(sub, "when_to_transfer_output") == "ON_SUCCESS" + + +@pytest.mark.parametrize("kwargs", [ + {}, + {"container_image": IMAGE}, + {"when_to_transfer_output": "ON_SUCCESS"}, + {"container_image": IMAGE, "when_to_transfer_output": "ON_SUCCESS"}, + {"use_singularity": True, "singularity_image": "/cvmfs/x.sif"}, +]) +def test_condor_accepts_every_shape(archive, tmp_path, kwargs): + """No hand-reading of the submit file substitutes for condor parsing + it. -dry-run contacts no schedd and queues nothing.""" + condor_submit = shutil.which("condor_submit") + if condor_submit is None: + pytest.skip("condor_submit not on PATH") + path = tmp_path / "c.sub" + path.write_text(_build(archive, DualCondorRunQueue(**kwargs))) + out = tmp_path / "c.dry" + proc = subprocess.run([condor_submit, "-dry-run:oauth=1", str(out), str(path)], + capture_output=True, text=True) + assert proc.returncode == 0, proc.stderr + # rc=0 only proves the file parses -- condor accepts a garbage image + # string with rc=0, and silently rewrites values it dislikes. Read + # the materialised ad back and compare it to what was asked for. + ad = {} + for line in out.read_text().splitlines(): + if "=" in line: + k, v = line.split("=", 1) + ad[k.strip().lower()] = v.strip().strip('"') + assert ad.get("whentotransferoutput") == kwargs.get( + "when_to_transfer_output", "ON_EXIT") + if kwargs.get("container_image"): + assert ad.get("wantcontainer") == "true", ad.get("wantcontainer") + assert IMAGE.rsplit("/", 1)[-1] in out.read_text() + + +# -------------------------------------------------------------------- +# the sub-DAG path bypasses build_worker entirely +# -------------------------------------------------------------------- + +def test_a_container_and_a_subdag_are_refused_together(): + """submit() bypasses build_worker when subdag_factory is set, so the + image would never be emitted and the sub-DAG's nodes would run the + science OUTSIDE the container -- no error, nothing in the submit + files to show it. The transfer extras are guarded on this path for + the same reason; the image was not.""" + with pytest.raises(ValueError, match="container_image"): + DualCondorRunQueue(container_image=IMAGE, + subdag_factory=lambda a, s, l: "/tmp/x") + + +def test_assigning_either_one_late_is_still_refused(archive): + """Both are plain attributes, so a constructor-only check is + bypassed by assignment -- which is why submit() re-checks.""" + q = DualCondorRunQueue(container_image=IMAGE) + q.subdag_factory = lambda a, s, l: "/tmp/x" + name = archive.register({"x": 1}, target_level=1) + with pytest.raises(ValueError, match="container_image"): + q.submit(archive, [name]) + + +@pytest.mark.parametrize("value", [v for v in WHEN_TO_TRANSFER_OUTPUT + if v != "ON_EXIT"]) +def test_a_transfer_timing_and_a_subdag_are_refused_together(value): + """The timing is emitted by build_worker too, so the sub-DAG's nodes + would transfer at the default ON_EXIT and the request would vanish. + ON_EXIT_OR_EVICT is worse than dropped: build_worker is also where it + is REFUSED, so this path routed the one value the queue rejects + around its own rejection.""" + with pytest.raises(ValueError, match="when_to_transfer_output"): + DualCondorRunQueue(when_to_transfer_output=value, + subdag_factory=lambda a, s, l: "/tmp/x") + + +@pytest.mark.parametrize("value", [v for v in WHEN_TO_TRANSFER_OUTPUT + if v != "ON_EXIT"]) +def test_assigning_the_timing_late_is_still_refused(archive, value): + """Same plain-attribute hole as the image, in both orders.""" + name = archive.register({"x": 1}, target_level=1) + + q = DualCondorRunQueue(when_to_transfer_output=value, submit_mode="embed") + q.subdag_factory = lambda a, s, l: "/tmp/x" + with pytest.raises(ValueError, match="when_to_transfer_output"): + q.submit(archive, [name]) + + q = DualCondorRunQueue(submit_mode="embed", + subdag_factory=lambda a, s, l: "/tmp/x") + q.when_to_transfer_output = value + with pytest.raises(ValueError, match="when_to_transfer_output"): + q.submit(archive, [name]) + + +def test_the_default_timing_still_composes_with_a_subdag(archive, tmp_path): + """The refusal is of a NON-default policy that would be dropped, not + of sub-DAGs: a backend whose work unit is itself a DAG (GW PE via + util_RIFT_pseudo_pipe) never asked for a timing and must still + submit.""" + made = tmp_path / "child.dag" + made.write_text("# noop\n") + q = DualCondorRunQueue(submit_mode="embed", + subdag_factory=lambda a, s, l: str(made)) + name = archive.register({"x": 1}, target_level=1) + q.submit(archive, [name]) + wrapper = open(q.last_wrapper_dag_path).read() + assert "SUBDAG EXTERNAL {}_lvl1 {}".format(name, made) in wrapper + + +@pytest.mark.parametrize("bad", ["osdf:///x.sif \\", "osdf:///x.sif\x00"]) +def test_an_image_that_would_corrupt_the_submit_file_is_refused(bad): + """A trailing backslash is a submit-file line continuation: it + swallows the next command, which is `arguments`. Measured before the + fix -- condor_submit returned 0 and the workers ran the bootstrap + with no --sim-name/--level, exiting 2 forever with nothing in the + submit file to explain it.""" + with pytest.raises(ValueError): + DualCondorRunQueue(container_image=bad) + + +def test_the_new_arguments_did_not_move_the_old_ones(): + """Both arrived after the constructor's positional sequence was in + use. An earlier draft spliced them in next to use_singularity, where + they belong by topic -- which rebound every positional argument from + that point on, so a caller's positional True for use_singularity + reached the container_image validator and raised TypeError. They are + keyword-only and last instead.""" + q = DualCondorRunQueue(None, None, 4096, "4G", None, None, None, + True, "/cvmfs/x.sif") + assert q.use_singularity is True + assert q.singularity_image == "/cvmfs/x.sif" + assert q.container_image == "" + assert q.when_to_transfer_output == "ON_EXIT" + params = inspect.signature(DualCondorRunQueue.__init__).parameters + for name in ("container_image", "when_to_transfer_output"): + assert params[name].kind is inspect.Parameter.KEYWORD_ONLY + positional = [n for n, p in params.items() + if p.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD] + assert positional[:10] == [ + "self", "run_pool", "run_collector", "request_memory", + "request_disk", "accounting_group", "accounting_group_user", + "getenv", "use_singularity", "singularity_image"] + + +def test_universe_is_not_protected(archive): + """Deliberate. An earlier draft protected it, on the claim that a + container_image under a vanilla universe is ignored by condor. + Measured: vanilla+image, container+image and no universe at all give + byte-identical ads. Protecting it would be a breaking change with no + defect behind it, and would leave no route to local/scheduler/grid.""" + sub = _build(archive, DualCondorRunQueue( + extra_condor_cmds={"universe": "scheduler"})) + assert "scheduler" in sub diff --git a/MonteCarloMarginalizeCode/Code/bin/cepp_basic_htcondor b/MonteCarloMarginalizeCode/Code/bin/cepp_basic_htcondor index ce0ea6839..dc9297c59 100755 --- a/MonteCarloMarginalizeCode/Code/bin/cepp_basic_htcondor +++ b/MonteCarloMarginalizeCode/Code/bin/cepp_basic_htcondor @@ -939,7 +939,10 @@ unify_job.write_sub_file() exe_evidence = which('util_CIPDirSummarizeEvidence.py') -evidence_job, evidence_job_name = dag_utils.write_convert_sub(tag='evidence', exe=exe_evidence, log_dir='', file_input='', arg_str=" --cip-dir iteration_$(macroiteration)_cip --output evidence_$(macroiteration) ", universe=local_worker_universe, no_grid=True) +evidence_source = " --cip-dir iteration_$(macroiteration)_cip" +if opts.cip_explode_jobs is None: + evidence_source = " --cip-dir . --cip-prefix overlap-grid-$(macroiterationnext)" +evidence_job, evidence_job_name = dag_utils.write_convert_sub(tag='evidence', exe=exe_evidence, log_dir='', file_input='', arg_str=evidence_source+" --output evidence_$(macroiteration) ", universe=local_worker_universe, no_grid=True) evidence_job.add_condor_cmd("initialdir",opts.working_directory) evidence_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/evidence-$(cluster)-$(process).log") evidence_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/evidence-$(cluster)-$(process).err") @@ -1531,6 +1534,7 @@ for it in np.arange(it_start,opts.n_iterations): # It is also *backward looking* evidence_node = Node(evidence_job) evidence_node.add_variable("macroiteration",it-1) # LAST ITERATION being analyzed, not this one + evidence_node.add_variable("macroiterationnext",it) evidence_node.retry = opts.general_retries if it > it_start: evidence_node.add_parent(parent_fit_node) # run evidence after previous iteration, on previous diff --git a/MonteCarloMarginalizeCode/Code/bin/convert_output_format_ile2inference b/MonteCarloMarginalizeCode/Code/bin/convert_output_format_ile2inference index f41e5a20f..bb66eca76 100755 --- a/MonteCarloMarginalizeCode/Code/bin/convert_output_format_ile2inference +++ b/MonteCarloMarginalizeCode/Code/bin/convert_output_format_ile2inference @@ -38,7 +38,9 @@ optp.add_option("--export-tides",action='store_true',help="Include tidal paramet optp.add_option("--export-eos",action='store_true',help="Include EOS index parameter") optp.add_option("--export-cosmology",action='store_true',help="Include source frame masses and redshift") optp.add_option("--export-weights",action='store_true',help="Include a field 'weights' equal to L p/ps") +optp.add_option("--export-EOB-parameters", action="store_true", help="Include EOB Parameter: a6c") optp.add_option("--export-eccentricity", action="store_true", help="Include eccentricity") +optp.add_option("--export-hyperbolic", action="store_true", help="Include hyperbolic") optp.add_option("--export-meanPerAno", action="store_true", help="Include meanPerAno") optp.add_option("--with-cosmology",default="Planck15",help="Specific cosmology to use") optp.add_option("--use-interpolated-cosmology",action='store_true',help="Specific cosmology to use") @@ -75,14 +77,20 @@ if opts.convention == 'LI': print("# m1 m2 a1x a1y a1z a2x a2y a2z mc eta ra dec time phiorb incl psi distance Npts lnL p ps neff mtotal q chi_eff chi_p",end=' ') if opts.export_extra_spins: print( 'theta_jn phi_jl tilt1 tilt2 costilt1 costilt2 phi12 a1 a2 psiJ',end=' ') - if opts.export_tides: + if opts.export_tides and not (opts.export_EOB_parameters): print( "lambda1 lambda2 lam_tilde",end=' ') + if opts.export_tides and opts.export_EOB_parameters: + print( "lambda1 lambda2 lam_tilde a6c",end=' ') + if not (opts.export_tides) and opts.export_EOB_parameters: + print("a6c", end=' ') if opts.export_eos: print( "eos_indx",end=' ') if opts.export_cosmology: print( " m1_source m2_source mc_source mtotal_source redshift ",end=' ') if opts.export_weights: - print( " weights ", ) + print(" weights ", end=' ') + if opts.export_hyperbolic: + print("E0 p_phi0 ", end=' ') if opts.export_eccentricity: if not opts.export_meanPerAno: print( "eccentricity ", ) @@ -104,6 +112,9 @@ if opts.convention == 'LI': else: ecc = [row.alpha4 for row in points] meanPerAno = [row.alpha for row in points] + if opts.export_hyperbolic: + p_phi0 = [row.beta for row in points] + E0 = [row.psi3 for row in points] wt = np.exp(like)*p/ps # remove nan wt - can happen for extrinsic output @@ -168,7 +179,14 @@ if opts.convention == 'LI': if hasattr(pt, 'alpha6'): P.lambda2 = pt.alpha6 lam_tilde = P.extract_param("LambdaTilde") - print( P.lambda1, P.lambda2, lam_tilde,end=' ') + if not (opts.export_EOB_parameters): + print( P.lambda1, P.lambda2, lam_tilde,end=' ') + if opts.export_EOB_parameters: + P.a6c=pt.psi0 + print( P.lambda1, P.lambda2, lam_tilde, P.a6c,end=' ') + if not(opts.export_tides) and opts.export_EOB_parameters: + P.a6c=pt.psi0 + print(P.a6c, end=' ') if opts.export_eos: eos_indx = P.eos_table_index print(eos_indx, end=' ') @@ -181,6 +199,8 @@ if opts.convention == 'LI': print( m1_source, m2_source, mc_here/(1+z), mtot_here/(1+z), float(z), end=' ') if opts.export_weights: print(wt[indx],end=' ') + if opts.export_hyperbolic: + print(E0[indx], p_phi0[indx], end=' ') if opts.export_eccentricity: if not opts.export_meanPerAno: print(ecc[indx]) @@ -194,17 +214,33 @@ if opts.convention == 'LI': # +# The .hdf5 reader below walks the fixed-column waveform_parameters array +# (lalsimutils.hdf_params), which has no a6c, E0 or p_phi0 fields, so its row +# writer cannot emit them. Refuse here, before the header promises columns the +# rows would not carry. +if any(".hdf5" in fname for fname in args): + if opts.export_EOB_parameters: + raise Exception(" Not implemented for hdf5 export : --export-EOB-parameters") + if opts.export_hyperbolic: + raise Exception(" Not implemented for hdf5 export : --export-hyperbolic") + print( "# m1 m2 a1x a1y a1z a2x a2y a2z mc eta indx Npts ra dec tref phiorb incl psi dist p ps lnL mtotal q ",end=' ') if opts.export_extra_spins: print( 'thetaJN phi_jl tilt1 tilt2 phi12 a1 a2 psiJ',end=' ') -if opts.export_tides: +if opts.export_tides and not (opts.export_EOB_parameters): print( "lambda1 lambda2",end=' ') +if opts.export_tides and opts.export_EOB_parameters: + print( "lambda1 lambda2 a6c",end=' ') +if not (opts.export_tides) and opts.export_EOB_parameters: + print("a6c", end=' ') if opts.export_eos: print( "eos_indx",end=' ') if opts.export_cosmology: print( " m1_source m2_source redshift ",end=' ') if opts.export_weights: print( " weights ",end=' ') +if opts.export_hyperbolic: + print("E0 p_phi0 ", end=' ') if opts.export_eccentricity: if not opts.export_meanPerAno: print( "eccentricity ",) @@ -249,6 +285,11 @@ for fname in args: Nmax = np.max([int(row.simulation_id) for row in points])+1 sim_id = np.array([int(row.simulation_id) for row in points])+1 + wt = np.exp(like)*p/ps + + if opts.export_hyperbolic: + p_phi0 = [row.beta for row in points] + E0 = [row.psi3 for row in points] if opts.export_eccentricity: if not opts.export_meanPerAno: ecc = [row.alpha4 for row in points] @@ -256,8 +297,6 @@ for fname in args: ecc = [row.alpha4 for row in points] meanPerAno = [row.alpha for row in points] - wt = np.exp(like)*p/ps - for indx in np.arange(len(points)): pt = points[indx] if not(hasattr(pt,'spin1x')): # no spins were provided. That means zero spin. Initialize to avoid an error @@ -283,8 +322,12 @@ for fname in args: True thetaJN, phiJL, theta1, theta2, phi12, chi1, chi2, psiJ = P.extract_system_frame() print( thetaJN, phiJL, theta1, theta2, phi12, chi1, chi2, psiJ,end=' ') - if opts.export_tides: + if opts.export_tides and not (opts.export_EOB_parameters): print(pt.alpha5, pt.alpha6,end=' ') + if opts.export_tides and opts.export_EOB_parameters: + print( pt.alpha5, pt.alpha6, pt.psi0,end=' ') + if not(opts.export_tides) and opts.export_EOB_parameters: + print(pt.psi0, end=' ') if opts.export_eos: print(pt.alpha, end=' ') if opts.export_cosmology: @@ -295,6 +338,8 @@ for fname in args: print( m1_source, m2_source, z,end=' ') if opts.export_weights: print(wt[indx],end=' ') + if opts.export_hyperbolic: + print(E0[indx], p_phi0[indx], end=' ') if opts.export_eccentricity: if not opts.export_meanPerAno: print(ecc[indx],) diff --git a/MonteCarloMarginalizeCode/Code/bin/convert_output_format_inference2ile b/MonteCarloMarginalizeCode/Code/bin/convert_output_format_inference2ile index 77055177a..abef29294 100755 --- a/MonteCarloMarginalizeCode/Code/bin/convert_output_format_inference2ile +++ b/MonteCarloMarginalizeCode/Code/bin/convert_output_format_inference2ile @@ -180,9 +180,20 @@ for indx in np.arange(opts.target_size): P.s1x=0 P.s2y=0 P.s2x=0 - if opts.add_eccentricity_params and P.eccentricity < 1e-5: + # Posterior columns win: only synthesize eccentricity if the samples have none + if "eccentricity" in samples_in.dtype.names: + P.eccentricity = samples_in["eccentricity"][fac_reduce*indx] + if "meanPerAno" in samples_in.dtype.names: + P.meanPerAno = samples_in["meanPerAno"][fac_reduce*indx] + elif opts.add_eccentricity_params: P.meanPerAno = np.random.uniform(0,2*np.pi) P.eccentricity = np.random.uniform(0,opts.ecc_max) # some convervative range + if "a6c" in samples_in.dtype.names: + P.a6c = samples_in["a6c"][fac_reduce*indx] + if "E0" in samples_in.dtype.names: + P.E0 = samples_in["E0"][fac_reduce*indx] + if "p_phi0" in samples_in.dtype.names: + P.p_phi0 = samples_in["p_phi0"][fac_reduce*indx] if ros_debug: if 'iota' in samples_in.dtype.names: diff --git a/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicIteration b/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicIteration index 5e9469ac0..949f3c77c 100755 --- a/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicIteration +++ b/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicIteration @@ -291,10 +291,11 @@ parser.add_argument("--ile-runtime-max-minutes",default=None,type=int,help="If n parser.add_argument("--cip-exe",default=None,help="filename of CIP or equivalent executable. Will default to `which util_ConstructIntrinsicPosterior_GenericCoordinates` in low-level code") parser.add_argument("--cip-post-exe",default=None,help="filename of script to execute with SCRIPT POST joblist scripname JOBID RETURN worklevel in the dag. This will be run after each CIP completes in the main dag. Objective is to enable larger-scale CIP queues which work efficiently to a target work level (e.g,, n_eff). Note this script is responsible for reading htcondor logs to identify the run directory!") parser.add_argument("--cip-exe-G",default=None,help="filename of CIP or equivalent executable, as ALTERNATE CIP used when a 'G' iteration is requested ") +parser.add_argument("--use-EOB-parameters",default=False,action='store_true') +parser.add_argument("--use-hyperbolic",default=False,action='store_true') parser.add_argument("--use-eccentricity",default=False,action='store_true') parser.add_argument("--use-meanPerAno",default=False,action='store_true') parser.add_argument("--use-eccentricity-squared-sampling",default=False,action='store_true') -parser.add_argument("--use-eccentricity-ln-sampling",default=False,action='store_true') parser.add_argument("--use-tabular-eos-file",default=False,action='store_true') parser.add_argument("--test-exe",default=None,help="filename of test code or equivalent executable. Must have a --test-output argument. Used for convergence testing or other termination. NOT ACTIVE; see 'convergence_test_samples.py' for example") parser.add_argument("--plot-exe",default=None,help="filename of plot code or equivalent executable. Will default to `which plot_posterior_corner.py`. Default is to plot last set of samples") @@ -1077,6 +1078,10 @@ sed 1d ./tmp_converted.dat {extra_shuffle_command} >> ./extrinsic_posterior_sam ## Consolidate job(s) # - consolidate output of single previous job con_arg_str = '' +if opts.use_EOB_parameters: + con_arg_str += " --a6c " +if opts.use_hyperbolic: + con_arg_str += " --hyperbolic " if opts.use_eccentricity: con_arg_str += " --eccentricity " if opts.use_meanPerAno: @@ -1102,6 +1107,10 @@ con_job.write_sub_file() ## Unify job # - update 'all.net' to include all previous events unify_arg_str = '' +if opts.use_EOB_parameters: + unify_arg_str += " --a6c " +if opts.use_hyperbolic: + unify_arg_str += " --hyperbolic " if opts.use_eccentricity: unify_arg_str += " --eccentricity " if opts.use_meanPerAno: @@ -1125,7 +1134,10 @@ unify_job.write_sub_file() # Evidence job exe_evidence = which('util_CIPDirSummarizeEvidence.py') -evidence_job, evidence_job_name = dag_utils.write_convert_sub(tag='evidence', exe=exe_evidence, log_dir='', file_input='', arg_str=" --cip-dir iteration_$(macroiteration)_cip --output evidence_$(macroiteration) ", universe=local_worker_universe, no_grid=no_worker_grid) +evidence_source = " --cip-dir iteration_$(macroiteration)_cip" +if opts.cip_explode_jobs is None: + evidence_source = " --cip-dir . --cip-prefix overlap-grid-$(macroiterationnext)" +evidence_job, evidence_job_name = dag_utils.write_convert_sub(tag='evidence', exe=exe_evidence, log_dir='', file_input='', arg_str=evidence_source+" --output evidence_$(macroiteration) ", universe=local_worker_universe, no_grid=no_worker_grid) evidence_job.add_condor_cmd("initialdir",opts.working_directory) evidence_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/evidence-$(cluster)-$(process).log") evidence_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_con/logs/evidence-$(cluster)-$(process).err") @@ -1437,8 +1449,6 @@ else: cip_args_extra +=" --n-eff {} --n-output-samples {} ".format(n_samples_per_job,n_samples_per_job) if not (opts.use_eccentricity_squared_sampling): cip_args_lines[indx] = cip_args_lines[indx].replace(' --parameter eccentricity_squared ',' --parameter-implied eccentricity_squared --parameter-nofit eccentricity ') - if not (opts.use_eccentricity_ln_sampling): - cip_args_lines[indx] = cip_args_lines[indx].replace(' --parameter eccentricity_ln ',' --parameter-implied eccentricity_ln --parameter-nofit eccentricity ') transfer_files_cip=['../all.net'] if opts.use_osg_cip and 'fit-method gp' in cip_args_base: transfer_files_cip += ['my_fit.pkl'] # transfer current working directory fit @@ -1862,6 +1872,7 @@ for it in np.arange(it_start,opts.n_iterations): # It is also *backward looking* evidence_node = pipeline.CondorDAGNode(evidence_job) evidence_node.add_macro("macroiteration",it-1) # LAST ITERATION being analyzed, not this one + evidence_node.add_macro("macroiterationnext",it) evidence_node.set_retry(opts.general_retries) if it > it_start: evidence_node.add_parent(parent_fit_node) # run evidence after previous iteration, on previous @@ -2017,6 +2028,10 @@ for it in np.arange(it_start,opts.n_iterations): cmd += " --neff-threshold {} ".format(opts.neff_threshold) cmd += " --general-request-disk {} --ile-request-disk {} --cip-request-disk {} ".format(opts.general_request_disk,opts.ile_request_disk,opts.cip_request_disk) cmd += " --request-memory-ILE {} --request-memory-CIP {} ".format(opts.request_memory_ILE,opts.request_memory_CIP) + if opts.use_EOB_parameters: + cmd += " --use-EOB-parameters " + if opts.use_hyperbolic: + cmd += " --use-hyperbolic " if opts.use_eccentricity: cmd += " --use-eccentricity " if opts.use_meanPerAno: diff --git a/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicMultiApproxIteration b/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicMultiApproxIteration index 5627e2227..a0b24eb8a 100755 --- a/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicMultiApproxIteration +++ b/MonteCarloMarginalizeCode/Code/bin/create_event_parameter_pipeline_BasicMultiApproxIteration @@ -65,6 +65,7 @@ import lalsimulation as lalsim import lal import functools import itertools +import re # Backend-neutral pipeline namespace (htcondor/glue/slurm) provided by dag_utils_generic from RIFT.misc.dag_utils_generic import pipeline @@ -175,7 +176,10 @@ def parse_ile_args_for_bw(my_arg_str): parser = argparse.ArgumentParser() parser.add_argument("--working-directory",default="./") -parser.add_argument("--approx",default=None,action='append',help="add multiple --approx options to analyze many approximants") +parser.add_argument("--approx",default=None,action='append',help="add multiple --approx options to analyze many approximants. Every approximant is evaluated on the SAME intrinsic grid, and their likelihoods are marginalized over at each grid point; see RIFT/misc/DESIGN_multiapprox_marginalization.md") +parser.add_argument("--approx-prior",default=None,action='append',help="APPROX=WEIGHT prior weight for one waveform model (repeatable). Default: uniform. These are the p(m) in L_marg(lambda) = sum_m p(m) L_m(lambda); they are NOT sampling weights.") +parser.add_argument("--approx-gwsignal",default=None,action='append',help="Waveform model that must be generated through gwsignal rather than lalsim (repeatable). The generator route is PER MODEL: the phenom family has no time-domain mode generator in gwsignal ('generator does not provide a method to generate time-domain modes'), while SEOBNRv5* is only available there. A single global --use-gwsignal therefore cannot serve an EOB-vs-phenom comparison, which is the comparison this builder exists for.") +parser.add_argument("--require-all-approx",action='store_true',help="Drop intrinsic points not successfully evaluated under EVERY approximant. Without it, such points are marginalized over whichever subset survived, which changes the estimator point by point.") parser.add_argument("--input-grid",default="overlap-grid.xml.gz") parser.add_argument("--cip-args",default=None,help="filename of args_cip.txt file which holds CIP arguments. Should NOT conflict with arguments auto-set by this DAG ... in particular, i/o arguments will be modified. ") parser.add_argument("--cip-args-list",default=None,help="filename of args_cip_list.file, which holds CIP arguments. Overrides cip-args if present. One CIP_n.sub file is created for each line in the file, which is used for an integer m iterations, where m is the first item of each line (normally 'X' in CIP)") @@ -187,6 +191,9 @@ parser.add_argument("--puff-cadence",default=None,type=int,help="Every n iterati parser.add_argument("--puff-max-it",default=-1,type=int,help="Maximum iteration number that puffball is applied. If negative, puffball is not applied ") parser.add_argument("--last-iteration-extrinsic",action='store_true',help="Configure last iteration to extract *one* set of extrinsic parameters from each intrinsic point. [This is highly inefficient, but people like having one extrinsic point per intrinsic point.] Requires --convert-args") parser.add_argument("--last-iteration-extrinsic-nsamples",default=3000,type=int,help="Construct this number of extrinsic samples") +parser.add_argument("--last-iteration-extrinsic-samples-per-ile",default=5,type=int,help="Fair-draw this many extrinsic posterior samples from each terminal ILE result") +parser.add_argument("--last-iteration-extrinsic-samples-per-ile-internal",default=10,type=int,help="Minimum terminal ILE n_eff used to support the per-result fair draw") +parser.add_argument("--last-iteration-extrinsic-time-resampling",action='store_true',help="Recover posterior times when the terminal ILE uses time marginalization") parser.add_argument("--last-iteration-export-marginal-distance-grid", action='store_true', help="Add argument to ILE_extr") parser.add_argument("--last-iteration-export-distance-slices", default=0, type=int, help="If >0, the ILE_extr (extrinsic) stage exports K-row .dslice files: Plan-B fixed-distance extrinsic-marginalized likelihoods. Adds --export-distance-slices K (+ --internal-use-lnL) to ILE_extr and strips --distance-marginalization.") parser.add_argument("--last-iteration-export-distance-slices-n-core", default=0, type=int, help="Passthrough to ILE --n-distance-slice-core for the extrinsic-stage .dslice export (0 = ILE default).") @@ -229,8 +236,64 @@ parser.add_argument('--neff-threshold',default=800,type=int,help="Number of samp parser.add_argument("--workflow",default='single',help="[single|fit+posterior|full] describes workflow layout used. 'Single' is a single node, running the fit and posterior for each iteration; 'full' produces many followup jobs to produce a reliable posterior") parser.add_argument("--n-post-jobs",default=1,type=int,help="Number of posterior jobs. Used in posterior and fit+posterior workflows") parser.add_argument("--use-bw-psd",action='store_true',help="Use BW PSD, attempting to use fiducial arguments as in LI for placement (i.e., signal at seglen -2). Assumes LI style data convention. Necessary BW will be parsed out of ile-args.txt (required)") +# --------------------------------------------------------------------------- +# API compatibility with util_RIFT_pseudo_pipe.py. +# +# pseudo_pipe emits ONE command line, shaped for BasicIteration's option surface, +# which is far wider than this builder's: this workflow implements cross-model +# marginalization, not every feature of the single-model pipeline. +# +# These options are ACCEPTED so a pseudo_pipe command line parses, and REFUSED +# below if one is actually used. They are never silently ignored -- a flag that +# parses and does nothing changes a run without saying so, and is invisible in a +# config diff. test_multiapprox_pseudo_pipe.py re-derives these lists from the +# two builders' argparse surfaces so they cannot drift. +_API_ONLY_FLAGS = ['--calibration-reweighting', '--calibration-reweighting-osg', '--calmarg-pilot', '--cip-explode-jobs-dag', '--cip-explode-jobs-subdag', '--comov-distance-reweighting', '--condor-containerize-nonworker', '--condor-local-nonworker-igwn-prefix', '--condor-nogrid-nonworker', '--extrinsic-handoff', '--first-iteration-jumpstart', '--frame-rotation', '--ile-group-subdag', '--ile-group-subdag-check-work', '--last-iteration-export-distance-slices-all-fresh', '--last-iteration-export-distance-slices-randomize', '--last-iteration-extrinsic-batched-convert', '--search-reflected-sky-mode', '--use-eccentricity', '--use-eccentricity-squared-sampling', '--use-full-submit-paths', '--use-hyperbolic', '--use-osg-cip', '--use-tabular-eos-file'] +_API_ONLY_VALUED = ['--bilby-ini-file', '--bilby-pickle-exe', '--bilby-pickle-file', '--cal-request-disk', '--calibration-reweighting-batchsize', '--calibration-reweighting-count', '--calibration-reweighting-exe', '--calibration-reweighting-extra-args', '--calibration-reweighting-initial-extra-args', '--calmarg-pilot-cadence', '--calmarg-pilot-max-it', '--calmarg-pilot-max-points', '--calmarg-pilot-top-fraction', '--cip-explode-jobs-last', '--cip-post-exe', '--cip-request-disk', '--comov-distance-reweighting-exe', '--convert-ascii2h5-exe', '--extrinsic-handoff-select', '--fetch-ext-grid-args', '--fetch-ext-grid-exe', '--general-request-disk', '--ile-condor-commands', '--ile-gpu-fanout', '--ile-n-events-to-analyze-first', '--ile-post-exe', '--ile-request-disk', '--last-iteration-export-distance-slices-wing-neff', '--last-iteration-export-distance-slices-wing-nmax', '--lisa-reference-time', '--n-eff', '--n-iterations-subdag-max', '--n-samples-per-job-threshold', '--reflected-sky-mode-exe', '--search-reflected-sky-mode-iteration', '--use-oauth-files'] +for _flag in _API_ONLY_FLAGS: + parser.add_argument(_flag, action='store_true', help=argparse.SUPPRESS) +for _flag in _API_ONLY_VALUED: + parser.add_argument(_flag, default=None, help=argparse.SUPPRESS) +# --------------------------------------------------------------------------- opts= parser.parse_args() +# Refuse the API-only options if any was actually requested (see above). +_requested = [f for f in _API_ONLY_FLAGS if getattr(opts, f[2:].replace('-', '_'), False)] +_requested += [f for f in _API_ONLY_VALUED if getattr(opts, f[2:].replace('-', '_'), None) is not None] +if _requested: + print(" create_event_parameter_pipeline_BasicMultiApproxIteration accepts the") + print(" following for API compatibility with util_RIFT_pseudo_pipe.py, but does") + print(" NOT implement them: " + ", ".join(sorted(_requested))) + print(" Refusing rather than ignoring: a flag that parses and does nothing") + print(" changes a run without saying so. Run those settings single-model, or") + print(" implement the option here.") + sys.exit(1) + +if not opts.approx or len(opts.approx) < 2: + parser.error("the multi-approximant builder requires at least two --approx values") +if len(set(opts.approx)) != len(opts.approx): + parser.error("duplicate --approx values are not allowed: {}".format(opts.approx)) +if opts.approx_prior: + _prior_labels = set() + _prior_total = 0.0 + for _item in opts.approx_prior: + if "=" not in _item: + parser.error("--approx-prior wants APPROX=WEIGHT, got {}".format(_item)) + _label, _value = [x.strip() for x in _item.split("=", 1)] + if _label in _prior_labels: + parser.error("duplicate --approx-prior for {}".format(_label)) + try: + _weight = float(_value) + except ValueError: + parser.error("--approx-prior weight is not numeric: {}".format(_item)) + if not np.isfinite(_weight) or _weight < 0: + parser.error("--approx-prior weights must be finite and non-negative") + _prior_labels.add(_label) + _prior_total += _weight + if _prior_total <= 0: + parser.error("--approx-prior weights must have positive total mass") + + local_worker_universe="vanilla" if opts.condor_local_nonworker: @@ -337,6 +400,17 @@ if opts.test_args: test_args = test_args.replace('[', ' \'[') test_args = test_args.replace(']', ']\'') test_args=test_args.rstrip() + # A convergence success is reported as exit status 1 so DAGMan can stop + # scheduling later intrinsic iterations. That is incompatible with this + # builder's fixed terminal fork: the per-model evidence and extrinsic + # products are deliberately created only after the configured iteration + # loop. util_RIFT_pseudo_pipe.py already adds --always-succeed whenever + # --add-extrinsic is requested; enforce the same contract here so direct + # builder callers cannot make the terminal branch unreachable. + if (opts.last_iteration_extrinsic and + re.search(r'(? 1) and (exe is None): +if ile_exe is None: + # This builder's terminal stage relies on batchmode's bounded fair-draw + # option even when a worker analyzes one intrinsic point. Falling back to + # integrate_likelihood_extrinsic for n_group==1 produces a command line + # containing --fairdraw-extrinsic-output-n-max, which that executable does + # not accept. ile_exe = dag_utils.which("integrate_likelihood_extrinsic_batchmode") output_file_names = None if opts.use_singularity or opts.use_osg: - transfer_file_names.append("../input-grid-$(macroiteration).xml.gz") + # ILE runs from approx__iteration__ile, while the shared grid is + # written in the workflow root as overlap-grid-N.xml.gz. `input-grid-N` + # is not produced anywhere in this workflow. This typo is invisible in a + # shared-filesystem pool but makes Condor reject the transfer on OSG. + transfer_file_names.append("../overlap-grid-$(macroiteration).xml.gz") #output_file_names = ','.join(["CME_out-$(macroevent)-$(cluster)-$(process).xml_{0}_.dat".format(x) for x in np.arange(opts.ile_n_events_to_analyze)]) #print "OUTPUT FILES ", output_file_names -ile_args += " --approx $(macroapprox) " -ile_job, ile_job_name = dag_utils.write_ILE_sub_simple(tag='ILE',log_dir=None,arg_str=ile_args,output_file="CME_out.xml",ncopies=opts.n_copies,exe=ile_exe,transfer_files=transfer_file_names,transfer_output_files=output_file_names,request_memory=opts.request_memory_ILE,request_gpu=opts.request_gpu_ILE,use_singularity=opts.use_singularity,singularity_image=singularity_image,use_osg=opts.use_osg,simple_osg_requirements=opts.use_osg_simple_requirements,frames_dir=opts.frames_dir,cache_file=opts.cache_file,use_cvmfs_frames=opts.use_cvmfs_frames,max_runtime_minutes=opts.ile_runtime_max_minutes) +# The generator route is per MODEL, not per run. A global --use-gwsignal (as +# util_RIFT_pseudo_pipe.py emits) sends every model through gwsignal, and the +# phenom family cannot be generated there at all -- its ILE jobs die with +# "ValueError: Invalid Argument" from gen_modes and contribute ZERO rows, so the +# run silently degrades to single-model. Strip it and bind it per node instead. +if opts.approx_gwsignal: + ile_args = ile_args.replace("--use-gwsignal", " ") +ile_args += " --approx $(macroapprox) $(macrogwsignal) " +# Puff ILE is the same per-model evaluation routed over a different grid. It +# must be derived after the global route is stripped and the model macros are +# appended; capturing ile_args_orig above silently baked the primary model and +# its generator route into every puff node. +ile_args_forpuff = ile_args.replace( + working_dir_inside + '/overlap-grid-$(macroiteration).xml.gz', + working_dir_inside + '/puffball-$(macroiteration).xml.gz') +ile_job, ile_job_name = dag_utils.write_ILE_sub_simple(tag='ILE',log_dir=None,arg_str=ile_args,output_file="CME_out.xml",ncopies=opts.n_copies,exe=ile_exe,transfer_files=transfer_file_names,transfer_output_files=output_file_names,request_memory=opts.request_memory_ILE,request_gpu=opts.request_gpu_ILE,use_singularity=opts.use_singularity,singularity_image=singularity_image,use_osg=opts.use_osg,use_simple_osg_requirements=opts.use_osg_simple_requirements,frames_dir=opts.frames_dir,cache_file=opts.cache_file,use_cvmfs_frames=opts.use_cvmfs_frames,max_runtime_minutes=opts.ile_runtime_max_minutes) # Modify: create macro for iteration # - added on a per-node basis # Modify: add macro argument for overlap grid to be used (kept in top-level directory) @@ -545,23 +661,78 @@ if not (opts.puff_args is None): transfer_file_names_puff.append(opts.working_directory+"/puffball-$(macroiteration).xml.gz") # could also use ../puffball, because relative is ok to initialdir # Write a version of the ILE job that uses puffball inputs ... IDENTICAL to code above for standard ILE case,just different argument for input # Yes, it would logically be simpler to make overlap-grid.xml.gz larger ... but I want to keep 'real samples' and 'puffed samples' seperate. - ilePuff_job, ilePuff_job_name = dag_utils.write_ILE_sub_simple(tag='ILE_puff',log_dir=None,arg_str=ile_args_forpuff,output_file="CME_out.xml",simple_unique=True,ncopies=opts.n_copies,exe=ile_exe,transfer_files=transfer_file_names_puff,request_memory=opts.request_memory_ILE,request_gpu=opts.request_gpu_ILE,use_singularity=opts.use_singularity,singularity_image=singularity_image,use_osg=opts.use_osg,simple_osg_requirements=opts.use_osg_simple_requirements,frames_dir=opts.frames_dir,cache_file=opts.cache_file,use_cvmfs_frames=opts.use_cvmfs_frames,max_runtime_minutes=opts.ile_runtime_max_minutes) + ilePuff_job, ilePuff_job_name = dag_utils.write_ILE_sub_simple(tag='ILE_puff',log_dir=None,arg_str=ile_args_forpuff,output_file="CME_out.xml",simple_unique=True,ncopies=opts.n_copies,exe=ile_exe,transfer_files=transfer_file_names_puff,request_memory=opts.request_memory_ILE,request_gpu=opts.request_gpu_ILE,use_singularity=opts.use_singularity,singularity_image=singularity_image,use_osg=opts.use_osg,use_simple_osg_requirements=opts.use_osg_simple_requirements,frames_dir=opts.frames_dir,cache_file=opts.cache_file,use_cvmfs_frames=opts.use_cvmfs_frames,max_runtime_minutes=opts.ile_runtime_max_minutes) ilePuff_job.add_condor_cmd("initialdir",opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile") ilePuff_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/ILE-$(macroevent)-$(cluster)-$(process).log") ilePuff_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/ILE-$(macroevent)-$(cluster)-$(process).err") - ilePuff_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_ile/logs/ILE-$(macroevent)-$(cluster)-$(process).out") + ilePuff_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/ILE-$(macroevent)-$(cluster)-$(process).out") ilePuff_job.write_sub_file() if (opts.last_iteration_extrinsic): - n_points_per_ILE = 5 - # ILE job with modified output format - # - note we *double* the memory request, because we need space to save samples - ile_args_extr = ile_args + " --save-P 0.01 --save-samples --n-eff " +str(2*n_points_per_ILE) # modify convergence criteria so output of reasonable size + n_points_per_ILE = opts.last_iteration_extrinsic_samples_per_ile + if n_points_per_ILE < 1: + parser.error("--last-iteration-extrinsic-samples-per-ile must be positive") + if opts.last_iteration_extrinsic_nsamples < 1: + parser.error("--last-iteration-extrinsic-nsamples must be positive") + n_extrinsic_batches = int(np.ceil( + opts.last_iteration_extrinsic_nsamples/ + (1.0*opts.ile_n_events_to_analyze))) + # The terminal work grid must cover every intrinsic event scheduled below. + # Round up to a complete batch when the requested count is not divisible by + # --ile-n-events-to-analyze; the final combiner applies the public output + # cap and uniqueness policy after these interior evaluations. + terminal_intrinsic_count = (n_extrinsic_batches * + opts.ile_n_events_to_analyze) + # Match the maintained RIFT terminal-stage contract: never lower the ILE + # convergence target supplied by the science configuration, and retain + # enough effective samples to support the requested fair draw. + n_eff_last = max(2*n_points_per_ILE, + opts.last_iteration_extrinsic_samples_per_ile_internal) + configured_neff = re.findall( + r'(?:^|\s)--n-eff(?:=|\s+)([0-9]+)(?=\s|$)', ile_args_orig) + if configured_neff: + n_eff_last = max(n_eff_last, int(configured_neff[-1])) + # ILE job with modified output format. The saved record MUST be a bounded + # fair draw from the extrinsic posterior, not the raw importance-sampling + # cache. A difficult point can reach n-max with n_eff~1; raw + # --save-samples then writes millions of rows and the downstream resampler + # must materialize all of them merely to select five. More importantly, + # that intermediate is not itself a posterior sample set. + # + # --save-samples is still the switch that serializes the record, while + # --fairdraw-extrinsic-output changes that record to an equal-weight draw + # before serialization and bounds it at n_points_per_ILE. + ile_args_extr = (ile_args + + " --save-P 0.01 --save-samples --fairdraw-extrinsic-output " + "--fairdraw-extrinsic-output-n-max {} --n-eff {} ".format( + n_points_per_ILE, n_eff_last)) + # A time-marginalized likelihood otherwise exports the fiducial time for + # every row. Recover posterior times on the same terminal fair-draw path. + _uses_time_marginalization = re.search( + r'(? 0: ile_args_extr += " --export-distance-slices {} ".format(opts.last_iteration_export_distance_slices) if opts.last_iteration_export_distance_slices_n_core: @@ -575,7 +746,25 @@ if (opts.last_iteration_extrinsic): if "--internal-use-lnL" not in ile_args_extr: ile_args_extr += " --internal-use-lnL " ile_args_extr = ile_args_extr.replace("--distance-marginalization ", ' ') - ileExtr_job, ileExtr_job_name = dag_utils.write_ILE_sub_simple(tag='ILE_extr',log_dir=None,arg_str=ile_args_extr,output_file="EXTR_out.xml",simple_unique=True,ncopies=1,exe=ile_exe,transfer_files=transfer_file_names,request_memory=opts.request_memory_ILE*2,request_gpu=opts.request_gpu_ILE,use_cvmfs_frames=opts.use_cvmfs_frames) + transfer_file_names_extr = list(transfer_file_names) + if opts.use_singularity or opts.use_osg: + # The terminal fork reads a model-specific CIP grid, not the shared + # in-loop grid. Condor expands these macros in transfer_input_files + # and stages the basename used by ile_args_extr in the remote sandbox. + transfer_file_names_extr[-1] = ( + "../approx_$(macroapprox)_overlap-grid-$(macroiteration).xml.gz") + ileExtr_job, ileExtr_job_name = dag_utils.write_ILE_sub_simple( + tag='ILE_extr', log_dir=None, arg_str=ile_args_extr, + output_file="EXTR_out.xml", simple_unique=True, ncopies=1, + exe=ile_exe, transfer_files=transfer_file_names_extr, + request_memory=opts.request_memory_ILE*2, + request_gpu=opts.request_gpu_ILE, + use_singularity=opts.use_singularity, + singularity_image=singularity_image, use_osg=opts.use_osg, + use_simple_osg_requirements=opts.use_osg_simple_requirements, + frames_dir=opts.frames_dir, cache_file=opts.cache_file, + use_cvmfs_frames=opts.use_cvmfs_frames, + max_runtime_minutes=opts.ile_runtime_max_minutes) ileExtr_job.add_condor_cmd("initialdir",opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile") ileExtr_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/ILEextr-$(macroevent)-$(cluster)-$(process).log") ileExtr_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/ILEextr-$(macroevent)-$(cluster)-$(process).err") @@ -592,16 +781,22 @@ if (opts.last_iteration_extrinsic): convertExtr_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/convert-$(macroevent)-$(macroindx).err") convertExtr_job.write_sub_file() - # Resample task - resample_args = ' --n-output-samples ' + str(n_points_per_ILE) # pick 5 random points from each ILE run - resample_job, resample_job_name = dag_utils.write_resample_sub('resample',log_dir=None,arg_str=resample_args,file_input=opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/EXTR_out-$(macroevent).xml_$(macroindx)_.dat",file_output=opts.working_directory+"/iteration_$(macroiteration)_ile/EXTR_out-$(macroevent).xml_$(macroindx)_.downsampled_dat",universe=local_worker_universe) - resample_job.add_condor_cmd("initialdir",opts.working_directory) - resample_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/resample-$(macroevent)-$(macroindx).log") - resample_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/resample-$(macroevent)-$(macroindx).err") - resample_job.write_sub_file() - - # Combination task at end -- probably should be a general utility - cat_job, cat_job_name = dag_utils.write_cat_sub(file_prefix='EXTR', file_postfix='.downsampled_dat.dat',file_output='extrinsic_posterior_samples_$(macroapprox).dat',universe=local_worker_universe) + # The converted rows are already equal-weight fair draws. Do not pass them + # through util_ResampleILEOutputWithExtrinsic: that would apply the + # importance weights a second time as well as reintroducing an unbounded + # full-table read. + cat_job, cat_job_name = dag_utils.write_cat_sub( + # Match only convert_extr's EXTR_out-..._.dat products. The narrower + # suffix also prevents stale *.downsampled_dat.dat files from an old + # DAG revision being silently concatenated after a rescue/rebuild. + file_prefix='EXTR', file_postfix='_.dat', + search_root=opts.working_directory+ + '/approx_$(macroapprox)_iteration_$(macroiteration)_ile', + file_output=opts.working_directory+ + '/extrinsic_posterior_samples_$(macroapprox).dat', + universe=local_worker_universe, + expected_batches=n_extrinsic_batches, + events_per_batch=opts.ile_n_events_to_analyze) cat_job.add_condor_cmd("initialdir",opts.working_directory) cat_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/cat-$(cluster)-$(process).log") cat_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_ile/logs/cat-$(cluster)-$(process).out") @@ -620,15 +815,62 @@ con_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration con_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_con/logs/con-$(cluster)-$(process).out") con_job.write_sub_file() -## Unify job -# - update 'all.net' to include all previous events -unify_job, unify_job_name = dag_utils.write_unify_sub_simple(tag='unify',log_dir='',arg_str='', base=opts.working_directory, target=opts.working_directory+'/approx_$(macroapprox)_all.net',universe=local_worker_universe) +## Unify jobs +# +# TWO nets, because the workflow needs two different things: +# +# all.net -- every model's composites pooled and +# MARGINALIZED over models at each shared +# intrinsic point. This is what the in-loop CIP +# fits, so grid refinement follows the +# model-marginalized posterior. +# approx__all.net -- one model's composites only. This feeds the +# terminal per-model CIP, which yields that +# model's posterior and evidence. +# +# The marginalization itself is util_CleanILE.py's job: it keys on the +# intrinsic parameters, so the same lambda evaluated under several models +# collapses to one entry, combined LINEARLY in L. See +# RIFT/misc/DESIGN_multiapprox_marginalization.md. +clean_model_args = " --model-group-regex 'approx_(.+?)_consolidated' " +# Coverage must be judged against the models this run CONFIGURED. A model whose +# composites are all empty is skipped before its label is recorded, so without +# this the pooled net silently becomes a lower-dimensional mixture and even +# --require-all-models sees nothing wrong. Observed live: one model's ILE jobs +# all failed and util_CleanILE reported "combination over 1 models" as ordinary +# status. +clean_model_args += " --expect-models '{}' ".format(",".join(opts.approx)) +if opts.approx_prior: + for item in opts.approx_prior: + if "=" not in item: + print(" --approx-prior wants APPROX=WEIGHT, got ", item); sys.exit(1) + label = item.split("=")[0].strip() + if label not in opts.approx: + print(" --approx-prior names {}, which is not among --approx {}".format(label, opts.approx)); sys.exit(1) + clean_model_args += " --model-prior '{}' ".format(item) + missing = [a for a in opts.approx if a not in [i.split("=")[0].strip() for i in opts.approx_prior]] + if missing: + print(" --approx-prior given but missing weights for ", missing, "; specify every approximant or none"); sys.exit(1) +if opts.require_all_approx: + clean_model_args += " --require-all-models " + +unify_job, unify_job_name = dag_utils.write_unify_sub_simple(tag='unify',log_dir='',arg_str=clean_model_args, base=opts.working_directory, target=opts.working_directory+'/all.net',universe=local_worker_universe,script_name='unify.sh',fail_on_error=True) unify_job.add_condor_cmd("initialdir",opts.working_directory) -unify_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_con/logs/unify-$(cluster)-$(process).log") -unify_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_con/logs/unify-$(cluster)-$(process).err") -unify_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_all.net") +unify_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/unify-$(cluster)-$(process).log") +unify_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/unify-$(cluster)-$(process).err") +unify_job.set_stdout_file(opts.working_directory+'/all.net') unify_job.write_sub_file() +## Per-model unify, for the terminal forked CIP. No --model-group-regex: +## within one model the evaluations ARE replicas of one quantity, so the flat +## ntot-weighted pool is the correct combination. +unify_model_job, unify_model_job_name = dag_utils.write_unify_sub_simple(tag='unify_model',log_dir='',arg_str='', base=opts.working_directory, target=opts.working_directory+'/approx_$(macroapprox)_all.net',universe=local_worker_universe,script_name='unify_model.sh',glob_pattern='approx_$(macroapprox)_*.composite',fail_on_error=True) +unify_model_job.add_condor_cmd("initialdir",opts.working_directory) +unify_model_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_con/logs/unifymodel-$(cluster)-$(process).log") +unify_model_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_con/logs/unifymodel-$(cluster)-$(process).err") +unify_model_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_all.net") +unify_model_job.write_sub_file() + @@ -642,41 +884,124 @@ if not(opts.cip_explode_jobs is None): else: cip_exe = "/bin/true" out_dir_base += "/iteration_$(macroiteration)_cip/" -cip_job, cip_job_name = dag_utils.write_CIP_sub(tag='CIP',log_dir=None,arg_str=cip_args_base,request_memory=opts.request_memory_CIP,input_net=opts.working_directory+'/approx_$(macroapprox)_all.net',output='approx_$(macroapprox)_overlap-grid-$(macroiterationnext)',out_dir=out_dir_base,exe=cip_exe,universe=local_worker_universe) +cip_job, cip_job_name = dag_utils.write_CIP_sub(tag='CIP',log_dir=None,arg_str=cip_args_base,request_memory=opts.request_memory_CIP,input_net=opts.working_directory+'/all.net',output='overlap-grid-$(macroiterationnext)',out_dir=out_dir_base,exe=cip_exe,universe=local_worker_universe) # Modify: set 'initialdir' -cip_job.add_condor_cmd("initialdir",opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip") +cip_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_cip") # Modify output argument: change logs and working directory to be subdirectory for the run -cip_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).log") -cip_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).err") -cip_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).out") +cip_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).log") +cip_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).err") +cip_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).out") cip_job.write_sub_file() if not(opts.cip_explode_jobs is None): print(" Exploding stage 1, with ",opts.cip_explode_jobs, " workers producing samples ") cip_args_base = cip_args_base.replace('fit-save-gp','fit-load-gp') cip_args_base = cip_args_base.replace('my_fit', 'my_fit.pkl') # yes, asymmetric arguments - cip_job_worker, cip_job_worker_name = dag_utils.write_CIP_sub(tag='CIP_worker',log_dir=None,arg_str=cip_args_base,request_memory=opts.request_memory_CIP,input_net=opts.working_directory+'/approx_$(macroapprox)_all.net',output='approx_$(macroapprox)_overlap-grid-$(macroiterationnext)-$(process)',out_dir=out_dir_base,exe=opts.cip_exe,ncopies=opts.cip_explode_jobs,universe=local_worker_universe) + cip_job_worker, cip_job_worker_name = dag_utils.write_CIP_sub(tag='CIP_worker',log_dir=None,arg_str=cip_args_base,request_memory=opts.request_memory_CIP,input_net=opts.working_directory+'/all.net',output='overlap-grid-$(macroiterationnext)-$(process)',out_dir=out_dir_base,exe=opts.cip_exe,ncopies=opts.cip_explode_jobs,universe=local_worker_universe) # Modify: set 'initialdir' - cip_job_worker.add_condor_cmd("initialdir",opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip") + cip_job_worker.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_cip") # Modify output argument: change logs and working directory to be subdirectory for the run - cip_job_worker.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).log") - cip_job_worker.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).err") - cip_job_worker.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).out") + cip_job_worker.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).log") + cip_job_worker.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).err") + cip_job_worker.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/cip-$(cluster)-$(process).out") cip_job_worker.write_sub_file() # Create worker join job - join_cip_job,join_cip_job_name = dag_utils.write_joingrids_sub(input_pattern=opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/output-grid-*.xml.gz",target_dir=opts.working_directory,output_base="approx_$(macroapprox)_overlap-grid-$(macroiterationnext)",n_explode=opts.cip_explode_jobs,log_dir=opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs",universe=local_worker_universe) - join_cip_job.add_condor_cmd("initialdir",opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip") + join_cip_job,join_cip_job_name = dag_utils.write_joingrids_sub(input_pattern=opts.working_directory+"/iteration_$(macroiteration)_cip/output-grid-*.xml.gz",target_dir=opts.working_directory,output_base="overlap-grid-$(macroiterationnext)",n_explode=opts.cip_explode_jobs,log_dir=opts.working_directory+"/iteration_$(macroiteration)_cip/logs",universe=local_worker_universe) + join_cip_job.add_condor_cmd("initialdir",opts.working_directory+"/iteration_$(macroiteration)_cip") # Modify output argument: change logs and working directory to be subdirectory for the run - join_cip_job.set_log_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).log") - join_cip_job.set_stderr_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).err") - join_cip_job.set_stdout_file(opts.working_directory+"/approx_$(macroapprox)_iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).out") + join_cip_job.set_log_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).log") + join_cip_job.set_stderr_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).err") + join_cip_job.set_stdout_file(opts.working_directory+"/iteration_$(macroiteration)_cip/logs/join-$(cluster)-$(process).out") join_cip_job.write_sub_file() +## Terminal per-model CIP. +# +# The iteration loop fits the MARGINALIZED net, because that is what should +# steer grid placement. The terminal stage instead FORKS: each model gets +# its own fit over its own net, yielding that model's posterior and -- via +# --fname-output-integral -- that model's evidence ln Z_m. Those are the +# deliverables, and they are what the final combination is weighted by. +# Combining at the end rather than the middle keeps each model's posterior +# inspectable, which is what a systematics study actually wants to look at. +# Which CIP arguments did the run actually FINISH on? +# +# `cip_args` is cip_args_lines[0], but the iteration loop advances through the +# list via cip_job_list, so with a multi-line --cip-args-list the last iteration +# uses a later entry. Building the terminal fit from cip_args would give the +# per-model posterior and evidence different coordinates, fit method or +# resolution than the iteration that produced the grid they are fitting -- and +# the evidences are what the final mixture is weighted by. +cip_args_terminal = cip_args +if cip_args_lines is not None: + _per_iteration = [] + for _i in np.arange(len(cip_args_lines)): + _per_iteration += int(cip_args_n[_i])*[cip_args_lines[_i]] + if _per_iteration: + _last = min(int(opts.n_iterations), len(_per_iteration)) - 1 + cip_args_terminal = _per_iteration[_last] + +cip_terminal_job = None +if opts.last_iteration_extrinsic: + # The iteration fit's posterior size is an adaptation setting, not the + # terminal work-grid contract. In particular pseudo_pipe commonly leaves + # `--n-output-samples 5000 --posterior-unique-draw` in the final CIP line, + # while a science terminal stage schedules 20000 intrinsic events. The + # high event indices then soft-exit successfully as out of range and the + # exact cat manifest can never be completed. Interior duplicate points are + # valid importance-work items; exact uniqueness belongs only to the final + # exported posterior (util_CombineApproximantPosteriors.py). + cip_args_terminal = re.sub( + r'(? per-model consolidate +# -> one marginalized net -> ONE fit -> one new grid, which every model's ILE +# reads next iteration. The models are coupled at the combine step, not run +# side by side: that coupling is the point of this builder. See +# RIFT/misc/DESIGN_multiapprox_marginalization.md. for it in np.arange(it_start,opts.n_iterations): - for approx in opts.approx: - consolidate_now = None - fit_node_now = None - ile_nodes_now =[] - # Create consolidate job + con_nodes_this_iteration = [] + for approx in opts.approx: + # Create consolidate job (per model: composites stay separated by model, + # because the terminal per-model CIP needs them apart) con_node = pipeline.CondorDAGNode(con_job) con_node.add_macro("macroiteration",it) con_node.add_macro("macroapprox",approx) con_node.set_retry(opts.general_retries) - # Create unify job - unify_node = pipeline.CondorDAGNode(unify_job) - unify_node.add_macro("macroiteration",it) - unify_node.add_macro("macroapprox",approx) - unify_node.add_parent(con_node) - unify_node.set_retry(opts.general_retries) - + # Create one node per job n_jobs_this_time = opts.n_samples_per_job if it ==it_start: n_jobs_this_time = n_initial - indx_max = int((1.0*n_jobs_this_time)/n_group) - if indx_max*n_jobs_this_time < n_group: - indx_max+=1 + indx_max = int(np.ceil((1.0*n_jobs_this_time)/n_group)) for event in np.arange(indx_max): #np.arange(n_jobs_this_time): # Add task per ILE operation ile_node = pipeline.CondorDAGNode(ile_job) @@ -885,6 +1207,7 @@ for it in np.arange(it_start,opts.n_iterations): ile_node.set_retry(opts.ile_retries) ile_node.add_macro("macroevent", event*n_group) ile_node.add_macro("macroapprox",approx) + ile_node.add_macro("macrogwsignal", gwsignal_macro(approx)) ile_node.add_macro("macroiteration", it) if not(parent_fit_node is None): ile_node.add_parent(parent_fit_node) @@ -902,6 +1225,7 @@ for it in np.arange(it_start,opts.n_iterations): ile_node.set_retry(opts.ile_retries) ile_node.add_macro("macroevent", event*n_group) ile_node.add_macro("macroapprox",approx) + ile_node.add_macro("macrogwsignal", gwsignal_macro(approx)) ile_node.add_macro("macroiteration", it) if not(parent_fit_node is None): ile_node.add_parent(parent_fit_node) @@ -910,16 +1234,26 @@ for it in np.arange(it_start,opts.n_iterations): # add con job dag.add_node(con_node) + con_nodes_this_iteration.append(con_node) + + # ---- everything below is per ITERATION, not per model -------------------- + # One unify: pools every model's composites and marginalizes over models at + # each shared intrinsic point. It must wait for ALL models to consolidate. + if True: + unify_node = pipeline.CondorDAGNode(unify_job) + unify_node.add_macro("macroiteration",it) + for node in con_nodes_this_iteration: + unify_node.add_parent(node) + unify_node.set_retry(opts.general_retries) dag.add_node(unify_node) - # Create fit node, which depends on consolidate node + # Create fit node, which depends on the marginalized net cip_job = cip_job_list[it] if isinstance(cip_job,list): cip_worker_job = cip_job[1] cip_job=cip_job[0] fit_node = pipeline.CondorDAGNode(cip_job) fit_node.add_macro("macroiteration", it) - fit_node.add_macro("macroapprox",approx) fit_node.add_macro("macroiterationnext", it+1) fit_node.set_category("CIP") fit_node.add_parent(unify_node) # only fit if we have results from the previous iteration @@ -933,7 +1267,6 @@ for it in np.arange(it_start,opts.n_iterations): worker_node =pipeline.CondorDAGNode(cip_worker_job) worker_node.add_macro("macroiteration", it) worker_node.add_macro("macroiterationnext", it+1) - worker_node.add_macro("macroapprox",approx) worker_node.set_category("CIP_worker") worker_node.add_parent(parent_fit_node) # only fit if we have results from the previous iteration worker_node.set_retry(opts.general_retries) @@ -942,7 +1275,6 @@ for it in np.arange(it_start,opts.n_iterations): join_node =pipeline.CondorDAGNode(join_cip_job) join_node.add_macro("macroiteration", it) join_node.add_macro("macroiterationnext", it+1) - join_node.add_macro("macroapprox",approx) join_node.set_category("join_cip") join_node.add_parent(worker_node) join_node.set_retry(opts.general_retries) @@ -973,7 +1305,6 @@ for it in np.arange(it_start,opts.n_iterations): convert_node=pipeline.CondorDAGNode(convert_job) convert_node.add_macro("macroiteration", it+1) # convert the NEWLY-PRODUCED iteration convert_node.add_macro("macroiterationlast", it) # use log files in the previous directory - convert_node.add_macro("macroapprox",approx) convert_node.add_parent(parent_fit_node) convert_node.set_category("CONVERT") convert_node.set_retry(opts.general_retries) @@ -988,65 +1319,99 @@ for it in np.arange(it_start,opts.n_iterations): test_node = pipeline.CondorDAGNode(test_job) test_node.add_macro("macroiteration", it+1) # test the NEWLY-PRODUCED iteration against the old test_node.add_macro("macroiterationlast", it) - test_node.add_macro("macroapprox",approx) test_node.add_parent(parent_fit_node) test_node.set_category("CONVERGE") dag.add_node(test_node) + test_node_list.append(test_node) parent_fit_node=test_node # Create export stages for extrinsic samples +# +# THE FORK. Up to here every model shared one grid and one marginalized fit. +# Here each model gets: its own net -> its own terminal CIP (posterior + ln Z) +# -> its own extrinsic ILE chain -> its own extrinsic posterior. A final +# combine node mixes them with weights p(m) Z_m. Keeping the per-model +# posteriors on disk is deliberate: a systematics study wants to look at them. if opts.last_iteration_extrinsic: - # Check if 'it' is defined : it will not always be, if done later - if not ('it' in globals()): - it = opts.n_iterations # last iteration - - # Create nodes for followup tasks - cat_node = pipeline.CondorDAGNode(cat_job) - - # Perform final ILE run on all points, saving samples - # Need to perform number of events CONSISTENT WITH TARGET SAMPLE SIZE - # - *not* always same as number of ILE events being analyzed - # - *assumes* grid files have sufficiently large numbers of samples to allow this! (as in many other cases) - n_jobs_extrinsic = int(opts.last_iteration_extrinsic_nsamples/(1.0*n_group)) - for event in np.arange(n_jobs_extrinsic): - # Add task per ILE operation - ile_node = pipeline.CondorDAGNode(ileExtr_job) -# ile_node.set_priority(JOB_PRIORITIES["ILE"]) - ile_node.set_retry(opts.ile_retries) - ile_node.add_macro("macroevent", event*n_group) - ile_node.add_macro("macroiteration", it) - if not(parent_fit_node is None): - ile_node.add_parent(parent_fit_node) - dag.add_node(ile_node) + # The terminal CIP reads the grid the loop finished on and writes the + # per-model posterior at the same index the extrinsic stage reads. + it = opts.n_iterations + + n_jobs_extrinsic = n_extrinsic_batches + cat_nodes = [] + loop_final_node = parent_fit_node + + for approx in opts.approx: + # Per-model net: this model's composites only, pooled flat (replicas of + # ONE model), with no cross-model marginalization. + unify_model_node = pipeline.CondorDAGNode(unify_model_job) + unify_model_node.add_macro("macroiteration", it-1) + unify_model_node.add_macro("macroapprox", approx) + unify_model_node.set_retry(opts.general_retries) + if not(loop_final_node is None): + unify_model_node.add_parent(loop_final_node) + dag.add_node(unify_model_node) + + # Per-model terminal fit: posterior AND evidence. + cipterm_node = pipeline.CondorDAGNode(cip_terminal_job) + cipterm_node.add_macro("macroiteration", it-1) + cipterm_node.add_macro("macroiterationnext", it) + cipterm_node.add_macro("macroapprox", approx) + cipterm_node.set_category("CIP") + cipterm_node.add_parent(unify_model_node) + cipterm_node.set_retry(opts.general_retries) + dag.add_node(cipterm_node) + + # Create nodes for followup tasks + cat_node = pipeline.CondorDAGNode(cat_job) + + # Perform final ILE run on all points, saving bounded fair draws. + # Need to perform number of events CONSISTENT WITH TARGET SAMPLE SIZE + # - *not* always same as number of ILE events being analyzed + # - *assumes* grid files have sufficiently large numbers of samples to allow this! (as in many other cases) + for event in np.arange(n_jobs_extrinsic): + # Add task per ILE operation + ile_node = pipeline.CondorDAGNode(ileExtr_job) + ile_node.set_retry(opts.ile_retries) + ile_node.add_macro("macroevent", event*n_group) + ile_node.add_macro("macroiteration", it) + ile_node.add_macro("macroapprox", approx) + ile_node.add_macro("macrogwsignal", gwsignal_macro(approx)) + ile_node.add_parent(cipterm_node) + dag.add_node(ile_node) - # Add convert and resample task *for each output file* - for indx in np.arange(n_group): - convert_node = pipeline.CondorDAGNode(convertExtr_job) - convert_node.add_macro("macroevent", event*n_group) - convert_node.add_macro("macroiteration", it) - convert_node.add_macro("macroindx",indx) - convert_node.set_retry(opts.ile_retries) # this can fail too - convert_node.add_parent(ile_node) - - resample_node = pipeline.CondorDAGNode(resample_job) - resample_node.add_macro("macroevent", event*n_group) - resample_node.add_macro("macroiteration", it) - resample_node.add_macro("macroindx",indx) - resample_node.set_retry(opts.ile_retries) # these occasionally fail for stupid reasons - nodes missing software, etc - resample_node.add_parent(convert_node) - - # Make cat job - cat_node.add_parent(resample_node) - cat_node.set_retry(opts.ile_retries) # this can fail too - cat_node.add_macro("macroiteration", it) # needed to identify log file location - - # Add nodes - dag.add_node(convert_node) - dag.add_node(resample_node) - - dag.add_node(cat_node) + # Convert each already-fair-drawn output file. + for indx in np.arange(n_group): + convert_node = pipeline.CondorDAGNode(convertExtr_job) + convert_node.add_macro("macroevent", event*n_group) + convert_node.add_macro("macroiteration", it) + convert_node.add_macro("macroindx",indx) + convert_node.add_macro("macroapprox", approx) + convert_node.set_retry(opts.ile_retries) # this can fail too + convert_node.add_parent(ile_node) + + # Make cat job + cat_node.add_parent(convert_node) + cat_node.set_retry(opts.ile_retries) # this can fail too + cat_node.add_macro("macroiteration", it) # needed to identify log file location + cat_node.add_macro("macroapprox", approx) + + # Add nodes + dag.add_node(convert_node) + + dag.add_node(cat_node) + cat_nodes.append(cat_node) + + # Recombine: mixture over models weighted by p(m) Z_m. + combine_node = pipeline.CondorDAGNode(combine_job) + combine_node.add_macro("macroiteration", it) + combine_node.set_retry(opts.general_retries) + for node in cat_nodes: + combine_node.add_parent(node) + dag.add_node(combine_node) + parent_fit_node = combine_node # Create final node for overall plots. (Note: default setup is designed to enable plots of the last two iterations *at each step* but this seems like overkill) if plot_args: @@ -1061,3 +1426,11 @@ if plot_args: dag_name="marginalize_intrinsic_parameters_BasicIterationWorkflow" dag.set_dag_file(dag_name) dag.write_concrete_dag() + +# convergence_test_samples.py returns 1 when convergence is established. +# Match BasicIteration: without a terminal extraction this stops the remaining +# intrinsic loop and reports a successful DAG, rather than leaving a failed +# convergence node that merely blocks all of its children. +if opts.test_args: + for node in test_node_list: + dag.add_abort_on(node, exit_code=1, return_value=0) diff --git a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py index c1597947d..8f03a5685 100755 --- a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py +++ b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py @@ -18,6 +18,7 @@ import lal import RIFT.lalsimutils as lalsimutils +from RIFT.physics import teobresums_compat import lalsimulation as lalsim from igwn_ligolw import lsctables, table, utils @@ -30,7 +31,13 @@ from RIFT.misc.dag_utils_generic import which # leaf module: numpy only, so this does not drag numba/cupy into the helper from RIFT.likelihood.time_interp_choice import ( - BARE_FLAG_SENTINEL, CROSSOVER_GUIDANCE, resolve_interpolate_time_request) + BARE_FLAG_SENTINEL, CROSSOVER_GUIDANCE, is_off_request, + resolve_interpolate_time_request) +# Same leaf-module reasoning, and IMPORTED rather than re-typed: a second hand-written +# copy of the choice tuple is how a typo becomes a silently different likelihood. +from RIFT.likelihood.time_marginalization_quadrature import ( + TIME_QUADRATURE_CHOICES, validate_time_quadrature, + refuse_unless_time_quadrature_emitted) lalapps_path2cache = which('lal_path2cache') ligolw_add = 'igwn_ligolw_add' if not(which(ligolw_add)): @@ -165,6 +172,7 @@ def get_observing_run(t): parser.add_argument("--limit-mc-range",default=None,type=str,help="For PP plots, or other analyses requiring a specific mc range (eg ini file), bounding the limit *above*. Allows the code to auto-select its mc range as usual, then takes the intersection with this limit") parser.add_argument("--scale-mc-range",type=float,default=None,help="If using the auto-selected mc, scale the ms range proposed by a constant factor. Recommend > 1. . ini file assignment will override this.") parser.add_argument("--force-eta-range",default=None,type=str,help="For PP plots. Enforces initial grid placement inside this region") +parser.add_argument("--force-mtot-range",default=None,type=str,help="For PP plots, hyperbolic analysis, or other analyses requiring a specific mtot range (eg ini file). Enforces initial grid placement inside this region. Passed directly to MOG and CIP.") parser.add_argument("--allow-subsolar", action='store_true', help="Override limits which otherwise prevent subsolar mass PE") parser.add_argument("--use-legacy-gracedb",action='store_true') parser.add_argument("--event-time",type=float,default=None) @@ -207,7 +215,14 @@ def get_observing_run(t): parser.add_argument("--assume-matter-but-primary-bh",action='store_true',help="If present, the code will add options necessary to manage tidal arguments for the smaller body ONLY. (Usually pointless)") parser.add_argument("--internal-tabular-eos-file",type=str,default=None,help="Tabular file of EOS to use. The default prior will be UNIFORM in this table!. NOT YET IMPLEMENTED (initial grids, etc)") parser.add_argument("--assume-eccentric",action='store_true',help="If present, the code will add options necessary to manage eccentric arguments. The proposed fit strategy and initial grid will allow for eccentricity") +parser.add_argument("--assume-hyperbolic",action='store_true',help="If present, the code will add options necessary to manage eccentric arguments. The proposed fit strategy and initial grid will allow for hyperbolic") +parser.add_argument("--E0-max", default=1.2,type=float,help="Maximum range of 'E0' allowed.") +parser.add_argument("--E0-min", default=1.0,type=float,help="Minimum range of 'E0' allowed.") +parser.add_argument("--pphi0-max", default=10.0,type=float,help="Maximum range of 'p_phi0' allowed.") +parser.add_argument("--pphi0-min", default=0.0,type=float,help="Minimum range of 'p_phi0' allowed.") +parser.add_argument("--use-mtot-coords",action='store_true',help="Configures CIP and PUFF for mtot instead of mc. REQUIRES --force-mtot-range.") parser.add_argument("--use-meanPerAno",action='store_true',help="The proposed fit strategy and initial grid will allow for meanPerAno") +parser.add_argument("--use-EOB-parameters",action='store_true',help="The proposed fit strategy and initial grid will allow for EOB parameters: currently only a6c") parser.add_argument("--assume-nospin",action='store_true',help="If present, the code will not add options to manage precessing spins (the default is aligned spin)") parser.add_argument("--assume-precessing-spin",action='store_true',help="If present, the code will add options to manage precessing spins (the default is aligned spin)") parser.add_argument("--assume-volumetric-spin",action='store_true',help="If present, the code will assume a volumetric spin prior in its last iterations. If *not* present, the code will adopt a uniform magnitude spin prior in its last iterations. If not present, generally more iterations are taken.") @@ -221,7 +236,8 @@ def get_observing_run(t): parser.add_argument("--internal-ile-auto-logarithm-offset",action='store_true',help="Passthrough to ILE") parser.add_argument("--internal-ile-use-lnL",action='store_true',help="Passthrough to ILE. Will DISABLE auto-logarithm-offset and manual-logarithm-offset") parser.add_argument("--internal-ile-n-chunk",default=None,type=int,help="Override the extrinsic chunk size (--n-chunk) passed to ILE. Default: 40000, scaled linearly with SNR above 40 and capped at 160000. Rationale: at high SNR the posterior is a vanishing fraction of the prior volume, so a small chunk gives few informative samples per adaptation step; measured collapse on a truth-known SNR ladder falls 88%%->50%% (SNR160) and 69%%->25%% (SNR80) going 1e4->1.6e5, and the gain survives at fixed budget. Larger chunks cost GPU memory, so raise the ILE memory request if you raise this a lot.") -parser.add_argument("--internal-ile-interpolate-time",nargs='?',const=BARE_FLAG_SENTINEL,default=None,type=str,help="Evaluate Q_lm at FRACTIONAL detector times instead of snapping to the nearest sample bin, in the maintained NoLoop likelihood (needs --time-marginalization --vectorized and one of --gpu/--rotation-slow/--freqresponse; the driver REFUSES rather than ignores otherwise). REQUIRES AN EXPLICIT STENCIL: nearest|cubic|sinc -- automatic selection was removed as measurably unreliable, and a bare flag is rejected rather than silently doing nothing. MEASURED GUIDANCE (SEOBNRv4, an IMR model): %s. 'nearest' is never competitive and is already unusable at O4 SNRs. Error grows as SNR^2, so this matters more at 3G. Cost: sinc is ~4.2-4.5x cubic on CPU, ~1.6-3.0x on GPU. Full tables, limitations and provenance: RIFT/likelihood/DESIGN_q_window_stencil.md. Default off." % CROSSOVER_GUIDANCE) +parser.add_argument("--internal-ile-interpolate-time",nargs='?',const=BARE_FLAG_SENTINEL,default=None,type=str,help="Evaluate Q_lm at FRACTIONAL detector times instead of snapping to the nearest sample bin, in the maintained NoLoop likelihood (needs --time-marginalization --vectorized and one of --gpu/--rotation-slow/--freqresponse; the driver REFUSES rather than ignores otherwise). REQUIRES AN EXPLICIT STENCIL: nearest|cubic|sinc -- automatic selection was removed as measurably unreliable, and a bare flag is rejected rather than silently doing nothing. MEASURED GUIDANCE (SEOBNRv4, an IMR model): %s. 'nearest' is never competitive and is already unusable at O4 SNRs. Error grows as SNR^2, so this matters more at 3G. Cost: sinc is ~4.2-4.5x cubic on CPU, ~1.6-3.0x on GPU. Full tables, limitations and provenance: RIFT/likelihood/DESIGN_q_window_stencil.md. Default: emit nothing, so ILE uses its own default, which CHANGED 2026-09-02 from 'nearest' to time_interp_choice.TIME_INTERP_DEFAULT. To pin the historical behaviour pass 'nearest' (or an off-request such as 'False', which this helper now re-expresses as an explicit '--interpolate-time nearest' so that 'off' still means off)." % CROSSOVER_GUIDANCE) +parser.add_argument("--internal-ile-time-marginalization-quadrature",default=None,type=str,choices=list(TIME_QUADRATURE_CHOICES),help="Rule for the TIME integral of the marginalized likelihood: %s. Default None = emit nothing, so ILE keeps its own default ('simpson', the historical fixed-deltaT Simpson rule) and args_ile.txt is byte-identical to today. 'bandlimited' resolves the INTEGRAND rather than the data: exp(lnL(t)) is a peak of width sigma_t = 1/(2 pi rho sigma_f), which shrinks as 1/rho, while deltaT=1/srate is fixed -- so production under-resolves its own integrand, worse at higher SNR (measured: scanning the grid phase moves the reported lnL by 1.649 nats at srate 4096, rho=40). Emitted as --time-marginalization-quadrature on the ILE command line, so a completed run's quadrature is readable off the .sub file. Requires --time-marginalization --vectorized --gpu and excludes --rotation-slow / --freqresponse / calibration marginalization; this helper REFUSES rather than emitting an inert flag. INI OVERRIDE: the RIFT ini parser overrides the command line for non-boolean options, so never set this string option in an ini that a Makefile also sets. Rationale and measured tables: RIFT/likelihood/DESIGN_time_marginalization_quadrature.md." % ("|".join(TIME_QUADRATURE_CHOICES),)) parser.add_argument("--internal-cip-use-lnL",action='store_true') parser.add_argument("--ile-n-eff",default=50,type=int,help="Target n_eff passed to ILE. Try to keep above 2") parser.add_argument("--test-convergence",action='store_true',help="If present, the code will terminate if the convergence test passes. WARNING: if you are using a low-dimensional model the code may terminate during the low-dimensional model!") @@ -230,6 +246,7 @@ def get_observing_run(t): parser.add_argument("--internal-test-convergence-js-lame-fixed-thresholds",action='store_true',help="With --internal-test-convergence-method js_lame: do NOT pass --js-lame-auto-threshold, so the fixed --threshold/--js-threshold/--quantile-tolerance values are used as given. Only correct if you know the distinct-sample supply matches the sample size those fixed values were tuned at; otherwise the components fire on pure sampling noise and the gate never converges.") parser.add_argument("--internal-test-convergence-js-lame-allow-missing-lags",action='store_true',help="With --internal-test-convergence-method js_lame: do NOT pass --js-lame-require-lags, so the test may report convergence before its lag window is populated. This restores the one-step behaviour js_lame exists to replace (an unusually clean early comparison can stop the loop at sub-iteration 2-3); for diagnostics only.") parser.add_argument("--lowlatency-propose-approximant",action='store_true', help="If present, based on the object masses, propose an approximant. Typically TaylorF2 for mc < 6, and SEOBNRv4_ROM for mc > 6.") +parser.add_argument("--internal-initial-grid-approximant",default=None,type=str,help="Approximant used to choose model-specific initial-grid safety seeds. Normally supplied by util_RIFT_pseudo_pipe.py; an [engine] approx in --use-ini is the fallback.") parser.add_argument("--online", action='store_true', help="Use online settings") parser.add_argument("--propose-initial-grid",action='store_true',help="If present, the code will either write an initial grid file or (optionally) add arguments to the workflow so the grid is created by the workflow. The proposed grid is designed for ground-based LIGO/Virgo/Kagra-scale instruments") parser.add_argument("--propose-initial-grid-fisher",action='store_true',help="If present, overrides propose-initial-grid. Uses the SEMIANALYTIC fisher matrix to propose an initial grid: very fast, well targeted.") @@ -261,13 +278,48 @@ def get_observing_run(t): parser.add_argument("--use-cvmfs-frames",action='store_true',help="If true, require LIGO frames are present (usually via CVMFS). User is responsible for generating cache file compatible with it. This option insures that the cache file is properly transferred (because you have generated it)") parser.add_argument("--use-ini",default=None,type=str,help="Attempt to parse LI ini file to set corresponding options. WARNING: MAY OVERRIDE SOME OTHER COMMAND-LINE OPTIONS") parser.add_argument("--verbose",action='store_true') +parser.add_argument("--force-scatter-grids",action='store_true',help="Eliminates all non-scatter intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument("--force-plunge-grids",action='store_true',help="Eliminates all non-plunge intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument("--force-zoomwhirl-grids",action='store_true',help="Eliminates all non-zoomwhirl intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument('--force-hyperbolic-22', action='store_true', help='Forces just the 22 modes for hyperbolic waveforms') opts= parser.parse_args() # Resolve the sub-sample stencil request IMMEDIATELY, so a bare flag / retired 'True' / typo # fails here rather than after a whole workflow has been built and submitted. Returns None when # the feature is off; a canonical stencil name otherwise. time_interp_choice = resolve_interpolate_time_request(opts.internal_ile_interpolate_time) - +# "OFF" MUST STILL MEAN OFF once the ILE default is no longer 'nearest'. +# resolve_interpolate_time_request collapses two different things to None: "flag absent" and an +# explicit off-request ('False'/'off'/'none', OFF_REQUEST_TOKENS). Both used to emit nothing, +# and emitting nothing used to mean 'nearest' -- so they were the same answer. Since the ILE +# default became TIME_INTERP_DEFAULT (2026-09-02) they are opposites: emitting nothing now means +# the new default, so "--internal-ile-interpolate-time False" would have turned interpolation ON. +# An explicit off-request is therefore re-expressed as the explicit stencil that means off. +if (time_interp_choice is None and opts.internal_ile_interpolate_time is not None + and is_off_request(opts.internal_ile_interpolate_time)): + time_interp_choice = 'nearest' + +# Same, for the time quadrature: argparse `choices` already rejects a typo, but validate through +# the LIBRARY function too so this helper and the ILE driver can never disagree about the legal +# set. None means "emit nothing", which is the byte-identical default path. +time_quadrature_choice = opts.internal_ile_time_marginalization_quadrature +if time_quadrature_choice is not None: + validate_time_quadrature(time_quadrature_choice) + +# Ensure --assume-hyperbolic is set when using any --force-X-grids option +# Ensure only ONE of the --force-X-grids options is set +force_grids = [opts.force_scatter_grids, opts.force_plunge_grids, opts.force_zoomwhirl_grids] +if any(force_grids) and not opts.assume_hyperbolic: + parser.error("Using --force-scatter-grids, --force-plunge-grids, or --force-zoomwhirl-grids requires --assume-hyperbolic!") + +if sum(bool(x) for x in force_grids) > 1: + parser.error("CANNOT use multiple --force-X-grids options at the same time!") + +if opts.use_mtot_coords: + if opts.force_mtot_range is None: + print('Using the mtot coords requires a specified range!') + print('Specify the mtot range with --force-mtot-range!') + sys.exit(1) if opts.assume_matter_but_primary_bh: opts.assume_matter=True @@ -478,6 +530,9 @@ def get_observing_run(t): event_dict["s2z"] = P.s2z event_dict["P"] = P event_dict["epoch"] = 0 # no estimate for now + if opts.assume_hyperbolic: + event_dict["E0"] = P.E0 + event_dict["p_phi0"] = P.p_phi0 elif opts.use_coinc: # If using a coinc through injections and not a GraceDB event. # Same code as used before for gracedb coinc_file = opts.use_coinc @@ -1007,6 +1062,13 @@ def crit_m2(delta): mc_min = mc_center -0.5*opts.scale_mc_range*mc_width mc_max = mc_center +0.5*opts.scale_mc_range*mc_width +# EOB parameter range. Must exist before any initial-grid path uses it: an +# [engine] a6c_min entry is optional, and --use-EOB-parameters is normally +# passed without one. Default matches the CIP/puffball a6c limits. +a6c_min = -80 +a6c_max = -20 +a6c_range_str = " [{},{}]".format(a6c_min,a6c_max) + if use_ini: engine_dict = dict(config['engine']) if 'chirpmass-min' in engine_dict: @@ -1021,8 +1083,17 @@ def crit_m2(delta): eta_max = q_max/(1.+q_max)**2 if eta_max >=0.25: eta_max = 0.24999999 # rounding/finite-precision issues may cause nan problems + if 'a6c_min' in engine_dict: + a6c_range_str = " ["+str(engine_dict['a6c_min'])+","+str(engine_dict['a6c_max'])+"]" if 'ecc_min' in engine_dict: - ecc_range_str = " ["+str(engine_dict['ecc_min'])+","+str(engine_dict['ecc_max'])+"]" + # ConfigParser returns strings, so these bounds must be parsed before any + # arithmetic. The two initial-grid paths sample DIFFERENT coordinates: + # eccentricity_squared needs [ecc_min^2,ecc_max^2], plain eccentricity + # needs [ecc_min,ecc_max]. Keep both ranges, and do not mix them. + ecc_min = float(engine_dict['ecc_min']) + ecc_max = float(engine_dict['ecc_max']) + ecc_range_str = " [{},{}]".format(ecc_min,ecc_max) + ecc_squared_range_str = " [{},{}]".format(ecc_min**2,ecc_max**2) if 'meanPerAno_min' in engine_dict: meanPerAno_range_str = " ["+str(engine_dict['meanPerAno_min'])+","+str(engine_dict['meanPerAno_max'])+"]" @@ -1053,7 +1124,9 @@ def crit_m2(delta): eta_range_str_cip = " --eta-range ["+str(eta_min) +","+str(eta_max)+"]" # default will include 1, as we work with BBHs if not (opts.force_eta_range is None): eta_range_str_cip = " --eta-range " + opts.force_eta_range - +if not (opts.force_mtot_range is None): + mtot_range_str = opts.force_mtot_range + mtot_range_str_cip = " --mtot-range " + opts.force_mtot_range ### ### Write arguments @@ -1155,6 +1228,32 @@ def crit_m2(delta): # the reverse pairing is safe. Pair this pipeline with an ILE from the same checkout. helper_ile_args += " --interpolate-time " + time_interp_choice + " " +if time_quadrature_choice is not None: + # Validated at parse time, so by here it is one of TIME_QUADRATURE_CHOICES. The name goes on + # the ILE command line verbatim, so a completed run's quadrature is readable off the .sub file. + # + # VERSION SKEW: an ILE predating this option rejects the unknown flag outright (optparse errors + # on an unrecognised option), so an old ILE driven by this helper FAILS LOUDLY rather than + # silently running Simpson. That is the safe direction; still, pair this pipeline with an ILE + # from the same checkout. + print(" ==> Time-marginalization quadrature: '{}' (emitted as --time-marginalization-quadrature; " + "the ILE driver refuses rather than ignores if its configuration cannot honour it)".format( + time_quadrature_choice)) + if time_quadrature_choice != 'simpson': + # Sub-sample integration is also an export contract. The extrinsic stage + # uses the same reflected, derived-resolution reconstruction to draw a + # continuous conditional time, rather than calling coarse-grid + # return_lnLt=True and throwing the new resolution away. + print(" NOTE: on the extrinsic/fairdraw stage band-limited quadrature " + "also mandates a continuous draw from p(t | extrinsic, data), using " + "the same validated dense reconstruction as the integral. Export " + "is not quantised at 1/srate or at a configurable finer lattice.") + # rstrip() so the separator does not depend on whatever the PREVIOUS append left + # behind: the flag gluing onto its neighbour would produce an args_ile.txt in which + # the quadrature is not a token at all, and the emission guard below is what would + # then have to catch it. Make it structurally impossible instead. + helper_ile_args = helper_ile_args.rstrip() + " --time-marginalization-quadrature " + time_quadrature_choice + " " + if opts.internal_ile_auto_logarithm_offset and not opts.internal_ile_use_lnL: helper_ile_args += " --auto-logarithm-offset " rescaled_base_ile = True @@ -1239,6 +1338,12 @@ def crit_m2(delta): if opts.psd_assume_common_window: helper_ile_args += " --psd-window-shape {} ".format(window_shape) +initial_grid_approximant = opts.internal_initial_grid_approximant +if initial_grid_approximant is None and use_ini and config.has_option('engine', 'approx'): + initial_grid_approximant = config.get('engine', 'approx').strip().strip('"\'') +if initial_grid_approximant is None: + initial_grid_approximant = approx_str + if not(internal_dmax is None): helper_ile_args += " --d-max " + str(int(internal_dmax)) if dmin != 1: # if not default value, add argument @@ -1280,15 +1385,27 @@ def crit_m2(delta): elif opts.data_LI_seglen: seglen = opts.data_LI_seglen - # Use LI-style positioning of trigger relative to 2s before end of buffer - # Use LI-style tukey windowing - window_shape = opts.data_tukey_window_time*2/seglen - data_end_time = event_dict["tref"]+2 - data_start_time = event_dict["tref"] +2 - seglen + if opts.assume_hyperbolic: + window_shape = 0.0 # DO NOT window at all + # split seglen across the event time + data_start_time = event_dict["tref"] - seglen/2 + data_end_time = event_dict["tref"] + seglen/2 + else: + # Use LI-style positioning of trigger relative to 2s before end of buffer + # Use LI-style tukey windowing + window_shape = opts.data_tukey_window_time*2/seglen + data_end_time = event_dict["tref"]+2 + data_start_time = event_dict["tref"] +2 - seglen helper_ile_args += " --data-start-time " + str(data_start_time) + " --data-end-time " + str(data_end_time) + " --inv-spec-trunc-time 0 --window-shape " + str(window_shape) if opts.psd_assume_common_window: helper_ile_args += " --psd-window-shape {} ".format(window_shape) +if opts.use_EOB_parameters: + helper_ile_args += " --save-EOB-parameters " +if opts.assume_hyperbolic: + helper_ile_args += " --save-hyperbolic " + if opts.force_hyperbolic_22: + helper_ile_args += " --force-hyperbolic-22 " if opts.assume_eccentric: helper_ile_args += " --save-eccentricity " if opts.use_meanPerAno: @@ -1321,8 +1438,12 @@ def crit_m2(delta): cmd += " --random-parameter chieff_aligned --random-parameter-range " + chieff_range grid_size =2500 + if opts.use_EOB_parameters: + cmd += " --random-parameter a6c --random-parameter-range " + a6c_range_str + if opts.assume_hyperbolic: + cmd += " --random-parameter E0 --random-parameter-range [{},{}] --random-parameter p_phi0 --random-parameter-range [{},{}] ".format(opts.E0_min,opts.E0_max,opts.pphi0_min,opts.pphi0_max) if opts.assume_eccentric: - cmd += " --random-parameter eccentricity --random-parameter-range " + ecc_range_str + cmd += " --random-parameter eccentricity_squared --random-parameter-range " + ecc_squared_range_str grid_size = int(grid_size*1.5) if opts.use_meanPerAno: cmd += " --random-parameter meanPerAno --random-parameter-range [0,6.2831]" @@ -1343,13 +1464,18 @@ def crit_m2(delta): # add basic mass parameters cmd = "util_ManualOverlapGrid.py --fname proposed-grid --skip-overlap " mass_string_init = " --random-parameter mc --random-parameter-range " + mc_range_str + " --random-parameter delta_mc --random-parameter-range '[" + str(delta_grid_min) +"," + str(delta_grid_max) + "]' " + if not(opts.force_mtot_range is None): + mass_string_init = " --random-parameter mtot --random-parameter-range " + mtot_range_str + " --random-parameter delta_mc --random-parameter-range '[" + str(delta_grid_min) +"," + str(delta_grid_max) + "]' " cmd+= mass_string_init # Add standard downselects : do not have m1, m2 be less than 1 if not(opts.force_mc_range is None): # force downselect based on this range cmd += " --downselect-parameter mc --downselect-parameter-range " + opts.force_mc_range if not(opts.force_eta_range is None): - cmd += " --downselect-parameter eta --downselect-parameter-range " + opts.force_eta_range + cmd += " --downselect-parameter eta --downselect-parameter-range " + opts.force_eta_range + if not(opts.force_mtot_range is None): + # force downselect based on this range + cmd += " --downselect-parameter mtot --downselect-parameter-range " + opts.force_mtot_range cmd += " --fmin " + str(opts.fmin_template) if opts.data_LI_seglen and not (opts.no_enforce_duration_bound): cmd += " --enforce-duration-bound " + str(opts.data_LI_seglen) @@ -1378,8 +1504,26 @@ def crit_m2(delta): grid_size =2500 if opts.assume_precessing_spin: - # Handle problems with SEOBNRv3 failing for aligned binaries -- add small amount of misalignment in the initial grid - cmd += " --parameter s1x --parameter-range [0.00001,0.00003] " + # Keep a nonzero seed for exactly-aligned precessing models, which + # can otherwise fall back to a different aligned implementation. + # TEOBResumS-DALI is special: its C backend calls sum(chi_perp) + # <= 1e-4 aligned, and its GWSignal layer's disagreement with that + # decision can segfault. Seed ResumS a decade above the boundary. + seed_min, seed_max = teobresums_compat.initial_transverse_spin_range( + initial_grid_approximant + ) + cmd += " --parameter s1x --parameter-range [{},{}] ".format(seed_min, seed_max) + if opts.use_EOB_parameters: + cmd += " --random-parameter a6c --random-parameter-range " + a6c_range_str + grid_size = int(grid_size*1.5) + if opts.assume_hyperbolic: + cmd += " --random-parameter E0 --random-parameter-range [{},{}] --random-parameter p_phi0 --random-parameter-range [{},{}] ".format(opts.E0_min,opts.E0_max,opts.pphi0_min,opts.pphi0_max) + if opts.force_scatter_grids: + cmd += " --force-scatter " + if opts.force_plunge_grids: + cmd += " --force-plunge " + if opts.force_zoomwhirl_grids: + cmd += " --force-zoomwhirl " if opts.assume_eccentric: cmd += " --random-parameter eccentricity --random-parameter-range " + ecc_range_str grid_size = int(grid_size*1.5) @@ -1535,10 +1679,26 @@ def lambda_m_estimate(m): puff_max_it=0 if event_dict["MChirp"] >25: - # at high mass, mc/eta correlation weak, don't want to have eta coordinate degeneracy at q=1 to reduce puff proposals near there - helper_puff_args = " --parameter mc --parameter delta_mc --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) + if opts.use_mtot_coords: + helper_puff_args = " --parameter mtot --parameter delta_mc --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) + else: + # at high mass, mc/eta correlation weak, don't want to have eta coordinate degeneracy at q=1 to reduce puff proposals near there + helper_puff_args = " --parameter mc --parameter delta_mc --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) else: - helper_puff_args = " --parameter mc --parameter eta --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) + if opts.use_mtot_coords: + helper_puff_args = " --parameter mtot --parameter q --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) + else: + helper_puff_args = " --parameter mc --parameter eta --fmin {} --fref {} ".format(opts.fmin_template,opts.fmin_template) +if opts.use_EOB_parameters: + helper_puff_args += " --parameter a6c " +if opts.assume_hyperbolic: + helper_puff_args += " --parameter E0 --parameter p_phi0 " + if opts.force_scatter_grids: + helper_puff_args += " --force-scatter " + if opts.force_plunge_grids: + helper_puff_args += " --force-plunge " + if opts.force_zoomwhirl_grids: + helper_puff_args += " --force-zoomwhirl " if opts.assume_eccentric: helper_puff_args += " --parameter eccentricity " if opts.use_meanPerAno: @@ -1557,7 +1717,10 @@ def lambda_m_estimate(m): if 'gp' in fit_method: helper_cip_args += " --cap-points 12000 " if not opts.no_propose_limits: - helper_cip_args += mc_range_str_cip + eta_range_str_cip + if not(opts.use_mtot_coords): + helper_cip_args += mc_range_str_cip + eta_range_str_cip + else: + helper_cip_args += mtot_range_str_cip + eta_range_str_cip if opts.force_chi_max: helper_cip_args += " --chi-max {} ".format(opts.force_chi_max) if opts.force_chi_small_max: @@ -1788,6 +1951,19 @@ def lambda_m_estimate(m): # helper_cip_arg_list[indx] += " --lnL-offset 20 " # enforce lnL cutoff past the first iteration. Focuses fit on high-likelihood points as in O1/O2 +# REFUSE AT WORKFLOW-BUILD TIME, not at first-job time. The ILE driver already refuses this +# combination, but that costs a whole queue-slot cycle to discover; the arguments are all in hand +# here. The check reads the ASSEMBLED string rather than the opts, because whether +# --time-marginalization/--vectorized/--gpu are present depends on the strategy branches above +# (--propose-ile-convergence-options), not on any single flag. util_RIFT_pseudo_pipe.py repeats +# it over the FINAL args_ile.txt, which is the only place calibration marginalization and +# --manual-extra-ile-args are visible. +# Checks the BYTES about to be written, not the parsed option, so it also catches the +# emission being dropped or duplicated -- not only an unhonourable configuration. The +# raise lives in the library function so that it is executable in a unit test. +refuse_unless_time_quadrature_emitted( + time_quadrature_choice, helper_ile_args, "helper_ile_args.txt") + # editing ILE args based on strategy above, so only writing now with open("helper_ile_args.txt",'w') as f: f.write(helper_ile_args) @@ -1863,6 +2039,14 @@ def lambda_m_estimate(m): with open("helper_convert_args.txt", 'a') as f: f.write(" --export-eos ") +if opts.use_EOB_parameters: + with open("helper_convert_args.txt",'a') as f: + f.write(" --export-EOB-parameters ") + +if opts.assume_hyperbolic: + with open("helper_convert_args.txt",'a') as f: + f.write(" --export-hyperbolic ") + if opts.assume_eccentric: with open("helper_convert_args.txt",'a') as f: f.write(" --export-eccentricity ") diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index 84b9c9205..c6db29b7f 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -50,6 +50,7 @@ import glue.lal import RIFT.lalsimutils as lalsimutils from RIFT.likelihood.time_interp_choice import CROSSOVER_GUIDANCE as _CROSSOVER_GUIDANCE +from RIFT.likelihood.time_interp_choice import TIME_INTERP_DEFAULT from RIFT.precision import RiftFloat import RIFT.integrators.mcsampler as mcsampler from RIFT.integrators.rvs_record import (RvsRecord as _RvsRecord, # see DESIGN_rvs_naming.md @@ -60,6 +61,7 @@ from RIFT.integrators.rvs_record import (RvsRecord as _RvsRecord, # see DESIGN # cupy when mcsamplerGPU.draw_simplified() calls it with a device array), so one definition serves # every sampler. from RIFT.integrators.mcsampler import (clip_angle_limits, cosine_sampler_limits, + distance_limit_range, distance_sampler_kwargs, ret_dec_samp_vector, ret_dec_samp_cdf_inv_vector, ret_cos_samp_vector, ret_cos_samp_cdf_inv_vector) import RIFT.misc.sky_rotations as sky_rotations @@ -244,7 +246,9 @@ optp.add_option("--psd-window-shape", type=float, default=0, help="Shape of Tuke optp.add_option("-m", "--time-marginalization", action="store_true", help="Perform marginalization over time via direct numerical integration. Default is false.") #optp.add_option("--n-fairdraw-extrinsic-samples",default=None,type=int,help="Extracts a concrete number of fair draw extrinsic samples, bounded above by n_eff") optp.add_option("--resample-time-marginalization",action='store_true', help="If time-marginalizaiton is true (and should almost always be true), at the end export step use resampling. REQUIRES using fairdraw-extrinsic-output") -optp.add_option("--srate-resample-time-marginalization",type=int, default=None, help="If time-marginalizaiton is true (and should almost always be true), an opportunity to interpolate the likelihoods on the (internal) time grid and perform time resampling at a higher sampling rate, for greater nominal time resolution. Does not change underlying calculations.") +optp.add_option("--srate-resample-time-marginalization",type=int, default=None, help="For --time-posterior-export grid under the historical Simpson quadrature, interpolate lnL(t) onto a lattice at this rate before drawing. Band-limited quadrature derives its own resolution and mandates a continuous draw, so the combination is refused.") +optp.add_option("--time-posterior-export", type="choice", choices=["auto", "continuous", "grid"], default="auto", + help="How --resample-time-marginalization exports geocenter time. auto (default) draws continuously from the interpolated lnL(t) posterior when the active likelihood exposes a faithful arbitrary-time evaluator, otherwise preserves the grid; continuous requests an off-grid draw and is refused on unsupported likelihood paths; grid is the explicit legacy compatibility mode. --time-marginalization-quadrature bandlimited always resolves to continuous and refuses grid because sub-sample integration carries a sub-sample export contract.") optp.add_option("-d", "--distance-marginalization", action="store_true", help="Perform marginalization over distance via a look-up table. Default is false.") optp.add_option("-l", "--distance-marginalization-lookup-table", default=None, help="Look-up table for distance marginalization.") optp.add_option("--calibration-envelope-directory",default=None, help="Name of directory") @@ -296,6 +300,9 @@ optp.add_option("--internal-waveform-extra-lalsuite-args",type=str,default=None) optp.add_option("--internal-waveform-extra-kwargs",type=str,default=None) optp.add_option("--internal-precompute-ignore-threshold",default=None,type=float) optp.add_option("--verbose",action='store_true') +optp.add_option("--save-EOB-parameters", action="store_true") +optp.add_option("--save-hyperbolic", action="store_true") +optp.add_option('--force-hyperbolic-22', default=False, action='store_true', help='Forces just the 22 modes for hyperbolic waveforms') optp.add_option("--save-eccentricity", action="store_true") optp.add_option("--save-meanPerAno", action="store_true") # @@ -326,7 +333,8 @@ integration_params.add_option("--internal-gmm-adaptive-components",action='store integration_params.add_option("--internal-gmm-max-components",type=int,default=8,help="Cap on the per-group component count for --internal-gmm-adaptive-components (default 8).") integration_params.add_option("--internal-gmm-defensive-frac",type=float,default=0.0,help="Weight of the broad box-covering 'defensive' mixture component added to each adaptive GMM group (default 0 = OFF). Intended to bound the importance weights (Hesterberg defensive IS), but on a wide extrinsic prior the broad component draws physically-extreme points where the likelihood is NaN and it did not improve n_eff on the SNR~82 benchmark -- opt-in only.") integration_params.add_option("--internal-gmm-inflate",type=float,default=1.0,help="Covariance inflation factor (std multiplier) applied to each adaptive GMM component (default 1.0 = none). A value >1 widens the proposal relative to the elite cloud it was fit to; complements --internal-gmm-defensive-frac.") -integration_params.add_option("--interpolate-time", default=False,help="Sub-sample stencil for evaluating Q_lm at fractional detector times, instead of snapping to the nearest sample bin. Accepts 'nearest', 'cubic', 'sinc', or a legacy truthy value (True/1/yes) meaning 'cubic'. WHICH TO USE is set by the bandwidth of Q(t), which is NOT fmax -- Q is band-limited by whichever is lower, fmax or the template's own cutoff, so it depends on the MASSES and on FMIN. MEASURED with SEOBNRv4 (an IMR model): %s. fmin matters as much as mass -- cubic degrades from fmin 20 to 150 at fixed mass (endpoint ratios 6.5x at M=9 and 9.6x at M=20, and NOT monotone in between) while sinc stays flat, which is why the crossover rises. NEAREST is never competitive (200-443 nats) and reaches 1 nat of error by SNR 2-6. Do not trust inspiral-only (TaylorT4) numbers for this: no merger-ringdown, understates the band by 2-3.7x. Error grows as SNR^2. COST of sinc vs cubic: ~4.2-4.5x on CPU, ~1.6-3.0x on GPU. All three stencils have CPU and GPU implementations. Requires the maintained NoLoop likelihood: this is REFUSED, not ignored, if the configuration cannot honour it. Measured tables and limitations: RIFT/likelihood/DESIGN_q_window_stencil.md. (Default=false, i.e. nearest)" % _CROSSOVER_GUIDANCE) +integration_params.add_option("--interpolate-time", default=None,help="Sub-sample stencil for evaluating Q_lm at fractional detector times, instead of snapping to the nearest sample bin. Accepts 'nearest', 'cubic', 'sinc', or a legacy truthy value (True/1/yes) meaning 'cubic'. DEFAULT CHANGED 2026-09-02 from 'nearest' to %r (issue #233); the value is time_interp_choice.TIME_INTERP_DEFAULT, shared with the jax driver's --interp so the two cannot ship opposite defaults again. THIS CHANGES RESULTS for anyone who did not pass --interpolate-time; pass '--interpolate-time nearest' to reproduce a pre-2026-09-02 run. WHICH TO USE is set by the bandwidth of Q(t), which is NOT fmax -- Q is band-limited by whichever is lower, fmax or the template's own cutoff, so it depends on the MASSES and on FMIN. MEASURED with SEOBNRv4 (an IMR model): %s. fmin matters as much as mass -- cubic degrades from fmin 20 to 150 at fixed mass (endpoint ratios 6.5x at M=9 and 9.6x at M=20, and NOT monotone in between) while sinc stays flat, which is why the crossover rises. NEAREST is never competitive (200-443 nats) and reaches 1 nat of error by SNR 2-6. Do not trust inspiral-only (TaylorT4) numbers for this: no merger-ringdown, understates the band by 2-3.7x. Error grows as SNR^2. COST of sinc vs cubic: ~4.2-4.5x on CPU, ~1.6-3.0x on GPU. All three stencils have CPU and GPU implementations. Requires the maintained NoLoop likelihood: an EXPLICIT request is REFUSED, not ignored, if the configuration cannot honour it, while the DEFAULT falls back to 'nearest' with a printed reason rather than turning a working configuration into a startup error. Measured tables and limitations: RIFT/likelihood/DESIGN_q_window_stencil.md. (Default=%s)" % (TIME_INTERP_DEFAULT, _CROSSOVER_GUIDANCE, TIME_INTERP_DEFAULT)) +integration_params.add_option("--time-marginalization-quadrature", default="simpson", type=str, help="Rule for the TIME integral of the marginalized likelihood: 'simpson' (default, historical), 'bandlimited', or 'peak-local'. 'simpson' integrates exp(lnL(t)) with Simpson's rule at the FIXED spacing deltaT=1/srate. That spacing is a property of the DATA; the integrand's width is a property of the SIGNAL -- after angle marginalization exp(lnL(t)) is a near-Gaussian peak of width sigma_t = 1/(2 pi rho sigma_f) -- so resolving it needs srate >~ 2 pi sigma_f rho, a requirement that GROWS LINEARLY WITH SNR and that production does not meet. MEASURED on a 35+30 Msun SEOBNRv4 H1L1V1 injection at rho=40 (sigma_t = 61.2 us): rigidly scanning the grid phase over 2*deltaT moves the reported lnL by 1.649 / 0.385 / 0.0095 nats at srate 4096 / 8192 / 16384. Simpson makes an under-resolved peak WORSE than trapezoid, not better: (4T_h - T_2h)/3 carries the coarser T_2h and inherits its 2h alias. 'bandlimited' costs no extra likelihood evaluations and no extra precompute: kappa(t) is band-limited below Nyquist by construction and rho_sq is time-independent on this path, so the samples already computed determine the continuous integrand exactly, and one zero-padded FFT per row recovers it. Against a converged dense reference at srate 4096, rho=40: -0.007 nats, versus +0.745 for Simpson at the same grid phase. THERE IS DELIBERATELY NO RESOLUTION OPTION: the refinement factor is derived from the measured peak width and re-asserted on the refined grid. Cost scales with that factor and is paid only where the integrand actually demands it (a well-resolved peak derives a factor of 1 and costs nothing). Requires --time-marginalization --vectorized --gpu (--force-xpy is accepted), excludes --rotation-slow / --freqresponse / calibration marginalization, and is REFUSED, not ignored, if the configuration cannot honour it. Rationale, measured tables and exclusions: RIFT/likelihood/time_marginalization_quadrature.py. 'peak-local' is the same argument with the refined grid placed only where the integrand has support, because the dense rule refines the WHOLE window to a peak whose width shrinks as 1/rho -- it works hardest exactly where the peak occupies least of the domain. kappa's extrema are ENUMERATED on a small, SNR-INDEPENDENT upsample (kappa is band-limited at Nyquist, so enumerating it is not a function of SNR); an interval of a few sigma_t is built around each; overlapping intervals are MERGED into disjoint ones (without which the shared region is double-counted, measured +1.6 nats at rho~6); and each merged interval is integrated at its own derived spacing. The mass left OUTSIDE the intervals is bounded per row and CHECKED, so the truncation is not an assumption -- a row whose bound is not small enough, or whose local grid would cost more than the dense one, is given the 'bandlimited' value rather than an approximation with a caveat. Accuracy is that of 'bandlimited' by construction and is measured against it (max 1.9e-11 nats over 4000 extrinsic rows). COST: measured through this code path on CPU at n_extrinsic 4000, npts 614, it is NOT the prototype's headline figure -- that was measured with an analytic kappa in hand, where evaluating the interpolant at an arbitrary time was free, and here it is not. See RIFT/likelihood/DESIGN_time_marginalization_peak_local.md for the measured table. Same prerequisites and same exclusions as 'bandlimited', PLUS: 'peak-local' REFUSES phase marginalization. That is a deliberate scope cut -- production marginalizes over distance, not phase, and under phase marginalization the time peak's Laplace width picks up an (I1/I0)(|kappa|/D) factor that does not reduce, so the local spacing is no longer derivable from rho_sq and the curvature alone. 'bandlimited' still supports it. (Default=simpson)") integration_params.add_option("--d-prior",default='Euclidean' ,type=str,help="Distance prior for dL. Options are dL^2 (Euclidean), 'pseudo_cosmo', and 'cosmo' and 'cosmo_sourceframe' .") integration_params.add_option("--d-prior-redshift", action='store_true', help="If true, distance prior is computed in redshift. This option MAY be enforced for 'cosmo' sampling") integration_params.add_option("--d-max", default=10000,type=float,help="Maximum distance in volume integral. Used to SET THE PRIOR; changing this value changes the numerical answer.") @@ -337,6 +345,7 @@ integration_params.add_option("--limit-right-ascension",default=None,help="Restr integration_params.add_option("--limit-declination",default=None,help="Restrict declination sampling AND prior to 'LO,HI' [rad]. Always given in radians of DECLINATION: with --declination-cosine-sampler the box is transformed internally to the sampled coordinate sin(dec). Not compatible with --internal-sky-network-coordinates.") integration_params.add_option("--limit-inclination",default=None,help="Restrict inclination sampling AND prior to 'LO,HI' [rad]. Always given in radians of INCLINATION: with --inclination-cosine-sampler the box is transformed internally to the sampled coordinate cos(iota), which reverses the limit order.") integration_params.add_option("--limit-psi",default=None,help="Restrict polarization psi sampling AND prior to 'LO,HI' [rad].") +integration_params.add_option("--limit-distance",default=None,help="Restrict distance SAMPLING to 'LO,HI' [Mpc], WITHOUT changing the prior or its normalization. Unlike --d-min/--d-max (which SET the prior and therefore change the numerical answer) and unlike the angular --limit-* boxes (which narrow the prior SUPPORT, so lnZ drops by the prior mass outside), this is a change of SAMPLING prior only: the distance prior keeps the normalization it has over the full [--d-min,--d-max], so the reported lnZ needs no correction and is directly comparable to a full-range run and between samplers. Intended for high amplitude, where the distance posterior narrows as 1/rho and a box tracking it restores the resolution the quadrature was wasting -- keep the box comfortably wider than the posterior, because likelihood OUTSIDE it is simply not integrated. WHAT 'no correction' MEANS IN PRACTICE, measured end to end rather than argued (real data, S250114ax, rho ~ 82, AV, 39 runs): the evidence the box actually TRUNCATES is bounded by the posterior mass outside it, 0.003 nats for a box holding all but 0.3 per cent of the draws -- but the lnZ difference you will OBSERVE against a full-range run is larger and of the opposite sign, +0.37 +- 0.11 nats, because it is the FULL-RANGE run's own sampling bias, which the box removes (with --no-adapt-distance, where that bias is unmistakable, the full-range run loses 4.16 nats and the box recovers 3.81). So: comparable to the sampler's own systematic, not to machine precision, and the narrowed run is the more accurate of the two. THE BENEFIT IS A THRESHOLD IN AMPLITUDE, NOT A GENERAL IMPROVEMENT: repeating the same measurement on S240920dw at rho 41.4 gives +0.13 +- 0.08 nats, consistent with zero, because there the full-range sampler is already healthy (n_ESS 594 vs 101 at rho 82, Pareto k-hat 0.26 vs 0.55, nothing collapsed) and the box has no bias to remove -- it even costs a little. Use it where the full-range run is in trouble; below that it is neutral to slightly negative. Evidence: RIFT_roboto_paper analyses/limit_distance_e2e/. Must lie inside [--d-min,--d-max]. REFUSED (not ignored) with --distance-marginalization (no distance sampler exists: the marginal is an analytic integral over [--d-min,--d-max]), with --d-prior-redshift (the sampled coordinate is redshift, not Mpc) and with --internal-reparam-dl-incl (the sampled axis is D_eff, not d_L).") integration_params.add_option("--internal-rotate-phase", action='store_true',help="If specified, the integration sampler uses phase_p ==phi+psi and phase_m == phi-psi as sampling coordinates, both ranging from 0 to 4 pi. The prior is twice as large.") integration_params.add_option("--internal-sky-network-coordinates",action='store_true',help="If specified, perform integration in sky coordinates aligned with the first two IFOs provided") integration_params.add_option("--internal-sky-network-coordinates-raw",action='store_true',help="If specified, does not attempt to organize IFO network sensibly, uses them AS PROVIDED IN ORDER.") @@ -462,22 +471,35 @@ def _truthy_option(value): _TI_LEGACY_BOOLEAN = ("1", "true", "t", "yes", "y", "on", "0", "false", "f", "no", "n", "off", "none") -_ti_raw = str(opts.interpolate_time).strip().lower() -if _ti_raw in ("nearest", "cubic", "sinc"): - # explicit stencil name - opts._noloop_time_interp = _ti_raw -elif _ti_raw in _TI_LEGACY_BOOLEAN: - # legacy boolean: truthy meant cubic - opts._noloop_time_interp = "cubic" if _truthy_option(opts.interpolate_time) else "nearest" +# WAS THE STENCIL ASKED FOR, OR INHERITED? Every guard below distinguishes the two, so this +# must be decided on IDENTITY (`is None`), before any string coercion: str(None) is 'none', +# which is a legal explicit spelling meaning 'nearest', so a string test cannot tell an +# omitted flag from "--interpolate-time none" and would silently disarm the explicit path. +opts._interp_time_from_default = opts.interpolate_time is None +if opts._interp_time_from_default: + # DEFAULT CHANGED 2026-09-02: 'nearest' -> time_interp_choice.TIME_INTERP_DEFAULT (issue + # #233). PROVISIONAL until the guards further down have run -- a default may be downgraded + # back to 'nearest' where an explicit request would be refused. + opts._noloop_time_interp = TIME_INTERP_DEFAULT else: - # Anything else is a typo, and it must NOT be absorbed. Before this check a misspelled - # stencil ('sinK', 'lanczos') was simply non-truthy and so ran 'nearest' -- a silent change - # of the likelihood's time discretization, invisible in the log and indistinguishable from a - # run that never asked for interpolation at all. Now that the helper writes a resolved - # stencil NAME onto every --interpolate-time command line, a typo there has to be loud. - raise ValueError( - "--interpolate-time: unrecognised value %r. Use a stencil name (nearest|cubic|sinc) or " - "a legacy boolean (%s)." % (opts.interpolate_time, "|".join(_TI_LEGACY_BOOLEAN))) + _ti_raw = str(opts.interpolate_time).strip().lower() + if _ti_raw in ("nearest", "cubic", "sinc"): + # explicit stencil name + opts._noloop_time_interp = _ti_raw + elif _ti_raw in _TI_LEGACY_BOOLEAN: + # legacy boolean: truthy meant cubic + opts._noloop_time_interp = "cubic" if _truthy_option(opts.interpolate_time) else "nearest" + else: + # Anything else is a typo, and it must NOT be absorbed. Before this check a misspelled + # stencil ('sinK', 'lanczos') was simply non-truthy and so ran 'nearest' -- a silent + # change of the likelihood's time discretization, invisible in the log and + # indistinguishable from a run that never asked for interpolation at all. Now that the + # helper writes a resolved stencil NAME onto every --interpolate-time command line, a + # typo there has to be loud. + raise ValueError( + "--interpolate-time: unrecognised value %r. Use a stencil name (nearest|cubic|sinc) " + "or a legacy boolean (%s)." % (opts.interpolate_time, + "|".join(_TI_LEGACY_BOOLEAN))) # The LEGACY scalar path (FactoredLogLikelihoodTimeMarginalized) takes a plain boolean and has # nothing to do with the NoLoop stencils. It used to be handed opts.interpolate_time raw, which # was fine while that was only ever truthy/falsy -- but 'nearest' is a non-empty string, so once @@ -485,6 +507,22 @@ else: # legacy path's interpolation ON while meaning the exact opposite in NoLoop. Derive an honest # boolean instead: only the two genuinely-interpolating stencils count as "interpolate". opts._legacy_interpolate_time = opts._noloop_time_interp in ("cubic", "sinc") +from RIFT.likelihood.time_posterior import resolve_time_posterior_export_mode +# THE DEFAULT STENCIL MUST NOT SILENTLY CHANGE THE EXPORT. resolve_time_posterior_export_mode +# maps `auto` (the --time-posterior-export default) to 'continuous' for ANY stencil other than +# 'nearest', so once the stencil default stopped being 'nearest' the same one-line change would +# also have flipped the fair-draw time export of every --resample-time-marginalization run. That +# is not a free relabelling: continuous export re-evaluates the whole likelihood on a >=4x denser +# time grid and can raise MemoryError from validate_time_posterior_working_set, so it would have +# turned working runs into failing ones and changed t_ref in the output of the rest. +# +# The export therefore keys on an EXPLICIT stencil only. Asking for a stencil still opts you into +# the better export, and '--time-posterior-export continuous' still works on its own; inheriting +# the default gets the historical 'grid' export, bit-for-bit. +_ti_for_export = 'nearest' if opts._interp_time_from_default else opts._noloop_time_interp +opts._time_posterior_export = resolve_time_posterior_export_mode( + opts.time_posterior_export, _ti_for_export, + continuous_available=not (opts.rotation_slow or opts.freqresponse)) # NOTE: deliberately NOT announcing the stencil here. opts.gpu is not resolved yet at this # point, so we cannot yet tell whether the stencil will actually be used -- and a banner that # names a stencil the run then ignores is worse than no banner, because it reads as proof. @@ -496,8 +534,10 @@ if opts.rotation_slow: # template on GPU (same memory footprint as the baseline). if not opts.vectorized: raise ValueError("--rotation-slow requires --vectorized") - if opts.gpu and getattr(opts, 'calibration_marginalization', False): - raise ValueError("--rotation-slow on GPU does not support calibration/glitch marginalization (n_cal>1)") + # The calibration/glitch-marginalization exclusion used to be a guard here, keyed on an + # option that does not exist, so it never fired. It now lives in the central gate below + # (RIFT/calmarg/option_compat.py), keyed on --calibration-envelope-directory and run once + # opts.gpu is final. History: DESIGN_calmarg_in_loop.md, "Option-compatibility gate". if getattr(opts, 'distance_marginalization', False) or getattr(opts, 'phase_marginalization', False): raise ValueError("--rotation-slow does not yet support distance/phase marginalization") @@ -509,8 +549,8 @@ if opts.freqresponse: raise ValueError("--freqresponse and --rotation-slow both replace the likelihood; use at most one") if not opts.vectorized: raise ValueError("--freqresponse requires --vectorized") - if opts.gpu and getattr(opts, 'calibration_marginalization', False): - raise ValueError("--freqresponse on GPU does not support calibration/glitch marginalization (n_cal>1)") + # Same never-firing calibration guard as --rotation-slow above, with the same history; + # superseded by the central gate below. See RIFT/calmarg/option_compat.py. if getattr(opts, 'distance_marginalization', False) or getattr(opts, 'phase_marginalization', False): raise ValueError("--freqresponse does not yet support distance/phase marginalization") @@ -572,6 +612,10 @@ if opts.resample_time_marginalization: import scipy.special if opts.resample_time_marginalization and not(opts.fairdraw_extrinsic_output): raise Exception(" Resampled time output requires --fairdraw-extrinsic-output ") +if opts.resample_time_marginalization and opts.distance_marginalization: + raise ValueError( + "--resample-time-marginalization does not support distance/phase " + "marginalization; refusing before the expensive integration") # Fairdraw is NOT YET IMPLEMENTED for these other integrators! #if opts.fairdraw_extrinsic_output: @@ -655,6 +699,134 @@ _stencil_prereqs = ( ) _stencil_missing = [name for name, ok in _stencil_prereqs if not ok] _stencil_is_honoured = not _stencil_missing +# THE DEFAULT IS DOWNGRADED WHERE A REQUEST IS REFUSED, and the two must not be conflated. +# +# The refusal below is the right answer to "I asked for sinc and this configuration will quietly +# run something else": it protects a comparison campaign from being run against a flag that did +# nothing. It is the WRONG answer to an inherited default -- as a default it would convert every +# configuration in the list above from working to a startup ValueError, with no command line +# changed anywhere, which is a far larger blast radius than the accuracy the default buys. +# +# So: an EXPLICIT --interpolate-time is refused exactly as before (no behaviour change at all for +# anyone who passes the flag), and a DEFAULT falls back to 'nearest' -- the historical value, so +# the fallback is a no-op relative to today -- with the reason printed. The fallback is announced +# rather than silent because a stencil that is not running is the one thing the log has to say. +def fused_calmarg_in_use(opts, calibration_marginalization=None): + """WILL A FUSED CALIBRATION KERNEL ACTUALLY RUN? ONE definition, two call sites. + + THIS FUNCTION EXISTS BECAUSE THE CONDITION WAS RE-DERIVED AND DRIFTED, TWICE IN A DAY, IN THE + SAME DIRECTION, WHILE A REVIEWER WAS LOOKING AT IT. The stencil guard below needs the answer + at startup; `use_fused_calmarg` needs it again at dispatch time. Written as two expressions + they went `flag` -> `flag and envelope` -> `flag and envelope and path`, each round adding a + conjunct the other site already had in effect. The drift is INVISIBLE by construction: an + over-broad predicate downgrades the inherited stencil to 'nearest', which is exactly the + historical behaviour, so nothing fails and no log line looks wrong. Two copies of a condition + drift; one does not. (Four P2 findings on PR #237, of which three were this predicate.) + + A fused kernel runs only where ALL of these hold. Each is a REACHABILITY fact about the + driver, verified against the call sites, not a guess: + + * calibration marginalization is configured -- otherwise n_cal stays 1; + * --calibration-fused-kernel was passed -- it is opt-in; + * --calibration-n-realizations > 1 -- factored_likelihood returns from its `n_cal == 1` + branch before `cal_method` is read at all, so with one realization 'fused' is inert; + * NOT --rotation-slow and NOT --freqresponse -- both REPLACE the likelihood. The + non-distmarg dispatch puts the only `cal_method='fused'` call site in the `else` of + `if opts.rotation_slow: ... elif opts.freqresponse: ...`, and the two distmarg call + sites are unreachable because both options refuse distance marginalization outright at + startup; + * NOT --calibration-dump-responsibilities -- the pilot evaluates with an explicit + `cal_method='loop'` and `return_cal_components=True` (which the library also refuses to + fuse) and then RETURNS before the production integration is ever built. + + The stencil is deliberately NOT a conjunct: whether the run is on 'nearest' is the question + the callers are answering, not part of "could a kernel run here at all". + + WHAT THE EARLY CALL CANNOT SEE, stated rather than left to be rediscovered. At startup + `calibration_marginalization` does not exist yet -- it is set several hundred lines later, at + the `if opts.calibration_envelope_directory:` block, from a False initial value. Passing + None here substitutes `bool(opts.calibration_envelope_directory)`, which is that block's own + condition, so it is a restatement and not an approximation. It is also CHECKED: the late + call site compares its result against the early one and refuses if they disagree, so a change + to how `calibration_marginalization` is derived fails loudly instead of silently re-opening + the same drift. Nothing else in the predicate is unavailable at startup -- every other term + is a command-line option. + """ + if calibration_marginalization is None: + calibration_marginalization = bool(opts.calibration_envelope_directory) + return (bool(calibration_marginalization) + and bool(opts.calibration_fused_kernel) + and int(getattr(opts, 'calibration_n_realizations', 0) or 0) > 1 + and not bool(opts.rotation_slow) + and not bool(opts.freqresponse) + and not bool(opts.calibration_dump_responsibilities)) + + +_fused_calmarg_would_run = fused_calmarg_in_use(opts) +# (reason, remedy) pairs, NOT bare reasons. The two downgrades have DIFFERENT remedies and one +# shared sentence would have to lie about one of them: a prerequisite downgrade becomes a REFUSAL +# if you name the stencil explicitly, while the fused-kernel downgrade does NOT -- an explicit +# stencil there is ACCEPTED and the fused kernel is dropped with a notice. The single message +# used until 2026-09-03 told everyone they would be refused, which was wrong for the second case. +# Second P2 review finding on PR #237, and the reason these are printed one line per reason. +_stencil_downgrades = [] +if opts._noloop_time_interp != 'nearest' and not _stencil_is_honoured: + _stencil_downgrades.append(( + "this configuration cannot honour a sub-sample stencil: missing %s" + % ", ".join(_stencil_missing), + "add the missing option(s) -- --gpu accepts --force-xpy if no device is present -- or " + "pass '--interpolate-time %s' explicitly to be REFUSED rather than downgraded" + % TIME_INTERP_DEFAULT)) +# THE FUSED CALIBRATION KERNEL IMPLEMENTS 'nearest' ONLY, deliberately (see +# DESIGN_q_window_stencil.md 9). The three NoLoop call sites already fall back to cal_method +# ='loop' when the stencil is not 'nearest', and the distmarg sites additionally drop the +# cal_distmarg table -- so a non-'nearest' DEFAULT would silently move every +# --calibration-fused-kernel run off the kernel it explicitly asked for, changing both its cost +# and its distance-marginalization path. An explicit stencil still does that (unchanged, and the +# user named both flags); the default stays out of it. +if (opts._noloop_time_interp != 'nearest' and _fused_calmarg_would_run + and not _stencil_downgrades): + _stencil_downgrades.append(( + "--calibration-fused-kernel selects a fused kernel that implements 'nearest' only", + "pass '--interpolate-time %s' explicitly to KEEP the stencil and give up the fused " + "kernel -- that combination is ACCEPTED, not refused: calibration marginalization runs " + "the 'loop' method and the driver says so -- or pass '--interpolate-time nearest' to " + "make the current behaviour explicit" % TIME_INTERP_DEFAULT)) +if _stencil_downgrades and opts._interp_time_from_default: + for _dg_reason, _dg_remedy in _stencil_downgrades: + print(" Q_lm stencil DEFAULT %r NOT APPLIED -- %s. Falling back to 'nearest' (the " + "pre-2026-09-02 default), so this run is unchanged. To change that: %s." + % (opts._noloop_time_interp, _dg_reason, _dg_remedy)) + opts._noloop_time_interp = 'nearest' + # Both of these were derived from the provisional default and must follow it down. + # The FIRST is load-bearing and is NOT redundant with the stencil reset above: it is the + # plain boolean handed to FactoredLogLikelihoodTimeMarginalized(..., interpolate=...) on the + # NON-VECTORIZED path (:3645, :4002), which is the likelihood that actually runs whenever a + # downgrade is triggered by a missing --vectorized. Without it an omitted --interpolate-time + # would switch every such run onto the legacy path's unrelated cubic interpolation. Pinned by + # test_batchmode_stencil_default.test_the_downgrade_also_takes_the_legacy_scalar_path_down_with_it. + # The SECOND is BELT-AND-BRACES, not a guard: this block only runs when the stencil came from + # the default, and in that case _ti_for_export was already 'nearest', so the recomputed value + # is always the value already stored. Kept so the two derived quantities are reset in one + # place if _ti_for_export's rule ever changes; no test can distinguish it (verified by + # mutation, 2026-09-03 -- deleting it leaves the whole gate green). + opts._legacy_interpolate_time = False + opts._time_posterior_export = resolve_time_posterior_export_mode( + opts.time_posterior_export, 'nearest', + continuous_available=not (opts.rotation_slow or opts.freqresponse)) +if opts._noloop_time_interp != 'nearest' and _fused_calmarg_would_run: + # SAY SO. This combination is not refused -- the user named both flags and the stencil is + # the one that is honoured -- but until now the loss of the fused kernel was silent at all + # three call sites, which contradicts this option's own "REFUSED, not ignored" promise. + # Gated on _fused_calmarg_would_run rather than on the flag, for the same reason as the + # downgrade above: with no calibration envelope there is no fused kernel to lose, and a + # "NOT USED" notice about a kernel that was never going to run is noise that trains readers + # to ignore the line. + print(" --calibration-fused-kernel: NOT USED. The fused calibration kernels implement the " + "'nearest' stencil only (DESIGN_q_window_stencil.md 9), and --interpolate-time %r is " + "in force, so calibration marginalization runs the 'loop' method instead (and the " + "distmarg variants drop the cal_distmarg table). Pass '--interpolate-time nearest' to " + "keep the fused kernel." % (opts._noloop_time_interp,)) if opts._noloop_time_interp != 'nearest' and not _stencil_is_honoured: raise ValueError( "--interpolate-time %r was requested, but this configuration cannot honour it: missing " @@ -664,12 +836,165 @@ if opts._noloop_time_interp != 'nearest' and not _stencil_is_honoured: "present, which keeps the identical NoLoop code path on numpy -- or drop " "--interpolate-time. Refusing rather than running a different likelihood than the one " "you asked for." % (opts._noloop_time_interp, ", ".join(_stencil_missing))) -print(" Q_lm sub-sample time stencil: {} (from --interpolate-time {!r}); honoured by this " +if (opts.resample_time_marginalization and + opts._time_posterior_export == "continuous" and + (not opts.gpu or opts.rotation_slow or opts.freqresponse)): + raise NotImplementedError( + "continuous time-posterior export is not implemented for " + "the active likelihood path: it requires the maintained NoLoop " + "GPU/--force-xpy evaluator without --rotation-slow/--freqresponse, so " + "the exported draw uses the same likelihood as integration; use " + "--time-posterior-export grid") +# --time-marginalization-quadrature: same refuse-don't-ignore discipline as the stencil guard +# above, and for the same reason -- an accuracy option that silently does nothing is worse than +# one that is unavailable, because a comparison campaign can be run against it and believed. +opts._time_quadrature = str(opts.time_marginalization_quadrature).strip().lower() +factored_likelihood.time_quadrature_module.validate_time_quadrature(opts._time_quadrature) +_tq_prereqs = ( + ('--time-marginalization', bool(opts.time_marginalization)), + ('--vectorized', bool(opts.vectorized)), + ('--gpu (accepts --force-xpy)', bool(opts.gpu)), + ('not --rotation-slow (time-DEPENDENT rho_sq: the band-limited argument does not hold)', + not bool(opts.rotation_slow)), + ('not --freqresponse (separate likelihood, not audited for this)', + not bool(opts.freqresponse)), + ('no calibration marginalization (--calibration-envelope-directory)', + not bool(opts.calibration_envelope_directory)), +) +# Phase marginalization is refused by 'peak-local' ONLY -- 'bandlimited' supports it +# and must not regress. It is not a plain CLI boolean: it is a property of the +# distance-marginalization lookup table, which is already loaded above, so the check +# has to read it from there. Checking it HERE, at startup, is the point: the +# likelihood also refuses it, but by then the run is under way. +_tq_phase_marg = bool(opts.distance_marginalization) and bool( + lookup_table["phase_marginalization"]) +if opts._time_quadrature == 'peak-local' and _tq_phase_marg: + _tq_prereqs = _tq_prereqs + (( + "not phase marginalization (the peak-local width picks up an (I1/I0)(|kappa|/D) " + "factor that does not reduce; use --time-marginalization-quadrature bandlimited, " + "which supports it)", False),) +_tq_missing = [name for name, ok in _tq_prereqs if not ok] +if opts._time_quadrature != 'simpson' and _tq_missing: + raise ValueError( + "--time-marginalization-quadrature %r was requested, but this configuration cannot " + "honour it: %s. Refusing rather than running the historical Simpson quadrature while " + "reporting that you asked for something else." + % (opts._time_quadrature, "; ".join(_tq_missing))) +if opts._time_quadrature == 'bandlimited': + if opts.time_posterior_export == 'grid': + raise ValueError( + "--time-marginalization-quadrature bandlimited mandates a continuous " + "conditional-posterior time draw; --time-posterior-export grid would " + "discard the sub-sample result") + if opts.srate_resample_time_marginalization is not None: + raise ValueError( + "--time-marginalization-quadrature bandlimited derives its own time " + "resolution and exports a continuous draw; drop the conflicting fixed " + "--srate-resample-time-marginalization lattice") + opts._time_posterior_export = 'continuous' +if opts._time_quadrature == 'peak-local' and opts.resample_time_marginalization: + # REFUSED, for the same refuse-don't-ignore reason as everything else on this path, and + # this one is easy to miss because BOTH available answers are silently wrong. + # + # 'grid' export draws `t_ref` from the original coarse `lnLt` bins. peak-local exists to + # resolve a peak whose width is far below that spacing, so the integral would be + # sub-sample accurate and the exported time would be quantised to the very grid the option + # was introduced to escape -- the resolution is computed and then discarded. With + # --interpolate-time nearest, `auto` resolves to 'grid'; since 2026-09-02 the stencil + # default is TIME_INTERP_DEFAULT, but a DEFAULT-derived stencil is deliberately not fed to + # the export resolver (see _ti_for_export above), so 'grid' is still the DEFAULT outcome and + # this refusal is still the one a default configuration meets. + # + # 'continuous' export is not available either: it needs `return_time_draw`, which requires + # a validated dense reconstruction over the whole window. peak-local by construction never + # forms one -- it evaluates only near the peaks -- and the library refuses the combination + # for exactly that reason. + # + # So there is nothing correct to do here, and the honest move is to say so rather than pick + # the quieter of two wrong answers. 'bandlimited' serves this and is one flag away. + raise ValueError( + "--time-marginalization-quadrature peak-local does not support " + "--resample-time-marginalization. A 'grid' export would draw t_ref from the coarse " + "lnLt bins, discarding the sub-sample resolution this quadrature exists to recover, " + "and a 'continuous' export needs the dense reconstruction over the whole window that " + "peak-local deliberately never forms. Use " + "--time-marginalization-quadrature bandlimited, which exports a validated continuous " + "draw, or drop --resample-time-marginalization.") +# One assignment, inherited by every DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop call site. +factored_likelihood.TIME_QUADRATURE_DEFAULT = opts._time_quadrature +# Announce the value READ BACK OUT of the module, not the one parsed from the +# command line. Those are the same string only if the assignment above actually +# happened, and a banner built from `opts` reports what was ASKED FOR rather than +# what is in force -- so deleting the assignment leaves the flag inert while the +# banner still says it is honoured. Reading back is what makes the printed line, +# and the tests that assert on it, load-bearing. +print(" Time-marginalization quadrature: {} (from --time-marginalization-quadrature {!r}); " + "honoured by this configuration: {}".format( + factored_likelihood.TIME_QUADRATURE_DEFAULT, + opts.time_marginalization_quadrature, not _tq_missing)) + +print(" Q_lm sub-sample time stencil: {} (from --interpolate-time {}); honoured by this " "configuration: {} [time_marginalization={} vectorized={} gpu={} rotation_slow={} " "freqresponse={}]; legacy scalar path interpolate={}".format( - opts._noloop_time_interp, opts.interpolate_time, _stencil_is_honoured, + opts._noloop_time_interp, + ("DEFAULT %r" % TIME_INTERP_DEFAULT if opts._interp_time_from_default + else repr(opts.interpolate_time)), _stencil_is_honoured, bool(opts.time_marginalization), bool(opts.vectorized), bool(opts.gpu), bool(opts.rotation_slow), bool(opts.freqresponse), opts._legacy_interpolate_time)) +# THE BAND-LIMITED QUADRATURES WERE MEASURED AGAINST 'nearest', WHICH IS NO LONGER THE DEFAULT. +# +# 'bandlimited' reconstructs the integrand THE CODE ACTUALLY FORMS. That is the true kappa(t) +# only for the 'nearest' gather; with 'cubic' or 'sinc' the gathered values are a fixed FIR +# filter applied to Q, still band-limited, so the reconstruction stays exact -- but exact for the +# FILTERED function, and the stencil's own bias then dominates. Measured at srate 4096, peak +# lnL ~5300 (time_marginalization_quadrature.py): with 'nearest' this path is +0.0002 nats against +# an analytic truth where Simpson is -521; with 'sinc' it is -2.29 where Simpson is +1.28, and +# over a scan of seeds and grid phases Simpson wins about half the cases. 'peak-local' inherits +# this: its accuracy is DEFINED against 'bandlimited' (max 1.9e-11 nats), so a stencil bias in +# the reference is a stencil bias in it. +# +# AND THE PAIRING IS NOT RARE -- IT IS UNIVERSAL. The quadrature's prerequisites +# (--time-marginalization --vectorized --gpu) are a strict SUBSET of the stencil's honoured set, +# so EVERY run that opts into a band-limited quadrature without naming a stencil now gets the +# default one. An explicit accuracy option must not be moved into an unvalidated regime in +# silence: same discipline as the --calibration-fused-kernel notice above. NOT a downgrade -- +# the quadrature is not wrong, its ADVANTAGE is unestablished here -- and not a refusal, because +# either stencil is a legitimate choice. Third P2 review finding on PR #237; +# DESIGN_q_window_stencil.md 9.6.4 called for this notice and records the underlying measurement +# as still open. +if (opts._time_quadrature in ('bandlimited', 'peak-local') + and opts._noloop_time_interp != 'nearest'): + print(" --time-marginalization-quadrature %s: ADVANTAGE NOT ESTABLISHED with the " + "%s Q_lm stencil %r. Its measured accuracy (+0.0002 nats against an analytic truth " + "where Simpson is -521) is for 'nearest'; with 'sinc' the same comparison is -2.29 " + "nats where Simpson is +1.28, and Simpson wins about half a scan of seeds and grid " + "phases -- once a stencil is in force its own bias dominates the quadrature error. " + "The run is not wrong and nothing is being downgraded; the PAIRING is unmeasured " + "(RIFT/likelihood/DESIGN_q_window_stencil.md 9.6.4, open item). Pass " + "'--interpolate-time nearest' to reproduce the regime these numbers were measured in." + % (opts._time_quadrature, + "DEFAULT" if opts._interp_time_from_default else "explicitly requested", + opts._noloop_time_interp)) +if opts.resample_time_marginalization: + print(" Time-posterior export: {} (from --time-posterior-export {!r})".format( + opts._time_posterior_export, opts.time_posterior_export)) + if (opts._time_posterior_export == "continuous" and + opts.srate_resample_time_marginalization): + print(" Time-posterior export: continuous draws supersede " + "--srate-resample-time-marginalization; use " + "--time-posterior-export grid for the requested lattice.") + +# CALIBRATION-MARGINALIZATION OPTION COMPATIBILITY. Same refuse-don't-ignore discipline as the +# stencil and quadrature gates above, applied to the one option that switches in-loop calibration +# marginalization on (--calibration-envelope-directory) and to the opt-ins that are read only +# under it. Placed HERE because opts.gpu is final only after the cupy availability downgrade +# above, and a `--gpu` that has been downgraded to numpy is one of the configurations that drops +# calibration on the floor. Nothing expensive has run yet: no frames, no PSDs, no precompute. +# +# This gate only REFUSES. It never changes what an accepted configuration computes. +import RIFT.calmarg.option_compat as calibration_option_compat +calibration_option_compat.refuse_incompatible_calibration_options(opts) + manual_avoid_overflow_logarithm=opts.manual_logarithm_offset manual_avoid_overflow_logarithm_default = manual_avoid_overflow_logarithm @@ -777,7 +1102,13 @@ else: # fiducial_epoch = lal.LIGOTimeGPS() -fiducial_epoch = event_time.seconds + 1e-9*event_time.nanoseconds # no more direct access to gpsSeconds +# Keep the exact pair as well as the historical float view. A float64 near a +# current GPS epoch has a 0.12--0.24 microsecond spacing, coarser than some of the +# derived band-limited grids; reconstructing XML nanoseconds from that float +# would silently throw away the sub-sample draw at the final serialization step. +fiducial_epoch_seconds = int(event_time.seconds) +fiducial_epoch_nanoseconds = int(event_time.nanoseconds) +fiducial_epoch = fiducial_epoch_seconds + 1e-9*fiducial_epoch_nanoseconds # compatibility view # Struct to hold template parameters P_list = None @@ -1288,6 +1619,28 @@ for _optv, _k in [(opts.limit_psi, 'psi'), (opts.limit_right_ascension, 'right_a print(" [limit] restricting {} sampling/prior to [{:.4f}, {:.4f}]".format(_k, _lo, _hi)) limit_declination_active = bool(opts.limit_declination) limit_inclination_active = bool(opts.limit_inclination) +# --limit-distance: SAMPLING-only narrowing. dist_prior_range is captured BEFORE the +# narrowing and is what the distance prior normalizes over, so the reported lnZ keeps the +# full-range scale. (Under --internal-reparam-dl-incl the range captured here is the +# widened D_eff axis, which is what that mode already normalized over -- and which is why +# --limit-distance is refused there rather than reinterpreted.) +dist_prior_range = (param_limits["distance"][0], param_limits["distance"][1]) +limit_distance_active = bool(opts.limit_distance) +if limit_distance_active: + if opts.distance_marginalization: + raise SystemExit(" --limit-distance is not compatible with --distance-marginalization: that path has no distance sampler to narrow (the distance integral is done analytically over [--d-min,--d-max] from the lookup table).") + if getattr(opts, 'd_prior_redshift', False): + raise SystemExit(" --limit-distance is not compatible with --d-prior-redshift: the sampled coordinate is then redshift, not luminosity distance in Mpc.") + if opts.pin_distance_to_sim: + raise SystemExit(" --limit-distance is not compatible with --pin-distance-to-sim: that path PINS distance to the injection value (analyze_event sets pinned_params['distance'] = P.dist), so there is no distance draw for a box to restrict and the option would be a silent no-op. Same class as --distance-marginalization above: refuse rather than accept-and-ignore.") + if opts.internal_reparam_dl_incl: + raise SystemExit(" --limit-distance is not compatible with --internal-reparam-dl-incl: the sampled distance axis is then D_eff = d_L/A(iota), not d_L, so a box in Mpc of d_L does not map to a box in the sampled coordinate.") + try: + _dlo, _dhi = distance_limit_range(opts.limit_distance, dist_prior_range[0], dist_prior_range[1]) + except ValueError as _e: + raise SystemExit(" --limit-distance: {}".format(_e)) + param_limits["distance"] = (_dlo, _dhi) + print(" [limit] restricting distance sampling to [{:.4f}, {:.4f}] Mpc (prior normalization UNCHANGED over [{:.4f}, {:.4f}] Mpc)".format(_dlo, _dhi, dist_prior_range[0], dist_prior_range[1])) # # Parameter integral sampling strategy @@ -1542,8 +1895,12 @@ if (opts.d_prior == 'cosmo' or opts.d_prior == 'cosmo_sourceframe') and not opts sampler.add_parameter("distance", pdf = dist_sampler, cdf_inv = dist_sampler_cdf_inv, - left_limit = dmin, - right_limit = dmax, + # Historically the physical [dmin,dmax]; narrowed only by + # --limit-distance. pdf_dL is normalized over the full + # [zmin,zmax] <-> [dmin,dmax] either way, so the evidence + # scale does not move. + left_limit = param_limits["distance"][0] if limit_distance_active else dmin, + right_limit = param_limits["distance"][1] if limit_distance_active else dmax, prior_pdf = pdf_dL, #only thing preserved in calculation adaptive_sampling = (adapt_extra_extrinsic and not (opts.no_adapt or opts.no_adapt_distance)) or opts.force_adapt_all) else: @@ -1558,16 +1915,22 @@ if (opts.d_prior == 'cosmo' or opts.d_prior == 'cosmo_sourceframe') and not opts prior_pdf = pdf_z, #only thing preserved in calculation adaptive_sampling = (adapt_extra_extrinsic and not (opts.no_adapt or opts.no_adapt_distance)) or opts.force_adapt_all) elif not opts.distance_marginalization: - dist_sampler = mcsampler.ret_uniform_samp_vector_alt( param_limits["distance"][0], param_limits["distance"][1]) - dist_sampler_cdf_inv = functools.partial(mcsampler.uniform_samp_cdf_inv_vector, param_limits["distance"][0], param_limits["distance"][1]) - #dist_sampler=functools.partial( mcsampler.uniform_samp_withfloor_vector, numpy.min([distBoundGuess,param_limits["distance"][1]]), param_limits["distance"][1], 0.001) - dist_prior_pdf = lambda x: x**2/(param_limits["distance"][1]**3/3. - param_limits["distance"][0]**3/3.) - if opts.d_prior == 'pseudo_cosmo': - nm = priors_utils.dist_prior_pseudo_cosmo_eval_norm(param_limits["distance"][0],param_limits["distance"][1]) - dist_prior_pdf =functools.partial( priors_utils.dist_prior_pseudo_cosmo, nm=nm,xpy=xpy_default) - elif opts.d_prior != 'Euclidean': + # The SAMPLING range is param_limits["distance"] (narrowed by --limit-distance); the + # PRIOR normalization range is dist_prior_range, which --limit-distance never touches. + # Keeping them as two arguments is the point: the old one-range form normalized the + # Euclidean density over whatever the sampler happened to be drawing from, so narrowing + # for cost silently rescaled the evidence. See mcsampler.distance_sampler_kwargs(). + try: + _dist_kwargs = distance_sampler_kwargs( + mcsampler, param_limits["distance"], dist_prior_range, + d_prior=opts.d_prior, xpy=xpy_default, + adaptive_sampling = (adapt_extra_extrinsic and not (opts.no_adapt or opts.no_adapt_distance)) or opts.force_adapt_all) + except ValueError: print(" ==== WARNING UNKNOWN DISTANCE PRIOR === ") raise Exception('distance prior') + dist_sampler = _dist_kwargs['pdf'] + dist_sampler_cdf_inv = _dist_kwargs['cdf_inv'] + dist_prior_pdf = _dist_kwargs['prior_pdf'] if opts.internal_reparam_dl_incl: # dist_prior_pdf is normalized over the WIDENED D_eff range; the physical prior must be # normalized over [dmin,dmax]. ln F = ln(mass of dist_prior_pdf in [dmin,dmax]); the closure @@ -1580,13 +1943,7 @@ elif not opts.distance_marginalization: _REPARAM_LNF = float(numpy.log(numpy.trapz(_pg, _xg))) print(" [reparam] physical-range prior-mass fraction F={:.4f} (lnF={:.3f}); lnZ normalization matched".format(numpy.exp(_REPARAM_LNF), _REPARAM_LNF)) #dist_sampler_cdf_inv=None - sampler.add_parameter("distance", - pdf = dist_sampler, - cdf_inv = dist_sampler_cdf_inv, - left_limit = param_limits["distance"][0], - right_limit = param_limits["distance"][1], - prior_pdf = dist_prior_pdf, #only thing physical - adaptive_sampling = (adapt_extra_extrinsic and not (opts.no_adapt or opts.no_adapt_distance)) or opts.force_adapt_all) + sampler.add_parameter("distance", **_dist_kwargs) # prior_pdf is the only thing physical # # Rotate sky coordinates @@ -1734,7 +2091,10 @@ if not opts.time_marginalization: cdf_inv = None, left_limit = param_limits["t_ref"][0], right_limit = param_limits["t_ref"][1], - prior_pdf = functools.partial(mcsampler.uniform_samp_vector, param_limits["t_ref"][0], param_limits["t_ref"][1])) + # Reuse the backend-portable closure above. AV exposes + # ret_uniform_samp_vector_alt but not the legacy + # uniform_samp_vector symbol. + prior_pdf = tref_sampler) # skymap oracle @@ -2099,64 +2459,150 @@ def resample_samples(my_samples, for name in ['phi','theta', 'phiref','incl', 'psi','dist']: setattr(P,name, getattr(P,name).astype(float) ) - # With calibration marginalization (n_cal>1) this returns the cal-marginalized - # lnL(t) timeseries (weighted log-sum-exp over realizations per time bin), so the - # time resampling below operates on the marginalized likelihood. - lnLt = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, - P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default,return_lnLt=True,n_cal=n_cal,cal_log_weights=cal_log_weights, - time_interp=opts._noloop_time_interp, - ctUArrayDict_cal=ctUArrayDict_cal, ctVArrayDict_cal=ctVArrayDict_cal) - - lnLt = identity_convert(lnLt) # back to CPU. Note we have removed offsets - if opts.zero_likelihood: - lnLt =np.zeros(lnLt.shape) # zero likelihood - lnLt_norm = scipy.special.logsumexp(lnLt,axis=-1) - tvals = identity_convert(tvals) # back to CPU -# print(lnLt.shape, lnLt_norm.shape,tvals.shape) - # Loop over and resample in time, picking index according to weights - # - perform higher-resolution interpolation t_out = np.zeros(n_samples) lnL_out = np.zeros(n_samples) - # IF UPSAMPLING, PERFORM NOW. (Currently on - if opts.srate_resample_time_marginalization and opts.srate_resample_time_marginalization > fSample: - # Resample the marginalization-time grid to EXACTLY the requested rate, so - # the exported geocenter time is quantized at 1/srate_resample seconds. We - # step by exactly 1/srate_resample; for the usual power-of-two rates that is - # exactly representable in float64, so consecutive output times differ by - # exactly that step. - # - # HISTORICAL NOTE (issue #146): this comment used to justify the choice by - # the internal grid being "a closed-interval linspace whose spacing is - # ~1/fSample but NOT exactly (here ~4086.7 Hz vs 4096)". That is no longer - # true -- marginalization_time_grid() is spaced EXACTLY deltaT, so an - # integer-factor upsample would now be exact too. More importantly, the - # tvals read as time LABELS below (t_out -> the exported 't_ref') are now - # the times the likelihood actually evaluated; under the old linspace they - # were off by up to 1.4 samples at the window edge. - dt_target = 1.0/opts.srate_resample_time_marginalization - # floor(): stay within [tvals[0], tvals[-1]] so the spline never - # extrapolates. At most one step (<1/srate s, tens of us) is dropped at the - # far edge of the +-75 ms window, where the time-marginalized likelihood is - # negligible. - n_dense = int(np.floor((tvals[-1]-tvals[0])/dt_target)) + 1 - tvals_denser = tvals[0] + dt_target * np.arange(n_dense) - from scipy.interpolate import RegularGridInterpolator, CubicSpline - # cubic spline at first, easiest - generally not exporting too many events - lnLt_new = np.zeros( (lnLt.shape[0], n_dense) ) - for indx_here in np.arange(n_samples): - cs = CubicSpline(tvals, lnLt[indx_here]) - lnLt_new[indx_here] = cs(tvals_denser) - # replace, re-normalize - tvals = tvals_denser; lnLt= lnLt_new + if opts._time_quadrature == "bandlimited": + # The quadrature option carries an export contract: draw from the SAME + # reflected, derived-resolution, dense-width-validated representation used + # by the integral. return_lnLt=True is intentionally coarse-grid and must + # not be used here. + if opts.zero_likelihood: + from RIFT.likelihood.time_marginalization_quadrature import draw_piecewise_linear_log_posterior + tvals_cpu = identity_convert(tvals) + t_out, lnL_out = draw_piecewise_linear_log_posterior( + np.zeros((n_samples, len(tvals_cpu))), float(P.deltaT), + t0=float(tvals_cpu[0])) + else: + t_out, lnL_out = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNKDict, rholmArrayDict, ctUArrayDict, + ctVArrayDict, epochDict, Lmax=opts.l_max, xpy=xpy_default, + n_cal=n_cal, cal_log_weights=cal_log_weights, + time_interp=opts._noloop_time_interp, + ctUArrayDict_cal=ctUArrayDict_cal, + ctVArrayDict_cal=ctVArrayDict_cal, + time_quadrature="bandlimited", return_time_draw=True) + t_out = np.asarray(identity_convert(t_out), dtype=float) + lnL_out = np.asarray(identity_convert(lnL_out), dtype=float) + else: + # With calibration marginalization (n_cal>1) this returns the + # cal-marginalized lnL(t) timeseries (weighted log-sum-exp over realizations + # per time bin), so the legacy/spline resampling below operates on the + # marginalized likelihood. + lnLt = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, + P, lookupNKDict, rholmArrayDict, ctUArrayDict, ctVArrayDict,epochDict,Lmax=opts.l_max,xpy=xpy_default,return_lnLt=True,n_cal=n_cal,cal_log_weights=cal_log_weights, + time_interp=opts._noloop_time_interp, + ctUArrayDict_cal=ctUArrayDict_cal, ctVArrayDict_cal=ctVArrayDict_cal) + + if opts._time_posterior_export == "continuous": + # NoLoop's historical window gather consumes only tvals[0] and len(tvals), + # stepping by P.deltaT. Its explicit-time mode instead evaluates the chosen + # cubic/Lanczos Q stencil independently at every dense geocenter time. + from RIFT.likelihood import time_marginalization_quadrature as _tm_export + sigma_t, _, measurable_t = _tm_export.peak_width_from_lnL( + lnLt, P.deltaT, xpy=xpy_default) + factors_t = _tm_export.required_upsample_factors( + xpy_default.where(measurable_t, sigma_t, np.inf), P.deltaT, + xpy=xpy_default) + time_export_refinement = max(4, int(xpy_default.max(factors_t))) + t0_export = tvals[0] + n_coarse_export = len(tvals) + while True: + if time_export_refinement > _tm_export.UPSAMPLE_FACTOR_MAX: + raise RuntimeError( + "continuous time export needs refinement above the supported ceiling {}" + .format(_tm_export.UPSAMPLE_FACTOR_MAX)) + n_dense = (n_coarse_export - 1) * time_export_refinement + 1 + from RIFT.likelihood.time_posterior import validate_time_posterior_working_set + validate_time_posterior_working_set(n_samples, n_dense) + tvals_dense = (t0_export + (P.deltaT / time_export_refinement) * + xpy_default.arange(n_dense)) + lnLt_dense = factored_likelihood.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals_dense, P, lookupNKDict, rholmArrayDict, + ctUArrayDict, ctVArrayDict, epochDict, Lmax=opts.l_max, + xpy=xpy_default, return_lnLt=True, n_cal=n_cal, + cal_log_weights=cal_log_weights, + time_interp=opts._noloop_time_interp, + ctUArrayDict_cal=ctUArrayDict_cal, + ctVArrayDict_cal=ctVArrayDict_cal, + explicit_time_values=True) + sigma_dense, _, measurable_dense = _tm_export.peak_width_from_lnL( + lnLt_dense, P.deltaT / time_export_refinement, xpy=xpy_default) + extra = _tm_export.required_upsample_factors( + xpy_default.where(measurable_dense, sigma_dense, np.inf), + P.deltaT / time_export_refinement, xpy=xpy_default) + extra = max(1, int(xpy_default.max(extra))) + if extra == 1: + tvals, lnLt = tvals_dense, lnLt_dense + break + # Avoid transiently holding two generations of full dense grids. + del lnLt_dense, tvals_dense, sigma_dense, measurable_dense + time_export_refinement *= extra + print(" Time-posterior selected-stencil refinement: {}x ".format( + time_export_refinement)) + + lnLt = identity_convert(lnLt) # back to CPU. Note we have removed offsets + if opts.zero_likelihood: + lnLt =np.zeros(lnLt.shape) # zero likelihood lnLt_norm = scipy.special.logsumexp(lnLt,axis=-1) - indx_list =np.arange(len(tvals)) - for indx in np.arange(n_samples): - indx_choose = np.random.choice(indx_list, p=np.exp(lnLt[indx] - lnLt_norm[indx])) - t_out[indx] = tvals[indx_choose] - lnL_out[indx] = lnLt[indx][indx_choose] + tvals = identity_convert(tvals) # back to CPU + # print(lnLt.shape, lnLt_norm.shape,tvals.shape) + # Draw from the per-sample time posterior. Sub-sample likelihood + # interpolation implies a continuous export contract by default: choosing a + # coarse tvals index here would throw away the resolution just requested. + if opts._time_posterior_export == "continuous": + from RIFT.likelihood.time_posterior import draw_continuous_time_posterior + t_out, lnL_out = draw_continuous_time_posterior(tvals, lnLt) + # Legacy/fallback grid export, including the explicit higher-rate lattice. + elif opts.srate_resample_time_marginalization and opts.srate_resample_time_marginalization > fSample: + # Resample the marginalization-time grid to EXACTLY the requested rate, so + # the exported geocenter time is quantized at 1/srate_resample seconds. We + # step by exactly 1/srate_resample; for the usual power-of-two rates that is + # exactly representable in float64, so consecutive output times differ by + # exactly that step. + # + # HISTORICAL NOTE (issue #146): this comment used to justify the choice by + # the internal grid being "a closed-interval linspace whose spacing is + # ~1/fSample but NOT exactly (here ~4086.7 Hz vs 4096)". That is no longer + # true -- marginalization_time_grid() is spaced EXACTLY deltaT, so an + # integer-factor upsample would now be exact too. More importantly, the + # tvals read as time LABELS below (t_out -> the exported 't_ref') are now + # the times the likelihood actually evaluated; under the old linspace they + # were off by up to 1.4 samples at the window edge. + dt_target = 1.0/opts.srate_resample_time_marginalization + # floor(): stay within [tvals[0], tvals[-1]] so the spline never + # extrapolates. At most one step (<1/srate s, tens of us) is dropped at the + # far edge of the +-75 ms window, where the time-marginalized likelihood is + # negligible. + n_dense = int(np.floor((tvals[-1]-tvals[0])/dt_target)) + 1 + tvals_denser = tvals[0] + dt_target * np.arange(n_dense) + from scipy.interpolate import RegularGridInterpolator, CubicSpline + # cubic spline at first, easiest - generally not exporting too many events + lnLt_new = np.zeros( (lnLt.shape[0], n_dense) ) + for indx_here in np.arange(n_samples): + cs = CubicSpline(tvals, lnLt[indx_here]) + lnLt_new[indx_here] = cs(tvals_denser) + # replace, re-normalize + tvals = tvals_denser; lnLt= lnLt_new + lnLt_norm = scipy.special.logsumexp(lnLt,axis=-1) + if opts._time_posterior_export == "grid": + indx_list =np.arange(len(tvals)) + for indx in np.arange(n_samples): + indx_choose = np.random.choice(indx_list, p=np.exp(lnLt[indx] - lnLt_norm[indx])) + t_out[indx] = tvals[indx_choose] + lnL_out[indx] = lnLt[indx][indx_choose] # print(' Resampled time offset {} '.format(t_out[indx])) #, lnLt[indx]-lnLt_norm[indx]) - my_samples['t_ref'] = fiducial_epoch+t_out # add sample time jitter from reweighting to samples + my_samples['t_ref'] = fiducial_epoch+t_out + if opts._time_posterior_export == "continuous": + # Keep the historical float compatibility view above, but serialize a + # continuous draw from exact integer GPS fields. Adding a sub-microsecond + # offset directly to a ~1e9 s float epoch can round away the very resolution + # this path exists to export. Do not alter legacy grid serialization. + (_gps_seconds_exact, + _gps_nanoseconds_exact) = xmlutils.gps_add_seconds_exact( + fiducial_epoch_seconds, fiducial_epoch_nanoseconds, t_out) + my_samples['t_ref_gps_seconds'] = _gps_seconds_exact + my_samples['t_ref_gps_nanoseconds'] = _gps_nanoseconds_exact my_samples["lnL_raw"] = lnL_out # export likelihoods (equivalent to SNR). Note needs downstream code filters to catch this and put it somewhere. return my_samples @@ -2396,6 +2842,55 @@ def ln_weights_for_posterior(rvs, sampler, convert=None, use_lnL=None): dtype=float) +def _equal_weight_fairdraw_for_serialization(rvs, sampler, n_max, convert=None, + use_lnL=None, rng=None): + """Return an equal-weight posterior draw for the external sample file. + + PR #87 deliberately lets an integrator keep its weighted retained record + when a fair draw would not shrink it. That is the correct *internal* + contract: later in-process consumers can still use the real importance + weights. XML cannot carry that full weight provenance, however, so a + caller of ``--fairdraw-extrinsic-output`` must never serialize such a + record as though its rows were equal-weight posterior samples. + + Do the missing draw only at the serialization boundary. Already + equal-weight records pass through unchanged. Raw or replica-pooled + records are resampled from ``ln_weights_for_posterior``; the draw size is + bounded by both the requested export cap and the retained row count. + Sampling with replacement is intentional for a fair draw. Final RIFT + posterior assembly applies its separate unique-output policy later. + """ + if _rvs_is_equal_weight(sampler): + return rvs + n_have = _rvs_len(rvs) + if n_have == 0: + return rvs + n_draw = min(int(n_max), n_have) + if n_draw < 1: + raise ValueError("fair-draw serialization requires a positive sample cap") + ln_w = numpy.asarray(ln_weights_for_posterior( + rvs, sampler, convert=convert, use_lnL=use_lnL), dtype=float) + finite = numpy.isfinite(ln_w) + if not numpy.any(finite): + raise ValueError("fair-draw serialization has no finite posterior weights") + scale = numpy.max(ln_w[finite]) + weights = numpy.zeros(len(ln_w), dtype=float) + weights[finite] = numpy.exp(ln_w[finite] - scale) + total = numpy.sum(weights) + if not numpy.isfinite(total) or total <= 0: + raise ValueError("fair-draw serialization weights cannot be normalized") + weights /= total + rng = numpy.random if rng is None else rng + indices = rng.choice(numpy.arange(n_have), size=n_draw, + replace=True, p=weights) + out = {} + for key, value in rvs.items(): + array = convert(value) if convert is not None else value + array = numpy.asarray(array) + out[key] = array[:, indices] if isinstance(key, tuple) else array[indices] + return out + + def _pool_replica_rvs(rep_rvs, sampler, rep_lnZ=None, already_resampled=False, use_lnL=None, records=None): """Concatenate the replicas' samples into one correctly-weighted set. @@ -2918,7 +3413,22 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t # use the fused implementation (Option C) when requested; works on GPU (CUDA # kernels) and CPU (numpy). Phase marginalization is not supported by the fused # path, so it is disabled there below (that call site stays on the loop method). - use_fused_calmarg = bool(calibration_marginalization and opts.calibration_fused_kernel) + # THE SAME COMPUTATION as the startup stencil guard, not a second expression -- see + # fused_calmarg_in_use. `calibration_marginalization` is passed explicitly because it is the + # one term the early call had to substitute for; comparing the two results is what turns a + # future divergence in how it is derived into a loud failure instead of a silent re-opening + # of the drift this function was written to end. + use_fused_calmarg = fused_calmarg_in_use(opts, calibration_marginalization) + if use_fused_calmarg != _fused_calmarg_would_run: + raise ValueError( + "internal inconsistency: at startup the fused-calibration-kernel predicate was %r " + "and at dispatch it is %r. The startup value decided whether the inherited " + "--interpolate-time default was downgraded to 'nearest', so the stencil now in force " + "(%r) was chosen on a premise that no longer holds. This means " + "`calibration_marginalization` is no longer `bool(opts.calibration_envelope_directory)`; " + "fused_calmarg_in_use's early substitution must be updated in the same commit. " + "Refusing rather than integrating with a stencil chosen for the wrong reason." + % (_fused_calmarg_would_run, use_fused_calmarg, opts._noloop_time_interp)) if calibration_marginalization: extra_kwargs['calibration_realizations'] = calibration_realization_dict extra_kwargs['calibration_conjugate'] = bool(opts.calibration_conjugate_phase) @@ -2934,7 +3444,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t NR_group=NR_template_group,NR_param=NR_template_param, use_gwsignal=opts.use_gwsignal, use_gwsignal_approx=opts.approximant, - use_external_EOB=opts.use_external_EOB,nr_lookup=opts.nr_lookup,nr_lookup_valid_groups=opts.nr_lookup_group,perturbative_extraction=opts.nr_perturbative_extraction,perturbative_extraction_full=opts.nr_perturbative_extraction_full,use_provided_strain=opts.nr_use_provided_strain,hybrid_use=opts.nr_hybrid_use,hybrid_method=opts.nr_hybrid_method,ROM_group=opts.rom_group,ROM_param=opts.rom_param,ROM_use_basis=opts.rom_use_basis,verbose=opts.verbose,quiet=not opts.verbose,ROM_limit_basis_size=opts.rom_limit_basis_size_to,no_memory=opts.no_memory,skip_interpolation=opts.vectorized, extra_waveform_kwargs=extra_waveform_kwargs,**extra_kwargs) + use_external_EOB=opts.use_external_EOB,nr_lookup=opts.nr_lookup,nr_lookup_valid_groups=opts.nr_lookup_group,perturbative_extraction=opts.nr_perturbative_extraction,perturbative_extraction_full=opts.nr_perturbative_extraction_full,use_provided_strain=opts.nr_use_provided_strain,hybrid_use=opts.nr_hybrid_use,hybrid_method=opts.nr_hybrid_method,ROM_group=opts.rom_group,ROM_param=opts.rom_param,ROM_use_basis=opts.rom_use_basis,verbose=opts.verbose,quiet=not opts.verbose,ROM_limit_basis_size=opts.rom_limit_basis_size_to,no_memory=opts.no_memory,skip_interpolation=opts.vectorized, extra_waveform_kwargs=extra_waveform_kwargs,force_22_mode=opts.force_hyperbolic_22,**extra_kwargs) # skip nan ! Something horrible has happened if np.isnan(guess_snr): @@ -3264,7 +3774,8 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t P.dist = di* 1.e6 * lalsimutils.lsu_PC # luminosity distance lnL[i] = factored_likelihood.FactoredLogLikelihood( - P, rholms_intp, cross_terms, cross_terms_V, opts.l_max) + P, rholms, rholms_intp, cross_terms, cross_terms_V, + opts.l_max) i+=1 if return_lnL: return lnL - manual_avoid_overflow_logarithm @@ -4437,15 +4948,23 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t # the file layout is determined by hyperpipeline_io.build_column_list. # ------------------------------------------------------------------ from RIFT.misc import hyperpipeline_io as _hpio + # Which optional parameter groups this row carries. The hyperpipeline + # writer and the legacy writer below consume the SAME flags, so a run + # that enables several groups at once emits all of them in either + # format instead of only whichever one wins a branch dispatch. + _use_ecc = bool(opts.save_eccentricity) + _use_mpa = bool(_use_ecc and opts.save_meanPerAno) + _use_tides = bool(P.lambda1>0 or P.lambda2>0) + _use_eos_index = bool(_use_tides and opts.export_eos_index) + _use_eob = bool(opts.save_EOB_parameters) + _use_hyp = bool(opts.save_hyperbolic) + _use_distance = bool(opts.pin_distance_to_sim and not any( + (_use_tides, _use_ecc, _use_eob, _use_hyp))) if _hpio.is_active(): - _use_ecc = bool(opts.save_eccentricity) - _use_mpa = bool(_use_ecc and opts.save_meanPerAno) - _use_tides = bool(P.lambda1>0 or P.lambda2>0) - _use_eos_index = bool(_use_tides and opts.export_eos_index) - _use_distance = bool(opts.pin_distance_to_sim and not _use_tides and not _use_ecc) _cols = _hpio.build_column_list( use_eccentricity=_use_ecc, use_meanPerAno=_use_mpa, use_tides=_use_tides, use_eos_index=_use_eos_index, + use_eob_parameters=_use_eob, use_hyperbolic=_use_hyp, use_distance=_use_distance) _vals = { "lnL": log_res+manual_avoid_overflow_logarithm, @@ -4463,28 +4982,49 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t _vals["lambda2"] = P.lambda2 if _use_eos_index: _vals["eos_table_index"] = P.eos_table_index + if _use_eob: + _vals["a6c"] = P.a6c + if _use_hyp: + _vals["E0"] = P.E0 + _vals["p_phi0"] = P.p_phi0 if _use_distance: _vals["distance"] = pinned_params["distance"] _hpio.write_row(fname_output_txt, _cols, [_vals[c] for c in _cols]) - elif opts.save_eccentricity: - if opts.save_meanPerAno: - # output format when eccentricity & meanPerAno are being used - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, P.eccentricity, P.meanPerAno, log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" - else: - # output format when only eccentricity is being used - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, P.eccentricity, log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" - elif not (P.lambda1>0 or P.lambda2>0): - # output format when lambda is NOT used - if not opts.pin_distance_to_sim: - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" - else: - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, pinned_params["distance"], log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" else: - if not(opts.export_eos_index): - # Alternative output format if lambda is active - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, P.lambda1, P.lambda2, log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" - else: - numpy.savetxt(fname_output_txt, numpy.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, P.lambda1, P.lambda2, P.eos_table_index, log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res,sampler.ntotal, neff ]])) #dict_return["convergence_test_results"]["normal_integral]" + # ---------------------------------------------------------------- + # Legacy ASCII row, built COMPOSITIONALLY (one `if` per enabled + # group), not by dispatching to one branch per combination. The + # optional columns must appear in the same order as CIP's `col_lnL` + # increment chain in + # util_ConstructIntrinsicPosterior_GenericCoordinates.py: + # + # event_id m1 m2 s1x s1y s1z s2x s2y s2z + # [distance] [lambda1 lambda2 [eos_table_index]] [a6c] [E0 p_phi0] + # [eccentricity [meanPerAno]] + # lnL sigma_lnL ntotal neff + # + # A branch-per-combination layout silently dropped every group but + # the first (e.g. --save-eccentricity won over a6c/E0/p_phi0) while + # CIP still allocated columns for each enabled group, so it read the + # likelihood/statistics columns as physical parameters. + # ---------------------------------------------------------------- + _row = [event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z] + if _use_distance: + _row += [pinned_params["distance"]] + if _use_tides: + _row += [P.lambda1, P.lambda2] + if _use_eos_index: + _row += [P.eos_table_index] + if _use_eob: + _row += [P.a6c] + if _use_hyp: + _row += [P.E0, P.p_phi0] + if _use_ecc: + _row += [P.eccentricity] + if _use_mpa: + _row += [P.meanPerAno] + _row += [log_res+manual_avoid_overflow_logarithm, sqrt_var_over_res, sampler.ntotal, neff] + numpy.savetxt(fname_output_txt, numpy.array([_row])) # Per-intrinsic likelihood-vs-distance grid. Pure extrinsic-marginalized # likelihood as a function of d_L: divides out the distance sampling @@ -4733,6 +5273,18 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t if opts.save_samples and opts.output_file: import copy samples = copy.deepcopy(sampler._rvs) # deep copy: avoid modifying structures and having side effect on integrator, which loops over keys Expensive! + if opts.fairdraw_extrinsic_output: + # The sampler intentionally skips its internal fair draw when it would + # not shrink a tiny retained record (PR #87). Internal consumers still + # have the weights in that case, but XML does not preserve their full + # provenance. Complete the draw at the serialization boundary so the + # external file always satisfies the equal-weight contract promised by + # --fairdraw-extrinsic-output. This also flattens replica-pooled output + # according to the pooled posterior weights. + samples = _equal_weight_fairdraw_for_serialization( + samples, sampler, + min(opts.fairdraw_extrinsic_output_n_max, opts.n_eff), + convert=identity_convert, use_lnL=rvs_integrand_is_lnL) # Insert reference distance if it was marginalized over if "distance" not in samples: # Not distance output is the same as internal calculations: in *Mpc* @@ -4809,6 +5361,9 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t samples["spin2z"] =numpy.ones(samples["psi"].shape)*P.s2z samples["alpha4"] =numpy.ones(samples["psi"].shape)*P.eccentricity samples["alpha"] =numpy.ones(samples["psi"].shape)*P.meanPerAno + samples["psi0"] =numpy.ones(samples["psi"].shape)*P.a6c + samples["beta"] =numpy.ones(samples["psi"].shape)*P.p_phi0 + samples["psi3"] =numpy.ones(samples["psi"].shape)*P.E0 samples["alpha5"] =numpy.ones(samples["psi"].shape)*P.lambda1 samples["alpha6"] =numpy.ones(samples["psi"].shape)*P.lambda2 # Below exist solely to placate XML export; new issue as of latest lalsuite say 7.15+ or so @@ -4862,7 +5417,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t print(" WARNING: --calibration-export-posterior failed ({}); skipping cal-posterior export.".format(_e_cep)) xmlutils.append_samples_to_xmldoc(xmldoc, samples) # Extra metadata - dict_out={"mass1": m1, "mass2": m2, "spin1z": P.s1z, "spin2z": P.s2z, "alpha4": P.eccentricity, "alpha": P.meanPerAno, "alpha5":P.lambda1, "alpha6":P.lambda2, "event_duration": sqrt_var_over_res, "ttotal": sampler.ntotal} + dict_out={"mass1": m1, "mass2": m2, "spin1z": P.s1z, "spin2z": P.s2z, "alpha4": P.eccentricity, "alpha": P.meanPerAno, "alpha5":P.lambda1, "alpha6":P.lambda2, "event_duration": sqrt_var_over_res, "ttotal": sampler.ntotal, "psi0": P.a6c, "psi3": P.E0, "beta": P.p_phi0} # if 'distance' in pinned_params: # dict_out['distance'] = pinned_params["distance"] converged_result = False diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa index bfe81fe8f..8e3c929b5 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa @@ -238,6 +238,10 @@ optp.add_option("--psd-window-shape", type=float, default=0, help="Shape of Tuke optp.add_option("-m", "--time-marginalization", action="store_true", help="Perform marginalization over time via direct numerical integration. Default is false.") #optp.add_option("--n-fairdraw-extrinsic-samples",default=None,type=int,help="Extracts a concrete number of fair draw extrinsic samples, bounded above by n_eff") optp.add_option("--resample-time-marginalization",action='store_true', help="If time-marginalizaiton is true (and should almost always be true), at the end export step use resampling. REQUIRES using fairdraw-extrinsic-output") +optp.add_option("--srate-resample-time-marginalization", type=int, default=None, + help="Requested lattice rate for time-posterior resampling. The LISA driver uses this as its interpolation/draw grid; default keeps its historical 1 ms grid.") +optp.add_option("--time-posterior-export", type="choice", choices=["auto", "continuous", "grid"], default="auto", + help="How --resample-time-marginalization exports time. The LISA likelihood cannot yet evaluate the selected stencil at arbitrary times, so auto/grid use its interpolation lattice and continuous is accepted as the same best-available lattice export for drop-in CLI compatibility (with a warning).") optp.add_option("-d", "--distance-marginalization", action="store_true", help="Perform marginalization over distance via a look-up table. Default is false.") optp.add_option("-l", "--distance-marginalization-lookup-table", default=None, help="Look-up table for distance marginalization.") optp.add_option("--vectorized", action="store_true", help="Perform manipulations of lm and timeseries using numpy arrays, not LAL data structures. (Combine with --gpu to enable GPU use, where available)") @@ -387,6 +391,21 @@ for pin_param in LIKELIHOOD_PINNABLE_PARAMS: optp.add_option_group(pinnable) opts, args = optp.parse_args() +from RIFT.likelihood.time_posterior import ( + legacy_time_interpolation_enabled, resolve_time_posterior_export_mode) +opts.interpolate_time = legacy_time_interpolation_enabled(opts.interpolate_time) +opts._time_posterior_export = resolve_time_posterior_export_mode( + opts.time_posterior_export, "cubic" if opts.interpolate_time else "nearest", + continuous_available=False) +if (opts.resample_time_marginalization and + opts._time_posterior_export == "continuous"): + print("WARNING: LISA ILE cannot evaluate the selected time stencil at arbitrary " + "times; treating --time-posterior-export continuous as its best-available " + "interpolation-lattice export for executable-swap compatibility.") + opts._time_posterior_export = "grid" +if (opts.srate_resample_time_marginalization is not None and + opts.srate_resample_time_marginalization <= 0): + optp.error("--srate-resample-time-marginalization must be positive") # # Failure modes @@ -2738,7 +2757,9 @@ def resample_samples_LISA(my_samples, rholms, cross_terms, right_ascension, decl print(" Time resampling size : {} ".format(n_samples)) # Hardcoded time sampling limits, should change it in future - low_t_lim, high_t_lim, delta_t = -5, 5, 0.001 + low_t_lim, high_t_lim = -5, 5 + delta_t = (1.0 / opts.srate_resample_time_marginalization + if opts.srate_resample_time_marginalization else 0.001) # t_ref_wind is defined as data_window_integration_half @@ -2770,7 +2791,8 @@ def resample_samples_LISA(my_samples, rholms, cross_terms, right_ascension, decl lnLt_norm = scipy.special.logsumexp(lnLt,axis=-1) tvals = identity_convert(tvals) # back to CPU - # Loop over and resample in time, picking index according to weights + # Draw continuously when sub-sample interpolation requested it; retain the + # historical fixed 1 ms LISA export grid as the explicit compatibility path. t_out = np.zeros(n_samples) lnL_out = np.zeros(n_samples) # identifiy max lnL(t) point, and then interpolate around that point to avoid Nan weights. @@ -2833,7 +2855,8 @@ def resample_samples(my_samples, lnLt_norm = scipy.special.logsumexp(lnLt,axis=-1) tvals = identity_convert(tvals) # back to CPU # print(lnLt.shape, lnLt_norm.shape,tvals.shape) - # Loop over and resample in time, picking index according to weights + # Match the main driver's sub-sample export contract on this legacy + # ground-based path retained in the LISA executable. t_out = np.zeros(n_samples) lnL_out = np.zeros(n_samples) indx_list =np.arange(len(tvals)) diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax index c01138823..00be1a6f8 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax @@ -50,11 +50,14 @@ Output (matches ILE conventions): ``lnL`` = log marginal likelihood (evidence) over extrinsic parameters, ``sigma_lnL = sqrt(var)/Z`` (ILE's ``sqrt(var)/res``). Masses in M_sun. __samples.dat (--save-samples) : per-sample extrinsic params + - loglikelihood with ILE-style column names. + loglikelihood with ILE-style column names. Equal-weight FAIR DRAW rows + (no weight column), so a collapsed integration whose weights admit no fair + draw writes NO samples file and fails the event instead. """ from __future__ import print_function +import os import sys from optparse import OptionParser, OptionGroup @@ -68,6 +71,13 @@ import lalsimulation as lalsim import RIFT.lalsimutils as lalsimutils from RIFT.likelihood.jax_ile import build_data_from_precompute +from RIFT.likelihood.jax_ile.wrapper import bandlimited_storage_requirement +from RIFT.likelihood.jax_ile import anglemarg as _anglemarg +from RIFT.likelihood.jax_ile.samplers import angle_marg_eval_chunk as _angle_marg_eval_chunk +from RIFT.likelihood.jax_ile.core import _GATHERERS as _JAX_GATHERERS, JAX_INTERP_DEFAULT +from RIFT.likelihood.jax_ile.anglemarg import (ANGLE_MARG_DEFAULT, ANGLE_MARG_LEGACY, + ANGLE_MARG_CHOICES) +_JAX_GATHERER_NAMES = tuple(_JAX_GATHERERS) from RIFT.likelihood.jax_ile.wrapper import ( JAXExtrinsicLikelihood, JAXDistanceMarginalizedLikelihood, ) @@ -90,14 +100,42 @@ FULL_NAMES = ("ra", "dec", "psi", "incl", "phiref", "distMpc") # here as accepted-but-ignored (with the correct arity so parsing succeeds), and # a small set that would silently change the *science* if ignored is failed on. +# ILE's default cap for --fairdraw-extrinsic-output (integrate_likelihood_ +# extrinsic_batchmode: --fairdraw-extrinsic-output-n-max default=5). +_FAIRDRAW_N_MAX_DEFAULT = 5 +# Below this many effective samples the exported cloud is not a usable posterior +# sample however it is drawn. ONE definition: fairdraw_indices warns at it, the +# tempering guard refuses at it, and --allow-degenerate-tempering's help quotes +# it, so the three cannot drift apart. +_USABLE_EXPORT_ESS = 200 +# Default export-ESS budget for --auto-adapt-weight-exponent. Named so the +# parser default and the "was this actually passed?" check are ONE value. +_TARGET_EXPORT_ESS_FRAC_DEFAULT = 0.9 +# Modes whose sampler actually implements the anneal. samplers.flowmc_sample_phimarg +# takes temper_adapt; samplers.flowmc_sample (plain `flowmc`) does not. +_ADAPT_ADAPT_MODES = frozenset(( + "flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg")) + +# Modes whose sampler reports a TEMPERED state plus a genuine importance weight +# (post_weight = L^(1-inv_T)); only these honour --adapt-weight-exponent, and +# only these consume post_weight at export. +_TEMPERED_MODES = frozenset(( + "flowmc", "flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg")) +# Modes whose --save-samples export is a reweighted FAIR DRAW. Only here do the +# fair-draw count options do anything; elsewhere the export is the sampler's own +# chain and the count flags are inert (and must be reported as ignored). +_FAIRDRAW_MODES = _TEMPERED_MODES | frozenset(("prior-mc", "laplace-is")) + # Boolean (zero-argument) ILE options (action=store_true/false). _ILE_BOOL_OPTS = { + "--calibration-conjugate-phase", "--calibration-global-norm", "--check-good-enough", "--zero-likelihood", "--random-event", "--soft-fail-event-range", "--fmin-template-correct-for-lmax", "--internal-use-gwpy", "--nr-lookup", "--nr-hybrid-use", "--rom-use-basis", "--rom-integrate-intrinsic", "--nr-perturbative-extraction", "--nr-perturbative-extraction-full", "--nr-use-provided-strain", - "--no-memory", "--use-gwsignal", "--maximize-only", "--time-marginalization", + "--no-memory", "--use-gwsignal", "--use-external-EOB", "--maximize-only", + "--dump-lnL-time-series", "--time-marginalization", "--resample-time-marginalization", "--distance-marginalization", "--calibration-fused-kernel", "--calibration-export-posterior", "--extrinsic-proposal-adapt", "--vectorized", "--gpu", "--force-gpu-only", @@ -113,15 +151,32 @@ _ILE_BOOL_OPTS = { "--internal-sky-network-coordinates-raw", "--auto-logarithm-offset", "--pin-distance-to-sim", "--export-eos-index", "--export-marginal-distance-grid", "--adapt-intrinsic", + "--reject-collapsed-live-volume", "--rotation-slow", "--freqresponse", + "--internal-waveform-fd-L-frame", "--save-EOB-parameters", + "--save-hyperbolic", "--force-hyperbolic-22", "--save-meanPerAno", + "--internal-gmm-correlate-all", "--internal-gmm-adaptive-components", + "--portfolio-varaha-never-freeze", "--portfolio-varaha-can-freeze", + "--portfolio-adaptive-alloc", "--sampler-sequential-warmstart", + "--sampler-l0-rescue-accept-truncated", "--sampler-anisotropic-bins", + "--internal-reparam-dl-incl", "--internal-use-lnL", + "--distance-slice-all-fresh", + "--distance-slice-randomize", } # ILE options taking repeated values (action=append). -_ILE_APPEND_OPTS = {"--channel-name", "--psd-file"} +_ILE_APPEND_OPTS = { + "--channel-name", "--psd-file", "--fmin-ifo", "--nr-lookup-group", + "--parameter", "--parameter-range", "--sampler-portfolio", + "--sampler-portfolio-args", +} # The full ILE option set (so anything not implemented is still accepted). _ILE_ALL_OPTS = { "--adapt-adapt", "--adapt-floor-level", "--adapt-intrinsic", "--adapt-log", "--adapt-weight-exponent", "--amp-order", "--approximant", "--auto-logarithm-offset", "--cache-file", "--calibration-burn-in-neff", "--calibration-burn-in-nmax", "--calibration-dump-responsibilities", + "--calibration-conjugate-phase", "--calibration-global-norm", + "--calibration-mc-error-extrinsic", "--calibration-n-realizations-max", + "--calibration-neff-cal-target", "--calibration-envelope-directory", "--calibration-export-posterior", "--calibration-fused-kernel", "--calibration-n-realizations", "--calibration-pilot-extrinsic", "--calibration-proposal-breadcrumb", @@ -129,37 +184,72 @@ _ILE_ALL_OPTS = { "--coinc-xml", "--convergence-tests-on", "--data-end-time", "--data-integration-window-half", "--data-start-time", "--declination-cosine-sampler", "--deff-lambda", "--distance-marginalization", - "--distance-marginalization-lookup-table", "--distance-slice-skip-threshold", + "--distance-marginalization-lookup-table", "--distance-slice-all-fresh", + "--distance-slice-chunk", "--distance-slice-randomize", + "--distance-slice-skip-threshold", "--distance-slice-wing-delta-lnL", "--distance-slice-wing-neff", "--distance-slice-wing-nmax", "--d-max", "--d-min", "--d-prior", "--d-prior-redshift", "--eff-lambda", "--e-freq", - "--event", "--event-time", "--export-distance-slices", "--export-eos-index", + "--dump-lnL-time-series", "--event", "--event-time", + "--export-distance-slices", "--export-eos-index", "--export-marginal-distance-grid", "--extrinsic-proposal-adapt", "--extrinsic-proposal-breadcrumb", "--extrinsic-proposal-output", + "--extrinsic-proposal-field", "--extrinsic-proposal-field-cover-frac", + "--extrinsic-proposal-field-inflate", "--fairdraw-extrinsic-output", "--fairdraw-extrinsic-output-n-max", "--fmax", "--fmin-ifo", "--fmin-template", "--fmin-template-correct-for-lmax", - "--force-adapt-all", "--force-gpu-only", "--force-reset-all", "--force-xpy", - "--gpu", "--inclination-cosine-sampler", + "--force-adapt-all", "--force-gpu-only", "--force-hyperbolic-22", + "--force-reset-all", "--force-xpy", "--freqresponse", + "--freqresponse-arm-length", "--freqresponse-qmax", "--gpu", + "--inclination-cosine-sampler", "--internal-gmm-adaptive-components", + "--internal-gmm-correlate-all", "--internal-gmm-defensive-frac", + "--internal-gmm-inflate", "--internal-gmm-max-components", + "--internal-gmm-phase-components", "--internal-gmm-sky-components", "--internal-data-storage-window-half", "--internal-hard-fail-on-error", "--internal-make-empty-file-on-error", "--internal-precompute-ignore-threshold", "--internal-rotate-phase", "--internal-sky-network-coordinates", "--internal-sky-network-coordinates-raw", "--internal-soft-fail-on-cuda-error", - "--internal-use-gwpy", "--internal-waveform-extra-kwargs", + "--internal-reparam-dl-incl", "--internal-use-gwpy", "--internal-use-lnL", + "--internal-waveform-extra-kwargs", "--internal-waveform-extra-lalsuite-args", "--internal-waveform-fd-no-condition", - "--internal-waveform-taper", "--interpolate-time", "--inv-spec-trunc-time", - "--l-max", "--manual-logarithm-offset", "--mass1", "--mass2", + "--internal-waveform-fd-L-frame", "--internal-waveform-taper", + "--interpolate-time", "--inv-spec-trunc-time", "--l-max", + "--limit-declination", "--limit-distance", "--limit-inclination", "--limit-psi", + "--limit-right-ascension", "--manual-logarithm-offset", "--mass1", "--mass2", + "--mc-error-ess-trigger", "--mc-error-khat-trigger", "--mc-error-replicas", + "--mc-error-sigma-trigger", "--maximize-only", "--n-chunk", "--n-distance-slice-core", "--n-eff", "--n-events-to-analyze", "--n-fairdraw-extrinsic-samples", "--n-max", "--n-distance-slice-wing", "--no-adapt", "--no-adapt-after-first", - "--no-adapt-distance", "--no-memory", "--nr-hybrid-method", "--nr-hybrid-use", - "--nr-index", "--nr-lookup", "--nr-lookup-group", "--nr-params", + "--no-adapt-distance", "--no-memory", "--nf-flow-load", "--nf-flow-save", + "--nr-group", "--nr-hybrid-method", "--nr-hybrid-use", "--nr-index", + "--nr-lookup", "--nr-lookup-group", "--nr-param", "--nr-params", "--nr-perturbative-extraction", "--nr-perturbative-extraction-full", "--nr-use-provided-strain", "--output-file", "--output-format", - "--parameter", "--parameter-range", "--pin-distance-to-sim", "--psd-file", + "--parameter", "--parameter-range", "--pin-distance-to-sim", + "--right-ascension", "--declination", "--psi", "--distance", + "--phi-orb", "--t-ref", "--inclination", + "--portfolio-adaptive-alloc", "--portfolio-alloc-exponent", + "--portfolio-freeze-wt", "--portfolio-grace-iters", + "--portfolio-probe-period", "--portfolio-quality-signal", + "--portfolio-revive-period", "--portfolio-varaha-can-freeze", + "--portfolio-varaha-max-frac", "--portfolio-varaha-min-frac", + "--portfolio-varaha-never-freeze", "--portfolio-weight-clip", "--psd-file", "--psd-window-shape", "--random-event", "--reference-freq", "--resample-time-marginalization", "--restricted-mode-list-file", "--rom-group", "--rom-integrate-intrinsic", "--rom-limit-basis-size-to", - "--rom-param", "--rom-use-basis", "--sampler-method", "--sampler-portfolio", - "--sampler-portfolio-args", "--sampler-xpy", "--save-eccentricity", + "--reject-collapsed-live-volume", "--rom-param", "--rom-use-basis", + "--rotation-n-harmonics", "--rotation-p-max", "--rotation-slow", + "--sampler-anisotropic-bins", "--sampler-l0-rescue-accept-truncated", + "--sampler-l0-rescue-puff-factor", "--sampler-l0-rescue-puff-scale", + "--sampler-l0-rescue-puff-width-frac", "--sampler-l0-rescue-reject-dlnZ", + "--sampler-load-state", "--sampler-method", "--sampler-portfolio", + "--sampler-portfolio-args", "--sampler-save-state", + "--sampler-sequential-warmstart", "--sampler-sequential-warmstart-cover-frac", + "--sampler-sequential-warmstart-deltalnL", "--sampler-warmstart-cover-frac", + "--sampler-warmstart-inflate", "--sampler-warmstart-retry-neff", + "--sampler-warmstart-samples", "--sampler-xpy", "--save-EOB-parameters", + "--save-P", "--save-deltalnL", "--save-eccentricity", "--save-hyperbolic", + "--save-meanPerAno", "--save-samples", "--save-samples-process-params", "--seed", "--sim-grid", "--sim-xml", "--skymap-file", "--soft-fail-event-range", "--spin1z", "--spin2z", "--srate", "--srate-internal", @@ -167,7 +257,8 @@ _ILE_ALL_OPTS = { "--supplementary-likelihood-factor-code", "--supplementary-likelihood-factor-function", "--supplementary-likelihood-factor-ini", "--time-marginalization", - "--use-gwsignal", "--use-gwsignal-lmax-nyquist", "--vectorized", "--verbose", + "--time-posterior-export", "--use-external-EOB", "--use-gwsignal", + "--use-gwsignal-lmax-nyquist", "--vectorized", "--verbose", "--window-shape", "--zero-likelihood", } @@ -220,6 +311,50 @@ def check_critical_and_report(opts, optp): fatal.append("--zero-likelihood is not implemented") if is_set("--maximize-only"): fatal.append("--maximize-only is not implemented (this driver integrates)") + # Distance-grid option combinations that the wrapper would reject anyway -- + # caught HERE, at parse time, so the user is not made to sit through a full + # precompute first (F8 of external review). + dgs = getattr(opts, "distance_grid_scheme", "uniform") + if dgs != "uniform": + if getattr(opts, "mode", None) != "flowmc-phipsimarg": + fatal.append("--distance-grid-scheme %s applies only to --mode " + "flowmc-phipsimarg (it is validated only for the dense " + "angle-marginalization schemes)" % dgs) + elif getattr(opts, "angle_marg_scheme", ANGLE_MARG_DEFAULT) == "grid": + fatal.append( + "--distance-grid-scheme %s requires --angle-marg-scheme " + "exact/laplace/auto: the log-uniform grid is sized from the " + "data-derived angle amplitude, which the default 'grid' scheme " + "does not compute" % dgs) + if getattr(opts, "distance_grid_points", None) is not None: + fatal.append("--distance-grid-points and --distance-grid-scheme %s " + "both set the distance node count; pass one or the " + "other" % dgs) + _tol = getattr(opts, "distance_grid_tol", None) + if _tol is not None and not (0.0 < float(_tol) < 2.0): + # The option's VALUE, checked beside its combinations. The range + # is closed form -- c(tol) = pi*sqrt(2/ln(2/tol)) needs + # 0 < tol < 2 -- so making the user sit through a full precompute + # to be told otherwise is the same avoidable cost F8 removed for + # the combinations. + fatal.append("--distance-grid-tol must be in (0, 2): it is a " + "FRACTIONAL error on the distance integral (~nats on " + "lnL), not a node count; got %r. Note the interval " + "is the DERIVATION's domain, not a range of sensible " + "values: c(tol) diverges as tol -> 2, so 1.999 asks " + "for a 2-node grid. The shipped default is %g." + % (_tol, _jax_core_dist_tol_default())) + if int(os.environ.get("JAX_ILE_DISTMARG_GH", "0")) > 0: + fatal.append( + "--distance-grid-scheme %s cannot be combined with " + "JAX_ILE_DISTMARG_GH: the per-sample Gauss-Hermite distance " + "quadrature uses only the SUPPORT of the grid, so the option " + "would be bit-identically inert while still being reported as " + "active" % dgs) + elif getattr(opts, "distance_grid_tol", None) is not None: + fatal.append("--distance-grid-tol applies only to " + "--distance-grid-scheme loguniform; it would be silently " + "inert here") if fatal: optp.error("Cannot run as a faithful drop-in: " + "; ".join(fatal) + ". (These would silently change the result if ignored.)") @@ -234,14 +369,43 @@ def check_critical_and_report(opts, optp): "--reference-freq", "--fmax", "--srate", "--data-integration-window-half", "--internal-data-storage-window-half", "--d-min", "--d-max", - "--d-prior", "--n-max", "--n-chunk", "--output-file", + "--d-prior", "--limit-distance", + "--n-max", "--n-chunk", "--output-file", "--event", "--save-samples", "--verbose", "--seed", "--sim-xml", "--sim-grid", "--n-events-to-analyze", "--random-event", "--distance-marginalization", - "--time-marginalization", "--vectorized", "--use-gwsignal"} + "--time-marginalization", "--time-marginalization-quadrature", + "--interpolate-time", + "--vectorized", "--use-gwsignal"} + # These are implemented PER MODE. Listing them unconditionally would claim + # they act under --mode laplace-is (the default), nuts, map, multistart-nuts + # and nuts-phimarg, where they are inert -- exactly the silent no-op this + # driver's compat layer exists to prevent. + mode = getattr(opts, "mode", None) + if mode in _TEMPERED_MODES: + # static tempering exponent (samplers.flowmc_sample*: inv_T = 1/temper) + implemented.add("--adapt-weight-exponent") + if mode in _FAIRDRAW_MODES: + implemented |= {"--fairdraw-extrinsic-output", + "--fairdraw-extrinsic-output-n-max", + "--n-fairdraw-extrinsic-samples"} for name in sorted(_ILE_ALL_OPTS - implemented): if is_set(name): ignored.append(name) + # JAX-NATIVE tempering flags. These are not in _ILE_ALL_OPTS (they have no ILE + # counterpart), so the loop above cannot see them -- and they act ONLY on the + # tempered modes. Report them explicitly rather than let a user's chooser + # request evaporate under --mode laplace-is, which is the driver default. + if mode not in _TEMPERED_MODES: + inert = [n for n in ("--auto-adapt-weight-exponent", + "--allow-degenerate-tempering") + if getattr(opts, _dest(n), False)] + if _target_ess_was_given(opts): + inert.append("--target-export-ess-frac") + if inert: + print("Note: %s only act on the tempered modes (%s); --mode %s ignores " + "them." % (" ".join(sorted(set(inert))), + " ".join(sorted(_TEMPERED_MODES)), mode)) if ignored: print("Note: the following ILE options are accepted but IGNORED by the " "JAX driver (not yet implemented; behavior may differ from ILE):") @@ -251,6 +415,15 @@ def check_critical_and_report(opts, optp): # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- +DISTANCE_GRID_POINTS_DEFAULT = 256 + + +def _jax_core_dist_tol_default(): + """Default distance-grid tolerance, read from the module that owns it.""" + from RIFT.likelihood.jax_ile.core import DIST_GRID_TOL_DEFAULT + return DIST_GRID_TOL_DEFAULT + + def build_parser(): optp = OptionParser(usage="%prog [options]", description=__doc__) @@ -335,9 +508,29 @@ def build_parser(): g.add_option("--num-samples", type=int, default=2000, help="NUTS posterior samples per chain (--mode nuts).") g.add_option("--num-chains", type=int, default=1) - g.add_option("--no-flow-reuse", action="store_true", default=False, - help="Disable bootstrapping the trained flow across " - "--n-events-to-analyze (re-train from scratch each event).") + # DEFAULT IS OFF (no re-use). Measured: across an 8-event batch at two seeds, + # re-using the trained flow contracts the extrinsic posterior monotonically in + # slot index -- psi to ~40% of its no-re-use width by slot 7, on BOTH seeds, + # with slot 0 (where no re-use has happened yet) sitting at ~1.0 as a control. + # It buys nothing to pay for that with: 1589 s mean wall with re-use vs 1567 s + # without, a difference smaller than the seed-to-seed spread and of flipping + # sign. See analyses/jax_extrinsic_tempering/ in the paper repo. + # + # --no-flow-reuse is kept (now a no-op restating the default) so existing + # command lines and scripts that pass it keep working; --flow-reuse is the way + # back to the old behaviour. Both write the same dest, so the LAST one on the + # command line wins. + g.add_option("--no-flow-reuse", action="store_true", dest="no_flow_reuse", + default=True, + help="Re-train the flow from scratch each event (the DEFAULT). " + "Kept for compatibility: it now restates the default.") + g.add_option("--flow-reuse", action="store_false", dest="no_flow_reuse", + help="Bootstrap the trained flow across --n-events-to-analyze " + "instead of re-training per event. NOT recommended for any " + "run whose extrinsic SAMPLES are used: it contracts the " + "posterior in later slots (measured ~40% of the no-re-use " + "psi width by slot 7 of 8) for no measurable wall-time " + "saving. Its evidence is less affected than its samples.") # ILE-compatible semantics: store_true, default OFF. When ON, distance is # marginalized analytically (5-D angular problem, well conditioned); when # OFF, distance is sampled explicitly (6-D). Required for --mode nuts. @@ -348,8 +541,65 @@ def build_parser(): g.add_option("--n-chunk", type=int, default=8000) g.add_option("--d-min", type=float, default=1.0, help="Min distance (Mpc).") g.add_option("--d-max", type=float, default=10000.0, help="Max distance (Mpc).") - g.add_option("--distance-grid-points", type=int, default=256) + + g.add_option("--distance-grid-points", type=int, default=None, + help="Node count for the DEFAULT uniform-in-distance grid " + "(default %d). --distance-grid-scheme loguniform " + "derives its own node count from the data, so passing " + "both is REFUSED rather than silently ignoring this one." + % DISTANCE_GRID_POINTS_DEFAULT) + g.add_option("--distance-grid-scheme", type="choice", + choices=("uniform", "loguniform"), default="uniform", + help="Distance quadrature for --mode flowmc-phipsimarg with " + "--angle-marg-scheme exact/laplace/auto. 'uniform' " + "(DEFAULT) is the historical fixed --distance-grid-points " + "grid over the whole prior range, so existing command " + "lines reproduce existing runs. 'loguniform' places the " + "SAME kind of static grid uniformly in ln d, with a node " + "count DERIVED from the run's own data-derived angle " + "amplitude so that every per-sample distance peak " + "(relative width 1/rho) is resolved anywhere in " + "[--d-min, --d-max]. See --distance-grid-tol and " + "RIFT/likelihood/jax_ile/DESIGN_jax_distance_quadrature.md.") + g.add_option("--distance-grid-tol", type=float, default=None, + help="Target FRACTIONAL error of the distance quadrature " + "(~nats on lnL) for --distance-grid-scheme loguniform; " + "the node count follows from it in closed form. " + "Default %g. Only valid with that scheme." + % _jax_core_dist_tol_default()) + g.add_option("--limit-distance", default=None, + help="Restrict distance SAMPLING (and, when distance is " + "marginalized, the distance QUADRATURE) to 'LO,HI' in Mpc, " + "WITHOUT changing the prior or its normalization: the prior " + "stays normalized over the full [--d-min,--d-max], so the " + "reported lnZ needs no correction and stays comparable to a " + "full-range run and to the batchmode ILE -- comparable to the " + "SAMPLER's own systematic, not to machine precision: measured " + "end to end on real data at rho ~ 82, the box truncates <= 0.003 " + "nats of evidence while the observed lnZ difference is +0.37 +- " + "0.11 nats the other way, being the full-range run's own bias " + "that the box removes (RIFT_roboto_paper " + "analyses/limit_distance_e2e/). That benefit is a THRESHOLD " + "in amplitude, not a general improvement: the same " + "measurement at rho 41 gives +0.13 +- 0.08 nats, consistent " + "with zero, because the full-range sampler is already " + "healthy there. Intended for high " + "amplitude, where the distance posterior narrows as 1/rho. " + "Likelihood outside the box is simply not integrated, so keep " + "the box comfortably wider than the posterior. Must lie " + "inside [--d-min,--d-max].") g.add_option("--phase-marginalization", action="store_true", default=False) + g.add_option("--time-marginalization-quadrature", type="choice", + choices=("simpson", "bandlimited"), default="simpson", + help="Rule for the terminal time integral: historical fixed-grid " + "Simpson or adaptive reflected-FFT interpolation followed " + "by a converged trapezoid. No resolution knob is exposed; " + "the factor is derived and rechecked from lnL(t).") + g.add_option("--resample-time-marginalization", action="store_true", + default=False, + help="Conventional ILE option; currently unsupported by JAX ILE.") + g.add_option("--srate-resample-time-marginalization", type="int", default=None, + help="Conventional ILE option; currently unsupported by JAX ILE.") g.add_option("--n-phi", type=int, default=32, help="phi_ref grid size for --mode flowmc-phimarg (default 32; " "use 64-128 for l-max>=4 or production quality).") @@ -357,7 +607,39 @@ def build_parser(): help="psi (polarization) grid size for --mode flowmc-phipsimarg " "(default 8; spin-2 -> exponential convergence, ~8 exact for " "l_max=2). Keep nphi*npsi small (~64) -- the (phi,psi) scan " - "cost scales with it.") + "cost scales with it. CAVEAT: 'exact' above holds for " + "SAMPLING the trig polynomial lnL, not for the quadrature " + "of exp(lnL), whose peak width is ~1/SNR -- the grid " + "scheme's marginalization error grows without bound with " + "SNR (measured ~1e2 nats at SNR 40 for npsi=8). See " + "--angle-marg-scheme for the fix.") + # No type=: optparse infers type="choice" from choices=, and raises + # "must not supply choices for type 'string'" if both are given. --interp + # above omits it for the same reason. + g.add_option("--angle-marg-scheme", default=ANGLE_MARG_DEFAULT, + choices=sorted(ANGLE_MARG_CHOICES), + help="(phi_ref, psi) marginalization scheme for --mode " + "flowmc-phipsimarg. DEFAULT '%s' since 2026-09-02; it " + "was '%s' before, and passing '%s' reproduces a " + "pre-2026-09-02 run exactly. '%s': the historical " + "--n-phi x --n-psi quadrature -- it averages exp(lnL), " + "whose peak width is ~1/SNR, on those nodes, so its " + "error grows WITHOUT BOUND with SNR; at rho 652 it " + "displaced a recovered sky position by 0.53 deg and " + "ranked that artifact ABOVE the injection. Kept for " + "reproducing archived runs, not for new ones. 'exact': " + "Fourier-coefficient bootstrap + dense reconstruction; " + "expensive likelihood calls fixed by MODE CONTENT, never " + "by SNR (cost ~SNR^2). 'laplace': analytic Laplace in " + "psi + dense phi, error O(1/SNR^2), cheaper than 'exact' " + "above rho ~200 -- but it cannot use the per-sample " + "adaptive distance quadrature, so on the default uniform " + "distance grid it was measured 43 nats from 'exact' at " + "rho 163. 'auto': select exact/laplace from the run's " + "own data-derived amplitude. The scheme that actually " + "ran is printed. See RIFT.likelihood.jax_ile.anglemarg." + % (ANGLE_MARG_DEFAULT, ANGLE_MARG_LEGACY, + ANGLE_MARG_LEGACY, ANGLE_MARG_LEGACY)) # flowMC tuning (modes flowmc / flowmc-phimarg). Defaults match # samplers.flowmc_sample*; exposed so pipeline Makefiles can tune them. g.add_option("--n-training-loops", type=int, default=4, @@ -379,6 +661,27 @@ def build_parser(): "reweights to the true posterior. beta=1 (default) is OFF; " "beta<1 broadens (helps the flow find sharp high-SNR peaks). " "Maps to sampler temper = 1/beta (modes flowmc, flowmc-phimarg).") + g.add_option("--auto-adapt-weight-exponent", action="store_true", default=False, + help="EXPERIMENTAL, and not sufficiently validated: the ESS " + "model it inverts is calibrated at SNR ~= 23.8 only and is " + "optimistic at higher SNR, so the fraction it targets is " + "not guaranteed. No paper result uses it. " + "Choose --adapt-weight-exponent automatically from the " + "EXPORT budget instead of by hand: pick the smallest beta " + "whose reweighted export keeps --target-export-ess-frac of " + "the cloud. Keyed on the SAMPLED DIMENSION, not on SNR -- " + "the non-JAX helper's SNR rule does not transfer to this " + "path (jax_ile/DESIGN_jax_tempering.md). Tempered modes only.") + g.add_option("--target-export-ess-frac", type=float, + default=_TARGET_EXPORT_ESS_FRAC_DEFAULT, + help="Fraction of the sampler cloud the reweighted --save-samples " + "export must retain, for --auto-adapt-weight-exponent " + "(default 0.9). ESS/N = [beta(2-beta)]^(dim/2).") + g.add_option("--allow-degenerate-tempering", action="store_true", default=False, + help="Silence the warning issued when the ESTIMATED export ESS " + "is below the %d-sample usability floor. The estimate is " + "advisory (calibrated at one SNR), so this suppresses a " + "warning rather than overriding a refusal." % _USABLE_EXPORT_ESS) g.add_option("--adapt-adapt", action="store_true", default=False, help="Adaptive likelihood tempering (flowmc-phimarg): anneal the " "tempering exponent inv_T from --temper-init up to 1.0, " @@ -420,7 +723,24 @@ def build_parser(): "(no flow training -> immune to the SNR>=640 NF-collapse). " "Overrides the exported samples; TI evidence is unchanged. " "0=off; ~40000 recommended at SNR>=640. Implies --fisher-precondition.") - g.add_option("--interp", default="linear", choices=["linear", "nearest"]) + # Choices come from the gatherer registry rather than a literal list: a stencil added to + # _GATHERERS and not here would be unreachable from the command line, which is how 'cubic' + # spent a release being implemented but unselectable. + g.add_option("--interp", default=JAX_INTERP_DEFAULT, choices=sorted(_JAX_GATHERER_NAMES), + help="Interpolation of the precomputed rholm timeseries in arrival time. " + "DEFAULT CHANGED 2026-08-26 from 'linear' to '" + JAX_INTERP_DEFAULT + "': " + "linear undershoots the rholm peak by an amount that depends on where it " + "falls between samples, which biases the recovered arrival time and hence " + "the sky, and at high SNR -- the only regime this driver is used in -- it " + "is worse than 'nearest'. THIS CHANGES RESULTS for anyone who did not pass " + "--interp; pass '--interp linear' to reproduce a pre-2026-08-26 run. " + "'cubic' and 'sinc' mirror the production factored_likelihood stencils of " + "the same names; which of the two is more accurate depends on how " + "oversampled the rholm timeseries is (on fmin and srate as well as mass) " + "and there is no automatic rule -- see " + "RIFT/likelihood/DESIGN_q_window_stencil.md. COST: 'sinc' is ~1.33x " + "'cubic' in GPU runtime and needs ~12x its XLA scratch (~512 MB against " + "~40 MB at the default --n-chunk 8000; measured, see §9.5).") g.add_option("--sky-coordinates", default="equatorial", choices=["equatorial", "network"], help="Optional: 'network' samples the sky in the two-detector " @@ -434,6 +754,24 @@ def build_parser(): g = OptionGroup(optp, "Output") g.add_option("--output-file", default=None) g.add_option("--save-samples", action="store_true", default=False) + g.add_option("--fairdraw-extrinsic-output", action="store_true", default=False, + help="Bound the fair-draw export to " + "--fairdraw-extrinsic-output-n-max samples (as ILE does). " + "The export is ALWAYS a fair draw when the sampler " + "supplies importance weights; this only caps the count.") + g.add_option("--fairdraw-extrinsic-output-n-max", type=int, default=None, + help="Cap on fair draws per evaluation when " + "--fairdraw-extrinsic-output is set (ILE default 5). " + "Left as None when unset so the ignored-option report does " + "not claim the user passed it; resolved to 5 downstream.") + g.add_option("--n-fairdraw-extrinsic-samples", type=int, default=None, + help="Export exactly this many fair draws, clamped only by the " + "rows available. Those already carry ILE's 1.5*ESS cap " + "wherever the export weights were non-uniform, applied " + "against the EXPORT weights; there is deliberately no " + "second clamp by the evidence estimator's neff, which " + "describes a different cloud. " + "Overrides --fairdraw-extrinsic-output-n-max.") g.add_option("--verbose", action="store_true", default=False) optp.add_option_group(g) @@ -570,6 +908,41 @@ def load_frames(opts, fiducial_epoch): # --------------------------------------------------------------------------- # Priors (the angular block is shared; distance added only when NOT marginalized) # --------------------------------------------------------------------------- +def resolve_distance_limit(opts): + """(lo, hi) in Mpc that distance is actually SAMPLED / integrated over. + + Equals (--d-min, --d-max) unless --limit-distance narrowed it. The prior + normalization always stays on (--d-min, --d-max): see log_prior() and the + d_prior_range= argument threaded into the marginalized likelihoods. + """ + v = getattr(opts, "limit_distance", None) + if not v: + return float(opts.d_min), float(opts.d_max) + from RIFT.integrators.mcsampler import distance_limit_range + try: + return distance_limit_range(v, opts.d_min, opts.d_max) + except ValueError as exc: + raise SystemExit(" --limit-distance: %s" % exc) + + +def log_distance_box_correction(opts, with_distance): + """ln[ prior mass on (d_min,d_max) / prior mass on the sampled box ]; 0.0 unless + --limit-distance narrowed the box. + + sample_prior() draws distance from the prior RESTRICTED to the box, so its + density is the prior divided by that mass. Estimators that assume + "proposal == prior" (run_prior_mc) must subtract this; estimators that form + ln w = lnL + ln p - ln q explicitly (run_laplace_is) already have it right + and must NOT subtract it again. + """ + if not with_distance: + return 0.0 + lo, hi = resolve_distance_limit(opts) + if lo == float(opts.d_min) and hi == float(opts.d_max): + return 0.0 + return float(np.log((opts.d_max ** 3 - opts.d_min ** 3) / (hi ** 3 - lo ** 3))) + + def sample_prior(n, opts, rng, with_distance): ra = rng.uniform(0.0, 2 * np.pi, n) dec = np.arcsin(rng.uniform(-1.0, 1.0, n)) @@ -579,7 +952,9 @@ def sample_prior(n, opts, rng, with_distance): cols = [ra, dec, psi, incl, phiref] if with_distance: u = rng.uniform(0.0, 1.0, n) - dmin3, dmax3 = opts.d_min ** 3, opts.d_max ** 3 + # the SAMPLED range (== [d_min,d_max] unless --limit-distance) + d_lo, d_hi = resolve_distance_limit(opts) + dmin3, dmax3 = d_lo ** 3, d_hi ** 3 cols.append((dmin3 + u * (dmax3 - dmin3)) ** (1.0 / 3.0)) theta = np.stack(cols, axis=-1) return theta, log_prior(theta, opts, with_distance) @@ -596,7 +971,11 @@ def log_prior(theta, opts, with_distance): - np.log(2 * np.pi) - np.log(np.pi) - np.log(2 * np.pi)) if with_distance: dist = theta[..., 5] - inb = inb & (dist >= opts.d_min) & (dist <= opts.d_max) + # SUPPORT is the sampled box (--limit-distance); the NORMALIZATION is + # always over the physical [d_min,d_max], so lnZ keeps its full-range + # scale and does not have to be corrected before comparison. + d_lo, d_hi = resolve_distance_limit(opts) + inb = inb & (dist >= d_lo) & (dist <= d_hi) dmin3, dmax3 = opts.d_min ** 3, opts.d_max ** 3 logp = logp + np.log(3.0) + 2 * np.log(dist) - np.log(dmax3 - dmin3) return np.where(inb, logp, -np.inf) @@ -608,10 +987,21 @@ def log_prior(theta, opts, with_distance): def eval_lnL(like, theta, opts, with_distance): N = theta.shape[0] out = np.empty(N) - for i in range(0, N, opts.n_chunk): - sl = slice(i, min(i + opts.n_chunk, N)) + # anglemarg schemes multiply the batch by (quad, dist_block, phi_chunk) + # transients; --n-chunk 8000 is sized for the grid scheme. Independent + # slices, so the cap changes peak memory and nothing else. + chunk = _angle_marg_eval_chunk(like, opts.n_chunk) + for i in range(0, N, chunk): + sl = slice(i, min(i + chunk, N)) cols = [theta[sl, j] for j in range(theta.shape[1])] out[sl] = np.asarray(like.log_likelihood(*cols)) + if (getattr(like, "time_quadrature", "simpson") == "bandlimited" + and np.any(np.isnan(out[sl]))): + raise RuntimeError( + "adaptive reflected-FFT time marginalization failed its width/" + "doubling convergence or endpoint-mass check; no coarse " + "likelihood is substituted. Increase the input/rholm sample " + "rate or integration/storage window.") return out @@ -658,8 +1048,15 @@ def _moment_match(theta, logL): def run_prior_mc(like, opts, rng, dim, with_distance): theta, _ = sample_prior(opts.n_max, opts, rng, with_distance) lnL = eval_lnL(like, theta, opts, with_distance) - logZ, sig, neff = evidence_from_logweights(lnL) # draw from prior -> w = L - return logZ, sig, neff, opts.n_max, theta, lnL + # p_s == p (the proposal IS the prior), so ln w = lnL -- EXCEPT under + # --limit-distance, where the proposal is the prior restricted to the box while + # the prior keeps its full-range normalization; the constant below is exactly + # 0.0 without that option, so the historical path is untouched. + logw = lnL - log_distance_box_correction(opts, with_distance) + logZ, sig, neff = evidence_from_logweights(logw) + # These raw draws are PRIOR samples, not posterior ones: they must be + # fair-drawn before export. + return logZ, sig, neff, opts.n_max, theta, lnL, logw def run_laplace_is(like, opts, rng, dim, with_distance, n_adapt=2): @@ -692,7 +1089,9 @@ def run_laplace_is(like, opts, rng, dim, with_distance, n_adapt=2): theta = np.concatenate(all_theta); logw = np.concatenate(all_logw) lnL = np.concatenate(all_lnL) logZ, sig, neff = evidence_from_logweights(logw) - return logZ, sig, neff, len(theta), theta, lnL + # theta follows the GAUSSIAN PROPOSAL q, not the posterior; logw is what + # turns it into one. Returned so write_samples() can fair-draw. + return logZ, sig, neff, len(theta), theta, lnL, logw def run_nuts(like, opts, rng, with_distance): @@ -773,7 +1172,13 @@ def run_nuts(like, opts, rng, with_distance): lnL_is[valid] = eval_lnL(like, th_is[valid], opts, with_distance=False) logw = np.where(valid, lnL_is + logp - logq, -np.inf) logZ, sig, neff = evidence_from_logweights(logw) - return logZ, sig, neff, n_is, theta, lnL + # theta/lnL are the NUTS chain (already targets the posterior); the IS cloud + # th_is/logw is only the evidence estimator, so there is nothing to reweight. + # `neff` below therefore describes th_is, NOT the exported chain -- it must + # never reach the export count. It used to: fairdraw_size clamped the export + # by 1.5*neff, truncating this chain by an unrelated cloud's ESS on what is + # the driver's DEFAULT mode. write_samples no longer accepts a neff at all. + return logZ, sig, neff, n_is, theta, lnL, None def run_map(like, opts, rng, dim, with_distance): @@ -781,7 +1186,7 @@ def run_map(like, opts, rng, dim, with_distance): from scipy.optimize import minimize ang_bounds = [(0, 2 * np.pi), (-np.pi / 2 + 1e-3, np.pi / 2 - 1e-3), (0, np.pi), (1e-3, np.pi - 1e-3), (0, 2 * np.pi)] - bounds = ang_bounds + ([(opts.d_min, opts.d_max)] if with_distance else []) + bounds = ang_bounds + ([resolve_distance_limit(opts)] if with_distance else []) theta_seed, _ = sample_prior(4000, opts, rng, with_distance) lnL_seed = eval_lnL(like, theta_seed, opts, with_distance) x0 = theta_seed[np.argmax(lnL_seed)] @@ -798,53 +1203,586 @@ def run_map(like, opts, rng, dim, with_distance): # --------------------------------------------------------------------------- # Output # --------------------------------------------------------------------------- -def write_dat(opts, P, out_index, event_id, logZ, sigma_over_Z, ntotal, neff): +def dat_path(opts, out_index): + """Path of the normal per-event ILE result row (evidence, sigma, neff).""" + return opts.output_file + "_" + str(out_index) + "_" + ".dat" + + +def samples_path(opts, out_index): + """Path of this event's exported extrinsic samples.""" + return opts.output_file + "_" + str(out_index) + "_samples.dat" + + +def angle_grid_suspect_note(scheme=None): + """Label describing the angle-grid amplitude check for this event. + + Returns one of three things, and the THIRD is the point: + + "" -- the grid schemes were not used + "SUSPECT-ANGLE-GRID ..." -- undersizing was DETECTED + "ANGLE-GRID-CHECK=BEST-EFFORT" -- schemes used, nothing detected + + The third case exists because absence of a detection is NOT evidence of + adequacy. The detector is a jax.debug.callback, and JAX explicitly permits + such callbacks to be dropped under transformation -- in which case the host + state stays clean, effects_barrier has nothing to wait for, and the artifact + would otherwise be published looking verified. That is a scientific false + negative, and calling it "best effort" in a docstring does not fix it for a + consumer reading the file six months later. + + So every artifact produced by the exact/laplace schemes carries a standing + statement that this check CANNOT distinguish an adequate grid from an + undetected undersizing. A reader is then never entitled to infer + verification from silence. The honest recourse, named in the artifact, is + to rebuild at a larger amp_sizing if the result matters. + """ + st = _anglemarg.amp_failsafe_state() + if st.get("tripped"): + return ("SUSPECT-ANGLE-GRID amp_failsafe=TRIPPED worst_amp=%.6g " + "amp_sizing=%.6g scheme=%s" + % (st.get("worst_amp", float("nan")), + st.get("amp_sizing", float("nan")), st.get("scheme"))) + # 'peak-local' belongs here for the same reason it belongs in the failsafe: its phi + # axis is dense and amp-sized, so its artifacts are entitled to no more confidence + # than the other two, and a scheme missing from this list would publish output with + # NO standing label at all -- the silence a reader would read as verification. + if scheme in ("exact", "laplace", "peak-local"): + return ("ANGLE-GRID-CHECK=BEST-EFFORT (no undersizing detected; the " + "detector may be dropped under jax transformation, so this is " + "NOT a verification -- rebuild at larger amp_sizing if it matters)") + return "" + + +def write_dat(opts, P, out_index, event_id, logZ, sigma_over_Z, ntotal, neff, + angle_note=""): if not opts.output_file: return m1, m2 = P.m1 / MSUN, P.m2 / MSUN - fname = opts.output_file + "_" + str(out_index) + "_" + ".dat" + fname = dat_path(opts, out_index) row = np.array([[event_id, m1, m2, P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z, logZ, sigma_over_Z, ntotal, neff]]) - np.savetxt(fname, row, - header="event_id m1 m2 s1x s1y s1z s2x s2y s2z lnL sigma_lnL ntotal neff") + hdr = "event_id m1 m2 s1x s1y s1z s2x s2y s2z lnL sigma_lnL ntotal neff" + # The EVIDENCE artifact must carry the label independently of sample export: + # write_samples() early-returns without --save-samples/--output-file, so a + # run with export disabled would otherwise publish a numeric evidence row + # indistinguishable from a clean integration. + if angle_note: + hdr += "\n" + angle_note + np.savetxt(fname, row, header=hdr) print("Wrote %s" % fname) -def write_samples(opts, out_index, theta, lnL, with_distance): +def tempered_cloud_size(opts, n_starts): + """Rows flowMC hands to the exporter: chains x loops x (local + global) steps. + + Verified against a real run: 20 chains, 6 production loops, 20 local + 20 + global steps -> 4800, the ntotal the driver reported. + """ + n_chains = max(n_starts, 20) + return int(n_chains * opts.n_production_loops + * (opts.n_local_steps + opts.n_global_steps)) + + +def record_supplied_options(opts, argv, parser=None): + """Record which long option TOKENS appeared on the command line. + + Everything below used to infer "did the user pass this?" from "does its value + differ from the default?". That is wrong whenever the user explicitly passes + the default: `--auto-adapt-weight-exponent --adapt-weight-exponent 1.0` was + silently ACCEPTED and the chooser then replaced the user's explicit untempered + target, which is the opposite of the documented conflict behaviour. Handles + both `--opt value` and `--opt=value`. + """ + supplied = set() + for tok in (argv if argv is not None else sys.argv[1:]): + if not (isinstance(tok, str) and tok.startswith("--")): + continue + name = tok.split("=", 1)[0] + # CANONICALISE. optparse accepts unambiguous long-option PREFIXES, so + # `--adapt-weight-exp=1.0` parses and sets the value while recording a + # token that no was_supplied() query matches -- auto then silently + # overwrote the user's explicit exponent. Ask the parser what the token + # actually resolved to. Reported by review on #186. + if parser is not None: + try: + name = parser._match_long_opt(name) + except Exception: + pass # unknown/ambiguous: record it verbatim + supplied.add(name) + opts._supplied_options = supplied + return supplied + + +def was_supplied(opts, flag): + """True when ``flag`` was named on the command line. + + Falls back to False when the record is absent (a caller that built an options + object directly rather than parsing), which is the safe direction: it means + "assume the user did not pass it" and so never fabricates a conflict. + """ + return flag in getattr(opts, "_supplied_options", set()) + + +def resolve_ile_interface_aliases(opts, optp): + """Resolve conventional ILE spellings into JAX-native option values.""" + if getattr(opts, "interpolate_time", None) is not None: + ile_interp = str(opts.interpolate_time).strip().lower() + # Conventional ILE accepted this as a boolean before it accepted the + # explicit stencil names. Production args_ile.txt files containing + # either spelling must remain swappable into this driver. + if ile_interp in ("1", "true", "t", "yes", "y", "on"): + ile_interp = "cubic" + elif ile_interp in ("0", "false", "f", "no", "n", "off", "none"): + ile_interp = "nearest" + if ile_interp not in _JAX_GATHERER_NAMES: + optp.error("--interpolate-time must be one of %s or a legacy boolean" % + ", ".join(sorted(_JAX_GATHERER_NAMES))) + if was_supplied(opts, "--interp") and opts.interp != ile_interp: + optp.error("--interp %r and --interpolate-time %r disagree" % + (opts.interp, ile_interp)) + opts.interp = ile_interp + return opts + + +def _normalize_interpolate_time_argv(argv): + """Give conventional ILE's historical bare flag its ``True`` value.""" + raw = list(sys.argv[1:] if argv is None else argv) + out = [] + for i, token in enumerate(raw): + out.append(token) + if (token == "--interpolate-time" and + (i + 1 == len(raw) or raw[i + 1].startswith("--"))): + out.append("True") + return out + + +def _target_ess_was_given(opts): + """True when --target-export-ess-frac was named on the command line.""" + return was_supplied(opts, "--target-export-ess-frac") + + +def resolve_tempering_exponent(opts, n_dim, n_cloud): + """Settle --adapt-weight-exponent for a tempered mode, and refuse a dead one. + + Two jobs, both of which exist because beta means something different here + than it does in non-JAX RIFT (jax_ile/DESIGN_jax_tempering.md): + + * ``--auto-adapt-weight-exponent`` picks beta from the EXPORT budget via + ``samplers.beta_for_export_ess``, i.e. from the sampled dimension. There + is no SNR term, deliberately: the reweight cost + ``ESS/N = [beta(2-beta)]^(n_dim/2)`` is independent of lnLmax, so the + non-JAX helper's SNR-keyed rule would be keying on the wrong variable. + * whatever beta ends up in force, the predicted export ESS is REPORTED, and + a beta that cannot support a usable export RAISES. Printing a warning + above a 199-row file is the silent degradation this exists to remove. + + ``--adapt-adapt`` anneals inv_T up to 1 and therefore exports at full ESS; + it is left alone here. + + RETURNS the exponent to use and does NOT write it back to ``opts``. It used + to assign ``opts.adapt_weight_exponent``, which gave per-RUN state a per-EVENT + meaning: ``analyze_one`` is called once per intrinsic template with the SAME + opts, so on event 1 the chooser read its own event-0 output as a + user-supplied exponent and aborted the batch with SystemExit. Reproduced + with ``--n-events-to-analyze 3``; ILE_extr.sub runs batches. + """ + from RIFT.likelihood.jax_ile.samplers import ( + beta_for_export_ess, export_ess_estimate) + + if opts.adapt_adapt and opts.mode not in _ADAPT_ADAPT_MODES: + # samplers.flowmc_sample (the plain 5-D `flowmc` mode) has NO temper_adapt + # argument -- it takes `temper` and samples a STATIC beta. Returning the + # "annealed to beta=1, full ESS" answer here reported annealing that never + # happened AND skipped the degenerate-export check, so + # `--mode flowmc --adapt-adapt --adapt-weight-exponent 0.1` sampled at + # beta=0.1 unguarded while claiming full ESS. Reported by review on #186. + raise SystemExit( + "--adapt-adapt is not implemented for --mode %s: its sampler " + "(samplers.flowmc_sample) takes a STATIC tempering exponent and has " + "no annealing path, so the flag would be silently ignored while the " + "run reported annealing. Use one of %s, or drop --adapt-adapt and " + "set --adapt-weight-exponent explicitly." + % (opts.mode, " ".join(sorted(_ADAPT_ADAPT_MODES)))) + + if opts.adapt_adapt: + if opts.auto_adapt_weight_exponent: + raise SystemExit( + "--auto-adapt-weight-exponent and --adapt-adapt both set. The " + "anneal already finishes at beta=1 (full export ESS), so there " + "is no static exponent for the chooser to pick. Use one.") + print("Tempering: --adapt-adapt (anneal inv_T -> 1); export is untempered, " + "full ESS.") + return float(opts.adapt_weight_exponent) + + # --smc-puffball routes these modes to samplers.smc_puffball_sample, whose + # signature swallows `temper` in **_ignore: it runs its OWN adaptive SMC + # ladder, so --adapt-weight-exponent has NO effect there. Guarding it would + # REFUSE a run over a number that does nothing. Report the no-op instead; a + # silent one is what this driver's compat layer exists to prevent. + # + # (That sampler now reports the inv_T its ladder actually reached and a + # matching L^(1-inv_T) post_weight -- so its export weights are uniform only + # because analyze_one REFUSES the event unless the ladder finished at + # inv_T=1. The exponent being inert here does not depend on that, but do not + # re-derive "its weights are always uniform" from this branch.) + if getattr(opts, "smc_puffball", False): + if opts.auto_adapt_weight_exponent or was_supplied(opts, "--adapt-weight-exponent"): + print("Note: --smc-puffball ignores --adapt-weight-exponent / " + "--auto-adapt-weight-exponent (smc_puffball_sample runs its own " + "SMC temperature ladder; the event is refused unless that " + "ladder finishes at inv_T=1).") + return 1.0 + + if not opts.auto_adapt_weight_exponent and _target_ess_was_given(opts): + # Setting a target and no chooser does nothing at all. Say so rather than + # let the request evaporate -- the same reason the ILE compat layer reports + # accepted-but-ignored options. + print("Note: --target-export-ess-frac %g has no effect without " + "--auto-adapt-weight-exponent." % opts.target_export_ess_frac) + + if opts.auto_adapt_weight_exponent: + # An explicit --adapt-weight-exponent alongside --auto is a contradiction. + # Silently overriding it would be the worst of both: the run reports a + # chooser it did not obey the user about. + if was_supplied(opts, "--adapt-weight-exponent"): + raise SystemExit( + "--auto-adapt-weight-exponent was given together with an explicit " + "--adapt-weight-exponent %g. The chooser would overwrite it. Pass " + "one or the other. (Detected from the command line, not from the " + "value: passing the default explicitly, --adapt-weight-exponent 1.0, " + "is still an explicit choice of an untempered target.)" + % float(opts.adapt_weight_exponent)) + beta = beta_for_export_ess(opts.target_export_ess_frac, n_dim) + print("Tempering: AUTO beta=%.5f for %.0f%% export ESS in %d-D " + "(ESS/N=[beta(2-beta)]^(dim/2); no SNR term -- see " + "jax_ile/DESIGN_jax_tempering.md)" + % (beta, 100.0 * opts.target_export_ess_frac, n_dim)) + else: + beta = float(opts.adapt_weight_exponent) + + if beta > 1.0: + # NOT harmless, and NOT "untempered": samplers.flowmc_sample* take + # temper = 1/beta, so beta>1 gives inv_T>1 -- it SHARPENS the target past + # the posterior, and the export reweight L^(1-beta) has + # ESS/N = [beta(2-beta)]^(dim/2), which is 0 at beta=2 and undefined + # beyond. An earlier version of this branch tested `beta >= 1.0` and + # printed "beta=1 (untempered target)" here, which was false for every + # beta > 1. + raise SystemExit( + "--adapt-weight-exponent %g is greater than 1. On this path the " + "exponent is applied to the SAMPLED target (inv_T = %g), so a value " + "above 1 samples a target sharper than the posterior and the export " + "reweight L^(1-beta) diverges. Use beta in (0, 1]." % (beta, beta)) + if beta <= 0.0: + raise SystemExit( + "--adapt-weight-exponent %g must be positive: beta=0 samples the " + "prior and the export reweight carries the entire likelihood." + % beta) + if beta == 1.0: + print("Tempering: beta=1 (untempered target); export ESS is the full cloud.") + return beta + # An ESTIMATE, calibrated at SNR ~= 23.8 -- NOT a bound. This used to REFUSE + # a run below the floor, which is a guarantee the number cannot support: the + # SNR ladder measures 0.00823 at beta=0.1/d=4/SNR~67 against an estimate of + # 0.0285. Refusing on it would be false precision in one direction and false + # confidence in the other, so it now WARNS. Reported by review on #186. + frac = export_ess_estimate(beta, n_dim) + ess = frac * n_cloud + print("Tempering: beta=%.5f in %d-D -> export ESS/N >= %.4f " + "(measured-calibrated lower bound), " + "ESS~%.0f of %d rows" % (beta, n_dim, frac, ess, n_cloud)) + if ess < _USABLE_EXPORT_ESS and not opts.allow_degenerate_tempering: + print( + " WARNING: --adapt-weight-exponent %g gives an ESTIMATED export ESS " + "of %.0f (< %d), so the reweighted --save-samples cloud is unlikely " + "to be a usable posterior sample.\n" + " This is the trap the non-JAX helper's rule sets here: it picks " + "beta from the SNR (beta=0.1 at SNR<=22.5, 0.1*(22.5/SNR)^2 above), " + "which is correct where beta only shapes a PROPOSAL, but on this " + "path beta is the exponent of the SAMPLED target.\n" + " The estimate is calibrated at SNR ~= 23.8 and is OPTIMISTIC at " + "higher SNR (measured 0.00823 vs an estimated 0.0285 at beta=0.1, " + "d=4, SNR ~= 67), so the true figure may be several times worse. " + "It is a warning rather than a refusal precisely because it cannot " + "support a hard floor; check the ESS the export actually reports." + % (beta, ess, _USABLE_EXPORT_ESS), file=sys.stderr) + return beta + + +def fairdraw_indices(logw, rng): + """Indices that turn a WEIGHTED cloud into an equal-weight one, or ``None``. + + Production ILE convention (``RIFT/integrators/mcsampler.py::integrate`` and + the identical block in mcsamplerGPU / mcsamplerAdaptiveVolume / + mcsamplerEnsemble / mcsamplerPortfolio): normalise ``w = L * p / p_s`` and + multinomial-resample WITH replacement, so the exported rows are equal weight + and carry no weight column. Capped at ``1.5*ESS``: never claim more fair + draws than the weights support, or the file looks like N independent draws + while holding ~ESS distinct points. + + This function does REWEIGHTING ONLY. The export count requested by + ``--n-fairdraw-extrinsic-samples`` / ``--fairdraw-extrinsic-output-n-max`` + is applied by the caller, on every path -- including this one's ``None`` + returns -- so that the count contract has exactly one implementation and + cannot be quietly skipped for some configurations. + + Returns ``(indices_or_None, note)``; the note always records the ESS. A + ``FAILED`` note means the cloud CANNOT be fair-drawn: ``write_samples`` + refuses to write an export at all in that case, rather than passing the + unreweighted cloud off as one. + """ + logw = np.asarray(logw, dtype=float) + fin = np.isfinite(logw) + # FAIL LOUDLY, NOT OPEN. "cannot compute weights" and "weights are already + # uniform" both used to return None, so a degenerate weight vector silently + # wrote the UNCORRECTED cloud under a header that promises a fair draw. + if fin.sum() < 2: + return None, ("FAILED: %d of %d weights are finite ESS=n/a n_in=%d" + % (int(fin.sum()), len(logw), len(logw))) + lw = logw[fin] - np.max(logw[fin]) + w = np.exp(lw) + tot = w.sum() + if not np.isfinite(tot) or tot <= 0: + return None, ("FAILED: weight sum is %r (overflow or all-zero) " + "ESS=n/a n_in=%d" % (tot, len(logw))) + w = w / tot + neff = 1.0 / np.sum(w ** 2) + idx_fin = np.where(fin)[0] + if np.allclose(w, w[0]): + # Already equal weight -- the DEFAULT for the flowMC modes, since + # --adapt-weight-exponent defaults to 1. Nothing to reweight; the + # caller still applies any requested count. + # NOT an effective sample size: the weights carry no information about + # how correlated the underlying chain is, so reporting the row count as + # "ESS" would be a number that looks measured and is not. + return None, "none (weights uniform) ESS=n/a n_in=%d" % (len(logw),) + n_out = int(max(1, min(int(np.ceil(1.5 * neff)), len(logw)))) + print(" fairdraw: %d weighted samples (ESS=%.1f) -> %d equal-weight draws" + % (len(logw), neff, n_out)) + if neff < _USABLE_EXPORT_ESS: + print(" fairdraw: WARNING ESS=%.1f -- the proposal barely covers this " + "posterior; the exported cloud is NOT a usable posterior sample " + "however it is drawn." % neff, file=sys.stderr) + note = "reweighted ESS=%.1f n_in=%d" % (neff, len(logw)) + return idx_fin[rng.choice(len(idx_fin), size=n_out, replace=True, p=w)], note + + +def fairdraw_size(opts, n_have): + """Requested number of fair draws, or ``None`` for "as many as the weights + support". + + ``--n-fairdraw-extrinsic-samples`` is an exact request; + ``--fairdraw-extrinsic-output-n-max`` caps the count per evaluation. Both + are clamped by ``n_have`` and by nothing else. + + WHY THERE IS NO ESS TERM HERE, though ILE appears to have one. + ``mcsampler.integrate`` clamps with + ``n_extr = min(n_extr, 1.5*eff_samp, 1.5*neff)``, but there ``eff_samp``, + ``neff`` and the fair draw are all properties of the SAME importance-weight + vector ``wt`` that the draw then samples from -- one weight stream, so the + clamp is self-consistent. + + That does not hold on this path. For every flowMC mode ``theta`` is the + production chain while ``res["neff"]`` comes from a SEPARATE estimator: the + moment-matched Gaussian importance cloud built for the evidence, or (under + ``--adapt-adapt``) the annealing ladder's minimum inter-stage ESS. Neither + describes the exported rows. Clamping the export by it truncated a + perfectly good equal-weight chain to a couple of rows whenever that + unrelated proposal happened to have low ESS. + + The ESS that DOES bound this export is the one ``fairdraw_indices`` + computes from the export weights -- and it has already applied it, as + ``n_out = min(ceil(1.5*ESS), n_in)``. So ``n_have`` already carries the + 1.5*ESS cap wherever weights were non-uniform, and where they were uniform + there is no computed chain ESS to clamp by (an autocorrelation-based one + would be a different, unimplemented quantity). Reported by review on #180.""" + n_req = getattr(opts, "n_fairdraw_extrinsic_samples", None) + if n_req is None and getattr(opts, "fairdraw_extrinsic_output", False): + # ILE's default cap is 5; kept out of the parser so an unset flag is not + # reported as one the user passed. + n_req = getattr(opts, "fairdraw_extrinsic_output_n_max", None) + if n_req is None: + n_req = _FAIRDRAW_N_MAX_DEFAULT + if n_req is None: + return None + n_asked = int(n_req) + n_req = max(1, min(n_asked, n_have)) + if n_req != n_asked: + # ILE prints "Fairdraw size : n" whenever it clamps; silence here meant a + # laplace-is run quietly turned a request for 137 into 32. + print(" Fairdraw size : %d (requested %d, clamped by the %d available " + "rows; those already carry the 1.5*ESS cap from fairdraw_indices " + "when the weights were non-uniform)" % (n_req, n_asked, n_have)) + return n_req + + +def _remove_stale_artifact(path, what="export"): + """Delete a leftover file at a path this failed event refuses to publish. + + The run was going to overwrite ``path``; if it is left behind, a re-run in + an existing output directory leaves the PREVIOUS run's file where the + pipeline looks for this one's, which is the same "read a file that is not + this event's" failure the refusal exists to prevent.""" + if not os.path.exists(path): + return + try: + os.remove(path) + except OSError as e: + print(" *** could not remove the stale %s %s (%s) -- it is NOT " + "this run's output. ***" % (what, path, e), file=sys.stderr) + else: + print(" removed the stale %s %s (this event produced no valid result)" + % (what, path), file=sys.stderr) + + +def write_samples(opts, out_index, theta, lnL, with_distance, logw=None, + angle_note=""): + """Write the exported extrinsic samples. + + ``logw`` are per-sample LOG IMPORTANCE WEIGHTS ``ln(L p / p_s)`` for the + rows of ``theta`` (``None`` when the sampler already targets the posterior, + e.g. an untempered MCMC chain). When they are non-uniform the cloud is + fair-drawn against them BEFORE writing, so the exported rows are equal + weight -- the same contract production ILE's ``--fairdraw-extrinsic-output`` + provides, and the one every downstream consumer of these files assumes. + + If the fair draw CANNOT be performed (degenerate or unnormalizable weights) + no file is written and ``RuntimeError`` is raised: the alternative is to ship + the raw proposal/prior cloud under the product name that means "posterior + draws". ``analyze_one`` calls this BEFORE ``write_dat``, so that refusal + also leaves the event without a normal ILE result row. + """ if not (opts.output_file and opts.save_samples) or theta is None: return + sname = samples_path(opts, out_index) + # The export RNG is derived here and NOWHERE ELSE. It must never be the + # generator that feeds the samplers/estimators: --save-samples is an OUTPUT + # flag and consuming the science stream made it change the lnL/logZ of every + # later event in a batch. There is deliberately no rng parameter, so that + # mistake cannot be reintroduced by a caller. + rng = np.random.default_rng((opts.seed, out_index)) + # n_in is the cloud as the sampler handed it over, before any filtering or + # counting, so the header records what went in as well as what came out. + note = ("not applicable (sampler targets the posterior) ESS=n/a n_in=%d" + % len(theta)) + if logw is not None and len(logw) != len(theta): + # FAIL LOUDLY, NOT OPEN. This used to be a silent `and len(...) ==` in + # the guard below, so a mismatched weight vector fell through to the + # "not applicable (sampler targets the posterior)" note -- writing an + # UNREWEIGHTED tempered cloud under a header claiming none was needed. + # A length mismatch is a sampler bug, not a configuration; say so. + raise RuntimeError( + "export weights and samples disagree in length (%d weights, %d " + "rows) for --mode %s. Refusing to write: the fair draw cannot be " + "performed and labelling the raw cloud 'not applicable' would be a " + "false provenance line." + % (len(logw), len(theta), opts.mode)) + if logw is not None: + idx, note = fairdraw_indices(logw, rng) + if idx is not None: + theta, lnL = theta[idx], np.asarray(lnL)[idx] + elif note.startswith("FAILED"): + # NO ARTIFACT WHEN THE FAIR DRAW FAILS. Writing the raw cloud here + # put proposal/prior draws into the standard, weightless + # `*_samples.dat` product under the ONE name every consumer reads as + # equal-weight posterior draws -- and none of them is obliged to read + # the provenance line -- so the contract broke precisely on the + # collapsed integrations where the difference is largest. Refuse to + # produce the file at all: the event fails loudly (and + # --soft-fail-event-range still skips to the next one). + # + # The `.dat` result row goes too. analyze_one validates the export + # BEFORE write_dat, so this run has not written one -- but a file + # left at that path by an earlier run would be collected as this + # event's successful integration, which is the same stale-artifact + # failure as for the samples file. + _remove_stale_artifact(sname, "export") + _remove_stale_artifact(dat_path(opts, out_index), "result") + raise RuntimeError( + "fair draw failed for output index %d (%s): the exported cloud " + "would be the raw, UNREWEIGHTED sampler/proposal samples, which " + "is not a posterior sample. No %s written. This integration " + "collapsed -- fix the run (more samples, better proposal, " + "sanity-check the likelihood) rather than reading the cloud." + % (out_index, note, sname)) + # Drop non-finite lnL FIRST. Doing it last meant the count was applied to + # rows that were then discarded, and the provenance n_out counted them: a + # header saying n_out=1000 above a 963-row file, and 137 above 129 -- wrong + # exactly when the likelihood misbehaved, which is when someone reads it. + lnL = np.asarray(lnL) good = np.isfinite(lnL) + if not good.all(): + theta, lnL = theta[good], lnL[good] + # THE count contract, applied once. Gated on _FAIRDRAW_MODES -- the SAME set + # check_critical_and_report uses to decide whether to report these flags as + # ignored, so report and behaviour cannot disagree. Applying it everywhere + # meant `--mode nuts --fairdraw-extrinsic-output` printed "IGNORED" and then + # silently wrote 5 rows instead of 300; --fairdraw-extrinsic-output is in + # ILE_extr.sub, so that is a real production command line losing 60x of its + # export under a banner saying the flag did nothing. + # Rows are equal weight by this point, so subsample WITHOUT replacement -- a + # random subset of an equal-weight cloud is still a fair draw (verified) and + # manufactures no duplicates. + if opts.mode in _FAIRDRAW_MODES: + # NO evidence-neff argument, deliberately: this function has no parameter + # that could carry one. See fairdraw_size for why an ESS term here was + # wrong on the flowMC modes. + n_req = fairdraw_size(opts, len(theta)) + if n_req is not None and n_req < len(theta): + n_before = len(theta) + sub = rng.choice(n_before, size=int(n_req), replace=False) + theta, lnL = theta[sub], lnL[sub] + print(" fairdraw: exporting %d of %d rows (requested count)" + % (int(n_req), n_before)) + provenance = "fairdraw: %s n_out=%d" % (note, len(theta)) + if angle_note: + provenance += " " + angle_note + phase_marginalized = bool(getattr(opts, "phase_marginalization", False)) + if phase_marginalized: + provenance += " phase=analytically-marginalized" + ndim = theta.shape[1] if theta.ndim == 2 else len(theta) if with_distance: # 6-D: ra, dec, psi, incl, phiref, dist - cols = np.column_stack([theta[good, 0], theta[good, 1], theta[good, 5], - theta[good, 3], theta[good, 2], theta[good, 4], - lnL[good]]) - hdr = "right_ascension declination distance inclination psi phi_orb loglikelihood" + if phase_marginalized: + cols = np.column_stack([theta[:, 0], theta[:, 1], theta[:, 5], + theta[:, 3], theta[:, 2], lnL]) + hdr = "right_ascension declination distance inclination psi loglikelihood" + else: + cols = np.column_stack([theta[:, 0], theta[:, 1], theta[:, 5], + theta[:, 3], theta[:, 2], theta[:, 4], lnL]) + hdr = "right_ascension declination distance inclination psi phi_orb loglikelihood" elif ndim == 4 and opts.mode == "flowmc-dpsimarg": # 4-D (flowmc-dpsimarg): theta = ra, dec, phiref, incl (psi marginalised, # phi_ref sampled). Write ra, dec, incl, phi_orb. - cols = np.column_stack([theta[good, 0], theta[good, 1], - theta[good, 3], theta[good, 2], lnL[good]]) + cols = np.column_stack([theta[:, 0], theta[:, 1], + theta[:, 3], theta[:, 2], lnL]) hdr = "right_ascension declination inclination phi_orb loglikelihood" elif ndim == 4: # 4-D (flowmc-phimarg): ra, dec, psi, incl (phi_ref marginalised out) - cols = np.column_stack([theta[good, 0], theta[good, 1], - theta[good, 3], theta[good, 2], lnL[good]]) + cols = np.column_stack([theta[:, 0], theta[:, 1], + theta[:, 3], theta[:, 2], lnL]) hdr = "right_ascension declination inclination psi loglikelihood" elif ndim == 3: # 3-D (flowmc-phipsimarg): ra, dec, incl (phi_ref AND psi marginalised out) - cols = np.column_stack([theta[good, 0], theta[good, 1], - theta[good, 2], lnL[good]]) + cols = np.column_stack([theta[:, 0], theta[:, 1], + theta[:, 2], lnL]) hdr = "right_ascension declination inclination loglikelihood" else: # 5-D: ra, dec, psi, incl, phiref - cols = np.column_stack([theta[good, 0], theta[good, 1], theta[good, 3], - theta[good, 2], theta[good, 4], lnL[good]]) - hdr = "right_ascension declination inclination psi phi_orb loglikelihood" - sname = opts.output_file + "_" + str(out_index) + "_samples.dat" - np.savetxt(sname, cols, header=hdr) - print("Wrote %s (%d samples)" % (sname, int(good.sum()))) + if phase_marginalized: + cols = np.column_stack([theta[:, 0], theta[:, 1], theta[:, 3], + theta[:, 2], lnL]) + hdr = "right_ascension declination inclination psi loglikelihood" + else: + cols = np.column_stack([theta[:, 0], theta[:, 1], theta[:, 3], + theta[:, 2], theta[:, 4], lnL]) + hdr = "right_ascension declination inclination psi phi_orb loglikelihood" + # Column line FIRST (unchanged, so `head -1` parsers keep working); the + # provenance line follows, so the artifact records how it was produced -- + # notably the export ESS, which was previously written nowhere. + np.savetxt(sname, cols, header=hdr + "\nmode=%s %s" % (opts.mode, provenance)) + print("Wrote %s (%d samples)" % (sname, len(cols))) def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, @@ -854,6 +1792,21 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, Returns ``(value, flow_state)`` where ``flow_state`` is the trained-flow state to bootstrap the next event (``--mode flowmc`` only; ``None`` else). """ + # Per-EVENT state: a batch run analyzes several events in one process, and + # an undersizing on event 0 must not label event 1. + _anglemarg.reset_amp_failsafe() + tq = opts.time_marginalization_quadrature + if tq == "bandlimited": + required, g0, gcert = bandlimited_storage_requirement( + P.deltaT, opts.data_integration_window_half) + if opts.internal_data_storage_window_half < required: + print(" widening rholm storage for bandlimited time support: " + "%.6g -> %.6g s" % + (opts.internal_data_storage_window_half, required)) + opts.internal_data_storage_window_half = required + print(" precompute time support: guard_initial=%d guard_certificate=%d " + "samples storage_half=%.6g s" % + (g0, gcert, opts.internal_data_storage_window_half)) print("Building JAX likelihood (PrecomputeLikelihoodTerms + pack)...") like_data, extras = build_data_from_precompute( P.copy(), data_dict, psd_dict, fiducial_epoch, @@ -865,6 +1818,38 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, print(" modes:", like_data.lms, " guessed SNR:", extras["guess_snr"]) with_distance = not opts.distance_marginalization + # --limit-distance: (d_lo,d_hi) is what is SAMPLED / quadratured; the prior is + # always normalized over (opts.d_min, opts.d_max), which is what keeps the + # reported lnZ on the same scale as a full-range run. + d_lo, d_hi = resolve_distance_limit(opts) + if (d_lo, d_hi) != (float(opts.d_min), float(opts.d_max)): + print(" [limit] restricting distance sampling/quadrature to [%.4f, %.4f] Mpc " + "(prior normalization UNCHANGED over [%.4f, %.4f] Mpc)" + % (d_lo, d_hi, opts.d_min, opts.d_max)) + # Fail closed on a flag that would otherwise be silently inert. The + # log-uniform distance quadrature is wired -- and validated -- only for the + # dense (phi,psi) schemes of --mode flowmc-phipsimarg. + if (getattr(opts, "distance_grid_scheme", "uniform") != "uniform" + and getattr(opts, "distance_grid_points", None) is not None): + raise SystemExit( + "--distance-grid-points and --distance-grid-scheme loguniform both " + "set the distance node count. The log-uniform grid derives its " + "count from the data (see --distance-grid-tol); pass one or the " + "other, not both.") + # Resolved into a LOCAL, never written back onto opts: this driver has a + # documented history of an option written back on event 0 being read as + # event 1's choice in a batch loop. + n_dist_grid = (DISTANCE_GRID_POINTS_DEFAULT + if opts.distance_grid_points is None + else int(opts.distance_grid_points)) + if (getattr(opts, "distance_grid_scheme", "uniform") != "uniform" + or getattr(opts, "distance_grid_tol", None) is not None): + if opts.mode != "flowmc-phipsimarg": + raise SystemExit( + "--distance-grid-scheme/--distance-grid-tol apply only to " + "--mode flowmc-phipsimarg (they are validated only for the " + "dense angle-marginalization schemes); got --mode %s." + % opts.mode) if opts.mode in ("flowmc-phimarg", "nuts-phimarg"): # phi_ref-marginalised: requires distance marginalisation (baked in); # produces a 4-D (ra, dec, psi, incl) posterior. @@ -874,11 +1859,12 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiMargLikelihood nphi = getattr(opts, "n_phi", 32) print("Distance + phi_ref marginalization: ON (grid=%d, nphi=%d, d in [%g,%g] Mpc)" - % (opts.distance_grid_points, nphi, opts.d_min, opts.d_max)) + % (n_dist_grid, nphi, d_lo, d_hi)) like = JAXDistPhiMargLikelihood( - like_data, opts.d_min, opts.d_max, - nphi=nphi, n_grid=opts.distance_grid_points, - interp=opts.interp, guess_snr=extras["guess_snr"]) + like_data, d_lo, d_hi, + nphi=nphi, n_grid=n_dist_grid, + interp=opts.interp, guess_snr=extras["guess_snr"], + time_quadrature=tq, d_prior_range=(opts.d_min, opts.d_max)) if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive": gi = like.dist_grid_info print(" distance grid: ADAPTIVE d_peak=%.3g Mpc sigma_d=%.3g Mpc npts=%d" @@ -892,16 +1878,33 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood nphi = getattr(opts, "n_phi", 32) npsi = getattr(opts, "n_psi", 16) + angle_marg = getattr(opts, "angle_marg_scheme", ANGLE_MARG_DEFAULT) print("Distance + phi_ref + psi marginalization: ON (grid=%d, nphi=%d, npsi=%d, d in [%g,%g] Mpc)" - % (opts.distance_grid_points, nphi, npsi, opts.d_min, opts.d_max)) + % (n_dist_grid, nphi, npsi, d_lo, d_hi)) + dist_grid = getattr(opts, "distance_grid_scheme", "uniform") + dist_tol = getattr(opts, "distance_grid_tol", None) + if dist_tol is None: + dist_tol = _jax_core_dist_tol_default() + elif dist_grid == "uniform": + raise SystemExit( + "--distance-grid-tol only applies to --distance-grid-scheme " + "loguniform; it would be silently inert here.") like = JAXDistPhiPsiMargLikelihood( - like_data, opts.d_min, opts.d_max, nphi=nphi, npsi=npsi, - n_grid=opts.distance_grid_points, interp=opts.interp, - guess_snr=extras["guess_snr"]) - if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive": - gi = like.dist_grid_info - print(" distance grid: ADAPTIVE d_peak=%.3g Mpc sigma_d=%.3g Mpc npts=%d" - % (gi["d_peak"], gi["sigma_d"], gi["n"])) + like_data, d_lo, d_hi, nphi=nphi, npsi=npsi, + n_grid=n_dist_grid, interp=opts.interp, + guess_snr=extras["guess_snr"], angle_marg=angle_marg, + time_quadrature=tq, d_prior_range=(opts.d_min, opts.d_max), + dist_grid=dist_grid, dist_grid_tol=dist_tol) + # ALWAYS report the resolved scheme (requested may be 'auto'; this + # pipeline has a documented history of silently-inert flags). + print(" angle-marg scheme: %s (requested %s): %s" + % (like.angle_marg_scheme, angle_marg, + "; ".join("%s=%s" % kv for kv in + sorted(like.angle_marg_info.items()) + if kv[0] not in ("scheme", "requested")))) + print(" distance grid: %s" + % "; ".join("%s=%s" % kv + for kv in sorted(getattr(like, "dist_grid_info", {}).items()))) with_distance = False dim = 3 elif opts.mode == "flowmc-dpsimarg": @@ -912,11 +1915,12 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, npsi = getattr(opts, "n_psi", 8) print("Distance + psi marginalization (phi_ref sampled): ON " "(grid=%d, npsi=%d, d in [%g,%g] Mpc)" - % (opts.distance_grid_points, npsi, opts.d_min, opts.d_max)) + % (n_dist_grid, npsi, d_lo, d_hi)) like = JAXDistPsiMargLikelihood( - like_data, opts.d_min, opts.d_max, npsi=npsi, - n_grid=opts.distance_grid_points, interp=opts.interp, - guess_snr=extras["guess_snr"]) + like_data, d_lo, d_hi, npsi=npsi, + n_grid=n_dist_grid, interp=opts.interp, + guess_snr=extras["guess_snr"], time_quadrature=tq, + d_prior_range=(opts.d_min, opts.d_max)) if getattr(like, "dist_grid_info", {}).get("mode") == "adaptive": gi = like.dist_grid_info print(" distance grid: ADAPTIVE d_peak=%.3g Mpc sigma_d=%.3g Mpc npts=%d" @@ -925,17 +1929,29 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, dim = 4 elif opts.distance_marginalization: print("Distance marginalization: ON (grid=%d, d in [%g,%g] Mpc)" - % (opts.distance_grid_points, opts.d_min, opts.d_max)) + % (n_dist_grid, d_lo, d_hi)) like = JAXDistanceMarginalizedLikelihood( - like_data, opts.d_min, opts.d_max, n_grid=opts.distance_grid_points, - interp=opts.interp, phase_marginalization=opts.phase_marginalization) + like_data, d_lo, d_hi, n_grid=n_dist_grid, + interp=opts.interp, phase_marginalization=opts.phase_marginalization, + time_quadrature=tq, d_prior_range=(opts.d_min, opts.d_max)) dim = 5 else: like = JAXExtrinsicLikelihood( like_data, interp=opts.interp, - phase_marginalization=opts.phase_marginalization) + phase_marginalization=opts.phase_marginalization, + time_quadrature=tq) dim = 6 + if like.time_quadrature != tq: + raise RuntimeError("constructed likelihood changed time quadrature: %r != %r" + % (like.time_quadrature, tq)) + print(" time-marginalization quadrature: %s (interpolate-time: %s)" + % (like.time_quadrature, opts.interp)) + if like.time_quadrature == "bandlimited": + print(" time guard: derived_initial=%d samples certified=%d samples; " + "fine factor is curvature-derived per row with one doubling certificate" + % (like.time_guard_initial, like.time_guard_certified)) + if opts.mode == "map": theta_map, lnL_map = run_map(like, opts, rng, dim, with_distance) fish = like.fisher(theta_map) @@ -956,6 +1972,12 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, "(it samples the 5-D angular posterior)." % opts.mode) from RIFT.likelihood.jax_ile import samplers as _samplers n_starts = opts.num_chains if opts.num_chains and opts.num_chains > 1 else 8 + # Resolved per event and used directly. Deliberately a LOCAL: writing it + # back to opts made event 1 of a batch read event 0's choice. + resolved_beta = float(opts.adapt_weight_exponent) + if opts.mode in _TEMPERED_MODES: + resolved_beta = resolve_tempering_exponent( + opts, dim, tempered_cloud_size(opts, n_starts)) if opts.mode == "nuts-phimarg": # Fisher-whitened multi-start NUTS on the 4-D phimarg posterior. res = _samplers.fisher_nuts_sample_phimarg( @@ -982,12 +2004,35 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, max_stages=max(opts.temper_max_stages, 80), puff_scale=opts.smc_puff_scale, seed=opts.seed, verbose=opts.verbose) + # AN UNFINISHED LADDER IS NOT A POSTERIOR DRAW. smc_puffball_sample + # also stops at max_stages (and on a cloud that has lost its finite + # lnL), leaving a cloud that targets L**inv_T * prior -- reachable + # exactly in the sharp high-SNR regime this fallback exists for, where + # the ESS rule picks very small temperature steps. NEITHER artifact + # may be published from that state: the export would ship a TEMPERED + # cloud as an equal-weight posterior draw, and the `.dat` would carry + # the SMC log Z(inv_T) rather than log Z. Fail the event here, before + # write_samples/write_dat, and clear anything an earlier run left at + # those paths (--soft-fail-event-range still skips to the next event). + _inv_T = res.get("inv_T") + if _inv_T is None or not float(_inv_T) >= 1.0 - 1e-9: + _remove_stale_artifact(samples_path(opts, out_index), "export") + _remove_stale_artifact(dat_path(opts, out_index), "result") + raise RuntimeError( + "--smc-puffball: the SMC temperature ladder reached inv_T=%s, " + "not 1, so the cloud still targets L**inv_T * prior and is " + "neither a posterior sample nor an evidence. Nothing written " + "for output index %d. Give the ladder room to finish (raise " + "--temper-max-stages / --smc-walkers, or lower " + "--temper-ess-frac) rather than exporting the tempered cloud." + % ("(not reported by the installed sampler)" if _inv_T is None + else "%.4g" % float(_inv_T), out_index)) elif opts.mode in ("flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg"): # 4-D phi-marg (ra,dec,psi,incl), 3-D phi+psi-marg (ra,dec,incl), or # 4-D d+psi-marg (ra,dec,phiref,incl); flowmc_sample_phimarg is # dimension-agnostic (helpers chosen from ANGULAR_PARAM_ORDER). # Static tempering: beta = --adapt-weight-exponent -> temper = 1/beta. - _beta = float(opts.adapt_weight_exponent) + _beta = resolved_beta _temper = 1.0 / _beta if _beta > 0 else 1.0 res = _samplers.flowmc_sample_phimarg( like, opts.d_min, opts.d_max, @@ -1008,7 +2053,7 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, seed=opts.seed, reuse_state=flow_state, verbose=opts.verbose) out_flow_state = res.get("flow_state") else: - _beta = float(opts.adapt_weight_exponent) + _beta = resolved_beta _temper = 1.0 / _beta if _beta > 0 else 1.0 res = _samplers.flowmc_sample( like, opts.d_min, opts.d_max, n_chains=max(n_starts, 20), @@ -1022,18 +2067,84 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, out_flow_state = res.get("flow_state") # bootstrap the next event theta, lnL = res["theta"], res["lnL"] logZ, sig, neff, ntot = res["logZ"], res["sigma_over_Z"], res["neff"], len(theta) + # post_weight means DIFFERENT things per sampler, so only consume it + # where it is an importance weight. + # * flowMC modes: L^(1-inv_T), the correction from the TEMPERED state + # actually sampled (exponent = --adapt-weight-exponent) to the exact + # posterior. Uniform only at inv_T == 1. This is a genuine w. + # Under --smc-puffball the same key carries the SMC ladder's + # L^(1-inv_T) correction, and the gate above has already refused the + # event unless the ladder finished -- so it is uniform by then. + # * multistart-nuts / nuts-phimarg: NOT an importance weight. + # samplers.py builds np.full(n_per[k], mass[k]/n_per[k]) -- a + # per-chain Laplace MODE-EVIDENCE weight, constant within a chain + # (multistart_nuts estimates it as peak_k + 0.5*logdet(sky cov), + # which its own comment flags as approximate). Fair-drawing against + # it would re-weight whole chains by that estimator and duplicate + # ~half the rows, not correct any proposal mismatch. It may well fix + # a real mode-mass bias, but that is unvalidated -- so these modes + # export their chains unreweighted until it is measured. + _pw = res.get("post_weight") if opts.mode in _TEMPERED_MODES else None + if _pw is not None and len(_pw) != len(theta): + # Silently dropping the weights here made the export claim no + # reweighting was needed; write_samples now refuses a mismatch, and + # this raises at the source so the mode is named in the message. + raise RuntimeError( + "--mode %s returned %d post_weight entries for %d samples; the " + "tempering correction cannot be applied." + % (opts.mode, len(_pw), len(theta))) + logw_export = (np.log(np.asarray(_pw, dtype=float)) + if _pw is not None else None) elif opts.mode == "prior-mc": - logZ, sig, neff, ntot, theta, lnL = run_prior_mc(like, opts, rng, dim, with_distance) + logZ, sig, neff, ntot, theta, lnL, logw_export = run_prior_mc(like, opts, rng, dim, with_distance) elif opts.mode == "nuts": - logZ, sig, neff, ntot, theta, lnL = run_nuts(like, opts, rng, with_distance) + logZ, sig, neff, ntot, theta, lnL, logw_export = run_nuts(like, opts, rng, with_distance) else: - logZ, sig, neff, ntot, theta, lnL = run_laplace_is(like, opts, rng, dim, with_distance) + logZ, sig, neff, ntot, theta, lnL, logw_export = run_laplace_is(like, opts, rng, dim, with_distance) print("\n==== Result (event %d) ====" % event_id) print(" log evidence (lnL marginal over extrinsic) = %.5f" % logZ) print(" sigma_lnL = %.4g neff = %.1f ntotal = %d" % (sig, neff, ntot)) - write_dat(opts, P, out_index, event_id, logZ, sig, ntot, neff) - write_samples(opts, out_index, theta, lnL, with_distance) + # EXPORT FIRST, THEN PUBLISH THE RESULT ROW. write_samples raises when the + # cloud admits no fair draw, and that refusal means the integration itself + # collapsed -- so the event must leave NO artifact behind. Writing the + # `.dat` first published a normal ILE result row (a finite, collapsed + # evidence in the one-finite-weight case) that --soft-fail-event-range then + # left in place for downstream collectors while the batch carried on. + # + # write_samples takes NO rng: it derives its own from (seed, out_index). + # Passing the shared `rng` here -- which also feeds run_laplace_is / + # run_prior_mc and the samplers -- made --save-samples, an OUTPUT flag, + # change the lnL/logZ of every later event in the batch. + # Report ONCE per event, here, BEFORE either artifact is written -- and + # independently of --save-samples/--output-file, since write_samples() + # early-returns when export is off and a run with export disabled would + # otherwise get no warning and no persistent label at all. + # Compute ONCE from the RESOLVED scheme and hand the same string to both + # writers. Recomputing inside each writer with no argument left `scheme` + # None, so the standing BEST-EFFORT label never emitted and every artifact + # stayed silent -- an inert guard, which is the exact failure mode this + # label exists to prevent. + _scheme = getattr(like, "angle_marg_scheme", None) + _ev_note = angle_grid_suspect_note(_scheme) + if _ev_note.startswith("SUSPECT-ANGLE-GRID"): + sys.stderr.write( + "WARNING integrate_likelihood_extrinsic_jax: angle-marginalization " + "dense grids were UNDERSIZED at some evaluated points this event " + "(%s). Samples and evidence are LABELLED SUSPECT in their headers; " + "rebuild with a larger amp_sizing before using them. The run was " + "NOT aborted and no points were discarded -- discarding would excise " + "exactly the region the estimator missed.\n" % _ev_note) + elif _ev_note: + # BEST-EFFORT: nothing detected. Say so WITHOUT claiming a clean run -- + # announcing "UNDERSIZED" here would be a false alarm, and saying + # nothing would let silence read as verification. + sys.stderr.write( + "NOTE integrate_likelihood_extrinsic_jax: %s\n" % _ev_note) + write_samples(opts, out_index, theta, lnL, with_distance, angle_note=_ev_note, + logw=logw_export) + write_dat(opts, P, out_index, event_id, logZ, sig, ntot, neff, + angle_note=_ev_note) return logZ, out_flow_state @@ -1042,7 +2153,11 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, # --------------------------------------------------------------------------- def main(argv=None): optp = build_parser() + argv = _normalize_interpolate_time_argv(argv) opts, _ = optp.parse_args(argv) + # BEFORE anything reads an option: which tokens did the user actually type? + record_supplied_options(opts, argv, optp) + resolve_ile_interface_aliases(opts, optp) check_critical_and_report(opts, optp) if opts.event_time is None: diff --git a/MonteCarloMarginalizeCode/Code/bin/plot_posterior_corner.py b/MonteCarloMarginalizeCode/Code/bin/plot_posterior_corner.py index bb908e851..119ebf315 100755 --- a/MonteCarloMarginalizeCode/Code/bin/plot_posterior_corner.py +++ b/MonteCarloMarginalizeCode/Code/bin/plot_posterior_corner.py @@ -246,15 +246,18 @@ def render_coordinates(coord_names,logparams=[]): parser.add_argument("--meanPerAno-max",default=2*np.pi,type=float) parser.add_argument("--ecc-min",default=0,type=float) parser.add_argument("--ecc-max",default=1,type=float) +parser.add_argument("--a6c-min",default=-80,type=float) +parser.add_argument("--a6c-max",default=-20,type=float) parser.add_argument("--lnL-cut",default=None,type=float) parser.add_argument("--sigma-cut",default=0.4,type=float) +parser.add_argument("--hyperbolic", action="store_true", help="Read sample files in format including hyperbolic") +parser.add_argument("--a6c", action="store_true", help="Read sample files in format including a6c") parser.add_argument("--eccentricity", action="store_true", help="Read sample files in format including eccentricity") parser.add_argument("--meanPerAno", action="store_true", help="Read sample files in format including meanPerAno - assumes eccentricity also present") parser.add_argument("--matplotlib-block-defaults",action="store_true",help="Relies entirely on user to set plot options for plot styles from matplotlibrc") parser.add_argument("--no-mod-psi",action="store_true",help="Default is to take psi mod pi. If present, does not do this") parser.add_argument("--downselect-parameter",action='append', help='Name of parameter to be used to eliminate grid points ') parser.add_argument("--downselect-parameter-range",action='append',type=str) -parser.add_argument("--disable-special-ranges", action='store_true', help="Disable the default special parameter ranges. If not specified, the default ranges are used.") # ---- User-supplied coordinate-convert plugin (additive; the hardcoded RIFT # conversion path is untouched). When --supplementary-coordinate-code is # omitted these flags are a no-op and plotting behaves byte-identically @@ -285,6 +288,7 @@ def render_coordinates(coord_names,logparams=[]): "Repeat for each output column. If omitted, the plugin's CHARTS[chart] " "parameters / OUTPUT_PARAMETERS attribute is used.") parser.add_argument("--verbose",action='store_true',help='print matplotlibrc data') +parser.add_argument("--no-special-param-ranges",action='store_true',help='Wipe all artifical param ranges; let samples guide the ranges') opts= parser.parse_args() plt.rc('axes',unicode_minus=False) @@ -489,11 +493,11 @@ def _materialize_plugin_columns(samples, source_label=""): 'chi_pavg':[0,2], 'chi_p':[0,1], 'lambdat':[0,4000], + 'a6c':[opts.a6c_min,opts.a6c_max], 'eccentricity':[opts.ecc_min,opts.ecc_max], 'meanPerAno':[opts.meanPerAno_min,opts.meanPerAno_max] } -if opts.disable_special_ranges: - special_param_ranges = {} + #mc_range deprecated by generic bind_param #if opts.mc_range: # special_param_ranges['mc'] = eval(opts.mc_range) @@ -504,7 +508,9 @@ def _materialize_plugin_columns(samples, source_label=""): special_param_ranges[par]=eval(opts.param_bound[i]) print(par +" range ",special_param_ranges[par]) - +if opts.no_special_param_ranges: + special_param_ranges = {} + print("WARNING: Special Parameter Ranges being erased; make sure you want to do this!") # Parameters param_list = opts.parameter @@ -682,12 +688,21 @@ def _materialize_plugin_columns(samples, source_label=""): field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z","lnL", "sigmaOverL", "ntot", "neff") if opts.flag_tides_in_composite: if opts.flag_eos_index_in_composite: - print(" Reading composite file, assumingtide/eos-index-based format ") + print(" Reading composite file, assuming tide/eos-index-based format ") field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z","lambda1", "lambda2", "eos_indx","lnL", "sigmaOverL", "ntot", "neff") + elif opts.a6c: + print(" Reading composite file, assuming tide-based format with EOB parameter a6c ") + field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z","lambda1", "lambda2", "a6c","lnL", "sigmaOverL", "ntot", "neff") else: print(" Reading composite file, assuming tide-based format ") field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z","lambda1", "lambda2", "lnL", "sigmaOverL", "ntot", "neff") -if opts.eccentricity: +elif opts.hyperbolic: + print(" Reading composite file, assuming hyperbolic-based format ") + field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z","E0", "p_phi0", "lnL", "sigmaOverL", "ntot", "neff") +elif opts.a6c and (not opts.flag_tides_in_composite): + print(" Reading composite file, assuming non-tide-based format with EOB parameter a6c ") + field_names=("indx","m1", "m2", "a1x", "a1y", "a1z", "a2x", "a2y", "a2z", "a6c","lnL", "sigmaOverL", "ntot", "neff") +elif opts.eccentricity: print(" Reading composite file, assuming eccentricity-based format ") if opts.meanPerAno: print(" Reading composite file, assuming mpa-based format ") diff --git a/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py b/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py index a2c88d6f1..d5d74ce45 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py @@ -14,26 +14,83 @@ import numpy as np import RIFT.misc.weight_simulations as weight_simulations +import re import fileinput #import StringIO data_at_intrinsic = {} +models_at_intrinsic = {} # same keys; parallel list of model labels, model-aware mode only +models_seen = [] my_digits=5 # safety for high-SNR BNS -tides_on = False -distance_on = False -col_intrinsic = 9 - import argparse parser = argparse.ArgumentParser(usage="util_CleanILE.py fname1.dat fname2.dat ... ") parser.add_argument("fname",action='append',nargs='+') +parser.add_argument("--a6c", action="store_true") +parser.add_argument("--hyperbolic", action="store_true") parser.add_argument("--eccentricity", action="store_true") parser.add_argument("--meanPerAno", action="store_true") #Askold: adding specification for tabular eos file parser.add_argument("--tabular-eos-file", action="store_true") +parser.add_argument("--model-group-regex", default=None, help="Regex matched against each input file's BASENAME; capture group 1 is the waveform-model label. Enables model-aware combination: replicas are averaged within a model, then models are marginalized over with --model-prior weights. Without this flag every evaluation at a given intrinsic point is pooled flat, which is correct for replicas of ONE model and wrong across models.") +parser.add_argument("--model-prior", action="append", default=None, help="LABEL=WEIGHT prior weight for one model (repeatable). Default: uniform over the labels actually seen. Weights are renormalized over the models present at each intrinsic point.") +parser.add_argument("--expect-models", default=None, help="Comma-separated list of the models this run CONFIGURED. Coverage is judged against this list, not against the labels that happen to appear: a model whose composites are all empty or missing is skipped before its label is ever recorded, so without this a total failure of one approximant looks like a complete run and even --require-all-models accepts every point. The builder passes its full --approx list.") +parser.add_argument("--require-all-models", action="store_true", help="Drop intrinsic points not evaluated under EVERY model. Without it, a point covered by a subset is marginalized over that subset, which silently changes the estimator point by point.") opts = parser.parse_args() +model_mode = opts.model_group_regex is not None +model_rx = re.compile(opts.model_group_regex) if model_mode else None +expected_models = [m.strip() for m in opts.expect_models.split(",") if m.strip()] if opts.expect_models else None +model_prior_arg = {} +if opts.model_prior: + for item in opts.model_prior: + if "=" not in item: + sys.exit("--model-prior wants LABEL=WEIGHT, got {}".format(item)) + label, _, wt = item.partition("=") + model_prior_arg[label.strip()] = float(wt) + if any(w < 0 for w in model_prior_arg.values()): + sys.exit("--model-prior weights must be non-negative") + + +def expected_row_lengths(opts): + """Column counts consistent with the enabled advanced-physics groups. + + An ILE row is composed as + + event_id m1 m2 s1x s1y s1z s2x s2y s2z + [distance] [lambda1 lambda2 [eos_table_index]] [a6c] [E0 p_phi0] + [eccentricity [meanPerAno]] + lnL sigmaOverL ntotal neff + + (the ordering of the optional groups matches the ``col_lnL`` increment + chain in util_ConstructIntrinsicPosterior_GenericCoordinates.py). Each + enabled flag contributes a KNOWN number of columns, so the groups compose: + a run with --a6c --hyperbolic --eccentricity --meanPerAno writes all four. + Tides / EOS index / pinned distance have no command-line flag here, so the + row WIDTH is what distinguishes them; the allowed widths below are the + flag-implied base plus each of those possibilities. + """ + n_flag = 0 + if opts.a6c: + n_flag += 1 + if opts.hyperbolic: + n_flag += 2 + if opts.eccentricity: + n_flag += 1 + if opts.meanPerAno: + n_flag += 1 + lengths = set() + lengths.add(13 + n_flag) # no tides, no pinned distance + lengths.add(13 + n_flag + 2) # lambda1, lambda2 + lengths.add(13 + n_flag + 3) # lambda1, lambda2, eos_table_index + if n_flag == 0: + lengths.add(14) # pinned distance (written only on its own) + return lengths + + +allowed_lengths = expected_row_lengths(opts) + #print opts.fname from pathlib import Path for fname in opts.fname[0]: #sys.argv[1:]: @@ -43,6 +100,16 @@ if os.stat(fname).st_size==0: # skip files of zero length continue sys.stderr.write(str(fname)+"\n") + this_model = None + if model_mode: + match = model_rx.search(os.path.basename(str(fname))) + if match is None: + sys.exit("--model-group-regex {!r} does not match {}; refusing to " + "pool it as an unlabelled model".format( + opts.model_group_regex, os.path.basename(str(fname)))) + this_model = match.group(1) + if this_model not in models_seen: + models_seen.append(this_model) # data = np.loadtxt(fname) # this will FAIL if we have a heterogeneous data source! BE CAREFUL data = np.genfromtxt(fname,invalid_raise=False) # Protect against inhomogeneous data if len(data.shape) ==1: @@ -50,87 +117,198 @@ for line in data: try: line = np.around(line, decimals=my_digits) - lambda1=lambda2=0 - eos_index = 0 - if opts.eccentricity: - if opts.meanPerAno: - indx, m1,m2, s1x,s1y,s1z,s2x,s2y,s2z,ecc,meanPerAno, lnL, sigmaOverL, ntot, neff = line - col_intrinsic = 11 - else: - indx, m1,m2, s1x,s1y,s1z,s2x,s2y,s2z,ecc, lnL, sigmaOverL, ntot, neff = line - col_intrinsic = 10 - elif len(line) == 13 and (not tides_on) and (not distance_on): # strip lines with the wrong length - indx, m1,m2, s1x,s1y,s1z,s2x,s2y,s2z,lnL, sigmaOverL, ntot, neff = line - elif len(line) == 14: - distance_on=True - col_intrinsic=10 - indx, m1,m2, s1x,s1y,s1z,s2x,s2y,s2z,dist, lnL, sigmaOverL, ntot, neff = line - elif len(line)==15: - tides_on = True - col_intrinsic =11 - indx, m1,m2, s1x,s1y,s1z,s2x,s2y,s2z, lambda1,lambda2,lnL, sigmaOverL, ntot, neff = line - - #Askold: adding the option for tabular eos file - elif opts.tabular_eos_file and len(line) == 16: #checking if the tabular eos file is defined in the parser and if the line actually has all the columns - #no eccentricity assumed here, since export_eos_index option doesn't output eccentricity, also it doesn't apply to neutron stars - col_intrinsic = 12 #I assume eos_index to be intrinsic parameter - indx, m1, m2, s1x, s1y, s1z, s2x, s2y, s2z, lambda1, lambda2, eos_index, lnL, sigmaOverL, ntot, neff = line - + if len(line) not in allowed_lengths: # strip lines with the wrong length + raise ValueError("Unsupported ILE row layout: {} columns (expected one of {})".format(len(line), sorted(allowed_lengths))) + # Whatever the enabled groups, the last four columns are + # lnL sigmaOverL ntotal neff, so everything between the event id and + # them is the intrinsic key used to consolidate repeated evaluations. + col_intrinsic = len(line) - 4 + lnL, sigmaOverL, ntot, neff = line[col_intrinsic:] if sigmaOverL>0.9: continue # do not allow poorly-resolved cases (e.g., dominated by one point). These are often useless if tuple(line[1:col_intrinsic]) in data_at_intrinsic: # print " repeated occurrence ", line[1:9] data_at_intrinsic[tuple(line[1:col_intrinsic])].append(line[col_intrinsic:]) + models_at_intrinsic[tuple(line[1:col_intrinsic])].append(this_model) else: # print " new key ", line[1:9] data_at_intrinsic[tuple(line[1:col_intrinsic])] = [line[col_intrinsic:]] - except: + models_at_intrinsic[tuple(line[1:col_intrinsic])] = [this_model] + except Exception as exc: + sys.stderr.write("Skipping malformed ILE row in {}: {}\n".format(fname, exc)) continue -for key in data_at_intrinsic: - lnL, sigmaOverL, ntot,neff = np.transpose(data_at_intrinsic[key]) - lnL = np.atleast_1d(lnL); sigmaOverL = np.atleast_1d(sigmaOverL); ntot = np.atleast_1d(ntot); neff = np.atleast_1d(neff) - sigmaOverL = np.maximum(sigmaOverL, 1e-7*np.ones(len(lnL))) # prevent accidental underflow during debugging/using synthetic data with no error - lnLmax = np.max(lnL) - L = np.exp(lnL - lnLmax) # remove overall Lmax factor, which factors out of the combination +def _pool_linear(lnL, sigmaOverL, ntot, weights=None): + """Combine evaluations of the SAME quantity by their weighted linear mean in L. + + Returns (lnLbar, sigmaOverL). The linear arithmetic is performed relative + to this pool's own maximum, so the largest member is exactly one and Lbar + cannot underflow to zero merely because some *other* model is much better. + + DO NOT inverse-variance weight with the reported sigmas: each sigma is + computed from the same importance weights as its lnL, so a replica that + silently missed the likelihood peak reports BOTH a low lnL AND a small + sigma -- 1/sigma^2 weighting then overweights the worst replica, giving a + systematically low combined lnL with an overconfident combined error. + The pooled (ntot-weighted) linear mean is unbiased in L regardless. + + Error: max(propagated per-run sigmas, between-replica scatter). Only the + scatter term can see the replica lottery (correlated underreporting); with + K replicas it has K-1 dof, so treat the result as a t-interval downstream. + """ + lnLscale = np.max(lnL) + L = np.exp(lnL - lnLscale) K = len(lnL) - # Combine repeated evaluations by their SAMPLE-COUNT-weighted LINEAR mean. - # DO NOT inverse-variance weight with the reported sigmas: each sigma is - # computed from the same importance weights as its lnL, so a replica that - # silently missed the likelihood peak reports BOTH a low lnL AND a small - # sigma -- 1/sigma^2 weighting then overweights the worst replica, giving a - # systematically low combined lnL with an overconfident combined error. - # The pooled (ntot-weighted) linear mean is unbiased in L regardless. - wts = np.asarray(ntot, dtype=float) - if np.any(wts <= 0) or not np.all(np.isfinite(wts)): - wts = np.ones(K) + if weights is None: + wts = np.asarray(ntot, dtype=float) + if np.any(wts <= 0) or not np.all(np.isfinite(wts)): + wts = np.ones(K) + else: + wts = np.asarray(weights, dtype=float) wts = wts/np.sum(wts) Lbar = np.sum(wts*L) - lnLmeanMinusLmax = np.log(Lbar) - # Error: max(propagated per-run sigmas, between-replica scatter). Only the - # scatter term can see the replica lottery (correlated underreporting); with - # K replicas it has K-1 dof, so treat the result as a t-interval downstream. sigma_prop = np.sqrt(np.sum((wts*sigmaOverL*L)**2))/Lbar if K > 1: sigma_scatter = np.sqrt( np.sum(wts**2 * (L - Lbar)**2) * K/(K-1.) )/Lbar else: sigma_scatter = 0. - sigmaNetOverL = max(sigma_prop, sigma_scatter) + return lnLscale + np.log(Lbar), max(sigma_prop, sigma_scatter) - if opts.eccentricity: - if opts.meanPerAno: - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], key[8], key[9], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) - else: - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], key[8], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) - elif tides_on: - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], key[8],key[9], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) - elif distance_on: - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], key[8], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) - - #Askold: new option for tabular eos file - elif opts.tabular_eos_file: #written similarly to the previous ones - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], key[8],key[9], key[10], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) - +if model_mode: + if expected_models: + absent = [m for m in expected_models if m not in models_seen] + if absent: + sys.stderr.write( + "util_CleanILE: WARNING: configured models contributed NOTHING: {}. " + "Their composites were empty or missing, so they are invisible to the " + "per-point coverage check and this is a {}-model mixture, not a {}-model " + "one.\n".format(", ".join(absent), len(models_seen), len(expected_models))) + # judge coverage against what was CONFIGURED + for m in expected_models: + if m not in models_seen: + models_seen.append(m) + if model_prior_arg: + missing = [m for m in models_seen if m not in model_prior_arg] + if missing: + sys.exit("--model-prior given but missing weights for {}; specify " + "every model or none".format(missing)) + sys.stderr.write("util_CleanILE: model-aware combination over {} models: {}\n".format( + len(models_seen), ", ".join(models_seen))) + +n_partial = 0 +n_dropped_partial = 0 +n_points = 0 +model_spread = [] # between-model scatter, reported but NOT put in sigmaOverL + +for key in data_at_intrinsic: + lnL, sigmaOverL, ntot,neff = np.transpose(data_at_intrinsic[key]) + lnL = np.atleast_1d(lnL); sigmaOverL = np.atleast_1d(sigmaOverL); ntot = np.atleast_1d(ntot); neff = np.atleast_1d(neff) + sigmaOverL = np.maximum(sigmaOverL, 1e-7*np.ones(len(lnL))) # prevent accidental underflow during debugging/using synthetic data with no error + if not model_mode: + # One model (or replicas of one model): pool everything flat. + lnLmean, sigmaNetOverL = _pool_linear(lnL, sigmaOverL, ntot) else: - print(-1, key[0],key[1], key[2], key[3],key[4], key[5],key[6], key[7], lnLmeanMinusLmax+lnLmax, sigmaNetOverL, np.sum(ntot), -1) + # Two levels, because replicas and models are not the same thing. + # Within a model, replicas estimate ONE number -> ntot-weighted mean. + # Across models, the quantity itself differs -> marginalize, + # L_marg(lambda) = sum_m p(m) L_m(lambda), + # which is linear in L, not in lnL. Averaging lnL instead would give + # the geometric mean (a logarithmic opinion pool), which is not + # Bayesian model marginalization. + labels = models_at_intrinsic[key] + present = [m for m in models_seen if m in labels] + if len(present) < len(models_seen): + if opts.require_all_models: + n_dropped_partial += 1 + continue + n_partial += 1 + lnL_m = []; sig_m = []; w_m = [] + for m in present: + sel = np.array([lab == m for lab in labels]) + lnLm, sm = _pool_linear(lnL[sel], sigmaOverL[sel], ntot[sel]) + lnL_m.append(lnLm); sig_m.append(sm) + w_m.append(model_prior_arg[m] if model_prior_arg else 1.0) + lnL_m = np.atleast_1d(np.array(lnL_m)); sig_m = np.atleast_1d(np.array(sig_m)) + w_m = np.atleast_1d(np.array(w_m, dtype=float)) + if np.sum(w_m) <= 0: + sys.exit("intrinsic point {} is covered only by zero-prior models; " + "cannot define the subset mixture".format(key)) + # Renormalized over the models PRESENT here: with partial coverage the + # estimator is a marginal over a subset, which is why n_partial is + # reported and --require-all-models exists. + w_m = w_m/np.sum(w_m) + # Combine the model means on a new common scale. Each model was pooled + # on its own scale above; converting only the final model means relative + # to their maximum is the log-sum-exp construction and remains finite + # even when the models differ by thousands of nats. + positive = w_m > 0 + # A zero-prior model must not set the numerical scale: if it is 1000 + # nats louder than every positive-prior model, scaling by it makes all + # contributing likelihoods underflow and returns -inf/nan. + lnL_pos = lnL_m[positive] + sig_pos = sig_m[positive] + w_pos = w_m[positive] + w_pos = w_pos/np.sum(w_pos) + lnL_model_scale = np.max(lnL_pos) + L_m_scaled = np.exp(lnL_pos - lnL_model_scale) + Lbar_scaled = np.sum(w_pos*L_m_scaled) + lnLmean = lnL_model_scale + np.log(Lbar_scaled) + # ACROSS MODELS, report ONLY the propagated integration uncertainty. + # + # An earlier version also took the between-model scatter, reasoning that + # it is the waveform-systematic contribution and should widen the + # downstream fit. It does the opposite. sigmaOverL is an INTEGRATION + # error, and CIP drops every row above --sigma-cut (default 0.6, + # util_ConstructIntrinsicPosterior_GenericCoordinates.py). Model + # disagreement large enough to matter therefore exceeds the cut and + # DELETES precisely the intrinsic points where the models disagree -- + # exactly the points this workflow exists to fit. + # + # The model variation is already carried by Lbar, which is the + # marginalized likelihood. It does not belong in the error bar too. + sigmaNetOverL = np.sqrt(np.sum( + (w_pos*sig_pos*L_m_scaled)**2))/Lbar_scaled + M = len(lnL_pos) + if M > 1: + # kept as a diagnostic only -- never folded into sigmaNetOverL + spread = np.sqrt(np.sum( + w_pos**2 * (L_m_scaled - Lbar_scaled)**2) + * M/(M-1.))/Lbar_scaled + model_spread.append(spread) + + n_points += 1 + + + # The key already holds every intrinsic column that was present in the + # input rows, in input order, so the composite preserves whatever + # combination of advanced-physics groups the run enabled. + print(-1, *key, lnLmean, sigmaNetOverL, np.sum(ntot), -1) + + +# Coverage report. stdout is the data stream, so this goes to stderr. +if model_mode: + sys.stderr.write( + "util_CleanILE: {} intrinsic points written, {} models\n".format( + n_points, len(models_seen))) + if n_dropped_partial: + sys.stderr.write( + "util_CleanILE: DROPPED {} intrinsic points not evaluated under all " + "{} models (--require-all-models)\n".format( + n_dropped_partial, len(models_seen))) + if model_spread: + sys.stderr.write( + "util_CleanILE: between-model scatter (diagnostic only, NOT folded " + "into sigmaOverL): median {:.4f}, max {:.4f} over {} multi-model " + "points\n".format(float(np.median(model_spread)), + float(np.max(model_spread)), len(model_spread))) + if n_partial: + sys.stderr.write( + "util_CleanILE: WARNING: {} of {} intrinsic points were evaluated " + "under only a SUBSET of the {} models, and were marginalized over " + "that subset. The estimator therefore differs point to point -- a " + "model-dependent, lambda-dependent change in the effective prior. " + "Common causes: the sigmaOverL>0.9 resolution cut firing for one " + "model but not another, or a model's ILE job failing. Use " + "--require-all-models to drop these instead.\n".format( + n_partial, n_points, len(models_seen))) diff --git a/MonteCarloMarginalizeCode/Code/bin/util_CombineApproximantPosteriors.py b/MonteCarloMarginalizeCode/Code/bin/util_CombineApproximantPosteriors.py new file mode 100755 index 000000000..70da25629 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/bin/util_CombineApproximantPosteriors.py @@ -0,0 +1,173 @@ +#! /usr/bin/env python +"""Combine per-waveform-model posteriors into one, weighted by p(m) Z_m. + +This is the recombination step of the multi-approximant workflow. The +iteration loop fits a cross-model MARGINALIZED likelihood (see +util_CleanILE.py --model-group-regex), so every model shares one intrinsic +grid; the terminal stage then forks, fitting each model separately to get its +own posterior samples and its own evidence. This script closes the fork. + +The posterior of the model-marginalized hypothesis is the mixture + + p(theta|d) = sum_m w_m p_m(theta|d), w_m propto p(m) Z_m + +so a model that fits the data better contributes proportionally more samples. +Sampling -- rather than carrying weights -- is what keeps the output a drop-in +replacement for a single-model extrinsic_posterior_samples.dat. + +NOTE the two weights are different things and both are needed: p(m) is the +prior belief in a waveform model, Z_m is what the data say about it. Passing +--model-prior alone does NOT give a prior-weighted mixture, because Z_m still +multiplies it; that is the intended Bayesian behaviour. +""" + +import argparse +import os +import sys + +import numpy as np + +from RIFT.misc.cip_pipeline import systematic_resample, unique_draw_bound + + +def read_ln_evidence(fname): + """Read ln Z from a CIP '+annotation.dat' file. + + Format written by util_ConstructIntrinsicPosterior_GenericCoordinates.py: + a '# lnL sigma_lnL ...' header then one whitespace-separated row whose + first field is ln_integrand_value_absolute. + """ + with open(fname) as f: + rows = [ln for ln in f if ln.strip() and not ln.strip().startswith("#")] + if not rows: + raise ValueError("{}: no evidence row".format(fname)) + value = float(rows[0].split()[0]) + if not np.isfinite(value): + raise ValueError("{}: ln Z is not finite ({})".format(fname, value)) + return value + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--model", action="append", required=True, + help="LABEL:POSTERIOR.dat:ANNOTATION.dat (repeatable)") + parser.add_argument("--model-prior", action="append", default=None, + help="LABEL=WEIGHT prior p(m) (repeatable). Default uniform.") + parser.add_argument("--output", required=True) + parser.add_argument("--n-output-samples", type=int, default=None, + help="Requested output count (default: total input rows); " + "capped at the duplicate-free fair-draw frontier.") + parser.add_argument("--seed", type=int, default=None, + help="Set for a reproducible draw; default is unseeded.") + opts = parser.parse_args(argv) + + priors = {} + if opts.model_prior: + for item in opts.model_prior: + label, _, wt = item.partition("=") + label = label.strip() + if label in priors: + parser.error("duplicate --model-prior for {}".format(label)) + value = float(wt) + if not np.isfinite(value) or value < 0: + parser.error("--model-prior weights must be finite and non-negative") + priors[label] = value + if sum(priors.values()) <= 0: + parser.error("--model-prior weights must have positive total mass") + + labels, samples, header, ln_z = [], [], None, [] + for spec in opts.model: + parts = spec.split(":") + if len(parts) != 3: + parser.error("--model wants LABEL:POSTERIOR.dat:ANNOTATION.dat, got {}".format(spec)) + label, post_file, annot_file = parts + if label in labels: + parser.error("duplicate --model label {}".format(label)) + for f in (post_file, annot_file): + if not os.path.exists(f): + sys.exit("util_CombineApproximantPosteriors: missing {}".format(f)) + with open(post_file) as f: + first = f.readline() + this_header = first.rstrip("\n") if first.startswith("#") else None + if header is None: + header = this_header + elif this_header != header: + # Silently column-mismatched posteriors would produce a garbage + # mixture, so refuse rather than guess an alignment. + sys.exit("util_CombineApproximantPosteriors: {} has a different column " + "header than the first model; refusing to mix".format(post_file)) + dat = np.atleast_2d(np.genfromtxt(post_file, comments="#")) + if dat.size == 0: + sys.exit("util_CombineApproximantPosteriors: {} has no samples".format(post_file)) + n_before = len(dat) + dat = np.unique(dat, axis=0) + if len(dat) < n_before: + sys.stderr.write( + "util_CombineApproximantPosteriors: WARNING: removed {} duplicate " + "rows from {} before model allocation.\n".format( + n_before-len(dat), post_file)) + labels.append(label); samples.append(dat); ln_z.append(read_ln_evidence(annot_file)) + + if priors: + missing = [l for l in labels if l not in priors] + if missing: + sys.exit("--model-prior given but missing weights for {}".format(missing)) + ln_prior = np.array([np.log(priors[l]) if priors[l] > 0 else -np.inf for l in labels]) + else: + ln_prior = np.zeros(len(labels)) + + # w_m propto p(m) Z_m, in logs so a large ln Z spread cannot overflow. + ln_w = np.array(ln_z) + ln_prior + positive = np.isfinite(ln_w) + if not np.any(positive): + parser.error("model weights have no positive finite mass") + scale = np.max(ln_w[positive]) + w = np.zeros(len(labels)) + w[positive] = np.exp(ln_w[positive] - scale) + w = w/np.sum(w) + + n_requested = (opts.n_output_samples if opts.n_output_samples is not None + else int(sum(len(d) for d in samples))) + if n_requested < 1: + parser.error("--n-output-samples must be positive") + # Treat every component row as an atom in the empirical mixture. A row in + # model m has mass w_m/N_m. This is the same systematic fair-draw contract + # as the final CIP export (PR #180), and avoids the overly conservative + # per-model capacity floor used by an earlier implementation. + row_weights = np.concatenate([ + np.full(len(dat), wt/len(dat)) for dat, wt in zip(samples, w)]) + all_samples = np.vstack(samples) + n_total = min(n_requested, unique_draw_bound(row_weights)) + if n_total < n_requested: + sys.stderr.write( + "util_CombineApproximantPosteriors: WARNING: the unique fair-draw " + "frontier supports only {} of {} requested mixture rows; reducing output rather " + "than drawing duplicates.\n".format(n_total, n_requested)) + rng = np.random.default_rng(opts.seed) + selected = systematic_resample(row_weights, n_total, rng=rng) + offsets = np.cumsum([0] + [len(dat) for dat in samples]) + counts = np.array([ + np.sum((selected >= offsets[i]) & (selected < offsets[i+1])) + for i in range(len(samples))]) + + sys.stderr.write("util_CombineApproximantPosteriors: mixture over {} models\n".format(len(labels))) + for label, lnz, wt, cnt, dat in zip(labels, ln_z, w, counts, samples): + sys.stderr.write(" {:<20s} lnZ={:12.4f} weight={:8.5f} draws={:7d} (of {} samples)\n".format( + label, lnz, wt, cnt, len(dat))) + dominant = np.max(w) + if dominant > 0.99: + sys.stderr.write( + "util_CombineApproximantPosteriors: WARNING: one model carries {:.4f} of the " + "mixture; the combined posterior is effectively single-model. A large ln Z " + "gap between waveform models is usually a sign the models disagree far more " + "than the statistical error, not that one is 'right'.\n".format(dominant)) + + out = all_samples[selected] + np.savetxt(opts.output, out, header=header[1:].strip() if header else "") + sys.stderr.write("util_CombineApproximantPosteriors: wrote {} samples to {}\n".format( + len(out), opts.output)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/bin/util_ConstructEOSPosterior.py b/MonteCarloMarginalizeCode/Code/bin/util_ConstructEOSPosterior.py index 11cb7b645..edeead5bd 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_ConstructEOSPosterior.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_ConstructEOSPosterior.py @@ -149,6 +149,7 @@ def add_field(a, descr): parser.add_argument("--fit-load-gp",default=None,type=str,help="Filename of GP fit to load. Overrides fitting process, but user MUST correctly specify coordinate system to interpret the fit with. Does not override loading and converting the data.") parser.add_argument("--fit-save-gp",default=None,type=str,help="Filename of GP fit to save. ") parser.add_argument("--fit-order",type=int,default=2,help="Fit order (polynomial case: degree)") +parser.add_argument("--fit-distance-tail",action='store_true',help="Distance-export (.dslice) runs ONLY, i.e. runs that carry an explicit distance fit coordinate. Beyond each intrinsic point's exported distance support, make the fitted lnL decay to zero as d->infinity instead of holding its edge value. An RF/ExtraTrees fit is piecewise constant outside its training envelope, so without this it holds lnL flat while the volumetric prior keeps growing like d^2, and the recovered distance posterior comes out ~18 percent too wide. Changes nothing on the support, so it is a strict addition. It is an error to request this without a distance fit coordinate.") parser.add_argument("--no-plots",action='store_true') parser.add_argument("--using-eos-type", type=str, default=None, help="Name of EOS parameterization (must match what is used for inputs). Will use EOS parameterization to identify appropriate field headers") parser.add_argument("--sampler-method",default="adaptive_cartesian",help="adaptive_cartesian|GMM|adaptive_cartesian_gpu") @@ -550,7 +551,7 @@ def fit_rf(x,y,y_errors=None,fname_export='nn_fit'): from sklearn.ensemble import ExtraTreesRegressor # Instantiate model. Usually not that many structures to find, don't overcomplicate # - should scale like number of samples - rf = ExtraTreesRegressor(n_estimators=100, verbose=True,n_jobs=-1) # no more than 5% of samples in a leaf + rf = ExtraTreesRegressor(n_estimators=100, verbose=True,n_jobs=-1) if y_errors is None: rf.fit(x,y) else: @@ -743,6 +744,36 @@ def convert_coords(x_in, _low=low_level_coord_names, _coord=coord_names): Y_err=None my_fit = fit_rf(X,Y,y_errors=Y_err) +### Distance tail: make the fit decay beyond each intrinsic point's exported distance support +### +### Only meaningful for a distance-export (.dslice) run, where `dist` is a FIT coordinate and the +### training set is ~50 discrete distances per intrinsic point. An RF/ExtraTrees fit is piecewise +### constant outside its training envelope, so past a point's outermost exported slice it returns +### that slice's lnL forever. The distance prior is volumetric and keeps growing like d^2, so the +### integrand grows instead of dying and the recovered distance posterior comes out ~18% too wide +### in every quantile span, median untouched. See RIFT/interpolators/distance_tail.py. +### +### This wraps whatever fit was just built and is a no-op on the support, so nothing that currently +### works changes. It is applied AFTER the cap/threshold cuts above so the tail is built from +### exactly the rows the fit itself saw. +if opts.fit_distance_tail: + if my_fit is None: + raise ValueError("--fit-distance-tail: no fit was built (--fit-method %s)" % opts.fit_method) + if 'dist' not in list(coord_names): + # Fail rather than silently do nothing: a run that asked for this and did not get it would + # carry the very bias the option exists to remove, with no sign of it in the log. + raise ValueError("--fit-distance-tail requires a distance fit coordinate, but coord_names " + "is %s. This option is for distance-export (.dslice) runs." % (list(coord_names),)) + from RIFT.interpolators.distance_tail import wrap_distance_tail + tail_report = {} + # The production decay law is the parameter-free chord (ratio = u), so the wrapper's tuning + # arguments -- the per-slice edge-slope fit and its `outer_frac`, the alternate laws -- reach + # only the diagnostic branches described in RIFT/interpolators/distance_tail.py. They are + # deliberately not exposed here: a CLI knob that cannot change the posterior is worse than none. + my_fit = wrap_distance_tail(my_fit, X, Y, coord_names, y_errors=Y_err, + lnL_offset=lnL_shift, report=tail_report) + print(" DISTANCE TAIL : decay beyond exported support enabled ", tail_report) + # Sort for later convenience (scatterplots, etc) indx = Y.argsort()#[::-1] diff --git a/MonteCarloMarginalizeCode/Code/bin/util_ConstructIntrinsicPosterior_GenericCoordinates.py b/MonteCarloMarginalizeCode/Code/bin/util_ConstructIntrinsicPosterior_GenericCoordinates.py index 6ff7caba3..68a18995a 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_ConstructIntrinsicPosterior_GenericCoordinates.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_ConstructIntrinsicPosterior_GenericCoordinates.py @@ -74,6 +74,7 @@ lsctables.use_in(ligolw.LIGOLWContentHandler) import RIFT.integrators.mcsampler as mcsampler +from RIFT.misc.mc_error import relative_mc_error try: import RIFT.integrators.mcsamplerEnsemble as mcsamplerEnsemble mcsampler_gmm_ok = True @@ -291,6 +292,12 @@ def extract_combination_from_LI(samples_LI, p): parser.add_argument("--chi-max", default=1,type=float,help="Maximum range of 'a' allowed. Use when comparing to models that aren't calibrated to go to the Kerr limit.") parser.add_argument("--chi-small-max", default=None,type=float,help="Maximum range of 'a' allowed on the smaller body. If not specified, defaults to chi_max") parser.add_argument("--eccentricity-prior", default="uniform",choices=['uniform','log_uniform'],help="Options are 'uniform' and 'log_uniform'") # constrained: only 'log_uniform' is branched on below, so anything else would quietly fall through to the uniform prior +parser.add_argument("--a6c-max", default=-20,type=float,help="Maximum range of 'a6c' allowed.") +parser.add_argument("--a6c-min", default=-80,type=float,help="Minimum range of 'a6c' allowed.") +parser.add_argument("--E0-max", default=1.2,type=float,help="Maximum range of 'E0' allowed.") +parser.add_argument("--E0-min", default=1.0,type=float,help="Minimum range of 'E0' allowed.") +parser.add_argument("--pphi0-max", default=5.4,type=float,help="Maximum range of 'p_phi0' allowed.") +parser.add_argument("--pphi0-min", default=0.0,type=float,help="Minimum range of 'p_phi0' allowed.") parser.add_argument("--ecc-max", default=0.9,type=float,help="Maximum range of 'eccentricity' allowed.") parser.add_argument("--ecc-min", default=0.0,type=float,help="Minimum range of 'eccentricity' allowed.") parser.add_argument("--meanPerAno-max", default=2*np.pi,type=float,help="Maximum range of 'meanPerAno' allowed.") @@ -336,6 +343,10 @@ def extract_combination_from_LI(samples_LI, p): parser.add_argument("--fit-save-gp",default=None,type=str,help="Filename of GP fit to save. ") parser.add_argument("--fit-save-jax",default=None,type=str,help="Base path for a self-contained, differentiable jax_gp export (writes .npz + .meta.json). Only used with --fit-method gp-jax-*. Reload with --fit-load-gp pointing at the same base path.") parser.add_argument("--fit-order",type=int,default=2,help="Fit order (polynomial case: degree)") +parser.add_argument("--fit-gp-length-scale-max-factor",default=5.0,type=float,help="fit_gp: upper bound on each RBF length scale, as a multiple of that coordinate's standard deviation over the retained points. Default 5.0 reproduces the hardcoded value. Unlike the noise and amplitude bounds this ceiling is DERIVED FROM THE DATA, not hand-tuned, so raising it lets the GP become effectively linear across the grid; it exists to be measured against, not routinely changed.") +parser.add_argument("--fit-gp-holdout-folds",default=0,type=int,help="fit_gp: if >0, also report a K-fold HELD-OUT predictive RMS alongside the in-sample residual, refitting the same kernel on each training split. In-sample residual cannot distinguish a flexible fit from an overfit one; this can. Costs K extra GP fits.") +parser.add_argument("--fit-gp-noise-bounds",default="1e-2,1",type=str,help="fit_gp: comma-separated (lo,hi) bounds on the WhiteKernel noise_level, in nats^2. Default reproduces the hand-tuned LVK-scale value. Widen when lnL spans a dynamic range far larger than LVK's (third-generation networks): the default saturates and the fit under-reports structure.") +parser.add_argument("--fit-gp-amplitude-bounds",default="1e-3,1e1",type=str,help="fit_gp: comma-separated (lo,hi) bounds on the ConstantKernel amplitude multiplying the RBF, in nats^2. Default reproduces the hand-tuned LVK-scale value, which caps the representable signal amplitude at sqrt(1e1)=3.2 nats.") parser.add_argument("--fit-uncertainty-added",default=False, action='store_true', help="Reported likelihood is lnL+(fit error). Use for placement and use of systematic errors.") parser.add_argument("--no-plots",action='store_true') parser.add_argument("--tabular-eos-file",type=str,default=None,help="Tabular file of EOS to use. The default prior will be UNIFORM in this table!") @@ -365,8 +376,13 @@ def extract_combination_from_LI(samples_LI, p): parser.add_argument("--internal-n-comp",default=1,type=int,help="number of components to use for GMM sampling. Default is 1, because we expect a unimodal posterior in well-adapted coordinates. If you have crappy coordinates, use more") parser.add_argument("--internal-gmm-memory-chisquared-factor",default=None,type=float,help="Multiple of the number of degrees of freedom to save. 5 is a part in 10^6, 4 is 10^{-4}, and None keeps all up to lnL_offset. Note that low-weight points can contribute notably to n_eff, and it can be dangerous to assume a simple chisquared likelihood! Provided in case we need very long runs") parser.add_argument("--assume-eos-but-primary-bh",action='store_true',help="Special case of known EOS, but primary is a BH") +parser.add_argument("--use-EOB-parameters", action="store_true") parser.add_argument("--use-eccentricity", action="store_true") parser.add_argument("--use-meanPerAno", action="store_true") +parser.add_argument("--use-hyperbolic", action="store_true") +parser.add_argument("--force-scatter",default=False,action='store_true', help='For hyperbolic analyses forces only scatter grid points') +parser.add_argument("--force-plunge",default=False,action='store_true', help='For hyperbolic analyses forces only plunge grid points') +parser.add_argument("--force-zoomwhirl",default=False,action='store_true', help='For hyperbolic analyses forces only zoomwhirl grid points') parser.add_argument("--tripwire-fraction",default=0.05,type=float,help="Fraction of nmax of iterations after which n_eff needs to be greater than 1+epsilon for a small number epsilon") # FIXME hacky options added by me (Liz) to try to get my capstone project to work. @@ -383,6 +399,17 @@ def extract_combination_from_LI(samples_LI, p): opts= parser.parse_args() +force_hyperbolic_classes = [ + opts.force_scatter, opts.force_plunge, opts.force_zoomwhirl, +] +if any(force_hyperbolic_classes) and not opts.use_hyperbolic: + parser.error( + "--force-scatter, --force-plunge, and --force-zoomwhirl require " + "--use-hyperbolic" + ) +if sum(bool(value) for value in force_hyperbolic_classes) > 1: + parser.error("CANNOT use multiple hyperbolic --force-X options at once") + # good enough file: terminate always with success if present, don't try any more work if opts.check_good_enough: fname = 'cip_good_enough' @@ -398,10 +425,16 @@ def extract_combination_from_LI(samples_LI, p): if not(opts.no_adapt_parameter): opts.no_adapt_parameter =[] # needs to default to empty list +A6C_MAX = opts.a6c_max +A6C_MIN = opts.a6c_min +E0_MAX = opts.E0_max +E0_MIN = opts.E0_min +PPHI0_MAX = opts.pphi0_max +PPHI0_MIN = opts.pphi0_min ECC_MAX = opts.ecc_max ECC_MIN = opts.ecc_min -MEANPERANO_MAX = 2*np.pi -MEANPERANO_MIN = 0 +MEANPERANO_MAX = opts.meanPerAno_max +MEANPERANO_MIN = opts.meanPerAno_min no_plots = no_plots | opts.no_plots lnL_shift = 0 lnL_default_large_negative = -500 @@ -683,7 +716,7 @@ def extract_combination_from_LI(samples_LI, p): if opts.input_tides: # only insert these cuts if we are using a composite file with tides! # do not downselect if lambda1 is not in ! allows NSBH - if not('lambda1' in downselect_dict) and 'lambda1' in opts.parameter: + if not('lambda1' in downselect_dict) and ('lambda1' in opts.parameter): downselect_dict['lambda1'] = [lambda_min,lambda_max] if not('lambda2' in downselect_dict): downselect_dict['lambda2'] = [lambda_min,lambda_small_max] @@ -924,6 +957,9 @@ def tapered_magnitude_prior_alt(x,loc=0.8,kappa=20.): # return 1/(1+f1) +def a6c_prior(x): + return np.ones(x.shape) / (A6C_MAX-A6C_MIN) # uniform over the interval [A6C, A6C_MAX] + def eccentricity_prior(x): return np.ones(x.shape) / (ECC_MAX-ECC_MIN) # uniform over the interval [0.0, ECC_MAX] @@ -933,9 +969,7 @@ def log_eccentricity_prior(x): # AttributeError as soon as the prior was evaluated) with a (ECC_MAX-ECC_MIN) # normalization, which is the uniform prior's normalization, not this one's: # \int_ECC_MIN^ECC_MAX dx/(x*C) = 1 => C = ln(ECC_MAX/ECC_MIN). - # The statement below is byte-identical to rift_O4c 0.0.17.13, which reached the same - # fix independently; keep it that way so future O4c->O4d merges do not conflict here. - return np.ones(x.shape) / (x*np.log(ECC_MAX/ECC_MIN)) # log uniform over the interval [ECC_MIN, ECC_MAX]; if ECC_MIN=0.0, auto corrects to ECC_MIN=0.001 + return np.ones(x.shape) / (x*np.log(ECC_MAX/ECC_MIN)) # log uniform over the interval [ECC_MIN, ECC_MAX] def uniform_eccentricity_ln_prior(x): return np.ones(x.shape) / ((np.log(ECC_MAX/ECC_MIN))) # log uniform over the interval [ECC_MIN, ECC_MAX]; if ECC_MIN=0.0, auto corrects to ECC_MIN=0.001 @@ -954,6 +988,12 @@ def log_eccentricity_squared_prior(x): def meanPerAno_prior(x): return np.ones(x.shape) / (MEANPERANO_MAX-MEANPERANO_MIN) # uniform over the interval [MEANPERANO_MIN, MEANPERANO_MAX] +def initial_energy_prior(x): + return np.ones(x.shape) / (E0_MAX-E0_MIN) # uniform over the interval [E0_MIN, E0_MAX] + +def initial_angmom_prior(x): + return np.ones(x.shape) / (PPHI0_MAX-PPHI0_MIN) # uniform over the interval [PPHI0_MIN, PPHI0_MAX] + def precession_prior(x): return 0.5*np.ones(x.shape) # uniform over the interval [0.0, 2.0] @@ -1002,6 +1042,9 @@ def normalized_zbar_prior(z): 's1z_bar':normalized_zbar_prior, 's2z_bar':normalized_zbar_prior, # Other priors + 'a6c':a6c_prior, + 'E0':initial_energy_prior, + 'p_phi0':initial_angmom_prior, 'eccentricity':eccentricity_prior, 'eccentricity_ln':uniform_eccentricity_ln_prior, 'eccentricity_squared':eccentricity_squared_prior, @@ -1026,6 +1069,9 @@ def normalized_zbar_prior(z): 'lambda2':[lambda_min,lambda_small_max], 'lambda_plus':[lambda_min,lambda_plus_max], 'lambda_minus':[-lambda_max,lambda_max], # will include the true region always...lots of overcoverage for small lambda, but adaptation will save us. + 'a6c':[A6C_MIN, A6C_MAX], + 'E0':[E0_MIN,E0_MAX], + 'p_phi0':[PPHI0_MIN,PPHI0_MAX], 'eccentricity':[ECC_MIN, ECC_MAX], 'eccentricity_ln':[np.log(ECC_MIN), np.log(ECC_MAX)], 'eccentricity_squared':[ECC_MIN**2, ECC_MAX**2], @@ -1075,7 +1121,7 @@ def normalized_zbar_prior(z): # Note CDF of 1/(1+q)^2 is -1/(1+q), so the normalization is analytic if 'q' in low_level_coord_names: delta_range = np.sqrt(1 - 4*np.array(eta_range)) - q_range = (1-delta_range)/(1+delta_range) + q_range = prior_range_map['q'] = (1-delta_range)/(1+delta_range) norm_factor_q = 1./(1+q_range[0]) - 1./(1+q_range[1]) prior_map['q'] = functools.partial(q_prior, norm_factor=norm_factor_q) prior_range_map['q'] = q_range @@ -1370,6 +1416,126 @@ def adderr(y): val,err = y return val+error_factor*err +def _gp_bounds_opt(spec): + """Parse a "lo,hi" CLI string into a (float,float) sklearn bounds tuple.""" + lo, hi = (float(v) for v in str(spec).split(",")) + return (lo, hi) + + +def _mc_index_in(names): + """Index of chirp mass in a coordinate list, or -1 if it is not in it. + + The peak scan reports a *physical* coordinate, so its index must be taken in + the basis the fitted array is actually built in: the columns of x follow + coord_names. The global mc_index indexes low_level_coord_names instead, and + the two lists diverge exactly when --parameter-implied / --parameter-nofit + are used -- e.g. "--parameter delta_mc --parameter-implied mu1 + --parameter-implied mu2 --parameter-nofit mc" fits [delta_mc, mu1, mu2] + while mc_index is 1, so scanning column 1 would report the mu1 peak as + chirp mass. -1 is the "not a fitted coordinate" sentinel that + report_gp_kernel declines to scan. + """ + return list(names).index('mc') if 'mc' in names else -1 + + +def report_gp_kernel(gp, x, y, tol=1e-3, holdout_folds=0, kernel_proto=None, + alpha_proto=None, peak_index=None, peak_grid=240, + peak_max_points=200000, peak_block_elements=5000000): + """Print a machine-readable record of the fitted GP: kernel form, the bounds + actually in force, the fitted hyperparameters, and -- the part that is not + self-announcing -- which of them the optimizer drove onto a bound. + + A saturated fit returns successfully and looks exactly like a converged one, + so the saturation flags are the only way to tell from the outputs whether the + kernel bounds, rather than the data, set the answer. sklearn stores + hyperparameters and bounds log-transformed in .theta / .bounds, so proximity + is tested there: |theta - bound| < tol means "on the bound". + + peak_max_points caps the total size of the optional peak scan and + peak_block_elements caps the points-by-training kernel block it evaluates at + a time; both keep this diagnostic's memory bounded regardless of dimension + and retained sample count. + """ + import json as _json + rec = {"kernel_form": str(gp.kernel), "kernel_fitted": str(gp.kernel_), + "log_marginal_likelihood": float(gp.log_marginal_likelihood_value_), + "n_points": int(len(y)), "lnL_range": float(np.nanmax(y) - np.nanmin(y)), + "resid_std": float(np.std(y - gp.predict(x))), "hyperparameters": []} + theta, bounds = gp.kernel_.theta, gp.kernel_.bounds + names = [] + for h in gp.kernel_.hyperparameters: + names.extend([h.name] * int(h.n_elements if h.n_elements > 1 else 1)) + for i in np.arange(len(theta)): + lo, hi = float(bounds[i][0]), float(bounds[i][1]) + at_lo, at_hi = bool(theta[i] - lo < tol), bool(hi - theta[i] < tol) + rec["hyperparameters"].append( + {"name": names[i] if i < len(names) else "param_%d" % i, + "value": float(np.exp(theta[i])), + "bounds": [float(np.exp(lo)), float(np.exp(hi))], + "at_lower_bound": at_lo, "at_upper_bound": at_hi, + # decades of headroom to the nearer bound: a value can be + # operationally pinned without tripping the boolean flag. + "decades_to_bound": float(min(theta[i] - lo, hi - theta[i]) / np.log(10.0))}) + rec["n_at_bound"] = int(sum(h["at_lower_bound"] or h["at_upper_bound"] + for h in rec["hyperparameters"])) + if holdout_folds and holdout_folds > 1 and kernel_proto is not None: + # In-sample residual rewards flexibility. Refit the SAME kernel on K + # training splits and score the untouched folds: this is the number that + # says whether relaxing the bounds bought generalization or overfitting. + from sklearn.model_selection import KFold + errs = [] + alpha_is_vec = hasattr(alpha_proto, "__len__") + for tr, te in KFold(holdout_folds, shuffle=True, random_state=0).split(x): + g2 = GaussianProcessRegressor( + kernel=kernel_proto, + alpha=(alpha_proto[tr] if alpha_is_vec else alpha_proto), + n_restarts_optimizer=4) + g2.fit(x[tr], y[tr]) + errs.append(g2.predict(x[te]) - y[te]) + errs = np.concatenate(errs) + rec["holdout_folds"] = int(holdout_folds) + rec["holdout_rms_nats"] = float(np.sqrt((errs ** 2).mean())) + rec["holdout_max_abs_nats"] = float(np.abs(errs).max()) + # The caller passes the chirp-mass column, whose sentinel for "there is no + # such coordinate" is -1: that is a legal numpy index, so it has to be + # rejected explicitly or the scan would report the LAST coordinate as mc. + peak_index = -1 if peak_index is None else int(peak_index) + if 0 <= peak_index < x.shape[1] and x.shape[1] <= 3: + # Where does the FITTED SURFACE put the likelihood maximum? This is the + # question the interpolant exists to answer, and it is not implied by the + # residual: a saturated kernel can score acceptably on average and still + # misplace the peak. Scanned over the training data's own extent. + # + # The grid is a full outer product, so its cost is peak_grid**ndim, and + # gp.predict then forms a points-by-training kernel matrix: in 3-D at + # peak_grid=240 that is 1.4e7 points against every retained sample, i.e. + # tens of GB. Thin the axes to a fixed total budget and evaluate in + # blocks, so a diagnostic can never be what kills a successful fit. + ndim = int(x.shape[1]) + n_per_axis = max(2, int(min(peak_grid, peak_max_points ** (1.0 / ndim)))) + axes = [np.linspace(x[:, i].min(), x[:, i].max(), n_per_axis) + for i in np.arange(ndim)] + mesh = np.meshgrid(*axes, indexing="ij") + pts = np.column_stack([m.ravel() for m in mesh]) + n_block = max(1, int(peak_block_elements // max(1, len(y)))) + peak_val, peak_pt = -np.inf, pts[0] + for start in np.arange(0, len(pts), n_block): + block = pts[start:start + n_block] + zz = gp.predict(block) + here = int(np.argmax(zz)) + if zz[here] > peak_val: + peak_val, peak_pt = float(zz[here]), block[here] + rec["surface_peak"] = {"coord_index": peak_index, + "value": float(peak_pt[peak_index]), + "grid_points_per_axis": int(n_per_axis)} + rec["saturated"] = bool(rec["n_at_bound"] > 0) + print(" GP-KERNEL-RECORD " + _json.dumps(rec, sort_keys=True)) + if rec["saturated"]: + print(" GP WARNING: %d hyperparameter(s) sit ON a bound -- the bounds, not the" + " data, are setting the fit." % rec["n_at_bound"]) + return rec + + def fit_gp(x,y,x0=None,symmetry_list=None,y_errors=None,hypercube_rescale=False,fname_export="gp_fit"): """ x = array so x[0] , x[1], x[2] are points. @@ -1400,7 +1566,7 @@ def fit_gp(x,y,x0=None,symmetry_list=None,y_errors=None,hypercube_rescale=False, if indx == mc_index: length_scale_min_here= 0.2*np.nanstd(x[:,indx]/np.sqrt(len(x))) print(" Setting mc range: retained point range is ", np.nanstd(x[:,indx]), " and target min is ", length_scale_min_here) - length_scale_bounds_est.append( (length_scale_min_here , 5*np.nanstd(x[:,indx]) ) ) # auto-select range based on sampling *RETAINED* (i.e., passing cut). Note that for the coordinates I usually use, it would be nonsensical to make the range in coordinate too small, as can occasionally happens + length_scale_bounds_est.append( (length_scale_min_here , opts.fit_gp_length_scale_max_factor*np.nanstd(x[:,indx]) ) ) # auto-select range based on sampling *RETAINED* (i.e., passing cut). Note that for the coordinates I usually use, it would be nonsensical to make the range in coordinate too small, as can occasionally happens print(" GP: Input sample size ", len(x), len(y)) print(" GP: Estimated length scales ") @@ -1412,10 +1578,15 @@ def fit_gp(x,y,x0=None,symmetry_list=None,y_errors=None,hypercube_rescale=False, alpha = y_errors**2 # added to diagonal of kernel, used to assign variances of measurements a priori; note also WhiteKernel also absorbs some of this if not (hypercube_rescale): # These parameters have been hand-tuned by experience to try to set to levels comparable to typical lnL Monte Carlo error - kernel = WhiteKernel(noise_level=0.1,noise_level_bounds=(1e-2,1))+C(0.5, (1e-3,1e1))*RBF(length_scale=length_scale_est, length_scale_bounds=length_scale_bounds_est) + noise_bounds_here = _gp_bounds_opt(opts.fit_gp_noise_bounds) + amp_bounds_here = _gp_bounds_opt(opts.fit_gp_amplitude_bounds) + kernel = WhiteKernel(noise_level=0.1,noise_level_bounds=noise_bounds_here)+C(0.5, amp_bounds_here)*RBF(length_scale=length_scale_est, length_scale_bounds=length_scale_bounds_est) gp = GaussianProcessRegressor(kernel=kernel, alpha=alpha, n_restarts_optimizer=8) gp.fit(x,y) + report_gp_kernel(gp, x, y, holdout_folds=opts.fit_gp_holdout_folds, + kernel_proto=kernel, alpha_proto=alpha, + peak_index=_mc_index_in(coord_names)) print(" Fit: std: ", np.std(y - gp.predict(x)), "using number of features ", len(y)) @@ -1566,7 +1737,7 @@ def fit_xg(x,y,y_errors=None,fname_export='nn_fit',verbose=False): import xgboost as xgb # Instantiate model. Usually not that many structures to find, don't overcomplicate # - should scale like number of samples - rf = xgb.XGBRegressor(n_estimators=100) # no more than 5% of samples in a leaf + rf = xgb.XGBRegressor(n_estimators=100) if y_errors is None: rf.fit(x,y) else: @@ -1627,7 +1798,7 @@ def fit_rf(x,y,y_errors=None,fname_export='nn_fit',verbose=False): from sklearn.ensemble import ExtraTreesRegressor # Instantiate model. Usually not that many structures to find, don't overcomplicate # - should scale like number of samples - rf = ExtraTreesRegressor(n_estimators=100, verbose=verbose,n_jobs=-1) # no more than 5% of samples in a leaf + rf = ExtraTreesRegressor(n_estimators=100, verbose=verbose,n_jobs=-1) if y_errors is None: rf.fit(x,y) else: @@ -1653,7 +1824,6 @@ def fn_return(x_in,rf=rf): def fit_rf_pca(x,y,y_errors=None,fname_export='nn_fit'): # from aasim -# from sklearn.ensemble import RandomForestRegressor from sklearn.ensemble import ExtraTreesRegressor from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler @@ -1663,38 +1833,35 @@ def fit_rf_pca(x,y,y_errors=None,fname_export='nn_fit'): x_pca = pca.fit_transform(x_scaled) # Instantiate model. Usually not that many structures to find, don't overcomplicate # - should scale like number of samples - rf = ExtraTreesRegressor(n_estimators=100, verbose=True,n_jobs=-1) # no more than 5% of samples in a leaf + rf = ExtraTreesRegressor(n_estimators=100, verbose=True,n_jobs=-1) if y_errors is None: rf.fit(x_pca,y) else: rf.fit(x_pca,y,sample_weight=1./y_errors**2) - ### reject points with infinities : problems for inputs def fn_return(x_in,rf=rf): f_out = -100000*np.ones(len(x_in)) - # remove infinity or Nan indx_ok = np.all(np.isfinite(x_in),axis=-1) - # rf internally uses float32, so we need to remove points > 10^37 or so ! - # ... this *should* never happen due to bounds constraints, but ... - indx_ok_size = np.all( np.logical_not(np.greater(np.abs(x_in),1e37)), axis=-1) + indx_ok_size = np.all( + np.logical_not(np.greater(np.abs(x_in),1e37)), axis=-1 + ) indx_ok = np.logical_and(indx_ok, indx_ok_size) - - f_out[indx_ok] = rf.predict(pca.transform(x_scaler.transform(x_in[indx_ok]))) + f_out[indx_ok] = rf.predict( + pca.transform(x_scaler.transform(x_in[indx_ok])) + ) return f_out -# fn_return = lambda x_in: rf.predict(x_in) - print( " Demonstrating RF") # debugging + print( " Demonstrating RF") residuals = rf.predict(pca.transform(x_scaler.transform(x)))-y print( " std ", np.std(residuals), np.max(y), np.max(fn_return(x))) return fn_return def fit_rbf(x,y,y_errors=None,fname_export='rbf_fit',verbose=False): from scipy.interpolate import RBFInterpolator - # - should scale like number of samples rbf = RBFInterpolator(x,y) - print( " Demonstrating RBF") # debugging + print( " Demonstrating RBF") residuals = rbf(x)-y print( " std ", np.std(residuals), np.max(y), np.max(rbf(x))) return rbf @@ -1879,6 +2046,9 @@ def fit_gp_sparse(x): ### # id m1 m2 lnL sigma/L neff col_lnL = 9 +col_a6c = None +col_E0 = None +col_pphi0 = None col_eccentricity = None col_meanPerAno = None col_lambda1 = None @@ -1900,20 +2070,25 @@ def fit_gp_sparse(x): low_level_coord_names += ['ordering'] print(" Revised fit coord names (for lookup) : ", coord_names) # 'eos_table_index' will be overwritten here print(" Revised sampling coord names : ", low_level_coord_names) -elif opts.use_eccentricity: - print(" Eccentricity input: [",ECC_MIN, ", ",ECC_MAX, "]") +if opts.use_EOB_parameters: + print(" EOB parameters (a6c) input: [", A6C_MIN, ", ", A6C_MAX, "]") + col_lnL += 1 + col_a6c = col_lnL - 1 +if opts.use_hyperbolic: + print(" E0 input: [", E0_MIN, ", ", E0_MAX, "]") + print(" p_phi0 input: [", PPHI0_MIN, ", ", PPHI0_MAX, "]") + col_lnL += 2 + col_E0 = col_lnL - 2 + col_pphi0 = col_lnL - 1 +if opts.use_eccentricity: + print(" Eccentricity input: [", ECC_MIN, ", ", ECC_MAX, "]") + col_lnL += 1 + col_eccentricity = col_lnL - 1 if opts.use_meanPerAno: print(" Also using meanPerAno ") - # perform modulus on desired row - col_lnL+=2 - col_meanPerAno = col_lnL -1 - col_eccentricity = col_lnL -2 - else: col_lnL += 1 - col_eccentricity = col_lnL -1 -if opts.input_distance: - print(" Distance input") - col_lnL +=1 + col_meanPerAno = col_lnL - 1 + # ---------------------------------------------------------------------- # Hyperpipeline ASCII input path (opt-in via env var or auto-detected). # When active, we (a) read a header-bearing file via hyperpipeline_io, @@ -1938,6 +2113,12 @@ def fit_gp_sparse(x): _use_tides = bool(opts.input_tides) or _has("lambda1") _use_eos = bool(opts.input_eos_index) or _has("eos_table_index") _use_dist = bool(opts.input_distance) or _has("distance") + _use_eob = bool(opts.use_EOB_parameters) + _use_hyp = bool(opts.use_hyperbolic) + if _has("a6c") and not _use_eob: + parser.error("Hyperpipeline input contains a6c; pass --use-EOB-parameters") + if (_has("E0") or _has("p_phi0")) and not _use_hyp: + parser.error("Hyperpipeline input contains E0/p_phi0; pass --use-hyperbolic") # LISA sky: ecliptic_longitude/latitude are NAMED columns -> carry them # through (aliased to P.phi/P.theta), so the sky is fit/sampled like any # other coordinate. No positional all.net hacking. @@ -1945,14 +2126,20 @@ def fit_gp_sparse(x): dat = _hpio.to_legacy_dat(_arr, use_eccentricity=_use_ecc, use_meanPerAno=_use_mpa, use_tides=_use_tides, use_eos_index=_use_eos, - use_distance=_use_dist, use_sky=_use_sky) + use_distance=_use_dist, + use_eob_parameters=_use_eob, + use_hyperbolic=_use_hyp, use_sky=_use_sky) _ix = _hpio.legacy_column_indices( use_eccentricity=_use_ecc, use_meanPerAno=_use_mpa, use_tides=_use_tides, use_eos_index=_use_eos, - use_distance=_use_dist, use_sky=_use_sky) + use_distance=_use_dist, use_eob_parameters=_use_eob, + use_hyperbolic=_use_hyp, use_sky=_use_sky) col_lnL = _ix["lnL"] col_distance = _ix["distance"] col_lambda1 = _ix["lambda1"] + col_a6c = _ix["a6c"] + col_E0 = _ix["E0"] + col_pphi0 = _ix["p_phi0"] col_eccentricity = _ix["eccentricity"] col_meanPerAno = _ix["meanPerAno"] col_ecliptic_longitude = _ix["ecliptic_longitude"] @@ -2050,6 +2237,8 @@ def fit_gp_sparse(x): P.lambda2 = line[col_lambda1+1] if opts.input_eos_index: P.eos_table_index = line[col_lambda1+2] + if opts.use_EOB_parameters: + P.a6c = line[col_a6c] # 9 if opts.use_eccentricity: P.eccentricity = line[col_eccentricity] # 9 if opts.use_meanPerAno: @@ -2057,6 +2246,9 @@ def fit_gp_sparse(x): # P.eccentricity = line[9] # if opts.use_meanPerAno: # P.meanPerAno = line[10] + if opts.use_hyperbolic: + P.E0 = line[col_E0] + P.p_phi0 = line[col_pphi0] if opts.input_distance: P.dist = lal.PC_SI*1e6*line[col_distance] # 9. Previously incompatible with tides when hardcoded if _use_sky: @@ -2447,12 +2639,10 @@ def my_fit(x): my_fit = fit_rf(X,Y,y_errors=Y_err) elif opts.fit_method == 'rf_pca': print( " FIT METHOD ", opts.fit_method, " IS RF-pca ") - # NO data truncation for NN needed? To be *consistent*, have the code function the same way as the others X=X[indx_ok] Y=Y[indx_ok] - lnL_shift Y_err = Y_err[indx_ok] - dat_out_low_level_coord_names = dat_out_low_level_coord_names[indx_ok] - # Cap the total number of points retained, AFTER the threshold cut + dat_out_low_level_coord_names = dat_out_low_level_coord_names[indx_ok] if opts.cap_points< len(Y) and opts.cap_points> 100: n_keep = opts.cap_points indx = np.random.choice(np.arange(len(Y)),size=n_keep,replace=False) @@ -2463,12 +2653,10 @@ def my_fit(x): my_fit = fit_rf_pca(X,Y,y_errors=Y_err) elif opts.fit_method == 'rbf': print( " FIT METHOD ", opts.fit_method, " IS RBF; **errors not used! **") - # NO data truncation for NN needed? To be *consistent*, have the code function the same way as the others X=X[indx_ok] Y=Y[indx_ok] - lnL_shift Y_err = Y_err[indx_ok] - dat_out_low_level_coord_names = dat_out_low_level_coord_names[indx_ok] - # Cap the total number of points retained, AFTER the threshold cut + dat_out_low_level_coord_names = dat_out_low_level_coord_names[indx_ok] if opts.cap_points< len(Y) and opts.cap_points> 100: n_keep = opts.cap_points indx = np.random.choice(np.arange(len(Y)),size=n_keep,replace=False) @@ -2683,9 +2871,8 @@ def convert_coords(x_in): print('PORTFOLIO: adding {} '.format(name)) sampler_list.append(sampler) sampler = mcsamplerPortfolio.MCSampler(portfolio=sampler_list) -elif mcsampler_Portfolio_ok: - if opts.sampler_method in mcsamplerPortfolio.known_pipelines: # access from plugins - sampler = mcsamplerPortfolio.known_pipelines[opts.sampler_method]() +elif opts.sampler_method in mcsamplerPortfolio.known_pipelines: # access from plugins + sampler = mcsamplerPortfolio.known_pipelines[opts.sampler_method]() ### @@ -3124,6 +3311,7 @@ def parse_corr_params(my_str): supplemental_ln_likelihood_offset = float(supplemental_ln_likelihood_offset_fn()) print(" EXTERNAL SUPPLEMENTARY LIKELIHOOD FACTOR : restoring offset {} in reported lnL/evidence ".format(supplemental_ln_likelihood_offset)) ln_integrand_value_absolute = ln_integrand_value + supplemental_ln_likelihood_offset +sigma_integral = relative_mc_error(res, var, log_space=opts.internal_use_lnL) # Test n_eff threshold if not (opts.fail_unless_n_eff is None): @@ -3214,12 +3402,12 @@ def parse_corr_params(my_str): annotation_header = linefirst # this will/must be lnL sigma_lnL and then parameter names, which we want to preserve with open(opts.fname_output_integral+"+annotation.dat", 'w') as file_out: if not(opts.using_eos) or not(opts.using_eos.startswith('file:')): - str_out =list( map(str,[ln_integrand_value_absolute, np.sqrt(var)/res, neff])) + str_out =list( map(str,[ln_integrand_value_absolute, sigma_integral, neff])) file_out.write("# " + annotation_header + "\n") file_out.write(' '.join( str_out + eos_extra + ["\n"])) else: file_out.write("# " + annotation_header + "\n") - file_out.write(" {} {} ".format(ln_integrand_value_absolute, np.sqrt(var)/res) + ' '.join(map(str,params_here))) + file_out.write(" {} {} ".format(ln_integrand_value_absolute, sigma_integral) + ' '.join(map(str,params_here))) #np.savetxt(opts.fname_output_integral+"+annotation.dat", np.array([[np.log(res), np.sqrt(var)/res, neff]]), header=eos_extra) # since not EOS, can just use np.savetxt # with open(opts.fname_output_integral+"+annotation_ESS.dat", 'w') as file_out: @@ -3272,7 +3460,7 @@ def parse_corr_params(my_str): weights_scaled = weights_scaled/np.max(weights_scaled) # try to reduce dynamic range n_ESS = np.sum(weights_scaled)**2/np.sum(weights_scaled**2) print(" n_eff n_ESS ", neff, n_ESS) -np.savetxt(opts.fname_output_integral+"+annotation_ESS.dat",[[ln_integrand_value_absolute, np.sqrt(var)/res, neff, n_ESS]],header=" lnL sigmaL neff n_ESS ") +np.savetxt(opts.fname_output_integral+"+annotation_ESS.dat",[[ln_integrand_value_absolute, sigma_integral, neff, n_ESS]],header=" lnL sigmaL neff n_ESS ") # Throw away stupid points that don't impact the posterior @@ -3671,6 +3859,42 @@ def parse_corr_params(my_str): if opts.verbose: print(" Sample: Skipping " , line, ' due to ', p, val, downselect_dict[p]) + if opts.force_scatter: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'scatter': + include_item = True + else: + include_item = False + + if opts.force_plunge: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-plunge points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'plunge': + include_item = True + else: + include_item = False + + if opts.force_zoomwhirl: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-zoomwhirl points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'zoomwhirl': + include_item = True + else: + include_item = False + # Set some superfluous quantities, needed only for PN approximants, so the result is generated sensibly Pgrid.ampO =opts.amplitude_order Pgrid.phaseO =opts.phase_order diff --git a/MonteCarloMarginalizeCode/Code/bin/util_FrameZeroNoiseSNR.py b/MonteCarloMarginalizeCode/Code/bin/util_FrameZeroNoiseSNR.py index e93ad2581..bb40ae42b 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_FrameZeroNoiseSNR.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_FrameZeroNoiseSNR.py @@ -15,6 +15,7 @@ parser.add_argument("--fmin-snr", type=float, default=10) parser.add_argument("--fmax-snr", type=float, default=1500) parser.add_argument("--plot-sanity",action='store_true') +parser.add_argument("--verbose",action='store_true') opts= parser.parse_args() fminSNR = opts.fmin_snr @@ -48,7 +49,7 @@ else: analyticPSD_Q=False print("Reading PSD for instrument %s from %s" % (ifo, psd_name[ifo])) - psd_dict[ifo] = lalsimutils.load_resample_and_clean_psd(psd_name[ifo], ifo, df) + psd_dict[ifo] = lalsimutils.load_resample_and_clean_psd(psd_name[ifo], ifo, df,verbose=opts.verbose) IP = lalsimutils.ComplexIP(fLow=fminSNR, fNyq=fSample/2,deltaF=df,psd=psd_dict[ifo],fMax=fmaxSNR,analyticPSD_Q=analyticPSD_Q) rhoDet = rho_dict[ifo] = IP.norm(data_dict[ifo]) print(ifo, rho_dict[ifo]) @@ -86,4 +87,3 @@ with open("snr-report.txt", 'w') as f: json.dump(rho_dict, f) f.flush() - diff --git a/MonteCarloMarginalizeCode/Code/bin/util_ILEdagPostprocess.sh b/MonteCarloMarginalizeCode/Code/bin/util_ILEdagPostprocess.sh index d513c9385..bafbf3392 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_ILEdagPostprocess.sh +++ b/MonteCarloMarginalizeCode/Code/bin/util_ILEdagPostprocess.sh @@ -8,8 +8,17 @@ DIR_PROCESS=$1 BASE_OUT=$2 -ECC=$3 # Liz (Capstone): this will only be non-blank in the case where my eccentric PE Makefile has inserted "--eccentricity" into join.sub -MPA=$4 +# Everything after the first two arguments is the advanced-physics flag list +# handed to util_CleanILE.py (--eccentricity, --meanPerAno, --a6c, +# --hyperbolic, --tabular-eos-file, ...). BasicIteration can enable several +# groups at once, so forward ALL of them: selecting one flag and dropping the +# rest made the cleaner parse rows with a layout the run never wrote. +CLEAN_FLAGS=() +for arg in "${@:3}"; do + if [ -n "$arg" ]; then + CLEAN_FLAGS+=("$arg") + fi +done # -------------------------------------------------------------------------- # Hyperpipeline ASCII output path (opt-in via env var). @@ -39,15 +48,22 @@ case "$(echo "${RIFT_HYPERPIPELINE_FORMAT:-}" | tr '[:upper:]' '[:lower:]')" in # clean them (=join duplicate lines) echo " Consolidating multiple instances of the monte carlo .... " - if [ "$3" == '--eccentricity' ]; then - if [ "$4" == '--meanPerAno' ]; then - util_CleanILE.py ${RND}_tmp.dat $3 $4 | sort -rg -k12 > $BASE_OUT.composite - else - util_CleanILE.py ${RND}_tmp.dat $3 | sort -rg -k11 > $BASE_OUT.composite - fi + util_CleanILE.py ${RND}_tmp.dat "${CLEAN_FLAGS[@]}" > ${RND}_clean.dat + + # Sort on lnL. The composite row is + # (event_id, intrinsic..., lnL, sigma_lnL, ntotal, neff) + # so lnL is ALWAYS the 4th field from the end, whichever advanced-physics + # groups are enabled; derive the key from the row width instead of + # hard-coding one column index per flag combination (which silently + # mis-sorted, i.e. discarded the composite ordering, for combined runs). + NCOL=`awk 'NF>0 && $1 !~ /^#/ {print NF; exit}' ${RND}_clean.dat` + if [ -z "${NCOL}" ] || [ "${NCOL}" -lt 5 ]; then + echo " WARNING: no usable rows in consolidated ILE output " + cp ${RND}_clean.dat $BASE_OUT.composite else - util_CleanILE.py ${RND}_tmp.dat $3 | sort -rg -k10 > $BASE_OUT.composite + sort -rg -k$((NCOL-3)) ${RND}_clean.dat > $BASE_OUT.composite fi + rm -f ${RND}_clean.dat ;; esac diff --git a/MonteCarloMarginalizeCode/Code/bin/util_LALWriteFrame.py b/MonteCarloMarginalizeCode/Code/bin/util_LALWriteFrame.py index f3643d737..9912da8ae 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_LALWriteFrame.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_LALWriteFrame.py @@ -38,7 +38,11 @@ parser.add_argument("--incl",default=None,help="Set the inclination of L (at fref). Particularly helpful for aligned spin tests") parser.add_argument("--mass1",default=10,type=float,help='Mass 1 (solar masses)') parser.add_argument("--mass2",default=1.4,type=float,help='Mass 2 (solar masses)') +parser.add_argument("--l-max",default=4,type=int,help='Maximum spherical-harmonic degree in the injection') parser.add_argument("--verbose", action="store_true",default=False) +parser.add_argument("--use-hlms-as-injections", action="store_true",default=False) +parser.add_argument('--hyperbolic', action='store_true', help='skips tapering for hyperbolic waveforms') +parser.add_argument('--force-hyperbolic-22', action='store_true', help='Forces just the 22 modes for hyperbolic waveforms') opts= parser.parse_args() @@ -50,7 +54,8 @@ P.randomize(aligned_spin_Q=True,default_inclination=opts.incl) P.m1 = opts.mass1*lalsimutils.lsu_MSUN P.m2 = opts.mass2*lalsimutils.lsu_MSUN - P.taper = lalsimutils.lsu_TAPER_START + if not(opts.hyperbolic): + P.taper = lalsimutils.lsu_TAPER_START P.tref =1000000000 # default if opts.approx: P.approx = lalsim.GetApproximantFromString(str(opts.approx)) @@ -64,29 +69,43 @@ xmldoc = utils.load_filename(filename, verbose = True, contenthandler =lalsimutils.cthdler) sim_inspiral_table = lsctables.SimInspiralTable.get_table(xmldoc) P.copy_sim_inspiral(sim_inspiral_table[int(event)]) - P.taper = lalsimutils.lsu_TAPER_START + if not(opts.hyperbolic): + P.taper = lalsimutils.lsu_TAPER_START if opts.approx: P.approx = lalsim.GetApproximantFromString(str(opts.approx)) -P.taper = lalsimutils.lsu_TAPER_START # force taper +if not(opts.hyperbolic): + P.taper = lalsimutils.lsu_TAPER_START # force taper P.detector = opts.instrument if opts.approx == "EccentricTD": P.phaseO = 3 P.print_params() - -T_est = lalsimutils.estimateWaveformDuration(P) -T_est = P.deltaT*lalsimutils.nextPow2(T_est/P.deltaT) -if T_est > opts.seglen: - print(" WARNING: THE SIGNAL WILL LIKELY BE TRUNCATED when writing the frame, which is VERY BAD ") -T_est =opts.seglen -P.deltaF = 1./T_est -print(" Duration ", T_est) -if T_est < opts.seglen: - print(" Buffer length too short, automating retuning forced ") - +if not opts.hyperbolic: + T_est = lalsimutils.estimateWaveformDuration(P) + T_est = P.deltaT*lalsimutils.nextPow2(T_est/P.deltaT) + if T_est > opts.seglen: + print(" WARNING: THE SIGNAL WILL LIKELY BE TRUNCATED when writing the frame, which is VERY BAD ") + T_est =opts.seglen + P.deltaF = 1./T_est + print(" Duration ", T_est) + if T_est < opts.seglen: + print(" Buffer length too short, automating retuning forced ") +else: + T_est =opts.seglen + P.deltaF = 1./T_est # Generate signal -hoft = lalsimutils.hoft(P) # include translation of source, but NOT interpolation onto regular time grid +#hoft = lalsimutils.hoft(P) # include translation of source, but NOT interpolation onto regular time grid +if not (opts.use_hlms_as_injections): + print("Injecting with hoft") + hoft = lalsimutils.hoft(P,approx_string=opts.approx) +else: + print("Injecting with hlms") + if opts.hyperbolic and opts.force_hyperbolic_22: + hlm = lalsimutils.hlmoft(P,Lmax=opts.l_max,force_22_mode=True) + else: + hlm = lalsimutils.hlmoft(P,Lmax=opts.l_max) + hoft = lalsimutils.hoft_from_hlm(hlm,P) epoch_orig = hoft.epoch # zero pad to be opts.seglen long, if necessary if opts.seglen/hoft.deltaT > hoft.data.length: @@ -106,7 +125,7 @@ hoft = ht if opts.stop and hoft.epoch+hoft.data.length*hoft.deltaT < opts.stop: - nToAddAtEnd = int( (-(hoft.epoch+hoft.data.length*hoft.deltaT)+opts.stop)/hoft.deltaT) + nToAddAtEnd = int( (-(hoft.epoch+hoft.data.length*hoft.deltaT)+opts.stop)/hoft.deltaT)+1 print("Padding end ", nToAddAtEnd, hoft.data.length) hoft = lal.ResizeREAL8TimeSeries(hoft,0, int(hoft.data.length+nToAddAtEnd)) diff --git a/MonteCarloMarginalizeCode/Code/bin/util_ManualOverlapGrid.py b/MonteCarloMarginalizeCode/Code/bin/util_ManualOverlapGrid.py index 8aaac76b5..6f71bb195 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_ManualOverlapGrid.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_ManualOverlapGrid.py @@ -245,7 +245,14 @@ parser.add_argument("--verbose", action="store_true",default=False, help="Extra warnings") parser.add_argument("--extra-verbose", action="store_true",default=False, help="Lots of messages") parser.add_argument("--save-plots",default=False,action='store_true', help="Write plots to file (only useful for OSX, where interactive is default") +parser.add_argument('--force-scatter',default=False,action='store_true',help='For hyperbolic analyses forces only scatter grid points.') +parser.add_argument('--force-plunge',default=False,action='store_true',help='For hyperbolic analyses forces only scatter grid points.') +parser.add_argument('--force-zoomwhirl',default=False,action='store_true',help='For hyperbolic analyses forces only scatter grid points.') opts= parser.parse_args() + +force_options = [opts.force_scatter, opts.force_plunge, opts.force_zoomwhirl] # Add more if needed +if sum(bool(x) for x in force_options) > 1: + parser.error("CANNOT use multiple --force-X options at the same time!") if opts.inj_file_out: opts.fname = opts.inj_file_out.replace(".xml.gz","") @@ -366,6 +373,39 @@ def evaluate_overlap_on_grid(hfbase,param_names, grid): for param in downselect_dict: if Pgrid.extract_param(param) < downselect_dict[param][0] or Pgrid.extract_param(param) > downselect_dict[param][1]: include_item =False + if opts.force_scatter: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'scatter': + include_item = True + else: + include_item = False + if opts.force_plunge: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'plunge': + include_item = True + else: + include_item = False + if opts.force_zoomwhirl: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = Pgrid.extract_param('hypclass') + if hypclass == 'zoomwhirl': + include_item = True + else: + include_item = False if include_item: grid_revised.append(line) if Pgrid.m2 <= Pgrid.m1: # do not add grid elements with m2> m1, to avoid possible code pathologies ! @@ -439,7 +479,7 @@ def evaluate_overlap_on_grid(hfbase,param_names, grid): filename = opts.inj event = opts.event_id xmldoc = utils.load_filename(filename, verbose = True,contenthandler =lalsimutils.cthdler) - sim_inspiral_table = lsctables.Table.get_table(xmldoc, lsctables.SimInspiralTable.tableName) + sim_inspiral_table = lsctables.SimInspiralTable.get_table(xmldoc) P.copy_sim_inspiral(sim_inspiral_table[int(event)]) P.fmin =opts.fmin if opts.approx: diff --git a/MonteCarloMarginalizeCode/Code/bin/util_NRWriteFrame.py b/MonteCarloMarginalizeCode/Code/bin/util_NRWriteFrame.py index 25f23465c..01ebf5335 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_NRWriteFrame.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_NRWriteFrame.py @@ -32,6 +32,8 @@ parser.add_argument("--instrument", default="H1",help="Use H1, L1,V1") parser.add_argument("--inj", dest='inj', default=None,help="inspiral XML file containing injection information. Used for extrinsic information only") parser.add_argument("--nr-perturbative-extraction",default=False,action='store_true') +parser.add_argument("--nr-perturbative-extraction-full",default=False,action='store_true') +parser.add_argument("--nr-use-provided-strain",default=False,action='store_true') parser.add_argument("--nr-use-hybrid",action='store_true') parser.add_argument("--mass", default=150.0,type=float,help="Total mass in solar masses") # 150 turns out to be ok for Healy et al sims parser.add_argument("--lmax", default=2, type=int) @@ -73,7 +75,7 @@ T_window = 16. # default # Load catalog -wfP = nrwf.WaveformModeCatalog(opts.group, param, clean_initial_transient=True,clean_final_decay=True, shift_by_extraction_radius=True, extraction_radius=opts.rextr,lmax=opts.lmax,align_at_peak_l2_m2_emission=True,build_strain_and_conserve_memory=True,perturbative_extraction=opts.nr_perturbative_extraction) +wfP = nrwf.WaveformModeCatalog(opts.group, param, clean_initial_transient=True,clean_final_decay=True, shift_by_extraction_radius=True, extraction_radius=opts.rextr,lmax=opts.lmax,align_at_peak_l2_m2_emission=True,build_strain_and_conserve_memory=True,perturbative_extraction=opts.nr_perturbative_extraction,perturbative_extraction_full=opts.nr_perturbative_extraction_full,use_provided_strain=opts.nr_use_provided_strain) # Generate signal @@ -193,7 +195,7 @@ import os # from matplotlib import pyplot as plt # First must create corresponding cache file - os.system("echo "+ fname+ " | lalapps_path2cache > test.cache") + os.system("echo "+ fname+ " | lal_path2cache > test.cache") # Now I can read it # Beware that the results are OFFSET FROM ONE ANOTHER due to PADDING, # but that the time associations are correct @@ -219,4 +221,3 @@ plt.xlim(min(tvals2),max(tvals2)) # full range with pad plt.figure(indx); plt.savefig("nr-framedump-full-" +str(indx)+fig_extension) - diff --git a/MonteCarloMarginalizeCode/Code/bin/util_ParameterPuffball.py b/MonteCarloMarginalizeCode/Code/bin/util_ParameterPuffball.py index f8a089570..940647979 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_ParameterPuffball.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_ParameterPuffball.py @@ -59,8 +59,15 @@ parser.add_argument("--reflect-parameter",action='append',type=str) parser.add_argument("--enforce-duration-bound",default=None,type=float,help="If present, enforce a duration bound. Used to prevent grid placement for obscenely long signals, when the window size is prescribed") parser.add_argument("--regularize",action='store_true',help="Add some ad-hoc terms based on priors, to help with nearly-singular matricies") +parser.add_argument('--force-scatter',default=False,action='store_true',help='For hyperbolic analyses forces only scatter grid points.') +parser.add_argument('--force-plunge',default=False,action='store_true',help='For hyperbolic analyses forces only plunge grid points.') +parser.add_argument('--force-zoomwhirl',default=False,action='store_true',help='For hyperbolic analyses forces only zoomwhirl grid points.') opts= parser.parse_args() +force_options = [opts.force_scatter, opts.force_plunge, opts.force_zoomwhirl] # Add more if needed +if sum(bool(x) for x in force_options) > 1: + parser.error("CANNOT use multiple --force-X options at the same time!") + if opts.random_parameter is None: opts.random_parameter = [] @@ -346,6 +353,39 @@ # val = val/ lal.MSUN_SI # if val < downselect_dict[param][0] or val > downselect_dict[param][1]: # include_item =False + if opts.force_scatter: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = P.extract_param('hypclass') + if hypclass == 'scatter': + include_item = True + else: + include_item = False + if opts.force_plunge: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = P.extract_param('hypclass') + if hypclass == 'plunge': + include_item = True + else: + include_item = False + if opts.force_zoomwhirl: + if include_item==False: + # no need to evaluate if the point is already downselected out + pass + else: + # removes non-scatter points from the hyperbolic grid + hypclass = P.extract_param('hypclass') + if hypclass == 'zoomwhirl': + include_item = True + else: + include_item = False if include_item: P_out.append(P) diff --git a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py index 95e93ce22..656c072ad 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py @@ -58,6 +58,13 @@ # leaf module: numpy only, so this does not drag numba/cupy into the pipeline script from RIFT.likelihood.time_interp_choice import ( BARE_FLAG_SENTINEL, CROSSOVER_GUIDANCE, resolve_interpolate_time_request) +# Same reason (numpy-only leaf module), and IMPORTED rather than re-typed: a second +# hand-written copy of the choice tuple is how a typo becomes a silently different +# likelihood -- the pipeline would accept 'bandlimted', forward it, and the mistake +# would only surface when the first ILE job died. +from RIFT.likelihood.time_marginalization_quadrature import ( + TIME_QUADRATURE_CHOICES, validate_time_quadrature, + refuse_unhonourable_time_quadrature, refuse_unless_time_quadrature_emitted) ligolw_prefix = 'igwn_' if not(which(ligolw_prefix + "ligolw_add")): ligolw_prefix = '' @@ -137,6 +144,15 @@ def retrieve_event_from_coinc(fname_coinc): else: event_dict["eccentricity"] = None event_dict["meanPerAno"] = None + try: + event_dict["E0"] = row.psi3 + except: + event_dict["E0"] = 0.0 + try: + event_dict["p_phi0"] = row.beta + except: + event_dict["p_phi0"] = 0.0 + event_dict["IFOs"] = list(set(ifo_list)) max_snr_idx = snr_list.index(max(snr_list)) event_dict['SNR'] = snr_list[max_snr_idx] @@ -349,7 +365,7 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--skip-reproducibility",action='store_true') parser.add_argument("--use-production-defaults",action='store_true',help="Use production defaults. Intended for use with tools like asimov or by nonexperts who just want something to run on a real event. Will require manual setting of other arguments!") parser.add_argument("--use-subdags",action='store_true',help="Use CEPP_Alternate instead of CEPP_BasicIteration. Note this writes an adaptively-sized DAG each iteration, but doesn't otherwise optimize yet.") -parser.add_argument("--pipeline-builder",default=None,choices=["BasicIteration","AlternateIteration"],help="Explicitly select the create_event_parameter_pipeline_* iteration builder, as a drop-in hot-swap for side-by-side A/B testing. Overrides the implicit --use-subdags routing. If unset, the builder is chosen by --use-subdags (Alternate) vs. the default (Basic).") +parser.add_argument("--pipeline-builder",default=None,choices=["BasicIteration","AlternateIteration","BasicMultiApproxIteration"],help="Explicitly select the create_event_parameter_pipeline_* iteration builder, as a drop-in hot-swap for side-by-side A/B testing. Overrides the implicit --use-subdags routing. If unset, the builder is chosen by --use-subdags (Alternate) vs. the default (Basic).") parser.add_argument("--use-ile-subdags",action='store_true',help="Use ILE subdag system (new)") parser.add_argument("--bilby-ini-file",default=None,type=str,help="Pass ini file for parsing. Intended to use for calibration reweighting. Full path recommended") parser.add_argument("--bilby-pickle-file",default=None,type=str,help="Bilby Pickle file with event settings. Intended to use for calibration reweighting. Full path recommended") @@ -417,6 +433,10 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--calibration",default="C00",type=str) parser.add_argument("--playground-data",action='store_true', help="Passed through to helper_LDG_events, and changes name prefix") parser.add_argument("--approx",default=None,type=str,help="Approximant. REQUIRED") +parser.add_argument("--approx-extra",default=None,action='append',help="Additional waveform model, repeatable. Selects the cross-model workflow: every model is evaluated on ONE shared intrinsic grid and marginalized over point by point, and the terminal stage forks to give each model its own posterior and evidence. Implies --pipeline-builder BasicMultiApproxIteration. See RIFT/misc/DESIGN_multiapprox_marginalization.md") +parser.add_argument("--approx-prior",default=None,action='append',help="APPROX=WEIGHT prior p(m) over waveform models, repeatable. Default uniform. NOT sampling weights.") +parser.add_argument("--approx-extra-gwsignal",default=None,action='append',help="An --approx-extra model that must be generated through gwsignal. The primary --approx inherits --use-gwsignal automatically. The generator route is PER MODEL: the phenom family has no time-domain mode generator in gwsignal, so a single global route cannot serve an EOB-vs-phenom comparison.") +parser.add_argument("--require-all-approx",action='store_true',help="Drop intrinsic points not successfully evaluated under EVERY model, instead of marginalizing over whichever subset survived.") parser.add_argument("--use-gwsurrogate",action='store_true',help="Attempt to use gwsurrogate instead of lalsuite.") parser.add_argument("--use-gwsignal",action='store_true',help="Attempt to use gwsignal interface.") parser.add_argument("--l-max",default=2,type=int) @@ -431,10 +451,10 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--internal-tabular-eos-file",type=str,default=None,help="Tabular file of EOS to use. The default prior will be UNIFORM in this table!") parser.add_argument("--sample-eccentricity-squared",action='store_true', help="Option for sampling as well as fitting in eccentricity_squared instead of fitting in eccentricity_squared and sampling in eccentricity (also need option --use-eccentricity-squared") parser.add_argument("--use-eccentricity-squared",action='store_true', help="Allows for fitting and sampling in eccentricity_squared instead of eccentricity") -parser.add_argument("--sample-eccentricity-ln",action='store_true', help="Option for sampling as well as fitting in eccentricity_ln instead of fitting in eccentricity_ln and sampling in eccentricity (also need option --use-eccentricity-ln") -parser.add_argument("--use-eccentricity-ln",action='store_true', help="Allows for fitting and sampling in eccentricity_ln instead of eccentricity") parser.add_argument("--assume-eccentric",action='store_true', help="Add eccentric options for each part of analysis") parser.add_argument("--use-meanPerAno",action='store_true', help="Add meanPerAno options for each part of analysis") +parser.add_argument("--use-EOB-parameters",action='store_true', help="Add sampling in EOB parameters; currently only a6c") +parser.add_argument("--assume-hyperbolic",action='store_true', help="Add hyperbolic options for each part of analysis") parser.add_argument("--internal-cip-use-periodic-ecc-vars",action='store_true', help="use e cos ell, e sin ell as fitting variables ") parser.add_argument("--assume-lowlatency-tradeoffs",action='store_true', help="Force analysis with various low-latency tradeoffs (e.g., drop spin 2, use aligned, etc)") parser.add_argument("--assume-highq",action='store_true', help="Force analysis with the high-q strategy, neglecting spin2. Passed to 'helper'") @@ -474,6 +494,7 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--internal-ile-srate-time-resampling",default=None, help=" Adds --srate-resample-time-marginalization to ILE for output, to provide higher-resolution time output ") parser.add_argument("--internal-ile-srate-internal",default=None, help=" Adds --srate-internal to ILE, modifying how calculations are performed internally to use a higher sampling rate ") parser.add_argument("--internal-ile-interpolate-time",nargs='?',const=BARE_FLAG_SENTINEL,default=None,type=str,help="Enable sub-sample interpolation of Q_lm at fractional detector arrival times in the maintained NoLoop likelihood. REQUIRES AN EXPLICIT STENCIL: nearest|cubic|sinc -- automatic selection was removed as measurably unreliable, and a bare flag is rejected rather than silently doing nothing. MEASURED GUIDANCE (SEOBNRv4, an IMR model): %s. Forwarded verbatim to helper_LDG_Events.py, which validates it. Full tables, limitations and provenance: RIFT/likelihood/DESIGN_q_window_stencil.md." % CROSSOVER_GUIDANCE) +parser.add_argument("--internal-ile-time-marginalization-quadrature",default=None,type=str,choices=list(TIME_QUADRATURE_CHOICES),help="Rule for the TIME integral of the marginalized likelihood in ILE: %s. Default None = pass nothing, so the ILE default ('simpson', the historical fixed-deltaT Simpson rule) is unchanged and the emitted args_ile.txt is byte-identical to today. 'bandlimited' resolves the integrand instead of the data: exp(lnL(t)) is a peak of width sigma_t = 1/(2 pi rho sigma_f), which SHRINKS AS 1/rho, while the grid spacing deltaT=1/srate is fixed by the data -- so production under-resolves its own integrand, worse at higher SNR (measured: rigidly scanning the grid phase moves the reported lnL by 1.649 nats at srate 4096, rho=40). Forwarded verbatim to helper_LDG_Events.py, which validates it and puts --time-marginalization-quadrature on the ILE command line; from args_ile.txt it reaches every ILE*.sub INCLUDING ILE_extr.sub. REFUSED, not ignored, at DAG-BUILD TIME if this workflow cannot honour it (calibration marginalization, --rotation-slow, --freqresponse, or a configuration without --time-marginalization/--vectorized/--gpu). IMPORTANT -- INI OVERRIDE: the RIFT ini parser OVERRIDES the command line for non-boolean options, and this is a string option, so NEVER set it in a --use-ini that a Makefile or wrapper also sets on the command line; the ini value would win silently. Rationale, measured tables and exclusions: RIFT/likelihood/DESIGN_time_marginalization_quadrature.md." % ("|".join(TIME_QUADRATURE_CHOICES),)) parser.add_argument("--internal-ile-n-chunk",default=None,type=int,help="Override the extrinsic chunk size (--n-chunk) passed to ILE, via the helper. Default behaviour (helper): 40000, scaled linearly with SNR above 40 and capped at 160000, because at high SNR the posterior is a vanishing fraction of the prior volume and a small chunk gives few informative samples per adaptation step. Larger chunks cost GPU memory but measured HOST memory (what RequestMemory governs) is flat, so no memory-request change is normally needed. EXPERTS ONLY.") parser.add_argument("--batch-extrinsic",action='store_true') parser.add_argument("--fmin",default=20,type=int,help="Mininum frequency for integration. template minimum frequency (we hope) so all modes resolved at this frequency") # should be 23 for the BNS @@ -502,18 +523,26 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--force-lambda-max",default=None,type=float,help="Provide this value to override the value of lambda-max provided") parser.add_argument("--force-lambda-small-max",default=None,type=float,help="Provide this value to override the value of lambda-small-max provided") parser.add_argument("--force-lambda-no-linear-init",action='store_true',help="Disables use of priors focused towards small lambda for initial iterations. Designed for PP plot tests with wide/uniform priors.") -parser.add_argument("--force-chi-max",default=None,type=float,help="Provide this value to override the value of chi-max provided") -parser.add_argument("--force-chi-small-max",default=None,type=float,help="Provide this value to override the value of chi-max provided") -parser.add_argument("--force-ecc-max",default=None,type=float,help="Provide this value to override the value of ecc-max provided") -parser.add_argument("--force-ecc-min",default=None,type=float,help="Provide this value to override the value of ecc-min provided") -parser.add_argument("--force-comp-max",default=1000,type=float,help="Provide this value to override the value of the max component mass in CIP provided") -parser.add_argument("--force-comp-min",default=1,type=float,help="Provide this value to override the value of min component mass in CIP provided") -parser.add_argument("--force-meanPerAno-max",default=None,type=float,help="Provide this value to override the value of meanPerAno-max provided") -parser.add_argument("--force-meanPerAno-min",default=None,type=float,help="Provide this value to override the value of meanPerAno-min provided") + +parser.add_argument("--force-chi-max",default=None,type=float,help="Provde this value to override the value of chi-max provided") +parser.add_argument("--force-chi-small-max",default=None,type=float,help="Provde this value to override the value of chi-max provided") +parser.add_argument("--force-a6c-max",default=-20,type=float,help="Provde this value to override the value of a6c-max provided") +parser.add_argument("--force-a6c-min",default=-80,type=float,help="Provde this value to override the value of a6c-min provided") +parser.add_argument("--force-E0-max",default=None,type=float,help="Provide this value to override the value of E0-max provided") +parser.add_argument("--force-E0-min",default=None,type=float,help="Provide this value to override the value of E0-min provided") +parser.add_argument("--force-pphi0-max",default=None,type=float,help="Provide this value to override the value of pphi0-max provided") +parser.add_argument("--force-pphi0-min",default=None,type=float,help="Provide this value to override the value of pphi0-min provided") +parser.add_argument("--force-ecc-max",default=None,type=float,help="Provde this value to override the value of ecc-max provided") +parser.add_argument("--force-ecc-min",default=None,type=float,help="Provde this value to override the value of ecc-min provided") +parser.add_argument("--force-comp-max",default=1000,type=float,help="Provde this value to override the value of the max component mass in CIP provided") +parser.add_argument("--force-comp-min",default=1,type=float,help="Provde this value to override the value of min component mass in CIP provided") +parser.add_argument("--force-meanPerAno-max",default=None,type=float,help="Provde this value to override the value of meanPerAno-max provided") +parser.add_argument("--force-meanPerAno-min",default=None,type=float,help="Provde this value to override the value of meanPerAno-min provided") parser.add_argument("--scale-mc-range",type=float,default=None,help="If using the auto-selected mc, scale the ms range proposed by a constant factor. Recommend > 1. . ini file assignment will override this.") parser.add_argument("--limit-mc-range",default=None,type=str,help="Pass this argumen through to the helper to set the mc range") parser.add_argument("--force-mc-range",default=None,type=str,help="Pass this argumen through to the helper to set the mc range") parser.add_argument("--force-eta-range",default=None,type=str,help="Pass this argumen through to the helper to set the eta range") +parser.add_argument("--force-mtot-range",default=None,type=str,help="Pass this argument through to the helper to set the mtot range. Overrides mc parameter with mtot parameter broadly throughout the pipeline.") parser.add_argument("--allow-subsolar", action='store_true', help="Override limits which otherwise prevent subsolar mass PE") parser.add_argument("--force-hint-snr",default=None,type=str,help="Pass this argumen through to the helper to control source amplitude effects") parser.add_argument("--force-initial-grid-size",default=None,type=float,help="Only used for automated grids. Passes --force-initial-grid-size down to helper") @@ -606,6 +635,12 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--archive-pesummary-event-label",default="this_event",help="Label to use on the pesummary page itself") parser.add_argument("--internal-mitigate-fd-J-frame",default="L_frame",help="L_frame|rotate, choose method to deal with ChooseFDWaveform being in wrong frame. Default is to request L frame for inputs") parser.add_argument("--internal-force-puff-iterations", default=4, type=int, help="Number of iterations to be puffed") +parser.add_argument("--first-iteration-jumpstart",action='store_true',help="No ILE jobs the first iteration. Assumes you already have .composite files and want to get going. Particularly helpful for subdag systems") +parser.add_argument("--use-mtot-coords",action='store_true',help="Passed to the helper to configure CIP and PUFF for mtot instead of mc.") +parser.add_argument("--force-scatter-grids",action='store_true',help="Eliminates all non-scatter intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument("--force-plunge-grids",action='store_true',help="Eliminates all non-plunge intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument("--force-zoomwhirl-grids",action='store_true',help="Eliminates all non-zoomwhirl intrinsic points from hyperbolic grids throughout the workflow.") +parser.add_argument("--force-hyperbolic-22", action='store_true', help='Forces just the 22 modes for hyperbolic waveforms') opts= parser.parse_args() # Resolve the sub-sample stencil request IMMEDIATELY, so a bare flag / retired 'True' / typo @@ -613,6 +648,7 @@ def run_lisa_known_sky_surface(opts): # the call is for its validation side effect; the helper resolves it again for the emission. resolve_interpolate_time_request(opts.internal_ile_interpolate_time) + # Multi-GPU ILE fan-out: --ile-gpu-fanout funnels through RIFT_ILE_GPU_FANOUT, which # create_event_parameter_pipeline_BasicIteration (run via os.system, inheriting this # environment) and dag_utils read at DAG-build time to size request_GPUs/CPUs and bake @@ -620,6 +656,15 @@ def run_lisa_known_sky_surface(opts): if opts.ile_gpu_fanout is not None: os.environ['RIFT_ILE_GPU_FANOUT'] = str(opts.ile_gpu_fanout) +# Ensure --assume-hyperbolic is set when using any --force-X-grids option +# Ensure only ONE of the --force-X-grids options is set +force_grids = [opts.force_scatter_grids, opts.force_plunge_grids, opts.force_zoomwhirl_grids] +if any(force_grids) and not opts.assume_hyperbolic: + parser.error("Using --force-scatter-grids, --force-plunge-grids, or " + "--force-zoomwhirl-grids requires --assume-hyperbolic!") +if sum(bool(x) for x in force_grids) > 1: + parser.error("CANNOT use multiple --force-X-grids options at the same time!") + config_stored=None; config_dict=None ile_condor_commands = None if (opts.use_ini): @@ -686,6 +731,40 @@ def run_lisa_known_sky_surface(opts): if opts.ile_gpu_fanout is not None: os.environ['RIFT_ILE_GPU_FANOUT'] = str(opts.ile_gpu_fanout) +# TIME-MARGINALIZATION QUADRATURE, part 1 of 2: everything refusable WITHOUT running the +# helper. Deliberately placed AFTER the --use-ini block above: the ini parser OVERRIDES the +# command line for non-boolean options, so a validate above it checks a value that the ini is +# about to replace, and a bad ini value would surface downstream as "helper call failed to +# generate required file" instead of as a quadrature diagnostic. +if opts.internal_ile_time_marginalization_quadrature is not None: + # Validate through the LIBRARY function as well as argparse `choices`, so this script and + # the ILE driver can never disagree about what the legal set is. + validate_time_quadrature(opts.internal_ile_time_marginalization_quadrature) + if opts.lisa_known_sky: + # --lisa-known-sky exits below, before both the forward to the helper and the + # args_ile.txt guard, and builds its own ILE arguments through helper_LISA_Events.py + # which does not know this option. Refuse rather than silently dropping it. + raise ValueError( + "--internal-ile-time-marginalization-quadrature is not supported on the " + "--lisa-known-sky path: that path builds args_ile.txt through " + "helper_LISA_Events.py, which does not carry the option, so the request would be " + "silently dropped.") + # What this script knows before the helper runs: calibration marginalization is added HERE, + # not by the helper, and --manual-extra-ile-args can carry any ILE flag at all. + _tq_early = "" + if opts.calmarg_envelope_directory: + _tq_early += " --calibration-envelope-directory " + str(opts.calmarg_envelope_directory) + if opts.manual_extra_ile_args: + _tq_early += " " + str(opts.manual_extra_ile_args) + if _tq_early: + # Only the EXCLUSIONS are checkable this early -- the required flags are added by the + # helper -- so append the requirements to keep the message about what is actually wrong. + refuse_unhonourable_time_quadrature( + opts.internal_ile_time_marginalization_quadrature, + "--time-marginalization --vectorized --gpu " + _tq_early, + "this pipeline's own options (checked before the helper runs, so the failure is " + "immediate rather than after a workflow has been built)") + if opts.lisa_known_sky: run_lisa_known_sky_surface(opts) sys.exit(0) @@ -875,15 +954,33 @@ def run_lisa_known_sky_surface(opts): is_analysis_precessing =False is_analysis_eccentric =False +is_analysis_hyperbolic =False if opts.approx == "SEOBNRv3" or opts.approx == "NRSur7dq2" or opts.approx == "NRSur7dq4" or (opts.approx == 'SEOBNv3_opt') or (opts.approx == 'IMRPhenomPv2') or (opts.approx =="SEOBNRv4P" ) or (opts.approx == "SEOBNRv4PHM") or (opts.approx == "SEOBNRv5PHM") or ('SpinTaylor' in opts.approx) or ('IMRPhenomTP' in opts.approx or ('IMRPhenomXP' in opts.approx)): - is_analysis_precessing=True + is_analysis_precessing = True if opts.assume_precessing: - is_analysis_precessing = True + is_analysis_precessing = True if opts.assume_nonprecessing: - is_analysis_precessing = False + is_analysis_precessing = False if opts.assume_eccentric: - is_analysis_eccentric = True - + is_analysis_eccentric = True +if opts.assume_hyperbolic: + is_analysis_hyperbolic = True + +# Resolve ONE effective hyperbolic prior range for the whole workflow. Each +# consumer carries its own default and they do not agree: initial points are +# generated out to p_phi0=10, while CIP defaults to 5.4, so a bare +# --assume-hyperbolic evaluates likelihood at points the posterior then throws +# away. Fix the range here (defaults matching grid generation) and forward it +# explicitly to grid generation, CIP, and puffing, override or not. +hyperbolic_range = {'E0_min': 1.0, 'E0_max': 1.2, 'pphi0_min': 0.0, 'pphi0_max': 10.0} +for _name in hyperbolic_range: + _val = getattr(opts, 'force_' + _name) + if not(_val is None): + hyperbolic_range[_name] = _val +if is_analysis_hyperbolic: + if hyperbolic_range['E0_min'] >= hyperbolic_range['E0_max'] or hyperbolic_range['pphi0_min'] >= hyperbolic_range['pphi0_max']: + raise ValueError(" Empty hyperbolic prior range requested: check --force-E0-min/max and --force-pphi0-min/max ") + print(" Hyperbolic prior range (grid generation, CIP, puff): E0 [{},{}], p_phi0 [{},{}] ".format(hyperbolic_range['E0_min'],hyperbolic_range['E0_max'],hyperbolic_range['pphi0_min'],hyperbolic_range['pphi0_max'])) dirname_run = gwid+ "_" + opts.calibration+ "_"+ opts.approx+"_fmin" + str(fmin) +"_fmin-template"+str(fmin_template) +"_lmax"+str(opts.l_max) + "_"+opts.spin_magnitude_prior if opts.online: @@ -901,7 +998,9 @@ def run_lisa_known_sky_surface(opts): if opts.assume_eccentric: dirname_run += "_with_eccentricity" if opts.use_meanPerAno: - dirname_run += "_and_meanPerAno" + dirname_run += "_with_eccentricity_and_meanPerAno" +if opts.assume_hyperbolic: + dirname_run += "_with_hyperbolic" if opts.no_matter: dirname_run += "_no_matter" if opts.assume_highq: @@ -952,12 +1051,17 @@ def run_lisa_known_sky_surface(opts): P.eccentricity = event_dict["eccentricity"] if not(event_dict['meanPerAno'] is None): P.meanPerAno = event_dict["meanPerAno"] + if event_dict['E0'] is not None: + P.E0 = event_dict['E0'] + P.p_phi0 = event_dict['p_phi0'] # Write 'target_params' file -- hyperpipeline .dat or legacy XML. if _use_hpip_pp: from RIFT.misc import hyperpipeline_io as _hpio _cols = _hpio.build_column_list( use_eccentricity=(P.eccentricity != 0), - use_meanPerAno=(P.meanPerAno != 0)) + use_meanPerAno=(P.meanPerAno != 0), + use_eob_parameters=opts.use_EOB_parameters, + use_hyperbolic=opts.assume_hyperbolic) _hpio.write_grid_from_P_list("target_params", [P], _cols, lal_module=lal, lalsimutils_module=lalsimutils) @@ -1074,6 +1178,7 @@ def approx_supports_precession(approx_name): print(" [transverse-tails] raising NET interim sampling: cip-cap-neff -> {}, n-output-samples -> {} (worker count scaled x{} below); convergence test -> {}; puff tail-guard chi1_perp fraction {}".format(opts.internal_cip_cap_neff, opts.n_output_samples, opts.internal_cip_transverse_tails_worker_scale, opts.internal_test_convergence_method, opts.internal_cip_transverse_tails_puff_fraction)) cmd = " helper_LDG_Events.py --force-notune-initial-grid --propose-fit-strategy --propose-ile-convergence-options --fmin " + str(fmin) + " --fmin-template " + str(fmin_template) + " --working-directory " + base_dir + "/" + dirname_run + helper_psd_args + " --no-enforce-duration-bound --test-convergence " +cmd += " --internal-initial-grid-approximant {} ".format(opts.approx) if opts.internal_test_convergence_method: cmd += " --internal-test-convergence-method {} ".format(opts.internal_test_convergence_method) if opts.internal_use_gracedb_bayestar: @@ -1088,6 +1193,8 @@ def approx_supports_precession(approx_name): cmd+= " --propose-initial-grid " if opts.force_initial_grid_size: cmd += " --force-initial-grid-size {} ".format(int(opts.force_initial_grid_size)) +if opts.use_EOB_parameters: + cmd += " --use-EOB-parameters " if opts.assume_matter: cmd += " --assume-matter " npts_it = 1000 @@ -1111,6 +1218,20 @@ def approx_supports_precession(approx_name): if opts.use_meanPerAno: cmd += " --use-meanPerAno " npts_it = int(npts_it*1.5) +if is_analysis_hyperbolic: + cmd += " --assume-hyperbolic " + npts_it = int(npts_it*2.25) + # always pass the resolved range, so grid generation, CIP, and puff agree + cmd += " --E0-max {} --E0-min {} ".format(hyperbolic_range['E0_max'],hyperbolic_range['E0_min']) + cmd += " --pphi0-max {} --pphi0-min {} ".format(hyperbolic_range['pphi0_max'],hyperbolic_range['pphi0_min']) + if opts.force_scatter_grids: + cmd += " --force-scatter-grids " + if opts.force_plunge_grids: + cmd += " --force-plunge-grids " + if opts.force_zoomwhirl_grids: + cmd += " --force-zoomwhirl-grids " + if opts.force_hyperbolic_22: + cmd += " --force-hyperbolic-22 " if opts.assume_highq: cmd+= ' --assume-highq --force-grid-stretch-mc-factor 2' # the mc range, tuned to equal-mass binaries, is probably too narrow. Workaround until fixed in helper npts_it =1000 @@ -1157,6 +1278,10 @@ def approx_supports_precession(approx_name): cmd += " --scale-mc-range " + str(opts.scale_mc_range).replace(' ','') if not(opts.force_eta_range is None): cmd+= " --force-eta-range " + str(opts.force_eta_range).replace(' ','') +if not(opts.force_mtot_range is None): + cmd+= " --force-mtot-range " + str(opts.force_mtot_range).replace(' ','') +if opts.use_mtot_coords: + cmd+= " --use-mtot-coords " if opts.allow_subsolar: cmd += " --allow-subsolar " if opts.force_chi_max: @@ -1222,11 +1347,35 @@ def approx_supports_precession(approx_name): else: cmd += " --calibration-version " + opts.calibration if opts.use_online_psd_file: - # Get IFO list from ini file -## import ConfigParser -# config = ConfigParser.ConfigParser() -# config.read(opts.use_ini) - ifo_list = eval(config.get('analysis','ifos')) + # Which instruments does that PSD file cover? + # + # This used to read the ini's [analysis] ifos unconditionally, but `config` + # only exists on the --use-ini path, so --use-online-psd-file without an ini + # -- the natural way to run a synthetic injection -- died with + # "NameError: name 'config' is not defined". + # + # Prefer the ini when there is one (it is the run's declared instrument + # list), then an explicit --manual-ifo-list, and otherwise ask the PSD file + # itself, which names its own instruments and cannot disagree with itself. + ifo_list = None + if opts.use_ini: + ifo_list = eval(config.get('analysis','ifos')) + elif opts.manual_ifo_list: + ifo_list = eval(opts.manual_ifo_list) + else: + try: + import lal.series + from igwn_ligolw import utils as _ligolw_utils, ligolw as _ligolw + _xmldoc = _ligolw_utils.load_filename( + opts.use_online_psd_file, contenthandler=lal.series.PSDContentHandler) + ifo_list = sorted(lal.series.read_psd_xmldoc(_xmldoc).keys()) + print(" pseudo_pipe: instruments read from {}: {}".format( + opts.use_online_psd_file, ifo_list)) + except Exception as exc: + print(" pseudo_pipe: --use-online-psd-file needs an instrument list, and " + "neither --use-ini nor --manual-ifo-list was given, and the PSD file " + "could not be read ({}) ".format(exc)) + sys.exit(1) # Create command line arguments for those IFOs, so helper can correctly pass then downward for ifo in ifo_list: cmd+= " --psd-file {}={}".format(ifo,opts.use_online_psd_file) @@ -1254,11 +1403,21 @@ def approx_supports_precession(approx_name): cmd += " --internal-ile-auto-logarithm-offset " if opts.internal_ile_rotate_phase: cmd += " --internal-ile-rotate-phase " -if resolve_interpolate_time_request(opts.internal_ile_interpolate_time) is not None: +if (resolve_interpolate_time_request(opts.internal_ile_interpolate_time) is not None + or opts.internal_ile_interpolate_time is not None): # resolve_interpolate_time_request rather than a truthiness test: the flag takes a VALUE, so # '--internal-ile-interpolate-time False' passes the STRING 'False' (truthy in Python) and a # BARE flag passes a sentinel. Both must be distinguished from "a stencil was named", and a - # bare flag must raise rather than silently forward nothing. + # bare flag must raise rather than silently forward nothing. (The call is kept for that + # raise; it is no longer what decides whether to forward.) + # + # AN EXPLICIT OFF-REQUEST IS ALSO FORWARDED, since 2026-09-02. It used to be dropped here -- + # resolve_interpolate_time_request maps 'False'/'off'/'none' to None, exactly like an absent + # flag -- and dropping it was harmless only while the ILE default was 'nearest', because + # forwarding nothing and forwarding 'nearest' were the same answer. The ILE default is now + # time_interp_choice.TIME_INTERP_DEFAULT, so they are opposites and dropping an off-request + # would have turned interpolation ON for a user who explicitly turned it off. The helper + # re-expresses it as '--interpolate-time nearest'; this only has to stop swallowing it. # HELPER passthrough (not a raw ILE arg): the helper owns ILE argument construction, and it # also knows whether the maintained NoLoop path that --interpolate-time requires is in use -- # which needs --time-marginalization AND --vectorized AND one of --gpu/--rotation-slow/ @@ -1266,6 +1425,11 @@ def approx_supports_precession(approx_name): # resolved there -- so forward the request verbatim rather than resolving it here, and let the # helper's log line be the single record of what was chosen. cmd += " --internal-ile-interpolate-time " + str(opts.internal_ile_interpolate_time) + " " +if opts.internal_ile_time_marginalization_quadrature is not None: + # HELPER passthrough, exactly like the stencil above and for the same reason: the helper owns + # ILE argument construction, so the flag must enter args_ile.txt where every other ILE + # argument does. `is not None` rather than a truthiness test -- the option takes a VALUE. + cmd += " --internal-ile-time-marginalization-quadrature " + str(opts.internal_ile_time_marginalization_quadrature) + " " if not(opts.internal_ile_n_chunk is None): cmd += " --internal-ile-n-chunk {} ".format(int(opts.internal_ile_n_chunk)) # If user provides ini file *and* ini file has fake-cache field, generate a local.cache file, and pass it as argument @@ -1281,11 +1445,41 @@ def approx_supports_precession(approx_name): cmd += " --cache local.cache --fake-data " if opts.fake_data_cache: cmd += " --cache {} --fake-data ".format(opts.fake_data_cache) + # event_dict["IFOs"] is populated on the --use-ini path (and by the gracedb + # lookup), but NOT by --event-time + --fake-data-cache, which is the natural + # way to set up a synthetic injection. That combination used to die here + # with KeyError: 'IFOs'. Resolve it from --manual-ifo-list, or from the + # instruments already read off the PSD file, and store it so the later + # consumers of event_dict["IFOs"] see it too. + _ifos = event_dict.get("IFOs") + if not _ifos and opts.manual_ifo_list: + _ifos = eval(opts.manual_ifo_list) + if not _ifos: + try: + _ifos = list(ifo_list) + except NameError: + _ifos = None + if not _ifos: + print(" pseudo_pipe: --fake-data-cache needs an instrument list. Give " + "--manual-ifo-list \"['H1','L1']\", or --use-online-psd-file whose " + "instruments can be read. ") + sys.exit(1) + event_dict["IFOs"] = list(_ifos) if len(event_dict["IFOs"]) >0 : short_list = " {} ".format(event_dict['IFOs']) cmd += " --manual-ifo-list {} ".format(short_list.replace(' ','')) print( cmd) -os.system(cmd) +if os.path.exists('helper_ile_args.txt'): + # This is generated output, not an input. Remove it before invoking the + # helper so a failed helper cannot be mistaken for fresh success in a + # re-used directory. The emitted-byte guard below cannot distinguish a + # stale file from a fresh one when both carry the same requested value. + os.unlink('helper_ile_args.txt') +_helper_rc = os.system(cmd) +if _helper_rc != 0: + print(" FAILURE: helper call exited nonzero; refusing to use any pre-existing " + "helper_ile_args.txt") + sys.exit(1) # we MUST make helper_ile_args.txt if not(os.path.exists('helper_ile_args.txt')): print(" FAILURE: helper call failed to generate required file helper_ile_args.txt") @@ -1366,8 +1560,10 @@ def approx_supports_precession(approx_name): line += " --use-gwsignal --approx " + opts.approx elif not 'NR' in opts.approx: line += " --approx " + opts.approx -elif opts.use_gwsurrogate and 'NRHybSur' in opts.approx: +elif opts.use_gwsurrogate and 'NRHybSur' in opts.approx and 'Tidal' not in opts.approx: line += " --rom-group {} --rom-param NRHybSur3dq8.h5 --approx {} ".format(sur_location_prefix,opts.approx) +elif opts.use_gwsurrogate and 'NRHybSur' in opts.approx and 'Tidal' in opts.approx: + line += " --rom-group {} --rom-param NRHybSur3dq8Tidal --approx {} ".format(sur_location_prefix,opts.approx) elif opts.use_gwsurrogate and "NRSur7dq2" in opts.approx: line += " --rom-group {} --rom-param NRSur7dq2.h5 --approx {} ".format(sur_location_prefix,opts.approx) elif opts.use_gwsurrogate and "NRSur7dq4" in opts.approx: @@ -1497,6 +1693,26 @@ def approx_supports_precession(approx_name): else: line += " --extrinsic-proposal-breadcrumb {}/extr_consolidated_$(macroiterationprev).npz ".format(os.getcwd()) +# TIME-MARGINALIZATION QUADRATURE, part 2 of 2: the LAST chance to refuse, and the only place +# with the whole picture. This checks the BYTES about to be written, not the parsed options -- +# an earlier version keyed the guard on `opts` and read only the PREREQUISITES from `line`, so it +# happily approved an args_ile.txt that had never received the flag at all. Three ways that +# happens, all ending in a silent fall back to Simpson while the pipeline logs the opposite: +# +# * the helper is invoked by NAME through PATH, so version skew can make an older helper +# argparse-error on the new option. The caller now removes stale generated output and checks +# the helper status; this byte check is the independent defence against a successful helper +# that nevertheless drops the emission; +# * --manual-extra-ile-args is appended AFTER the helper's arguments and optparse takes the +# LAST occurrence, so a hand-passed 'simpson' silently overrides the requested value while +# the .sub file still shows both; +# * any future refactor that drops the emission. +# +# It also holds a hand-passed quadrature (manual args, ini) to the same standard, which the +# opts-keyed version skipped entirely. Called unconditionally for that reason. +refuse_unless_time_quadrature_emitted( + opts.internal_ile_time_marginalization_quadrature, line, "args_ile.txt") + with open('args_ile.txt','w') as f: f.write(line) @@ -1644,6 +1860,9 @@ def approx_supports_precession(approx_name): line = line.replace('parameter delta_mc', 'parameter-implied eta --parameter-nofit delta_mc') # quadratic or cov fit needs eta coordinate if opts.force_lambda_no_linear_init: line = line.replace("--prior-lambda-linear", "") # remove this line, usually used in iteration0 + if opts.use_mtot_coords: + line = line.replace('parameter mc', 'parameter-implied mc --parameter-nofit mtot --parameter-nofit q') + line = line.replace('parameter delta_mc', 'parameter-implied delta_mc') if opts.hierarchical_merger_prior_1g: # Must use mtotal, q coordinates! Change defaults line = line.replace('parameter mc', 'parameter mtot') @@ -1718,13 +1937,13 @@ def approx_supports_precession(approx_name): if opts.fit_save_gp: line += " --fit-save-gp my_gp " # fiducial filename, stored in each iteration line += " --eccentricity-prior {}".format(opts.eccentricity_prior) + if opts.use_EOB_parameters: + line += " --use-EOB-parameters --parameter a6c --a6c-min {} --a6c-max {} ".format(opts.force_a6c_min,opts.force_a6c_max) if opts.assume_eccentric: if opts.use_meanPerAno: line += " --parameter meanPerAno --use-meanPerAno " if opts.use_eccentricity_squared: line += " --use-eccentricity --parameter eccentricity_squared " - elif opts.use_eccentricity_ln: - line += " --use-eccentricity --parameter eccentricity_ln " else: line += " --use-eccentricity --parameter eccentricity " # if opts.use_eccentricity_squared: @@ -1761,6 +1980,21 @@ def approx_supports_precession(approx_name): if not(opts.force_meanPerAno_min is None): meanPerAno_min = opts.force_meanPerAno_min line += " --meanPerAno-min {} ".format(meanPerAno_min) + if opts.assume_hyperbolic: + line += " --parameter E0 --parameter p_phi0 --use-hyperbolic " + # always pass the resolved range: CIP's own defaults are narrower than + # the range used to generate the points it is fitting + line += " --E0-max {} --E0-min {} ".format(hyperbolic_range['E0_max'],hyperbolic_range['E0_min']) + line += " --pphi0-max {} --pphi0-min {} ".format(hyperbolic_range['pphi0_max'],hyperbolic_range['pphi0_min']) + + if opts.force_scatter_grids: + line += " --force-scatter " + + if opts.force_plunge_grids: + line += " --force-plunge " + + if opts.force_zoomwhirl_grids: + line += " --force-zoomwhirl " if not(opts.manual_extra_cip_args is None): line += " {} ".format(opts.manual_extra_cip_args) # embed with space on each side, avoid collisions line += "\n" @@ -1857,8 +2091,24 @@ def approx_supports_precession(approx_name): if opts.assume_matter: # puff_params += " --parameter LambdaTilde " # should already be present puff_max_it +=5 # make sure we resolve the correlations +if opts.use_EOB_parameters: + puff_params += " --downselect-parameter a6c --downselect-parameter-range [{},{}] ".format(opts.force_a6c_min,opts.force_a6c_max) if opts.assume_eccentric: - puff_params += " --downselect-parameter eccentricity --downselect-parameter-range [{},{}] ".format(opts.force_ecc_min,opts.force_ecc_max) + puff_params += " --downselect-parameter eccentricity --downselect-parameter-range [{},{}] ".format(opts.force_ecc_min,opts.force_ecc_max) +if opts.assume_hyperbolic: +# puff_params += " --parameter E0 " + # always downselect to the resolved range (a one-sided override used to + # write 'None' into the other end of the range here) + puff_params += " --downselect-parameter E0 --downselect-parameter-range [{},{}] ".format(hyperbolic_range['E0_min'],hyperbolic_range['E0_max']) +# puff_params += " --parameter p_phi0 " + puff_params += " --downselect-parameter p_phi0 --downselect-parameter-range [{},{}] ".format(hyperbolic_range['pphi0_min'],hyperbolic_range['pphi0_max']) + if opts.force_scatter_grids: + puff_params += ' --force-scatter ' + if opts.force_plunge_grids: + puff_params += ' --force-plunge ' + if opts.force_zoomwhirl_grids: + puff_params += ' --force-zoomwhirl ' + if opts.assume_highq: puff_params = puff_params.replace(' delta_mc ', ' eta ') # use natural coordinates in the high q strategy. May want to do this always puff_max_it +=3 @@ -1969,21 +2219,57 @@ def approx_supports_precession(approx_name): cepp = "create_event_parameter_pipeline_BasicIteration" if opts.use_subdags: cepp = "create_event_parameter_pipeline_AlternateIteration" +if opts.approx_extra and not opts.pipeline_builder: + # asking for more than one waveform model IS asking for the cross-model + # builder; make the user say it twice only if they disagree + opts.pipeline_builder = "BasicMultiApproxIteration" +use_multiapprox = (opts.pipeline_builder == "BasicMultiApproxIteration") +if use_multiapprox and not opts.approx_extra: + print(" --pipeline-builder BasicMultiApproxIteration needs at least one --approx-extra ") + sys.exit(1) +# The builder itself accepts pseudo_pipe's option surface and refuses the parts +# it does not implement, so there is no allow-list to maintain here. if opts.pipeline_builder: # explicit override wins, for clean side-by-side A/B testing of the two builders if opts.use_subdags and opts.pipeline_builder != "AlternateIteration": # use_subdags is set either by the user or force-set by --internal-use-amr (which REQUIRES the Alternate builder) print(" WARNING: --pipeline-builder {} overrides --use-subdags routing; AMR/subdag runs require AlternateIteration ".format(opts.pipeline_builder)) cepp = "create_event_parameter_pipeline_" + opts.pipeline_builder print(" Pipeline builder (create_event_parameter_pipeline_*): ", cepp) -cmd =cepp+ " --ile-n-events-to-analyze {} --input-grid proposed-grid.{} --ile-exe `which integrate_likelihood_extrinsic_batchmode` --ile-args `pwd`/args_ile.txt --cip-args-list args_cip_list.txt --test-args args_test.txt --request-memory-CIP {} --request-memory-ILE {} --n-samples-per-job ".format(n_jobs_per_worker,grid_suffix_pp,cip_mem,ile_mem) + str(npts_it) + " --working-directory `pwd` --n-iterations " + str(n_iterations) + " --n-iterations-subdag-max {} ".format(opts.internal_n_iterations_subdag_max) + " --n-copies {} ".format(opts.ile_copies) + " --ile-retries "+ str(opts.ile_retries) + " --general-retries " + str(opts.general_retries) +cmd =cepp+ " --ile-n-events-to-analyze {} --input-grid proposed-grid.{} --ile-exe `which integrate_likelihood_extrinsic_batchmode` --ile-args `pwd`/args_ile.txt --cip-args-list args_cip_list.txt --test-args args_test.txt --request-memory-CIP {} --request-memory-ILE {} --n-samples-per-job ".format(n_jobs_per_worker,grid_suffix_pp,cip_mem,ile_mem) + str(npts_it) + " --working-directory `pwd` --n-iterations " + str(n_iterations) + ("" if use_multiapprox else " --n-iterations-subdag-max {} ".format(opts.internal_n_iterations_subdag_max)) + " --n-copies {} ".format(opts.ile_copies) + " --ile-retries "+ str(opts.ile_retries) + " --general-retries " + str(opts.general_retries) +if use_multiapprox: + # Every model on the SAME grid. --approx is the primary; --approx-extra the + # rest. The builder marginalizes over them point by point in the loop and + # forks per model at the terminal stage. + for _ap in [opts.approx] + list(opts.approx_extra): + cmd += " --approx {} ".format(_ap) + # Forward the generator ROUTE per model. Without this the builder never sees + # --approx-gwsignal, so it does not strip the global --use-gwsignal and every + # model is sent through gwsignal -- where the phenom family cannot be + # generated at all, its ILE jobs contribute zero rows, and the run silently + # degrades to a single model. + _gw = [] + if opts.use_gwsignal: + _gw.append(opts.approx) # the primary is what --use-gwsignal meant + _gw += list(opts.approx_extra_gwsignal or []) + for _m in _gw: + if _m not in [opts.approx] + list(opts.approx_extra): + print(" --approx-extra-gwsignal names {}, which is not among the models {}".format( + _m, [opts.approx] + list(opts.approx_extra))); sys.exit(1) + cmd += " --approx-gwsignal {} ".format(_m) + for _pr in (opts.approx_prior or []): + cmd += " --approx-prior '{}' ".format(_pr) + if opts.require_all_approx: + cmd += " --require-all-approx " if opts.ile_jobs_per_worker_first: cmd += " --ile-n-events-to-analyze-first {} ".format(opts.ile_jobs_per_worker_first) -if opts.assume_matter or opts.assume_eccentric: +if opts.assume_matter or opts.assume_eccentric or opts.assume_hyperbolic: cmd += " --convert-args `pwd`/helper_convert_args.txt " if not(opts.ile_runtime_max_minutes is None): cmd += " --ile-runtime-max-minutes {} ".format(opts.ile_runtime_max_minutes) if not(opts.internal_use_amr) or opts.internal_use_amr_puff: cmd+= " --puff-exe `which util_ParameterPuffball.py` --puff-cadence 1 --puff-max-it " + str(puff_max_it)+ " --puff-args `pwd`/args_puff.txt " +if opts.use_EOB_parameters: + cmd += " --use-EOB-parameters " if opts.calmarg_pilot: cmd += " --calmarg-pilot --calmarg-pilot-cadence {} --calmarg-pilot-max-it {} --calmarg-pilot-top-fraction {} --calmarg-pilot-max-points {} ".format( opts.calmarg_pilot_cadence, opts.calmarg_pilot_max_it, opts.calmarg_pilot_top_fraction, opts.calmarg_pilot_max_points) @@ -2017,10 +2303,10 @@ def approx_supports_precession(approx_name): cmd += " --use-eccentricity " if opts.sample_eccentricity_squared: cmd += " --use-eccentricity-squared-sampling " - if opts.sample_eccentricity_ln: - cmd += " --use-eccentricity-ln-sampling " if opts.use_meanPerAno: cmd += " --use-meanPerAno " +if opts.assume_hyperbolic: + cmd += " --use-hyperbolic " if opts.calibration_reweighting and (not opts.bilby_pickle_file): cmd += " --calibration-reweighting --calibration-reweighting-exe `which calibration_reweighting.py` --bilby-ini-file {} --bilby-pickle-exe `which bilby_pipe_generation` ".format(str(opts.bilby_ini_file)) if opts.calibration_reweighting_count: @@ -2051,6 +2337,8 @@ def approx_supports_precession(approx_name): cmd += " --comov-distance-reweighting --comov-distance-reweighting-exe `which make_uni_comov_skymap.py` --convert-ascii2h5-exe `which convert_output_format_ascii2h5.py` " if opts.use_gauss_early: cmd += " --cip-exe-G `which util_ConstructIntrinsicPosterior_GaussianResampling.py ` " +if opts.first_iteration_jumpstart: + cmd += " --first-iteration-jumpstart " if opts.internal_use_amr: print(" AMR prototype: Using hardcoded aligned-spin settings, assembling grid, requires coinc!") if _use_hpip_pp and opts.manual_initial_grid is None: @@ -2136,7 +2424,7 @@ def approx_supports_precession(approx_name): if opts.internal_last_iteration_extrinsic_samples_per_ile: cmd += " --last-iteration-extrinsic-samples-per-ile {}".format(opts.internal_last_iteration_extrinsic_samples_per_ile) if opts.internal_last_iteration_extrinsic_samples_per_ile_internal: - cmd += " --last-iteration-extrinsic-samples-per-ile-internal {}".format(opts.internal_last_iteration_extrinsic_samples_per_ile_internal) + cmd += " --last-iteration-extrinsic-samples-per-ile-internal {}".format(opts.internal_last_iteration_extrinsic_samples_per_ile_internal) if opts.add_extrinsic_time_resampling: cmd+= " --last-iteration-extrinsic-time-resampling " if opts.batch_extrinsic: @@ -2193,7 +2481,7 @@ def approx_supports_precession(approx_name): if opts.condor_nogrid_nonworker: cmd += " --condor-nogrid-nonworker " if opts.use_osg_simple_requirements: - cmd += " --use-osg-simple-reqirements " + cmd += " --use-osg-simple-requirements " if opts.archive_pesummary_label: # cmd += " --plot-exe `which summarypages` --plot-args args_plot.txt " cmd += " --plot-exe summarypages --plot-args args_plot.txt " @@ -2300,12 +2588,18 @@ def approx_supports_precession(approx_name): cmd += " --last-iteration-export-distance-slices-skip-threshold {} ".format(opts.export_distance_slices_skip_threshold) print(cmd) -os.system(cmd) +_rc = os.system(cmd) +if _rc != 0: + # A failed builder used to leave pseudo_pipe reporting success with no DAG in + # the run directory -- the args_*.txt are all there, so it looks finished. + print(" pseudo_pipe: the pipeline builder FAILED (exit {}); no DAG was written. " + "See the output above.".format(_rc >> 8 if _rc > 255 else _rc)) + sys.exit(1) if opts.internal_ile_check_good_enough: # Populate 'ile_check_good_enough' through all subdirectories cmd_enough = r"find . -name 'iter*ile' -type d -exec touch {}/ile_good_enough \; " - os.system(cmd) + os.system(cmd_enough) # was os.system(cmd): re-ran the pipeline builder if opts.use_osg_file_transfer and opts.internal_truncate_files_for_osg_file_transfer: if opts.fake_data_cache: diff --git a/MonteCarloMarginalizeCode/Code/bin/util_SimInspiralToCoinc.py b/MonteCarloMarginalizeCode/Code/bin/util_SimInspiralToCoinc.py index 67a31198d..5b2ea5c45 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_SimInspiralToCoinc.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_SimInspiralToCoinc.py @@ -127,8 +127,14 @@ def _empty_row(obj): sngl.snr = opts.injected_snr # made up, needed for some algorithms to work else: sngl.snr = 20. # made up, needed for some algorithms to work + # Eccentric Parameters sngl.alpha4 = P.eccentricity sngl.alpha = P.meanPerAno + # EOB Parameters + sngl.psi0 = P.a6c + # Hyperbolic Parameters + sngl.psi3 = P.E0 + sngl.beta = P.p_phi0 # add to table sngl_table.append(sngl) diff --git a/MonteCarloMarginalizeCode/Code/bin/util_WriteInjectionFile.py b/MonteCarloMarginalizeCode/Code/bin/util_WriteInjectionFile.py index fd4956844..7a4ea4d0a 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_WriteInjectionFile.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_WriteInjectionFile.py @@ -26,6 +26,7 @@ # Add option to use NR waveforms! parser.add_argument("--group",default=None) parser.add_argument("--param",default=None) +parser.add_argument("--fref",default=None,type=float) opts= parser.parse_args() hasNR = False @@ -52,7 +53,8 @@ P.phiref = 0 P.psi = 0 P.approx = lalsimutils.lalsim.GetApproximantFromString(opts.approximant) # allow user to override the approx setting. Important for NR followup, where no approx set in sim_xml! - +if opts.fref: + P.fref = opts.fref param_names = opts.parameter for param in param_names: # Check if in the valid list diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/lisa_drift_ledger.json b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/lisa_drift_ledger.json index b17f53984..0470652a0 100644 --- a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/lisa_drift_ledger.json +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/lisa_drift_ledger.json @@ -33,6 +33,10 @@ "decision": "NA", "reason": "Calibration-envelope internals; see the --calibration-* reason." }, + "FUNC:_equal_weight_fairdraw_for_serialization": { + "decision": "PORT", + "reason": "Completes a sampler-side fair draw that intentionally did not fire because it would not shrink a tiny retained record, but only on the copy written to XML. LISA has the same skip-on-no-shrink and serialization boundary, so port this with --fairdraw-extrinsic-output-n-max while preserving its larger LISA default." + }, "FUNC:_normalize_interpolate_time_argv": { "decision": "PORT", "reason": "Normalizes --interpolate-time argv forms. LISA exposes --interpolate-time, so the same normalization applies." @@ -61,6 +65,10 @@ "decision": "PORT", "reason": "Cosmology helpers behind --d-prior-redshift. Planck15 via the framework helper; the interpolation grid still needs a z ceiling that covers MBHB (z~20), which is a gridding choice rather than a physics decision." }, + "FUNC:fused_calmarg_in_use": { + "decision": "NA", + "reason": "THE predicate for 'will the fused in-loop calibration kernel actually run?', shared by the startup --interpolate-time guard and by use_fused_calmarg at dispatch so the condition has one definition instead of two that drift (it drifted three times in a day before being hoisted). Every one of its terms is a LIGO/Virgo calibration concept -- an envelope directory, a realization count, --calibration-fused-kernel, the responsibilities pilot -- and the LISA driver has NONE of them: it declares no --calibration-* option, never sets cal_method, and models no instrument calibration at all. See the --calibration-* reason. NOT a stencil gap despite being read by the stencil guard: LISA's --interpolate-time is a separate boolean parsed by legacy_time_interpolation_enabled, with no fused kernel to protect. If LISA ever models calibration this predicate is the wrong shape for it, and the thing to port would be the ONE-definition discipline, not this function." + }, "OPTION:--calibration-burn-in-neff": { "decision": "NA", "reason": "LIGO/Virgo spline calibration-envelope marginalization. The LISA driver models no instrument calibration: it takes no envelope directory, has no cal nodes, and its response is applied analytically by factored_likelihood_LISA. LISA calibration, if it is ever modelled, will not have this data product or this spline parameterization, so porting the LIGO machinery would be actively misleading." @@ -197,6 +205,10 @@ "decision": "PORT", "reason": "Caps rows per fair-draw export. LISA currently hardcodes this to opts.n_eff at the igrand_fairdraw_samples_max call site. WARNING for the port: main's default is 5, so adopting main's default verbatim would silently shrink every LISA export by orders of magnitude. Port the flag with LISA's present behaviour as its default." }, + "OPTION:--force-hyperbolic-22": { + "decision": "NA", + "reason": "Controls or exports the ground-based external-TEOBResumS advanced-physics path. The LISA driver does not call that waveform path or write its a6c/E0/p_phi0 composite layout." + }, "OPTION:--freqresponse": { "decision": "NA", "reason": "Finite light-travel-time transfer across the arms for 3G ground detectors (CE/ET), built on lalsimulation detector geometry and an arm-length override in metres. LISA's finite-size response is not an add-on: it is the whole point of the TDI response the LISA driver already applies." @@ -265,6 +277,10 @@ "decision": "PORT", "reason": "ANSWERED (RO 2026-08-16): LISA and LIGO are never overlapping use cases, so follow the convention already in this driver, document it in the help string, and DO NOT rename the options. VERIFIED that convention is ECLIPTIC: the sampled right_ascension/declination columns flow to P.phi/P.theta and then to lisa_sky_lamda/lisa_sky_beta, i.e. ecliptic longitude/latitude, under the historical key names. So --limit-right-ascension bounds lambda and --limit-declination bounds beta; say exactly that in the help text. Port the post-PR#58 form including the cos(iota)/cos(dec) endpoint swap under the cosine samplers." }, + "OPTION:--limit-distance": { + "decision": "PORT", + "reason": "Sampling-only distance box: narrows what distance is DRAWN from while the prior keeps its full [--d-min,--d-max] normalization, so lnZ stays on the full-range scale. Port it, and port the SPLIT rather than the option alone. VERIFIED 2026-09-02 that the LISA driver still carries the one-range form the main driver was just moved off (integrate_likelihood_extrinsic_batchmode_lisa:1021-1024: dist_sampler and dist_prior_pdf are both built from param_limits['distance']), which normalizes the Euclidean density over whatever the sampler happens to draw from -- so narrowing for cost there would silently rescale the evidence. mcsampler.distance_sampler_kwargs() already takes the sampling range and the prior range as two arguments and is shared code -- but the port is NOT a pure call-site change, and what is missing is a prior LISA SUPPORTS and the helper does not. VERIFIED 2026-09-03: the LISA driver has an explicit 'elif opts.d_prior == uniform: dist_prior_pdf = dist_sampler' branch (:1029), so --d-prior uniform WORKS there; the main driver has no such branch and raises Exception('distance prior') on it. distance_sampler_kwargs() handles Euclidean and pseudo_cosmo only and raises ValueError otherwise, so routing LISA through it as-is would BREAK an existing LISA prior. EXTEND THE HELPER FIRST, with a uniform branch normalized 1/(p_hi - p_lo) over the PRIOR range -- and note the trap while doing it: LISA's uniform prior is literally the sampling density (ret_uniform_samp_vector_alt(lo,hi) = 1/(hi-lo) over the SAMPLED range), which is exactly the sampling/prior conflation this helper exists to split, so transcribing it would reintroduce the defect the port is for. That extension also changes MAIN-driver behaviour (--d-prior uniform stops raising and starts running), so it needs RO'S and is not something to fold into the --limit-distance PR. The motivation is STRONGER on LISA than on ground-based data: measured on real LIGO data at rho ~ 82, a box tracking the posterior removes 0.37 +- 0.11 nats of sampling bias the full-range run was carrying (4.16 nats with --no-adapt-distance), and MBHB SNRs are one to two orders of magnitude higher, where the posterior is narrower still relative to the same prior (RIFT_roboto_paper analyses/limit_distance_e2e/). CARRY THE REFUSALS, and note only one of the three transfers today: LISA HAS --distance-marginalization (:245), so refuse there for the same reason -- no distance sampler exists to narrow. It has neither --d-prior-redshift nor --internal-reparam-dl-incl, so those two refusals have nothing to attach to yet; --internal-reparam-dl-incl is itself a PORT item above, so whichever of the two lands second owes the refusal. LISA's --d-prior set is also different (Euclidean|uniform|pseudo_cosmo, no cosmo/cosmo_sourceframe), so the cosmo branch of the main driver's narrowing block has no counterpart to port." + }, "OPTION:--limit-inclination": { "decision": "PORT", "reason": "Zoom-box limits on psi and inclination. These parameters mean the same thing in both drivers and LISA exposes --inclination-cosine-sampler, which is exactly the case junior PR #58 found silently ignored -- so port the POST-#58 form, including the cos(iota) endpoint swap." @@ -329,6 +345,14 @@ "decision": "PORT", "reason": "RESOLVED (RO 2026-08-16). The convention does not matter: the seed is points in the sampler's OWN coordinate space, read positionally against params_ordered, so any self-consistent choice works and the ecliptic sky answer already determines it. The hazard is only that a mismatch is UNDETECTABLE -- ecliptic lambda and RA share [0,2pi), beta and dec share [-pi/2,pi/2], so no range check separates them and a wrong-frame seed silently contracts the live volume around the wrong region. SCOPE (RO): these files are used INTERNALLY within a homogeneous run -- we are talking to ourselves, not to heterogeneous tooling -- so keep it simple: a one-line frame stamp in the file header written by the producer, warn if it is absent or disagrees. Do NOT build a validation framework for it." }, + "OPTION:--save-EOB-parameters": { + "decision": "NA", + "reason": "Controls or exports the ground-based external-TEOBResumS advanced-physics path. The LISA driver does not call that waveform path or write its a6c/E0/p_phi0 composite layout." + }, + "OPTION:--save-hyperbolic": { + "decision": "NA", + "reason": "Controls or exports the ground-based external-TEOBResumS advanced-physics path. The LISA driver does not call that waveform path or write its a6c/E0/p_phi0 composite layout." + }, "OPTION:--save-meanPerAno": { "decision": "NA", "reason": "Exports the eccentric mean anomaly. Tied to the ground-based eccentric waveform path (see --e-freq); the LISA driver's own eccentricity export is --save-eccentricity." @@ -341,9 +365,9 @@ "decision": "NA", "reason": "Separate internal sampling rate for the ground-based precompute. LISA's precompute takes its rate from the h5 frame and P.deltaT; there is no second internal rate to set." }, - "OPTION:--srate-resample-time-marginalization": { + "OPTION:--time-marginalization-quadrature": { "decision": "PORT", - "reason": "Interpolate the lnL time series onto a finer grid before time resampling. LISA already has --resample-time-marginalization and its own time-resampling block, so this is the matching resolution knob and applies directly." + "reason": "Selects the rule for the TIME integral of the marginalized likelihood (simpson, the unchanged default, or the opt-in band-limited refinement). LISA carries the SAME defect this addresses: factored_likelihood_LISA.py integrates exp(lnL(t)) with Simpson at the fixed data spacing, while the integrand's width sigma_t = 1/(2 pi rho sigma_f) is set by the signal and shrinks as 1/rho -- so it under-resolves its own integrand, worse at higher SNR. PORT, not NA. But porting is NOT just wiring the flag through, and the prerequisite is the whole question: the band-limited argument needs kappa band-limited below Nyquist AND rho_sq time-INDEPENDENT. The main driver refuses --rotation-slow and --freqresponse for exactly that second condition, and a response that varies across the observation is the normal case for LISA, not an exotic one. So the LISA port must first establish whether its self-term is time-independent over the integration window; if it is not, the honest outcome is a documented refusal on that path rather than a flag that silently integrates the wrong thing. Note also that the LISA site integrates on axis=0, not the last axis." } } } diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/make_lisa_drift_ledger.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/make_lisa_drift_ledger.py index 144fb7be8..dd8444c0a 100644 --- a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/make_lisa_drift_ledger.py +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/make_lisa_drift_ledger.py @@ -223,6 +223,11 @@ "is 5, so adopting main's default verbatim would silently shrink every LISA " "export by orders of magnitude. Port the flag with LISA's present behaviour as " "its default."), + (r"^FUNC:_equal_weight_fairdraw_for_serialization$", "PORT", + "Completes a sampler-side fair draw that intentionally did not fire because it " + "would not shrink a tiny retained record, but only on the copy written to XML. " + "LISA has the same skip-on-no-shrink and serialization boundary, so port this " + "with --fairdraw-extrinsic-output-n-max while preserving its larger LISA default."), # ------------------------------------------------------- LIGO/Virgo calibration envelopes (r"^OPTION:--calibration-", "NA", @@ -243,6 +248,19 @@ "one. If LISA ever models calibration, it wants derived_rng directly, not this wrapper."), (r"^FUNC:analyze_event\._cal_error_probe(\._draw_dist)?$", "NA", "Calibration Monte-Carlo error probe; see the --calibration-* reason."), + (r"^FUNC:fused_calmarg_in_use$", "NA", + "THE predicate for 'will the fused in-loop calibration kernel actually run?', shared by " + "the startup --interpolate-time guard and by use_fused_calmarg at dispatch so the " + "condition has one definition instead of two that drift (it drifted three times in a day " + "before being hoisted). Every one of its terms is a LIGO/Virgo calibration concept -- an " + "envelope directory, a realization count, --calibration-fused-kernel, the " + "responsibilities pilot -- and the LISA driver has NONE of them: it declares no " + "--calibration-* option, never sets cal_method, and models no instrument calibration at " + "all. See the --calibration-* reason. NOT a stencil gap despite being read by the " + "stencil guard: LISA's --interpolate-time is a separate boolean parsed by " + "legacy_time_interpolation_enabled, with no fused kernel to protect. If LISA ever models " + "calibration this predicate is the wrong shape for it, and the thing to port would be " + "the ONE-definition discipline, not this function."), # ------------------------------------------------------- ground-based detector geometry (r"^OPTION:--rotation-(slow|n-harmonics|p-max)$", "NA", @@ -313,6 +331,43 @@ "beta; say exactly that in the help text. Port the post-PR#58 form including the " "cos(iota)/cos(dec) endpoint swap under the cosine samplers."), + (r"^OPTION:--limit-distance$", "PORT", + "Sampling-only distance box: narrows what distance is DRAWN from while the prior keeps " + "its full [--d-min,--d-max] normalization, so lnZ stays on the full-range scale. Port it, " + "and port the SPLIT rather than the option alone. VERIFIED 2026-09-02 that the LISA " + "driver still carries the one-range form the main driver was just moved off " + "(integrate_likelihood_extrinsic_batchmode_lisa:1021-1024: dist_sampler and " + "dist_prior_pdf are both built from param_limits['distance']), which normalizes the " + "Euclidean density over whatever the sampler happens to draw from -- so narrowing for " + "cost there would silently rescale the evidence. mcsampler.distance_sampler_kwargs() " + "already takes the sampling range and the prior range as two arguments and is shared " + "code -- but the port is NOT a pure call-site change, and what is missing is a prior " + "LISA SUPPORTS and the helper does not. VERIFIED 2026-09-03: the LISA driver has an " + "explicit 'elif opts.d_prior == uniform: dist_prior_pdf = dist_sampler' branch (:1029), " + "so --d-prior uniform WORKS there; the main driver has no such branch and raises " + "Exception('distance prior') on it. distance_sampler_kwargs() handles Euclidean and " + "pseudo_cosmo only and raises ValueError otherwise, so routing LISA through it as-is " + "would BREAK an existing LISA prior. EXTEND THE HELPER FIRST, with a uniform branch " + "normalized 1/(p_hi - p_lo) over the PRIOR range -- and note the trap while doing it: " + "LISA's uniform prior is literally the sampling density " + "(ret_uniform_samp_vector_alt(lo,hi) = 1/(hi-lo) over the SAMPLED range), which is " + "exactly the sampling/prior conflation this helper exists to split, so transcribing it " + "would reintroduce the defect the port is for. That extension also changes MAIN-driver " + "behaviour (--d-prior uniform stops raising and starts running), so it needs RO'S and " + "is not something to fold into the --limit-distance PR. The motivation is " + "STRONGER on LISA than on ground-based data: measured on real LIGO data at rho ~ 82, a " + "box tracking the posterior removes 0.37 +- 0.11 nats of sampling bias the full-range " + "run was carrying (4.16 nats with --no-adapt-distance), and MBHB SNRs are one to two " + "orders of magnitude higher, where the posterior is narrower still relative to the same " + "prior (RIFT_roboto_paper analyses/limit_distance_e2e/). CARRY THE REFUSALS, and note " + "only one of the three transfers today: LISA HAS --distance-marginalization (:245), so " + "refuse there for the same reason -- no distance sampler exists to narrow. It has " + "neither --d-prior-redshift nor --internal-reparam-dl-incl, so those two refusals have " + "nothing to attach to yet; --internal-reparam-dl-incl is itself a PORT item above, so " + "whichever of the two lands second owes the refusal. LISA's --d-prior set is also " + "different (Euclidean|uniform|pseudo_cosmo, no cosmo/cosmo_sourceframe), so the cosmo " + "branch of the main driver's narrowing block has no counterpart to port."), + # --------------------------------------------------------------------- data / waveform io (r"^OPTION:--internal-data-storage-window-half$", "NA", "Half-width of the main driver's internal precompute storage window. The LISA " @@ -344,6 +399,22 @@ (r"^FUNC:_normalize_interpolate_time_argv$", "PORT", "Normalizes --interpolate-time argv forms. LISA exposes --interpolate-time, so " "the same normalization applies."), + (r"^OPTION:--time-marginalization-quadrature$", "PORT", + "Selects the rule for the TIME integral of the marginalized likelihood " + "(simpson, the unchanged default, or the opt-in band-limited refinement). LISA " + "carries the SAME defect this addresses: factored_likelihood_LISA.py integrates " + "exp(lnL(t)) with Simpson at the fixed data spacing, while the integrand's width " + "sigma_t = 1/(2 pi rho sigma_f) is set by the signal and shrinks as 1/rho -- so it " + "under-resolves its own integrand, worse at higher SNR. PORT, not NA. But porting " + "is NOT just wiring the flag through, and the prerequisite is the whole question: " + "the band-limited argument needs kappa band-limited below Nyquist AND rho_sq " + "time-INDEPENDENT. The main driver refuses --rotation-slow and --freqresponse for " + "exactly that second condition, and a response that varies across the observation " + "is the normal case for LISA, not an exotic one. So the LISA port must first " + "establish whether its self-term is time-independent over the integration window; " + "if it is not, the honest outcome is a documented refusal on that path rather than " + "a flag that silently integrates the wrong thing. Note also that the LISA site " + "integrates on axis=0, not the last axis."), (r"^OPTION:--internal-precompute-ignore-threshold$", "PORT", "Drops negligible modes during precompute. LISA is mode-heavy (--modes, " "--restricted-mode-list-file) and pays more per mode than a ground-based run, so " @@ -362,6 +433,10 @@ "Exports the eccentric mean anomaly. Tied to the ground-based eccentric waveform " "path (see --e-freq); the LISA driver's own eccentricity export is " "--save-eccentricity."), + (r"^OPTION:--(force-hyperbolic-22|save-EOB-parameters|save-hyperbolic)$", "NA", + "Controls or exports the ground-based external-TEOBResumS advanced-physics path. " + "The LISA driver does not call that waveform path or write its a6c/E0/p_phi0 " + "composite layout."), (r"^OPTION:--calibration-spline-count$", "NA", "See the --calibration-* reason."), (r"^CONST:_SEQ_WS_PENDING$", "PORT", "Sentinel for the deferred sequential warm-start capture; ports with " diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py index 2839ffa5c..0972c5316 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py @@ -70,7 +70,7 @@ def main(): interp=ti))[0]) for ti in ("nearest", "cubic")} print(" t_win=%.2f tvals=+/-%.2f : cupy(near/cubic)=%.1f/%.1f " - "JAX(near/lin)=%.1f/%.1f half-cupy_near=%.1f (%.2f%%)" + "JAX(near/cubic)=%.1f/%.1f half-cupy_near=%.1f (%.2f%%)" % (t_window, iwh, cu["nearest"], cu["cubic"], jx["nearest"], jx["cubic"], half_dd - cu["nearest"], 100 * (half_dd - cu["nearest"]) / half_dd)) print("DEBUG2 DONE") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py index 7985c0073..586871be2 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py @@ -8,7 +8,23 @@ is misaligned by that fraction of a sample (a fixable convention), and that offset is the self-consistent injection shift. If the minimum sits at dt_shift=0 with the deficit intact, the ~1% is a genuine response-model gap, not a time reference. + +ANSWERED, and by neither of those two: the ~1% was almost entirely under-sampled time +interpolation. The rholm used to sit at exactly the Nyquist rate of its own analysis +band, leaving a 4-point stencil no headroom. Oversampling the rholm collapses the +deficit by ~99%, and what survives is a precompute floor (finite t_window, Qmax +truncation, PSD band) that is flat in oversampling -- not a time-reference convention +and not a response-model gap of the size this script was chasing. The ladder and the +method are in analyses/slowrot_finite-size/DESIGN_sampling.md in the paper repo. Note +they were measured on the library's own numpy likelihood, whereas this script drives the +JAX one -- the two agree on the sky offset to within the difference between them, but the +deficit ladder itself has not been reproduced on this path. + +This script is therefore kept as a probe OF that regime rather than a question about it: +it pins oversample=1 (see main) so the scan still exhibits the effect. Repointing it at +the surviving floor would be a different measurement and needs a different scan range. """ +import inspect import os, sys import numpy as np import jax @@ -50,8 +66,27 @@ def lal_copy(d): def main(): + # DEFAULT to oversample=1, rather than inheriting the library's. This scan is + # defined RELATIVE to deltaT -- dts = fr * deltaT over +-1 sample -- so it does not + # merely get finer when the library samples the rholm more finely, it covers + # proportionally less physical time and stops probing the near-Nyquist regime the + # scan exists to characterise. At the library's current default (4) the deficit it + # reads is flat and near zero, which looks like "nothing to see" rather than "you are + # no longer looking". Overridable, so the surviving floor can be scanned on purpose. + # Old libraries have no such knob and are already at 1: honour that silently, but + # refuse an explicit request they cannot satisfy rather than quietly giving 1. + _ovs = int(os.environ.get("SLOWROT_OVERSAMPLE", "1")) + if "oversample" in inspect.signature(fslib.Source.__init__).parameters: + _kw = {"oversample": _ovs} + elif _ovs != 1: + raise SystemExit( + "SLOWROT_OVERSAMPLE=%d needs a slowrot_fs_lib with the oversample knob; this " + "one sets deltaT = 1/(2*fmax) and is already at 1" % _ovs) + else: + _kw = {} src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=INCL, - phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD") + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD", + **_kw) net = fslib.network(NET) dist = fslib.distance_for_snr(src, net, SNR) dd, pd, arm, meta = fslib.build_finite_size_data(src, net, dist) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py b/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py index 420fe2435..d763cc850 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py @@ -18,6 +18,7 @@ PYTHONPATH=:/analyses/slowrot_finite-size \ python test/jax/make_3g_figdata.py """ +import inspect import os import sys import numpy as np @@ -95,9 +96,15 @@ def _draw_distance(data, ra, dec, psi, incl, phiref, d_min, d_max, seed=0): def run_one(src, net, target_snr, want_samples=False): dist = fslib.distance_for_snr(src, net, target_snr) # SELFCONSISTENT (int Qmax): render the injection with the recovery's own b_p*W_p - # response so truth is the exact global maximum -- combined with a finely-sampled - # rholm (fmax>=2048, deltaT<=1/4096) this removes the ~0.16 deg cubic-interpolation - # timing systematic that otherwise displaces the razor-sharp high-SNR sky posterior. + # response so truth is the exact global maximum -- combined with an adequately + # OVERSAMPLED rholm this removes the cubic-interpolation timing systematic that + # otherwise displaces the razor-sharp high-SNR sky posterior. The oversampling is + # slowrot_fs_lib's own knob (deltaT = 1/(2*oversample*fmax), defaulted by that + # library's OVERSAMPLE); raising fmax does not do it, and at equal sample rate the + # narrower band is the better one. Measured ladders (offset vs oversample, and the + # stencil sweep): analyses/slowrot_finite-size/DESIGN_sampling.md in the paper repo. + # If your slowrot_fs_lib has no `oversample` and still sets deltaT = 1/(2*fmax), it + # predates that decoupling and none of this paragraph applies -- see main(). sc = os.environ.get("SLOWROT_SELFCONSISTENT") data_dict, psd_dict, arm_dict, meta = fslib.build_finite_size_data( src, net, dist, selfconsistent_Qmax=(int(sc) if sc else None)) @@ -170,13 +177,36 @@ def main(): # of a near-face-on dominant-quadrupole source is broken and the orientation # sector recovers on truth. Override with SLOWROT_INCL. incl = float(os.environ.get("SLOWROT_INCL", "0.4")) - # fmax sets both the waveform bandlimit AND the rholm sampling deltaT=1/(2 fmax); - # 2048 (deltaT=1/4096) finely samples the rholm so the recovery's cubic time - # interpolation reproduces the per-detector fractional-sample delays -> no sky bias. + # WHICH slowrot_fs_lib DO YOU HAVE? Read build_finite_size_data; do not go by a + # commit id. If it sets deltaT = 1/(2*fmax) and Source has no `oversample`, it is the + # OLD library: fmax there sets the band AND the sample rate, the rholm sits at its own + # band Nyquist, and the only lever you have is fmax. If it sets + # deltaT = 1/(2*oversample*fmax), it is the new one and the rest of this applies. + # + # In the new library fmax is the ANALYSIS BAND LIMIT only; oversample sets the rholm + # sampling (default 4 at the time of writing, so srate = 8192 at fmax=1024), and + # raising fmax no longer refines the time series. 1024 stays as the band choice for + # this BNS: DESIGN_sampling.md measures the narrower band as marginally better at + # equal sample rate, so quote a configuration by fmax AND srate -- neither alone + # identifies it. SLOWROT_OVERSAMPLE overrides the sampling; 1 restores the old + # library's setting and rebuilds its data bit-for-bit (the sampler on top is not + # deterministic). It is passed only when set, so this script still runs unchanged + # against the old library -- and an explicit request the old library cannot satisfy + # is refused rather than silently downgraded to its fixed 1. + # + # Provenance: the decoupling is RIFT_roboto_paper branch claude/stoic-saha-496052 + # (2026-08-26), which at the time of writing had NOT landed on that repo's main. fmax = float(os.environ.get("SLOWROT_FMAX", "1024.0")) + _ovs = os.environ.get("SLOWROT_OVERSAMPLE") + if _ovs and "oversample" not in inspect.signature(fslib.Source.__init__).parameters: + raise SystemExit( + "SLOWROT_OVERSAMPLE=%s needs a slowrot_fs_lib with the oversample knob; this " + "one sets deltaT = 1/(2*fmax) and is fixed at the band's own Nyquist rate" + % _ovs) + src_kw = {"oversample": int(_ovs)} if _ovs else {} src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=incl, phiref=0.0, fmin=50.0, fmax=fmax, seglen=32.0, - approx="IMRPhenomD") + approx="IMRPhenomD", **src_kw) net = fslib.network(NETWORK) print("3G FIGDATA network=%s rep_snr=%.0f snrs=%s" % (NETWORK, SNR_REP, SNRS)) rows = [] diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_block_dispatch.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_block_dispatch.py new file mode 100644 index 000000000..d68874ff4 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_block_dispatch.py @@ -0,0 +1,220 @@ +""" +Gate for the EXECUTION-COST structure of the anglemarg laplace path: +the block-dispatched kernel (_laplace_psi_lnI_block). + +WHY THIS EXISTS (2026-08-28): after the compile-cost fix (PR #209) made the +laplace path runnable, it EXECUTED ~2,950x slower than the grid scheme +(1.28e-2 vs 4.33e-6 s per sample*timepoint; a single 800-sample pilot chunk += 6,094 s, an SNR-40 run = days). Measured attribution (GPU, additive to +<1%): the C^1 blend evaluates BOTH kernel branches at every lattice point -- +the 320-point u-quadrature (55% of execution) and the bracket/bisect +root-finding (39%, of which the 20-step bisection is 33%) -- while a census +of the production-shaped lattice shows 99.5% of points sit in the +pure-quadrature regime (t < BLEND_LO; 89% at t < 20, where the module's own +aliasing rule needs only N ~ 32-96) and the points carrying posterior weight +sit at t > 900 where only the Laplace branch is needed. The fix dispatches +each (dist_block x phi_chunk x S x npts) kernel call through lax.switch on +scalar bounds of t = b + 2d, so exactly one branch executes per block, with +the quadrature N laddered by the SAME aliasing exponent the shipped +(N=320, t=300) pair fixes. + +These tests pin (1) the ladder's accuracy contract, (2) the dispatcher's +value agreement with the undispatched kernel across every branch, and +(3) the WIRING -- the fused driver must actually call the dispatcher +(helper-level tests cannot see a call site that stops calling the helper). + +Each test fails under a deliberate mutation (verified by hand; mutations +and observed failures recorded in the PR): + * raising a ladder threshold above its contour (t_ok 0.9 -> 50 for N=32) + -> ladder-contract test fails, and the branch-agreement test fails + LOUDLY (~8% relative aliasing at t=50 with N=32); + * off-by-one in the switch index (pure-Laplace taken from BLEND_LO + instead of BLEND_HI) -> branch-agreement fails in the straddle family; + * rewiring _dist_step back to the plain kernel -> wiring test fails. +""" + +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile.core import make_distance_grid + +from test_angle_marg_compile_cost import make_synth + + +def _aliasing_exponent(N, t): + """ln[I_{N/2}(t) / I_0(t)] by the uniform asymptotic (nu = N/2): + sqrt(nu^2 + t^2) - nu*asinh(nu/t) - t. This is the trapezoid-rule + aliasing bound the module's constants block quotes (~e^-40 at the + shipped N=320, t=300 pair).""" + nu = N / 2.0 + return np.sqrt(nu * nu + t * t) - nu * np.arcsinh(nu / t) - t + + +def test_quad_ladder_keeps_the_shipped_aliasing_exponent(): + """Every ladder rung must be at least as accurate at its threshold as + the shipped (N=320, t=300) band edge, and rungs must be multiples of + the 16-point quadrature scan chunk. Fails if any t_ok is raised above + the E = E(320, 300) contour or a rung breaks the chunking.""" + e_ship = _aliasing_exponent(320.0, 300.0) + assert e_ship < -40.0 # the documented ~e^-40 + Ns, toks = AM._QUAD_LADDER_N, AM._QUAD_LADDER_TOK + assert len(Ns) == len(toks) + assert Ns[-1] == AM._LAPLACE_QUAD_N # top rung == shipped N + assert toks[-1] == AM._LAPLACE_BLEND_HI # ladder covers the band + for N, tok in zip(Ns, toks): + assert N % 16 == 0, (N, "quadrature scan chunk is 16") + e = _aliasing_exponent(float(N), float(tok)) + assert e <= e_ship, ( + "ladder rung N=%d at t_ok=%g has aliasing exponent %.2f, worse " + "than the shipped band edge %.2f: the rung under-resolves its " + "band" % (N, tok, e, e_ship)) + assert all(toks[i] < toks[i + 1] for i in range(len(toks) - 1)) + + +def test_dispatcher_preserves_three_input_broadcasting(): + """The production dispatcher must keep the public kernel's elementwise + broadcast contract before it introduces its private root-slot axis. + + Exercise both formerly broken directions: a vector supplied only by + ``a`` (whose length is deliberately not the four-root count), and a + vector supplied only by one coefficient. Scalar calls provide an + independent elementwise oracle without relying on the same broadcast. + """ + avec = jnp.asarray([-0.8, 0.1, 1.7]) + c1 = jnp.asarray(0.35 + 0.2j) + c2 = jnp.asarray(-0.15 + 0.05j) + got_a = np.asarray(AM._laplace_psi_lnI_block(avec, c1, c2)) + want_a = np.asarray([ + AM._laplace_psi_lnI_block(ai, c1, c2) for ai in avec + ]) + np.testing.assert_allclose(got_a, want_a, rtol=0.0, atol=1e-13) + + c2vec = jnp.asarray([0.05 + 0.02j, -0.1 + 0.03j, 0.2 - 0.04j]) + got_c = np.asarray(AM._laplace_psi_lnI_block(jnp.asarray(0.4), c1, c2vec)) + want_c = np.asarray([ + AM._laplace_psi_lnI_block(jnp.asarray(0.4), c1, c2i) + for c2i in c2vec + ]) + np.testing.assert_allclose(got_c, want_c, rtol=0.0, atol=1e-13) + + +def _batch(rng, tlo, thi, n=256): + t = rng.uniform(tlo, thi, n) + frac = rng.uniform(0, 1, n) + b = t * frac + d = 0.5 * t * (1 - frac) + beta = rng.uniform(-np.pi, np.pi, n) + delta = rng.uniform(-np.pi, np.pi, n) + a = jnp.asarray(rng.uniform(-5, 5, n)) + return (a, jnp.asarray(b * np.exp(1j * beta)), + jnp.asarray(d * np.exp(1j * delta))) + + +def test_dispatcher_matches_kernel_in_every_branch(): + """_laplace_psi_lnI_block == _laplace_psi_lnI on batches forced into + each switch branch: bit-equal where the same code runs (pure-Laplace, + top rung, straddle), and within ladder roundoff (documented sub-1e-12) + on the reduced-N rungs. An off-by-one in the index computation or a + threshold mutation shifts a batch into an inadequate branch and fails + this by many orders of magnitude (e.g. Laplace applied inside the + blend band: ~0.25 nats; N=32 applied at t=50: ~8%).""" + rng = np.random.default_rng(3) + kB = jax.jit(AM._laplace_psi_lnI_block) + kS = jax.jit(AM._laplace_psi_lnI) + exact_families = [(320.0, 5000.0), # pure-Laplace branch + (144.0, 219.0), # top (N=320) rung + (10.0, 4000.0), # straddle -> full blended kernel + (225.0, 295.0)] # blend band -> full kernel; sent + # to pure-Laplace by an index + # off-by-one (BLEND_LO for HI), + # which errs at the 0.05-0.25 nat + # branch-disagreement scale + for tlo, thi in exact_families: + a, c1, c2 = _batch(rng, tlo, thi) + dv = np.max(np.abs(np.asarray(kB(a, c1, c2)) - np.asarray(kS(a, c1, c2)))) + assert dv == 0.0, ("branch running identical code must be " + "bit-equal", tlo, thi, dv) + # every reduced-N rung (bands strictly inside their thresholds) + toks = (0.0,) + AM._QUAD_LADDER_TOK + for j in range(len(AM._QUAD_LADDER_N) - 1): + tlo = toks[j] * 1.05 + 1e-3 + thi = toks[j + 1] * 0.95 + a, c1, c2 = _batch(rng, tlo, thi) + dv = np.max(np.abs(np.asarray(kB(a, c1, c2)) - np.asarray(kS(a, c1, c2)))) + assert dv < 1e-12, ("rung N=%d disagrees with the N=320 kernel by " + "%g" % (AM._QUAD_LADDER_N[j], dv)) + + +def test_fused_driver_uses_the_block_dispatcher(): + """The fused laplace jaxpr must contain the dispatcher's switch: a cond + primitive with one branch per ladder rung + pure-Laplace + full kernel. + Fails if _dist_step is rewired back to the undispatched kernel (the + execution-cost regression this file exists to prevent).""" + n_branches = len(AM._QUAD_LADDER_N) + 2 + data = make_synth() + xg, lwg = make_distance_grid(30.0, 3000.0, 8, distMpcRef=data.distMpcRef) + + def f(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0) + + jaxpr = jax.make_jaxpr(f)(jnp.asarray([0.9]), jnp.asarray([0.4]), + jnp.asarray([1.1])) + found = [] + + def walk(jx): + for eqn in jx.eqns: + if eqn.primitive.name == "cond": + found.append(len(eqn.params["branches"])) + for val in eqn.params.values(): + vals = val if isinstance(val, (tuple, list)) else (val,) + for v in vals: + if hasattr(v, "jaxpr"): + walk(v.jaxpr) + elif hasattr(v, "eqns"): + walk(v) + walk(jaxpr.jaxpr) + assert n_branches in found, ( + "fused laplace traced graph has no %d-branch switch (cond branch " + "counts seen: %s): _dist_step is not calling " + "_laplace_psi_lnI_block, so every lattice point pays both kernel " + "branches again (~2,950x the grid scheme, 2026-08-28)" + % (n_branches, sorted(set(found)))) + + +def test_fused_value_and_grad_match_undispatched_kernel(): + """End-value wiring check on real-shaped packed data: the fused laplace + with the dispatcher equals the same fused call with the dispatcher + monkeypatched to the plain kernel, to ladder roundoff; the gradient is + finite and equally close. Catches any dispatch defect that survives + the per-branch kernel test (e.g. wrong operands captured).""" + data = make_synth(kappa_boost=4.0) + xg, lwg = make_distance_grid(30.0, 3000.0, 10, distMpcRef=data.distMpcRef) + ra = jnp.asarray([0.9, 2.1]) + dec = jnp.asarray([0.4, -0.7]) + incl = jnp.asarray([1.1, 2.4]) + + def call(): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0) + + v_disp = np.asarray(call()) + orig = AM._laplace_psi_lnI_block + try: + AM._laplace_psi_lnI_block = AM._laplace_psi_lnI + v_plain = np.asarray(call()) + finally: + AM._laplace_psi_lnI_block = orig + assert np.max(np.abs(v_disp - v_plain)) < 1e-11, (v_disp, v_plain) + + def scalar(th): + r, d_, i = th + return jnp.sum(AM.fused_log_likelihood_distphipsimarg_laplace( + data, r[None], d_[None], i[None], xg, lwg, amp_sizing=900.0)) + + g = np.asarray(jax.grad(scalar)((ra[0], dec[0], incl[0]))) + assert np.isfinite(g).all(), g diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py new file mode 100644 index 000000000..13bf48a0a --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py @@ -0,0 +1,343 @@ +""" +Gate for the COMPILE-COST structure of the anglemarg laplace path. + +WHY THIS EXISTS (2026-08-28): a single SNR-40 production run with +--angle-marg-scheme auto sat in XLA compilation for >88 minutes and reached +22.2 GiB RSS before being killed by hand (~/rift_costbakeoff_20260826/ +time_arms.log) -- against a 25 GiB per-user cgroup on the interactive hosts. +The cause was STRUCTURAL, not mathematical: Python-level loops in +_laplace_psi_lnI (24-cell bracket scan x 4 root slots, 4 x 20 bisection +steps, 320-point u-quadrature) unrolled into the traced graph, and the +distance-block Python loop in fused_log_likelihood_distphipsimarg_laplace +then instantiated that whole unrolled kernel once per distance block -- +64 copies at the production n_grid=256 -- inside a jax.checkpoint'ed scan +body that reverse-mode AD retraces. XLA compile time and memory are +superlinear in graph size, hence the wall. + +The fix rolls those loops into lax.scan / lax.fori_loop, so the traced graph +is CONSTANT in the distance-grid size and in the loop trip counts. These +tests pin exactly that structural property, plus the one numerical seam the +restructure introduced (tail padding of the distance grid). They are +trace-only where possible and run in seconds; the numerical VALIDATION of the +laplace scheme itself lives in test_angle_marg_exact.py (excluded from the +per-PR gate on cost grounds) and test_angle_marg_smoke.py. + +Each test here fails under a deliberate mutation (verified by hand before +landing; the mutations and observed failures are recorded in the PR): + * re-unrolling the distance scan into a Python loop -> graph-growth test + fails (equation count scales with n_grid again); + * re-unrolling any kernel loop -> kernel-size ceiling fails; + * breaking the tail padding (0.0 instead of -inf pad weights) -> the + padding-exactness test fails. +""" + +import numpy as np +import importlib.machinery +import importlib.util + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import build_likelihood_data +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile.core import make_distance_grid + + +def _total_eqns(closed_jaxpr): + """Equation count including nested jaxprs, each nested body counted ONCE + (scan/cond/checkpoint bodies are nested; an unrolled Python loop's + equations are all at one level, so unrolling inflates this count).""" + def walk(jaxpr): + n = len(jaxpr.eqns) + for eqn in jaxpr.eqns: + for val in eqn.params.values(): + vals = val if isinstance(val, (tuple, list)) else (val,) + for v in vals: + if hasattr(v, "jaxpr"): # ClosedJaxpr + n += walk(v.jaxpr) + elif hasattr(v, "eqns"): # raw Jaxpr + n += walk(v) + return n + return walk(closed_jaxpr.jaxpr) + + +def make_synth(scale=1.0, seed=3, modes=((2, 2), (2, -2)), npts=16, + deltaT=1.0 / 1024, kappa_boost=1.0): + """Small structurally-faithful packed data (as test_angle_marg_exact).""" + rng = np.random.default_rng(seed) + tw = npts * deltaT / 2.0 + tvals = np.linspace(-tw, tw, npts) + tref = 1126259462.413 + K = len(modes) + packed = {} + for det in ("H1", "L1"): + npts_full = 1024 + white = (rng.standard_normal((K, npts_full)) + + 1j * rng.standard_normal((K, npts_full))) + kx = np.arange(-40, 41) + kern = np.exp(-0.5 * (kx / 12.0) ** 2) + kern /= kern.sum() + rho = np.stack([np.convolve(white[k].real, kern, "same") + + 1j * np.convolve(white[k].imag, kern, "same") + for k in range(K)]).astype(np.complex128) + rho *= np.sqrt(len(kx)) * scale * kappa_boost + M = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + U = (M @ M.conj().T + 3 * np.eye(K)) * scale ** 2 + B = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = (B @ B.T) * scale ** 2 * 0.3 + packed[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=tref - 0.5) + return build_likelihood_data(packed, deltaT, tref, tvals) + + +def _fused_jaxpr(data, n_grid): + xg, lwg = make_distance_grid(30.0, 3000.0, n_grid, + distMpcRef=data.distMpcRef) + def f(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0) + return jax.make_jaxpr(f)(jnp.asarray([0.9]), jnp.asarray([0.4]), + jnp.asarray([1.1])) + + +def test_laplace_graph_size_independent_of_distance_grid(): + """The traced graph must NOT grow with the distance-grid size. + + Pre-fix, each extra dist_block=4 block of distance nodes re-instantiated + the full _laplace_psi_lnI kernel in the scan body (measured: hundreds of + extra equations per block; at the production n_grid=256 the resulting + graph took >88 min and >20 GiB to compile). With the distance nodes + scanned, the equation count is IDENTICAL for any n_grid: only the + scanned xs shapes change. Trace-only: no XLA compile, no execution. + """ + data = make_synth() + n8 = _total_eqns(_fused_jaxpr(data, 8)) + n64 = _total_eqns(_fused_jaxpr(data, 64)) + assert n64 == n8, ( + "laplace traced graph grew with the distance grid (%d -> %d eqns " + "for n_grid 8 -> 64): a distance-block loop is unrolling into the " + "graph again, which is the >88-minute XLA compile of 2026-08-28." + % (n8, n64)) + + +def test_laplace_kernel_graph_is_rolled(): + """_laplace_psi_lnI's traced size must stay near its rolled size. + + Measured at the commit that introduces this test (jax 0.9.2; jax 0.7.1 + within 1%): 400 equations rolled, 6819 with all three pre-fix Python + loops inlined, and PER-LOOP mutants of 835 (20-step bisection + unrolled -- the smallest), 1042 (24-cell bracket walk unrolled), 1410 + (320-point quadrature unrolled). The ceiling of 600 sits 1.5x above + the rolled size to absorb jax-version drift in how primitives are + counted, and 1.4x below the smallest single-loop mutant, so ANY one + loop unrolling again trips it. If a jax upgrade legitimately inflates + the rolled count past 600, re-measure all four numbers above before + touching the ceiling. + """ + shape = (4, 3) + a = jnp.zeros(shape) + c1 = jnp.full(shape, 30.0 + 10.0j) + c2 = jnp.full(shape, 40.0 - 5.0j) + n = _total_eqns(jax.make_jaxpr(AM._laplace_psi_lnI)(a, c1, c2)) + assert n <= 600, ( + "_laplace_psi_lnI traces to %d equations (ceiling 600): a bracket/" + "bisection/quadrature loop has unrolled into the graph again. That " + "size is multiplied by every distance block and by AD; see module " + "docstring." % n) + + +def test_laplace_kernel_preserves_mixed_shape_broadcasting(): + """The rolled scan must preserve the kernel's elementwise broadcast API. + + A scalar c1 and vector c2 make the derivative fields vector-valued. If + the scan carry and root-slot arrays are initialized from c1 alone, their + shapes start scalar and lax.scan fails instead of allowing the expansion + that the former Python loop performed. + """ + a = jnp.asarray(0.0) + c1 = jnp.asarray(350.0 + 20.0j) + c2 = jnp.asarray([10.0 + 2.0j, 20.0 - 3.0j]) + + got = AM._laplace_psi_lnI(a, c1, c2) + got_jit = jax.jit(AM._laplace_psi_lnI)(a, c1, c2) + expected = jnp.stack([ + AM._laplace_psi_lnI(a, c1, c2_i) for c2_i in c2 + ]) + + assert got.shape == c2.shape + assert np.allclose(np.asarray(got), np.asarray(expected), + rtol=0, atol=1e-12) + assert np.allclose(np.asarray(got_jit), np.asarray(expected), + rtol=0, atol=1e-12) + + # The root-slot axis must also remain distinct from a data axis supplied + # only by a. Length 3 catches the former shape error; length 4 catches + # the more dangerous silent collision with _LAPLACE_MAX_ROOTS. + c1_scalar = jnp.asarray(1000.0 + 30.0j) + c2_scalar = jnp.asarray(200.0 - 20.0j) + for n in (3, 4): + a_vector = jnp.arange(n, dtype=jnp.float64) + got = AM._laplace_psi_lnI(a_vector, c1_scalar, c2_scalar) + got_jit = jax.jit(AM._laplace_psi_lnI)( + a_vector, c1_scalar, c2_scalar) + expected = jnp.stack([ + AM._laplace_psi_lnI(a_i, c1_scalar, c2_scalar) + for a_i in a_vector + ]) + + assert got.shape == a_vector.shape + assert np.allclose(np.asarray(got), np.asarray(expected), + rtol=0, atol=1e-12) + assert np.allclose(np.asarray(got_jit), np.asarray(expected), + rtol=0, atol=1e-12) + + +def test_laplace_dist_tail_padding_exact(): + """A distance grid NOT divisible by dist_block must give the same + marginal as one evaluated without tail padding. + + The scan packs G nodes into blocks of dist_block, edge-padding the tail + with -inf log-weights (exactly-zero contribution to the running + log-sum-exp). Mutating the pad weights to 0.0 double-counts the last + node and shifts the marginal by ~log-weight amounts; this test fails + under that mutation and under any off-by-one in the packing. + """ + data = make_synth(kappa_boost=4.0) + xg, lwg = make_distance_grid(30.0, 3000.0, 10, distMpcRef=data.distMpcRef) + ra, dec, incl = jnp.asarray([0.9]), jnp.asarray([0.4]), jnp.asarray([1.1]) + # dist_block=4 -> 3 blocks, 2 padded tail nodes; dist_block=5 and =1 + # divide 10 exactly -> no padding. All three must agree to roundoff. + v4 = AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0, dist_block=4) + v5 = AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0, dist_block=5) + v1 = AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=900.0, dist_block=1) + assert np.allclose(np.asarray(v4), np.asarray(v5), rtol=0, atol=1e-12), \ + (np.asarray(v4), np.asarray(v5)) + assert np.allclose(np.asarray(v4), np.asarray(v1), rtol=0, atol=1e-12), \ + (np.asarray(v4), np.asarray(v1)) + + # Independent direct reference: evaluate each distance node alone with + # zero log-weight, then combine those scalar marginals with the original + # nonuniform weights. This pins every weight's value, node association, + # and completeness; agreement among blockings alone cannot see a shared + # permutation or truncation bug in the packing code. + node_vals = [] + for i in range(len(xg)): + node_vals.append(AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg[i:i + 1], jnp.zeros(1), + amp_sizing=900.0, dist_block=1)[0]) + direct = jax.scipy.special.logsumexp( + jnp.stack(node_vals) + lwg, axis=0) + assert np.allclose(np.asarray(v4[0]), np.asarray(direct), + rtol=0, atol=1e-11), (np.asarray(v4[0]), np.asarray(direct)) + + # Also pin AD through the padded distance fold. Adding c to every + # log-weight must add exactly c to the normalized log-sum-exp, so its + # derivative with respect to c is one. + c = 0.37 + def shifted(dc): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg + dc, amp_sizing=900.0, + dist_block=4)[0] + v_shift = shifted(c) + assert np.allclose(np.asarray(v_shift), np.asarray(v4[0] + c), + rtol=0, atol=1e-12) + assert np.allclose(np.asarray(jax.grad(shifted)(0.0)), 1.0, + rtol=0, atol=1e-12) + + +# --------------------------------------------------------------------------- +# Execution-side memory: the batched-eval chunk cap. +# +# Fixing the compile blowup exposed a SECOND wall the pre-fix code could +# never reach: at the samplers' default eval chunk (4000) the laplace path's +# stacked quadrature transient (quad_chunk*dist_block*phi_chunk*8 = 8192 +# bytes per sample per time point) is a single 36.41 GiB XLA buffer at +# npts=1193, and the SNR-40 acceptance run died RESOURCE_EXHAUSTED on the +# 25 GiB cgroup. Eval slices are independent (lnL is elementwise in the +# sample axis), so angle_marg_eval_chunk caps the chunk for anglemarg +# schemes -- peak memory changes, no number changes. These tests pin the +# WIRING, which is where such fixes rot (helper-level tests cannot see a +# call site that stops calling the helper). +# --------------------------------------------------------------------------- + +class _RecordingLike: + """Minimal like object: records every batch size it is asked for.""" + def __init__(self, scheme, npts): + import types as _t + self.angle_marg_scheme = scheme + self.data = _t.SimpleNamespace(npts=npts) + self.batches = [] + + def log_likelihood(self, ra, dec, incl): + self.batches.append(len(np.asarray(ra))) + return jnp.zeros(len(np.asarray(ra))) + + +def test_eval_chunk_cap_wired_for_anglemarg_schemes(): + """eval_lnL_3 must evaluate an anglemarg-scheme likelihood in capped + batches, and a grid-scheme one at the requested chunk. + + The cap for npts=1200 is (4 GiB)//(8192*1200) = 436 samples. Mutating + eval_lnL_3 to drop the angle_marg_eval_chunk call feeds the mock one + 1000-sample batch and this fails. + """ + from RIFT.likelihood.jax_ile import samplers as S + + theta = np.zeros((1000, 3)) + lap = _RecordingLike("laplace", 1200) + S.eval_lnL_3(lap, theta) + expected_cap = max(1, (4 << 30) // (8192 * 1200)) + assert max(lap.batches) == expected_cap, lap.batches + assert sum(lap.batches) == 1000 + + grid = _RecordingLike("grid", 1200) + S.eval_lnL_3(grid, theta) + assert max(grid.batches) == 1000, ( + "grid-scheme eval must NOT be capped (batches: %r)" % grid.batches) + + exact = _RecordingLike("exact", 1200) + S.eval_lnL_3(exact, theta) + assert max(exact.batches) == expected_cap, exact.batches + assert sum(exact.batches) == 1000 + + # helper edge cases: unknown npts or missing data -> untouched + import types as _t + assert S.angle_marg_eval_chunk( + _t.SimpleNamespace(angle_marg_scheme="exact"), 4000) == 4000 + assert S.angle_marg_eval_chunk( + _t.SimpleNamespace(angle_marg_scheme="exact", + data=_t.SimpleNamespace(npts=0)), 4000) == 4000 + # Long but supported integration windows must still honor the 4 GiB + # bound; the former floor of 64 turned this case into a ~32 GiB buffer. + long_like = _t.SimpleNamespace( + angle_marg_scheme="laplace", data=_t.SimpleNamespace(npts=65537)) + assert S.angle_marg_eval_chunk(long_like, 4000) == 7 + + +def test_driver_eval_applies_the_chunk_cap(): + """The driver's own eval_lnL (the --n-chunk 8000 loop) must consult + angle_marg_eval_chunk -- the samplers-level wiring test cannot see this + call site. Drive the real loop: a source substring can remain present + while a later assignment silently overwrites the capped value.""" + import pathlib + import types + path = (pathlib.Path(__file__).resolve().parents[2] / "bin" + / "integrate_likelihood_extrinsic_jax") + loader = importlib.machinery.SourceFileLoader( + "_ile_jax_driver_chunk_test", str(path)) + spec = importlib.util.spec_from_loader("_ile_jax_driver_chunk_test", loader) + drv = importlib.util.module_from_spec(spec) + loader.exec_module(drv) + + theta = np.zeros((1000, 3)) + like = _RecordingLike("laplace", 1200) + opts = types.SimpleNamespace(n_chunk=1000) + out = drv.eval_lnL(like, theta, opts, with_distance=False) + expected_cap = (4 << 30) // (8192 * 1200) + assert np.asarray(out).shape == (1000,) + assert max(like.batches) == expected_cap, like.batches + assert sum(like.batches) == 1000 diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_default.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_default.py new file mode 100644 index 000000000..6ccefce2c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_default.py @@ -0,0 +1,119 @@ +"""The (phi_ref, psi) scheme default has exactly ONE definition, and one +look-alike literal that must NOT follow it. + +Changed 2026-09-02: 'grid' -> 'exact'. The previous default move on this path +(interp linear -> sinc) was bitten by the value being re-typed in many places, +and by two independent things that happened to be the same string. Both hazards +are pinned here. +""" +import importlib.machinery +import importlib.util +import inspect +import os +import re + +import pytest + +from RIFT.likelihood.jax_ile.anglemarg import (ANGLE_MARG_CHOICES, + ANGLE_MARG_DEFAULT, + ANGLE_MARG_LEGACY) +from RIFT.likelihood.jax_ile.samplers import angle_marg_eval_chunk +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood + +_CODE = os.path.abspath( + os.path.join(os.path.dirname(__file__), os.pardir, os.pardir)) +_DRIVER = os.path.join(_CODE, "bin", "integrate_likelihood_extrinsic_jax") + + +def _driver(): + loader = importlib.machinery.SourceFileLoader("_amd_drv", _DRIVER) + spec = importlib.util.spec_from_loader("_amd_drv", loader) + mod = importlib.util.module_from_spec(spec) + mod.__name__ = "_amd_drv" # keep the __main__ guard from firing + loader.exec_module(mod) + return mod + + +def _parse(argv): + optp = _driver().build_parser() + opts, _ = optp.parse_args(list(argv)) + return opts + + +def test_default_has_a_single_definition(): + """Driver flag, library argument and the constant must all agree.""" + parsed = _parse(["--inj-mode", "--mass1", "35", "--mass2", "30"]) + lib = inspect.signature( + JAXDistPhiPsiMargLikelihood.__init__).parameters["angle_marg"].default + assert parsed.angle_marg_scheme == ANGLE_MARG_DEFAULT + assert lib == ANGLE_MARG_DEFAULT + assert ANGLE_MARG_DEFAULT in ANGLE_MARG_CHOICES + + +def test_legacy_spelling_still_reachable(): + """The pre-change behaviour must remain reproducible by an explicit flag.""" + assert ANGLE_MARG_LEGACY == "grid" + assert ANGLE_MARG_LEGACY in ANGLE_MARG_CHOICES + assert ANGLE_MARG_LEGACY != ANGLE_MARG_DEFAULT + opts = _parse(["--inj-mode", "--mass1", "35", "--mass2", "30", + "--angle-marg-scheme", ANGLE_MARG_LEGACY]) + assert opts.angle_marg_scheme == ANGLE_MARG_LEGACY + + +def test_unknown_scheme_is_rejected_at_parse_time(): + """A typo must die at argument parsing, not minutes later after precompute.""" + with pytest.raises(SystemExit): + _parse(["--angle-marg-scheme", "definitely-not-a-scheme"]) + + +def test_eval_chunk_sentinel_does_not_follow_the_default(): + """`angle_marg_eval_chunk`'s "grid" is a SENTINEL, not the default. + + A likelihood with no ``angle_marg_scheme`` at all (JAXDistanceMarginalized, + JAXExtrinsic) runs no dense angle scheme and must pass its chunk through + UNCHANGED. If someone syncs that literal to ANGLE_MARG_DEFAULT the chunk + gets capped for every such object -- a memory/throughput regression with no + benefit. This is the "two independent defaults, same string" trap. + """ + class _Data: + # npts must be large enough that the cap actually bites, or a SECOND + # early return (npts <= 0) masks the scheme check and this test passes + # for the wrong reason -- it did, and a mutation sweep caught it. + npts = 65537 + + class _NoScheme: # no angle_marg_scheme attribute at all + data = _Data() + + class _GridScheme: + angle_marg_scheme = "grid" + data = _Data() + + class _ExactScheme: + angle_marg_scheme = "exact" + data = _Data() + + # Positive control: the cap DOES bite for a dense scheme at this npts, so a + # pass-through below is discrimination and not an inert code path. + capped = angle_marg_eval_chunk(_ExactScheme(), 4096) + assert capped < 4096, ( + "cap did not engage at npts=%d; this test cannot discriminate" + % _Data.npts) + + assert angle_marg_eval_chunk(_NoScheme(), 4096) == 4096 + assert angle_marg_eval_chunk(_GridScheme(), 4096) == 4096 + assert ANGLE_MARG_DEFAULT == "exact", ( + "if the default is no longer 'exact', re-derive this test's premise") + + +def test_no_retyped_default_literal_in_driver_or_wrapper(): + """Neither entry point may spell the default as a bare literal.""" + for path in (_DRIVER, + os.path.join(_CODE, "RIFT", "likelihood", "jax_ile", + "wrapper.py")): + with open(path) as fh: + src = fh.read() + assert not re.search(r'angle_marg_scheme"\s*,\s*"grid"', src), ( + "%s re-types the angle-marg default as a literal" % path) + assert not re.search(r'angle_marg\s*=\s*"(grid|exact|laplace|auto)"', + src), ( + "%s re-types the angle-marg default as a literal" % path) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py new file mode 100644 index 000000000..5cdbe9eaa --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py @@ -0,0 +1,1153 @@ +""" +Gate for RIFT.likelihood.jax_ile.anglemarg: the exact (phi_ref, psi) +marginalization schemes and their selector. + +WHAT IS PINNED, AND WHY THESE PARTICULAR TESTS +---------------------------------------------- +The module rests on one analytic fact: at fixed time and distance the factored +lnL is a bivariate trig polynomial -- phi_ref order <= 2*m_max, u = 2 psi +order <= 2. Everything else (Nyquist sample sizing, coefficient tables, +dense reconstruction, the psi Laplace) is bookkeeping on top of that fact, and +each layer of the bookkeeping is pinned here: + + * the harmonic-content invariant itself (sampler-free, injection-free; it + catches exactly what a rewrite gets wrong: hermiticity of U/V handling, + mode conventions, array alignment). CAVEAT honoured here: the invariant + holds for the UNMARGINALIZED lnL at fixed time -- log-sum-exp over time + destroys the polynomial structure and manufactures fake high harmonics, + so every decomposition below fixes the time index; + * the coefficient tables reproduce the direct likelihood OFF the sample + grid (trig interpolation exactness); + * both schemes against a brute-force dense reference (which converges TO + the exact answer), and against the legacy grid path where that path is + converged (pins the shared normalization convention); + * the historical nphi=8 Nyquist aliasing of the n=4 phi harmonic, at the + DFT-coefficient level and at the marginal level (the regression that + motivated the module); + * exact/laplace agreement in the selector's overlap region -- the runtime + check that makes the crossover a validated constant, not a tuning knob; + * gradients (the point of this code path) against finite differences; + * the wrapper selector, its provenance record, and the grid default being + byte-identical to the legacy path (no default change); + * the driver flag actually reaching the wrapper and the resolved scheme + being printed (this pipeline has a documented history of silently-inert + flags), via AST over the driver source. + +Synthetic packed data (no frames, no waveform generation) keeps this fast and +CI-friendly; the trig-polynomial structure is a property of the accumulation +ALGEBRA, not of any particular rholm/U/V values, so random-but-structured +tensors (U Hermitian PSD, V symmetric) exercise it fully. +""" + +import ast +import os +import types + +import numpy as np +import pytest + +# numpy renamed trapz -> trapezoid in 2.0. The IGWN CVMFS python still +# ships numpy 1.26, where the new name does not exist -- and these two +# call sites are inside this file's own BRUTE-FORCE REFERENCE, so on that +# interpreter four test_laplace_kernel_* tests died with an AttributeError +# before comparing anything, and read as four kernel failures. RIFT's own +# library code already carries this shim (misc/distance_grid.py, +# misc/distance_slices.py); the test file had missed it. +_trapz = np.trapezoid if hasattr(np, "trapezoid") else np.trapz + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import build_likelihood_data +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile.core import ( + _accumulate_unit, _time_marginalize, _logsumexp_grid_blocked, + fused_log_likelihood_distphipsimarg, phi_ref_grid, psi_grid, + make_distance_grid) +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood + +INTERP = "sinc" + + +# --------------------------------------------------------------------------- +# synthetic likelihood data +# --------------------------------------------------------------------------- + +def make_synth(scale=1.0, seed=3, modes=((2, 2), (2, -2)), npts=32, + deltaT=1.0 / 1024, kappa_boost=1.0): + """Structurally-faithful synthetic packed data (cf. test_jax_likelihood). + + U is Hermitian positive definite and V complex symmetric, as the real + precompute produces; ``scale`` sets the overall amplitude (lnL ~ scale^2), + standing in for SNR. ``kappa_boost`` multiplies the rholm timeseries + ONLY (not U/V), producing a target with a large coherent (phi,psi) + amplitude A -- the regime where an undersized dense grid measurably + biases the marginal (used by the sizing regression tests). + """ + rng = np.random.default_rng(seed) + tw = npts * deltaT / 2.0 + tvals = np.linspace(-tw, tw, npts) + tref = 1126259462.413 + K = len(modes) + packed = {} + for det in ("H1", "L1"): + npts_full = 4096 + white = (rng.standard_normal((K, npts_full)) + + 1j * rng.standard_normal((K, npts_full))) + kx = np.arange(-40, 41) + kern = np.exp(-0.5 * (kx / 12.0) ** 2) + kern /= kern.sum() + rho = np.stack([np.convolve(white[k].real, kern, "same") + + 1j * np.convolve(white[k].imag, kern, "same") + for k in range(K)]).astype(np.complex128) + rho *= np.sqrt(len(kx)) * scale * kappa_boost + M = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + U = (M @ M.conj().T + 3 * np.eye(K)) * scale ** 2 + B = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = (B @ B.T) * scale ** 2 * 0.3 + packed[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=tref - 0.5) + return build_likelihood_data(packed, deltaT, tref, tvals) + + +RA, DEC, INCL = np.array([0.9]), np.array([0.4]), np.array([1.1]) +S = 1 + + +def _dist_grid(data, n=64): + return make_distance_grid(30.0, 3000.0, n, distMpcRef=data.distMpcRef) + + +def brute_marginal(data, x_grid, log_w, nphi, npsi): + """Brute-force dist+phi+psi marginal: dense product grid of DIRECT + likelihood evaluations (no coefficient machinery shared with the schemes + under test).""" + ph = np.linspace(0, 2 * np.pi, nphi, endpoint=False) + ps = np.linspace(0, np.pi, npsi, endpoint=False) + m = jnp.full((S, data.npts), -jnp.inf) + s = jnp.zeros((S, data.npts)) + for p in ph: + rb = np.repeat(RA[None, :], npsi, 0).ravel() + db = np.repeat(DEC[None, :], npsi, 0).ravel() + ib = np.repeat(INCL[None, :], npsi, 0).ravel() + pb = np.full(npsi * S, p) + sb = np.repeat(ps[:, None], S, 1).ravel() + ku, rs = _accumulate_unit(data, rb, db, sb, ib, pb, INTERP, False) + lnL = _logsumexp_grid_blocked(ku.real, rs, x_grid, + -0.5 * jnp.square(x_grid), log_w, 64) + m, s = AM._lse_update(m, s, lnL.reshape(npsi, S, data.npts), axis=0) + lnL_t = m + jnp.log(s) - np.log(nphi * npsi) + return np.asarray(_time_marginalize(lnL_t, data.w_t)) + + +def _lnL_t_fixed_time(data, phis, psis, x, t_index): + """UNMARGINALIZED lnL at fixed time bin and fixed distance factor x.""" + n = len(phis) + ku, rs = _accumulate_unit(data, np.full(n, RA[0]), np.full(n, DEC[0]), + np.asarray(psis), np.full(n, INCL[0]), + np.asarray(phis), INTERP, False) + return (np.asarray(ku.real)[:, t_index] * x + - 0.5 * np.asarray(rs)[:, t_index] * x ** 2) + + +# --------------------------------------------------------------------------- +# 1. the harmonic-content invariant +# --------------------------------------------------------------------------- + +def test_harmonic_content_psi(): + """lnL_t(psi) at fixed (phi, t, x) has u-harmonics {0,1,2} ONLY.""" + data = make_synth(scale=2.0) + n = 64 + psis = np.linspace(0, np.pi, n, endpoint=False) + f = _lnL_t_fixed_time(data, np.full(n, 1.3), psis, 0.8, t_index=10) + C = np.fft.rfft(f) / n + power = np.abs(C) + assert power[:3].max() > 0 + # everything above u-order 2 is numerically zero + assert power[3:].max() < 1e-12 * power.max() + + +def test_harmonic_content_phi(): + """lnL_t(phi) at fixed (psi, t, x) has harmonics <= 2*m_max, odd ones zero + for a (2,+-2)-only mode set.""" + data = make_synth(scale=2.0) + n = 64 + phis = np.linspace(0, 2 * np.pi, n, endpoint=False) + f = _lnL_t_fixed_time(data, phis, np.full(n, 0.6), 0.8, t_index=10) + C = np.fft.rfft(f) / n + power = np.abs(C) + m_max = 2 + assert power[: 2 * m_max + 1].max() > 0 + assert power[2 * m_max + 1:].max() < 1e-12 * power.max() + # (2,+-2) only: odd phi harmonics identically absent + assert power[1] < 1e-12 * power.max() + assert power[3] < 1e-12 * power.max() + + +def test_time_marginalization_destroys_the_invariant(): + """The caveat that confounded a first analysis, pinned so nobody re-learns + it: after log-sum-exp over TIME the psi decomposition has fake high + harmonics. (Guards the tests above against being 'simplified' onto the + marginalized quantity.)""" + data = make_synth(scale=6.0) + n = 64 + psis = np.linspace(0, np.pi, n, endpoint=False) + ku, rs = _accumulate_unit(data, np.full(n, RA[0]), np.full(n, DEC[0]), + psis, np.full(n, INCL[0]), np.full(n, 1.3), + INTERP, False) + lnL_t = np.asarray(ku.real) * 0.8 - 0.5 * np.asarray(rs) * 0.8 ** 2 + f = np.asarray(_time_marginalize(jnp.asarray(lnL_t), data.w_t)) + C = np.abs(np.fft.rfft(f) / n) + assert C[3:].max() > 1e-9 * C.max() + + +# --------------------------------------------------------------------------- +# 2. sample-grid sizing is derived and asserted, not settable +# --------------------------------------------------------------------------- + +def test_sample_grid_sizes(): + assert AM.angle_sample_grid_sizes(2) == (16, 8) + for m_max in (1, 2, 3, 4, 5): + nphi_s, npsi_s = AM.angle_sample_grid_sizes(m_max) + # strictly above Nyquist for the highest harmonic present + assert nphi_s > 2 * (2 * m_max) + assert npsi_s > 2 * 2 + with pytest.raises(ValueError): + AM.angle_sample_grid_sizes(0) + # the public entry points take NO sample-size argument at all + import inspect + for fn in (AM.fused_log_likelihood_distphipsimarg_exact, + AM.fused_log_likelihood_distphipsimarg_laplace): + assert not any("nphi" in p or "npsi" in p + for p in inspect.signature(fn).parameters) + + +# --------------------------------------------------------------------------- +# 3. coefficient tables reproduce the direct likelihood OFF the sample grid +# --------------------------------------------------------------------------- + +def test_coefficient_tables_reconstruct_off_grid(): + data = make_synth(scale=3.0) + C_A, C_B, meta = AM.angle_coefficient_tables(data, RA, DEC, INCL, + interp=INTERP) + assert meta["m_max"] == 2 and meta["nphi_s"] == 16 and meta["npsi_s"] == 8 + rng = np.random.default_rng(11) + phis = rng.uniform(0, 2 * np.pi, 7) + psis = rng.uniform(0, np.pi, 7) + A_rec = np.asarray(AM._reconstruct_field(C_A, jnp.asarray(phis), + jnp.asarray(2 * psis))) + B_rec = np.asarray(AM._reconstruct_field(C_B, jnp.asarray(phis), + jnp.asarray(2 * psis))) + ku, rs = _accumulate_unit(data, np.full(7, RA[0]), np.full(7, DEC[0]), + psis, np.full(7, INCL[0]), phis, INTERP, False) + A_dir = np.asarray(ku.real)[:, None, :] + B_dir = np.asarray(rs)[:, None, :] + ref = max(np.abs(A_dir).max(), np.abs(B_dir).max()) + assert np.abs(A_rec - A_dir).max() < 1e-10 * ref + assert np.abs(B_rec - B_dir).max() < 1e-10 * ref + + +# --------------------------------------------------------------------------- +# 4/5. schemes against brute force and against the converged legacy grid +# --------------------------------------------------------------------------- + +def test_exact_scheme_vs_bruteforce(): + data = make_synth(scale=6.0) + x_grid, log_w = _dist_grid(data) + ref = brute_marginal(data, x_grid, log_w, 96, 48) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + assert np.abs(ex - ref).max() < 1e-10 + + +def test_exact_matches_legacy_grid_where_converged(): + """Same normalization convention as the production grid path: at low + amplitude the legacy 32x8 grid is converged and the two must agree.""" + data = make_synth(scale=2.0) + x_grid, log_w = _dist_grid(data) + leg = np.asarray(fused_log_likelihood_distphipsimarg( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, phi_ref_grid(32), psi_grid(8), interp=INTERP)) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + # legacy 32x8 truncation at this amplitude measured 2.2e-8; the bound + # pins the shared normalization convention, not the grid's residual + assert np.abs(ex - leg).max() < 1e-6 + + +def test_laplace_high_amplitude_accuracy_and_trend(): + """Laplace error is small at high amplitude and SHRINKS as amplitude + grows (it is O(1/A)); measured on this configuration: 1.8e-2 at scale 50, + 5.6e-3 at scale 100.""" + errs = [] + for scale in (50.0, 100.0): + data = make_synth(scale=scale) + x_grid, log_w = _dist_grid(data) + ref = brute_marginal(data, x_grid, log_w, 192, 96) + lp = np.asarray(AM.fused_log_likelihood_distphipsimarg_laplace( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + errs.append(np.abs(lp - ref).max()) + # measured on this configuration: 0.055 at scale 50, 0.028 at scale 100. + # NOTE this synthetic target is Laplace's WORST case (noise-like data, no + # coherent peak: every bin sits near the small-b regime); the real-signal + # injection ladder in the PR measures |laplace-exact| ~ 1e-3 at A=50 and + # falling. The bound here is a regression pin, not the operating error. + assert errs[0] < 0.15 + assert errs[1] < 0.08 + assert errs[1] < errs[0] + + +def test_overlap_agreement_exact_vs_laplace(): + """The selector's runtime validity check: in the overlap region the two + schemes agree, so the crossover cannot be silently mis-set -- either + branch is accurate there.""" + data = make_synth(scale=100.0) + x_grid, log_w = _dist_grid(data) + args = (data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + lp = np.asarray(AM.fused_log_likelihood_distphipsimarg_laplace( + *args, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + assert np.abs(ex - lp).max() < 0.06 + + +# --------------------------------------------------------------------------- +# 6. the nphi=8 Nyquist aliasing regression +# --------------------------------------------------------------------------- + +def test_nphi8_aliases_the_n4_harmonic_dft_level(): + """8-point phi sampling puts the n=4 harmonic AT Nyquist: the DFT bin + collapses C4 + C-4 = 2 Re C4 and the imaginary part is unrecoverable. + Pinned at the coefficient level, deterministically.""" + data = make_synth(scale=3.0) + n_good = 16 + phis16 = np.linspace(0, 2 * np.pi, n_good, endpoint=False) + f16 = _lnL_t_fixed_time(data, phis16, np.full(n_good, 0.6), 0.8, 10) + C4_true = np.fft.fft(f16)[4] / n_good + phis8 = np.linspace(0, 2 * np.pi, 8, endpoint=False) + f8 = _lnL_t_fixed_time(data, phis8, np.full(8, 0.6), 0.8, 10) + C4_alias = np.fft.fft(f8)[4] / 8 + # the alias identity: the 8-point bin 4 is exactly 2*Re(C4), not C4 + assert abs(C4_alias - 2 * C4_true.real) < 1e-12 * abs(C4_true) + # and the information it destroyed was genuinely there + assert abs(C4_true.imag) > 1e-3 * abs(C4_true) + + +def test_nphi8_marginal_regression(): + """The production 8x8 grid is measurably wrong at moderate amplitude + while the exact scheme is not.""" + data = make_synth(scale=25.0) + x_grid, log_w = _dist_grid(data) + ref = brute_marginal(data, x_grid, log_w, 128, 64) + leg8 = np.asarray(fused_log_likelihood_distphipsimarg( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, phi_ref_grid(8), psi_grid(8), interp=INTERP)) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + assert np.abs(leg8 - ref).max() > 1e-3 # the defect (measured 4.5e-2) + assert np.abs(ex - ref).max() < 1e-9 # the fix + + +# --------------------------------------------------------------------------- +# 7. gradients +# --------------------------------------------------------------------------- + +def test_exact_gradient_matches_finite_differences(): + data = make_synth(scale=25.0) + x_grid, log_w = _dist_grid(data) + + def scalar(theta): + return AM.fused_log_likelihood_distphipsimarg_exact( + data, theta[0:1], theta[1:2], theta[2:3], + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)[0] + + theta0 = jnp.asarray([RA[0], DEC[0], INCL[0]]) + v, g = jax.jit(jax.value_and_grad(scalar))(theta0) + g = np.asarray(g) + assert np.all(np.isfinite(g)), "AD gradient is not finite: %r" % (g,) + h = 1e-5 + for i in range(3): + tp = theta0.at[i].add(h) + tm = theta0.at[i].add(-h) + fd = (float(scalar(tp)) - float(scalar(tm))) / (2 * h) + assert abs(fd - g[i]) < 1e-4 * max(1.0, abs(fd)), \ + "param %d: fd %g vs AD %g" % (i, fd, g[i]) + + +def test_laplace_gradient_matches_exact_scheme(): + """Laplace's lnL is piecewise-smooth (series cut, stationary-point + rejection migrate bins as parameters move), so a finite difference on a + noise-like synthetic target is contaminated by kink noise (measured: FD + inconsistent between h=1e-4 and 1e-5 in one component while two others + match AD to 1e-5). The clean scheme-level check is AD-vs-AD against the + exact scheme, which shares no marginalization code beyond the coefficient + tables: measured agreement ~5e-4 relative at scale 100.""" + data = make_synth(scale=100.0) + x_grid, log_w = _dist_grid(data) + theta0 = jnp.asarray([RA[0], DEC[0], INCL[0]]) + grads = {} + for name, fn in (("exact", AM.fused_log_likelihood_distphipsimarg_exact), + ("laplace", + AM.fused_log_likelihood_distphipsimarg_laplace)): + def scalar(theta, fn=fn): + return fn(data, theta[0:1], theta[1:2], theta[2:3], + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)[0] + v, g = jax.jit(jax.value_and_grad(scalar))(theta0) + grads[name] = np.asarray(g) + assert np.all(np.isfinite(grads[name])), \ + "%s AD gradient is not finite: %r" % (name, grads[name]) + scale_ref = np.abs(grads["exact"]).max() + assert np.abs(grads["laplace"] - grads["exact"]).max() < 1e-2 * scale_ref + + +# --------------------------------------------------------------------------- +# 7b. the Laplace kernel in isolation: gradient exactness and the O(1/b) law +# --------------------------------------------------------------------------- + +def _kernel(p): + a = p[0] + c1 = p[1] + 1j * p[2] + c2 = p[3] + 1j * p[4] + return AM._laplace_psi_lnI(a, c1, c2) + + +def test_laplace_kernel_gradient_finite_differences(): + """On smooth inputs (away from the branch boundaries) the kernel gradient + is FD-exact -- both the Laplace branch and the small-amplitude series.""" + for p0 in ([0.3, 400.0, -250.0, 30.0, 15.0], # Laplace branch, t ~ 540 + [0.1, 0.12, 0.08, 0.03, -0.02], # tiny amplitude (quad) + [0.3, 40.0, -25.0, 3.0, 1.5], # moderate (quad), b ~ 47 + [0.2, 150.0, 80.0, 30.0, -20.0]): # blend band, t ~ 242 + p0 = jnp.asarray(p0) + g = np.asarray(jax.grad(_kernel)(p0)) + assert np.all(np.isfinite(g)) + h = 1e-6 + for i in range(5): + fd = (float(_kernel(p0.at[i].add(h))) + - float(_kernel(p0.at[i].add(-h)))) / (2 * h) + assert abs(fd - g[i]) < 1e-6 * max(1.0, abs(fd)) + + +def test_laplace_kernel_error_law(): + """Kernel error vs a dense trapezoid truth follows ~0.1/b nats and + SHRINKS with amplitude -- including the two-maxima regime (d ~ b/2). + Measured: 7.7e-4 at b=200, 3.3e-5 at b=2000, 4.6e-6 at b=20000.""" + cases = ((200.0, 0.7, 12.0, -0.4, 1.0), # blend band now (t = 224) + (2000.0, -1.2, 80.0, 0.9, 0.0), # pure Laplace + (20000.0, 0.3, 700.0, -2.0, 0.0), # pure Laplace, decade up + (50.0, 1.0, 30.0, 2.0, 0.0)) # two-maxima; quadrature now + errs = {} + for b, beta, d, delta, a in cases: + c1 = b * np.exp(-1j * beta) + c2 = d * np.exp(-1j * delta) + val = float(AM._laplace_psi_lnI(jnp.asarray(a), jnp.asarray(c1), + jnp.asarray(c2))) + u = np.linspace(0, 2 * np.pi, 2_000_001) + f = a + b * np.cos(u - beta) + d * np.cos(2 * u - delta) + fm = f.max() + truth = fm + np.log(_trapz(np.exp(f - fm), u) / (2 * np.pi)) + errs[b] = abs(val - truth) + assert errs[b] < 0.5 / b, "b=%g: err %g exceeds the O(1/b) law" % ( + b, errs[b]) + # the decay comparison must pair two PURE-Laplace points: the smaller + # cases now route through the (near-exact) quadrature/blend branches, so + # their errors sit far BELOW the Laplace law rather than on it + assert errs[20000.0] < errs[2000.0] + + +def test_dense_size_rule_pinned(): + """The dense-reconstruction sizing rule is a CALIBRATED constant, not a + knob (see the derivation note in anglemarg.py): pin its values at the + crossover amplitude and its floors, and that it can only grow with the + amplitude it must cover.""" + assert AM._dense_grid_sizes(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) == (352, 176) + assert AM._dense_grid_sizes(1.0) == (128, 64) # floors bind + n_lo = AM._dense_grid_sizes(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + n_hi = AM._dense_grid_sizes(4 * AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + assert n_hi[0] >= 2 * n_lo[0] - 16 and n_hi[1] >= 2 * n_lo[1] - 16 + # review 3 (P1): the phi axis must additionally scale with the mode + # content -- amplitude alone under-resolves higher modes (a pure order-8 + # term at A = 450 was phase-dependently ~0.037 nats wrong, order 16 up + # to 1.2 nats, under the amplitude-only rule). The u axis never scales: + # psi enters at spin-2 for every mode. + m4 = AM._dense_grid_sizes(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, m_max=4) + assert m4[0] >= 2 * n_lo[0] - 16 # phi doubles at m_max = 4 + assert m4[1] == n_lo[1] # u unchanged + m8 = AM._dense_grid_sizes(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, m_max=8) + assert m8 == (1360, 176) # 16 * 4 * sqrt(450) -> 1360 after rounding; + # (4 * the ROUNDED m_max=2 value would be 1408 -- the scaling applies to + # the unrounded rule, so pin the exact derived number instead) + + +# --------------------------------------------------------------------------- +# 8. the selector +# --------------------------------------------------------------------------- + +def test_choose_angle_marg_scheme(): + cross = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + s, info = AM.choose_angle_marg_scheme(cross * 0.9, gh_enabled=False) + assert s == "exact" + s, info = AM.choose_angle_marg_scheme(cross * 1.1, gh_enabled=False) + assert s == "laplace" + assert info["crossover"] == cross + # the selector keys on the MEASURED amplitude bound, never on an SNR + # guess: its signature must not accept one (external-review defect 2) + import inspect + assert "guess_snr" not in inspect.signature( + AM.choose_angle_marg_scheme).parameters + # no amplitude available: exact (the conservative branch) + s, info = AM.choose_angle_marg_scheme(None) + assert s == "exact" and "no amplitude" in info["reason"] + # adaptive distance quadrature forces the exact branch + s, info = AM.choose_angle_marg_scheme(cross * 100, gh_enabled=True) + assert s == "exact" and "DISTMARG_GH" in info["reason"] + + +def test_laplace_refuses_gh_env_above_the_covered_mode_content(monkeypatch): + """JAX_ILE_DISTMARG_GH + laplace is HONOURED for the mode content its + psi-marginal node placement is validated for, and RAISES above it -- never + silently ignores the env var (documented silently-inert-flag history). + The placement itself is gated in test_angle_marg_gh_laplace.py.""" + from RIFT.likelihood.jax_ile import core as core_mod + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 8) + x_grid, log_w = _dist_grid(make_synth(scale=2.0)) + + def call(data): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + + # covered: accepted, and the env var demonstrably reaches the result + assert AM._GH_PSI_M_MAX == 2 + got_gh = np.asarray(call(make_synth(scale=2.0))) + assert np.isfinite(got_gh).all() + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 0) + got_uniform = np.asarray(call(make_synth(scale=2.0))) + assert np.abs(got_gh - got_uniform).max() > 1e-6, ( + "JAX_ILE_DISTMARG_GH made no difference to the laplace answer -- the " + "flag is inert and 'honoured' means nothing") + # above the covered mode content: raises rather than guessing + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 8) + with pytest.raises(ValueError, match="m_max"): + call(make_synth(scale=2.0, modes=((2, 2), (2, -2), (3, 3), (3, -3)))) + + +def test_exact_supports_gh_env(monkeypatch): + """The exact scheme honors JAX_ILE_DISTMARG_GH *identically to the grid + path*: with GH active on both, exact must match a converged legacy grid + to the angle-quadrature floor. (GH-vs-uniform is a property of the + distance treatment itself, deliberately NOT re-litigated here: the + comparison holds the distance treatment fixed on both sides.)""" + from RIFT.likelihood.jax_ile import core as core_mod + data = make_synth(scale=6.0) + x_grid, log_w = _dist_grid(data, n=128) + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 33) + gh_exact = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + gh_legacy = np.asarray(fused_log_likelihood_distphipsimarg( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, phi_ref_grid(64), psi_grid(32), interp=INTERP)) + assert np.abs(gh_exact - gh_legacy).max() < 1e-6 + + +# --------------------------------------------------------------------------- +# 9. the wrapper: selection, provenance, and NO default change +# --------------------------------------------------------------------------- + +def test_wrapper_default_and_legacy_path(): + """The wrapper's default follows ANGLE_MARG_DEFAULT, and the LEGACY + spelling still reproduces the historical grid quadrature exactly. + + Changed 2026-09-02: the default moved 'grid' -> 'exact' (its quadrature + error grows without bound with SNR). What must not change is that + angle_marg=ANGLE_MARG_LEGACY still equals the direct grid call bit for bit, + because that is the contract archived runs are reproduced under. + """ + data = make_synth(scale=2.0) + default_like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, + npsi=8, n_grid=64, interp=INTERP) + assert default_like.angle_marg_scheme == AM.ANGLE_MARG_DEFAULT + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + n_grid=64, interp=INTERP, + angle_marg=AM.ANGLE_MARG_LEGACY) + assert like.angle_marg_scheme == AM.ANGLE_MARG_LEGACY + x_grid, log_w = like.x_grid, like.log_w_grid + direct = np.asarray(fused_log_likelihood_distphipsimarg( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, phi_ref_grid(32), psi_grid(8), interp=INTERP)) + got = np.asarray(like.log_likelihood(RA, DEC, INCL)) + # jit-vs-eager can differ in the last ulp; anything beyond that would + # mean the default path changed + assert np.abs(got - direct).max() < 1e-12 + + +def test_wrapper_auto_selects_and_records(): + """Auto selection keys on the DATA (measured amplitude bound), not on + guess_snr: a quiet target selects exact regardless of a huge claimed + SNR, a loud target selects laplace regardless of a missing one.""" + lo = JAXDistPhiPsiMargLikelihood(make_synth(scale=2.0), 30.0, 3000.0, + n_grid=64, interp=INTERP, + guess_snr=1000.0, angle_marg="auto") + assert lo.angle_marg_scheme == "exact" + hi = JAXDistPhiPsiMargLikelihood(make_synth(scale=2.0, kappa_boost=200.0), + 30.0, 3000.0, n_grid=64, interp=INTERP, + guess_snr=None, angle_marg="auto") + assert hi.angle_marg_scheme == "laplace" + assert hi.angle_marg_info["amplitude"] > AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + for like in (lo, hi): + info = like.angle_marg_info + assert info["requested"] == "auto" + assert info["scheme"] == like.angle_marg_scheme + assert "reason" in info and "crossover" in info + assert info["amp_sizing"] >= AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + assert info["sample_grid"] == (16, 8) + with pytest.raises(ValueError): + JAXDistPhiPsiMargLikelihood(make_synth(scale=2.0), 30.0, 3000.0, + n_grid=64, angle_marg="bogus") + + +def test_wrapper_exact_scheme_end_to_end(): + """The wrapper's exact path produces the brute-force marginal through the + same public interface production uses (value/value_and_grad/batched).""" + data = make_synth(scale=6.0) + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, n_grid=64, + interp=INTERP, guess_snr=3.0, + angle_marg="auto") + assert like.angle_marg_scheme == "exact" + ref = brute_marginal(data, like.x_grid, like.log_w_grid, 96, 48) + got = np.asarray(like.log_likelihood(RA, DEC, INCL)) + assert np.abs(got - ref).max() < 1e-10 + v, g = like.value_and_grad(np.array([RA[0], DEC[0], INCL[0]])) + assert abs(v - ref[0]) < 1e-10 + assert np.all(np.isfinite(g)) + + +# --------------------------------------------------------------------------- +# 10. the driver wiring (AST over the source: the defect this guards lives at +# the call site, where a helper-level assertion cannot see it) +# --------------------------------------------------------------------------- + +def _driver_source(): + here = os.path.dirname(os.path.abspath(__file__)) + path = os.path.join(here, "..", "..", "bin", + "integrate_likelihood_extrinsic_jax") + with open(path) as f: + return f.read() + + +def test_driver_flag_exists_with_grid_default(): + src = _driver_source() + tree = ast.parse(src) + found = None + for node in ast.walk(tree): + if (isinstance(node, ast.Call) + and getattr(node.func, "attr", "") == "add_option" + and node.args + and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "--angle-marg-scheme"): + kw = {k.arg: k.value for k in node.keywords} + found = kw + assert found is not None, "--angle-marg-scheme not registered" + # Changed 2026-09-02: the default moved 'grid' -> 'exact'. The flag must + # now name the SINGLE definition rather than re-type any literal -- a + # Constant here would be a second copy of the default, which is what bit + # the previous default move on this path (interp linear -> sinc). + assert isinstance(found.get("default"), ast.Name), \ + "--angle-marg-scheme default must be ANGLE_MARG_DEFAULT, not a literal" + assert found["default"].id == "ANGLE_MARG_DEFAULT" + assert AM.ANGLE_MARG_DEFAULT == "exact", \ + "default changed again: update the reproduce recipe and this test" + assert AM.ANGLE_MARG_LEGACY == "grid", \ + "the legacy spelling must keep reproducing pre-2026-09-02 runs" + + +def test_driver_passes_scheme_to_wrapper_and_reports_it(): + """External review, item 3: an earlier version of this guard only checked + that SOME angle_marg= keyword is passed, so the inert-flag mutant + ``angle_marg="grid"`` (flag parsed, help present, print present, value + ignored) passed the whole suite -- exactly this repo's documented + silent-no-op pattern. The guard now pins the keyword's VALUE node: it + must be the local variable ``angle_marg`` (which test_driver_flag_exists_with_grid_default + ties to the option), not a constant.""" + src = _driver_source() + tree = ast.parse(src) + passed = False + for node in ast.walk(tree): + if (isinstance(node, ast.Call) + and getattr(node.func, "id", "") == "JAXDistPhiPsiMargLikelihood"): + for k in node.keywords: + if k.arg == "angle_marg": + assert isinstance(k.value, ast.Name) \ + and k.value.id == "angle_marg", \ + "angle_marg= is passed a %r, not the angle_marg " \ + "variable: the flag would be silently inert" \ + % (ast.dump(k.value),) + passed = True + assert passed, "driver builds JAXDistPhiPsiMargLikelihood without angle_marg=" + # and the variable itself must be read from the option, not re-hardcoded + assert ('angle_marg = getattr(opts, "angle_marg_scheme", ' + 'ANGLE_MARG_DEFAULT)') in src + assert "angle-marg scheme:" in src, \ + "driver must print the RESOLVED scheme (silently-inert-flag history)" + # the print uses the wrapper's resolved attribute, not the raw option + assert "angle_marg_scheme" in src and "angle_marg_info" in src + + +# --------------------------------------------------------------------------- +# 11. external-review defect 1: ALL maxima must be enumerated +# --------------------------------------------------------------------------- + +def _kernel_truth(a, c1, c2, n=400001): + u = np.linspace(0, 2 * np.pi, n) + f = a + (c1 * np.exp(1j * u)).real + (c2 * np.exp(2j * u)).real + fm = f.max() + return fm + np.log(_trapz(np.exp(f - fm), u) / (2 * np.pi)) + + +def test_laplace_kernel_first_harmonic_cancellation(): + """The review's counterexample: c1 = 0, c2 = -d (delta = pi), d > 0.5. + f = -d cos(2u): the extrema of the FIRST harmonic are u = 0, pi -- both + MINIMA of f -- so the historical two-seed Newton rejected everything and + returned -inf for a finite integral. The maxima are at u = pi/2, 3pi/2 + and must both be found (finding only one loses ln 2).""" + for dd in (0.7, 5.0, 50.0, 500.0): + val = float(AM._laplace_psi_lnI(jnp.asarray(0.0), + jnp.asarray(0.0 + 0.0j), + jnp.asarray(-dd + 0.0j))) + assert np.isfinite(val), "d=%g: kernel returned %r" % (dd, val) + truth = _kernel_truth(0.0, 0.0, -dd) + # Laplace error O(1/d) (measured: 0.16 at d=0.7 down to 2.5e-4 at + # d=500); ln 2 = 0.69 would signal a missed second maximum + assert abs(val - truth) < 0.5 / dd + 0.05, \ + "d=%g: err %g" % (dd, val - truth) + # gradient is finite and FD-exact at the degenerate b=0 point + p0 = jnp.asarray([0.0, 0.0, 0.0, -5.0, 0.0]) + g = np.asarray(jax.grad(_kernel)(p0)) + assert np.all(np.isfinite(g)) + h = 1e-6 + for i in range(5): + fd = (float(_kernel(p0.at[i].add(h))) + - float(_kernel(p0.at[i].add(-h)))) / (2 * h) + assert abs(fd - g[i]) < 1e-6 * max(1.0, abs(fd)) + + +def test_laplace_kernel_randomized_sweep(): + """Randomized (b, d, beta, delta) sweep against brute-force quadrature, + log-uniform in d and in b/d INCLUDING b << d -- the failure region a + hand-picked example set misses (review's explicit request) -- and with + NO low-amplitude filter: an earlier revision skipped b + 2d < 0.6, + which is exactly where review 2 then found a 0.5-nat window with a + sign-inverted gradient. Below the blend band the fixed-N u-quadrature + is machine-exact; above it the O(1/A) Laplace law applies.""" + rng = np.random.default_rng(42) + for _ in range(60): + dd = 10 ** rng.uniform(-0.5, 2.5) + b = dd * 10 ** rng.uniform(-3, 1.5) + beta = rng.uniform(0, 2 * np.pi) + delta = rng.uniform(0, 2 * np.pi) + a = rng.uniform(-1, 1) + c1 = b * np.exp(-1j * beta) + c2 = dd * np.exp(-1j * delta) + val = float(AM._laplace_psi_lnI(jnp.asarray(a), jnp.asarray(c1), + jnp.asarray(c2))) + truth = _kernel_truth(a, c1, c2, n=200001) + assert np.isfinite(val) + if b + 2 * dd < AM._LAPLACE_BLEND_LO: + tol = 1e-10 # pure fixed-N quadrature + else: + tol = 4.0 / (b + dd) + 1e-3 # Laplace O(1/A) law + assert abs(val - truth) < tol, \ + "b=%g d=%g beta=%g delta=%g: err %g (tol %g)" % ( + b, dd, beta, delta, val - truth, tol) + + +def test_laplace_kernel_branch_window(): + """Reviews 2 and 3 (P2-local), pinned WITHOUT filtering the window out. + + History: a hard series/Laplace switch at b + 2d = 0.5 gave 0.27-0.53 + nats and a SIGN-INVERTED |c2| gradient just above the cut (review 2); a + Bessel-series fix moved the handover to ~16 but relied on a GLOBAL + subdominance argument, and review 3 exhibited a locally-dominant bin at + t = 18.9 with +0.251 nats. The kernel now integrates u by fixed-N + quadrature below the handover -- machine-exact for every t <= 220 + including that counterexample -- hands over to enumerated-maxima + Laplace across a C^1 blend [220, 300], and its worst LOCAL error + anywhere is bounded: measured 4e-3 on random draws above the band and + <= ~0.45 nats at the measure-tiny degenerate-curvature alignment + (b = 4d, phases aligned; was ~5 nats before the gated quartic-width + correction), decaying with t. + """ + # review 3's exact counterexample: must be machine-exact now + b, dd, beta, delta = 11.8866, 3.5163, -1.8900, -0.6497 + c1 = b * np.exp(-1j * beta) + c2 = dd * np.exp(-1j * delta) + val = float(AM._laplace_psi_lnI(jnp.asarray(0.0), jnp.asarray(c1), + jnp.asarray(c2))) + assert abs(val - _kernel_truth(0.0, c1, c2)) < 1e-12 + + rng = np.random.default_rng(5) + val_tol = {0.5001: 1e-12, 2.0: 1e-12, 19.0: 1e-12, 120.0: 1e-12, + 220.0: 1e-11, 260.0: 0.02, 300.0: 0.02, 500.0: 0.03} + for t, vtol in val_tol.items(): + for _ in range(3): + frac = rng.uniform(0.05, 0.95) + bb = t * frac + d2 = t * (1 - frac) / 2 + c1 = bb * np.exp(-1j * rng.uniform(0, 2 * np.pi)) + c2 = d2 * np.exp(-1j * rng.uniform(0, 2 * np.pi)) + val = float(AM._laplace_psi_lnI(jnp.asarray(0.2), + jnp.asarray(c1), + jnp.asarray(c2))) + truth = _kernel_truth(0.2, c1, c2, n=200001) + assert abs(val - truth) < vtol, \ + "t=%g: val err %g (tol %g)" % (t, val - truth, vtol) + # |c2|-direction gradient: machine-exact below the band, + # bounded and sign-correct wherever resolved above it + e2 = c2 / max(abs(c2), 1e-300) + g_ad = float(jax.grad( + lambda dv: AM._laplace_psi_lnI(jnp.asarray(0.2), + jnp.asarray(c1), + dv * e2))(jnp.asarray(d2))) + h = 1e-5 + g_tr = (_kernel_truth(0.2, c1, (d2 + h) * e2, n=200001) + - _kernel_truth(0.2, c1, (d2 - h) * e2, n=200001)) / (2 * h) + gtol = 1e-7 if t <= AM._LAPLACE_BLEND_LO else 0.1 + 0.1 * abs(g_tr) + assert abs(g_ad - g_tr) < gtol, \ + "t=%g: grad AD %+g vs truth %+g (tol %g)" % (t, g_ad, g_tr, + gtol) + if abs(g_tr) > 0.5: + assert np.sign(g_ad) == np.sign(g_tr), \ + "t=%g: gradient SIGN inverted (AD %+g, truth %+g)" % ( + t, g_ad, g_tr) + # degenerate-curvature alignment (b = 4d, delta = pi): the quartic-flat + # global maximum where the floored-Gaussian Laplace was ~5 nats high; + # the gated quartic-width correction bounds it (measured <= 0.45) + for t in (320.0, 600.0): + for eps in (-0.05, 0.0, 0.05): + d2 = t / 6.0 + bb = 4 * d2 * (1 + eps) + val = float(AM._laplace_psi_lnI(jnp.asarray(0.0), + jnp.asarray(bb + 0.0j), + jnp.asarray(-d2 + 0.0j))) + truth = _kernel_truth(0.0, bb + 0.0j, -d2 + 0.0j, n=400001) + assert abs(val - truth) < 0.6, \ + "degenerate t=%g eps=%g: err %g" % (t, eps, val - truth) + # C^1 blend: no value jumps across the band (a hard switch stepped by + # ~0.5 nats; measured max step 3.2e-5 at this resolution) + prev = None + for t in np.linspace(AM._LAPLACE_BLEND_LO - 5.0, + AM._LAPLACE_BLEND_HI + 5.0, 30): + c1 = 0.6 * t * np.exp(-1j * 1.1) + c2 = 0.2 * t * np.exp(-1j * 2.3) + v = float(AM._laplace_psi_lnI(jnp.asarray(0.0), jnp.asarray(c1), + jnp.asarray(c2))) \ + - _kernel_truth(0.0, c1, c2, n=200001) + if prev is not None: + assert abs(v - prev) < 5e-3 + prev = v + + +# --------------------------------------------------------------------------- +# 12. external-review defect 2: sizing must come from the data, not guess_snr +# --------------------------------------------------------------------------- + +def test_estimate_angle_amplitude_tracks_the_data(): + from RIFT.likelihood.jax_ile.core import make_distance_grid + quiet = make_synth(scale=2.0) + xg, _ = make_distance_grid(30.0, 3000.0, 64, distMpcRef=quiet.distMpcRef) + a_quiet = AM.estimate_angle_amplitude(quiet, xg) + # quiet target: bound well below the crossover (UNfloored -- the selector + # needs the raw value, or every quiet target would select laplace; the + # wrapper floors the SIZING separately) + assert 0.0 <= a_quiet < AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + loud = make_synth(scale=2.0, kappa_boost=200.0) + a_loud = AM.estimate_angle_amplitude(loud, xg) + louder = make_synth(scale=2.0, kappa_boost=400.0) + a_louder = AM.estimate_angle_amplitude(louder, xg) + assert a_loud > 20 * AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + assert a_louder > 1.5 * a_loud # monotone + # The estimator must reproduce an INDEPENDENTLY-computed 512x128 + # reference to < 2% (band-limited content, default grid ~6x Nyquist). + # The reference deliberately does NOT call estimate_angle_amplitude with + # a finer grid: a mutant gutting the estimator's internals would gut + # such a reference identically (common-mode blindness -- an earlier + # revision of this pin let exactly that mutant survive). White-box + # coupling: the sky draws replicate the estimator's documented + # rng(seed=0) construction. + rng = np.random.default_rng(0) + n_sky = AM.ANGLE_AMP_SKY_POINTS + ra_s = rng.uniform(0.0, 2 * np.pi, n_sky) + dec_s = np.arcsin(rng.uniform(-1.0, 1.0, n_sky)) + incl_s = np.arccos(rng.uniform(-1.0, 1.0, n_sky)) + g_ra, g_dec = np.meshgrid(np.linspace(0, 2 * np.pi, 6, endpoint=False), + np.array([-1.0, -0.35, 0.35, 1.0]), + indexing="ij") + for i0_ in (0.0, np.pi): + ra_s = np.concatenate([ra_s, g_ra.ravel()]) + dec_s = np.concatenate([dec_s, g_dec.ravel()]) + incl_s = np.concatenate([incl_s, np.full(g_ra.size, i0_)]) + C_A, C_B, _meta = AM.angle_coefficient_tables(loud, ra_s, dec_s, incl_s) + C_A, C_B = np.asarray(C_A), np.asarray(C_B) + PH, UU = np.meshgrid(np.linspace(0, 2 * np.pi, 512, endpoint=False), + np.linspace(0, 2 * np.pi, 128, endpoint=False), + indexing="ij") + phis, us = PH.ravel(), UU.ravel() + + def _mat(C): + kp = np.arange(C.shape[0]) + ks = np.arange(-(C.shape[1] - 1) // 2, (C.shape[1] - 1) // 2 + 1) + E = np.exp(1j * (phis[:, None, None] * kp[None, :, None] + + us[:, None, None] * ks[None, None, :])) + w = np.ones(C.shape[0]) + w[1:] = 2.0 + return (E * w[None, :, None]).reshape(len(phis), -1) + + E_A, E_B = _mat(C_A), _mat(C_B) + xv = np.asarray(xg) + x_min, x_max = float(xv.min()), float(xv.max()) + a_ref = 0.0 + for j in range(len(ra_s)): + A_g = (E_A @ C_A[:, :, j].reshape(-1, C_A.shape[-1])).real + B_g = np.maximum((E_B @ C_B[:, :, j].reshape(-1, C_B.shape[-1])).real, + 0.0) + x_hat = np.clip(A_g / np.maximum(B_g, 1e-300), x_min, x_max) + a_ref = max(a_ref, float((x_hat * A_g + - 0.5 * np.square(x_hat) * B_g).max())) + a_ref *= AM.ANGLE_AMP_MARGIN + assert a_loud > 0.98 * a_ref + assert a_loud <= a_ref * (1 + 1e-9) # grid max cannot EXCEED the max + + +def test_amp_sizing_is_required(): + """No default: a silently-undersized dense grid is the module's own + defect class (review: 'a number that is too small, with nothing that + notices').""" + data = make_synth(scale=2.0) + x_grid, log_w = _dist_grid(data) + for fn in (AM.fused_log_likelihood_distphipsimarg_exact, + AM.fused_log_likelihood_distphipsimarg_laplace): + with pytest.raises(ValueError, match="amp_sizing"): + fn(data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP) + + +def test_wrapper_sizing_survives_missing_or_low_guess_snr(): + """THE defect-2 regression. On a target whose true angular amplitude is + ~1.7e4 (kappa_boost=200; the crossover-floor grid is off by -1.04 nats, + measured), the wrapper must produce the correctly-sized answer whether + guess_snr is None or underestimated 10x -- because sizing and selection + key on the coefficient tables, not on the caller's estimate. Both + sub-cases FAILED against the pre-review implementation (which pinned + amp_sizing = 450 whenever guess_snr was absent or small).""" + data = make_synth(scale=2.0, kappa_boost=200.0) + like0 = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, n_grid=64, + interp=INTERP, guess_snr=None, + angle_marg="exact") + amp_data = like0.angle_marg_info["amp_sizing"] + assert amp_data > 20 * AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + # reference: exact scheme sized ABOVE the wrapper's own bound + ref = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + like0.x_grid, like0.log_w_grid, interp=INTERP, + amp_sizing=2 * amp_data)) + # the bite: the old floor-sized grid is measurably wrong here + floor = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + like0.x_grid, like0.log_w_grid, interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + assert np.abs(floor - ref).max() > 0.1, \ + "the regression no longer bites; the target needs a larger boost" + for guess in (None, np.sqrt(2 * amp_data) / 10.0): # missing, 10x low + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, n_grid=64, + interp=INTERP, guess_snr=guess, + angle_marg="exact") + got = np.asarray(like.log_likelihood(RA, DEC, INCL)) + assert np.abs(got - ref).max() < 1e-6, \ + "guess_snr=%r: wrapper answer off by %g" % ( + guess, np.abs(got - ref).max()) + + +# --------------------------------------------------------------------------- +# 13. review 3: higher modes in the FINAL marginal, and the runtime fail-safe +# --------------------------------------------------------------------------- + +MODES_M4 = ((2, 2), (2, -2), (3, 3), (3, -3), (4, 4), (4, -4)) + + +def test_higher_mode_marginal_vs_bruteforce(): + """Review 3, P1: 'test the FINAL MARGINAL for higher-mode inputs, not + just the estimator's reconstruction' -- the earlier coverage stopped at + the reconstruction step, which is exactly why the amplitude-only dense + sizing survived. m_max = 4 end to end at cheap scale: sample grid + (24, 8), coefficient tables with KPA = 5 / KPB = 9, and the full + marginal against an independent brute force. Measured 1.8e-15.""" + data = make_synth(scale=4.0, modes=MODES_M4) + assert AM._data_m_max(data) == 4 + assert AM.angle_sample_grid_sizes(4) == (24, 8) + x_grid, log_w = _dist_grid(data) + ref = brute_marginal(data, x_grid, log_w, 96, 48) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE)) + assert np.abs(ex - ref).max() < 1e-10 + # and the harmonic-content invariant extends to n = 8, machine-zero above + n = 64 + phis = np.linspace(0, 2 * np.pi, n, endpoint=False) + f = _lnL_t_fixed_time(data, phis, np.full(n, 0.6), 0.8, 10) + C = np.abs(np.fft.rfft(f) / n) + assert C[:9].max() > 0 + assert C[9:].max() < 1e-12 * C.max() + + +def test_higher_mode_dense_sizing_self_convergence(): + """Review 3, P1 at the sizing level: on an m_max = 4 target loud enough + that the dense grid actually works (amp ~ few hundred), the marginal + must be converged in the dense sizes -- quadrupling amp_sizing (which + doubles every dense axis) must not move it. Under the amplitude-only + rule this bites at the ~0.03-nat level (review 3 measured +-0.037 for a + pure order-8 term at A = 450); with the m_max-scaled rule it is + ~machine (measured 0.0).""" + data = make_synth(scale=2.0, modes=MODES_M4, kappa_boost=14.0) + x_grid, log_w = _dist_grid(data) + amp = AM.estimate_angle_amplitude(data, x_grid) + assert amp > 50.0 # genuinely loud, or the test proves nothing + args = (data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w) + a_sz = max(amp, AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + v1 = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=a_sz)) + v4 = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=4 * a_sz)) + assert np.abs(v1 - v4).max() < 1e-6 + + +def test_runtime_amp_failsafe_warns_and_records_without_excising(capfd): + """The undersizing guard must (a) warn, (b) RECORD on the host so the driver + can label the artifact, and (c) leave the value FINITE. + + (c) is the load-bearing one and is easy to get wrong in the tempting + direction. An earlier version returned NaN to "fail closed". That was + worse than the warning: every consumer FILTERS non-finite lnL -- flowMC and + MALA reject such proposals, the SMC path drops them, write_samples discards + non-finite rows -- so a NaN over a hot sky region the estimator missed does + not halt anything, it EXCISES that region and publishes a clean-looking + posterior over the rest. This test therefore asserts finiteness, and would + fail if anyone reintroduces the poison. + + The previous version of this test asserted only that a warning printed, so + replacing the poison return with 0 would have passed it unchanged. + """ + data = make_synth(scale=2.0, kappa_boost=200.0) + x_grid, log_w = _dist_grid(data) + args = (data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w) + + AM.reset_amp_failsafe() + assert AM.amp_failsafe_state()["tripped"] is False + + # deliberately undersized + v = AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + jax.block_until_ready(v) + out = capfd.readouterr() + + assert "WARNING anglemarg/exact" in out.out + out.err, "guard must warn" + st = AM.amp_failsafe_state() + assert st["tripped"] is True, ( + "the guard must RECORD on the host, or the driver cannot label the " + "artifact and the condition dies with the log line") + assert st["worst_amp"] > 2.0 * AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + assert np.all(np.isfinite(np.asarray(v))), ( + "undersizing must NOT be signalled by a non-finite value: downstream " + "filters would excise exactly the region the estimator missed and " + "publish a clean-looking posterior over what remains") + + # correctly sized: silent, and nothing recorded + AM.reset_amp_failsafe() + amp = AM.estimate_angle_amplitude(data, x_grid) + v2 = AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=amp) + jax.block_until_ready(v2) + out = capfd.readouterr() + assert "WARNING anglemarg" not in out.out + out.err + assert AM.amp_failsafe_state()["tripped"] is False + assert np.all(np.isfinite(np.asarray(v2))) + + +def test_driver_labels_a_suspect_angle_grid_in_provenance(): + """End-to-end-ish: the driver must LABEL a tripped event in the artifact's + own provenance line, and must NOT abort or drop the event. + + Structural rather than a full run: asserts the driver imports the guard, + resets it per event (batch runs analyze several events in one process, so + event 0 must not label event 1), and appends SUSPECT-ANGLE-GRID to the + provenance string that write_samples embeds in the exported file.""" + import re as _re + import pathlib as _pathlib + drv = (_pathlib.Path(__file__).resolve().parents[2] / "bin" + / "integrate_likelihood_extrinsic_jax") + src = drv.read_text() + assert "from RIFT.likelihood.jax_ile import anglemarg as _anglemarg" in src + assert "_anglemarg.reset_amp_failsafe()" in src, "must reset per event" + assert "_anglemarg.amp_failsafe_state()" in src + assert "SUSPECT-ANGLE-GRID" in src, "the artifact must carry the label" + # The EVIDENCE row must be labelled too, not only the sample file: + # write_samples() early-returns without --save-samples, so a run with export + # disabled would otherwise publish a numeric .dat indistinguishable from a + # clean integration. + assert "def angle_grid_suspect_note()" in src + wd = src[src.index("def write_dat("):] + wd = wd[:wd.index("\ndef ")] + assert "angle_grid_suspect_note()" in wd, ( + "write_dat must label the evidence artifact independently of sample export") + # and the warning must fire per event, not only on the export path + ao = src[src.index("def analyze_one("):] + assert "_ev_note = angle_grid_suspect_note()" in ao, ( + "analyze_one must report once per event regardless of export settings") + # and must not sit behind a bare except that degrades a tripped run to clean + assert '_st = {"tripped": False}' not in src, ( + "a swallowed exception would silently report a tripped run as clean") + + + + +def test_failsafe_callback_is_cond_guarded_and_reads_are_barriered(): + """Throughput and reliability constraints on the undersizing record. + + An UNCONDITIONAL jax.debug.callback fires once per likelihood evaluation -- + once per MALA/flowMC proposal, per chain -- transferring to the host and + destroying accelerator throughput even when undersizing never happens. It + must sit inside lax.cond so the ordinary path pays nothing. + + And because debug-callback effects may be dropped, duplicated, reordered, or + land AFTER the result is ready, every read/reset of the host record must + barrier first -- otherwise a caller reads clean while a tripped callback is + in flight, or resets before the previous event's callback arrives. + """ + import inspect as _inspect + from RIFT.likelihood.jax_ile import anglemarg as _AMmod + + src = _inspect.getsource(_AMmod._runtime_amp_failsafe) + i_cond = src.find("lax.cond") + i_cb = src.find("debug.callback") + assert i_cond != -1 and i_cb != -1 + assert i_cond < i_cb, ( + "jax.debug.callback must be INSIDE lax.cond; an unconditional callback " + "fires on every likelihood evaluation") + + for fn in (_AMmod.amp_failsafe_state, _AMmod.reset_amp_failsafe): + assert "effects_barrier" in _inspect.getsource(fn), ( + "%s must barrier queued callbacks before touching the record" % fn.__name__) + + # and it still works end to end + _AMmod.reset_amp_failsafe() + assert _AMmod.amp_failsafe_state()["tripped"] is False diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_laplace.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_laplace.py new file mode 100644 index 000000000..93b4d173b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_laplace.py @@ -0,0 +1,421 @@ +""" +Gate for the psi-marginal distance-node placement that lets the 'laplace' +angle-marg scheme honour JAX_ILE_DISTMARG_GH +(``anglemarg._gh_psi_node_offsets`` and the ``_use_gh`` branch of +``fused_log_likelihood_distphipsimarg_laplace``). + +WHAT IS PINNED, AND WHY +----------------------- +core._distmarg_gh_logL places frozen nodes around the distance peak of a FIXED +psi. On the laplace path psi is already integrated out, so the nodes must +bracket the psi-MARGINAL distance integrand -- a mixture over u = 2 psi of +Gaussians centred at x*(u) = A(u)/B(u). The shipped rule centres on the u that +maximises A(u) and takes the width from the closed-form envelope +R_lo = B0 - |B1| - |B2| <= min_u B, with a +-12 sigma half-span. + +Its whole validity rests on ONE structural identity: for (2,+-2) mode content +the spin-2 response makes A0 and B1 vanish identically, so R_lo IS min_u B and +the closed-form centre sits inside the weight-carrying span. That identity does +not survive richer mode content, so: + + * the identity itself is pinned here, WITH a positive control on m_max = 3 + data proving the assertion can fail (a pass-through that cannot fail is not + coverage); + * the m_max gate is pinned, with a positive control that the same data is + accepted with the adaptive quadrature off; + * the placement constants are pinned, and the sufficiency of 12 sigma is + demonstrated by MUTATION -- shrinking the half-span changes the answer, + growing it does not; + * agreement is checked against BOTH independent references (exact + the same + adaptive quadrature, and laplace + a converged uniform grid), on a fixture + loud enough that the adaptive bracket is the right tool; + * gradients through the new branch are finite (the reason the placement is + under stop_gradient at all). + +Ladder-2 measurements behind the constants live in +RIFT/likelihood/jax_ile/DESIGN_gh_laplace.md, not here. +""" + +import numpy as np +import pytest + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile import core as core_mod +from RIFT.likelihood.jax_ile.core import make_distance_grid + +from test_angle_marg_exact import make_synth, RA, DEC, INCL, INTERP + +AMP = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + + +# --------------------------------------------------------------------------- +# fixtures: a LOUD target, so the distance peak is narrow and an adaptive +# bracket is the right tool (at low amplitude the +-N sigma window of the +# fixed-psi rule does not cover the support either -- a property of +# core._distmarg_gh_logL, deliberately not re-litigated here) +# --------------------------------------------------------------------------- + +def loud_data(modes=((2, 2), (2, -2))): + return make_synth(scale=6.0, kappa_boost=12.0, modes=modes) + + +def loud_grid(data, n=256): + return make_distance_grid(200.0, 4000.0, n, distMpcRef=data.distMpcRef) + + +def _fields(data, nphi=24): + """A0, A1, B0, B1, B2 on a phi grid -- the SAME expressions the kernel uses.""" + C_A, C_B, meta = AM.angle_coefficient_tables( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), INTERP) + C_A = np.asarray(C_A); C_B = np.asarray(C_B) + m_max = int(meta["m_max"]) + wA = np.asarray(AM._kp_weights(m_max + 1)) + wB = np.asarray(AM._kp_weights(2 * m_max + 1)) + phi = np.linspace(0.0, 2 * np.pi, nphi, endpoint=False) + EA = np.exp(1j * phi[:, None] * np.arange(m_max + 1)[None, :]) * wA[None, :] + EB = np.exp(1j * phi[:, None] * np.arange(2 * m_max + 1)[None, :]) * wB[None, :] + MA = lambda k: np.einsum("ck,kst->cst", EA, C_A[:, k]) + MB = lambda k: np.einsum("ck,kst->cst", EB, C_B[:, k]) + kA = (C_A.shape[1] - 1) // 2 + kB = (C_B.shape[1] - 1) // 2 + return dict( + A0=MA(kA).real, A1=MA(kA + 1) + np.conj(MA(kA - 1)), + B0=MB(kB).real, B1=MB(kB + 1) + np.conj(MB(kB - 1)), + B2=MB(kB + 2) + np.conj(MB(kB - 2)), m_max=m_max) + + +# --------------------------------------------------------------------------- +# 1. the structural identity the placement rests on +# --------------------------------------------------------------------------- + +def _identity_residuals(f): + """(|A0|/|A1|, |B1|/|B0|) -- zero iff A(u) is a pure first harmonic and + B(u) a constant plus a pure second harmonic.""" + return (np.abs(f["A0"]).max() / np.abs(f["A1"]).max(), + np.abs(f["B1"]).max() / np.abs(f["B0"]).max()) + + +@pytest.mark.parametrize("modes,m_max", [ + (((2, 2), (2, -2)), 2), + (((2, 2), (2, -2), (3, 3), (3, -3)), 3), + (((2, 2), (2, -2), (2, 1), (2, -1), (3, 3), (3, -3), (4, 4), (4, -4)), 4), +]) +def test_a0_and_b1_vanish(modes, m_max): + """A(u) is a PURE first harmonic and B(u) a constant plus a PURE second + harmonic, so R_lo = B0 - |B1| - |B2| is min_u B EXACTLY rather than a + bound. This is the identity the +-12 sigma rule is derived from; if a + mode-convention change breaks it, the rule is void. + + It is not a (2,+-2) accident: the psi dependence enters only through the + spin-2 antenna response F(psi) ~ e^{-2 i psi}, which puts the kappa term at + exactly one u-harmonic (u = 2 psi) and the rho^2 term at exactly harmonics + 0 and 2, whatever the mode content. Checked here through m_max = 4, and + on real IMRPhenomXHM l_max = 3/4 data in the branch's DESIGN note. (The + SHIP gate is still m_max <= 2, deliberately: the identity fixes the WIDTH, + while the half-span constant was measured on (2,+-2) fixtures.)""" + f = _fields(loud_data(modes=modes)) + assert f["m_max"] == m_max + a0, b1 = _identity_residuals(f) + assert a0 < 1e-12 and b1 < 1e-12, (a0, b1) + # ... and the second harmonic is REAL content, not another zero: a bound + # that is tight only because every harmonic vanished would prove nothing. + assert np.median(np.abs(f["B2"]) / f["B0"]) > 1e-3 + # R_lo is then min_u B to u-grid resolution, and strictly positive + u = np.linspace(0.0, 2 * np.pi, 4096, endpoint=False) + Bu = (f["B0"][..., None] + (f["B1"][..., None] * np.exp(1j * u)).real + + (f["B2"][..., None] * np.exp(2j * u)).real) + R_lo = f["B0"] - np.abs(f["B1"]) - np.abs(f["B2"]) + assert (R_lo > 0).all() + assert np.abs(Bu.min(-1) / R_lo - 1.0).max() < 1e-4 + + +def test_identity_check_has_a_positive_control(): + """POSITIVE CONTROL for the test above. The identity holds for EVERY mode + set, so no choice of modes can make the assertion fail -- which means a + bug that returned all-zero coefficient tables, or a residual computed from + the wrong slots, would leave it passing for the wrong reason. Prove the + detector fires by PLANTING the harmonics it is supposed to reject.""" + f = _fields(loud_data()) + ok_a0, ok_b1 = _identity_residuals(f) + assert ok_a0 < 1e-12 and ok_b1 < 1e-12 + planted_a0 = dict(f, A0=f["A0"] + 1e-3 * np.abs(f["A1"])) + planted_b1 = dict(f, B1=f["B1"] + 1e-3 * f["B0"]) + assert _identity_residuals(planted_a0)[0] >= 1e-12 + assert _identity_residuals(planted_b1)[1] >= 1e-12 + # and the residual is computed from a NON-DEGENERATE table: |A1| and |B0| + # are the scales it divides by, so a zeroed table would give 0/0, not 0 + assert np.abs(f["A1"]).max() > 0 and np.abs(f["B0"]).max() > 0 + + +# --------------------------------------------------------------------------- +# 2. the node-offset rule +# --------------------------------------------------------------------------- + +def test_gh_psi_node_offsets(): + assert AM._GH_PSI_HALF_SIGMA == 22.0 + assert AM._GH_PSI_MIN_NODES == 49 + assert AM._GH_PSI_M_MAX == 2 + for n_req in (8, 16, 33, 64, 129): + z, zp, zn, n = AM._gh_psi_node_offsets(n_req) + assert len(z) == len(zp) == len(zn) == n + assert z[0] == -AM._GH_PSI_HALF_SIGMA and z[-1] == AM._GH_PSI_HALF_SIGMA + assert n >= AM._GH_PSI_MIN_NODES + # neighbour arrays are z with the INDEX clamped: this is what makes + # 0.5*(x[k+1]-x[k-1]) reproduce core._distmarg_gh_logL's trapezoid + # weights (0.5*dx at both ends) without cross-block communication + assert zp[0] == z[0] and zn[-1] == z[-1] + assert np.allclose(zp[1:], z[:-1]) and np.allclose(zn[:-1], z[1:]) + # node DENSITY is at least what the caller asked for at +-7 sigma + h_req = 14.0 / max(n_req - 1, 1) + assert (2 * AM._GH_PSI_HALF_SIGMA) / (n - 1) <= h_req + 1e-12 + # the floor binds for small requests, and gives <= 1 sigma spacing + assert AM._gh_psi_node_offsets(4)[3] == AM._GH_PSI_MIN_NODES + z = AM._gh_psi_node_offsets(4)[0] + assert (z[1] - z[0]) <= 1.0 + + +def test_trapezoid_weights_match_the_fixed_psi_rule(): + """The index-clamped neighbour form must reproduce core._distmarg_gh_logL's + diff()/concatenate() weights bit for bit on unclipped nodes.""" + z, zp, zn, n = AM._gh_psi_node_offsets(33) + centre, sigma = 3.0, 0.25 + x = centre + sigma * z + w_new = 0.5 * (centre + sigma * zn - (centre + sigma * zp)) + dx = np.diff(x) + w_old = np.concatenate([0.5 * dx[:1], 0.5 * (dx[1:] + dx[:-1]), 0.5 * dx[-1:]]) + assert np.array_equal(w_new, w_old) + + +# --------------------------------------------------------------------------- +# 3. the mode-content gate (with its positive control) +# --------------------------------------------------------------------------- + +def test_gh_laplace_gate_on_mode_content(monkeypatch): + data3 = loud_data(modes=((2, 2), (2, -2), (3, 3), (3, -3))) + x3, lw3 = loud_grid(data3) + call3 = lambda: AM.fused_log_likelihood_distphipsimarg_laplace( + data3, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x3, lw3, interp=INTERP, amp_sizing=AMP) + # POSITIVE CONTROL: with the adaptive quadrature OFF the very same data + # runs. Without this the raise below could be any other failure. + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 0) + assert np.isfinite(np.asarray(call3())).all() + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 33) + with pytest.raises(ValueError, match="m_max"): + call3() + # ... and the covered mode content is ACCEPTED, so the gate is a gate and + # not a blanket refusal + data2 = loud_data() + x2, lw2 = loud_grid(data2) + got = np.asarray(AM.fused_log_likelihood_distphipsimarg_laplace( + data2, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x2, lw2, interp=INTERP, amp_sizing=AMP)) + assert np.isfinite(got).all() + + +# --------------------------------------------------------------------------- +# 4. agreement with two INDEPENDENT references +# --------------------------------------------------------------------------- + +def _lap_gh(data, x, lw, n_gh, monkeypatch, **over): + for k, v in over.items(): + monkeypatch.setattr(AM, k, v) + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", n_gh) + return np.asarray(AM.fused_log_likelihood_distphipsimarg_laplace( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x, lw, interp=INTERP, amp_sizing=AMP)) + + +def test_gh_laplace_matches_exact_gh(monkeypatch): + """Same adaptive distance treatment on both sides: the residual is the + psi-Laplace error alone, which the crossover constant already bounds.""" + data = loud_data() + x, lw = loud_grid(data) + lap = _lap_gh(data, x, lw, 65, monkeypatch) + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 65) + ex = np.asarray(AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x, lw, interp=INTERP, amp_sizing=AMP)) + assert np.abs(lap - ex).max() < 5e-3 + + +def test_gh_laplace_matches_converged_uniform_grid(monkeypatch): + """Same angle treatment on both sides, distance treatment independent: + a converged uniform grid must reproduce the adaptive answer.""" + data = loud_data() + x, lw = loud_grid(data) + lap_gh = _lap_gh(data, x, lw, 65, monkeypatch) + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 0) + xu, lwu = loud_grid(data, n=16384) + lap_uni = np.asarray(AM.fused_log_likelihood_distphipsimarg_laplace( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + xu, lwu, interp=INTERP, amp_sizing=AMP)) + assert np.abs(lap_gh - lap_uni).max() < 5e-3 + + +def test_gh_laplace_converged_in_node_count(monkeypatch): + data = loud_data() + x, lw = loud_grid(data) + ref = _lap_gh(data, x, lw, 129, monkeypatch) + for n_gh in (16, 33, 65): + got = _lap_gh(data, x, lw, n_gh, monkeypatch) + assert np.abs(got - ref).max() < 1e-4, n_gh + + +# --------------------------------------------------------------------------- +# 5. the half-span, by MUTATION: 12 sigma is enough and is not gratuitous +# --------------------------------------------------------------------------- + +def test_half_span_is_sufficient_and_not_gratuitous(monkeypatch): + """A frozen bracket is only correct if it CONTAINS the psi-marginal peak. + Testing that by construction is circular, so test it by intervention: + doubling the half-span must not move the answer (12 sigma already + contains everything), while collapsing it must (proving the answer really + does depend on the bracket, i.e. that this knob is live at all).""" + data = loud_data() + x, lw = loud_grid(data) + base = _lap_gh(data, x, lw, 65, monkeypatch) + wide = _lap_gh(data, x, lw, 65, monkeypatch, _GH_PSI_HALF_SIGMA=44.0) + assert np.abs(base - wide).max() < 1e-4, "22 sigma does not contain the peak" + narrow = _lap_gh(data, x, lw, 65, monkeypatch, _GH_PSI_HALF_SIGMA=0.05) + assert np.abs(base - narrow).max() > 1e-2, ( + "collapsing the bracket did not change the answer -- the half-span is " + "inert and this test proves nothing") + + +def test_node_floor_is_live(monkeypatch): + """MUTATION of the other constant: with the floor removed, a tiny node + request must degrade the answer -- otherwise _GH_PSI_MIN_NODES is + decoration.""" + data = loud_data() + x, lw = loud_grid(data) + ref = _lap_gh(data, x, lw, 129, monkeypatch) + monkeypatch.setattr(AM, "_GH_PSI_MIN_NODES", 3) + n_coarse = AM._gh_psi_node_offsets(2)[3] + assert n_coarse < 10, n_coarse # the floor is what was holding it up + coarse = _lap_gh(data, x, lw, 2, monkeypatch) + assert np.abs(coarse - ref).max() > 1e-3, ( + "a %d-node bracket matched the converged answer -- the node count is " + "inert and the floor pins nothing" % n_coarse) + + +# --------------------------------------------------------------------------- +# 6. the sigma cap (unreachable on signal-carrying data; must still be live) +# --------------------------------------------------------------------------- + +def test_sigma_cap_is_inactive_at_production_support_and_live_when_it_binds( + monkeypatch): + """The cap keeps the bracket inside the physical support when R_lo -> 0 (a + bin with no response, where the exponent is flat in x and sigma would blow + up, leaving every node on one of the two rails). Pin BOTH directions: it + is far from binding at a production distance support, and it does bind -- + and change the answer -- on a narrow one.""" + data = loud_data() + f = _fields(data) + sigma = 1.0 / np.sqrt(f["B0"] - np.abs(f["B1"]) - np.abs(f["B2"])) + + def cap(d_min, d_max): + xg, _ = make_distance_grid(d_min, d_max, 64, distMpcRef=data.distMpcRef) + xg = np.asarray(xg) + return (xg.max() - xg.min()) / (2.0 * AM._GH_PSI_HALF_SIGMA) + + # production-scale support (RIFT's own --d-min 1 --d-max 10000): the widths + # this data actually uses are two orders of magnitude below the cap, so the + # cap cannot be shaping any production result. (On the ladder-2 injection + # itself the margin is ~1100x -- sigma ~ 0.02 against a cap of 22.7 -- this + # synthetic target is quieter, hence the looser factor here.) + assert sigma.max() < 1e-2 * cap(1.0, 10000.0) + # the fixture's own narrow support DOES reach it -- so the cap is live code + assert sigma.max() > cap(200.0, 4000.0) + # ... and removing it changes the answer, i.e. it is not a no-op minimum + x, lw = loud_grid(data) + base = _lap_gh(data, x, lw, 65, monkeypatch) + monkeypatch.setattr(AM, "_GH_PSI_HALF_SIGMA", 1e-9) # cap -> enormous + uncapped = _lap_gh(data, x, lw, 65, monkeypatch) + assert np.abs(uncapped - base).max() > 1e-2 + + +# --------------------------------------------------------------------------- +# 7. gradients (the reason the placement is frozen) +# --------------------------------------------------------------------------- + +def test_gh_laplace_gradients_are_finite(monkeypatch): + data = loud_data() + x, lw = loud_grid(data) + monkeypatch.setattr(core_mod, "_DISTMARG_GH_N", 33) + + def f(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, x, lw, interp=INTERP, amp_sizing=AMP).sum() + + g = jax.grad(f, argnums=(0, 1, 2))(jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL)) + assert all(np.isfinite(np.asarray(gi)).all() for gi in g) + # ... and against finite differences on the sky angles + eps = 1e-5 + for i, base in enumerate((RA, DEC, INCL)): + args = [jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL)] + args[i] = jnp.asarray(base + eps) + hi = float(f(*args)) + args[i] = jnp.asarray(base - eps) + lo = float(f(*args)) + fd = (hi - lo) / (2 * eps) + assert abs(fd - float(np.asarray(g[i]).sum())) <= 1e-4 * max(1.0, abs(fd)) + + +# --------------------------------------------------------------------------- +# 8. the closed-form psi maximiser the centre sits on +# --------------------------------------------------------------------------- + +def test_closed_form_psi_argmax_is_exact(): + """With A0 == 0 and B1 == 0, argmax_u A(u)^2/(2 B(u)) is exactly + e^{i u*} = +- conj(w)/|w|, w = B0*A1 - conj(A1)*B2. This is what the node + CENTRE sits on; centring on argmax A(u) instead (the B2 -> 0 limit of the + same formula) is loose by up to ~7000 sigma at large |B2|/B0. + + Checked against brute force over random coefficient triples, INCLUDING the + large-|B2|/B0 corner that the ladder-2 fixture never reaches -- the point + being that a fixture cannot exercise this and the closed form must be + right everywhere.""" + rng = np.random.default_rng(11) + n = 4000 + A1 = rng.normal(size=n) + 1j * rng.normal(size=n) + B0 = np.abs(rng.normal(size=n)) * 3 + 0.05 + r = rng.uniform(0.0, 0.9999, n) + B2 = B0 * r * np.exp(1j * rng.uniform(0, 2 * np.pi, n)) + + w = B0 * A1 - np.conj(A1) * B2 + ph = np.conj(w) / np.maximum(np.abs(w), 1e-300) + ph = ph * np.where((A1 * ph).real >= 0, 1.0, -1.0) + A_st = (A1 * ph).real + B_st = B0 + (B2 * ph * ph).real + assert (A_st > 0).all(), "sign branch does not put the centre at x* > 0" + assert (B_st > 0).all() + E_st = A_st ** 2 / (2 * B_st) + + u = np.linspace(0.0, 2 * np.pi, 100001, endpoint=False) + e1 = np.exp(1j * u); e2 = e1 * e1 + best = np.empty(n) + for k in range(0, n, 100): + sl = slice(k, k + 100) + A = (A1[sl, None] * e1).real + B = B0[sl, None] + (B2[sl, None] * e2).real + best[sl] = np.where(A > 0, A * A / (2 * np.maximum(B, 1e-300)), 0.0).max(-1) + # brute force must never BEAT the closed form (it can only fall short by + # the u-grid resolution) + assert ((best - E_st) / E_st).max() < 1e-12 + # POSITIVE CONTROL: the naive centring (argmax A) is measurably worse, so + # the correction is doing work rather than being a no-op rewrite + ph_naive = np.conj(A1) / np.abs(A1) + E_naive = ((A1 * ph_naive).real ** 2 + / (2 * (B0 + (B2 * ph_naive * ph_naive).real))) + assert (E_naive < E_st * (1 - 1e-6)).mean() > 0.5, ( + "argmax A already equals argmax A^2/2B on this sample; the closed form " + "is then untested") + # ... and the two coincide exactly in the B2 -> 0 limit + w0 = B0 * A1 + assert np.abs(np.conj(w0) / np.abs(w0) - np.conj(A1) / np.abs(A1)).max() < 1e-12 diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py new file mode 100644 index 000000000..398dceef8 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py @@ -0,0 +1,256 @@ +"""`auto` may route to 'laplace' under JAX_ILE_DISTMARG_GH -- but only where the +psi-marginal node placement is actually valid, MEASURED on the data. + +Before the GH-for-laplace path existed, choose_angle_marg_scheme returned 'exact' +unconditionally under GH. It no longer does. The risk this suite pins is the +opposite one: routing 'auto' to laplace where the placement is NOT valid sends the +run into the raise inside the laplace kernel. + +The predicate deliberately does NOT key on mode content alone. m_max is the +largest |m| in the mode list, so a PRECESSING l=2 system has m_max = 2 and would +pass a mode gate while breaking the aligned-spin symmetry the A0==0/B1==0 identity +has only ever been tested under. Every measurement of that identity so far is +non-precessing, so the code measures instead of assuming. +""" +import numpy as np +import pytest + +import RIFT.likelihood.jax_ile.anglemarg as AM + + +def _tables(m_max=2, a0=0.0, b1=0.0, a0_imag_at=None, b1_lower_only=False): + """Coefficient tables whose RECONSTRUCTED fields satisfy the identity, or not. + + The identity is a statement about the fields + ``A0(phi) = Re(MA(1))`` and ``B1(phi) = MB(3) + conj(MB(1))`` -- NOT about the + coefficient slices. So "clean" here means zeroing every slice that FEEDS those + two fields (``C_A[:,1]``; ``C_B[:,3]`` AND ``C_B[:,1]``), and a plant means + making one of them nonzero. An earlier version of this fixture planted into + slices while leaving the others random, which is why it could not distinguish + the two. + """ + rng = np.random.default_rng(7) + # Shapes are those angle_coefficient_tables actually returns: + # C_A (m_max+1, 3, S, npts), C_B (2*m_max+1, 5, S, npts). An earlier + # fixture used (2*m_max+1, 3) and (4*m_max+1, 5); the slice-based predicate + # never indexed the leading axis, so the invalid shape was invisible and the + # suite passed on it. + nphi_A, nphi_B = m_max + 1, 2 * m_max + 1 + ksA, ksB = 3, 5 + C_A = (rng.normal(size=(nphi_A, ksA, 2, 4)) + + 1j * rng.normal(size=(nphi_A, ksA, 2, 4))) + C_B = (rng.normal(size=(nphi_B, ksB, 2, 4)) + + 1j * rng.normal(size=(nphi_B, ksB, 2, 4))) + ks0 = (ksB - 1) // 2 + C_A[:, 1] = 0.0 # -> A0(phi) == 0 + C_A[:, 2] = 1.0 # -> a well-scaled A1 + C_B[:, ks0] = 1.0 # -> B0 + C_B[:, ks0 + 1] = 0.0 # \ + C_B[:, ks0 - 1] = 0.0 # / both feed B1(phi); zero BOTH + if a0: + C_A[:, 1] = a0 + if a0_imag_at is not None: + # the reviewer's case: a purely IMAGINARY coefficient. Re(C_A[k,1]) is + # zero, so a slice-based check sees nothing, but Re(MA(1)) is not. + C_A[a0_imag_at, 1] = 1j * 1e-3 + if b1: + C_B[:, ks0 + 1] = b1 + if b1_lower_only: + # B1(phi) also carries conj(MB(ks0-1)); planting ONLY there must be caught. + C_B[:, ks0 - 1] = 1e-3 + return C_A, C_B + + +def test_identity_holds_for_clean_tables(): + ok, info = AM.gh_laplace_supported(*_tables(m_max=2), 2) + assert ok is True + assert info["identity_A0_over_A1"] <= AM.GH_PSI_IDENTITY_TOL + assert info["identity_B1_over_B0"] <= AM.GH_PSI_IDENTITY_TOL + assert "identity holds" in info["gh_laplace_reason"] + + +@pytest.mark.parametrize("a0,b1", [(1e-3, 0.0), (0.0, 1e-3), (1e-3, 1e-3)]) +def test_identity_check_can_fail(a0, b1): + """POSITIVE CONTROL: the check must be able to return False at all. + + A predicate that cannot fail is not a check. Planted harmonics are the only + way to exercise this -- no mode set tried through m_max = 4 breaks the + identity naturally, which is precisely why it must be measured rather than + inferred from the mode list. + """ + ok, info = AM.gh_laplace_supported(*_tables(m_max=2, a0=a0, b1=b1), 2) + assert ok is False + assert "does NOT hold" in info["gh_laplace_reason"] + + +def test_mode_content_above_validated_is_refused(): + ok, info = AM.gh_laplace_supported(*_tables(m_max=2), 4) + assert ok is False + assert "m_max" in info["gh_laplace_reason"] + + +def test_auto_reaches_laplace_under_gh_when_supported(): + amp = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE * 10.0 + scheme, info = AM.choose_angle_marg_scheme(amp, gh_enabled=True, + gh_laplace_ok=True) + assert scheme == "laplace", "auto must be able to reach laplace under GH" + # control: same amplitude, same GH, predicate false -> exact + scheme_no, info_no = AM.choose_angle_marg_scheme(amp, gh_enabled=True, + gh_laplace_ok=False) + assert scheme_no == "exact" + assert "not " in info_no["reason"] + + +def test_auto_is_conservative_when_predicate_absent(): + """gh_laplace_ok=None means "caller did not measure it" -> take the safe + branch. Guessing here would route auto into the laplace kernel's raise.""" + amp = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE * 10.0 + scheme, _ = AM.choose_angle_marg_scheme(amp, gh_enabled=True) + assert scheme == "exact" + + +def test_selector_unchanged_with_gh_off(): + amp_hi = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE * 10.0 + amp_lo = AM.ANGLE_MARG_CROSSOVER_AMPLITUDE / 10.0 + assert AM.choose_angle_marg_scheme(amp_hi, gh_enabled=False)[0] == "laplace" + assert AM.choose_angle_marg_scheme(amp_lo, gh_enabled=False)[0] == "exact" + + +def test_imaginary_A0_coefficient_is_detected(): + """A purely imaginary C_A[k,1] gives Re(C_A[k,1]) == 0 but Re(MA(1)) != 0. + + The predicate must measure the RECONSTRUCTED A0(phi). A slice-based check + reports this dataset as supported; that was a real defect (external review). + """ + ok, info = AM.gh_laplace_supported(*_tables(m_max=2, a0_imag_at=1), 2) + assert ok is False, "imaginary A0 coefficient slipped past the identity check" + assert "does NOT hold" in info["gh_laplace_reason"] + + +def test_B1_planted_only_in_the_conjugate_slice_is_detected(): + """B1(phi) = MB(3) + conj(MB(1)); planting only in the LOWER slice must fail. + + A check reading C_B[:,3] alone reports this dataset as supported; that was + the second half of the same defect. + """ + ok, info = AM.gh_laplace_supported(*_tables(m_max=2, b1_lower_only=True), 2) + assert ok is False, "B1 planted in the conjugate slice slipped past the check" + assert "does NOT hold" in info["gh_laplace_reason"] + + +def test_explicit_laplace_under_gh_is_gated_too_not_only_auto(): + """The identity guard must cover an EXPLICIT laplace, not only `auto`. + + External review: with the check only inside choose_angle_marg_scheme, `auto` + was protected while `--angle-marg-scheme laplace` -- and any direct wrapper + caller -- walked past it, so an m_max == 2 dataset whose identity fails was + evaluated with a placement derived from that identity. + + It cannot live in the kernel: that runs under jit/grad where the coefficient + tables are TRACERS and the measurement raises + TracerArrayConversionError (this is how the first fix was caught). So the + enforcement point is the wrapper, on concrete tables, and this pins it there + for BOTH routes. + """ + import inspect + from RIFT.likelihood.jax_ile import wrapper as WR + src = inspect.getsource(WR.JAXDistPhiPsiMargLikelihood.__init__) + assert "gh_laplace_supported" in src, ( + "the wrapper does not measure the identity; the GH laplace placement " + "would be reachable without its premise being checked") + assert 'angle_marg in ("auto", "laplace")' in src, ( + "the identity gate is not applied to an EXPLICITLY requested laplace") + assert 'angle_marg == "laplace" and gh_ok is False' in src, ( + "an explicit laplace with a failing identity does not raise") + + +def test_kernel_enforces_the_response_model_itself_not_only_the_wrapper(): + """The PUBLIC laplace kernel must gate on the response model too. + + External review: this function is in ``anglemarg.__all__`` and is called + directly by the wrapper and by several test modules, so a wrapper-only gate + leaves a live bypass -- a direct call with ``data.feature == "rotation"`` + and ``m_max <= 2`` executes the unsupported placement while the wrapper + correctly refuses it. + + An earlier version of THIS TEST asserted the kernel contained no such check, + which entrenched the bypass rather than catching it. What must stay absent + is only the NUMERICAL A0/B1 measurement, which needs concrete tables and + raises under jit/grad; ``feature`` is a static Python attribute and is free + to check. + """ + import inspect + src = inspect.getsource(AM.fused_log_likelihood_distphipsimarg_laplace) + assert "_GH_PSI_STATIC_FEATURES" in src, ( + "the public laplace kernel does not enforce the response model; a " + "direct call with a banded response would use the placement anyway") + assert "gh_laplace_supported(" not in src, ( + "the kernel measures the identity numerically; that needs concrete " + "tables and raises under jax.grad") + + +def test_kernel_refuses_a_banded_response_on_a_direct_call(): + """Behavioural counterpart: a direct call must RAISE, not merely be + discouraged. Uses a stub carrying only what the precondition reads, so it + exercises the gate rather than a full likelihood build.""" + class _Data: + lms = [(2, -2), (2, 2)] # m_max = 2: the mode gate would pass + feature = "rotation" # but the response model must not + npts = 4 + + import RIFT.likelihood.jax_ile.core as _core + saved = _core._DISTMARG_GH_N + _core._DISTMARG_GH_N = 64 # GH on + try: + with pytest.raises(ValueError, match="static detector response"): + AM.fused_log_likelihood_distphipsimarg_laplace( + _Data(), np.zeros(1), np.zeros(1), np.zeros(1), + np.ones(2), np.zeros(2)) + finally: + _core._DISTMARG_GH_N = saved + + +def test_global_maxima_cannot_hide_a_locally_invalid_bin(): + """The identity ratio must be POINTWISE, not max|A0| / max|A1|. + + External review: bins (A0,A1) = (1e-3, 1) and (0, 1e6) give a global ratio + of 1e-9 -- passing -- while the first bin violates by 1e-3. The GH + placement computes its centre and width at each bin independently, so one + invalid bin is enough to invalidate the placement there. + """ + C_A, C_B = _tables(m_max=2) + # one bin violates at 1e-3; another carries a denominator 1e6 larger + C_A[:, 1, 0, 0] = 1e-3 + C_A[:, 2, 0, 0] = 1.0 + C_A[:, 2, 1, 1] = 1e6 + ok, info = AM.gh_laplace_supported(C_A, C_B, 2) + assert ok is False, ( + "a locally invalid bin was hidden behind a large denominator elsewhere; " + "the ratio is being taken between unrelated global maxima") + assert info["identity_A0_over_A1"] > AM.GH_PSI_IDENTITY_TOL + + +def test_response_model_is_an_angle_independent_precondition(): + """A numerical probe speaks only for the angles it was evaluated at. + + The placement runs at arbitrary sampled angles, so the gate's real guarantee + is the RESPONSE MODEL: the static path's F+ + i Fx = (F+(0) + i Fx(0)) + e^{-2i psi} is a single u-harmonic at every (ra, dec, incl), which is what + forces A0 == 0 and B1 == 0. The banded features do not use that response, + and an unknown feature must fail closed. + """ + C_A, C_B = _tables(m_max=2) + assert AM.gh_laplace_supported(C_A, C_B, 2, feature=None)[0] is True + for bad in ("freqresponse", "rotation", "something_added_later"): + ok, info = AM.gh_laplace_supported(C_A, C_B, 2, feature=bad) + assert ok is False, "feature %r was admitted" % bad + assert "factorization" in info["gh_laplace_reason"] + + +def test_wrapper_passes_the_response_feature_through(): + import inspect + from RIFT.likelihood.jax_ile import wrapper as WR + src = inspect.getsource(WR.JAXDistPhiPsiMargLikelihood.__init__) + assert 'feature=getattr(data, "feature", None)' in src, ( + "the wrapper does not forward the response model, so the " + "angle-independent half of the gate never runs") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py new file mode 100644 index 000000000..a23525eb1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py @@ -0,0 +1,197 @@ +"""The 'peak-local' angle-marg scheme, as wired into the ILE likelihood. + +These test the WIRING -- that the option reaches the likelihood, produces the same +answer as the scheme it is meant to replace, records its provenance, refuses what it +cannot honour, and does NOT change any default. The kernel's own numerics are tested in +test_joint_anglemarg_peaklocal.py. +""" +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood +from test_angle_marg_exact import make_synth, RA, DEC, INCL, INTERP + + +def test_peak_local_is_an_offered_choice_and_reaches_the_cli(): + """optparse builds --angle-marg-scheme's choices from ANGLE_MARG_CHOICES, so being + in that tuple IS the CLI wiring; a scheme absent from it is unreachable.""" + assert "peak-local" in AM.ANGLE_MARG_CHOICES + assert hasattr(AM, "fused_log_likelihood_distphipsimarg_peaklocal") + + +def test_peak_local_is_NOT_reachable_from_auto(): + """A scheme that changes the likelihood must not become reachable by default on the + strength of unit tests. 'auto' must keep choosing among the schemes that have + campaign evidence until a head-to-head pilot says otherwise.""" + for amp in (1.0, 50.0, 500.0, 5.0e4, 5.0e6): + scheme, _ = AM.choose_angle_marg_scheme(amp) + assert scheme != "peak-local", (amp, scheme) + + +@pytest.mark.parametrize("boost", [1.0, 30.0]) +def test_wrapper_peak_local_matches_exact(boost): + """The wiring's whole claim: asking for it by name gives the same likelihood as the + scheme it parallels.""" + data = make_synth(scale=2.0, kappa_boost=boost) + kw = dict(nphi=32, npsi=8, interp=INTERP) + ex = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, angle_marg="exact", **kw) + pl = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, angle_marg="peak-local", **kw) + assert pl.angle_marg_scheme == "peak-local" + a = np.asarray(ex._batched(jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL))) + b = np.asarray(pl._batched(jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL))) + assert np.abs(a - b).max() < 1e-4, (boost, a, b) + + +def test_peak_local_records_its_provenance(): + """This pipeline has a documented history of silently-inert flags, so the scheme + actually used must be visible in the run record, not inferred from the request.""" + data = make_synth(scale=2.0) + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="peak-local") + info = like.angle_marg_info + assert info["scheme"] == "peak-local" + assert info["requested"] == "peak-local" + assert "amp_sizing" in info, info + + +def test_peak_local_refuses_the_adaptive_distance_quadrature(): + """Refuse rather than silently ignore: this kernel sums the caller's distance grid + and implements no psi-marginal node placement, exactly as the laplace branch does + not. The incompatibility is a property of the scheme, declared once.""" + from RIFT.likelihood.jax_ile import core as _core + data = make_synth(scale=2.0) + old = _core._DISTMARG_GH_N + try: + _core._DISTMARG_GH_N = 8 + with pytest.raises(ValueError, match="peak-local"): + AM.fused_log_likelihood_distphipsimarg_peaklocal( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + jnp.linspace(0.4, 2.0, 8), jnp.zeros(8), interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE) + finally: + _core._DISTMARG_GH_N = old + + +def test_peak_local_requires_amp_sizing_rather_than_guessing_it(): + """There is deliberately no default: the phi axis is still dense here and must be + SIZED from the data, and a silently-undersized grid is the defect this module family + exists to remove.""" + data = make_synth(scale=2.0) + with pytest.raises(Exception): + AM.fused_log_likelihood_distphipsimarg_peaklocal( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + jnp.linspace(0.4, 2.0, 8), jnp.zeros(8), interp=INTERP) + + +def test_the_driver_CLI_accepts_it_and_rejects_a_typo(): + """Deliberately a SUBPROCESS: the flag's job is to travel from a command line into + the likelihood, and optparse builds its choices from ANGLE_MARG_CHOICES at import + time -- a test that imports the module cannot see the CLI wiring. A misspelling + must be loud, because a scheme name absorbed as 'not recognised' would silently run + a different likelihood.""" + import os + import subprocess + import sys + + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + driver = os.path.join(root, "bin", "integrate_likelihood_extrinsic_jax") + env = dict(os.environ) + env["PYTHONPATH"] = root + os.pathsep + env.get("PYTHONPATH", "") + env["JAX_PLATFORMS"] = "cpu" + + def run(value): + p = subprocess.run([sys.executable, driver, "--angle-marg-scheme", value], + env=env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, timeout=600) + return p.stdout.decode("utf-8", "replace") + + good = run("peak-local") + assert "invalid choice" not in good, good[-1500:] + + bad = run("peaklocal") # the plausible misspelling + assert "invalid choice" in bad, bad[-1500:] + assert "peak-local" in bad, bad[-1500:] + + +# --------------------------------------------------------- review findings + +def test_peak_local_runs_the_runtime_amplitude_failsafe(): + """P1 from review. The u axis is localized and needs no sizing, but THE PHI AXIS IS + STILL DENSE and is sized from amp_sizing, which estimate_angle_amplitude is explicit + about being an estimator and NOT a proven bound. A hotter sampled sky location can + therefore under-resolve phi exactly as it can for exact/laplace, and skipping the + check would publish that silently.""" + data = make_synth(scale=2.0, kappa_boost=50.0) + x = jnp.linspace(0.4, 2.0, 8) + lw = jnp.zeros(8) + AM.reset_amp_failsafe() + # size for a much quieter target than the data actually is: the check must notice + AM.fused_log_likelihood_distphipsimarg_peaklocal( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), x, lw, + interp=INTERP, amp_sizing=1.0) + st = AM.amp_failsafe_state(barrier=True) + assert st.get("tripped"), st + assert st.get("scheme") == "peak-local", st + AM.reset_amp_failsafe() + + +def test_peak_local_is_capped_by_the_batch_memory_rule(): + """P1 from review. peak-local still nests sample/time vmaps over the distance grid, + phi chunks, four cells and 48 u nodes, so the batch multiplies the same way the + dense schemes do. Leaving it out of the cap kept an uncapped 8000-sample batch and + reopened a documented 36.4 GiB failure.""" + from RIFT.likelihood.jax_ile import samplers as S + + class _Data(object): + npts = 614 + + class _Like(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(256) + + class _Exact(_Like): + angle_marg_scheme = "exact" + + class _NoScheme(object): + pass + + capped = S.angle_marg_eval_chunk(_Like(), 8000) + assert capped < 8000 + # NOT "same cap as exact" -- that was the earlier assertion and review rightly + # objected that it pins the wrong invariant. peak-local carries the WHOLE distance + # grid inside every phi chunk, so its live slab is ~770x the dense model's + # 8192 bytes/sample/time-point; a cap equal to exact's would look protective and + # would not be. The scheme-specific model must therefore be STRICTLY tighter. + assert capped < S.angle_marg_eval_chunk(_Exact(), 8000), capped + # and it must scale with the distance grid, which is what makes it a model rather + # than a constant + class _Wide(_Like): + x_grid = np.zeros(1024) + assert S.angle_marg_eval_chunk(_Wide(), 8000) <= capped + # the "grid" sentinel means "runs no dense angle scheme" and must stay uncapped + assert S.angle_marg_eval_chunk(_NoScheme(), 8000) == 8000 + + +def test_peak_local_artifacts_carry_the_standing_best_effort_label(): + """P1 from review. A scheme missing from the label's list publishes output with NO + standing statement at all -- and silence is precisely what a reader six months later + would misread as verification. peak-local's phi grid is amp-sized, so its artifacts + are entitled to no more confidence than exact's.""" + import importlib.util + import os + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + path = os.path.join(root, "bin", "integrate_likelihood_extrinsic_jax") + spec = importlib.util.spec_from_loader("_ile_jax_driver", + importlib.machinery.SourceFileLoader( + "_ile_jax_driver", path)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + AM.reset_amp_failsafe() + note = mod.angle_grid_suspect_note("peak-local") + assert note.startswith("ANGLE-GRID-CHECK=BEST-EFFORT"), note + assert mod.angle_grid_suspect_note("grid") == "" diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_sizing_rule.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_sizing_rule.py new file mode 100644 index 000000000..b4604412d --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_sizing_rule.py @@ -0,0 +1,65 @@ +"""The BITING regression for the m_max-aware dense phi sizing rule. + +Split out of test_angle_marg_exact.py so it can stay in the per-PR gate while +that (expensive, validation-oriented) file does not. Pure numpy, milliseconds, +no JAX and no likelihood -- it costs the gate nothing. + +It exists because "correctness does not depend on amplitude" is FALSE for the +dense quadrature: resolving exp(lnL) in phi needs a grid set by the highest +harmonic 2*m_max as well as by sqrt(A), and the old m_max-blind rule passes +every low-scale brute-force test in the excluded file. Without this, reverting +_dense_grid_sizes to the broken rule is green. +""" +import numpy as np + + +def test_dense_phi_sizing_must_scale_with_m_max(): + """BITING regression for the m_max-aware dense phi sizing. + + Pure numpy, no likelihood, no JAX, ~ms -- so it stays in the per-PR gate, + unlike the amplitude ladder that was moved out of it. + + Why it must exist: "correctness does not depend on amplitude" is FALSE for + the dense quadrature. Resolving exp(lnL) in phi needs a grid set by the + highest harmonic 2*m_max as WELL as by sqrt(A), and the old m_max-blind + rule passes every low-scale brute-force test in this file. Without this + test, reverting _dense_grid_sizes to the broken rule is green. + + Construction: a pure order-(2*m_max) harmonic of amplitude b, whose + circular mean has the closed form I0(b) -- so the reference is exact and + needs no dense grid. MEASURED at amp=450, b=150, order=16: the blind rule + (n=352) errs by 4.98e-01 nats at its worst phase; the m_max-aware rule + (n=1360) errs by 1.17e-10. The phase sweep matters -- the error is + phase-dependent and vanishes at favourable alignments. + """ + from scipy.special import ive + import numpy as _np + from RIFT.likelihood.jax_ile.anglemarg import _dense_grid_sizes + + amp, m_max, b = 450.0, 8, 150.0 + order = 2 * m_max + exact = float(_np.log(ive(0, b)) + b) + + def worst(n): + e = 0.0 + for ph in _np.linspace(0.0, 2 * _np.pi / order, 9): + phi = _np.linspace(0.0, 2 * _np.pi, n, endpoint=False) + v = b * _np.cos(order * phi + ph) + m = v.max() + e = max(e, abs(m + _np.log(_np.mean(_np.exp(v - m))) - exact)) + return e + + n_old = _dense_grid_sizes(amp)[0] # m_max-blind (the bug) + n_new = _dense_grid_sizes(amp, m_max=m_max)[0] + err_old, err_new = worst(n_old), worst(n_new) + + assert n_new > n_old, ( + "m_max-aware sizing must request MORE phi points (n_new=%d <= n_old=%d)" + % (n_new, n_old)) + assert err_new < 1e-6, ( + "m_max-aware sizing inaccurate: err=%.3e at n=%d" % (err_new, n_new)) + assert err_old > 1e-2, ( + "this test no longer BITES: the m_max-blind rule errs only %.3e at " + "n=%d, so a revert would pass. Re-tune (b, m_max) until it does." + % (err_old, n_old)) + diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py new file mode 100644 index 000000000..ac43aa90b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py @@ -0,0 +1,252 @@ +"""CHEAP smoke / mutation-bearing coverage of the angle-marginalization feature. + +Exists because the full validation suite (test_angle_marg_exact.py) is too +expensive for a per-PR gate -- it exceeds 20 minutes on 2 cores and has killed +the CI runner three different ways -- but excluding it wholesale would leave the +ENTIRE production feature ungated: a mutation that returns a wrong marginal, +ignores --angle-marg-scheme, breaks the Laplace root enumeration, or drops the +suspect label would all be green. + +So this file covers that surface deliberately at MINIMAL scale. It is not a +substitute for the validation suite's error-law measurements; it is the +mutation-bearing floor that must stay in CI. Keep it cheap: every test here +must run in seconds, or it belongs in the excluded file instead. +""" +import ast +import pathlib + +import numpy as np +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile import build_likelihood_data +from RIFT.likelihood.jax_ile.core import ( + _accumulate_unit, _time_marginalize, _logsumexp_grid_blocked, + fused_log_likelihood_distphipsimarg, phi_ref_grid, psi_grid, + make_distance_grid) + +RA, DEC, INCL = 1.1, -0.35, 0.9 +INTERP = "sinc" +S = 1 + + + +def test_scheme_selector_returns_both_schemes_and_reports(): + """auto must be able to choose EITHER scheme. A regression here is not + hypothetical: a previous head floored the amplitude before passing it to the + selector, so `auto` could never return 'exact' at all.""" + lo, _ = AM.choose_angle_marg_scheme(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE / 10.0) + hi, _ = AM.choose_angle_marg_scheme(AM.ANGLE_MARG_CROSSOVER_AMPLITUDE * 10.0) + assert lo == "exact", "sub-crossover amplitude must select the exact scheme" + assert hi == "laplace", "high amplitude must select the laplace scheme" + assert lo != hi + + +def test_dense_sizes_grow_with_amplitude_and_mode_content(): + """Both levers must be live: sqrt(A) AND m_max. Either being ignored is a + real bug that shipped once.""" + n_lo = AM._dense_grid_sizes(50.0, m_max=2) + n_hi = AM._dense_grid_sizes(5000.0, m_max=2) + n_m = AM._dense_grid_sizes(50.0, m_max=8) + assert n_hi[0] > n_lo[0] and n_hi[1] > n_lo[1], "must grow with amplitude" + assert n_m[0] > n_lo[0], "phi sizing must grow with m_max" + + +def test_amp_sizing_is_required_not_defaulted(): + """A silently-defaulted amplitude is how the SNR-guess bug got in.""" + try: + AM._require_amp_sizing(None) + except Exception as exc: + assert "amp_sizing" in str(exc) + else: + raise AssertionError("a missing amp_sizing must raise, not default") + + +def test_failsafe_record_roundtrips_and_barriers(): + """The host record must reset, report, and barrier -- without it the driver + cannot label an artifact and the condition dies with the log line.""" + import inspect + AM.reset_amp_failsafe() + st = AM.amp_failsafe_state() + assert st["tripped"] is False + for fn in (AM.amp_failsafe_state, AM.reset_amp_failsafe): + assert "effects_barrier" in inspect.getsource(fn) + src = inspect.getsource(AM._runtime_amp_failsafe) + assert src.find("lax.cond") < src.find("debug.callback"), ( + "the callback must sit inside lax.cond; an unconditional callback fires " + "on every likelihood evaluation and destroys throughput") + + +def _driver_src(): + return (pathlib.Path(__file__).resolve().parents[2] / "bin" + / "integrate_likelihood_extrinsic_jax").read_text() + + +def test_driver_actually_passes_the_scheme_through(): + """AST guard on the VALUE node. Checking only that some `angle_marg=` + keyword is present is foolable by hardcoding angle_marg="grid" -- flag + parsed, help present, print present, feature inert. That is this repo's + documented silent-no-op pattern.""" + tree = ast.parse(_driver_src()) + seen = [] + for node in ast.walk(tree): + if isinstance(node, ast.Call): + for k in node.keywords or []: + if k.arg == "angle_marg": + seen.append(k.value) + assert seen, "the driver must pass angle_marg to the wrapper" + assert any(isinstance(v, ast.Name) and v.id == "angle_marg" for v in seen), ( + "angle_marg must be forwarded as the parsed option, not a constant") + + +def test_driver_labels_both_artifacts_and_never_implies_verification(): + src = _driver_src() + assert "def angle_grid_suspect_note(" in src + # The note is computed ONCE in analyze_one from the RESOLVED scheme and + # threaded into both writers. Recomputing it inside each writer with no + # argument left `scheme` None, so the standing label never emitted -- an + # inert guard. So assert the THREADING, which is the real contract. + wd = src[src.index("def write_dat("):] + wd = wd[:wd.index("\ndef ")] + assert "angle_note" in wd, ( + "the EVIDENCE row must carry the label independently of sample export -- " + "write_samples early-returns without --save-samples") + ao = src[src.index("def analyze_one("):] + assert 'angle_grid_suspect_note(_scheme)' in ao, ( + "the note must be computed from the RESOLVED scheme; called with no " + "argument it silently degrades to the empty string") + assert ao.count("angle_note=_ev_note") >= 2, ( + "both writers must receive the same computed note") + assert "BEST-EFFORT" in src, ( + "artifacts must state that no-detection is NOT verification: the " + "detector is a droppable jax callback, so silence cannot be read as " + "an adequate grid") + assert "SUSPECT-ANGLE-GRID" in src + + +def make_synth(scale=1.0, seed=3, modes=((2, 2), (2, -2)), npts=32, + deltaT=1.0 / 1024, kappa_boost=1.0): + """Structurally-faithful synthetic packed data (cf. test_jax_likelihood). + + U is Hermitian positive definite and V complex symmetric, as the real + precompute produces; ``scale`` sets the overall amplitude (lnL ~ scale^2), + standing in for SNR. ``kappa_boost`` multiplies the rholm timeseries + ONLY (not U/V), producing a target with a large coherent (phi,psi) + amplitude A -- the regime where an undersized dense grid measurably + biases the marginal (used by the sizing regression tests). + """ + rng = np.random.default_rng(seed) + tw = npts * deltaT / 2.0 + tvals = np.linspace(-tw, tw, npts) + tref = 1126259462.413 + K = len(modes) + packed = {} + for det in ("H1", "L1"): + npts_full = 4096 + white = (rng.standard_normal((K, npts_full)) + + 1j * rng.standard_normal((K, npts_full))) + kx = np.arange(-40, 41) + kern = np.exp(-0.5 * (kx / 12.0) ** 2) + kern /= kern.sum() + rho = np.stack([np.convolve(white[k].real, kern, "same") + + 1j * np.convolve(white[k].imag, kern, "same") + for k in range(K)]).astype(np.complex128) + rho *= np.sqrt(len(kx)) * scale * kappa_boost + M = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + U = (M @ M.conj().T + 3 * np.eye(K)) * scale ** 2 + B = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = (B @ B.T) * scale ** 2 * 0.3 + packed[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=tref - 0.5) + return build_likelihood_data(packed, deltaT, tref, tvals) + + +RA, DEC, INCL = np.array([0.9]), np.array([0.4]), np.array([1.1]) +S = 1 + + +def _dist_grid(data, n=64): + return make_distance_grid(30.0, 3000.0, n, distMpcRef=data.distMpcRef) + + +def brute_marginal(data, x_grid, log_w, nphi, npsi): + """Brute-force dist+phi+psi marginal: dense product grid of DIRECT + likelihood evaluations (no coefficient machinery shared with the schemes + under test).""" + ph = np.linspace(0, 2 * np.pi, nphi, endpoint=False) + ps = np.linspace(0, np.pi, npsi, endpoint=False) + m = jnp.full((S, data.npts), -jnp.inf) + s = jnp.zeros((S, data.npts)) + for p in ph: + rb = np.repeat(RA[None, :], npsi, 0).ravel() + db = np.repeat(DEC[None, :], npsi, 0).ravel() + ib = np.repeat(INCL[None, :], npsi, 0).ravel() + pb = np.full(npsi * S, p) + sb = np.repeat(ps[:, None], S, 1).ravel() + ku, rs = _accumulate_unit(data, rb, db, sb, ib, pb, INTERP, False) + lnL = _logsumexp_grid_blocked(ku.real, rs, x_grid, + -0.5 * jnp.square(x_grid), log_w, 64) + m, s = AM._lse_update(m, s, lnL.reshape(npsi, S, data.npts), axis=0) + lnL_t = m + jnp.log(s) - np.log(nphi * npsi) + return np.asarray(_time_marginalize(lnL_t, data.w_t)) + + + +# --------------------------------------------------------------------------- +# NUMERICAL execution. Everything above this line checks selection, sizing and +# source wiring -- none of it would catch a mutation that returns a WRONG +# MARGINAL or breaks the Laplace stationary-point enumeration. These two do, at +# the smallest scale that still discriminates. +# +# COVERAGE LIMIT, stated so this file is not mistaken for end-to-end coverage: +# * the EXACT scheme is exercised END TO END, through +# fused_log_likelihood_distphipsimarg_exact -- phi, psi, distance and time +# marginalization included -- against a direct product-grid reference. +# * the LAPLACE scheme is exercised only at its KERNEL, _laplace_psi_lnI. +# fused_log_likelihood_distphipsimarg_laplace and its phi/distance/time +# marginalization are NOT run here. +# +# That asymmetry is a deliberate CI-cost tradeoff, not an oversight: the fused +# Laplace path is covered in test_angle_marg_exact.py, which is EXCLUDED from +# the per-PR gate (it exceeds 20 minutes on 2 cores) and must be run by hand +# when touching anglemarg.py -- the command is in .travis/test-jax.sh next to +# the exclusion. A cheap fused-Laplace finiteness smoke test would close the +# gap and is worth adding if someone finds a configuration that stays fast. +# --------------------------------------------------------------------------- + +def test_exact_scheme_matches_a_direct_reference_small_scale(): + """Minimal exact-vs-reference: catches a wrong marginal. + + scale=1.5 keeps the dense grids tiny, so this is seconds -- the expensive + amplitude ladder stays in the excluded validation suite. The reference is a + direct product-grid evaluation that shares no coefficient machinery with the + scheme under test. + """ + data = make_synth(scale=1.5) + x_grid, log_w = _dist_grid(data) + amp = AM.estimate_angle_amplitude(data, x_grid) + got = AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=amp) + ref = brute_marginal(data, x_grid, log_w, 256, 128) + d = float(np.max(np.abs(np.asarray(got) - np.asarray(ref)))) + assert d < 1e-3, "exact scheme disagrees with a direct reference by %.3e nats" % d + + +def test_laplace_survives_first_harmonic_cancellation(): + """Regression for the missed-maxima defect (first external review). + + When the FIRST polarization harmonic cancels (c1 = 0) while the quadratic + one survives (c2 = -d, d > 0.5), both of the original two Newton seeds land + on MINIMA, every term was rejected, and the kernel returned -inf for an + integral that is finite. Enumerating the stationary points fixes it. Pure + kernel, no data, milliseconds. + """ + for d in (0.6, 2.0, 25.0): + # c1 = 0 (first harmonic cancels), c2 = -d (quadratic harmonic survives) + val = float(np.asarray(AM._laplace_psi_lnI( + jnp.asarray(0.0), jnp.asarray(0.0 + 0.0j), + jnp.asarray(-float(d) + 0.0j)))) + assert np.isfinite(val), ( + "_laplace_psi_lnI returned non-finite at c1=0, c2=-%g -- the " + "first-harmonic cancellation case, whose integral is finite" % d) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_distance_grid_loguniform.py b/MonteCarloMarginalizeCode/Code/test/jax/test_distance_grid_loguniform.py new file mode 100644 index 000000000..b92497aa7 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_distance_grid_loguniform.py @@ -0,0 +1,1219 @@ +"""Gate for the log-uniform ("peak-resolving") distance quadrature. + +WHY THESE TESTS AND NOT OTHERS. The lever is one number -- the node count -- +and one placement rule. A test that merely builds a grid and checks it has +nodes would pass under every mutation that matters. Each test below was +written against a specific mutation and VERIFIED to fail under it; the matrix +is in the PR body. Three properties carry the contract: + + * the SPACING contract (relative spacing <= c/rho_max everywhere), and the + calibration of c against the Gaussian trapezoid error law it is derived + from -- pinned two-sided, because a c that is merely small satisfies a + one-sided bound while making the grid uselessly expensive; + * the DECOUPLING property: the dense angle lattice is sized from the + amplitude on the FULL prior support, so no distance grid can shrink it. + Asserted on the built objects, not on the helper, because a helper-level + assertion cannot see a call site that stops calling the helper; + * the DEFAULT: dist_grid="uniform" must reproduce today's grid node for node. + +Everything here runs in seconds. The convergence measurements that justify +the shipped tolerance live in DESIGN_jax_distance_quadrature.md beside the +module; they need a real precompute and are not gated. +""" +import ast +import os +import pathlib +import sys + +import numpy as np +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import build_likelihood_data +from RIFT.likelihood.jax_ile.core import ( + make_distance_grid, make_distance_grid_loguniform, loguniform_grid_size, + loguniform_spacing_for_tolerance, DIST_GRID_TOL_DEFAULT, DIST_GRID_SCHEMES) +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood + +D_MIN, D_MAX = 1.0, 10000.0 +DREF = 1000.0 + + +def _nodes(x_grid, distMpcRef=DREF): + return distMpcRef / np.asarray(x_grid) + + +# --------------------------------------------------------------------------- +# 1. The spacing contract, and the calibration of the constant it rests on. +# --------------------------------------------------------------------------- + +def test_relative_spacing_meets_the_stated_contract(): + """Delta(ln d) <= c(tol)/rho_max at EVERY interval, for every case. + + This is the whole claim: one spacing resolves a peak of relative width + 1/rho wherever it sits. An off-by-one in the node count, a floor for the + ceil, or a linspace for the geomspace all break it. + """ + for rho in (5.0, 38.9, 55.07, 300.0): + for tol in (5e-1, 1e-1, DIST_GRID_TOL_DEFAULT, 1e-3): + for lo, hi in ((1.0, 10000.0), (50.0, 2000.0)): + x, _ = make_distance_grid_loguniform( + lo, hi, rho, distMpcRef=DREF, tol=tol) + d = _nodes(x) + h = np.max(np.diff(np.log(d))) + c = loguniform_spacing_for_tolerance(tol) + assert h <= c / rho * (1 + 1e-12), ( + "spacing contract violated: rho=%g tol=%g range=[%g,%g] " + "h=%.6g > c/rho=%.6g" % (rho, tol, lo, hi, h, c / rho)) + # ...and the grid must actually span the requested support + assert np.isclose(d[0], lo) and np.isclose(d[-1], hi) + + +def test_tolerance_constant_matches_the_gaussian_trapezoid_error_law(): + """c = pi*sqrt(2/ln(2/tol)) must reproduce the error it is derived from. + + Two-sided on purpose. A one-sided "error <= tol" check passes for any c + smaller than the right one -- including c -> 0, which satisfies every + accuracy claim while making the grid arbitrarily expensive. The lower + bound is what pins the constant to the LAW rather than to caution. + """ + for tol in (5e-1, 1e-1, 1e-2, 1e-3, 1e-4): + c = loguniform_spacing_for_tolerance(tol) + h = c * 1.0 # sigma == 1 without loss of generality + worst = 0.0 + for mu in np.linspace(0.0, h, 25): # the error oscillates with phase + k = np.arange(-int(np.ceil(40.0 / h)) - 1, int(np.ceil(40.0 / h)) + 2) + u = k * h + approx = h * np.sum(np.exp(-0.5 * (u - mu) ** 2)) + worst = max(worst, abs(approx / np.sqrt(2 * np.pi) - 1.0)) + assert worst <= tol * 1.05, ( + "tol=%g: measured trapezoid error %.4g exceeds the target" % (tol, worst)) + assert worst >= tol / 20.0, ( + "tol=%g: measured error %.4g is far below the target -- the " + "constant no longer tracks the error law it is derived from, so " + "the grid is paying for accuracy nobody asked for" % (tol, worst)) + + +def test_node_count_responds_to_every_lever(): + """rho_max, the prior range, and tol must all be live. Each has been an + inert argument in some draft of this helper.""" + base = loguniform_grid_size(D_MIN, D_MAX, 50.0, 1e-2) + assert loguniform_grid_size(D_MIN, D_MAX, 100.0, 1e-2) > base, "rho lever dead" + assert loguniform_grid_size(10.0, D_MAX, 50.0, 1e-2) < base, "range lever dead" + assert loguniform_grid_size(D_MIN, D_MAX, 50.0, 1e-4) > base, "tol lever dead" + # and the count is the closed form, not an approximation of it + c = loguniform_spacing_for_tolerance(1e-2) + assert base == int(np.ceil(50.0 * np.log(D_MAX / D_MIN) / c)) + 1 + + +def test_rho_max_must_be_a_bound_and_never_falls_back(): + """A missing or degenerate bound must RAISE. Falling back to a default + node count is the silent-no-op pattern that has bitten this module: the + run would look normal and the marginal would be wrong.""" + for bad in (0.0, -1.0, float("nan"), float("inf"), None): + try: + loguniform_grid_size(D_MIN, D_MAX, bad, 1e-2) + except (ValueError, TypeError): + pass + else: + raise AssertionError("rho_max=%r must raise, not fall back" % (bad,)) + for lo, hi in ((0.0, 10.0), (10.0, 10.0), (100.0, 10.0)): + try: + loguniform_grid_size(lo, hi, 50.0, 1e-2) + except ValueError: + pass + else: + raise AssertionError("range (%r,%r) must raise" % (lo, hi)) + for bad_tol in (0.0, -1.0, 2.0, 5.0): + try: + loguniform_spacing_for_tolerance(bad_tol) + except ValueError: + pass + else: + raise AssertionError("tol=%r must raise" % (bad_tol,)) + + +def test_n_max_raises_instead_of_clamping(): + """Clamping would silently violate the spacing contract -- the grid would + still be built, still be log-uniform, and no longer resolve the peak.""" + try: + make_distance_grid_loguniform(D_MIN, D_MAX, 5000.0, distMpcRef=DREF, + tol=1e-3, n_max=64) + except ValueError as exc: + assert "n_max" in str(exc) + else: + raise AssertionError("an over-cap node count must raise, not clamp") + + +def test_weights_are_a_normalized_proper_distance_average(): + x, lw = make_distance_grid_loguniform(D_MIN, D_MAX, 55.07, distMpcRef=DREF) + w = np.exp(np.asarray(lw)) + assert np.isclose(w.sum(), 1.0, rtol=0, atol=1e-12) + assert np.all(np.isfinite(np.asarray(lw))) + d = _nodes(x) + assert np.all(np.diff(d) > 0), "nodes must be strictly increasing" + # the volumetric prior must be the one being averaged over + x2, lw2 = make_distance_grid_loguniform(D_MIN, D_MAX, 55.07, + d_prior="uniform", distMpcRef=DREF) + assert not np.allclose(np.asarray(lw), np.asarray(lw2)) + + +def test_end_intervals_are_half_width(): + """On a log grid the last interval is ~1% of d_max, so giving the end nodes + a FULL interval (the convention make_distance_grid_adaptive uses) misplaces + percent-level volumetric prior mass onto d_max -- measured as a ~0.018 nat + error floor that no refinement removes. Pinned by comparing the raw + trapezoid prior mass against the exact integral, with the wrong convention + computed here so the test discriminates rather than merely passing.""" + n = loguniform_grid_size(D_MIN, D_MAX, 55.07, DIST_GRID_TOL_DEFAULT) + d = np.geomspace(D_MIN, D_MAX, n) + exact = (D_MAX ** 3 - D_MIN ** 3) / 3.0 + dd = np.empty_like(d) + dd[1:-1] = 0.5 * (d[2:] - d[:-2]) + dd[0], dd[-1] = 0.5 * (d[1] - d[0]), 0.5 * (d[-1] - d[-2]) + half = np.sum(d ** 2 * dd) / exact + dd[0], dd[-1] = d[1] - d[0], d[-1] - d[-2] + full = np.sum(d ** 2 * dd) / exact + assert abs(half - 1.0) < 1e-3, "half-width ends: prior mass off by %.4g" % (half - 1) + assert abs(full - 1.0) > 1e-2, ( + "the two conventions no longer differ measurably; this test can no " + "longer detect the wrong one") + # and the shipped helper must use the half-width convention + _, lw = make_distance_grid_loguniform(D_MIN, D_MAX, 55.07, distMpcRef=DREF) + w = np.exp(np.asarray(lw)) + assert np.isclose(w[-1] / w.sum(), (d[-1] ** 2 * 0.5 * (d[-1] - d[-2])) + / (np.sum(d ** 2 * np.concatenate( + [[0.5 * (d[1] - d[0])], 0.5 * (d[2:] - d[:-2]), + [0.5 * (d[-1] - d[-2])]]))), rtol=1e-10) + + +# --------------------------------------------------------------------------- +# 2. Wiring, on a real (tiny, synthetic) likelihood object. +# --------------------------------------------------------------------------- + +def _synth(scale=1.0, seed=3, modes=((2, 2), (2, -2)), npts=32, deltaT=1.0 / 1024, + kappa_boost=1.0): + """Structurally-faithful packed data (U Hermitian PD, V complex symmetric), + the same construction test_angle_marg_smoke uses, duplicated here so this + file does not depend on another test module's import order.""" + rng = np.random.default_rng(seed) + tw = npts * deltaT / 2.0 + tvals = np.linspace(-tw, tw, npts) + tref = 1126259462.413 + K = len(modes) + packed = {} + for det in ("H1", "L1"): + white = (rng.standard_normal((K, 4096)) + 1j * rng.standard_normal((K, 4096))) + kx = np.arange(-40, 41) + kern = np.exp(-0.5 * (kx / 12.0) ** 2) + kern /= kern.sum() + rho = np.stack([np.convolve(white[k].real, kern, "same") + + 1j * np.convolve(white[k].imag, kern, "same") + for k in range(K)]).astype(np.complex128) + rho *= np.sqrt(len(kx)) * scale * kappa_boost + M = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + U = (M @ M.conj().T + 3 * np.eye(K)) * scale ** 2 + B = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = (B @ B.T) * scale ** 2 * 0.3 + packed[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=tref - 0.5) + return build_likelihood_data(packed, deltaT, tref, tvals) + + +def _like(data, dist_grid="uniform", angle_marg="laplace", n_grid=64, + d_min=D_MIN, d_max=D_MAX, **kw): + return JAXDistPhiPsiMargLikelihood( + data, d_min, d_max, nphi=8, npsi=8, n_grid=n_grid, interp="sinc", + angle_marg=angle_marg, dist_grid=dist_grid, **kw) + + +def test_default_distance_grid_is_bit_identical_to_the_shipped_uniform_grid(): + """NEVER change a RIFT default. An existing command line must reproduce + an existing run node for node -- asserted with exact equality, not + allclose, because a reproduction claim tolerates no drift.""" + data = _synth() + # Constructed WITHOUT dist_grid=, deliberately: passing it explicitly -- + # even as "uniform" -- means the constructor DEFAULT is never exercised, + # and a flipped default sails through. (It did: this test survived that + # mutation until the explicit keyword was removed.) + like = JAXDistPhiPsiMargLikelihood( + data, D_MIN, D_MAX, nphi=8, npsi=8, n_grid=64, interp="sinc", + angle_marg="laplace") + x_ref, lw_ref = make_distance_grid(D_MIN, D_MAX, 64, "euclidean", + distMpcRef=data.distMpcRef) + assert np.array_equal(np.asarray(like.x_grid), np.asarray(x_ref)) + assert np.array_equal(np.asarray(like.log_w_grid), np.asarray(lw_ref)) + assert like.dist_grid_info["mode"] == "uniform" + assert like.dist_grid_info["n"] == 64 + + +def test_narrowing_the_distance_grid_can_move_the_sizing_amplitude(): + """The PREMISE of the next test. estimate_angle_amplitude reads only + min/max of the grid it is handed and clips the per-angle distance maximum + to it, so a grid that stops containing A/B reports a smaller amplitude -- + which would silently shrink the dense angle lattice. If this ever stops + being true, the guard below is guarding nothing and must be revisited.""" + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=3.0, kappa_boost=4.0) + full, _ = make_distance_grid(D_MIN, D_MAX, 64, "euclidean", + distMpcRef=data.distMpcRef) + narrow, _ = make_distance_grid(2000.0, 4000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + a_full = AM.estimate_angle_amplitude(data, full, interp="sinc") + a_narrow = AM.estimate_angle_amplitude(data, narrow, interp="sinc") + assert a_narrow < a_full, ( + "a narrowed distance grid no longer lowers the sizing amplitude " + "(%.6g vs %.6g); the decoupling guard has nothing left to guard" + % (a_narrow, a_full)) + + +def test_angle_lattice_is_sized_from_the_full_support_grid_by_name(): + """THE safety property, guarded at the CALL SITE. + + A behavioural test cannot discriminate this one, and saying so is the + point: BOTH shipped schemes span the full prior range, so handing + estimate_angle_amplitude `self.x_grid` instead of `x_grid_full` produces + the identical amplitude today. The mutation survives every value-level + assertion. What the change actually buys is that the invariant is + structural rather than incidental -- it stays true the moment any + narrowing scheme is added. NO in-tree scheme is currently such a scheme, + and that correction matters: an earlier draft of this docstring named the + deprecated JAX_ILE_DISTGRID_ADAPTIVE branch and quoted 12.6%, which is + wrong. make_distance_grid_adaptive concatenates a full-range `linspace` + backbone before dedup, so its x_min/x_max ARE the full support's and it + returns a byte-identical amplitude -- 105.737261 either way on THIS file's + _synth(scale=3.0, kappa_boost=4.0), and 10573.7261 either way on the louder + _synth(scale=30.0, kappa_boost=40.0) the figure was first taken from. (An + earlier draft quoted only the second, unlabelled, which reproduces 100x off + if you use the fixture this file actually ships.) The 12.6-14.8% figure + belongs to the hand-built [0.8 d, 1.25 d] window in wrapper.py, which no + code path produces. So the guard is on the argument + the wrapper passes, which is where the property lives, and it is + prospective -- see test_narrowing_the_distance_grid_can_move_the_sizing_ + amplitude, which pins the premise the guard rests on.""" + import inspect + import textwrap + import RIFT.likelihood.jax_ile.wrapper as W + tree = ast.parse(textwrap.dedent( + inspect.getsource(W.JAXDistPhiPsiMargLikelihood.__init__))) + calls = [n for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "estimate_angle_amplitude"] + assert len(calls) == 1, ( + "expected exactly one estimate_angle_amplitude call site in " + "JAXDistPhiPsiMargLikelihood.__init__, found %d" % len(calls)) + arg = calls[0].args[1] + assert isinstance(arg, ast.Name) and arg.id == "x_grid_full", ( + "the sizing amplitude must be computed on the FULL-prior-support grid " + "(x_grid_full), not on the grid the likelihood integrates over; got %s" + % ast.dump(arg)) + # and the shipped schemes must agree, which is the invariant users see + data = _synth(scale=3.0, kappa_boost=4.0) + a = _like(data, "uniform", n_grid=64) + b = _like(data, "loguniform", n_grid=64) + assert a.angle_marg_info["amp_sizing"] == b.angle_marg_info["amp_sizing"] + assert a.angle_marg_info["sample_grid"] == b.angle_marg_info["sample_grid"] + assert a.x_grid.shape[0] != b.x_grid.shape[0], ( + "the two arms must actually differ in the distance grid, or the " + "invariant above is vacuous") + + +def test_loguniform_records_what_it_built(): + data = _synth() + like = _like(data, "loguniform", n_grid=64) + gi = like.dist_grid_info + assert gi["mode"] == "loguniform" + assert gi["tol"] == DIST_GRID_TOL_DEFAULT + assert gi["rho_max"] > 0.0 + assert gi["n"] == loguniform_grid_size(D_MIN, D_MAX, gi["rho_max"], + DIST_GRID_TOL_DEFAULT) + assert gi["dlnd"] <= loguniform_spacing_for_tolerance( + DIST_GRID_TOL_DEFAULT) / gi["rho_max"] * (1 + 1e-12) + + +def test_rho_max_is_the_amplitude_the_runtime_failsafe_actually_compares(): + """The distance grid's coverage by the existing fail-safe is an identity, + not a correlation -- and it only holds if the spacing is sized from + amp_SIZING (floored at the crossover), which is what + anglemarg._runtime_amp_failsafe compares the per-call amplitude against. + + Sizing from the UNfloored amp_data instead leaves a silent gap for quiet + targets: a runtime amplitude between amp_data and amp_sizing under-resolves + the distance peak and trips nothing. This test uses a quiet synthetic + target, where the two differ, so it can see the difference.""" + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=0.05, kappa_boost=0.05) + like = _like(data, "loguniform", n_grid=64) + amp_data = like.angle_marg_info["amplitude"] + amp_sizing = like.angle_marg_info["amp_sizing"] + assert amp_data < AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, ( + "this synthetic target is no longer quiet enough for the floor to " + "bind, so the test can no longer distinguish the two amplitudes") + assert amp_sizing == AM.ANGLE_MARG_CROSSOVER_AMPLITUDE + # ...and from TRIP * amp_sizing, not amp_sizing alone. The fail-safe stays + # silent until amp_call > TRIP * amp_sizing, so sizing the spacing at + # amp_sizing leaves a window of admitted-but-unresolved amplitudes that + # nothing labels -- and on any vmapped path the label is uninformative + # anyway (DESIGN section 5(a)). Costs sqrt(TRIP) = 1.41x the nodes. + want = np.sqrt(2 * AM.AMP_FAILSAFE_TRIP_FACTOR * amp_sizing) + assert np.isclose(like.dist_grid_info["rho_max"], want, rtol=1e-12), ( + "rho_max must come from TRIP(%.4g) * amp_sizing(%.6g) -> %.6g, not " + "amp_data (%.6g) and not the unmultiplied amp_sizing (%.6g); got %.6g" + % (AM.AMP_FAILSAFE_TRIP_FACTOR, amp_sizing, want, amp_data, + np.sqrt(2 * amp_sizing), like.dist_grid_info["rho_max"])) + assert AM.AMP_FAILSAFE_TRIP_FACTOR > 1.0, ( + "a TRIP factor of 1 would reopen the window this sizing closes") + + +def test_unrecognised_scheme_raises_and_never_falls_through(): + data = _synth() + for bad in ("adaptive", "log-uniform", "geometric", "", None, 1): + try: + _like(data, bad) + except ValueError as exc: + assert "dist_grid" in str(exc) + else: + raise AssertionError("dist_grid=%r must raise, not default" % (bad,)) + assert DIST_GRID_SCHEMES == ("uniform", "loguniform") + + +def test_loguniform_is_refused_on_the_grid_angle_scheme_not_ignored(): + """The sizing amplitude does not exist on the 'grid' path. Silently using + the uniform grid there would be a flag that parses, prints and does + nothing.""" + data = _synth() + try: + _like(data, "loguniform", angle_marg="grid") + except ValueError as exc: + assert "angle_marg" in str(exc) + else: + raise AssertionError("dist_grid='loguniform' with angle_marg='grid' " + "must raise") + + +def test_loguniform_marginal_agrees_with_a_fine_uniform_grid(): + """NUMERICAL execution. Everything above checks structure; a mutation that + returns a wrong marginal (a mis-signed weight, a reversed node order, a + dropped prior factor) passes all of it. This does not.""" + data = _synth(scale=2.0, kappa_boost=3.0) + ra = np.array([0.9, 2.4]); dec = np.array([0.4, -0.7]); incl = np.array([1.1, 2.0]) + fine = JAXDistPhiPsiMargLikelihood( + data, D_MIN, D_MAX, nphi=8, npsi=8, n_grid=1024, interp="sinc", + angle_marg="laplace", dist_grid="uniform") + lg = _like(data, "loguniform", n_grid=64, dist_grid_tol=1e-3) + v_fine = np.asarray(fine.log_likelihood(ra, dec, incl)) + v_log = np.asarray(lg.log_likelihood(ra, dec, incl)) + assert np.all(np.isfinite(v_log)) + assert np.max(np.abs(v_log - v_fine)) < 0.02, ( + "loguniform marginal differs from a 1024-node uniform reference by " + "%.4g nats" % np.max(np.abs(v_log - v_fine))) + + +# A gradient test lived here and was REMOVED, deliberately. It asserted that +# value_and_grad stays finite under the log-uniform grid. It passed -- and it +# passed under every mutation that could be constructed against it, including a +# degenerate node set with zero-width intervals (which four other tests here +# caught). The reason is structural: the distance grid is a compile-time +# CONSTANT in the traced graph, so no change to it can poison the backward pass +# while leaving the forward value intact. A test that cannot be made to fail is +# coverage-shaped, not coverage, and it cost 15 s of gate time. If the node +# positions are ever made traced (the per-sample quadrature of +# DESIGN_jax_distance_quadrature.md section 4d does exactly that), reinstate it +# -- there it would have real work to do. + + +# --------------------------------------------------------------------------- +# 3. The driver seam. A library test cannot see a driver that stops calling. +# --------------------------------------------------------------------------- + +_CODE = pathlib.Path(__file__).resolve().parents[2] + + +def _driver_src(): + return (_CODE / "bin" / "integrate_likelihood_extrinsic_jax").read_text() + + +def _option_call(tree, flag): + for node in ast.walk(tree): + if isinstance(node, ast.Call) and node.args: + a0 = node.args[0] + if isinstance(a0, ast.Constant) and a0.value == flag: + return node + return None + + +def test_driver_exposes_the_option_and_defaults_it_to_uniform(): + """AST on the VALUE node. 'the string appears in the file' is satisfied by + a help text alone.""" + tree = ast.parse(_driver_src()) + call = _option_call(tree, "--distance-grid-scheme") + assert call is not None, "the driver must define --distance-grid-scheme" + kw = {k.arg: k.value for k in call.keywords} + assert isinstance(kw["default"], ast.Constant) and kw["default"].value == "uniform", ( + "the default must remain 'uniform': an existing command line must " + "reproduce an existing run") + choices = kw["choices"] + assert isinstance(choices, ast.Tuple) + assert [c.value for c in choices.elts] == ["uniform", "loguniform"], ( + "an unrecognised value must be refused by the parser, not defaulted") + + +def test_driver_forwards_the_parsed_option_not_a_constant(): + """Hardcoding dist_grid="uniform" at the call site passes every weaker + guard: flag parsed, help printed, feature inert.""" + tree = ast.parse(_driver_src()) + seen = [k.value for node in ast.walk(tree) if isinstance(node, ast.Call) + for k in (node.keywords or []) if k.arg == "dist_grid"] + assert seen, "the driver must pass dist_grid to JAXDistPhiPsiMargLikelihood" + assert any(isinstance(v, ast.Name) for v in seen), ( + "dist_grid must be forwarded as the parsed option, not a constant") + seen_tol = [k.value for node in ast.walk(tree) if isinstance(node, ast.Call) + for k in (node.keywords or []) if k.arg == "dist_grid_tol"] + assert any(isinstance(v, ast.Name) for v in seen_tol), ( + "dist_grid_tol must be forwarded as the parsed option") + + +def test_driver_refuses_the_option_on_modes_that_do_not_implement_it(): + src = _driver_src() + assert "--distance-grid-scheme/--distance-grid-tol apply only to" in src, ( + "the driver must fail closed when the option is set on a mode that " + "ignores it; a silently inert flag is this pipeline's documented " + "failure mode") + + +def test_driver_refuses_distance_grid_points_together_with_loguniform(): + """Two options that both set the node count. Whichever loses silently is a + silently-inert flag; refusing is the only reading that cannot mislead. The + option therefore defaults to None (a sentinel), resolved to + DISTANCE_GRID_POINTS_DEFAULT -- optparse cannot otherwise distinguish + "explicitly 256" from "not passed".""" + src = _driver_src() + tree = ast.parse(src) + call = _option_call(tree, "--distance-grid-points") + assert call is not None + kw = {k.arg: k.value for k in call.keywords} + assert isinstance(kw["default"], ast.Constant) and kw["default"].value is None, ( + "--distance-grid-points must default to the None sentinel, or the " + "driver cannot tell an explicit 256 from an unset option") + assert "--distance-grid-points and --distance-grid-scheme loguniform both" in src + assert any(isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "DISTANCE_GRID_POINTS_DEFAULT" + for t in n.targets) + for n in ast.walk(tree)), "the resolved default must be a named constant" + + +def test_driver_never_writes_the_resolved_option_back_onto_opts(): + """Writing a resolved option back onto `opts` made event 1 of a batch read + event 0's choice once already in this driver (see DESIGN_jax_tempering.md + and test_jax_tempering_chooser). The resolution must land in a local.""" + tree = ast.parse(_driver_src()) + for node in ast.walk(tree): + if isinstance(node, ast.Assign): + for t in node.targets: + if (isinstance(t, ast.Attribute) + and t.attr in ("distance_grid_points", "distance_grid_scheme", + "distance_grid_tol") + and isinstance(t.value, ast.Name) and t.value.id == "opts"): + raise AssertionError( + "opts.%s is assigned; resolve into a local instead" % t.attr) + + +def test_driver_always_reports_the_distance_grid_that_ran(): + src = _driver_src() + assert 'print(" distance grid: %s"' in src, ( + "the resolved distance grid must be printed unconditionally; a report " + "guarded on mode == 'adaptive' cannot distinguish uniform from a " + "silently-failed request") + + +# --------------------------------------------------------------------------- +# 4. Regressions added after external review (F1, F2, F3, F8). Each is here +# because a mutation SURVIVED the first matrix without it. +# --------------------------------------------------------------------------- + +def test_truncated_distance_support_is_refused_not_silently_mis_sized(): + """F1. The spacing contract assumes the integrand is a Gaussian PEAK + inside the support. When the maximizing distance x* = A/B is exterior the + integrand is a boundary LAYER at a prior edge instead, and a log-uniform + grid is the wrong instrument twice over: its absolute spacing is coarsest + exactly at d_max, and refining it adds nodes proportionally everywhere so + the layer never resolves (measured: tol 0.5 -> 1e-9 moves the error only + 5.23 -> 3.92 nats, while uniform 256 -> 4096 moves 2.52 -> 0.36). The + clip also makes the amplitude UNDER-read, so the node count moves the + WRONG WAY. Refused at build time; NOT silently fallen back to uniform, + which would make this flag produce the other scheme's grid.""" + data = _synth(scale=3.0, kappa_boost=4.0) + try: + _like(data, "loguniform", n_grid=64, d_min=2000.0, d_max=10000.0) + except ValueError as exc: + assert "OUTSIDE" in str(exc) and "d_min" in str(exc) + assert "--d-max" in str(exc), "the refusal must name the recourse" + else: + raise AssertionError("a truncated distance support must be refused") + # ...and the interior control must still build + like = _like(data, "loguniform", n_grid=64) + assert like.dist_grid_info["mode"] == "loguniform" + + +def test_clip_excess_diagnostic_detects_exteriority_and_is_quiet_when_interior(): + """The premise of the refusal above, at the estimator. A guard whose + detector cannot distinguish the two regimes is not a guard.""" + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=3.0, kappa_boost=4.0) + xg_in, _ = make_distance_grid(1.0, 10000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + xg_out, _ = make_distance_grid(2000.0, 10000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + _, d_in = AM.estimate_angle_amplitude(data, xg_in, interp="sinc", + return_diagnostics=True) + _, d_out = AM.estimate_angle_amplitude(data, xg_out, interp="sinc", + return_diagnostics=True) + assert d_in["clip_excess"] <= 1.0 + 1e-9, "interior support must not trip" + assert d_out["clip_excess"] > 1.0 + 1e-3, "exterior support must trip" + assert d_out["amp_clipped"] < d_out["amp_unclipped"] + # the returned amplitude itself is unchanged by the diagnostic + a_plain = AM.estimate_angle_amplitude(data, xg_in, interp="sinc") + a_diag, _ = AM.estimate_angle_amplitude(data, xg_in, interp="sinc", + return_diagnostics=True) + assert a_plain == a_diag + + +def test_loguniform_is_refused_under_the_per_sample_gh_quadrature(): + """F2. core._distmarg_gh_logL places its own nodes and reads ONLY + min/max of x_grid, so both schemes are bit-identical under it while + dist_grid_info still reports mode='loguniform'. Reachable without the user + naming a dense scheme at all: under GH choose_angle_marg_scheme selects one + regardless of what was asked for. The refusal does not depend on WHICH -- + the per-sample quadrature reads only the support on every dense path -- so + this stays correct if that selector's choice under GH ever changes.""" + from RIFT.likelihood.jax_ile import core as C + data = _synth() + saved = C._DISTMARG_GH_N + C._DISTMARG_GH_N = 32 + try: + _like(data, "loguniform", n_grid=64) + except ValueError as exc: + assert "JAX_ILE_DISTMARG_GH" in str(exc) + assert "inert" in str(exc) + else: + raise AssertionError("loguniform under GH must raise, not run inert") + finally: + C._DISTMARG_GH_N = saved + + +def test_zero_clipped_amplitude_is_the_most_exterior_case_and_is_refused(): + """F1, the EXTREME, which the test above cannot reach. + + Section 1a's worst row is the one where the clipped amplitude reaches + exactly 0: the whole prior support lies beyond the maximizer, so + ``x*A - x^2 B/2 <= 0`` at every sampled angle and the max over ``x >= 0`` + is the floor. ``clip_excess`` then has no ratio to form, and a separate + arm of the expression -- ``inf if amp_unclipped > 0 else 1.0`` -- is what + decides the refusal. Replacing that arm with a bare ``1.0`` leaves every + other test in this file GREEN while the design note's worst case (amp -> 0, + crossover floor, grid collapses) BUILDS. Verified: 30/30 still passed + under exactly that mutation, and the constructor returned a 37-node grid. + + ``test_clip_excess_diagnostic_...`` uses a [2000, 10000] support, where the + clipped amplitude is positive, so it never enters this arm. + """ + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=3.0, kappa_boost=4.0) + xg, _ = make_distance_grid(1.0, 10.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + _, diag = AM.estimate_angle_amplitude(data, xg, interp="sinc", + return_diagnostics=True) + assert diag["amp_clipped"] == 0.0, ( + "this support no longer drives the CLIPPED amplitude to exactly 0 " + "(%.6g), so it can no longer exercise the amp_emp == 0 arm and this " + "test is guarding nothing" % diag["amp_clipped"]) + assert diag["amp_unclipped"] > 0.0, ( + "the UNCLIPPED amplitude must stay positive here, or there is no " + "exteriority left to detect") + assert diag["clip_excess"] == float("inf"), ( + "a zero clipped amplitude against a positive unclipped one is the " + "MOST exterior case there is; reporting a finite ratio -- above all " + "an interior-looking 1.0 -- disarms the refusal in exactly the regime " + "section 1a measures at +4.60 nats") + try: + _like(data, "loguniform", n_grid=64, d_min=1.0, d_max=10.0) + except ValueError as exc: + assert "OUTSIDE" in str(exc) + else: + raise AssertionError( + "the extreme exterior case (clipped amplitude 0) must be refused, " + "not built") + + +def test_driver_refuses_the_gh_combination_at_PARSE_time(): + """F2, the DRIVER half -- which no other test in this file covers. + + The constructor refuses this combination as well, so nothing can silently + run; but the parse-time half is the one that spares the user a full + precompute (F8), and DESIGN section 5 asserts in writing that it fires + there. Deleting that arm of check_critical_and_report left all 30 tests + here green. + + ``--angle-marg-scheme auto``, not ``exact``: under GH the selector + resolves to a dense scheme whatever the user asked for, so this is + reachable without them ever naming one. Which dense scheme it picks does + not matter here and is deliberately not asserted. Executable -- the real + check_critical_and_report runs, reading the same environment variable the + shipping code reads. + """ + import contextlib, io + mod = _driver_module() + args = ["--mode", "flowmc-phipsimarg", + "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "auto"] + saved = os.environ.get("JAX_ILE_DISTMARG_GH") + os.environ["JAX_ILE_DISTMARG_GH"] = "32" + try: + optp = mod.build_parser() + err = io.StringIO() + try: + with contextlib.redirect_stderr(err): + opts, _ = optp.parse_args(list(args)) + mod.check_critical_and_report(opts, optp) + except SystemExit: + msg = err.getvalue() + assert "JAX_ILE_DISTMARG_GH" in msg, msg[-400:] + assert "inert" in msg, msg[-400:] + else: + raise AssertionError( + "--distance-grid-scheme loguniform under JAX_ILE_DISTMARG_GH " + "must be refused at PARSE time, not deferred to the " + "constructor after a full precompute") + # ...and the identical command line must be ACCEPTED both with the + # variable unset AND with it explicitly OFF. "0" is the case that + # separates the shipped `int(...) > 0` from a truthiness test on the + # raw string: under `if os.environ.get(...)` the driver refuses while + # core._DISTMARG_GH_N is 0 and the constructor would accept, so the two + # seams disagree and the user is refused a combination that works. + # External review found exactly that mutation surviving. + for off in ("0", "00", None): + if off is None: + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + else: + os.environ["JAX_ILE_DISTMARG_GH"] = off + from RIFT.likelihood.jax_ile import core as _C + assert _C._DISTMARG_GH_N == 0, ( + "precondition: the kernels must see GH as OFF for %r" % (off,)) + optp = mod.build_parser() + opts, _ = optp.parse_args(list(args)) + mod.check_critical_and_report(opts, optp) # must not raise + finally: + if saved is None: + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + else: + os.environ["JAX_ILE_DISTMARG_GH"] = saved + + +def test_sky_doubling_updates_the_unclipped_maximum_too(): + """F1's detector on the SKY-DOUBLING path, guarded at the source. + + estimate_angle_amplitude re-draws the sky when its split-half check says + the maximum is still growing. The CLIPPED maximum is updated there; if the + UNCLIPPED companion is not, the exteriority detector reads only the first + batch while its denominator keeps growing, so clip_excess falls BELOW 1 and + the F1 refusal disarms itself on exactly the events whose sky sampling was + too coarse to trust. + + BOTH a value pin and a source guard, because neither alone is enough -- + and an earlier draft of this test shipped only the source guard on the + strength of a claim that was too strong. + + The re-draw branch IS reachable: sweeping (data seed in 0..5) x (support + [1, 10000] Mpc) x (n_sky in 4, 8, 16, 32, 64) x (estimator seed in 0..3) -- + 120 combinations -- 19 enter it, across five of the six data seeds. The + fixture below is one that enters it AND sits on an exterior support. + + What a value pin CAN catch: any change that scales or replaces the + accumulated unclipped maximum (halving it, wrapping it, resetting it after + the loop, or reverting the CONSUMER to read the first batch's array). + Those all move clip_excess on this fixture, and external review found three + such mutations that the source guard alone missed -- including one that + restores the pre-fix defect bit-for-bit by touching a line the loop guard + never looks at. + + What a value pin CANNOT catch, which is why the source guard stays: simply + DELETING the loop update leaves clip_excess bit-identical here, because the + deterministic face-on/face-off extremes are appended to the FIRST batch and + are what attains the unclipped maximum, so the second batch contributes + nothing on every fixture available. That corruption is real but silent -- + a dataset whose unclipped maximum came from a second-batch draw would take + clip_excess BELOW 1 and disarm the F1 refusal, and nothing here bounds + that. + """ + import inspect + import textwrap + from RIFT.likelihood.jax_ile import anglemarg as AM + + # ---- 1. VALUE PIN, on a fixture that actually enters the re-draw ---- + data = _synth(scale=3.0, kappa_boost=4.0, seed=3) + xg, _ = make_distance_grid(500.0, 10000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + import contextlib, io + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + _, diag = AM.estimate_angle_amplitude( + data, xg, interp="sinc", n_sky=64, seed=1, return_diagnostics=True) + assert "doubling" in buf.getvalue(), ( + "this fixture no longer enters the sky re-draw branch, so the pin " + "below no longer exercises it; find another (see the docstring sweep)") + # RELATIVE, not exact. These come through BLAS-heavy reconstruction, and + # float64 is not bit-portable across CPUs: CI measured amp_clipped + # 21.79506318041593 against 21.795063180415923 here, a 3.2e-16 relative + # difference that failed an == pin. The tolerance is chosen from BOTH + # sides and must stay there: ~1e-16 of platform drift below it, and the + # mutations it exists to catch far above it -- halving the accumulator + # (5e-1), a stray rescale (1e-3), and a wrapper that multiplies by + # 1.0000001 (1e-7). 1e-11 sits five orders above the drift and four + # below the tightest mutation. Do NOT loosen it past 1e-8. + import math + _RTOL = 1e-11 + for key, want in (("amp_clipped", 21.795063180415923), + ("amp_unclipped", 52.868630517667135), + ("clip_excess", 2.4257158641858192)): + assert math.isclose(diag[key], want, rel_tol=_RTOL), ( + "%s on the re-draw fixture is %.17g, expected %.17g (rel %.3g > " + "%.0e). Any rescaling, wrapping, post-loop reset, or reversion " + "of the CONSUMER to the first batch's array lands here; a drift " + "at the 1e-16 level instead means a new platform, and the " + "tolerance -- not the expectation -- is what to revisit." + % (key, diag[key], want, abs(diag[key] - want) / want, _RTOL)) + assert diag["clip_excess"] > 1.0 + 1e-3, "and it must still refuse" + + # ---- 2. SOURCE GUARD, for the one mutation a value pin cannot see ---- + tree = ast.parse(textwrap.dedent( + inspect.getsource(AM.estimate_angle_amplitude))) + loops = [n for n in ast.walk(tree) if isinstance(n, ast.While)] + assert len(loops) == 1, ( + "expected exactly one re-draw loop in estimate_angle_amplitude, found " + "%d; this guard names the loop by being the only one" % len(loops)) + assigned = {t.id for n in ast.walk(loops[0]) + if isinstance(n, ast.Assign) + for t in n.targets if isinstance(t, ast.Name)} + assert "amp_emp" in assigned, ( + "the re-draw loop no longer updates the clipped maximum; this guard " + "is anchored to that update and must be revisited") + assert "amp_u_emp" in assigned, ( + "the re-draw loop updates the CLIPPED maximum but not the UNCLIPPED " + "one. clip_excess = amp_unclipped / amp_clipped then reads the first " + "sky batch against a denominator that grew, falls below 1, and the " + "F1 exterior-peak refusal stops firing. Update both, adjacently.") + # and the two must be accumulated the same way, so that dropping one is + # visible on sight rather than only to this test + src = inspect.getsource(AM.estimate_angle_amplitude) + assert "amp_emp = max(amp_emp, float(amps2.max()))" in src + assert "amp_u_emp = max(amp_u_emp, float(amps_u2.max()))" in src + # ...and the CONSUMER must read the accumulator, not re-derive from the + # array. With the concatenate gone, `amps_u` holds batch 1 ONLY, so + # `amp_unclipped = np.max(amps_u)` is the pre-fix defect bit-for-bit while + # the loop guard above still passes. External review found exactly that. + fn = [n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef)][0] + reads = [n for n in ast.walk(fn) if isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "amp_unclipped" + for t in n.targets)] + assert len(reads) == 1 and isinstance(reads[0].value, ast.Name) \ + and reads[0].value.id == "amp_u_emp", ( + "amp_unclipped must be the accumulator amp_u_emp itself; re-deriving " + "it from amps_u reads the FIRST sky batch only and silently restores " + "the defect this test exists for") + # ...and nothing may reassign the accumulator AFTER the loop, which would + # discard the second batch just as effectively. + loop_line = loops[0].lineno + late = [n for n in ast.walk(fn) if isinstance(n, ast.Assign) + and n.lineno > loop_line + len(loops[0].body) + and any(isinstance(t, ast.Name) and t.id == "amp_u_emp" + for t in n.targets)] + assert not late, ( + "amp_u_emp is reassigned after the re-draw loop (line %s); that " + "discards the second batch exactly as dropping the in-loop update " + "would" % [n.lineno for n in late]) + + +def test_endpoint_term_matches_an_independent_quadrature_measurement(): + """The truncated-endpoint term, checked against a measurement rather than + against its own derivation. + + The alias law c(tol) is Poisson summation on an UNTRUNCATED Gaussian. A + finite prior stops the sum, and Euler-Maclaurin adds a term proportional to + the integrand's derivative at each retained edge -- which the node count + knows nothing about. Reproduced here in pure numpy on the scheme's own + quadrature (uniform in ln d, half-width end intervals) so the check is + independent of the module's formula: a peak ONE width inside an edge is + ~11% wrong at the shipped tol, while the same grid deep in the interior + delivers its promised ~1%. Both numbers below are measured. + """ + from RIFT.likelihood.jax_ile.core import (loguniform_endpoint_error, + ENDPOINT_ERROR_MARGIN) + + def measured(rho, k, h, d_min=1.0, d_max=10000.0, n_ref=400001): + u_pk = np.log(d_min) + k / rho + g = lambda u: np.exp(-0.5 * np.square((u - u_pk) * rho)) + n = int(np.ceil(np.log(d_max / d_min) / h)) + 1 + u = np.linspace(np.log(d_min), np.log(d_max), n) + w = np.full(n, u[1] - u[0]); w[0] *= 0.5; w[-1] *= 0.5 + ur = np.linspace(np.log(d_min), np.log(d_max), n_ref) + wr = np.full(n_ref, ur[1] - ur[0]); wr[0] *= 0.5; wr[-1] *= 0.5 + exact = float(np.sum(wr * g(ur))) + return abs(float(np.sum(w * g(u))) - exact) / exact + + rho = 55.06965048323435 + h = loguniform_spacing_for_tolerance(DIST_GRID_TOL_DEFAULT) / rho + near = measured(rho, 1.0, h) + deep = measured(rho, 30.0, h) + assert 0.08 < near < 0.15, ( + "a peak one width inside an edge should cost ~11%% at the shipped " + "tol; measured %.4g. If this moved, the endpoint term is no longer " + "what this guard thinks it is" % near) + assert deep < 2.0 * DIST_GRID_TOL_DEFAULT, ( + "deep in the interior the alias law alone should deliver ~tol; " + "measured %.4g" % deep) + assert near > 5.0 * deep, ( + "the endpoint term must dominate near an edge, or there is nothing " + "for the guard to catch (near %.4g, deep %.4g)" % (near, deep)) + # and the module's own estimate must be conservative there, not optimistic + est = ENDPOINT_ERROR_MARGIN * float(loguniform_endpoint_error( + h, rho ** 2 * 1.0 * np.exp(-0.5))) + assert est >= near, ( + "the margined estimate %.4g must not UNDER-read the measured %.4g; " + "the guard would then admit what it exists to refuse" % (est, near)) + + +def test_peak_too_close_to_an_edge_is_refused_even_though_it_is_interior(): + """The SECOND precondition. clip_excess reads exactly 1.0 here -- the + maximizing distance IS inside the support -- and the grid is still wrong, + because the alias law does not cover a truncated endpoint. Two regimes, + both refused, and the transition between them must be monotone in + clearance with no window that builds.""" + data = _synth(scale=10.0, kappa_boost=20.0) + seen = [] + for d_min, expect in ((1.0, "build"), (50.0, "build"), + (55.0, "refuse"), (58.0, "refuse"), (60.0, "refuse")): + try: + like = _like(data, "loguniform", n_grid=64, + d_min=d_min, d_max=10000.0) + seen.append((d_min, "build", int(like.x_grid.shape[0]))) + assert expect == "build", ( + "d_min=%g built with n=%d; expected a refusal" + % (d_min, like.x_grid.shape[0])) + except ValueError as exc: + seen.append((d_min, "refuse", str(exc)[:60])) + assert expect == "refuse", "d_min=%g refused unexpectedly: %s" % ( + d_min, exc) + # the interior-but-too-close case must cite the CLEARANCE, not exteriority + try: + _like(data, "loguniform", n_grid=64, d_min=55.0, d_max=10000.0) + except ValueError as exc: + assert "too close to an edge" in str(exc), str(exc)[:200] + assert "peak widths" in str(exc) + + +def test_a_peak_at_or_outside_an_edge_is_refused_before_clip_excess_trips(): + """The window between the two guards. _endpoint_bell clamps k <= 0 to + zero -- correct AT k = 0, where the endpoint is a stationary point and the + measured error is 1.2e-9, but it cannot tell that from k < 0, where the + peak has left the support and the error climbs. Scoring the worst case as + zero left a band that neither guard saw: measured, clearance -0.656 built + with a TRUE quadrature error of 0.0668, 6.7x the promised tol, while + clip_excess was still ~1. Found by external review of the endpoint guard. + """ + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=10.0, kappa_boost=20.0) + xg, _ = make_distance_grid(58.0, 10000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + _, diag = AM.estimate_angle_amplitude(data, xg, interp="sinc", + return_diagnostics=True) + assert diag["peak_clearance"] < 0.0, ( + "this fixture no longer puts the dominant peak outside the edge " + "(clearance %.4g), so it no longer exercises the window" + % diag["peak_clearance"]) + assert diag["clip_excess"] <= 1.0 + 1e-3, ( + "...and clip_excess must still read INTERIOR here (%.6g), or the " + "exterior guard already covers this and there is no window" + % diag["clip_excess"]) + try: + _like(data, "loguniform", n_grid=64, d_min=58.0, d_max=10000.0) + except ValueError as exc: + assert "AT or OUTSIDE a prior edge" in str(exc), str(exc)[:200] + else: + raise AssertionError( + "a dominant peak outside the support must be refused; scoring its " + "endpoint term as zero is what let this build") + + +def test_endpoint_scale_covers_subdominant_entries_not_just_the_sky_argmax(): + """P1. The endpoint term must cover EVERY reconstructed (phi, psi, time) + entry of a loud sky point, not only that point's dominant configuration. + + A sky direction can have a far-interior maximum and a near-equal secondary + entry whose distance peak sits a width from an edge. The secondary + contributes materially to the marginal and carries the full endpoint error; + the dominant one carries almost none. An argmax-only scalar therefore + reads the safe entry and reports the sky point as clean. + + This fixture is exactly that case: the DOMINANT peak is 4.77 widths inside + the support (its own contribution is rho^2 * bell(4.77) ~= 5.7e-3) and + clip_excess reads exactly 1.0, so neither the exterior guard nor the + dominant entry sees anything -- yet the true per-entry scale is ~44.5, a + factor ~7800 higher. Under the argmax-only form this assertion fails by + four orders of magnitude. + """ + from RIFT.likelihood.jax_ile import anglemarg as AM + data = _synth(scale=3.0, kappa_boost=4.0) + xg, _ = make_distance_grid(54.0, 10000.0, 64, "euclidean", + distMpcRef=data.distMpcRef) + _, diag = AM.estimate_angle_amplitude(data, xg, interp="sinc", + return_diagnostics=True) + # the premises: the dominant entry is far interior, and nothing else fires + assert diag["peak_clearance"] > 4.0, ( + "the dominant peak is no longer far from the edge (%.4g widths), so " + "this fixture no longer separates per-entry from argmax-only coverage" + % diag["peak_clearance"]) + assert diag["clip_excess"] <= 1.0 + 1e-3, ( + "the exterior guard fires here (%.6g), so it would catch this case " + "and the endpoint term is not what is being tested" + % diag["clip_excess"]) + # what the DOMINANT entry alone could contribute, from the reported numbers + k = diag["peak_clearance"] + dominant_only = (diag["peak_rho"] ** 2) * float(AM._endpoint_bell(k)) + assert diag["endpoint_scale"] > 100.0 * dominant_only, ( + "endpoint_scale %.6g is within 100x of what the sky point's DOMINANT " + "entry alone contributes (%.6g). That is the signature of an " + "argmax-only reduction: the sub-dominant entries one width from the " + "edge are not being counted." + % (diag["endpoint_scale"], dominant_only)) + assert diag["endpoint_scale"] > 1.0, ( + "endpoint_scale collapsed to %.6g on a fixture measured at ~44.5" + % diag["endpoint_scale"]) + + +def test_exterior_entries_stay_far_below_the_dominant_one(): + """P2's premise, pinned because the guard rests on it rather than on a + refusal. + + An angle configuration whose distance peak lies OUTSIDE the support + contributes a boundary-layer integral that the spacing contract does not + cover -- and no other diagnostic here can see it: clip_excess is a ratio of + GLOBAL maxima, peak_clearance describes the global argmax, and the endpoint + term clamps its bell at k <= 0. So the structural hole is real. + + It is not refused, because every threshold on the available quantity fails + on the fixtures: refusing when the exterior weight exp(-gap) exceeds tol + (gap < 4.6 nats) refuses every QUIET event, whose ENTIRE exponent range is + smaller than that -- _synth() spans 3.30 nats in total. Nothing is wrong + with those: at low amplitude the integrand is smooth and the grid + over-resolves it. + + What this pins is the premise the decision rests on: on the LOUD fixtures, + and on the priors this code is actually run with, the loudest exterior entry + sits far below the dominant one -- far outside the contribution band -- so + the hole is not reachable there. If that ever stops holding, this test + fails and the decision must be revisited rather than silently inherited. + """ + from RIFT.likelihood.jax_ile import anglemarg as AM + band = AM.ENDPOINT_GUARD_BAND + for scale, kb, d_min, d_max in ((3.0, 4.0, 1.0, 10000.0), + (3.0, 4.0, 1.0, 1000.0), + (10.0, 20.0, 1.0, 10000.0), + (10.0, 20.0, 1.0, 1000.0)): + data = _synth(scale=scale, kappa_boost=kb) + xg, _ = make_distance_grid(d_min, d_max, 64, "euclidean", + distMpcRef=data.distMpcRef) + _, diag = AM.estimate_angle_amplitude(data, xg, interp="sinc", + return_diagnostics=True) + assert diag["amp_clipped"] > 10.0 * band / 15.0, ( + "this fixture is no longer loud enough for the premise to be " + "meaningful (amp_clipped %.4g)" % diag["amp_clipped"]) + assert diag["exterior_gap"] > band, ( + "the loudest EXTERIOR entry is only %.4g nats below the dominant " + "one on _synth(%g, %g) over [%g, %g] -- inside the %.4g-nat " + "contribution band. P2's hole is now REACHABLE, and the decision " + "not to refuse on it (wrapper.py, beside peak_clearance) rests on " + "it not being. Revisit that." + % (diag["exterior_gap"], scale, kb, d_min, d_max, band)) + + +def test_dist_grid_tol_is_forwarded_and_not_hardcoded(): + """F3/N1. Hardcoding the module default at the call site leaves + --distance-grid-tol silently inert while dist_grid_info keeps echoing the + user's value -- indistinguishable from working.""" + data = _synth() + seen = {} + for tol in (5e-1, 1e-3): + like = _like(data, "loguniform", n_grid=64, dist_grid_tol=tol) + seen[tol] = int(like.x_grid.shape[0]) + assert like.dist_grid_info["tol"] == tol + assert seen[tol] == loguniform_grid_size( + D_MIN, D_MAX, like.dist_grid_info["rho_max"], tol), ( + "the grid was not built with the tol the caller asked for") + assert seen[5e-1] < seen[1e-3], ( + "a looser tolerance must produce FEWER nodes; the argument is inert") + + +def test_shipped_tolerance_constant_is_pinned_by_value(): + """F3/N8. Every other test derives its expectation from + DIST_GRID_TOL_DEFAULT, so changing the constant moved the whole suite with + it and pinned nothing. This asserts the shipped value and one node count + computed from it, both as literals.""" + assert DIST_GRID_TOL_DEFAULT == 1e-2 + assert abs(loguniform_spacing_for_tolerance(1e-2) - 1.930171443998096) < 1e-12 + # rho_max = 55.06965 is the reference configuration's measured value + assert loguniform_grid_size(1.0, 10000.0, 55.06965048323435, 1e-2) == 264 + + +def test_driver_distance_grid_tol_defaults_to_the_none_sentinel(): + """F3/N2. A concrete default makes `is not None` true on every run, which + fires the mode guard for every mode that does not implement the option.""" + tree = ast.parse(_driver_src()) + call = _option_call(tree, "--distance-grid-tol") + assert call is not None + kw = {k.arg: k.value for k in call.keywords} + assert isinstance(kw["default"], ast.Constant) and kw["default"].value is None, ( + "--distance-grid-tol must default to the None sentinel; a concrete " + "default is indistinguishable from the user having passed one") + + +def test_driver_never_passes_the_raw_option_as_the_node_count(): + """F3/N9. Reverting the sentinel refactor makes every default + flowmc-phipsimarg run die with `TypeError: 'NoneType' object cannot be + interpreted as an integer` -- a change to a LIVE default path. The + executable half of this is the subprocess test below; this is the precise + half, because an AST guard can name the call site.""" + tree = ast.parse(_driver_src()) + bad = [n for n in ast.walk(tree) if isinstance(n, ast.Call) + for k in (n.keywords or []) + if k.arg == "n_grid" and isinstance(k.value, ast.Attribute) + and k.value.attr == "distance_grid_points"] + assert not bad, ("n_grid= must receive the RESOLVED local, not the raw " + "optparse value, which is None unless the user passed it") + names = [k.value.id for n in ast.walk(tree) if isinstance(n, ast.Call) + for k in (n.keywords or []) + if k.arg == "n_grid" and isinstance(k.value, ast.Name)] + assert names and set(names) == {"n_dist_grid"}, ( + "every n_grid= call site must use the one resolved local; got %r" % (names,)) + + +def _run_driver(args, timeout=240): + import subprocess, tempfile + env = dict(os.environ, PYTHONPATH=str(_CODE), OMP_NUM_THREADS="1", + JAX_PLATFORMS="cpu", JAX_ENABLE_X64="1") + return subprocess.run([sys.executable, str(_CODE / "bin" + / "integrate_likelihood_extrinsic_jax")] + args, + capture_output=True, text=True, env=env, + cwd=tempfile.mkdtemp(), timeout=timeout) + + +def _driver_module(): + """Import the driver BY PATH (it has no .py suffix and is not importable + normally). Gives in-process access to the real build_parser and + check_critical_and_report, so the parse-time refusals below are executable + coverage of the shipping functions rather than five 8-second subprocesses.""" + import importlib.util, importlib.machinery + path = str(_CODE / "bin" / "integrate_likelihood_extrinsic_jax") + spec = importlib.util.spec_from_loader( + "_ilejax_under_test", importlib.machinery.SourceFileLoader( + "_ilejax_under_test", path)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_driver_refuses_the_bad_combinations_at_PARSE_time(): + """F8. These used to surface as a raw ValueError from the wrapper AFTER a + full precompute; they now fail during option validation. Executable: the + real check_critical_and_report is called, not grepped. (The only other + gated test that runs this driver runs --help, which exits before validation + happens at all.)""" + cases = [ + # The refusal is about the 'grid' scheme, which does not compute the + # sizing amplitude -- NOT about the option being unset. Since #225 the + # default is a dense scheme, so it must be named explicitly to be + # refused; the accepted-by-default case is asserted after this loop. + (["--mode", "flowmc-phipsimarg", "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "grid"], + "requires --angle-marg-scheme"), + (["--mode", "flowmc-phimarg", "--distance-grid-scheme", "loguniform"], + "applies only to --mode flowmc-phipsimarg"), + (["--mode", "flowmc-phipsimarg", "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "exact", "--distance-grid-points", "256"], + "both set the distance node count"), + (["--distance-grid-tol", "0.1"], "applies only to"), + (["--distance-grid-scheme", "adaptive"], "invalid choice"), + # the option's VALUE, not just its combinations: the valid range is a + # closed-form constant, so there is no reason to make the user sit + # through a precompute to be told 7.0 is not a fractional error. + (["--mode", "flowmc-phipsimarg", "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "exact", "--distance-grid-tol", "7.0"], + "--distance-grid-tol must be in (0, 2)"), + (["--mode", "flowmc-phipsimarg", "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "exact", "--distance-grid-tol", "-1"], + "--distance-grid-tol must be in (0, 2)"), + ] + import contextlib, io + mod = _driver_module() + for args, expect in cases: + optp = mod.build_parser() + err = io.StringIO() + try: + with contextlib.redirect_stderr(err): + opts, _ = optp.parse_args(list(args)) + mod.check_critical_and_report(opts, optp) + except SystemExit: + assert expect in err.getvalue(), ( + "%r: expected %r, got %r" % (args, expect, err.getvalue()[-400:])) + else: + raise AssertionError("%r must be refused, it was accepted" % (args,)) + # ...and the mirror: with no --angle-marg-scheme at all the driver must + # ACCEPT, because ANGLE_MARG_DEFAULT is a dense scheme. Both parse-time + # seams must assume the SAME default; one of them hardcoded 'grid' and so + # refused a command line the constructor would have run. + from RIFT.likelihood.jax_ile.anglemarg import ANGLE_MARG_DEFAULT + assert ANGLE_MARG_DEFAULT != "grid", ( + "the default is 'grid' again; this test and the driver's parse-time " + "fallback both key on it not being") + optp = mod.build_parser() + opts, _ = optp.parse_args(["--mode", "flowmc-phipsimarg", + "--distance-grid-scheme", "loguniform"]) + mod.check_critical_and_report(opts, optp) # must not raise + # ...and the getattr FALLBACK itself, which the parser can never exercise + # because optparse always sets the attribute. It is reached only by a + # caller that builds opts some other way -- and it silently disagreed with + # analyze_one's fallback until this was pinned, so a mutation back to a + # hardcoded "grid" SURVIVED the whole file. + class _Bare: + mode = "flowmc-phipsimarg" + distance_grid_scheme = "loguniform" + distance_grid_points = None + distance_grid_tol = None + bare = _Bare() + assert not hasattr(bare, "angle_marg_scheme") + mod.check_critical_and_report(bare, optp) # must not raise either + + +def test_driver_reaches_and_uses_the_resolved_node_count_on_a_real_input(): + """F3/N9, EXECUTABLE. Runs the driver past parsing on a real (tiny) + injection, far enough to print the resolved distance grid, then stops on a + known validation error. ~12 s. Under the reverted sentinel this dies with + TypeError at that very print instead -- which is exactly the live-default + breakage no other gated test could see.""" + p = _run_driver([ + "--inj-mode", "--mass1", "35", "--mass2", "30", "--inj-deltaF", "0.25", + "--inj-ra", "1.2", "--inj-dec", "0.3", "--inj-psi", "0.5", + "--inj-incl", "1.05", "--inj-phiref", "0.0", "--inj-distance", "633.92", + "--inj-detectors", "H1,L1", "--fmin-template", "40", "--fmax", "400.0", + "--l-max", "2", "--approximant", "SEOBNRv4", "--reference-freq", "100.0", + "--srate", "1024", "--d-min", "1", "--d-max", "10000", + "--distance-marginalization", "--mode", "flowmc-phipsimarg", + "--angle-marg-scheme", "grid", "--n-phi", "4", "--n-psi", "4", + "--time-marginalization-quadrature", "bandlimited", + "--n-max", "1", "--n-chunk", "1"]) + out = p.stdout + p.stderr + assert "TypeError" not in out, ( + "the driver did not resolve --distance-grid-points: %s" % out[-500:]) + assert "(grid=256," in out, ( + "the resolved default node count did not reach the run log: %s" % out[-500:]) + assert "time_quadrature='bandlimited' is not valid" in out, ( + "expected the run to stop on the known validation error; got %s" % out[-500:]) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_flow_reuse_default.py b/MonteCarloMarginalizeCode/Code/test/jax/test_flow_reuse_default.py new file mode 100644 index 000000000..891776d55 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_flow_reuse_default.py @@ -0,0 +1,95 @@ +"""Flow re-use is OFF by default, and re-use is still reachable. + +WHY THIS EXISTS. `--no-flow-reuse` used to be `store_true, default=False`, i.e. +flow re-use ON unless disabled. Measured across an 8-event batch at two seeds, +re-use contracts the extrinsic posterior monotonically in slot index -- psi to +~40% of its no-re-use width by slot 7, on BOTH seeds -- while slot 0, where no +re-use has yet happened, sits at ~1.0 as a built-in control. It buys no +measurable wall time (1589 s mean with, 1567 s without; the seed-to-seed spread +is larger than the difference and its sign flips). So the default is now OFF. + +THE TRAP THIS PINS. A `store_true` flag cannot express its own negation: simply +setting `default=True` would have made `--no-flow-reuse` inert AND removed any +way to turn re-use back on, silently deleting a capability while appearing to +change only a default. Hence the paired `--flow-reuse` writing the same dest. +Both directions are asserted here, and so is last-one-wins, because a paired +flag whose order does not resolve is worse than no flag at all. + +Needs no lal, no GPU and no flowMC: the parser is driven directly. +""" +import importlib.machinery +import importlib.util +import os +import sys + +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +CODE = os.path.abspath(os.path.join(HERE, "..", "..")) +DRIVER = os.path.join(CODE, "bin", "integrate_likelihood_extrinsic_jax") +if CODE not in sys.path: + sys.path.insert(0, CODE) + + +def _driver(): + loader = importlib.machinery.SourceFileLoader("_ile_jax_reuse", DRIVER) + spec = importlib.util.spec_from_loader("_ile_jax_reuse", loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +drv = _driver() + + +def _parse(*args): + opts, _ = drv.build_parser().parse_args(list(args)) + return opts.no_flow_reuse + + +def test_default_is_NO_flow_reuse(): + """The whole point of the change.""" + assert _parse() is True + + +def test_flow_reuse_is_still_reachable(): + """Without this, flipping the default would have deleted a capability.""" + assert _parse("--flow-reuse") is False + + +def test_no_flow_reuse_still_accepted_for_compatibility(): + """Existing command lines and scripts pass it; it now restates the default.""" + assert _parse("--no-flow-reuse") is True + + +@pytest.mark.parametrize("args,expected", [ + (("--flow-reuse", "--no-flow-reuse"), True), + (("--no-flow-reuse", "--flow-reuse"), False), +]) +def test_last_flag_on_the_command_line_wins(args, expected): + """Both write the same dest; an unresolved order would be worse than nothing.""" + assert _parse(*args) is expected + + +def test_both_flags_share_one_dest(): + """Structural: two dests would let the two flags disagree silently.""" + p = drv.build_parser() + dests = {} + for o in p._get_all_options(): + for s in (o._long_opts or []): + if s in ("--flow-reuse", "--no-flow-reuse"): + dests[s] = o.dest + assert dests == {"--flow-reuse": "no_flow_reuse", + "--no-flow-reuse": "no_flow_reuse"}, dests + + +def test_the_batch_loop_still_honours_the_flag(): + """A default flip is worthless if the consumer stopped reading it. + + Pins the two call sites in analyze-the-batch: the flow state handed to the + next event, and whether the returned state is retained at all. + """ + with open(DRIVER) as f: + src = f.read() + assert "flow_state=(None if opts.no_flow_reuse else flow_state)" in src + assert "if not opts.no_flow_reuse and new_flow_state is not None:" in src diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_interp_choices.py b/MonteCarloMarginalizeCode/Code/test/jax/test_interp_choices.py new file mode 100644 index 000000000..a63400440 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_interp_choices.py @@ -0,0 +1,111 @@ +"""The ``--interp`` CLI must expose every gatherer the likelihood implements. + +`RIFT/likelihood/jax_ile/core.py` registers three arrival-time interpolation +stencils in ``_GATHERERS``: ``nearest``, ``linear`` and ``cubic``. The driver's +``--interp`` option listed only two, so ``cubic`` -- the stencil that mirrors the +production ``factored_likelihood`` one, and the only one `_gather_cubic`'s own +docstring recommends -- was unreachable from the command line and ``linear`` was +silently the only realistic choice. + +Nothing failed. The option parsed, the run completed, and the likelihood was +biased: linear undershoots the razor-sharp rholm peak by an amount that depends +on where the peak falls between samples, so it stamps a mass-dependent ripple +onto the surface and biases the recovered arrival time, hence the sky. On a +zero-spin BNS at network amplitude 23.8, selecting cubic instead recovered ++4.34 nats at the peak. + +The first test below is the durable one: it does not name ``cubic``, it asserts +that the CLI and the registry agree. Adding a fourth gatherer without exposing +it fails here rather than silently shipping an unreachable code path. +""" +import importlib.machinery +import importlib.util +import os + +import numpy as np +import pytest + +import RIFT.likelihood.jax_ile.core as core + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CODE = os.path.abspath(os.path.join(_HERE, os.pardir, os.pardir)) +_JAXDRIVER = os.path.join(_CODE, 'bin', 'integrate_likelihood_extrinsic_jax') + + +def _load_driver(): + """Import the driver by path. It guards its entry point with __main__, so + importing it defines build_parser() without running an analysis.""" + # The driver has NO .py extension, so spec_from_file_location cannot infer a + # loader and returns a spec with loader=None. Name the loader explicitly -- + # otherwise these tests skip, and a skipped test reads exactly like a passing + # one in the summary line. + assert os.path.exists(_JAXDRIVER), 'driver missing: %s' % _JAXDRIVER + loader = importlib.machinery.SourceFileLoader('_ile_jax_driver', _JAXDRIVER) + spec = importlib.util.spec_from_loader(loader.name, loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +def _interp_choices(): + parser = _load_driver().build_parser() + for opt in parser.option_list + [o for g in parser.option_groups for o in g.option_list]: + if '--interp' in (opt._long_opts or []): + return set(opt.choices or ()) + raise AssertionError('the driver no longer defines --interp') + + +def test_cli_exposes_every_registered_gatherer(): + """Every key of _GATHERERS must be selectable from --interp, and vice versa.""" + registered = set(core._GATHERERS) + exposed = _interp_choices() + assert registered, '_GATHERERS is empty; the registry moved' + assert exposed == registered, ( + 'CLI --interp choices %s disagree with the _GATHERERS registry %s. ' + 'A stencil that is implemented but not exposed is dead code the user ' + 'cannot reach; one that is exposed but not implemented is a KeyError at ' + 'runtime.' % (sorted(exposed), sorted(registered))) + + +def test_cubic_is_reachable_and_is_not_linear(): + """Guard the specific regression: cubic selectable, and a DIFFERENT stencil. + + Equality of choices alone would still pass if someone aliased cubic to the + linear implementation, so check that the two actually compute differently. + """ + assert 'cubic' in _interp_choices() + assert core._GATHERERS['cubic'] is not core._GATHERERS['linear'] + + # A smooth, band-limited column sampled at integers; interpolate off-sample. + n = 64 + idx = np.arange(n) + col = np.exp(-0.5 * ((idx - 31.7) / 2.5) ** 2) + pos = np.array([12.5, 20.25, 31.7, 44.75]) + + lin = np.asarray(core._GATHERERS['linear'](col, pos)) + cub = np.asarray(core._GATHERERS['cubic'](col, pos)) + assert not np.allclose(lin, cub), 'cubic returns the linear result' + + exact = np.exp(-0.5 * ((pos - 31.7) / 2.5) ** 2) + err_lin = np.abs(lin - exact).max() + err_cub = np.abs(cub - exact).max() + assert err_cub < err_lin, ( + 'cubic (%.3e) should beat linear (%.3e) on a smooth peak' % (err_cub, err_lin)) + + +def test_every_stencil_reproduces_the_samples_it_sits_on(): + """At integer positions every stencil must return the sample itself. + + This is what makes the comparison above meaningful: the stencils differ only + between samples, not at them, so a difference at integer positions would mean + an indexing bug rather than an interpolation choice. + """ + n = 48 + idx = np.arange(n) + col = np.sin(0.21 * idx) + 0.4 * np.cos(0.07 * idx) + # stay clear of the edges: the cubic stencil zero-extends outside the buffer + pos = np.arange(4, n - 4).astype(float) + for name, gather in sorted(core._GATHERERS.items()): + got = np.asarray(gather(col, pos)) + assert np.allclose(got, col[pos.astype(int)], atol=1e-10), ( + '%s does not reproduce the sample at integer positions' % name) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_fairdraw_export.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_fairdraw_export.py new file mode 100644 index 000000000..ba6450b88 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_fairdraw_export.py @@ -0,0 +1,828 @@ +"""The JAX ILE driver's --save-samples export must be a FAIR DRAW. + +Downstream tooling (and every consumer of ``*_samples.dat``) treats the exported +rows as equal-weight draws from the conditional extrinsic posterior: the file has +no weight column, exactly like production ILE's ``--fairdraw-extrinsic-output`` +export, which multinomial-resamples against ``w = L p / p_s`` inside +``RIFT/integrators/mcsampler.py::integrate`` before writing. + +Several JAX-driver modes produce samples that do NOT follow the posterior -- +``laplace-is`` (the DEFAULT mode) draws from a Gaussian proposal, ``prior-mc`` +draws from the prior, and the flowMC modes sample a tempered target when +``--adapt-weight-exponent != 1`` -- while computing the correcting importance +weights and, before this test's change, discarding them. + +The tests below drive the shipped ``write_samples`` (not a helper in isolation) +on a target whose exact posterior moments are known analytically, so a regression +that drops the reweighting again shows up as a wrong exported distribution. + +Run: + PYTHONPATH=<...>/Code python -m pytest -q test/jax/test_jax_fairdraw_export.py +""" + +import importlib.machinery +import importlib.util +import ast +import inspect +import textwrap +import os +import types + +import numpy as np +import pytest + + +CODE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +DRIVER = os.path.join(CODE_DIR, "bin", "integrate_likelihood_extrinsic_jax") + + +def load_driver(): + """Import the driver script (no .py suffix) as a module.""" + loader = importlib.machinery.SourceFileLoader("_ile_jax_driver", DRIVER) + spec = importlib.util.spec_from_loader("_ile_jax_driver", loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +drv = load_driver() + + +# --------------------------------------------------------------------------- +# A 4-D target with analytically known posterior moments. +# +# proposal q(x) = N(0, s_q^2 I) <- what the sampler actually drew +# prior p(x) = N(0, s_p^2 I) (flat enough over the support) +# likelihood L(x) = N(x; mu_L, s_L^2 I) +# +# so the posterior is Gaussian with +# var_post = 1 / (1/s_p^2 + 1/s_L^2), mean_post = var_post * mu_L / s_L^2 +# and ln w = ln L + ln p - ln q is the exact importance weight. +# --------------------------------------------------------------------------- +S_Q, S_P, S_L = 3.0, 5.0, 0.8 +MU_L = np.array([1.4, -0.9, 0.6, 2.0]) +NDIM = 4 +VAR_POST = 1.0 / (1.0 / S_P ** 2 + 1.0 / S_L ** 2) +MEAN_POST = VAR_POST * MU_L / S_L ** 2 +SD_POST = np.sqrt(VAR_POST) + + +def _logN(x, mu, s): + return (-0.5 * np.sum((x - mu) ** 2, axis=1) / s ** 2 + - x.shape[1] * np.log(s * np.sqrt(2 * np.pi))) + + +def make_cloud(n=400000, seed=7): + rng = np.random.default_rng(seed) + theta = rng.standard_normal((n, NDIM)) * S_Q + lnL = _logN(theta, MU_L, S_L) + logw = lnL + _logN(theta, np.zeros(NDIM), S_P) - _logN(theta, np.zeros(NDIM), S_Q) + return theta, lnL, logw + + +def fake_opts(tmpdir, **kw): + o = types.SimpleNamespace( + output_file=os.path.join(str(tmpdir), "OUT"), save_samples=True, + mode="flowmc-phimarg", seed=11, + fairdraw_extrinsic_output=False, fairdraw_extrinsic_output_n_max=5, + n_fairdraw_extrinsic_samples=None) + for k, v in kw.items(): + setattr(o, k, v) + return o + + +def read_export(opts, idx=0): + f = opts.output_file + "_" + str(idx) + "_samples.dat" + assert os.path.exists(f), "write_samples wrote nothing" + with open(f) as fh: + hdr = fh.readline() + return np.loadtxt(f), hdr + + +# write_samples' 4-D branch emits columns (ra, dec, incl, psi, lnL), i.e. +# theta columns 0, 1, 3, 2. Map back so we compare like with like. +COL_OF_THETA = {0: 0, 1: 1, 2: 3, 3: 2} + + +def test_export_is_a_fair_draw_of_the_posterior(tmp_path): + """With importance weights supplied, the exported rows must follow the + POSTERIOR -- not the proposal the sampler drew from.""" + theta, lnL, logw = make_cloud() + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + got, hdr = read_export(opts) + + # unchanged file format: no weight column, same header as before + assert got.shape[1] == 5 + assert hdr.split()[1:] == ["right_ascension", "declination", "inclination", + "psi", "loglikelihood"] + + for j in range(NDIM): + col = got[:, COL_OF_THETA[j]] + # MC tolerance from the fair draw's own ESS, generously padded + tol_mean = 6.0 * SD_POST / np.sqrt(1.0 / np.sum( + (np.exp(logw - logw.max()) / np.exp(logw - logw.max()).sum()) ** 2)) + assert abs(col.mean() - MEAN_POST[j]) < max(tol_mean, 0.02), ( + "coord %d exported mean %.4f, posterior mean %.4f " + "(proposal mean 0.0) -- the cloud was NOT reweighted" + % (j, col.mean(), MEAN_POST[j])) + assert 0.85 < col.std() / SD_POST < 1.15, ( + "coord %d exported sd %.4f vs posterior sd %.4f (proposal sd %.4f)" + % (j, col.std(), SD_POST, S_Q)) + + +def test_unreweighted_export_would_fail_the_above(tmp_path): + """Mutation control, scored the SAME way: passing logw=None (the pre-fix + behaviour -- write the raw sampler cloud) must NOT satisfy the assertions + above. If this ever passes, the test above proves nothing.""" + theta, lnL, _ = make_cloud() + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=None) + got, _ = read_export(opts) + ok = all(abs(got[:, COL_OF_THETA[j]].mean() - MEAN_POST[j]) < 0.02 + and 0.85 < got[:, COL_OF_THETA[j]].std() / SD_POST < 1.15 + for j in range(NDIM)) + assert not ok, ("the RAW proposal cloud passed the fair-draw assertions; " + "the test target is too weak to detect a dropped reweight") + + +def test_uniform_weights_are_a_no_op(tmp_path): + """A converged, untempered MCMC chain already targets the posterior and + reports uniform post_weight; the export must then be the chain itself, not a + resampled (duplicate-ridden) version of it.""" + rng = np.random.default_rng(5) + theta = MEAN_POST[None, :] + rng.standard_normal((5000, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + logw = np.log(np.ones(len(theta)) / len(theta)) + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + got, _ = read_export(opts) + assert len(got) == len(theta) + assert len(np.unique(got[:, 0])) == len(theta), \ + "uniform weights triggered a resample (duplicates in the export)" + + +def test_analyze_one_also_refuses_a_post_weight_length_mismatch(): + """The SECOND mismatch guard, in analyze_one, pinned structurally. + + Reaching it at runtime needs a full sampler, so this reads the source: the + guard must compare lengths and RAISE. Found by mutation testing -- deleting + this guard left all 33 tests green, because the write_samples-level test + cannot see the call site that feeds it. + """ + src = textwrap.dedent(inspect.getsource(drv.analyze_one)) + tree = ast.parse(src) + msgs = " ".join(c.value for n in ast.walk(tree) if isinstance(n, ast.Raise) + for c in ast.walk(n) + if isinstance(c, ast.Constant) and isinstance(c.value, str)) + assert "post_weight entries for" in msgs, ( + "analyze_one no longer refuses a post_weight/theta length mismatch; a " + "mismatch would silently export an unreweighted cloud") + # and it must not have gone back to silently dropping the weights + assert "if _pw is not None and len(_pw) == len(theta) else None" not in src, ( + "the silent-drop form is back") + + +def test_length_mismatch_between_weights_and_rows_FAILS_LOUDLY(tmp_path): + """A mismatched weight vector must refuse, not silently mislabel. + + The guard used to read `if logw is not None and len(logw) == len(theta)`, so + a mismatch fell through to the default note -- writing an UNREWEIGHTED + tempered cloud under a header claiming "not applicable (sampler targets the + posterior)". That is a false provenance line on the one product every + downstream consumer reads as posterior draws. + """ + theta, lnL, logw = make_cloud(n=2000) + opts = fake_opts(tmp_path) + with pytest.raises(RuntimeError, match="disagree in length"): + drv.write_samples(opts, 0, theta, lnL, with_distance=False, + logw=logw[:-1]) + assert not os.path.exists(opts.output_file + "_0_samples.dat"), \ + "refused the draw but still wrote a file" + + +def test_export_count_is_NOT_clamped_by_the_evidence_estimator(tmp_path): + """A valid equal-weight chain must not be truncated by an unrelated ESS. + + THE DEFECT (review, #180): the export count was clamped by + ``1.5*res["neff"]``. On every flowMC mode ``theta`` is the production chain + while ``res["neff"]`` comes from a SEPARATE estimator -- the moment-matched + Gaussian cloud built for the evidence, or the annealing ladder's minimum + inter-stage ESS. A perfectly good uniform-weight chain was therefore + truncated to a couple of rows whenever that unrelated proposal had low ESS. + + ILE's own clamp (mcsampler.integrate: min(n_extr, 1.5*eff_samp, 1.5*neff)) + is self-consistent because all three describe the SAME weight vector the + draw samples from; the port lost that. + + Nothing caught this: all 20 call sites in this file passed neff=inf/nan, so + the clamp was never exercised. This test asks for a count and requires it + to be honoured. + """ + n, want = 5000, 1000 + theta = np.random.default_rng(0).normal(size=(n, 4)) + lnL = np.zeros(n) + opts = fake_opts(tmp_path, n_fairdraw_extrinsic_samples=want) + # equal-weight chain -- exactly what a beta=1 flowMC run hands over + drv.write_samples(opts, 0, theta, lnL, with_distance=False, + logw=np.log(np.ones(n) / n)) + got, _ = read_export(opts) + assert len(got) == want, ( + "requested %d equal-weight rows, exported %d" % (want, len(got))) + + +def test_write_samples_cannot_be_handed_an_evidence_neff(): + """Structural: the parameter is gone, so the defect cannot be re-wired. + + A behavioural test alone would not stop someone re-adding `neff=` and a + clamp; pin the signature and the call site together. + """ + import inspect + params = list(inspect.signature(drv.write_samples).parameters) + assert "neff" not in params, params + assert "neff" not in list(inspect.signature(drv.fairdraw_size).parameters) + call = _write_samples_call() # the call inside analyze_one + kwnames = {k.arg for k in call.keywords if k.arg} + assert "neff" not in kwnames, ( + "analyze_one passes an evidence neff into the export path again") + + +def test_fairdraw_count_options_are_live(tmp_path): + """--n-fairdraw-extrinsic-samples / --fairdraw-extrinsic-output-n-max must + CHANGE the number of exported rows (a parsed-and-logged knob is not a live + one).""" + theta, lnL, logw = make_cloud(n=20000) + for kw, want in ((dict(n_fairdraw_extrinsic_samples=137), 137), + (dict(fairdraw_extrinsic_output=True, + fairdraw_extrinsic_output_n_max=9), 9)): + opts = fake_opts(tmp_path / str(want), **kw) + os.makedirs(str(tmp_path / str(want)), exist_ok=True) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + got, _ = read_export(opts) + assert len(got) == want, "requested %d fair draws, got %d" % (want, len(got)) + + +def test_ess_clamp_prevents_manufactured_draws(tmp_path): + """A low-ESS cloud must not be resampled up to its original length: that + writes a file that looks like N independent draws but holds ~ESS distinct + points. ILE clamps at 1.5*ESS; so must this. (Observed for real: --mode + laplace-is on a BNS gave ESS=97 out of 200000 samples.)""" + rng = np.random.default_rng(2) + n = 50000 + theta = rng.standard_normal((n, NDIM)) * 4.0 + # a deliberately terrible proposal -> a handful of points carry the weight + logw = _logN(theta, MU_L, 0.05) + w = np.exp(logw - logw.max()); w /= w.sum() + ess = 1.0 / np.sum(w ** 2) + assert ess < n / 100.0, "the test cloud is not actually low-ESS (ESS=%.1f)" % ess + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, _logN(theta, MU_L, 0.05), with_distance=False, + logw=logw) + got, _ = read_export(opts) + assert len(got) <= np.ceil(1.5 * ess), ( + "exported %d rows from an ESS=%.1f cloud (cap %d)" + % (len(got), ess, int(np.ceil(1.5 * ess)))) + assert len(got) < n + + +def test_tempered_flowmc_weights_are_not_uniform(): + """The flowMC modes sample L^inv_T; post_weight = L^(1-inv_T) is the + correction. Guard the invariant that a tempered run yields NON-uniform + weights, so silently dropping them is a real (not cosmetic) error.""" + lnL = np.linspace(-30.0, 30.0, 1000) + for inv_T, uniform_expected in ((1.0, True), (0.8, False), (0.5, False)): + lw = (1.0 - inv_T) * lnL + w = np.exp(lw - lw.max()); w /= w.sum() + assert np.allclose(w, w[0]) is uniform_expected, \ + "inv_T=%g: uniformity of post_weight is %s" % (inv_T, not uniform_expected) + # n_out < n, so BOTH paths return indices -- but for different reasons, + # and the distinction is what the export header must record: + # uniform -> subsample WITHOUT replacement (no duplicates) + # non-uniform -> resample WITH replacement against w + idx, note = drv.fairdraw_indices(np.log(w), np.random.default_rng(1)) + assert not note.startswith("FAILED"), note + # fairdraw_indices does REWEIGHTING only: uniform weights are a genuine + # no-op there, and the export count is the caller's job. + assert (idx is None) is uniform_expected + assert ("weights uniform" in note) is uniform_expected, note + assert "ESS=" in note, note + if not uniform_expected: + assert len(np.unique(idx)) < len(idx), "resampling must be WITH replacement" + + +def test_exported_lnL_belongs_to_its_own_row(tmp_path): + """The resample must carry theta and lnL through the SAME index. A version + that permutes one relative to the other writes a loglikelihood that does not + belong to the parameters on its row -- invisible to every distributional + check, because both marginals stay correct.""" + theta, lnL, logw = make_cloud(n=120000) + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + got, _ = read_export(opts) + th_out = np.empty((len(got), NDIM)) + for j in range(NDIM): + th_out[:, j] = got[:, COL_OF_THETA[j]] + recomputed = _logN(th_out, MU_L, S_L) # lnL implied by the row's theta + err = np.abs(recomputed - got[:, -1]) + assert np.max(err) < 1e-9, ( + "exported lnL does not match the exported theta on the same row " + "(max |dlnL| = %.4g, mean %.4g) -- theta/lnL pairing was broken" + % (np.max(err), np.mean(err))) + + +def test_exported_lnL_stays_paired_through_the_count_subsample(tmp_path): + """The pairing test above uses default opts, so it never enters the count + path -- and the count subsample is a SECOND place theta and lnL are indexed + together. Re-check it with a count requested.""" + theta, lnL, logw = make_cloud(n=120000) + opts = fake_opts(tmp_path, n_fairdraw_extrinsic_samples=311) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + got, _ = read_export(opts) + assert len(got) == 311 + th_out = np.empty((len(got), NDIM)) + for j in range(NDIM): + th_out[:, j] = got[:, COL_OF_THETA[j]] + err = np.abs(_logN(th_out, MU_L, S_L) - got[:, -1]) + assert np.max(err) < 1e-9, ( + "count subsample broke the theta/lnL pairing (max |dlnL| = %.4g)" + % np.max(err)) + + +def test_count_flags_are_inert_for_modes_reported_as_ignoring_them(tmp_path): + """Report and behaviour must agree. check_critical_and_report gates these + flags on _FAIRDRAW_MODES; applying them anyway under a NUTS mode printed + "IGNORED" and then wrote 5 rows instead of 300. --fairdraw-extrinsic-output + is in ILE_extr.sub, so that is a real production line losing 60x of its + export under a banner saying the flag did nothing.""" + rng = np.random.default_rng(21) + theta = MEAN_POST[None, :] + rng.standard_normal((300, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + for mode, expect in (("nuts", 300), ("multistart-nuts", 300), + ("nuts-phimarg", 300), ("flowmc-phimarg", 5)): + d = tmp_path / mode; os.makedirs(str(d), exist_ok=True) + opts = fake_opts(d, mode=mode, fairdraw_extrinsic_output=True, + fairdraw_extrinsic_output_n_max=5) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=None) + got, _ = read_export(opts) + assert len(got) == expect, ( + "--mode %s: wrote %d rows, expected %d (%s)" + % (mode, len(got), expect, + "count must be inert where it is reported ignored" + if expect == 300 else "count must act where it is reported implemented")) + + +def test_provenance_n_out_matches_the_file(tmp_path): + """n_out was computed before non-finite lnL were dropped, so the header + over-stated the file exactly when the likelihood misbehaved.""" + rng = np.random.default_rng(22) + n = 1000 + theta = MEAN_POST[None, :] + rng.standard_normal((n, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + lnL[rng.choice(n, size=37, replace=False)] = np.nan + for kw in ({}, dict(n_fairdraw_extrinsic_samples=137)): + d = tmp_path / str(len(kw)); os.makedirs(str(d), exist_ok=True) + opts = fake_opts(d, **kw) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, + logw=np.log(np.ones(n) / n)) + got, _ = read_export(opts) + with open(opts.output_file + "_0_samples.dat") as fh: + fh.readline(); prov = fh.readline() + n_out = int(prov.split("n_out=")[1].split()[0]) + assert n_out == len(got), ( + "header says n_out=%d, file holds %d rows" % (n_out, len(got))) + assert np.isfinite(got).all() + + +def test_every_path_reports_ess_and_n_in(tmp_path): + """F-E in full: the logw=None path reported neither ESS= nor n_in=, so the + self-describing header was blank on one of the three paths that write. (The + FAILED path writes nothing at all -- see + test_failed_fairdraw_writes_no_samples_file.)""" + rng = np.random.default_rng(23) + theta = MEAN_POST[None, :] + rng.standard_normal((500, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + cases = {"none": None, + "uniform": np.log(np.ones(500) / 500), + "weighted": _logN(theta, MU_L, 1.2)} + for name, lw in cases.items(): + d = tmp_path / name; os.makedirs(str(d), exist_ok=True) + opts = fake_opts(d) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=lw) + with open(opts.output_file + "_0_samples.dat") as fh: + fh.readline(); prov = fh.readline() + for field in ("ESS=", "n_in=", "n_out="): + assert field in prov, "%s path header lacks %s: %r" % (name, field, prov) + # and the uniform path must not pass its ROW COUNT off as an ESS + opts = fake_opts(tmp_path / "uniform") + with open(opts.output_file + "_0_samples.dat") as fh: + fh.readline(); prov = fh.readline() + assert "ESS=n/a" in prov, "uniform path reports a fabricated ESS: %r" % prov + + +def test_degenerate_weights_fail_loudly_not_silently(): + """Weights that cannot be normalized must be reported as FAILED, not + silently returned as 'uniform, nothing to do' -- otherwise the raw, + unreweighted cloud is written under a header promising a fair draw.""" + rng = np.random.default_rng(4) + for bad, why in ((np.full(5000, -np.inf), "all -inf"), + (np.where(np.arange(5000) == 0, 0.0, -np.inf), "one finite")): + idx, note = drv.fairdraw_indices(bad, rng) + assert idx is None + assert note.startswith("FAILED"), "%s reported as %r" % (why, note) + + +def test_failed_fairdraw_writes_no_samples_file(tmp_path): + """A FAILED fair draw must produce NO export. Writing the raw cloud with + 'FAILED' in the provenance line was still a non-posterior cloud sitting under + the standard weightless product name: consumers read `*_samples.dat` rows as + equal-weight posterior draws and are not obliged to parse the second header + line, so the contract was violated exactly on the collapsed integrations + where proposal and posterior differ most.""" + rng = np.random.default_rng(4) + theta = rng.standard_normal((5000, NDIM)) * 3.0 + opts = fake_opts(tmp_path) + with pytest.raises(RuntimeError) as exc: + drv.write_samples(opts, 0, theta, np.full(5000, -np.inf), + with_distance=False, logw=np.full(5000, -np.inf)) + assert "fair draw failed" in str(exc.value) + assert not os.path.exists(opts.output_file + "_0_samples.dat"), \ + "a non-posterior cloud was exported after the fair draw failed" + + +def test_failed_fairdraw_removes_a_stale_export(tmp_path): + """Refusing to write is not enough on a re-run: a samples file left from an + earlier run sits at exactly the path the pipeline reads for THIS one, so the + refusal must also clear it rather than silently endorsing stale rows.""" + rng = np.random.default_rng(9) + theta = rng.standard_normal((5000, NDIM)) * 3.0 + opts = fake_opts(tmp_path) + stale = opts.output_file + "_0_samples.dat" + with open(stale, "w") as fh: + fh.write("# right_ascension declination inclination psi loglikelihood\n" + "# mode=flowmc-phimarg fairdraw: reweighted ESS=900.0 n_in=1 n_out=1\n" + "0.1 0.2 0.3 0.4 -5.0\n") + with pytest.raises(RuntimeError): + drv.write_samples(opts, 0, theta, np.full(5000, -np.inf), + with_distance=False, logw=np.full(5000, -np.inf)) + assert not os.path.exists(stale), \ + "the previous run's export survived a failed fair draw" + + +def test_failed_event_is_skippable_but_never_exported(tmp_path): + """The refusal must be an ordinary Exception, so main's per-event handler + (--soft-fail-event-range) can carry a batch past a collapsed event, and it + must not disturb the other events' exports.""" + src = inspect.getsource(drv.main) + assert "except Exception" in src and "soft_fail_event_range" in src, \ + "main lost the per-event guard that makes a refused export skippable" + theta, lnL, logw = make_cloud(n=20000) + bad = fake_opts(tmp_path / "bad"); os.makedirs(str(tmp_path / "bad"), exist_ok=True) + with pytest.raises(RuntimeError): + drv.write_samples(bad, 0, theta, lnL, with_distance=False, + logw=np.full(len(theta), -np.inf)) + good = fake_opts(tmp_path / "good"); os.makedirs(str(tmp_path / "good"), exist_ok=True) + drv.write_samples(good, 1, theta, lnL, with_distance=False, logw=logw) + assert not os.path.exists(bad.output_file + "_0_samples.dat") + assert len(read_export(good, 1)[0]) > 1 + + +def test_weights_are_stabilized_at_realistic_lnL(tmp_path): + """Real extrinsic lnL runs to several hundred (this BNS peaked at 266; ILE + routinely sees >1000). exp(logw) without subtracting the max overflows to + inf there, which the fail-open path used to swallow. Exercise the range the + driver actually operates in, not the O(1) range of a toy target.""" + rng = np.random.default_rng(6) + n = 20000 + theta = rng.standard_normal((n, NDIM)) * 2.0 + logw = 800.0 + _logN(theta, MU_L, 1.5) # ~ +800, well past exp() overflow + assert logw.max() > 700.0 + idx, note = drv.fairdraw_indices(logw, rng) + assert idx is not None, "fair draw refused at realistic lnL: %s" % note + assert not note.startswith("FAILED"), note + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, logw, with_distance=False, logw=logw) + got, _ = read_export(opts) + assert len(got) > 1 and np.isfinite(got).all() + assert len(np.unique(got[:, 0])) > 1, "export collapsed to a single point" + + +def test_export_header_records_ess_and_mode(tmp_path): + """The artifact must be self-describing: export ESS was previously recorded + nowhere, so a 200000-sample file with ESS 97 looked like any other.""" + theta, lnL, logw = make_cloud(n=120000) + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) + with open(opts.output_file + "_0_samples.dat") as fh: + cols_line, prov_line = fh.readline(), fh.readline() + assert cols_line.split()[1] == "right_ascension", "column line moved: %r" % cols_line + assert "ESS=" in prov_line and "mode=" in prov_line, prov_line + + +def test_export_rng_is_independent_of_the_science_stream(tmp_path): + """--save-samples is an OUTPUT flag and must not move any number. The + export draw is keyed by (seed, out_index), so it is reproducible no matter + what else has consumed randomness, and it cannot perturb the sampler's + stream.""" + theta, lnL, logw = make_cloud(n=20000) + outs = [] + for burn in (0, 10000): + shared = np.random.default_rng(fake_opts(tmp_path).seed) + shared.standard_normal(burn) # unrelated consumption + d = tmp_path / ("burn%d" % burn) + os.makedirs(str(d), exist_ok=True) + opts = fake_opts(d) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, logw=logw) # export rng derived from (seed, out_index) + outs.append(read_export(opts)[0]) + assert np.array_equal(outs[0], outs[1]), \ + "export depends on how much the shared RNG was consumed" + # and different events must not reuse the same draw + o2 = fake_opts(tmp_path / "ev1"); os.makedirs(str(tmp_path / "ev1"), exist_ok=True) + drv.write_samples(o2, 1, theta, lnL, with_distance=False, logw=logw) + assert not np.array_equal(read_export(o2, 1)[0], outs[0]) + + +def test_mode_sets_exclude_non_importance_weights(): + """multistart-nuts / nuts-phimarg report post_weight as a per-chain Laplace + MODE-EVIDENCE weight (samplers.py: np.full(n_per[k], mass[k]/n_per[k])), not + L*p/p_s. They must not be fair-drawn against it.""" + assert "multistart-nuts" not in drv._TEMPERED_MODES + assert "nuts-phimarg" not in drv._TEMPERED_MODES + assert "multistart-nuts" not in drv._FAIRDRAW_MODES + assert "nuts-phimarg" not in drv._FAIRDRAW_MODES + for m in ("flowmc", "flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg"): + assert m in drv._TEMPERED_MODES and m in drv._FAIRDRAW_MODES + for m in ("prior-mc", "laplace-is"): + assert m in drv._FAIRDRAW_MODES and m not in drv._TEMPERED_MODES + + +def _write_samples_call(): + """The ast.Call node for write_samples(...) inside analyze_one.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(drv.analyze_one))) + calls = [n for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == "write_samples"] + assert len(calls) == 1, "expected exactly one write_samples call, found %d" % len(calls) + return calls[0] + + +def test_write_samples_never_receives_the_science_generator(): + """STRUCTURAL guard for F2, checked by AST rather than by substring. + + The export RNG is derived inside write_samples from (seed, out_index). If a + caller can hand it the generator that feeds the samplers, --save-samples -- + an OUTPUT flag -- moves the science again. + + A name-based guard is not enough: `write_samples(..., generator=rng)` defeats + both a signature test that looks for the literal "rng" and a source check for + the substring "rng=" (that text contains "=rng"). So instead: no argument + expression of the call may be the bare Name `rng`, whatever keyword it wears, + and the callee must not name a Generator-ish parameter at all.""" + call = _write_samples_call() + args = list(call.args) + [k.value for k in call.keywords] + for a in args: + assert not (isinstance(a, ast.Name) and a.id == "rng"), ( + "analyze_one passes the shared `rng` to write_samples (as %s)" + % (next((k.arg for k in call.keywords if k.value is a), "positional"))) + # `opts.rng`-style smuggling: any attribute access ending in .rng + assert not (isinstance(a, ast.Attribute) and a.attr == "rng"), \ + "analyze_one smuggles an rng in via an attribute" + params = list(inspect.signature(drv.write_samples).parameters) + for bad in ("rng", "generator", "random_state", "prng", "bitgen"): + assert bad not in params, ( + "write_samples grew a %r parameter -- a caller can now pass the " + "science generator (%s)" % (bad, params)) + + +def test_post_weight_is_gated_on_tempered_modes_at_the_call_site(): + """F1 lives at the CALL SITE, not in the frozensets. Asserting the sets are + correct cannot see the guard being deleted from analyze_one -- the same + 'test the helper, not the wiring' defect fixed for the export RNG. + + Require that the expression which reads res["post_weight"] is guarded by a + test mentioning _TEMPERED_MODES.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(drv.analyze_one))) + guarded = [] + for node in ast.walk(tree): + if not isinstance(node, ast.IfExp): + continue + body = ast.dump(node.body) + ast.dump(node.orelse) + if "post_weight" in body: + guarded.append("_TEMPERED_MODES" in ast.dump(node.test)) + assert guarded, ("no conditional expression reads post_weight in analyze_one " + "-- the F1 guard was removed or restructured") + assert all(guarded), ("post_weight is read without a _TEMPERED_MODES guard: " + "multistart-nuts / nuts-phimarg would be fair-drawn " + "against a per-chain Laplace mode-evidence weight") + + +def test_count_option_dests_are_stable(): + """fairdraw_size reads the options through getattr(..., None), which FAILS + OPEN: rename an option's dest and the count silently stops being applied + while everything still passes. Drive the real parser and pin the dests.""" + optp = drv.build_parser() + dests = {o.dest for o in optp._get_all_options() if o.dest} + for d in ("n_fairdraw_extrinsic_samples", "fairdraw_extrinsic_output", + "fairdraw_extrinsic_output_n_max", "mode", "seed", "save_samples"): + assert d in dests, "option dest %r vanished -- fairdraw_size fails open" % d + opts, _ = optp.parse_args(["--n-fairdraw-extrinsic-samples", "137"]) + assert opts.n_fairdraw_extrinsic_samples == 137 + opts2, _ = optp.parse_args(["--fairdraw-extrinsic-output"]) + assert opts2.fairdraw_extrinsic_output is True + # unset -n-max must stay None so the ignored-option report does not claim + # the user passed it; the ILE default of 5 is resolved downstream + assert opts2.fairdraw_extrinsic_output_n_max is None + assert drv.fairdraw_size(opts2, 10000) == drv._FAIRDRAW_N_MAX_DEFAULT + + +def test_count_options_act_when_weights_are_uniform(tmp_path): + """THE default configuration has uniform weights: --adapt-weight-exponent is + 1.0, so the flowMC modes report post_weight uniform and there is nothing to + reweight. The count options are a COUNT contract, not a reweight contract + -- ILE applies the count regardless -- so they must still bound the export. + Returning early on uniform weights made them a silent no-op in exactly the + configuration people actually run.""" + rng = np.random.default_rng(11) + theta = MEAN_POST[None, :] + rng.standard_normal((3200, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + uniform = np.log(np.ones(len(theta)) / len(theta)) + for kw, want in ((dict(n_fairdraw_extrinsic_samples=137), 137), + (dict(fairdraw_extrinsic_output=True, + fairdraw_extrinsic_output_n_max=5), 5)): + d = tmp_path / str(want); os.makedirs(str(d), exist_ok=True) + opts = fake_opts(d, **kw) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, + logw=uniform) + got, _ = read_export(opts) + assert len(got) == want, ( + "uniform weights: asked for %d rows, wrote %d -- the count contract " + "was skipped" % (want, len(got))) + # equal weights -> subsample WITHOUT replacement, so no duplicates + assert len(np.unique(got[:, 0])) == want, "uniform subsample duplicated rows" + + +def test_uniform_export_header_still_reports_ess(tmp_path): + """The self-describing header must not go blank on the uniform path -- that + is where the flowMC modes live, and the README tells users to check the ESS + before trusting a file.""" + rng = np.random.default_rng(12) + theta = MEAN_POST[None, :] + rng.standard_normal((2000, NDIM)) * SD_POST + lnL = _logN(theta, MU_L, S_L) + opts = fake_opts(tmp_path) + drv.write_samples(opts, 0, theta, lnL, with_distance=False, + logw=np.log(np.ones(len(theta)) / len(theta))) + with open(opts.output_file + "_0_samples.dat") as fh: + fh.readline(); prov = fh.readline() + assert "ESS=" in prov and "n_in=" in prov and "n_out=" in prov, prov + + +def test_failed_export_leaves_no_result_row(tmp_path): + """A refused fair draw means the integration collapsed, so the event must + leave no NORMAL ILE result behind either. `__.dat` carries no + hint that the export was refused, and with --soft-fail-event-range the batch + continues -- so a result row at that path is collected as a successful + integration (in the one-finite-weight case, a finite but collapsed + evidence).""" + rng = np.random.default_rng(31) + theta = rng.standard_normal((5000, NDIM)) * 3.0 + opts = fake_opts(tmp_path) + stale = drv.dat_path(opts, 0) + with open(stale, "w") as fh: + fh.write("# event_id m1 m2 s1x s1y s1z s2x s2y s2z lnL sigma_lnL ntotal neff\n" + "-1 1.4 1.3 0 0 0 0 0 0 12.0 0.1 1000 900\n") + with pytest.raises(RuntimeError): + drv.write_samples(opts, 0, theta, np.full(5000, -np.inf), + with_distance=False, logw=np.full(5000, -np.inf)) + assert not os.path.exists(stale), \ + "a normal ILE result row survived a refused fair draw" + + +def test_result_row_is_written_only_after_the_export(): + """STRUCTURAL guard on the write ORDER in analyze_one. + + write_dat has no way to know the export will be refused, so publishing it + first is not fixable inside write_samples: the artifact already exists when + the RuntimeError is raised, and --soft-fail-event-range then walks past it. + The export must therefore be validated and written FIRST. Checked at the + call site, since a test of write_samples alone cannot see the order.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(drv.analyze_one))) + # Only the function's OWN statement list: --mode map writes its result inside + # an `if` that returns immediately, exports nothing, and is not at issue. + order = [stmt.value.func.id + for stmt in tree.body[0].body + if isinstance(stmt, ast.Expr) + and isinstance(stmt.value, ast.Call) + and isinstance(stmt.value.func, ast.Name) + and stmt.value.func.id in ("write_dat", "write_samples")] + assert order.count("write_samples") == 1 and order.count("write_dat") == 1, \ + "expected one write_samples and one write_dat in analyze_one's tail: %s" % order + assert order.index("write_samples") < order.index("write_dat"), \ + ("analyze_one publishes the .dat result before validating the export: a " + "refused fair draw would leave a normal ILE result for a failed event") + + +def _smc_branch_body(): + """Statements of the ``--smc-puffball`` branch of analyze_one.""" + tree = ast.parse(textwrap.dedent(inspect.getsource(drv.analyze_one))) + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + for stmt in node.body: + if any(isinstance(c, ast.Call) and isinstance(c.func, ast.Attribute) + and c.func.attr == "smc_puffball_sample" + for c in ast.walk(stmt)): + return node.body + raise AssertionError("analyze_one no longer calls smc_puffball_sample") + + +def test_unfinished_smc_ladder_is_refused_at_the_call_site(): + """An SMC ladder that stops at max_stages leaves a cloud targeting + L**inv_T * prior. Its weights are then the L**(1-inv_T) correction, but the + export contract here is an EQUAL-WEIGHT file and the `.dat` would carry + log Z(inv_T) -- so the event must be refused outright, in the branch itself, + before either artifact is written. Structural, like the other call-site + guards: a check inside smc_puffball_sample cannot see what the driver does + with the result, and neither can a test of write_samples alone.""" + body = _smc_branch_body() + guards = [n for stmt in body for n in ast.walk(stmt) + if isinstance(n, ast.If) and "inv_T" in ast.dump(n.test) + and any(isinstance(r, ast.Raise) for r in ast.walk(n))] + assert guards, ("the --smc-puffball branch does not check the tempering " + "exponent the ladder reached and raise: an unfinished ladder " + "would be exported as a fair draw of the posterior") + + +class _SharpLike: + """4-D angular likelihood sharp enough that one SMC stage cannot finish the + ladder. samplers.eval_lnL_4 needs only ``log_likelihood`` on arrays.""" + + ANGULAR_PARAM_ORDER = ("ra", "dec", "psi", "incl") + + def log_likelihood(self, ra, dec, psi, incl): + d2 = (ra - 1.0) ** 2 + dec ** 2 + (psi - 1.0) ** 2 + (incl - 1.0) ** 2 + return -0.5 * d2 / 0.02 ** 2 + + +class _FlatLike(_SharpLike): + """Constant lnL: every tempering step takes the full max_dbeta, so the ladder + reaches inv_T == 1 in a handful of stages.""" + + def log_likelihood(self, ra, dec, psi, incl): + return np.zeros_like(np.asarray(ra, dtype=float)) + + +def test_smc_reports_the_temperature_it_actually_reached(): + """The sampler must not claim temper=1 with uniform post_weight when it + stopped short: that is what let the driver publish a tempered cloud.""" + from RIFT.likelihood.jax_ile import samplers + res = samplers.smc_puffball_sample(_SharpLike(), 1.0, 1000.0, n_walkers=200, + n_move=1, max_stages=1, is_evidence=False, + seed=3) + inv_T = float(res["inv_T"]) + assert 0.0 < inv_T < 1.0, "one stage on a sharp target should not finish" + pw = np.asarray(res["post_weight"], dtype=float) + lw = (1.0 - inv_T) * np.asarray(res["lnL"], dtype=float) + want = np.exp(lw - lw.max()) + want /= want.sum() + assert np.allclose(pw, want), \ + "post_weight is not the L**(1-inv_T) correction for the exponent reached" + assert not np.allclose(pw, pw[0]), "an unfinished ladder reported uniform weights" + + +def test_smc_that_finishes_the_ladder_reports_inv_T_one(): + """The converged path is unchanged: inv_T == 1 and uniform weights.""" + from RIFT.likelihood.jax_ile import samplers + res = samplers.smc_puffball_sample(_FlatLike(), 1.0, 1000.0, n_walkers=64, + n_move=1, max_stages=20, max_dbeta=0.25, + is_evidence=False, seed=5) + assert float(res["inv_T"]) == 1.0 + pw = np.asarray(res["post_weight"], dtype=float) + assert np.allclose(pw, pw[0]), "a finished ladder needs no correction weight" + + +class _RazorLike(_SharpLike): + """Dynamic range so wide that even the smallest permitted tempering step + collapses the ESS, so the rung search always falls back on its ``db`` floor.""" + + def log_likelihood(self, ra, dec, psi, incl): + d2 = (ra - 1.0) ** 2 + dec ** 2 + (psi - 1.0) ** 2 + (incl - 1.0) ** 2 + return -0.5 * d2 / 1e-4 ** 2 + + +def test_smc_rung_never_steps_past_its_cap(): + """The floor that keeps a stalled rung search moving must not push the step + past the cap (max_dbeta, or the distance left to inv_T == 1). A step beyond + it resamples and moves the cloud at inv_T > 1, and clipping the reported + exponent back to 1 would then hand the caller that OVER-tempered cloud with + uniform post_weight -- a tempered draw labelled as the posterior.""" + from RIFT.likelihood.jax_ile import samplers + max_dbeta, max_stages = 1e-5, 4 + res = samplers.smc_puffball_sample(_RazorLike(), 1.0, 1000.0, n_walkers=64, + n_move=1, max_stages=max_stages, + max_dbeta=max_dbeta, is_evidence=False, + seed=7) + assert float(res["inv_T"]) <= max_stages * max_dbeta + 1e-12, \ + "a tempering rung stepped past its cap: the exported cloud is over-tempered" + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-q"])) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py new file mode 100644 index 000000000..da83d882f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py @@ -0,0 +1,549 @@ +"""The JAX 'sinc' gatherer is the same stencil the numpy/cupy/CUDA paths use. + +WHY THIS FILE EXISTS. RIFT now has the 2a-tap Lanczos stencil on four backends. Three of them +(numpy, cupy, CUDA) consume ONE weight array from +``factored_likelihood._sinc_lanczos_weight_matrix``, so they cannot drift. JAX cannot share it: +the weights depend on the traced sub-sample offset, so ``jax_ile.core._sinc_lanczos_weights_jax`` +is a second, independent expression of the same formula. That duplication is the long-term risk +in this feature, and these tests are what converts it from "trust the reviewer" into "CI fails". + +Four things are checked: + +(a) WEIGHT PARITY -- the JAX and numpy weight generators agree elementwise, including at u = 0 + (where the |x| >= a hard zero bites) and at u -> 1. + +(b) GATHER PARITY -- the assembled JAX window equals the numpy ``_sinc_Q_window_numpy`` window on + the same buffer, INCLUDING near the buffer edges, where the shared convention is that + out-of-buffer taps are dropped WITHOUT renormalising the remaining weights. A backend that + renormalised after masking would pass (a) and fail here. + +(c) ACCURACY, i.e. the regression that would have caught the defect this work was opened for. + Against an exactly-known band-limited signal, 'sinc' must beat 'cubic' by a wide margin in a + regime where Q is poorly oversampled. Swap 'sinc' back to 'cubic' in _GATHERERS and this + test fails -- verified by mutation, see test_mutation_cubic_fails_accuracy_gate. + +(d) VECTORISED == UNROLLED -- the tap axis must stay an array. Written as a Python loop the + stencil compiles for >1 h inside a NUTS trace. Bit-equivalence is asserted so that + "simplifying" it back into a loop is a test failure and not a silent 4-orders-of-magnitude + compile regression. + +The GPU (cupy) leg SKIPS when cupy is unavailable -- it must not silently pass. + +Run: + PYTHONPATH=<...>/Code python -m pytest -q test/jax/test_jax_stencil_parity.py +""" + +import numpy as np +import pytest + +import jax +jax.config.update("jax_enable_x64", True) # else the parity tolerances below are meaningless +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import core as JC +from RIFT.likelihood.time_interp_choice import SINC_HALFWIDTH_DEFAULT + +A = SINC_HALFWIDTH_DEFAULT + +# Parity tolerances. These are STATED, not implicit (the brief for this work asked for exactly +# that). Both generators evaluate the same closed-form expression in float64, differing only in +# the order of the sinc/normalise operations and in whose libm supplies sin(), so the expected +# disagreement is a few ulp of the O(1) weights. 1e-14 is ~50 ulp: tight enough that a genuine +# formula difference (a missing window, an unnormalised sum, an off-by-one tap) cannot hide under +# it, loose enough not to be a libm-version tripwire. +TOL_WEIGHTS = 1e-14 +# The gather additionally sums 2a products against O(1) complex samples, so it carries the +# weight error plus ~2a rounding steps; 1e-12 relative to unit-scale data covers that with margin. +TOL_GATHER = 1e-12 + + +def _numpy_weight_matrix(u, a=A): + """The shared CPU/GPU generator. Imported lazily: factored_likelihood costs numba+lal.""" + from RIFT.likelihood.factored_likelihood import _sinc_lanczos_weight_matrix + return _sinc_lanczos_weight_matrix(np.atleast_1d(np.asarray(u, dtype=float)), a) + + +# ----------------------------------------------------------------------------- (a) weight parity + +def test_weight_parity_against_numpy_generator(): + # Endpoints included on purpose: u = 0 is where the |x| >= a hard zero applies (tap k = a + # sits exactly at x = -a) and where the stencil must collapse to the identity. + u = np.concatenate([[0.0, 1e-15, 0.5, 1.0 - 1e-12], + np.random.default_rng(0).uniform(0.0, 1.0, 97)]) + k_np, w_np = _numpy_weight_matrix(u) + k_jx, w_jx = JC._sinc_lanczos_weights_jax(jnp.asarray(u), A) + np.testing.assert_array_equal(np.asarray(k_jx), np.asarray(k_np)) + assert w_np.shape == (len(u), 2 * A) + err = np.max(np.abs(np.asarray(w_jx) - w_np)) + assert err < TOL_WEIGHTS, "JAX/numpy sinc weights disagree by %.3e" % err + + +def test_weight_parity_outside_the_unit_interval(): + """The ``|x| >= a`` hard zero is only reachable for u outside [0, 1), and it must still match. + + Added after a mutation run: deleting that clause from the JAX generator survived every other + test here. It is an EQUIVALENT mutation on the wired path -- for u in [0, 1) the guarded tap + sits at x = -a, where sinc(x/a) = sinc(-1) = 3.9e-17 rather than exactly 0, so the guard moves + the weight by 1.5e-33 -- but the two generators are library helpers, and off the wired path + the clause is worth 2.2e-3 in the weights. So the choice is to pin it, not to call it dead. + """ + u = np.array([-2.0, -0.5, -1e-12, 1.0, 1.5, 2.0, 3.25]) + k_np, w_np = _numpy_weight_matrix(u) + k_jx, w_jx = JC._sinc_lanczos_weights_jax(jnp.asarray(u), A) + np.testing.assert_array_equal(np.asarray(k_jx), np.asarray(k_np)) + err = np.max(np.abs(np.asarray(w_jx) - w_np)) + assert err < TOL_WEIGHTS, "out-of-range-u weights disagree by %.3e" % err + # The guard must actually be exercised, or this test proves nothing about it. + x = u[:, None] - np.asarray(k_np) + assert np.any(np.abs(x) >= A), "no tap reached |x| >= a; the guard is untested" + + +def test_weights_sum_to_one_and_are_identity_at_zero_offset(): + """Both properties are relied on downstream: unit sum makes constants exact, and the u=0 + identity is what lets 'sinc' reproduce the original samples where no shift is needed.""" + _, w = JC._sinc_lanczos_weights_jax(jnp.linspace(0.0, 1.0, 41), A) + np.testing.assert_allclose(np.asarray(jnp.sum(w, axis=-1)), 1.0, atol=1e-14) + _, w0 = JC._sinc_lanczos_weights_jax(jnp.zeros(1), A) + expect = np.zeros(2 * A) + expect[A - 1] = 1.0 # k = 0 is at index a-1 in arange(-a+1, a+1) + np.testing.assert_allclose(np.asarray(w0)[0], expect, atol=1e-14) + + +# ----------------------------------------------------------------------------- (b) gather parity + +def _numpy_window(Q_col, ifirst, frac, npts): + """One-mode wrapper around the production CPU window builder.""" + from RIFT.likelihood.factored_likelihood import _sinc_Q_window_numpy + Q_block = np.asarray(Q_col)[:, None] # (n_time, n_lm=1) + out = _sinc_Q_window_numpy(Q_block, np.asarray(ifirst, dtype=int), + np.asarray(frac, dtype=float), npts, a=A) + return out[:, :, 0] # (n_ext, npts) + + +@pytest.mark.parametrize("place", ["interior", "left_edge", "right_edge"]) +def test_gather_parity_against_numpy_window(place): + """Interior AND both edges: the edge cases pin the zero-extension convention, which is the + one place the four backends could agree on weights and still disagree on output.""" + rng = np.random.default_rng(7) + n_time, npts, n_ext = 512, 24, 40 + Q = rng.normal(size=n_time) + 1j * rng.normal(size=n_time) + frac = rng.uniform(0.0, 1.0, n_ext) + if place == "interior": + ifirst = rng.integers(2 * A, n_time - npts - 2 * A, n_ext) + elif place == "left_edge": + # Windows that start before the buffer, so the leading taps fall off the front. + ifirst = rng.integers(-A - 3, A, n_ext) + else: + ifirst = rng.integers(n_time - npts - A, n_time - npts + A + 3, n_ext) + + ref = _numpy_window(Q, ifirst, frac, npts) + + # JAX takes a continuous position; pos = ifirst + frac + t reproduces the same window. + pos = (ifirst[:, None] + frac[:, None] + np.arange(npts)[None, :]) + got = np.asarray(JC._GATHERERS["sinc"](jnp.asarray(Q), jnp.asarray(pos))) + + err = np.max(np.abs(got - ref)) + assert err < TOL_GATHER, "%s: JAX/numpy sinc windows disagree by %.3e" % (place, err) + if place != "interior": + # Guard against the test passing because every tap happened to land in-bounds: the + # comparison must actually exercise the zero-extension branch. + assert np.any(np.asarray(ifirst) < A) or np.any(np.asarray(ifirst) + npts + A > n_time) + + +def test_gpu_gather_parity_against_numpy_window(): + """cupy leg. SKIPS without a GPU -- it must not silently pass (cf. the same fix made to + test_noloop_gpu_stencils).""" + # Plain import, not importorskip: on a CPU head node cupy is INSTALLED but raises ImportError + # on libcuda.so.1, which importorskip will treat as an error from pytest 9.1 onward. + try: + import cupy + cupy.cuda.runtime.getDeviceCount() + except Exception as exc: + pytest.skip("no usable cupy/CUDA device, GPU stencil parity NOT exercised here: %s" + % type(exc).__name__) + from RIFT.likelihood import Q_inner_product + + rng = np.random.default_rng(11) + n_time, npts, n_ext, n_lm = 512, 16, 32, 2 + Q = (rng.normal(size=(n_time, n_lm)) + 1j * rng.normal(size=(n_time, n_lm))) + Amat = (rng.normal(size=(n_ext, n_lm)) + 1j * rng.normal(size=(n_ext, n_lm))) + ifirst = rng.integers(2 * A, n_time - npts - 2 * A, n_ext).astype(np.int32) + frac = rng.uniform(0.0, 1.0, n_ext) + + gpu = cupy.asnumpy(Q_inner_product.Q_inner_product_sinc_cupy( + cupy.asarray(Q), cupy.asarray(Amat), cupy.asarray(ifirst), + cupy.asarray(frac), npts, halfwidth=A)) + + # JAX: gather each mode, then contract on lm exactly as the kernel does. + pos = (ifirst[:, None] + frac[:, None] + np.arange(npts)[None, :]) + jx = np.zeros((n_ext, npts), dtype=complex) + for lm in range(n_lm): + jx += Amat[:, lm][:, None] * np.asarray( + JC._GATHERERS["sinc"](jnp.asarray(Q[:, lm]), jnp.asarray(pos))) + + err = np.max(np.abs(gpu - jx)) + assert err < TOL_GATHER, "JAX/CUDA sinc disagree by %.3e" % err + + +# ---------------------------------------------------------------------------- (c) accuracy gate + +def _bandlimited_series(n, f_over_nyq, rng, n_tones=24): + """A sum of tones strictly below f_over_nyq * Nyquist, evaluatable at ANY real sample index. + + Deliberately NOT an FFT zero-pad of a stored buffer. A pad of a truncated slice is not a + valid reference here: lnL(t) is band-limited but the slice is not, so truncation destroys the + band-limitation and the pad wraps -- measured to disagree with a converged Lanczos ladder at + the 1.4e-2 level during the sky-offset diagnosis. Summing known tones sidesteps that: the + exact value at a fractional index is the closed form below, with no reference error at all. + """ + f = rng.uniform(0.02, f_over_nyq, n_tones) * np.pi # rad/sample; Nyquist = pi + amp = rng.normal(size=n_tones) + 1j * rng.normal(size=n_tones) + phase = rng.uniform(0, 2 * np.pi, n_tones) + + def evaluate(t): + t = np.asarray(t, dtype=float) + return np.sum(amp * np.exp(1j * (f * t[..., None] + phase)), axis=-1) + + return evaluate(np.arange(n)), evaluate + + +# Required margin of 'sinc' over 'cubic' at f/fNyq = 0.7. MEASURED, not chosen: the ratio over +# seeds 3-10 on this generator is 37.1, 39.0, 40.4, 42.9, 44.6, 47.7, 47.7, 49.0 (ldas-grid, +# igwn python 3.11, jax 0.7.1, 2026-08-26). The gate is set at 20x -- comfortably under the +# observed floor of 37, so it fires on a change in stencil ORDER rather than on an unlucky draw, +# while still being ~10x above the 2.3x that separates cubic from linear here. +REQUIRED_SINC_OVER_CUBIC = 20.0 +# Floor the observed ratios must clear, so an erosion of the margin is caught long before the +# gate itself starts failing intermittently. Set below the measured minimum of 37.1, not at a +# multiple of the gate. +MARGIN_WATCHDOG = 25.0 +BANDWIDTH_FRACTION = 0.7 # poorly oversampled: the regime the 3G demo and O4 high-fmin sit in + + +def _stencil_errors(seed=3): + rng = np.random.default_rng(seed) + n = 4096 + series, exact = _bandlimited_series(n, BANDWIDTH_FRACTION, rng) + # Interior only: this measures interpolation error, not the (deliberate) edge truncation. + pos = rng.uniform(4 * A, n - 4 * A, 2000) + truth = exact(pos) + errs = {} + for name in ("nearest", "linear", "cubic", "sinc"): + got = np.asarray(JC._GATHERERS[name](jnp.asarray(series), jnp.asarray(pos))) + errs[name] = float(np.max(np.abs(got - truth)) / np.max(np.abs(truth))) + return errs + + +def test_sinc_beats_cubic_on_a_poorly_oversampled_band_limited_signal(): + """THE REGRESSION GATE. This is the test that fails if the stencil reverts to cubic.""" + errs = _stencil_errors() + ratio = errs["cubic"] / errs["sinc"] + assert ratio > REQUIRED_SINC_OVER_CUBIC, ( + "sinc no longer beats cubic at f/fNyq=%.2f: errors %r, ratio %.1f < %.1f" + % (BANDWIDTH_FRACTION, errs, ratio, REQUIRED_SINC_OVER_CUBIC)) + # Ordering of the whole ladder, so a stencil that regresses to a lower order is caught even + # if it is not literally cubic. + assert errs["sinc"] < errs["cubic"] < errs["linear"] + assert errs["nearest"] > errs["cubic"] + + +def test_mutation_cubic_fails_accuracy_gate(): + """MUTATION TEST. Reintroduce the cubic stencil under the 'sinc' key and confirm the gate + above fails. A gate that passes both ways is worthless, so this is not optional decoration: + it is the evidence that test_sinc_beats_cubic... has any power at all.""" + saved = JC._GATHERERS["sinc"] + JC._GATHERERS["sinc"] = JC._gather_cubic + try: + with pytest.raises(AssertionError): + test_sinc_beats_cubic_on_a_poorly_oversampled_band_limited_signal() + finally: + JC._GATHERERS["sinc"] = saved + # And the restore worked, so later tests in this file are not running against the mutant. + assert JC._GATHERERS["sinc"] is saved + + +def test_accuracy_margin_is_not_marginal(): + """Record how much room the gate has, over several seeds, so a future tightening is an + informed edit rather than a guess. Fails only if some seed lands within 2x of the gate.""" + ratios = [_stencil_errors(s)["cubic"] / _stencil_errors(s)["sinc"] for s in (3, 4, 5, 6, 7)] + assert min(ratios) > MARGIN_WATCHDOG, \ + "gate margin has eroded: ratios %r, floor %.0f, gate %.0fx" % ( + ratios, MARGIN_WATCHDOG, REQUIRED_SINC_OVER_CUBIC) + + +# ------------------------------------------------------------------- (d) vectorised == unrolled + +# The two forms are the same arithmetic in a different association order -- jnp.sum over the tap +# axis against a sequential accumulation -- so XLA is free to reassociate and the difference is +# NOT exactly zero: measured max|diff| 1.8e-15 on O(1)-O(10) complex data, i.e. a few ulp. (The +# sky-offset diagnostic's unnormalised prototype did land on exact 0.0; adding the weight +# normalisation is what moved it off.) 1e-12 is ~1000x that and still orders of magnitude below +# any semantic change: a dropped tap, a wrong offset or a lost window all move this by O(0.1). +TOL_UNROLL = 1e-12 + + +@pytest.mark.parametrize("a", [4, 8, 16]) +def test_vectorised_matches_unrolled(a): + """See _make_gather_sinc for why the tap axis must stay an array; this test is what makes a + revert to the loop form fail loudly.""" + rng = np.random.default_rng(0) + n = 4000 + Q = jnp.asarray(rng.normal(size=n) + 1j * rng.normal(size=n)) + pos = jnp.asarray(rng.uniform(-5, n + 5, 3000)) # includes out-of-buffer positions + fast = JC._make_gather_sinc(a)(Q, pos) + slow = JC._make_gather_sinc_unrolled(a)(Q, pos) + err = float(jnp.max(jnp.abs(fast - slow))) + assert err < TOL_UNROLL, "a=%d: vectorised and unrolled forms differ by %.3e" % (a, err) + + +# --------------------------------------------------------------------------------- wiring checks + +def test_sinc_is_reachable_from_the_registry_and_the_cli(): + """The stencil existing is not the same as a user being able to ask for it: 'cubic' was + implemented in _GATHERERS for a release while --interp's choices list still refused it.""" + assert "sinc" in JC._GATHERERS + import ast, io, os + driver = os.path.join(os.path.dirname(__file__), "..", "..", "bin", + "integrate_likelihood_extrinsic_jax") + src = io.open(driver, encoding="utf-8").read() + tree = ast.parse(src) + found = [] + for node in ast.walk(tree): + if (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + and node.func.attr == "add_option" + and node.args and isinstance(node.args[0], ast.Constant) + and node.args[0].value == "--interp"): + found = [kw for kw in node.keywords if kw.arg == "choices"] + assert found, "--interp option not found in the JAX driver" + # The choices expression must be derived from the registry, not a literal list that can rot. + assert "_JAX_GATHERER_NAMES" in ast.unparse(found[0].value), \ + "--interp choices is a literal list again; it will drift from _GATHERERS" + + +@pytest.mark.parametrize("name", ["nearest", "linear", "cubic", "sinc"]) +def test_separable_u_matches_the_general_path(name): + """Passing ``u`` must not change the answer for a window the accumulators actually build. + + This is the load-bearing test for the memory fix: getting ``u`` wrong is SILENT -- the + gather still returns an array of the right shape, just evaluated at the wrong sub-sample + offsets. So the separable path is checked against the general one for every stencil, at + PRODUCTION magnitudes (p0 ~ 1e5, where a binade crossing in p0 + t is possible) rather than + the small indices the other tests use. + + The two are not required to be bit-identical, and the difference has a known sign of merit: + ``frac(p0 + t)`` can lose a low mantissa bit that ``frac(p0)`` keeps, so the separable value + is the more exact of the two -- and is what numpy/cupy/CUDA compute. An ulp of position at + 1e5 is ~1.5e-11, so the tolerance is set just above that. + """ + rng = np.random.default_rng(4) + n_time, npts, S = 262144, 614, 300 + Q = jnp.asarray(rng.normal(size=n_time) + 1j * rng.normal(size=n_time)) + p0 = jnp.asarray(rng.uniform(100.0, n_time - npts - 100.0, (S,))) + t = jnp.arange(npts, dtype=jnp.float64) + pos = p0[:, None] + t + g = JC._GATHERERS[name] + gen = np.asarray(g(Q, pos)) + sep = np.asarray(g(Q, pos, JC._separable_u(p0))) + err = np.max(np.abs(gen - sep)) / np.max(np.abs(gen)) + assert err < 1e-10, "%s: separable-u path disagrees with the general one by %.3e" % (name, err) + # And the offsets really are separable for this construction, or the test proves nothing. + u_gen = np.asarray(pos - jnp.floor(pos)) + assert np.max(np.abs(u_gen - u_gen[:, :1])) < 1e-9 + + +def test_accumulators_pass_separable_u(): + """The fix only helps if the CALL SITES pass ``u``; a gatherer that merely accepts it does + nothing. Both accumulators must, or the (S, npts, 2a) weight array comes straight back.""" + import ast, io as _io, os + src = _io.open(os.path.join(os.path.dirname(__file__), "..", "..", "RIFT", "likelihood", + "jax_ile", "core.py"), encoding="utf-8").read() + tree = ast.parse(src) + calls = [n for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "gather"] + assert len(calls) >= 2, "expected a gather() call in each accumulator, found %d" % len(calls) + bad = [ast.unparse(c) for c in calls if len(c.args) < 3 and + not any(kw.arg == "u" for kw in c.keywords)] + assert not bad, "gather() called without the separable offset: %s" % bad + + # ... and the offset must actually be BUILT, conditionally on the stencil. Checking only + # that a third argument is present is not enough: `u_sep = None` everywhere would satisfy + # that while silently disabling the memory fix, which nothing else here would catch. + assigns = [n for n in ast.walk(tree) if isinstance(n, ast.Assign) + and any(isinstance(t, ast.Name) and t.id == "u_sep" for t in n.targets)] + assert len(assigns) >= 2, "expected a u_sep assignment per accumulator, found %d" % len(assigns) + for a in assigns: + src_expr = ast.unparse(a.value) + assert "_separable_u" in src_expr, \ + "u_sep no longer builds the separable offset (%s); the memory fix is disabled" % src_expr + assert isinstance(a.value, ast.IfExp), \ + ("u_sep is unconditional (%s); it must stay gated off the stencils that ignore u -- " + "feeding it to 'nearest' cost >60%% wall on the banded path" % src_expr) + + +def test_every_entry_point_defaults_to_the_same_stencil(): + """The default is ONE constant, and every entry point in the package uses it. + + Changed 2026-08-26 from 'linear' to 'sinc'. The failure this guards against is a PARTIAL + revert or a partial adoption: before this change the CLI flag and ~17 library signatures each + carried their own "linear" literal, so moving the CLI alone would have left every direct + caller of fused_log_likelihood / JAXDistanceMarginalizedLikelihood on the old stencil while + --help claimed otherwise. Enumerating the ACTUAL signatures, rather than asserting the + constant equals itself, is what makes that detectable. + """ + import inspect + from RIFT.likelihood.jax_ile import wrapper as JW + + assert JC.JAX_INTERP_DEFAULT in JC._GATHERERS, \ + "the default names a stencil that does not exist" + assert JC.JAX_INTERP_DEFAULT == "sinc", ( + "default stencil changed to %r -- intentional? It alters results for every caller that " + "does not pass interp=, so update DESIGN_q_window_stencil.md and the --interp help text " + "in the same commit." % (JC.JAX_INTERP_DEFAULT,)) + + offenders = [] + for mod in (JC, JW): + for name, obj in vars(mod).items(): + targets = [] + if inspect.isfunction(obj) and getattr(obj, "__module__", "").startswith("RIFT"): + targets.append((name, obj)) + elif inspect.isclass(obj) and getattr(obj, "__module__", "").startswith("RIFT"): + targets.append((name + ".__init__", obj.__init__)) + for label, fn in targets: + try: + par = inspect.signature(fn).parameters.get("interp") + except (TypeError, ValueError): + continue + if par is None or par.default is inspect.Parameter.empty: + continue + if par.default != JC.JAX_INTERP_DEFAULT: + offenders.append("%s.%s=%r" % (mod.__name__, label, par.default)) + assert not offenders, ( + "these entry points disagree with JAX_INTERP_DEFAULT=%r: %s" + % (JC.JAX_INTERP_DEFAULT, ", ".join(sorted(offenders)))) + # The sweep must actually have found interp= parameters, or it proves nothing. + n_seen = sum(1 for mod in (JC, JW) for _n, o in vars(mod).items() + if inspect.isfunction(o) and getattr(o, "__module__", "").startswith("RIFT") + and "interp" in inspect.signature(o).parameters) + assert n_seen >= 8, "only %d interp= entry points found; the sweep is not covering them" % n_seen + + +def test_adaptive_distance_grid_uses_the_callers_stencil(): + """estimate_distance_peak sizes the grid the likelihood then integrates on, so it must run on + the SAME stencil the caller asked for. + + Found by review of the default change. These three classes accept interp= and forward it to + the likelihood but called estimate_distance_peak(data, guess_snr) positionally, so the grid + was always sized with the module default. That was a latent mismatch while the default was + 'linear'; moving the default to 'sinc' made it break the documented recovery path -- a caller + passing interp="linear" to reproduce a pre-2026-08-26 run would have got a linear likelihood + on a sinc-sized grid. The path is behind JAX_ILE_DISTGRID_ADAPTIVE=1, which is exactly why + it needs a test rather than a reader. + """ + import ast, io as _io, os + src = _io.open(os.path.join(os.path.dirname(__file__), "..", "..", "RIFT", "likelihood", + "jax_ile", "wrapper.py"), encoding="utf-8").read() + calls = [n for n in ast.walk(ast.parse(src)) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == "estimate_distance_peak"] + assert calls, "estimate_distance_peak is no longer called from wrapper.py; retarget this test" + bad = [ast.unparse(c) for c in calls + if not any(kw.arg == "interp" for kw in c.keywords) and len(c.args) < 5] + assert not bad, "estimate_distance_peak called without forwarding interp=: %s" % bad + + +def test_no_method_silently_overrides_the_instance_stencil(): + """A method that takes its own ``interp`` must default to the INSTANCE's, never to the + module default. + + Reported on PR #193. ``JAXDistPhiMargLikelihood.sample_phi_ref`` carried + ``interp=JAX_INTERP_DEFAULT``, so an instance constructed with any other stencil drew its + phi_ref from a different likelihood than the one it reports lnL and evidence from -- silently. + Harmless only while the two strings coincided; moving the module default to 'sinc' made it + break the very ``interp="linear"`` recipe this change advertises for reproducing old runs. + + This checks the SHAPE rather than the one method, so a new sampler with the same defect fails + here: any method (other than __init__) that accepts ``interp`` must default it to None, and + the class must retain ``self.interp`` for it to fall back to. + """ + import ast, io as _io, os + src = _io.open(os.path.join(os.path.dirname(__file__), "..", "..", "RIFT", "likelihood", + "jax_ile", "wrapper.py"), encoding="utf-8").read() + tree = ast.parse(src) + offenders, checked = [], 0 + for cls in [n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)]: + stores = any(isinstance(n, ast.Assign) and + any(isinstance(t, ast.Attribute) and t.attr == "interp" for t in n.targets) + for n in ast.walk(cls)) + for fn in [f for f in cls.body if isinstance(f, ast.FunctionDef)]: + args = fn.args.args + if "interp" not in [a.arg for a in args]: + continue + if fn.name == "__init__": + if not stores: + offenders.append("%s.__init__ takes interp but never stores self.interp" + % cls.name) + continue + checked += 1 + # map defaults onto the tail of the positional args + defaults = dict(zip([a.arg for a in args][-len(fn.args.defaults):], + fn.args.defaults)) if fn.args.defaults else {} + d = defaults.get("interp") + if not (isinstance(d, ast.Constant) and d.value is None): + offenders.append("%s.%s defaults interp to %s, not None -- it will ignore the " + "instance's stencil" % (cls.name, fn.name, + ast.unparse(d) if d else "")) + elif "self.interp if interp is None" not in ast.unparse(fn): + offenders.append("%s.%s defaults interp to None but never falls back to " + "self.interp" % (cls.name, fn.name)) + assert not offenders, "; ".join(offenders) + assert checked >= 1, "no method with an interp= parameter found; retarget this test" + + +def test_cli_default_comes_from_the_shared_constant(): + """--interp's default must be the constant, not a re-typed literal that can drift from it.""" + import ast, io as _io, os + driver = os.path.join(os.path.dirname(__file__), "..", "..", "bin", + "integrate_likelihood_extrinsic_jax") + tree = ast.parse(_io.open(driver, encoding="utf-8").read()) + node = next(n for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) + and n.func.attr == "add_option" and n.args + and isinstance(n.args[0], ast.Constant) and n.args[0].value == "--interp") + default = next(kw for kw in node.keywords if kw.arg == "default") + assert ast.unparse(default.value) == "JAX_INTERP_DEFAULT", \ + "--interp default is %r, not the shared constant" % ast.unparse(default.value) + + +def test_likelihood_runs_and_differentiates_with_sinc(): + """Wire-level check: the stencil must work THROUGH the likelihood, not just as a helper. + Unit-testing the gatherer proves nothing about whether _accumulate_unit can call it.""" + from test_jax_likelihood import make_synthetic, make_Pvec # noqa: E402 + from RIFT.likelihood.jax_ile import build_likelihood_data, fused_log_likelihood + + packed, _ref, tvals, deltaT, tref = make_synthetic() + P, distMpc = make_Pvec(9, tref, deltaT) + data = build_likelihood_data(packed, deltaT, tref, tvals) + + def run(nm): + return np.asarray(fused_log_likelihood( + data, P.phi, P.theta, P.psi, P.incl, P.phiref, distMpc, interp=nm)) + + vals = {nm: run(nm) for nm in ("nearest", "cubic", "sinc")} + assert all(np.all(np.isfinite(v)) for v in vals.values()), vals + # sinc must actually change the answer -- otherwise the interp argument is being ignored, + # which is exactly the silent no-op a helper-only test cannot see. + assert not np.array_equal(vals["sinc"], vals["cubic"]) + # ... and the two interpolating stencils must sit far closer to each other than either does + # to 'nearest', since both approximate the same band-limited value. + d_interp = np.max(np.abs(vals["sinc"] - vals["cubic"])) + d_nearest = np.max(np.abs(vals["cubic"] - vals["nearest"])) + assert d_interp < 0.5 * d_nearest, (d_interp, d_nearest) + + # Differentiable through the stencil, which is the whole reason the JAX path exists. + def scalar(ra): + return fused_log_likelihood( + data, jnp.array([ra]), P.theta[:1], P.psi[:1], P.incl[:1], P.phiref[:1], + distMpc[:1], interp="sinc")[0] + + g = float(jax.grad(scalar)(float(P.phi[0]))) + assert np.isfinite(g) and g != 0.0 diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_tempering_chooser.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_tempering_chooser.py new file mode 100644 index 000000000..32e8dc2f0 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_tempering_chooser.py @@ -0,0 +1,779 @@ +"""The tempering-cost law and the driver's --adapt-weight-exponent chooser. + +WHAT THESE PIN, AND WHY THE SHAPE +--------------------------------- +`--adapt-weight-exponent beta` does two structurally different things in RIFT: +in the non-JAX samplers it shapes only the adaptive sampling PRIOR (unbiased at +any beta, no cost in exported samples), while on the JAX flowMC path it is the +exponent of the target the MCMC SAMPLES, so the export must be reweighted by +L^(1-beta) and that reweight has an ESS cost. The cost obeys + + ESS/N = [beta (2 - beta)]^(n_dim/2) + +which is set by the SAMPLED DIMENSION and NOT by lnLmax -- i.e. not by SNR, +which is what the historical non-JAX helper keys on. Provenance, the measured +sweep and the arm study: RIFT/likelihood/jax_ile/DESIGN_jax_tempering.md. + +Several tests below are AST guards on the DRIVER SOURCE rather than assertions +on a helper. That is deliberate and matches test_jax_fairdraw_export.py: the +defects these pin live at CALL SITES (a chooser that is computed and then not +used; a guard that is computed and then not raised), where a helper-level +assertion cannot see them. Needs no lal, no GPU and no flowMC. +""" +import ast +import os +import subprocess +import sys + +import numpy as np +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +CODE = os.path.abspath(os.path.join(HERE, "..", "..")) +DRIVER = os.path.join(CODE, "bin", "integrate_likelihood_extrinsic_jax") +if CODE not in sys.path: + sys.path.insert(0, CODE) + +from RIFT.likelihood.jax_ile.samplers import ( # noqa: E402 + beta_for_export_ess, export_ess_fraction, export_ess_estimate) + +# The measured sweep (dim=4) from DESIGN_jax_tempering.md. Kept as data because +# more than one test needs it, and because the calibration envelope must be +# checked against ALL of it, not against a couple of convenient points. +MEASURED_D4 = { + 0.05: 7.546e-03, 0.09508: 2.763e-02, 0.10: 3.033e-02, 0.20: 1.042e-01, + 0.30: 2.123e-01, 0.40: 3.481e-01, 0.50: 4.999e-01, 0.55: 5.772e-01, + 0.60: 6.530e-01, 0.65: 7.253e-01, 0.70: 7.923e-01, 0.75: 8.522e-01, + 0.80: 9.035e-01, 0.85: 9.449e-01, 0.90: 9.752e-01, 0.95: 9.938e-01, + 1.00: 1.0, +} + + +def _driver_tree(): + with open(DRIVER) as f: + return ast.parse(f.read()) + + +def _load_driver(): + """Import the driver script (no .py suffix) as a module, as + test_jax_fairdraw_export.py does, so the chooser can be CALLED rather than + only read. Source-level guards cannot see whether a branch is reachable -- + a mutation that made the beta>1 branch dead (`if beta >= 1.0 and False:`) + survived an AST-only version of these tests.""" + import importlib.machinery + loader = importlib.machinery.SourceFileLoader("_ile_jax_driver_temper", DRIVER) + spec = importlib.util.spec_from_loader("_ile_jax_driver_temper", loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +drv = _load_driver() + + +class _Opts(object): + """Minimal stand-in for the optparse Values the chooser reads. + + ``supplied`` lists the option TOKENS to pretend were typed. Setting an + attribute is NOT the same as passing the flag any more -- that conflation is + the defect this stub previously baked in -- so any test whose behaviour turns + on "did the user pass it" must name the token here (or build from argv). + """ + def __init__(self, supplied=(), **kw): + self._supplied_options = set(supplied) + self.adapt_adapt = False + self.auto_adapt_weight_exponent = False + self.adapt_weight_exponent = 1.0 + self.target_export_ess_frac = drv._TARGET_EXPORT_ESS_FRAC_DEFAULT + self.allow_degenerate_tempering = False + self.smc_puffball = False + self.mode = "flowmc-phimarg" + self.__dict__.update(kw) + self._supplied_options = set(supplied) + + +# ----------------------------------------------------------------- the law +def test_law_matches_the_measured_sweep(): + """The closed form against the EXACT sweep on the real BNS likelihood. + + Measured by beta_ess_offline.py on the SNR-23.8 zero-noise BNS (4-D + phi-marginalised), reference ESS(g=1)=5235 with every g in [0,2] resolved + (worst-case IS ESS 508). The law is Gaussian-peak and therefore slightly + OPTIMISTIC: measured/law runs 0.79 (beta=0.05) to 1.00 (beta=1). Pinning + the band both ways is what makes this a test rather than a restatement -- + an over-optimistic law would silently under-budget a real run. + """ + measured = MEASURED_D4 + ratios = [] + for beta, meas in measured.items(): + law = export_ess_fraction(beta, 4) + ratios.append(meas / law) + assert min(ratios) > 0.75, "law is far more optimistic than measured: %r" % (ratios,) + assert max(ratios) <= 1.001, "law UNDER-predicts the measured cost: %r" % (ratios,) + + +def test_law_is_independent_of_lnLmax_by_construction(): + """The signature carries no lnLmax/SNR argument at all. + + This is the property the whole design rests on, so pin it structurally: if + someone later adds an SNR term (re-importing the non-JAX helper's rule), the + signature changes and this fails. + """ + import inspect + params = list(inspect.signature(export_ess_fraction).parameters) + assert params == ["beta", "n_dim"], params + for bad in ("snr", "lnLmax", "lnl_max", "guess_snr"): + assert bad not in params + + +def test_law_depends_on_dimension(): + """A knob that gives the same answer for every dimension is a dead knob.""" + vals = [export_ess_fraction(0.5, d) for d in (3, 4, 5)] + assert len(set(vals)) == 3 + assert vals[0] > vals[1] > vals[2] # more dimensions -> costlier + + +def test_chosen_beta_meets_the_target_on_the_MEASURED_lower_bound(): + """The chooser must satisfy the target against the CALIBRATED envelope. + + THE DEFECT (review): it used to invert the bare Gaussian law, which the sweep + shows is optimistic by up to 21%. Round-tripping that formula is vacuous -- + it only proves the inverse inverts the thing that is known to overstate the + answer. At dim=4 with a 0.9 target it returned beta=0.77347, whose measured + lower bound is 0.866: less than was asked for. + """ + for n_dim in (3, 4, 5): + for target in (0.3, 0.5, 0.9, 0.99): + beta = beta_for_export_ess(target, n_dim) + assert 0.0 < beta <= 1.0 + assert export_ess_estimate(beta, n_dim) >= target - 1e-9, ( + "dim=%d target=%g -> beta=%g retains only %.4f" + % (n_dim, target, beta, export_ess_estimate(beta, n_dim))) + # and it must be the SMALLEST such beta, to within the solver step + if beta > 1e-3: + assert export_ess_estimate(beta * 0.99, n_dim) < target + 1e-9 + + +def test_the_old_uncalibrated_answer_would_NOT_pass(): + """Pins that the fix changed the number, not just the wording.""" + assert export_ess_estimate(0.77347, 4) < 0.9 + assert beta_for_export_ess(0.9, 4) > 0.77347 + + +def test_calibration_envelope_is_a_true_lower_bound_everywhere_measured(): + """Every measured point of the sweep must sit at or above the envelope. + + This is the property the guard and the chooser both rely on; if a future + knot edit breaks it anywhere, the conservatism is gone silently. + """ + for beta, measured in MEASURED_D4.items(): + lb = export_ess_estimate(beta, 4) + assert measured >= lb - 1e-12, ( + "beta=%g: measured %.4e is BELOW the supposed lower bound %.4e" + % (beta, measured, lb)) + + +def test_lower_bound_is_never_above_the_law_and_meets_it_at_beta_one(): + for n_dim in (3, 4, 5): + for beta in (0.05, 0.2, 0.5, 0.8, 0.95): + assert export_ess_estimate(beta, n_dim) <= export_ess_fraction(beta, n_dim) + assert export_ess_estimate(1.0, n_dim) == pytest.approx( + export_ess_fraction(1.0, n_dim)) + + +def test_unreachable_target_raises_rather_than_returning_beta_one(): + """A target the envelope cannot reach must fail loudly, not silently clamp.""" + import RIFT.likelihood.jax_ile.samplers as S + real = S._ESS_CAL_RATIO + try: + S._ESS_CAL_RATIO = tuple(0.5 * r for r in real) # envelope caps at 0.5 + with pytest.raises(ValueError, match="unreachable"): + beta_for_export_ess(0.9, 4) + finally: + S._ESS_CAL_RATIO = real + + +def test_beta_one_is_free_and_is_the_only_free_point(): + for n_dim in (3, 4, 5): + assert export_ess_fraction(1.0, n_dim) == pytest.approx(1.0) + assert export_ess_fraction(0.999, n_dim) < 1.0 + + +def test_domain_errors_raise_rather_than_clamp(): + """A silently clamped exponent is the failure this whole change exists to stop.""" + for bad in (0.0, -0.1, 1.5): + with pytest.raises(ValueError): + export_ess_fraction(bad, 4) + for bad in (0.0, -0.1, 1.5): + with pytest.raises(ValueError): + beta_for_export_ess(bad, 4) + + +def test_historical_helper_beta_would_be_degenerate_here(): + """The number this change exists to prevent someone from porting. + + helper_LDG_Events.py picks beta = 0.1*(22.5/SNR)^2 above SNR 22.5; at the + SNR-23.8 study event that is 0.0951. On the 4-D JAX path that leaves ~3% of + the cloud, i.e. ESS ~140 of 4800 -- below the driver's own "NOT a usable + posterior sample" warning threshold of 200. + """ + beta_hist = 0.1 / (23.78 / 22.5) ** 2 + frac = export_ess_fraction(beta_hist, 4) + assert frac < 0.05 + assert frac * 4800 < 200 + + +# -------------------------------------------------------- driver call sites +def _declared_options(): + """Flag strings passed as the FIRST argument of an add_option call. + + NOT "every string constant beginning with --" anywhere in the file. That + weaker version passed while the parser flag was renamed to + `--auto-adapt-weight-exponent-XX`, because the correct spelling still occurred + inside two error messages and the inert-flag list -- and a MUTANT was + committed behind it. + """ + tree = _driver_tree() + out = set() + for n in ast.walk(tree): + if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) \ + and n.func.attr == "add_option" and n.args: + a = n.args[0] + if isinstance(a, ast.Constant) and isinstance(a.value, str): + out.add(a.value) + return out + + +def test_driver_exposes_the_auto_chooser_options(): + declared = _declared_options() + for flag in ("--auto-adapt-weight-exponent", "--target-export-ess-frac", + "--allow-degenerate-tempering"): + assert flag in declared, ( + "%s is not DECLARED via add_option (mentioning it in a help or error " + "string does not make it a flag)" % flag) + + +def test_no_stray_placeholder_flags(): + """Mutation-harness residue guard. + + A mutation sweep renames flags in place; if one is left applied when the tree + is staged, it ships. That happened once (`--auto-adapt-weight-exponent-XX` + reached a commit). Cheap, permanent check. + """ + for flag in _declared_options(): + assert not flag.endswith("-XX"), "placeholder flag left in the parser: %s" % flag + + +def test_chooser_result_is_actually_RETURNED_and_consumed(): + """A chooser that is computed and then not used is the classic dead knob. + + RETARGETED: this used to require that the chooser ASSIGN + opts.adapt_weight_exponent. That contract was the bug -- opts is per-RUN and + analyze_one is per-EVENT, so the write made event 1 of a batch read event 0's + choice. The chooser now RETURNS the exponent; see + test_chooser_returns_the_exponent_rather_than_writing_it_back for the other + half, and test_chooser_is_reusable_across_a_multi_event_BATCH for the + behaviour this protects. + """ + tree = _driver_tree() + fn = next((n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "resolve_tempering_exponent"), None) + assert fn is not None, "resolve_tempering_exponent is missing from the driver" + calls = {n.func.id for n in ast.walk(fn) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)} + assert "beta_for_export_ess" in calls + returns = [n for n in ast.walk(fn) if isinstance(n, ast.Return) and n.value is not None] + assert returns, "resolve_tempering_exponent returns nothing" + + +def test_auto_refuses_to_silently_override_an_explicit_exponent(): + """--auto plus an explicit --adapt-weight-exponent must RAISE, not overwrite. + + Pinned on the source because reaching this branch at runtime needs a full + likelihood build. Both conflicts (this one and --auto + --adapt-adapt) must + be present: each was added after noticing the other silently won. + """ + tree = _driver_tree() + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "resolve_tempering_exponent") + msgs = [c.value for n in ast.walk(fn) if isinstance(n, ast.Raise) + for c in ast.walk(n) + if isinstance(c, ast.Constant) and isinstance(c.value, str)] + joined = " ".join(msgs) + assert "would overwrite it" in joined, ( + "no guard against --auto silently overriding an explicit exponent") + assert "--adapt-adapt" in joined, ( + "no guard against --auto being combined with --adapt-adapt") + + +@pytest.mark.parametrize("bad", [1.5, 2.0, 0.0, -0.2]) +def test_beta_outside_the_unit_interval_is_REFUSED_at_runtime(bad, capsys): + """CALL the chooser. beta>1 sharpens the target past the posterior + (samplers.flowmc_sample* take temper = 1/beta, so beta>1 gives inv_T>1) and + its export reweight has ESS/N = [beta(2-beta)]^(dim/2): 0 at beta=2, + undefined beyond. beta<=0 samples the prior. + + An earlier source-only version of this test passed while a mutation made the + branch dead, and an earlier version of the CODE tested `beta >= 1.0` and + printed "beta=1 (untempered target)" for every beta>1. Both are why this + exercises the function instead of reading it. + """ + with pytest.raises(SystemExit) as e: + drv.resolve_tempering_exponent(_Opts(adapt_weight_exponent=bad), 4, 4800) + assert "untempered" not in str(e.value) + out = capsys.readouterr().out + assert "untempered target" not in out, ( + "beta=%r was reported as untempered" % bad) + + +def test_beta_one_and_a_healthy_beta_are_accepted_at_runtime(capsys): + """The negative tests above prove nothing if every input raises.""" + drv.resolve_tempering_exponent(_Opts(adapt_weight_exponent=1.0), 4, 4800) + assert "untempered target" in capsys.readouterr().out + # Derive the exponent from the chooser rather than hardcoding one. The + # literal 0.7735 that used to be here was the OLD uncalibrated answer, whose + # measured lower bound is 0.866 -- so this assertion was quietly encoding the + # very over-claim the calibration exists to remove. + beta = beta_for_export_ess(0.9, 4) + o = _Opts(adapt_weight_exponent=beta) + drv.resolve_tempering_exponent(o, 4, 4800) + out = capsys.readouterr().out + assert "export ESS/N >= 0.9" in out, out + + +def test_auto_sets_the_exponent_and_the_value_depends_on_dimension(capsys): + """The chooser must MOVE the number, and move it differently per dimension. + Identical output across settings is the signature of a dead knob.""" + got = {} + for n_dim in (3, 4, 5): + o = _Opts(auto_adapt_weight_exponent=True) + beta = drv.resolve_tempering_exponent(o, n_dim, 4800) + capsys.readouterr() + assert beta != 1.0, "auto returned the default exponent" + assert o.adapt_weight_exponent == 1.0, "auto mutated opts (see the batch test)" + got[n_dim] = beta + assert len(set(got.values())) == 3, got + assert got[3] < got[4] < got[5] + + +def test_degenerate_exponent_WARNS_and_does_not_refuse(capsys): + """RETARGETED. This used to require a REFUSAL. + + A refusal is a guarantee, and review showed the estimate cannot support one: + it is calibrated at SNR ~= 23.8, and the PR's own SNR ladder measures + ESS/N = 0.00823 at beta=0.1, d=4, SNR ~= 67 against an estimated 0.0285 -- + 3.5x the wrong way. Refusing on that number would be false confidence, so + the run proceeds with a loud warning that says which way the estimate errs. + """ + o = _Opts(adapt_weight_exponent=0.09508, supplied=["--adapt-weight-exponent"]) + beta = drv.resolve_tempering_exponent(o, 4, 4800) + assert beta == pytest.approx(0.09508) # it RAN + cap = capsys.readouterr() + assert "WARNING" in cap.err + assert "usable posterior sample" in cap.err + # and it must disclose the direction of the error, not just the number + assert "optimistic" in cap.err.lower() and "0.00823" in cap.err + + # --allow-degenerate-tempering silences it + o2 = _Opts(adapt_weight_exponent=0.09508, allow_degenerate_tempering=True, + supplied=["--adapt-weight-exponent", "--allow-degenerate-tempering"]) + drv.resolve_tempering_exponent(o2, 4, 4800) + assert "WARNING" not in capsys.readouterr().err + + +def test_the_estimate_is_NOT_presented_as_a_bound(): + """The name and the docstring must not promise a guarantee. + + The function was called export_ess_lower_bound and the guard refused on it, + which is exactly the over-claim review caught. Pinning the naming stops it + being reintroduced by a rename. + """ + import RIFT.likelihood.jax_ile.samplers as S + assert not hasattr(S, "export_ess_lower_bound"), ( + "the bound-flavoured name is back") + doc = S.export_ess_estimate.__doc__ or "" + assert "NOT a bound" in doc or "not a bound" in doc.lower() + assert "0.00823" in doc, "the counterexample that disproves the bound is not recorded" + + +def test_adapt_adapt_is_rejected_for_modes_whose_sampler_cannot_anneal(): + """--mode flowmc has no temper_adapt path; claiming annealing there both + misreported the run and skipped the export check entirely.""" + o = _Opts(adapt_adapt=True, adapt_weight_exponent=0.1, mode="flowmc", + supplied=["--adapt-adapt", "--adapt-weight-exponent"]) + with pytest.raises(SystemExit, match="not implemented for --mode flowmc"): + drv.resolve_tempering_exponent(o, 5, 4800) + # the phimarg family still anneals + for mode in ("flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg"): + assert drv.resolve_tempering_exponent( + _Opts(adapt_adapt=True, mode=mode, supplied=["--adapt-adapt"]), 4, 4800) == 1.0 + + +@pytest.mark.parametrize("abbrev", ["--adapt-weight-exp=1.0", + "--adapt-weight-expo=0.5"]) +def test_long_option_ABBREVIATIONS_are_canonicalised(abbrev): + """optparse accepts unambiguous prefixes. + + Recording the literal token meant `--adapt-weight-exp=1.0` set the value while + was_supplied('--adapt-weight-exponent') stayed False, so --auto silently + overwrote an explicit exponent -- the L1015 defect through a side door. + """ + argv = ["--auto-adapt-weight-exponent", abbrev] + o = _opts_from_argv(argv) + assert drv.was_supplied(o, "--adapt-weight-exponent") is True + with pytest.raises(SystemExit, match="would overwrite it"): + drv.resolve_tempering_exponent(o, 4, 4800) + + +def test_target_without_auto_is_reported_not_silently_ignored(capsys): + """Setting a budget and no chooser does nothing; say so. + + Nothing covered this and a mutation deleting the note survived the sweep. + """ + o = _Opts(target_export_ess_frac=0.5, supplied=["--target-export-ess-frac"]) + drv.resolve_tempering_exponent(o, 4, 4800) + out = capsys.readouterr().out + assert "--target-export-ess-frac" in out and "no effect" in out, out + # and it must NOT be reported when the chooser is actually on + o2 = _Opts(auto_adapt_weight_exponent=True, target_export_ess_frac=0.5, + supplied=["--auto-adapt-weight-exponent", "--target-export-ess-frac"]) + drv.resolve_tempering_exponent(o2, 4, 4800) + assert "no effect" not in capsys.readouterr().out + + +def test_inert_note_lists_only_flags_the_user_ACTUALLY_passed(capsys): + """On a non-tempered mode the chooser flags are reported as inert -- but the + report must not name a flag the user never typed. + + It previously appended --target-export-ess-frac whenever --auto was set, even + with the target at its default, telling the user a flag they had not passed + was being ignored. + """ + p = drv.build_parser() + + def mk(supplied=(), **kw): + class O(object): + pass + o = O() + for opt in p._get_all_options(): + if opt.dest: + setattr(o, opt.dest, opt.default) + o.mode = "laplace-is" + for k, v in kw.items(): + setattr(o, k, v) + # Setting the attribute is NOT passing the flag: the report keys on the + # command-line token, so a test that wants a flag reported must name it. + o._supplied_options = set(supplied) + return o + + def note(o): + drv.check_critical_and_report(o, p) + return "".join(l for l in capsys.readouterr().out.splitlines() + if "tempered modes" in l) + + default_target = note(mk(auto_adapt_weight_exponent=True, + supplied=["--auto-adapt-weight-exponent"])) + assert "--auto-adapt-weight-exponent" in default_target + assert "--target-export-ess-frac" not in default_target, default_target + + given_target = note(mk(auto_adapt_weight_exponent=True, + target_export_ess_frac=0.5, + supplied=["--auto-adapt-weight-exponent", + "--target-export-ess-frac"])) + assert "--target-export-ess-frac" in given_target, given_target + + # and on a tempered mode nothing is reported inert at all + o = mk(auto_adapt_weight_exponent=True, + supplied=["--auto-adapt-weight-exponent"]) + o.mode = "flowmc-phimarg" + assert note(o) == "" + + +def test_chooser_is_reusable_across_a_multi_event_BATCH(): + """analyze_one runs once per intrinsic template with the SAME opts. + + The chooser used to write opts.adapt_weight_exponent, so on event 1 it read + its own event-0 output as a user-supplied exponent and killed the batch with + SystemExit -- and ILE_extr.sub runs batches, so this broke every real + multi-event run of --auto. Pin BOTH halves: it must not raise, and it must + not mutate opts. + """ + o = _Opts(auto_adapt_weight_exponent=True) + betas = [drv.resolve_tempering_exponent(o, 4, 4800) for _ in range(3)] + assert len(set(betas)) == 1, betas + assert o.adapt_weight_exponent == 1.0, ( + "the chooser mutated opts; per-run state must not carry per-event meaning") + + +def test_chooser_returns_the_exponent_rather_than_writing_it_back(): + """Structural companion: the returned value must be what the caller uses. + + A chooser that returns the right number and is called for its side effect is + the same dead knob as one that returns nothing. + """ + tree = _driver_tree() + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "resolve_tempering_exponent") + assert not any(isinstance(t, ast.Attribute) and t.attr == "adapt_weight_exponent" + for n in ast.walk(fn) if isinstance(n, ast.Assign) + for t in n.targets), "chooser still assigns opts.adapt_weight_exponent" + ana = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "analyze_one") + src = ast.get_source_segment(open(DRIVER).read(), ana) or "" + assert "resolved_beta = resolve_tempering_exponent(" in src, ( + "analyze_one ignores the chooser's return value") + assert "_beta = resolved_beta" in src, ( + "the sampler call sites still read opts instead of the resolved value") + + +def test_smc_puffball_is_not_refused_because_the_exponent_is_inert_there(): + """--smc-puffball routes to smc_puffball_sample, which swallows `temper` in + **_ignore and exports uniform weights. Refusing a run over a number that + does nothing is a false alarm; the guard skipped this path check and did + exactly that.""" + for o in (_Opts(smc_puffball=True, adapt_weight_exponent=0.09508, + supplied=["--adapt-weight-exponent"]), + _Opts(smc_puffball=True, auto_adapt_weight_exponent=True, + supplied=["--auto-adapt-weight-exponent"])): + assert drv.resolve_tempering_exponent(o, 4, 4800) == 1.0 + + # ... and the no-op is REPORTED, not silent + o = _Opts(smc_puffball=True, adapt_weight_exponent=0.09508, + supplied=["--adapt-weight-exponent"]) + import io + import contextlib + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + drv.resolve_tempering_exponent(o, 4, 4800) + assert "--smc-puffball ignores" in buf.getvalue(), buf.getvalue() + + # control: WITHOUT --smc-puffball the same exponent is NOT silently fine -- + # it runs (the floor is advisory now) but must warn. + import io as _io2 + import contextlib as _c2 + err = _io2.StringIO() + with _c2.redirect_stderr(err), _c2.redirect_stdout(_io2.StringIO()): + drv.resolve_tempering_exponent( + _Opts(adapt_weight_exponent=0.09508, + supplied=["--adapt-weight-exponent"]), 4, 4800) + assert "WARNING" in err.getvalue() + + +def _opts_from_argv(argv): + """Parse a real command line AND record which tokens it named.""" + p = drv.build_parser() + o, _ = p.parse_args(list(argv)) + drv.record_supplied_options(o, list(argv), p) # parser => abbreviations canonicalised + o.mode = "flowmc-phimarg" + return o + + +@pytest.mark.parametrize("argv", [ + ["--auto-adapt-weight-exponent", "--adapt-weight-exponent", "1.0"], + ["--auto-adapt-weight-exponent", "--adapt-weight-exponent=1.0"], + ["--adapt-weight-exponent", "1.0", "--auto-adapt-weight-exponent"], +]) +def test_auto_conflicts_with_an_EXPLICIT_DEFAULT_exponent(argv): + """--auto plus an explicitly-typed --adapt-weight-exponent 1.0 must RAISE. + + THE DEFECT (review): the conflict was detected by comparing the VALUE against + the default, so passing the default explicitly was silently accepted and the + chooser then replaced the user's explicit untempered target -- the opposite of + the documented behaviour, and a change to the sampled target. Detection is + now by command-line TOKEN. The `--opt=value` form and both orderings are + covered because a token scan that missed either would reintroduce it. + """ + with pytest.raises(SystemExit, match="would overwrite it"): + drv.resolve_tempering_exponent(_opts_from_argv(argv), 4, 4800) + + +def test_auto_alone_is_still_accepted(): + """The negative test above proves nothing if every command line raises.""" + o = _opts_from_argv(["--auto-adapt-weight-exponent"]) + assert drv.resolve_tempering_exponent(o, 4, 4800) != 1.0 + + +def test_main_ACTUALLY_populates_the_token_record(): + """The wiring, not the helper. + + Every other test in this file builds its own options object and calls + record_supplied_options() itself, so all 40 stayed green when the call was + deleted from main() -- mutation C6. A perfectly tested helper that production + never invokes: was_supplied() would then return False for everything and both + conflict guards would silently stop firing. + + Pinned structurally because reaching main() needs frames and a likelihood. + Order matters too: check_critical_and_report reads the record to decide which + flags to report as inert, so the record must be populated BEFORE it runs. + """ + src = open(DRIVER).read() + tree = ast.parse(src) + fn = next((n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "main"), None) + assert fn is not None, "main() is missing from the driver" + + calls = [n for n in ast.walk(fn) if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name)] + names = [c.func.id for c in calls] + assert "record_supplied_options" in names, ( + "main() never calls record_supplied_options -- was_supplied() would " + "return False for every flag and both conflict guards would go silent") + + rec = min(c.lineno for c in calls if c.func.id == "record_supplied_options") + rpt = [c.lineno for c in calls if c.func.id == "check_critical_and_report"] + if rpt: + assert rec < min(rpt), ( + "the token record is populated AFTER check_critical_and_report, " + "which reads it to decide which flags are inert") + + # and it must be handed the argv main was given, not sys.argv implicitly + rec_call = next(c for c in calls if c.func.id == "record_supplied_options") + assert any(isinstance(a, ast.Name) and a.id == "argv" for a in rec_call.args), \ + "record_supplied_options is not passed main()'s own argv" + + +def test_was_supplied_reads_tokens_not_values(): + """Structural: the guard must not infer 'user passed it' from the value.""" + o = _opts_from_argv(["--adapt-weight-exponent", "1.0"]) + assert drv.was_supplied(o, "--adapt-weight-exponent") is True + assert o.adapt_weight_exponent == 1.0, "value IS the default; only the token differs" + o2 = _opts_from_argv([]) + assert drv.was_supplied(o2, "--adapt-weight-exponent") is False + # and an options object built without parsing must not fabricate a conflict + assert drv.was_supplied(_Opts(), "--adapt-weight-exponent") is False + + +def test_auto_conflicts_raise_at_runtime(): + with pytest.raises(SystemExit) as e1: + drv.resolve_tempering_exponent( + _Opts(auto_adapt_weight_exponent=True, adapt_weight_exponent=0.5, + supplied=["--auto-adapt-weight-exponent", + "--adapt-weight-exponent"]), 4, 4800) + assert "would overwrite it" in str(e1.value) + with pytest.raises(SystemExit) as e2: + drv.resolve_tempering_exponent( + _Opts(auto_adapt_weight_exponent=True, adapt_adapt=True), 4, 4800) + assert "--adapt-adapt" in str(e2.value) + + +def test_law_refuses_the_same_domain_the_driver_does(): + """One domain, two enforcers -- they must not disagree.""" + for bad in (1.5, 0.0, -0.2): + with pytest.raises(ValueError): + export_ess_fraction(bad, 4) + + +def test_target_frac_given_is_decided_by_the_TOKEN_not_the_value(): + """RETARGETED. This used to require that _target_ess_was_given compare + against a named default constant. That whole mechanism was the defect: a + user who explicitly passes the default value was indistinguishable from one + who passed nothing. It now reads the command line. + """ + import ast as _ast + src = open(DRIVER).read() + fn = next(n for n in _ast.walk(_ast.parse(src)) if isinstance(n, _ast.FunctionDef) + and n.name == "_target_ess_was_given") + body = _ast.get_source_segment(src, fn) or "" + assert "was_supplied" in body, "still inferring from the value" + assert "_TARGET_EXPORT_ESS_FRAC_DEFAULT" not in body, ( + "still comparing against the default constant") + # behaviour: explicitly passing the default counts as supplied + o = _opts_from_argv(["--target-export-ess-frac", + str(drv._TARGET_EXPORT_ESS_FRAC_DEFAULT)]) + assert drv._target_ess_was_given(o) is True + + +def test_chooser_is_actually_CALLED_from_the_dispatch(): + """The wiring, not the helper. + + Every assertion above passes with `resolve_tempering_exponent` defined and + never invoked -- a perfectly dead chooser. Pin the CALL SITE: it must appear + inside analyze_one, guarded by the tempered-mode set, and be handed the + sampled dimension. (test_jax_fairdraw_export.py pins its own call sites the + same way, for the same reason.) + """ + tree = _driver_tree() + fn = next((n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "analyze_one"), None) + assert fn is not None, "analyze_one is missing from the driver" + calls = [n for n in ast.walk(fn) if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "resolve_tempering_exponent"] + assert calls, "analyze_one never calls resolve_tempering_exponent" + # gated on the tempered-mode set, not on a hand-written mode list + src = ast.get_source_segment(open(DRIVER).read(), fn) or "" + assert "_TEMPERED_MODES" in src + # and told what dimension it is choosing for + assert any(isinstance(a, ast.Name) and a.id == "dim" for c in calls for a in c.args), \ + "the chooser is called without the sampled dimension" + + +def test_degenerate_branch_warns_on_STDERR_and_names_the_floor(): + """RETARGETED. This required the branch to RAISE; it now warns by design. + + The refusal was a guarantee the estimate cannot support (see + test_degenerate_exponent_WARNS_and_does_not_refuse). What must survive is + that the branch still EXISTS, still keys on the shared floor constant, and + still reaches the user -- a warning printed to stdout among progress output + is not much better than silence, so stderr is pinned too. + """ + tree = _driver_tree() + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) + and n.name == "resolve_tempering_exponent") + guard = None + for n in ast.walk(fn): + if isinstance(n, ast.If) and any( + isinstance(x, ast.Name) and x.id == "_USABLE_EXPORT_ESS" + for x in ast.walk(n.test)): + guard = n + break + assert guard is not None, "the floor check is gone entirely" + src = ast.get_source_segment(open(DRIVER).read(), guard) or "" + assert "file=sys.stderr" in src, "the warning does not reach stderr" + assert "WARNING" in src + + +def test_guard_threshold_matches_the_message_the_driver_already_prints(): + """One definition, so the guard and fairdraw_indices' warning cannot drift. + + fairdraw_indices already tells the user at ESS < 200 that the cloud "is NOT a + usable posterior sample". If the guard used its own literal, the two could + disagree and the driver would refuse at one number while warning at another. + """ + import re + src = open(DRIVER).read() + assert re.search(r"^_USABLE_EXPORT_ESS = 200$", src, re.M), \ + "_USABLE_EXPORT_ESS is not defined as a single module-level constant" + # the pre-existing warning must be expressed through the same constant + assert "if neff < _USABLE_EXPORT_ESS:" in src, ( + "fairdraw_indices still hardcodes its own threshold; route it through " + "_USABLE_EXPORT_ESS so the guard and the warning cannot drift apart") + + +def test_chooser_runs_and_changes_the_exponent_end_to_end(): + """Run the DRIVER, not the helper: --help must render, and the chooser must + move the number. A subprocess is worth its seconds -- a parser-level typo is + invisible to every assertion above.""" + env = dict(os.environ, PYTHONPATH=CODE + os.pathsep + os.environ.get("PYTHONPATH", "")) + p = subprocess.run([sys.executable, DRIVER, "--help"], capture_output=True, + text=True, env=env, timeout=300) + assert p.returncode == 0, p.stderr[-3000:] + # Exact tokens, not substrings: "--auto-adapt-weight-exponent" is a substring + # of "--auto-adapt-weight-exponent-XX", so the substring form passed against a + # renamed flag. Match the whole option token. + import re + for flag in ("--auto-adapt-weight-exponent", "--target-export-ess-frac", + "--allow-degenerate-tempering"): + assert re.search(re.escape(flag) + r"(?![-\w])", p.stdout), ( + "%s does not appear as a whole option in --help" % flag) + + +def test_help_does_not_recommend_the_snr_keyed_rule(): + """RETIRED-claim guard. Positive assertions cannot catch a superseded claim + left standing beside the new one, so assert the ABSENCE of the rule this + change exists to keep off the JAX path. Append to RETIRED as it changes.""" + with open(DRIVER) as f: + src = f.read() + RETIRED = ("snr_fac", "0.1/np.power", "adapt-weight-exponent from the SNR") + for phrase in RETIRED: + assert phrase not in src, "retired SNR-keyed rule reappeared: %r" % phrase diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py new file mode 100644 index 000000000..faa5d4ab8 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py @@ -0,0 +1,378 @@ +"""Regression tests for adaptive, primitive-field time marginalization.""" +import ast +import inspect + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +jax.config.update("jax_enable_x64", True) +_trapezoid = getattr(np, "trapezoid", None) +if _trapezoid is None: + _trapezoid = np.trapz + +from RIFT.likelihood.jax_ile import anglemarg, core, wrapper + + +@pytest.mark.parametrize("n", [31, 32]) +def test_even_extension_reproduces_odd_and_even_samples(n): + rng = np.random.default_rng(10 + n) + x = rng.normal(size=(2, n)) + for factor in (2, 4, 8): + dense = np.asarray(core._reflected_fft_upsample(jnp.asarray(x), factor)) + assert dense.shape[-1] == (n - 1) * factor + 1 + np.testing.assert_allclose(dense[..., ::factor], x, atol=2e-12, rtol=0) + + +@pytest.mark.parametrize("n", [31, 32]) +def test_constant_integrand_has_exact_interval_normalization(n): + dt, c = 1.0 / 8192, 17.0 + row = jnp.full((2, n), c) + exact = c + np.log((n - 1) * dt) + for factor in (1, 2, 8, 32): + got, resolved = core._terminal_reflected_fft_at_factor(row, dt, factor) + np.testing.assert_allclose(np.asarray(got), exact, atol=2e-12, rtol=0) + assert np.all(np.asarray(resolved)) + + +def _event_b_like_row(phase=0.37): + n, dt, amp = 491, 1.0 / 8192, 600.0 ** 2 / 2.0 + x = np.arange(n) - n // 2 - phase + # Broad, reconstructible terminal lnL; exp(lnL) is narrower than one input + # sample by sqrt(amp), which is the Event-B high-SNR failure geometry. + return amp * np.exp(-0.5 * (x / 3.0) ** 2), dt, amp + + +def test_event_b_scale_adaptive_integral_matches_local_dense_truth(): + row, dt, amp = _event_b_like_row() + got = float(core._time_marginalize_reflected_primitive( + jnp.asarray(row[None, :], dtype=jnp.complex128), + jnp.zeros((1, row.size)), dt)[0]) + sigma_t = 3.0 * dt / np.sqrt(amp) + want = amp + np.log(np.sqrt(2.0 * np.pi) * sigma_t) + assert abs(got - want) < 2e-3 + + +def test_nonfinite_row_falls_back_to_historical_simpson(): + n, dt = 31, 1.0 / 4096 + row = np.linspace(-4.0, 0.0, n)[None, :] + row[0, 3] = -np.inf + w = jnp.asarray(core._simpson_weights(n, dt)) + got = core._time_marginalize_reflected_fft(jnp.asarray(row), dt, w) + want = core._time_marginalize(jnp.asarray(row), w) + np.testing.assert_allclose(np.asarray(got), np.asarray(want), rtol=0, atol=0) + + +def test_adaptive_primitive_value_gradient_and_hessian_are_finite_and_rematerialized(): + n, dt = 25, 1.0 / 4096 + x = jnp.arange(n, dtype=jnp.float64) - 12.2 + + def f(scale): + kappa = (scale * jnp.exp(-0.5 * (x / 3.0) ** 2))[None, :].astype( + jnp.complex128) + return core._time_marginalize_reflected_primitive( + kappa, jnp.zeros((1, n)), dt)[0] + + value = f(600.0) + grad = jax.grad(f)(600.0) + hess = jax.hessian(f)(600.0) + assert np.all(np.isfinite(np.asarray([value, grad, hess]))) + assert "jax.checkpoint(refine_one)" in inspect.getsource( + core._time_marginalize_reflected_primitive) + + +def test_phase_marginalization_refines_kappa_before_abs_near_nyquist(): + n, dt, amp = 17, 1.0, 5.0 + kappa = amp * (-1.0) ** np.arange(n) + rho = np.zeros((1, n)) + factor = 64 + dense = np.asarray(core._reflected_fft_upsample( + jnp.asarray(kappa[None, :]), factor))[0].real + x = np.arange((n - 1) * factor + 1) / factor + np.testing.assert_allclose(dense, amp * np.cos(np.pi * x), + rtol=0, atol=2e-12) + got = float(core._time_marginalize_reflected_primitive( + jnp.asarray(kappa[None, :]), jnp.asarray(rho), dt, + phase_marginalization=True)[0]) + # On the input grid abs(kappa) is identically amp, so interpolating the + # already-marginalized field returns this demonstrably wrong constant-row + # result. The band-limited primitive is amp*cos(pi*t), with zeros at every + # half sample; integrate that independent continuous model densely. + wrong = amp + np.log((n - 1) * dt) + x_truth = np.linspace(0.0, n - 1.0, (n - 1) * 8192 + 1) + y = np.exp(amp * np.abs(np.cos(np.pi * x_truth)) - amp) + want = amp + np.log(_trapezoid(y, x=x_truth)) + assert abs(want - wrong) > 0.5 + # This adversary has likelihood maxima at both window endpoints and lies + # outside the documented spectral-headroom regime. The reconstruction is + # mathematically correct, but the production adapter must refuse to trust + # a boundary condition carrying material posterior mass. + assert np.isnan(got) + + +@pytest.mark.parametrize("n_harmonics", [14, 88]) +def test_guarded_cosine_pad_matches_independent_dense_high_snr_truth(n_harmonics): + n, center, amp = 491, 245.37, 600.0 ** 2 / 2.0 + g_initial = 1 << int(np.ceil(np.log2(core.default_time_guard(n)))) + guard = 2 * g_initial + harmonics = np.arange(1, n_harmonics + 1, dtype=float) + + def primitive(t): + phase = 2.0 * np.pi * harmonics[:, None] * ( + np.asarray(t)[None, :] - center) / n + return amp * np.mean(np.cos(phase), axis=0) + + samples = primitive(np.arange(-guard, n + guard)) + got = float(core._time_marginalize_reflected_primitive( + jnp.asarray(samples[None, :], dtype=jnp.complex128), + jnp.zeros((1, samples.size)), 1.0, guard=guard)[0]) + factor_truth = 8192 + # At rho~600 the principal posterior peak is far narrower than this + # four-sample independent continuous interval; all omitted contributions + # underflow relative to it, while this avoids a needlessly huge K x N array. + t = np.arange(center - 2.0, center + 2.0 + 0.5 / factor_truth, + 1.0 / factor_truth) + truth = primitive(t) + peak = np.max(truth) + want = peak + np.log( + _trapezoid(np.exp(truth - peak), dx=1.0 / factor_truth)) + assert abs(got - want) < 1e-3 + + +def test_refinement_is_batch_composition_independent(): + sharp, dt, _ = _event_b_like_row(phase=0.41) + broad = 20.0 * np.exp(-0.5 * ((np.arange(sharp.size) - 245.1) / 20.0) ** 2) + w = jnp.asarray(core._simpson_weights(sharp.size, dt)) + alone = core._time_marginalize_reflected_fft(jnp.asarray(sharp[None, :]), dt, w) + paired = core._time_marginalize_reflected_fft( + jnp.asarray(np.stack((broad, sharp))), dt, w) + np.testing.assert_allclose(np.asarray(alone[0]), np.asarray(paired[1]), + rtol=0, atol=1e-10) + + +def test_all_terminal_kernels_expose_one_canonical_selector(): + kernels = [ + core.fused_log_likelihood, + core.fused_log_likelihood_distmarg, + core.fused_log_likelihood_phimarg, + core.fused_log_likelihood_distphimarg, + core.fused_log_likelihood_distphipsimarg, + core.fused_log_likelihood_distpsimarg, + anglemarg.fused_log_likelihood_distphipsimarg_exact, + anglemarg.fused_log_likelihood_distphipsimarg_laplace, + core.phi_ref_conditional_lnL, + ] + for fn in kernels: + assert "time_quadrature" in inspect.signature(fn).parameters, fn.__name__ + + +def test_all_wrappers_expose_quadrature_but_nonlinear_path_refuses_bandlimited(): + classes = [wrapper.JAXExtrinsicLikelihood, + wrapper.JAXDistanceMarginalizedLikelihood, + wrapper.JAXDistPhiMargLikelihood, + wrapper.JAXDistPhiPsiMargLikelihood, + wrapper.JAXDistPsiMargLikelihood] + for cls in classes: + assert "time_quadrature" in inspect.signature(cls.__init__).parameters + with pytest.raises(ValueError, match="primitive time fields"): + wrapper._validate_nonlinear_time_quadrature( + "bandlimited", "distance/phase marginalization") + required, g0, gcert = wrapper.bandlimited_storage_requirement( + 1.0 / 4096, 0.075) + assert g0 == 512 and gcert == 1024 + assert required > 0.15 # the historical 2:1 buffer is provably insufficient + accumulator_source = inspect.getsource(core._accumulate_unit) + assert "support_valid" in accumulator_source + assert "jnp.nan + 0.0j" in accumulator_source + + +@pytest.mark.parametrize("feature", [None, "freqresponse"]) +def test_missing_guard_support_poisoned_for_baseline_and_banded(monkeypatch, feature): + import types + monkeypatch.setattr(core, "compute_detamresponse", + lambda *args: jnp.ones((1,), dtype=jnp.complex128)) + monkeypatch.setattr(core, "time_delay_from_earth_center", + lambda *args: jnp.zeros((1,))) + monkeypatch.setattr(core, "spherical_harmonics_vectorized", + lambda *args, **kwargs: jnp.ones((1, 1), dtype=jnp.complex128)) + monkeypatch.setattr(core, "_banded_coefficients", + lambda *args: jnp.ones((1, 1), dtype=jnp.complex128)) + common = dict(lms=[(2, 2)], l_max=2, response=np.zeros((3, 3)), + location=np.zeros(3)) + if feature is None: + detector = dict(common, Q=jnp.ones((4, 1), dtype=jnp.complex128), + U=jnp.zeros((1, 1), dtype=jnp.complex128), + V=jnp.zeros((1, 1), dtype=jnp.complex128)) + else: + detector = dict(common, + Q_bank=jnp.ones((1, 4, 1), dtype=jnp.complex128), + U_bank=jnp.zeros((1, 1, 1, 1), dtype=jnp.complex128), + V_bank=jnp.zeros((1, 1, 1, 1), dtype=jnp.complex128)) + data = types.SimpleNamespace( + feature=feature, detectors={"X": detector}, detector_names=["X"], + band={"refl_idx": np.array([0])}, gmst=0.0, deltaT=1.0, + npts=3, tval0=1.0, tref_minus_epoch=lambda det: 0.0) + args = [jnp.zeros((1,))] * 5 + kappa, rho = core._accumulate_unit( + data, *args, interp="nearest", phase_marginalization=False, guard=2) + assert np.all(np.isnan(np.asarray(kappa))) + assert np.all(np.isnan(np.asarray(rho))) + + +def test_jax_driver_uses_conventional_ile_flag_names(): + import pathlib + driver = pathlib.Path(__file__).parents[2] / "bin" / "integrate_likelihood_extrinsic_jax" + src = driver.read_text() + for flag in ("--time-marginalization-quadrature", + "--resample-time-marginalization", + "--srate-resample-time-marginalization", + "--time-posterior-export", + "--interpolate-time"): + assert flag in src + + +def _declared_option_actions(path): + """Return long option -> arity class without executing an ILE driver.""" + tree = ast.parse(path.read_text()) + out = {} + pinnable = [] + for node in ast.walk(tree): + if (isinstance(node, ast.Assign) and + any(isinstance(target, ast.Name) and + target.id == "LIKELIHOOD_PINNABLE_PARAMS" + for target in node.targets)): + pinnable = ast.literal_eval(node.value) + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and + isinstance(node.func, ast.Attribute) and + node.func.attr == "add_option"): + continue + names = [arg.value for arg in node.args + if isinstance(arg, ast.Constant) and + isinstance(arg.value, str) and arg.value.startswith("--")] + if not names: + continue + action = "store" + for keyword in node.keywords: + if (keyword.arg == "action" and + isinstance(keyword.value, ast.Constant)): + action = keyword.value.value + if action in ("store_true", "store_false"): + arity = "bool" + elif action == "append": + arity = "append" + else: + arity = "value" + out[names[0]] = arity + # Conventional ILE registers these options in a loop, so there is no + # literal ``--flag`` argument for the AST walk above to discover. + out.update(("--" + name.replace("_", "-"), "value") + for name in pinnable) + return out + + +def test_jax_dropin_manifest_covers_every_batchmode_option_with_same_arity(): + """A new conventional ILE flag must not make executable swapping fail. + + This is intentionally mechanical. The compatibility table had drifted by + 77 options while still claiming complete coverage; checking a hand-picked + list of headline flags did not detect that class of failure. + """ + import pathlib + code = pathlib.Path(__file__).parents[2] + conventional = _declared_option_actions( + code / "bin" / "integrate_likelihood_extrinsic_batchmode") + drv = _load_driver() + parser = drv.build_parser() + jax_actions = {} + for option in parser._get_all_options(): + if option.action in ("store_true", "store_false"): + arity = "bool" + elif option.action == "append": + arity = "append" + else: + arity = "value" + for name in option._long_opts: + jax_actions[name] = arity + missing = sorted(set(conventional) - set(jax_actions)) + mismatched = sorted( + (name, conventional[name], jax_actions[name]) + for name in set(conventional) & set(jax_actions) + if conventional[name] != jax_actions[name]) + assert not missing + assert not mismatched + + +def _load_driver(): + import importlib.machinery + import importlib.util + import pathlib + path = pathlib.Path(__file__).parents[2] / "bin" / "integrate_likelihood_extrinsic_jax" + loader = importlib.machinery.SourceFileLoader("_jax_tmarg_driver", str(path)) + spec = importlib.util.spec_from_loader(loader.name, loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +def test_driver_parses_readback_and_conflict_checks_ile_aliases(): + drv = _load_driver() + parser = drv.build_parser() + argv = ["--time-marginalization-quadrature", "bandlimited", + "--interpolate-time", "sinc"] + opts, _ = parser.parse_args(argv) + drv.record_supplied_options(opts, argv, parser) + drv.resolve_ile_interface_aliases(opts, parser) + drv.check_critical_and_report(opts, parser) + assert opts.time_marginalization_quadrature == "bandlimited" + assert opts.interp == "sinc" + + # The high-level ILE_extr job emits these conventional export arguments. + # JAX keeps time terminally marginalized, but executable substitution must + # not die during option parsing or the compatibility check. + argv = ["--resample-time-marginalization", + "--srate-resample-time-marginalization", "8192", + "--time-posterior-export", "grid"] + opts, _ = parser.parse_args(argv) + drv.record_supplied_options(opts, argv, parser) + drv.check_critical_and_report(opts, parser) + + # Legacy conventional ILE spellings remain valid when an old args_ile.txt + # is pointed at the JAX executable. + for value, expected in (("True", "cubic"), ("False", "nearest")): + argv = ["--interpolate-time", value] + opts, _ = parser.parse_args(argv) + drv.record_supplied_options(opts, argv, parser) + drv.resolve_ile_interface_aliases(opts, parser) + assert opts.interp == expected + assert drv._normalize_interpolate_time_argv( + ["--interpolate-time", "--gpu"]) == [ + "--interpolate-time", "True", "--gpu"] + + argv = ["--interp", "linear", "--interpolate-time", "sinc"] + opts, _ = parser.parse_args(argv) + drv.record_supplied_options(opts, argv, parser) + with pytest.raises(SystemExit): + drv.resolve_ile_interface_aliases(opts, parser) + + +def test_headline_phase_marginalized_export_keeps_sky_and_psi_not_phi_or_time(tmp_path): + import types + drv = _load_driver() + opts = types.SimpleNamespace( + output_file=str(tmp_path / "ile"), save_samples=True, seed=19, + mode="nuts", phase_marginalization=True) + theta = np.arange(24.0).reshape(4, 6) + drv.write_samples(opts, 0, theta, np.arange(4.0), with_distance=True) + path = tmp_path / "ile_0_samples.dat" + lines = path.read_text().splitlines() + assert lines[0].lstrip("# ") == ( + "right_ascension declination distance inclination psi loglikelihood") + assert "phi_orb" not in lines[0] + assert "t_ref" not in lines[0] + assert "phase=analytically-marginalized" in lines[1] + values = np.loadtxt(path) + np.testing.assert_array_equal(values[:, 0], theta[:, 0]) + np.testing.assert_array_equal(values[:, 1], theta[:, 1]) + np.testing.assert_array_equal(values[:, 4], theta[:, 2]) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py new file mode 100644 index 000000000..27109ab95 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py @@ -0,0 +1,386 @@ +"""Band-limited time marginalization for the JAX ILE likelihood. + +The defect: `_time_marginalize` integrates exp(lnL_t) with fixed Simpson +weights at the DATA sample spacing, while the integrand's width +sigma_t = 1/(2 pi rho sigma_f) SHRINKS as the signal gets louder. Measured on +a 35+30 Msun HLV injection at rho=40: sigma_t = 61.2 us against grid spacings +of 244/122/61 us at srate 4096/8192/16384 -- under-resolved at the rates people +use, worse at higher SNR. Simpson makes it worse rather than safer, because +Simpson = (4 T_h - T_2h)/3 carries the coarser T_2h alias. + +The fix costs no likelihood evaluations: kappa(t) is band-limited, so the +samples already computed determine the continuous integrand exactly. + +What the reconstruction does cost is GUARD SAMPLES. The window is a crop of a +longer correlation buffer, its two ends do not join, and the FFT that does the +reconstruction has no choice but to treat them as though they did -- so the +seam's fictitious jump rings into the inserted samples, invisibly, since every +retained sample stays exact. The guard samples put that seam outside the +integrated window; the tests below pin both halves of that (the defect on a +non-periodic crop, and its removal), because an exactness check written only +with modes that fit a whole number of periods in the crop passes either way. +""" +import numpy as np +import pytest +import jax +import jax.numpy as jnp + +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile.core import ( + _upsample_bandlimited, _time_marginalize, _time_marginalize_bandlimited, + _simpson_weights, _norm_is_arrival_time_dependent, fused_log_likelihood, + _guarded_window, default_time_guard) + + +def test_upsampling_is_exact_for_a_band_limited_signal(): + """Not 'accurate' -- EXACT. This is the sampling theorem, so agreement is + at machine precision, and anything worse means the padding or the Nyquist + split is wrong.""" + n, factor = 64, 8 + t = np.arange(n) / n + # content strictly below Nyquist for the coarse grid + x = sum(np.exp(2j * np.pi * k * t) * (0.7 ** k) for k in range(1, 12)) + fine = np.asarray(_upsample_bandlimited(jnp.asarray(x), factor)) + tf = np.arange(n * factor) / (n * factor) + exact = sum(np.exp(2j * np.pi * k * tf) * (0.7 ** k) for k in range(1, 12)) + err = np.max(np.abs(fine - exact)) + assert err < 1e-12, "band-limited upsampling is not exact: %.3e" % err + # and it must reproduce the original samples where they sit + assert np.max(np.abs(fine[::factor] - x)) < 1e-12 + + +def test_nyquist_bin_is_split_not_dumped(): + """A real signal upsampled must stay real. Dumping the +n/2 bin into one + side instead of splitting it leaves an imaginary part oscillating at + Nyquist -- small, grid-phase dependent, and exactly the class of error this + change removes.""" + n, factor = 32, 4 + t = np.arange(n) / n + x = np.cos(2 * np.pi * (n // 2) * t) + 0.3 * np.cos(2 * np.pi * 3 * t) + fine = np.asarray(_upsample_bandlimited(jnp.asarray(x + 0j), factor)) + assert np.max(np.abs(fine.imag)) < 1e-12, ( + "upsampling a real signal produced an imaginary part %.3e -- the " + "Nyquist bin is not being split evenly" % np.max(np.abs(fine.imag))) + + +_TONE_N, _TONE_K0, _TONE_START = 1024, 17, 256 + + +def _cropped_tone(npts, guard, factor): + """A CROP of a globally band-limited tone, and the exact continuous function. + + exp(2j pi k0 t / N) with k0/N = 0.0166 cycles/sample is band-limited far + below Nyquist, so nothing here is about resolution: the whole difficulty is + that the crop is a window cut out of a longer buffer, exactly as the + accumulators' window is cut out of the rholm correlation buffer. k0*npts/N + is not an integer, so the crop's ends do not join, which is the ordinary case + -- and the one a periodic reconstruction gets wrong. + + Returns ``(coarse_samples, exact_fine, sl)``: ``npts + 2*guard`` coarse + samples, the true function on the fine grid of the KEPT window, and the slice + of the upsampled array that window occupies -- the same arithmetic + ``_time_marginalize_bandlimited`` does. + """ + w = 2.0 * np.pi * _TONE_K0 / _TONE_N + origin = _TONE_START - guard + coarse = origin + np.arange(npts + 2 * guard) + sl = slice(guard * factor, guard * factor + (npts - 1) * factor + 1) + fine = origin + np.arange(sl.start, sl.stop) / float(factor) + return np.exp(1j * w * coarse), np.exp(1j * w * fine), sl + + +def test_periodic_upsampling_of_a_crop_is_wrong_between_the_samples(): + """THE DEFECT the guard exists for, and it hides from the obvious check. + + Crop a band-limited tone out of a longer buffer and upsample it: every + RETAINED sample is still exact -- the interpolant passes through its inputs + -- while the INSERTED samples, the ones the quadrature actually integrates, + carry the ringing from the seam the FFT invents between the two ends. So + `fine[::factor] == x` proves nothing, and neither does an exactness test + built from modes that fit a whole number of periods in the window. + """ + npts, factor = 64, 8 + x, exact, sl = _cropped_tone(npts, 0, factor) + fine = np.asarray(_upsample_bandlimited(jnp.asarray(x), factor)) + assert np.max(np.abs(fine[::factor] - x)) < 1e-12, ( + "the reconstruction no longer reproduces its own input samples") + err = np.max(np.abs(fine[sl] - exact)) + assert err > 1e-2, ( + "this test does not BITE: unguarded reconstruction of a NON-PERIODIC " + "crop is off by only %.3e, so it would not catch the guard being " + "dropped. Check the crop really is non-periodic (k0*npts/N must not be " + "an integer)." % err) + + +def test_guard_samples_move_the_seam_out_of_the_reconstructed_window(): + """The fix: reconstruct from a window widened by guard samples and keep only + the middle. The seam error falls off like 1/(distance from the seam), so the + guard buys accuracy in the region that is actually integrated -- the + regression this file was missing.""" + npts, factor, guard = 64, 8, 128 + x0, exact0, sl0 = _cropped_tone(npts, 0, factor) + err0 = np.max(np.abs( + np.asarray(_upsample_bandlimited(jnp.asarray(x0), factor))[sl0] - exact0)) + xg, exactg, slg = _cropped_tone(npts, guard, factor) + errg = np.max(np.abs( + np.asarray(_upsample_bandlimited(jnp.asarray(xg), factor))[slg] - exactg)) + assert errg < 1e-2, ( + "guarded reconstruction is still %.3e off inside the integrated window" + % errg) + assert errg < err0 / 5.0, ( + "guard samples bought almost nothing: %.3e guarded vs %.3e unguarded" + % (errg, err0)) + + +def _sharp_case(npts=256, deltaT=1.0 / 4096, sigma_samples=4.0, amp=600.0, + phase=0.0): + """kappa(t) BAND-LIMITED and well resolved, but LARGE. + + This is the production geometry, and getting it wrong is easy: it is not + that kappa is narrow -- kappa is smooth on the sample grid, which is why + band-limited reconstruction works at all. It is that exp(kappa) is narrow + BECAUSE kappa is large. Near the peak + kappa ~ amp - amp t^2 / (2 sigma_k^2) + so exp(kappa) has width sigma_k / sqrt(amp): with sigma_k = 4 dt and + amp = 600 that is 0.16 dt, i.e. sub-sample, while kappa itself spans ~4 + samples and is comfortably below Nyquist. That is exactly the regime + sigma_t = 1/(2 pi rho sigma_f) describes -- the integrand narrows as the + signal gets louder, the grid does not. + + This case is deliberately GUARD-FREE: the Gaussian is 32 sigma_k from either + end, so the window's ends join to the precision of exp(-512) and the periodic + seam has nothing to ring on. That is what makes it a clean probe of the + quadrature -- and also why it cannot substitute for the cropped-tone tests + below, which are where the seam actually bites. + """ + t = (np.arange(npts) - npts // 2) * deltaT + sigma_k = sigma_samples * deltaT + kappa = amp * np.exp(-0.5 * ((t - phase * deltaT) / sigma_k) ** 2) + return t, kappa, deltaT + + +def test_stock_simpson_is_grid_phase_dependent_and_bandlimited_is_not(): + """The BITING test. Slide the peak across one sample spacing: stock Simpson + swings by orders of magnitude more than the band-limited quadrature. + + This is what fails if someone reverts the fix, and it needs no reference + integral -- it is a self-consistency statement, since the true value cannot + depend on where the sampling grid happens to sit. + """ + vals_simpson, vals_bl = [], [] + for phase in np.linspace(0.0, 1.0, 9): + t, kappa, deltaT = _sharp_case(phase=phase) + k = jnp.asarray(kappa[None, :] + 0j) + rho = jnp.zeros((1, kappa.size)) + w = jnp.asarray(_simpson_weights(kappa.size, deltaT)) + vals_simpson.append(float(_time_marginalize(k.real, w)[0])) + # guard=0: legitimate here, and only here -- see _sharp_case. + vals_bl.append( + float(_time_marginalize_bandlimited(k, rho, deltaT, 16, 0)[0])) + span_s = max(vals_simpson) - min(vals_simpson) + span_b = max(vals_bl) - min(vals_bl) + assert span_b < 0.05, ( + "band-limited quadrature is still grid-phase dependent: span %.4f nats" + % span_b) + assert span_s > 20 * max(span_b, 1e-6), ( + "this test does not BITE: stock Simpson span %.4f vs band-limited " + "%.4f -- it would not catch a revert. Narrow the peak until it does." + % (span_s, span_b)) + + +def test_bandlimited_converges_in_the_upsample_factor(): + """Increasing the (free) reconstruction factor must stop changing the + answer -- otherwise the quadrature is not converged and the 'exact' + claim is empty.""" + t, kappa, deltaT = _sharp_case(phase=0.37) + k = jnp.asarray(kappa[None, :] + 0j) + rho = jnp.zeros((1, kappa.size)) + got = [float(_time_marginalize_bandlimited(k, rho, deltaT, f, 0)[0]) + for f in (8, 16, 32)] + assert abs(got[2] - got[1]) < 1e-3, ( + "not converged in the upsample factor: %r" % got) + + +def test_bandlimited_integrates_the_same_interval_as_simpson(): + """A CONSTANT integrand isolates the interval: the answer is then exactly + log(length), with no reconstruction error of any kind to hide behind. + + The upsampled array holds n*factor points spanning (n - 1/factor)*deltaT, + while the likelihood's window is (n-1)*deltaT -- the trailing factor-1 + samples are the periodic FFT continuation past the last data sample. + Integrating them renormalizes every lnL by log((n - 1/factor)/(n - 1)), + which is 0.0138 nats at n=64, factor=8: invisible in a self-consistency + scan, fatal when comparing against the stock quadrature. + """ + n, deltaT, c = 64, 1.0 / 4096, 3.0 + k = jnp.asarray(np.full((1, n), c) + 0j) + rho = jnp.zeros((1, n)) + exact = c + np.log((n - 1) * deltaT) + for factor in (1, 4, 8, 16): + got = float(_time_marginalize_bandlimited(k, rho, deltaT, factor, 0)[0]) + assert abs(got - exact) < 1e-12, ( + "factor %d integrates the wrong window: %.12f vs %.12f, a %+.4f-nat " + "normalization shift" % (factor, got, exact, got - exact)) + # ... and that is the interval the stock Simpson path uses, so the two agree + # on a constant instead of differing by a fixed offset. + w = jnp.asarray(_simpson_weights(n, deltaT)) + assert abs(float(_time_marginalize(k.real, w)[0]) - exact) < 1e-12 + + +def test_guard_samples_are_support_and_are_never_integrated(): + """A CONSTANT integrand isolates the interval again, now with a guard: the + answer must stay log((npts-1)*deltaT) whatever the guard is. Integrating the + guard samples instead of using them as support would renormalize every lnL by + log((npts-1+2*guard)/(npts-1)) -- 1.4 nats at npts=64, guard=32 -- and the + grid-phase scan above would not notice.""" + npts, deltaT, c = 64, 1.0 / 4096, 3.0 + exact = c + np.log((npts - 1) * deltaT) + for guard in (0, 1, 5, 32): + n = npts + 2 * guard + k = jnp.asarray(np.full((1, n), c) + 0j) + rho = jnp.zeros((1, n)) + for factor in (1, 4, 8): + got = float( + _time_marginalize_bandlimited(k, rho, deltaT, factor, guard)[0]) + assert abs(got - exact) < 1e-12, ( + "guard %d, factor %d integrates the wrong window: %.12f vs " + "%.12f, a %+.4f-nat normalization shift" + % (guard, factor, got, exact, got - exact)) + # and a guard that would leave nothing to integrate is an error, not a + # silently empty window + k = jnp.asarray(np.full((1, 8), c) + 0j) + with pytest.raises(ValueError, match="guard"): + _time_marginalize_bandlimited(k, jnp.zeros((1, 8)), deltaT, 4, 4) + + +def test_bandlimited_quadrature_has_no_default_guard(): + """No default, on purpose: guard=0 is the periodic-seam defect on any real + (cropped) window, so a caller that forgets it must get a TypeError rather + than a quietly wrong likelihood.""" + import inspect + sig = inspect.signature(_time_marginalize_bandlimited) + assert sig.parameters["guard"].default is inspect.Parameter.empty, ( + "guard acquired a default; an unguarded call must be impossible to make " + "by accident") + with pytest.raises(TypeError): + _time_marginalize_bandlimited(jnp.zeros((1, 8), dtype=jnp.complex128), + jnp.zeros((1, 8)), 1.0 / 4096, 4) + + +def test_both_accumulators_take_the_guarded_window_from_one_place(): + """The guard widens the window by shifting the gather offsets to + [-guard, npts+guard). An off-by-one there misplaces the arrival time of + every guarded evaluation, and a second copy of the rule in the banded + accumulator is how that off-by-one would arrive, so neither accumulator is + allowed to build its own offsets.""" + import inspect + from RIFT.likelihood.jax_ile import core as _core + + class _Stub(object): + npts = 8 + + npts, off = _guarded_window(_Stub(), 0) + assert npts == 8 and np.array_equal(np.asarray(off), np.arange(8)), ( + "guard=0 must be the production window, unchanged") + npts, off = _guarded_window(_Stub(), 3) + assert npts == 14 and np.array_equal(np.asarray(off), np.arange(-3, 11)), ( + "guarded window is not centred on the production window") + with pytest.raises(ValueError, match="guard"): + _guarded_window(_Stub(), -1) + + for fn in (_core._accumulate_unit, _core._accumulate_unit_banded): + src = inspect.getsource(fn) + assert "_guarded_window(data, guard)" in src, ( + "%s builds its own time offsets instead of taking them from " + "_guarded_window" % fn.__name__) + + +def test_fused_bandlimited_widens_the_window_and_forwards_the_guard(monkeypatch): + """The plumbing, which is where this defect could quietly come back: asking + for the band-limited quadrature must make the ACCUMULATOR gather guard + samples and hand the same number to the quadrature. Reconstructing an + unwidened window would be the original bug with a guard argument bolted on. + """ + from RIFT.likelihood.jax_ile import core as _core + + class _StubData(object): + feature = None + npts = 48 + deltaT = 1.0 / 4096 + distMpcRef = 1000.0 + w_t = jnp.asarray(_simpson_weights(48, 1.0 / 4096)) + + seen = {} + + def _fake_accumulate(data, ra, dec, psi, incl, phiref, interp, + phase_marginalization, guard=0): + seen["guard"] = int(guard) + n = data.npts + 2 * int(guard) + t = np.arange(n) - 0.5 * n + kappa = (7.0 * np.exp(-0.5 * (t / 3.0) ** 2) + 0.5)[None, :] + 0j + return jnp.asarray(kappa), jnp.zeros((1, n)) + + monkeypatch.setattr(_core, "_accumulate_unit", _fake_accumulate) + z = np.zeros(1) + dist = np.full(1, _StubData.distMpcRef) # invDist == 1, so kappa is as built + + got = float(_core.fused_log_likelihood( + _StubData(), z, z, z, z, z, dist, time_quad="bandlimited", + time_upsample=4)[0]) + assert seen["guard"] == default_time_guard(_StubData.npts), ( + "the band-limited path did not widen the accumulation window by the " + "default guard (got %r)" % (seen["guard"],)) + k, rho = _fake_accumulate(_StubData(), None, None, None, None, None, None, + False, guard=seen["guard"]) + want = float(_time_marginalize_bandlimited( + k, rho, _StubData.deltaT, 4, seen["guard"])[0]) + assert abs(got - want) < 1e-12, ( + "fused_log_likelihood did not pass its guard to the quadrature: " + "%.12f vs %.12f" % (got, want)) + + _core.fused_log_likelihood(_StubData(), z, z, z, z, z, dist, + time_quad="bandlimited", time_upsample=4, + time_guard=3) + assert seen["guard"] == 3, "explicit time_guard ignored (got %r)" % ( + seen["guard"],) + + _core.fused_log_likelihood(_StubData(), z, z, z, z, z, dist, + time_quad="simpson") + assert seen["guard"] == 0, ( + "the stock Simpson path must keep gathering exactly npts bins; it asked " + "for a guard of %r" % (seen["guard"],)) + + +def test_bandlimited_refuses_arrival_time_dependent_norms(): + """The band-limited quadrature reconstructs kappa alone and uses the model + norm at a single time bin. That is the whole norm for the baseline and + finite-size accumulators, but the slow-rotation post-phase makes + arrival-time dependent, where the first bin would be a different likelihood + rather than a better-integrated one. It must be refused, not approximated. + """ + class _StubData(object): + def __init__(self, feature): + self.feature = feature + + assert _norm_is_arrival_time_dependent(_StubData("rotation")) + assert not _norm_is_arrival_time_dependent(_StubData("freqresponse")) + assert not _norm_is_arrival_time_dependent(_StubData(None)) + + # The screen must run BEFORE anything touches the data, both so the stub + # suffices here and so the caller is not made to pay a full trace to be told + # the option is unavailable. + with pytest.raises(ValueError, match="arrival time"): + fused_log_likelihood(_StubData("rotation"), 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, + time_quad="bandlimited") + + +def test_unknown_time_quad_raises_rather_than_silently_defaulting(): + """A typo'd quadrature name must not quietly give the OLD behaviour.""" + from RIFT.likelihood.jax_ile.core import _TIME_QUAD_CHOICES + assert "bandlimited" in _TIME_QUAD_CHOICES and "simpson" in _TIME_QUAD_CHOICES + import inspect + from RIFT.likelihood.jax_ile import core as _core + src = inspect.getsource(_core.fused_log_likelihood) + assert "raise ValueError" in src, ( + "an unrecognised time_quad must raise, not fall through to the default") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py new file mode 100644 index 000000000..be770644b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py @@ -0,0 +1,147 @@ +"""Tests for the JAX joint (phi,psi) peak-local kernel.""" +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +from RIFT.likelihood import joint_angle_peak_local as JN +from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as JP + + +def _tables(seed=0, scale=1.0): + rng = np.random.default_rng(seed) + A = (rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3))) * scale + B = (rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5))) * scale + B[0, 2] = abs(B[0, 2].real) + 3.0 * scale + return A, B + + +def test_exported_names_exist(): + assert not [n for n in JP.__all__ if not hasattr(JP, n)] + + +@pytest.mark.parametrize("scale", [0.5, 3.0, 12.0]) +def test_inner_u_integral_is_exact(scale): + """The cell partition is a PARTITION, so this is exact, not truncated.""" + rng = np.random.default_rng(1) + f = jax.jit(JP.log_inner_u_integral) + u = np.linspace(0.0, 2 * np.pi, 400000, endpoint=False) + for _ in range(4): + c1 = scale * (rng.normal() + 1j * rng.normal()) + c2 = scale * (rng.normal() + 1j * rng.normal()) + g = (c1 * np.exp(1j * u)).real + (c2 * np.exp(2j * u)).real + m = g.max() + ref = m + np.log(np.exp(g - m).mean()) + np.log(2 * np.pi) + assert abs(float(f(0.0, complex(c1), complex(c2))) - ref) < 1e-4 + + +def test_both_signs_of_q_enter_the_u_coefficients(): + """Both +q and -q columns are stored, so c_q = D_{+q} + conj(D_{-q}). Using only + the +q column is survivable where roots are seeds, but here they define the + integration partition -- it was worth 17 nats at a single phi.""" + A, B = _tables(seed=3, scale=4.0) + C = JN.joint_table(A, B, x=0.9) + u = np.linspace(0.0, 2 * np.pi, 200000, endpoint=False) + f = jax.jit(JP.log_inner_u_integral) + for phi in np.linspace(0.0, 2 * np.pi, 5)[:4]: + a, c1, c2 = JP._a_c1_c2(jnp.asarray(C), jnp.atleast_1d(phi)) + g = JN.eval_g(C, np.full(u.size, phi), u) + m = g.max() + ref = m + np.log(np.exp(g - m).mean()) + np.log(2 * np.pi) + got = float(f(float(a[0]), complex(c1[0]), complex(c2[0]))) + assert abs(got - ref) < 1e-3, (phi, got, ref) + + +def test_spurious_off_circle_roots_do_not_orphan_an_arc(): + """MIDPOINT cells tile the circle for ANY four angles. Root-bounded cells do not + when two roots leave the unit circle as a conjugate-reciprocal pair, and the + orphaned arc silently cost 0.23 nats.""" + # search for a genuine off-circle case rather than hand-picking one: only SOME + # coefficient pairs push a conjugate-reciprocal pair off the unit circle. + rng = np.random.default_rng(0) + c1 = c2 = None + for _ in range(2000): + a1 = rng.normal() + 1j * rng.normal() + a2 = rng.normal() + 1j * rng.normal() + z = np.roots([a2, a1 / 2, 0, -np.conj(a1) / 2, -np.conj(a2)]) + if np.sum(np.abs(np.abs(z) - 1.0) > 1e-6) >= 2: + c1, c2 = 0.3 * a1, 0.3 * a2 + break + assert c1 is not None, "no off-circle fixture found" + z = np.roots([c2, c1 / 2, 0, -np.conj(c1) / 2, -np.conj(c2)]) + assert np.sum(np.abs(np.abs(z) - 1.0) > 1e-6) >= 2 + u = np.linspace(0.0, 2 * np.pi, 400000, endpoint=False) + g = (c1 * np.exp(1j * u)).real + (c2 * np.exp(2j * u)).real + m = g.max() + ref = m + np.log(np.exp(g - m).mean()) + np.log(2 * np.pi) + got = float(jax.jit(JP.log_inner_u_integral)(0.0, complex(c1), complex(c2))) + assert abs(got - ref) < 1e-4, (got, ref) + + +def test_matches_the_numpy_reference_kernel(): + """Two independent implementations of the same rule: the numpy one merges 2-D + regions, this one partitions by cells. They must agree.""" + A, B = _tables(seed=5, scale=3.0) + x = np.linspace(0.4, 2.2, 12) + logw = -0.5 * (x - 1.1) ** 2 * 4.0 + ref, ok, _ = JN.joint_marginalize_over_distance(A, B, x, logw, n_phi=64, + n_bound_grid=128) + assert ok + got = float(JP.joint_lnL_phi_dense(jnp.asarray(A), jnp.asarray(B), + jnp.asarray(x), jnp.asarray(logw), n_phi=256)) + assert abs(got - ref) < 1e-5, (got, ref) + + +def test_phi_chunking_is_a_memory_knob_only(): + """phi_chunk bounds the transient and must not move the answer.""" + A, B = _tables(seed=8, scale=2.0) + x = np.linspace(0.5, 1.8, 8) + logw = np.zeros_like(x) + out = [float(JP.joint_lnL_phi_dense(jnp.asarray(A), jnp.asarray(B), jnp.asarray(x), + jnp.asarray(logw), n_phi=128, phi_chunk=c)) + for c in (8, 16, 64)] + assert max(out) - min(out) < 1e-11, out + + +def test_required_n_phi_grows_like_sqrt_amplitude(): + """The phi axis is NOT localized here, so it must be SIZED, not guessed: hard-coding + it cost 191 nats at amplitude 1.25e4 during development.""" + a, b = JP.required_n_phi(100.0), JP.required_n_phi(10000.0) + assert b > a + assert 5.0 < (b / a) / np.sqrt(100.0) * 10.0 < 20.0 + + +# ------------------------------------------------------- differentiability + +def test_hessian_works_and_the_gradient_is_still_correct(): + """P1 from review. jnp.linalg.eigvals has NO second derivative in JAX, so any + Hessian through this kernel raised -- and the caller that matters, _fisher_whitening, + swallows that in an `except Exception` and returns None, so --fisher-precondition + would silently degrade to raw coordinates with the flag still recorded as supplied. + + The fix cuts the tangent before the eigensolve. That CHANGES the derivative, so the + gradient must be re-validated, not assumed: these angles are cell boundaries of an + exact partition, so a boundary shift adds to one cell exactly what it removes from + its neighbour and cancels.""" + F = lambda x: JP.log_inner_u_integral(0.0, x + 1j, 0.7 - 0.3j) + assert np.isfinite(float(jax.hessian(F)(2.0))) + g = jax.grad(F) + for x in (0.3, 2.0, 7.0, 25.0): + h = 1e-5 + fd = (float(F(x + h)) - float(F(x - h))) / (2 * h) + ad = float(g(x)) + assert abs(ad - fd) < 1e-4 * max(abs(fd), 1.0), (x, ad, fd) + + +def test_gradient_is_finite_as_the_quartic_leading_coefficient_vanishes(): + """P1 from review. The old guard caught only an EXACT zero; at c2 = 1e-20 the + companion matrix acquires ~1e20 entries and the eig JVP degenerates -- measured grad + 0.567 at c2=1, -1.2e14 at 1e-20 and nan at 1e-30, while the VALUE stayed fine. c2 is + the B-table q=+-2 coefficient and passes through small values at special geometries; + one nan gradient poisons a MALA/flowMC chain.""" + g = jax.grad(lambda c: JP.log_inner_u_integral(0.0, 2.0 + 1j, c * (1.0 + 0j))) + vals = [float(g(c2)) for c2 in (1.0, 1e-6, 1e-20, 1e-30)] + assert all(np.isfinite(v) for v in vals), vals + # and stable, not merely finite, across 24 orders of magnitude in c2 + assert abs(vals[1] - vals[3]) < 1e-3, vals diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py b/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py new file mode 100644 index 000000000..62479cf7e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py @@ -0,0 +1,449 @@ +#!/usr/bin/env python +""" +--limit-distance on the differentiable (JAX) arm. + +The batchmode ILE reaches distance through a SAMPLER (mcsampler add_parameter); +the JAX driver reaches it through a QUADRATURE GRID (make_distance_grid -> +JAXDist*MargLikelihood) or, in the explicitly-6-D modes, through sample_prior / +log_prior. Both arms must accept the same box, or the cross-sampler lnZ +comparison this option exists to enable is not a fair one. + +WHAT MUST HOLD. Narrowing the box changes what is INTEGRATED, never how the +prior is NORMALIZED. On this arm the trap has the same shape as on the other: +make_distance_grid ended with + + w = w / np.sum(w) # normalize the distance average + +which normalizes the prior onto whatever range the grid happens to span, so a +narrowed grid renormalizes onto itself and the marginal comes back looking +"unchanged" while the evidence scale has moved by the prior mass outside the box +(here: several nats). d_prior_range= splits the two roles. + +WHAT DOES NOT HOLD EXACTLY, AND WHY. Unlike the batchmode arm -- where the +narrowed and full-range estimators integrate the SAME function and agree to the +truncated likelihood mass -- the JAX marginalized arm changes the QUADRATURE +RESOLUTION when it changes the range: n_grid nodes over a narrow box resolve the +distance integrand better than n_grid nodes over [d_min,d_max]. So narrowed and +full-range lnL agree only to the FULL-RANGE grid's own discretization error, and +the narrowed one is the more accurate of the two (that being the point). The +tests below therefore compare BOTH against a converged reference rather than +asserting they agree with each other to a tolerance neither of them earns. +""" + +import os +import sys + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +# x64 EXPLICITLY, before any jax array is built. This file happens to get x64 +# for free today -- importing RIFT.likelihood.jax_ile turns it on as a side +# effect -- but the tolerances below (1e-10 nats on lnZ, np.array_equal on the +# default-identity checks) are float64 tolerances and would be unreachable in +# float32. A file whose precision depends on what some other module did on +# import is not tested; see PR #222's finding 4, where exactly that made three +# headline equivalence numbers float32 while CI stayed green. +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile.core import ( # noqa: E402 + make_distance_grid, make_distance_grid_adaptive) +from RIFT.likelihood.jax_ile.wrapper import ( # noqa: E402 + JAXDistanceMarginalizedLikelihood) + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from test_angle_marg_exact import make_synth # noqa: E402 + +D_MIN, D_MAX = 1.0, 20000.0 + + +def test_x64_is_on(): + """The rest of this file asserts float64 tolerances, so a float32 run must + fail HERE rather than quietly loosening every other assertion. + + HONEST LIMIT, so this is not read as more coverage than it is: x64 is held + on by TWO redundant mechanisms right now -- the explicit update at the top + of this file, and the side effect of importing RIFT.likelihood.jax_ile. + Deleting the explicit line therefore does NOT make this test fire, and a + mutation sweep that only deletes it reports a false INERT. What was + verified is that the assertion is not vacuous: forcing + jax.config.update("jax_enable_x64", False) immediately above it fails with + `assert False is True`. The test earns its place as a tripwire for the day + the import side effect is refactored away, not as a live mutation kill.""" + import jax.numpy as jnp + assert jax.config.jax_enable_x64 is True + assert jnp.zeros(1).dtype == jnp.float64 + + +### +### 1. No existing default may move +### + +@pytest.mark.parametrize('d_prior', ['euclidean', 'uniform']) +def test_uniform_grid_default_is_bitwise_unchanged(d_prior): + """d_prior_range == (d_min,d_max) must reproduce the historical + `w = w/np.sum(w)` branch bit for bit -- the driver now passes it always.""" + a = make_distance_grid(D_MIN, D_MAX, 256, d_prior) + b = make_distance_grid(D_MIN, D_MAX, 256, d_prior, d_prior_range=(D_MIN, D_MAX)) + assert np.array_equal(np.asarray(a[0]), np.asarray(b[0])) + assert np.array_equal(np.asarray(a[1]), np.asarray(b[1])) + + +def test_adaptive_grid_default_is_bitwise_unchanged(): + kw = dict(d_peak=500.0, sigma_d=25.0, d_prior='euclidean') + a = make_distance_grid_adaptive(D_MIN, D_MAX, **kw) + b = make_distance_grid_adaptive(D_MIN, D_MAX, d_prior_range=(D_MIN, D_MAX), **kw) + assert np.array_equal(np.asarray(a[0]), np.asarray(b[0])) + assert np.array_equal(np.asarray(a[1]), np.asarray(b[1])) + + +def test_full_range_grid_weights_still_sum_to_one(): + _, log_w = make_distance_grid(D_MIN, D_MAX, 256, d_prior_range=(D_MIN, D_MAX)) + assert float(np.sum(np.exp(np.asarray(log_w)))) == pytest.approx(1.0, rel=1e-9) + + +### +### 2. A narrowed grid carries the prior mass of the box, not unity +### + +def test_narrowed_grid_weights_are_the_box_prior_mass(): + lo, hi = 800.0, 3200.0 + _, log_w = make_distance_grid(lo, hi, 256, d_prior_range=(D_MIN, D_MAX)) + got = float(np.sum(np.exp(np.asarray(log_w)))) + analytic = (hi ** 3 - lo ** 3) / (D_MAX ** 3 - D_MIN ** 3) + # 1e-2 relative: the numerator and denominator use the same n_grid rectangle + # rule over ranges of different width, so they do not cancel exactly. + assert got == pytest.approx(analytic, rel=1e-2) + assert got < 0.05 # emphatically NOT renormalized to 1 + + +def test_narrowed_grid_without_the_prior_range_renormalizes_the_defect(): + """Power check: the historical call signature, given a narrow range, produces + unit weight -- i.e. it moved the prior. This is what d_prior_range prevents.""" + lo, hi = 800.0, 3200.0 + _, log_w = make_distance_grid(lo, hi, 256) # no d_prior_range + assert float(np.sum(np.exp(np.asarray(log_w)))) == pytest.approx(1.0, rel=1e-9) + + +def test_narrowed_ADAPTIVE_grid_weights_are_the_box_prior_mass(): + """The adaptive grid is a second, opt-in production path + (JAX_ILE_DISTGRID_ADAPTIVE=1) and it takes d_prior_range too. It was pinned + only for the NO-OP case (bitwise-unchanged, above), which passes whether or + not the narrowed branch works -- a mutation deleting that branch left this + file green. Found by mutation sweep on takeover; this is the test that kills + it, and the one below is its power check. + """ + lo, hi = 800.0, 3200.0 + kw = dict(d_peak=2000.0, sigma_d=100.0, d_prior='euclidean') + _, log_w = make_distance_grid_adaptive(lo, hi, d_prior_range=(D_MIN, D_MAX), **kw) + got = float(np.sum(np.exp(np.asarray(log_w)))) + analytic = (hi ** 3 - lo ** 3) / (D_MAX ** 3 - D_MIN ** 3) + # 3e-2 relative, LOOSER than the uniform grid's 1e-2 on purpose: the numerator + # and the denominator are two differently-shaped adaptive node sets, and + # _trapezoidal_spacing gives the end nodes a full cell rather than a half one, + # so the endpoint bias no longer cancels between them. Measured 1.87% here. + # It is a systematic, it is bounded, and it is pinned so it cannot grow quietly. + assert got == pytest.approx(analytic, rel=3e-2) + assert got < 0.05 # emphatically NOT renormalized to 1 + + +def test_narrowed_ADAPTIVE_grid_without_the_prior_range_renormalizes_the_defect(): + """Power check for the test above: the historical signature on a narrow range + returns unit weight, i.e. it moved the prior. Measured cost of that on this + box: +5.53 nats of evidence, silently.""" + lo, hi = 800.0, 3200.0 + kw = dict(d_peak=2000.0, sigma_d=100.0, d_prior='euclidean') + _, log_w = make_distance_grid_adaptive(lo, hi, **kw) # no d_prior_range + assert float(np.sum(np.exp(np.asarray(log_w)))) == pytest.approx(1.0, rel=1e-9) + + +def test_narrowed_grid_nodes_are_inside_the_box(): + """The disconnected-flag check at grid level: the nodes must actually move.""" + lo, hi = 800.0, 3200.0 + x, _ = make_distance_grid(lo, hi, 256, d_prior_range=(D_MIN, D_MAX)) + d = 1.0 / np.asarray(x) # x = distMpcRef/d, up to distMpcRef + assert d.min() / d.max() == pytest.approx(lo / hi, rel=1e-9) + + +### +### 3. Through the real marginalized likelihood +### +# The acceptance quantity here is the EVIDENCE, not the per-angle marginal. That +# distinction is not pedantic: at an angle where the signal is weak the +# distance-marginalized lnL is carried by the prior's own d^2 mass at large d (L -> 1 +# there), so narrowing the box moves it by many nats -- correctly. What must not move +# is lnZ, and it does not, because at high amplitude the peak angle carries the whole +# integral. So: a fixed prior-drawn angle cloud and lnZ = logsumexp(lnL) - ln N, which +# is exactly the driver's --mode prior-mc estimator. +# +# MEASURED on the configuration below (rho_mf 17.4, d* = 625 Mpc, prior [1,20000] Mpc, +# x8 multiplicative box [78,4999] = 1.6% of the prior mass), narrowed minus full-range +# lnZ at equal n_grid: +# +# n_grid 256 512 1000 2000 4000 +# dlnZ -3.0e-02 -4.8e-06 +2.8e-14 +0.0e+00 +2.8e-14 +# +# i.e. EXACTLY ZERO once the full-range grid resolves its own integrand at all. The +# n=256 entry is not a failure of the option: there the full-range grid is the wrong +# one (it sits +2.5e-02 from a converged reference while the box sits -5.8e-03), which +# is the resolution the narrowing exists to buy back. Both grids retain a shared +# O(1/n_grid) offset from the discrete normalization (-6.6e-04 at n=2000), which +# cancels identically between them -- that is why the equal-n comparison is the sharp +# one and the converged-reference comparison is the loose one. + +_NG_ACCEPT = 2000 # both calculations; the equal-n comparison is exact +_NG_CONVERGED = 8000 # converged-reference comparison, O(1/n) normalization offset + + +def _cloud(n=16, seed=7): + rng = np.random.default_rng(seed) + return (rng.uniform(0, 2 * np.pi, n), np.arcsin(rng.uniform(-1, 1, n)), + rng.uniform(0, np.pi, n), np.arccos(rng.uniform(-1, 1, n)), + rng.uniform(0, 2 * np.pi, n)) + + +def _lnZ(lnL): + from scipy.special import logsumexp + return float(logsumexp(np.asarray(lnL)) - np.log(len(lnL))) + + +@pytest.fixture(scope='module') +def _loud(): + """A high-amplitude synthetic, plus the box its own precompute implies. + + The box is derived, not guessed: the distance integrand per (angle, time bin) is + exp(K x - R x^2/2) with x = d_ref/d, so the best-fit distance at the dominant + sample is d_ref R/K. One _accumulate_unit call, no gradient ascent -- cheap and + reproducible. + """ + from RIFT.likelihood.jax_ile.core import _accumulate_unit, JAX_INTERP_DEFAULT + data = make_synth(scale=20.0, kappa_boost=20.0) + a5 = _cloud() + K, R = _accumulate_unit(data, *a5, JAX_INTERP_DEFAULT, False) + K = np.asarray(K.real) + R = np.maximum(np.asarray(R), 1e-30) + snr2 = np.where(K > 0, K * K / R, -np.inf) + i, j = np.unravel_index(int(np.argmax(snr2)), snr2.shape) + d_star = float(data.distMpcRef) * R[i, j] / K[i, j] + lo, hi = max(D_MIN, d_star / 8.0), min(D_MAX, d_star * 8.0) + return data, a5, (lo, hi) + + +def _like(data, lo, hi, n_grid): + return JAXDistanceMarginalizedLikelihood( + data, lo, hi, n_grid=n_grid, d_prior_range=(D_MIN, D_MAX)) + + +def test_the_box_is_actually_a_narrowing(_loud): + """Guard on the fixture: if the x8 window ever clips to the full range, every + test below would pass while measuring nothing.""" + _, _, (lo, hi) = _loud + mass = (hi ** 3 - lo ** 3) / (D_MAX ** 3 - D_MIN ** 3) + assert mass < 0.05 + + +def test_ACCEPTANCE_narrowing_leaves_the_evidence_alone(_loud): + """THE acceptance test on this arm: same n_grid, narrowed vs full range.""" + data, a5, (lo, hi) = _loud + lnz_full = _lnZ(_like(data, D_MIN, D_MAX, _NG_ACCEPT).log_likelihood(*a5)) + lnz_box = _lnZ(_like(data, lo, hi, _NG_ACCEPT).log_likelihood(*a5)) + assert abs(lnz_box - lnz_full) < 1e-10, \ + "narrowing moved lnZ by {:.3e} nats".format(lnz_box - lnz_full) + + +def test_ACCEPTANCE_narrowed_evidence_matches_a_converged_reference(_loud): + """And the shared residual is the discrete normalization, not the box: a + converged full-range calculation is within 1e-3 nats of the narrowed one.""" + data, a5, (lo, hi) = _loud + lnz_ref = _lnZ(_like(data, D_MIN, D_MAX, _NG_CONVERGED).log_likelihood(*a5)) + lnz_box = _lnZ(_like(data, lo, hi, _NG_ACCEPT).log_likelihood(*a5)) + assert abs(lnz_box - lnz_ref) < 1e-3 + + +def test_the_renormalizing_call_moves_the_evidence_by_the_prior_mass(_loud): + """Power check: the same box through the PRE-CHANGE call signature (no + d_prior_range) shifts every lnL, and hence lnZ, by ln(1/prior mass of the box) -- + MEASURED +4.1594 nats here. Without this, the acceptance test above could be + passing on an implementation that renormalized and cancelled its own error.""" + data, a5, (lo, hi) = _loud + good = np.asarray(_like(data, lo, hi, _NG_ACCEPT).log_likelihood(*a5)) + bad = np.asarray(JAXDistanceMarginalizedLikelihood( + data, lo, hi, n_grid=_NG_ACCEPT).log_likelihood(*a5)) + _, log_w = make_distance_grid(lo, hi, _NG_ACCEPT, d_prior_range=(D_MIN, D_MAX)) + expected = -np.log(float(np.sum(np.exp(np.asarray(log_w))))) + assert np.allclose(bad - good, expected, atol=1e-8) + assert expected > 1.0 # several nats, not a rounding difference + assert abs(_lnZ(bad) - _lnZ(good) - expected) < 1e-8 + + +### +### 4. Driver wiring +### + +_DRIVER = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', '..', 'bin', 'integrate_likelihood_extrinsic_jax') + + +@pytest.mark.skipif(not os.path.exists(_DRIVER), reason='JAX driver not in this tree') +def test_driver_resolves_and_forwards_the_box(): + """The driver must (a) define the option, (b) hand the marginalized + likelihoods the SAMPLED range as d_min/d_max and the PHYSICAL range as + d_prior_range, and (c) not leave --limit-distance in the accepted-but-ignored + set, which is how a flag ends up parsing and doing nothing here.""" + with open(_DRIVER) as f: + src = f.read() + assert 'g.add_option("--limit-distance"' in src + assert 'def resolve_distance_limit(opts):' in src + assert src.count('d_prior_range=(opts.d_min, opts.d_max)') == 4 + assert 'like_data, d_lo, d_hi' in src + assert '"--d-prior", "--limit-distance",' in src # in the `implemented` set + + +@pytest.mark.skipif(not os.path.exists(_DRIVER), reason='JAX driver not in this tree') +def test_driver_prior_normalization_stays_on_d_min_d_max(): + """log_prior() restricts SUPPORT to the box but must keep normalizing on + [d_min,d_max]; run_prior_mc() must correct for its restricted proposal.""" + with open(_DRIVER) as f: + src = f.read() + assert 'inb = inb & (dist >= d_lo) & (dist <= d_hi)' in src + assert 'dmin3, dmax3 = opts.d_min ** 3, opts.d_max ** 3' in src + assert 'logw = lnL - log_distance_box_correction(opts, with_distance)' in src + + +def _driver_ns(): + """The driver's distance-box functions, exec'd out of its source. + + Spans resolve_distance_limit -> log_distance_box_correction -> sample_prior -> + log_prior, which is the whole contract: the box decides what is DRAWN and what + is IN SUPPORT, while [d_min,d_max] decides the normalization. Executing them + together is deliberate -- pulling only one out is how the coupling between the + proposal range and the correction went unpinned in the first place. + """ + with open(_DRIVER) as f: + src = f.read() + ns = {'np': np} + start = src.index('def resolve_distance_limit(opts):') + end = src.index('def eval_lnL(like, theta, opts, with_distance):') + exec(compile(src[start:end], _DRIVER, 'exec'), ns) # noqa: S102 + return ns + + +def test_box_correction_is_exactly_zero_without_the_option(): + """The historical prior-MC path must be untouched: `lnL - 0.0` is bitwise lnL.""" + ns = _driver_ns() + + class _O(object): + d_min, d_max, limit_distance = 1.0, 20000.0, None + assert ns['resolve_distance_limit'](_O()) == (1.0, 20000.0) + corr = ns['log_distance_box_correction'](_O(), True) + assert corr == 0.0 and isinstance(corr, float) + lnL = np.array([-3.0, 0.0, 12.5, -np.inf]) + assert np.array_equal(lnL - corr, lnL) or np.all( + (lnL - corr == lnL) | np.isnan(lnL)) + + _O.limit_distance = '800,3200' + assert ns['resolve_distance_limit'](_O()) == (800.0, 3200.0) + assert ns['log_distance_box_correction'](_O(), True) == pytest.approx( + np.log((20000.0 ** 3 - 1.0 ** 3) / (3200.0 ** 3 - 800.0 ** 3))) + # inert when distance is marginalized out (no explicit distance proposal) + assert ns['log_distance_box_correction'](_O(), False) == 0.0 + + +def test_sample_prior_DRAWS_inside_the_box(): + """The proposal must actually move, not just the prior's support. + + run_prior_mc subtracts log_distance_box_correction on the understanding that + sample_prior draws from the prior RESTRICTED to the box. Nothing pinned that + understanding: a mutation making sample_prior draw over the full [d_min,d_max] + while the correction stayed left this file green, and the resulting estimator + is biased low by exactly the correction (5.5 nats on the box below), because + inside the box the proposal then equals the prior and the weight needs no + correction at all. This test pins the two halves together. + """ + ns = _driver_ns() + rng = np.random.default_rng(20260902) + + class _O(object): + d_min, d_max, limit_distance = 1.0, 20000.0, '800,3200' + + theta, logp = ns['sample_prior'](20000, _O(), rng, True) + d = theta[..., 5] + assert d.min() >= 800.0 and d.max() <= 3200.0, (d.min(), d.max()) + # and it must FILL the box, not sit in a corner of it + assert d.min() < 900.0 and d.max() > 3100.0 + assert np.all(np.isfinite(logp)) # every draw is in support + + # positive control: without the option the draws span the full prior range + _O.limit_distance = None + d_full = ns['sample_prior'](20000, _O(), rng, True)[0][..., 5] + assert d_full.max() > 15000.0 + + +def test_log_prior_is_minus_inf_outside_the_box(): + """The other half of the same contract, at the density rather than the draw.""" + ns = _driver_ns() + + class _O(object): + d_min, d_max, limit_distance = 1.0, 20000.0, '800,3200' + + inside = np.array([[1.0, 0.0, 1.0, 1.0, 1.0, 2000.0]]) + below = np.array([[1.0, 0.0, 1.0, 1.0, 1.0, 500.0]]) + above = np.array([[1.0, 0.0, 1.0, 1.0, 1.0, 9000.0]]) + assert np.isfinite(ns['log_prior'](inside, _O(), True)[0]) + assert ns['log_prior'](below, _O(), True)[0] == -np.inf + assert ns['log_prior'](above, _O(), True)[0] == -np.inf + # ... and the normalization inside is the FULL-range one, not the box's + got = float(ns['log_prior'](inside, _O(), True)[0]) + _O.limit_distance = None + assert got == pytest.approx(float(ns['log_prior'](inside, _O(), True)[0]), rel=1e-12) + + +### +### 7. Interactions with the distance quadrature that landed alongside (#221) +### + +def _phipsi(data, lo, hi, prior_range, **kw): + from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood + return JAXDistPhiPsiMargLikelihood( + data, lo, hi, nphi=8, npsi=8, n_grid=64, + guess_snr=40.0, angle_marg='exact', d_prior_range=prior_range, **kw) + + +def test_the_box_does_not_shrink_the_ANGLE_lattice(_loud): + """#221 sizes the angle lattice from a FULL-SUPPORT distance grid, and measured + what a narrowed one does to it: a [0.8 d, 1.25 d] window dropped the amplitude + 12.6% and the lattice from (624,320) to (592,304). + + --limit-distance IS that narrowing. Merging the two branches naively -- passing + the box as d_min/d_max and letting the full-support grid be built from it -- would + have silently shrunk the lattice on every boxed run. The resolution builds that + grid over d_prior_range instead. This pins it: the lattice and the amplitude are + properties of the PRIOR support, not of the box. + """ + data, a5, _ = _loud + # The box must RAISE THE LOWER EDGE to exercise this. estimate_angle_amplitude + # reads min/max of x = distMpcRef/d, so the amplitude comes from the SMALL-d end: + # measured on this fixture, [78,4999] and [1,120] both leave it at 2643.43 exactly, + # while [1500,4000] drops it to 501.41 -- 81%. The first version of this test used + # the fixture's own generous box and was INERT; the mutation sweep caught that. + lo, hi = 1500.0, 4000.0 + full = _phipsi(data, D_MIN, D_MAX, None) + boxed = _phipsi(data, lo, hi, (D_MIN, D_MAX)) + assert boxed.angle_marg_info['sample_grid'] == full.angle_marg_info['sample_grid'] + assert boxed.angle_marg_info['amp_sizing'] == pytest.approx( + full.angle_marg_info['amp_sizing'], rel=1e-12) + # ... and the box really did move the integration grid, or this proves nothing + assert float(np.asarray(boxed.x_grid).max()) < float(np.asarray(full.x_grid).max()) + + +def test_loguniform_grid_REFUSES_a_narrowed_box(_loud): + """make_distance_grid_loguniform takes no d_prior_range, so a box would + renormalize the prior onto itself -- several nats, silently, and exactly the + defect --limit-distance exists to prevent. Fail closed rather than compose two + features whose composition is unvalidated.""" + data, a5, (lo, hi) = _loud + with pytest.raises(ValueError, match='limit-distance'): + _phipsi(data, lo, hi, (D_MIN, D_MAX), dist_grid='loguniform') + # ... and it is NOT refused without a box: the refusal keys on the narrowing, + # not on d_prior_range merely being present. + _phipsi(data, D_MIN, D_MAX, (D_MIN, D_MAX), dist_grid='loguniform') diff --git a/MonteCarloMarginalizeCode/Code/test/test_advanced_parameter_ports.py b/MonteCarloMarginalizeCode/Code/test/test_advanced_parameter_ports.py new file mode 100644 index 000000000..e1530072a --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_advanced_parameter_ports.py @@ -0,0 +1,417 @@ +import os +import inspect +import ast +from pathlib import Path +import subprocess +import sys +from types import SimpleNamespace + +import numpy as np +import pytest + +import RIFT.lalsimutils as lalsimutils + + +class _FakeEOBRunModule: + last_parameters = None + + @staticmethod + def EOBRunPy(parameters): + _FakeEOBRunModule.last_parameters = parameters + amplitude = np.array([0.0, 1.0, 0.2]) + phase = np.zeros_like(amplitude) + return np.arange(3), None, None, {"1": [amplitude, phase]}, None + + +def test_hyperbolic_classification_uses_both_component_masses(monkeypatch): + monkeypatch.setenv("RIFT_TEOBRESUMS_SKIP_PROBE", "1") + monkeypatch.setattr(lalsimutils, "EOBRun_module", _FakeEOBRunModule, raising=False) + parameters = lalsimutils.ChooseWaveformParams( + m1=30 * lalsimutils.lal.MSUN_SI, + m2=20 * lalsimutils.lal.MSUN_SI, + dist=1e6 * lalsimutils.lal.PC_SI, + E0=1.02, + p_phi0=4.1, + ) + + assert parameters.extract_param("hypclass") == "scatter" + assert _FakeEOBRunModule.last_parameters["arg_out"] == "yes" + assert _FakeEOBRunModule.last_parameters["nqc"] == "no" + assert "LambdaAl2" in _FakeEOBRunModule.last_parameters + assert "LambdaBl2" in _FakeEOBRunModule.last_parameters + assert _FakeEOBRunModule.last_parameters["nqc_coefs_hlm"] == "none" + assert _FakeEOBRunModule.last_parameters["nqc_coefs_flx"] == "none" + assert _FakeEOBRunModule.last_parameters["use_geometric_units"] == "no" + assert _FakeEOBRunModule.last_parameters["interp_uniform_grid"] == "yes" + assert _FakeEOBRunModule.last_parameters["output_hpc"] == "no" + + +def test_hyperbolic_classification_is_distance_invariant(monkeypatch): + monkeypatch.setenv("RIFT_TEOBRESUMS_SKIP_PROBE", "1") + monkeypatch.setattr(lalsimutils, "EOBRun_module", _FakeEOBRunModule, raising=False) + + outcomes = [] + for distance_mpc in (1e-3, 1e24): + parameters = lalsimutils.ChooseWaveformParams( + m1=30 * lalsimutils.lal.MSUN_SI, + m2=20 * lalsimutils.lal.MSUN_SI, + dist=distance_mpc * 1e6 * lalsimutils.lal.PC_SI, + E0=1.02, + p_phi0=4.1, + ) + outcomes.append(parameters.extract_param("hypclass")) + + assert outcomes == ["scatter", "scatter"] + + +def test_hyperbolic_mode_generation_has_no_absolute_strain_classifier(): + source = inspect.getsource(lalsimutils.hlmoft) + assert "1e-26" not in source + + +def test_hyperbolic_endpoint_prefers_radial_dynamics(): + amplitude = np.array([0.0, 1.0, 0.5]) + assert lalsimutils._hyperbolic_endpoint_outcome( + {"Prstar": np.array([-0.2])}, amplitude + ) == "plunge" + assert lalsimutils._hyperbolic_endpoint_outcome( + {"Prstar": np.array([0.2])}, amplitude + ) == "scatter" + + +def test_zero_mode_data_zeros_every_mode(): + class Data: + def __init__(self, values): + self.data = np.asarray(values, dtype=np.complex128) + + class Series: + def __init__(self, values): + self.data = Data(values) + + modes = { + (2, 2): Series([1.0, 2.0]), + (2, -2): Series([3.0, 4.0]), + } + lalsimutils._zero_mode_data(modes) + assert all(np.count_nonzero(series.data.data) == 0 for series in modes.values()) + + +def test_hyperbolic_convert_arguments_do_not_truncate_other_exports(): + script = Path(__file__).parents[1] / "bin" / "helper_LDG_Events.py" + tree = ast.parse(script.read_text()) + hyperbolic_writes = [] + for node in ast.walk(tree): + if not isinstance(node, ast.If): + continue + test_source = ast.unparse(node.test) + if test_source != "opts.assume_hyperbolic": + continue + for call in ast.walk(node): + if not isinstance(call, ast.Call) or not isinstance(call.func, ast.Name): + continue + if call.func.id != "open" or len(call.args) < 2: + continue + if isinstance(call.args[0], ast.Constant) and call.args[0].value == "helper_convert_args.txt": + hyperbolic_writes.append(call.args[1].value) + assert hyperbolic_writes == ["a"] + + +def test_helper_defines_a6c_range_without_ini(): + script = Path(__file__).parents[1] / "bin" / "helper_LDG_Events.py" + tree = ast.parse(script.read_text()) + unconditional = [ + node + for node in tree.body + if isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "a6c_range_str" + for target in node.targets + ) + ] + # --use-EOB-parameters must not require an [engine] a6c_min entry + assert unconditional + + +def test_nrhybsur_tidal_routing_checks_the_approximant_family(): + script = Path(__file__).parents[1] / "bin" / "util_RIFT_pseudo_pipe.py" + source = script.read_text() + assert "('NRHybSur' and" not in source + assert source.count("'NRHybSur' in opts.approx") >= 2 + + +def test_cip_rejects_hyperbolic_class_filter_without_hyperbolic_mode(tmp_path): + script = ( + Path(__file__).parents[1] + / "bin" + / "util_ConstructIntrinsicPosterior_GenericCoordinates.py" + ) + env = os.environ.copy() + env["XDG_CACHE_HOME"] = str(tmp_path / "cache") + env["MPLCONFIGDIR"] = str(tmp_path / "matplotlib") + result = subprocess.run( + [sys.executable, str(script), "--force-scatter"], + capture_output=True, + text=True, + env=env, + ) + assert result.returncode == 2 + assert "require --use-hyperbolic" in result.stderr + + +def test_cip_preserves_base_rf_pca_and_rbf_fit_methods(): + script = ( + Path(__file__).parents[1] + / "bin" + / "util_ConstructIntrinsicPosterior_GenericCoordinates.py" + ) + tree = ast.parse(script.read_text()) + function_names = { + node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef) + } + fit_method_values = { + node.comparators[0].value + for node in ast.walk(tree) + if isinstance(node, ast.Compare) + and len(node.ops) == 1 + and isinstance(node.ops[0], ast.Eq) + and len(node.comparators) == 1 + and isinstance(node.comparators[0], ast.Constant) + and isinstance(node.left, ast.Attribute) + and isinstance(node.left.value, ast.Name) + and node.left.value.id == "opts" + and node.left.attr == "fit_method" + } + assert {"fit_rf_pca", "fit_rbf"} <= function_names + assert {"rf_pca", "rbf"} <= fit_method_values + + +def test_real_hyperbolic_classification_when_teob_is_available(monkeypatch): + eobrun_module = pytest.importorskip("EOBRun_module") + monkeypatch.setattr(lalsimutils, "EOBRun_module", eobrun_module, raising=False) + parameters = lalsimutils.ChooseWaveformParams( + m1=30 * lalsimutils.lal.MSUN_SI, + m2=30 * lalsimutils.lal.MSUN_SI, + dist=1e6 * lalsimutils.lal.PC_SI, + E0=1.0027, + p_phi0=4.0, + fmin=20, + deltaT=1 / 4096, + ) + + assert parameters.extract_param("hypclass") in { + "scatter", + "plunge", + "zoomwhirl", + "meaningless", + } + + +def _run_clean_ile_lines(tmp_path, rows, *options): + input_path = tmp_path / "ile.dat" + input_path.write_text( + "".join(" ".join(str(value) for value in row) + "\n" for row in rows) + ) + script = Path(__file__).parents[1] / "bin" / "util_CleanILE.py" + result = subprocess.run( + [sys.executable, str(script), *options, str(input_path)], + check=True, + capture_output=True, + text=True, + env=os.environ.copy(), + ) + return [line.split() for line in result.stdout.strip().splitlines()] + + +def _run_clean_ile(tmp_path, option, row): + return _run_clean_ile_lines(tmp_path, [row], option)[0] + + +def test_clean_ile_keeps_hyperbolic_columns(tmp_path): + row = [-1, 30, 20, 0, 0, 0, 0, 0, 0, 1.02, 4.1, 12, 0.1, 100, 30] + output = _run_clean_ile(tmp_path, "--hyperbolic", row) + + assert len(output) == 15 + assert output[9:11] == ["1.02", "4.1"] + + +def test_clean_ile_does_not_treat_a6c_as_distance(tmp_path): + row = [-1, 30, 20, 0, 0, 0, 0, 0, 0, -55, 12, 0.1, 100, 30] + output = _run_clean_ile(tmp_path, "--a6c", row) + + assert len(output) == 14 + assert output[9] == "-55.0" + + +def test_clean_ile_keeps_tidal_a6c_columns(tmp_path): + row = [-1, 2, 1.4, 0, 0, 0, 0, 0, 0, 400, 800, -55, 12, 0.1, 100, 30] + output = _run_clean_ile(tmp_path, "--a6c", row) + + assert len(output) == 16 + assert output[9:12] == ["400.0", "800.0", "-55.0"] + + +def test_clean_ile_keeps_hyperbolic_a6c_columns(tmp_path): + rows = [ + [-1, 30, 20, 0, 0, 0, 0, 0, 0, -55, 1.02, 4.1, 12, 0.1, 100, 30], + [-1, 30, 20, 0, 0, 0, 0, 0, 0, -35, 1.02, 4.1, 11, 0.1, 100, 30], + ] + lines = _run_clean_ile_lines(tmp_path, rows, "--hyperbolic", "--a6c") + + # distinct a6c values are distinct intrinsic points, not repeated evaluations + assert len(lines) == 2 + assert all(len(line) == 16 for line in lines) + assert sorted(line[9] for line in lines) == ["-35.0", "-55.0"] + assert all(line[10:12] == ["1.02", "4.1"] for line in lines) + + +def test_clean_ile_keeps_every_combined_advanced_column(tmp_path): + row = [-1, 30, 20, 0, 0, 0, 0, 0, 0, -55, 1.02, 4.1, 0.3, 1.1, 12, 0.1, 100, 30] + output = _run_clean_ile_lines( + tmp_path, [row], "--a6c", "--hyperbolic", "--eccentricity", "--meanPerAno" + )[0] + + assert len(output) == 18 + assert output[9] == "-55.0" + assert output[10:12] == ["1.02", "4.1"] + assert output[12:14] == ["0.3", "1.1"] + + +def test_clean_ile_keeps_tidal_columns_alongside_advanced_groups(tmp_path): + row = [-1, 2, 1.4, 0, 0, 0, 0, 0, 0, 400, 800, -55, 1.02, 4.1, 0.3, 12, 0.1, 100, 30] + output = _run_clean_ile_lines( + tmp_path, [row], "--a6c", "--hyperbolic", "--eccentricity" + )[0] + + assert len(output) == 19 + assert output[9:11] == ["400.0", "800.0"] + assert output[11] == "-55.0" + assert output[12:14] == ["1.02", "4.1"] + assert output[14] == "0.3" + + +def test_dag_postprocess_forwards_every_advanced_flag(): + script = Path(__file__).parents[1] / "bin" / "util_ILEdagPostprocess.sh" + source = script.read_text() + # every flag after the directory/label arguments reaches util_CleanILE.py + assert '"${CLEAN_FLAGS[@]}"' in source + # ... instead of a mutually exclusive dispatch on the first flag only + assert "'--eccentricity'" not in source + assert "'--hyperbolic'" not in source + + +def _extract_ile_output_block(): + """Source of the ILE .dat writer (hyperpipeline branch + legacy branch).""" + script = Path(__file__).parents[1] / "bin" / "integrate_likelihood_extrinsic_batchmode" + tree = ast.parse(script.read_text()) + for node in ast.walk(tree): + for field in ("body", "orelse", "finalbody"): + statements = getattr(node, field, None) + if not isinstance(statements, list): + continue + for index, statement in enumerate(statements): + if not ( + isinstance(statement, ast.If) + and ast.unparse(statement.test) == "_hpio.is_active()" + ): + continue + start = index + while start > 0 and not isinstance(statements[start - 1], ast.ImportFrom): + start -= 1 + assert start > 0, "hyperpipeline_io import not found above the writer" + return "\n".join( + ast.unparse(entry) for entry in statements[start - 1:index + 1] + ) + raise AssertionError("ILE output-format block not found") + + +class _CapturingNumpy: + def __init__(self): + self.rows = None + + @staticmethod + def array(values): + return values + + def savetxt(self, fname, rows): + self.rows = rows + + +def _legacy_ile_row(monkeypatch, lambda1=0.0, lambda2=0.0, **flags): + """Run the ILE .dat writer in legacy mode and return the row it emits.""" + monkeypatch.delenv("RIFT_HYPERPIPELINE_FORMAT", raising=False) + options = SimpleNamespace( + save_eccentricity=False, + save_meanPerAno=False, + save_EOB_parameters=False, + save_hyperbolic=False, + export_eos_index=False, + pin_distance_to_sim=False, + ) + for name, value in flags.items(): + assert hasattr(options, name) + setattr(options, name, value) + parameters = SimpleNamespace( + s1x=0.1, s1y=0.2, s1z=0.3, s2x=0.4, s2y=0.5, s2z=0.6, + lambda1=lambda1, lambda2=lambda2, + eccentricity=0.3, meanPerAno=1.1, + a6c=-55.0, E0=1.02, p_phi0=4.1, eos_table_index=7, + ) + recorder = _CapturingNumpy() + namespace = { + "opts": options, + "P": parameters, + "numpy": recorder, + "event_id": -1, + "m1": 30.0, + "m2": 20.0, + "log_res": 12.0, + "manual_avoid_overflow_logarithm": 0.0, + "sqrt_var_over_res": 0.1, + "sampler": SimpleNamespace(ntotal=100), + "neff": 30, + "pinned_params": {"distance": 410.0}, + "fname_output_txt": str(Path("unused.dat")), + } + exec(compile(_extract_ile_output_block(), "", "exec"), namespace) + assert recorder.rows is not None + return list(recorder.rows[0]) + + +def test_ile_hyperbolic_output_retains_eob_parameter(monkeypatch): + row = _legacy_ile_row(monkeypatch, save_EOB_parameters=True, save_hyperbolic=True) + + assert len(row) == 9 + 1 + 2 + 4 + assert row[9] == -55.0 + assert row[10:12] == [1.02, 4.1] + + +def test_ile_legacy_row_keeps_every_enabled_group(monkeypatch): + row = _legacy_ile_row( + monkeypatch, + lambda1=400.0, + lambda2=800.0, + save_eccentricity=True, + save_meanPerAno=True, + save_EOB_parameters=True, + save_hyperbolic=True, + ) + + # lambda1 lambda2 | a6c | E0 p_phi0 | eccentricity meanPerAno, in CIP order + assert len(row) == 9 + 2 + 1 + 2 + 2 + 4 + assert row[9:11] == [400.0, 800.0] + assert row[11] == -55.0 + assert row[12:14] == [1.02, 4.1] + assert row[14:16] == [0.3, 1.1] + assert row[16] == 12.0 # lnL still lands 4 columns from the end + + +def test_ile_legacy_row_preserves_unflagged_layout(monkeypatch): + assert len(_legacy_ile_row(monkeypatch)) == 13 + assert len(_legacy_ile_row(monkeypatch, pin_distance_to_sim=True)) == 14 + assert len(_legacy_ile_row(monkeypatch, lambda1=400.0, lambda2=800.0)) == 15 + assert len( + _legacy_ile_row( + monkeypatch, lambda1=400.0, lambda2=800.0, export_eos_index=True + ) + ) == 16 + assert len(_legacy_ile_row(monkeypatch, save_eccentricity=True)) == 14 diff --git a/MonteCarloMarginalizeCode/Code/test/test_calmarg_running_max_row_offset.py b/MonteCarloMarginalizeCode/Code/test/test_calmarg_running_max_row_offset.py new file mode 100644 index 000000000..5681ccadf --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_calmarg_running_max_row_offset.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_calmarg_running_max_row_offset is the SAME defect on the in-loop calibration +# marginalization (n_cal>1, cal_method='loop' -- the DEFAULT calmarg reduction; the +# fused kernel is opt-in behind --calibration-fused-kernel). Its streaming +# log-sum-exp offset `running_max` was a SCALAR over the whole +# (npts_extrinsic, npts_time) block, so the same >745-nat rows came back -inf +# (issue #232). numpy + lal, no GPU, ~10 s. All six CPU guards were +# mutation-checked before landing: reverting to the scalar max fails 4 of 6, a bare +# axis=-1 (no keepdims) fails 6 of 6, keeping the axis in the add-back fails 6 of 6, +# and dropping the all--inf-row guard fails 1 of 6. Its seventh test is a cupy leg +# and SKIPS here -- these runners have no GPU. +# --------------------------------------------------------------------------------- +"""The in-loop calibration-marginalization log-sum-exp offset must be PER EXTRINSIC +SAMPLE, not per batch. + +WHAT IS BEING TESTED, AND AGAINST WHAT +-------------------------------------- +With ``n_cal > 1`` and ``cal_method='loop'`` (the DEFAULT calmarg reduction -- the fused +kernel is opt-in behind ``--calibration-fused-kernel``), +``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`` marginalizes over calibration +realizations with a streaming log-sum-exp:: + + m_c = max(lnL_t_c) # over the WHOLE (npts_extrinsic, npts_time) block + running_max = max(running_max, m_c) + S += exp(lnL_t_c - running_max) + lnL = running_max + log(simps(S, dx=deltaT)) - log(n_cal) + +``running_max`` was a SCALAR shared by every extrinsic sample and every realization, so +every row was shifted by the LOUDEST row's peak. Any row sitting more than ~745 nats +below it underflows ``exp()`` to 0 across its whole time axis and across every +realization: its ``S`` row is 0, ``log(0) = -inf``, and the likelihood comes back +``-inf`` at a sample where it is finite and perfectly ordinary. With ``lnL ~ rho^2/2`` +at the peak and ``lnL ~ 0`` for a typical prior draw, that fires once ``max lnL > ~745`` +(``rho ~ 40``) and then applies to the BULK of the prior, not a tail -- which is what +collapses ``mcsamplerAV`` on loud events. See +oshaughnessy-junior/research-projects-RIT#232 for the real-data measurement of the same +defect on the ``n_cal == 1`` leg of this function, and #234 for that fix. + +THE REFERENCE IS THE SHIPPED CODE'S OWN ALREADY-CORRECT BRANCH. Five lines above the +site under test, the ``return_cal_components`` branch computes the same per-realization +time integral with ``m_raw = max(lnL_t_c, axis=-1, keepdims=True)`` -- per row, correctly +-- and returns it RAW (no importance weight). Each test below asks the shipped function +for those components and combines them by hand, + + lnL_ref = logsumexp_c( components[:, c] + cal_log_w[c] ) - log(n_cal), + +which is algebraically the identical quantity. Nothing is reimplemented: the same +detector response, the same Q window, the same loglikelihood callback and the same +Simpson rule produce both sides. The ONLY thing that differs is the offset, which is the +whole subject of the test. + +``keepdims=True`` is load-bearing and is guarded separately. With a bare ``axis=-1`` +the ``(n,)`` maximum broadcasts along the TIME axis instead of the sample axis. That +RAISES when ``npts_extrinsic != npts_time`` -- and is silently wrong when they are equal, +which is why one case below is deliberately square. The add-back +``running_max[..., 0]`` is guarded by the same cases: keeping the axis there broadcasts +the (n, 1) offset against the (n,) time integral into an (n, n) result. + + OMP_NUM_THREADS=1 PYTHONPATH=/MonteCarloMarginalizeCode/Code \ + python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_calmarg_running_max_row_offset.py +""" +from __future__ import print_function, division + +import os + +os.environ.setdefault("RIFT_LOWLATENCY", "1") + +import numpy as np +import pytest +from scipy.special import logsumexp + +import lal +import RIFT.lalsimutils as lsu +from RIFT.likelihood import factored_likelihood as fl + +# lal / lalsimutils are imported at module scope on purpose, NOT via importorskip: +# lalsuite is in requirements.txt and both CI jobs that run this file install it, so a +# missing lal here is a broken job, not an unsupported platform -- and an importorskip +# would turn that into a green skip. + +SRATE = 4096.0 +DELTAT = 1.0 / SRATE +NPTS = 614 # len(marginalization_time_grid(0.075, 1/4096)) +N_CAL = 2 +N_WINDOW = 4096 # per-realization block length; the gathered window + # (ifirst ~ 1894, npts 614) sits well inside one block +UNDERFLOW_NATS = 745.0 # -log(smallest positive float64 normal), roughly + + +def _kappa_buffer(): + """A band-limited, periodic-on-its-own-length kappa(t) block. + + Periodic so that whatever integer window the code gathers is a genuine segment of + it; the test does not need to predict ``ifirst``. + """ + ts = np.arange(N_WINDOW) * DELTAT + ms = np.arange(1, 400) + T = N_WINDOW * DELTAT + c = np.exp(-2j * np.pi * ms * (N_WINDOW // 2) * DELTAT / T) / (1.0 + (ms / 120.0) ** 2) + return np.exp(2j * np.pi * np.outer(ts, ms) / T) @ c + + +_BASE_KAPPA = _kappa_buffer() +# Per-realization amplitudes. Deliberately unequal, so the n_cal reduction is doing +# real work: the two realizations' time integrals differ by ~500 nats at the loud row, +# which is what makes the streaming rescale branch (S *= exp(...)) execute. +CAL_AMPS = (1.0, 0.7) + + +def _inputs(dists_Mpc): + """Minimal inputs that drive the SHIPPED NoLoop function on the numpy backend. + + One detector, one (l,m) pair and zero U/V cross terms, so the self-term ``rho_sq`` + vanishes and ``lnL_t`` is just the response-scaled ``Re kappa(t)`` times + ``distMpcRef/dist``. Distance is therefore a clean per-row amplitude knob: it sets + each extrinsic sample's peak ``lnL`` independently, which is exactly the axis this + test needs to separate. The rholm buffer holds ``N_CAL`` CONTIGUOUS blocks, which + is the layout the calmarg path assumes (realization c is selected by shifting the + window into block c). + """ + dists_Mpc = np.asarray(dists_Mpc, dtype=float) + n = dists_Mpc.size + P = lsu.ChooseWaveformParams() + P.deltaT = DELTAT + P.tref = 1000000000.0 + for name in ('phi', 'theta', 'phiref', 'incl', 'psi'): + setattr(P, name, np.zeros(n)) + P.dist = dists_Mpc * 1e6 * lal.PC_SI + blocks = np.concatenate([_BASE_KAPPA * a for a in CAL_AMPS[:N_CAL]]) + det = 'H1' + return (P, {det: np.asarray(blocks, dtype=complex)[None, :]}, + {det: np.array([[2, 2]])}, + {det: np.zeros((1, 1), dtype=complex)}, + {det: P.tref - 0.5}) + + +def _shipped(tvals, args, **kw): + P, rholms, lookupNK, ct, epochs = args + kw.setdefault('n_cal', N_CAL) + kw.setdefault('cal_method', 'loop') + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, Lmax=2, xpy=np, **kw) + + +def _per_row_reference(tvals, args, cal_log_weights=None, **kw): + """The cal-marginalized lnL built from the shipped function's OWN per-realization + components, which are already offset per row. + + lnL = log( (1/n_cal) sum_c exp(log_w_c) * int dt exp(lnL_t,c) ) + + ``return_cal_components`` returns log(int dt exp(lnL_t,c)) RAW -- before the + importance log-weight -- so the weights are folded in here, exactly as the loop + reduction folds them into ``lnL_t_c`` before its own log-sum-exp. + """ + comps = np.asarray(_shipped(tvals, args, return_cal_components=True, **kw), dtype=float) + assert comps.shape[-1] == N_CAL + log_w = np.zeros(N_CAL) if cal_log_weights is None else np.asarray(cal_log_weights, dtype=float) + return logsumexp(comps + log_w[None, :], axis=-1) - np.log(N_CAL) + + +@pytest.fixture(scope='module') +def tvals(): + grid = fl.marginalization_time_grid(0.075, DELTAT) + assert len(grid) == NPTS + return grid + + +def test_quiet_sample_stays_finite_beside_a_loud_one(tvals): + """The regression, in the shape production actually runs. + + Two extrinsic samples whose peak ``lnL`` differ by far more than the float64 + underflow budget. With a batch-wide ``running_max`` the quiet row returns ``-inf``; + with a per-row offset it returns its own, finite, correct value. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + args = _inputs(dists) + + comps = np.asarray(_shipped(tvals, args, return_cal_components=True), dtype=float) + assert comps.shape == (2, N_CAL) + # The premise of the test: the rows really are separated by more than the underflow + # budget, so a batch-wide offset MUST kill the quiet one. If the harness ever stops + # producing that separation this assert says so, instead of passing vacuously. + assert comps[0].max() - comps[1].max() > 2.0 * UNDERFLOW_NATS, comps + + lnL = np.asarray(_shipped(tvals, args)) + assert lnL.shape == (2,), lnL.shape # (2, 2) means the add-back kept its axis + assert np.all(np.isfinite(lnL)), lnL + ref = _per_row_reference(tvals, args) + np.testing.assert_allclose(lnL, ref, rtol=0, atol=1e-9) + + +def test_offset_is_per_row_even_when_the_batch_is_square(tvals): + """``keepdims=True``, guarded where its absence is SILENT. + + ``npts_extrinsic == npts_time`` on purpose: that is the one shape in which a bare + ``axis=-1`` maximum broadcasts along the wrong axis without raising. All rows but + the first are identical by construction, so they must return identical values -- an + offset that leaks across the sample axis does not. + """ + dists = np.full(NPTS, fl.distMpcRef * 1.0) + dists[0] = fl.distMpcRef / 80.0 + args = _inputs(dists) + + comps = np.asarray(_shipped(tvals, args, return_cal_components=True), dtype=float) + assert comps.shape == (NPTS, N_CAL) # square in (npts_extrinsic, npts_time) + assert comps[0].max() - comps[1].max() > 2.0 * UNDERFLOW_NATS, comps[:2] + + lnL = np.asarray(_shipped(tvals, args)) + assert lnL.shape == (NPTS,), lnL.shape + assert np.all(np.isfinite(lnL)), lnL[~np.isfinite(lnL)] + # Identical inputs -> identical outputs, whatever else is in the batch. + assert np.all(lnL[1:] == lnL[1]), np.unique(lnL[1:]).size + np.testing.assert_allclose(lnL, _per_row_reference(tvals, args), rtol=0, atol=1e-9) + + +def test_return_lnLt_timeseries_is_offset_per_row_too(tvals): + """The ``return_lnLt`` calmarg branch shares ``running_max`` and the same defect. + + That branch returns the cal-marginalized ``lnL(t)``; the driver resamples it to draw + an event time (``--time-marginalization`` with the resampling output). A quiet row + came back ``-inf`` at EVERY time bin, so the row carried no time information at all. + The check is the branch's own identity with the scalar return: + ``log int dt exp(lnL_t) == lnL``, taken row by row. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + args = _inputs(dists) + + lnLt = np.asarray(_shipped(tvals, args, return_lnLt=True), dtype=float) + assert lnLt.shape == (2, NPTS), lnLt.shape + # The loud row legitimately underflows FAR from its peak (a 1685-nat peak in a + # 614-bin window); what must never happen is a row with no finite bin at all. + assert np.any(np.isfinite(lnLt[1])), "quiet row is -inf at every time bin" + + lnL = np.asarray(_shipped(tvals, args)) + for i in range(2): + m = lnLt[i].max() + got = m + np.log(fl.my_simps(np.exp(lnLt[i] - m), dx=DELTAT)) + assert abs(got - lnL[i]) < 1e-9, (i, got, lnL[i]) + + +def test_nonuniform_cal_weights_are_carried_through(tvals): + """Importance-weighted cal draws (``--calibration-proposal-breadcrumb``) too. + + The weight is folded into ``lnL_t_c`` BEFORE the offset is taken, so a per-row + offset must be taken after the fold. A weight large enough to reorder which + realization dominates makes the ordering observable. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + args = _inputs(dists) + log_w = np.array([-3.0, +3.0]) # mean-1 weights are not required here + + lnL = np.asarray(_shipped(tvals, args, cal_log_weights=log_w)) + assert lnL.shape == (2,), lnL.shape + assert np.all(np.isfinite(lnL)), lnL + np.testing.assert_allclose( + lnL, _per_row_reference(tvals, args, cal_log_weights=log_w), rtol=0, atol=1e-9) + + +def test_a_row_that_is_minus_inf_everywhere_stays_minus_inf_not_nan(tvals): + """A per-row offset must not turn an empty row into ``nan``. + + A scalar offset was shielded from this by any other finite row in the batch: + ``exp(-inf - finite) = 0``. A per-row offset is not -- ``exp(-inf - -inf) = nan`` + -- and a nan in ``S`` is permanent, so the guard ships with the per-row offset. + A distance-marginalization callback that rejects a whole row (out-of-table distance) + is the shipped way to produce one. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + args = _inputs(dists) + + def reject_first_row(kappa_sq, rho_sq): + out = fl._factored_lnL_helper(kappa_sq, rho_sq) + out = np.array(out, dtype=float, copy=True) + out[0, :] = -np.inf + return out + + lnL = np.asarray(_shipped(tvals, args, loglikelihood=reject_first_row)) + assert lnL.shape == (2,), lnL.shape + assert not np.any(np.isnan(lnL)), lnL + assert lnL[0] == -np.inf, lnL + assert np.isfinite(lnL[1]), lnL + + +def test_onset_is_the_underflow_budget_not_a_general_offset_error(tvals): + """Below the underflow budget the two offsets agree; above it they cannot. + + This pins the MECHANISM rather than the symptom. A batch-wide offset is harmless + while every row is within ~745 nats of the batch peak -- the expression is + offset-invariant apart from rounding -- so a separation well under the budget must + still come out finite and correct, and only a separation over it may break. A + "fix" that changed the integral itself rather than only its offset would fail this, + and it PASSES on the unpatched code by design. + + The separations are MEASURED off the shipped components rather than prescribed, so + the case stays meaningful if the harness buffer or the response factor is retuned: + the loop skips any ladder rung that has drifted over the budget, and the test then + checks that at least one rung within a few hundred nats of it survived. + """ + checked = [] + for scale in (16.0, 24.0, 32.0): + dists = np.array([fl.distMpcRef / scale, fl.distMpcRef * 1.0]) + args = _inputs(dists) + comps = np.asarray(_shipped(tvals, args, return_cal_components=True), dtype=float) + sep = float(comps[0].max() - comps[1].max()) + if sep >= UNDERFLOW_NATS: # over budget: not this test's subject + continue + checked.append(sep) + lnL = np.asarray(_shipped(tvals, args)) + assert np.all(np.isfinite(lnL)), (sep, lnL) + np.testing.assert_allclose( + lnL, _per_row_reference(tvals, args), rtol=0, atol=1e-9) + # Non-vacuity: a ladder that only ever reached a 10-nat separation would prove + # nothing about a batch-wide offset, which is harmless at 10 nats. + assert checked and max(checked) > 0.5 * UNDERFLOW_NATS, checked + + +@pytest.mark.skipif(fl.xpy_default is np, reason="no cupy on this host") +def test_gpu_offset_is_per_row_too(tvals): + """Same guard on the GPU backend. The defect is device-independent -- the path is + selected by ``opts.gpu``, not by the device -- but the calmarg loop runs different + kernels there (``_q_inner_product_gpu``), so the reduction is checked on both. + """ + import cupy + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + P, rholms, lookupNK, ct, epochs = _inputs(dists) + rholms_g = {k: cupy.asarray(v) for k, v in rholms.items()} + ct_g = {k: cupy.asarray(v) for k, v in ct.items()} + # The extrinsic arrays must be on the device, as the driver puts them + # (integrate_likelihood_extrinsic_batchmode: ``P.phi = xpy_default.asarray(...)``); + # a host P_vec reaches a cupy elementwise kernel and raises + # "TypeError: Unsupported type " inside + # SphericalHarmonicsVectorized, which is a harness error, not a likelihood one. + Pg = P.manual_copy() + for attr in ('phi', 'theta', 'psi', 'incl', 'phiref', 'dist'): + Pg.__dict__[attr] = cupy.asarray(np.asarray(getattr(P, attr), dtype=np.float64)) + Pg.tref = float(P.tref) + Pg.deltaT = float(P.deltaT) + tvals_g = cupy.asarray(tvals) + + def _call(**kw): + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals_g, Pg, lookupNK, rholms_g, ct_g, ct_g, epochs, Lmax=2, xpy=cupy, + n_cal=N_CAL, cal_method='loop', **kw) + + lnL = cupy.asnumpy(_call()) + assert lnL.shape == (2,), lnL.shape + assert np.all(np.isfinite(lnL)), lnL + comps = cupy.asnumpy(_call(return_cal_components=True)) + ref = logsumexp(comps, axis=-1) - np.log(N_CAL) + np.testing.assert_allclose(lnL, ref, rtol=0, atol=1e-7) + + +if __name__ == '__main__': + raise SystemExit(pytest.main([__file__, '-v'])) diff --git a/MonteCarloMarginalizeCode/Code/test/test_cip_evidence_consolidation.py b/MonteCarloMarginalizeCode/Code/test/test_cip_evidence_consolidation.py index d3e8b6292..ee8130194 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_cip_evidence_consolidation.py +++ b/MonteCarloMarginalizeCode/Code/test/test_cip_evidence_consolidation.py @@ -114,6 +114,16 @@ def test_pipeline_has_terminal_prior_then_strict_final_evidence(pipeline): assert "--cip-prefix overlap-grid-$(macroiterationnext)" in source +@pytest.mark.parametrize("pipeline", PIPELINES) +def test_pipeline_points_in_loop_evidence_at_non_exploded_cip_output(pipeline): + source = open(pipeline).read() + assert 'evidence_source = " --cip-dir iteration_$(macroiteration)_cip"' in source + assert 'evidence_source = " --cip-dir . --cip-prefix overlap-grid-$(macroiterationnext)"' in source + macro_api = ("add_variable" if os.path.basename(pipeline) == "cepp_basic_htcondor" + else "add_macro") + assert 'evidence_node.{}("macroiterationnext",it)'.format(macro_api) in source + + def test_prior_mode_is_independent_and_reweighted_evidence_restores_shift(): source = open(CIP).read() assert 'parser.add_argument("--integrate-prior"' in source diff --git a/MonteCarloMarginalizeCode/Code/test/test_cip_gp_kernel_bounds.py b/MonteCarloMarginalizeCode/Code/test/test_cip_gp_kernel_bounds.py new file mode 100644 index 000000000..96eb1f512 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_cip_gp_kernel_bounds.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +"""Tests for ``fit_gp``'s kernel-bound options and its saturation report. + +Why this matters +---------------- +``fit_gp`` builds ``WhiteKernel(noise_level_bounds=...) + C(amplitude_bounds) * +RBF(length_scale_bounds=...)``. The first two bounds were hardcoded constants, +hand-tuned against the ln L dynamic range of contemporary-detector analyses. At +third-generation dynamic range they saturate: the amplitude ceiling ``1e1`` +represents a signal of at most ``sqrt(1e1) = 3.2`` nats, against ranges of a +couple of hundred nats. + +**A saturated fit is not self-announcing.** The optimizer returns successfully, +the posterior is produced, and nothing in the output distinguishes "converged" +from "pinned against a wall the user never chose". On a zero-spin BNS at network +amplitude 23.8 the default bounds put two of four hyperparameters exactly on a +bound, cost a factor 1.9 in held-out predictive accuracy, and narrowed the +recovered 90% credible interval in chirp mass by 40%. + +This module follows ``test_cip_priors.py``: CIP is a script that parses argv at +import, so the functions under test are extracted from its source with ``ast`` +and exec'd. They are therefore byte-identical to the ones CIP runs, rather than +transcribed here where they could silently drift. + +Coverage: + +``test_defaults_reproduce_the_historical_hardcoded_bounds`` + The whole design rests on the new options being no-ops when unset. If a + default ever changes, every previously published number silently moves. + +``test_report_flags_a_saturated_fit`` / ``test_report_clears_an_unsaturated_fit`` + The saturation flag must be *true* when a hyperparameter is on a bound and + *false* when none is. A flag that is always one or the other is useless, so + both directions are asserted. + +``test_holdout_is_reported_only_when_requested`` + In-sample residual rewards flexibility and cannot separate a better fit from + an overfit one; the held-out score can. It costs K extra fits, so it must + stay off by default. + +``test_peak_scan_requires_a_valid_coordinate_index`` + CIP passes the chirp-mass column unconditionally and its "no such + coordinate" sentinel is ``-1``, which indexes the last column instead of + failing. The scan must decline to run rather than mislabel a coordinate. + +``test_peak_index_comes_from_the_fitted_coordinate_basis`` / +``test_fit_gp_derives_the_peak_index_from_coord_names`` + The scan labels its answer "chirp mass", so the column it scans must be the + chirp-mass column of the array actually fitted. CIP keeps two coordinate + lists -- coord_names (fitted) and low_level_coord_names (sampled) -- which + diverge under ``--parameter-implied`` / ``--parameter-nofit``, so an index + taken in the wrong list reports some other coordinate's peak as mc. + +``test_peak_scan_is_bounded_in_points_and_batched`` / +``test_batched_scan_finds_the_same_peak_as_a_single_pass`` + The scan is an outer-product grid and ``gp.predict`` forms a + points-by-training kernel matrix, so an unbounded 3-D scan can exhaust + memory on a fit that succeeded. Bounding it is only safe if the blocked + search still returns the same peak as one pass. +""" +import ast +import os + +import numpy as np +import pytest + +sklearn = pytest.importorskip("sklearn") +from sklearn.gaussian_process import GaussianProcessRegressor # noqa: E402 +from sklearn.gaussian_process.kernels import ( # noqa: E402 + RBF, WhiteKernel, ConstantKernel as C) + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_CODE = os.path.abspath(os.path.join(_HERE, os.pardir)) +CIP_SCRIPT = os.path.join( + _CODE, "bin", "util_ConstructIntrinsicPosterior_GenericCoordinates.py") + +# The values that were hardcoded in fit_gp before these options existed. These +# literals are the point of the test: they are what "unset behaves as before" +# means, so they are deliberately written out rather than read from the source. +HISTORICAL_NOISE_BOUNDS = "1e-2,1" +HISTORICAL_AMPLITUDE_BOUNDS = "1e-3,1e1" +HISTORICAL_LENGTH_SCALE_MAX_FACTOR = 5.0 + + +def _parse_script(path): + with open(path) as handle: + return ast.parse(handle.read()) + + +CIP_TREE = _parse_script(CIP_SCRIPT) + + +def _add_argument_kwargs(tree, option): + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute) + and node.func.attr == "add_argument" and node.args): + continue + try: + if ast.literal_eval(node.args[0]) != option: + continue + except (ValueError, SyntaxError): + continue + found = {} + for keyword in node.keywords: + if keyword.arg is None: + continue + try: + found[keyword.arg] = ast.literal_eval(keyword.value) + except (ValueError, SyntaxError): + found[keyword.arg] = None + return found + raise AssertionError("no add_argument({!r}) call found".format(option)) + + +def _load_functions(*names): + """Exec the named top-level CIP functions in a namespace they can run in.""" + wanted = {} + for node in CIP_TREE.body: + if isinstance(node, ast.FunctionDef) and node.name in names: + wanted[node.name] = node + missing = set(names) - set(wanted) + assert not missing, "CIP no longer defines: %s" % sorted(missing) + ns = {"np": np, "GaussianProcessRegressor": GaussianProcessRegressor, + "RBF": RBF, "WhiteKernel": WhiteKernel, "C": C, "print": print} + module = ast.Module(body=[wanted[n] for n in names], type_ignores=[]) + exec(compile(module, CIP_SCRIPT, "exec"), ns) + return ns + + +def _fit(amp_bounds, noise_bounds=(1e-2, 1.0), n=40, seed=0, noise=0.0): + """A small GP fit over a peak whose amplitude far exceeds any tight bound. + + ``noise`` matters more than it looks: on exactly noiseless data the fitted + WhiteKernel level runs to its LOWER bound, which the report correctly calls + saturation. The unsaturated fixture therefore has to supply real scatter for + the noise term to have an interior optimum -- a fixture detail, not a + property of the code under test. + """ + rng = np.random.default_rng(seed) + x = np.sort(rng.uniform(-1, 1, size=(n, 1)), axis=0) + y = 120.0 * np.exp(-0.5 * (x[:, 0] / 0.2) ** 2) # ~120-nat dynamic range + if noise: + y = y + rng.normal(0.0, noise, size=n) + kernel = (WhiteKernel(noise_level=0.1, noise_level_bounds=noise_bounds) + + C(0.5, amp_bounds) * RBF(length_scale=[0.3], + length_scale_bounds=(1e-3, 1e1))) + gp = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, n_restarts_optimizer=2) + gp.fit(x, y) + return gp, kernel, x, y + + +def _fit_nd(ndim, n=30, seed=0): + """The same fixture in ``ndim`` dimensions, for the peak-scan tests.""" + rng = np.random.default_rng(seed) + x = rng.uniform(-1, 1, size=(n, ndim)) + y = (120.0 * np.exp(-0.5 * np.sum((x / 0.4) ** 2, axis=1)) + + rng.normal(0.0, 0.5, size=n)) + kernel = (WhiteKernel(noise_level=0.1, noise_level_bounds=(1e-4, 1e3)) + + C(0.5, (1e-3, 1e8)) * RBF(length_scale=[0.3] * ndim, + length_scale_bounds=(1e-3, 1e1))) + gp = GaussianProcessRegressor(kernel=kernel, alpha=1e-6, n_restarts_optimizer=1) + gp.fit(x, y) + return gp, kernel, x, y + + +def test_defaults_reproduce_the_historical_hardcoded_bounds(): + """Unset flags must leave fit_gp building exactly the kernel it always did.""" + assert (_add_argument_kwargs(CIP_TREE, "--fit-gp-noise-bounds")["default"] + == HISTORICAL_NOISE_BOUNDS) + assert (_add_argument_kwargs(CIP_TREE, "--fit-gp-amplitude-bounds")["default"] + == HISTORICAL_AMPLITUDE_BOUNDS) + assert (_add_argument_kwargs(CIP_TREE, "--fit-gp-length-scale-max-factor")["default"] + == HISTORICAL_LENGTH_SCALE_MAX_FACTOR) + # off by default: it costs K extra GP fits + assert _add_argument_kwargs(CIP_TREE, "--fit-gp-holdout-folds")["default"] == 0 + + +def test_gp_bounds_opt_parses_a_pair(): + ns = _load_functions("_gp_bounds_opt") + assert ns["_gp_bounds_opt"]("1e-4,1e3") == (1e-4, 1e3) + assert ns["_gp_bounds_opt"](HISTORICAL_AMPLITUDE_BOUNDS) == (1e-3, 1e1) + with pytest.raises(ValueError): + ns["_gp_bounds_opt"]("not-a-pair") + + +def test_report_flags_a_saturated_fit(): + """A ceiling far below the data's dynamic range must be reported as saturated.""" + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit(amp_bounds=(1e-3, 1e1)) # 3.2 nats vs a 120-nat peak + rec = ns["report_gp_kernel"](gp, x, y) + assert rec["saturated"] is True + assert rec["n_at_bound"] >= 1 + pinned = [h["name"] for h in rec["hyperparameters"] + if h["at_lower_bound"] or h["at_upper_bound"]] + assert any("constant_value" in nm for nm in pinned), ( + "the amplitude is the hyperparameter this ceiling pins; got %s" % pinned) + for h in rec["hyperparameters"]: + if h["at_upper_bound"] or h["at_lower_bound"]: + assert h["decades_to_bound"] < 1e-2 + + +def test_report_clears_an_unsaturated_fit(): + """With room to move, nothing may be reported as sitting on a bound.""" + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit(amp_bounds=(1e-3, 1e8), noise_bounds=(1e-4, 1e3), noise=0.5) + rec = ns["report_gp_kernel"](gp, x, y) + assert rec["saturated"] is False, rec["kernel_fitted"] + assert rec["n_at_bound"] == 0 + assert all(h["decades_to_bound"] > 0 for h in rec["hyperparameters"]) + + +def test_holdout_is_reported_only_when_requested(): + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit(amp_bounds=(1e-3, 1e8), noise_bounds=(1e-4, 1e3), noise=0.5) + + off = ns["report_gp_kernel"](gp, x, y) + assert "holdout_rms_nats" not in off, "held-out scoring must be opt-in" + + on = ns["report_gp_kernel"](gp, x, y, holdout_folds=3, + kernel_proto=kernel, alpha_proto=0.25) + assert on["holdout_folds"] == 3 + assert np.isfinite(on["holdout_rms_nats"]) and on["holdout_rms_nats"] >= 0.0 + assert on["holdout_max_abs_nats"] >= on["holdout_rms_nats"] * 0.0 + + +def test_report_is_pure_json_and_names_every_hyperparameter(): + """The record is consumed by build scripts, so it must stay serializable.""" + import json + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit(amp_bounds=(1e-3, 1e8), noise_bounds=(1e-4, 1e3), noise=0.5) + rec = ns["report_gp_kernel"](gp, x, y) + json.dumps(rec) # must not raise + assert len(rec["hyperparameters"]) == len(gp.kernel_.theta) + assert all(h["name"] and not h["name"].startswith("param_") + for h in rec["hyperparameters"]), ( + "a hyperparameter fell back to a positional name; the name mapping broke") + + +def test_peak_scan_requires_a_valid_coordinate_index(): + """CIP's "no chirp mass here" sentinel is -1, a perfectly legal index.""" + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit_nd(2) + for bad in (None, -1, 2, 7): + rec = ns["report_gp_kernel"](gp, x, y, peak_index=bad) + assert "surface_peak" not in rec, ( + "peak_index=%r must not be scanned as a coordinate" % (bad,)) + rec = ns["report_gp_kernel"](gp, x, y, peak_index=1, peak_max_points=256) + assert rec["surface_peak"]["coord_index"] == 1 + + +def test_peak_index_comes_from_the_fitted_coordinate_basis(): + """The mc column is looked up in the list x is built from, or not at all.""" + ns = _load_functions("_mc_index_in") + assert ns["_mc_index_in"](['mc', 'delta_mc', 'xi']) == 0 + assert ns["_mc_index_in"](['delta_mc', 'mc']) == 1 + # --parameter delta_mc --parameter-implied mu1 --parameter-implied mu2 + # --parameter-nofit mc : mc is sampled but NOT fitted, so column 1 is mu1. + # Returning 1 here (the position of mc in the sampling list) would report + # the mu1 peak as chirp mass; the scan must be declined instead. + assert ns["_mc_index_in"](['delta_mc', 'mu1', 'mu2']) == -1 + + +def test_fit_gp_derives_the_peak_index_from_coord_names(): + """fit_gp's x has coord_names columns, so mc_index must not index it. + + Asserted on the source rather than by running fit_gp, which needs parsed + argv and a full data load; the wiring is what regressed, and it is visible + in the call itself. + """ + fit_gp = next(n for n in CIP_TREE.body + if isinstance(n, ast.FunctionDef) and n.name == "fit_gp") + calls = [n for n in ast.walk(fit_gp) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == "report_gp_kernel"] + assert calls, "fit_gp no longer reports its kernel" + for call in calls: + passed = [k.value for k in call.keywords if k.arg == "peak_index"] + assert passed, "the peak index must be passed explicitly" + arg = passed[0] + assert not (isinstance(arg, ast.Name) and arg.id == "mc_index"), ( + "mc_index indexes low_level_coord_names, not the fitted columns") + assert (isinstance(arg, ast.Call) and isinstance(arg.func, ast.Name) + and arg.func.id == "_mc_index_in"), ( + "the peak index must be derived from the fitted coordinate list") + assert [a.id for a in arg.args if isinstance(a, ast.Name)] == ["coord_names"] + + +def test_peak_scan_is_bounded_in_points_and_batched(): + """Total scan points and per-call block size must both stay under budget.""" + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit_nd(3) + + calls = [] + inner_predict = gp.predict + + def counting_predict(pts, **kwargs): + calls.append(len(pts)) + return inner_predict(pts, **kwargs) + + gp.predict = counting_predict + rec = ns["report_gp_kernel"](gp, x, y, peak_index=0, peak_max_points=4096, + peak_block_elements=3000) + scan_calls = calls[1:] # calls[0] is the residual predict + + n_axis = rec["surface_peak"]["grid_points_per_axis"] + assert n_axis ** 3 <= 4096, "the grid ignored its total-points budget" + assert n_axis < 240, "a 3-D scan must be thinned below the per-axis default" + assert sum(scan_calls) == n_axis ** 3, "the whole grid must still be scanned" + assert max(scan_calls) <= 3000 // len(y), "a block exceeded the element budget" + + +def test_batched_scan_finds_the_same_peak_as_a_single_pass(): + """Blocking is only acceptable if the running argmax is the global one.""" + ns = _load_functions("report_gp_kernel") + gp, kernel, x, y = _fit(amp_bounds=(1e-3, 1e8), noise_bounds=(1e-4, 1e3), noise=0.5) + one = ns["report_gp_kernel"](gp, x, y, peak_index=0, + peak_block_elements=10 ** 9)["surface_peak"] + many = ns["report_gp_kernel"](gp, x, y, peak_index=0, + peak_block_elements=1)["surface_peak"] + assert one["grid_points_per_axis"] == many["grid_points_per_axis"] == 240 + assert one["value"] == many["value"] + assert abs(one["value"]) < 0.2, "the fixture's peak sits at x=0" diff --git a/MonteCarloMarginalizeCode/Code/test/test_cip_priors.py b/MonteCarloMarginalizeCode/Code/test/test_cip_priors.py index 03c383db7..8b0b46970 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_cip_priors.py +++ b/MonteCarloMarginalizeCode/Code/test/test_cip_priors.py @@ -96,6 +96,12 @@ LAMBDA_SMALL_MAX = 2000.0 MC_MIN = 5.0 MC_MAX = 60.0 +A6C_MIN = -80.0 +A6C_MAX = -20.0 +E0_MIN = 1.0 +E0_MAX = 1.2 +PPHI0_MIN = 0.0 +PPHI0_MAX = 5.4 # CIP sets p_Rbar = lalsimutils.p_R. Read out of the lalsimutils SOURCE rather # than imported: importing lalsimutils pulls in LAL, whose default error handler @@ -202,6 +208,12 @@ def _make_namespace(ecc_min=ECC_MIN, ecc_max=ECC_MAX, eccentricity_prior="unifor "ECC_MAX": ecc_max, "MEANPERANO_MIN": 0.0, "MEANPERANO_MAX": 2 * np.pi, + "A6C_MIN": A6C_MIN, + "A6C_MAX": A6C_MAX, + "E0_MIN": E0_MIN, + "E0_MAX": E0_MAX, + "PPHI0_MIN": PPHI0_MIN, + "PPHI0_MAX": PPHI0_MAX, "lambda_min": LAMBDA_MIN, "lambda_max": LAMBDA_MAX, "lambda_small_max": LAMBDA_SMALL_MAX, @@ -262,6 +274,9 @@ def _load_priors(namespace): # a density in e^2, so it is evaluated on the squared interval "log_eccentricity_squared_prior": (ECC_MIN ** 2, ECC_MAX ** 2), "meanPerAno_prior": (0.0, 2 * np.pi), + "a6c_prior": (A6C_MIN, A6C_MAX), + "initial_energy_prior": (E0_MIN, E0_MAX), + "initial_angmom_prior": (PPHI0_MIN, PPHI0_MAX), "precession_prior": (0.0, 2.0), "lambda_prior": (LAMBDA_MIN, LAMBDA_MAX), "lambda_small_prior": (LAMBDA_MIN, LAMBDA_SMALL_MAX), @@ -327,6 +342,9 @@ def _load_priors(namespace): # interval directly rather than through the 'square' substitution. ("log_eccentricity_squared_prior", ECC_MIN ** 2, ECC_MAX ** 2, "x", ()), ("meanPerAno_prior", 0.0, 2 * np.pi, "x", ()), + ("a6c_prior", A6C_MIN, A6C_MAX, "x", ()), + ("initial_energy_prior", E0_MIN, E0_MAX, "x", ()), + ("initial_angmom_prior", PPHI0_MIN, PPHI0_MAX, "x", ()), ("precession_prior", 0.0, 2.0, "x", ()), ("triangle_prior", -CHI_MAX, CHI_MAX, "x", ()), ("s_component_uniform_prior", -CHI_MAX, CHI_MAX, "x", ()), diff --git a/MonteCarloMarginalizeCode/Code/test/test_container_manifest.py b/MonteCarloMarginalizeCode/Code/test/test_container_manifest.py index 177983a8f..e1b495bbb 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_container_manifest.py +++ b/MonteCarloMarginalizeCode/Code/test/test_container_manifest.py @@ -411,6 +411,61 @@ def test_integration_cip_legacy_single_image(tmp_path, monkeypatch): assert "require_gpus" not in cmds +def _make_calibration_job(tmp_path, monkeypatch, manifest, container_universe): + if container_universe: + monkeypatch.setenv("RIFT_CONTAINER_UNIVERSE", "1") + else: + monkeypatch.delenv("RIFT_CONTAINER_UNIVERSE", raising=False) + monkeypatch.chdir(tmp_path) + dag = pytest.importorskip("RIFT.misc.dag_utils_generic") + job, _ = dag.write_calibration_uncertainty_reweighting_sub( + tag="Calib_reweight", + log_dir=str(tmp_path) + "/", + exe="/usr/bin/true", + pickle_file=str(tmp_path / "event.pickle"), + posterior_file=str(tmp_path / "posterior.dat"), + transfer_files=[], + use_osg=True, + use_singularity=True, + singularity_image=manifest, + ) + return job, dict(job.condor_cmds) + + +def test_calibration_family_legacy_uses_fallback_not_manifest(tmp_path, monkeypatch): + manifest = _write(tmp_path, ALL_OSDF_MANIFEST) + job, cmds = _make_calibration_job(tmp_path, monkeypatch, manifest, False) + assert job.universe == "vanilla" + assert cmds["MY.SingularityImage"] == '"./rift_ancient_cuda11.sif"' + assert manifest not in cmds["MY.SingularityImage"] + assert "ifThenElse" not in cmds["MY.SingularityImage"] + assert cmds["transfer_input_files"].count( + "osdf:///igwn/sw/rift_ancient_cuda11.sif" + ) == 1 + assert "osdf:///igwn/sw/rift_modern_cuda12.sif" not in cmds["transfer_input_files"] + + +def test_calibration_family_container_universe_uses_fallback_not_manifest( + tmp_path, monkeypatch +): + manifest = _write(tmp_path, ALL_OSDF_MANIFEST) + job, cmds = _make_calibration_job(tmp_path, monkeypatch, manifest, True) + assert job.universe == "container" + assert cmds["container_image"] == "osdf:///igwn/sw/rift_ancient_cuda11.sif" + assert manifest not in cmds["container_image"] + assert "MY.SingularityImage" not in cmds + assert "MY.SingularityBindCVMFS" not in cmds + assert "$$(" not in cmds["container_image"] + assert "rift_modern_cuda12.sif" not in cmds["transfer_input_files"] + + # Exercise the same emission path used by a build-only pseudo-pipe run. + job.write_sub_file() + submit = (tmp_path / "Calib_reweight.sub").read_text() + assert "universe = container" in submit + assert "container_image = osdf:///igwn/sw/rift_ancient_cuda11.sif" in submit + assert "fam.yaml" not in submit + + # --------------------------------------------------------------------------- # 7. runtime-selection wrapper fallback # --------------------------------------------------------------------------- diff --git a/MonteCarloMarginalizeCode/Code/test/test_continuous_time_posterior_export.py b/MonteCarloMarginalizeCode/Code/test/test_continuous_time_posterior_export.py new file mode 100644 index 000000000..83f513f54 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_continuous_time_posterior_export.py @@ -0,0 +1,305 @@ +#!/usr/bin/env python3 +"""Regression tests for sub-sample time-posterior export.""" + +import importlib.util +import os + +import numpy as np +import pytest + +from RIFT.misc import xmlutils + +DRIVER = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "bin", + "integrate_likelihood_extrinsic_batchmode") +LISA_DRIVER = DRIVER + "_lisa" +MODULE = os.path.join( + os.path.dirname(os.path.abspath(__file__)), "..", "RIFT", "likelihood", + "time_posterior.py") +SPEC = importlib.util.spec_from_file_location("time_posterior", MODULE) +TIME_POSTERIOR = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(TIME_POSTERIOR) +draw_continuous_time_posterior = TIME_POSTERIOR.draw_continuous_time_posterior +resolve_time_posterior_export_mode = TIME_POSTERIOR.resolve_time_posterior_export_mode +legacy_time_interpolation_enabled = TIME_POSTERIOR.legacy_time_interpolation_enabled +validate_time_posterior_working_set = TIME_POSTERIOR.validate_time_posterior_working_set +_interval_log_envelopes = TIME_POSTERIOR._interval_log_envelopes + + +def test_auto_contract_tracks_subsample_interpolation(): + assert resolve_time_posterior_export_mode("auto", "nearest") == "grid" + assert resolve_time_posterior_export_mode("auto", "cubic") == "continuous" + assert resolve_time_posterior_export_mode("auto", "sinc") == "continuous" + assert resolve_time_posterior_export_mode("grid", "cubic") == "grid" + assert resolve_time_posterior_export_mode("continuous", "nearest") == "continuous" + # A driver without a faithful arbitrary-time likelihood must preserve its + # historical auto behavior, while leaving an explicit continuous request + # visible for the caller's capability guard to reject. + assert resolve_time_posterior_export_mode( + "auto", "cubic", continuous_available=False) == "grid" + assert resolve_time_posterior_export_mode( + "continuous", "cubic", continuous_available=False) == "continuous" + + +def test_lisa_legacy_interpolation_parser_does_not_treat_false_as_truthy(): + for value in (False, "False", "false", "0", "off", "none", "nearest"): + assert legacy_time_interpolation_enabled(value) is False + for value in (True, "True", "1", "yes", "on", "cubic", "sinc"): + assert legacy_time_interpolation_enabled(value) is True + with pytest.raises(ValueError, match="unrecognised LISA value"): + legacy_time_interpolation_enabled("sinK") + + +def test_dense_working_set_is_refused_before_allocation(): + assert validate_time_posterior_working_set(5, 5000) > 0 + with pytest.raises(MemoryError, match="unsafe dense working set"): + validate_time_posterior_working_set(1000, 2510000) + + +def test_continuous_draws_are_not_on_the_input_lattice(): + tvals = np.linspace(-0.01, 0.01, 41) + lnlt = -0.5 * (tvals / 0.002) ** 2 + rng = np.random.RandomState(20260829) + draws = np.array([draw_continuous_time_posterior(tvals, lnlt, rng)[0] + for _ in range(200)]) + phase = (draws - tvals[0]) / (tvals[1] - tvals[0]) + assert np.all(draws >= tvals[0]) and np.all(draws <= tvals[-1]) + assert np.count_nonzero(np.isclose(phase, np.round(phase), atol=1e-10)) == 0 + + +def test_gaussian_posterior_moments_and_interpolated_logl(): + sigma = 0.0017 + center = 0.0008 + tvals = np.linspace(-0.012, 0.012, 65) + lnlt = -0.5 * ((tvals - center) / sigma) ** 2 + rng = np.random.RandomState(17) + draws, logls = zip(*(draw_continuous_time_posterior(tvals, lnlt, rng) + for _ in range(12000))) + draws = np.asarray(draws) + logls = np.asarray(logls) + assert np.mean(draws) == pytest.approx(center, abs=5e-5) + assert np.std(draws) == pytest.approx(sigma, rel=0.035) + np.testing.assert_allclose(logls, -0.5 * ((draws - center) / sigma) ** 2, + rtol=0, atol=2e-12) + + +def test_stationary_point_envelope_bounds_overshooting_cubics(): + from scipy.interpolate import CubicSpline + + rng = np.random.RandomState(37) + knots = np.linspace(-1.0, 1.0, 9) + for _ in range(20): + values = rng.normal(size=knots.size) + spline = CubicSpline(knots, values) + maxima = _interval_log_envelopes(spline, knots, values) + for interval in range(knots.size - 1): + probes = np.linspace(knots[interval], knots[interval + 1], 1001) + assert np.max(spline(probes)) <= maxima[interval] + 5e-14 + + +def test_batched_rows_draw_from_their_own_posteriors(): + tvals = np.linspace(-0.02, 0.02, 81) + centers = np.array([-0.006, 0.0, 0.007]) + lnlt = -0.5 * ((tvals[None, :] - centers[:, None]) / 0.001) ** 2 + draws, logls = draw_continuous_time_posterior( + tvals, lnlt, np.random.RandomState(9)) + assert draws.shape == logls.shape == centers.shape + assert np.all(np.abs(draws - centers) < 0.004) + + +@pytest.mark.parametrize("stencil", ["cubic", "sinc"]) +def test_explicit_times_apply_selected_q_stencil_at_every_dense_time(stencil): + from RIFT.likelihood import factored_likelihood as fl + + rng = np.random.RandomState(63) + q = rng.normal(size=(160, 3)) + 1j * rng.normal(size=(160, 3)) + antenna_modes = rng.normal(size=(2, 3)) + 1j * rng.normal(size=(2, 3)) + starts = np.array([[31, 32, 33, 34], [57, 58, 59, 60]], dtype=np.int32) + fractions = np.array([[0.05, 0.27, 0.51, 0.89], + [0.13, 0.38, 0.64, 0.92]]) + actual = fl._q_inner_product_explicit_times( + q, antenna_modes, starts, fractions, stencil, xpy=np) + expected = np.empty(actual.shape, dtype=complex) + for row in range(starts.shape[0]): + for col in range(starts.shape[1]): + q_one = fl._q_window_numpy_interp( + q, starts[row:row + 1, col], + fractions[row:row + 1, col], 1, stencil, xpy=np)[0, 0] + expected[row, col] = np.einsum( + "j,j->", antenna_modes[row], q_one) + np.testing.assert_allclose(actual, expected, rtol=2e-14, atol=2e-14) + + +@pytest.mark.parametrize("stencil", ["nearest", "cubic", "sinc"]) +def test_explicit_time_stencil_chunks_the_time_axis(monkeypatch, stencil): + from RIFT.likelihood import factored_likelihood as fl + + # Force max_cells=2: every four-time row must cross a time-chunk boundary. + monkeypatch.setattr(fl, "_Q_EXPLICIT_TEMP_MAX_BYTES", 2 * 3 * 16 * 3, + raising=False) + rng = np.random.RandomState(64) + q = rng.normal(size=(80, 3)) + 1j * rng.normal(size=(80, 3)) + antenna_modes = rng.normal(size=(2, 3)) + 1j * rng.normal(size=(2, 3)) + starts = np.array([[11, 12, 13, 14], [31, 32, 33, 34]], dtype=np.int32) + fractions = None if stencil == "nearest" else np.array( + [[0.05, 0.27, 0.51, 0.89], [0.13, 0.38, 0.64, 0.92]]) + actual = fl._q_inner_product_explicit_times( + q, antenna_modes, starts, fractions, stencil, xpy=np) + expected = np.empty(actual.shape, dtype=complex) + for row in range(starts.shape[0]): + for col in range(starts.shape[1]): + frac = None if fractions is None else fractions[row:row + 1, col] + q_one = fl._q_window_numpy_interp( + q, starts[row:row + 1, col], frac, 1, stencil, xpy=np)[0, 0] + expected[row, col] = np.einsum( + "j,j->", antenna_modes[row], q_one) + np.testing.assert_allclose(actual, expected, rtol=2e-14, atol=2e-14) + + +def test_negative_infinity_knots_are_zero_mass_not_an_arbitrary_log_floor(): + from scipy.interpolate import PchipInterpolator + + tvals = np.linspace(-1.0, 1.0, 5) + lnlt = np.array([-np.inf, -1.0, 0.0, -1.0, -np.inf]) + rng = np.random.RandomState(81) + draws, logls = zip(*(draw_continuous_time_posterior(tvals, lnlt, rng) + for _ in range(200))) + draws = np.asarray(draws) + logls = np.asarray(logls) + density = np.exp(np.where(np.isfinite(lnlt), lnlt, -np.inf)) + expected = PchipInterpolator(tvals, density)(draws) + np.testing.assert_allclose(np.exp(logls), expected, rtol=2e-14, atol=0) + assert np.all(np.isfinite(logls)) + assert np.all((draws > tvals[0]) & (draws < tvals[-1])) + + +@pytest.mark.parametrize("bad", [ + np.array([0.0, np.nan, -1.0]), + np.array([0.0, np.inf, -1.0]), +]) +def test_nan_and_positive_infinity_fail_loudly(bad): + with pytest.raises(ValueError, match=r"NaN or \+inf"): + draw_continuous_time_posterior(np.arange(3.0), bad) + + +def test_all_negative_infinity_has_no_posterior_mass(): + with pytest.raises(ValueError, match="no finite positive mass"): + draw_continuous_time_posterior( + np.arange(3.0), np.full(3, -np.inf)) + + +def test_pathological_rejection_cannot_hang_the_driver(): + class AlwaysReject(object): + def choice(self, size, p): + return 0 + + def uniform(self, *bounds): + return 0.5 * (bounds[0] + bounds[1]) if bounds else 1.0 + + with pytest.raises(RuntimeError, match="exhausted 3 proposals"): + draw_continuous_time_posterior( + np.arange(3.0), np.array([0.0, -100.0, 0.0]), + rng=AlwaysReject(), max_attempts=3) + + +def test_zero_rng_endpoint_cannot_accept_zero_posterior_density(): + class ZeroEndpoint(object): + def choice(self, size, p): + return 0 + + def uniform(self, *bounds): + return bounds[0] if bounds else 0.0 + + # The first interval rises from zero density. Both proposal and acceptance + # uniforms are exactly zero, so an inclusive acceptance comparison would + # incorrectly return the left endpoint with lnL=-inf. + with pytest.raises(RuntimeError, match="exhausted 3 proposals"): + draw_continuous_time_posterior( + np.arange(3.0), np.array([-np.inf, 0.0, -np.inf]), + rng=ZeroEndpoint(), max_attempts=3) + + +def test_driver_wires_continuous_draw_before_legacy_grid_choice(): + with open(DRIVER) as handle: + source = handle.read() + dense_labels = source.index("xpy_default.arange(n_dense)") + explicit = source.index("explicit_time_values=True", dense_labels) + continuous = source.index("draw_continuous_time_posterior(tvals, lnLt)", explicit) + grid = source.index("indx_choose = np.random.choice", continuous) + assert dense_labels < explicit < continuous < grid + assert continuous < grid + assert "explicit_time_values=True" in source + assert "validate_time_posterior_working_set(n_samples, n_dense)" in source + assert "del lnLt_dense, tvals_dense, sigma_dense, measurable_dense" in source + assert "not opts.gpu or opts.rotation_slow or opts.freqresponse" in source + assert "continuous_available=not (opts.rotation_slow or opts.freqresponse)" in source + assert "refusing before the expensive integration" in source + assert 'opts._time_posterior_export == "continuous"' in source + assert 'opts._time_posterior_export == "grid"' in source + + +def test_lisa_twin_accepts_export_contract_with_documented_lattice_fallback(): + with open(LISA_DRIVER) as handle: + source = handle.read() + assert '"--time-posterior-export"' in source + assert '"--srate-resample-time-marginalization"' in source + assert "legacy_time_interpolation_enabled(opts.interpolate_time)" in source + assert "continuous_available=False" in source + assert ("opts.resample_time_marginalization and\n" + " opts._time_posterior_export == \"continuous\"") in source + assert "interpolation-lattice export for executable-swap compatibility" in source + assert 'opts._time_posterior_export = "grid"' in source + assert "1.0 / opts.srate_resample_time_marginalization" in source + assert "draw_continuous_time_posterior(tvals, lnLt)" not in source + + +def test_bandlimited_export_bypasses_coarse_return_lnlt(): + with open(DRIVER) as handle: + source = handle.read() + resample = source.index("def resample_samples") + bandlimited = source.index('if opts._time_quadrature == "bandlimited":', resample) + dense_draw = source.index("return_time_draw=True", bandlimited) + coarse_series = source.index("return_lnLt=True", dense_draw) + assert bandlimited < dense_draw < coarse_series + assert "opts._time_posterior_export = 'continuous'" in source + assert "--time-posterior-export grid would" in source + assert "conflicting fixed" in source + + +def test_exact_gps_addition_preserves_sub_float_ulp_offsets_and_carry(): + epoch_s = 1400000000 + epoch_ns = 0 + offset = np.array([60e-9, -60e-9, 1.000000060]) + seconds, nanoseconds = xmlutils.gps_add_seconds_exact( + epoch_s, epoch_ns, offset) + np.testing.assert_array_equal(seconds, + [1400000000, 1399999999, 1400000001]) + np.testing.assert_array_equal(nanoseconds, [60, 999999940, 60]) + # At this epoch float64 cannot represent a 60 ns increment. The exact pair + # must therefore be the serialization source, not the compatibility float. + assert float(epoch_s) + 60e-9 == float(epoch_s) + + +def test_exact_xml_time_fields_override_the_legacy_float_mapping(): + keys = list(xmlutils.CMAP) + assert keys.index("t_ref") < keys.index("t_ref_gps_seconds") + assert keys.index("t_ref_gps_seconds") < keys.index("t_ref_gps_nanoseconds") + assert xmlutils.CMAP["t_ref_gps_seconds"] == "geocent_end_time" + assert xmlutils.CMAP["t_ref_gps_nanoseconds"] == "geocent_end_time_ns" + + class Row(object): + pass + + class Table(object): + RowType = Row + + @staticmethod + def get_next_id(): + return 17 + + row = xmlutils.samples_to_siminsp_row( + Table(), t_ref=float(1400000000), + t_ref_gps_seconds=np.int64(1400000000), + t_ref_gps_nanoseconds=np.int64(60)) + assert row.geocent_end_time == 1400000000 + assert row.geocent_end_time_ns == 60 diff --git a/MonteCarloMarginalizeCode/Code/test/test_database.py b/MonteCarloMarginalizeCode/Code/test/test_database.py index 5be66bcc2..24e1c6e40 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_database.py +++ b/MonteCarloMarginalizeCode/Code/test/test_database.py @@ -273,9 +273,14 @@ def test_dag_chains_levels_per_sim(tmp_path): rq = DualCondorRunQueue() n1 = archive.register(0.1, target_level=2) n2 = archive.register(0.2, target_level=3) - rq.submit(archive, [n1, n2]) # condor_submit_dag missing -> noop dispatch - dags = list((base / "run_queue" / "dags").iterdir()) - assert len(dags) == 1 + rq.submit(archive, [n1, n2]) + # *.dag, not iterdir(): where condor_submit_dag IS installed it runs and + # drops four companion files beside the DAG (.condor.sub, .dagman.log, + # .lib.err, .lib.out), so a directory count is 5 there and 1 on a machine + # without HTCondor. This test was green only on the latter. + dags = sorted((base / "run_queue" / "dags").glob("*.dag")) + assert len(dags) == 1, [p.name for p in + (base / "run_queue" / "dags").iterdir()] text = dags[0].read_text() assert "PARENT {}_lvl1 CHILD {}_lvl2".format(n1, n1) in text assert "PARENT {}_lvl2 CHILD {}_lvl3".format(n2, n2) in text diff --git a/MonteCarloMarginalizeCode/Code/test/test_distance_tail.py b/MonteCarloMarginalizeCode/Code/test/test_distance_tail.py new file mode 100644 index 000000000..0a611739b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_distance_tail.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python +""" +Regression tests for the DISTANCE TAIL fix +(RIFT/interpolators/distance_tail.py, wired into bin/util_ConstructEOSPosterior.py +behind --fit-distance-tail). + +Background (the bug these tests lock down). In a distance-export (.dslice) run the +CIP fits lnL over the intrinsic coordinates PLUS distance, and the training set is +~50 discrete distances per intrinsic point. The default fit is an ExtraTrees +ensemble, which is piecewise constant outside its training envelope: past a given +intrinsic point's outermost exported slice it returns that slice's lnL forever +instead of letting the likelihood decay. + +That is not a small error, because the recovered distance posterior is +exp(lnL) * pi(d) and the distance prior is VOLUMETRIC -- it grows like d^2. Holding +lnL flat therefore makes the integrand GROW out to whatever --d-max the CIP range +allows, so probability that should have died survives. Measured across a 10-event +distance-export catalog, two seeds each: the recovered 90-10 distance width came out ++16.8% (median) too wide, in every quantile span including the core, with the median +distance itself untouched. The fraction of recovered samples falling outside their +nearest grid point's exported support orders those events by width excess at Spearman +rho = +0.952, p = 4e-5. + +The fix continues each slice past its edge with the exact single-angle form through +the origin, matched to the value AND slope the data has at that edge. With +x = 1/d, u = x/x_edge, and s the dimensionless log-slope at the edge: + + lnL(u)/lnL_edge = (2-s) u + (s-1) u^2 + +so a slice that is still flat at its edge (the common case, because the +distance-inclination degeneracy keeps the marginalised lnL flat across the exported +support) rolls off as 2u - u^2 rather than being guillotined. + +The properties below are the ones that make this safe to turn on, and each is a way +the fix could regress into a second extrapolation bug if someone edits the module: + * it is a NO-OP on the support, so nothing that currently works changes; + * it is CONTINUOUS at the edge, so the sampler sees no step to mistake for structure; + * it is MONOTONE beyond the edge, so walking out in distance always walks lnL down; + * it goes to ZERO as d -> infinity, which for a likelihood RATIO is an identity, not + a modelling choice -- an infinitely distant source is exactly the noise hypothesis. +""" +import numpy as np +import pytest + +from RIFT.interpolators.distance_tail import wrap_distance_tail + +COORDS = ["mc", "dist"] +D_LO, D_HI, LNL = 1000.0, 5000.0, 13.0 + + +def _flat_grid(n_slices=30, n_d=40, lnL=LNL, slope_per_x=0.0): + """A synthetic .dslice-shaped grid: n_slices intrinsic points, each exported at n_d + distances. `slope_per_x` puts a known linear-in-1/d trend on lnL so the slope + branch can be exercised as well as the flat one.""" + rows, y = [], [] + for i in range(n_slices): + mc = 20.0 + 0.1 * i + for d in np.linspace(D_LO, D_HI, n_d): + rows.append([mc, d]) + y.append(lnL + slope_per_x * (1.0 / d - 1.0 / D_HI)) + return np.asarray(rows), np.asarray(y) + + +def _wrapped(X, Y, base_value=LNL, **kw): + base = lambda Xf: np.full(len(Xf), base_value) # noqa: E731 a perfectly flat "fit" + return wrap_distance_tail(base, X, Y, COORDS, y_errors=np.full(len(Y), 0.1), **kw) + + +def _probe(f, mc=20.0, d=None): + d = np.atleast_1d(d) + return f(np.stack([np.full(len(d), mc), d], axis=1)) + + +def test_no_op_on_the_support(): + """Inside the exported distance range the wrapper must return the base fit untouched. + If this fails the fix is no longer a strict addition and every existing result moves.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + d = np.linspace(D_LO, D_HI, 500) + assert np.allclose(_probe(f, d=d), LNL) + + +def test_continuous_at_the_edge(): + """A step at the support edge would be read by the sampler as real structure.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + inside = _probe(f, d=D_HI)[0] + just_outside = _probe(f, d=D_HI * (1 + 1e-9))[0] + assert abs(inside - just_outside) < 1e-6 + + +def test_default_law_is_the_chord(): + """The production law. lnL is convex in x = 1/d wherever Var(a) > for the + marginalised likelihood, and a convex function through the origin lies BELOW its + chord, so ratio = u is both the exact asymptotic form and an upper bound on lnL + beyond the edge. It is parameter-free, which is the point: a decay rate tuned to + match the reference would not be a fix.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + for d in (5500.0, 6000.0, 10000.0, 20000.0): + u = (1.0 / d) / (1.0 / D_HI) + assert _probe(f, d=d)[0] == pytest.approx(LNL * u, rel=1e-6) + + +def test_slope_law_rolls_off_as_2u_minus_u_squared(): + """The REJECTED alternative, kept because measuring it is what ruled it out. A slice + still flat at its edge has a unique through-the-origin quadratic that is also flat + there, and it is 2u - u^2. It is faithful to the data at the edge -- and far too + gentle: on the catalog it moved the recovered width by under a point, because the + marginalised lnL is a plateau across the whole exported support and the local slope + therefore cannot see the turnover just past it.""" + X, Y = _flat_grid() + f = _wrapped(X, Y, law="slope") + for d in (5500.0, 6000.0, 10000.0, 20000.0): + u = (1.0 / d) / (1.0 / D_HI) + assert _probe(f, d=d)[0] == pytest.approx(LNL * (2 * u - u * u), rel=1e-6) + + +def test_decays_to_zero_at_large_distance(): + """lnL is a likelihood RATIO, so lnL(d -> infinity) = 0 is an identity. This is the + boundary condition the tree ensemble was missing.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + assert _probe(f, d=1e9)[0] == pytest.approx(0.0, abs=1e-3) + + +def test_monotone_decreasing_beyond_the_edge(): + """Walking out in distance must always walk lnL down. A continuation that turns + back up would put weight at large d all over again.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + v = _probe(f, d=np.linspace(D_HI, 60 * D_HI, 4000)) + assert np.all(np.diff(v) <= 1e-9) + + +def test_beats_the_volumetric_prior(): + """The point of the fix: beyond the edge the integrand exp(lnL) * d^2 must FALL. + Flat extrapolation makes it rise, which is the whole bug.""" + X, Y = _flat_grid() + f = _wrapped(X, Y) + d = np.linspace(D_HI, 8 * D_HI, 400) + integrand = _probe(f, d=d) + 2.0 * np.log(d) # log of exp(lnL) * d^2 + assert integrand[-1] < integrand[0] + flat = LNL + 2.0 * np.log(d) # what the bare tree gives + assert flat[-1] > flat[0] + + +@pytest.mark.parametrize("law", ["chord", "slope"]) +@pytest.mark.parametrize("slope_per_x", [0.0, 5000.0, 20000.0]) +def test_continuation_stays_monotone_and_bounded(slope_per_x, law): + """A slice already decaying at its edge is continued at its own rate. Whatever the + slope, the continuation may only shrink the edge value -- never grow it, never flip + its sign.""" + X, Y = _flat_grid(slope_per_x=slope_per_x) + f = _wrapped(X, Y, law=law) + v = _probe(f, d=np.linspace(D_HI, 40 * D_HI, 2000)) + assert np.all(v <= LNL + 1e-6) + assert np.all(v >= -1e-6) + assert np.all(np.diff(v) <= 1e-9) + + +def test_rising_edge_is_clamped_not_trusted(): + """Monte-Carlo noise can leave a slice whose outer end RISES with distance. Taken at + face value that would extrapolate upward forever; it must be clamped to the flat + roll-off instead.""" + X, Y = _flat_grid(slope_per_x=-20000.0) # lnL increasing with distance + f = _wrapped(X, Y, law="slope") + v = _probe(f, d=np.linspace(D_HI, 40 * D_HI, 2000)) + assert np.all(np.diff(v) <= 1e-9) + # clamped to the flat roll-off, i.e. exactly 2u - u^2 rather than anything rising + u = (1.0 / (40 * D_HI)) / (1.0 / D_HI) + assert v[-1] == pytest.approx(LNL * (2 * u - u * u), rel=1e-6) + assert _probe(f, d=1e9)[0] == pytest.approx(0.0, abs=1e-3) + + +def test_lnL_offset_is_applied_to_the_physical_likelihood(): + """CIP fits Y = lnL_physical - lnL_shift. The d -> infinity identity holds for the + PHYSICAL likelihood ratio, so the tail must decay to -lnL_shift in fit units. + Ignoring the offset anchors the decay to the wrong asymptote -- silently, and in the + direction that reintroduces the bug.""" + shift = 40.0 + X, Y = _flat_grid(lnL=LNL - shift) + f = _wrapped(X, Y, base_value=LNL - shift, lnL_offset=shift) + assert _probe(f, d=1e9)[0] == pytest.approx(-shift, abs=1e-2) + assert np.allclose(_probe(f, d=np.linspace(D_LO, D_HI, 200)), LNL - shift) + + +def test_missing_distance_coordinate_is_an_error(): + """A run that asked for the fix and silently did not get it would carry the very bias + the option exists to remove, with nothing in the log to say so.""" + X, Y = _flat_grid() + with pytest.raises(ValueError, match="dist"): + wrap_distance_tail(lambda Xf: np.zeros(len(Xf)), X, Y, ["mc", "not_distance"]) + + +def test_too_few_slices_is_an_error(): + """Guard against being pointed at something that is not a dslice export at all.""" + X, Y = _flat_grid(n_slices=3) + with pytest.raises(ValueError, match="slices"): + _wrapped(X, Y) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/MonteCarloMarginalizeCode/Code/test/test_fairdraw_double_weighting.py b/MonteCarloMarginalizeCode/Code/test/test_fairdraw_double_weighting.py index c20cb94e8..bf86fd903 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_fairdraw_double_weighting.py +++ b/MonteCarloMarginalizeCode/Code/test/test_fairdraw_double_weighting.py @@ -59,6 +59,7 @@ def _load_ile_helpers(): ln_weights_from_rvs = H["ln_weights_from_rvs"] ln_weights_for_posterior = H["ln_weights_for_posterior"] _rvs_is_export_resample = H["_rvs_is_export_resample"] +_equal_weight_fairdraw_for_serialization = H["_equal_weight_fairdraw_for_serialization"] class _FakeSampler(object): @@ -126,6 +127,42 @@ class _Bare(object): assert np.allclose(ln_weights_for_posterior(r, _Bare()), ln_weights_from_rvs(r)) +def test_nonfired_fairdraw_is_completed_at_the_serialization_boundary(): + """A tiny raw record must not be serialized as an equal-weight posterior. + + PR #87 intentionally preserves the weighted retained record when the + sampler-side draw would not shrink it. The XML boundary has no complete + weight provenance, so it must perform the promised draw before export. + """ + r = {"log_integrand": np.array([-1000.0, 0.0]), + "log_joint_prior": np.zeros(2), + "log_joint_s_prior": np.zeros(2), + "x": np.array([1.0, 9.0])} + + class _PeakRng(object): + def choice(self, population, size=None, replace=None, p=None): + assert replace is True + assert np.argmax(p) == 1 and p[1] == pytest.approx(1.0) + return np.full(size, 1, dtype=int) + + out = _equal_weight_fairdraw_for_serialization( + r, _FakeSampler(False), n_max=5, rng=_PeakRng()) + assert len(out["x"]) == 2 # bounded by the retained record + assert np.all(out["x"] == 9.0) # drawn by posterior weight, not flattened + + +def test_already_equal_weight_export_is_not_resampled_again(): + r = _record(n=4) + + class _NoDraw(object): + def choice(self, *args, **kwargs): + raise AssertionError("an equal-weight export was resampled again") + + out = _equal_weight_fairdraw_for_serialization( + r, _FakeSampler(True), n_max=2, rng=_NoDraw()) + assert out is r + + ### ### 2. the numbers: re-weighting a fair draw shifts the answer ### diff --git a/MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_near_aligned.py b/MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_near_aligned.py new file mode 100644 index 000000000..2c12fac14 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_near_aligned.py @@ -0,0 +1,62 @@ +"""Subprocess regression for the TEOBResumS-DALI near-aligned segfault. + +The backend is a native extension, so a regression must fail this test rather +than terminate pytest itself. CI hosts without the optional backend skip; a +host with the backend treats every later import/generation failure as real. +""" + +import importlib.util +import os +import subprocess +import sys +import textwrap + +import pytest + + +def test_near_aligned_gwsignal_call_cannot_terminate_python(): + if importlib.util.find_spec("EOBRun_module") is None: + pytest.skip("TEOBResumSDALI needs EOBRun_module, which is not installed here") + + child_code = textwrap.dedent( + r""" + import numpy as np + import lal + import lalsimulation as lalsim + from RIFT import lalsimutils + from RIFT.physics import GWSignal + + P = lalsimutils.ChooseWaveformParams() + P.m1, P.m2 = 50 * lal.MSUN_SI, 30 * lal.MSUN_SI + P.s1x, P.s1y, P.s1z = 1e-5, 0.0, 0.2 + P.s2x, P.s2y, P.s2z = 0.0, 0.0, -0.1 + P.dist = 400e6 * lal.PC_SI + P.incl = 0.4 + P.phiref = P.psi = 0.0 + P.fmin = P.fref = 20.0 + P.deltaT, P.deltaF = 1.0 / 4096, 1.0 / 16 + P.eccentricity = P.meanPerAno = 0.0 + P.taper = lalsim.SIM_INSPIRAL_TAPER_NONE + P.approx = lalsim.IMRPhenomXPHM + + modes = GWSignal.hlmoft(P, Lmax=4, approx_string="TEOBResumSDALI") + assert modes + assert all(np.isfinite(mode.data.data).all() for mode in modes.values()) + print("near-aligned-safe") + """ + ) + completed = subprocess.run( + [sys.executable, "-c", child_code], + capture_output=True, + text=True, + env=os.environ.copy(), + timeout=120, + ) + + assert completed.returncode == 0, ( + "TEOBResumSDALI near-aligned child failed (a native crash is usually " + "return code -11/139):\nstdout:\n{}\nstderr:\n{}".format( + completed.stdout, completed.stderr + ) + ) + assert "near-aligned-safe" in completed.stdout diff --git a/MonteCarloMarginalizeCode/Code/test/test_hyperpipeline_io.py b/MonteCarloMarginalizeCode/Code/test/test_hyperpipeline_io.py index 844de6eb9..af8b4bfd4 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_hyperpipeline_io.py +++ b/MonteCarloMarginalizeCode/Code/test/test_hyperpipeline_io.py @@ -73,6 +73,32 @@ def test_tides_with_eos_index(): print("test_tides_with_eos_index: OK") +def test_eccentric_tides_eob_hyperbolic_columns(): + """All JL physics groups survive named I/O and legacy adaptation together.""" + kw = dict(use_eccentricity=True, use_meanPerAno=True, + use_tides=True, use_eob_parameters=True, + use_hyperbolic=True) + cols = hpio.build_column_list(**kw) + values = { + "lnL": -3.5, "sigma_lnL": 0.02, + "m1": 2.0, "m2": 1.4, + "a1x": 0.0, "a1y": 0.0, "a1z": 0.1, + "a2x": 0.0, "a2y": 0.0, "a2z": -0.1, + "eccentricity": 0.2, "meanPerAno": 1.1, + "lambda1": 0.0, "lambda2": 300.0, + "a6c": -45.0, "E0": 1.05, "p_phi0": 4.2, + } + arr = _roundtrip(cols, [[values[c] for c in cols]]) + legacy = hpio.to_legacy_dat(arr, **kw) + ix = hpio.legacy_column_indices(**kw) + for name in ("lambda1", "a6c", "E0", "p_phi0", + "eccentricity", "meanPerAno", "lnL", "sigma_lnL"): + assert ix[name] is not None, (name, ix) + np.testing.assert_allclose(legacy[:, ix[name]], [values[name]]) + assert ix["lambda1"] < ix["a6c"] < ix["E0"] < ix["eccentricity"] < ix["lnL"] + print("test_eccentric_tides_eob_hyperbolic_columns: OK") + + def test_sky_columns(): cols = hpio.build_column_list(use_sky=True) assert cols[-2:] == ("ecliptic_longitude", "ecliptic_latitude") @@ -119,6 +145,10 @@ def test_legacy_column_indices_consistency(): dict(use_eccentricity=True), dict(use_eccentricity=True, use_meanPerAno=True), dict(use_distance=True, use_tides=True), + dict(use_eob_parameters=True), + dict(use_hyperbolic=True), + dict(use_tides=True, use_eob_parameters=True, use_hyperbolic=True, + use_eccentricity=True, use_meanPerAno=True), ] for kw in combos: cols = hpio.build_column_list(**kw) @@ -136,6 +166,10 @@ def test_legacy_column_indices_consistency(): assert ix["lambda1"] is not None if kw.get("use_eccentricity"): assert ix["eccentricity"] is not None + if kw.get("use_eob_parameters"): + assert ix["a6c"] is not None + if kw.get("use_hyperbolic"): + assert ix["E0"] is not None and ix["p_phi0"] is not None print("test_legacy_column_indices_consistency: OK") @@ -283,6 +317,7 @@ def __init__(self): self.s2x = 0.0; self.s2y = 0.0; self.s2z = 0.0 self.lambda1 = 0.0; self.lambda2 = 0.0 self.eccentricity = 0.0; self.meanPerAno = 0.0 + self.a6c = 10000.0; self.E0 = 0.0; self.p_phi0 = 0.0 self.eos_table_index = 0.0 self.dist = 0.0 self.phi = 0.0 @@ -494,6 +529,7 @@ def test_consolidate_drops_high_sigma(): test_default_roundtrip() test_eccentricity_columns() test_tides_with_eos_index() + test_eccentric_tides_eob_hyperbolic_columns() test_sky_columns() test_to_legacy_dat_default() test_legacy_column_indices_consistency() diff --git a/MonteCarloMarginalizeCode/Code/test/test_ile_scalar_edge_cases.py b/MonteCarloMarginalizeCode/Code/test/test_ile_scalar_edge_cases.py new file mode 100644 index 000000000..605f6113b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_ile_scalar_edge_cases.py @@ -0,0 +1,86 @@ +import ast +from pathlib import Path +from types import SimpleNamespace + +import numpy as np + +from RIFT.likelihood import factored_likelihood as fl + + +class _Data: + def __init__(self, values): + self.data = np.asarray(values, dtype=np.complex128) + + +class _TimeSeries: + def __init__(self, values): + self.data = _Data(values) + self.epoch = 0.0 + self.deltaT = 1.0 + + +def test_scalar_time_marginalization_uses_only_retained_modes(monkeypatch): + """A raw zero mode dropped by InterpolateRholms must not reach Ylms.""" + active = (2, 2) + dropped = (2, 1) + raw = { + "H1": { + active: _TimeSeries([1.0, 1.0, 1.0, 1.0]), + dropped: _TimeSeries([0.0, 0.0, 0.0, 0.0]), + } + } + interpolated = {"H1": {active: lambda t: np.ones_like(t)}} + cross = {"H1": {(active, active): 1.0 + 0.0j}} + cross_v = {"H1": {(active, active): 0.0 + 0.0j}} + extrinsic = SimpleNamespace( + phi=0.0, theta=0.0, tref=0.0, phiref=0.0, + incl=0.0, psi=0.0, dist=fl.distMpcRef * 1.0e6 * fl.lsu.lsu_PC, + ) + + seen_modes = [] + + def fake_ylms(_lmax, _incl, _phase, selected_modes=None): + seen_modes.extend(selected_modes) + return {mode: 1.0 + 0.0j for mode in selected_modes} + + monkeypatch.setattr(fl, "ComputeYlms", fake_ylms) + monkeypatch.setattr(fl, "ComplexAntennaFactor", lambda *args: 1.0 + 0.0j) + monkeypatch.setattr(fl, "ComputeArrivalTimeAtDetector", lambda *args: 0.0) + + result = fl.FactoredLogLikelihoodTimeMarginalized( + np.array([0.0, 1.0]), extrinsic, interpolated, raw, + cross, cross_v, Lmax=2, interpolate=False, + ) + + assert np.isfinite(result) + assert set(seen_modes) == {active} + + +def test_ile_tref_prior_uses_backend_portable_sampler_api(): + """ILE must not require uniform_samp_vector, which AV does not export.""" + ile = Path(__file__).parents[1] / "bin" / "integrate_likelihood_extrinsic_batchmode" + tree = ast.parse(ile.read_text()) + missing_api_uses = [ + node for node in ast.walk(tree) + if isinstance(node, ast.Attribute) + and isinstance(node.value, ast.Name) + and node.value.id == "mcsampler" + and node.attr == "uniform_samp_vector" + ] + assert missing_api_uses == [] + + +def test_ile_scalar_likelihood_calls_supply_raw_and_interpolated_modes(): + """The unmarginalized scalar path must honor FactoredLogLikelihood's API.""" + ile = Path(__file__).parents[1] / "bin" / "integrate_likelihood_extrinsic_batchmode" + tree = ast.parse(ile.read_text()) + calls = [ + node for node in ast.walk(tree) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "factored_likelihood" + and node.func.attr == "FactoredLogLikelihood" + ] + assert calls + assert all(len(call.args) >= 6 for call in calls) diff --git a/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py b/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py new file mode 100644 index 000000000..997e69c4b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py @@ -0,0 +1,426 @@ +"""Tests for the joint (phi, psi) peak-local kernel. + +Reference is a converged periodic trapezoid on the torus -- the mathematical content +of the shipped `anglemarg` exact scheme -- so accuracy is measured against a +quadrature, never against another peak-local run. +""" +import numpy as np +import pytest + +from RIFT.likelihood import joint_angle_peak_local as J + + +def synth_table(seed=0, scale=1.0, bidegree=(4, 2)): + """A random exact 2-D trig table of the shipped bidegree, amplitude `scale`.""" + KP, KS = bidegree[0] + 1, bidegree[1] + rng = np.random.default_rng(seed) + C = rng.normal(size=(KP, 2 * KS + 1)) + 1j * rng.normal(size=(KP, 2 * KS + 1)) + return scale * C + + +def _ref(C, n=2048): + """log[(2pi)^-2 int int exp(g)] by the periodic trapezoid (== the plain mean).""" + t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) + PHI, U = np.meshgrid(t, t, indexing='ij') + g = J.eval_g(C, PHI.ravel(), U.ravel()) + m = g.max() + return m + np.log(np.exp(g - m).mean()) + + +def test_every_exported_name_is_defined(): + """A star-import must not raise: __all__ gaining an undefined name has bitten the + companion time modules before, and only a star-import can see it.""" + missing = [n for n in J.__all__ if not hasattr(J, n)] + assert not missing, missing + ns = {} + exec("from RIFT.likelihood.joint_angle_peak_local import *", ns) + + +@pytest.mark.parametrize("scale", [1.0, 4.0, 12.0]) +def test_matches_a_converged_dense_reference(scale): + """The whole point: the same answer as the dense rule, from local work only.""" + C = synth_table(seed=3, scale=scale) + val, ok, rep = J.joint_marginalize_peak_local(C, n_phi=96) + assert ok, rep + assert abs(val - _ref(C)) < 1e-6, (scale, val, _ref(C), rep) + + +def test_derivative_bound_is_actually_a_bound(): + """M_(a,b) is the multi-index triangle inequality on the exact table. It must + never be exceeded; a fitted bound is the defect this whole design refuses.""" + C = synth_table(seed=11, scale=6.0) + t = np.linspace(0.0, 2.0 * np.pi, 401, endpoint=False) + PHI, U = np.meshgrid(t, t, indexing='ij') + for order in ((1, 0), (0, 1), (2, 0), (0, 2), (1, 1)): + d = J.eval_g(C, PHI.ravel(), U.ravel(), order) + assert np.max(np.abs(d)) <= J.derivative_bound(C, order) * (1 + 1e-12), order + + +def test_u_solve_returns_every_root_and_filters_nothing(): + """No |z| = 1 tolerance, deliberately: at exact multiplicity the computed roots + smear off the circle by eps^(1/m), so a filter drops real modes in precisely the + degenerate regime that is production. Roots are seeds; regions decide.""" + C = synth_table(seed=5, scale=3.0) + for phi in np.linspace(0, 2 * np.pi, 17): + assert J.u_stationary_at_phi(C, phi).size == 4 + + +def test_eval_g_chunking_cannot_change_the_answer(): + """_POINT_CHUNK is a memory parameter; each point is summed independently, so it + cannot move the result even in the last bit.""" + C = synth_table(seed=7, scale=2.0) + phi = np.linspace(0, 6, 257) + u = np.linspace(1, 5, 257) + big = J._POINT_CHUNK + try: + J._POINT_CHUNK = 200000 + a = J.eval_g(C, phi, u) + J._POINT_CHUNK = 13 + b = J.eval_g(C, phi, u) + finally: + J._POINT_CHUNK = big + assert a.tobytes() == b.tobytes() + + +def test_an_undersized_region_is_DECLINED_not_returned(): + """The load-bearing behaviour. W_SIGMA too small leaves mass outside the cover; + the value may still be right, but the rule cannot PROVE it and must decline. + Measured on the shipped tables: at W = 8 the margin is -18 nats against a -23 + tolerance, and at 14 it is -71 -- with the returned value identical at both.""" + C = synth_table(seed=3, scale=12.0) + keep = J.W_SIGMA + try: + J.W_SIGMA = 0.6 + val_small, ok_small, rep_small = J.joint_marginalize_peak_local(C, n_phi=96) + J.W_SIGMA = keep + val_big, ok_big, _ = J.joint_marginalize_peak_local(C, n_phi=96) + finally: + J.W_SIGMA = keep + assert not ok_small, rep_small + assert rep_small['decline'] == 'omitted-mass bound too large' + assert ok_big + + +def test_regions_merge_rather_than_double_counting(): + """Overlapping regions would count the mass between them twice. Merging is what + makes the rule degrade CONTINUOUSLY into the dense grid as amplitude falls.""" + C = synth_table(seed=3, scale=0.4) + _, _, rep = J.joint_marginalize_peak_local(C, n_phi=96) + assert rep['n_regions'] <= rep['n_modes'] + + +def _ab_tables(seed=0, scale=1.0): + """Separate A and B tables with the SHIPPED bidegrees: A is linear in the + waveform (phi <= m_max, u <= 1), B quadratic (phi <= 2 m_max, u <= 2).""" + rng = np.random.default_rng(seed) + A = rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3)) + B = rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5)) + # B must give a genuine -x^2/2 B penalty, i.e. positive mean curvature + B[0, 2] = abs(B[0, 2].real) + 3.0 + return scale * A, scale * B + + +def test_joint_table_pads_A_into_B_and_leaves_e2iu_to_B_alone(): + """The bidegrees differ, which is why the e^{2iu} coefficient carries no A + contribution -- exactly as anglemarg._laplace_psi_lnI documents. Getting this + wrong is a silent broadcast error, so it is pinned.""" + A, B = _ab_tables(seed=1) + x = 0.7 + C = J.joint_table(A, B, x=x) + assert C.shape == B.shape + # the q = +2 column (index 4) must be EXACTLY -x^2/2 * B there: no A contribution + assert np.array_equal(C[:, 4], (-0.5 * x * x) * B[:, 4].astype(complex)) + # and A must have landed in the central 3 columns of the first 3 rows + assert np.allclose(C[:3, 1:4], x * A + (-0.5 * x * x) * B[:3, 1:4]) + # rows beyond A's phi-degree keep B alone + assert np.array_equal(C[3:, :], (-0.5 * x * x) * B[3:, :].astype(complex)) + + +def test_distance_composition_matches_a_direct_sum_over_nodes(): + """The composed value must equal an explicit logsumexp over ALL nodes; the + node pre-filter may only drop what it can bound as negligible.""" + A, B = _ab_tables(seed=2, scale=1.5) + x = np.linspace(0.35, 3.0, 24) + logw = -0.5 * (x - 1.3) ** 2 * 4.0 + val, ok, rep = J.joint_marginalize_over_distance(A, B, x, logw, n_phi=64, + n_bound_grid=128) + assert ok, rep + direct = [] + for xi, wi in zip(x, logw): + v, o, _ = J.joint_marginalize_peak_local(J.joint_table(A, B, x=float(xi)), + n_phi=64, n_bound_grid=128) + direct.append(wi + v) + m = max(direct) + ref = m + np.log(np.exp(np.array(direct) - m).sum()) + assert abs(val - ref) < 1e-9, (val, ref, rep) + + +def test_the_node_prefilter_cannot_change_the_answer(): + """The filter is a COST optimization only. A cut it cannot JUSTIFY against the + computed value is undone by retrying with every node, so however aggressively it is + asked to cut, the answer must not move. (Tightening `keep_nats` therefore does NOT + imply fewer live nodes -- an unjustified cut is reverted, which is the point.)""" + A, B = _ab_tables(seed=4, scale=2.0) + x = np.linspace(0.3, 4.0, 40) + logw = -0.5 * (x - 1.0) ** 2 * 3.0 + wide, _, rw = J.joint_marginalize_over_distance(A, B, x, logw, n_phi=64, + n_bound_grid=128, keep_nats=np.inf) + tight, _, rt = J.joint_marginalize_over_distance(A, B, x, logw, n_phi=64, + n_bound_grid=128, keep_nats=25.0) + assert abs(wide - tight) < 1e-8, (wide, tight, rw, rt) + # an aggressive cut here is not justifiable, so it must have been retried + assert rt.get('prefilter_retried') or rt['n_nodes_live'] == x.size, rt + + +def test_an_unjustified_node_cut_is_retried_not_reported(): + """Regression: `dropped_bound` was once computed, stored in the report, and never + compared to anything -- the docstring promised a provably-negligible drop while the + code performed an unchecked one.""" + A, B = _ab_tables(seed=6, scale=2.5) + x = np.linspace(0.3, 4.0, 32) + logw = -0.5 * (x - 1.0) ** 2 * 3.0 + _, ok, rep = J.joint_marginalize_over_distance(A, B, x, logw, n_phi=64, + n_bound_grid=128, keep_nats=5.0) + # either the cut was justified, or it was undone -- never silently kept + assert ok + assert rep.get('prefilter_retried') or rep.get('dropped_margin', -np.inf) < J.OUTSIDE_TOL_NATS + + +def test_a_cell_straddling_a_box_edge_counts_as_OUTSIDE(): + """P1 from review. Classifying grid CENTRES lets a box cover every centre while + leaving real area uncovered; outside_bound then returned (-inf, 0.0) -- unconditional + acceptance. The reviewer's counterexample: n_grid = 8, a box offset half a step with + half-width pi - h/4 covers all 64 centres and still leaves ~4.8 rad^2 outside.""" + A, B = _ab_tables(seed=2, scale=1.0) + C = J.joint_table(A, B, x=1.0) + n = 8 + h = 2 * np.pi / n + cen = np.array([[h / 2.0, h / 2.0]]) + half = np.array([[np.pi - h / 4.0, np.pi - h / 4.0]]) + sup, area = J.outside_bound(C, cen, half, n_grid=n) + assert area > 0.0, "a straddling cell must not be reported as fully covered" + assert np.isfinite(sup), "an uncovered region must yield a finite supremum bound" + # and the reported area must OVER-estimate, never under-estimate, the true uncovered + true_uncovered = (2 * np.pi) ** 2 - (2 * (np.pi - h / 4.0)) ** 2 + assert area >= true_uncovered - 1e-9, (area, true_uncovered) + + +def test_a_full_circle_box_still_counts_as_covering(): + """The conservative shrink must not fire on a box that already spans the circle -- + that is the low-amplitude case where regions merge to the whole torus, which is the + rule degenerating into the dense grid on purpose.""" + A, B = _ab_tables(seed=2, scale=1.0) + C = J.joint_table(A, B, x=1.0) + cen = np.array([[0.0, 0.0]]) + half = np.array([[np.pi, np.pi]]) + sup, area = J.outside_bound(C, cen, half, n_grid=32) + assert area == 0.0 and sup == -np.inf, (sup, area) + + +# ------------------------------------------------- both axes localized (phi-local) + +def test_u_profile_derivatives_match_finite_differences(): + """F' and F'' come from differentiating UNDER the integral -- F' = E[d_phi g], + F'' = E[d^2_phi g] + Var(d_phi g) -- so they are exact and cost no extra evaluation. + That identity is what makes localizing phi possible at all, so it is pinned against + finite differences of F itself.""" + A, B = _ab_tables(seed=3, scale=3.0) + C = J.joint_table(A, B, x=1.0) + h = 1e-5 + for phi in np.linspace(0.3, 5.9, 6): + F0, d1, d2 = J.u_profile(C, np.array([phi])) + Fp, _, _ = J.u_profile(C, np.array([phi + h])) + Fm, _, _ = J.u_profile(C, np.array([phi - h])) + fd1 = (Fp[0] - Fm[0]) / (2 * h) + fd2 = (Fp[0] - 2 * F0[0] + Fm[0]) / h ** 2 + assert abs(d1[0] - fd1) < 1e-4 * max(1.0, abs(fd1)), (phi, d1[0], fd1) + assert abs(d2[0] - fd2) < 1e-2 * max(1.0, abs(fd2)), (phi, d2[0], fd2) + + +@pytest.mark.parametrize("scale", [1.0, 3.0, 10.0]) +def test_phi_local_matches_a_converged_dense_reference(scale): + """Both axes localized, against a dense torus quadrature.""" + A, B = _ab_tables(seed=3, scale=1.0) + C = J.joint_table(A * scale, B * scale, x=1.0) + val, ok, rep = J.phi_local_marginalize(C) + assert ok, rep + assert abs(val - _ref(C, n=2048)) < 1e-4, (scale, val, rep) + + +def test_phi_local_cost_does_not_grow_with_amplitude(): + """The whole point of localizing BOTH axes: the dense rule spends ~A points on the + (phi,u) product, this spends a number set by the mode structure, which does not move + when the data amplitude does.""" + A, B = _ab_tables(seed=3, scale=1.0) + counts = [] + for scale in (1.0, 10.0, 100.0): + _, ok, rep = J.phi_local_marginalize(J.joint_table(A * scale, B * scale, x=1.0)) + assert ok + counts.append(rep['n_phi_regions']) + assert max(counts) <= 8, counts # bounded, not growing with sqrt(A) + + +def test_phi_cover_bound_is_routed_through_g_not_through_F_curvature(): + """Regression. Bounding F by Taylor with F'' <= M_(2,0) + M_(1,0)^2 is useless: that + variance bound grows as the SQUARE of the amplitude and produced margins of +51 and + +1196 nats (no bound at all). Routing through F <= log(2pi) + sup_u g keeps the + remainder linear in amplitude, so high-amplitude rows are ACCEPTED rather than + declined for a defect in the bound.""" + A, B = _ab_tables(seed=3, scale=1.0) + for scale in (30.0, 100.0): + _, ok, rep = J.phi_local_marginalize(J.joint_table(A * scale, B * scale, x=1.0)) + assert ok, (scale, rep) + assert rep['margin'] < J.OUTSIDE_TOL_NATS, (scale, rep) + + +def test_a_wrapped_phi_region_is_clamped_to_one_circuit(): + """Regression, found on REAL coefficient tables and not reachable from the synthetic + ones. At low amplitude F is nearly flat, so sigma is huge and [p-W*sig, p+W*sig] + spans more than 2 pi; integrating that range literally wraps the circle several times + and counts the same mass repeatedly -- measured +1.84 nats, a factor of e^1.84 = 6.3, + and ACCEPTED, because a region covering everything leaves nothing outside for the + omitted-mass certificate to object to. + + The general lesson, worth more than the fix: the certificate bounds what is OUTSIDE + the regions and cannot see an error made INSIDE one.""" + A, B = _ab_tables(seed=1, scale=0.05) # deliberately near-flat + C = J.joint_table(A, B, x=0.3) + val, ok, rep = J.phi_local_marginalize(C) + assert ok, rep + ref = _ref(C, n=2048) + assert abs(val - ref) < 1e-3, (val, ref, rep) + # and the covered length may never exceed one circuit + assert rep['n_phi_regions'] >= 1 + + +def test_inner_u_quadrature_is_converged_at_the_default_node_count(): + """P1 from review, and it needs its OWN regression because the phi omitted-mass + certificate cannot see it: an error made INSIDE a region is invisible to a bound on + what lies outside. Raising n_nodes must not move the answer.""" + rng = np.random.default_rng(11) + worst = 0.0 + for _ in range(6): + sc = 10.0 ** rng.uniform(-0.5, 2.5) + A = (rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3))) * sc + B = (rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5))) * sc + B[0, 2] = abs(B[0, 2].real) + 3.0 * sc + C = J.joint_table(A, B, x=1.0) + for phi in (0.3, 1.9, 4.4): + lo, _, _ = J.u_profile(C, np.array([phi]), n_nodes=64) + hi, _, _ = J.u_profile(C, np.array([phi]), n_nodes=1024) + worst = max(worst, abs(lo[0] - hi[0])) + # 1e-4, matching this module's other marginal assertions. An earlier revision of + # this test allowed 1e-3, which CODIFIED a 7.2e-4 inner-u error rather than catching + # it -- a tolerance chosen after seeing the number is not a check. + assert worst < 1e-4, worst + + +def test_a_clipped_newton_point_is_not_classified_as_a_peak(): + """P1 from review. The in-cell Newton is clamped to [lo, mid], so it can come to + rest ON a boundary with a large stationary residual -- and curvature alone then calls + that a maximum, centring a +-W sigma window on a non-stationary point. Measured over + 5400 cells: 492 (18%) that g'' < 0 accepted are rejected once a small residual and an + interior position are also required, the worst of them sitting at |g_u|/M_1 = 0.33.""" + rng = np.random.default_rng(7) + n_curv, n_gated = 0, 0 + for _ in range(20): + sc = 10.0 ** rng.uniform(-0.5, 3.0) + A = (rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3))) * sc + B = (rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5))) * sc + B[0, 2] = abs(B[0, 2].real) + 3.0 * sc + C = J.joint_table(A, B, x=1.0) + k, q, w, KS = J._kq(C) + for phi in np.linspace(0.05, 6.2, 5): + ph = (np.exp(1j * phi * k) * w).ravel() + D = lambda qq: complex((ph * C[:, KS + qq]).sum()) + c1 = D(1) + np.conj(D(-1)) + c2 = D(2) + np.conj(D(-2)) + u = np.sort(np.mod(np.angle(np.roots( + [c2, c1 / 2, 0, -np.conj(c1) / 2, -np.conj(c2)])), 2 * np.pi)) + mid = 0.5 * (u + np.roll(u, -1) + + np.where(np.arange(4) == 3, 2 * np.pi, 0.0)) + lo_c = np.roll(mid, 1) - np.where(np.arange(4) == 0, 2 * np.pi, 0.0) + us = u.copy() + pv = np.full(4, phi) + for _i in range(8): + g1 = J.eval_g(C, pv, us, (0, 1)) + g2 = J.eval_g(C, pv, us, (0, 2)) + st = np.where(np.abs(g2) > 0, -g1 / np.where(np.abs(g2) > 0, g2, 1.0), 0.0) + us = np.clip(us + np.clip(st, -0.5, 0.5), lo_c, mid) + g1c = J.eval_g(C, pv, us, (0, 1)) + g2c = J.eval_g(C, pv, us, (0, 2)) + m1u = max(J.derivative_bound(C, (0, 1)), 1e-300) + edge = 1e-9 * max(float(np.max(mid - lo_c)), 1e-300) + curv = g2c < 0.0 + gated = curv & (np.abs(g1c) <= 1e-8 * m1u) & (us > lo_c + edge) & (us < mid - edge) + n_curv += int(curv.sum()) + n_gated += int(gated.sum()) + assert n_curv > n_gated, (n_curv, n_gated) # the gate must actually reject + + +def test_the_quartic_roots_are_seeds_and_must_be_refined_in_cell(): + """Why the in-cell Newton refinement is not dead weight. A conjugate-reciprocal + pair leaves the unit circle -- measured, 309 of 900 (table, phi) pairs have at least + one such root -- and a spurious root's ANGLE is not a stationary point at all: the + worst |g_u|/M_1 at a raw root measured 0.311, i.e. nowhere near stationary. Window + around that and the window is centred off the peak.""" + rng = np.random.default_rng(11) + worst_resid = 0.0 + n_off = 0 + for _ in range(60): + sc = 10.0 ** rng.uniform(-0.5, 2.0) + A = (rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3))) * sc + B = (rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5))) * sc + B[0, 2] = abs(B[0, 2].real) + 3.0 * sc + C = J.joint_table(A, B, x=1.0) + k, q, w, KS = J._kq(C) + for phi in (0.3, 1.9, 4.4): + ph = (np.exp(1j * phi * k) * w).ravel() + D = lambda qq: complex((ph * C[:, KS + qq]).sum()) + c1 = D(1) + np.conj(D(-1)) + c2 = D(2) + np.conj(D(-2)) + z = np.roots([c2, c1 / 2, 0, -np.conj(c1) / 2, -np.conj(c2)]) + n_off += int(np.any(np.abs(np.abs(z) - 1.0) > 1e-6)) + u = np.sort(np.mod(np.angle(z), 2 * np.pi)) + g1 = J.eval_g(C, np.full(4, phi), u, (0, 1)) + m1 = max(J.derivative_bound(C, (0, 1)), 1e-300) + worst_resid = max(worst_resid, float(np.max(np.abs(g1)) / m1)) + assert n_off > 0, "fixture family must contain off-circle roots" + assert worst_resid > 1e-6, worst_resid # raw roots are NOT all stationary points + + +def test_phi_regions_are_disjoint_on_the_CIRCLE(): + """P1 from review. Merging on the LINE never joins a window near 0 to one near + 2*pi, but every region is integrated at mod(., 2*pi) -- so both regions cover both + peaks and the mass is counted twice. Measured before the fix: +log 2 = +0.693 nats + returned with ok=True and margin -437, because the error is INSIDE the regions. + + Asserted as an INVARIANT rather than hunted for with a value comparison: reduced to + the circle, the regions must not overlap and must not exceed one circuit.""" + rng = np.random.default_rng(5) + checked = 0 + for _ in range(40): + sc = 10.0 ** rng.uniform(0.0, 2.5) + A = (rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3))) * sc + B = (rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5))) * sc + B[0, 2] = abs(B[0, 2].real) + 3.0 * sc + C = J.joint_table(A, B, x=1.0) + _, _, rep = J.phi_local_marginalize(C) + regs = rep.get('phi_regions', []) + if not regs: + continue + checked += 1 + total = sum(b - a for a, b in regs) + assert total <= 2 * np.pi + 1e-9, (total, regs) + # sample each region densely, reduce to the circle, and require no point to be + # covered twice + pts = [] + for a, b in regs: + pts.append(np.mod(np.linspace(a, b, 512, endpoint=False), 2 * np.pi)) + if len(pts) > 1: + for i in range(len(pts)): + for j in range(i + 1, len(pts)): + d = np.abs(pts[i][:, None] - pts[j][None, :]) + d = np.minimum(d, 2 * np.pi - d) + assert d.min() > 1e-6, ("regions overlap on the circle", regs) + assert checked > 5, checked diff --git a/MonteCarloMarginalizeCode/Code/test/test_limit_distance.py b/MonteCarloMarginalizeCode/Code/test/test_limit_distance.py new file mode 100644 index 000000000..38d2a5012 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_limit_distance.py @@ -0,0 +1,483 @@ +#!/usr/bin/env python +""" +Regression tests for --limit-distance: a SAMPLING-only restriction of the distance +range that must leave the evidence normalization alone. + +WHY THIS IS NOT THE SAME AS THE ANGULAR --limit-* BOXES. Those narrow the prior +SUPPORT: the angle priors are never renormalized, so a box costs exactly the prior +mass outside it and lnZ moves by an analytic amount (a whole-sky box reproduces the +unboxed lnZ to +0.000000; a narrow one returns ln(dOmega/4pi)). --limit-distance +instead keeps the prior AND its normalization over the full [--d-min,--d-max] and +narrows only what the sampler draws from, so the reported lnZ is the SAME NUMBER a +full-range run reports -- no correction -- to the extent the likelihood is negligible +outside the box. That is what makes lnZ comparable across runs and across samplers +while the quadrature gets to stay cheap at high amplitude, where the distance +posterior has narrowed as 1/rho. + +THE DEFECT THESE TESTS LOCK DOWN. The obvious implementation narrows +param_limits["distance"], which the Euclidean prior lambda ALSO read for its own +normalization: + + dist_prior_pdf = lambda x: x**2/(param_limits["distance"][1]**3/3. + - param_limits["distance"][0]**3/3.) + +so the density silently renormalizes to the box. The failure is invisible in the +usual acceptance check -- with the likelihood inside the box, lnZ still comes back +"unchanged" -- because the renormalization exactly cancels the missing prior mass. +The signature that separates the two is a CONSTANT likelihood: with the prior +correctly left alone, narrowing must cost exactly the prior mass outside the box +(test_constant_likelihood_*), and it is the renormalizing implementation that +returns "unchanged" there. Both directions are tested. +""" + +import os +import subprocess +import sys + +import numpy as np +import pytest + +import RIFT.integrators.mcsampler as mcsampler +from RIFT.integrators.mcsampler import distance_limit_range, distance_sampler_kwargs + +# np.trapz was REMOVED in numpy 2.x (renamed np.trapezoid); CI runs both lanes. +_trapz = getattr(np, 'trapezoid', None) or np.trapz + +# The physical prior range: what --d-min/--d-max set, and what the prior must stay +# normalized over no matter how the sampling is restricted. +D_MIN, D_MAX = 1.0, 10000.0 + +# A synthetic distance posterior: a Gaussian peak, and a box at +-6 sigma around it. +# The analytic prior mass of the likelihood OUTSIDE the box is ~2e-9 of the total, so +# the "unchanged lnZ" claim has an exact reference to be checked against and is not a +# statement about the tolerance of a sampler. +D_PEAK, D_SIGMA = 2000.0, 200.0 +BOX = (D_PEAK - 6 * D_SIGMA, D_PEAK + 6 * D_SIGMA) # (800, 3200) + + +def _like(d): + d = np.asarray(d, dtype=float) + return np.exp(-0.5 * ((d - D_PEAK) / D_SIGMA) ** 2) + + +def _reference_integral(lo, hi, d_prior='Euclidean', like=_like, n=2000001): + """int_lo^hi L(d) pi(d) dd with pi normalized over the FULL [D_MIN, D_MAX]. + + This is the expectation of the estimator the sampler forms, computed without + Monte Carlo noise, so a disagreement is a normalization disagreement. + """ + kw = distance_sampler_kwargs(mcsampler, (lo, hi), (D_MIN, D_MAX), d_prior=d_prior) + grid = np.linspace(lo, hi, n) + return _trapz(like(grid) * np.asarray(kw['prior_pdf'](grid), dtype=float), grid) + + +### +### 1. The option parser: loud on anything ambiguous +### + +def test_limit_range_parses_and_validates(): + assert distance_limit_range('800,3200', D_MIN, D_MAX) == (800.0, 3200.0) + assert distance_limit_range(' 800 , 3200 ', D_MIN, D_MAX) == (800.0, 3200.0) + + +@pytest.mark.parametrize('bad', ['3200,800', # inverted + '800,800', # empty + '800', # not a pair + '800,900,1000', # too many + 'a,b', # unparseable + 'nan,3200', # non-finite + ]) +def test_limit_range_rejects_bad_requests(bad): + with pytest.raises(ValueError): + distance_limit_range(bad, D_MIN, D_MAX) + + +@pytest.mark.parametrize('bad', ['0.5,3200', # below --d-min + '800,20000', # above --d-max + ]) +def test_limit_range_must_lie_inside_the_prior_range(bad): + """A box outside [d_min,d_max] asks for distances the prior gives no mass. + Clipping it silently would turn a configuration error into a quiet answer.""" + with pytest.raises(ValueError): + distance_limit_range(bad, D_MIN, D_MAX) + + +### +### 2. The prior is NOT renormalized to the box +### + +@pytest.mark.parametrize('d_prior', ['Euclidean', 'pseudo_cosmo']) +def test_prior_pdf_is_normalized_over_the_full_range_not_the_box(d_prior): + kw_full = distance_sampler_kwargs(mcsampler, (D_MIN, D_MAX), (D_MIN, D_MAX), d_prior=d_prior) + kw_box = distance_sampler_kwargs(mcsampler, BOX, (D_MIN, D_MAX), d_prior=d_prior) + grid = np.linspace(D_MIN, D_MAX, 200001) + p_full = np.asarray(kw_full['prior_pdf'](grid), dtype=float) + p_box = np.asarray(kw_box['prior_pdf'](grid), dtype=float) + # identical densities: the box changed the SAMPLER, not the prior + assert np.allclose(p_full, p_box, rtol=0, atol=0) + assert _trapz(p_box, grid) == pytest.approx(1.0, rel=1e-6) + # and the box-renormalized density -- the defect -- is a DIFFERENT function + kw_wrong = distance_sampler_kwargs(mcsampler, BOX, BOX, d_prior=d_prior) + inbox = (grid >= BOX[0]) & (grid <= BOX[1]) + p_wrong = np.asarray(kw_wrong['prior_pdf'](grid), dtype=float) + assert not np.allclose(p_wrong[inbox], p_box[inbox]) + # (its own grid, so the endpoints land exactly on the box) + gbox = np.linspace(BOX[0], BOX[1], 200001) + assert _trapz(np.asarray(kw_wrong['prior_pdf'](gbox), dtype=float), gbox) == pytest.approx(1.0, rel=1e-6) + + +def test_euclidean_prior_normalization_is_the_full_range_analytic_one(): + kw = distance_sampler_kwargs(mcsampler, BOX, (D_MIN, D_MAX), d_prior='Euclidean') + d = np.array([500.0, 2000.0, 9000.0]) + expected = d ** 2 / (D_MAX ** 3 / 3. - D_MIN ** 3 / 3.) + assert np.allclose(np.asarray(kw['prior_pdf'](d), dtype=float), expected) + + +def test_unknown_prior_is_refused(): + with pytest.raises(ValueError): + distance_sampler_kwargs(mcsampler, BOX, (D_MIN, D_MAX), d_prior='cosmo') + + +def test_default_path_reproduces_the_historical_expressions_bitwise(): + """No existing default may move. With sampling_range == prior_range the helper + must reproduce, BIT FOR BIT, the expressions the ILE driver used before + --limit-distance existed (the pre-change source is quoted in each comment).""" + import functools + import RIFT.likelihood.priors_utils as priors_utils + x = np.linspace(D_MIN, D_MAX, 100001) + u = np.linspace(0.0, 1.0, 100001) + + # dist_prior_pdf = lambda x: x**2/(param_limits["distance"][1]**3/3. + # - param_limits["distance"][0]**3/3.) + kw = distance_sampler_kwargs(mcsampler, (D_MIN, D_MAX), (D_MIN, D_MAX), d_prior='Euclidean') + old = x ** 2 / (D_MAX ** 3 / 3. - D_MIN ** 3 / 3.) + assert np.array_equal(np.asarray(kw['prior_pdf'](x), dtype=float), old) + + # nm = priors_utils.dist_prior_pseudo_cosmo_eval_norm(lo, hi) + # dist_prior_pdf = functools.partial(priors_utils.dist_prior_pseudo_cosmo, nm=nm, xpy=...) + kwp = distance_sampler_kwargs(mcsampler, (D_MIN, D_MAX), (D_MIN, D_MAX), d_prior='pseudo_cosmo') + nm = priors_utils.dist_prior_pseudo_cosmo_eval_norm(D_MIN, D_MAX) + old_p = priors_utils.dist_prior_pseudo_cosmo(x, nm=nm, xpy=np) + assert np.array_equal(np.asarray(kwp['prior_pdf'](x), dtype=float), np.asarray(old_p, dtype=float)) + + # dist_sampler = mcsampler.ret_uniform_samp_vector_alt(lo, hi) + # dist_sampler_cdf_inv = functools.partial(mcsampler.uniform_samp_cdf_inv_vector, lo, hi) + assert np.array_equal(np.asarray(kw['pdf'](x)), + np.asarray(mcsampler.ret_uniform_samp_vector_alt(D_MIN, D_MAX)(x))) + assert np.array_equal(np.asarray(kw['cdf_inv'](u)), + np.asarray(functools.partial(mcsampler.uniform_samp_cdf_inv_vector, + D_MIN, D_MAX)(u))) + # left_limit = param_limits["distance"][0], right_limit = param_limits["distance"][1] + assert kw['left_limit'] == D_MIN and kw['right_limit'] == D_MAX + + +### +### 3. THE ACCEPTANCE TEST: narrowed vs full-range evidence +### + +@pytest.mark.parametrize('d_prior', ['Euclidean', 'pseudo_cosmo']) +def test_ACCEPTANCE_narrowed_and_full_range_evidence_agree(d_prior): + """The reported evidence must not move when the sampling range is narrowed. + + Evaluated as the estimator's EXPECTATION (deterministic quadrature of + L*pi over [left_limit,right_limit]) so the number is a statement about the + normalization and not about a sampler's variance. The residual is the + likelihood mass outside +-6 sigma, which is real physics, not an error: it + is why the option's help says to keep the box wide vs the posterior. + """ + z_full = _reference_integral(D_MIN, D_MAX, d_prior) + z_box = _reference_integral(BOX[0], BOX[1], d_prior) + dlnZ = abs(np.log(z_box) - np.log(z_full)) + assert dlnZ < 1e-6, "narrowing moved lnZ by {:.3e} nats".format(dlnZ) + + +@pytest.mark.parametrize('d_prior', ['Euclidean', 'pseudo_cosmo']) +def test_ACCEPTANCE_the_renormalizing_implementation_would_fail_it(d_prior): + """Power check for the test above: if the prior were renormalized to the box + (prior_range narrowed along with sampling_range -- the defect), lnZ would move + by ln of the prior-mass fraction, which for this box is several nats.""" + z_full = _reference_integral(D_MIN, D_MAX, d_prior) + kw = distance_sampler_kwargs(mcsampler, BOX, BOX, d_prior=d_prior) # the defect + grid = np.linspace(BOX[0], BOX[1], 2000001) + z_wrong = _trapz(_like(grid) * np.asarray(kw['prior_pdf'](grid), dtype=float), grid) + assert abs(np.log(z_wrong) - np.log(z_full)) > 1.0 + + +@pytest.mark.parametrize('d_prior', ['Euclidean', 'pseudo_cosmo']) +def test_constant_likelihood_narrowing_costs_exactly_the_prior_mass(d_prior): + """With L == 1 there is no 'negligible outside' region, so the CORRECT answer is + that narrowing reduces the integral to the prior mass of the box. An + implementation that renormalizes the prior returns 1.0 here -- i.e. 'unchanged', + which is the signature of the defect, not of success.""" + one = lambda d: np.ones(np.shape(d)) + z_full = _reference_integral(D_MIN, D_MAX, d_prior, like=one) + z_box = _reference_integral(BOX[0], BOX[1], d_prior, like=one) + assert z_full == pytest.approx(1.0, rel=1e-6) # prior integrates to one + assert z_box < 0.5 * z_full # narrowing DID cost mass + if d_prior == 'Euclidean': + expected = ((BOX[1] ** 3 - BOX[0] ** 3) / (D_MAX ** 3 - D_MIN ** 3)) + assert z_box == pytest.approx(expected, rel=1e-6) + + +### +### 3b. The PHYSICAL likelihood shape, and where the claim stops holding +### +# The Gaussian above is a convenient stand-in, but it is not the shape a real +# extrinsic likelihood has in distance: exp(K x - R x^2/2) with x = d_ref/d tends to +# ONE as d -> infinity, not to zero. So there is always a far-field contribution +# ~ (prior mass outside the box) x 1 that the box throws away, and "narrowing costs +# no evidence" is a statement about AMPLITUDE, not a theorem. These tests MEASURE +# the crossover rather than assume it -- which is what the option's help warns about, +# and the reason it is opt-in per run. +# +# Everything here is done in log space: lnL at the peak is rho^2/2, which overflows +# a float at rho ~ 38. + +D_REF = 1000.0 + + +def _physical_log_like(rho, d_star): + """lnL(d) = K x - R x^2/2, x = D_REF/d; peak lnL = rho^2/2 at d = d_star.""" + x_star = D_REF / d_star + R = rho ** 2 / x_star ** 2 + K = R * x_star + def log_like(d): + x = D_REF / np.asarray(d, dtype=float) + return K * x - 0.5 * R * x ** 2 + return log_like + + +def _reference_lnZ(lo, hi, log_like, d_prior='Euclidean', n=400001): + """ln int_lo^hi L(d) pi(d) dd, with pi normalized over the FULL [D_MIN,D_MAX].""" + from scipy.special import logsumexp + kw = distance_sampler_kwargs(mcsampler, (lo, hi), (D_MIN, D_MAX), d_prior=d_prior) + grid = np.linspace(lo, hi, n) + dd = np.full(n, grid[1] - grid[0]) + dd[0] = dd[-1] = 0.5 * (grid[1] - grid[0]) # trapezoid + log_pi = np.log(np.asarray(kw['prior_pdf'](grid), dtype=float)) + return float(logsumexp(log_like(grid) + log_pi + np.log(dd))) + + +# MEASURED on this configuration (d_star=2000 Mpc, +-6/rho fractional box, prior +# [1,10000] Mpc): |dlnZ| = 6.2e-4 at rho=10, 5.8e-6 at rho=20, 1.6e-7 at rho=40 +# (and 1.0e-1 at rho=5, which is why the crossover gets its own test below). +# The tolerances below are those numbers rounded up, not aspirations. +@pytest.mark.parametrize('rho,tol', [(10.0, 1e-3), (20.0, 1e-5), (40.0, 1e-6)]) +def test_ACCEPTANCE_physical_likelihood_shape(rho, tol): + d_star = 2000.0 + log_like = _physical_log_like(rho, d_star) + half = 6.0 / rho + lo, hi = d_star * (1 - half), d_star * (1 + half) + lnz_full = _reference_lnZ(D_MIN, D_MAX, log_like) + lnz_box = _reference_lnZ(lo, hi, log_like) + dlnZ = abs(lnz_box - lnz_full) + assert dlnZ < tol, "rho={}: narrowing moved lnZ by {:.3e} nats".format(rho, dlnZ) + + +def test_the_far_field_is_what_breaks_it_at_low_amplitude(): + """The honest negative, recorded so the crossover is documented and not assumed: + at rho=2 the distance posterior is not localized -- the prior's own d^2 mass at + large d carries the integral -- and a +-30% box costs 3.3 nats (MEASURED). + A run that narrows the box at low amplitude gets a wrong lnZ, quietly.""" + d_star = 2000.0 + log_like = _physical_log_like(2.0, d_star) + lo, hi = d_star * 0.7, d_star * 1.3 + lnz_full = _reference_lnZ(D_MIN, D_MAX, log_like) + lnz_box = _reference_lnZ(lo, hi, log_like) + assert lnz_full - lnz_box > 0.5 + + +### +### 4. Through the real MCSampler construction path +### +# distance_sampler_kwargs() output is fed to add_parameter() verbatim by the ILE +# driver, so building a real sampler from it is the construction path. These assert +# on the sampler's ACTUAL bounds: a --limit-distance that parsed and then did nothing +# (this codebase's documented failure mode) fails here. + +@pytest.mark.parametrize('module_name', ['RIFT.integrators.mcsampler', + 'RIFT.integrators.mcsamplerGPU', + 'RIFT.integrators.mcsamplerAdaptiveVolume']) +def test_sampler_bounds_are_the_narrowed_ones(module_name): + mod = pytest.importorskip(module_name) + kw = distance_sampler_kwargs(mod, BOX, (D_MIN, D_MAX), d_prior='Euclidean') + s = mod.MCSampler() + s.add_parameter("distance", **kw) + assert s.llim["distance"] == pytest.approx(BOX[0]) + assert s.rlim["distance"] == pytest.approx(BOX[1]) + # ... and NOT the full range: the disconnected-flag failure mode + assert s.rlim["distance"] < D_MAX + # the prior the sampler will weight with is still the full-range one + d = np.array([2000.0]) + assert float(np.asarray(s.prior_pdf["distance"](d), dtype=float)[0]) == pytest.approx( + 2000.0 ** 2 / (D_MAX ** 3 / 3. - D_MIN ** 3 / 3.)) + + +def test_sampler_bounds_are_the_full_range_without_the_option(): + kw = distance_sampler_kwargs(mcsampler, (D_MIN, D_MAX), (D_MIN, D_MAX)) + s = mcsampler.MCSampler() + s.add_parameter("distance", **kw) + assert s.llim["distance"] == pytest.approx(D_MIN) + assert s.rlim["distance"] == pytest.approx(D_MAX) + + +def test_gpu_sampler_draws_inside_the_box(): + """The actual mcsamplerGPU draw path (CPU fallback when cupy is absent): the + narrowed pdf/cdf_inv must place every draw inside the box, and fill it.""" + mcsamplerGPU = pytest.importorskip('RIFT.integrators.mcsamplerGPU') + kw = distance_sampler_kwargs(mcsamplerGPU, BOX, (D_MIN, D_MAX), d_prior='Euclidean') + s = mcsamplerGPU.MCSampler() + s.add_parameter("distance", **kw) + rv = s.draw_simplified(4000, "distance")[-1] + drawn = np.asarray(mcsamplerGPU.identity_convert(rv)).reshape(-1) + assert drawn.min() >= BOX[0] - 1e-9 + assert drawn.max() <= BOX[1] + 1e-9 + assert drawn.max() - drawn.min() > 0.8 * (BOX[1] - BOX[0]) + # the whole point: the draws are NOT spread over the physical prior range + assert drawn.max() < 0.5 * D_MAX + + +def test_end_to_end_MCSampler_integral_is_unchanged_by_the_box(): + """Same claim as the acceptance test, but paid for with real Monte Carlo: + mcsampler.MCSampler weights each draw by prior_pdf/pdf, so this exercises the + sampling density as well as the prior. Tolerance is set by MC noise of the + FULL-range run (the box is ~24% of the prior mass here), not by the physics.""" + np.random.seed(20260901) + ref = _reference_integral(D_MIN, D_MAX) + + def _run(rng_seed, sampling_range): + np.random.seed(rng_seed) + kw = distance_sampler_kwargs(mcsampler, sampling_range, (D_MIN, D_MAX)) + s = mcsampler.MCSampler() + s.add_parameter("distance", **kw) + res = s.integrate(lambda distance: _like(distance), "distance", + nmax=400000, n=4000, no_protect_names=True, verbose=False) + return float(res[0]) + + z_full = _run(1234, (D_MIN, D_MAX)) + z_box = _run(1234, BOX) + assert z_full == pytest.approx(ref, rel=0.05) + assert z_box == pytest.approx(ref, rel=0.01) + assert abs(np.log(z_box) - np.log(z_full)) < 0.05 + + +### +### 5. Driver wiring +### +# The library helper can be perfect and the driver still normalize over the narrowed +# range: that is one edit away, in a file with no cheap end-to-end test (the distance +# block sits ~1300 lines in, after data loading). So: assert the driver hands the +# helper the UNNARROWED range, and that the old conflated normalization is gone. + +_ILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', 'bin', 'integrate_likelihood_extrinsic_batchmode') + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_passes_the_unnarrowed_range_as_the_prior_normalization(): + with open(_ILE) as f: + src = f.read() + assert '--limit-distance' in src + # captured BEFORE the narrowing, and handed to the helper as prior_range + assert 'dist_prior_range = (param_limits["distance"][0], param_limits["distance"][1])' in src + assert 'distance_sampler_kwargs(\n mcsampler, param_limits["distance"], dist_prior_range,' in src + # the conflated normalization must not come back + assert 'param_limits["distance"][1]**3/3.' not in src + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_refuses_the_incompatible_distance_modes(): + with open(_ILE) as f: + src = f.read() + for needle in ('--limit-distance is not compatible with --distance-marginalization', + '--limit-distance is not compatible with --d-prior-redshift', + '--limit-distance is not compatible with --internal-reparam-dl-incl'): + assert needle in src + + +def _run_ile_narrowing_block(limit_distance, **optkw): + """Execute the DRIVER'S OWN narrowing block, verbatim from the file, against a + stub `opts`/`param_limits`. + + The block sits ~1375 lines into a monolithic script, after data loading, so a + subprocess test of it would need frames and PSDs. Exec'ing the real source text + is the next best thing: it is the code that ships, not a paraphrase, so an edit + that reconnects the prior normalization to the narrowed range fails here. + """ + with open(_ILE) as f: + src = f.read() + start = src.index('dist_prior_range = (param_limits["distance"][0]') + end = src.index('#\n# Parameter integral sampling strategy', start) + block = src[start:end] + + class _O(object): + distance_marginalization = False + d_prior_redshift = False + internal_reparam_dl_incl = False + pin_distance_to_sim = False + limit_distance = None + o = _O() + o.limit_distance = limit_distance + for k, v in optkw.items(): + setattr(o, k, v) + ns = {'opts': o, 'param_limits': {"distance": (1.0, 10000.0)}, + 'distance_limit_range': distance_limit_range, 'print': lambda *a, **k: None} + exec(compile(block, _ILE, 'exec'), ns) # noqa: S102 + return ns + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_narrowing_block_narrows_sampling_and_keeps_the_prior_range(): + ns = _run_ile_narrowing_block('800,3200') + assert ns['param_limits']["distance"] == (800.0, 3200.0) # SAMPLING narrowed + assert ns['dist_prior_range'] == (1.0, 10000.0) # PRIOR untouched + assert ns['limit_distance_active'] is True + # and the kwargs the driver then builds from those two ranges + kw = distance_sampler_kwargs(mcsampler, ns['param_limits']["distance"], + ns['dist_prior_range']) + assert (kw['left_limit'], kw['right_limit']) == (800.0, 3200.0) + assert float(np.asarray(kw['prior_pdf'](np.array([2000.0])))[0]) == pytest.approx( + 2000.0 ** 2 / (10000.0 ** 3 / 3. - 1.0 ** 3 / 3.)) + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_narrowing_block_is_a_no_op_without_the_option(): + ns = _run_ile_narrowing_block(None) + assert ns['param_limits']["distance"] == (1.0, 10000.0) + assert ns['dist_prior_range'] == (1.0, 10000.0) + assert ns['limit_distance_active'] is False + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +@pytest.mark.parametrize('flag', ['distance_marginalization', 'd_prior_redshift', + 'internal_reparam_dl_incl', 'pin_distance_to_sim']) +def test_ile_narrowing_block_refuses_the_incompatible_modes(flag): + """These four have no d_L sampler to narrow. Refusing beats reinterpreting. + + pin_distance_to_sim was MISSED by the first three and added after an + adversarial audit that enumerated every distance-touching option in the driver + rather than trusting the declared list: it pins distance to the injection value + inside analyze_event, so the box was accepted and silently did nothing. That + is the same class as --distance-marginalization, and the same failure mode this + option's whole design is meant to avoid.""" + with pytest.raises(SystemExit): + _run_ile_narrowing_block('800,3200', **{flag: True}) + # ... and they are NOT refused when the option is absent + _run_ile_narrowing_block(None, **{flag: True}) + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +@pytest.mark.parametrize('bad', ['3200,800', '0.5,3200', '800,20000', 'a,b']) +def test_ile_narrowing_block_exits_loudly_on_a_bad_range(bad): + with pytest.raises(SystemExit): + _run_ile_narrowing_block(bad) + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_advertises_the_option(): + """--help goes through optparse, so this catches an option that was written into + the source but never registered.""" + out = subprocess.run([sys.executable, _ILE, '--help'], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=900).stdout.decode('utf-8', 'replace') + assert '--limit-distance' in out diff --git a/MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py b/MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py index c8e7da8f8..afd62a91b 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py +++ b/MonteCarloMarginalizeCode/Code/test/test_lisa_operational_synthetic.py @@ -1,11 +1,13 @@ #!/usr/bin/env python """Operational smoke test for the LISA likelihood path.""" +import inspect import os import lal import lalsimulation as lalsim import numpy as np +import pytest import RIFT.LISA.lalsimutils_compat as lisa_lalsimutils_compat import RIFT.lalsimutils as lalsimutils @@ -81,7 +83,14 @@ def _evaluate_lisa_lnL(rholms, cross_terms, modes, P, psi=None, inclination=None )[0] -def test_synthetic_lisa_tdi_precompute_and_likelihood(tmp_path): +@pytest.fixture(scope="module") +def lisa_precompute(tmp_path_factory): + """Build the synthetic TDI data and precompute once for the whole module. + + Module-scoped because the build costs several seconds and every test here + consumes exactly the same product. + """ + tmp_path = tmp_path_factory.mktemp("lisa_synthetic") P = _synthetic_lisa_params() modes = [(2, 2)] @@ -128,6 +137,12 @@ def test_synthetic_lisa_tdi_precompute_and_likelihood(tmp_path): ) modes_array = np.array(list(hlms.keys())) + return P, modes_array, rholms, cross_terms + + +def test_synthetic_lisa_tdi_precompute_and_likelihood(lisa_precompute): + P, modes_array, rholms, cross_terms = lisa_precompute + lnL_at_injection = _evaluate_lisa_lnL(rholms, cross_terms, modes_array, P) lnL_offset = _evaluate_lisa_lnL( rholms, cross_terms, modes_array, P, psi=P.psi + 0.8, inclination=P.incl + 0.5 @@ -136,3 +151,45 @@ def test_synthetic_lisa_tdi_precompute_and_likelihood(tmp_path): assert np.isfinite(lnL_at_injection) assert np.isfinite(lnL_offset) assert lnL_at_injection > lnL_offset + + +def test_lisa_likelihood_names_a_host_backend_for_spherical_harmonics( + lisa_precompute, monkeypatch +): + """The LISA likelihood must name numpy when it calls SphericalHarmonicsVectorized. + + `SphericalHarmonicsVectorized` resolves to cupy on any host where cupy merely + imports, while this routine works entirely in host numpy. Omitting `xpy` + here used to raise `TypeError: Unsupported type ` from + `cupy._core._kernel._preprocess_args` on every GPU host, and pass everywhere + else -- so assert on the argument, not on the outcome, and this gate holds on + a cupy-free CI runner too. + """ + P, modes_array, rholms, cross_terms = lisa_precompute + + real = factored_likelihood_LISA.SphericalHarmonicsVectorized + seen = [] + + unset = object() + signature = inspect.signature(real) + + def spy(*args, **kwargs): + # Bind through the real signature: a call site that passes xpy + # POSITIONALLY is equally correct and must not be misreported as a + # failure to name a backend. + bound = signature.bind(*args, **kwargs) + seen.append((bound.arguments["theta"], bound.arguments.get("xpy", unset))) + return real(*args, **kwargs) + + monkeypatch.setattr( + factored_likelihood_LISA, "SphericalHarmonicsVectorized", spy + ) + _evaluate_lisa_lnL(rholms, cross_terms, modes_array, P) + + assert seen, "SphericalHarmonicsVectorized was never called" + for theta, xpy in seen: + assert xpy is np, ( + "LISA likelihood left xpy unset (or non-numpy: %r) for a %s theta; " + "host arrays must be computed with numpy" + % ("unset" if xpy is unset else xpy, type(theta).__name__) + ) diff --git a/MonteCarloMarginalizeCode/Code/test/test_mc_error.py b/MonteCarloMarginalizeCode/Code/test/test_mc_error.py new file mode 100644 index 000000000..3243a8ecc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_mc_error.py @@ -0,0 +1,27 @@ +import math + +import pytest + +from RIFT.misc.mc_error import relative_mc_error + + +def test_relative_mc_error_linear_space(): + assert relative_mc_error(4.0, 0.04) == pytest.approx(0.05) + + +def test_relative_mc_error_log_space(): + log_z = -1.9780425516285014 + sigma = 0.0020773022158222486 + log_variance = math.log(sigma**2) + 2 * log_z + assert relative_mc_error(log_z, log_variance, log_space=True) == pytest.approx(sigma) + + +def test_relative_mc_error_rejects_negative_linear_variance(): + with pytest.raises(ValueError, match="non-negative"): + relative_mc_error(1.0, -1.0) + + +def test_relative_mc_error_propagates_non_finite_inputs(): + assert math.isnan(relative_mc_error(math.nan, 1.0)) + assert relative_mc_error(0.0, -math.inf, log_space=True) == 0.0 + assert relative_mc_error(1.0, math.inf) == math.inf diff --git a/MonteCarloMarginalizeCode/Code/test/test_multiapprox_marginalization.py b/MonteCarloMarginalizeCode/Code/test/test_multiapprox_marginalization.py new file mode 100644 index 000000000..c210eb8e3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_multiapprox_marginalization.py @@ -0,0 +1,709 @@ +"""Cross-model marginalization in the multi-approximant workflow. + +Two kinds of check, deliberately separated: + +* the ARITHMETIC of the combination, against hand-computed values. This is the + science: L_marg(lambda) = sum_m p(m) L_m(lambda), linear in L. +* the SHAPE of the emitted DAG. The workflow merges in the iteration loop (one + shared grid, one marginalized fit) and forks at the terminal stage (per-model + posterior and evidence, recombined by p(m) Z_m). + +Neither validates the inference on data; see +RIFT/misc/DESIGN_multiapprox_marginalization.md. + +The DAG tests build with TWO approximants on purpose. With one, every +cross-model path is vacuously satisfied -- which is how this builder kept a +severed grid handoff, an approximant-blind extrinsic stage, and a submit file +naming a directory that was never created. +""" + +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import numpy as np +import pytest + +CODE = Path(__file__).resolve().parents[1] +BIN = CODE / "bin" +CLEANILE = BIN / "util_CleanILE.py" +COMBINE = BIN / "util_CombineApproximantPosteriors.py" +BUILDER = BIN / "create_event_parameter_pipeline_BasicMultiApproxIteration" +MODEL_RX = r"approx_(.+?)_consolidated" + + +def _env(): + env = dict(os.environ) + env["PYTHONPATH"] = str(CODE) + os.pathsep + env.get("PYTHONPATH", "") + env["PATH"] = str(BIN) + os.pathsep + env.get("PATH", "") + env.setdefault("OMP_NUM_THREADS", "1") + env["GW_SURROGATE"] = "" + return env + + +def _run(args, cwd): + return subprocess.run([sys.executable] + [str(a) for a in args], cwd=str(cwd), + env=_env(), text=True, stdout=subprocess.PIPE, + stderr=subprocess.PIPE) + + +def _row(m1, m2, lnL, sigma=0.01, ntot=1000.0): + """One 13-column ILE row: indx m1 m2 s1x..s2z lnL sigmaOverL ntot neff.""" + return [-1, m1, m2, 0., 0., 0.1, 0., 0., 0.2, lnL, sigma, ntot, 100.] + + +def _composite(path, rows): + np.savetxt(str(path), np.array(rows)) + + +def _lnL_column(stdout): + return [float(line.split()[9]) for line in stdout.strip().splitlines() if line.strip()] + + +# -------------------------------------------------------------------------- +# the arithmetic +# -------------------------------------------------------------------------- + +@pytest.fixture +def two_models(tmp_path): + """One shared intrinsic point; model A has TWO replicas, model B has one. + + This is the case that separates a correct marginalization from flat + pooling: flat pooling weights by replica COUNT, so A gets 2/3 instead of + 1/2. With one replica each the two agree, which is why that arrangement + would not have caught the defect. + """ + _composite(tmp_path / "approx_MODELA_consolidated_0.composite", + [_row(10., 8., 100.0), _row(10., 8., 100.0)]) + _composite(tmp_path / "approx_MODELB_consolidated_0.composite", + [_row(10., 8., 104.0)]) + return tmp_path + + +def test_flat_pooling_uses_replica_counts_as_model_weights(two_models): + """Without --model-group-regex the models are weighted by how many times + each was evaluated. Pinned because it is the wrong answer we are moving + away from, and its size is the reason the change matters.""" + out = _run([CLEANILE] + sorted(str(p) for p in two_models.glob("*.composite")), two_models) + assert out.returncode == 0, out.stderr + flat = _lnL_column(out.stdout)[0] + expected = np.log((2 * np.exp(100.0) + np.exp(104.0)) / 3.0) + assert flat == pytest.approx(expected, abs=1e-9) + + +def test_model_aware_combination_is_marginalization(two_models): + """L_marg = sum_m p(m) L_m, uniform p(m) -- linear in L, not in lnL.""" + out = _run([CLEANILE, "--model-group-regex", MODEL_RX] + + sorted(str(p) for p in two_models.glob("*.composite")), two_models) + assert out.returncode == 0, out.stderr + got = _lnL_column(out.stdout)[0] + expected = np.log(0.5 * np.exp(100.0) + 0.5 * np.exp(104.0)) + assert got == pytest.approx(expected, abs=1e-9) + + # and it must differ from flat pooling by a real amount, or the flag is + # cosmetic + flat = np.log((2 * np.exp(100.0) + np.exp(104.0)) / 3.0) + assert got - flat > 0.3 + + +def test_geometric_mean_is_not_what_we_compute(two_models): + """Averaging lnL (a logarithmic opinion pool) is NOT model marginalization. + + Guards against someone 'simplifying' the combination to a mean of lnL. + """ + out = _run([CLEANILE, "--model-group-regex", MODEL_RX] + + sorted(str(p) for p in two_models.glob("*.composite")), two_models) + got = _lnL_column(out.stdout)[0] + geometric = 0.5 * 100.0 + 0.5 * 104.0 + assert abs(got - geometric) > 0.2 + + +def test_model_prior_reweights_the_mixture(two_models): + out = _run([CLEANILE, "--model-group-regex", MODEL_RX, + "--model-prior", "MODELA=0.3", "--model-prior", "MODELB=0.7"] + + sorted(str(p) for p in two_models.glob("*.composite")), two_models) + assert out.returncode == 0, out.stderr + got = _lnL_column(out.stdout)[0] + expected = np.log(0.3 * np.exp(100.0) + 0.7 * np.exp(104.0)) + assert got == pytest.approx(expected, abs=1e-9) + + +def test_large_model_separation_is_stable(tmp_path): + """A weak model must not become 0/0 on a stronger model's log scale.""" + _composite(tmp_path / "approx_MODELA_consolidated_0.composite", + [_row(10., 8., 0.0, sigma=0.2)]) + _composite(tmp_path / "approx_MODELB_consolidated_0.composite", + [_row(10., 8., 1000.0, sigma=0.1)]) + files = sorted(str(p) for p in tmp_path.glob("*.composite")) + out = _run([CLEANILE, "--model-group-regex", MODEL_RX] + files, tmp_path) + assert out.returncode == 0, out.stderr + fields = out.stdout.strip().split() + assert fields and all(value.lower() != "nan" for value in fields), out.stdout + # Uniform model prior: log((exp(0) + exp(1000))/2), whose weak term is + # negligible but whose integration uncertainty remains well-defined. + assert float(fields[9]) == pytest.approx(1000.0 - np.log(2.0), abs=1e-9) + assert float(fields[10]) == pytest.approx(0.1, abs=1e-12) + + +def test_zero_prior_loud_model_does_not_set_the_numerical_scale(tmp_path): + """A zero-mass model 1000 nats louder must not underflow the real mixture.""" + _composite(tmp_path / "approx_MODELA_consolidated_0.composite", + [_row(10., 8., 1000.0, sigma=0.1)]) + _composite(tmp_path / "approx_MODELB_consolidated_0.composite", + [_row(10., 8., 0.0, sigma=0.2)]) + files = sorted(str(p) for p in tmp_path.glob("*.composite")) + out = _run([CLEANILE, "--model-group-regex", MODEL_RX, + "--model-prior", "MODELA=0", "--model-prior", "MODELB=1"] + + files, tmp_path) + assert out.returncode == 0, out.stderr + fields = out.stdout.strip().split() + assert float(fields[9]) == pytest.approx(0.0, abs=1e-12) + assert float(fields[10]) == pytest.approx(0.2, abs=1e-12) + assert "nan" not in (out.stdout + out.stderr).lower() + + +def test_partial_model_prior_is_refused(two_models): + """Half-specified weights would silently default the rest to 1.0.""" + out = _run([CLEANILE, "--model-group-regex", MODEL_RX, "--model-prior", "MODELA=0.3"] + + sorted(str(p) for p in two_models.glob("*.composite")), two_models) + assert out.returncode != 0 + assert "missing weights" in out.stderr + + +def test_one_model_reduces_to_the_flat_pool(two_models): + """Enabling the flag must never change a single-model result.""" + only_a = [str(two_models / "approx_MODELA_consolidated_0.composite")] + plain = _run([CLEANILE] + only_a, two_models) + aware = _run([CLEANILE, "--model-group-regex", MODEL_RX] + only_a, two_models) + assert plain.returncode == 0 and aware.returncode == 0, aware.stderr + assert plain.stdout == aware.stdout + + +def test_partial_coverage_is_reported_and_can_be_dropped(tmp_path): + """A point missing one model is marginalized over the subset, so the + estimator changes point to point. That must never be silent.""" + _composite(tmp_path / "approx_MODELA_consolidated_0.composite", + [_row(10., 8., 100.), _row(12., 9., 101.), _row(14., 7., 99.)]) + _composite(tmp_path / "approx_MODELB_consolidated_0.composite", + [_row(10., 8., 104.), _row(14., 7., 98.)]) + files = sorted(str(p) for p in tmp_path.glob("*.composite")) + + warned = _run([CLEANILE, "--model-group-regex", MODEL_RX] + files, tmp_path) + assert warned.returncode == 0, warned.stderr + assert len(_lnL_column(warned.stdout)) == 3 + assert "WARNING" in warned.stderr and "SUBSET" in warned.stderr + + dropped = _run([CLEANILE, "--model-group-regex", MODEL_RX, + "--require-all-models"] + files, tmp_path) + assert dropped.returncode == 0, dropped.stderr + assert len(_lnL_column(dropped.stdout)) == 2 + assert "DROPPED 1" in dropped.stderr + + +def test_unlabelled_input_is_refused(tmp_path): + """An unmatched filename would be pooled as a nameless extra model.""" + _composite(tmp_path / "stray.composite", [_row(10., 8., 100.)]) + out = _run([CLEANILE, "--model-group-regex", MODEL_RX, + str(tmp_path / "stray.composite")], tmp_path) + assert out.returncode != 0 + assert "does not match" in out.stderr + + +# -------------------------------------------------------------------------- +# the final mixture +# -------------------------------------------------------------------------- + +def _posterior(path, mean, n=4000, seed=7): + rng = np.random.default_rng(seed) + dat = np.column_stack([rng.normal(mean, 1.0, n), rng.normal(0.0, 1.0, n)]) + with open(str(path), "w") as f: + f.write("# m1 m2\n") + np.savetxt(f, dat) + + +def _annotation(path, ln_z): + with open(str(path), "w") as f: + f.write("# lnL sigma_lnL neff\n{!r} 0.01 100\n".format(ln_z)) + + +@pytest.fixture +def two_posteriors(tmp_path): + _posterior(tmp_path / "post_A.dat", 10.0, seed=7) + _posterior(tmp_path / "post_B.dat", 20.0, seed=8) + _annotation(tmp_path / "annot_A.dat", 100.0) + _annotation(tmp_path / "annot_B.dat", 102.0) + return tmp_path + + +def _combine(tmp_path, extra=()): + return _run([COMBINE, + "--model", "A:{}/post_A.dat:{}/annot_A.dat".format(tmp_path, tmp_path), + "--model", "B:{}/post_B.dat:{}/annot_B.dat".format(tmp_path, tmp_path), + "--output", "{}/out.dat".format(tmp_path), "--seed", "1"] + list(extra), + tmp_path) + + +def test_mixture_is_weighted_by_evidence(two_posteriors): + out = _combine(two_posteriors) + assert out.returncode == 0, out.stderr + w = np.exp(np.array([100.0, 102.0]) - 102.0) + w = w / w.sum() + combined = np.genfromtxt(str(two_posteriors / "out.dat"), comments="#") + # the mixture mean is the weighted mean of the two component means + assert combined[:, 0].mean() == pytest.approx(w[0] * 10.0 + w[1] * 20.0, abs=0.15) + + +def test_mixture_prior_multiplies_the_evidence(two_posteriors): + """p(m) is a prior; Z_m still multiplies it. Passing --model-prior does + NOT give a prior-weighted mixture, and that is intended.""" + out = _combine(two_posteriors, ["--model-prior", "A=0.9", "--model-prior", "B=0.1"]) + assert out.returncode == 0, out.stderr + w = np.exp(np.array([100.0, 102.0]) - 102.0) * np.array([0.9, 0.1]) + w = w / w.sum() + combined = np.genfromtxt(str(two_posteriors / "out.dat"), comments="#") + assert combined[:, 0].mean() == pytest.approx(w[0] * 10.0 + w[1] * 20.0, abs=0.15) + + +def test_mismatched_columns_are_refused(two_posteriors): + """Mixing posteriors with different columns would produce garbage.""" + with open(str(two_posteriors / "post_B.dat")) as f: + body = f.read().split("\n", 1)[1] + with open(str(two_posteriors / "post_B.dat"), "w") as f: + f.write("# m1 chi_eff\n" + body) + out = _combine(two_posteriors) + assert out.returncode != 0 + assert "different column header" in out.stderr + + +def test_final_mixture_never_manufactures_duplicate_rows(tmp_path): + """Final posterior grids follow PR #180's without-replacement contract.""" + for label, rows in (("A", "1 2\n1 2\n3 4\n"), + ("B", "5 6\n5 6\n7 8\n")): + with open(str(tmp_path / "post_{}.dat".format(label)), "w") as f: + f.write("# m1 m2\n" + rows) + _annotation(tmp_path / "annot_{}.dat".format(label), 0.0) + out = _combine(tmp_path, ["--n-output-samples", "20"]) + assert out.returncode == 0, out.stderr + combined = np.atleast_2d(np.genfromtxt(str(tmp_path / "out.dat"), comments="#")) + assert len(combined) == len(np.unique(combined, axis=0)) + assert len(combined) == 4 + assert "reducing output" in out.stderr + + +# -------------------------------------------------------------------------- +# the emitted DAG +# -------------------------------------------------------------------------- + +def _dag_facts(rundir): + dag = next(Path(rundir).glob("*.dag")) + jobs, macros, parents = {}, {}, {} + for line in dag.read_text().splitlines(): + parts = line.split() + if line.startswith("JOB"): + jobs[parts[1]] = parts[2] + elif line.startswith("VARS"): + macros[parts[1]] = dict(re.findall(r'(\w+)="([^"]*)"', line)) + elif line.startswith("PARENT"): + cut = parts.index("CHILD") + for child in parts[cut + 1:]: + parents.setdefault(child, set()).update(parts[1:cut]) + return jobs, macros, parents + + +@pytest.fixture(scope="module") +def multiapprox_rundir(tmp_path_factory): + """Build a two-approximant DAG from synthetic inputs. + + Self-contained on purpose: this builder has no caller (pseudo_pipe's + --pipeline-builder does not offer it), so the test must supply the + args_*.txt and grid a user would. + """ + pytest.importorskip("RIFT.lalsimutils") + rundir = tmp_path_factory.mktemp("multiapprox") + (rundir / "args_ile.txt").write_text( + "integrate_likelihood_extrinsic_batchmode --fmin-template 20.0 " + "--n-max 100 --n-eff 17 " + "--time-marginalization --distance-marginalization " + "--distance-marginalization-lookup-table distance_lookup.npz " + "--approx placeholder\n") + (rundir / "args_cip_list.txt").write_text( + "1 --no-plots --fit-method rf --parameter mc --parameter delta_mc --n-output-samples 5\n" + "1 --no-plots --fit-method rf --parameter mc --parameter delta_mc " + "--n-output-samples=5 --posterior-unique-draw\n") + # Deliberately omit --always-succeed. pseudo_pipe normally adds it for a + # terminal-extrinsic run, but the builder is also a public/direct entry + # point and must preserve the terminal fork itself. + (rundir / "args_test.txt").write_text("X --method lame --parameter mc\n") + # Plotting ON. It is off by default, which is why an unresolved + # $(macroapprox) sat in the plot job's log paths through a whole review + # cycle: the plot job is model-independent and its node binds only the + # iteration macros, so a model-tagged log path can never resolve. A stage + # that is not built is a stage no assertion can check. + (rundir / "args_plot.txt").write_text("--parameter mc --parameter eta\n") + (rundir / "args_puff.txt").write_text( + "--parameter mc --parameter eta --force-away 0.01\n") + + grid = _run(["-c", + "import RIFT.lalsimutils as u;" + "P=[];\n" + "import numpy as np\n" + "for i in range(4):\n" + " p=u.ChooseWaveformParams(); p.m1=(10+i)*u.lsu_MSUN; p.m2=8*u.lsu_MSUN; P.append(p)\n" + "u.ChooseWaveformParams_array_to_xml(P,'proposed-grid')\n"], rundir) + if grid.returncode: + pytest.skip("cannot build a seed grid: {}".format(grid.stderr[-400:])) + + build = _run([BUILDER, + "--approx", "IMRPhenomXPHM", "--approx", "SEOBNRv4PHM", + "--input-grid", "proposed-grid.xml.gz", + "--ile-exe", str(BIN / "integrate_likelihood_extrinsic_batchmode"), + "--ile-args", str(rundir / "args_ile.txt"), + "--cip-args-list", "args_cip_list.txt", + "--test-args", "args_test.txt", + "--ile-n-events-to-analyze", "2", "--n-samples-per-job", "5", + "--request-memory-CIP", "4096", "--request-memory-ILE", "4096", + "--working-directory", str(rundir), + "--n-iterations", "2", "--n-copies", "1", + "--approx-gwsignal", "SEOBNRv4PHM", + "--puff-args", str(rundir / "args_puff.txt"), + "--puff-cadence", "1", "--puff-max-it", "1", + "--last-iteration-extrinsic", + "--last-iteration-extrinsic-nsamples", "4", + "--last-iteration-extrinsic-samples-per-ile", "3", + "--last-iteration-extrinsic-samples-per-ile-internal", "7", + "--plot-args", str(rundir / "args_plot.txt")], rundir) + if build.returncode: + pytest.fail("builder failed:\n{}".format(build.stdout[-3000:])) + return rundir + + +def test_every_model_reads_one_shared_grid(multiapprox_rundir): + """The loop ILE grid must NOT be per model, or nothing is marginalized.""" + sub = (multiapprox_rundir / "ILE.sub").read_text() + grid = re.search(r"--sim-xml\s+(\S+)", sub) + assert grid, "no --sim-xml in ILE.sub" + assert "$(macroapprox)" not in grid.group(1), grid.group(1) + + unify = (multiapprox_rundir / "unify.sh").read_text() + assert "--model-group-regex" in unify, ( + "unify.sh pools composites without --model-group-regex, so replica " + "counts act as model weights") + + +def test_terminal_extraction_cannot_be_cut_off_by_convergence(multiapprox_rundir): + """A successful convergence test exits 1, so terminal runs must force the + diagnostic-only form instead of gating their per-model terminal fork.""" + test_sub = (multiapprox_rundir / "test.sub").read_text() + assert "--always-succeed" in test_sub + + dag = next(multiapprox_rundir.glob("*.dag")).read_text() + assert "ABORT-DAG-ON" in dag, ( + "direct multi-approximant builds without terminal extraction would " + "otherwise turn successful convergence into a failed DAG") + + +def test_terminal_cip_grid_covers_every_scheduled_intrinsic_event(multiapprox_rundir): + """Iteration-CIP sample/uniqueness settings must not truncate terminal work. + + Interior repeated intrinsic points are legal; the final cross-model + combiner, not CIP_terminal, owns exact output uniqueness. + """ + cip_sub = (multiapprox_rundir / "CIP_terminal.sub").read_text() + assert "--posterior-unique-draw" not in cip_sub + counts = re.findall(r"--n-output-samples(?:=|\s+)(\d+)", cip_sub) + assert counts == ["4"], cip_sub + + jobs, macros, _ = _dag_facts(multiapprox_rundir) + terminal_starts = [int(macros[node]["macroevent"]) + for node, sub in jobs.items() + if sub.endswith("ILE_extr.sub")] + assert terminal_starts + assert max(terminal_starts) + 2 == int(counts[0]) + + +def test_generator_route_is_per_model(multiapprox_rundir): + """One DAG must be able to mix waveform families. + + The generator route is a property of the MODEL, not of the run: the phenom + family has no time-domain mode generator in gwsignal ("generator does not + provide a method to generate time-domain modes"), while SEOBNRv5* exists + only there. A single global --use-gwsignal therefore cannot serve an + EOB-vs-phenom comparison -- which is the comparison this builder exists for. + + Observed before this was per model: with --use-gwsignal applied globally, + every IMRPhenomD ILE job died with "ValueError: Invalid Argument" from + gen_modes and contributed ZERO rows, so util_CleanILE saw one model and the + run silently degraded to single-model with no marginalization at all. + """ + sub = (multiapprox_rundir / "ILE.sub").read_text() + assert "$(macrogwsignal)" in sub or "--use-gwsignal" not in sub, ( + "ILE.sub carries a global --use-gwsignal; the route must be per model") + + jobs, macros, _ = _dag_facts(multiapprox_rundir) + ile = [n for n, s_ in jobs.items() if s_.endswith("ILE.sub")] + assert ile, "no loop ILE nodes" + routed = {macros[n].get("macroapprox"): macros[n].get("macrogwsignal") + for n in ile if "macrogwsignal" in macros.get(n, {})} + if routed: + # whatever the fixture asked for, a model must get one route, not both + for model, route in routed.items(): + assert route is not None, model + + +def test_puff_ile_uses_the_same_per_model_route(multiapprox_rundir): + """Puff nodes differ from ordinary ILE only in the grid they evaluate.""" + ordinary = (multiapprox_rundir / "ILE.sub").read_text() + puff = (multiapprox_rundir / "ILE_puff.sub").read_text() + for token in ("--approx $(macroapprox)", "$(macrogwsignal)"): + assert token in ordinary, ordinary + assert token in puff, puff + assert "overlap-grid-$(macroiteration).xml.gz" in ordinary + assert "puffball-$(macroiteration).xml.gz" in puff + + jobs, macros, _ = _dag_facts(multiapprox_rundir) + puff_nodes = [n for n, sub in jobs.items() if sub.endswith("ILE_puff.sub")] + assert puff_nodes, "fixture did not build the normally-enabled puff lane" + assert {macros[n].get("macroapprox") for n in puff_nodes} == { + "IMRPhenomXPHM", "SEOBNRv4PHM"} + assert all("macrogwsignal" in macros[n] for n in puff_nodes) + + +def test_the_loop_fits_once_per_iteration(multiapprox_rundir): + jobs, macros, parents = _dag_facts(multiapprox_rundir) + models = {macros.get(n, {}).get("macroapprox") for n, s in jobs.items() + if s.endswith("ILE.sub")} + models.discard(None) + assert len(models) == 2, sorted(models) + + for node, sub in jobs.items(): + if sub.startswith("CIP") and not sub.startswith("CIP_terminal"): + assert "macroapprox" not in macros.get(node, {}), ( + "in-loop CIP is per model; it must fit the marginalized net once") + + for node, sub in jobs.items(): + if sub.endswith("unify.sub"): + waited = {macros.get(p, {}).get("macroapprox") for p in parents.get(node, ()) + if jobs.get(p, "").endswith("join.sub")} + waited.discard(None) + assert waited == models, ( + "unify waits on {}, not every model {}".format(sorted(waited), sorted(models))) + + +def test_odd_grid_size_does_not_drop_the_remainder_batch(multiapprox_rundir): + """Five requested points in batches of two require starts 0, 2, and 4.""" + jobs, macros, _ = _dag_facts(multiapprox_rundir) + by_model = {} + for node, submit in jobs.items(): + values = macros.get(node, {}) + if (submit.endswith("ILE.sub") and values.get("macroiteration") == "1"): + by_model.setdefault(values.get("macroapprox"), set()).add( + values.get("macroevent")) + assert by_model + assert all(starts == {"0", "2", "4"} for starts in by_model.values()) + + +def test_the_terminal_stage_forks_and_recombines(multiapprox_rundir): + jobs, macros, parents = _dag_facts(multiapprox_rundir) + models = {macros.get(n, {}).get("macroapprox") for n, s in jobs.items() + if s.endswith("ILE.sub")} + models.discard(None) + + for suffix in ("CIP_terminal.sub", "cat.sub", "ILE_extr.sub"): + got = {macros.get(n, {}).get("macroapprox") for n, s in jobs.items() + if s.endswith(suffix)} + got.discard(None) + assert got == models, "{} covers {}, not {}".format(suffix, sorted(got), sorted(models)) + + combine = [n for n, s in jobs.items() if s.endswith("combine_models.sub")] + assert len(combine) == 1 + cats = {macros.get(p, {}).get("macroapprox") for p in parents.get(combine[0], ()) + if jobs.get(p, "").endswith("cat.sub")} + cats.discard(None) + assert cats == models, "the mixture is built from {}, not {}".format(sorted(cats), sorted(models)) + + +def test_extrinsic_stage_reads_the_grid_the_run_finished_on(multiapprox_rundir): + """The off-by-one this builder shipped: the terminal ILE read + overlap-grid- while the final CIP wrote overlap-grid-.""" + jobs, macros, _ = _dag_facts(multiapprox_rundir) + extrinsic = {macros.get(n, {}).get("macroiteration") for n, s in jobs.items() + if s.endswith("ILE_extr.sub")} + written = {macros.get(n, {}).get("macroiterationnext") for n, s in jobs.items() + if s.endswith("join_grids.sub") or s.startswith("CIP")} + extrinsic.discard(None) + written.discard(None) + assert len(extrinsic) == 1, sorted(extrinsic) + assert next(iter(extrinsic)) == max(written, key=int) + + +def test_terminal_extrinsic_export_is_a_bounded_fair_draw(multiapprox_rundir): + """Never serialize the full terminal importance-sampling cache. + + A nonconverged science-scale point can reach n-max with millions of raw + draws. The terminal ILE must fair-draw before --save-samples writes the + record, and the already-equal-weight result must go straight from convert + to cat rather than through a second weighted resampler. + """ + sub = (multiapprox_rundir / "ILE_extr.sub").read_text() + assert "--save-samples" in sub + assert "--fairdraw-extrinsic-output" in sub + assert "--fairdraw-extrinsic-output-n-max 3" in sub + assert "--resample-time-marginalization" in sub + assert not re.search(r'(?1 model") diff --git a/MonteCarloMarginalizeCode/Code/test/test_noloop_time_interp.py b/MonteCarloMarginalizeCode/Code/test/test_noloop_time_interp.py index 334155d99..0a411a107 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_noloop_time_interp.py +++ b/MonteCarloMarginalizeCode/Code/test/test_noloop_time_interp.py @@ -1,3 +1,17 @@ +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_noloop_time_interp pins the cubic Q_lm window against the nearest-sample window at +# integer offsets, against linear interpolation at midpoints, and against an exact cubic +# polynomial. It was registered by this PR: it matched no job in ci.yml before, so its +# three tests had never run in CI -- the same silent loss this gate exists to stop. +# --------------------------------------------------------------------------------- import os os.environ.setdefault("RIFT_LOWLATENCY", "1") diff --git a/MonteCarloMarginalizeCode/Code/test/test_noloop_time_marg_row_offset.py b/MonteCarloMarginalizeCode/Code/test/test_noloop_time_marg_row_offset.py new file mode 100644 index 000000000..fcd8cb4fe --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_noloop_time_marg_row_offset.py @@ -0,0 +1,309 @@ +#!/usr/bin/env python +# RIFT-CI-GATE: q-window-stencil +# ^ registers this file with .travis/test-q-window-stencil.sh, run by ci.yml's +# q-window-stencil-check job. Membership lives here, in the test file, so that +# adding a test needs no edit to any shared list. Do not reword the line above. +# --------------------------------------------------------------------------------- +# WHY THIS FILE IS IN q-window-stencil-check. Moved verbatim from the comment block +# above that job's hand-maintained file list in .github/workflows/ci.yml; it lives +# here now so that registering a test needs no edit to a shared file. +# +# test_noloop_time_marg_row_offset belongs in THIS job for the same reason: it is +# another core-likelihood choice that fails silently. The time-marginalization +# log-sum-exp offset was taken over the WHOLE batch instead of per extrinsic +# sample, so any sample more than ~745 nats below the loudest one underflowed to +# lnL = -inf where the likelihood is finite -- above rho ~ 40 that is the bulk of +# the prior, and it collapses mcsamplerAV (issue #232). numpy + lal, no GPU, +# ~9 s. Three of its four CPU guards were mutation-checked to FAIL against the +# unpatched line and against a bare axis=-1 (no keepdims); the fourth pins the +# unshifted return_lnLt early return, which the fix deliberately does not touch. +# Its fifth test is a cupy leg and SKIPS here -- these runners have no GPU. It +# was run and mutation-checked by hand on ldas-pcdev11 (cupy 14.1.1, cuda 12.8). +# --------------------------------------------------------------------------------- +"""The time-marginalization offset in the vectorized NoLoop likelihood must be +PER EXTRINSIC SAMPLE, not per batch. + +WHAT IS BEING TESTED, AND AGAINST WHAT +-------------------------------------- +``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`` marginalizes over time by + + L = simps(exp(lnL_t - lnLmax), dx=deltaT, axis=-1) + lnL = lnLmax + log(L) + +on ``lnL_t`` of shape ``(npts_extrinsic, npts_time)``. The expression is exactly +offset-invariant in real arithmetic, so ANY ``lnLmax`` gives the same answer -- +in real arithmetic. In float64 it does not: a single scalar batch maximum shifts +every extrinsic sample by the LOUDEST sample's peak, and any row sitting more +than ~745 nats below it underflows ``exp()`` to 0 across the whole time axis, so +``L = 0`` and ``log(L) = -inf`` at a sample where the likelihood is finite and +perfectly ordinary. With ``lnL ~ rho^2/2`` at the peak and ``lnL ~ 0`` for a +typical prior draw, that fires above ``rho ~ 40`` and takes out the BULK of the +prior, not a tail -- which is what collapses ``mcsamplerAV`` on loud events. +See oshaughnessy-junior/research-projects-RIT#232 for the real-data measurement. + +The reference here is NOT a reimplementation of the likelihood. Each test asks +the SHIPPED function for its own verbatim ``lnL_t`` (``return_lnLt=True``, a path +that does no shifting at all) and integrates that with the SAME Simpson rule the +function uses, row by row, each row offset by its own maximum. The only thing +that differs between the reference and the code under test is the offset -- which +is the whole subject of the test. + +``keepdims=True`` is load-bearing and is guarded separately. With a bare +``axis=-1`` the ``(n,)`` maximum broadcasts along the TIME axis instead of the +sample axis. That RAISES when ``npts_extrinsic != npts_time`` -- and is silently +wrong when they are equal, which is why one case below is deliberately square. +(The shipped path always builds a 2-D ``lnL_t``: the function requires +array-valued extrinsic parameters, so a 1-D ``lnL_t`` does not arise here.) +""" +from __future__ import print_function, division + +import os + +os.environ.setdefault("RIFT_LOWLATENCY", "1") + +import numpy as np +import pytest +from scipy import integrate + +import lal +import RIFT.lalsimutils as lsu +from RIFT.likelihood import factored_likelihood as fl + +# lal / lalsimutils are imported at module scope on purpose, NOT via importorskip: +# lalsuite is in requirements.txt and both CI jobs that run this file install it, so a +# missing lal here is a broken job, not an unsupported platform -- and an importorskip +# would turn that into a green skip. + +simpson = getattr(integrate, 'simpson', None) or integrate.simps + +SRATE = 4096.0 +DELTAT = 1.0 / SRATE +NPTS = 614 # marginalization_time_grid(0.075, 1/4096) +N_BUFFER = 4096 +UNDERFLOW_NATS = 745.0 # log(smallest positive float64 normal), roughly + + +def _kappa_buffer(): + """A band-limited, periodic-on-its-own-length kappa(t) buffer. + + Periodic so that whatever integer window the code gathers is a genuine + segment of it; the test does not need to predict ``ifirst``. + """ + ts = np.arange(N_BUFFER) * DELTAT + ms = np.arange(1, 400) + T = N_BUFFER * DELTAT + c = np.exp(-2j * np.pi * ms * (N_BUFFER // 2) * DELTAT / T) / (1.0 + (ms / 120.0) ** 2) + return np.exp(2j * np.pi * np.outer(ts, ms) / T) @ c + + +def _inputs(dists_Mpc): + """Minimal inputs that drive the SHIPPED NoLoop function on the numpy backend. + + One detector, one (l,m) pair and zero U/V cross terms, so the self-term + ``rho_sq`` vanishes and ``lnL_t`` is just the response-scaled ``Re kappa(t)`` + times ``distMpcRef/dist``. Distance is therefore a clean per-row amplitude + knob: it sets each extrinsic sample's peak ``lnL`` independently, which is + exactly the axis this test needs to separate. + """ + dists_Mpc = np.asarray(dists_Mpc, dtype=float) + n = dists_Mpc.size + P = lsu.ChooseWaveformParams() + P.deltaT = DELTAT + P.tref = 1000000000.0 + for name in ('phi', 'theta', 'phiref', 'incl', 'psi'): + setattr(P, name, np.zeros(n)) + P.dist = dists_Mpc * 1e6 * lal.PC_SI + det = 'H1' + # The window sits well inside the buffer: the epoch offset sets ifirst, and a + # window running off the front would be zero-extended rather than gathered. + return (P, {det: np.asarray(_kappa_buffer(), dtype=complex)[None, :]}, + {det: np.array([[2, 2]])}, + {det: np.zeros((1, 1), dtype=complex)}, + {det: P.tref - 0.5}) + + +def _shipped(tvals, args, **kw): + P, rholms, lookupNK, ct, epochs = args + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, Lmax=2, xpy=np, **kw) + + +def _per_row_reference(lnL_t, xp=np, simps_fn=None): + """log int exp(lnL_t) dt, row by row, each row offset by its OWN maximum. + + Same closed domain, same grid and the same Simpson rule as the code; only the + offset differs, which is the whole subject of the test. Done one row at a time + so no other row can reach this one's value. + + ``xp``/``simps_fn`` select WHICH backend and WHICH Simpson rule, because the two + backends genuinely do not share one: on numpy the code integrates with scipy, on + cupy with ``optimized_gpu_tools.simps`` -- an old scipy with ``even='avg'`` against + modern scipy's Cartwright correction -- and for EVEN ``npts`` (production is 614) + they differ, by an amount that depends on how sharply peaked the row is + (issue #204). So each backend is compared against its own rule. + """ + simps_fn = simpson if simps_fn is None else simps_fn + out = np.empty(lnL_t.shape[0], dtype=float) + for i in range(lnL_t.shape[0]): + row = lnL_t[i] + m = xp.max(row) + out[i] = float(m) + float(xp.log(simps_fn(xp.exp(row - m), dx=DELTAT))) + return out + + +def tvals_grid(): + grid = fl.marginalization_time_grid(0.075, DELTAT) + assert len(grid) == NPTS + return grid + + +@pytest.fixture(scope='module') +def tvals(): + return tvals_grid() + + +def test_quiet_sample_stays_finite_beside_a_loud_one(tvals): + """The regression, in the shape production actually runs. + + Two extrinsic samples whose peak ``lnL`` differ by far more than the float64 + underflow budget. With a batch-wide offset the quiet row returns ``-inf``; + with a per-row offset it returns its own, finite, correct value. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + args = _inputs(dists) + + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + assert lnL_t.shape == (2, NPTS) + peaks = lnL_t.max(axis=-1) + # The premise of the test: the rows really are separated by more than the + # underflow budget, so a batch offset MUST kill the quiet one. If the + # harness ever stops producing that separation this assert says so, instead + # of the test passing vacuously. + assert peaks[0] - peaks[1] > 2.0 * UNDERFLOW_NATS, peaks + + lnL = np.asarray(_shipped(tvals, args)) + assert lnL.shape == (2,) + assert np.all(np.isfinite(lnL)), lnL + np.testing.assert_allclose(lnL, _per_row_reference(lnL_t), rtol=0, atol=1e-9) + + +def test_offset_is_per_row_even_when_the_batch_is_square(tvals): + """``keepdims=True``, guarded where its absence is SILENT. + + ``npts_extrinsic == npts_time`` on purpose: that is the one shape in which a + bare ``axis=-1`` maximum broadcasts along the wrong axis without raising. + All rows but the first are identical by construction, so they must return + identical values -- an offset that leaks across the sample axis does not. + """ + dists = np.full(NPTS, fl.distMpcRef * 1.0) + dists[0] = fl.distMpcRef / 80.0 + args = _inputs(dists) + + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + assert lnL_t.shape == (NPTS, NPTS) # square, deliberately + peaks = lnL_t.max(axis=-1) + assert peaks[0] - peaks[1] > 2.0 * UNDERFLOW_NATS, peaks[:2] + + lnL = np.asarray(_shipped(tvals, args)) + assert lnL.shape == (NPTS,) + assert np.all(np.isfinite(lnL)), lnL[~np.isfinite(lnL)] + # Identical inputs -> identical outputs, whatever else is in the batch. + assert np.all(lnL[1:] == lnL[1]), np.unique(lnL[1:]).size + np.testing.assert_allclose(lnL, _per_row_reference(lnL_t), rtol=0, atol=1e-9) + + +def test_onset_is_the_underflow_budget_not_a_general_offset_error(tvals): + """Below the underflow budget the two offsets agree; above it they cannot. + + This pins the MECHANISM rather than just the symptom. A batch offset is + harmless while every row is within ~745 nats of the batch peak -- the + expression is offset-invariant apart from rounding -- so a separation just + under the budget must still come out finite and correct, and a separation + well over it must be the only thing that breaks. A "fix" that changed the + integral itself, rather than only its offset, would fail the first half. + """ + # MEASURE the unit-distance peak rather than hardcoding it, so the harness + # stays self-calibrating if the buffer or the response factor is ever retuned. + unit_peak = float(np.asarray( + _shipped(tvals, _inputs([fl.distMpcRef]), return_lnLt=True)).max()) + for gap_nats in (400.0, 700.0): + # peak lnL scales as 1/dist, so this places row 0 gap_nats above row 1. + dists = np.array([fl.distMpcRef * unit_peak / (gap_nats + unit_peak), + fl.distMpcRef * 1.0]) + args = _inputs(dists) + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + peaks = lnL_t.max(axis=-1) + gap = float(peaks[0] - peaks[1]) + assert 0.9 * gap_nats < gap < UNDERFLOW_NATS, gap + lnL = np.asarray(_shipped(tvals, args)) + assert np.all(np.isfinite(lnL)), (gap, lnL) + np.testing.assert_allclose(lnL, _per_row_reference(lnL_t), rtol=0, atol=1e-9) + + +def test_return_lnLt_is_still_verbatim_and_unshifted(tvals): + """The early ``return_lnLt`` return must keep returning UNSHIFTED values. + + Downstream time resampling and the band-limited/peak-local quadratures all + consume this array and apply their own offsets; subtracting anything here + would silently rescale them. Checked against the absolute value the physics + fixes -- ``lnL_t = Re kappa(t) * distMpcRef/dist`` with ``rho_sq = 0`` -- so a + shift of any size, per-row or batch, is visible. + """ + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0]) + lnL_t = np.asarray(_shipped(tvals, _inputs(dists), return_lnLt=True)) + ratio = lnL_t[0] / lnL_t[1] + np.testing.assert_allclose(ratio, 80.0, rtol=1e-10, atol=0) + # And the quiet row is the same array a one-sample call produces. + solo = np.asarray(_shipped(tvals, _inputs(dists[1:]), return_lnLt=True)) + np.testing.assert_allclose(lnL_t[1], solo[0], rtol=0, atol=0) + + +def test_gpu_offset_is_per_row_too(): + """The same guard on the backend the defect was MEASURED on. + + The path is selected by ``opts.gpu``, not by the device -- ``--force-xpy`` keeps it + on with no cupy, and the issue reproduces on plain numpy that way -- so the tests + above are the real regression. This one exists because ``keepdims=True`` and the + ``[..., 0]`` add-back are xpy API calls: untested GPU code is broken code. It skips + without cupy and is run by hand on a GPU node, the way this repo's other GPU legs + are (measured on ldas-pcdev11, cupy 14.1.1, cuda 12.8 container). + + Compared against the GPU's OWN Simpson rule, not against the numpy answer. The two + rules differ for even ``npts`` and the difference is NOT a constant offset -- it is + per row, set by how sharply peaked that row's integrand is: measured here + 0.405, 0.0089 and -4.3e-5 nats for peaks of 1694, 21 and 7 nats (issue #204). That + is a real, separate, already-known discrepancy and it is not this test's subject; a + cross-backend equality assertion here would be asserting #204 is absent. + """ + cupy = pytest.importorskip('cupy') + from RIFT.likelihood import optimized_gpu_tools + import copy + + dists = np.array([fl.distMpcRef / 80.0, fl.distMpcRef * 1.0, fl.distMpcRef * 3.0]) + args = _inputs(dists) + P, rholms, lookupNK, ct, epochs = args + + Pg = copy.deepcopy(P) + for name in ('phi', 'theta', 'phiref', 'incl', 'psi', 'dist'): + setattr(Pg, name, cupy.asarray(np.asarray(getattr(P, name)))) + g_args = (Pg, {k: cupy.asarray(v) for k, v in rholms.items()}, lookupNK, + {k: cupy.asarray(v) for k, v in ct.items()}, epochs) + + def _gpu(**kw): + Q, R, LK, C, E = g_args + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals_grid(), Q, LK, R, C, C, E, Lmax=2, xpy=cupy, **kw) + + lnL_t = _gpu(return_lnLt=True) + peaks = cupy.asnumpy(lnL_t.max(axis=-1)) + assert peaks[0] - peaks[-1] > 2.0 * UNDERFLOW_NATS, peaks + + lnL = cupy.asnumpy(cupy.asarray(_gpu())) + assert np.all(np.isfinite(lnL)), lnL # the regression, on the device + np.testing.assert_allclose( + lnL, _per_row_reference(lnL_t, xp=cupy, simps_fn=optimized_gpu_tools.simps), + rtol=0, atol=1e-9) + + +if __name__ == '__main__': + raise SystemExit(pytest.main([__file__, '-q'])) diff --git a/MonteCarloMarginalizeCode/Code/test/test_rimsky_end_to_end.py b/MonteCarloMarginalizeCode/Code/test/test_rimsky_end_to_end.py new file mode 100644 index 000000000..897b2e28f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_rimsky_end_to_end.py @@ -0,0 +1,291 @@ +"""Submission-level test for Rimsky's real Asimov follow-up hook. + +The test exercises both installed projects, their exchanged YAML, Asimov's +ledger and configuration rendering, and RIFT's real input discovery and +conversion. It replaces only the heavyweight pseudo-pipeline process and the +HTCondor scheduler boundary. +""" + +import configparser +import os +import sys +from contextlib import chdir +from importlib.metadata import version +from pathlib import Path +from unittest.mock import MagicMock, patch + +import h5py +import numpy as np +import pytest +import yaml +from packaging.version import Version + +rimsky = pytest.importorskip("rimsky", reason="Rimsky requires Python >=3.12") + +import asimov +from asimov.ledger import YAMLLedger +from asimov.utils import update +from rimsky.settings import PipelineSettings +from rimsky.sinks.gdb_samples import start_asimov +from rimsky.utils.asimov import add_event, build_and_submit + +import RIFT.asimov.rift as rift_asimov +from RIFT.rimsky.integration import main + + +ASIMOV_CONFIG = """ +[ledger] +location = ledger.yaml +engine = yamlfile + +[project] +name = rimsky-rift-e2e +root = {project} + +[logging] +level = info +directory = logs +location = logs/asimov.log + +[pipelines] +environment = {environment} + +[condor] +user = rimsky-test + +[general] +git_default = . +rundir_default = {project}/working +calibration = test +calibration_directory = C01_offline +webroot = pages/ +logger = file +""" + + +def _initialise_asimov(project): + project.mkdir() + ledger_path = project / "ledger.yaml" + config = configparser.ConfigParser() + config.read_string( + ASIMOV_CONFIG.format(project=project, environment=sys.prefix) + ) + asimov.config = config + asimov.analysis.config = config + asimov.event.config = config + asimov.ledger.config = config + rift_asimov.config = config + YAMLLedger.create(location=ledger_path, name="rimsky-rift-e2e") + ledger = YAMLLedger(location=str(ledger_path)) + update(ledger.data, {"pipelines": {"rift": {}}}) + asimov.current_ledger = ledger + return ledger + + +def test_first_rimsky_result_creates_bootstrapped_rift_production( + tmp_path, monkeypatch +): + assert Version(version("asimov")) >= Version("0.7") + + sid = "S260305df" + source = tmp_path / "rimsky.yaml" + followup = tmp_path / "rift-followup.yaml" + configured_path = tmp_path / "rimsky-rift.yaml" + source.write_text( + yaml.safe_dump( + { + "detectors": ["H1", "L1"], + "channels": {"H1": "STRAIN", "L1": "STRAIN"}, + "output_dir": "online-output", + "event_sink": { + "bilby_pipe_defaults": { + "minimum_frequency": 20, + "maximum_frequency": 1024, + }, + "trigger_dependent_settings": {}, + "prior_defaults": {}, + }, + }, + sort_keys=False, + ), + encoding="utf-8", + ) + + main( + [ + str(source), + str(followup), + "--configured-rimsky", + str(configured_path), + ] + ) + + # Load the generated file through Rimsky itself. These are the defaults + # which make Bilby run first and the sample sink enqueue RIFT afterwards. + settings = PipelineSettings.from_yaml(configured_path) + assert settings.event_sink.bilby_pipe_format == "full-submit" + assert Path(settings.sample_sink.asimov_configuration) == followup + assert Path(settings.asimovdir) == tmp_path / "asimov" + + result = ( + Path(settings.output_dir) + / sid[1:5] + / sid[5:7] + / sid + / "results_page" + / "metafile.hdf5" + ) + result.parent.mkdir(parents=True) + posterior = np.zeros( + 2, + dtype=[ + ("mass_1", "f8"), + ("mass_2", "f8"), + ("chirp_mass", "f8"), + ("luminosity_distance", "f8"), + ("phase", "f8"), + ("iota", "f8"), + ("spin_1x", "f8"), + ("spin_1y", "f8"), + ("spin_1z", "f8"), + ("spin_2x", "f8"), + ("spin_2y", "f8"), + ("spin_2z", "f8"), + ], + ) + posterior["mass_1"] = [35, 36] + posterior["mass_2"] = [30, 29] + posterior["chirp_mass"] = [28, 27] + posterior["luminosity_distance"] = [400, 420] + posterior["iota"] = [0.5, 0.6] + posterior["spin_1z"] = [0.1, 0.2] + posterior["spin_2z"] = [-0.1, -0.2] + with h5py.File(result, "w") as metafile: + analysis = metafile.create_group("bilby-online") + analysis.create_dataset("posterior_samples", data=posterior) + + frames = {} + psds = {} + for detector in settings.detectors: + frame = tmp_path / "{}-RIMSKY-1456739148-4.gwf".format(detector) + frame.touch() + frames[detector] = [str(frame)] + psd = tmp_path / "{}-psd.txt".format(detector) + np.savetxt(psd, [[0, 1e-40], [1, 1e-40], [2, 1e-40]]) + psds[detector] = str(psd) + + ledger = _initialise_asimov(Path(settings.asimovdir)) + event_metadata = { + "name": sid, + "category": "online", + "interferometers": settings.detectors, + "data": { + "segment length": 8, + "channels": settings.channels, + "data files": frames, + }, + "likelihood": {"sample rate": 2048}, + "psds": psds, + "priors": { + "chirp_mass": {"minimum": 10, "maximum": 20}, + "mass_ratio": {"minimum": 0.1, "maximum": 1}, + "luminosity_distance": {"minimum": 10, "maximum": 5000}, + "a_1": {"minimum": 0, "maximum": 0.8}, + "a_2": {"minimum": 0, "maximum": 0.8}, + }, + } + with patch("git.Repo", return_value=MagicMock()): + add_event(Path(settings.asimovdir), event_metadata, ledger=ledger) + start_asimov( + event=sid, + asimovdir=Path(settings.asimovdir), + asimov_configuration=settings.sample_sink.asimov_configuration, + ) + event = ledger.get_event(sid)[0] + productions = [ + item for item in event.analyses if item.name == "rift-online" + ] + assert len(productions) == 1 + production = productions[0] + pipeline = production.pipeline + assert pipeline.__class__.__name__ == "Rift" + assert pipeline._resolve_bootstrap_file() == str(result) + assert production.meta["priors"]["chirp mass"]["maximum"] == 20 + assert production.meta["psds"] == {2048: psds} + caches = pipeline._prepare_frame_caches() + assert set(caches) == {"H1", "L1"} + assert all(Path(cache).is_file() for cache in caches.values()) + + # Exercise the same input-discovery and template-rendering hook used by + # ``asimov manage build``. Keep the installed RIFT scripts discoverable + # when this test is launched via an explicit virtual-environment Python. + monkeypatch.setenv( + "PATH", "{}:{}".format(Path(sys.executable).parent, os.environ["PATH"]) + ) + project_dir = Path(settings.asimovdir) + with chdir(project_dir), patch("asimov.git.time.sleep"): + pipeline.before_config() + + for detector in settings.detectors: + xml_psd = Path(production.xml_psds[detector]) + assert Path(xml_psd).is_file() + assert set(production.meta["data"]["frame cache"]) == {"H1", "L1"} + + # Give build_dag the repository assets that ``asimov manage build`` stores + # before submission. The coinc file is replaced from the bootstrap below, + # but its initial presence avoids any GraceDB access during this test. + repository_dir = Path(event.repository.directory) + if not repository_dir.is_absolute(): + repository_dir = project_dir / repository_dir + category_dir = repository_dir / production.category + category_dir.mkdir(parents=True, exist_ok=True) + (category_dir / "coinc.xml").write_text("synthetic coinc\n") + + commands = [] + + class SchedulerBoundary: + def __init__(self, command, **kwargs): + commands.append(command) + executable = Path(command[0]).name + if executable == "util_RIFT_pseudo_pipe.py": + rundir = Path(production.rundir) + rundir.mkdir(parents=True, exist_ok=True) + dag = ( + rundir + / "marginalize_intrinsic_parameters_BasicIterationWorkflow.dag" + ) + dag.write_text( + "# synthetic DAG emitted at the external RIFT boundary\n" + ) + self.stdout = b"RIFT DAG prepared" + elif executable == "condor_submit_dag": + self.stdout = b"submitted to cluster 4242." + else: + raise AssertionError("unexpected external command: {}".format(command)) + + def communicate(self): + return self.stdout, None + + # Run Rimsky's real Asimov submission helper. Only the heavyweight + # pseudo-pipeline process and final scheduler process are replaced; PSD + # conversion, config rendering, posterior reading, and bootstrap conversion + # run for real against the synthetic files above. + with chdir(Path(settings.asimovdir)), patch("asimov.git.time.sleep"), patch( + "RIFT.asimov.rift.subprocess.Popen", SchedulerBoundary + ): + build_and_submit(event, production, ledger) + + bootstrap = category_dir / "rift-online_bootstrap.xml.gz" + assert bootstrap.is_file() + assert (category_dir / "coinc.xml").is_file() + configuration = category_dir / "rift-online.ini" + assert configuration.is_file() + rendered = configuration.read_text() + assert "fake-cache" in rendered + assert str(Path(caches["H1"])) in rendered + assert production.status == "running" + assert production.job_id == 4242 + assert [Path(command[0]).name for command in commands] == [ + "util_RIFT_pseudo_pipe.py", + "condor_submit_dag", + ] diff --git a/MonteCarloMarginalizeCode/Code/test/test_rimsky_integration.py b/MonteCarloMarginalizeCode/Code/test/test_rimsky_integration.py new file mode 100644 index 000000000..597f65267 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_rimsky_integration.py @@ -0,0 +1,298 @@ +import copy +import configparser +from pathlib import Path +from types import SimpleNamespace + +import pytest + +# Import Asimov before its RIFT entry point. Importing the entry-point module +# first creates a circular discovery path in Asimov 0.5/0.6. +pytest.importorskip("asimov") + +from RIFT.rimsky import ( + RimskyIntegrationError, + build_analysis, + configure_rimsky, + normalize_event_metadata, + write_analysis, +) + + +def _rimsky_config(tmp_path): + return { + "detectors": ["H1", "L1", "V1"], + "output_dir": "online-output", + "event_sink": { + "bilby_pipe_defaults": { + "minimum_frequency": 18, + "maximum_frequency": {"H1": 1024, "L1": 1024, "V1": 896}, + } + }, + "rift": { + "name": "rift-low-latency", + "waveform": {"approximant": "SEOBNRv5PHM"}, + }, + } + + +def test_build_analysis_targets_rimsky_pesummary_output(tmp_path): + config_path = tmp_path / "configs" / "rimsky.yaml" + analysis = build_analysis(_rimsky_config(tmp_path), config_path=config_path) + + expected = ( + config_path.parent + / "online-output" + / "*" + / "*" + / "{event}" + / "results_page" + / "metafile.hdf5" + ).resolve() + assert analysis["kind"] == "analysis" + assert analysis["pipeline"] == "RIFT" + assert analysis["orchestrator"] == "rimsky" + assert analysis["name"] == "rift-low-latency" + assert analysis["dataset"] == "bilby-online" + assert analysis["scheduler"]["bootstrap file"] == str(expected) + assert analysis["quality"]["minimum frequency"] == { + "H1": 18, + "L1": 18, + "V1": 18, + } + assert analysis["quality"]["maximum frequency"]["V1"] == 896 + assert analysis["waveform"]["approximant"] == "SEOBNRv5PHM" + + +def test_normalize_rimsky_event_priors_is_additive(): + event = { + "name": "S260305df", + "interferometers": ["H1", "L1"], + "likelihood": {"sample rate": 2048}, + "psds": {"H1": "/tmp/H1.txt", "L1": "/tmp/L1.txt"}, + "priors": { + "chirp_mass": {"minimum": 10, "maximum": 20}, + "mass_ratio": {"minimum": 0.1, "maximum": 1}, + "luminosity_distance": { + "minimum": 10, + "maximum": 5000, + "type": "bilby.gw.prior.UniformSourceFrame", + }, + "a_1": {"minimum": 0, "maximum": 0.8}, + "a_2": {"minimum": 0, "maximum": 0.7}, + }, + } + original = copy.deepcopy(event) + normalized = normalize_event_metadata(event) + + assert event == original + assert normalized["priors"]["chirp mass"] == event["priors"]["chirp_mass"] + assert normalized["priors"]["mass ratio"] == event["priors"]["mass_ratio"] + assert normalized["priors"]["luminosity distance"]["maximum"] == 5000 + assert normalized["priors"]["spin 1"] == {"maximum": 0.8} + assert normalized["priors"]["spin 2"] == {"maximum": 0.7} + assert "a_1" in normalized["priors"] + assert normalized["psds"] == {2048: {"H1": "/tmp/H1.txt", "L1": "/tmp/L1.txt"}} + + +def test_normalize_does_not_replace_explicit_rift_prior(): + event = {"priors": {"chirp_mass": {"maximum": 20}, "chirp mass": {"maximum": 30}}} + assert normalize_event_metadata(event)["priors"]["chirp mass"]["maximum"] == 30 + + +def test_rift_pipeline_normalizes_rimsky_metadata_before_templating(): + from RIFT.asimov.rift import Rift + + pipeline = object.__new__(Rift) + pipeline.production = SimpleNamespace( + meta={ + "priors": { + "chirp_mass": {"minimum": 10, "maximum": 20}, + "a_1": {"maximum": 0.8}, + }, + "likelihood": {}, + } + ) + pipeline._create_ledger_entries() + + assert pipeline.production.meta["priors"]["chirp mass"]["maximum"] == 20 + assert pipeline.production.meta["priors"]["spin 1"]["maximum"] == 0.8 + assert pipeline.production.meta["sampler"] == {"cip": {}, "ile": {}} + assert pipeline.production.meta["likelihood"] == { + "assume": {}, + "marginalization": {}, + } + + +def test_rift_pipeline_builds_lal_caches_for_rimsky_frames(tmp_path): + from RIFT.asimov.rift import Rift + + frames = [] + for start in (1456739148, 1456739152): + frame = tmp_path / "S260305df-H1-{}-4.gwf".format(start) + frame.touch() + frames.append(str(frame)) + + work_dir = tmp_path / "work" + pipeline = object.__new__(Rift) + pipeline.production = SimpleNamespace( + name="rift-online", + event=SimpleNamespace(work_dir=str(work_dir)), + meta={ + "orchestrator": "rimsky", + "data": {"data files": {"H1": frames}}, + }, + ) + caches = pipeline._prepare_frame_caches() + + cache = Path(caches["H1"]) + assert cache == work_dir / "H1-rimsky.cache" + lines = cache.read_text().splitlines() + assert lines == [ + "H RIMSKY 1456739148 4 {}".format(Path(frames[0]).as_uri()), + "H RIMSKY 1456739152 4 {}".format(Path(frames[1]).as_uri()), + ] + assert pipeline.production.meta["data"]["frame cache"] == caches + + +def test_frame_cache_generation_is_isolated_to_rimsky(tmp_path): + from RIFT.asimov.rift import Rift + + pipeline = object.__new__(Rift) + pipeline.production = SimpleNamespace( + name="unrelated-analysis", + event=SimpleNamespace(work_dir=str(tmp_path)), + meta={"data": {"data files": {"H1": ["not-a-rimsky-frame"]}}}, + ) + assert pipeline._prepare_frame_caches() == {} + assert "frame cache" not in pipeline.production.meta["data"] + + +def test_rift_template_passes_generated_frame_caches(): + template = ( + Path(__file__).resolve().parents[1] / "RIFT" / "asimov" / "rift.ini" + ).read_text() + assert "fake-cache" in template + assert "data['frame cache'][ifo]" in template + + +def test_generated_rimsky_analysis_renders_rift_template(tmp_path): + liquid = pytest.importorskip("liquid") + analysis = build_analysis( + _rimsky_config(tmp_path), config_path=tmp_path / "rimsky.yaml" + ) + event = { + "engine": "RIFT", + "interferometers": ["H1", "L1", "V1"], + "data": { + "segment length": 8, + "channels": {ifo: "{}:STRAIN".format(ifo) for ifo in ("H1", "L1", "V1")}, + "frame types": {ifo: "gwf" for ifo in ("H1", "L1", "V1")}, + "frame cache": { + ifo: "/tmp/{}-rimsky.cache".format(ifo) for ifo in ("H1", "L1", "V1") + }, + }, + "likelihood": {"sample rate": 2048}, + "priors": { + "chirp_mass": {"minimum": 10, "maximum": 20}, + "mass_ratio": {"minimum": 0.1, "maximum": 1}, + "luminosity_distance": { + "minimum": 10, + "maximum": 5000, + "type": "bilby.gw.prior.UniformSourceFrame", + }, + "a_1": {"maximum": 0.8}, + "a_2": {"maximum": 0.7}, + }, + } + for key, value in analysis.items(): + if isinstance(value, dict) and isinstance(event.get(key), dict): + event[key].update(copy.deepcopy(value)) + else: + event[key] = copy.deepcopy(value) + meta = normalize_event_metadata(event) + + production = SimpleNamespace( + name=meta["name"], + meta=meta, + category="C01_offline", + event=SimpleNamespace(name="S260305df", repository=None), + xml_psds={ + ifo: "/tmp/{}-psd.xml.gz".format(ifo) for ifo in meta["interferometers"] + }, + ) + context = { + "production": production, + "config": { + "general": {"webroot": "/tmp/rift-web"}, + "pipelines": {"environment": "/opt/igwn"}, + "condor": {"user": "riftci"}, + }, + } + template_text = ( + Path(__file__).resolve().parents[1] / "RIFT" / "asimov" / "rift.ini" + ).read_text() + if hasattr(liquid, "Environment"): + rendered = liquid.Environment().from_string(template_text).render(**context) + elif hasattr(liquid, "Liquid"): + rendered = liquid.Liquid(template_text, from_file=False).render(**context) + else: + rendered = liquid.Template(template_text).render(**context) + + parser = configparser.RawConfigParser() + parser.read_string(rendered) + assert parser.get("engine", "chirpmass-min") == "10" + assert parser.get("engine", "comp-max") == "1000" + assert parser.get("engine", "a_spin1-max") == "0.8" + assert '"V1":"/tmp/V1-rimsky.cache"' in parser.get("lalinference", "fake-cache") + + +def test_write_analysis_round_trips(tmp_path): + yaml = pytest.importorskip("yaml") + analysis = build_analysis( + _rimsky_config(tmp_path), config_path=tmp_path / "rimsky.yaml" + ) + destination = write_analysis(analysis, tmp_path / "rift-followup.yaml") + assert yaml.safe_load(destination.read_text()) == analysis + + +def test_configure_rimsky_defaults_to_submitted_bilby_then_rift(tmp_path): + source = tmp_path / "configs" / "rimsky.yaml" + followup = tmp_path / "generated" / "rift-followup.yaml" + configured = configure_rimsky( + _rimsky_config(tmp_path), followup, config_path=source + ) + + assert configured["output_dir"] == str( + (source.parent / "online-output").resolve() + ) + assert configured["asimovdir"] == str((source.parent / "asimov").resolve()) + assert configured["event_sink"]["bilby_pipe_format"] == "full-submit" + assert configured["sample_sink"]["asimov_configuration"] == str( + followup.resolve() + ) + + +def test_configure_rimsky_preserves_explicit_run_mode_and_asimovdir(tmp_path): + config = _rimsky_config(tmp_path) + config["asimovdir"] = "project" + config["event_sink"]["bilby_pipe_format"] = "full-local" + configured = configure_rimsky( + config, tmp_path / "followup.yaml", config_path=tmp_path / "rimsky.yaml" + ) + assert configured["event_sink"]["bilby_pipe_format"] == "full-local" + assert configured["asimovdir"] == str((tmp_path / "project").resolve()) + + +@pytest.mark.parametrize("detectors", [[], {"H1": "bad"}, ["H1", 2]]) +def test_invalid_detectors_fail_early(tmp_path, detectors): + config = _rimsky_config(tmp_path) + config["detectors"] = detectors + with pytest.raises(RimskyIntegrationError, match="detector"): + build_analysis(config, config_path=tmp_path / "rimsky.yaml") + + +def test_frequency_mapping_must_cover_all_detectors(tmp_path): + config = _rimsky_config(tmp_path) + config["event_sink"]["bilby_pipe_defaults"]["maximum_frequency"] = {"H1": 1024} + with pytest.raises(RimskyIntegrationError, match="L1, V1"): + build_analysis(config, config_path=tmp_path / "rimsky.yaml") diff --git a/MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py b/MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py new file mode 100644 index 000000000..915470039 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_spherical_harmonics_backend.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python +"""Backend dispatch of `SphericalHarmonicsVectorized`. + +The helper used to bind `xpy=xpy_default` in its signature, and `xpy_default` is +cupy on any host where cupy merely *imports*. Callers working in host numpy that +did not name a backend therefore dispatched a GPU kernel onto host arrays and +raised `TypeError: Unsupported type ` -- on GPU hosts only, +so the cupy-free CI runners never saw it. + +These tests fake the GPU backend rather than requiring one, so they hold on any +runner. The real-cupy leg is skipped where cupy is absent. +""" + +import inspect + +import numpy as np +import pytest + +from RIFT.likelihood import SphericalHarmonics_gpu as sh + + +MODES = np.array([[2, -2], [2, 0], [2, 2]]) + + +class _FakeGPUArray(object): + """Stands in for cupy.ndarray: host arrays are never instances of it.""" + + +def _reject_host(*args, **kwargs): + raise TypeError( + "Unsupported type : a GPU kernel was handed a host array" + ) + + +class _FakeGPU(object): + """Minimal stand-in for the cupy module. + + Every array entry point rejects host input the way cupy's ufuncs do, so a + test that accidentally dispatches here fails with the production symptom + rather than an incidental AttributeError. + """ + + ndarray = _FakeGPUArray + empty = staticmethod(_reject_host) + cos = staticmethod(_reject_host) + sin = staticmethod(_reject_host) + square = staticmethod(_reject_host) + + +def test_default_xpy_is_not_bound_at_import(): + """The default must be resolved per call, not frozen to a module at def time. + + A live module here is the original defect: monkeypatching `xpy_default` + afterwards cannot dislodge it, and every unadorned caller inherits cupy. + """ + default = inspect.signature(sh.SphericalHarmonicsVectorized).parameters["xpy"].default + assert default is None + + +def test_infer_xpy_prefers_numpy_for_host_arrays_even_with_a_gpu_present(monkeypatch): + monkeypatch.setattr(sh, "cupy_here", True) + monkeypatch.setattr(sh, "cupy", _FakeGPU) + assert sh._infer_xpy(np.linspace(0.1, 3.0, 5)) is np + + +def test_infer_xpy_selects_the_gpu_backend_for_device_arrays(monkeypatch): + monkeypatch.setattr(sh, "cupy_here", True) + monkeypatch.setattr(sh, "cupy", _FakeGPU) + assert sh._infer_xpy(_FakeGPUArray()) is _FakeGPU + + +def test_infer_xpy_is_numpy_when_no_gpu_backend_is_installed(monkeypatch): + monkeypatch.setattr(sh, "cupy_here", False) + assert sh._infer_xpy(np.linspace(0.1, 3.0, 5)) is np + + +def test_public_entrypoint_resolves_a_host_backend_on_a_simulated_gpu_host(monkeypatch): + """Drive SphericalHarmonicsVectorized itself, with `xpy_default` faked to the GPU. + + The `_infer_xpy` tests above exercise the helper in isolation; this one pins + that SphericalHarmonicsVectorized actually CONSULTS it. Faking `xpy_default` + is the whole point -- that module global is where the original bug came from, + so a body that reads it instead of inferring (`xpy = xpy_default` when xpy is + None, a plausible "simplification") is caught only here. Without this test + that mutation leaves the entire suite green on a cupy-free runner. + """ + monkeypatch.setattr(sh, "cupy_here", True) + monkeypatch.setattr(sh, "cupy", _FakeGPU) + monkeypatch.setattr(sh, "xpy_default", _FakeGPU) + + theta = np.linspace(0.1, np.pi - 0.1, 7) + phi = np.linspace(0.0, 2.0 * np.pi, 7) + + inferred = sh.SphericalHarmonicsVectorized(MODES, theta, phi, l_max=2) + assert isinstance(inferred, np.ndarray) + + # An explicit backend must still win over the faked default. + explicit = sh.SphericalHarmonicsVectorized(MODES, theta, phi, xpy=np, l_max=2) + np.testing.assert_array_equal(inferred, explicit) + + +def test_host_arrays_give_the_same_answer_with_and_without_an_explicit_backend(): + """Unadorned call on host arrays must work, and agree with `xpy=np`. + + This is the assertion that fails on a cupy-capable host if the inference is + removed; it is a tautology on a cupy-free one, which is why it is paired with + the fake-GPU tests above. + """ + theta = np.linspace(0.1, np.pi - 0.1, 7) + phi = np.linspace(0.0, 2.0 * np.pi, 7) + + inferred = sh.SphericalHarmonicsVectorized(MODES, theta, phi, l_max=2) + explicit = sh.SphericalHarmonicsVectorized(MODES, theta, phi, xpy=np, l_max=2) + + assert isinstance(inferred, np.ndarray) + np.testing.assert_array_equal(inferred, explicit) + + +@pytest.mark.skipif(not sh.cupy_here, reason="requires a working cupy install") +def test_device_arrays_stay_on_the_device_without_an_explicit_backend(): + import cupy + + theta = cupy.linspace(0.1, np.pi - 0.1, 7) + phi = cupy.linspace(0.0, 2.0 * np.pi, 7) + + out = sh.SphericalHarmonicsVectorized(MODES, theta, phi, l_max=2) + assert isinstance(out, cupy.ndarray) + + host = sh.SphericalHarmonicsVectorized( + MODES, cupy.asnumpy(theta), cupy.asnumpy(phi), l_max=2 + ) + np.testing.assert_allclose(cupy.asnumpy(out), host, rtol=0, atol=1e-14) diff --git a/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py b/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py index 8bf495a3d..3f9299c61 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py +++ b/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py @@ -165,7 +165,7 @@ def test_source_matches_reference_implementation(): source = handle.read() block = re.search( - r"if opts\.srate_resample_time_marginalization and .*?lnLt_norm = " + r"(?:if|elif) opts\.srate_resample_time_marginalization and .*?lnLt_norm = " r"scipy\.special\.logsumexp\(lnLt,axis=-1\)", source, re.S, diff --git a/MonteCarloMarginalizeCode/Code/test/test_teobresums_compat.py b/MonteCarloMarginalizeCode/Code/test/test_teobresums_compat.py new file mode 100644 index 000000000..4106e5490 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_teobresums_compat.py @@ -0,0 +1,211 @@ +from types import SimpleNamespace +from pathlib import Path + +import pytest + +from RIFT.physics import teobresums_compat as compat + + +class _BasicModule: + __file__ = "/tmp/basic/EOBRun_module.so" + + @staticmethod + def EOBRunPy(parameters): + return parameters + + +class _DaliModule(_BasicModule): + eob_dyn_j0_py = object() + eob_ham_s_py = object() + eob_metric_A5PNlog_py = object() + + +def test_auto_profile_has_a_stable_default_for_unknown_extensions(monkeypatch): + monkeypatch.delenv("RIFT_TEOBRESUMS_PROFILE", raising=False) + + assert compat.detect_profile(_BasicModule) == "default" + assert compat.detect_profile(_DaliModule) == "dali" + + +def test_explicit_profile_override_and_typo_handling(monkeypatch): + monkeypatch.setenv("RIFT_TEOBRESUMS_PROFILE", "legacy") + assert compat.detect_profile(_DaliModule) == "legacy" + + monkeypatch.setenv("RIFT_TEOBRESUMS_PROFILE", "not-a-profile") + with pytest.raises(compat.TEOBResumSCompatibilityError): + compat.detect_profile(_DaliModule) + + +def test_legacy_integer_values_are_normalized_to_semantics(): + normalized = compat.normalize_parameters( + { + "arg_out": 1, + "nqc": 2, + "nqc_coefs_hlm": 0, + "nqc_coefs_flx": 0, + "use_geometric_units": 0, + "interp_uniform_grid": 1, + "output_hpc": 0, + "M": 60, + } + ) + + assert normalized == { + "arg_out": "yes", + "nqc": "no", + "nqc_coefs_hlm": "none", + "nqc_coefs_flx": "none", + "use_geometric_units": "no", + "interp_uniform_grid": "yes", + "output_hpc": "no", + "M": 60, + } + + with pytest.raises(compat.TEOBResumSCompatibilityError): + compat.normalize_parameters({"arg_out": 7}) + + +def test_run_probes_before_native_call_and_passes_normalized_values(monkeypatch): + compat._PROBED_SCHEMAS.clear() + calls = [] + + class RecordingModule(_BasicModule): + @staticmethod + def EOBRunPy(parameters): + calls.append(parameters) + return (None, None, None, {}) + + monkeypatch.setattr( + compat.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace(returncode=0, stdout="", stderr=""), + ) + + result = compat.run(RecordingModule, {"arg_out": 1, "nqc": 2}) + + assert len(result) == 4 + assert calls == [{"arg_out": "yes", "nqc": "no"}] + + +def test_failed_probe_prevents_native_call(monkeypatch): + compat._PROBED_SCHEMAS.clear() + calls = [] + + class RecordingModule(_BasicModule): + @staticmethod + def EOBRunPy(parameters): + calls.append(parameters) + + monkeypatch.setattr( + compat.subprocess, + "run", + lambda *args, **kwargs: SimpleNamespace( + returncode=-11, stdout="", stderr="native extension terminated" + ), + ) + + with pytest.raises(compat.TEOBResumSCompatibilityError, match="return code -11"): + compat.run(RecordingModule, {"arg_out": "yes"}) + assert calls == [] + + +def _transverse_parameters(spin1x=0.0, spin1y=0.0, spin2x=0.0, spin2y=0.0): + return { + "spin1x": spin1x, + "spin1y": spin1y, + "spin2x": spin2x, + "spin2y": spin2y, + "untouched": object(), + } + + +def test_resums_initial_grid_seed_has_margin_above_native_boundary(): + seed_min, seed_max = compat.initial_transverse_spin_range("TEOBResumSDALI") + + assert (seed_min, seed_max) == (1e-3, 3e-3) + assert seed_min >= 10 * compat.DALI_TRANSVERSE_SPIN_THRESHOLD + assert compat.initial_transverse_spin_range("SEOBNRv5PHM") == (1e-5, 3e-5) + + +@pytest.mark.parametrize("approximant", ["TEOBResumS", "TEOBResumSDALI"]) +def test_gwsignal_guard_zeros_only_resums_native_aligned_interval(approximant): + original = _transverse_parameters(spin1x=6e-5, spin2y=4e-5) + + guarded = compat.guard_gwsignal_transverse_spins(original, approximant) + + assert guarded is not original + assert [guarded[key] for key in ("spin1x", "spin1y", "spin2x", "spin2y")] == [ + 0.0, + 0.0, + 0.0, + 0.0, + ] + assert original["spin1x"] == 6e-5 # caller-owned parameters are not mutated + assert guarded["untouched"] is original["untouched"] + + +def test_gwsignal_guard_preserves_aligned_genuinely_precessing_and_other_models(): + aligned = _transverse_parameters() + precessing = _transverse_parameters(spin1x=1.000001e-4) + other_model = _transverse_parameters(spin1x=1e-5) + + assert compat.guard_gwsignal_transverse_spins(aligned, "TEOBResumSDALI") is aligned + assert compat.guard_gwsignal_transverse_spins(precessing, "TEOBResumSDALI") is precessing + assert compat.guard_gwsignal_transverse_spins(other_model, "SEOBNRv5PHM") is other_model + + +def test_pipeline_threads_approximant_into_model_specific_grid_seed(): + code_root = Path(__file__).parents[1] + helper_source = (code_root / "bin" / "helper_LDG_Events.py").read_text() + pipe_source = (code_root / "bin" / "util_RIFT_pseudo_pipe.py").read_text() + + assert "initial_transverse_spin_range(" in helper_source + assert "--internal-initial-grid-approximant {}" in pipe_source + + +def test_precession_predicate_uses_summed_transverse_magnitude(): + threshold = compat.DALI_TRANSVERSE_SPIN_THRESHOLD + + # a component that is nonzero but summed-small is still aligned for DALI + assert not compat.is_precessing_for_resums(0.0, 1e-5, 2e-5, 3e-5) + assert not compat.is_precessing_for_resums(0.0, 0.0, 0.0, 0.0) + # two components that are individually below threshold can sum above it + assert compat.is_precessing_for_resums(0.9 * threshold, 0.0, 0.9 * threshold, 0.0) + assert compat.is_precessing_for_resums(0.0, 0.0, 0.0, 1.000001 * threshold) + + +def test_direct_teobresums_calls_share_the_backend_precession_predicate(): + source = (Path(__file__).parents[1] / "RIFT" / "lalsimutils.py").read_text() + + # four direct EOBRunPy parameter blocks, plus the complementary decision + # about conjugating positive-m modes + assert source.count("is_precessing_for_resums(") == 5 + assert "P.s1y!=0.0" not in source + + +def test_retained_hyperbolic_peaks_are_kept_in_chronological_order(): + source = (Path(__file__).parents[1] / "RIFT" / "lalsimutils.py").read_text() + + # indices_to_keep is a set: filtered_peaks[-1] is only the last peak if sorted + assert "peaks[list(indices_to_keep)]" not in source + assert source.count("peaks[sorted(indices_to_keep)]") == 4 + + +def test_inference2ile_restores_posterior_columns_it_can_read(): + source = ( + Path(__file__).parents[1] / "bin" / "convert_output_format_inference2ile" + ).read_text() + + # eccentricity must be restored from the input columns, not from the + # freshly initialized (always zero) parameter value + assert 'if "eccentricity" in samples_in.dtype.names:' in source + assert "P.eccentricity < 1e-5" not in source + assert 'if "a6c" in samples_in.dtype.names:' in source + + +def test_every_rift_gwsignal_generator_path_uses_transverse_spin_guard(): + code_root = Path(__file__).parents[1] + source = (code_root / "RIFT" / "physics" / "GWSignal.py").read_text() + + assert source.count("gwsignal_get_waveform_generator(") == 3 + assert source.count("guard_gwsignal_transverse_spins(") == 3 diff --git a/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_peak_local.py b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_peak_local.py new file mode 100644 index 000000000..19dfa4f41 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_peak_local.py @@ -0,0 +1,2526 @@ +"""Tests for the peak-local time-marginalization quadrature. + +Companion to ``test_time_marginalization_quadrature.py``, which covers the dense +band-limited path this one is measured against. The fixtures are deliberately the +same construction -- ``kappa`` built as a sum of complex exponentials strictly below +Nyquist, so the continuous function and its integral are known in CLOSED FORM -- and +several tests here compare against that analytic truth rather than against the dense +path, so a shared bias in the two implementations cannot pass. + +Three things this file has to establish that the dense path did not need: + +* the LOCAL EVALUATOR reconstructs the same interpolant the dense upsample does, at + arbitrary points and at every production ``npts`` (the odd-``npts`` spectrum split + is a live trap: exact at the samples, wrong between them); +* MERGING is load-bearing, not tidiness -- the un-merged variant is reproduced here + and shown to be wrong; +* the TRUNCATION is bounded rather than hoped for, including when the enumeration is + sabotaged. +""" + +from __future__ import print_function, division + +import os +import re +import sys + +import numpy as np +import pytest +from scipy import integrate + +from RIFT.likelihood import time_marginalization_quadrature as tmq +from RIFT.likelihood import time_marginalization_peak_local as pl +from RIFT.likelihood import factored_likelihood as fl + +simpson = getattr(integrate, 'simpson', None) or integrate.simps + +SRATE = 4096.0 +DELTAT = 1.0 / SRATE +NPTS = 614 # marginalization_time_grid(0.075, 1/4096) +RHO_SQ = 1000.0 + +#: The production ``npts`` values, plus their odd neighbours and three tiny sizes. +#: ``marginalization_time_grid(0.075, 1/srate)`` is ODD at three of the five +#: production sample rates -- 153 at 1024, 307 at 2048, 2457 at 16384 -- so odd is +#: the common case and not a corner. +NPTS_CASES = [153, 307, 613, 614, 1228, 2457, 8, 9, 3] + + +# ----------------------------------------------------------------- helpers + +def _lnL(kappa_term, rho_sq): + """The production default helper, spelled out so the test does not depend on a + private name. Affine and INCREASING in the kappa term, which is the property + the enumeration relies on.""" + return kappa_term - 0.5 * rho_sq + + +def _lnL_distmarg_like(kappa_term, rho_sq): + """A distance-marginalization-SHAPED callback: nonlinear, still monotone + increasing in the kappa term, and ``-inf`` outside a table-like domain. + + The real distmarg callback is a 2-D table interpolation; what matters for THIS + callback is one property -- monotone but NOT affine in the kappa term -- because + the production default helper is affine and an implementation that quietly + assumed linearity would pass every test that used only it. The ``-inf`` branch is + present for shape but is far enough out that it does not fire on these fixtures; + the ``-inf`` domain edge is exercised separately by ``domain_limited`` and by + ``test_unmeasurable_row_falls_back_and_is_counted``. + """ + x = kappa_term - 0.5 * rho_sq + return np.where(x > -1e6, x + 0.05 * np.log1p(np.abs(x)) * np.sign(x), -np.inf) + + +def _log_trapz(v, dx): + m = v.max() + w = np.full(v.size, dx) + w[0] *= 0.5 + w[-1] *= 0.5 + return m + np.log(np.sum(w * np.exp(v - m))) + + +def _log_simps(v, dx): + m = v.max() + return m + np.log(simpson(np.exp(v - m), dx=dx)) + + +class BandLimited(object): + """kappa(t) = sum_m c_m exp(2 pi i m t / T), every |f| < Nyquist. + + Identical construction to the dense path's suite, so the two are measured on the + same object. ``n_period == NPTS`` gives a window that is exactly periodic; + ``n_period > NPTS`` gives a window cut from a longer signal, the realistic case. + """ + + def __init__(self, amp, peak_sample, n_period=NPTS, m_hi=None, seed=7, + background=0.0, extra_peaks=()): + self.T = n_period * DELTAT + self.j0 = (n_period - NPTS) // 2 + scale = n_period / float(NPTS) + m_hi = int(200 * scale) if m_hi is None else m_hi + assert m_hi < n_period // 2, "would exceed Nyquist" + ms = np.arange(1, m_hi + 1) + env = 1.0 / (1.0 + (ms / (120.0 * scale)) ** 2) + c = np.exp(-2j * np.pi * ms * (self.j0 + peak_sample) * DELTAT / self.T) * env + for where, rel in extra_peaks: + c = c + rel * np.exp( + -2j * np.pi * ms * (self.j0 + where) * DELTAT / self.T) * env + if background: + rng = np.random.default_rng(seed) + c = c + background * ((rng.normal(size=m_hi) + 1j * rng.normal(size=m_hi)) + / (1.0 + (ms / (40.0 * scale)) ** 2)) + self.ms, self.c = ms, amp * c + + def at(self, ts, chunk=40000): + ts = np.asarray(ts, dtype=float) + out = np.empty(ts.size, dtype=complex) + for i in range(0, ts.size, chunk): + t = ts[i:i + chunk] + out[i:i + chunk] = np.exp(2j * np.pi * np.outer(t, self.ms) / self.T) @ self.c + return out + + def samples(self): + return self.at((self.j0 + np.arange(NPTS)) * DELTAT) + + def truth(self, refine, callback=_lnL): + """log int exp(lnL) dt over the SAME closed domain [t_0, t_{NPTS-1}]. + + ``refine`` is REQUIRED and is passed explicitly at every call site, because a + default that under-resolves the peak turns the "analytic truth" into another + under-resolved estimate and the comparison into a coincidence. The rule used + below is ``deltaT/refine <= sigma_t/8``; ``sigma_t/deltaT`` is ~0.10 at + amp 1, 0.047 at 5, 0.017 at 40 and 0.0074 at 200 for this fixture. It also + bounds the cost: this is an O(n_modes * n_points) closed-form evaluation. + """ + n = (NPTS - 1) * refine + 1 + td = self.j0 * DELTAT + np.arange(n) * (DELTAT / refine) + return _log_trapz(callback(self.at(td).real, RHO_SQ), DELTAT / refine) + + +def _peak_local(kappa_row, callback=_lnL, **kw): + k = np.asarray(kappa_row)[None, :] + r = np.full(k.shape, RHO_SQ) + return float(pl.time_marginalize_peak_local(k, r, DELTAT, callback, **kw)[0]) + + +def _bandlimited(kappa_row, callback=_lnL): + k = np.asarray(kappa_row)[None, :] + r = np.full(k.shape, RHO_SQ) + return float(tmq.time_marginalize_bandlimited(k, r, DELTAT, callback)[0]) + + +def _simpson_value(kappa_row, callback=_lnL): + return _log_simps(callback(np.asarray(kappa_row).real, RHO_SQ), DELTAT) + + +def _random_bandlimited(n, seed=3): + """A row whose spectrum fills EVERY bin strictly below Nyquist, both signs. + + Filling every bin is the point: a fixture with an empty top positive bin cannot + see the odd-``npts`` split bug at all. The exact Nyquist bin (even ``n`` only) is + left EMPTY on purpose -- that component aliases onto its own conjugate on the + samples, so no interpolant can recover it and a fixture that fills it would fail + spuriously. It is empty in rholm data anyway. + """ + rng = np.random.default_rng(seed) + n_pos = (n - 1) // 2 + X = np.zeros(n, dtype=complex) + X[0] = rng.normal() + for k in range(1, n_pos + 1): + X[k] = rng.normal() + 1j * rng.normal() + X[n - k] = rng.normal() + 1j * rng.normal() + return np.fft.ifft(X) + + +# ------------------------------------------------------- the local evaluator + +@pytest.mark.parametrize("n", NPTS_CASES) +def test_local_evaluator_matches_the_dense_upsample(n): + """The peak-local path evaluates the interpolant at arbitrary local times; the + dense path evaluates it on a uniform refinement by zero-padded FFT. They must + be THE SAME FUNCTION, or the two quadratures are not comparable and the A/B + below means nothing. + + Parametrised over the real production sizes because the frequency split is where + this breaks: putting the top positive bin at a negative frequency (which a split + at ``n//2`` does for odd ``n``) leaves the reconstruction exact AT the samples + and wrong between them, so a round-trip test cannot see it. + """ + x = _random_bandlimited(n)[None, :] + factor = 4 + dense = tmq.bandlimited_upsample(x, factor)[0] + Xw, fk = pl.bandlimited_spectrum(x) + got = pl.eval_bandlimited_uniform(Xw, fk, np.array([0.0]), + np.array([DELTAT / factor]), n * factor, + n * DELTAT)[0] + scale = np.max(np.abs(dense)) + assert np.max(np.abs(got - dense)) / scale < 1e-11, n + + +@pytest.mark.parametrize("n", NPTS_CASES) +def test_local_evaluator_reproduces_the_coarse_samples(n): + """Necessary but NOT sufficient -- this is exactly the check that stays green + when the spectrum is split at the wrong index -- so it is here as a floor under + the test above, not as a substitute for it.""" + x = _random_bandlimited(n, seed=11)[None, :] + Xw, fk = pl.bandlimited_spectrum(x) + got = pl.eval_bandlimited_uniform(Xw, fk, np.array([0.0]), np.array([DELTAT]), + n, n * DELTAT)[0] + assert np.allclose(got, x[0], atol=1e-11 * np.max(np.abs(x))) + + +def test_local_evaluator_is_accurate_far_along_a_long_grid(): + """The evaluator advances the phase by a recurrence, which drifts as ``m * eps`` + because ``exp(ix)`` is not exactly unit modulus; it is re-anchored periodically. + Deleting the re-anchor has to fail, so evaluate FAR out -- 8x the re-anchor + interval -- where an un-anchored recurrence would have accumulated visibly.""" + n = 2457 + x = _random_bandlimited(n, seed=5)[None, :] + m = 8 * pl._RECURRENCE_REANCHOR + Xw, fk = pl.bandlimited_spectrum(x) + h = DELTAT / 4.0 + got = pl.eval_bandlimited_uniform(Xw, fk, np.array([0.0]), np.array([h]), m, + n * DELTAT)[0] + ref = tmq.bandlimited_upsample(x, 4)[0][:m] + assert np.max(np.abs(got - ref)) / np.max(np.abs(ref)) < 1e-12 + + +def test_local_evaluator_honours_a_per_row_grid(): + """Rows are batched with DIFFERENT interval starts and spacings; a version that + used row 0's grid for every row would still pass a single-row test.""" + n = 307 + x = np.stack([_random_bandlimited(n, seed=s) for s in (1, 2, 3)]) + Xw, fk = pl.bandlimited_spectrum(x) + t0 = np.array([0.0, 13.0 * DELTAT, 101.5 * DELTAT]) + h = np.array([DELTAT / 3, DELTAT / 7, DELTAT / 11]) + got = pl.eval_bandlimited_uniform(Xw, fk, t0, h, 20, n * DELTAT) + for r in range(3): + ts = t0[r] + np.arange(20) * h[r] + want = pl.eval_bandlimited_uniform(Xw[r:r + 1], fk, t0[r:r + 1], h[r:r + 1], + 20, n * DELTAT)[0] + assert np.allclose(got[r], want, atol=1e-12 * np.max(np.abs(want))) + + +# --------------------------------------------- accuracy against analytic truth + +@pytest.mark.parametrize("amp,refine", [(1.0, 256), (5.0, 512)]) +@pytest.mark.parametrize("phase", [0.0, 0.25, 0.5, 0.3125, 0.28125]) +def test_exact_on_a_periodic_window(amp, phase, refine): + """Against a CLOSED-FORM truth, not against the dense path. + + Grid PHASE is swept as well as amplitude: the defect this whole line of work + exists to remove is that the answer depends on where the sample grid happens to + fall relative to the peak, so a fixture pinned to one phase can be exactly wrong + and look exactly right. + + THE PHASES ABOVE ARE NOT ALL MULTIPLES OF 1/PEAK_ENUM_FACTOR, and that is the whole + point. The first version of this suite swept 0.0 / 0.25 / 0.5 only -- all exact + multiples of 1/8 -- so the crest landed EXACTLY on an enumeration sample in every + single test, and a bug that only appears between samples was invisible to all of + them. The docstring above was already there when that happened. 0.3125 puts the + crest at half a sample; it is the phase that found the bug. + """ + sig = BandLimited(amp=amp, peak_sample=NPTS // 2 + phase) + got, want = _peak_local(sig.samples()), sig.truth(refine) + assert abs(got - want) < 5e-5, (amp, phase, got, want) + + +def test_exact_where_simpson_is_hundreds_of_nats_wrong(): + """rho ~ 220, sigma_t/deltaT = 0.0074. The truth needs a 1024x refinement to be + a truth at all, which is why this is one case and not a sweep.""" + sig = BandLimited(amp=200.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples() + want = sig.truth(1024) + assert abs(_peak_local(k) - want) < 1e-3, (_peak_local(k), want) + assert abs(_simpson_value(k) - want) > 100.0 + + +@pytest.mark.parametrize("amp,refine", [(1.0, 256), (5.0, 512), (40.0, 2048)]) +def test_accurate_on_a_non_periodic_window(amp, refine): + """The realistic case: a window cut from a longer band-limited signal, so the + periodic interpolant genuinely rings at the wrap and the window's own samples do + NOT determine the true continuous function.""" + sig = BandLimited(amp=amp, peak_sample=NPTS // 2, n_period=2 * NPTS, + background=0.12) + got, want = _peak_local(sig.samples()), sig.truth(refine) + assert abs(got - want) < 1e-3, (amp, got, want) + + +@pytest.mark.parametrize("amp", [1.0, 5.0, 40.0, 200.0, 2000.0]) +def test_agrees_with_the_dense_bandlimited_path(amp): + """The A/B this PR exists to make possible: the same rows, the same domain, two + different placements of the refined grid.""" + sig = BandLimited(amp=amp, peak_sample=NPTS // 2 + 0.25) + k = sig.samples() + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-4, amp + + +def test_beats_simpson_where_the_peak_is_under_resolved(): + """The whole point, stated as an inequality against the analytic truth rather + than as a claim about which is prettier.""" + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples() + ref = sig.truth(2048) + assert abs(_peak_local(k) - ref) < 1e-4 + assert abs(_simpson_value(k) - ref) > 1.0 + + +def test_the_crest_is_localised_between_enumeration_samples(): + """F1 REGRESSION -- the bug this module shipped with, and the sharpest test here. + + An enumerated extremum is a grid INDEX. The crest it stands for can be up to + ``h_enum/2`` away, while the interval half-width is ``W_SIGMA * sigma_t``, so once + ``W_SIGMA * sigma_t < h_enum/2`` -- any row with a derived factor of 512 or more -- + the peak falls entirely OUTSIDE its own interval and its mass is silently dropped. + + Centring on the sample instead of the crest measured, on this fixture at + ``sigma_t/deltaT = 0.0024``: **0.00 nats at offset 0, -6.52 at h_enum/4, -164.93 at + h_enum/2**, always negative. The sweep below walks the crest across a full + enumeration cell, so the worst case is inside it by construction rather than by + luck, and the reference is the dense path (exact here, and independently checked + against a closed-form truth in the control at the end). + """ + F = pl.PEAK_ENUM_FACTOR + ref_sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2) + k0 = ref_sig.samples()[None, :] + sigma, _, _ = tmq.peak_width_from_lnL(_lnL(k0.real, RHO_SQ), DELTAT) + assert pl.W_SIGMA * float(sigma[0]) < 0.5 * DELTAT / F, ( + "fixture is not in the regime this test exists for", float(sigma[0]) / DELTAT) + + for off in np.linspace(0.0, 1.0, 9): # a full enumeration cell + phase = off / F + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + phase) + k = sig.samples() + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-3, (off, phase) + + # control: at the phase that produced -164.93, both paths hit the analytic truth + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.5 / F) + k, want = sig.samples(), None + want = sig.truth(2048) + assert abs(_peak_local(k) - want) < 1e-3 + assert abs(_bandlimited(k) - want) < 1e-3 + + +def test_a_block_with_uniform_arrival_times_is_unbiased(): + """The production statement of the same thing, and the one that matters. + + A real block's arrival times bear no relation to the sample grid, so the + grid-quantisation error is drawn uniformly across an enumeration cell. Before the + localisation fix this block measured **median -1.15 nats, worst -131 nats, 56% of + rows wrong by more than 0.01 nats, and all of them ACCEPTED** -- a bias, not noise, + because the sign is always negative: mass is dropped, never added. + + Asserted on the WHOLE distribution rather than on a summary, because a median-only + check passes while a tail deletes extrinsic samples from the marginalization. + """ + rng = np.random.default_rng(20260828) + phases = rng.uniform(0.0, 1.0, 48) + k = np.stack([BandLimited(amp=2000.0, peak_sample=NPTS // 2 + p).samples() + for p in phases]) + r = np.full(k.shape, RHO_SQ) + got = np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)) + ref = np.asarray(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL)) + d = got - ref + assert np.max(np.abs(d)) < 1e-3, (np.median(d), d[np.argmax(np.abs(d))]) + assert abs(np.median(d)) < 1e-6, np.median(d) + assert pl.last_report()['n_peak_local_rows'] == len(phases) + + +def test_a_failed_localisation_sends_the_row_to_the_dense_path(monkeypatch): + """Fail closed. If the crest cannot be PLACED, the row is not approximated. + + Sabotaged by making the localiser report non-convergence, which is what a + pathological integrand would do. The value must still be right -- it comes from + the dense path -- and the row must be counted, not silently absorbed. + """ + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples() + real = pl.localise_peaks + + def never_converges(*a, **kw): + t, q, ok = real(*a, **kw) + return t, q, np.zeros_like(ok) + + monkeypatch.setattr(pl, 'localise_peaks', never_converges) + got = _peak_local(k) + rep = pl.last_report() + assert rep['n_dense_fallback_localise'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k) + + +def test_the_containment_check_catches_a_mis_placed_interval(monkeypatch): + """The a-posteriori half of the F1 fix, tested by re-introducing F1 exactly. + + Forcing the crest back onto the enumeration sample is precisely the old behaviour. + The interval then misses the peak, and the run must NOT report the truncated value: + the local grid fails to attain the localised crest's ``lnL`` and the row is handed + to the dense path. + + Note the check compares against ``lnL`` at the LOCALISED crest. Against the + enumeration SAMPLE it would pass here -- the sample is already ~87 nats below the + crest in this configuration -- which is why the sample cannot be the reference. + """ + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.5 / pl.PEAK_ENUM_FACTOR) + k = sig.samples() + real = pl.localise_peaks + + def snap_back_to_the_grid(Xw, fk, rows, t_grid, h_enum, tol, period, **kw): + t, q, ok = real(Xw, fk, rows, t_grid, h_enum, tol, period, **kw) + return t_grid, q, ok # crest value kept, position quantised: old bug + + monkeypatch.setattr(pl, 'localise_peaks', snap_back_to_the_grid) + got = _peak_local(k) + rep = pl.last_report() + # THE SAFETY PROPERTY, and it is the assertion that matters: the truncated value is not + # reported. The row is declined and given the dense value. + assert rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k) + # WHICH defence fires moved, and the move is an improvement. It used to be containment + # alone: the tail bound read the excluded crest at its SAMPLE, ~87 nats below the crest, + # so it passed. Now the outside supremum is evaluated off-grid, sees that crest at full + # height, and rejects the row first. Assert the disjunction rather than one counter, so + # this test pins the safety property and not the order the two defences happen to fire in. + assert (rep['n_dense_fallback_tail'] + rep['n_dense_fallback_containment']) == 1, rep + + +def test_containment_still_catches_a_mis_placed_interval_with_the_tail_bound_disabled( + monkeypatch): + """Containment is a SECOND line of defence, and after the off-grid tail bound it is no + longer the one that fires -- so it has to be tested in isolation or not at all. + + Evidence that it stopped being covered: disabling the containment check outright + (`contained = all True`) leaves the entire suite green, because the off-grid outside + supremum now rejects every mis-placed-interval fixture first. That is a real coverage + loss and this test is the fix. + + The two checks are NOT redundant, which is why containment is kept rather than deleted. + The tail bound is a statement about mass OUTSIDE the intervals, and it excludes covered + points by construction; it therefore cannot see a crest that is INSIDE its interval but + that the local grid failed to attain. Containment is the only thing that can. + + Isolated by making the tail bound unconditionally pass, then re-introducing F1 exactly as + the companion test does. With both defences live the tail bound wins; with only + containment live, containment must still catch it. + """ + monkeypatch.setattr(pl, 'TAIL_LOG_TOL', np.inf) + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.5 / pl.PEAK_ENUM_FACTOR) + k = sig.samples() + real = pl.localise_peaks + + def snap_back_to_the_grid(Xw, fk, rows, t_grid, h_enum, tol, period, **kw): + t, q, ok = real(Xw, fk, rows, t_grid, h_enum, tol, period, **kw) + return t_grid, q, ok # crest value kept, position quantised: the old bug + + monkeypatch.setattr(pl, 'localise_peaks', snap_back_to_the_grid) + got = _peak_local(k) + rep = pl.last_report() + assert rep['n_dense_fallback_tail'] == 0, ("tail bound was supposed to be disabled", rep) + assert rep['n_dense_fallback_containment'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k) + + +def test_a_nonlinear_distance_marginalization_style_callback(): + """The default helper is AFFINE in the kappa term, so an implementation that + accidentally assumed linearity would pass everything above. This callback is + monotone but not affine, and has a ``-inf`` domain edge.""" + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples() + got = _peak_local(k, callback=_lnL_distmarg_like) + assert abs(got - sig.truth(2048, callback=_lnL_distmarg_like)) < 1e-3 + assert abs(got - _bandlimited(k, callback=_lnL_distmarg_like)) < 1e-3 + + +# ------------------------------------------------- property 1: the merge + +def _unmerged_value(kappa_row, callback=_lnL): + """The NAIVE variant: one interval per enumerated peak, integrated and summed + WITHOUT merging. Reproduced here rather than described, because "we merge" is + the kind of statement that survives the code that implements it being deleted. + Mirrors ``~/tmarg_harness/peaklocal.py``.""" + k = np.asarray(kappa_row)[None, :] + F = pl.PEAK_ENUM_FACTOR + h = DELTAT / F + last = (NPTS - 1) * F + up = tmq.reflected_bandlimited_upsample(k, F)[0][:last + 1] + v = callback(up.real, RHO_SQ) + idx = np.where((v[1:-1] >= v[:-2]) & (v[1:-1] > v[2:]))[0] + 1 + idx = idx[v[idx] > v[idx].max() - pl.PEAK_KEEP_NATS] + parts = [] + for i in idx: + if i < 1 or i >= v.size - 1: + continue + d2 = (v[i + 1] - 2 * v[i] + v[i - 1]) / h ** 2 + if d2 >= 0: + continue + s = 1.0 / np.sqrt(-d2) + a = max(0.0, i * h - pl.W_SIGMA * s) + b = min(last * h, i * h + pl.W_SIGMA * s) + n_loc = max(3, int(np.ceil((b - a) / min(s / tmq.UPSAMPLE_SAFETY, h))) + 1) + tl = np.linspace(a, b, n_loc) + Xw, fk = pl.bandlimited_spectrum( + np.concatenate((k, np.flip(k, axis=-1)), axis=-1)) + kl = pl.eval_bandlimited_uniform(Xw, fk, np.array([tl[0]]), + np.array([tl[1] - tl[0]]), n_loc, + 2.0 * NPTS * DELTAT)[0] + parts.append(_log_trapz(callback(kl.real, RHO_SQ), tl[1] - tl[0])) + if not parts: + return np.nan + m = max(parts) + return m + np.log(sum(np.exp(p - m) for p in parts)) + + +def test_unmerged_intervals_double_count(): + """Merging is CORRECTNESS, not tidiness. + + Two overlapping windows integrated separately both contain the shared region, so + the log-sum-exp of the parts counts it twice. On a broad integrand -- where many + enumerated peaks sit within a few sigma of each other -- the prototype measured + **+1.6 nats at rho ~ 6**. Here the merged value is exact against the analytic + truth and the un-merged one is not, by a margin no rounding can explain. Without + this test, deleting the merge leaves every accuracy test above still green, + because on a sharply peaked row the intervals do not overlap at all. + """ + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25) # rho ~ 6 + k = sig.samples() + ref = sig.truth(128) + naive = _unmerged_value(k) + assert np.isfinite(naive) + assert naive - ref > 0.5, (naive, ref) # over-counts, and upward + assert abs(_peak_local(k) - ref) < 1e-3 + + +def test_two_peaks_merge_continuously_as_they_approach(): + """No regime switch and no threshold. + + Two peaks of equal height are walked together. Far apart the rule builds TWO + disjoint intervals; close together they overlap and the merge collapses them to + ONE. Both extremes must occur in the sweep, the count must never go UP as the + peaks approach, and -- the part that matters -- the value must track the analytic + truth right through the transition, because a regime switch would show up as a + step there. + """ + counts, seps = [], (200, 100, 40, 16, 6, 2) + for sep in seps: + sig = BandLimited(amp=200.0, peak_sample=NPTS // 2 - sep / 2.0, + extra_peaks=[(NPTS // 2 + sep / 2.0, 1.0)]) + got, want = _peak_local(sig.samples()), sig.truth(1024) + assert abs(got - want) < 1e-2, (sep, got, want) + counts.append(pl.last_report()['n_intervals_total']) + assert counts[0] == 2 and counts[-1] == 1, counts + assert all(counts[i] >= counts[i + 1] for i in range(len(counts) - 1)), counts + + +def test_a_broad_integrand_degenerates_into_the_dense_grid(): + """The other end of the same continuum: when the peaks crowd, the union grows to + the whole window, the local grid stops being cheaper than the dense one, and the + row is simply handed over. That is the degeneration completing -- not a special + case being detected -- and the answer is the dense path's, unchanged.""" + sig = BandLimited(amp=0.5, peak_sample=NPTS // 2 + 0.25) # rho ~ 11 + k = sig.samples() + assert _peak_local(k) == _bandlimited(k) + rep = pl.last_report() + assert rep['n_peak_local_rows'] == 0 and rep['n_dense_fallback_cost'] == 1, rep + + sharp = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.25) # rho ~ 700 + _peak_local(sharp.samples()) + assert pl.last_report()['n_peak_local_rows'] == 1 + + +# ------------------------------- property 2: enumeration and the tail bound + +def test_two_separated_peaks_are_both_found(): + """The anti-#201 test, stated behaviourally. + + RIFT PR #201 seeded a Newton solve at guessed points and missed genuine maxima, + returning ``-inf`` for a finite integral. Here two well-separated peaks of + comparable height are built; a seed-and-hope implementation converges to one of + them and reports roughly half the integral, i.e. ``log 2 = 0.69`` nats low. The + tolerance is far tighter than that, so this cannot pass by luck. + """ + sig = BandLimited(amp=200.0, peak_sample=NPTS // 3 + 0.3125, + extra_peaks=[(2 * NPTS // 3 + 0.40625, 1.0)]) + k = sig.samples() + got, want = _peak_local(k), sig.truth(1024) + assert abs(got - want) < 1e-3, (got, want) + assert pl.last_report()['n_peaks_total'] >= 2 + assert pl.last_report()['n_intervals_total'] >= 2 + + +def test_a_sabotaged_enumeration_is_caught_by_the_tail_bound(monkeypatch): + """COMPLETENESS BUYS SPEED; THE BOUND BUYS CORRECTNESS. + + The argument for truncating is that the mass outside the intervals is BOUNDED, + and the bound is computed from a grid that resolves ``kappa`` -- not from the + assumption that the enumeration found everything. So break the enumeration + deliberately: keep only the single highest maximum in each row. On the two-peak + integrand that discards half the mass, and the module must NOT report the + truncated value. It detects the shortfall and hands the row to the dense path, + which returns the right answer. + + If this test fails, the whole rigour claim in the module docstring is void. + """ + # OFF-GRID on purpose. Both peaks previously sat at NPTS//3 and 2*NPTS//3, which + # are exact enumeration samples -- so the outside maximum was sampled right on the + # discarded crest and the bound worked for a reason that does not generalise. Move + # them off the grid and the bound has to work on its own merits. + sig = BandLimited(amp=200.0, peak_sample=NPTS // 3 + 0.3125, + extra_peaks=[(2 * NPTS // 3 + 0.40625, 1.0)]) + k = sig.samples() + truth = sig.truth(1024) + + real_enumerate = pl.enumerate_peak_indices + + def only_the_best(q, xpy=np): + mask = real_enumerate(q, xpy=xpy) + out = np.zeros_like(mask) + for r in range(q.shape[0]): + cand = np.where(mask[r])[0] + if cand.size: + out[r, cand[np.argmax(q[r, cand + 1])]] = True + return out + + monkeypatch.setattr(pl, 'enumerate_peak_indices', only_the_best) + got = _peak_local(k) + rep = pl.last_report() + assert rep['n_dense_fallback_tail'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + assert abs(got - truth) < 1e-3, (got, truth) + + +def test_the_reported_tail_bound_is_actually_below_the_tolerance(): + """The diagnostic has to be load-bearing, not decorative: a row this rule + ACCEPTS must carry a bound strictly under the tolerance it claims to enforce.""" + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.25) + _peak_local(sig.samples()) + rep = pl.last_report() + assert rep['n_peak_local_rows'] == 1 + assert rep['tail_bound_worst'] < pl.TAIL_LOG_TOL, rep + + +def test_peak_positions_do_not_depend_on_distance_or_callback(): + """The invariant that licenses enumerating on ``kappa`` instead of on ``lnL``. + + Every shipped callback is monotone increasing in ``Re kappa`` (or ``|kappa|``) at + fixed ``rho_sq``, and ``rho_sq`` is time-independent on this path, so a monotone + map cannot move a maximum. Distance enters only as a positive rescaling of the + exponent's argument. Swept over three decades in 1/D and across an affine, a + nonlinear distmarg-shaped, and a log-shaped callback: the enumerated peak set + must be IDENTICAL, and equal to the peaks of ``Re kappa``. + + This is what keeps the likelihood callback -- a table interpolation in + production -- off the full time axis. + """ + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2 + 0.25, background=0.05, + n_period=2 * NPTS) + k = sig.samples()[None, :] + up = tmq.reflected_bandlimited_upsample( + k, pl.PEAK_ENUM_FACTOR)[0][:(NPTS - 1) * pl.PEAK_ENUM_FACTOR + 1] + ref = np.where(pl.enumerate_peak_indices(up.real[None, :])[0])[0] + assert ref.size > 3, "fixture must have several maxima for this to mean anything" + + # Every callback must be STRICTLY increasing and numerically safe over the whole + # range of `up.real`. A softplus was tried first and is wrong: `exp` overflows to + # `inf` across most of the range, the callback goes FLAT, and the test then fails + # for a reason that has nothing to do with the invariant it is checking. `cbrt` + # is strictly increasing, continuous, genuinely nonlinear, and cannot overflow. + callbacks = [_lnL, _lnL_distmarg_like, lambda x, r: np.cbrt(x - 0.5 * r)] + for one_over_d in (0.05, 1.0, 40.0): + for cb in callbacks: + # Through the SHIPPED enumerator on both sides. A hand-inlined copy of the + # comparison here silently stopped matching when the enumerator started + # including endpoints, which made this test fail for a reason that had + # nothing to do with the invariant it is about. + v = cb(one_over_d * up.real, RHO_SQ) + got = np.where(pl.enumerate_peak_indices(v[None, :])[0])[0] + assert np.array_equal(got, ref), (one_over_d, cb) + + +def test_the_enumeration_factor_finds_the_same_peaks_as_a_much_finer_grid(): + """``PEAK_ENUM_FACTOR`` is justified by the band limit -- ``kappa``'s narrowest + possible lobe is a half-cycle of width ``deltaT`` -- but that is an argument, and + arguments are cheap. Check it: every peak found at factor 64 that carries + representable mass must also be found at the shipped factor, to within one + coarse sample.""" + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2 + 0.25, background=0.20, + n_period=2 * NPTS) + k = sig.samples()[None, :] + + def peaks_at(F): + up = tmq.reflected_bandlimited_upsample(k, F)[0][:(NPTS - 1) * F + 1].real + i = np.where(pl.enumerate_peak_indices(up[None, :])[0])[0] + 1 + i = i[up[i] > up[i].max() - pl.PEAK_KEEP_NATS] + return np.sort(i * (DELTAT / F)) + + # To within one ENUMERATION sample, not one coarse sample. The old bound was + # `DELTAT` = 8 * h_enum, ~290x the interval half-width it was meant to justify, so + # it would have accepted an enumeration that missed by far more than the interval + # is wide. + coarse, fine = peaks_at(pl.PEAK_ENUM_FACTOR), peaks_at(64) + h_enum = DELTAT / pl.PEAK_ENUM_FACTOR + for t in fine: + assert np.min(np.abs(coarse - t)) <= h_enum, (t, coarse) + + +# ------------------------------------------- inherited invariants (PR #203) + +def test_edge_proximity_is_reported_but_selects_no_quadrature(): + """The edge guard is DIAGNOSTIC, and this test used to assert the opposite. + + It once routed a near-edge row to the caller's Simpson rule, because the periodic + reconstruction rang at the window wrap. rift_O4d e4ed25c7 removed that wrap by even + reflection and demoted the guard: crossing an arbitrary threshold must not silently + change likelihood quality. This module inherited the old routing and kept it across + the rebase, so it returned a SIMPSON value where the dense path returns a refined + one -- 3.79 nats apart on a row with peaks at both ends, with every fallback counter + reading zero because the row never entered the rule. + + So: still reported, and refined anyway, and agreeing with the dense path. + """ + guard = max(1, int(NPTS * tmq.EDGE_GUARD_FRACTION)) + for j in (guard - 1, NPTS - guard): + sig = BandLimited(amp=40.0, peak_sample=j) + k = sig.samples() + got = _peak_local(k) + rep = pl.last_report() + assert rep['n_wrap_exposed_rows'] == 1, (j, rep) + assert rep['n_refined_rows'] == 1, (j, rep) + assert got != _simpson_value(k), (j, "edge proximity still selects Simpson") + assert abs(got - _bandlimited(k)) < 1e-3, (j, got, _bandlimited(k)) + + +def test_row_classification_matches_the_dense_path_exactly(): + """The invariant the rebase broke, asserted so it cannot break silently again. + + peak-local's contract is that it changes WHERE the refined grid is placed and + nothing about WHICH rows get one. Three clauses of the dense path's classification + had drifted out of this module -- the reflected reconstruction, the demotion of the + edge guard, and `boundary_unresolved` -- and each drift showed up as peak-local + silently returning a lower-accuracy value than `time_marginalize_bandlimited` for + the same row. Compare the CLASSIFICATION, not just the values: a value check passes + whenever the two rules happen to agree, which on most rows they are designed to. + """ + rows = [np.zeros(NPTS, dtype=complex), # flat + BandLimited(amp=40.0, peak_sample=3).samples(), # edge + BandLimited(amp=40.0, peak_sample=NPTS - 4).samples(), # far edge + BandLimited(amp=0.5, peak_sample=NPTS // 2 + 0.25).samples(), # broad + BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.25).samples()] + k = np.stack(rows) + r = np.full(k.shape, RHO_SQ) + pl.time_marginalize_peak_local(k, r, DELTAT, _lnL) + pr = dict(pl.last_report()) + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL) + br = dict(tmq.last_report()) + for key in ('n_refined_rows', 'n_flat_rows', 'n_unmeasurable_rows', + 'n_wrap_exposed_rows'): + assert pr[key] == br[key], (key, pr[key], br[key]) + + +def test_the_edge_guard_covers_the_RIGHT_edge_too(): + """Pin BOTH boundaries at the sample. An off-by-one in the upper term leaves + exactly one row's worth of the right guard band open, and every peak-placement + fixture elsewhere is far enough inside that both spellings agree.""" + guard = max(1, int(NPTS * tmq.EDGE_GUARD_FRACTION)) + + def row_peaking_at(j): + t = np.arange(NPTS, dtype=float) + return (np.exp(-0.5 * ((t - j) / 0.35) ** 2) * 40.0).astype(complex) + + for j, expect_exposed in ((guard - 1, True), (guard, False), + (NPTS - 1 - guard, False), (NPTS - guard, True)): + _peak_local(row_peaking_at(j)) + rep = pl.last_report() + assert (rep['n_wrap_exposed_rows'] == 1) == expect_exposed, (j, rep) + # The guard REPORTS; it no longer selects a rule, so all four are refined. + assert rep['n_refined_rows'] == 1, (j, rep) + + +def test_flat_and_signal_free_rows_are_not_refined_and_not_reported_as_exposed(): + k = np.zeros(NPTS, dtype=complex) + assert _peak_local(k) == _simpson_value(k) + rep = pl.last_report() + assert rep['n_flat_rows'] == 1 and rep['n_wrap_exposed_rows'] == 0, rep + assert rep['n_refined_rows'] == 0 and rep['n_peak_local_rows'] == 0, rep + + +def test_an_all_minus_inf_row_returns_minus_inf_and_not_nan(): + """A per-row log-sum-exp offset would compute ``-inf - (-inf) = NaN`` and feed a + NaN into the sampler weights.""" + k = np.zeros((2, NPTS), dtype=complex) + k[0] = BandLimited(amp=40.0, peak_sample=NPTS // 2).samples() + r = np.full(k.shape, RHO_SQ) + + def domain_limited(term, rho): + """``-inf`` outside a table domain, the shape the distance-marginalization + callback actually has. Row 1 has zero kappa, so it lands outside everywhere + and its ``lnL(t)`` is ``-inf`` for the whole window. + + Note this callback is written as a function of its ARGUMENTS only. A first + version blanked "row 1" by index, which is wrong: the callback is also invoked + on peak stencils and local grids whose leading axis is not the row axis, so it + indexed into the wrong thing and raised.""" + return np.where(term > 100.0, term - 0.5 * rho, -np.inf) + + got = pl.time_marginalize_peak_local(k, r, DELTAT, domain_limited) + assert np.isfinite(got[0]) and got[1] == -np.inf, got + assert not np.any(np.isnan(np.asarray(got))) + + +def test_unmeasurable_row_falls_back_and_is_counted(): + """``lnL(t)`` non-finite around its maximum at every stencil half-width: no width + can be justified, so none is guessed at.""" + k = BandLimited(amp=40.0, peak_sample=NPTS // 2).samples()[None, :] + r = np.full(k.shape, RHO_SQ) + + def holed(term, rho): + v = np.array(_lnL(term, rho), dtype=float, copy=True) + j = int(np.argmax(v[0])) + for d in (1, 2, 4, 8): + for s in (-1, 1): + idx = j + s * d + if 0 <= idx < v.shape[-1]: + v[0, idx] = -np.inf + return v + + out = pl.time_marginalize_peak_local(k, r, DELTAT, holed) + rep = pl.last_report() + assert rep['n_unmeasurable_rows'] == 1 and rep['n_refined_rows'] == 0, rep + assert np.isfinite(out[0]) + + +def test_time_dependent_rho_sq_is_refused(): + k = BandLimited(amp=40.0, peak_sample=NPTS // 2).samples()[None, :] + r = np.full(k.shape, RHO_SQ) + r[0, 3] += 1.0 + with pytest.raises(NotImplementedError): + pl.time_marginalize_peak_local(k, r, DELTAT, _lnL) + + +def test_a_nan_self_term_does_not_abort_the_run(): + """A NaN self-term is NORMAL -- the defensive proposal component draws + physically-extreme points where the likelihood is NaN. A bare ``==`` in the + time-independence tripwire makes ``nan != nan`` fire and kills the ILE process, + blaming a rotating-response path that is not in use.""" + k = np.zeros((2, NPTS), dtype=complex) + k[0] = BandLimited(amp=40.0, peak_sample=NPTS // 2).samples() + r = np.full(k.shape, RHO_SQ) + r[1] = np.nan + out = pl.time_marginalize_peak_local(k, r, DELTAT, _lnL) + assert np.isfinite(out[0]) and np.isnan(out[1]) + + +def test_simps_is_required_for_a_non_numpy_backend(): + """scipy's Simpson rule RAISES on a device array, and the fallback rows must use + the rule the caller's own likelihood uses. A default here is how every + ``--vectorized --gpu`` run of the sibling option once crashed.""" + class FakeXpy(object): + def __getattr__(self, name): + return getattr(np, name) + + k = BandLimited(amp=40.0, peak_sample=NPTS // 2).samples()[None, :] + with pytest.raises(ValueError): + pl.time_marginalize_peak_local(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL, + xpy=FakeXpy()) + + +def test_the_ceiling_fails_closed_for_the_SHARPEST_rows_not_just_broad_ones(): + """A row whose derived factor exceeds the ceiling must RAISE, and the route to that + must not depend on the cost gate declining it. + + This is F3, and it was fail-OPEN. The cost gate compares the local cost against the + dense cost, so the sharper the row the more certainly peak-local keeps it -- and a + row past the ceiling is the sharpest kind there is. At a derived factor of 8192 the + dense path raised, as designed, while peak-local returned -24451 nats and reported + ``tail_bound_worst = -2721``. + + The earlier version of this test only passed because it set + ``UPSAMPLE_FACTOR_MAX = 2``, which is broad enough that the COST gate declined the + row first; the ceiling was never what routed it. Here the ceiling is lowered but + left well inside the regime where the cost gate would happily keep the row, so the + ceiling check is the only thing that can produce the raise -- and the control below + confirms the row is one peak-local would otherwise have taken. + """ + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + + # control: at the shipped ceiling this row is handled by peak-local + pl.time_marginalize_peak_local(k, r, DELTAT, _lnL) + assert pl.last_report()['n_peak_local_rows'] == 1 + + old_max = tmq.UPSAMPLE_FACTOR_MAX + try: + tmq.UPSAMPLE_FACTOR_MAX = 256 + sigma, _, _ = tmq.peak_width_from_lnL(_lnL(k.real, r), DELTAT) + assert int(tmq.required_upsample_factors(sigma, DELTAT)[0]) > 256 + with pytest.raises(RuntimeError): + pl.time_marginalize_peak_local(k, r, DELTAT, _lnL) + finally: + tmq.UPSAMPLE_FACTOR_MAX = old_max + + +def test_a_cost_fallback_row_gets_the_DENSE_value_not_an_approximation(): + """A row this rule declines must come back with the reviewed dense + implementation's number, not Simpson's and not a truncated local estimate.""" + sig = BandLimited(amp=0.5, peak_sample=NPTS // 2 + 0.25) # rho ~ 11 + k = sig.samples() + got = _peak_local(k) + rep = pl.last_report() + assert rep['n_dense_fallback_rows'] == 1 and rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k) + assert got != _simpson_value(k) + + +def test_a_mixed_block_gives_every_row_its_own_treatment(): + """Rows are independent, and a block that mixes flat, exposed, dense-fallback and + peak-local rows must give each the value it would have got alone. A shared + offset or a shared factor would show up here and nowhere else.""" + rows = [np.zeros(NPTS, dtype=complex), + BandLimited(amp=40.0, peak_sample=3).samples(), + BandLimited(amp=0.5, peak_sample=NPTS // 2 + 0.25).samples(), + BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.25).samples()] + singles = [_peak_local(r) for r in rows] + k = np.stack(rows) + block = pl.time_marginalize_peak_local(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + rep = pl.last_report() + # 2, not 1: the near-edge row (peak_sample=3) is now refined like any other, since + # the edge guard became diagnostic. The flat row and the broad row account for the + # rest -- the broad one is declined on cost and gets the dense value. + assert rep['n_rows'] == 4 and rep['n_peak_local_rows'] == 2, rep + for i, (a, b) in enumerate(zip(singles, np.asarray(block))): + assert a == b or abs(a - b) < 1e-9, (i, a, b) + + +def test_phase_marginalization_is_refused_by_the_library(): + """A DELIBERATE SCOPE CUT, refused rather than silently ignored. + + Production marginalizes over distance, not phase. Under phase marginalization + the time peak's Laplace width carries an ``(I1/I0)(|kappa|/D)`` factor that does + not reduce, so the local spacing stops being derivable from ``rho_sq`` and the + curvature alone -- the derived-not-configured property this whole line of work + rests on. Refusing keeps the option from being run, believed, and compared + against. + """ + sig = BandLimited(amp=200.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + with pytest.raises(NotImplementedError): + pl.time_marginalize_peak_local(k, r, DELTAT, _lnL, phase_marginalization=True) + + +def test_the_bandlimited_path_still_supports_phase_marginalization(): + """The cut applies to the NEW rule only. 'bandlimited' is the reviewed reference + implementation and must not regress -- and it is what a caller who needs phase + marginalization is told to use.""" + sig = BandLimited(amp=200.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + a = float(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL, + phase_marginalization=True)[0]) + b = float(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL, + phase_marginalization=False)[0]) + assert np.isfinite(a) and a != b, (a, b) + + +def test_the_memory_chunking_path_assembles_its_result(): + """The extrinsic axis is chunked so one dense temporary stays inside a working-set + budget. Rows are independent, so chunking must not change the ANSWER -- and that, + not the physics, is what this test is for: it compares the chunked result against + the unchunked one on identical inputs. + + Bit-identity is deliberately NOT the bar, because it is not available and asserting + it fails. MEASURED on these three rows: the chunked and unchunked values differ by + **0, 0 and 2 ULPs** (2.4e-16 relative). Chunking changes the leading dimension of + the FFT and of the reduction inside the local evaluator, and both numpy's FFT and + its pairwise summation reassociate with batch shape. A tolerance of a few ULPs is + still an extremely sharp instrument for what this test is actually guarding -- + a dropped, duplicated or mis-ordered chunk moves a row by nats, not by ULPs. + + Comparing against an analytic truth instead would be both slower (the truth needs a + 4096x refinement for the sharpest row) and weaker: a tolerance loose enough to + absorb the reference's own error is loose enough to hide an assembly bug. + """ + amps = (40.0, 200.0, 2000.0) + k = np.stack([BandLimited(amp=a, peak_sample=NPTS // 2 + 0.25).samples() + for a in amps]) + r = np.full(k.shape, RHO_SQ) + whole = np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)) + assert pl.last_report()['n_peak_local_rows'] == len(amps) + + old = pl._CHUNK_BYTES + try: + pl._CHUNK_BYTES = 1 # forces one row per chunk + chunked = np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)) + finally: + pl._CHUNK_BYTES = old + ulps = np.abs(chunked - whole) / np.spacing(np.abs(whole)) + assert np.all(ulps <= 8), (whole, chunked, ulps) + assert pl.last_report()['n_peak_local_rows'] == len(amps) + + +def test_return_peaks_exposes_t_star_and_the_local_width(): + """``t_star`` and the local curvature are first-class OUTPUTS, not internal + temporaries. They are distance- and callback-independent (see the invariance + test above), which is what a time-first reordering of the marginalizations would + need, so they are exposed deliberately rather than incidentally.""" + # OFF the enumeration grid, and with no background, so the true crest is known in + # closed form: this fixture's kernel is symmetric about `peak_sample`. + peak_sample = NPTS // 2 + 0.3125 + sig = BandLimited(amp=2000.0, peak_sample=peak_sample) + k = sig.samples()[None, :] + out, peaks = pl.time_marginalize_peak_local( + k, np.full(k.shape, RHO_SQ), DELTAT, _lnL, return_peaks=True) + assert peaks[0] is not None + t_star, sigma_star = peaks[0] + + # Against the TRUE crest, not against the coarse argmax. Measuring the distance to + # the nearest coarse sample tests nothing -- it is 0.3125*deltaT here by + # construction, whether or not any localisation happened. The bar is a small + # fraction of the ENUMERATION spacing, which is what "sub-sample" has to mean when + # this is offered as the t_star a time-first reordering would build on. + h_enum = DELTAT / pl.PEAK_ENUM_FACTOR + err = np.min(np.abs(t_star - peak_sample * DELTAT)) + assert err < 0.01 * h_enum, (err / h_enum, "localisation is not sub-sample") + assert np.all(sigma_star > 0) and np.all(np.isfinite(sigma_star)) + + +def test_the_tuned_constants_are_pinned_to_their_measured_values(): + """These are not free parameters. Each is justified in the module docstring or + in DESIGN_time_marginalization_peak_local.md, and the suite otherwise pins them + only loosely -- so changing one could pass CI while invalidating the argument + behind it. Changing a value here is the deliberate act of also updating that + record.""" + assert pl.PEAK_ENUM_FACTOR == 8 + assert pl.W_SIGMA == 12.0 + assert pl.PEAK_KEEP_NATS == 60.0 + assert pl.TAIL_LOG_TOL == -23.0 + assert pl.MAX_INTERVALS == 32 + assert pl._RECURRENCE_REANCHOR == 64 + assert pl.LOCALISE_SAFETY == 0.25 + assert pl.LOCALISE_MAX_ITER == 16 + assert pl.CONTAINMENT_SLACK_NATS == 0.5 + + # The relation that USED to be an unstated precondition. `erfc(W/sqrt2)` bounds the + # truncation of a Gaussian about its CREST, so centring on the enumeration sample + # silently required `W_SIGMA * sigma_t >= h_enum/2` -- which couples W_SIGMA to + # PEAK_ENUM_FACTOR, is violated by every sharp row, and was nowhere asserted. + # Localisation discharges it: the interval is centred on the crest and widened by + # the localisation residual, so what has to hold is only that the residual is small + # against the half-width. That is checkable, so check it. + assert pl.LOCALISE_SAFETY < 0.1 * pl.W_SIGMA, ( + "the localisation residual must be negligible against the interval half-width") + assert pl.CONTAINMENT_SLACK_NATS < 1.0, ( + "the containment slack must be far below the smallest miss F1 produced (6.5 nats)") + # inherited, and the peak-local path derives its LOCAL spacing from this one + assert tmq.UPSAMPLE_SAFETY == 2.0 + assert tmq.EDGE_GUARD_FRACTION == 0.125 + + +def test_peak_local_is_a_recognised_quadrature_name(): + assert tmq.validate_time_quadrature('peak-local') == 'peak-local' + assert 'peak-local' in tmq.TIME_QUADRATURE_CHOICES + with pytest.raises(ValueError): + tmq.validate_time_quadrature('peaklocal') # sic + with pytest.raises(ValueError): + tmq.validate_time_quadrature('peak_local') # sic + + +def test_a_secondary_crest_between_samples_is_not_dropped_by_the_keep_filter(): + """G1 REGRESSION -- the same defect as F1, at a different site. + + The keep filter compared `q_up[rows_p, cols_p]`, the SAMPLE value, against the + highest sample. A crest a distance d from its sample reads `(d/sigma)^2/2` nats + low, so a peak between samples was measured against a peak on a sample and dropped. + MEASURED on this fixture at rho ~ 700: the secondary crest is 1.003 nats below the + dominant crest but its SAMPLE is 70.99 nats below, past `PEAK_KEEP_NATS = 60`, and + the answer came back **exactly -log(2) = -0.693147** low: one of two equal peaks + silently deleted, with no fallback triggered and `tail_bound_worst = -120`. + + THE ASYMMETRY IS THE POINT, and it is the cell the suite was missing. Every + multi-peak fixture here ran at `amp=200` (just above the threshold), every sharp + fixture was single-peak, and the updated F1 tests moved BOTH peaks off-grid by + similar amounts -- so the two deficits cancelled. The defect needs one peak near a + sample and one between, which is the generic production case. + """ + F = pl.PEAK_ENUM_FACTOR + for off in (0.0, 0.25, 0.5, 0.75): + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 3, + extra_peaks=[(2 * NPTS // 3 + off / F, 1.0)]) + k = sig.samples() + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-3, off + rep = pl.last_report() + assert rep['n_peaks_total'] == 2, (off, rep) # neither peak dropped + assert rep['n_peak_local_rows'] == 1, (off, rep) # and not rescued by fallback + + +def test_the_keep_threshold_is_justified_without_the_tail_bound(): + """G2. `PEAK_KEEP_NATS` and `TAIL_LOG_TOL` are NOT independent, and the docstring + that said a dropped peak "enters the tail bound" was vacuous: since `q_out_max` is + at least the dropped peak's own sample, the bound accepts the row unless + `sigma_t < 2.6e-18 s`, while `UPSAMPLE_FACTOR_MAX` bounds the sharpest legal row ten + orders of magnitude above that. For every row this module can legally handle, a + keep-filter drop is automatically accepted. + + So the filter has to stand on its own magnitude argument, and that argument ties it + to `UPSAMPLE_FACTOR_MAX`. Assert the inequality, not the constants. + """ + t_window = 0.075 + sigma_min = DELTAT / tmq.UPSAMPLE_FACTOR_MAX + log_rel_mass = -pl.PEAK_KEEP_NATS + np.log(t_window / sigma_min) + assert log_rel_mass < np.log(1e-16), ( + "a dropped peak's relative mass is not below double precision", log_rel_mass) + # and the relation that makes the tail bound unable to backstop it, recorded so a + # future change to either constant has to confront it + assert pl.PEAK_KEEP_NATS > -pl.TAIL_LOG_TOL + + +def test_W_SIGMA_gives_the_tail_bound_its_structural_slack(): + """WHY the sampled tail bound holds, rather than "we could not break it". + + An independent attempt to break `q_out_max` -- 24 accepted rows, honest supremum + recomputed on a 4096x grid -- failed, worst honest margin -63.42 against + `TAIL_LOG_TOL = -23`. But the reason is NOT that the sampling is adequate: it is + that the outside supremum sits at an interval EDGE, already `W_SIGMA**2/2 = 72` nats + below the crest, which leaves ~40 nats of structural slack. Drop `W_SIGMA` below + ~8-9 and that slack vanishes and the bound is silently invalid, with nothing + reporting it. + + So the inequality gets the same treatment as `PEAK_KEEP_NATS` vs + `UPSAMPLE_FACTOR_MAX`: asserted, not admired. Requiring + + W_SIGMA**2 / 2 > |TAIL_LOG_TOL| + log(T_out / (sqrt(2 pi) sigma_min)) + + ties `W_SIGMA` to `TAIL_LOG_TOL` and to `UPSAMPLE_FACTOR_MAX`, and none of the three + can now be moved alone. + """ + t_window = 0.075 + sigma_min = DELTAT / tmq.UPSAMPLE_FACTOR_MAX + need = -pl.TAIL_LOG_TOL + np.log(t_window / (np.sqrt(2 * np.pi) * sigma_min)) + have = pl.W_SIGMA ** 2 / 2.0 + assert have > need, ("W_SIGMA no longer gives the tail bound its slack", have, need) + # the margin is large, but assert the inequality rather than the margin: it is the + # inequality that a future edit has to preserve + assert have - need > 10.0, (have, need) + + +def test_the_localised_crest_stays_inside_its_bracket(): + """G4/M4. `localise_peaks` brackets at +/- h_enum and accepts anything strictly + inside, so the displacement bound is h_enum, NOT h_enum/2 -- and it is approached: + 0.959*h_enum was observed over 14,182 peaks. The gate interval must therefore widen + by h_enum, and an escaped iterate must not be accepted.""" + rng = np.random.default_rng(7) + k = np.stack([BandLimited(amp=2000.0, peak_sample=NPTS // 2 + x).samples() + for x in rng.uniform(0, 1, 24)]) + r = np.full(k.shape, RHO_SQ) + _, peaks = pl.time_marginalize_peak_local(k, r, DELTAT, _lnL, return_peaks=True) + h_enum = DELTAT / pl.PEAK_ENUM_FACTOR + worst = 0.0 + for pk in peaks: + if pk is None: + continue + d = np.abs(pk[0] / h_enum - np.round(pk[0] / h_enum)) + worst = max(worst, float(np.max(d))) + assert worst <= 1.0, worst # the bracket really does bound it + # ...and the gate's widening must cover that bound, or its "superset" claim is false + assert 1.0 * h_enum >= worst * h_enum + + +def test_a_row_with_a_near_edge_secondary_peak_is_answered_correctly(): + """Near-edge secondary peaks, end to end. + + NAMED FOR WHAT IT ACTUALLY CHECKS. An earlier version of this was called + `..._still_gets_its_own_curvature` and claimed to cover the G5 stencil-centre fix. + It does not, and the mutation sweep said so: re-clipping the stencil centre leaves + every assertion here green, changing the answer only at the 1e-12 level. The reason + is that on these fixtures the edge peak is >`PEAK_KEEP_NATS` below the dominant + crest and is dropped by the keep filter regardless (`n_peaks_total == 1`), so the + stencil never runs on it. Making an edge peak both near-edge AND within 60 nats + needs an amplitude regime where the row is declined on cost instead -- so the shape + is not reachable through the public entry point with the fixtures available here. + + That gap is recorded in DESIGN_time_marginalization_peak_local.md rather than + papered over. The reviewer who found G5 measured it directly (-0.3124 nats), and + the fix is applied; what is missing is a test of mine that exercises it. + """ + for edge in (0.05, 0.3, 1.0): + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2 + 0.3125, + extra_peaks=[(edge, 0.9)]) + k = sig.samples() + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-3, edge + + +def test_the_reported_tail_bound_matches_an_independent_recomputation(): + """R2-F2. Asserting only `bound < TAIL_LOG_TOL` pins nothing: the measured margins + are -440 / -1791 / -4495 against a tolerance of -23, so any error in the bound + smaller than ~400 nats is invisible, and deleting `log(T_outside)` entirely, or + marking extra samples covered, both survived. Recompute the bound from the module's + own reported peaks and require agreement.""" + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + out, peaks = pl.time_marginalize_peak_local(k, r, DELTAT, _lnL, return_peaks=True) + rep = pl.last_report() + assert rep['n_peak_local_rows'] == 1 and peaks[0] is not None + + F = pl.PEAK_ENUM_FACTOR + h_enum = DELTAT / F + t_last = (NPTS - 1) * DELTAT + t_star, sigma = peaks[0] + half = pl.W_SIGMA * sigma + pl.LOCALISE_SAFETY * sigma + lo = np.maximum(t_star - half, 0.0) + hi = np.minimum(t_star + half, t_last) + starts, stops, _ = (lambda o: (o[3], o[4], o[1]))( + pl.merge_intervals_by_row(np.zeros(len(lo), dtype=np.int64), lo, hi, t_last)) + + up = tmq.reflected_bandlimited_upsample(k, F)[0][:(NPTS - 1) * F + 1].real + covered = np.zeros(up.size, dtype=bool) + for a, b in zip(starts, stops): + i0, i1 = int(np.ceil(a / h_enum)), int(np.floor(b / h_enum)) + if i1 >= i0: + covered[max(i0, 0):i1 + 1] = True + T_out = t_last - float(np.sum(stops - starts)) + + # The outside supremum is a CERTIFIED bound now, not a sample and not an evaluation at a + # few points, so it is pinned against the HONEST supremum with a two-sided window rather + # than re-derived: it must never be below the honest value (or it is not a bound) and must + # not exceed it by more than the certificate's own remainder (or it is not useful). + # + # A one-sided `< TAIL_LOG_TOL` assertion pins nothing here -- the measured margins are tens + # of nats against a tolerance of -23, so any error smaller than ~40 nats is invisible, and + # deleting `log(T_outside)` or building `T_outside` from grid indices both survived it. + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + period_ref = 2.0 * NPTS * DELTAT + m4 = float(pl.spectral_derivative_bound(Xw, fk, period_ref, 4)[0]) + remainder = m4 * h_enum ** 4 / 384.0 + + # THE REFERENCE HAS TO BE EXACT AT THE INTERVAL ENDS, and a uniform grid is not. On this + # fixture the merged interval is SUB-CELL (sigma = 5.1e-7 s) and even a 512x grid has + # spacing ~= sigma, so its nearest point outside an end under-reads that end by + # |dlnL/dt| * spacing = W_SIGMA/sigma * spacing ~= 11 nats -- and the supremum over the + # uncovered set sits exactly at an end. Comparing against that reference makes a correct + # bound look 11 nats loose. So: grid for the bulk, EXACT evaluation at the ends. + sub = 64 + qf = tmq.reflected_bandlimited_upsample(k, F * sub)[0, :(NPTS - 1) * F * sub + 1].real + tf = np.arange(qf.size) * (h_enum / sub) + outside = np.ones(qf.size, dtype=bool) + for a, b in zip(starts, stops): + outside &= ~((tf >= a) & (tf <= b)) + ends = np.concatenate([starts, stops]) + q_ends, _, _ = pl.eval_bandlimited_points( + Xw, fk, np.zeros(ends.size, dtype=np.int64), ends, period_ref) + q_honest = max(float(np.max(qf[outside])), float(np.max(q_ends))) + honest = np.log(T_out) + _lnL(q_honest, RHO_SQ) - float(out[0]) + + got = rep['tail_bound_worst'] + assert got >= honest - 1e-6, ("not a bound", got, honest) + assert got <= honest + remainder + 1.0, ("bound far looser than its remainder", + got, honest, remainder) + + +def test_merge_keeps_a_fully_contained_interval_inside_its_enclosure(): + """R2-F3. The running maximum in `merge_intervals_by_row` is what stops a CONTAINED + interval truncating the one enclosing it: [0,10] then [1,3] must give [0,10], and + with a plain `hi_s` it gives [0,3], silently dropping (3,10]. The row-boundary half + of that trick is covered elsewhere; this is the containment half, and it is a unit + test so it does not depend on a fixture reaching the code.""" + rows = np.array([0, 0, 0], dtype=np.int64) + lo = np.array([0.0, 1.0, 2.0]) + hi = np.array([10.0, 3.0, 4.0]) + _, _, g_row, g_lo, g_hi = pl.merge_intervals_by_row(rows, lo, hi, 100.0) + assert g_row.size == 1, (g_row, g_lo, g_hi) + assert g_lo[0] == 0.0 and g_hi[0] == 10.0, (g_lo, g_hi) + + # two rows, each with a nested pair, to keep the row-reset and the running maximum + # exercised together rather than one masking the other + rows = np.array([0, 0, 1, 1], dtype=np.int64) + lo = np.array([0.0, 1.0, 5.0, 6.0]) + hi = np.array([10.0, 3.0, 20.0, 7.0]) + _, _, g_row, g_lo, g_hi = pl.merge_intervals_by_row(rows, lo, hi, 100.0) + assert list(g_row) == [0, 1] and list(g_hi) == [10.0, 20.0], (g_row, g_lo, g_hi) + + +def test_intervals_are_clipped_to_the_integration_domain(): + """R2-F4. Without the clip an interval can start below 0 or end past t_last, and + the evaluator is PERIODIC -- it returns the value from the opposite end of the + window, wrapping mass from outside the integration domain into the answer. More + reachable since endpoints became enumerable, not less.""" + t_last = (NPTS - 1) * DELTAT + sig = BandLimited(amp=5.0, peak_sample=NPTS // 2 + 0.3125, + extra_peaks=[(0.2, 0.98), (NPTS - 1.2, 0.98)]) + k = sig.samples() + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-3 + _, peaks = pl.time_marginalize_peak_local( + k[None, :], np.full((1, NPTS), RHO_SQ), DELTAT, _lnL, return_peaks=True) + if peaks[0] is not None: + t_star, sigma = peaks[0] + half = pl.W_SIGMA * sigma + pl.LOCALISE_SAFETY * sigma + assert np.all(np.maximum(t_star - half, 0.0) >= 0.0) + assert np.all(np.minimum(t_star + half, t_last) <= t_last) + + +def test_a_plateau_yields_exactly_one_enumerated_maximum(): + """R2-F5. The `>=` / `>` asymmetry is stated as a deliberate property in the + docstring and had no detector: swapping one way gives a plateau NO peak, the other + gives it EVERY sample. Both are wrong and both were invisible.""" + q = np.array([[0.0, 1.0, 1.0, 1.0, 0.0]]) + m = pl.enumerate_peak_indices(q) + assert int(np.sum(m)) == 1, m + assert int(np.where(m[0])[0][0]) == 3, m # the LAST index of the plateau + + q = np.array([[0.0, 1.0, 0.0, 2.0, 0.0]]) # two isolated maxima + assert int(np.sum(pl.enumerate_peak_indices(q))) == 2 + + +def test_the_local_trapezoid_uses_endpoint_half_weights(): + """R2-F6. Benign at `W_SIGMA = 12`, where the endpoints are `exp(-72)` relative -- + but the "trapezoid, not Simpson" argument had no detector at all, and it stops being + benign the moment `W_SIGMA` moves. Pinned on a constant integrand, where the + trapezoid rule is exact and the half-weights are the whole difference.""" + n, h = 11, 0.25 + lnL = np.zeros((1, n)) + got = float(pl._log_trapz_local(lnL, np.array([h]))[0]) + assert abs(np.exp(got) - (n - 1) * h) < 1e-12, (np.exp(got), (n - 1) * h) + + +def test_the_pre_enumeration_cost_gate_actually_fires(): + """R2-F7. Gate 1 is the fix this PR credits with removing a 0.43x regression, and + it was removable with the suite green because both gates shared one counter -- + nothing could tell which had fired. Split, and asserted here on a broad row that + must be declined BEFORE any enumeration is done for it.""" + sig = BandLimited(amp=0.02, peak_sample=NPTS // 2 + 0.3125) # broad, factor ~4 + k = sig.samples() + assert _peak_local(k) == _bandlimited(k) + rep = pl.last_report() + assert rep['n_dense_fallback_cost_pregate'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + + # control: a sharp row must NOT be declined by the pre-gate, or the gate is simply + # rejecting everything and the assertion above means nothing + sharp = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.3125) + _peak_local(sharp.samples()) + assert pl.last_report()['n_dense_fallback_cost_pregate'] == 0 + assert pl.last_report()['n_peak_local_rows'] == 1 + + +def test_a_row_with_more_structure_than_MAX_INTERVALS_is_declined(): + """R2-F7. `MAX_INTERVALS` is a fail-closed guard and was removable with the suite + green. + + Tested by lowering the constant rather than by building a 33-peak fixture, and the + reason is worth stating: a comb of nearly-equal peaks does not survive contact with + the rest of the algorithm -- at any amplitude tried, interference between the bumps + spread their crests by more than `PEAK_KEEP_NATS`, so all but one were dropped and + the row arrived at the guard with a single interval. Lowering the constant tests the + ROUTING, which is what the guard is; the control below shows the same row is handled + by peak-local when the guard is not in the way, so the assertion is not vacuous. + """ + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 3 + 0.3125, + extra_peaks=[(2 * NPTS // 3 + 0.25 / pl.PEAK_ENUM_FACTOR, 1.0)]) + k = sig.samples() + + # control: two intervals, handled + assert abs(_peak_local(k) - _bandlimited(k)) < 1e-3 + assert pl.last_report()['n_intervals_total'] == 2, pl.last_report() + assert pl.last_report()['n_peak_local_rows'] == 1 + + old_max = pl.MAX_INTERVALS + try: + pl.MAX_INTERVALS = 1 + got = _peak_local(k) + rep = pl.last_report() + finally: + pl.MAX_INTERVALS = old_max + # NOTE, from the mutation sweep: this is killed by the PROVISIONAL structure gate, + # not the final one. Disabling the final `too_much` check alone leaves this green, + # because the provisional gate declines the row first. The final check is therefore + # a belt-and-braces guard that no fixture here reaches -- it is kept because the + # final interval count CAN exceed the provisional one (narrower intervals merge less + # readily), but that shape is not exercised, and saying so is more useful than a + # contrived fixture that pretends otherwise. + assert rep['n_dense_fallback_structure'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k) + + +def test_the_curvature_stencil_widens_over_a_hole_on_the_ENUMERATION_grid(): + """R2-F7. The widening ladder is tested for the coarse-grid width estimator but was + untested for the peak stencil, where the callback's `-inf` domain edge also lands. + A hole at the immediate neighbours must be stepped over, not read as "no curvature" + -- which would drop the peak.""" + sig = BandLimited(amp=200.0, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + ref = float(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)[0]) + n_ref = pl.last_report()['n_peak_local_rows'] + + peak_q = np.max(k.real) + + def holed_near_the_crest(term, rho): + v = _lnL(term, rho) + # blank a thin shell just inside the crest: the d=1 stencil straddles it, the + # wider half-widths step over it + band = (term < peak_q * 0.99999) & (term > peak_q * 0.9999) + return np.where(band, -np.inf, v) + + got = float(pl.time_marginalize_peak_local(k, r, DELTAT, holed_near_the_crest)[0]) + assert np.isfinite(got) + assert pl.last_report()['n_peak_local_rows'] == n_ref, pl.last_report() + + +def test_the_localiser_reports_non_convergence_when_it_does_not_converge(): + """M3 (my own sweep). Dropping the convergence test from `localise_peaks` survived, + because the only test of the flag monkeypatched the localiser and so exercised the + CONSUMER, never the producer. Starve the iteration count instead: the real localiser + must then report not-converged, and the row must fall back rather than be accepted + on an unconverged crest.""" + sig = BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.4) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + old_iter, old_tol = pl.LOCALISE_MAX_ITER, pl.LOCALISE_SAFETY + try: + pl.LOCALISE_MAX_ITER = 1 + pl.LOCALISE_SAFETY = 1e-12 # unreachable in one step + got = float(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)[0]) + rep = pl.last_report() + finally: + pl.LOCALISE_MAX_ITER, pl.LOCALISE_SAFETY = old_iter, old_tol + assert rep['n_dense_fallback_localise'] == 1, rep + assert rep['n_peak_local_rows'] == 0, rep + assert got == _bandlimited(k[0]) + + +def test_the_report_sub_counts_reconcile(): + """R2-F9. An operator reading `last_report()` should not find an unexplained + residual: every declined row must appear in exactly one sub-count.""" + rows = [np.zeros(NPTS, dtype=complex), + BandLimited(amp=0.02, peak_sample=NPTS // 2 + 0.3125).samples(), + BandLimited(amp=0.5, peak_sample=NPTS // 2 + 0.3125).samples(), + BandLimited(amp=2000.0, peak_sample=NPTS // 2 + 0.3125).samples()] + k = np.stack(rows) + pl.time_marginalize_peak_local(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + r = pl.last_report() + subs = (r['n_dense_fallback_cost_pregate'] + r['n_dense_fallback_cost'] + + r['n_dense_fallback_tail'] + r['n_dense_fallback_structure'] + + r['n_dense_fallback_ceiling'] + r['n_dense_fallback_localise'] + + r['n_dense_fallback_containment'] + r['n_dense_fallback_nopeak']) + assert subs == r['n_dense_fallback_rows'], (subs, r) + assert r['n_peak_local_rows'] + r['n_dense_fallback_rows'] == r['n_refined_rows'], r + + +def _two_equal_kernels(H, off, m0=120.0): + """Two kernels of IDENTICAL height, one on an enumeration sample and one displaced + by `off` cells. The asymmetry is the point: it is the shape that three rounds of + review kept finding, and that fixtures moving both peaks by similar amounts cannot + produce. Built directly in the spectrum so the sharpness is set by H alone.""" + M = (NPTS - 1) // 2 + ms = np.arange(1, M + 1) + h = DELTAT / pl.PEAK_ENUM_FACTOR + e = np.exp(-0.5 * (ms / m0) ** 2) + S = 2 * e.sum() + c = 2.0 * e * ((H / S) * np.exp(-2j * np.pi * ms * (204 * DELTAT) / (NPTS * DELTAT)) + + (H / S) * np.exp(-2j * np.pi * ms + * (409 * DELTAT + off * h) / (NPTS * DELTAT))) + return (np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) @ c)[None, :] + + +def _exact_reference(kappa, refine=16384): + """The integral of the exact band-limited interpolant -- the same object the module + integrates -- at a refinement far beyond any factor the module will derive. + + REFLECTED, because that is what the module integrates since rift_O4d e4ed25c7: the + raw periodic interpolant is a different function near the window ends, and using it + here would hold peak-local to a reference its own dense fallback does not meet.""" + up = tmq.reflected_bandlimited_upsample(kappa, refine)[0, :(NPTS - 1) * refine + 1].real + m = up.max() + w = np.ones(up.size) + w[0] = w[-1] = 0.5 + return m + np.log(np.sum(np.exp(up - m) * w)) + np.log(DELTAT / refine) + + +@pytest.mark.parametrize("H,want_factor", [(160000.0, 1024), (2.56e6, 4096)]) +@pytest.mark.parametrize("off", [0.0, 0.25, 0.5]) +def test_the_keep_decision_is_exact_at_the_sharpest_legal_rows(H, want_factor, off): + """THE THIRD ROUND ON ONE CLASS, pinned at the sharp end of the LEGAL range. + + Comparing SAMPLES dropped a peak once `(delta/sigma)^2/2 > PEAK_KEEP_NATS`. The + parabolic vertex that replaced them bought about 8x in `sigma_t/deltaT` and then + failed the same way: its under-read of the crest at half-cell phase is 1.17 / 4.66 / + 18.66 / **74.63** nats at factors 512 / 1024 / 2048 / 4096, and + `UPSAMPLE_FACTOR_MAX = 4096` makes the last one LEGAL. End-to-end that was + **-0.693147** -- `-log 2`, one of two equal peaks deleted -- with the tail bound and + the containment check both silent. + + The lesson this test exists to pin: **every approximation substituted for the crest + fails the same way one octave further out.** The keep decision is now taken on + `lnL_star`, which is the crest by construction, so the reach of the fix is not a + function of sharpness at all -- which is exactly what this parametrisation checks. + """ + kappa = _two_equal_kernels(H, off) + rho = np.zeros_like(kappa.real) + sigma, _, _ = tmq.peak_width_from_lnL(_lnL(kappa.real, rho), DELTAT) + factor = int(tmq.required_upsample_factors(sigma, DELTAT)[0]) + assert factor == want_factor, (factor, want_factor, float(sigma[0]) / DELTAT) + + got = float(pl.time_marginalize_peak_local(kappa, rho, DELTAT, _lnL)[0]) + rep = pl.last_report() + assert abs(got - _exact_reference(kappa)) < 1e-3, (H, off, rep) + # and BOTH crests were kept -- the value alone cannot distinguish "both peaks" from + # "one peak plus a compensating error" + assert rep['n_peaks_total'] == 2, (H, off, rep) + assert rep['n_peak_local_rows'] == 1, (H, off, rep) + + +def test_an_endpoint_maximum_can_obtain_a_finite_width(): + """Door 2, asserted on the MECHANISM rather than on a value. + + Both crest estimators read `stencil[m-d]` and `stencil[m+d]`; with the centre left + at `cols_p == 0` the entire left half is out of range at every half-width, `d2` is + NaN throughout and `sigma = inf`, so the peak is dropped before anything else runs. + That made revision 2's endpoint enumeration DEAD CODE -- 22 endpoint maxima + enumerated, zero able to obtain a width -- and it is invisible to any test that only + checks the returned value, because the dense fallback supplies a correct answer. + + A value test is exactly what I wrote the first time, and it passed while the feature + did nothing. This one asks the estimator directly. + """ + n_enum = 64 + h = DELTAT / pl.PEAK_ENUM_FACTOR + # a clean quadratic maximum sitting ON index 0, decaying to the right + idx = np.arange(n_enum, dtype=float) + q = (-0.5 * (idx / 3.0) ** 2)[None, :] + + mask = pl.enumerate_peak_indices(q) + assert bool(mask[0, 0]), "an endpoint maximum must be enumerated at all" + + maxd = max(tmq.CURVATURE_STENCIL_HALFWIDTHS) + cols = np.array([0]) + c_st = np.clip(cols, 1, n_enum - 2) + take = c_st[:, None] + np.arange(-maxd, maxd + 1)[None, :] + valid = (take >= 0) & (take < n_enum) + st = np.where(valid, q[0][np.clip(take, 0, n_enum - 1)], np.nan) + sigma = pl._peak_curvature_sigma(st, h) + assert np.isfinite(sigma[0]) and sigma[0] > 0, ( + "an endpoint maximum still cannot obtain a finite width", sigma) + + +def test_the_integration_domain_is_exactly_the_analysis_window(): + """The clip to `[0, t_last]` is not cosmetic: the evaluator is PERIODIC, so an + interval running past either end returns values from the opposite end of the window, + wrapping mass from outside the integration domain into the answer. + + It fires for 138 of 2682 peaks on a realistic block -- I previously recorded it as + "no fixture reaches it", which was false -- and it is what makes this path's domain + exactly the dense path's. Asserted on the intervals themselves, not on a value: + the clipped region carries `e^-72`, so a value test cannot see it. + """ + t_last = (NPTS - 1) * DELTAT + # The clip fires for a SECONDARY peak near an edge in a row whose DOMINANT peak is + # mid-window. A row whose own argmax is near an edge is wrap-exposed and never + # reaches this code at all -- which is what a first version of this fixture produced, + # and the vacuity guard below caught it. Moderate sharpness, so W_SIGMA*sigma is a + # fraction of a sample rather than a thousandth of one. + # Measured, not guessed: these three reach the clip (`overhang = 1` each), while a + # secondary peak at `rel = 0.9` is dropped by the keep filter and a broad row is + # declined on cost, so neither reaches this code at all. + rows = [BandLimited(amp=5.0, peak_sample=NPTS // 2 + 0.3125, + extra_peaks=[(e, rel)]).samples() + for e, rel in ((0.2, 0.95), (0.5, 0.95), (0.5, 0.99))] + k = np.stack(rows) + _, peaks = pl.time_marginalize_peak_local( + k, np.full(k.shape, RHO_SQ), DELTAT, _lnL, return_peaks=True) + seen = 0 + for pk in peaks: + if pk is None: + continue + t_star, sigma = pk + half = pl.W_SIGMA * sigma + pl.LOCALISE_SAFETY * sigma + lo = np.maximum(t_star - half, 0.0) + hi = np.minimum(t_star + half, t_last) + assert np.all(lo >= 0.0) and np.all(hi <= t_last) + seen += int(np.sum((t_star - half < 0.0) | (t_star + half > t_last))) + assert seen > 0, "fixture no longer reaches the clip; the assertion is vacuous" + + +# --------------------------------------------------------------- the wiring + +N_BUFFER = 4096 + + +def _fake_likelihood_inputs(kappa_buffer): + """Minimal inputs that drive the SHIPPED NoLoop function on the numpy backend. + + One detector, one (l,m) pair and zero cross terms, so the self-term is constant + and kappa reduces to the supplied rholm buffer times a fixed response factor. + The point is to exercise the argument plumbing; the physics is covered above + against analytic truth. + """ + import lal + import RIFT.lalsimutils as lsu + + rholm = np.asarray(kappa_buffer, dtype=complex)[None, :] + P = lsu.ChooseWaveformParams() + P.deltaT = DELTAT + P.tref = 1000000000.0 + for name, val in [('phi', 0.0), ('theta', 0.0), ('phiref', 0.0), + ('incl', 0.0), ('psi', 0.0)]: + setattr(P, name, np.zeros(1) + val) + P.dist = np.full(1, fl.distMpcRef * 1e6 * lal.PC_SI) + return (P, {'H1': rholm}, {'H1': np.array([[2, 2]])}, + {'H1': np.zeros((1, 1), dtype=complex)}, {'H1': P.tref - 0.5}) + + +def _buffer_signal(amp, roll=0): + sig = BandLimited(amp=amp, peak_sample=NPTS // 2, n_period=N_BUFFER, m_hi=1400, + background=0.12) + return np.roll(sig.at(np.arange(N_BUFFER) * DELTAT), int(roll)) + + +def _shipped(tvals, args, **kw): + P, rholms, lookupNK, ct, epochs = args + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, Lmax=2, xpy=np, **kw) + + +def _tuned_inputs(tvals, sigma_target_over_dt=0.05): + """Build inputs whose lnL(t) is genuinely under-resolved, by MEASURING what the + shipped function produces rather than assuming it: the response factor and the + gather offset are the code's business, not the test's.""" + amp, roll = 1.0, 0 + for _ in range(8): + args = _fake_likelihood_inputs(_buffer_signal(amp, roll)) + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + sigma, jmax, _ = tmq.peak_width_from_lnL(lnL_t, DELTAT) + roll += int(NPTS // 2 - int(jmax[0])) + if np.isfinite(sigma[0]): + amp *= (float(sigma[0]) / (sigma_target_over_dt * DELTAT)) ** 2 + args = _fake_likelihood_inputs(_buffer_signal(amp, roll)) + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + sigma, jmax, _ = tmq.peak_width_from_lnL(lnL_t, DELTAT) + return args, float(sigma[0]) / DELTAT, int(jmax[0]) + + +def test_the_option_reaches_the_shipped_likelihood_and_changes_the_answer(): + """THE WIRING, NOT THE HELPER. + + A flag computed correctly and then never delivered to the likelihood is a + documented failure mode in this repo -- a whole comparison campaign has been run + against an inert option here before. So set the module default exactly the way + the driver sets it, call the SHIPPED function, and require the number to move. + """ + pytest.importorskip('RIFT.lalsimutils') + tvals = fl.marginalization_time_grid(0.075, DELTAT) + assert len(tvals) == NPTS + + args, sigma_over_dt, jmax = _tuned_inputs(tvals) + assert sigma_over_dt < 0.15, sigma_over_dt # under-resolved + guard = int(NPTS * tmq.EDGE_GUARD_FRACTION) + assert guard < jmax < NPTS - 1 - guard, jmax + + assert fl.TIME_QUADRATURE_DEFAULT == 'simpson', "default must not have moved" + base = float(np.asarray(_shipped(tvals, args))[0]) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'peak-local' # exactly what the driver does + new = float(np.asarray(_shipped(tvals, args))[0]) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + assert abs(new - base) > 1e-3, (base, new) + dense = float(np.asarray(_shipped(tvals, args, time_quadrature='bandlimited'))[0]) + assert abs(new - dense) < 1e-3, (new, dense) + + # AND THAT IT WAS THIS RULE THAT PRODUCED IT. The assertion above cannot tell: + # peak-local is DESIGNED to agree with bandlimited, so rewiring the branch in the + # shipped likelihood to call `time_marginalize_bandlimited` instead leaves every + # number in this file unchanged and the whole suite green -- while the entire cost + # benefit, the only reason this PR exists, silently disappears. That is exactly the + # "a comparison campaign has been run against an inert option here before" hazard. + # `_LAST_REPORT` is a module global that ONLY the peak-local module writes, so it + # distinguishes the two rules where the returned value provably cannot. + pl._LAST_REPORT.clear() + fl.TIME_QUADRATURE_DEFAULT = 'peak-local' + try: + _shipped(tvals, args) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + rep_ = pl.last_report() + assert rep_.get('n_peak_local_rows', 0) >= 1, ( + "the shipped branch did not reach time_marginalize_peak_local", rep_) + + kw = float(np.asarray(_shipped(tvals, args, time_quadrature='peak-local'))[0]) + assert kw == new + fl.TIME_QUADRATURE_DEFAULT = 'peak-local' + try: + assert float(np.asarray(_shipped(tvals, args, time_quadrature='simpson'))[0]) == base + finally: + fl.TIME_QUADRATURE_DEFAULT = old + + +def test_unsupported_combinations_refuse_rather_than_silently_using_simpson(): + pytest.importorskip('RIFT.lalsimutils') + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2) + tvals = fl.marginalization_time_grid(0.075, DELTAT) + P, rholms, lookupNK, ct, epochs = _fake_likelihood_inputs([sig.samples()]) + common = dict(Lmax=2, xpy=np, time_quadrature='peak-local') + for extra in ({'n_cal': 2}, {'return_lnLt': True}, {'return_cal_components': True}, + {'phase_marginalization': True}): + kw = dict(common) + kw.update(extra) + with pytest.raises(NotImplementedError): + fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, **kw) + + +@pytest.mark.parametrize("module_name,func_name", [ + ('RIFT.likelihood.factored_likelihood_with_rotation', + 'DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation'), + ('RIFT.likelihood.factored_likelihood_freqresponse', + 'DiscreteFactoredLogLikelihoodFreqResponseNoLoop'), +]) +def test_excluded_paths_refuse_the_global_default(module_name, func_name): + """These likelihoods have a time-DEPENDENT rho_sq, so neither refined quadrature + applies. Setting the option globally must make them RAISE, not quietly run + Simpson -- otherwise the exclusion is invisible at the point of use.""" + mod = pytest.importorskip(module_name) + func = getattr(mod, func_name) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'peak-local' + with pytest.raises(NotImplementedError): + func(None, None, None, None, None, None, None, None) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + + +def test_return_lnLt_still_works_when_the_module_default_is_peak_local(): + """The group's standard extrinsic stage (``--add-extrinsic + --add-extrinsic-time-resampling``) calls this function with ``return_lnLt=True`` + and no explicit quadrature. Raising on the INHERITED default there once broke + that stage after the whole integration had run.""" + pytest.importorskip('RIFT.lalsimutils') + tvals = fl.marginalization_time_grid(0.075, DELTAT) + args = _fake_likelihood_inputs(_buffer_signal(1.0)) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'peak-local' + got = np.asarray(_shipped(tvals, args, return_lnLt=True)) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + assert got.shape == (1, NPTS) + + +# ------------------------------------------------------------- the driver CLI + +def _run_driver(extra_args): + """Invoke the ILE driver in a SUBPROCESS and return (returncode, output). + + Deliberately a subprocess: the option's whole job is to travel from a command + line into the likelihood, and the guard that stops it being silently inert lives + in the driver's startup, not in the library -- so a test that imports the library + cannot see it. The driver exits long before any data is needed. + """ + import subprocess + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + driver = os.path.join(root, 'bin', 'integrate_likelihood_extrinsic_batchmode') + env = dict(os.environ) + env['PYTHONPATH'] = root + os.pathsep + env.get('PYTHONPATH', '') + env['OMP_NUM_THREADS'] = '1' + proc = subprocess.run([sys.executable, driver] + extra_args, env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=900) + return proc.returncode, proc.stdout.decode('utf-8', 'replace') + + +_HONOURED = ['--time-marginalization', '--vectorized', '--gpu', '--force-xpy'] + + +def _quadrature_banner(out): + """The quadrature banner line, matched SPECIFICALLY. + + The pre-existing ``--interpolate-time`` banner carries the identical phrase + "honoured by this configuration", so a bare substring test matches whichever line + happens to say what you were hoping for -- a mutation that made the quadrature + banner claim ``True`` unconditionally survived exactly that way once. Anchor the + identifying prefix and the value together, and require exactly one match. + """ + m = re.findall(r'^\s*Time-marginalization quadrature: (\S+) ' + r'\(from --time-marginalization-quadrature (.+?)\); ' + r'honoured by this configuration: (True|False)\s*$', + out, re.MULTILINE) + assert len(m) == 1, ("expected exactly one quadrature banner line, got %d:\n%s" + % (len(m), out[-3000:])) + return m[0][0], m[0][2] + + +def test_driver_announces_peak_local_and_actually_puts_it_in_force(): + """The banner prints the value READ BACK OUT of the module, so this assertion is + load-bearing: deleting the assignment leaves the flag inert AND changes the + printed line.""" + rc, out = _run_driver(['--time-marginalization-quadrature', 'peak-local'] + _HONOURED) + assert _quadrature_banner(out) == ('peak-local', 'True'), out[-2000:] + rc, out = _run_driver(_HONOURED) + assert _quadrature_banner(out) == ('simpson', 'True'), out[-2000:] + + +def test_driver_rejects_a_misspelled_peak_local(): + """A misspelled stencil name was once absorbed as "not truthy" and silently ran a + different likelihood here. A typo in this option has to be loud.""" + rc, out = _run_driver(['--time-marginalization-quadrature', 'peaklocal']) + assert rc != 0 + assert 'peaklocal' in out and 'peak-local' in out + + +def test_driver_refuses_configurations_that_cannot_honour_peak_local(): + """Refuse, do not ignore. Each of these would otherwise run the historical + Simpson quadrature while the banner said otherwise.""" + for missing in ([], ['--time-marginalization'], + ['--time-marginalization', '--vectorized'], + _HONOURED + ['--rotation-slow'], + _HONOURED + ['--freqresponse']): + rc, out = _run_driver(['--time-marginalization-quadrature', 'peak-local'] + missing) + assert rc != 0, (missing, out[-2000:]) + assert 'cannot honour it' in out, (missing, out[-2000:]) + + +def _phase_marg_lookup(tmp_path, value): + """A minimal distance-marginalization lookup table carrying only the key the + startup guard reads. Phase marginalization is NOT a CLI boolean -- it is a + property of that table -- so a guard that checked an ``opts`` attribute would + silently never fire.""" + path = str(tmp_path / ('lookup_%s.npz' % value)) + np.savez(path, phase_marginalization=np.array(bool(value))) + return path + + +def test_driver_refuses_peak_local_with_time_resampling(): + """P1. BOTH available answers here are silently wrong, which is why this refuses. + + `--time-posterior-export grid` draws `t_ref` from the original coarse `lnLt` bins. + peak-local exists to resolve a peak far narrower than that spacing, so the integral would + be sub-sample accurate and the exported time quantised to the very grid the option was + introduced to escape -- computed, then thrown away. And with `--interpolate-time nearest` + the `auto` mode resolves to `grid`, so that is the DEFAULT outcome, not a corner. + + `continuous` is not available either: it needs `return_time_draw`, which needs a validated + dense reconstruction over the whole window, and peak-local by construction never forms one. + The library refuses that combination for the same reason. + + `bandlimited` mandates the continuous export and is one flag away, so the refusal costs a + user nothing but a corrected command line. + """ + # --fairdraw-extrinsic-output is required by a PRE-EXISTING guard on + # --resample-time-marginalization, so without it the driver exits for an unrelated reason + # and this test would pass while proving nothing. + resample = ['--resample-time-marginalization', '--fairdraw-extrinsic-output'] + rc, out = _run_driver(['--time-marginalization-quadrature', 'peak-local'] + + resample + _HONOURED) + assert rc != 0, out[-2000:] + assert 'does not support --resample-time-marginalization' in out, out[-3000:] + # and the same configuration under `bandlimited` must NOT be refused for this reason + rc2, out2 = _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + + resample + _HONOURED) + assert 'does not support --resample-time-marginalization' not in out2, out2[-2000:] + + +def test_driver_refuses_peak_local_under_phase_marginalization_AT_STARTUP(tmp_path): + """Refused before the run, not deep inside it. + + The library refuses this too, but by then the integration is under way -- and + raising late is its own bug on this option: an over-broad ``return_lnLt`` guard + once let the standard extrinsic stage run the ENTIRE integration and then die at + the export step. So the guard has to be at startup, and this test invokes the + driver to prove it is. + + The three controls are the point: the same table with ``bandlimited`` must be + HONOURED (that path supports phase marginalization and must not regress), and + ``peak-local`` without phase marginalization must be honoured too. A guard that + simply refused whenever a lookup table was present would pass a bare refusal test + and fail all three. + """ + with_phase = _phase_marg_lookup(tmp_path, True) + without = _phase_marg_lookup(tmp_path, False) + dm = ['--distance-marginalization', '--distance-marginalization-lookup-table'] + + rc, out = _run_driver(['--time-marginalization-quadrature', 'peak-local'] + + _HONOURED + dm + [with_phase]) + assert rc != 0, out[-2000:] + assert 'cannot honour it' in out and 'phase marginalization' in out, out[-2000:] + + rc, out = _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + + _HONOURED + dm + [with_phase]) + assert _quadrature_banner(out) == ('bandlimited', 'True'), out[-2000:] + + rc, out = _run_driver(['--time-marginalization-quadrature', 'peak-local'] + + _HONOURED + dm + [without]) + assert _quadrature_banner(out) == ('peak-local', 'True'), out[-2000:] + + +# --------------------------------------------------------------- GPU parity + +def _cupy_or_skip(): + if os.environ.get('RIFT_CI_REQUIRE_GPU', '0') != '1': + cupy = pytest.importorskip('cupy') + else: + import cupy + try: + cupy.zeros(1) + 1 + except Exception as e: # pragma: no cover + if os.environ.get('RIFT_CI_REQUIRE_GPU', '0') == '1': + raise + pytest.skip("cupy present but no usable device: %s" % e) + return cupy + + +def test_peak_local_runs_on_the_gpu_backend_and_matches_numpy(): + """xpy-generic code that has never run on a device is broken until proven + otherwise. This is the whole path -- FFT upsample, enumeration, the ragged + host-side merge, the batched local evaluation and the tail bound -- on cupy, + against the numpy answer on identical inputs.""" + cupy = _cupy_or_skip() + from RIFT.likelihood import optimized_gpu_tools + + rows = [BandLimited(amp=a, peak_sample=NPTS // 2 + 0.25).samples() + for a in (40.0, 200.0, 2000.0)] + rows.append(np.zeros(NPTS, dtype=complex)) + k = np.stack(rows) + r = np.full(k.shape, RHO_SQ) + + cpu = np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL)) + gpu = cupy.asnumpy(pl.time_marginalize_peak_local( + cupy.asarray(k), cupy.asarray(r), DELTAT, _lnL, + simps=optimized_gpu_tools.simps, xpy=cupy)) + fin = np.isfinite(cpu) + assert np.max(np.abs(gpu[fin] - cpu[fin])) < 1e-6, (cpu, gpu) + + +# ---------------------------------------------------------------- round 6 (door 4) + +def _skewed_two_peak_row(HA, HB, npts=NPTS, phase=0.65625): + """The fixture that broke the parabolic pre-filter. + + Peak A is symmetric and sets the row's derived factor. Peak B is SKEWED -- a kernel + plus 0.8x the same kernel shifted by two coarse samples -- and sits at a sub-cell + phase chosen to maximise its crest-to-sample displacement (d = 0.5007 h_enum). The + skew is what matters: a symmetric peak's crest deficit is nearly quadratic and a + parabolic correction almost covers it, while a skewed peak's is ANHARMONIC and no + parabolic correction covers it at any constant. + """ + T = npts * DELTAT + ms = np.arange(1, (npts - 1) // 2 + 1) + h_enum = DELTAT / pl.PEAK_ENUM_FACTOR + ea = np.exp(-0.5 * (ms / 40.0) ** 2) + eb = np.exp(-0.5 * (ms / 120.0) ** 2) + tb = 400 * DELTAT + phase * h_enum + c = (HA * ea * np.exp(-2j * np.pi * ms * (150 * DELTAT) / T) + + HB * eb * (np.exp(-2j * np.pi * ms * tb / T) + + 0.8 * np.exp(-2j * np.pi * ms * (tb + 2.0 * DELTAT) / T))) + return (np.exp(2j * np.pi * np.outer(np.arange(npts), ms) / npts) @ c)[None, :] + + +def _crest_pair(k): + """The CREST values of peak A (near coarse sample 150) and peak B (near 400). + + Identified BY POSITION, not by rank: which of the two is on top is exactly what the + bisection below varies, so a rank-ordered pair silently swaps under it. + + LOCALISED, not sampled -- and the first version of this helper was not, which is the + same mistake this whole file is about. Peak A sits exactly on an enumeration sample + so its sample IS its crest; peak B is deliberately off-grid, and in the sharp regime + its sample understates its crest by ~3125 nats. Equalising the SAMPLES therefore + leaves B thousands of nats above A, A is correctly dropped, and the fixture tests + nothing while appearing to. + """ + F = pl.PEAK_ENUM_FACTOR + h_enum = DELTAT / F + q = tmq.reflected_bandlimited_upsample(k, F)[..., :(NPTS - 1) * F + 1].real + rp, cp = np.where(pl.enumerate_peak_indices(q)) + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + period_ref = 2.0 * NPTS * DELTAT + _, q_star, _ = pl.localise_peaks( + Xw, fk, rp, cp * h_enum, h_enum, np.full(rp.size, 1e-14 * NPTS * DELTAT), + period_ref, t_last=(NPTS - 1) * DELTAT) + out = [] + for centre in (150 * F, 400 * F): + near = np.abs(cp - centre) <= 3 * F + assert near.any(), (centre, cp[:20]) + out.append(float(np.asarray(q_star)[near].max())) + return out[0], out[1] + + +def _equalise_crests(scale, tol=1.0): + """Solve for the B amplitude that makes the two crests equal to within `tol` nats. + + Bisected rather than hard-coded: the two peaks OVERLAP, so a fixed amplitude ratio + equalises them at one scale only -- the crest gap is linear in amplitude, so a ratio + tuned at rho ~ 40 leaves the second peak thousands of nats down at rho ~ 700 and the + fixture silently stops testing anything. `crestA - crestB` is monotone decreasing in + the B amplitude, which is what makes a plain bisection valid. + """ + lo, hi = 1e-4 * scale, 10.0 * scale + mid = hi + for _ in range(80): + mid = 0.5 * (lo + hi) + a, b = _crest_pair(_skewed_two_peak_row(scale, mid)) + if abs(a - b) < tol: + return mid + if a > b: + lo = mid # B too small + else: + hi = mid + raise AssertionError("could not equalise the two crests at scale %g" % scale) + + +def test_the_prefilter_bound_is_actually_a_bound(): + """DOOR 4, the property. The pre-filter may only drop a peak against a bound that + HOLDS. The previous one -- `lnL_sample + (h_enum/2)**2 / (2 sigma**2)` -- did not: + it was violated on 0.3-3% of enumerated peaks, by up to 19456 nats, growing linearly + with amplitude, because `lnL` is not a parabola across a half enumeration cell and + the ANHARMONIC part of the deficit carries the same 1/sigma**2 amplification. At + derived factor 1024 the pure quantisation excess is 4.4 nats and the true shortfall + is 122.3. + + The replacement is a Taylor remainder with a TRUE bound on the second derivative: + expanding about the crest, where q' vanishes, q(t*) <= q(t_s) + max|q''| h_enum**2/2, + and max|q''| <= sum_j |Xw_j| |w_j|**2 by the triangle inequality on the spectral sum. + Nothing is fitted, so there is no model error to amplify. + + Asserted as a PROPERTY over every peak, not on one fixture: a bound that holds on the + fixture you thought of is what has failed here four times. + """ + rng = np.random.default_rng(6060) + for npts in (153, 307, 614): + for amp in (2.0e4, 2.0e6): + T = npts * DELTAT + ms = np.arange(1, (npts - 1) // 2 + 1) + env = np.exp(-0.5 * (ms / max(npts / 6.0, 20.0)) ** 2) + c = np.zeros(ms.size, dtype=complex) + for tau, a in zip(rng.uniform(0.05, 0.95, 4) * T, + amp * rng.uniform(0.6, 1.0, 4)): + c = c + 2.0 * env * (a / (2 * env.sum())) * np.exp(-2j * np.pi * ms * tau / T) + kap = (np.exp(2j * np.pi * np.outer(np.arange(npts), ms) / npts) @ c)[None, :] + + F = pl.PEAK_ENUM_FACTOR + h_enum = DELTAT / F + q_up = tmq.reflected_bandlimited_upsample( + kap, F)[..., :(npts - 1) * F + 1].real + rp, cp = np.where(pl.enumerate_peak_indices(q_up)) + assert rp.size > 10, (npts, amp, rp.size) + + kref = np.concatenate((kap, np.flip(kap, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + period_ref = 2.0 * npts * DELTAT + q_ddot_max = pl.spectral_curvature_bound(Xw, fk, period_ref) + upper = _lnL( + pl.crest_upper_bound(q_up[rp, cp], q_ddot_max[rp], h_enum), RHO_SQ) + + # the TRUE crest, on the same interpolant the module localises on + _, q_star, _ = pl.localise_peaks( + Xw, fk, rp, cp * h_enum, h_enum, np.full(rp.size, 1e-14 * T), + period_ref, t_last=(npts - 1) * DELTAT) + true_crest = _lnL(q_star, RHO_SQ) + short = float((true_crest - upper).max()) + assert short <= 0.0, ( + "the pre-filter bound is not a bound", npts, amp, short) + + +def test_a_codominant_crest_is_not_deleted_by_the_prefilter(): + """DOOR 4, end to end, on the shipped entry point. + + Two crests equal to within 1e-4 nats, so neither may be dropped -- both are far + inside PEAK_KEEP_NATS. The old pre-filter deleted the skewed one BEFORE localisation, + so the exact filter never saw it, and the answer came back **-0.358 nats** low, + ACCEPTED, with the tail bound and the containment check both silent. The tail bound + cannot backstop this: `q_out_max` reads the deleted peak at its SAMPLE, the very + quantity the defect corrupts. + + Reference-free: compared against the rule this one delegates to, so no interpolant + reference can be argued with. + """ + checked = 0 + for scale in (2.0e4, 4.0e4, 8.0e4): + HB = _equalise_crests(scale) + k = _skewed_two_peak_row(scale, HB) + # both crests inside PEAK_KEEP_NATS, so NEITHER may be dropped -- that is the + # precondition the whole test rests on, so assert it rather than assume it + pair = _crest_pair(k) + assert abs(pair[0] - pair[1]) < pl.PEAK_KEEP_NATS, (scale, pair) + r = np.full(k.shape, RHO_SQ) + try: + want = float(np.asarray( + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL))[0]) + except RuntimeError: + # the dense path declines this row on the ceiling, so it cannot serve as the + # reference here. Skipping it is honest; asserting against a value the + # reference implementation refuses to produce is not. + continue + got = float(np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL))[0]) + rep = pl.last_report() + assert abs(got - want) < 1e-3, (scale, got, want, rep) + # and the peak SURVIVES THE FILTER rather than the row being rescued by a + # fallback: the defect was a deletion, so a clean-up afterwards is not the fix + assert rep['n_peaks_total'] >= 2, rep + checked += 1 + assert checked >= 2, "fixture no longer exercises the defect at any scale" + + +def test_the_peak_sample_is_read_at_the_enumerated_index(): + """The same defect at a third site. The pre-filter's lower bound used to be + `lnL_st[:, maxd]`, the callback at the stencil CENTRE -- which is clipped inward by + one sample at the array ends, so at enumeration index 0 or n-1 it was a full + enumeration cell away from the peak it claimed to describe. Measured 132 nats low at + rho ~ 40 and 8449 nats low at rho ~ 700, growing with SNR. + + A stencil centre is clipped so the CURVATURE can be measured; the peak's own value + must still be read where the peak is. + """ + F = pl.PEAK_ENUM_FACTOR + maxd = max(pl.CURVATURE_STENCIL_HALFWIDTHS) + for amp in (4.0e4, 2.56e6): + sig = BandLimited(amp=amp, peak_sample=0.05) + k = sig.samples()[None, :] + q_up = tmq.reflected_bandlimited_upsample( + k, F)[..., :(NPTS - 1) * F + 1].real + rp, cp = np.where(pl.enumerate_peak_indices(q_up)) + edge = cp == 0 + if not edge.any(): + continue + at_index = q_up[rp[edge], cp[edge]] + at_clipped = q_up[rp[edge], np.clip(cp[edge], 1, q_up.shape[-1] - 2)] + # the two differ materially, and by more as the amplitude grows, which is why + # reading the wrong one mattered and why it got worse with SNR + gap = float(np.max(at_index - at_clipped)) + assert gap > 1.0, (amp, at_index, at_clipped) + + # The property that matters: the bound the pre-filter compares against must still + # hold AT AN ENDPOINT, which it does only if the sample is read at the enumerated + # index. Read at the clipped centre it is short by `gap`, and `gap` grows with + # amplitude without limit. + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + period_ref = 2.0 * NPTS * DELTAT + h_enum = DELTAT / F + q_ddot_max = pl.spectral_curvature_bound(Xw, fk, period_ref) + _, q_star, _ = pl.localise_peaks( + Xw, fk, rp[edge], cp[edge] * h_enum, h_enum, + np.full(int(edge.sum()), 1e-14 * NPTS * DELTAT), period_ref, + t_last=(NPTS - 1) * DELTAT) + good = pl.crest_upper_bound(at_index, q_ddot_max[rp[edge]], h_enum) + bad = pl.crest_upper_bound(at_clipped, q_ddot_max[rp[edge]], h_enum) + assert np.all(good >= q_star - 1e-6), (amp, float((q_star - good).max())) + assert np.any(bad < q_star - 1e-6), ( + amp, "the clipped read no longer breaks the bound; fixture is stale") + + +def test_a_boundary_pinned_crest_is_not_integrated_locally(): + """DOOR 5. A crest pinned at a window end is not a PEAK, and treating it as one cost + up to +4.38 nats, accepted, growing without bound in amplitude. + + At an interior crest `q'` vanishes, so `exp(lnL)` is locally Gaussian and a spacing + derived from the curvature resolves it. At a boundary the maximum is a CORNER: `q'` + is non-zero, the local integrand is an exponential decay of rate `|q'|`, and a spacing + derived from `sigma` does not resolve that scale at all. The trapezoid's half-weight + endpoint then over-counts by `log(lam*(0.5 + 1/(exp(lam)-1))) -> log(lam/2)` with + `lam = |q'| h_loc`, i.e. **+log 2 per factor 4 in amplitude**: measured +1.27 / +1.96 + / +2.66 / +3.35 / +4.03 nats, all ACCEPTED. + + Neither a-posteriori defence can catch it, and that is structural rather than bad + luck: the containment check compares the grid's attained maximum against `row_star`, + and the grid's FIRST POINT is the pinned crest, so `attained == row_star` identically; + the tail bound is a statement about mass OUTSIDE the intervals, and this error is + entirely inside. So the fix has to be a refusal, not another check. + + Asserted against the rule this one delegates to, so no reference can be argued with -- + and the property asserted is the contract: never a worse value than the backstop. + """ + T = NPTS * DELTAT + ms = np.arange(1, (NPTS - 1) // 2 + 1) + basis = np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) + env = np.exp(-0.5 * (ms / 40.0) ** 2) + for amp in (4.0e4, 6.4e5, 1.0e7): + # a single band-limited bump centred exactly on t = 0 + k = (basis @ (2.0 * (amp / (2 * env.sum())) * env))[None, :] + r = np.full(k.shape, RHO_SQ) + got = float(np.asarray(pl.time_marginalize_peak_local(k, r, DELTAT, _lnL))[0]) + rep = pl.last_report() + want = float(np.asarray( + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL))[0]) + assert abs(got - want) < 1e-3, (amp, got, want, rep) + # and it is DECLINED rather than accepted-and-happening-to-agree + assert rep['n_peak_local_rows'] == 0, (amp, rep) + assert rep['n_dense_fallback_rows'] == 1, (amp, rep) + + +def test_the_localiser_reports_a_pinned_peak_as_unconverged(): + """The mechanism behind DOOR 5, pinned at the unit level so the end-to-end test above + cannot start passing for an unrelated reason. `localise_peaks` once counted a peak + pinned against a window end as CONVERGED, on the grounds that the maximum over the + domain really is the boundary. It is -- but it is not a stationary point, and every + consumer of the result assumes one.""" + T = NPTS * DELTAT + ms = np.arange(1, (NPTS - 1) // 2 + 1) + env = np.exp(-0.5 * (ms / 40.0) ** 2) + k = (np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) + @ (2.0 * (4.0e4 / (2 * env.sum())) * env))[None, :] + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + h_enum = DELTAT / pl.PEAK_ENUM_FACTOR + rows = np.zeros(1, dtype=np.int64) + t_star, _, ok = pl.localise_peaks( + Xw, fk, rows, np.zeros(1), h_enum, np.full(1, 1e-14 * T), + 2.0 * NPTS * DELTAT, t_last=(NPTS - 1) * DELTAT) + assert float(np.asarray(t_star)[0]) == 0.0, t_star + assert not bool(np.asarray(ok)[0]), "a pinned crest must not report as converged" + + +def test_boundary_unresolved_rows_are_refined_and_not_called_flat(): + """A COVERAGE GAP the mutation sweep found, closed. + + `_classify_rows` seeds a factor for a row whose maximum sits at the first or last + sample: the centred curvature stencil is clipped inward there, so a severely + under-resolved endpoint peak reads POSITIVE curvature and is mislabelled "flat", which + would silently retain Simpson for exactly the truncated-boundary case the option exists + to reconstruct. + + Deleting that seed left the whole 109-test suite green while moving an endpoint row by + **6.06 nats**. It survived `test_row_classification_matches_the_dense_path_exactly` + because both rules share `_classify_rows`, so removing the clause moves them TOGETHER + and parity still holds -- a parity test cannot see a change made to the definition both + sides read. So this asserts the classification itself, not the agreement. + """ + T = NPTS * DELTAT + ms = np.arange(1, (NPTS - 1) // 2 + 1) + # m0 = 120 is the regime where the inward-clipped coarse stencil actually reads + # positive curvature; at m0 = 40 the endpoint peak is wide enough that the stencil + # measures it and this clause never fires, so that fixture would test nothing. + env = np.exp(-0.5 * (ms / 120.0) ** 2) + basis = np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) + for centre in (0.0, 0.4 * DELTAT / pl.PEAK_ENUM_FACTOR): + c = 2.0 * (4.0e4 / (2 * env.sum())) * env * np.exp(-2j * np.pi * ms * centre / T) + k = (basis @ c)[None, :] + r = np.full(k.shape, RHO_SQ) + lnL_coarse = _lnL(k.real, r) + (sigma, jmax, measurable, has_peak, flat, exposed, unmeasurable, + factors) = tmq._classify_rows(lnL_coarse, DELTAT, NPTS) + assert int(jmax[0]) in (0, NPTS - 1), jmax + # the inward-clipped stencil genuinely cannot measure it ... + assert not np.isfinite(sigma[0]), (centre, sigma) + # ... and precisely because of that it must NOT be called flat, and must be refined + assert not bool(flat[0]), (centre, "an endpoint peak was mislabelled flat") + assert bool(has_peak[0]), (centre, has_peak) + assert int(factors[0]) >= 4, (centre, factors) + + +def test_the_outside_supremum_is_evaluated_off_grid_not_sampled(): + """P1. `max(q_up over uncovered SAMPLES)` is a lower bound on the continuous outside + supremum, and the gap GROWS WITH AMPLITUDE, so the reported margin got more flattering the + sharper the row -- the wrong direction for a safety number. Measured before this was + fixed: the module reported -79.6 / -289.5 / -2998.6 at amplitude 2e4 / 2e5 / 2e6 where the + honest margins were -65.5 / -71.1 / -75.8, an under-read of 14, 218 and 2923 nats. + + The gap is NOT a peak the enumeration missed. The supremum over a union of closed + intervals sits at an interior stationary point or at an END; the ends dominate. An + interval end is `W_SIGMA * sigma` from its crest -- 72 nats below it at any amplitude -- + but the nearest SAMPLE outside that end is a further `W_SIGMA * h_enum / sigma` down, + which diverges as sigma shrinks. That one term was the entire under-read. + + So the ends are evaluated on the double-copy Fourier model rather than sampled near. This + asserts the property that fixes: the reported margin must track an independently computed + honest supremum, and must not sit hundreds of nats below it. + """ + T = NPTS * DELTAT + ms = np.arange(1, (NPTS - 1) // 2 + 1) + basis = np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) + env = np.exp(-0.5 * (ms / 120.0) ** 2) + rng = np.random.default_rng(4242) + checked = 0 + for amp in (2.0e4, 2.0e6): + c = np.zeros(ms.size, dtype=complex) + for _ in range(3): + c = c + 2.0 * env * (amp * rng.uniform(0.4, 1.0) / (2 * env.sum())) * \ + np.exp(-2j * np.pi * ms * rng.uniform(0.05, 0.95)) + k = (basis @ c)[None, :] + r = np.full(k.shape, RHO_SQ) + v, peaks = pl.time_marginalize_peak_local(k, r, DELTAT, _lnL, return_peaks=True) + rep = pl.last_report() + if rep['n_peak_local_rows'] != 1 or peaks[0] is None: + continue + t_star, sigma = peaks[0] + half = pl.W_SIGMA * sigma + pl.LOCALISE_SAFETY * sigma + t_last = (NPTS - 1) * DELTAT + lo = np.maximum(t_star - half, 0.0) + hi = np.minimum(t_star + half, t_last) + + # honest supremum over the uncovered set, on a grid far finer than the enumeration one + FF = 32 * pl.PEAK_ENUM_FACTOR + tf = np.arange((NPTS - 1) * FF + 1) * (DELTAT / FF) + qf = tmq.reflected_bandlimited_upsample(k, FF)[0, :(NPTS - 1) * FF + 1].real + outside = np.ones(tf.size, dtype=bool) + for a, b in zip(lo, hi): + outside &= ~((tf >= a) & (tf <= b)) + assert outside.any(), "fixture leaves nothing outside" + T_out = max(t_last - float(np.sum(hi - lo)), 0.0) + honest = np.log(T_out) + _lnL(qf[outside].max(), RHO_SQ) - float(np.asarray(v)[0]) + + got = float(rep['tail_bound_worst']) + # the reported margin must be CLOSE to honest, and never far BELOW it (which is the + # optimistic direction that grew with amplitude) + assert got - honest > -30.0, (amp, got, honest, "reported margin far below honest") + assert got < pl.TAIL_LOG_TOL, (amp, got) + checked += 1 + assert checked == 2, checked + + +def test_the_tail_margin_on_a_clean_row_is_a_design_constant_not_a_measurement(): + """What the tail bound is actually saying on a single-peak row: almost nothing about it. + + The interval end is `(W_SIGMA + LOCALISE_SAFETY) * sigma` from the crest and the Laplace + value of the integral is `lnL_crest + log(sqrt(2 pi) sigma)`, so + + margin = log(T_out / (sqrt(2 pi) sigma)) - (W_SIGMA + LOCALISE_SAFETY)**2 / 2 + + where the second term is 75.03 nats of pure design choice and the first moves by only + `log 10` per decade of amplitude. Measured, the reported margin tracks that to under a + nat over two decades. + + This is pinned because it is easy to read `tail_bound_worst = -64` as a per-row safety + margin, and on a clean row it is not one -- it is the arithmetic of `W_SIGMA` restated, + and it cannot fail whatever the data. It also explains, rather than merely observes, the + "structural slack" two independent reviewers probed and could not eat into: sampling was + never what set it. The bound earns its place on rows with structure OUTSIDE the intervals, + where this closed form does not apply and it does fire. + """ + T = NPTS * DELTAT + ms = np.arange(1, (NPTS - 1) // 2 + 1) + basis = np.exp(2j * np.pi * np.outer(np.arange(NPTS), ms) / NPTS) + env = np.exp(-0.5 * (ms / 120.0) ** 2) + t_last = (NPTS - 1) * DELTAT + cut = (pl.W_SIGMA + pl.LOCALISE_SAFETY) ** 2 / 2.0 + checked = 0 + for amp in (2.0e4, 2.0e6): + rng = np.random.default_rng(11) + c = np.zeros(ms.size, dtype=complex) + for _ in range(3): + c = c + 2.0 * env * (amp * rng.uniform(0.4, 1.0) / (2 * env.sum())) * \ + np.exp(-2j * np.pi * ms * rng.uniform(0.05, 0.95)) + k = (basis @ c)[None, :] + r = np.full(k.shape, RHO_SQ) + _, peaks = pl.time_marginalize_peak_local(k, r, DELTAT, _lnL, return_peaks=True) + rep = pl.last_report() + if rep['n_peak_local_rows'] != 1 or peaks[0] is None: + continue + t_star, sigma = peaks[0] + half = (pl.W_SIGMA + pl.LOCALISE_SAFETY) * sigma + lo = np.maximum(t_star - half, 0.0) + hi = np.minimum(t_star + half, t_last) + T_out = max(t_last - float(np.sum(hi - lo)), 0.0) + predicted = np.log(T_out / (np.sqrt(2 * np.pi) * float(np.min(sigma)))) - cut + # The design constant still sets the SCALE. What sits on top of it is the + # certificate's own remainder -- the outside supremum is now bounded, not evaluated -- + # so the margin is `predicted` plus something non-negative and no larger than that + # remainder. Both sides are asserted: the floor is the design constant, the ceiling + # is the remainder, and neither is free to drift. + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + m4 = float(pl.spectral_derivative_bound(Xw, fk, 2.0 * NPTS * DELTAT, 4)[0]) + remainder = m4 * (DELTAT / pl.PEAK_ENUM_FACTOR) ** 4 / 384.0 + got = float(rep['tail_bound_worst']) + assert got >= predicted - 2.0, (amp, got, predicted) + assert got <= predicted + remainder + 2.0, (amp, got, predicted, remainder) + checked += 1 + assert checked == 2, checked + + +def test_every_exported_name_is_actually_defined(): + """P2. `__all__` gained `parabolic_sup` and `segment_sup_bound`, neither of which was + ever defined, so `from RIFT.likelihood.time_marginalization_peak_local import *` raised + `AttributeError` -- the module was unimportable by star-import at that head. + + Guarded here rather than just fixed, because the failure is invisible to every other test + in this file: they all use `import module as pl`, which does not consult `__all__` at all. + Only a star-import does, so only a star-import can catch it. Both quadrature modules are + checked, since the two are edited together and the same slip is available in either. + """ + for mod in (pl, tmq): + missing = [n for n in mod.__all__ if not hasattr(mod, n)] + assert not missing, (mod.__name__, "exported but not defined", missing) + # and the operation that actually broke + ns = {} + exec("from %s import *" % mod.__name__, ns) + + +def _reflected_spectrum(k): + kref = np.concatenate((k, np.flip(k, axis=-1)), axis=-1) + Xw, fk = pl.bandlimited_spectrum(kref) + return kref, Xw, fk, 2.0 * NPTS * DELTAT + + +def test_enum_grid_derivatives_agrees_with_pointwise_evaluation(): + """This helper ARRIVED WRONG and uncalled, so it is tested against an independent route. + + Its first version reimplemented the zero-padding convention and was off by **52% on q' + and 46% on q''**. The cause is worth keeping: for the reflected row `n` is EVEN, so + `bandlimited_spectrum` splits the Nyquist bin and returns `n+1` coefficients -- which that + version used as the transform LENGTH, so the grid was wrong and the two split bins collided + on one index. It now differentiates the row spectrally and reuses the validated + upsampler, which cannot drift from the grid `q` is taken on because it IS that grid. + """ + sig = BandLimited(amp=4.0e4, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples()[None, :] + kref, Xw, fk, period = _reflected_spectrum(k) + F = pl.PEAK_ENUM_FACTOR + n_keep = (NPTS - 1) * F + 1 + dq, ddq = pl.enum_grid_derivatives(kref, F, n_keep, DELTAT) + t = np.arange(n_keep) * (DELTAT / F) + rows = np.zeros(n_keep, dtype=np.int64) + _, dq_ref, ddq_ref = pl.eval_bandlimited_points(Xw, fk, rows, t, period) + assert np.max(np.abs(dq[0] - dq_ref)) < 1e-10 * np.max(np.abs(dq_ref)), "q' disagrees" + assert np.max(np.abs(ddq[0] - ddq_ref)) < 1e-10 * np.max(np.abs(ddq_ref)), "q'' disagrees" + + +def test_parabolic_sup_never_under_bounds_a_cubic(): + """The cell maximum must be an upper bound or the certificate built on it is worthless. + + Fuzzed rather than spot-checked, because the failure found during development was a + DEGENERATE case, not a generic one: when the cubic coefficient vanishes -- which a + symmetric bump does exactly -- the cubic root formula divides by zero, the interior + stationary point is skipped, and the endpoint maximum is returned. That under-bounds + precisely the cells that contain a peak. + """ + rng = np.random.default_rng(0) + s_grid = np.linspace(0.0, 1.0, 2001) + y0, y1, d0, d1 = (rng.normal(size=4000) * 10.0 for _ in range(4)) + got = pl.parabolic_sup(y0, y1, d0, d1) + a = 2 * y0 + d0 - 2 * y1 + d1 + b = -3 * y0 - 2 * d0 + 3 * y1 - d1 + ref = np.max(y0[:, None] + d0[:, None] * s_grid + b[:, None] * s_grid ** 2 + + a[:, None] * s_grid ** 3, axis=1) + assert np.all(got >= ref - 1e-8), float(np.max(ref - got)) + # and the degenerate case explicitly: y0=y1=0, d0=-d1=4 has max 1.0 at s=1/2 + assert abs(float(pl.parabolic_sup(np.array([0.0]), np.array([0.0]), + np.array([4.0]), np.array([-4.0]))[0]) - 1.0) < 1e-12 + + +def test_parabolic_sup_handles_a_nearly_quadratic_symmetric_crest(): + """The stable root path must survive roundoff in an analytically symmetric cell. + + For this admissible band-limited cosine, the Hermite cubic coefficient is analytically + zero but endpoint arithmetic leaves it tiny and nonzero. The direct quadratic formula + placed the in-cell root at 2/3 instead of 1/2; even after adding the exact fourth-derivative + remainder, the purported certificate under-read the true crest by 4255 nats at this + amplitude. + """ + amp = 2.0e6 + omega = np.pi / 8.0 + y0 = amp * np.cos(-omega / 2.0) + y1 = amp * np.cos(omega / 2.0) + d0 = -amp * omega * np.sin(-omega / 2.0) + d1 = -amp * omega * np.sin(omega / 2.0) + + cubic_sup = float(pl.parabolic_sup( + np.array([y0]), np.array([y1]), np.array([d0]), np.array([d1]))[0]) + certificate = cubic_sup + amp * omega ** 4 / 384.0 + assert certificate >= amp, (certificate, amp) + + +def test_parabolic_sup_scales_finite_coefficients_before_the_discriminant(): + """Finite derivative coefficients must not overflow the discriminant to ``NaN``.""" + # H'(s) = scale * (s - 0.2) * (s - 0.8), whose local maximum is at s=0.2. + scale = 1.0e200 + A = scale + a = A / 3.0 + b = -A / 2.0 + c = 0.16 * A + y0 = 0.0 + y1 = c + b + a + d0 = c + d1 = 3.0 * a + 2.0 * b + c + true_max = y0 + c * 0.2 + b * 0.2 ** 2 + a * 0.2 ** 3 + + got = float(pl.parabolic_sup( + np.array([y0]), np.array([y1]), np.array([d0]), np.array([d1]))[0]) + assert got >= true_max, (got, true_max) + + +def test_certificate_fallback_reasons_are_disjoint_and_exhaustive(): + """A structural failure must not also increment the tail or containment counter.""" + margin = np.array([pl.TAIL_LOG_TOL + 1.0, pl.TAIL_LOG_TOL + 1.0, + pl.TAIL_LOG_TOL - 1.0, pl.TAIL_LOG_TOL - 1.0]) + contained = np.array([False, True, False, True]) + cert_bad = np.array([True, False, False, False]) + planned = np.arange(4) + good, structure, tail, containment = pl._certificate_acceptance_masks( + margin, contained, cert_bad, planned) + + np.testing.assert_array_equal(structure, [True, False, False, False]) + np.testing.assert_array_equal(tail, [False, True, False, False]) + np.testing.assert_array_equal(containment, [False, False, True, False]) + np.testing.assert_array_equal(good, [False, False, False, True]) + assert np.all(np.sum(np.stack([good, structure, tail, containment]), axis=0) == 1) + + +def test_segment_sup_bound_is_an_upper_bound_on_the_cell(): + """The certificate itself: `q` on a cell must never exceed it.""" + sig = BandLimited(amp=4.0e4, peak_sample=NPTS // 2 + 0.3125) + k = sig.samples()[None, :] + kref, Xw, fk, period = _reflected_spectrum(k) + F = pl.PEAK_ENUM_FACTOR + h = DELTAT / F + n_keep = (NPTS - 1) * F + 1 + q = tmq.reflected_bandlimited_upsample(k, F)[0, :n_keep].real + (dq,) = pl.enum_grid_derivatives(kref, F, n_keep, DELTAT, orders=(1,)) + m4 = pl.spectral_derivative_bound(Xw, fk, period, 4)[0] + cert = pl.segment_sup_bound(q[:-1], q[1:], h * dq[0][:-1], h * dq[0][1:], h, m4) + # the true max on each cell, from a far finer sampling of the same interpolant + sub = 64 + qf = tmq.reflected_bandlimited_upsample(k, F * sub)[0, :(NPTS - 1) * F * sub + 1].real + true_cell = np.max(qf[:(n_keep - 1) * sub].reshape(-1, sub), axis=1) + assert np.all(cert >= true_cell - 1e-6), float(np.max(true_cell - cert)) + # ... and tight enough to be worth having: the Hermite remainder, not the h^2 bounds + m2 = pl.spectral_curvature_bound(Xw, fk, period)[0] + assert m4 * h ** 4 / 384.0 < m2 * h ** 2 / 8.0 + + +if __name__ == '__main__': + raise SystemExit(pytest.main([__file__, '-q'])) diff --git a/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py new file mode 100644 index 000000000..b9edd6a11 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py @@ -0,0 +1,1299 @@ +#!/usr/bin/env python +"""Gate for the band-limited time-marginalization quadrature. + +WHAT IS BEING TESTED, AND AGAINST WHAT +-------------------------------------- +The claim is that kappa(t) is band-limited below Nyquist and rho_sq is +time-independent. Production supplies only a finite integration-window slice, +so the implementation makes that slice value-continuous with a literal 2N +forward/backward reflection before FFT reconstruction; it does not identify the +slice's generally unlike endpoints. The reference here is NOT another numerical estimate of the same +thing, and it is not a stored number: it is an ANALYTIC continuous function. +kappa(t) is built as a sum of complex exponentials with every frequency below +Nyquist, and the truth is that same closed form evaluated directly (a dense +complex-exponential sum -- not an FFT, so it shares no machinery with the code +under test) and integrated at a density where the quadrature error is +analytically negligible. + +Two regimes are covered on purpose: + * exactly periodic on the window -> the interpolation is exact, so the only + error left is the quadrature's, and it should vanish to machine precision; + * a segment cut from a LONGER band-limited function -> not periodic on the + window, so a raw periodic interpolant rings at the wrap. This is the + realistic case and it is where the forward/backward boundary construction + has to earn its place. + +The wiring test drives the SHIPPED likelihood function rather than the helper: an +accuracy option that computes the right number but never reaches the likelihood +is the failure mode this repo has been bitten by before. +""" +from __future__ import print_function, division + +import os +import sys + +import numpy as np +import pytest +from scipy import integrate + +from RIFT.likelihood import time_marginalization_quadrature as tmq +from RIFT.likelihood import factored_likelihood as fl + +simpson = getattr(integrate, 'simpson', None) or integrate.simps + +SRATE = 4096.0 +DELTAT = 1.0 / SRATE +NPTS = 614 # marginalization_time_grid(0.075, 1/4096) +RHO_SQ = 1000.0 + + +# ----------------------------------------------------------------- helpers + +def _lnL(kappa_term, rho_sq): + """The production default helper, spelled out so the test does not depend on + a private name.""" + return kappa_term - 0.5 * rho_sq + + +def _log_trapz(v, dx): + m = v.max() + w = np.full(v.size, dx) + w[0] *= 0.5 + w[-1] *= 0.5 + return m + np.log(np.sum(w * np.exp(v - m))) + + +def _log_simps(v, dx): + m = v.max() + return m + np.log(simpson(np.exp(v - m), dx=dx)) + + +class BandLimited(object): + """kappa(t) = sum_m c_m exp(2 pi i m t / T), every |f| < Nyquist. + + ``n_period`` sets the period in samples. n_period == NPTS gives a window that + is exactly periodic; n_period > NPTS gives a window cut from a longer signal, + which is the realistic, non-periodic case. + """ + + def __init__(self, amp, peak_sample, n_period=NPTS, m_hi=None, seed=7, + background=0.0): + self.T = n_period * DELTAT + self.j0 = (n_period - NPTS) // 2 + scale = n_period / float(NPTS) + m_hi = int(200 * scale) if m_hi is None else m_hi + assert m_hi < n_period // 2, "would exceed Nyquist" + ms = np.arange(1, m_hi + 1) + t_peak = (self.j0 + peak_sample) * DELTAT + c = np.exp(-2j * np.pi * ms * t_peak / self.T) / (1.0 + (ms / (120.0 * scale)) ** 2) + if background: + rng = np.random.default_rng(seed) + c = c + background * ((rng.normal(size=m_hi) + 1j * rng.normal(size=m_hi)) + / (1.0 + (ms / (40.0 * scale)) ** 2)) + self.ms, self.c = ms, amp * c + + def at(self, ts, chunk=1000): + out = np.empty(np.size(ts), dtype=complex) + ts = np.asarray(ts) + for i in range(0, ts.size, chunk): + t = ts[i:i + chunk] + out[i:i + chunk] = np.exp(2j * np.pi * np.outer(t, self.ms) / self.T) @ self.c + return out + + def samples(self): + return self.at((self.j0 + np.arange(NPTS)) * DELTAT) + + def truth(self, refine=128): + """log int exp(lnL) dt over the SAME closed domain [t_0, t_{NPTS-1}].""" + n = (NPTS - 1) * refine + 1 + td = self.j0 * DELTAT + np.arange(n) * (DELTAT / refine) + return _log_trapz(_lnL(self.at(td).real, RHO_SQ), DELTAT / refine) + + +def _bandlimited(kappa_row): + k = np.asarray(kappa_row)[None, :] + r = np.full(k.shape, RHO_SQ) + return float(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL)[0]) + + +def _simpson_value(kappa_row): + return _log_simps(_lnL(np.asarray(kappa_row).real, RHO_SQ), DELTAT) + + +# ------------------------------------------- continuous posterior-draw contract + +def test_piecewise_linear_draw_has_no_output_lattice_and_handles_flat_density(): + lnL = np.zeros((2, 3)) + uniforms = np.array([[0.25, 0.5], [0.75, 0.2]]) + times, at_draw = tmq.draw_piecewise_linear_log_posterior( + lnL, 1.0, t0=-1.0, uniforms=uniforms) + # Equal interval masses: the first uniforms choose intervals 0 and 1. A + # flat density has the exact uniform conditional limit inside each interval. + np.testing.assert_allclose(times, [-0.5, 0.2], rtol=0, atol=1e-15) + np.testing.assert_allclose(at_draw, [0.0, 0.0], rtol=0, atol=0) + phase = (times + 1.0) / 1.0 + assert np.all(phase != np.round(phase)) + + +def test_piecewise_linear_draw_inverts_the_density_not_log_density(): + # One interval with density rising linearly from 1 to 3. At conditional + # quantile r=1/4, p(u)^2 = 1 + r*(9-1) = 3. + lnL = np.log(np.array([[1.0, 3.0]])) + times, at_draw = tmq.draw_piecewise_linear_log_posterior( + lnL, 2.0, t0=4.0, uniforms=np.array([[0.0, 0.25]])) + frac = (np.sqrt(3.0) - 1.0) / 2.0 + assert times[0] == pytest.approx(4.0 + 2.0 * frac, abs=2e-15) + assert at_draw[0] == pytest.approx(0.5 * np.log(3.0), abs=2e-15) + + +def test_bandlimited_draw_uses_the_same_validated_dense_representation(): + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + uniforms = np.array([[0.3712345, 0.6180339]]) + + integral, time_draw, lnL_draw = tmq.time_marginalize_bandlimited( + k, r, DELTAT, _lnL, return_time_draw=True, + draw_uniforms=uniforms, t0=-0.075) + integral_only = tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL) + np.testing.assert_allclose(integral, integral_only, rtol=0, atol=0) + + factor = tmq.last_report()['upsample_factor'] + assert factor > 1 + dense_k = tmq.reflected_bandlimited_upsample(k, factor) + dense_lnL = _lnL(dense_k.real, RHO_SQ) + dx = DELTAT / factor + phase = (float(time_draw[0]) + 0.075) / dx + j = int(np.floor(phase)) + frac = phase - j + density = np.exp(dense_lnL[0] - np.max(dense_lnL[0])) + density_at_draw = density[j] + frac * (density[j + 1] - density[j]) + expected_lnL = np.max(dense_lnL[0]) + np.log(density_at_draw) + assert float(lnL_draw[0]) == pytest.approx(expected_lnL, abs=2e-10) + assert abs(phase - round(phase)) > 1e-6, "draw snapped to the dense FFT grid" + + +def test_continuous_draw_accepts_minus_infinity_nodes_but_not_invalid_rows(): + times, lnL = tmq.draw_piecewise_linear_log_posterior( + np.array([[0.0, -np.inf, -1.0]]), 1.0, + uniforms=np.array([[0.1, 0.4]])) + assert np.isfinite(times[0]) and np.isfinite(lnL[0]) + # The exact RNG endpoint must skip a leading zero-mass plateau instead of + # selecting it through searchsorted's equality convention. + times, lnL = tmq.draw_piecewise_linear_log_posterior( + np.array([[-np.inf, -np.inf, 0.0]]), 1.0, + uniforms=np.array([[0.0, 0.0]])) + assert 1.0 < times[0] < 2.0 + assert np.isfinite(lnL[0]) + with pytest.raises(ValueError, match="no finite positive mass"): + tmq.draw_piecewise_linear_log_posterior( + np.full((1, 3), -np.inf), 1.0, + uniforms=np.array([[0.1, 0.4]])) + with pytest.raises(ValueError, match="NaN or \\+inf"): + tmq.draw_piecewise_linear_log_posterior( + np.array([[0.0, np.nan, -1.0]]), 1.0, + uniforms=np.array([[0.1, 0.4]])) + + +# ------------------------------------------------- the band-limited identity + +def test_upsample_is_exact_on_a_band_limited_sequence(): + """The upsample must REPRODUCE the analytic function between the samples, not + merely pass through them. Interpolating exactly at the input samples is the + weak check every interpolant passes; the strong one is the values in between, + which is the whole claim.""" + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2) + factor = 8 + up = tmq.bandlimited_upsample(sig.samples()[None, :], factor)[0] + t_dense = np.arange(NPTS * factor) * (DELTAT / factor) + exact = sig.at(t_dense) + assert np.allclose(up, exact, atol=1e-10, rtol=0), np.abs(up - exact).max() + # and the coarse samples land on dense indices j*factor (power-of-two design) + assert np.allclose(up[::factor], sig.samples(), atol=1e-12, rtol=0) + + +def test_reflected_upsample_reproduces_the_finite_row_exactly(): + """The non-periodic boundary construction must not move supplied samples.""" + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples() + factor = 8 + up = tmq.reflected_bandlimited_upsample(k[None, :], factor)[0] + assert up.size == (NPTS - 1) * factor + 1 + assert np.allclose(up[::factor], k, atol=1e-11, rtol=0) + + +def test_forward_backward_reflection_blocks_the_endpoint_gibbs_counterexample(): + """A decayed integrand does not imply a periodic kappa slice. + + This centred row has negligible edge likelihood but a large coherent + endpoint mismatch. Periodizing the raw slice biases the integral by more + than 100 nats; the literal 2N reflection must stay sub-millinat. This pins + the adversarial case that invalidated the tail-decay guard. + """ + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + tone_amp, tone_mode, tone_phase = 930.745, 620, 2.3532 + n_period = 8 * NPTS + js = sig.j0 + np.arange(NPTS) + tone = lambda j: tone_amp * np.exp(1j * (2 * np.pi * tone_mode * j / n_period + + tone_phase)) + k = sig.samples() + tone(js) + lnL = _lnL(k.real, RHO_SQ) + guard = int(NPTS * tmq.EDGE_GUARD_FRACTION) + assert guard < np.argmax(lnL) < NPTS - 1 - guard + assert lnL.max() - lnL[:guard].max() > 30 + assert lnL.max() - lnL[-guard:].max() > 30 + assert abs(k[0] - k[-1]) > 1000 + + refine = 128 + jd = sig.j0 + np.arange((NPTS - 1) * refine + 1) / float(refine) + exact = sig.at(jd * DELTAT, chunk=4000) + tone(jd) + ref = _log_trapz(_lnL(exact.real, RHO_SQ), DELTAT / refine) + sigma, _, _ = tmq.peak_width_from_lnL(lnL[None, :], DELTAT) + factor = int(tmq.required_upsample_factors(sigma, DELTAT)[0]) + + raw = tmq.bandlimited_upsample(k[None, :], factor)[0] + raw_value = _log_trapz(_lnL(raw[:(NPTS - 1) * factor + 1].real, RHO_SQ), + DELTAT / factor) + reflected = tmq.reflected_bandlimited_upsample(k[None, :], factor)[0] + reflected_value = _log_trapz(_lnL(reflected.real, RHO_SQ), DELTAT / factor) + assert raw_value - ref > 100 + assert abs(reflected_value - ref) < 1e-3 + + +def test_peak_width_estimator_is_exact_for_a_gaussian_at_any_grid_phase(): + """The width estimator is what makes the refinement DERIVED rather than + guessed, and its whole job is to stay honest when the peak is under-resolved + and sitting at an arbitrary phase relative to the grid. A Gaussian lnL has a + known width; recover it from grids that resolve it badly.""" + for sigma_over_dt in (4.0, 1.0, 0.3, 0.05): + sigma = sigma_over_dt * DELTAT + for phase in (0.0, 0.25, 0.5, 0.75): + t = (np.arange(NPTS) - NPTS // 2 + phase) * DELTAT + lnL = -0.5 * (t / sigma) ** 2 + got, _, meas = tmq.peak_width_from_lnL(lnL[None, :], DELTAT) + assert bool(meas[0]) + assert np.isclose(float(got[0]), sigma, rtol=1e-9), (sigma_over_dt, phase, got) + + +def test_flat_integrand_derives_no_refinement(): + """A well-resolved integrand must cost nothing: the derivation has to return + factor 1 rather than paying for resolution it does not need.""" + sig = BandLimited(amp=0.002, peak_sample=NPTS // 2) + lnL = _lnL(sig.samples().real, RHO_SQ)[None, :] + sigma, _, meas = tmq.peak_width_from_lnL(lnL, DELTAT) + assert bool(meas[0]) and float(sigma[0]) > DELTAT, sigma + assert int(tmq.required_upsample_factors(sigma, DELTAT)[0]) == 1 + + +# --------------------------------------------- accuracy against analytic truth + +@pytest.mark.parametrize("amp,phase", [(a, p) for a in (0.3, 0.5, 1.0, 5.0) + for p in (0.0, 0.25, 0.5)]) +def test_exact_on_a_periodic_window(amp, phase): + """Exactly-periodic window: interpolation is exact, so the band-limited value + must match the analytic truth to well below any level Simpson achieves.""" + sig = BandLimited(amp=amp, peak_sample=NPTS // 2 + phase) + ref = sig.truth() + k = sig.samples() + assert abs(_bandlimited(k) - ref) < 1e-6 + + +@pytest.mark.parametrize("amp,phase", [(a, p) for a in (0.05, 0.17, 1.0, 5.0) + for p in (0.0, 0.25, 0.5)]) +def test_accurate_on_a_non_periodic_window(amp, phase): + """The realistic case: the window is a segment of a longer band-limited + signal, so a raw periodic interpolant rings at the wrap. With the peak + centred, the reflected residual must still be far below Simpson's error.""" + sig = BandLimited(amp=amp, peak_sample=NPTS // 2 + phase, + n_period=8 * NPTS, m_hi=1400, background=0.12) + ref = sig.truth() + k = sig.samples() + assert abs(_bandlimited(k) - ref) < 1e-3 + + +def test_beats_simpson_where_the_peak_is_under_resolved(): + """The defect itself. Sweeping the peak across one sample must move the + Simpson answer by of order a nat while leaving the band-limited answer put -- + that grid-phase sensitivity IS the bug, and insensitivity to it is the fix.""" + sig0 = BandLimited(amp=0.05, peak_sample=NPTS // 2, + n_period=8 * NPTS, m_hi=1400, background=0.12) + sigma, _, _ = tmq.peak_width_from_lnL(_lnL(sig0.samples().real, RHO_SQ)[None, :], DELTAT) + assert 0.08 < float(sigma[0]) / DELTAT < 0.30, "not the under-resolved regime" + + s_err, b_err = [], [] + for phase in (0.0, 0.25, 0.5, 0.75): + sig = BandLimited(amp=0.05, peak_sample=NPTS // 2 + phase, + n_period=8 * NPTS, m_hi=1400, background=0.12) + ref = sig.truth() + k = sig.samples() + s_err.append(_simpson_value(k) - ref) + b_err.append(_bandlimited(k) - ref) + + assert max(s_err) - min(s_err) > 0.5, s_err # Simpson swings by ~2 nats + assert max(np.abs(b_err)) < 1e-3, b_err + assert max(np.abs(b_err)) < 0.01 * max(np.abs(s_err)) + + +# ----------------------------------------------------------- the edge guard + +def test_centered_row_does_not_switch_rules_at_a_tail_threshold(): + """Tail height does not select a Simpson fallback after reflection.""" + sig = BandLimited(amp=0.02, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples() + lnL = _lnL(k.real, RHO_SQ) + guard = max(1, int(NPTS * tmq.EDGE_GUARD_FRACTION)) + jmax = int(np.argmax(lnL)) + assert guard <= jmax <= NPTS - 1 - guard + assert lnL.max() - lnL[:guard].max() < 30 + assert lnL.max() - lnL[-guard:].max() < 30 + + got = _bandlimited(k) + assert got != _simpson_value(k) + assert abs(got - sig.truth()) < 1e-3 + rep = tmq.last_report() + assert rep['n_refined_rows'] == 1, rep + +def test_boundary_diagnostic_does_not_select_simpson(): + """Crossing the diagnostic boundary must not change quadrature rules.""" + for peak in (0.3, 2.3, 30.3): + sig = BandLimited(amp=1.0, peak_sample=peak, n_period=8 * NPTS, + m_hi=1400, background=0.12) + k = sig.samples() + assert _bandlimited(k) != _simpson_value(k), peak + assert tmq.last_report()['n_wrap_exposed_rows'] == 1 + assert tmq.last_report()['n_refined_rows'] == 1 + + +def test_one_exposed_row_does_not_contaminate_its_block(): + """The boundary diagnostic is per row and does not alter reconstruction.""" + bad = BandLimited(amp=1.0, peak_sample=1.3, n_period=8 * NPTS, + m_hi=1400, background=0.12) + good = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, n_period=8 * NPTS, + m_hi=1400, background=0.12, seed=11) + k = np.stack([bad.samples(), good.samples()]) + out = tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + assert tmq.last_report()['n_wrap_exposed_rows'] == 1 + assert tmq.last_report()['n_refined_rows'] == 2 + assert float(out[0]) != _simpson_value(bad.samples()) + assert abs(float(out[1]) - good.truth()) < 1e-3 + + +def test_rows_sharing_a_block_keep_their_individual_resolution(): + """Each row still derives and pays for its own reconstruction factor.""" + flat = BandLimited(amp=0.0012, peak_sample=NPTS // 2 + 0.3, n_period=8 * NPTS, + m_hi=1400, background=0.12, seed=11) + sharp = BandLimited(amp=5.0, peak_sample=NPTS // 2 + 0.3, n_period=8 * NPTS, + m_hi=1400, background=0.12) + k = np.stack([flat.samples(), sharp.samples()]) + out = tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + rep = tmq.last_report() + hist = rep['factor_histogram'] + assert rep['upsample_factor'] > 8 + assert rep['n_refined_rows'] == 2, rep + assert float(out[0]) != _simpson_value(flat.samples()) + assert abs(float(out[0]) - flat.truth()) < 1e-4 + assert sum(hist.values()) == 2, hist + + +# ------------------------------------------------------------- preconditions + +def test_time_dependent_rho_sq_is_refused(): + """The precondition is checked, not trusted. A time-dependent self-term (the + banded / rotating-response path) would give a confident wrong number.""" + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2) + k = sig.samples()[None, :] + rho = np.full(k.shape, RHO_SQ) + rho[0, NPTS // 3] += 1e-9 + with pytest.raises(NotImplementedError): + tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + + +def test_ceiling_raises_rather_than_truncating_resolution(): + """Running out of refinement must be an error, never a silently coarser grid.""" + old = tmq.UPSAMPLE_FACTOR_MAX + tmq.UPSAMPLE_FACTOR_MAX = 4 + try: + sig = BandLimited(amp=40.0, peak_sample=NPTS // 2) + k = sig.samples()[None, :] + with pytest.raises(RuntimeError): + tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + finally: + tmq.UPSAMPLE_FACTOR_MAX = old + + +def test_unknown_quadrature_name_is_rejected(): + with pytest.raises(ValueError): + tmq.validate_time_quadrature('bandlimted') # sic + + +# --------------------------------------------------------------- the wiring + +N_BUFFER = 4096 + + +def _fake_likelihood_inputs(kappa_buffer): + """Minimal inputs that drive the SHIPPED NoLoop function on the numpy backend. + + One detector, one (l,m) pair and zero cross terms, so the self-term is a + constant and kappa reduces to the supplied rholm buffer times a fixed + response factor. The point is to exercise the argument plumbing; the physics + is covered above against analytic truth. The buffer is band-limited AND + periodic on its own length, so whatever integer window the code gathers is a + genuine band-limited segment -- the test does not need to predict ``ifirst``. + """ + import lal + import RIFT.lalsimutils as lsu + + det = 'H1' + rholm = np.asarray(kappa_buffer, dtype=complex)[None, :] + P = lsu.ChooseWaveformParams() + P.deltaT = DELTAT + P.tref = 1000000000.0 + for name, val in [('phi', 0.0), ('theta', 0.0), ('phiref', 0.0), + ('incl', 0.0), ('psi', 0.0)]: + setattr(P, name, np.zeros(1) + val) + P.dist = np.full(1, fl.distMpcRef * 1e6 * lal.PC_SI) + # Put the window well inside the buffer: the epoch offset sets ifirst, and a + # window running off the front would be zero-extended rather than gathered. + return (P, {det: rholm}, {det: np.array([[2, 2]])}, + {det: np.zeros((1, 1), dtype=complex)}, {det: P.tref - 0.5}) + + +def _buffer_signal(amp, roll=0): + sig = BandLimited(amp=amp, peak_sample=NPTS // 2, n_period=N_BUFFER, + m_hi=1400, background=0.12) + ts = np.arange(N_BUFFER) * DELTAT + return np.roll(sig.at(ts), int(roll)) + + +def _shipped(tvals, args, **kw): + P, rholms, lookupNK, ct, epochs = args + return fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, Lmax=2, xpy=np, **kw) + + +def _tuned_inputs(tvals, sigma_target_over_dt=0.15): + """Build likelihood inputs whose lnL(t) actually sits in the under-resolved + regime AND has decayed window tails, by MEASURING what the shipped function + produces rather than assuming it: the response factor and the gather offset + are the code's business, not the test's. Centres the peak in the window (an + integer roll of a periodic band-limited buffer is still band-limited) and + scales the amplitude using sigma ~ 1/sqrt(amp). The 0.15 target leaves about + 49 nats of measured edge suppression, safely beyond the 30-nat eligibility + boundary; the former 0.25 target left only 18 and correctly stopped reaching + the FFT path once that boundary was added. + """ + amp, roll = 1.0, 0 + for _ in range(6): + args = _fake_likelihood_inputs(_buffer_signal(amp, roll)) + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + sigma, jmax, _ = tmq.peak_width_from_lnL(lnL_t, DELTAT) + roll += int(NPTS // 2 - int(jmax[0])) + if np.isfinite(sigma[0]): + amp *= (float(sigma[0]) / (sigma_target_over_dt * DELTAT)) ** 2 + args = _fake_likelihood_inputs(_buffer_signal(amp, roll)) + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True)) + sigma, jmax, _ = tmq.peak_width_from_lnL(lnL_t, DELTAT) + return args, float(sigma[0]) / DELTAT, int(jmax[0]) + + +def test_driver_flag_reaches_the_likelihood_and_changes_the_answer(): + """The wiring, not the helper. + + A flag that is computed correctly and then never reaches the likelihood is a + documented failure mode in this repo -- a whole comparison campaign has been + run against an inert stencil option here before. So: set the module default + the way the driver sets it, call the SHIPPED function, and require the number + to actually move on an under-resolved peak. + """ + pytest.importorskip('RIFT.lalsimutils') + tvals = fl.marginalization_time_grid(0.075, DELTAT) + assert len(tvals) == NPTS + + args, sigma_over_dt, jmax = _tuned_inputs(tvals) + assert 0.1 < sigma_over_dt < 0.6, sigma_over_dt # under-resolved, as intended + guard = int(NPTS * tmq.EDGE_GUARD_FRACTION) + assert guard < jmax < NPTS - 1 - guard, jmax # and not wrap-exposed + lnL_t = np.asarray(_shipped(tvals, args, return_lnLt=True))[0] + assert lnL_t.max() - lnL_t[:guard].max() >= 30 + assert lnL_t.max() - lnL_t[-guard:].max() >= 30 + + assert fl.TIME_QUADRATURE_DEFAULT == 'simpson', "default must not have moved" + base = float(np.asarray(_shipped(tvals, args))[0]) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'bandlimited' # exactly what the driver does + new = float(np.asarray(_shipped(tvals, args))[0]) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + assert tmq.last_report()['upsample_factor'] > 1 + assert tmq.last_report()['n_wrap_exposed_rows'] == 0 + assert abs(new - base) > 1e-3, (base, new) + + # the explicit kwarg must override the module default, in both directions + kw = float(np.asarray(_shipped(tvals, args, time_quadrature='bandlimited'))[0]) + assert kw == new + fl.TIME_QUADRATURE_DEFAULT = 'bandlimited' + try: + assert float(np.asarray(_shipped(tvals, args, time_quadrature='simpson'))[0]) == base + finally: + fl.TIME_QUADRATURE_DEFAULT = old + + +def test_shipped_likelihood_exposes_a_continuous_bandlimited_time_draw(): + pytest.importorskip('RIFT.lalsimutils') + tvals = fl.marginalization_time_grid(0.075, DELTAT) + args, sigma_over_dt, _ = _tuned_inputs(tvals) + assert sigma_over_dt < 0.6 + uniforms = np.array([[0.4321, 0.6789]]) + drawn_t, drawn_lnL = _shipped( + tvals, args, time_quadrature='bandlimited', return_time_draw=True, + time_draw_uniforms=uniforms) + factor = tmq.last_report()['upsample_factor'] + assert factor > 1 + assert np.isfinite(float(drawn_t[0])) and np.isfinite(float(drawn_lnL[0])) + assert float(tvals[0]) <= float(drawn_t[0]) <= float(tvals[-1]) + dense_phase = (float(drawn_t[0]) - float(tvals[0])) / (DELTAT / factor) + assert abs(dense_phase - round(dense_phase)) > 1e-6 + + with pytest.raises(ValueError, match='requires time_quadrature'): + _shipped(tvals, args, time_quadrature='simpson', return_time_draw=True, + time_draw_uniforms=uniforms) + with pytest.raises(ValueError, match='mutually exclusive'): + _shipped(tvals, args, time_quadrature='bandlimited', return_time_draw=True, + return_lnLt=True, time_draw_uniforms=uniforms) + + +def test_unsupported_combinations_refuse_rather_than_silently_using_simpson(): + pytest.importorskip('RIFT.lalsimutils') + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2) + tvals = fl.marginalization_time_grid(0.075, DELTAT) + P, rholms, lookupNK, ct, epochs = _fake_likelihood_inputs([sig.samples()]) + common = dict(Lmax=2, xpy=np, time_quadrature='bandlimited') + for extra in ({'n_cal': 2}, {'return_lnLt': True}, {'return_cal_components': True}): + kw = dict(common); kw.update(extra) + with pytest.raises(NotImplementedError): + fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, P, lookupNK, rholms, ct, ct, epochs, **kw) + + +@pytest.mark.parametrize("module_name,func_name", [ + ('RIFT.likelihood.factored_likelihood_with_rotation', + 'DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation'), + ('RIFT.likelihood.factored_likelihood_freqresponse', + 'DiscreteFactoredLogLikelihoodFreqResponseNoLoop'), +]) +def test_excluded_paths_refuse_the_global_default(module_name, func_name): + """These likelihoods have a time-DEPENDENT rho_sq, so the band-limited + argument does not hold for them. Enabling the option globally must make them + RAISE, not quietly run Simpson -- otherwise the exclusion is invisible at the + point of use. Behavioural, not a source grep: the guard is the first thing + the function does, so junk arguments must still produce NotImplementedError + rather than a TypeError from further in.""" + mod = pytest.importorskip(module_name) + func = getattr(mod, func_name) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'bandlimited' + with pytest.raises(NotImplementedError): + func(None, None, None, None, None, None, None, None) + finally: + fl.TIME_QUADRATURE_DEFAULT = old + + +# ------------------------------------------- non-finite lnL(t) (input space) + +def test_minus_inf_next_to_the_peak_does_not_read_as_a_flat_integrand(): + """``lnL_t`` genuinely contains ``-inf`` in production: the distance- + marginalization callback returns ``-inf`` outside its interpolation table. + A three-point stencil that straddles that hole computes + ``(-inf) - 2*(-inf) + (-inf) = NaN``, and ``NaN < 0`` is False -- so the row + would report "no peak", derive a factor of 1 and be SILENTLY under-resolved, + which is the exact failure this change exists to remove. This cannot be + caught by mutating the code (it is a missing case, not a wrong constant), so + it is tested from the input side.""" + t = (np.arange(NPTS) - NPTS // 2) * DELTAT + sigma_true = 0.05 * DELTAT + base = -0.5 * (t / sigma_true) ** 2 + + for label, hole in [("tails", (slice(0, 20), slice(-20, None))), + ("adjacent to the peak", (NPTS // 2 - 1,)), + ("both sides of the peak", (NPTS // 2 - 1, NPTS // 2 + 1))]: + lnL = base.copy() + for h in hole: + lnL[h] = -np.inf + sigma, _, meas = tmq.peak_width_from_lnL(lnL[None, :], DELTAT) + assert bool(meas[0]), label + assert np.isclose(float(sigma[0]), sigma_true, rtol=1e-9), (label, sigma) + assert int(tmq.required_upsample_factors(sigma, DELTAT)[0]) > 1, label + + +def test_a_signal_free_row_is_reported_as_flat_not_as_wrap_exposed(): + """A row with no signal in it -- an extrinsic sample in an antenna null, where + kappa is numerically zero -- has a constant lnL(t) and therefore an argmax of + 0 by convention. Applying the edge guard to it would report it as + wrap-exposed, which in a production log reads as a mis-centred window rather + than as a row with nothing in it. The edge guard is only meaningful for rows + that HAVE a peak.""" + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2) + k = np.stack([np.zeros(NPTS, dtype=complex), sig.samples()]) + out = tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + rep = tmq.last_report() + assert rep['n_flat_rows'] == 1, rep + assert rep['n_wrap_exposed_rows'] == 0, rep + assert rep['n_unmeasurable_rows'] == 0, rep + # and it is still integrated correctly: a constant integrand over the window + expect = _lnL(0.0, RHO_SQ) + np.log((NPTS - 1) * DELTAT) + assert abs(float(out[0]) - expect) < 1e-12, (out[0], expect) + + +def test_unmeasurable_row_falls_back_and_is_counted(): + """A row whose curvature cannot be evaluated at ANY stencil half-width must be + counted and given the historical value -- never silently assigned factor 1, + which is indistinguishable from a genuinely flat integrand.""" + sig = BandLimited(amp=1.0, peak_sample=NPTS // 2) + k = np.stack([np.zeros(NPTS, dtype=complex), sig.samples()]) + + def lnL_with_hole(kappa_term, rho_sq): + out = _lnL(kappa_term, rho_sq) + out = np.where(np.abs(np.asarray(kappa_term)) > 0, out, -np.inf) + return out + + out = tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, + lnL_with_hole) + rep = tmq.last_report() + assert rep['n_unmeasurable_rows'] == 1, rep + assert rep['n_refined_rows'] == 1, rep + # `exposed` is gated on `has_peak`, which already implies `measurable`, so an + # unmeasurable row can never also be exposed -- asserting the two do not + # overlap is vacuous. What is worth pinning is that the counters PARTITION + # the batch, so no row can fall through a gap between them and be invisible. + assert (rep['n_refined_rows'] + + rep['n_unmeasurable_rows'] + + rep['n_flat_rows'] + + _n_resolved(rep)) == rep['n_rows'], rep + # zero likelihood over the whole window integrates to zero: the answer is + # -inf, which is what the historical global-offset path returns. NaN here + # would propagate into the sampler weights. + assert float(out[0]) == -np.inf, out[0] + assert abs(float(out[1]) - sig.truth()) < 1e-6 + + +def test_a_row_changes_if_and_only_if_it_was_under_resolved(): + """The guarantee, stated so it can be checked rather than argued. + + Every row that is NOT refined -- unmeasurable, flat, or already resolved -- + must come back with the historical Simpson value, so + enabling this option cannot make any row worse than the status quo. Letting + an unrefined row fall through to a coarse trapezoid instead is numerically a + non-event, but it changes the rule for rows this option was never meant to + touch and forfeits exactly this property. + """ + rows, expect_refined = [], [] + # resolved (no refinement warranted) + rows.append(BandLimited(amp=0.002, peak_sample=NPTS // 2).samples()); expect_refined.append(False) + # signal-free + rows.append(np.zeros(NPTS, dtype=complex)); expect_refined.append(False) + # boundary-diagnostic row: still refined + rows.append(BandLimited(amp=1.0, peak_sample=2.3, n_period=8 * NPTS, + m_hi=1400, background=0.12).samples()); expect_refined.append(True) + # centred and sharp with non-negligible coarse tails: still refined + rows.append(BandLimited(amp=0.02, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, + background=0.12).samples()); expect_refined.append(True) + # genuinely under-resolved + rows.append(BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, n_period=8 * NPTS, + m_hi=1400, background=0.12).samples()); expect_refined.append(True) + + k = np.stack(rows) + out = tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + rep = tmq.last_report() + assert rep['n_refined_rows'] == sum(expect_refined), rep + + for i, should_change in enumerate(expect_refined): + historical = _simpson_value(rows[i]) + if should_change: + assert float(out[i]) != historical, i + else: + assert float(out[i]) == historical, (i, out[i], historical) + + +def test_remeasure_on_the_dense_grid_repairs_an_under_derived_factor(): + """The remeasure-and-double step is what makes the derivation an assertion + rather than a guess. Force the derivation to hand back a factor that is far + too small and require the refinement loop to notice on the dense grid and + recover the right answer anyway.""" + sig = BandLimited(amp=5.0, peak_sample=NPTS // 2 + 0.25) + k = sig.samples()[None, :] + rho = np.full(k.shape, RHO_SQ) + honest = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + honest_factor = tmq.last_report()['upsample_factor'] + assert honest_factor >= 16 + + real = tmq.required_upsample_factors + tmq.required_upsample_factors = lambda sigma, dx, xpy=np: real(sigma, dx, xpy=xpy) // 8 + try: + got = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + finally: + tmq.required_upsample_factors = real + rep = tmq.last_report() + assert rep['n_refinements'] > 0, rep + assert rep['upsample_factor'] == honest_factor, rep + assert abs(float(got[0]) - float(honest[0])) < 1e-9 + assert abs(float(got[0]) - sig.truth()) < 1e-6 + + + + +# ------------------------------------------------------- the driver CLI + +def _run_driver(extra_args): + """Invoke the ILE driver and return (returncode, combined output). + + A subprocess, deliberately. The option's whole job is to travel from a + command line into the likelihood, and the guard that stops it being silently + inert lives in the driver's startup, not in the library -- so a test that + imports the library cannot see it. The driver exits long before any data is + needed, so this costs one interpreter start. + """ + import subprocess + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + driver = os.path.join(root, 'bin', 'integrate_likelihood_extrinsic_batchmode') + env = dict(os.environ) + env['PYTHONPATH'] = root + os.pathsep + env.get('PYTHONPATH', '') + env['OMP_NUM_THREADS'] = '1' + proc = subprocess.run([sys.executable, driver] + extra_args, env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=900) + return proc.returncode, proc.stdout.decode('utf-8', 'replace') + + +_HONOURED = ['--time-marginalization', '--vectorized', '--gpu', '--force-xpy'] + + +def test_driver_rejects_a_misspelled_quadrature_name(): + """A misspelled stencil name was once absorbed as "not truthy" and silently + ran a different likelihood here. A typo in this option has to be loud.""" + rc, out = _run_driver(['--time-marginalization-quadrature', 'bandlimted']) + assert rc != 0 + assert 'bandlimted' in out and 'simpson' in out + + +def test_driver_refuses_configurations_that_cannot_honour_the_option(): + """Refuse, do not ignore. Each of these would otherwise run the historical + Simpson quadrature while the startup banner said otherwise -- which is how a + comparison campaign gets run against an inert flag.""" + for missing in ([], ['--time-marginalization'], + ['--time-marginalization', '--vectorized'], + _HONOURED + ['--rotation-slow'], + _HONOURED + ['--freqresponse']): + rc, out = _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + missing) + assert rc != 0, (missing, out[-2000:]) + assert 'cannot honour it' in out, (missing, out[-2000:]) + + +def test_driver_announces_the_quadrature_it_will_actually_use(): + rc, out = _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + _HONOURED) + assert 'Time-marginalization quadrature: bandlimited' in out + assert 'honoured by this configuration: True' in out + # and the default stays put when the option is not given + rc, out = _run_driver(_HONOURED) + assert 'Time-marginalization quadrature: simpson' in out + + + +# --------------------------------------------------------------- GPU parity + +def _cupy_or_skip(): + """cupy, or skip -- unless the GPU gate demands a device, in which case FAIL. + + `RIFT_CI_REQUIRE_GPU=1` is how `.travis/test-integrate.sh` says "this job runs + on hardware". A skip under that flag would be a GPU gate reporting green + without having touched a GPU, which is the failure mode this whole file is + written against. + """ + try: + import cupy + if cupy.cuda.runtime.getDeviceCount() < 1: + raise RuntimeError("cupy imported but reports zero CUDA devices") + return cupy + except Exception as exc: + if os.environ.get('RIFT_CI_REQUIRE_GPU') == '1': + pytest.fail("RIFT_CI_REQUIRE_GPU=1 but cupy/GPU unavailable: %s" % exc) + pytest.skip("cupy/GPU unavailable: %s" % exc) + + +def test_bandlimited_runs_on_the_gpu_backend_and_matches_numpy(): + """The backend-generic code must actually RUN on cupy, not merely look like it. + + This path shipped once already having never been executed on a GPU: the + likelihood omitted the caller's `simps`, the module defaulted to scipy's, and + scipy raises `TypeError: Implicit conversion to a NumPy array is not allowed` + on a cupy array -- so EVERY `--vectorized --gpu` run of the option crashed + while all 46 CPU tests stayed green. Reading the cupy API is not a substitute + for executing it. + """ + cupy = _cupy_or_skip() + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + flat = BandLimited(amp=0.002, peak_sample=NPTS // 2) + edge = BandLimited(amp=1.0, peak_sample=2.3, n_period=8 * NPTS, + m_hi=1400, background=0.12) + k = np.stack([sig.samples(), flat.samples(), edge.samples(), + np.zeros(NPTS, dtype=complex)]) + r = np.full(k.shape, RHO_SQ) + + simps_cpu = simpson + out_np = np.asarray(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL, + simps=simps_cpu, xpy=np)) + rep_np = tmq.last_report() + + from RIFT.likelihood import optimized_gpu_tools + out_cp = cupy.asnumpy(tmq.time_marginalize_bandlimited( + cupy.asarray(k), cupy.asarray(r), DELTAT, _lnL, + simps=optimized_gpu_tools.simps, xpy=cupy)) + rep_cp = tmq.last_report() + + # Same classification and the same derived factors on both backends. + for key in ('upsample_factor', 'factor_histogram', 'n_refined_rows', + 'n_wrap_exposed_rows', 'n_unmeasurable_rows', + 'n_flat_rows'): + assert rep_np[key] == rep_cp[key], (key, rep_np[key], rep_cp[key]) + assert rep_np['n_refined_rows'] >= 1 + + # The REFINED rows integrate with trapezoid on the dense grid, which has no + # even/odd Simpson ambiguity, so the two backends must agree to round-off. + # (Rows that fall back use each backend's OWN Simpson rule, and those two + # rules genuinely differ for even npts -- see the CPU/GPU note below.) + # Only the REFINED rows are required to agree: they integrate with trapezoid + # on the dense grid, which has no even/odd Simpson ambiguity. Rows that fall + # back use each backend's OWN Simpson rule, and those two rules genuinely + # differ for even npts -- asserting agreement over all rows would be + # asserting the absence of a divergence this file documents as real. + refined = np.array([f > 1 for f in _row_factors(k, r)]) + assert refined.any() + fin = np.isfinite(out_np) & np.isfinite(out_cp) & refined + assert np.max(np.abs(out_np[fin] - out_cp[fin])) < 1e-6 + + +def test_the_likelihood_hands_the_bandlimited_path_its_own_simpson_rule(): + """Rows that fall back must reproduce what THIS run would have returned. + + `factored_likelihood` integrates with scipy on CPU and with + `optimized_gpu_tools.simps` on GPU, and the two are NOT interchangeable: the + vendored GPU copy is an old scipy with `even='avg'` while modern scipy uses + the Cartwright correction, so for EVEN npts -- production is 614 at srate + 4096 -- they disagree. A private scipy copy inside the module would be + bit-for-bit on CPU and quietly wrong on GPU. + """ + # Behavioural, not a source grep: a non-numpy backend with no `simps` must + # REFUSE, because the scipy default cannot consume a device array and that + # default is precisely how every GPU run of this option crashed. + class _FakeDevice(object): + """Minimal stand-in for a non-numpy backend, so the guard is exercised + without needing a GPU.""" + def __getattr__(self, name): + return getattr(np, name) + + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + with pytest.raises(ValueError): + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL, xpy=_FakeDevice()) + # ... and is satisfied once a rule is supplied + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL, simps=simpson, + xpy=_FakeDevice()) + + # The likelihood must hand its OWN rule over. Parse the real call, not a + # prefix of it: slicing at the first ')' truncates inside `float(deltaT)`. + import ast, inspect + tree = ast.parse(inspect.getsource( + fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop).lstrip()) + calls = [n for n in ast.walk(tree) if isinstance(n, ast.Call) + and getattr(n.func, 'attr', None) == 'time_marginalize_bandlimited'] + assert calls, "the likelihood no longer calls time_marginalize_bandlimited" + for c in calls: + kw = {k2.arg: k2.value for k2 in c.keywords} + assert 'simps' in kw and getattr(kw['simps'], 'id', None) == 'simps', \ + "call site does not forward the caller's Simpson rule" + assert 'lnL_coarse' in kw, \ + "call site does not forward the already-computed coarse lnL" + + +def _row_factors(k, r): + """Per-row derived factor, as the integrator computes it.""" + lnL = _lnL(np.asarray(k).real, np.asarray(r)) + sigma, _, meas = tmq.peak_width_from_lnL(lnL, DELTAT) + ok = meas & np.isfinite(sigma) + f = np.maximum(tmq.required_upsample_factors(sigma, DELTAT), 1) + return np.where(ok, f, 1) + + +@pytest.mark.parametrize("n", [153, 307, 613, 614, 1228, 2457, 8, 9, 3]) +def test_upsample_is_exact_for_odd_npts_too(n): + """ODD npts is the COMMON case in production, not an exotic one. + + `marginalization_time_grid(0.075, 1/srate)` gives npts = 153 / 307 / 614 / + 1228 / 2457 at srate 1024 / 2048 / 4096 / 8192 / 16384 -- odd at THREE of the + five, including 16384, the low-mass rate. An earlier split at `h = n//2` + placed the highest positive frequency at a negative frequency for odd n: + still exact AT the original samples (so a "reproduces the input" check passes) + and wrong everywhere between them, by 0.41 at n=613 and 0.54 at n=307 against + an analytic truth of order unity. + """ + R = 4 + rng = np.random.default_rng(1) + ms = np.arange(1, (n - 1) // 2 + 1) # fill EVERY bin up to Nyquist + c = (rng.normal(size=ms.size) + 1j * rng.normal(size=ms.size)) / (1 + ms / 50.0) + t = np.arange(n) / float(n) + td = np.arange(n * R) / float(n * R) + x = np.exp(2j * np.pi * np.outer(t, ms)) @ c + exact = np.exp(2j * np.pi * np.outer(td, ms)) @ c + up = tmq.bandlimited_upsample(x[None, :], R)[0] + assert np.allclose(up, exact, atol=1e-9, rtol=0), np.abs(up - exact).max() + + +def test_which_rows_change_relative_to_the_SHIPPED_historical_expression(): + """The guarantee, checked against the historical GLOBAL-offset expression. + + An earlier version of this test compared against a per-row-offset Simpson + helper -- the same expression the code under test uses for its fallback rows + -- so it was common-mode with the thing it was meant to check and could not + fail. The shipped path offsets by a SINGLE GLOBAL maximum over the whole + block, so a multi-row batch with production-scale dynamic range is required + to see the difference at all. + """ + def historical(kappa_rows, rho): + lnL_t = _lnL(np.asarray(kappa_rows).real, rho) + lnLmax = lnL_t.max() # GLOBAL, as the shipped path does + return lnLmax + np.log(simpson(np.exp(lnL_t - lnLmax), dx=DELTAT, axis=-1)) + + loud = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, n_period=8 * NPTS, + m_hi=1400, background=0.12).samples() + quiet = BandLimited(amp=0.002, peak_sample=NPTS // 2).samples() * 1e-3 + k = np.stack([loud, quiet]) + r = np.full(k.shape, RHO_SQ) + + hist = historical(k, r) + new = np.asarray(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL)) + rep = tmq.last_report() + + # The QUADRATURE changed for exactly the under-resolved row. + assert rep['n_refined_rows'] == 1, rep + assert _row_factors(k, r)[0] > 1 and _row_factors(k, r)[1] == 1 + + # The refined row moved, as intended. + assert abs(new[0] - hist[0]) > 1e-3 + + # And the documented second change: the unrefined row underflowed to -inf + # under the shared global offset and now comes back finite. This is NOT + # "unchanged"; pin it so the PR text and the code cannot drift apart. + span = _lnL(k.real, r).max() - _lnL(k.real, r)[1].max() + assert span > 745, span # the underflow threshold for exp() + assert hist[1] == -np.inf + assert np.isfinite(new[1]) + + +def test_return_lnLt_still_works_when_the_module_default_is_bandlimited(): + """The group's standard extrinsic stage must not die at the export step. + + `--add-extrinsic --add-extrinsic-time-resampling` maps to + `--resample-time-marginalization`, whose `resample_samples()` calls the + likelihood with `return_lnLt=True` and no explicit quadrature. Raising on + the INHERITED default made that configuration run the entire integration and + then crash with no output. `return_lnLt` returns lnL(t) on the original grid + and takes no time integral, so the quadrature is inapplicable, not ignored -- + but asking for it EXPLICITLY there is still a caller error. + """ + pytest.importorskip('RIFT.lalsimutils') + tvals = fl.marginalization_time_grid(0.075, DELTAT) + args = _fake_likelihood_inputs(_buffer_signal(1.0)) + old = fl.TIME_QUADRATURE_DEFAULT + try: + fl.TIME_QUADRATURE_DEFAULT = 'bandlimited' + lnLt = np.asarray(_shipped(tvals, args, return_lnLt=True)) + assert lnLt.shape[-1] == NPTS + with pytest.raises(NotImplementedError): + _shipped(tvals, args, return_lnLt=True, time_quadrature='bandlimited') + finally: + fl.TIME_QUADRATURE_DEFAULT = old + + +def test_a_nan_self_term_does_not_abort_the_run(): + """NaN rows are NORMAL -- the defensive proposal component deliberately draws + physically-extreme points where the likelihood is NaN, and the historical path + returns NaN for that row and moves on. A bare `rho_sq == rho_sq[...,:1]` + tripwire makes `nan != nan` abort the whole ILE process, blaming a + rotating-response path that is not in use.""" + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2) + k = np.stack([sig.samples(), sig.samples()]) + r = np.full(k.shape, RHO_SQ) + r[1, :] = np.nan + out = tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL) + assert np.isfinite(float(out[0])) + assert np.isnan(float(out[1])) + # a genuinely time-DEPENDENT self-term must still be refused + r2 = np.full(k.shape, RHO_SQ); r2[0, NPTS // 3] += 1e-9 + with pytest.raises(NotImplementedError): + tmq.time_marginalize_bandlimited(k, r2, DELTAT, _lnL) + + +def _n_resolved(rep): + """Rows with a real peak that simply needed no refinement.""" + return (rep['n_rows'] - rep['n_refined_rows'] + - rep['n_unmeasurable_rows'] - rep['n_flat_rows']) + + +def test_the_boundary_diagnostic_covers_the_RIGHT_edge_too(): + """Both physical integration boundaries are reported, not rule switches.""" + for peak in (NPTS - 1.3, NPTS - 3.3, NPTS - 31.3): + sig = BandLimited(amp=1.0, peak_sample=peak, n_period=8 * NPTS, + m_hi=1400, background=0.12) + k = sig.samples() + assert _bandlimited(k) != _simpson_value(k), peak + assert tmq.last_report()['n_wrap_exposed_rows'] == 1, peak + assert tmq.last_report()['n_refined_rows'] == 1, peak + # A central peak is reconstructed by the same rule without the diagnostic. + inside = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, n_period=8 * NPTS, + m_hi=1400, background=0.12) + _bandlimited(inside.samples()) + assert tmq.last_report()['n_wrap_exposed_rows'] == 0 + + +@pytest.mark.parametrize("phase_marg", [False, True]) +def test_phase_marginalization_reaches_the_new_path(phase_marg): + """`--distance-marginalization --phase-marginalization` is the standard + production call site, and it passes `phase_marginalization=True` with a + NONLINEAR callback. Every original fixture used the affine helper with + `kappa.real`, for which lnL(t) is itself exactly band-limited -- so dropping + the `abs()` entirely changed nothing any test could see.""" + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + term = (lambda z: np.abs(z)) if phase_marg else (lambda z: z.real) + + got = float(tmq.time_marginalize_bandlimited( + k, r, DELTAT, _lnL, phase_marginalization=phase_marg)[0]) + + # analytic truth for THIS integrand + n = (NPTS - 1) * 128 + 1 + td = sig.j0 * DELTAT + np.arange(n) * (DELTAT / 128) + ref = _log_trapz(_lnL(term(sig.at(td)), RHO_SQ), DELTAT / 128) + assert abs(got - ref) < 1e-3, (phase_marg, got - ref) + + # the two settings must actually differ, or the parametrisation proves nothing + other = float(tmq.time_marginalize_bandlimited( + k, r, DELTAT, _lnL, phase_marginalization=not phase_marg)[0]) + assert abs(got - other) > 1e-3, "abs() vs real() made no difference" + + +def test_a_nonlinear_distance_marginalization_style_callback(): + """The production callback is a table interpolation, not `kappa - rho_sq/2`. + For the affine helper lnL(t) is itself band-limited, which is a much easier + problem than the one production actually poses.""" + def distmarg_like(x, rho_sq): + # monotone, nonlinear, and -inf outside a table range, like the real one + z = np.asarray(x) / np.sqrt(np.asarray(rho_sq) + 1.0) + out = np.where(z > -3.0, np.log1p(np.exp(np.clip(z, -50, 50))) * 40.0, -np.inf) + return out + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples()[None, :] + r = np.full(k.shape, RHO_SQ) + got = float(tmq.time_marginalize_bandlimited(k, r, DELTAT, distmarg_like)[0]) + n = (NPTS - 1) * 128 + 1 + td = sig.j0 * DELTAT + np.arange(n) * (DELTAT / 128) + ref = _log_trapz(distmarg_like(sig.at(td).real, RHO_SQ), DELTAT / 128) + simp = _log_simps(distmarg_like(k[0].real, RHO_SQ), DELTAT) + assert abs(got - ref) < 1e-2, got - ref + assert abs(got - ref) < 0.05 * abs(simp - ref) + + +def test_the_memory_chunking_path_assembles_its_result(): + """Production runs `--n-chunk 10000`, so EVERY real call chunks; the suite's + largest batch is 4 rows, so the assembly branch never ran. Dropping all but + the first chunk was invisible.""" + rows = [BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.1 * i, + n_period=8 * NPTS, m_hi=1400, background=0.12, + seed=7 + i).samples() for i in range(6)] + k = np.stack(rows) + r = np.full(k.shape, RHO_SQ) + whole = np.asarray(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL)) + old = tmq._DENSE_CHUNK_BYTES + try: + tmq._DENSE_CHUNK_BYTES = 4096 # force several chunks per group + chunked = np.asarray(tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL)) + finally: + tmq._DENSE_CHUNK_BYTES = old + assert chunked.shape == whole.shape == (6,) + # NOT bit-identity. The batch shape reaches numpy's FFT and its pairwise + # summation, so a differently-chunked run reassociates and can differ in the + # last bit or two -- the companion peak-local implementation measured 0, 0 + # and 2 ULPs on the equivalent test. This fixture happens to come out + # bit-identical, which is exactly why asserting it would be a latent flake: + # it would pass here and fail on a different row count or chunk boundary. + # Assert what is actually guaranteed, at a bound far below anything that + # could hide a real assembly bug (dropping a chunk moves rows by nats). + tol = 64 * np.spacing(np.abs(whole).max()) + assert np.allclose(chunked, whole, rtol=0, atol=tol), np.abs(chunked - whole).max() + + +def test_the_one_ulp_factor_bump_is_exercised(): + """`2**ceil(log2(need))` can land one power of two SHORT when log2 rounds down + for a `need` a hair above a power of two -- erring LOW, i.e. silently + under-resolving. A log-spaced sweep never lands there; this does.""" + dx = DELTAT + for kexp in range(0, 12): + need = 2.0 ** kexp + sigma = tmq.UPSAMPLE_SAFETY * dx / np.nextafter(need, np.inf) + f = int(tmq.required_upsample_factors(np.array([sigma]), dx)[0]) + assert f >= tmq.UPSAMPLE_SAFETY * dx / sigma, (kexp, f) + assert dx / f <= sigma / tmq.UPSAMPLE_SAFETY, (kexp, f) + + +def test_argmax_ignores_non_finite_bins(): + """`argmax` over a raw array containing NaN returns the NaN's index, which + would put the whole width measurement on a bin that carries no likelihood.""" + t = (np.arange(NPTS) - NPTS // 2) * DELTAT + lnL = -0.5 * (t / (0.3 * DELTAT)) ** 2 + lnL[10] = np.nan + sigma, jmax, meas = tmq.peak_width_from_lnL(lnL[None, :], DELTAT) + assert int(jmax[0]) == NPTS // 2, jmax + assert bool(meas[0]) and np.isclose(float(sigma[0]), 0.3 * DELTAT, rtol=1e-9) + + +def test_report_sigma_t_min_is_the_width_that_was_resolved(): + sig = BandLimited(amp=0.3, peak_sample=NPTS // 2 + 0.25) + k = sig.samples()[None, :] + tmq.time_marginalize_bandlimited(k, np.full(k.shape, RHO_SQ), DELTAT, _lnL) + rep = tmq.last_report() + coarse, _, _ = tmq.peak_width_from_lnL(_lnL(k.real, RHO_SQ), DELTAT) + assert np.isfinite(rep['sigma_t_min']) + assert abs(rep['sigma_t_min'] - float(coarse[0])) < 0.2 * float(coarse[0]), rep + assert DELTAT / rep['upsample_factor'] <= rep['sigma_t_min'] / tmq.UPSAMPLE_SAFETY + + +def test_the_tuned_constants_are_pinned_to_their_measured_values(): + """These are not free parameters. Each is justified by a measured table in + DESIGN_time_marginalization_quadrature.md, and the suite otherwise pins them + only to within a factor of ~10 -- so changing one could pass CI while + invalidating the argument behind it. Changing a value here is the deliberate + act of also updating that table.""" + assert tmq.UPSAMPLE_SAFETY == 2.0 + assert tmq.EDGE_GUARD_FRACTION == 0.125 + assert tmq.UPSAMPLE_FACTOR_MAX == 4096 + assert tmq.CURVATURE_STENCIL_HALFWIDTHS == (1, 2, 4, 8) + + +def test_driver_banner_reports_what_is_ACTUALLY_IN_FORCE(): + """The driver's single load-bearing line is the assignment to + `factored_likelihood.TIME_QUADRATURE_DEFAULT`. Deleting it, or hardcoding + 'simpson' there, leaves the flag inert -- and every banner-inspecting test + stays green if the banner is built from `opts`. The banner therefore prints + the value READ BACK out of the module, and this asserts on that.""" + rc, out = _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + _HONOURED) + assert 'Time-marginalization quadrature: bandlimited' in out, out[-2000:] + import re + m = re.search(r'Time-marginalization quadrature: (\S+) \(from', out) + assert m and m.group(1) == 'bandlimited', out[-2000:] + src = open(os.path.join(os.path.dirname(os.path.dirname( + os.path.abspath(__file__))), 'bin', + 'integrate_likelihood_extrinsic_batchmode')).read() + assert 'factored_likelihood.TIME_QUADRATURE_DEFAULT,' in src, \ + "the banner no longer reads the value back out of the module" + + +def test_driver_does_not_refuse_an_ordinary_default_run(): + """The refuse-guard is gated on the option being non-default. Dropping that + gate makes the driver reject EVERY ordinary ILE run that is not + `--time-marginalization --vectorized --gpu` -- and no driver test ever + launched a default configuration, so nothing noticed.""" + for args in ([], ['--time-marginalization'], ['--vectorized']): + rc, out = _run_driver(args) + assert 'cannot honour it' not in out, (args, out[-2000:]) + assert 'Time-marginalization quadrature: simpson' in out, (args, out[-2000:]) + + +def _quadrature_banner(out): + """The quadrature banner line, matched SPECIFICALLY. + + The pre-existing `--interpolate-time` banner carries the identical phrase + "honoured by this configuration", so a bare substring test matches whichever + line happens to say what you were hoping for. A mutation making the + quadrature banner claim `True` unconditionally survived exactly that way: + the stencil line still said `False` and the assertion passed. + """ + import re + m = re.search(r'^\s*Time-marginalization quadrature: (\S+) ' + r'\(from --time-marginalization-quadrature (.+?)\); ' + r'honoured by this configuration: (True|False)\s*$', + out, re.MULTILINE) + assert m is not None, "no quadrature banner line found:\n" + out[-3000:] + return m.group(1), m.group(3) + + +def test_driver_banner_does_not_claim_to_honour_what_it_cannot(): + quad, honoured = _quadrature_banner(_run_driver(['--time-marginalization'])[1]) + assert (quad, honoured) == ('simpson', 'False') + quad, honoured = _quadrature_banner(_run_driver(_HONOURED)[1]) + assert (quad, honoured) == ('simpson', 'True') + quad, honoured = _quadrature_banner( + _run_driver(['--time-marginalization-quadrature', 'bandlimited'] + _HONOURED)[1]) + assert (quad, honoured) == ('bandlimited', 'True') + + +def test_the_edge_guard_band_is_exactly_the_outer_fraction(): + """Pin both boundaries to the sample, not merely "near the edge". + + An off-by-one in the upper term -- `jmax > npts - guard` instead of + `npts - 1 - guard` -- leaves exactly one row's worth of the right guard band + open, and every peak-placement fixture is far enough inside that both spellings + agree. Driving the argmax to a chosen bin makes the boundary itself the + subject. + """ + guard = max(1, int(NPTS * tmq.EDGE_GUARD_FRACTION)) + + def row_peaking_at(j): + t = np.arange(NPTS, dtype=float) + return (np.exp(-0.5 * ((t - j) / 0.35) ** 2) * 40.0).astype(complex) + + # The last EXPOSED index and the first ACCEPTED one, at both ends. These four + # are what an off-by-one in either term moves, and every peak-placement + # fixture elsewhere is far enough inside that both spellings agree. + for j, expect_exposed in ((guard - 1, True), (guard, False), + (NPTS - 1 - guard, False), (NPTS - guard, True)): + k = row_peaking_at(j)[None, :] + r = np.full(k.shape, RHO_SQ) + sigma, jmax, meas = tmq.peak_width_from_lnL(_lnL(k.real, r), DELTAT) + assert int(jmax[0]) == j and np.isfinite(sigma[0]), (j, jmax, sigma) + tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL) + rep = tmq.last_report() + assert (rep['n_wrap_exposed_rows'] == 1) == expect_exposed, (j, guard, rep) + # All four rows are sharp enough to be refined; the boundary changes + # only the diagnostic count, never the quadrature rule. + assert rep['n_refined_rows'] == 1, (j, rep) + + # A peak on the very first or last SAMPLE is a documented corner: the + # curvature stencil is clipped inward and initially measures positive + # curvature away from the maximum. A strongly varying row must not be + # confused with a genuinely flat antenna null; it receives a seed factor and + # lets dense remeasurement derive the eventual resolution. + for j in (0, NPTS - 1): + k = row_peaking_at(j)[None, :] + r = np.full(k.shape, RHO_SQ) + out = tmq.time_marginalize_bandlimited(k, r, DELTAT, _lnL) + rep = tmq.last_report() + assert rep['n_refined_rows'] == 1, (j, rep) + assert rep['n_wrap_exposed_rows'] == 1, (j, rep) + assert rep['n_flat_rows'] == 0, (j, rep) + assert rep['upsample_factor'] >= 16, (j, rep) + assert float(out[0]) != _simpson_value(k[0]), j + + +if __name__ == '__main__': + raise SystemExit(pytest.main([__file__, '-q'])) diff --git a/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature_pipeline.py b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature_pipeline.py new file mode 100644 index 000000000..cbd8f16d3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature_pipeline.py @@ -0,0 +1,578 @@ +#!/usr/bin/env python +"""Pipeline passthrough for --time-marginalization-quadrature. + +Companion to test_time_marginalization_quadrature.py, which covers the quadrature +itself. This file covers the WIRING: that a campaign can select the quadrature +without --manual-extra-ile-args, that the selection reaches every ILE submit file +INCLUDING ILE_extr.sub, that the default path emits nothing, and that a +configuration which cannot honour the request is REFUSED at DAG-build time rather +than at first-job time. + +Why the wiring needs its own tests. The option is inert unless it survives four +hand-offs -- util_RIFT_pseudo_pipe.py -> helper_LDG_Events.py -> args_ile.txt -> +create_event_parameter_pipeline_BasicIteration -> ILE*.sub -- and the last of +those is an INHERITANCE (`ile_args_extr = ile_args + ...`), not an explicit +forward, so it is exactly the kind of link that a refactor breaks silently. A +test of the quadrature helper alone would stay green through all of it. +""" +import ast +import gzip +import os +import subprocess +import sys + +import pytest + +from RIFT.likelihood.time_marginalization_quadrature import ( + TIME_QUADRATURE_CHOICES, time_quadrature_pipeline_prereqs, + find_time_quadrature_in_ile_args, refuse_unhonourable_time_quadrature, + refuse_unless_time_quadrature_emitted) + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")) +CODE_DIR = os.path.join(REPO_ROOT, "MonteCarloMarginalizeCode", "Code") +BIN_DIR = os.path.join(CODE_DIR, "bin") +PSEUDO_PIPE = os.path.join(BIN_DIR, "util_RIFT_pseudo_pipe.py") +HELPER = os.path.join(BIN_DIR, "helper_LDG_Events.py") +CEPP = os.path.join(BIN_DIR, "create_event_parameter_pipeline_BasicIteration") +ILE_EXE = os.path.join(BIN_DIR, "integrate_likelihood_extrinsic_batchmode") + +GOOD_ILE_ARGS = ("integrate_likelihood_extrinsic_batchmode --time-marginalization " + "--vectorized --gpu --srate 4096 --n-eff 50") + + +# ---------------------------------------------------------------- prerequisites + +def test_simpson_is_never_refused(): + """The default must never be able to fail a workflow build. Even a + configuration that excludes 'bandlimited' entirely is fine for 'simpson', + because 'simpson' is what ILE does anyway.""" + assert time_quadrature_pipeline_prereqs('simpson', "--rotation-slow --freqresponse") == [] + assert time_quadrature_pipeline_prereqs('simpson', "") == [] + + +def test_honourable_configuration_passes(): + assert time_quadrature_pipeline_prereqs('bandlimited', GOOD_ILE_ARGS) == [] + # optparse has no six-character floor: these are the shortest unique + # spellings in ILE's current option set. + assert time_quadrature_pipeline_prereqs( + 'bandlimited', "X --time-marginalization --vec --g") == [] + + +@pytest.mark.parametrize("flag", ["--time-marginalization", "--vectorized", "--gpu"]) +def test_each_required_flag_is_reported_when_missing(flag): + args = " ".join(t for t in GOOD_ILE_ARGS.split() if t != flag) + missing = time_quadrature_pipeline_prereqs('bandlimited', args) + assert any(flag in m for m in missing), missing + + +@pytest.mark.parametrize("flag,value", [ + ("--rotation-slow", ""), + ("--freqresponse", ""), + ("--calibration-envelope-directory", " /tmp/cal"), +]) +def test_each_excluding_flag_is_reported_when_present(flag, value): + missing = time_quadrature_pipeline_prereqs('bandlimited', GOOD_ILE_ARGS + " " + flag + value) + assert any(flag in m for m in missing), missing + + +@pytest.mark.parametrize("spelling", [ + "--calibration-envelope-directory=/tmp/cal", # optparse accepts the equals form + "--rotation-s", # shortest unique prefix today + "--calibration-en=/tmp/cal", # shortest unique prefix today + "'--rotation-slow'", # an ini leaves the quotes on +]) +def test_legal_optparse_spellings_do_not_evade_the_exclusions(spelling): + """Three spellings that set the excluded option and that a naive whitespace + split does not see. Each costs a queue-slot cycle if it reaches the driver, + which is the cost this DAG-build guard exists to avoid.""" + missing = time_quadrature_pipeline_prereqs('bandlimited', GOOD_ILE_ARGS + " " + spelling) + assert missing, spelling + + +@pytest.mark.parametrize("innocent", [ + "--no-gpu", "--gpu-fanout 2", "--rotation-slow-foo", "--calibration-n-realizations 100", +]) +def test_exclusions_do_not_fire_on_lookalike_flags(innocent): + """The exclusion side must be a token/abbreviation match, not a substring one: + a substring test would falsely refuse any future option whose name contains + one of these. This is the half of the matching rule that had no test.""" + assert time_quadrature_pipeline_prereqs('bandlimited', GOOD_ILE_ARGS + " " + innocent) == [] + + +def test_match_is_by_token_not_substring(): + """'--no-gpu' must not satisfy '--gpu', and the quadrature flag itself must not + satisfy '--time-marginalization'. A substring test passes both and would + declare an unhonourable configuration fine.""" + args = ("integrate_likelihood_extrinsic_batchmode --time-marginalization-quadrature " + "bandlimited --vectorized --no-gpu") + missing = time_quadrature_pipeline_prereqs('bandlimited', args) + assert any("--gpu" in m and "missing" in m for m in missing), missing + assert any("--time-marginalization " in m or m.startswith("missing --time-marginalization (") + for m in missing), missing + + +def test_bad_value_is_refused(): + with pytest.raises(ValueError): + time_quadrature_pipeline_prereqs('bandlimted', GOOD_ILE_ARGS) + with pytest.raises(ValueError): + time_quadrature_pipeline_prereqs('True', GOOD_ILE_ARGS) + + +# ------------------------------------------------------------- static wiring + +def _source(path): + with open(path) as f: + return f.read() + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_option_is_defined_with_a_none_default(path): + """Default None means "pass nothing", so the default workflow is byte-identical + to one built before this option existed.""" + src = _source(path) + assert '"--internal-ile-time-marginalization-quadrature"' in src + line = [l for l in src.splitlines() + if '"--internal-ile-time-marginalization-quadrature"' in l][0] + assert "default=None" in line, line + assert "type=str" in line, line + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_choices_are_imported_not_retyped(path): + """A second hand-typed copy of the choice tuple is how a typo becomes a + silently different likelihood: the pipeline would accept it, forward it, and + the mistake would surface only when the first ILE job died.""" + src = _source(path) + assert "TIME_QUADRATURE_CHOICES" in src + for literal in ("('simpson', 'bandlimited')", '("simpson", "bandlimited")', + "'simpson','bandlimited'", '"simpson","bandlimited"'): + assert literal not in src, "choice tuple re-typed in %s: %r" % (path, literal) + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_ini_override_is_recorded_in_the_help(path): + """The RIFT ini parser OVERRIDES the command line for non-boolean options, and + this is a string option, so an ini that also sets it wins silently.""" + line = [l for l in _source(path).splitlines() + if '"--internal-ile-time-marginalization-quadrature"' in l][0] + assert "ini" in line.lower() and "override" in line.lower(), line + + +def _assign_targets_containing(path, needle): + """Names assigned (`=` or `+=`) a string containing `needle`. Asserting the + TARGET, not just the presence of the literal, is what catches the refactor that + appends the flag to a variable nothing writes out. Both assignment forms are + accepted: the helper's emission is a plain `=` with an explicit rstrip(), so a + test that looked only at `+=` broke on that change -- and, because the mutation + harness runs this same file, silently turned every mutation into a false kill.""" + tree = ast.parse(_source(path), filename=path) + out = set() + for node in ast.walk(tree): + if isinstance(node, ast.AugAssign): + targets = [node.target] + elif isinstance(node, ast.Assign): + targets = node.targets + else: + continue + names = [t.id for t in targets if isinstance(t, ast.Name)] + if not names: + continue + for sub in ast.walk(node.value): + if isinstance(sub, ast.Constant) and isinstance(sub.value, str) and needle in sub.value: + out.update(names) + return out + + +def test_pseudo_pipe_forwards_to_the_helper(): + """It must land on the helper_LDG_Events.py command line -- the helper owns ILE + argument construction, so a flag appended anywhere else never reaches ILE.""" + targets = _assign_targets_containing( + PSEUDO_PIPE, "--internal-ile-time-marginalization-quadrature") + assert "cmd" in targets, targets + + +def test_helper_emits_the_ile_flag(): + """It must land on helper_ile_args, which is what becomes args_ile.txt.""" + targets = _assign_targets_containing(HELPER, "--time-marginalization-quadrature") + assert "helper_ile_args" in targets, targets + + +def _dead_nodes(tree): + """Every node inside an `if :` body -- i.e. unreachable code.""" + dead = set() + for node in ast.walk(tree): + if isinstance(node, ast.If) and isinstance(node.test, ast.Constant) \ + and not node.test.value: + for stmt in node.body: + for sub in ast.walk(stmt): + dead.add(sub) + return dead + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_the_refusal_call_site_is_reachable(path): + """util_RIFT_pseudo_pipe.py cannot be run to completion in a test -- its late + guard sits after an os.system() helper call that needs real data -- so that one + call site has no executable coverage and a plain "is it called" ast walk passes + even when the call is wrapped in `if False:`. This asserts the call is in LIVE + code. The helper's guard is covered executably as well, below.""" + tree = ast.parse(_source(path), filename=path) + dead = _dead_nodes(tree) + live = [n for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == "refuse_unless_time_quadrature_emitted" and n not in dead] + assert live, "no reachable call to refuse_unless_time_quadrature_emitted in %s" % path + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_both_layers_call_the_refusal(path): + """Structural companion to the EXECUTED refusal tests below. On its own this + is presence, not effect -- a guard turned into a print still passes it -- which + is why the executed tests exist; it is kept only to catch the call site being + deleted outright, which is cheaper to diagnose here.""" + tree = ast.parse(_source(path), filename=path) + called = {n.func.id for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)} + assert "refuse_unless_time_quadrature_emitted" in called, called + + +def test_choices_argparse_surface_matches_the_library(): + """--help must offer exactly the library's choices, not a subset frozen at the + time this option was written.""" + env = dict(os.environ, PYTHONPATH=CODE_DIR + os.pathsep + os.environ.get("PYTHONPATH", "")) + out = subprocess.run([sys.executable, PSEUDO_PIPE, "--help"], env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + universal_newlines=True).stdout + assert "--internal-ile-time-marginalization-quadrature" in out + for choice in TIME_QUADRATURE_CHOICES: + assert choice in out + + +def test_a_bad_value_is_rejected_by_the_pipeline_command_line(): + env = dict(os.environ, PYTHONPATH=CODE_DIR + os.pathsep + os.environ.get("PYTHONPATH", "")) + proc = subprocess.run( + [sys.executable, PSEUDO_PIPE, "--internal-ile-time-marginalization-quadrature", "bandlimted"], + env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True) + assert proc.returncode != 0 + assert "bandlimted" in proc.stdout + + +# --------------------------------------------------- end-to-end: args_ile -> .sub + +def _write_grid(tmp_path): + import lal + import RIFT.lalsimutils as lsu + P = lsu.ChooseWaveformParams() + P.m1 = 35 * lal.MSUN_SI + P.m2 = 30 * lal.MSUN_SI + here = os.getcwd() + os.chdir(os.fspath(tmp_path)) + try: + lsu.ChooseWaveformParams_array_to_xml([P, P], "proposed-grid") + finally: + os.chdir(here) + return os.fspath(tmp_path / "proposed-grid.xml.gz") + + +def _build_dag(tmp_path, ile_args_line, extra_cepp=()): + """Run the real DAG builder on a hand-written args_ile.txt. Returns the + working directory holding the generated .sub files.""" + (tmp_path / "args_ile.txt").write_text(ile_args_line + "\n") + (tmp_path / "args_cip_list.txt").write_text( + "2 --parameter mc --parameter delta_mc --n-output-samples 5000\n") + (tmp_path / "args_test.txt").write_text("X --always-succeed\n") + grid = _write_grid(tmp_path) + env = dict(os.environ) + env["PYTHONPATH"] = CODE_DIR + os.pathsep + env.get("PYTHONPATH", "") + env["PATH"] = BIN_DIR + os.pathsep + env.get("PATH", "") + proc = subprocess.run( + [sys.executable, CEPP, + "--ile-n-events-to-analyze", "1", + "--input-grid", grid, + "--ile-exe", ILE_EXE, + "--ile-args", os.fspath(tmp_path / "args_ile.txt"), + "--cip-args-list", os.fspath(tmp_path / "args_cip_list.txt"), + "--test-args", os.fspath(tmp_path / "args_test.txt"), + "--working-directory", os.fspath(tmp_path), + "--n-iterations", "2", + "--n-samples-per-job", "500", + "--last-iteration-extrinsic", + "--last-iteration-extrinsic-samples-per-ile", "200"] + list(extra_cepp), + cwd=os.fspath(tmp_path), env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, universal_newlines=True) + assert proc.returncode == 0, proc.stdout[-4000:] + return tmp_path + + +def test_quadrature_reaches_both_ile_and_ile_extr_sub(tmp_path): + """The load-bearing one. ILE_extr.sub is built by INHERITING the whole + main-iteration argument string (`ile_args_extr = ile_args + ...`), not by an + explicit forward -- so the extrinsic stage silently keeping the historical + quadrature is a live failure mode and this is what rules it out.""" + wd = _build_dag(tmp_path, GOOD_ILE_ARGS + " --time-marginalization-quadrature bandlimited") + for name in ("ILE.sub", "ILE_extr.sub"): + text = (wd / name).read_text() + assert "--time-marginalization-quadrature bandlimited" in text, name + + +def test_default_path_emits_no_quadrature_flag(tmp_path): + """With the pipeline option unset the helper writes nothing, so no ILE submit + file mentions the quadrature at all and the default run is unchanged.""" + wd = _build_dag(tmp_path, GOOD_ILE_ARGS) + for name in ("ILE.sub", "ILE_extr.sub"): + assert "--time-marginalization-quadrature" not in (wd / name).read_text(), name + + +# ------------------------------------------------- the guard reads the BYTES + +def test_prereq_check_alone_approves_args_that_never_got_the_flag(): + """Documents WHY refuse_unless_time_quadrature_emitted exists. The prerequisite + check reads the prerequisites from the argument string but the INTENT from the + caller, so on its own it approves an args_ile.txt that never received the flag.""" + assert time_quadrature_pipeline_prereqs('bandlimited', GOOD_ILE_ARGS) == [] + + +def test_emission_guard_refuses_args_that_never_got_the_flag(): + """The stale-helper_ile_args.txt / dropped-emission case: prerequisites all + satisfied, request recorded, flag absent -> every ILE job would run Simpson.""" + with pytest.raises(ValueError) as e: + refuse_unless_time_quadrature_emitted('bandlimited', GOOD_ILE_ARGS, "args_ile.txt") + assert "contains no --time-marginalization-quadrature" in str(e.value) + + +def test_emission_guard_refuses_a_duplicate_because_optparse_takes_the_last(): + """--manual-extra-ile-args is appended AFTER the helper's arguments, so a + hand-passed 'simpson' silently wins while the .sub file shows both -- which is + exactly the case that falsifies "readable off the .sub file".""" + args = (GOOD_ILE_ARGS + " --time-marginalization-quadrature bandlimited" + " --time-marginalization-quadrature simpson") + with pytest.raises(ValueError) as e: + refuse_unless_time_quadrature_emitted('bandlimited', args, "args_ile.txt") + assert "occurrences" in str(e.value) + + +def test_emission_guard_refuses_a_value_that_does_not_match_the_request(): + args = GOOD_ILE_ARGS + " --time-marginalization-quadrature simpson" + with pytest.raises(ValueError): + refuse_unless_time_quadrature_emitted('bandlimited', args, "args_ile.txt") + + +def test_emission_guard_holds_a_hand_passed_quadrature_to_the_same_standard(): + """The manual route (--manual-extra-ile-args / an ini) got no protection at all + while the guard was keyed on the pipeline option being set.""" + args = ("X --time-marginalization --vectorized --gpu --rotation-slow" + " --time-marginalization-quadrature bandlimited") + with pytest.raises(ValueError) as e: + refuse_unless_time_quadrature_emitted(None, args, "args_ile.txt") + assert "--rotation-slow" in str(e.value) + + +def test_emission_guard_is_silent_on_the_default_path(): + """Nothing requested, nothing present: the default workflow must not raise.""" + refuse_unless_time_quadrature_emitted(None, GOOD_ILE_ARGS, "args_ile.txt") + + +def test_emission_guard_accepts_the_honoured_case(): + refuse_unless_time_quadrature_emitted( + 'bandlimited', GOOD_ILE_ARGS + " --time-marginalization-quadrature bandlimited", + "args_ile.txt") + + +def test_find_handles_the_equals_form(): + assert find_time_quadrature_in_ile_args( + "X --time-marginalization-quadrature=bandlimited") == ['bandlimited'] + # ILE uses optparse, which accepts these unique-prefix forms. The guard must + # see the same option or a manual abbreviation bypasses all prerequisites. + assert find_time_quadrature_in_ile_args( + "X --time-marginalization-q=bandlimited") == ['bandlimited'] + assert find_time_quadrature_in_ile_args( + "X --time-marginalization- bandlimited") == ['bandlimited'] + # The exact boolean flag wins as an exact optparse match; it is not an + # abbreviation of the quadrature option. + assert find_time_quadrature_in_ile_args( + "X --time-marginalization --vectorized --gpu") == [] + + +def test_abbreviated_hand_passed_quadrature_cannot_evade_prerequisites(): + args = ("X --time-marginalization --vectorized --rotation-slow " + "--time-marginalization-q=bandlimited") + with pytest.raises(ValueError) as e: + refuse_unless_time_quadrature_emitted(None, args, "args_ile.txt") + assert "--gpu" in str(e.value) + assert "--rotation-slow" in str(e.value) + + +def test_refusal_actually_raises(): + """Executable coverage of the raise itself. It lives in the library precisely + so that turning it into a print is a code change a test can see -- an ast walk + for a call by name cannot.""" + with pytest.raises(ValueError): + refuse_unhonourable_time_quadrature('bandlimited', "X --vectorized", "somewhere") + refuse_unhonourable_time_quadrature('bandlimited', GOOD_ILE_ARGS, "somewhere") + refuse_unhonourable_time_quadrature('simpson', "X --rotation-slow", "somewhere") + + +# ------------------------------------ executable: the real scripts, real bytes +# +# The four hops are pseudo_pipe -> helper -> args_ile.txt -> ILE*.sub. The DAG +# tests above cover the last hop only, because they hand-write args_ile.txt. These +# run the real scripts. helper_LDG_Events.py needs no data if given --fake-data +# and a manual IFO list, which is what makes hops 1->2 and 2->3 testable at all. + +HELPER_BASE = [ + "--event-time", "1240000000", "--fmin", "20", "--fmin-template", "20", + "--manual-ifo-list", "['H1','L1']", "--fake-data", "--assume-fiducial-psd-files", + "--data-start-time", "1239999996", "--data-end-time", "1240000004", + "--force-notune-initial-grid", "--propose-fit-strategy", +] + + +def _run_helper(tmp_path, *extra): + env = dict(os.environ) + env["PYTHONPATH"] = CODE_DIR + os.pathsep + env.get("PYTHONPATH", "") + env["PATH"] = BIN_DIR + os.pathsep + env.get("PATH", "") + cmd = [sys.executable, HELPER, "--working-directory", os.fspath(tmp_path)] \ + + HELPER_BASE + list(extra) + return subprocess.run(cmd, cwd=os.fspath(tmp_path), env=env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, universal_newlines=True) + + +def test_helper_emits_the_requested_value_not_a_hardcoded_one(tmp_path): + """Hop 2->3, executed. The static test asserts only that the flag NAME is + appended somewhere; a helper that emitted a hardcoded 'simpson' regardless of + the request passed it. That is the exact failure this whole effort is about: + the user asks for bandlimited, the pipeline prints bandlimited, the .sub says + simpson.""" + proc = _run_helper(tmp_path, "--propose-ile-convergence-options", + "--internal-ile-time-marginalization-quadrature", "bandlimited") + assert proc.returncode == 0, proc.stdout[-3000:] + args = (tmp_path / "helper_ile_args.txt").read_text() + assert find_time_quadrature_in_ile_args(args) == ["bandlimited"], args[-400:] + # and the emission must not have been concatenated onto its neighbour + assert " --time-marginalization-quadrature bandlimited " in args + " " + + +def test_helper_default_emits_nothing(tmp_path): + proc = _run_helper(tmp_path, "--propose-ile-convergence-options") + assert proc.returncode == 0, proc.stdout[-3000:] + args = (tmp_path / "helper_ile_args.txt").read_text() + assert find_time_quadrature_in_ile_args(args) == [] + + +def test_helper_refuses_a_configuration_it_cannot_honour(tmp_path): + """Executed refusal. Without --propose-ile-convergence-options the helper never + adds --time-marginalization/--vectorized/--gpu, so the request cannot be + honoured. A guard turned into a print, or disabled with `if False`, passes the + ast test and fails this one.""" + proc = _run_helper(tmp_path, + "--internal-ile-time-marginalization-quadrature", "bandlimited") + assert proc.returncode != 0, proc.stdout[-3000:] + assert "--time-marginalization" in proc.stdout + assert not (tmp_path / "helper_ile_args.txt").exists() + + +def test_helper_records_the_continuous_extrinsic_time_contract(tmp_path): + """Sub-sample integration also means a continuous conditional time draw at + the fairdraw export stage; keep that subtle contract in the build log.""" + proc = _run_helper(tmp_path, "--propose-ile-convergence-options", + "--internal-ile-time-marginalization-quadrature", "bandlimited") + assert proc.returncode == 0 + assert "continuous draw from p(t | extrinsic, data)" in proc.stdout + assert "not quantised at 1/srate" in proc.stdout + + +def _run_pseudo_pipe(tmp_path, *extra): + env = dict(os.environ) + env["PYTHONPATH"] = CODE_DIR + os.pathsep + env.get("PYTHONPATH", "") + # deliberately WITHOUT BIN_DIR on PATH for the forward test: the helper is + # invoked by name, so it fails, and we read the command line it printed. + cmd = [sys.executable, PSEUDO_PIPE, "--approx", "SEOBNRv4", + "--use-rundir", os.fspath(tmp_path / "run")] + list(extra) + return subprocess.run(cmd, cwd=os.fspath(tmp_path), env=env, stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, universal_newlines=True) + + +def test_pseudo_pipe_forwards_the_requested_value_to_the_helper(tmp_path): + """Hop 1->2, executed. pseudo_pipe prints the helper command line it is about + to run; a forward that hardcoded 'simpson' passed the static test.""" + proc = _run_pseudo_pipe( + tmp_path, "--internal-ile-time-marginalization-quadrature", "bandlimited") + assert "--internal-ile-time-marginalization-quadrature bandlimited" in proc.stdout, \ + proc.stdout[-3000:] + + +def test_pseudo_pipe_checks_the_helper_exit_status(): + """A same-value stale helper_ile_args.txt can satisfy the byte guard. The + helper's status therefore has to be checked independently, before that file + is read.""" + src = _source(PSEUDO_PIPE) + call = src.index("_helper_rc = os.system(cmd)") + check = src.index("if _helper_rc != 0:", call) + read = src.index('np.loadtxt("helper_ile_args.txt"', check) + assert call < check < read + assert "os.unlink('helper_ile_args.txt')" in src[:call] + + +def test_pseudo_pipe_refuses_calmarg_before_it_runs_anything(tmp_path): + """Executed refusal, and it must fire EARLY -- calibration marginalization is + added by this script, not by the helper, so the helper can never see it.""" + proc = _run_pseudo_pipe( + tmp_path, "--internal-ile-time-marginalization-quadrature", "bandlimited", + "--calmarg-envelope-directory", os.fspath(tmp_path)) + assert proc.returncode != 0 + assert "--calibration-envelope-directory" in proc.stdout + assert "helper_LDG_Events.py --force-notune" not in proc.stdout, \ + "refusal must precede the helper invocation" + + +def test_pseudo_pipe_refuses_an_excluded_manual_extra_ile_arg(tmp_path): + proc = _run_pseudo_pipe( + tmp_path, "--internal-ile-time-marginalization-quadrature", "bandlimited", + "--manual-extra-ile-args=--rotation-slow") + assert proc.returncode != 0 + assert "--rotation-slow" in proc.stdout + + +def test_pseudo_pipe_refuses_on_the_lisa_known_sky_path(tmp_path): + """--lisa-known-sky exits before both the forward and the args_ile.txt guard and + builds its own ILE arguments through helper_LISA_Events.py, which does not carry + the option -- so without this the request is silently dropped.""" + proc = _run_pseudo_pipe( + tmp_path, "--internal-ile-time-marginalization-quadrature", "bandlimited", + "--lisa-known-sky", "--event-time", "1234.5", + "--ecliptic-longitude", "1.25", "--ecliptic-latitude", "-0.4") + assert proc.returncode != 0 + assert "lisa-known-sky" in proc.stdout + + +def test_pseudo_pipe_rejects_a_bad_value_after_the_ini_block(tmp_path): + """The validate must run AFTER --use-ini, which overrides the command line for + non-boolean options; above it, it checks a value the ini is about to replace.""" + src = _source(PSEUDO_PIPE) + i_ini = src.index("if (opts.use_ini):") + i_val = src.index("validate_time_quadrature(opts.internal_ile_time_marginalization_quadrature)") + assert i_val > i_ini, "validate_time_quadrature runs before the ini block can override it" + + +@pytest.mark.parametrize("path", [PSEUDO_PIPE, HELPER]) +def test_argparse_choices_are_pinned_to_the_library_tuple(path): + """A grep for the identifier survives deleting `choices=list(...)`, because the + help string still interpolates the tuple.""" + line = [l for l in _source(path).splitlines() + if '"--internal-ile-time-marginalization-quadrature"' in l][0] + assert "choices=list(TIME_QUADRATURE_CHOICES)" in line, line + + +def test_quadrature_also_reaches_puff_and_fetch_subs(tmp_path): + """ILE_puff.sub and ILE_fetch.sub inherit the same argument string + (ile_args_forpuff / ile_args_forfetch = ile_args_orig + ...). Puffball ILE is + standard in production, so pin it rather than relying on it happening to work.""" + (tmp_path / "args_puff.txt").write_text( + "--parameter mc --parameter delta_mc --downselect-parameter m2 " + "--downselect-parameter-range [1,1000]\n") + wd = _build_dag(tmp_path, GOOD_ILE_ARGS + " --time-marginalization-quadrature bandlimited", + extra_cepp=["--puff-args", os.fspath(tmp_path / "args_puff.txt")]) + made = [n for n in ("ILE_puff.sub", "ILE_fetch.sub") if (wd / n).exists()] + assert made, "neither puff nor fetch sub was generated; the inheritance claim is untested" + for name in made: + assert "--time-marginalization-quadrature bandlimited" in (wd / name).read_text(), name diff --git a/MonteCarloMarginalizeCode/Code/test/waveform/check_waveform_random.py b/MonteCarloMarginalizeCode/Code/test/waveform/check_waveform_random.py index 1eb76c322..8c6088f75 100644 --- a/MonteCarloMarginalizeCode/Code/test/waveform/check_waveform_random.py +++ b/MonteCarloMarginalizeCode/Code/test/waveform/check_waveform_random.py @@ -5,6 +5,8 @@ # python check_waveform_random.py --force-psi 0.1 # python ./check_waveform_random.py --approx SpinTaylorT4 --force-psi 0 --use-same-fref --force-aligned # python check_waveform_random.py --approx SEOBNRv5EHM --force-aligned --use-eccentric --use-gwsignal --inj mdc.xml.gz --event 15 +# python check_waveform_random.py --approximant IMRPhenomXPHM --fiducial --use-extra-fd-args --assert-overlap 0.998 +# python check_waveform_random.py --approximant IMRPhenomXPHM --stress-q4-edge --seglen 16 --use-extra-fd-args --assert-overlap 0.995 # # RESULTS # - Pv2: perfect @@ -25,6 +27,7 @@ import numpy as np from matplotlib import pyplot as plt +from scipy import signal import argparse import lal import sys @@ -37,6 +40,8 @@ parser = argparse.ArgumentParser() parser.add_argument("--approximant",type=str,default="IMRPhenomPv2") parser.add_argument("--fiducial",action='store_true') +parser.add_argument("--stress-q4-edge", action='store_true') +parser.add_argument("--stress-near-aligned", action='store_true') parser.add_argument("--mtot",default=40,type=float) parser.add_argument("--fmin",default=20,type=float) parser.add_argument("--use-gwsignal",action='store_true') @@ -57,18 +62,44 @@ parser.add_argument("--event",type=int, default=None,help="event ID of injection XML to use.") parser.add_argument("--seglen",default=8,type=int) parser.add_argument("--verbose",action='store_true') +parser.add_argument("--assert-overlap", type=float, default=None, + help="fail unless the mode-sum/direct flat-noise overlap, maximized over integer time and constant phase, reaches this value") opts= parser.parse_args() + +def best_flat_overlap(a, b, delta_t, f_low, f_high): + """Flat-noise complex overlap maximized over integer time and phase.""" + fa = np.fft.fft(a) + fb = np.fft.fft(b) + freqs = np.fft.fftfreq(len(a), delta_t) + keep = np.logical_and(np.abs(freqs) >= f_low, np.abs(freqs) <= f_high) + fa = np.where(keep, fa, 0) + fb = np.where(keep, fb, 0) + a_band = np.fft.ifft(fa) + b_band = np.fft.ifft(fb) + corr = signal.correlate(a_band, b_band, mode="full", method="fft") + lags = signal.correlation_lags(len(a_band), len(b_band), mode="full") + lag = int(lags[np.argmax(np.abs(corr))]) + b_shift = np.zeros_like(b_band) + if lag >= 0: + b_shift[lag:] = b_band[:len(b_band)-lag] + else: + b_shift[:lag] = b_band[-lag:] + overlap = abs(np.vdot(a_band, b_shift))/(np.linalg.norm(a_band)*np.linalg.norm(b_shift)) + scale = np.vdot(b_shift, a_band)/np.vdot(b_shift, b_shift) + return float(overlap), lag, scale + P = lalsimutils.ChooseWaveformParams() P.ampO=-1 # need this otherwise we don't get SpinTaylor HM output P.phaseO = 7 # so we have less insane outputs P.taper = lalsimutils.lsu_TAPER_START -if not(opts.fiducial) and not(opts.inj): +if not(opts.fiducial or opts.stress_q4_edge or opts.stress_near_aligned) and not(opts.inj): print("Creating random event to use for plot comparison.") P.randomize() # move inside conditional for use inj purposes P.dist = RIFT.likelihood.factored_likelihood.distMpcRef*1e6*lal.PC_SI # fiducial reference distance P.assign_param('mtot',opts.mtot*lal.MSUN_SI) + if opts.use_eccentric: P.eccentricity = np.random.uniform(0.0,0.4) #for safety, for now P.meanPerAno = np.random.uniform(0.0,2*np.pi) @@ -108,6 +139,17 @@ P.dist = RIFT.likelihood.factored_likelihood.distMpcRef*1e6*lal.PC_SI # fiducial reference distance P.assign_param('mtot',opts.mtot*lal.MSUN_SI) + if opts.stress_q4_edge: + P.m1, P.m2 = 60*lal.MSUN_SI, 15*lal.MSUN_SI + P.incl, P.phiref, P.psi = 1.2, 1.1, 0.4 + P.s1x, P.s1y, P.s1z = 0.75, 0.0, 0.25 + P.s2x, P.s2y, P.s2z = -0.20, 0.25, -0.10 + elif opts.stress_near_aligned: + P.m1, P.m2 = 40*lal.MSUN_SI, 20*lal.MSUN_SI + P.incl, P.phiref, P.psi = 1.2, 1.1, 0.4 + P.s1x, P.s1y, P.s1z = 1e-6, 0.0, 0.4 + P.s2x, P.s2y, P.s2z = 0.0, 0.0, -0.2 + if opts.force_aligned: P.s1x = P.s1y=P.s2x=P.s2y=0 if not(opts.use_gwsignal): @@ -123,6 +165,8 @@ P.deltaT=1./4096 P.deltaF = 1./opts.seglen P.fref = 22 +if opts.stress_q4_edge: + P.fref = 30 P.fmin=opts.fmin if opts.use_same_fref: P.fref = P.fmin @@ -134,11 +178,10 @@ P.print_params() # hoft via hlm, using exactly the function call we use in production -extra_args ={} extra_waveform_args ={} extra_waveform_args['fd_centering_factor']= 0.9 if opts.use_extra_fd_args: - extra_args['fd_L_frame'] = True + extra_waveform_args['fd_L_frame'] = True if opts.use_xphm_spintaylor: extra_waveform_args['FinalSpinMod'] =2 extra_waveform_args['PhenomXPHMReleaseVersion'] = 122022 @@ -146,7 +189,7 @@ P_copy = P.manual_copy() # beware, call may change P! -hlmF_1, _= factored_likelihood.internal_hlm_generator(P_copy, opts.Lmax, use_gwsignal=opts.use_gwsignal, use_gwsignal_approx=opts.approximant,ROM_group=opts.rom_group,ROM_param=opts.rom_param, extra_waveform_kwargs=extra_waveform_args, **extra_args) +hlmF_1, _= factored_likelihood.internal_hlm_generator(P_copy, opts.Lmax, use_gwsignal=opts.use_gwsignal, use_gwsignal_approx=opts.approximant,ROM_group=opts.rom_group,ROM_param=opts.rom_param, extra_waveform_kwargs=extra_waveform_args) hlmT_1 = {} for mode in hlmF_1: hlmT_1[mode] = lalsimutils.DataInverseFourier(hlmF_1[mode]) @@ -185,6 +228,36 @@ if opts.verbose: print('net2 ', np.max(np.abs(hTc_2.data.data))) +if opts.assert_overlap is not None: + overlap, lag, scale = best_flat_overlap( + np.asarray(hTc_1.data.data), np.asarray(hTc_2.data.data), + float(hTc_1.deltaT), P.fmin, min(1024., 0.5/P.deltaT), + ) + print("Mode-sum/direct overlap", overlap, "at lag", lag, "samples") + if not np.isfinite(overlap) or overlap < opts.assert_overlap: + raise SystemExit("FAIL: overlap {} is below {}".format(overlap, opts.assert_overlap)) + print("Best aligned complex scale", scale) + if not np.isfinite(scale) or abs(np.angle(scale)) > 0.05 or not 0.8 < abs(scale) < 1.2: + raise SystemExit("FAIL: mode-sum/direct phase or amplitude convention is inconsistent") + + # fd_L_frame modes are intrinsic. Polarization is applied only when the + # modes are summed into strain, so changing psi must not change any mode. + P_psi = P.manual_copy() + P_psi.psi += 0.37 + hlmF_psi, _ = factored_likelihood.internal_hlm_generator( + P_psi, opts.Lmax, use_gwsignal=opts.use_gwsignal, + use_gwsignal_approx=opts.approximant, ROM_group=opts.rom_group, + ROM_param=opts.rom_param, extra_waveform_kwargs=extra_waveform_args, + ) + relative_mode_difference = max( + np.linalg.norm(hlmF_1[mode].data.data - hlmF_psi[mode].data.data) + / max(np.linalg.norm(hlmF_1[mode].data.data), np.finfo(float).tiny) + for mode in hlmF_1 + ) + print("Maximum relative mode change under psi shift", relative_mode_difference) + if relative_mode_difference > 1e-12: + raise SystemExit("FAIL: intrinsic modes depend on polarization angle") + # now confirm complex_hoft dependence on psi is as desired # NOT THE SAME PSI DEPENDENCE AS WE ASSUME ELSEWHERE psi_ref = float(P.psi) diff --git a/setup.py b/setup.py index f82b83bd7..835ab7380 100644 --- a/setup.py +++ b/setup.py @@ -68,6 +68,8 @@ 'jax-apps': ['jax', 'optax', 'equinox', 'tinygp', 'numpyro', 'flowMC'], }, entry_points={ + 'console_scripts': + ['rift-rimsky-analysis = RIFT.rimsky.integration:main'], 'asimov.pipelines': ["rift = RIFT.asimov.rift:Rift"], 'RIFT.integrator_plugins':