diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d4167d924..d71cb5b84 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -301,6 +301,47 @@ jobs: - name: Census every test file against the CI roster run: python .travis/test-ci-roster.py + roster-verify-check: + needs: install + runs-on: ubuntu-latest + # Companion to ci-roster-check, and the answer to a gap that job cannot close: it enforces + # that every ungated test file carries a REASON, but not that the reason is TRUE. That was + # not hypothetical -- the roster asserted "skips cleanly without them" for two jax_gp files + # while one ERRORED without jax and the other had no guard at all (PR #248), and it carried + # two OPTDEP entries whose prose said "unverified on a runner" / "belongs in another job" + # when both in fact collect and pass completely. A reason nobody re-checks is prose, which + # is the thing this whole census exists to stop being mistaken for coverage. + # + # So each status now has a falsifiable predicate and this job runs it: LEGACY must fail to + # import, HANDRUN must collect nothing, EXPENSIVE must collect but pass nothing without + # RIFT_RUN_EXPENSIVE, and OPTDEP must name its dependencies (`needs:`) which + # requirements.txt must not install. See .travis/test-roster-verify.py for what is + # deliberately NOT checked and why. + # + # SEPARATE from ci-roster-check on purpose: that one is stdlib-only with no `needs: install` + # and must stay that way, so it still reports when the install matrix is broken. This one + # imports RIFT. ~3.5 min measured on CIT (52 entries, one pytest collection each). + timeout-minutes: 25 + 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 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: Verify each roster reason still holds + env: + OMP_NUM_THREADS: 1 + run: python .travis/test-roster-verify.py + core-unit-check: needs: install runs-on: ubuntu-latest @@ -320,8 +361,12 @@ jobs: # 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. + # Python 3.10 to match the sibling numpy jobs. timeout-minutes is a runaway backstop, + # not a budget -- but the margin is no longer large: the gate measures 350 s on CIT + # (was 49 s when this comment first claimed 'an order of magnitude'), because the + # per-file collection loop spawns one RIFT-importing interpreter per manifest entry and + # grows with the manifest. 20 min is ~3.4x that, not 20x. Re-measure before adding + # many more files, and see .travis/test-core-units.sh for the breakdown. timeout-minutes: 20 steps: - uses: actions/checkout@v4 @@ -393,6 +438,15 @@ jobs: jax-ile-check: needs: install runs-on: ubuntu-latest + # Sharded 2026-09-08: the suite outgrew the 60-minute cap below. Only the + # execute step splits; the FILES manifest, the collection floor against + # EXPECTED_TESTS and the deselect-resolution check run in full in every + # shard, so a shard cannot go green on a partial view. Raise `shard` to + # add capacity; the split is round-robin by index in test-jax.sh. + strategy: + fail-fast: false + matrix: + shard: [1, 2, 3] # test/jax/ was run by NOTHING in this workflow until this job landed (the file had # zero matches for "jax"), and two real defects survived a month each behind that # gap. See .travis/test-jax.sh for why the gate counts tests instead of just @@ -456,11 +510,46 @@ jobs: python -m pip install --upgrade pip --break-system-packages python -m pip install -r requirements.txt coverage pytest "jax[cpu]" numpyro --break-system-packages python -m pip install --editable . --break-system-packages + # WHY THIS EXISTS. jax is installed UNPINNED above, so this job's stack is + # whatever pip resolved the day it ran, and no developer environment + # necessarily has that version. Without a record, a red gate six weeks old + # cannot be reproduced -- you cannot even tell which jax produced it. Two + # failures during the #214 landing were chased locally before a + # pristine-base control showed they were the environment, not the branch + # (issue #292). The gate prints jax and numpyro; this captures the rest, + # jaxlib and lal included. + - name: Record the resolved jax stack + run: | + python -m pip freeze > pip-freeze-jax-ile.txt + echo "== resolved stack (full list in the jax-ile-resolved-stack artifact) ==" + # importlib.metadata, not a grep of pip freeze: freeze reports a URL + # rather than a version for anything not installed from PyPI, so the + # grep prints a wheel path in exactly the environments someone is most + # likely to be comparing against. + python -c " + import importlib.metadata as md + for n in ('jax', 'jaxlib', 'numpyro', 'numpy', 'scipy', 'numba', 'lal'): + try: + print('%-10s %s' % (n, md.version(n))) + except Exception: + print('%-10s (absent)' % n) + " - name: Run jax_ile CPU regression gate env: JAX_PLATFORMS: cpu OMP_NUM_THREADS: 1 + JAX_GATE_SHARDS: 3 + JAX_GATE_SHARD: ${{ matrix.shard }} run: bash .travis/test-jax.sh + # always(): a FAILING gate is the case the record is for, so it must upload + # after a red run too, not only after a green one. + - name: Upload the resolved jax stack + if: always() + uses: actions/upload-artifact@v4 + with: + name: jax-ile-resolved-stack + path: pip-freeze-jax-ile.txt + if-no-files-found: warn # 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 @@ -715,6 +804,20 @@ jobs: bash .travis/test-run.sh bash .travis/test-run-alts.sh bash .travis/test-build.sh + - name: Run ILE executable selection tests (--use-jax-ile / --ile-exe) + # Full subprocess DAG builds against the same reference ini/coinc as + # test-build.sh, so the pseudo_pipe CLI wiring is what is exercised: + # --use-jax-ile and --ile-exe land in ILE.sub/ILE_puff.sub/ILE_extr.sub, + # and --use-jax-ile is refused at DAG-build time with --calmarg-envelope-directory. + run: | + python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_jax_ile_selectable.py + - name: Run q-time-pregrid DAG-build test (--internal-ile-q-time-pregrid-factor) + # Full subprocess DAG build against the same reference ini/coinc, same reason as + # the step above (PR #281 follow-up review, MAJOR #1): asserts + # --internal-ile-q-time-pregrid-factor 8 lands in ILE.sub, ILE_extr.sub, and + # ILE_puff.sub, not only in helper_ile_args.txt / args_ile.txt. + run: | + python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_q_time_pregrid_dag.py - name: Upload test logs on failure if: failure() uses: actions/upload-artifact@v4 diff --git a/.travis/PRECOMMIT.md b/.travis/PRECOMMIT.md new file mode 100644 index 000000000..737d87ce3 --- /dev/null +++ b/.travis/PRECOMMIT.md @@ -0,0 +1,47 @@ +# CI tiers + +`.github/workflows/ci.yml` gates every push/PR with ~22 jobs. That is already large, and +covering every backend and hardware corner in it would grow it without bound: no GitHub +runner has a GPU, and the JAX suite already costs real time. This tiers the existing suite; +nothing here removes or shortens a job. + +## Tier 1 -- fast CI, every push/PR (unchanged) + +Everything in `ci.yml` runs automatically: install, help-check, import-check, +dependency-compat-check, sim-manager-check, integrator-gate-accounting-check, +q-window-stencil-check, ci-roster-check, roster-verify-check, core-unit-check, slowrot-check, +jax-ile-check, lisa-check, calmarg-check, integration-check, asimov-integration, +rimsky-integration, test-run, container-dep-canary, container-swig-canary, docs. Plus +`.gitlab-ci.yml`'s default pipeline (coord/integrate/posterior/run/run-alts/build). + +RECOMMENDATION for RO'S: `jax-ile-check` runs on every push, and its cost argues for tier 2 or +a schedule instead. `.travis/test-jax.sh:494` measures the gated 139-test baseline at 13m53s +(833s), against a 60-minute timeout; the suite has since grown to 577 tests +(`.travis/test-jax.sh:766`). `test_angle_marg_exact.py` (36 tests, excluded from this gate at +`.travis/test-jax.sh:484`) does not add to that per-push cost; run it by hand per the comment +there. Moving `jax-ile-check` to a nightly schedule would return the 833s to every push, since +nothing else in tier 1 imports jax. This PR leaves it in place: moving a job is a behavior +change and belongs in its own PR. + +## Tier 2 -- strongly recommended before merging into rift_O4d, by hand on a GPU/big node + +No GitHub runner has a GPU, so these either never run, or run a cupy-less CPU parity check +that cannot see a real device. Run `.travis/precommit-recommended.sh` before merging +anything touching the ILE likelihood, the NoLoop stencils, calibration marginalization, or +the JAX driver: + +- `RIFT/likelihood/test_q_window_interp_gpu.py`, `test_noloop_gpu_stencils.py` (real cupy) +- `.travis/test-calmarg-gpu.sh` (fused calibration CUDA kernels) +- `test/jax/test_angle_marg_exact.py`, and the full `.travis/test-jax.sh` +- `test/expensive_before_merging/` (`RIFT_RUN_EXPENSIVE=1`) + +The script exits 1 when cupy/CUDA or the jax stack is absent; it does not skip. + +## Tier 3 -- scheduled or manual only + +`.gitlab-ci.yml`'s `gpu_integration` job: a real GPU runner, running `test-integrate.sh` + +`test-calmarg-gpu.sh` + `test-lisa-gpu.sh` under CUDA. Its `rules:` block +(`.gitlab-ci.yml:180-184`) has two branches: web-triggered `when: manual`, and +`$CI_PIPELINE_SOURCE == "schedule"` with `when: on_success`. So it also runs automatically on a +scheduled pipeline. A nightly `jax-ile-check` would belong on the same schedule if RO'S takes +the tier-1 recommendation above. diff --git a/.travis/ci_roster.txt b/.travis/ci_roster.txt index 8456d86e7..a6475abd4 100644 --- a/.travis/ci_roster.txt +++ b/.travis/ci_roster.txt @@ -16,7 +16,6 @@ # 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 @@ -69,22 +68,31 @@ MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp_pinned.py LE MonteCarloMarginalizeCode/Code/test/integrators/test_mcsampler_gpu.py LEGACY imports ourio # --------------------------------------------------------------------------------------- -# HANDRUN -- quantitative studies with real internal gates, run by hand. +# HANDRUN -- scripts that are not pytest targets. All collect ZERO items, so pytest exits 5 +# ("no tests ran") on each, which reads as a pass; that is why they are recorded here rather +# than pointed at a job. # -# 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 FIVE STUDIES BELOW ARE NOW RUN IN CI ANYWAY. test/integrators/test_integrator_studies.py +# invokes each as a subprocess and requires exit 0, and core-unit-check gates that wrapper. They +# still belong in this file because they remain non-collectable themselves. # -# 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 +# The earlier version of this entry called them "expensive" and proposed hiding them behind +# RIFT_RUN_EXPENSIVE. That was asserted, not measured, and it was wrong: on CIT with the IGWN +# python (OMP_NUM_THREADS=1) they take 5, 2, 14, 4 and 4 seconds -- 29 s for all five. Each ends +# in `raise SystemExit(1)` on failure, and all five seed numpy explicitly (RandomState(0/1/3), +# np.random.seed), so they are deterministic rather than merely lucky. +# +# THE WRAPPER PASSES --as-test, AND THAT IS LOAD-BEARING. Every one of these keeps its +# scientific comparisons AND its SystemExit(1) behind `if args.as_test`, so without the flag +# a biased result still prints and exits 0. The first version of the wrapper omitted it and +# therefore gated only crashes -- an inert guard, caught in review of #251, and not by the +# three clean runs I had cited as evidence: those were guaranteed to pass. The timings here +# are WITH the flag, and are larger than without it precisely because the gates then run. +MonteCarloMarginalizeCode/Code/test/integrators/test_AV_bootstrap.py HANDRUN AV warm-start bias/efficiency gate; 0 collected, run by test_integrator_studies.py --as-test (8 s) +MonteCarloMarginalizeCode/Code/test/integrators/test_AV_warmstart_safety.py HANDRUN anti-bias guard for reusing a proposal across problems; 0 collected, run by test_integrator_studies.py --as-test (4 s) +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_adaptive_alloc.py HANDRUN portfolio draw-allocation vs standalone AV; 0 collected, run by test_integrator_studies.py --as-test (19 s) +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_balance_heuristic.py HANDRUN portfolio safety under a decoy member; 0 collected, run by test_integrator_studies.py --as-test (7 s) +MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_oracle.py HANDRUN needle-target oracle study; 0 collected, run by test_integrator_studies.py --as-test (5 s) 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 @@ -95,42 +103,55 @@ MonteCarloMarginalizeCode/Code/demo/rift/export_likelihoods/head_to_head/run_tes # 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 +MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py OPTDEP needs:jax -- skips cleanly without it, 2 tests pass with it; belongs in jax-ile-check, which already installs a CPU jax stack +# test_interpolators.py was rostered OPTDEP claiming "10 tests where both are installed". +# That was never run. With jax 0.9.2 + optax 0.2.8 (~/.cache/jaxci_venv on CIT, taskset -c 0-7, +# JAX_ENABLE_X64=1) it collects 12 and FAILS 2 on accuracy: exact rmse 0.7696 vs tol 0.05, svgp +# 0.9087 vs tol 0.3. Not under-training -- exact plateaus at 0.7697 for n_opt_steps 150/600/2000 +# while rff reaches 0.0045 on the identical target, so it is converged and wrong. Both backends +# are user-selectable as CIP --fit-method gp-jax-exact / gp-jax-svgp. Reported, not fixed here. +MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py BROKEN needs jax+optax, and with them 2 of 12 fail on accuracy (exact, svgp) -- see the note above # 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/hyperpipe/tests/test_hydra_integration.py OPTDEP needs:hydra,omegaconf -- skips cleanly without them +MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_mode_sign.py OPTDEP needs:EOBRun_module -- the gwsignal TEOB route; 2 collected, 1 passes and 1 skips without it +MonteCarloMarginalizeCode/Code/test/test_gwsignal_teob_near_aligned.py OPTDEP needs:EOBRun_module -- 1 collected, skips without it +# test_rimsky_integration.py is the worked example of why OPTDEP entries must DECLARE their +# dependencies rather than describe them. Its old reason was prose ("belongs in the +# rimsky-integration job"), test-roster-verify.py rightly refused it, and I resolved that by +# gating the file -- which was wrong: it collects 15 on CIT and 0 on a runner, because CIT's +# IGWN environment happens to carry asimov. core-unit-check's per-file collection floor caught +# it. With the deps named, the same check now keys off their real presence and stays quiet +# here while remaining meaningful on a runner. Its proper home is a job that installs asimov +# (rimsky-integration already runs test_rimsky_end_to_end.py); that move needs a collection +# floor there so a skip cannot pass silently, and is not attempted in this PR. +MonteCarloMarginalizeCode/Code/test/test_rimsky_integration.py OPTDEP needs:asimov,liquid -- 15 collected where asimov is installed, 0 without it +MonteCarloMarginalizeCode/Code/test/integrators/test_NF_reuse.py OPTDEP needs:nflows -- the normalizing-flow store; collection errors without it +MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp_vegas.py OPTDEP needs:vegas -- 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_eosmanager_misc.py OPTDEP needs:env:LALSIMULATION_DATADIR -- 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 +# test_mcsampler_foridiots.py is HANDRUN rather than BROKEN because it is a demo script with no +# test functions -- it was never going to be gated. But it fails for a reason that is NOT the +# demo's: it dies in RIFT/integrators/mcsamplerGPU.py:1324, inside integrate(), on +# +# weights_alt = int_vals**tempering_exp # NameError: int_vals is not defined +# +# the `not save_intg` branch of the adaptation weighting. `int_vals` exists nowhere in that +# scope; the sibling branches use self._rvs["integrand"][-n_history:], and the local holding the +# same values when nothing is being saved is `fval` (a commented-out line two above prints it), +# so `fval**tempering_exp` is the near-certain intent. This branch cannot ever have run. +# +# NOT FIXED HERE: that is core sampler code, and guessing the intended expression is exactly the +# kind of change that should be RO'S call rather than a side effect of a test-hygiene PR. +# Reported instead. Reachable on CPU -- this demo hit it with no GPU involved. +MonteCarloMarginalizeCode/Code/test/test_mcsampler_foridiots.py HANDRUN plotting demo with no test functions; dies in mcsamplerGPU.integrate on an undefined int_vals (see note above) # --------------------------------------------------------------------------------------- # 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/precommit-recommended.sh b/.travis/precommit-recommended.sh new file mode 100755 index 000000000..1d1feb40b --- /dev/null +++ b/.travis/precommit-recommended.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Tier-2 gate: "strongly recommended before merging into rift_O4d" (see +# .travis/PRECOMMIT.md). GPU + JAX suites no GitHub runner can exercise, plus the +# expensive_before_merging/ regression tests. Run BY HAND on a GPU/big node before +# merging anything touching the ILE likelihood, the NoLoop stencils, calibration +# marginalization, or the JAX driver. +# +# REFUSES, NOT SKIPS. A tier-2 run that quietly no-ops when cupy or the jax stack is +# missing is worse than no run at all: it produces the same "ran clean" impression as a +# real pass. Every prerequisite below is checked before anything runs, the same way +# .travis/test-jax.sh guards its own interpreter/jax/numpyro imports and +# .travis/test-calmarg-gpu.sh assumes a real GPU because it only ever runs on one. +set -uo pipefail +cd "$(dirname "$0")/.." || { echo "precommit-recommended.sh: cannot cd to repo root" >&2; exit 1; } + +PYTHON_BIN="${PYTHON:-python}" +command -v "${PYTHON_BIN}" >/dev/null 2>&1 || PYTHON_BIN="$(command -v python3)" +command -v "${PYTHON_BIN}" >/dev/null 2>&1 || { + echo "precommit-recommended.sh: no python interpreter found" >&2; exit 1; } + +echo "== environment guard: GPU / cupy ==" +# Device(0).compute_capability alone is NOT a guard: it reads a static property and +# returns cleanly even when the device is busy/unavailable for real work (measured: +# passes, then the first actual allocation raises cudaErrorDevicesUnavailable). Force +# a real allocation + kernel + sync so a device that cannot currently run anything +# fails HERE, not partway through a tier-2 test file with a confusing traceback. +"${PYTHON_BIN}" -c ' +import cupy +x = cupy.arange(1024) +y = int((x * 2).sum()) +cupy.cuda.Stream.null.synchronize() +assert y == 2 * 1024 * 1023 // 2 +' || { + echo "precommit-recommended.sh: cupy could not run a real op on a CUDA device in ${PYTHON_BIN}." >&2 + echo " This gate requires a real, currently-available GPU -- it has no CPU fallback." >&2 + echo " Run it on a GPU/big node (see infra-atlas: pcdev11/13) or GitLab's gpu_integration"\ + "runner, and check nvidia-smi if the device looks present but busy." >&2 + exit 1 +} +echo "cupy + CUDA device: OK (real allocation + kernel + sync)" + +echo "== environment guard: jax stack ==" +"${PYTHON_BIN}" -c 'import pytest' || { + echo "precommit-recommended.sh: pytest unavailable in ${PYTHON_BIN}" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import jax, jaxlib; print("jax", jax.__version__)' || { + echo "precommit-recommended.sh: jax unavailable in ${PYTHON_BIN}" >&2; exit 1; } +"${PYTHON_BIN}" -c 'import numpyro; print("numpyro", numpyro.__version__)' || { + echo "precommit-recommended.sh: numpyro unavailable in ${PYTHON_BIN}" >&2; exit 1; } +echo "jax + jaxlib + numpyro: OK" + +CODE="MonteCarloMarginalizeCode/Code" +fail=0 + +echo "== RIFT/likelihood/test_q_window_interp_gpu.py + test_noloop_gpu_stencils.py (real cupy) ==" +"${PYTHON_BIN}" -m pytest -q \ + "${CODE}/RIFT/likelihood/test_q_window_interp_gpu.py" \ + "${CODE}/RIFT/likelihood/test_noloop_gpu_stencils.py" || fail=1 + +echo "== .travis/test-calmarg-gpu.sh ==" +bash .travis/test-calmarg-gpu.sh || fail=1 + +echo "== test/jax/test_angle_marg_exact.py ==" +JAX_PLATFORMS="${JAX_PLATFORMS:-cpu}" OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" \ + "${PYTHON_BIN}" -m pytest -q "${CODE}/test/jax/test_angle_marg_exact.py" || fail=1 + +echo "== full .travis/test-jax.sh (measured ~859s+~600s locally as of its own comments;"\ + "grows with the suite) ==" +JAX_PLATFORMS="${JAX_PLATFORMS:-cpu}" OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}" \ + bash .travis/test-jax.sh || fail=1 + +echo "== test/expensive_before_merging/ (RIFT_RUN_EXPENSIVE=1) ==" +RIFT_RUN_EXPENSIVE=1 "${PYTHON_BIN}" -m pytest -q "${CODE}/test/expensive_before_merging/" \ + || fail=1 + +if [ "${fail}" -ne 0 ]; then + echo "precommit-recommended.sh: one or more tier-2 gates FAILED" >&2 + exit 1 +fi +echo "precommit-recommended.sh: all tier-2 gates PASSED" diff --git a/.travis/test-ci-roster.py b/.travis/test-ci-roster.py index 06e8d241f..91290e5fe 100755 --- a/.travis/test-ci-roster.py +++ b/.travis/test-ci-roster.py @@ -203,8 +203,6 @@ def _live_gates(live_cfg): "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", } @@ -432,49 +430,14 @@ def main(): # 2. A roster entry for a file that IS now reachable is stale -- it records a decision that # has been overtaken, and leaving it invites the next reader to trust it. - # - # PENDING used to be UNCONDITIONALLY exempt from this, which made it the one status that - # could sit in the roster for ever: its whole point was to stay legal before AND after the - # companion PR landed, so nothing ever forced its removal. That bought merge-order - # independence at the price of a status with no expiry, which is the rot this file exists - # to prevent. It now carries an ENFORCEABLE condition instead of a promise: the reason - # must name the gate it waits on as `gate:`, 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)) + "-- delete it." % (ROSTER, f, status, reachable[f])) 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)) diff --git a/.travis/test-core-units.sh b/.travis/test-core-units.sh index c82fbb8d9..812f19c0b 100755 --- a/.travis/test-core-units.sh +++ b/.travis/test-core-units.sh @@ -11,8 +11,11 @@ # 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. +# The original manifest was run file by file 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. Later +# entries are verified by this gate itself, which collects every file individually before the +# combined run, so an addition that collects nothing or fails is caught here rather than +# trusted on a quoted number. # # 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 @@ -56,13 +59,21 @@ FILES=( "$C/test/test_calmarg_calibration.py" # -- likelihood dispatch "$C/RIFT/likelihood/test_td_dispatch_epoch.py" + "$C/RIFT/likelihood/test_precompute_crossterm_batching.py" "$C/test/test_ile_scalar_edge_cases.py" + "$C/test/test_mcsamplerGPU_cdf_inverse_scalar_probe.py" "$C/test/test_srate_resample_time_marginalization.py" + "$C/test/test_vectorized_lal_tools_split.py" + "$C/test/test_noloop_accumulator_shapes.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" + # Wraps the five integrator studies as subprocesses (29 s). They collect nothing + # themselves -- pytest exits 5 on each -- so this is how their gates reach CI at all. + "$C/test/integrators/test_integrator_studies.py" + "$C/test/integrators/test_replica_pooling.py" "$C/test/integrators/test_rvs_weight_derivation.py" "$C/test/integrators/test_seeding_public_paths.py" "$C/test/integrators/test_seeding_reproducibility.py" @@ -77,7 +88,21 @@ FILES=( "$C/test/hyperpipe/tests/test_config.py" "$C/test/hyperpipe/tests/test_coords.py" "$C/test/hyperpipe/tests/test_drivers.py" + "$C/test/hyperpipe/tests/test_marg_list.py" "$C/test/test_hyperpipeline_io.py" + # -- promoted out of the roster after roster-verify-check caught its reason being false ON + # THE RUNNER: it was OPTDEP needs:glue,htcondor, and with htcondor absent there it still + # collected 15 and passed 15. Confirmed locally with BOTH blocked via a sys.meta_path + # finder: 15/15. Its `import htcondor` / `from glue import pipeline` are capability + # probes inside the tests, not requirements. CIT has both, which is exactly why CIT could + # not see this and a runner could. + "$C/test/backends/test_backends_lowlevel.py" + # -- promoted out of the roster: it was OPTDEP on prose ("unverified on a runner") and + # .travis/test-roster-verify.py caught it collecting and passing COMPLETELY with nothing + # missing. (test_rimsky_integration.py was promoted alongside it and REVERTED: it + # importorskips asimov, which CIT has and this job does not, so it collected 15 here and + # 0 on the runner. The per-file collection floor below caught that -- see the roster.) + "$C/test/test_teobresums_compat.py" # -- packaging / config contracts / waveform conventions "$C/test/test_advanced_parameter_ports.py" "$C/test/test_container_manifest.py" @@ -94,6 +119,86 @@ for f in "${FILES[@]}"; do done [ "$missing" -eq 0 ] || { echo " Fix the manifest or restore the file; left as is it covers nothing." >&2; 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: 319 collected, +# 307 passed, 12 skipped (11 pytest.skip + 1 xfail). +# +# COST: 363 s total on CIT for the 347-test manifest, of which the pytest run is ~145 s. The +# rest is the per-file collection loop below -- one interpreter per manifest entry, each +# importing RIFT (lal, numpy, numba), so it grows linearly with the manifest and now dominates. +# That is the price of the exit-5 defence and it is worth paying, but it is why this job's +# timeout-minutes is 20 rather than something tight: measure before trimming the budget. +# +# History. Both branches of a merge have now moved these numbers twice, and the arithmetic is +# NOT the way to combine them -- this gate collects every file individually before the combined +# run, so the merged floor is RE-MEASURED, never added up: +# 278/266 original manifest +# 296/284 + test_replica_pooling.py, test_marg_list.py (rostered BROKEN until fixed) +# 299/287 + test_vectorized_lal_tools_split.py (from rift_O4d, then 3 -> 7 tests: its +# original comparison pitted the combined wrapper against the composition of its own +# two halves, which after the split IS the wrapper, so it could not fail. Rewritten +# against a FROZEN copy of the pre-split bodies, parametrized per detector -- #256) +# 307/295 + test_noloop_accumulator_shapes.py (from rift_O4d, then 8 -> 9 tests: the added +# one asserts the synthetic inputs actually exercise the data term, after the first +# version put the Q window entirely outside the buffer and left kappa_sq identically +# zero -- an assertion that held for the wrong reason) +# 312/300 rift_O4d #258, a pure FLOOR correction: the growth above had already landed in the +# files while the floors still said 299/287, so 312 collected was clearing a 307 +# floor -- the silent under-coverage this gate exists to prevent, appearing in the +# gate's own bookkeeping. Worth noting how the two branches differed here: #258 +# reconstructed 312 by arithmetic (299 + 4 + 9) and this branch measured 347 with +# those same files already grown, because it never adds. Merging changed this +# manifest's numbers by ZERO. +# + + test_teobresums_compat.py (OPTDEP on prose until test-roster-verify.py caught it +# passing completely with nothing missing; its companion test_rimsky_integration.py +# was promoted alongside and REVERTED -- see the note by the manifest entry) +# + + test_integrator_studies.py, wrapping the five integrator studies that collect +# nothing themselves (+43 s, with --as-test) +# + + test_backends_lowlevel.py (OPTDEP needs:glue,htcondor until roster-verify-check +# found it passing 15/15 on a runner with htcondor absent) +# 358/346 + test_mcsamplerGPU_cdf_inverse_scalar_probe.py (11 tests: mcsamplerGPU.cdf_inverse +# fed odeint's float probe to len(x) pdfs; the t_ref wiring in all three ILE drivers) +# +# RAISE these when files are added: a floor left at the old value passes while covering less, +# which is the failure this gate exists to catch. +# DO NOT RAISE THESE TO THE RUNNER'S NUMBERS. The GitHub runner reports 350 collected / 338 +# passed for this same manifest, CIT reports 347 / 335, and the underlying results are IDENTICAL: +# both say "335 passed, 11 skipped, 1 xfailed". The gap is pytest-subtests, which is in the +# runner's dependency closure and not in CIT's IGWN environment; with it, the three `subTest` +# blocks in test/backends/test_backends_lowlevel.py are counted as separate cases in the junit +# XML this gate parses. Per-FILE collection is 347 on both, file by file. +# +# So the floors are pinned to the PLUGIN-FREE count. That is the robust choice in the only +# direction that matters: 350 >= 347 passes today, and if pytest-subtests ever leaves the +# runner's closure the count falls back to 347 and still passes. Pinning 350 would turn an +# unrelated dependency change into a red gate. +EXPECTED_TESTS=370 +# 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. +EXPECTED_PASSED=358 +MAX_SKIPPED=12 + +# The floors must be INTEGERS, and this is checked rather than assumed. `[ 347 -lt FOO ]` does +# not fail the build: bash prints "integer expression expected", returns 2, and the `if` is +# simply FALSE -- so a malformed floor silently disables the check it looks like it performs, +# and the gate goes green having compared nothing. Observed: a placeholder left in during a +# merge resolution produced exactly that, and only the stderr line gave it away. `set -u` +# catches a MISSING floor; nothing caught a malformed one until this did. +for _f in EXPECTED_TESTS EXPECTED_PASSED MAX_SKIPPED; do + eval "_v=\${${_f}}" + case "${_v}" in + ''|*[!0-9]*) + echo "test-core-units.sh: ${_f}='${_v}' is not a non-negative integer." >&2 + echo " A non-numeric floor makes its comparison a silent no-op -- the gate would pass" >&2 + echo " while checking nothing. Set it to the MEASURED value." >&2 + exit 1 + ;; + esac +done + # 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. @@ -111,23 +216,6 @@ for f in "${FILES[@]}"; do 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 ==" diff --git a/.travis/test-integrate.sh b/.travis/test-integrate.sh index 9f4a3cd5a..14398c304 100755 --- a/.travis/test-integrate.sh +++ b/.travis/test-integrate.sh @@ -44,6 +44,13 @@ python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_cosine_sample python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_distance.py python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_mcsampler_ensemble_log_contract.py +# --psi-marginalization: analytic polarization-angle marginalization made reachable on +# the legacy scalar likelihood path (factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized +# was previously dead code, unreachable from any driver and untested by any importable +# test). Covers the analytic marginal against a brute-force quadrature, the driver's +# refuse-don't-ignore prerequisite checks, and a real subprocess run on synthetic data. +python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_psi_marginalization.py + # Supplementary-likelihood plugin hook: the NAL reader/evaluator (pure numpy, no data) and the # static guard on the drivers' prepare-hook wiring, which is what makes the plugin receive the # SAMPLING basis at all. Both are seconds-long and protect a silent-wrong-answer path. @@ -72,7 +79,7 @@ _TMARG_TESTS=( # 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_EXPECTED=171 _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 @@ -137,11 +144,14 @@ fi # 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. +# against the computed value, and that an undersized region is routed to the finite +# dense fallback rather than returned locally. The algebraic follow-up also pins +# the BKK/resultant enumerator on co-dominant, near-annihilating, exactly degenerate, +# and amplitude-scaled systems, requires inside-cover convergence even after a +# complete enumeration, and keeps the NumPy fallback independent of optional JAX. _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_EXPECTED=37 _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 @@ -152,3 +162,39 @@ 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 + +# Q_lm pregrid factor (PR #261), pipeline passthrough. --q-time-pregrid-factor had NO +# helper/pseudo_pipe wiring at all until this option was added -- it was reachable only +# through --manual-extra-ile-args, which RO'S directive 2026-09-08 says is too easy to get +# wrong for the time-stencil/time-quadrature family. Same discipline as the +# time-marginalization-quadrature gate above: an unlisted test file is simply never run, so +# wiring the file in is part of shipping the wiring. What these files protect: the +# driver-mirroring prerequisite check (--vectorized required; --rotation-slow/--freqresponse/ +# calibration marginalization excluded), the forced-cubic-stencil conflict (factor 8 refuses +# an explicit --interpolate-time other than cubic, with the driver's OWN wording), the +# two-stage refuse-not-ignore emission guard, and that the option actually reaches +# helper_ile_args.txt / args_ile.txt rather than being inert. test_q_time_pregrid_driver_ +# parity.py (PR #281 follow-up review, MAJOR #2) adds the piece those two files left +# untested: it EXECUTES bin/integrate_likelihood_extrinsic_batchmode as a subprocess for +# every prerequisite above and asserts the builder refuses exactly when the driver refuses. +# The real DAG-build regression (--internal-ile-q-time-pregrid-factor reaching ILE.sub / +# ILE_extr.sub / ILE_puff.sub, PR #281 review MAJOR #1) is test_q_time_pregrid_dag.py, +# registered in .github/workflows/ci.yml's test-run job next to test_jax_ile_selectable.py +# rather than here: it is a full subprocess DAG build, not a fast unit gate. test-run is +# matrixed over TWO lanes (legacy py3.9, modern py3.12), so this step runs twice per push, +# measured at ~236s/lane -- ~8 minutes total, not ~3 (PR #291 review, NOTE #5: the single-run +# figure this comment used to state undercounted the per-lane doubling every step in that +# job already pays). +_QPREGRID_TESTS=( + MonteCarloMarginalizeCode/Code/test/test_q_time_pregrid.py + MonteCarloMarginalizeCode/Code/test/test_q_time_pregrid_pipeline.py + MonteCarloMarginalizeCode/Code/test/test_q_time_pregrid_driver_parity.py +) +# Raise EXPECTED by RUNNING collection, never by arithmetic. +_QPREGRID_EXPECTED=77 +_QPREGRID_FOUND=$(python -m pytest -q --collect-only "${_QPREGRID_TESTS[@]}" 2>/dev/null | grep -c '::' || true) +if [ "$_QPREGRID_FOUND" -ne "$_QPREGRID_EXPECTED" ]; then + echo "q-time-pregrid gate: collected $_QPREGRID_FOUND tests, expected $_QPREGRID_EXPECTED" >&2 + exit 1 +fi +python -m pytest -q "${_QPREGRID_TESTS[@]}" diff --git a/.travis/test-jax.sh b/.travis/test-jax.sh old mode 100755 new mode 100644 index 96b4efe5a..5cfc2ef77 --- a/.travis/test-jax.sh +++ b/.travis/test-jax.sh @@ -53,7 +53,8 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # test_jax_endtoend.py 1 full precompute -> pack -> JAX vs the numpy # NoLoop on a real injection (fixed by #144) # test_jax_slowrot_coeffs.py 2 rotation + freqresponse response coefficients -# against their numpy references +# against their numpy references; the compound +# algebra is folded into the freqresponse test # test_jax_slowrot_wrapper.py 1 the one-call build_*_data_from_precompute path # test_jax_slowrot.py 3 rotation Path A (p_max=0), Path B (p_max=1) # and freqresponse: NoLoop parity + AD/jit/ @@ -147,16 +148,21 @@ 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 +# test_angle_marg_smoke.py 12 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 +# amp_sizing, synchronous output-cloud +# recording with training-call exclusion, +# that an amp-sized scheme which recorded +# NOTHING is labelled NOT-PERFORMED and +# never OUTPUT-CLOUD-PASS (the composite +# policy leaves the recorder unwired), +# 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. +# artifacts carry the deterministic checked +# scope. 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 @@ -172,6 +178,22 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # cap must stay WIRED in samplers and the # driver. Each fails under a verified # mutation (see the PR). Seconds. +# test_jax_cache.py 30 the shipped ILE selects a stable +# compatibility namespace, Condor uses +# scratch by default, unwritable caches +# fail open, and transferred bundles +# round-trip while rejecting profile, +# runtime, checksum, and archive-member +# mismatches; concurrent manifest +# writers and imported-entry readers +# cannot race; import provenance survives +# later startup; accelerator +# plugin identity is recorded; a device +# probe that raises disables the cache +# instead of escaping driver import; +# and two fresh real JAX processes +# prove an actual persistent-cache +# reuse. # 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 @@ -233,6 +255,39 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # is the only gated check that distinguishes # the corrected sizing. The rest of the # angle-marg suite is EXCLUDED; see below. +# test_is_proposal_jitter.py 29 issue #227: a Gaussian IS proposal must be +# SCORED under the matrix it was DRAWN from. +# Seven sites drew from cov + 1e-12*I and +# scored under bare cov; the DEFAULT --mode +# laplace-is returned lnZ = 5.85e9 on real O4 +# data and exited 0. Pure numpy synthetic +# likelihood -- no frames, no PSDs, ~8 s. +# 14 of the 24 FAIL on the parent commit +# (4c4f6492). The 10 that pass there are the +# AST detector's own self-test and the nine +# healthy-regime reference comparisons, which +# must pass on BOTH sides by construction -- +# they are the gate on the fix (and on the +# collapse guard) not disturbing the regime +# where this estimator actually works. +# test_jax_phase_marg_mode_order.py 14 phase marginalization must accept EITHER +# packed order of the (2,+-2) pair. +# _accumulate_unit hardcoded column 0 = (2,2) +# and raised NotImplementedError otherwise -- +# but the column order comes from a dict's +# iteration order in the precompute, not from +# the caller, so a correctly configured +# --phase-marginalization run died on valid, +# complete data. U and V carry the mode index +# on BOTH axes, so a half-permutation is a +# silent wrong answer; one test asserts the +# fixture can SEE each single-axis mistake, or +# the equality tests would not gate it. Two +# tests defend the ordering that already works +# by making _permute_modes fatal: the canonical +# order must take the untouched path, not an +# identity permutation. Synthetic packed data, +# no frames, no PSDs, ~60 s. # 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 @@ -320,6 +375,135 @@ JAXDIR="MonteCarloMarginalizeCode/Code/test/jax" # honest phase-marginalized sky/psi export, # K=14/K=88 independent guarded references, # and executable baseline/banded support refusal. +# test_all_axis_peaklocal.py 30 fail-closed four-axis peak-local prototype: +# U,V-guided time ranking and algebraic angular +# starts, JAX refinement, fixed-shape multimode +# quadrature, exact selected-time reconstruction, +# explicit omitted-mass/time-reconstruction +# warrants, geometry/capacity refusal, and outer +# jit/grad/hessian/vmap transform compatibility, +# two-guard primitive support/convergence, +# harmonic-order U,V/Q starts, and the +# empirical enrichment/exact-reserve +# disposition gate. +# test_multipeak_planner.py 15 opt-in U,V,Q-guided four-axis multi-peak +# planner: exact symmetry expansion, strict +# stationary refinement, two-tier empirical +# convergence, overlap ownership and finite +# reserve, plus FOUR refinement-stall guards: +# the bounded step's ascent contract, the +# step-bound sweep, the max_step check the +# rescale requires, and the symmetry-orbit +# invariant a campaign write-up misread as +# degeneracy. Three of the four use a +# narrow-time-peak fixture (the ascent contract +# needs no table): the older _synthetic_tables +# puts its maximum ON the targeting lattice, so +# the Newton loop was never exercised and a +# fixed point in it passed this gate for a +# month while declining every row of the +# 2026-09-07 ladder campaign. CPU-only; no +# lal, cupy, or GPU required. The file defines +# 19 tests; the four real-table oracle +# regressions need external validation packets +# that no fixture in this repository provides, +# so they are DESELECTED here -- see +# DESELECTED_TESTS -- and 15 are gated. +# test_direct_marginalization_policy.py +# 18 opt-in cross-axis policy WIRING: choices and +# refusals, measure conversion on both distance +# paths against the exact scheme, decline to a +# warranted band-limited reserve that keeps the +# sample, ledger completeness, wrapper end to +# end on real synthetic tables with a finite +# gradient, and the driver CLI (subprocess). +# test_jax_q_time_pregrid.py 21 opt-in reflected Q time pregrid on the JAX +# arm: factor-1 bit identity (same array object, +# positions bit-identical to the pre-pregrid +# expressions), the 2n-vs-2(n-1) reflection +# choice measured against an exact-period +# oracle, refined-grid position scaling, the +# fail-closed length/factor checks, 'nearest' +# refusal, wrapper/driver forwarding, and the +# merge interaction with #272's phase-marginalized +# mode permutation. +# Synthetic fixtures; no lal frames, no GPU. +# test_multipeak_fallback_visibility.py +# 13 the multi-peak planner's fallback must not +# read as a policy decline: an exception-driven +# fallback is reported once per call and carries +# decline_kind/fault in the record, a +# budget-driven decline does neither, +# fail_on_fallback is fatal on the first and +# inert on the second, and the default path is +# value- and provenance-identical to the +# pre-change module. Synthetic tables; the +# tier faults are injected at the +# _run_structural_tier seam. CPU-only. +# Two of these pin properties that the first +# version of the change got wrong. The record +# is a 13-element tuple, tested by UNPACKING it +# (13 defaulted-field CONSTRUCTION kept working +# while `a, ..., m = result` had started to +# raise, so a construction test could not see +# it). And the fault report goes to a logger, +# tested under `-W error::RuntimeWarning`, where +# warnings.warn had made the DEFAULT +# fail_on_fallback=False path raise. +# test_jax_bandlimited_distmarg.py 20 time_quadrature="bandlimited" on the +# DISTANCE-marginalized wrapper: agreement at +# two amplitudes with an independently +# reconstructed fine-grid reference (plain +# periodic FFT + numpy reduction + numpy +# trapezoid, converged in its own guard and +# factor), the sample-rate ladder closing on +# that value, the reduce-then-refine order +# being a different number, the refusal set +# still refusing, the two fail-closed doors +# (return_lnLt, rotation norms), the single +# definitions of the guard pair and the +# distance reduction, and one subprocess run +# of the driver through --mode flowmc +# --distance-marginalization. Real +# precompute; needs lal, no GPU. +# test_distance_gh_nodes_cli.py 27 --distance-gh-nodes: makes the per-sample +# Gauss-Hermite distance quadrature (previously +# reachable only via JAX_ILE_DISTMARG_GH) an ILE +# argument, and the warn-not-silently-ignore +# compatibility notes for --phase-marginalization +# on the four phi_ref-analytic modes, +# --sky-coordinates on the modes that do not +# implement it, and --d-prior (always volumetric). +# Parse-time CLI/env resolution and refusal run the +# real check_critical_and_report in-process (no +# subprocess), including that a CLI/env conflict on +# DIFFERENT nonzero values is REFUSED rather than +# reconciled, and that a refused command line never +# mutates core._DISTMARG_GH_N. BLOCKER fix (external +# review, same day): the option now defaults to None, +# not 0, so an explicit --distance-gh-nodes 0 is +# distinguishable from not-passed; four tests pin +# this -- explicit 0 against a nonzero env refuses, +# explicit 16 against agreeing env 16 is accepted, +# env 16 alone resolves and is named in the banner, +# and env 16 against CLI 32 refuses (the mutation +# target: a reversed CLI/env priority passes every +# other test in this file, since most cases here +# exercise only one of the two knobs). A numeric liveness +# check on cheap synthetic packed data (no lal, no +# frames) pins that the resolved count actually +# changes the constructed likelihood's VALUE, not +# just an echoed CLI flag, with a fresh jax.jit +# closure per setting so no stale trace can mask the +# difference; and that 0 reproduces the untouched +# legacy grid bit-for-bit. Two subprocess checks +# against the real driver entry point (--inj-mode, +# tiny budget, stopped at the same known +# post-construction validation error +# test_distance_grid_loguniform.py's own subprocess +# test relies on) confirm the resolved count reaches +# the run log end to end. Needs no lal beyond what +# build_likelihood_data already requires; no GPU. FILES=( "${JAXDIR}/test_jax_time_quadrature.py" @@ -332,6 +516,7 @@ FILES=( "${JAXDIR}/test_jax_slowrot_cauchy_schwarz.py" "${JAXDIR}/test_network_coords.py" "${JAXDIR}/test_nuts_phimarg.py" + "${JAXDIR}/test_jax_av.py" "${JAXDIR}/test_jax_fairdraw_export.py" "${JAXDIR}/test_jax_tempering_chooser.py" "${JAXDIR}/test_tvals_grid_convention.py" @@ -339,6 +524,7 @@ FILES=( "${JAXDIR}/test_jax_stencil_parity.py" "${JAXDIR}/test_flow_reuse_default.py" "${JAXDIR}/test_angle_marg_sizing_rule.py" + "${JAXDIR}/test_anglemarg_buffer_cap.py" "${JAXDIR}/test_angle_marg_smoke.py" "${JAXDIR}/test_angle_marg_compile_cost.py" "${JAXDIR}/test_angle_marg_block_dispatch.py" @@ -348,7 +534,25 @@ FILES=( "${JAXDIR}/test_angle_marg_gh_selection.py" "${JAXDIR}/test_joint_anglemarg_peaklocal.py" "${JAXDIR}/test_angle_marg_peaklocal_wiring.py" + "${JAXDIR}/test_angle_marg_multipeak_wiring.py" "${JAXDIR}/test_limit_distance_jax.py" + "${JAXDIR}/test_direct_marginalization_planner.py" + "${JAXDIR}/test_time_first_peaklocal.py" + "${JAXDIR}/test_all_axis_peaklocal.py" + "${JAXDIR}/test_is_proposal_jitter.py" + "${JAXDIR}/test_multipeak_planner.py" + "${JAXDIR}/test_multipeak_fallback_visibility.py" + "${JAXDIR}/test_jax_phase_marg_mode_order.py" + "${JAXDIR}/test_jax_q_time_pregrid.py" + "${JAXDIR}/test_direct_marginalization_policy.py" + "${JAXDIR}/test_reserve_pair_selection.py" + "${JAXDIR}/test_angle_marg_laplace_table.py" + "${JAXDIR}/test_distance_gh_nodes_cli.py" + "${JAXDIR}/test_jax_cache.py" + "${JAXDIR}/test_direct_marginalization_policy_cli.py" + "${JAXDIR}/test_jax_bandlimited_distmarg.py" + "${JAXDIR}/test_jax_bandlimited_6d_blind.py" + "${JAXDIR}/test_policy_peaklocal_reserve.py" ) # EXCLUDED: files in JAXDIR matching test_*.py that are deliberately NOT gated. The @@ -356,7 +560,16 @@ FILES=( # 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=( + # importorskip("flowMC"): flowMC is deliberately not installed in jax-ile-check + # (ci.yml), and the OUTCOME check below rejects a skip. The prior-mc and + # laplace-is driver tests in the same file are the executable coverage that + # runs here; run the flowMC one by hand where flowMC is installed. + "${JAXDIR}/test_jax_bandlimited_distmarg.py::test_driver_runs_flowmc_distance_marginalized_bandlimited" "${JAXDIR}/test_jax_stencil_parity.py::test_gpu_gather_parity_against_numpy_window" + "${JAXDIR}/test_multipeak_planner.py::test_hm_second_mode_survives_unsafe_proxy_gap" + "${JAXDIR}/test_multipeak_planner.py::test_hm_two_tier_integral_matches_overcomplete_oracle" + "${JAXDIR}/test_multipeak_planner.py::test_real_low_snr_declines_to_finite_reserve" + "${JAXDIR}/test_multipeak_planner.py::test_real_high_snr_two_tier_path_matches_overcomplete_oracle" ) EXCLUDED=( # test_angle_marg_exact.py -- the angle-marginalization VALIDATION suite. @@ -400,6 +613,25 @@ EXCLUDED=( # 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. +# +# test_multipeak_planner.py::test_hm_second_mode_survives_unsafe_proxy_gap +# test_multipeak_planner.py::test_hm_two_tier_integral_matches_overcomplete_oracle +# test_multipeak_planner.py::test_real_low_snr_declines_to_finite_reserve +# test_multipeak_planner.py::test_real_high_snr_two_tier_path_matches_overcomplete_oracle +# The four real-table oracle regressions of the multi-peak planner. Each is +# skipif-guarded on an external validation packet -- a saved (C_A, C_B) coefficient +# table from a real analysis -- and NOTHING in this repository or in the CI setup +# supplies one, so on this runner all four skip. A skip is precisely what the +# post-run junit check below refuses, so leaving them selected would redden the +# gate on every PR while asserting nothing. They cannot be made to run from a +# synthetic fixture either: they pin numbers measured on those tables (mode +# spacings, oracle log-integrals) to ~1e-8, which is a property of the real +# tables and not of any stand-in this repo could ship. +# The 15 remaining tests in that file are self-contained and stay gated; they +# carry the planner's structural coverage (symmetry expansion, strict stationary +# refinement, two-tier convergence, overlap ownership, reserve fallback). +# RUN THE FOUR BY HAND, with the packets present, when touching +# multipeak_planner.py, and record the numbers in the PR per records-protocol. DESELECT=() for t in "${DESELECTED_TESTS[@]}"; do DESELECT+=( --deselect "$t" ); done @@ -477,7 +709,365 @@ fi # 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 +# +# The u-FALLBACK branch adds 2 in test_joint_anglemarg_peaklocal.py (required_u_nodes +# is derived and follows the sqrt-A law under a cap, and a whole-cell integration sized +# by it agrees with a 4x finer one). The floor is 310, READ FROM THIS JOB'S OWN LOG. +# Two wrong numbers preceded it, failing in opposite directions: +# 308 -- by adding 2 to the previous 306, which is exactly what the paragraph above +# says not to do. The base is 309 after #239 merged, so 308 would still have +# PASSED while silently under-promising three tests. +# 311 -- by running the collection on a dev host. Wrong by exactly one, because the +# harness sliced this script by line number to reuse FILES and stopped before +# the loop that populates DESELECT from DESELECTED_TESTS -- so it counted +# test_gpu_gather_parity_against_numpy_window, which THIS job deselects. +# Arithmetic lands below the truth and passes; a mis-set-up local collection lands above +# it and fails. Read the floor off this job's "collected N tests from 27 files" line -- +# the only source that is not a guess. +# The production-policy follow-up adds one mutation-bearing streaming test; this job's +# own collection reports 312. +# PR #250 adds test_anglemarg_buffer_cap.py, test_direct_marginalization_planner.py and +# test_time_first_peaklocal.py; 247 adds four tests to test_joint_anglemarg_peaklocal.py +# and REMOVES test_joint_angle_algebraic.py with the duplicate enumerator it covered. +# Neither branch guessed well: 250 derived a provisional 408 from arithmetic and said to +# replace it with a real collection, and 247 measured 329 against a different file set. +# This number is the MERGED collection, run over this job own FILES/DESELECT with the +# DESELECT loop actually applied (the run reports "424/425 tests collected (1 deselected)"). +# +# 250 also inferred a standing "this environment collects one more than CI" offset and +# subtracted it. There is no such offset: the 311 case documented above was wrong by one +# because a harness sliced this script by line number and never ran the DESELECT loop, so +# it counted the one test this job deselects. That was a one-off setup bug, not a property +# of the environment, and subtracting for it would under-promise by one -- which is the +# failure direction this whole comment exists to warn about, because a low floor PASSES. +# +# The #227 IS-proposal branch then adds the 29 pins in test_is_proposal_jitter.py (27, +# plus the two external review's P1 required: the Markov floor and the inflated-pilot +# regression case). The FILES array above takes the UNION of every side that has +# touched it. +# +# The phase-marginalization mode-order branch then adds the 14 pins in +# test_jax_phase_marg_mode_order.py. Its number was NOT derived by adding 14 to the +# constant above -- that shortcut is what the paragraphs below warn about. It was read +# off this job's own line after rebasing on rift_O4d: "collected 475 tests from 32 +# files", with the DESELECT loop applied (the script itself prints it, so there is no +# way to run this and get the half-configured count). +# +# THIS BRANCH HAS NOW HIT THIS CONFLICT THREE TIMES, on three consecutive days, and +# every one of its own numbers was read off a real collection run when written: +# +# 2026-09-05 339 stale within a day (main had moved 293 -> 312) +# 2026-09-06 451 stale within a day (main had moved 312 -> 424) +# 2026-09-06 453 stale by the next merge (main had moved 424 -> 432) +# +# So the constant does not go stale because someone was careless. It goes stale +# because rift_O4d moves faster than any single branch can hold a global count, and +# that is a property of the counter, not of the people using it. Note that the +# arithmetic would have been RIGHT all three times (312+27 = 339 after the two review +# tests landed later, 424+27 = 451, 432+29 = 461). That is exactly what makes it an +# unreliable shortcut rather than a safe one: it is nearly always right, so the once it +# is wrong there is no habit of checking left to catch it. The number below is READ +# OFF this job's own collection line after the #227 merge: 461/462 collected, +# 1 deselected, 31 files. PR #270 adds 15 multipeak tests but deselects the four +# real-table regressions whose external packets CI does not provide. The merged +# gate therefore adds 11 self-contained tests. Confirmed from the merged +# collection: 472/477 collected, 5 deselected, 32 files. +# +# The phase-marginalization mode-order merge added the 14 pins in +# test_jax_phase_marg_mode_order.py on top of #270, and hit the same conflict a +# fourth time: each side of it carried a number the other side had already +# invalidated (475 vs 472). Read off the job's own line then: 486 from 33 files. +# +# FIFTH time, on the Q time-pregrid branch (this merge). Both sides were stale +# again -- 480 on the branch, 486 on rift_O4d -- for the same reason, and the +# resolution is again a MEASUREMENT, not the sum. Read off this job's own line +# after resolving, DESELECT loop applied, on the merged tree: +# "collected 505 tests from 34 files". The arithmetic (486 + the 19 in +# test_jax_q_time_pregrid.py) agrees, and is again not where the number came +# from. Independently recollected on citlogin6 with ~/.cache/jaxci_venv +# (jax 0.9.2, numpyro 0.21.0) during the landing review: the same 505 from the +# same 34 files. +# +# 507, not 505, and the gap is two tests added after that measurement: the P2 +# review's no-metadata refusal, and one for a path the merge creates that +# neither side covers (#272's phase-marginalized mode permutation acting on a +# REFINED Q grid). The P2 commit left this constant at 505, which the >= floor +# accepts silently -- exactly the drift this comment exists to stop. Recollected +# after both: "507/512 tests collected (5 deselected)", +# "collected 507 tests from 34 files". +# +# SIXTH time, on the full-circuit phi-region branch (this merge). Both sides were +# stale in the usual way -- 487 on the branch, 507 on rift_O4d -- and the resolution +# is again a MEASUREMENT. This branch adds ONE test, +# test_a_full_circuit_phi_window_is_one_region_at_every_peak_location. Recollected on +# the merged tree, citlogin6, ~/.cache/jaxci_venv, DESELECT loop applied: +# "collected 508 tests from 34 files". +# SEVENTH time, on the four-axis peak-local branch (#268, this merge). Both +# sides stale again: 453 on the branch, 507 and then 508 on rift_O4d. +# Resolution is again a measurement on the merged tree with the DESELECT loop +# applied, read off this +# job's own collection line on citlogin6 (~/.cache/jaxci_venv, jax 0.9.2): +# "542/547 tests collected (5 deselected)", gate-style count 542 from 35 +# files. Def-count arithmetic (508 + 30 in test_all_axis_peaklocal.py + 1 in +# test_angle_marg_exact.py + 2 in test_time_first_peaklocal.py = 541) does NOT +# reproduce it, which is one more reason the constant is measured. +# +# EIGHTH: the cross-axis policy wiring adds test_direct_marginalization_policy.py +# (15 tests). Measured on this tree with the DESELECT loop applied, citlogin6, +# ~/.cache/jaxci_venv (jax 0.9.2): "557/562 tests collected (5 deselected)", +# gate-style count 557 from 36 files. Independently recollected with the CVMFS +# igwn python on ldas-pcdev11 during the same landing: same 557 from 36 files. +# +# NINTH, on the multi-peak refinement-stall branch (this change). It adds FOUR +# tests to test_multipeak_planner.py and touches no other test file. The +# branch measured 546 against a base of 542; #278 has since taken the base to +# 557, so 546 is stale and 557+4 would be the arithmetic this comment forbids. +# Re-measured on the merged tree, DESELECT loop applied, read off the gate's +# own collection line: +# "561/566 tests collected (5 deselected)", gate-style count 561 from 36 files. +# +# TENTH, on the multi-peak fallback-visibility branch (#277, this merge). It +# adds ONE file, test_multipeak_fallback_visibility.py (13 tests), and touches +# no existing test file. The branch measured 555 against a base of 542; #278 +# and #279 have since taken the base to 561, so 555 is stale and neither side +# nor their sum is usable. Re-measured on the merged tree, ldas-grid, +# ~/.cache/jaxci_venv (jax 0.9.2, numpyro 0.21.0), DESELECT loop applied, read +# off this job own collection line: +# "574/579 tests collected (5 deselected)", gate-style count 574 from 37 +# files. +# +# The policy follow-up PR (this branch, numbered NINTH on its own side before the +# merge) adds three wiring tests (guard preflight, +# operating-point defaults, and a parametrized fail-closed case). Measured on +# the follow-up tree with the DESELECT loop applied, ldas-grid, CVMFS igwn +# python: "560/565 tests collected (5 deselected)", gate-style count 560 from +# 36 files. +# Both sides stale again after PR #279 (multipeak planner Newton step) landed +# under this branch: re-measured on the merged tree, ldas-grid, CVMFS igwn +# python, DESELECT applied: "564/569 tests collected (5 deselected)", gate-style +# count 564 from 36 files. +# +# ELEVENTH, reconciling the policy follow-ups with #277 (this merge). Both sides +# were stale in the usual way -- 564 on the branch, 574 on rift_O4d -- and neither +# number nor their difference describes the merged tree, because each side counted +# a file set the other had already changed. The FILES array is again the UNION, +# now 37 files. Re-measured on the merged tree with the DESELECT loop applied, +# ldas-grid, CVMFS igwn python +# (/cvmfs/software.igwn.org/conda/envs/igwn/bin/python), read off this job's own +# collection line: "577/582 tests collected (5 deselected)", gate-style count 577 +# from 37 files. +# +# TWELFTH, adding test_distance_gh_nodes_cli.py (--distance-gh-nodes, this branch). +# 20 test_* entry points, one parametrized x4, so 23 collected; none deselected. +# FILES is now 38 files, EXPECTED_TESTS raised by exactly that: 577 + 23 = 600. +# +# THIRTEENTH, same branch, same day: adversarial review found a BLOCKER (the +# 0-default made an explicit --distance-gh-nodes 0 indistinguishable from +# not-passed, so a nonzero JAX_ILE_DISTMARG_GH silently won). Fixed with a +# None default and four new tests pinning CLI-given-including-0 wins, plus a +# mutation-target regression test for the reversed-priority case. 24 test_* +# entry points now, one parametrized x4, so 27 collected; none deselected. +# EXPECTED_TESTS raised by exactly that: 600 + 4 = 604. +# +# FOURTEENTH, on rift_O4d (#285, not this branch): the jax 0.9.2 empty-pool cap fix +# touches only test_anglemarg_buffer_cap.py: removes 1 test +# (test_a_zero_largest_free_block_is_a_known_full_device, whose "0 means full" premise +# was the bug) and adds 6, net +5, none parametrized. 577 + 5 = 582. +# +# FIFTEENTH, same file, the review-MAJOR follow-up (forced probe allocation before +# reading memory_stats(), plus the on-demand-allocator bound): adds 5 tests, none +# parametrized, no removals. 582 + 5 = 587. +# +# SIXTEENTH, reconciling TWELFTH/THIRTEENTH (this branch, test_distance_gh_nodes_cli.py, +# +27 off the 577 base) with FOURTEENTH/FIFTEENTH (rift_O4d #285, +# test_anglemarg_buffer_cap.py, +10 off the same 577 base) at this merge. The two +# deltas land in disjoint files, so unlike the earlier reconciliations in this history +# the sum is exact, not a guess: 577 + 27 + 10 = 614. FILES is 38 (37 + this branch's +# one new file; #285 added no file). +# THIRTEENTH, the review-MAJOR follow-up (forced probe allocation before reading +# memory_stats(), plus the on-demand-allocator bound) again touches only +# test_anglemarg_buffer_cap.py: adds 5 tests, none parametrized, no removals. +# 582 + 5 = 587. +# FOURTEENTH, on the JAX persistent/transferable compilation cache (#214, this +# merge). It adds ONE file, test_jax_cache.py, and changes no test count in an +# existing file: the amplitude failsafe changes HOW it reports -- a returned +# value instead of a host callback, which is what makes the angle-marg graph +# eligible for JAX's persistent cache at all -- but not how many pins cover it. +# The branch opened carrying 189 against a base of 171, both months stale, was +# then measured at 600 against a base of 574, and pushed 603 against a base of +# 577. rift_O4d has since moved to 587 (#285 and its review follow-up) and +# again with #284, so every number either side carries is stale, for the +# fourteenth time running. Re-measured on the merged tree, DESELECT loop +# applied, read off this job's own collection line: +# "621/626 tests collected (5 deselected)", gate-style count 621 from 38 +# files. Independently recollected on citlogin6 and on ldas-grid, same +# interpreter, same 621/626. +# +# The +8 over the branch's own 613 are ALL from this landing, not from the +# branch. Three pin defects found reviewing it (the NOT-PERFORMED label, the +# policy-composite wiring, and a device probe that raised out of driver +# import), and five close mutation survivors: both --jax-cache-dir spellings, +# both cache opt-outs separately, a member declaring zero compressed size, and +# the two refusals that keep an unsized scheme from inventing a metric. +# test_jax_cache.py collects 30, not the 17 the branch's per-file line claimed. +# +# READ THIS BEFORE TREATING A LOCAL RED AS A BRANCH DEFECT. jax and numpyro +# are installed UNPINNED here (see ci.yml for why), so CI and your shell can be +# on different jax versions at the same time, and which one is newer changes +# over time -- do not infer it from this comment. Landing #214, two failures +# reproduced in a local venv on the branch AND on its pristine base while CI +# was green on the whole gate: a trend assertion at the noise floor, and a +# full-suite abort (134/139) in a file the branch never touched. Both were the +# environment, and finding that out cost two runs. +# +# So: check the jax version each side actually ran (the gate prints it as its +# second line), and reproduce any local failure on the PRISTINE BASE in the SAME +# environment before believing it. The collected COUNT has been stable across +# versions; pass/fail has not. Issue #292 tracks the environment spread and +# what to do about it. +# +# One practical note for whoever hits this next, because it cost a wasted run: +# PYTHONPATH must be pinned to the tree under test before collecting. The +# conda environment on the CIT interactive hosts resolves RIFT to a DIFFERENT +# checkout (~/RIFT_ralph), and collection then fails on imports that have +# nothing to do with the branch. +# +# +# FIFTEENTH, the four-axis policy row-batching branch (this merge). It adds +# SEVEN tests to test_direct_marginalization_policy.py (the batched/sequential +# equivalence test, five parametrized validate_batch_rows cases, and the driver +# knob test) and adds no file. Its own side measured 581 against a base of 574; +# rift_O4d has since reached 587, so neither number nor their sum describes the +# merged tree. Re-measured by running THIS script on the merged tree, +# ldas-grid, ~/.cache/jaxci_venv, DESELECT loop applied, read off its own +# collection line: "collected 594 tests from 37 files". +# Both sides were stale in the usual way: 594 on this branch against a base of +# 587, and 621 on rift_O4d, and neither number nor their difference describes +# the merged tree because each counted a file set the other had changed. +# Re-measured on the MERGED tree by running this script and reading its own +# line: "collected 628 tests from 38 files" (ldas-grid, ~/.cache/jaxci_venv, +# DESELECT loop applied). +# +# NEXT, the policy observability branch (this change). It adds ONE file, +# test_direct_marginalization_policy_cli.py, with twenty tests: seventeen +# driver-seam refusals and three that pin the return arity of +# direct_marginalization_policy_note. File count 38 -> 39. Measured on the +# REBASED tree with /scratch/richard.oshaughnessy/envs/jaxci-py311 (python +# 3.11.13, jax 0.10.2) by running the collection and reading its own line, not +# by adding 20 to 628: "648/653 tests collected (5 deselected)". +# FIFTEENTH, the peak-local phi-scan reduction (joint_lnL_phi_dense reduces into its +# lax.scan carry instead of stacking the phi axis; RIFT PR #295). Touches only +# test_angle_marg_peaklocal_wiring.py: replaces +# test_peak_local_model_includes_streamed_body_and_scan_output (2 params) with +# test_peak_local_model_is_flat_in_n_phi_because_the_scan_reduces (the same 2 params) +# and adds test_peak_local_model_does_not_grow_with_the_phi_axis. The file-local delta +# is exact at +1. +# +# NOT 628 + 1, for the reason this block has now recorded four times. Re-measured by +# running THIS script on the merged tree and reading its own collection line: +# "collected 629 tests from 38 files". +# SIXTEENTH, the bandlimited distance-marginalization branch (this merge). It +# adds ONE file, test_jax_bandlimited_distmarg.py, which collects 20 and has one +# test DESELECTED here (its flowMC driver run would importorskip, and a skip +# fails the OUTCOME check), so +19 over the merged base. Its own side carried +# 596 against a base of 577; rift_O4d reached 628 meanwhile. Re-measured by +# running this script on the MERGED tree, ldas-grid, ~/.cache/jaxci_venv, +# DESELECT loop applied, read off its own collection line: +# "647/653 tests collected (6 deselected)", gate-style count 647 from 39 files +# (the new file alone: "19/20 tests collected (1 deselected)"). +# +# SEVENTEENTH, the fixed-distance blind-draw follow-up (endpoint gap off on the +# 6-D field; parse-time window refusal). ONE new file, +# test_jax_bandlimited_6d_blind.py, nothing deselected. Re-measured by running +# this script on this tree, ldas-grid, ~/.cache/jaxci_venv, DESELECT loop +# applied, read off its own collection line: +# "collected 657 tests from 40 files" (the gate's own line; 647 + 10) +# SEVENTEENTH, merging rift_O4d (#214's test_jax_cache.py, gate count 621 from 38 +# files) into this branch (test_distance_gh_nodes_cli.py, +27): the two deltas +# land in disjoint files, so 621 + 27 = 648 from 39 files. Re-measured on the +# merged tree with the CI-equivalent /scratch jaxci-py311 interpreter before +# this commit. + +# EIGHTEENTH, the YOLO integration merge of 2026-09-08 (RIFT PRs #286, #295, +# #297, #298, #299 -- #299 carries #288 -- merged onto rift_O4d after #294). +# Every block above was measured on its own tree, so none of their numbers nor +# their sum describes this one. Re-measured by running THIS script on the merged +# tree (ldas-grid, ~/.cache/jaxci_venv, jax 0.9.2, DESELECT loop applied) and +# reading its own collection line: "collected 705 tests from 42 files". This +# assignment is the one that binds; the earlier ones are kept as provenance. +# ONE ASSIGNMENT ONLY. This file carried FIVE consecutive unconditional +# EXPECTED_TESTS= assignments (648, 629, 657, 648, 705) accumulated by parallel +# merges, separated only by their comment blocks. Bash keeps the LAST, so the +# four above it were dead while reading as authoritative -- the same shadowing +# that the driver's duplicate add_option produced, one file over. The comment +# history is kept; the dead assignments are not. Re-derive by running this +# script and reading its own collection line, never by adding a delta. +# Measured on this branch: 731 collected, 44 files, 6 deselected -- superseded by +# the twentieth block below, which is the one that binds. + +# NINETEENTH, RIFT PR #301 (preset local-plan capacities) merged with rift_O4d +# at 43918b22. #301 adds four tests to +# test/jax/test_direct_marginalization_policy.py, a file already in FILES. +# The eighteenth block measured 705 on the integration tree, so the tempting +# number here is 705 + 4 = 709. That is wrong: the gate's own line on THIS +# tree reads "collected 712 tests from 42 files" (ldas-grid, +# /scratch/$USER/envs/jaxci-py311, jax 0.10.2, DESELECT loop applied). The two +# trees are not the same tree, which is the whole reason this file says to +# measure and never to add. +# +# TWENTIETH, this branch (PR #305) rebased onto #301 -- first onto its head +# f30b6a06 while it was open, then onto rift_O4d c0654025 once #301 merged. +# The second rebase was CLEAN and the count did not move; it was re-measured +# anyway, because a clean rebase is when this file's one-assignment property is +# most likely to have been quietly undone. The nineteenth block and +# the eighteenth-plus-mine block were BOTH left in the file by that rebase, in +# that order, and git merged them without a conflict because they touch +# different lines. Bash keeps the last, so #301's 712 silently replaced the 731 +# this branch had measured -- the same shadowing this file was already collapsed +# once to remove, reintroduced by a clean rebase rather than by an edit. A +# textual merge cannot see that two assignments to one name are in conflict, so +# ONE ASSIGNMENT ONLY is a property this file has to be re-checked for after +# every merge, not one it keeps on its own. +# Re-measured by running this script on the rebased tree (ldas-grid, +# /scratch/$USER/envs/jaxci-py311, jax 0.10.2, PYTHONPATH pinned to THIS +# checkout, DESELECT loop applied): "collected 733 tests from 45 files", and +# 737 after the four reserve-roster/refusal tests added later in the branch, and +# 738 once the wrapper's source-text gate test became two behaviour tests, and +# 740 with the two reserve-resolution tests. +# This assignment is the one that binds. +# Plus the peak-local time reserve branch (PR #304): test_policy_peaklocal_reserve.py, +# one file. Re-measured by running this script on the rebased tree, ldas-grid, +# ~/.cache/jaxci_venv, DESELECT loop applied, read off its own collection line: +# "754/760 tests collected (6 deselected)", gate-style count 754 from 46 files +# (2026-09-09). This assignment is the one that binds. +# +# SIXTEENTH, --angle-marg-scheme multipeak (the four-axis controller wired into the +# driver). Adds ONE file, test_angle_marg_multipeak_wiring.py, 7 tests, none +# parametrized, no removals, and touches no existing test count. 754 + 7 = 761, +# measured by running this script. +# (superseded assignment removed 2026-09-09; see the single EXPECTED_TESTS= below) +# 2026-09-09 (laplace reserve kernel keyword fix): +1 test in test_policy_peaklocal_reserve.py +# (the resolved kernel through anglemarg's REAL Laplace function). Read off this script's +# own collection line on ldas-pcdev12 (~/.cache/jaxci_venv, CPU): "collected 755 tests from 45 files". +# (superseded assignment removed 2026-09-09; see the single EXPECTED_TESTS= below) + +# Simultaneous rotation + finite response adds cheap analytic coefficient parity +# to an existing collected test. Real waveform precompute, JIT/grad, the one-call +# wrapper, and scaling profiles remain explicit manual checks in the same files: +# they are too expensive for the already runner-limited per-PR JAX gate. + +# 2026-09-09, #313 rebased over #312 (laplace keyword fix, 755) and the base's +# test_limit_distance_jax tightening: neither 761 + 1 nor 755 + 7 is the number. +# Re-measured by running THIS script on the merged tree (ldas-grid, ~/.cache/jaxci_venv, +# jax 0.9.2, CPU, DESELECT applied): "collected 762 tests from 46 files". Binding. +# 2026-09-09 (locator search sizing, stacked on #312 + #313): +1 test in +# test_policy_peaklocal_reserve.py (the sized locator on a rho-632 carrier). Read off this +# script's own collection line on ldas-pcdev12 (~/.cache/jaxci_venv, CPU) after rebasing on +# rift_O4d 336f86133: "collected 763 tests from 46 files". The #312/#313 merges had left +# three EXPECTED_TESTS= assignments (761, 755, 762; last wins); this is the single one. +# 2026-09-10: +22 value-only AV/portfolio, prior-window, wrapper, and driver +# contract tests in test_jax_av.py. +EXPECTED_TESTS=785 echo "== collection floor check (expect >= ${EXPECTED_TESTS} tests) ==" collect_out="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${DESELECT[@]}" "${FILES[@]}" 2>&1)" @@ -525,8 +1115,60 @@ 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}" "${DESELECT[@]}" "${FILES[@]}" +# SHARDING. Only the EXECUTE step is split. Everything above -- the FILES +# manifest, the EXCLUDED accounting, the collection floor against +# EXPECTED_TESTS, and the deselect-resolution check -- runs in full in every +# shard, so no shard can pass on a partial view of the suite and the counts +# stay one number rather than N. +# +# Why: the suite outgrew the 60-minute job cap. Measured 2026-09-08, the base +# suite ran 3277 s of pytest against a ~3518 s budget, and one new test in #290 +# added ~850 s, so every run was ~609 s over and jax-ile-check was cancelled at +# 1h00m17s with the suite at 91% and ZERO failures. Trimming was costed at +# ~529 s across five changes and does not clear the overrun on its own. Three +# shards leave each well inside the cap with room for the next test. +# +# Round-robin by index, not a hand-tuned split: a cost table in this file would +# go stale exactly the way the EXPECTED_TESTS comments above record every other +# hardcoded number going stale. +JAX_GATE_SHARDS="${JAX_GATE_SHARDS:-1}" +JAX_GATE_SHARD="${JAX_GATE_SHARD:-1}" +if ! [ "${JAX_GATE_SHARDS}" -ge 1 ] 2>/dev/null || ! [ "${JAX_GATE_SHARD}" -ge 1 ] 2>/dev/null \ + || [ "${JAX_GATE_SHARD}" -gt "${JAX_GATE_SHARDS}" ]; then + echo "test-jax.sh: bad shard ${JAX_GATE_SHARD}/${JAX_GATE_SHARDS}" >&2; exit 1 +fi +if [ "${JAX_GATE_SHARDS}" -gt 1 ]; then + SHARD_FILES=() + for i in "${!FILES[@]}"; do + if [ $(( i % JAX_GATE_SHARDS )) -eq $(( JAX_GATE_SHARD - 1 )) ]; then + SHARD_FILES+=( "${FILES[$i]}" ) + fi + done + if [ "${#SHARD_FILES[@]}" -eq 0 ]; then + echo "test-jax.sh: shard ${JAX_GATE_SHARD}/${JAX_GATE_SHARDS} got 0 files" >&2 + exit 1 + fi + echo "== running shard ${JAX_GATE_SHARD}/${JAX_GATE_SHARDS}: ${#SHARD_FILES[@]} of ${#FILES[@]} files ==" +else + SHARD_FILES=( "${FILES[@]}" ) + echo "== running ==" +fi +# The OUTCOME floor below must be the count THIS invocation was asked to run. With +# JAX_GATE_SHARDS>1 that is the shard's own collection, not EXPECTED_TESTS: the +# whole-suite floor has already been asserted above on the full FILES list, and +# holding one shard to it fails every shard that passes ("ran 257 tests, expected +# at least 705", 2026-09-08, first run of the three-way split). Collect the shard's +# files the same way, so a shard still cannot go green on a partial view of ITS files. +if [ "${JAX_GATE_SHARDS}" -gt 1 ]; then + shard_collect="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${DESELECT[@]}" "${SHARD_FILES[@]}" 2>&1)" || { + printf '%s\n' "${shard_collect}"; echo "test-jax.sh: shard collection failed" >&2; exit 1; } + RUN_EXPECTED="$(printf '%s\n' "${shard_collect}" | grep -cE '^[^[:space:]]+\.py::')" + echo "shard ${JAX_GATE_SHARD}/${JAX_GATE_SHARDS} collects ${RUN_EXPECTED} tests from ${#SHARD_FILES[@]} files" + if [ "${RUN_EXPECTED}" -lt 1 ]; then echo "test-jax.sh: shard collected 0 tests" >&2; exit 1; fi +else + RUN_EXPECTED="${EXPECTED_TESTS}" +fi +"${PYTHON_BIN}" -m pytest -q -p no:cacheprovider --durations=0 --junit-xml="${junit}" "${DESELECT[@]}" "${SHARD_FILES[@]}" rc=$? if [ "${rc}" -ne 0 ]; then # rc 5 == "no tests ran"; it is a FAILURE here, not a pass. @@ -538,7 +1180,7 @@ fi # 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. That is # the very shape this script exists to prevent, so assert what the RUN did. -"${PYTHON_BIN}" - "${junit}" "${EXPECTED_TESTS}" <<'PYCHECK' +"${PYTHON_BIN}" - "${junit}" "${RUN_EXPECTED}" <<'PYCHECK' import sys, xml.etree.ElementTree as ET path, expected = sys.argv[1], int(sys.argv[2]) root = ET.parse(path).getroot() @@ -562,4 +1204,4 @@ if bad: PYCHECK if [ $? -ne 0 ]; then exit 1; fi -echo "jax_ile CPU regression gate: PASS (${n_collected} tests)" +echo "jax_ile CPU regression gate: PASS (${n_collected} tests collected; this invocation ran ${RUN_EXPECTED})" diff --git a/.travis/test-q-window-stencil.sh b/.travis/test-q-window-stencil.sh index 8aad3b979..5326a3fd5 100755 --- a/.travis/test-q-window-stencil.sh +++ b/.travis/test-q-window-stencil.sh @@ -136,9 +136,21 @@ SCOPE_GLOBS=( # 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. +# +# test_noloop_accumulator_ Belongs to another job, not to a GPU. It matches the +# shapes.py test_noloop_* pattern by NAME but not by subject: it pins +# NoLoop's rho_sq and kappa_sq ACCUMULATOR shapes against a +# reference, and its time-integral test is about the +# quadrature rule, not about sub-sample interpolation of +# Q_lm. It is registered with core-unit-check, whose FILES +# manifest carries it and whose floors count it. Listed +# here rather than renamed so the decision is recorded where +# the next such file will hit it: renaming to dodge a +# manifest is how these gates quietly stop covering things. EXCLUDED=( "${CODEDIR}/RIFT/likelihood/test_q_window_interp_gpu.py" "${CODEDIR}/RIFT/likelihood/test_noloop_gpu_stencils.py" + "${CODEDIR}/test/test_noloop_accumulator_shapes.py" ) echo "== registered files (marker: ${MARKER}) ==" @@ -210,15 +222,15 @@ fi # 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 +EXPECTED_TESTS=78 +EXPECTED_PASSED=75 # 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 +MAX_SKIPS=3 # 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 diff --git a/.travis/test-roster-verify.py b/.travis/test-roster-verify.py new file mode 100755 index 000000000..a20162510 --- /dev/null +++ b/.travis/test-roster-verify.py @@ -0,0 +1,304 @@ +#!/usr/bin/env python3 +"""Check that each roster entry's STATUS is still true, not merely present. + +WHY THIS EXISTS. .travis/test-ci-roster.py enforces that every ungated test file carries a +reason. It cannot tell whether the reason is CORRECT, and that gap is not theoretical: the +roster asserted "skips cleanly without them" for two jax_gp files while one of them ERRORED +without jax and the other had no guard at all (PR #248). A reason nobody re-checks is prose, +and prose is what this whole census exists to stop being mistaken for coverage. + +So each status carries a FALSIFIABLE, direction-checked predicate, and this job runs it: + + LEGACY must FAIL to import, and only a collection/import ERROR shows that. A clean + collection satisfies nothing -- neither one that finds tests nor one that finds + none -- because both mean the pre-package module names it supposedly needs are + resolving, and the file is a candidate for real gating. + HANDRUN must collect NO tests. If it collects some, it is a pytest suite wearing the + wrong label -- and one that no job runs. + EXPENSIVE must collect tests AND, without RIFT_RUN_EXPENSIVE, SKIP them in a run that exits + cleanly. Passing none is NOT the predicate: a suite that fails or errors passes + none too, and a broken suite is not a guarded one. Catches both a suite that + stopped collecting and an opt-in guard that stopped guarding. + OPTDEP must declare its dependencies as `needs:[,]` or `needs:env:VAR`, none of + which may appear in requirements.txt -- if CI installs it, it is not optional. The + behavioural half is keyed off whether those deps are ACTUALLY present in the running + environment, because that differs between CIT and a runner: with one absent the file + must not collect-and-fully-pass; with all present it may. An earlier version of this + check simply flagged "collects and all pass here", which fired on CIT purely because + CIT has jax -- a check that reported the environment rather than the claim. + +DELIBERATELY NOT CHECKED: which dependency an OPTDEP file wants, and whether a HANDRUN study's +internal gate still holds. Both need the missing stack or a long run; claiming to check them +would be the same overreach this file exists to catch. The predicates above are the part that +is decidable HERE, and the docstrings say so. + +This is a SEPARATE job from ci-roster-check on purpose: that one is stdlib-only with no +`needs: install`, and must stay that way so it reports even when the install matrix is broken. +This one needs RIFT importable. +""" + +import os +import re +import subprocess +import sys + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +ROSTER = os.path.join(".travis", "ci_roster.txt") +CODE = os.path.join("MonteCarloMarginalizeCode", "Code") +TIMEOUT = 300 +CLEAN_RCS = (0, 5) # pytest: 0 = collected/ran without error, 5 = imported fine, no tests here + + +def _read_roster(): + out = [] + for n, raw in enumerate(open(ROSTER, errors="replace"), 1): + if raw.lstrip().startswith("#") or not raw.strip(): + continue + bits = raw.rstrip("\n").split(None, 2) + if len(bits) >= 3: + out.append((n, bits[0], bits[1], bits[2].strip())) + return out + + + +def _declared_deps(reason): + """Modules / env vars an OPTDEP entry claims, from `needs:a,b` or `needs:env:VAR`.""" + m = re.search(r"needs:([A-Za-z0-9_.,:]+)", reason) + return [d for d in m.group(1).split(",") if d] if m else [] + + +# Import names that installed metadata cannot resolve to their distribution. Kept SHORT and +# justified: on CIT, `lalsuite` is a conda metapackage whose dist-info lists no top-level modules +# at all (its files() shows only __pycache__ and the dist-info), so neither packages_distributions +# nor a file scan can learn that `lal` comes from it. A pip-installed lalsuite wheel does declare +# them, so this table is a fallback for the environment, not a replacement for the lookup. +# Add an entry only when the two mechanisms below genuinely cannot answer. +KNOWN_ALIASES = { + "lal": "lalsuite", "lalsimulation": "lalsuite", "lalframe": "lalsuite", + "lalmetaio": "lalsuite", "lalburst": "lalsuite", "lalinspiral": "lalsuite", + "lalpulsar": "lalsuite", "lalinference": "lalsuite", +} + + +def _distributions_for(mod): + """Distribution names that provide this IMPORT name, e.g. sklearn -> {scikit-learn}. + + Three mechanisms, cheapest first: the metadata index, a scan of each distribution's files + for a top-level `/` or `.py`, and finally KNOWN_ALIASES for distributions whose + metadata declares nothing. + """ + top = mod.split(".")[0] + out = set() + try: + from importlib.metadata import packages_distributions, distributions, files + except ImportError: # pragma: no cover - py<3.10 + return {KNOWN_ALIASES[top]} if top in KNOWN_ALIASES else set() + try: + out |= set(packages_distributions().get(top, [])) + except Exception: # pragma: no cover - defensive + pass + if not out: + try: + for d in distributions(): + name = (d.metadata["Name"] or "") + if not name: + continue + for f in (files(name) or []): + parts = str(f).split("/") + if parts[0] == top or parts[0] == top + ".py": + out.add(name) + break + except Exception: # pragma: no cover - defensive + pass + if not out and top in KNOWN_ALIASES: + out.add(KNOWN_ALIASES[top]) + return out + + +def _in_requirements(mod): + """True if requirements.txt installs this module -- in which case it is not optional. + + THE IMPORT NAME IS NOT THE DISTRIBUTION NAME, and comparing them directly made this check + fail open: `needs:sklearn` sailed past a requirements.txt that says `scikit-learn`, and so + did `needs:lal` against `lalsuite` -- both importable in CI, both therefore NOT optional, + both silently accepted as OPTDEP. Reproduced before this fix for sklearn, lal, + lalsimulation and skimage. + + So the import name is resolved to the distributions that provide it, and any of those + matching a requirements line counts. LIMIT, stated because it is real: the mapping comes + from installed metadata, so a dependency absent from the CHECKING environment cannot be + resolved and falls back to the bare name comparison. In this job requirements.txt is + installed, which is exactly the case that matters. + """ + try: + req = open(os.path.join(REPO, "requirements.txt"), errors="replace").read() + except OSError: + return False + want = {n.lower().replace("-", "_") for n in ({mod} | _distributions_for(mod))} + for line in req.splitlines(): + line = line.split("#", 1)[0].strip() + if not line: + continue + name = re.split(r"[<>=\[]", line)[0].strip().lower().replace("-", "_") + if name in want: + return True + return False + + +def _dep_present(dep): + """Is this declared dependency actually available in the environment running the check?""" + if dep.startswith("env:"): + return bool(os.environ.get(dep[4:])) + pr = subprocess.run([sys.executable, "-c", "import %s" % dep], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + return pr.returncode == 0 + + +def _pytest(path, extra_env=None, collect_only=True): + """Return (rc, n, n_passed), or (None, None, None) on timeout. + + n is the collected count under --collect-only and the SKIPPED count for a real run, since + that is what tells a guarded suite apart from a broken one. n_passed is -1 when not run. + """ + env = dict(os.environ) + env["PYTHONPATH"] = os.path.join(REPO, CODE) + os.pathsep + env.get("PYTHONPATH", "") + env.setdefault("OMP_NUM_THREADS", "1") + env.setdefault("MPLBACKEND", "Agg") + # A None value REMOVES the variable. The EXPENSIVE predicate needs a run that genuinely + # lacks RIFT_RUN_EXPENSIVE; inheriting it from the caller inverts the whole check -- a + # correct guard then runs its tests and is reported as broken, and an inverted guard skips + # and is reported as fine. Reproduced with the variable exported. + for k, v in (extra_env or {}).items(): + if v is None: + env.pop(k, None) + else: + env[k] = v + cmd = [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider"] + if collect_only: + cmd.append("--collect-only") + cmd.append(path) + try: + pr = subprocess.run(cmd, env=env, cwd=REPO, timeout=TIMEOUT, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + except subprocess.TimeoutExpired: + return None, None, None + text = pr.stdout.decode("utf-8", "replace") + if collect_only: + # -q prints one `path::id` line per collected test; count lines, not `::` occurrences, + # so a class-based id does not count twice. + return pr.returncode, sum(1 for ln in text.splitlines() if "::" in ln), -1 + m = re.search(r"(\d+) passed", text) + s = re.search(r"(\d+) skipped", text) + return pr.returncode, (int(s.group(1)) if s else 0), (int(m.group(1)) if m else 0) + + +def main(): + os.chdir(REPO) + errs, checked = [], {} + for lineno, path, status, reason in _read_roster(): + if not os.path.exists(path): + continue # ci-roster-check owns that error + if status == "LEGACY": + rc, n, _ = _pytest(path) + if rc is None: + errs.append("%s:%d: %s timed out during collection." % (ROSTER, lineno, path)) + elif n > 0: + errs.append("%s:%d: %s is LEGACY (\"cannot be imported\") but COLLECTS %d " + "tests.\n Whatever it needed now resolves. Re-check the reason: " + "it is probably gateable." % (ROSTER, lineno, path, n)) + elif rc in CLEAN_RCS: + errs.append("%s:%d: %s is LEGACY (\"cannot be imported\") but pytest collected it " + "WITHOUT error (exit %d) and found no tests.\n Collecting nothing " + "is not evidence of a failed import -- only a collection error is -- " + "so the reason is stale: whatever it needed now resolves. Re-check it; " + "the file is HANDRUN at most, and probably gateable." + % (ROSTER, lineno, path, rc)) + elif status == "HANDRUN": + rc, n, _ = _pytest(path) + if rc is None: + errs.append("%s:%d: %s timed out during collection." % (ROSTER, lineno, path)) + elif n > 0: + errs.append("%s:%d: %s is HANDRUN (\"not a pytest target\") but COLLECTS %d " + "tests.\n It is a real suite that no job runs -- gate it, or " + "correct the status." % (ROSTER, lineno, path, n)) + elif status == "EXPENSIVE": + # Both calls drop RIFT_RUN_EXPENSIVE: collection too, since a module-level skip may + # key off it and change what is collected at all. + no_optin = {"RIFT_RUN_EXPENSIVE": None} + rc, n, _ = _pytest(path, extra_env=no_optin) + if rc is None or n == 0: + errs.append("%s:%d: %s is EXPENSIVE but collects nothing.\n The opt-in " + "suite is gone or stopped importing." % (ROSTER, lineno, path)) + else: + rc2, skipped, passed = _pytest(path, collect_only=False, extra_env=no_optin) + if rc2 is None: + errs.append("%s:%d: %s is EXPENSIVE but the run WITHOUT RIFT_RUN_EXPENSIVE " + "timed out after %ds.\n Opting out should cost nothing, so " + "something is executing: the guard is not holding." + % (ROSTER, lineno, path, TIMEOUT)) + elif passed > 0: + errs.append("%s:%d: %s is EXPENSIVE (\"skips unless RIFT_RUN_EXPENSIVE=1\") " + "but %d test(s) PASSED without it.\n The opt-in guard stopped " + "guarding." % (ROSTER, lineno, path, passed)) + elif rc2 != 0 or skipped < n: + errs.append("%s:%d: %s is EXPENSIVE (\"skips unless RIFT_RUN_EXPENSIVE=1\") " + "but the run WITHOUT it exited %d with %d of %d collected test(s) " + "skipped.\n Passing nothing is not the same as being guarded: " + "a suite that fails or errors passes nothing either. Opting out " + "must be a CLEAN skip of every collected test." + % (ROSTER, lineno, path, rc2, skipped, n)) + elif status == "OPTDEP": + deps = _declared_deps(reason) + if not deps: + errs.append("%s:%d: %s is OPTDEP but names no dependency.\n" + " Write `needs:[,]` or `needs:env:VAR` in the " + "reason so the claim can be checked instead of believed. Two entries " + "carrying only prose here turned out to collect and pass completely, " + "and belonged in a job." % (ROSTER, lineno, path)) + for d in deps: + if not d.startswith("env:") and _in_requirements(d): + errs.append("%s:%d: %s is OPTDEP on %r, which requirements.txt DOES install.\n" + " Then it is not optional -- gate the file." + % (ROSTER, lineno, path, d)) + missing = [d for d in deps if not _dep_present(d)] + if missing: + rc, n, _ = _pytest(path) + if rc is None: + # A timeout is a verification that did NOT happen. Counting it as checked + # is the same silent pass this file exists to remove: with TIMEOUT dropped + # to 1 s every OPTDEP subprocess timed out and the run still reported + # "OPTDEP 8 checked ... PASS". + errs.append("%s:%d: %s is OPTDEP and its collection TIMED OUT after %ds.\n" + " Nothing was verified; a timeout is not a pass. Raise " + "TIMEOUT if the file is legitimately slow, or fix the hang." + % (ROSTER, lineno, path, TIMEOUT)) + elif n > 0: + rc2, _, passed = _pytest(path, collect_only=False) + if rc2 is None: + errs.append("%s:%d: %s is OPTDEP and its RUN timed out after %ds with " + "%s missing.\n Nothing was verified; a timeout is not a " + "pass." + % (ROSTER, lineno, path, TIMEOUT, ",".join(missing))) + elif rc2 == 0 and passed == n: + errs.append("%s:%d: %s is OPTDEP on missing %s, yet collects %d tests and " + "ALL PASS.\n It does not actually need what it claims; gate " + "it, or correct the reason." + % (ROSTER, lineno, path, ",".join(missing), n)) + else: + continue + checked[status] = checked.get(status, 0) + 1 + + print("test-roster-verify: predicates checked per status") + for s in sorted(checked): + print(" %-10s %3d" % (s, checked[s])) + if errs: + print("\ntest-roster-verify: FAIL", file=sys.stderr) + for e in errs: + print(" " + e, file=sys.stderr) + return 1 + print("test-roster-verify: PASS -- every checkable roster reason still holds.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/.travis/test-run-alts.sh b/.travis/test-run-alts.sh index 3ce77a9b8..72587bc35 100755 --- a/.travis/test-run-alts.sh +++ b/.travis/test-run-alts.sh @@ -1,19 +1,24 @@ #! /bin/bash -# GETTING_STARTED.md example +# GETTING_STARTED.md example. Same NoLoop lane as test-run.sh +# (--time-marginalization --vectorized --gpu, --force-xpy for the CPU +# fallback), plus --resample-time-marginalization --fairdraw-extrinsic-output +# on top -- see test-run.sh for why the flags are explicit and what the +# banner assertion below is checking. No separate legacy-scalar lane here; +# test-run.sh already carries that one. -if [ ! -d ILE-GPU-Paper ]; then +if [ ! -d ILE-GPU-Paper ]; then git clone https://github.com/oshaughn/ILE-GPU-Paper.git -fi +fi cd ILE-GPU-Paper/demos/ if [ -d test_workflow_batch_gpu_lowlatency ]; then echo " Deleting test directory !" rm -rf test_workflow_batch_gpu_lowlatency -fi +fi make test_workflow_batch_gpu_lowlatency cd test_workflow_batch_gpu_lowlatency -# force standard code path -switcheroo '--maximize-only ' ' --force-xpy ' command-single.sh +# Exercise the maintained NoLoop path explicitly (see test-run.sh). +switcheroo '--maximize-only ' ' --vectorized --gpu --force-xpy ' command-single.sh # Reduce number of analyses for this worker to 1, to reduce runtime switcheroo '\$\(macrongroup\)' 1 command-single.sh alias macrongroup='echo 1' @@ -22,5 +27,16 @@ echo 'echo 1' > macrongroup; chmod a+x macrongroup; PATH=${PATH}:`pwd`# # ... and save-samples switcheroo 'n-max 2000000' 'n-max 50000 --save-samples --output-file my_stuff ' command-single.sh switcheroo '--save-samples ' '--save-samples --resample-time-marginalization --fairdraw-extrinsic-output ' command-single.sh -./command-single.sh +./command-single.sh 2>&1 | tee run_command_single.log +status=${PIPESTATUS[0]} +if [ "$status" -ne 0 ]; then + echo "command-single.sh FAILED (exit $status)" + exit "$status" +fi +if ! grep -q 'Q_lm sub-sample time stencil.*vectorized=True gpu=True' run_command_single.log; then + echo "the run did NOT take the NoLoop path -- vectorized=True gpu=True not in the startup banner:" + grep 'Q_lm sub-sample time stencil' run_command_single.log + exit 1 +fi +echo "NoLoop path confirmed (vectorized=True gpu=True in run_command_single.log)" diff --git a/.travis/test-run.sh b/.travis/test-run.sh index 1d389893b..e62b2ce30 100755 --- a/.travis/test-run.sh +++ b/.travis/test-run.sh @@ -1,6 +1,27 @@ #! /bin/bash -# GETTING_STARTED.md example +# GETTING_STARTED.md example. Two lanes, from one build: +# +# 1. NoLoop (production path): --time-marginalization --vectorized --gpu, +# the exact option set helper_LDG_Events.py emits for every real run +# (--propose-ile-convergence-options). --force-xpy takes the identical +# NoLoop code (DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop) on +# plain numpy when no cupy device is present, instead of silently +# downgrading --gpu to a no-op. This is asserted below by grepping the +# run's own startup banner ("Q_lm sub-sample time stencil ... [... ]") +# for vectorized=True gpu=True, rather than trusting the flags alone: +# command-single.sh already carries an incidental --vectorized --gpu +# from create_event_parameter_pipeline_BasicIteration's --request-gpu-ILE +# handling, so passing the flags here is about making this script +# correct on its own (and future-proof if that incidental append ever +# goes away), and the banner check is what actually proves the path ran. +# 2. legacy-scalar (sanity only, NOT what production runs): the same build, +# with --vectorized/--gpu/--force-xpy stripped back out and n-max/n-chunk +# cut so it costs seconds. FactoredLogLikelihoodTimeMarginalized already +# has unit coverage (test_ile_scalar_edge_cases.py, +# factored_likelihood_test.py); this lane exists because unit tests miss +# the driver/CLI seam, and only checks that the legacy branch still runs +# at all, not that its answer is right. #if command -v apt ; then #apt install lalsuite=7.22 # problem with 7.23 and newer lalapps_path2cache, workaround @@ -12,13 +33,55 @@ git clone https://github.com/oshaughn/ILE-GPU-Paper.git cd ILE-GPU-Paper/demos/ make test_workflow_batch_gpu_lowlatency cd test_workflow_batch_gpu_lowlatency -# force standard code path -switcheroo '--maximize-only ' ' --force-xpy ' command-single.sh -# Reduce number of analyses for this worker to 1, to reduce runtime +# Exercise the maintained NoLoop path explicitly (see banner above). +switcheroo '--maximize-only ' ' --vectorized --gpu --force-xpy ' command-single.sh +# Reduce number of analyses for this worker to 1, to reduce runtime. +# create_event_parameter_pipeline_BasicIteration replaces $(macrongroup) with the literal +# '5' at command-single.sh GENERATION time (bin/create_event_parameter_pipeline_ +# BasicIteration, arg_list.replace('$(macrongroup)','5')), before this script ever runs -- +# so the line below targeted a token this file no longer carries, a silent no-op (PR #283 +# review, MINOR). Target the literal value actually written instead. +switcheroo '--n-events-to-analyze 5 ' '--n-events-to-analyze 1 ' command-single.sh +# Backstop for an older create_event_parameter_pipeline_BasicIteration that still leaves +# $(macrongroup) as a literal token: bash would then treat it as command substitution +# (running a command named macrongroup) when command-single.sh executes, so make that +# resolve to 1 too rather than failing with "macrongroup: command not found". switcheroo '\$\(macrongroup\)' 1 command-single.sh -# new format for n-events-to-analyze. Backstop alias macrongroup='echo 1' echo 'echo 1' > macrongroup; chmod a+x macrongroup; PATH=${PATH}:`pwd` # Reduce the number of points investigated by x100 switcheroo 'n-max 2000000' 'n-max 50000' command-single.sh -./command-single.sh + +./command-single.sh 2>&1 | tee run_command_single_noloop.log +status=${PIPESTATUS[0]} +if [ "$status" -ne 0 ]; then + echo "NoLoop lane: command-single.sh FAILED (exit $status)" + exit "$status" +fi +if ! grep -q 'Q_lm sub-sample time stencil.*vectorized=True gpu=True' run_command_single_noloop.log; then + echo "NoLoop lane: the run did NOT take the NoLoop path -- vectorized=True gpu=True not in the startup banner:" + grep 'Q_lm sub-sample time stencil' run_command_single_noloop.log + exit 1 +fi +echo "NoLoop lane: confirmed (vectorized=True gpu=True in run_command_single_noloop.log)" + +# --- legacy-scalar sanity lane (cheap; NOT the production path) ---------- +cp command-single.sh command-single-legacy-scalar.sh +switcheroo ' --vectorized' '' command-single-legacy-scalar.sh +switcheroo ' --gpu' '' command-single-legacy-scalar.sh +switcheroo ' --force-xpy' '' command-single-legacy-scalar.sh +switcheroo 'n-max 50000' 'n-max 500' command-single-legacy-scalar.sh +switcheroo 'n-chunk 10000' 'n-chunk 500' command-single-legacy-scalar.sh + +./command-single-legacy-scalar.sh 2>&1 | tee run_command_single_legacy_scalar.log +status=${PIPESTATUS[0]} +if [ "$status" -ne 0 ]; then + echo "legacy-scalar lane: command-single-legacy-scalar.sh FAILED (exit $status)" + exit "$status" +fi +if ! grep -q 'Q_lm sub-sample time stencil.*vectorized=False gpu=False' run_command_single_legacy_scalar.log; then + echo "legacy-scalar lane: the run did NOT take the legacy scalar path -- vectorized=False gpu=False not in the startup banner:" + grep 'Q_lm sub-sample time stencil' run_command_single_legacy_scalar.log + exit 1 +fi +echo "legacy-scalar lane: confirmed (vectorized=False gpu=False in run_command_single_legacy_scalar.log)" diff --git a/.travis/test-slowrot.sh b/.travis/test-slowrot.sh index b960c6b70..a028da8d3 100755 --- a/.travis/test-slowrot.sh +++ b/.travis/test-slowrot.sh @@ -75,6 +75,7 @@ FILES=( "${SLOWDIR}/test_slowrot_pathB.py" "${SLOWDIR}/test_slowrot_precompute_integration.py" "${SLOWDIR}/test_slowrot_response.py" + "${SLOWDIR}/test_slowrot_rotating_freqresponse.py" ) # DESELECTED, and EXPECTED_TESTS is one lower because of it. @@ -169,7 +170,7 @@ fi # Re-derive with `pytest --collect-only -q` over FILES; never lower it without saying why # in the commit message. A bare `pytest ${SLOWDIR}` would sweep up files that collect 0, # and a partial loss still exits 0, which is what this pins against. -EXPECTED_TESTS=43 +EXPECTED_TESTS=50 DESELECT_ARGS=() for d in "${DESELECT[@]}"; do DESELECT_ARGS+=(--deselect "${d}"); done diff --git a/CHANGES.rst b/CHANGES.rst index ab383a7b6..316eebe29 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -16,6 +16,14 @@ development tree is rift_O4d; PRs refer to oshaughn/research-projects-RIT. normalization, distance tails and waveform handling; strengthen calibration and CI checks (PRs #159, #173, #178, #188). JAX flow reuse is now off by default (opt-in via --flow-reuse); JAX --save-samples now exports fair posterior draws. + - (rc4 pending) JAX ILE gains an opt-in, fail-closed four-axis direct-marginalization policy, + peak-local time/distance/angle planning with exact or band-limited reserve rules, and + device-aware batching/pregridding. Add value-only JAX adaptive-volume and portfolio sampling, + phase-rotated coordinates, persistent compilation caching, and Asimov selection. Correct + importance-proposal regularization and refuse non-finite or implausible evidence before result + publication. Extend slow-rotation/finite-size response support and cross-term batching; expand + CPU/JAX regression and CI-roster coverage (fork PRs #214, #245, #247, #255, #268, #270, + #274, #280--#285, #294, #301--#315, #319). 0.0.17.12 --------- diff --git a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini index ccc3a225a..c5a6088fa 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini +++ b/MonteCarloMarginalizeCode/Code/RIFT/asimov/rift.ini @@ -202,6 +202,23 @@ internal-cip-use-lnL=True {% if sampler contains 'ile' %} {% if sampler['ile'] contains 'rotate phase' %} internal-ile-rotate-phase={{ sampler['ile']['rotate phase'] }} {% endif %} {% endif %} +# +# Time-stencil / time-quadrature ledger keys (RO'S directive 2026-09-08): these are +# complicated enough that people should not need --manual-extra-ile-args for them. All +# three default to ABSENT -- omitted below unless the ledger sets them -- so an existing +# production is byte-identical unless one of these keys is added. See +# bin/helper_LDG_Events.py / bin/util_RIFT_pseudo_pipe.py --help for the corresponding +# pipeline options and RIFT/likelihood/{time_interp_choice,time_marginalization_quadrature, +# q_time_pregrid}.py for the validation each one goes through. +{% if sampler contains 'ile' %} {% if sampler['ile'] contains 'interpolate time' %} +internal-ile-interpolate-time='{{ sampler['ile']['interpolate time'] }}' +{% endif %} {% endif %} +{% if sampler contains 'ile' %} {% if sampler['ile'] contains 'time marginalization quadrature' %} +internal-ile-time-marginalization-quadrature='{{ sampler['ile']['time marginalization quadrature'] }}' +{% endif %} {% endif %} +{% if sampler contains 'ile' %} {% if sampler['ile'] contains 'q time pregrid factor' %} +internal-ile-q-time-pregrid-factor={{ sampler['ile']['q time pregrid factor'] }} +{% endif %} {% endif %} # # Assume settings @@ -248,6 +265,11 @@ ile-n-eff= {{ sampler['ile']['n eff'] | default: 10 }} ile-copies = {{ sampler['ile']['copies'] | default: 1}} ile-sampler-method='{{ sampler['ile']['sampling method'] | default: "AV" }}' internal-ile-freezeadapt={{ sampler['ile']['freezeadapt'] | default: False }} +# Select `which integrate_likelihood_extrinsic_jax` instead of the default batchmode +# ILE driver. See util_RIFT_pseudo_pipe.py --help (--use-jax-ile) for what that +# driver does not implement; it is REFUSED at DAG-build time together with in-loop +# calibration marginalization. +use-jax-ile={{ sampler['ile']['use jax ile'] | default: False }} # {%- if sampler['ile'] contains "manual extra args" %} # manual-extra-ile-args="{% for arg in sampler['ile']['manual extra args'] %} {{ arg }} {% endfor %}" # {%- endif %} diff --git a/MonteCarloMarginalizeCode/Code/RIFT/hyperpipe/marg_list.py b/MonteCarloMarginalizeCode/Code/RIFT/hyperpipe/marg_list.py index 8ee27d864..d54a4ddae 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/hyperpipe/marg_list.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/hyperpipe/marg_list.py @@ -155,19 +155,29 @@ def _stage_event_file( base_dir: str, run_dir: str, ) -> Tuple[str, bool]: - """Materialize this entry's event file at base_dir/event-.net. + """Materialize this entry's event file at run_dir/event-.net. Returns ``(abs_path, is_empty_sentinel)``. If the entry has no ``event-file`` set, we write a sentinel file with the single token ``empty_event_file`` so the downstream pipeline still sees a well-formed input. + + Sources resolve against ``base_dir`` (where the user's config paths are + relative to); the staged copy is written to ``run_dir``. That split is + what :func:`assemble_marg_list` documents, and what the exe staging a + few lines below already does. The destination used to be ``base_dir``, + with ``run_dir`` accepted and unused: under hydra those are different + directories -- ``base_dir`` is the ORIGINAL cwd the user launched from, + ``run_dir`` the per-run output dir -- so the staged files landed in the + launch directory, and two runs started from one directory overwrote each + other's ``event-.net``. """ src = None if hasattr(entry, "get"): src = entry.get("event-file") or entry.get("event_file") elif "event-file" in entry: src = entry["event-file"] - dest = os.path.join(base_dir, f"event-{indx}.net") + dest = os.path.join(run_dir, f"event-{indx}.net") if src: src = os.path.expanduser(src) if not os.path.isabs(src): diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py index a1926465f..5c84bde78 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py @@ -412,10 +412,25 @@ def cdf_inverse(self, param): Numerically determine the inverse CDF from a given sampling PDF. If the PDF itself is not normalized, the class will keep an internal record of the normalization and adjust the PDF values as necessary. Returns a function object which is the interpolated CDF inverse. """ # Solve P'(x) == p(x), with P[lower_boun] == 0 + pdf = self.pdf[param] + # odeint probes with a python float. Scalar-style pdfs (uniform_samp, + # numpy.vectorize, uniform_samp_withfloor_vector) take it as-is; this + # module's vectorized helpers (ones(len(x)), xpy.sin(x)) need a length-1 + # backend array. Decide once by probing, then reduce whatever comes back + # (float, numpy scalar, 0-d or length-1 array on either backend) to a float. + try: + pdf(float(0.5*(self.llim[param]+self.rlim[param]))) + as_array = False + except TypeError: + as_array = True + def _scalar(val): + return float(val.ravel()[0]) if hasattr(val, 'ravel') else float(val) def dP_cdf(p, x): if x > self.rlim[param] or x < self.llim[param]: return 0 - return self.pdf[param](x) + if as_array: + return _scalar(pdf(xpy_default.asarray([x], dtype=numpy.float64))) + return _scalar(pdf(x)) x_i = numpy.linspace(self.llim[param], self.rlim[param], 1000) # Integrator needs to have a step size which doesn't step over the # probability mass diff --git a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/__init__.py b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/__init__.py index b8287a9af..21c5e5172 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/__init__.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/__init__.py @@ -23,15 +23,43 @@ """ from __future__ import annotations -import jax as _jax +# The docstring above calls this subpackage OPTIONAL, but importing it used to require jax +# unconditionally -- so merely TOUCHING the package died with ModuleNotFoundError on a machine +# without the stack. Pytest touches it: collecting any test module in this directory imports +# this __init__ first, which is why a skip guard inside test_interpolators.py could never fire. +# +# Only the ABSENCE is tolerated here. When jax is present the sequence below is unchanged and +# stays EAGER on purpose: three callers import this package for no reason but its side effect +# (applications/compare.py, applications/jax_cip.py, applications/export_at_scale.py all say +# "enables float64"), and x64 must be set before any submodule builds a jax array. Deferring it +# into get_interpolator() would leave those three silently in float32, which is a wrong-gradient +# bug that raises nothing. +try: + import jax as _jax +except ImportError: # pragma: no cover - exercised only where the jax stack is absent + _jax = None +else: + if not _jax.config.read("jax_enable_x64"): + _jax.config.update("jax_enable_x64", True) -if not _jax.config.read("jax_enable_x64"): - _jax.config.update("jax_enable_x64", True) - -from .interface import BaseInterpolator # noqa: E402 + from .interface import BaseInterpolator # noqa: E402 __all__ = ["BaseInterpolator"] + +def __getattr__(name): + """Re-raise the real ImportError for the eager exports when jax is missing. + + Without this the jax-absent case reports a bare AttributeError, which reads like a typo + rather than a missing dependency. Unknown names still raise AttributeError, so + ``from RIFT.interpolators.jax_gp import export`` (a SUBMODULE) keeps working -- the import + machinery falls back to importing the submodule when this returns AttributeError. + """ + if name in __all__: + from . import interface # raises ModuleNotFoundError naming the missing package + return getattr(interface, name) + raise AttributeError("module {!r} has no attribute {!r}".format(__name__, name)) + # Method classes are imported lazily by name to avoid importing every backend # (and its heavier deps, e.g. tinygp) when only one is needed. def get_interpolator(name): diff --git a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py index 78d6e67eb..e178ab637 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_coordinates.py @@ -7,8 +7,26 @@ """ from __future__ import annotations +import sys as _sys + import numpy as np -import jax + +# jax is not in requirements.txt, so SKIP rather than let the ImportError escape at +# collection: an import error there reports as a FAILING suite, and a suite that fails for +# environmental reasons is one people learn to ignore. +# +# Skip only when actually running UNDER pytest. `import pytest` succeeding is not that test +# -- pytest is installed nearly everywhere -- and using it as one makes the direct +# `python -m ...` run this file's docstring advertises die with a pytest `Skipped` exception +# instead of the real ImportError. Under pytest, pytest is already in sys.modules by the +# time it imports this module; under a direct run it is not. +try: + import jax +except ImportError as _exc: # pragma: no cover - environment probe + _pytest = _sys.modules.get("pytest") + if _pytest is None: + raise + _pytest.skip("jax_gp coordinates need jax: %s" % _exc, allow_module_level=True) from . import coordinates as C diff --git a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py index 3e9d00de2..a0f8c8076 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/interpolators/jax_gp/test_interpolators.py @@ -15,8 +15,28 @@ import os +import sys as _sys + import numpy as np +# The jax stack (jax for the models, optax for their optimisers) is not in requirements.txt, so SKIP rather than let the ImportError escape at +# collection: an import error there reports as a FAILING suite, and a suite that fails for +# environmental reasons is one people learn to ignore. +# +# Skip only when actually running UNDER pytest. `import pytest` succeeding is not that test +# -- pytest is installed nearly everywhere -- and using it as one makes the direct +# `python -m ...` run this file's docstring advertises die with a pytest `Skipped` exception +# instead of the real ImportError. Under pytest, pytest is already in sys.modules by the +# time it imports this module; under a direct run it is not. +try: + import jax # noqa: F401 + import optax # noqa: F401 +except ImportError as _exc: # pragma: no cover - environment probe + _pytest = _sys.modules.get("pytest") + if _pytest is None: + raise + _pytest.skip("jax_gp interpolators need jax and optax: %s" % _exc, allow_module_level=True) + def _target(X): # smooth, anisotropic quadratic bowl -- exactly representable-ish, known grad diff --git a/MonteCarloMarginalizeCode/Code/RIFT/jax_cache.py b/MonteCarloMarginalizeCode/Code/RIFT/jax_cache.py new file mode 100644 index 000000000..4d34a491f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/jax_cache.py @@ -0,0 +1,383 @@ +"""Persistent, transferable JAX compilation-cache support for RIFT ILE. + +JAX includes compiler options and argument shapes in its cache keys. RIFT adds +an outer compatibility namespace so cache bundles are never mixed across the +JAX/JAXLIB/backend/device combinations that matter on heterogeneous GPU pools. +""" + +from __future__ import annotations + +import hashlib +import importlib.metadata +import json +import os +import platform +import shutil +import sys +import tempfile +import zipfile +from pathlib import Path, PurePosixPath + + +MANIFEST_NAME = "rift-jax-cache-manifest.json" +IMPORT_MANIFEST_NAME = "rift-jax-cache-import.json" +IMPORT_MANIFEST_PREFIX = "rift-jax-cache-import-" +FORMAT_VERSION = 1 +MAX_BUNDLE_FILES = 100_000 +MAX_BUNDLE_MEMBER_BYTES = 4 * 1024**3 +MAX_BUNDLE_TOTAL_BYTES = 16 * 1024**3 +MAX_BUNDLE_COMPRESSION_RATIO = 10_000 +MAX_MANIFEST_BYTES = 16 * 1024**2 +_ACCELERATOR_PLUGIN_PACKAGES = ( + "jax-cuda13-plugin", "jax-cuda12-plugin", "jax-cuda11-plugin", + "jax-rocm7-plugin", "jax-rocm60-plugin", "jax-rocm-plugin", "jax-metal", +) + + +def _package_version(name): + try: + return importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + return None + + +def runtime_compatibility(jax_module=None): + """Return the conservative runtime identity used for cache transfer.""" + if jax_module is None: + import jax as jax_module + backend = jax_module.default_backend() + devices = list(jax_module.devices(backend)) + device = devices[0] if devices else None + client = getattr(device, "client", None) + capability = getattr(device, "compute_capability", None) + if callable(capability): + capability = capability() + if isinstance(capability, (tuple, list)): + capability = ".".join(str(part) for part in capability) + accelerator_plugins = { + name: version for name in _ACCELERATOR_PLUGIN_PACKAGES + if (version := _package_version(name)) is not None + } + return { + "python": platform.python_version(), + "jax": getattr(jax_module, "__version__", _package_version("jax")), + "jaxlib": _package_version("jaxlib"), + "accelerator_plugins": accelerator_plugins, + "backend": backend, + "platform_version": str(getattr(client, "platform_version", None)), + "device_kind": str(getattr(device, "device_kind", None)), + "compute_capability": capability, + } + + +def compatibility_key(compatibility): + raw = json.dumps(compatibility, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:20] + + +def _argv_cache_controls(argv): + root = None + disabled = False + argv = list(argv or ()) + for i, token in enumerate(argv): + if token == "--no-jax-persistent-cache": + disabled = True + elif token.startswith("--jax-cache-dir="): + root = token.split("=", 1)[1] + elif token == "--jax-cache-dir" and i + 1 < len(argv): + root = argv[i + 1] + return root, disabled + + +def argv_option(argv, name): + """Return the last ``--name value``/``--name=value`` occurrence.""" + value = None + argv = list(argv or ()) + for i, token in enumerate(argv): + if token.startswith(name + "="): + value = token.split("=", 1)[1] + elif token == name and i + 1 < len(argv): + value = argv[i + 1] + return value + + +def default_cache_root(): + explicit = os.environ.get("RIFT_JAX_CACHE_ROOT") + if explicit: + return Path(explicit).expanduser() + scratch = os.environ.get("_CONDOR_SCRATCH_DIR") + if scratch: + return Path(scratch) / ".rift_cache" / "jax" + xdg = os.environ.get("XDG_CACHE_HOME") + base = Path(xdg).expanduser() if xdg else Path.home() / ".cache" + return base / "rift" / "jax" + + +def _write_manifest(directory, compatibility, extra=None): + manifest = { + "format_version": FORMAT_VERSION, + "compatibility": compatibility, + "compatibility_key": compatibility_key(compatibility), + } + if extra: + manifest.update(extra) + return _write_json_atomic(directory / MANIFEST_NAME, manifest) + + +def _write_import_manifest(directory, compatibility, bundle_manifest): + """Persist one immutable record per contributing bundle manifest.""" + manifest_raw = json.dumps( + bundle_manifest, sort_keys=True, separators=(",", ":")).encode("utf-8") + manifest_digest = hashlib.sha256(manifest_raw).hexdigest() + record = { + "format_version": FORMAT_VERSION, + "compatibility": compatibility, + "compatibility_key": compatibility_key(compatibility), + "bundle_manifest_sha256": manifest_digest, + "imported_profile": bundle_manifest.get("profile"), + "static_shapes": bundle_manifest.get("static_shapes", {}), + } + target = directory / (IMPORT_MANIFEST_PREFIX + manifest_digest + ".json") + return _write_json_atomic(target, record) + + +def _is_provenance_file(path): + """Return whether *path* is runtime/import metadata, not compiler data.""" + return (path.name in (MANIFEST_NAME, IMPORT_MANIFEST_NAME) + or (path.name.startswith(IMPORT_MANIFEST_PREFIX) + and path.name.endswith(".json"))) + + +def _write_json_atomic(target, value): + target.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=".%s." % target.name, suffix=".tmp", dir=str(target.parent)) + try: + with os.fdopen(fd, "w", encoding="utf-8") as stream: + stream.write(json.dumps(value, indent=2, sort_keys=True) + "\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, target) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + return value + + +def _publish_file_atomic(source, target): + """Copy one validated entry without exposing a partial target to readers.""" + target.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp( + prefix=".%s." % target.name, suffix=".tmp", dir=str(target.parent)) + os.close(fd) + try: + shutil.copy2(source, temporary_name) + sync_fd = os.open(temporary_name, os.O_RDONLY) + try: + os.fsync(sync_fd) + finally: + os.close(sync_fd) + os.replace(temporary_name, target) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + +def configure_persistent_cache(jax_module, argv=None): + """Enable RIFT's cache before any ILE JIT is constructed. + + ``--jax-cache-dir`` and ``RIFT_JAX_CACHE_ROOT`` name a cache *root*. A + compatibility-keyed child is selected automatically. The standard + ``JAX_COMPILATION_CACHE_DIR`` remains supported as an exact expert override. + """ + cli_root, disabled = _argv_cache_controls(argv) + if disabled or os.environ.get("RIFT_DISABLE_JAX_CACHE") == "1": + jax_module.config.update("jax_enable_compilation_cache", False) + return None + + # runtime_compatibility() PROBES THE DEVICE, and that probe can fail on a + # perfectly ordinary execute node: jax.default_backend()/jax.devices() force + # backend init, so unloadable CUDA libraries or a card that is busy for every + # tenant raise here. This function runs at driver IMPORT, before the option + # parser exists, so an escaping exception turns "no usable accelerator" into + # a driver that cannot even print --help -- a performance optimization + # failing a scientific run, which is the thing the OSError handler below + # exists to prevent. Catch it in the same place and for the same reason. + try: + compatibility = runtime_compatibility(jax_module) + except Exception as exc: + print("WARNING: disabling JAX persistent cache, cannot identify the " + "JAX runtime/device: %s" % exc, file=sys.stderr) + try: + jax_module.config.update("jax_enable_compilation_cache", False) + except Exception: + pass + return None + exact = os.environ.get("JAX_COMPILATION_CACHE_DIR") + if exact and not cli_root: + directory = Path(exact).expanduser() + else: + root = Path(cli_root).expanduser() if cli_root else default_cache_root() + directory = root / compatibility_key(compatibility) + try: + directory.mkdir(parents=True, exist_ok=True) + _write_manifest(directory, compatibility) + except OSError as exc: + # A read-only/missing home must not turn a performance optimization into + # a failed scientific run. Condor normally avoids this via its scratch + # fallback above; unusual sites can still opt in explicitly. + print("WARNING: disabling JAX persistent cache: %s" % exc, file=sys.stderr) + jax_module.config.update("jax_enable_compilation_cache", False) + return None + os.environ["JAX_COMPILATION_CACHE_DIR"] = str(directory.resolve()) + # JAX 0.9.2's auxiliary per-fusion GPU autotune cache embeds its absolute + # directory in CompileOptions, but does not exclude that field from the + # persistent-executable cache key. A bundle imported under a different + # absolute root would therefore miss every executable it contains. Keep + # the portable executable cache enabled, but disable only that auxiliary + # path-valued cache. Deliberately let config.update fail loudly if a future + # JAX stops accepting a setting it advertises: silently restoring + # non-portable keys would make a successful-looking transfer useless. + # Older supported JAX (for example 0.4.24) predates this auxiliary cache and + # does not advertise the option, so there is nothing path-valued to disable. + if hasattr(jax_module.config, "jax_persistent_cache_enable_xla_caches"): + jax_module.config.update("jax_persistent_cache_enable_xla_caches", "") + jax_module.config.update("jax_enable_compilation_cache", True) + jax_module.config.update("jax_compilation_cache_dir", str(directory.resolve())) + return directory.resolve() + + +def _file_hash(path): + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def export_bundle(cache_dir, output, compatibility, profile=None, static_shapes=None): + """Create a self-describing zip bundle from an already-warmed cache.""" + cache_dir = Path(cache_dir) + output = Path(output) + if not cache_dir.is_dir(): + raise ValueError("cache directory does not exist: %s" % cache_dir) + files = {} + total_size = 0 + for path in sorted(cache_dir.rglob("*")): + is_manifest_temp = (path.name.startswith(".rift-jax-cache-") + and path.name.endswith(".tmp")) + is_provenance = _is_provenance_file(path) + if path.is_file() and not is_provenance and not is_manifest_temp: + rel = path.relative_to(cache_dir).as_posix() + size = path.stat().st_size + if size > MAX_BUNDLE_MEMBER_BYTES: + raise ValueError("cache member exceeds the bundle size limit: %s" % rel) + total_size += size + if total_size > MAX_BUNDLE_TOTAL_BYTES: + raise ValueError("cache exceeds the total bundle size limit") + files[rel] = _file_hash(path) + if len(files) > MAX_BUNDLE_FILES: + raise ValueError("cache has too many files to bundle safely") + manifest = { + "format_version": FORMAT_VERSION, + "compatibility": compatibility, + "compatibility_key": compatibility_key(compatibility), + "profile": profile, + "static_shapes": static_shapes or {}, + "files": files, + } + output.parent.mkdir(parents=True, exist_ok=True) + with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(MANIFEST_NAME, json.dumps(manifest, indent=2, sort_keys=True) + "\n") + for rel in files: + archive.write(cache_dir / rel, "cache/" + rel) + return manifest + + +def _validate_member(info, *, manifest=False): + limit = MAX_MANIFEST_BYTES if manifest else MAX_BUNDLE_MEMBER_BYTES + if info.file_size > limit: + raise ValueError("cache bundle member exceeds the size limit: %s" % info.filename) + if info.file_size and not info.compress_size: + raise ValueError("cache bundle member has an invalid compressed size: %s" % info.filename) + if info.compress_size and info.file_size / info.compress_size > MAX_BUNDLE_COMPRESSION_RATIO: + raise ValueError("cache bundle member exceeds the compression-ratio limit: %s" % info.filename) + + +def _read_limited(archive, info, limit): + with archive.open(info, "r") as stream: + data = stream.read(limit + 1) + if len(data) > limit: + raise ValueError("cache bundle member exceeds the size limit: %s" % info.filename) + return data + + +def import_bundle(bundle, cache_root, compatibility, expected_profile=None, + destination=None): + """Validate and merge a bundle into this runtime's cache namespace.""" + bundle = Path(bundle) + with zipfile.ZipFile(bundle, "r") as archive: + infos = archive.infolist() + names = [info.filename for info in infos] + if len(names) != len(set(names)): + raise ValueError("cache bundle contains duplicate archive members") + if len(names) > MAX_BUNDLE_FILES + 1: + raise ValueError("cache bundle contains too many archive members") + if MANIFEST_NAME not in names: + raise ValueError("bundle has no %s" % MANIFEST_NAME) + info_by_name = {info.filename: info for info in infos} + _validate_member(info_by_name[MANIFEST_NAME], manifest=True) + manifest = json.loads(_read_limited( + archive, info_by_name[MANIFEST_NAME], MAX_MANIFEST_BYTES)) + if manifest.get("format_version") != FORMAT_VERSION: + raise ValueError("unsupported cache bundle format") + if manifest.get("compatibility") != compatibility: + raise ValueError("cache bundle is incompatible with this JAX runtime/device") + if expected_profile is not None and manifest.get("profile") != expected_profile: + raise ValueError("cache bundle warmup profile does not match --expect-profile") + declared = manifest.get("files", {}) + expected_names = {"cache/" + rel for rel in declared} + actual_names = {name for name in names if name.startswith("cache/") and not name.endswith("/")} + if actual_names != expected_names: + raise ValueError("cache bundle contents do not match its manifest") + if set(names) != expected_names | {MANIFEST_NAME}: + raise ValueError("cache bundle contains unexpected archive members") + total_size = 0 + for name in expected_names: + info = info_by_name[name] + _validate_member(info) + total_size += info.file_size + if total_size > MAX_BUNDLE_TOTAL_BYTES: + raise ValueError("cache bundle exceeds the total size limit") + with tempfile.TemporaryDirectory(prefix="rift-jax-cache-") as temp: + temp_root = Path(temp) + for rel, expected_hash in declared.items(): + pure = PurePosixPath(rel) + if pure.is_absolute() or ".." in pure.parts: + raise ValueError("unsafe cache bundle path: %s" % rel) + target = temp_root / rel + target.parent.mkdir(parents=True, exist_ok=True) + digest = hashlib.sha256() + written = 0 + with archive.open(info_by_name["cache/" + rel], "r") as source, target.open("wb") as output: + for block in iter(lambda: source.read(1024 * 1024), b""): + written += len(block) + if written > MAX_BUNDLE_MEMBER_BYTES: + raise ValueError("cache bundle member exceeds the size limit: %s" % rel) + digest.update(block) + output.write(block) + if digest.hexdigest() != expected_hash: + raise ValueError("cache bundle checksum mismatch: %s" % rel) + destination = (Path(destination).expanduser() if destination is not None + else Path(cache_root).expanduser() / compatibility_key(compatibility)) + destination.mkdir(parents=True, exist_ok=True) + for source in temp_root.rglob("*"): + if source.is_file(): + target = destination / source.relative_to(temp_root) + _publish_file_atomic(source, target) + _write_import_manifest(destination, compatibility, manifest) + return destination diff --git a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py index d90fc5a0c..0c2f48111 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py @@ -2340,18 +2340,20 @@ def __init__(self, fLow=10., fMax=None, fNyq=2048., deltaF=1./8., else: # if we get here psd must be an array fPSD = (len(psd) - 1) * self.deltaF # -1 b/c start at f=0 assert self.fMax <= fPSD - # ivals = np.arange(self.minIdx,self.maxIdx) - # ivals_ok = psd[ivals]>0 - # extra_weight=np.ones(len(self.weights)) - # if waveform_is_psi4: - # extra_weight[ivals_ok] = 1./(2*np.pi*ivals[ivals_ok]*deltaF)**2 - # self.weights[ivals_ok] = 1./psd[ivals_ok] * extra_weight[ivals_ok] - for i in range(self.minIdx,self.maxIdx): - if psd[i] != 0.: - extra_weight=1.0 - if waveform_is_psi4: - extra_weight=1.0/(2*np.pi*i*deltaF)/(2*np.pi*i*deltaF) - self.weights[i] = 1./psd[i]*extra_weight + # Vectorised form of the per-bin loop this replaces: 22.2 ms -> 1.9 ms at + # 128512 in-band bins (seglen 128 s, srate 8192, ldas-grid). The driver + # passes a REAL8FrequencySeries, so this branch is off that path. + # Bit-identical: the mask stays `!= 0` as in the loop, since the `> 0` used + # by the REAL8FrequencySeries branch above would drop negative bins here, + # and the psi4 weight keeps the `1/x/x` association rather than `1/x**2`. + ivals = np.arange(self.minIdx, self.maxIdx) + ivals_ok = ivals[psd[self.minIdx:self.maxIdx] != 0.] + if waveform_is_psi4: + _x = 2*np.pi*ivals_ok*deltaF + extra_weight = 1.0/_x/_x + else: + extra_weight = 1.0 + self.weights[ivals_ok] = 1./psd[ivals_ok]*extra_weight else: raise ValueError("analyticPSD_Q must be either True or False") @@ -2372,8 +2374,11 @@ def __init__(self, fLow=10., fMax=None, fNyq=2048., deltaF=1./8., WFD.data.data[:] = np.sqrt(self.weights) # W_FD is 1/sqrt(S_n(f)) WFD.data.data[0] = WFD.data.data[-1] = 0. # zero 0, f_Nyq bins lal.REAL8FreqTimeFFT(WTD, WFD, revplan) # IFFT to TD - for i in range(int(N_spec/2), self.len2side - int(N_spec/2)): - WTD.data.data[i] = 0. # Zero all but T_spec/2 ends of W_TD + # Zero all but T_spec/2 ends of W_TD. Slice assignment, bit-identical to the + # per-element loop it replaces -- which was ~1e6 SWIG element writes at seglen + # 128 s / srate 8192, ~0.3-0.7 s, and inverse spectrum truncation is ON by + # default (--inv-spec-trunc-time 8), so it was paid once per ComputeModeCrossTermIP. + WTD.data.data[int(N_spec/2) : self.len2side - int(N_spec/2)] = 0. lal.REAL8TimeFreqFFT(WFD, WTD, fwdplan) # FFT back to FD WFD.data.data[0] = WFD.data.data[-1] = 0. # zero 0, f_Nyq bins # Square to get trunc. inv. PSD @@ -2386,6 +2391,17 @@ def __init__(self, fLow=10., fMax=None, fNyq=2048., deltaF=1./8., self.weights2side[:len(self.weights)] = self.weights[::-1] self.weights2side[len(self.weights)-1:] = self.weights[0:-1] + # Contiguous support of the band weights. Everything outside [band_lo, band_hi) + # multiplies by exactly zero, so a batched inner product may skip it; interior + # zeros (dead PSD bins) stay inside the range and are still multiplied through. + # Used only by the opt-in batched path -- self.ip() still integrates the full array. + _nz2 = self.weights2side != 0 + if _nz2.any(): + self.band_lo2side = int(np.argmax(_nz2)) + self.band_hi2side = len(_nz2) - int(np.argmax(_nz2[::-1])) + else: + self.band_lo2side, self.band_hi2side = 0, 0 + def ip(self, h1, h2): """ Compute inner product between two COMPLEX16Frequency Series @@ -2496,14 +2512,55 @@ def ip(self, h1, h2,include_epoch_differences=False): assert abs(h1.deltaF-h2.deltaF) <= TOL_DF\ and abs(h1.deltaF-self.deltaF) <= TOL_DF val = 0. - factor_shift = np.ones( len(h1.data.data)) if include_epoch_differences: fvals = evaluate_fvals(h1) factor_shift = np.exp(-1j* (float(h1.epoch) - float(h2.epoch))*fvals*2*np.pi) # exp( i omega( t_2 - t_1) ) - val = np.sum( np.conj(h1.data.data)*h2.data.data*factor_shift*self.weights2side ) + val = np.sum( np.conj(h1.data.data)*h2.data.data*factor_shift*self.weights2side ) + else: + # Bit-identical to multiplying by an all-ones factor_shift (x*1.0 == x in IEEE + # 754), but skips a len2side float64 allocation + one full-length complex + # multiply per call. At 1e6 bins that was ~0.5 ms of allocation and ~20% of the + # arithmetic, times O(10^5) calls in a higher-mode rotation precompute. + val = np.sum( np.conj(h1.data.data)*h2.data.data*self.weights2side ) val *= 2. * self.deltaF return val + def ip_matrix(self, listA, listB, chunk=1<<18): + r"""Batched form of ip(): returns the matrix M[a,b] = self.ip(listA[a], listB[b]). + + listA, listB are sequences of COMPLEX16FrequencySeries on the same 2-sided grid. + The double loop over (a,b) is a single matrix product, + + M = 2 df . conj(A) . (B . W)^T , A[a,f] = listA[a](f), B[b,f] = listB[b](f) + + which turns O(Na.Nb) separate full-length reductions into one pass over the data + plus a GEMM, and skips the frequency bins where the band weight is exactly zero. + Accumulated in chunks of `chunk` bins so the working set stays bounded. + + NOT bit-identical to the loop: the reduction order changes (pairwise np.sum over + the full array vs. blocked GEMM accumulation over the band). Measured deviation is + ~1e-15 relative to max|M| -- see DESIGN_precompute_crossterm_batching.md. Callers + that need the shipped rounding must keep using ip(). + """ + lo, hi = self.band_lo2side, self.band_hi2side + na, nb = len(listA), len(listB) + out = np.zeros((na, nb), dtype=np.complex128) + if hi <= lo: + return out + for m in listA: + assert m.data.length == self.len2side + for m in listB: + assert m.data.length == self.len2side + colsA = [m.data.data for m in listA] + colsB = [m.data.data for m in listB] + for s in range(lo, hi, chunk): + e = min(s+chunk, hi) + A = np.conj(np.stack([c[s:e] for c in colsA], axis=0)) + B = np.stack([c[s:e] for c in colsB], axis=0) * self.weights2side[s:e] + out += A @ B.T + out *= 2. * self.deltaF + return out + def norm(self, h): """ Compute norm of a COMPLEX16Frequency Series diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_bandlimited_retained_fft.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_bandlimited_retained_fft.md new file mode 100644 index 000000000..09b52cc7c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_bandlimited_retained_fft.md @@ -0,0 +1,295 @@ +# Retained-grid FFT for ordinary-ILE band-limited time marginalization + +## Scope and status + +This note covers only the dense `bandlimited` time-marginalization implementation +in `time_marginalization_quadrature.py`. It does not change the Q_lm time +stencil (`sinc` remains the ordinary-ILE default in the benchmark), the +peak-local implementation, the method selector, or the frozen paper benchmark. + +The code is a production-safe optimization candidate: supported CuPy complex128 +inputs use a retained-grid chirp-z evaluation, as do NumPy inputs at factor 8 and +above. NumPy factors 2 and 4 intentionally retain the established full FFT +below a conservative measured CPU crossover. Every declined or failed optimized +transform also retries the full-padding reconstruction. Cost selection and +failure retry are separately recorded by `last_report()`. An optimization +decline is therefore not reported as a waveform/likelihood failure and does not +by itself remove an AV sample. + +The numerical-identity, focused-kernel, and matched production-AV claims below +are verified for the enumerated 22 and higher-mode cells. This note does not +turn the microbenchmark into an evidence claim or extrapolate the production +result beyond those configurations. + +## Exact mismatch at production window sizes + +Let the gathered integration window have `n` coarse samples and let `f` be the +derived power-of-two refinement factor. The boundary construction forms the +literal reflected period + +``` +[x[0], ..., x[n-1], x[n-1], ..., x[0]] +``` + +of length `N = 2n`. The reference implementation zero-pads its spectrum to +`N f`, takes the entire inverse FFT, and retains only +`m = (n - 1) f + 1` forward-window samples. + +The two representative NCHUNK=40,000 shapes are: + +| cell | n | f | reflected N | reference IFFT Nf | consumed m | factorization | +|---|---:|---:|---:|---:|---:|---| +| 22, srate 4096 | 614 | 64 | 1228 | 78,592 | 39,233 | 78,592 = 256 x 307 | +| Lmax=4, srate 8192 | 1228 | 32 | 2456 | 78,592 | 39,265 | 78,592 = 256 x 307 | + +Thus roughly half of the explicitly generated inverse-FFT outputs are discarded. +More importantly, reflection leaves the prime factor 307 in every power-of-two +refinement length. The exact vendor-library implementation of that nonsmooth +FFT is not assumed here; measured cost, rather than a claim about proprietary +cuFFT internals, is the performance evidence below. + +## Retained-grid identity + +After the length-N FFT, arrange `N+1` coefficients at consecutive signed +frequencies `k=-N/2,...,+N/2`. As in the reference implementation, split the +even-period Nyquist coefficient equally between the two endpoints. The desired +sample `j` is then + +``` +y[j] = exp(-i pi j/f) / N + sum(q=0..N) C[q] exp(2 pi i q j/(N f)), j=0,...,m-1. +``` + +The sum is a uniform unit-circle chirp-z transform. Bluestein convolution +evaluates only the requested `m` points. Its compatible FFT lengths are 40,500 +for the 22 cell and 42,000 for the Lmax=4 cell, versus 78,592 in the reference +path. Chirp phases are reduced exactly modulo `2 N f` in int64 before conversion +to complex128; this avoids the accumulated unit-circle drift of repeatedly +raising one rounded complex root to high powers. + +All arrays, coefficient rearrangement, chirps, and FFTs use the caller's `xpy` +backend. `scipy.fft.next_fast_len` computes one host integer; it does not move +data off a GPU. Independent rows remain batched. Chirp plans are reused across +chunks and factors within one marginalization call, then released rather than +held in a process-global GPU cache. + +## Why cost still grows with SNR + +For the near-Gaussian time peak, +`sigma_t = 1/(2 pi rho sigma_f)`. The certified resolution requires +`deltaT/f <= sigma_t/2`, so the derived `f` grows approximately linearly with +SNR (in power-of-two steps). Both the reconstructed grid and the nonlinear +distance/phase likelihood callback contain `m ~ n f` points per refined row. +Consequently the irreducible callback/reduction work grows approximately as +rho, while the reference transform grows as roughly `rho log rho` and also pays +for the discarded reflected half and the nonsmooth FFT length. + +NCHUNK=40,000 is not itself an accuracy parameter. It supplies many rows to the +dense stage, which is divided into about 128-MiB working chunks. Commit +`70599f1f` already prevents a rare unresolved row from doubling the factor for +the whole group; each row now retires at its own certified factor. A larger AV +chunk still means proportionally more row transforms/callback evaluations and +can contain more high-factor rows. This optimization reduces the transform +constant, but intentionally does not alter the SNR-dependent resolution rule or +the number of `sinc`/likelihood evaluations. + +## GPU benchmark + +Hardware and software: NVIDIA RTX PRO 4000 Blackwell SFF (24,026.7 MiB), CUDA +12.8 runtime, cuFFT 11.3.3, CuPy 14.1.1, SciPy 1.15.3. Source base was the +immutable ordinary-ILE benchmark commit `476145cb`; candidate source was an +isolated clone based on that commit. The time-quadrature source in the HM +snapshot `50f470f8` was byte-identical to `476145cb`. Each arm processed 40,000 +row transforms in production-sized batches of 26. The reported wall interval +excludes Python and RIFT import but includes optimized-plan construction. Each +arm ran in a fresh process; RSS therefore includes the same RIFT/container +import baseline. + +The committed reproducer is `Code/test/benchmark_bandlimited_retained_fft.py`. +Two independent executions gave the wall ranges below; memory columns are from +the committed-reproducer execution. + +| cell | retained outputs | full wall (s) | retained wall (s) | paired speedup | host max RSS full/new (MiB) | CuPy pool full/new (MiB) | device delta full/new (MiB) | +|---|---:|---:|---:|---:|---:|---:|---:| +| 22, n=614, f=64 | 1,569,320,000 | 6.03--6.79 | 1.77--3.13 | 2.17--3.40x | 483.2 / 483.1 | 277.4 / 55.2 | 296 / 60 | +| Lmax=4, n=1228, f=32 | 1,570,600,000 | 6.12--6.69 | 2.28--2.89 | 2.32--2.69x | 480.7 / 484.1 | 285.1 / 66.8 | 306 / 72 | + +The host RSS difference is noise at an import-dominated baseline. The device +figures demonstrate that the explicit retained-grid transform does not hide a +larger chirp/workspace or CPU transfer: its CuPy-pool footprint is 20--23% of +the full-padding arm in these cells. + +A pre-commit sweep over every factor 2, 4, 8, 16, 32, and 64 processed 40,000 +rows at each of `n=614` and `n=1228`, using the same 128-MiB-derived batches. +The retained path was faster in all 12 cells; the smallest measured speedup was +1.53x (n=1228, factor 2). Thus applying it to every supported refinement factor +does not hide a measured low-SNR crossover on this device. + +The corresponding four-worker CPU sweep did have a small-grid crossover. In +balanced repeats the retained factor-2 transform cost 1.03--1.9 times the full +FFT for `n=614,1228,2457`, and factor 4 was 1.14 times slower at `n=2457` +(although faster for the prime-307 lengths). Since these grids are cheap and +not the high-SNR bottleneck, NumPy conservatively selects the full transform at +both factors 2 and 4. This selection is telemetry, not a failed optimization; +CuPy continues to use retained evaluation because its measured crossover is +below factor 2. + +Fixed-input parity used 32 full-band complex rows and a smooth nonlinear map +`100 logaddexp(0, Re(kappa))` before trapezoidal time integration and a +log-sum-exp evidence-like reduction: + +| cell | max abs delta kappa | max abs delta row lnL | delta aggregate lnZ | +|---|---:|---:|---:| +| 22, n=614, f=64 | 7.71e-15 | 2.27e-13 nat | -1.14e-13 nat | +| Lmax=4, n=1228, f=32 | 8.04e-15 | 3.98e-13 nat | +5.68e-14 nat | + +CPU tests also compare random Nyquist-populated rows at `n=614,1228,2457` +against the full-padding reference. The largest observed complex discrepancy +in the wider diagnostic sweep (`n=3` through 2457, factors 2 through 64) was +`5.7e-15`. + +A matched bounded 22 ordinary-ILE integration smoke used `bandlimited+sinc`, SNR +label 160, seed 99002, and `NMAX=NCHUNK=4000`. Both arms completed 4000 AV +evaluations. Full/new wall was 19.88/19.67 s, host max RSS was +1491.6/1493.5 MiB, and the reported log integral differed by `7.3e-12` nat +(13224.475448007970 versus 13224.475448007977). The deliberately tiny run had +ESS 1.73 and Pareto k-hat 11 in both arms, so it is an integration smoke, not +acceptable evidence or a throughput benchmark. The Lmax=4 claim remains the +fixed-shape kernel/parity result above; the converged production runs below are +the separate stochastic promotion gate. + +## Matched production AV validation + +The production test was registered before inspecting candidate results. It +replayed the exact argv and seed from each accepted immutable baseline record, +changing only the output prefix and RIFT tree. Every run used AV, +`bandlimited+sinc`, `NCHUNK=40000`, `NEFF=100`, `NMAX=4000000`, physical GPU 2, +and the same Apptainer image (SHA256 +`1367a60df7037a20927337f00175dfa72cf54e8dd843a424334a70ce7faf3427`). +All baseline-record input hashes were rechecked successfully after the runs. +Candidate source was the clean frozen commit +`c78c39ad25540cce0b2fadc95a5eb2c5735915d9`. The 22 reference is +`476145cbe9c1fb4e8c5621fdf3b11eebf97bcf47`; the HM reference is +`50f470f8a9355187387b0446d800fdb72bc2534c`. Between those two reference +commits, `factored_likelihood.py` and `time_marginalization_quadrature.py` are +byte-identical. The only ordinary-ILE driver change is in the 22-only direct +phase-marginalization guard, which the HM argv does not enter. + +A run was rejected, rather than interpreted, for a nonzero exit, source drift, +AV live-volume collapse, ESS below 100, Pareto k-hat at or above 0.7, or an +unverified CUDA backend. The optimization claim additionally required complete +telemetry, the requested transform route, zero transform fallback rows, and zero +failed marginalization calls. The harness records sampler acceptance separately +from optimization validation, so an optimization decline can remain a finite +full-sinc likelihood point without being mislabeled as either a waveform failure +or a successful retained-path validation. + +All candidate/control rows below passed both gates; the immutable rows had +already passed the same sampler gate. Delta-lnZ is candidate minus the +same-seed immutable reference; the final column divides it by the quadrature sum +of the two reported Monte Carlo errors. + +The HM result has two independent candidate seeds. Each 22 priority cell has +one candidate seed matched to an accepted reference seed, so no new 22 +candidate seed-scatter estimate is claimed here. + +| model/SNR | arm | seed | lnZ +/- sigma | ESS | k-hat | evaluations | delta lnZ | delta/combined sigma | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| 22/40 | immutable full reference | 1001 | 795.5928947372948 +/- 0.05992296 | 370.26 | -0.3513 | 80,014 | -- | -- | +| 22/40 | explicit full control at `c78c39ad` | 1001 | 795.5928947372948 +/- 0.05992296 | 370.26 | -0.3513 | 80,014 | 0 | 0 | +| 22/40 | retained at `c78c39ad` | 1001 | 795.5928947372953 +/- 0.05992296 | 370.26 | -0.3513 | 80,014 | +4.55e-13 | 5.37e-12 | +| 22/640 | immutable full reference | 1001 | 212588.68445187947 +/- 0.08341426 | 481.52 | -0.1458 | 321,015 | -- | -- | +| 22/640 | retained at `c78c39ad` | 1001 | 212588.68445187970 +/- 0.08341426 | 481.52 | -0.1458 | 321,015 | +2.33e-10 | 1.97e-9 | +| HM/51 | immutable full reference | 1001 | 1283.8369908486260 +/- 0.08481719 | 710.72 | 0.2726 | 1,461,820 | -- | -- | +| HM/51 | retained at `c78c39ad` | 1001 | 1283.8369908486268 +/- 0.08481719 | 710.72 | 0.2726 | 1,461,820 | +6.82e-13 | 5.69e-12 | +| HM/51 | immutable full reference | 1003 | 1283.8278744619759 +/- 0.08517261 | 693.57 | 0.2594 | 1,292,458 | -- | -- | +| HM/51 | retained at `c78c39ad` | 1003 | 1283.8278744619765 +/- 0.08517261 | 693.57 | 0.2594 | 1,292,458 | +6.82e-13 | 5.66e-12 | + +The wall and memory measurements are end-to-end process maxima, not the focused +kernel allocations reported above. Speedup is relative to the same-seed +immutable baseline. The baseline HM seed-1001 GPU monitor was incomplete, so +that cell has no baseline GPU-memory value; seed 1003 provides the matched HM +memory comparison. + +| model/SNR | arm | seed | wall (s) | speedup | host max RSS (MiB) | GPU peak (MiB) | retained/full transform rows | fallback/failed calls | +|---|---|---:|---:|---:|---:|---:|---:|---:| +| 22/40 | immutable full reference | 1001 | 27.93 | 1.000x | 1603.5 | 9874 | n/a | n/a | +| 22/40 | explicit full control | 1001 | 24.40 | 1.145x | 1478.6 | 9874 | 0 / 81,307 | 0 / 0 | +| 22/40 | retained | 1001 | 23.55 | 1.186x | 1479.0 | 9868 | 81,307 / 0 | 0 / 0 | +| 22/640 | immutable full reference | 1001 | 481.56 | 1.000x | 1514.5 | 17,422 | n/a | n/a | +| 22/640 | retained | 1001 | 344.50 | 1.398x | 1508.2 | 17,266 | 325,678 / 0 | 0 / 0 | +| HM/51 | immutable full reference | 1001 | 192.04 | 1.000x | 1690.6 | -- | n/a | n/a | +| HM/51 | retained | 1001 | 132.94 | 1.445x | 1684.7 | 18,006 | 1,466,334 / 0 | 0 / 0 | +| HM/51 | immutable full reference | 1003 | 172.54 | 1.000x | 1675.4 | 17,808 | n/a | n/a | +| HM/51 | retained | 1003 | 117.11 | 1.473x | 1662.4 | 17,228 | 1,296,699 / 0 | 0 / 0 | + +The explicit full control is validation-only instrumentation around the unchanged +reference helper; there is no new production switch or default. It reproduced +the immutable SNR-40 evidence and AV diagnostics exactly. Retained evaluation +was 1.036x faster than that same-commit full control in this low-SNR cell. At +SNR 640, where 287,898 of 321,015 refined rows required factor 256, the retained +path reduced end-to-end wall time by 28.5%. The focused transform's large memory +reduction is diluted by waveform, sampler, and likelihood allocations in a full +job: measured end-to-end GPU peaks decreased by 156 MiB at 22/SNR640 and by 580 +MiB in the matched HM seed-1003 run. + +Transform routing was fully observed, not inferred from the selected option: + +| model/SNR/seed | successful calls | refined-row factor histogram | max reference/retained-plan FFT | retained rows | selected-full rows | fallback rows | failed calls | +|---|---:|---|---:|---:|---:|---:|---:| +| 22/40/1001 retained | 2 | 2:98, 4:1,581, 8:74,816, 16:3,510 | 19,648 / 11,088 | 81,307 | 0 | 0 | 0 | +| 22/40/1001 full control | 2 | 2:98, 4:1,581, 8:74,816, 16:3,510 | 19,648 / -- | 0 | 81,307 | 0 | 0 | +| 22/640/1001 retained | 8 | 8:2, 16:13, 32:56, 64:369, 128:32,677, 256:287,898 | 314,368 / 158,400 | 325,678 | 0 | 0 | 0 | +| HM/51/1001 retained | 37 | 2:22,099, 4:48,048, 8:77,241, 16:1,311,948, 32:312 | 78,592 / 42,000 | 1,466,334 | 0 | 0 | 0 | +| HM/51/1003 retained | 32 | 2:22,196, 4:48,193, 8:76,719, 16:1,142,753, 32:296 | 78,592 / 42,000 | 1,296,699 | 0 | 0 | 0 | + +The separate low-factor CPU router witness passed with exact array equality: +factor 4 selected the established full transform for two rows, recorded two +`full_fft_selected_rows`, and recorded zero fallback rows. The adjacent forced +transform-decline and likelihood-exception tests also passed, establishing that a +finite full-sinc retry and a genuine likelihood failure remain distinct outcomes. + +## Failure and telemetry contract + +The retained path is certified only for NumPy/CuPy, complex128 spectra, even +reflected periods, and power-of-two factors above one whose modular chirp indices +fit exactly in int64. Other combinations, plan-construction failures, and +transform exceptions enter the full-padding reference path. A RuntimeWarning is +emitted once per reason per call when warning policy permits it; warnings promoted +to exceptions are contained so diagnostics cannot drop the point. + +`last_report()` records: + +- `bandlimited_fft_strategy`: retained, full selected, full fallback, mixed, or + unused; +- retained/selected/fallback batch and row-transform counts; +- a reason map for an intentional full-FFT cost selection; +- the fallback exception/reason map; +- reference full length, retained-grid length, compatible convolution length, + largest factor, and number of per-call plans. + +The likelihood callback is invoked outside the guarded transform helper. Its +exception is therefore not swallowed or relabeled as an FFT decline. Tests pin +both directions: forced optimized failure returns the finite full-sinc result +with provenance, while a forced callback failure retains its original identity. + +## Validation and promotion gate + +The focused suite passes on the actual CuPy backend, including GPU/CPU parity, +unsupported-factor fallback, warnings-as-errors, and callback-failure identity. +The complete `test_time_marginalization_quadrature.py` gate passed 90 tests. The +matched production gate now also passes for 22/SNR40, 22/SNR640, and two HM/SNR51 +seeds: every run converged without collapse, every same-seed delta-lnZ was below +`2e-9` of the combined Monte Carlo uncertainty, and no optimized transform +declined or failed. This promotes the implementation for those ordinary-ILE +configurations. It does not validate a different time stencil, backend, +marginalization method, or model family; the fail-safe full-sinc route remains +required outside the certified transform contract. + +The committed production reproducer is +`Code/test/run_bandlimited_retained_ile_validation.py`; its companion driver +`Code/test/telemetry_bandlimited_retained_ile.py` aggregates per-call reports and +provides the explicitly labeled full-FFT control. Raw scientific products stay +outside the repository under `/tmp/rift-retained-production-validation/runs`. +Only compact `validation_record.json` files there are needed to audit the tables +above; no posterior samples, logs, or run directories are committed. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_freqresponse_vectorized_coefficients.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_freqresponse_vectorized_coefficients.md new file mode 100644 index 000000000..bf6be2b7c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_freqresponse_vectorized_coefficients.md @@ -0,0 +1,114 @@ +# Finite-size response coefficients: one block, not one sample at a time + +`DiscreteFactoredLogLikelihoodFreqResponseNoLoop` is the production finite-size +(`--freqresponse`) likelihood. Everything in it was vectorized over the Monte Carlo block +except the response coefficients `b_p(det, RA, DEC, psi)`, which were built by a Python loop +calling the scalar `response_coefficients` once per sample per detector. That loop cost more +than the rest of the likelihood. + +## What the loop was doing + +Per sample, per detector, `response_coefficients` called: + +| call | depends on the sample? | +|---|---| +| `lal.GreenwichMeanSiderealTime(tref)` | no, `tref` is one scalar for the whole block | +| `slowrot_freqresponse.detector_geometry(det, L_arm)` | no, detector and arm length only | +| `_triad(dec, psi, g)` | yes | +| `_lwl_response(response, X, Y)` | yes | +| arm projections `X.x_arm`, `nhat.x_arm`, ... | yes | +| `finite_size_beta(geom, Qmax)` | yes | + +The two sample-independent rows were recomputed for every sample. In a profiled five-detector +CE+ET+K run (282,000 samples) that was 1,409,085 LAL detector lookups for five distinct +answers. + +The loop body also read + +```python +bvec.setdefault(p, np.zeros(npts_ex, dtype=complex)) +``` + +Python evaluates `setdefault`'s second argument whether or not the key is present, so this +allocated an `npts_ex`-long complex array on every `(sample, p)` iteration: 8,455,844 +allocations in that run, 83 s of a 293 s integration. + +## What it does now + +`slowrot_freqresponse` gained two functions: + +- `detector_geometry_cached(det, L_arm)` memoizes `detector_geometry` on `(det, L_arm)` and + returns read-only arrays, so a caller that writes to one fails instead of corrupting the + next caller. `detector_geometry` itself is unchanged and uncached. +- `finite_size_geometry_vector(det, ra, dec, psi, gmst, L_arm)` evaluates the triad, the + long-wavelength contraction and the arm projections for a whole block. `finite_size_beta` + already accepted arrays and takes its result unchanged. + +`factored_likelihood_freqresponse` gained `response_coefficients_vector`, which returns +`{p: (npts_ex,) complex}`. The NoLoop calls it once per detector. The scalar +`response_coefficients` is still shipped and is still the definition of `b_p`; the block +form is checked against it. + +This is the same structure `factored_likelihood_with_rotation.rotation_coefficients_vector` +already had for the sidereal-harmonic coefficients, which is why the `--rotation-slow` row +cost 5.4x the long-wavelength baseline while `--freqresponse` cost 137x for six basis +elements. + +## Does lnL move? + +`finite_size_beta` now squares `zx` through `_csquare`, which writes out CPython's own +complex multiply. numpy's `complex128 ** 2` rounds differently from Python's `complex ** 2` +in about 29% of random samples, at one ulp; `_csquare` reproduces the scalar value bit for +bit on both paths. + +After that, `b_0` through `b_3` are bit-identical to the scalar loop for every detector +tested. `b_4` and `b_5` carry `a_x**3` and `a_x**4`, where CPython's `float.__pow__` calls +libm `pow` and numpy's power loop does not; those agree to one ulp on about 18% of samples, +worst relative difference 5e-16. No reassociation of the physics is involved and no +alternative array expression matches libm `pow` without a per-element Python call. + +The consequence for the likelihood was measured, not argued. `analyses/slowrot_finite-size` +(RIFT_roboto_paper) builds a self-consistent finite-size injection where truth is the exact +global maximum. Evaluating the NoLoop on a fixed block of extrinsic samples under the two +trees: + +| configuration | samples x times | `lnL_t` differing | max abs difference | +|---|---|---|---| +| CE+ET+K, Qmax=4, SNR 30.1, seed 4242 | 2000 x 64 | 0 of 128,000 | 0 | +| CE-ET, Qmax=6, SNR 27.1, seed 909 | 1500 x 64 | 0 of 96,000 | 0 | + +lnL is bit-identical on both. The ulp differences in `b_4`, `b_5` are below the rounding of +the sums those coefficients enter. + +## Cost + +`analyses/response_cost_scaling/run_one.py --response finite --lmax 2 --device gpu --stencil +sinc`, CE+ET+K, IMRPhenomXPHM, seglen 128 s, srate 8192, Qmax 4, on an RTX PRO 4000 Blackwell +(ldas-pcdev11). Stage attribution by that directory's `parse_profile.py`. Arms alternate per +seed. + +| seed | baseline `tau_it,like` | block `tau_it,like` | speedup | `tau_it,rest` | `N_it` both arms | +|---|---|---|---|---|---| +| 1002 | 1706.83 us | 37.03 us | 46.1x | 7.2 us | 281,817 | +| 2002 | 1778.03 us | 71.09 us | 25.0x | 16.3 us | 321,687 | +| 3002 | 1755.94 us | 73.08 us | 24.0x | 16.6 us | 403,083 | + +The spread is node contention. `tau_it,rest` is the sampler outside the likelihood and +matches within each pair; the block arm tracks it at 4.4 to 5.2x, while the baseline arm has +sd 2.1% across the three because it is bound by Python call overhead. + +Provenance, correctness table and the one row that was killed by the per-UID memory cgroup: +RIFT_roboto_paper +`analyses/finite_response_vectorization/RESULTS_2026-09-09_block_response_coefficients.md`. + +## Not done + +The GPU path still brings `RA`, `DEC`, `psi` to the host (`_h()`) and computes `b_p` in +numpy, then copies `b_p` back. That is one transfer of three float64 blocks per likelihood +call, not one per sample, and it was already there before this change. Moving the geometry +onto the device would need `lal.ComputeDetAMResponse` reimplemented in cupy, and the +long-wavelength baseline `F0` is the one part of the response the module keeps exact against +LAL. The remaining host cost is a fixed number of array operations per block. + +PR #307 (`codex/combined-slowrot-freqresponse`) adds a combined rotation-and-finite-size +module with its own block-form coefficients. It does not touch either file changed here. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_noloop_per_detector_glue.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_noloop_per_detector_glue.md new file mode 100644 index 000000000..38dd42e61 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_noloop_per_detector_glue.md @@ -0,0 +1,233 @@ +# NoLoop: what the detector loop was recomputing, and why the split is bitwise exact + +Scope: `DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop`, the maintained GPU +likelihood (`--vectorized --gpu`). This note records *why* the source geometry was +lifted out of the detector loop, and the constraint that decided the implementation. + +## The measurement that motivated it + +Stage attribution inside NoLoop, RTX PRO 4000 Blackwell (sm_120), cupy 14.1.1 / +CUDA 12.9, ILE-GPU-Paper demo, `--interpolate-time nearest`, `--n-chunk 10000`, +1000 calls, each stage device-synced (which inflates the total by 2.5%): + +| stage | share of NoLoop | +|---|---| +| `simps` | 32.6% | +| `SphericalHarmonicsVectorized` | 22.2% | +| `ComputeDetAMResponse` | 17.1% | +| residual (`kappa_sq`, `rho_sq` einsums, `exp`/`log`, allocation) | 18.8% | +| `TimeDelayFromEarthCenter` | 3.6% | +| `Q_inner_product_cupy` (the CUDA kernel) | **5.7%** | + +The hand-written kernel is a twentieth of the cost; the rest is cupy glue. The three +geometry stages total ~43% and act on `(n_extrinsic,)` arrays — a few hundred KB. Time +spent there is therefore kernel-launch and op-count bound, not bandwidth bound: long +chains of small elementwise operations. Each was being rebuilt once **per detector** +although none of them depends on the detector. + +Independent confirmation that the per-call cost is launch-bound: sweeping `--n-chunk` +on an RTX 3080 fits `cost = 5.7 ms + 0.61 us x n_chunk`, i.e. at `--n-chunk 10000` +roughly half of every call does no more work for a larger batch. + +## What is actually per-detector + +Only the contraction with the interferometer's own constants: + +- `ComputeDetAMResponse` — six trig evaluations and twelve elementwise combinations + build the `(X, Y)` polarization basis from RA/DEC/psi/GMST. Only the two `inner` + contractions against `detector_response_matrix` are per-detector. +- `TimeDelayFromEarthCenter` — `ehat_src`, the unit vector towards the source, is + source-only. Only the `inner` against `detector_earthfixed_xyz_metres` is not. +- `SphericalHarmonicsVectorized` — depends on `(modes, incl, phiref)`. Detectors share + a mode list in practice, since the modes come from one waveform. +- `DetectorPrefixToLALDetector` plus two host-to-device transfers were also being + redone every call, for values fixed for the lifetime of the process. + +## The constraint: bitwise, not approximately + +This is a likelihood behind published results, so the split had to leave lnL +*bit-identical*, which rules out the obvious vectorization. `ComputeDetAMResponse`'s +docstring advertises a leading detector axis, but that path does not actually work — +`X * xpy.inner(X, R)` fails to broadcast for `(n_ex, 3)` against `(n_det, 3, 3)`. The +natural fix, one batched `einsum` over stacked detectors, reassociates the contraction +and agrees only to ~4e-16. Fewer launches, but not the same number. + +So the per-detector halves keep the identical `inner` calls in the identical order and +only the source-only prologue is shared. `test/test_vectorized_lal_tools_split.py` +pins that with `array_equal`, not a tolerance, on three real interferometer geometries. + +**Against a FROZEN COPY of the pre-split bodies, not against the wrapper.** The first +version of that test compared `ComputeDetAMResponse(...)` to +`ComputeDetAMResponsePrecomputed(SourcePolarizationBasis(...))` -- but after the split +the wrapper *is* that composition, so the comparison was tautological and could not +fail. An adversarial review demonstrated it passing with a sign flipped in the +source-only half, with the response matrix doubled in the per-detector half, and with +the speed of light wrong by 0.1%. All three now fail. The lesson generalizes: **when a +refactor splits a function, the two halves are not an independent check on each other** +-- freeze what was replaced, or compare against an outside implementation. + +## Sharing hazards, and how they are handled + +- **The phase-marginalization branch mutates `Ylms_vec` in place** (`[:, 1] = conj(...)`), + and `rho_sq_det` above it must see the un-conjugated array. A shared array would leak + one detector's conjugation into the next detector's self-term. Each detector gets a + copy when `phase_marginalization` is on; a copy of `(n_extrinsic, n_lms)` is still far + cheaper than rebuilding the harmonics. +- **`lookupNKDict[det]` may be a device array**, so comparing mode lists per call would + force a synchronization. `_mode_list_key` memoizes a hashable host key on the array + *object*, keeping a reference so `id()` cannot be recycled. Detectors with genuinely + different mode lists therefore get a correct, merely unshared, result. +- `TimeDelayFromEarthCenterPrecomputed` divides in place into the result of `inner`, + which is a fresh array — not into the shared `ehat_src`. The test pins that too. + +## Measured effect + +Same captured NoLoop arguments replayed through both trees (Blackwell, `nearest`, +`--n-chunk 10000`, 100 calls per timing, 3 repetitions), output bitwise identical: + +| configuration | before | after | | +|---|---|---|---| +| H1 L1 (2 detectors) | 12.41 ms/call | 10.26 ms/call | **-17.4%** | +| H1 L1 V1 (3 detectors) | 16.81 ms/call | 13.06 ms/call | **-22.3%** | + +The saving scales with detector count, as it should: the shared prologue is paid once +instead of `n_det` times. The CPU (`xpy=numpy`) path is unchanged within noise — it is +dominated by the `(n_extrinsic, npts, n_lms)` window build, not by this glue. + +## What this deliberately does NOT do + +- `simps`, the single largest stage, is untouched. It is a fixed linear functional, so + it could be one `gemv` against precomputed weights — which is what the fused calmarg + path already does via `w_t = simps(eye(npts))`. That changes summation order and so + is not bitwise; it belongs in its own change with its own accuracy argument. +- The post-kernel reduction is untouched. Routing `n_cal == 1` through the existing + `Q_fused_calmarg` kernel measured a further ~24%, agreeing within Monte Carlo error + but not bitwise. Also a separate change. + +--- + +# Round 2: the accumulators, and a per-operation cost table + +After the hoist above, stage attribution became misleading: it device-syncs after every +wrapped call, so a function called once per *detector* is charged three times the sync +penalty of one called once per *likelihood call*, and the mode inflated the total by 26%. +The numbers below come instead from timing each operation in a tight loop with a single +sync (`bench/micro_ops.py` in the profiling archive), at production shapes +`n_extrinsic = 10000`, `npts = 614`, three detectors, on an RTX PRO 4000 Blackwell. They +sum to 11.87 ms against a measured 12.10 ms/call, i.e. they account for 98% of the +function. + +| operation | ms/op | x per call | ms per NoLoop call | +|---|---|---|---| +| `kappa_sq += Q_prod * invDist` | 1.519 | 3 | **4.556** | +| `ComputeDetAMResponsePrecomputed` | 0.628 | 3 | 1.885 | +| `simps` over `(10000, 614)` | 1.765 | 1 | 1.765 | +| `Q_inner_product_cupy` | 0.391 | 3 | 1.173 | +| `kappa.real - 0.5*rho` (stride-0 view) | 0.619 | 1 | 0.619 | +| `exp` in place | 0.499 | 1 | 0.499 | +| `SphericalHarmonicsVectorized` | 0.401 | 1 | 0.401 | +| `SourcePolarizationBasis` | 0.367 | 1 | 0.367 | +| `max(axis=-1, keepdims)` | 0.264 | 1 | 0.264 | +| `TimeDelayFromEarthCenterPrecomputed` | 0.062 | 3 | 0.187 | +| `SourcePropagationDirection` | 0.127 | 1 | 0.127 | +| `rho_sq` vector accumulate | 0.008 | 3 | 0.024 | + +The data term dominates, and it dominates because of how it is *stored*, not what is +computed into it. + +## rho_sq was 49 MB of duplicated scalars + +`rho_sq` is the `` term. Every detector contributes `rho_sq_det` of shape +`(npts_extrinsic,)` — it has no time dependence at all — and that was being broadcast +into a dense `(npts_extrinsic, npts)` accumulator: a 49 MB zero-fill, then one 49 MB +read-modify-write per detector, to store `npts` identical copies of each value. + +It is now summed as a vector and exposed as a stride-0 `broadcast_to` view. Measured: +dense accumulate 0.121 ms/detector against 0.008 for the vector, and the downstream +`kappa.real - 0.5*rho` drops from 0.758 ms to 0.619 ms because the subtrahend now fits +in cache. The calibration path already did exactly this for `rho_sq_cal`; this brings +the ordinary path in line. + +Consumers that need real backing memory go through `_dense_rho_sq()` and pay what they +paid before. There are two classes: the fused calmarg CUDA kernels, which index raw +device pointers and would read garbage from a stride-0 view, and the non-Simpson +quadrature helpers, which are free to write into what they are handed. + +## kappa_sq did not need to start at zero + +`kappa_sq` is 98 MB of complex128. It was zero-filled, then for each detector the +distance scaling allocated another full-size temporary and the result was accumulated in +— so a three-detector network paid one 98 MB fill, three 98 MB temporaries, and three +98 MB read-modify-writes. It now scales the Q kernel's own freshly allocated output +buffer in place and takes the first detector's buffer as the accumulator. + +The one arithmetic caveat: `0.0 + x` is exactly `x` for every finite `x`, and for Inf and +NaN, but `0.0 + (-0.0)` is `+0.0` while starting from the buffer preserves `-0.0`. A +signed zero in `kappa_sq` is unobservable downstream — it survives `.real`, and +`exp(-0.0) == exp(+0.0) == 1.0` — so this is noted for completeness rather than as a +behavioural difference. + +## Measured, cumulative, all bitwise + +Same captured NoLoop arguments replayed through each tree, H1 L1 V1, `nearest`, +`n_chunk 10000`, 100 calls per timing: + +| tree | ms/call | vs base | +|---|---|---| +| `rift_O4d` | 17.06 | — | +| \+ hoist source-only geometry | 13.08 | −23.3% | +| \+ `rho_sq` as a vector | 12.10 | −29.1% | +| \+ `kappa_sq` in-place | 11.12 | **−34.8%** | + +`test/test_noloop_accumulator_shapes.py` pins both accumulators against a reference +implementation written the original way, with `array_equal` rather than a tolerance, at +one, two and three detectors. The reference passes against the unpatched tree as well, +which is what makes it a check on the change rather than a transcription of it. + +## Round 3: the time integral, and the one change that is not bitwise + +`simps` was 1.765 ms/call. It is a fixed linear functional at fixed `dx`, so it equals a +matrix-vector product against precomputed weights — measured at **0.049 ms**, a 36x +saving. The fused calmarg path already built exactly those weights by hand with +`w_t = simps(eye(npts))`; that is now a single cached helper, `_simps_weights`, so the +tree carries one definition of the equivalence instead of two. + +A `gemv` reassociates the summation, so unlike everything above this is **not** bitwise. +It is the same RULE: the weights come from the very `simps` implementation the call site +would otherwise have used, so the `even='avg'`-versus-Cartwright distinction that +separates the vendored GPU copy from scipy's is preserved exactly. Only the order of the +additions changes. + +**Measured discrepancy**, over 10 000 real extrinsic samples spanning lnL from +-2.2e6 to +116: + +| | | +|---|---| +| max abs difference | **2.8e-14 nats** | +| median abs difference | exactly 0 | +| max relative difference | 7.1e-14 | +| float64 rounding scale of the values themselves (`eps x max abs lnL`) | 4.9e-10 | + +The difference is below the rounding scale of the quantities being compared, and both +paths are deterministic run to run. For physical scale, the errors already present in +this integral are between eleven and sixteen orders of magnitude larger: the two `simps` +variants in this tree disagree by **0.405 nats** on an under-resolved peak, and the +`nearest` time stencil costs **200-443 nats at SNR 100** (`--interpolate-time` help text, +issue #233). Simpson's real accuracy limit here is sub-sample resolution of a peak whose +width is set by the signal rather than by the sample rate — which is what the +`time_quadrature` and stencil work addresses — not the order of its additions. + +`test/test_noloop_accumulator_shapes.py` splits the two guarantees rather than blurring +them: the accumulators are checked with `array_equal` at `return_lnLt=True`, before the +integral, and the quadrature is checked separately against `simps` at a tolerance far +tighter than anything physical. A failure of the second means the rule changed, not that +rounding drifted. + +## Where the remaining time goes + +After all three rounds, at three detectors and `n_chunk 10000`, no single item dominates: +the Q kernel (~1.2 ms), the detector-response contraction (~1.9 ms), and the +`exp`/`max`/subtract reduction (~1.4 ms) are the three largest, and none has an obvious +order-preserving win left. The response contraction is the best remaining candidate — +four `inner` calls per detector against a 3x3 matrix — but batching it over stacked +detectors reassociates, for a much smaller payoff than this round bought. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md index b90d8b8f1..9f7d34d3d 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_peak_local_framework.md @@ -610,52 +610,65 @@ the symmetry can be broken at roundoff. A symmetry assumed exact when it is 1e- 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. +### The 2-D enumerator is a finite resultant, not a φ grid + +The earlier hybrid in this section solved the degree-four u polynomial at 64 sampled φ +values. That cost was amplitude-independent, but it was still GRID SEEDING and therefore +was not enumeration of the known finite stationary set. Higher-mode likelihoods know both +orders: the exponent is a real Laurent polynomial of bidegree `(K,Q)=(2 m_max,2)`. There is +no reason to replace that information by an angular resolution. + +`bivariate_trig_stationary.py` now implements the host reference construction. Expand the +stored half-table into its full Hermitian Laurent table and form + +``` +F(z,w) = partial_phi g, G(z,w) = partial_u g, +z = exp(i phi), w = exp(i u). +``` + +After clearing negative powers these are ordinary bivariate polynomials. A coordinate +resultant is a poor numerical choice because several real modes commonly have exactly the +same φ (or u), making the hidden root multiple. Instead choose a generic affine hidden +variable `t = z + alpha w`, substitute `z=t-alpha w`, and eliminate `w` with the Sylvester +matrix polynomial `S(t)`. A block companion linearization turns `det S(t)=0` into one +generalized eigenproblem. The Newton polygons give the exact mixed-volume budget; for the +full rectangular derivative supports it is `8 K Q` (64 at `(4,2)`). This is the finite +object that exhausts the isolated complex stationary set. + +The numerical certificate has four gates, all fail closed: + +1. recover the mixed-volume number of verified roots in `(C*)^2`; +2. require nonsingular, adequately conditioned stationary Jacobians and a backward-stable + generalized eigenproblem; +3. classify torus roots with the Laurent system's reciprocal-conjugate involution, not an + `abs(|z|-1)/dev/null` + is indistinguishable from a clean result; + - a *sweep that hides the defect from itself* — grouping numerals by string reports + `7.069` and `7.07` as unrelated, so the tool written to find multi-spelling reports + clean on a file that has it. + The third is the worst of the three, because running it converts "unchecked" into + "checked and clean" without touching the code. Before trusting any check, ask what + input would make it FAIL; if you cannot name one, it is decoration. +* **A COMPRESSION of verified facts is a NEW claim, and does not inherit their + verification.** The same shape as the rule above, from the opposite end: one is a check + that cannot fail, this is a claim nobody checked *because its parts were checked*. + Measured on this work: four per-axis defaults were each independently verified from the + code and each held, and the one-sentence summary of them was still false — it asserted a + pattern that one of the four axes is a counterexample to, because the default there had + deliberately been moved to the accurate scheme. Every input was true and the summary was + not. Verifying the parts is the step that makes checking the whole feel unnecessary, + which is exactly when it is required. +* **Do not put a broad `except` around a certificate call, in shipped code OR in a + harness.** An error filter converts a bug into a result, and the result looks clean. + Measured while sizing this note's own acceptance table: a broad `except Exception` + around `joint_marginalize_peak_local` caught a tuple-unpack error and scored it as a + DECLINE, reporting a flat 0% acceptance at every amplitude — a uniform, plausible, + entirely fabricated headline that was caught only because it contradicted a number + already in hand. A decline must come from the ledger, never from an exception. * **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. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_precompute_crossterm_batching.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_precompute_crossterm_batching.md new file mode 100644 index 000000000..8d97e4ba3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_precompute_crossterm_batching.md @@ -0,0 +1,195 @@ +# Batching the ILE mode cross terms + +## The problem + +`PrecomputeLikelihoodTerms` and its two response variants build the mode cross-term +banks `U` and `V` before any Monte Carlo sampling starts. For the slow-rotation path at +`--l-max 4` the precompute is 94% of ILE wall clock. + +Profiled on `ldas-pcdev2`, one intrinsic point, CE+ET+K (5 detectors), IMRPhenomXPHM, +seglen 128 s, srate 8192, RTX PRO 4000 Blackwell, container +`rift_o4d_cc90-120_cuda128_20260717.sif`, RIFT `43918b22a`. Records: +`/scratch/richard.oshaughnessy/response_cost/campaign_main/`. + +| response | modes | precompute (s) | integration (s) | wall (s) | +|---|---|---|---|---| +| long-wavelength | l<=2 | 26.5 | 28 | 54.6 | +| long-wavelength | l<=4 | 61.4 | 32 | 93.0 | +| slow rotation | l<=2 | 335 | 38 | 373 | +| slow rotation | l<=4 | 1318 | 86 | 1404 | + +`cProfile` on the l<=4 rotation run splits the 1318 s as follows. All three response +paths reach the same function, `factored_likelihood.ComputeModeCrossTermIP`. + +| site | calls | cumulative (s) | share of precompute | +|---|---|---|---| +| `ComputeModeCrossTermIP` | 260 | 1151.6 | 87% | +| ... of which `ComplexIP.ip` | 112560 | 929.8 | 71% | +| ... of which `InnerProduct.__init__` | 260 | 221.0 | 17% | +| `ComputeModeIPTimeSeries` (data-term overlaps, FFTs) | 25 | ~130 | 10% | +| waveform generation | 2 | ~21 | 2% | + +Waveform generation is 2% of the precompute. The FFTs are 10%. The cost is the cross +terms, and inside them it is 112560 separate reductions over a 1048576-point array. + +## Structure + +For one detector and one pair of response-basis elements, + + U[a][b] = 2 df sum_f conj(hA_a[f]) hB_b[f] W(f) + +with `W` the two-sided `1/S(f)` band weight. The shipped code evaluates this with a +Python loop over `(a,b)`, calling `ComplexIP.ip` once per pair. Each call allocates four +full-length temporaries and reduces them. + +This is one matrix product. Stack the modes into `A[a,f]` and `B[b,f]` and the whole +block is `2 df . conj(A) . (B W)^T`: one pass over the data plus a GEMM, instead of +`Na*Nb` passes. + +The counts multiply. At `--l-max 4` there are 21 modes, so 441 pairs per block. The +slow-rotation path has 5 basis elements, giving 25 blocks per detector for `U` and 25 +for `V`, over 5 detectors: 250 blocks, 110250 inner products. `--freqresponse` +(`Qmax+2 = 6` elements) and the combined path added by PR #307 +(`factored_likelihood_rotating_freqresponse.py`, `6 x 5 = 30` elements, 900 blocks per +detector per bank) scale as the square of the basis size. All four call +`ComputeModeCrossTermIP`, so all four take the batched path when it is enabled. + +## What is not the answer + +**More cores.** The inner product is memory-bandwidth bound, not compute bound. Raising +`OMP_NUM_THREADS` from 1 to 8 or 16 made every variant slower, including the GEMM. +Per-block times on `ldas-pcdev2` in the container, 21 modes, 1048576 bins, 4 replicates: + +| variant | OMP=1 (s) | OMP=8 (s) | +|---|---|---| +| shipped loop | 6.94 | 7.17 | +| batched GEMM | 0.478 | 0.502 | + +**The GPU.** The device sits at 3-9% mean utilization through these runs, but it is not +what is missing. A cupy version of the same GEMM takes 0.037 s against 0.118 s on the +CPU, and it still needs the host-side stacking that dominates the batched path. The +card's FP64 throughput is a small fraction of its FP32 throughput, so the arithmetic +never becomes the constraint. The 26x from batching on the CPU is available without a +device transfer, and the remaining time is host memory traffic that a GPU does not +remove. + +## Changes + +Four, in `lalsimutils.py` and `factored_likelihood.py`. No response-path module is +modified. + +1. `ComplexIP.ip` no longer builds an all-ones `factor_shift` array when + `include_epoch_differences` is false. Bit-identical: `x * 1.0 == x` in IEEE 754. It + removes a `len2side` allocation and one full-length complex multiply per call, 55.5 s + of allocation alone in the profiled run. + +2. `InnerProduct.__init__` zeroes the inverse-spectrum-truncation window by slice + assignment instead of a per-element loop over ~1e6 SWIG elements. Bit-identical. + Construction drops from 0.287 s to 0.0455 s at campaign shapes (`ldas-grid`, IGWN + CVMFS python). This is the whole 221 s above. It applies only when + `--inv-spec-trunc-time` is nonzero, which is the driver default but is not the + production setting; see "Provenance of the motivating numbers". + +3. The array-PSD weight fill is vectorized, preserving the loop's `!= 0` mask: 22.2 ms + to 1.9 ms at 128512 in-band bins. The driver passes a `REAL8FrequencySeries`, so this + branch is off the measured path. + +4. `ComplexIP.ip_matrix(listA, listB)` computes a whole block as a chunked GEMM over the + nonzero support of the weights. `ComputeModeCrossTermIP` uses it when + `RIFT_PRECOMPUTE_BATCHED_CROSSTERMS=1` or `batched=True`. Default is off. + +Changes 1-3 are on by default because they alter no output bit. Change 4 is opt-in +because it does. + +## Numerics of the batched path + +The reduction order changes: pairwise `np.sum` over the full array becomes blocked GEMM +accumulation over the band. Measured against the shipped loop, worst element relative to +`max|U|` over the block, 21 modes, 1048576 bins: + +| weight support | deviation | +|---|---| +| band only, 41% support | 1.4e-15 | +| full support (inverse spectrum truncation on) | 2.7e-15 | + +`lnL` is a contraction of these matrices, so a fractional error of 3e-15 on entries whose +scale sets `lnL ~ rho^2/2` gives `|d lnL| ~ 1e-12` nats at `rho = 40`. End-to-end +confirmation is in the results section below. + +The batched path re-imposes the `same_waveform_Q` mirror explicitly, so the exact +Hermitian and transpose relations the shipped path guarantees still hold element for +element rather than to rounding. + +## Known behaviour difference + +Frequency bins where the weight is exactly zero are skipped. Their contribution is +`0.0` in the shipped path, so the value is unaffected, but a non-finite template sample +outside the band would propagate to `NaN` in the shipped path and be dropped here. +Inverse spectrum truncation smears the weights to full support, in which case nothing is +skipped and `band_lo2side, band_hi2side` span the whole array; the support is read off +the weights rather than assumed from `(fmin, fMax)`. + +## Provenance of the motivating numbers + +The `response_cost` campaign ran without `--inv-spec-trunc-time`, taking the driver +default of 8 s. Production runs set it to 0 (RO, 2026-09-09). Two consequences for the +table at the top of this file: about 190 s of the 1318 s precompute is a term production +does not pay, and the band weights in production are zero outside `[fmin, fMax]`, which +production runs. + +## Results + +Interleaved arms, `ldas-pcdev2`, container `rift_o4d_cc90-120_cuda128_20260717.sif`, +one intrinsic point, CE+ET+K, IMRPhenomXPHM, seglen 128 s, srate 8192, `--rotation-slow +--l-max 4`, `--inv-spec-trunc-time 0` (production), seed 1002, 2 replicates. Run dirs: +`/scratch/richard.oshaughnessy/precompute_speed/ab/`. + +| stage | baseline (s) | batched (s) | speedup | +|---|---|---|---| +| `ComputeModeCrossTermIP`, 260 calls | 679.9, 659.6 | 25.9, 25.7 | 26.0 | +| slow-rotation precompute | 777.8, 758.4 | 137.0, 137.2 | 5.6 | +| long-wavelength precompute | 43.4, 43.0 | 29.4, 29.8 | 1.5 | +| `ComputeModeIPTimeSeries`, 30 calls | 94.4, 95.0 | 93.6, 94.4 | 1.0 | +| `InnerProduct.__init__`, 290 calls | 2.1, 1.0 | 1.1, 1.0 | 1.5 | +| ILE wall | 917.3, 889.3 | 253.8, 253.0 | 3.57 | + +Replicate spread is 3.1% on the baseline arm and 0.3% on the batched arm, so the ratio +is not a single-run artefact. + +`lnL` was 769.5527137530971 in both baseline arms and 769.5527137530969 in both batched +arms: a difference of 2.3e-13 nats against a reported `sigma_lnL` of 0.1297. The sampler +took an identical path in all four runs, with the same `ntotal` (523199) and the same +`n_ESS` (73.189198442), so the difference is the precompute and nothing downstream of it. + +At production settings `InnerProduct.__init__` costs 2 s, not the 221 s in the campaign +profile. That 221 s was the inverse-spectrum-truncation window, which the campaign paid +because it took the driver default. + +The data-term overlaps in `ComputeModeIPTimeSeries` are now the largest remaining term, +at 94 s. They are FFT-bound and untouched here. + +## Enabling it + + export RIFT_PRECOMPUTE_BATCHED_CROSSTERMS=1 + +Submit files inherit it: RIFT's condor jobs default to `getenv = *`, the same route +`RIFT_ILE_GPU_FANOUT` uses. There is no CLI flag yet; adding one touches the ILE driver, which this change stays +out of. + +## Tried and rejected + +- Caching the stacked mode matrix across blocks. Stacking is 0.46 s of the 0.577 s + batched block, so a cache would give a further 4x, but one entry is 146 MB at l<=4 and + the combined path in PR #307 would hold 60 of them. +- One GEMM per detector over all basis elements at once, rather than per block. Same + arithmetic, fewer calls, but it requires changing every response module's caller + instead of the one function they share. +- A bit-identical batched form. Floating-point multiplication does not reassociate, so + no reordering of the reduction reproduces `np.sum` exactly. + +## Tests + +`test_precompute_crossterm_batching.py`, wired into `.travis/test-core-units.sh`. Three +of the four changes claim bit identity, so each is checked against a replay of the exact +code it replaced rather than a tolerance. The default-off behaviour is pinned by a +counter the batched path increments. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md index 5e8654894..2cb8aa084 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_q_window_stencil.md @@ -741,3 +741,103 @@ not a shared checkout, so a branch switch could not move code mid-run. Its fmin- reproduces #97's shipped numbers bit-for-bit, and the analysis code was validated by re-deriving #97's published bracket from the original 9 points alone. No row is reference-limited (per-stencil reference floors ≥ 400× below the smallest measured error; M→2M reference checks ≤ 5.7e-5 nats). + + +### 9.7 The JAX arm gained `--q-time-pregrid-factor` (2026-09-07) + +Refining the stored Q once, at build time, beats every choice of local stencil on +the coarse grid. PR #262 made the JAX driver refuse any factor but 1, because at +that point the arm had no refined-Q path. It has one now. + +`build_q_time_pregrid` calls `factored_likelihood.build_reflected_q_pregrid`, the +host-side builder #261 ships. Both arms therefore answer with the same refined Q, +and this one inherits that function's round-trip guard. + +Factor 1 remains the default and is bit-identical to base `bec19ad5`: 52 arrays +on a toy likelihood (four stencils, phase marginalization on and off, guard 0 and +8, over `fused_log_likelihood`, its `return_lnLt` form, both accumulator outputs +and `fused_log_likelihood_distmarg`) and 35 arrays from a rebuilt production +likelihood, all SHA256-equal. + +**Accuracy against an exact oracle.** `test/jax/test_jax_q_time_pregrid.py` +builds a band-limited series with a known finite Fourier sum, crops it as +`ComputeModeIPTimeSeries` crops `rhoTS`, and evaluates the truth at arbitrary +real times by direct summation. Production geometry: 1229-sample buffer, +positions about 300 samples clear of its ends. + +| stencil | relative max error | +|---|---| +| `nearest`, coarse | 4.88e-1 | +| `linear`, coarse | 1.98e-1 | +| `cubic`, coarse | 1.25e-1 | +| `sinc` a=8, coarse (production default) | 1.88e-2 | +| `cubic`, pregrid 2 | 1.09e-2 | +| `cubic`, pregrid 4 | 8.0e-4 | +| `cubic`, pregrid 8 | 4.62e-5 | +| `cubic`, pregrid 16 | 4.88e-6 | +| `cubic`, pregrid 32 | 3.46e-6 | +| `sinc` a=8, pregrid 8 | 4.86e-4 | + +Three decisions follow from that table. + +1. Cubic, not the arm's `sinc` default. A fixed 2a-tap Lanczos window does not + gain from a finer grid the way a fourth-order stencil does. On the same + factor-8 grid cubic is 10x more accurate at a quarter of the taps. The driver + selects `cubic` with the pregrid and refuses a different explicit stencil, as + conventional ILE does (#261). +2. Factor 8. The error falls about 16x per doubling (13.5x for 2 to 4, 17.4x for + 4 to 8), then saturates: 9.5x for 8 to 16, 1.4x for 16 to 32. Past about 8 the + residual is the reflection boundary condition, which no factor reduces. +3. Not the default. As in #261, promotion is a separate discussion. + +**The residual is a boundary error.** Error against clearance from the buffer +end, `cubic` on pregrid 8: 2.5e-3 at 8 samples, 8.3e-4 at 16, 2.4e-4 at 32, +1.0e-4 at 64, 4.7e-5 at 128, 4.0e-5 at 256. A fixture that gathers near the ends +measures the reflection while every assertion in it still passes. + +**Which reflection.** `jax_ile.core._reflected_fft_upsample` omits the duplicate +turning samples, period `2(n-1)`. `reflected_bandlimited_upsample` duplicates +them, period `2n`. The two docstrings each assert their own convention is +correct, and each is right about its own problem: `2(n-1)` reconstructs `kappa` +on the terminal integration window, where a series at Nyquist must keep +reconstructing `cos(pi t)`. For a cropped Q the `2n` form wins against the oracle +by 15.0x in the interior and 6.7x near the ends, matching the conventional arm's +independent finding in `DESIGN_time_marginalization_quadrature.md`. Routed +through `_reflected_fft_upsample` the factor-8 pregrid saturates at 7.8e-4 and +stops improving at factor 16 (7.6e-4), a 17x worse floor with no convergence. +`test_duplicated_reflection_is_the_right_one_for_a_crop` fails if a later change +reroutes it. + +**`nearest` is refused with a pregrid.** It would gather correctly. +`_accumulate_unit_banded` reconstructs the arrival time its post-phase applies as +`rint(p0)` in coarse samples, which is no longer the sample a refined-grid +nearest gather reads. The data term and the model norm would then disagree by up +to half a coarse bin. + +**Real rows.** Phase-marginalized JAX endpoint, 4096 Hz, SEOBNRv4 35+30, max over +6 rows, in nats, against a converged reference (see the PR). + +| rho | stencil, `sinc` a=8 coarse | stencil, `cubic` pregrid 8 | Simpson quadrature | +|---|---:|---:|---:| +| 40.77 | 0.15 | 3.2e-5 | 1.62 | +| 163.08 | 3.03 | 0.0005 | 33.96 | +| 652.31 | 27.31 | 0.0088 | 105.73 | + +The stencil term is removed. The quadrature term is untouched, because the +pregrid refines how Q is interpolated and not what is integrated. Above rho about +100 the time integral is limited by the quadrature rule alone. + +**Cost.** The pregrid is one host-side FFT, 0.03 s for a 3-detector 2-mode bank, +and it multiplies only the stored Q, 0.112 to 0.899 MiB here. On GPU it is faster +than the production default, because four taps replace sixteen. Matched at +S=20000, npts=614, interleaved, on an RTX PRO 4000 Blackwell: + +| case | s/eval | GPU peak in use | +|---|---:|---:| +| factor 1, `sinc` a=8 | 0.0823 | 1836.5 MiB | +| factor 1, `cubic` | 0.0160 | 257.3 MiB | +| factor 8, `cubic` | 0.0160 | 258.3 MiB | + +On CPU the ordering reverses and the pregrid costs 1.09x the production default: +0.2319, 0.2020 and 0.2520 s/eval at S=4000, from the strided gather's cache +behaviour. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_rotating_freqresponse.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_rotating_freqresponse.md new file mode 100644 index 000000000..def41eda1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_rotating_freqresponse.md @@ -0,0 +1,89 @@ +# Simultaneous slow rotation and finite-arm response + +## Scope + +`integrate_likelihood_extrinsic_batchmode --rotation-slow --freqresponse` now selects +one compound response model. It is intended for long, loud BNS-like signals, including +strongly precessing systems and higher modes. The response operators act on each full +inertial-frame mode, so the implementation makes no per-mode stationary-phase or +time-frequency-track approximation. + +This is not the economical path for short eccentric BBH mergers. Use `--freqresponse` +alone there unless Earth rotation is independently relevant. + +## Factorization + +The finite-arm response already has the form + +```text +F(f,t) = sum_b beta_b(t) W_b(f). +``` + +`beta_0` is the exact LAL long-wavelength response and `beta_(1+q)` contains an +arm-projection polynomial of order `q`. Under Earth rotation these coefficients have +finite sidereal half-widths 2 and `q+2`, respectively. A small exact DFT recovers their +Fourier coefficients. Composing the slow-delay expansion gives elementary templates + +```text +chi_(b,p,n) = M_n d_t^p [W_b h_lm]. +``` + +The coefficient half-width is `width(beta_b)+p`. Conjugation reflects only the sidereal +index, `(b,p,n) -> (b,p,-n)`, because each `W_b` is Hermitian. + +## Cost and controls + +The number of compound elements is + +```text +sum_b sum_(p=0)^pmax [2 (width(beta_b)+p) + 1]. +``` + +At the current defaults (`Qmax=4`, `pmax=0`) this is 50 elements and 2500 ordered U/V +pairs per detector. At `pmax=1` it is 112 elements and 12544 pairs. The driver prints +both counts before integration. Start with `--rotation-p-max 0` and the smallest +`--freqresponse-qmax` justified by a response-convergence check. + +## JAX implementation + +The JAX library now supports `feature="rotation_freqresponse"`. It consumes the same +production precompute and dense U/V bank as conventional ILE, evaluates the compound +coefficients in `jax.numpy`, reflects `(b,p,n) -> (b,p,-n)`, and applies the arrival-time +post-phase to both likelihood terms. The production JAX driver selects the individual or +compound banded builder from `--rotation-slow` and `--freqresponse`; these flags are no +longer compatibility no-ops. + +For the compound norm, JAX contracts one sidereal-difference bucket at a time. This avoids +materializing several full `(A,A,S)` arrays: peak pair scratch scales with the largest +harmonic bucket rather than all `A^2` pairs. The driver reports `A`, `A^2`, and persistent +device-bank storage before sampling. + +## Initial usability profile + +`test/jax/profile_rotating_freqresponse.py` profiles both conventional and JAX contractions +using a one-detector IMRPhenomD `(2,+/-2)` model. On one constrained CPU core, 32 extrinsic +samples, `pmax=0`, the measured progression was: + +| Qmax | A | pairs | precompute (s) | conventional eval (s) | JAX compile (s) | warm samples/s | +|---:|---:|---:|---:|---:|---:|---:| +| 0 | 10 | 100 | 0.30 | 0.006 | 4.26 | 77,100 | +| 1 | 17 | 289 | 0.82 | 0.006 | 4.90 | 44,800 | +| 2 | 26 | 676 | 1.93 | 0.010 | 6.65 | 28,000 | +| 3 | 37 | 1369 | 3.50 | 0.015 | 8.60 | 18,300 | +| 4 | 50 | 2500 | 6.51 | 0.021 | 13.0 | 9,640 | + +The `Qmax=4` warm rate is after memory bucketing; it trades about 15% CPU throughput for a +substantially smaller production-chunk temporary footprint. With delay drift, `pmax=1` +gave `A=24`/576 pairs/24,300 samples/s at `Qmax=0`, and `A=40`/1600 pairs/11,500 samples/s +at `Qmax=1`. These are microbenchmark numbers, not a detector-count or BNS-duration cost +model. The first science runs should use `pmax=0` and establish Qmax convergence; enable +`pmax=1` only after showing that delay drift changes the target high-SNR likelihood. + +## Validation + +`test_slowrot_rotating_freqresponse.py` checks the exact compound basis roster, sidereal +reconstruction of every finite-response coefficient, delay-order band support, and the +full precompute/NoLoop reduction to the existing finite-response likelihood at zero +sidereal rate, plus the Cauchy--Schwarz bound at zero and physical sidereal rates. The +existing `test_slowrot_noloop.py` remains the regression for the shared rotation +contraction. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md index c337106c1..52d1b2d59 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/DESIGN_time_marginalization_quadrature.md @@ -211,18 +211,23 @@ this rate exceeds the 11 GB card, which is itself worth knowing: | 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 +**The ratio triples between the measured 4,000 and the production 40,000, and the original +implementation explains why.** It is not only that the GPU baseline is nearly free. That +implementation re-doubled the refinement factor until the criterion held for the group minimum. +Ten times as many rows reach ten times deeper into the tail of that minimum, so the whole group +paid an extra octave: the factor histogram at the worst rung moved 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. +`_integrate_group` now retires each row as soon as its dense-grid remeasurement satisfies the +resolution criterion and doubles only the unresolved active set. The accuracy criterion and +128-MiB working-memory chunk remain unchanged; the returned factor histogram records the actual +per-row factors. The tables above predate that fix and are retained as the performance problem +the new production-SNR benchmark must remeasure, not as its expected post-fix cost. + **The affine, n=4,000 table, kept because it is what the CPU table compares against.** Same device, `--callback affine`, `rho_sq = 0`: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/Q_inner_product.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/Q_inner_product.py index 5b715944c..b0bbf9110 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/Q_inner_product.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/Q_inner_product.py @@ -54,7 +54,8 @@ def Q_inner_product_cupy(Q, A, start_indices, window_size): return out -def Q_inner_product_cubic_cupy(Q, A, start_indices, fractional_offsets, window_size): +def Q_inner_product_cubic_cupy(Q, A, start_indices, fractional_offsets, window_size, + time_stride=1): """Cubic-interpolated Q inner product for fractional detector-time offsets. ``start_indices`` are the integer floor indices of the first requested time @@ -98,7 +99,7 @@ def Q_inner_product_cubic_cupy(Q, A, start_indices, fractional_offsets, window_s 0, ) args = ( - Q, A, start_indices, fractional_offsets, window_size, + Q, A, start_indices, fractional_offsets, window_size, int(time_stride), num_time_points, num_extrinsic_samples, num_lms, out, ) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SLOWROT_HANDOFF.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SLOWROT_HANDOFF.md index 41d761f9a..5259f84c0 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SLOWROT_HANDOFF.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/SLOWROT_HANDOFF.md @@ -105,8 +105,9 @@ Path A + Path B are implemented, validated, and wired into the ILE. Next work i a finite-size injection the long-wavelength NoLoop deficit 2.71 -> finite-size 0.558 (converged Qmax=2), residual ~= peak-resolution floor; V3 Cauchy-Schwarz respected. Scalar companion agrees with the NoLoop to 0.156 (interp-vs-nearest floor). - REMAINING: ILE wiring (a --freqresponse flag mirroring --rotation-slow); full precessing+HM - (route a pinned-sky); a value demo (the finite-size effect only bites for CE/3G). + ILE wiring is complete for `--freqresponse` alone and for its composition with + `--rotation-slow`. REMAINING: a full long-duration precessing+HM science campaign; the + compound zero-rate reduction and Cauchy--Schwarz checks pass on a short waveform fixture. AUDIT NOTE (2026-07): all slowrot likelihoods route through the maintained NoLoop; there are NO SingleDetectorLogLikelihood calls; the ILE wires only NoLoop paths. The scalar FactoredLogLikelihood*/SingleDetectorLogLikelihood are used only as secondary references and @@ -179,8 +180,15 @@ End-to-end ILE head-to-head (ILE-GPU-Paper demo data), baseline vs rotation vs f integrate_likelihood_extrinsic_batchmode --vectorized --rotation-slow --rotation-p-max 1 ... # Path B integrate_likelihood_extrinsic_batchmode --vectorized --freqresponse \ --freqresponse-arm-length 40000 --freqresponse-qmax 6 ... # Path D (finite-size) + integrate_likelihood_extrinsic_batchmode --vectorized --rotation-slow --freqresponse \ + --rotation-p-max 0 --freqresponse-arm-length 40000 --freqresponse-qmax 4 ... # combined A+D # --interpolate-time selects cubic sub-bin time interpolation for all of the above (default nearest). +The combined path is intentionally scoped to long, loud BNS-like signals. Its compound +`(frequency basis, delay order, sidereal harmonic)` bank has 50 elements / 2500 ordered U/V +pairs at the current `Qmax=4,pmax=0` defaults and 112 / 12544 at `pmax=1`; see +`DESIGN_rotating_freqresponse.md` before increasing either order. + ## Validation status (all PASSING) - Response harmonics vs LAL: ~1e-16. FD ops vs LAL round trips: ~1e-13. - Path A scalar: V1a (Omega=0 vs baseline) 2.7e-12; V1b (real vs brute force) 2.6e-9. @@ -215,6 +223,15 @@ End-to-end ILE head-to-head (ILE-GPU-Paper demo data), baseline vs rotation vs f - Path D (finite-size, --freqresponse): response Sum_p b_p W_p == antenna_response_fd to 6e-11 on both +/-f; likelihood L->0 reduces to baseline NoLoop 3e-9; Cauchy-Schwarz respected; V4 positive control asserts finite-size beats LWL by +38.9 nats (15+13 Msun, fmax=2000, CE 40km). +- Combined Path A/B+D: every finite-response sky coefficient reconstructs under sidereal + evolution to 2.6e-12; the full compound precompute/NoLoop likelihood reduces to Path D at + zero sidereal rate to <1e-8 and respects the Cauchy--Schwarz bound at zero and physical + sidereal rates. Guarded by `test_slowrot_rotating_freqresponse.py`. +- JAX combined Path A/B+D: JAX-native compound coefficients agree with numpy to 2.7e-16; + the packed `Qmax=0,pmax=0` likelihood agrees with conventional NoLoop to 1.2e-14 + relative and executes under JIT/grad. The production JAX driver now wires the + rotation-only, frequency-response-only, and combined selections. Initial CPU basis + scaling is recorded in `DESIGN_rotating_freqresponse.md`. - Cubic time-interp (from calmarg_in_loop, --interpolate-time): both slow-response NoLoops now support time_interp='nearest'|'cubic'. Cubic exposed+fixed a sub-bin GPS-cancellation bug in the time reference; head-to-head regression floor 1.6e-3 -> 4.5e-13, test_slowrot_noloop 3.6e-12. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/bivariate_trig_stationary.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/bivariate_trig_stationary.py new file mode 100644 index 000000000..aaba1d2f2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/bivariate_trig_stationary.py @@ -0,0 +1,604 @@ +"""Algebraic stationary-point enumeration for real bivariate trig polynomials. + +This is the NumPy/SciPy reference enumerator for the finite ``(phi, 2 psi)`` +polynomials used by the higher-mode angle likelihood. It does not sample an +angle grid. Both derivatives are Laurent polynomials; after clearing the +Laurent powers and making a generic affine projection + + t = z + alpha w, z = exp(i phi), w = exp(i u), + +their common zeros are the zeros of a Sylvester resultant in ``t``. The +resultant is solved as a generalized polynomial eigenproblem. Every isolated +stationary point is therefore in a finite algebraic candidate set whose size is +fixed by the harmonic support, not by the likelihood amplitude. + +The input uses the RIFT coefficient-table convention: ``C[k, q + Q]`` stores +non-negative phi harmonics only and the real field is + + Re sum_{k=0..K,q=-Q..Q} (2 if k else 1) C[k,q] exp(i(k phi+q u)). + +Floating-point algebra cannot make an unconditional exact-root promise. The +enumeration certificate is consequently fail closed. A certified result requires: + +* the mixed-volume (BKK) number of isolated roots in ``(C*)^2``; +* nonsingular complex stationary Jacobians; +* an unambiguous unit-torus classification for every algebraic root; and +* identical torus stationary sets from two independent affine projections. + +An exact or near degeneracy, a projection collision, a root close enough to the +unit torus that its membership is numerically ambiguous, or a root-count deficit +returns ``ok=False``. Definite candidates remain available for a downstream +outside-cover bound; if that cannot certify their omitted impact the caller must +take its dense fallback. No tolerance silently discards a possible real mode or +likelihood sample. A certified call returns all isolated local maxima, including +exactly co-dominant symmetry-related maxima. + +This module is deliberately host-side. Generalized QZ, variable finite-root +counts, cross-projection matching, and fail-closed conditioning diagnostics do +not have an honest static-shape JAX transcription yet. A JAX adapter should +consume a host-built fixed-capacity plan; it must not replace this solve with a +sampled phi grid and call that enumeration. +""" + +from dataclasses import dataclass +from math import factorial + +import numpy as np +from scipy import linalg +from scipy.optimize import linear_sum_assignment + + +__all__ = [ + "StationaryPointEnumeration", + "canonical_laurent_table", + "stationary_mixed_volume", + "enumerate_torus_maxima", +] + + +def _binomial(n, k): + """Small exact binomial coefficient (keeps the reference Python-3.6 safe).""" + return factorial(n) // (factorial(k) * factorial(n - k)) + + +@dataclass(frozen=True) +class StationaryPointEnumeration: + """Result of :func:`enumerate_torus_maxima`. + + ``stationary_points`` contains every verified isolated stationary candidate + on the torus. ``points`` is its negative-definite-Hessian subset. When + ``ok`` is false these arrays may be partial and are targeting data only: a + caller may use them if an independent outside-cover bound passes, otherwise + it must take a dense fallback. ``report`` contains no fallback policy. + """ + + points: np.ndarray + hessians: np.ndarray + values: np.ndarray + stationary_points: np.ndarray + ok: bool + report: dict + + +def canonical_laurent_table(C): + """Return the full Hermitian Laurent table represented by a RIFT table. + + The returned array has shape ``(2*K+1, 2*Q+1)`` and indices ``(k+K,q+Q)``. + It obeys ``A[-k,-q] = conj(A[k,q])`` by construction, including the + potentially overlapping ``k=0`` contributions. + """ + C = np.asarray(C, dtype=np.complex128) + if C.ndim != 2 or C.shape[0] < 2 or C.shape[1] < 3 or C.shape[1] % 2 != 1: + raise ValueError("C must have shape (K+1, 2*Q+1) with K,Q >= 1") + if not np.all(np.isfinite(C.real) & np.isfinite(C.imag)): + raise ValueError("C must be finite") + K = C.shape[0] - 1 + Q = (C.shape[1] - 1) // 2 + A = np.zeros((2 * K + 1, 2 * Q + 1), dtype=np.complex128) + for k in range(K + 1): + weight = 1.0 if k == 0 else 2.0 + for iq, q in enumerate(range(-Q, Q + 1)): + a = 0.5 * weight * C[k, iq] + A[k + K, q + Q] += a + A[-k + K, -q + Q] += np.conj(a) + return A + + +def _convex_hull(points): + """Integer monotone-chain hull, without a numerical geometry tolerance.""" + pts = sorted(set(tuple(map(int, p)) for p in points)) + if len(pts) <= 1: + return pts + + def cross(o, a, b): + return ((a[0] - o[0]) * (b[1] - o[1]) + - (a[1] - o[1]) * (b[0] - o[0])) + + lower = [] + for p in pts: + while len(lower) >= 2 and cross(lower[-2], lower[-1], p) <= 0: + lower.pop() + lower.append(p) + upper = [] + for p in reversed(pts): + while len(upper) >= 2 and cross(upper[-2], upper[-1], p) <= 0: + upper.pop() + upper.append(p) + return lower[:-1] + upper[:-1] + + +def _twice_polygon_area(points): + hull = _convex_hull(points) + if len(hull) < 3: + return 0 + return abs(sum( + hull[i][0] * hull[(i + 1) % len(hull)][1] + - hull[(i + 1) % len(hull)][0] * hull[i][1] + for i in range(len(hull)))) + + +def _derivative_tables(A): + K = (A.shape[0] - 1) // 2 + Q = (A.shape[1] - 1) // 2 + k = np.arange(-K, K + 1)[:, None] + q = np.arange(-Q, Q + 1)[None, :] + return 1j * k * A, 1j * q * A + + +def stationary_mixed_volume(C): + """BKK count for the two stationary Laurent equations. + + This is the exact integer mixed volume of their Newton polygons. It is the + number of isolated roots in ``(C*)^2`` for a non-degenerate system, counted + with multiplicity, and an upper bound otherwise. + """ + A = canonical_laurent_table(C) + F, G = _derivative_tables(A) + K = (A.shape[0] - 1) // 2 + Q = (A.shape[1] - 1) // 2 + exponents = [(k, q) for k in range(-K, K + 1) + for q in range(-Q, Q + 1)] + sf = [p for p, c in zip(exponents, F.ravel()) if c != 0.0] + sg = [p for p, c in zip(exponents, G.ravel()) if c != 0.0] + # A derivative may have a one-dimensional Newton polytope without making + # the JOINT system one-dimensional: the separable field + # cos(m phi)+cos(n u) has two transverse segments and 4mn isolated roots. + if len(sf) < 2 or len(sg) < 2: + return 0 + hf = _convex_hull(sf) + hg = _convex_hull(sg) + summed = [(a[0] + b[0], a[1] + b[1]) for a in hf for b in hg] + twice = (_twice_polygon_area(summed) + - _twice_polygon_area(hf) - _twice_polygon_area(hg)) + if twice < 0 or twice % 2: + raise RuntimeError("stationary mixed volume was not a non-negative integer") + return twice // 2 + + +def _projected_polynomial(D, alpha): + """Coefficients in ``w,t`` after ``z=t-alpha*w`` and Laurent clearing.""" + K = (D.shape[0] - 1) // 2 + Q = (D.shape[1] - 1) // 2 + # After multiplying by z^K w^Q, z-degree is <=2K and w-degree <=2Q. + # Substitution can transfer all z degree to w. + out = np.zeros((2 * (K + Q) + 1, 2 * K + 1), dtype=np.complex128) + for iz in range(2 * K + 1): + for iw in range(2 * Q + 1): + c = D[iz, iw] + if c == 0.0: + continue + for it in range(iz + 1): + out[iw + iz - it, it] += ( + c * _binomial(iz, it) * ((-alpha) ** (iz - it))) + nz = np.nonzero(np.any(out != 0.0, axis=1))[0] + if nz.size == 0: + return np.zeros((0, 0), dtype=np.complex128) + out = out[nz[0]:nz[-1] + 1] + scale = np.max(np.abs(out)) + return out / scale if scale > 0.0 else out + + +def _sylvester_matrix_polynomial(F, G): + """Return ``S[j]`` for the Sylvester matrix polynomial ``sum t^j S[j]``.""" + if F.size == 0 or G.size == 0: + raise ValueError("an identically-zero stationary equation is degenerate") + m = F.shape[0] - 1 + n = G.shape[0] - 1 + if m < 1 or n < 1: + raise ValueError("projection produced an equation independent of the eliminated variable") + degree = max(F.shape[1], G.shape[1]) - 1 + size = m + n + S = np.zeros((degree + 1, size, size), dtype=np.complex128) + for shift in range(n): + for j in range(m + 1): + S[:F.shape[1], shift, shift + j] = F[j] + for shift in range(m): + for j in range(n + 1): + S[:G.shape[1], n + shift, shift + j] = G[j] + nz = np.nonzero(np.any(S != 0.0, axis=(1, 2)))[0] + if nz.size < 2: + raise ValueError("constant or zero resultant pencil") + return S[:nz[-1] + 1] + + +def _linearize_matrix_polynomial(S): + """First companion linearization ``L0 - t L1`` of ``sum S[j] t^j``.""" + degree = S.shape[0] - 1 + size = S.shape[1] + L0 = np.zeros((degree * size, degree * size), dtype=np.complex128) + L1 = np.zeros_like(L0) + eye = np.eye(size, dtype=np.complex128) + for i in range(degree - 1): + L0[i * size:(i + 1) * size, (i + 1) * size:(i + 2) * size] = eye + L1[i * size:(i + 1) * size, i * size:(i + 1) * size] = eye + last = slice((degree - 1) * size, degree * size) + for j in range(degree): + L0[last, j * size:(j + 1) * size] = -S[j] + L1[last, (degree - 1) * size:degree * size] = S[degree] + return L0, L1 + + +def _eval_laurent(D, z, w): + K = (D.shape[0] - 1) // 2 + Q = (D.shape[1] - 1) // 2 + zp = z ** np.arange(-K, K + 1) + wp = w ** np.arange(-Q, Q + 1) + return np.einsum("ij,i,j->", D, zp, wp) + + +def _laurent_scale(D, z, w): + K = (D.shape[0] - 1) // 2 + Q = (D.shape[1] - 1) // 2 + zp = np.abs(z) ** np.arange(-K, K + 1) + wp = np.abs(w) ** np.arange(-Q, Q + 1) + return float(np.einsum("ij,i,j->", np.abs(D), zp, wp)) + + +def _laurent_order(A, a, b): + K = (A.shape[0] - 1) // 2 + Q = (A.shape[1] - 1) // 2 + k = np.arange(-K, K + 1)[:, None] + q = np.arange(-Q, Q + 1)[None, :] + return ((1j * k) ** int(a)) * ((1j * q) ** int(b)) * A + + +def _laurent_newton(A, z, w, iterations=60): + """Newton in complex angle coordinates, avoiding cleared-power scaling.""" + Dp = _laurent_order(A, 1, 0) + Du = _laurent_order(A, 0, 1) + Dpp = _laurent_order(A, 2, 0) + Dpu = _laurent_order(A, 1, 1) + Duu = _laurent_order(A, 0, 2) + for _ in range(int(iterations)): + gradient = np.array([_eval_laurent(Dp, z, w), + _eval_laurent(Du, z, w)]) + H = np.array([[_eval_laurent(Dpp, z, w), + _eval_laurent(Dpu, z, w)], + [_eval_laurent(Dpu, z, w), + _eval_laurent(Duu, z, w)]]) + if not np.all(np.isfinite(H)) or np.linalg.cond(H) > 1e16: + return z, w, np.inf, 0.0, False + try: + step = np.linalg.solve(H, -gradient) + except np.linalg.LinAlgError: + return z, w, np.inf, 0.0, False + if not np.all(np.isfinite(step)) or np.max(np.abs(step)) > 4.0: + return z, w, np.inf, 0.0, False + z *= np.exp(1j * step[0]) + w *= np.exp(1j * step[1]) + if (not np.isfinite(z) or not np.isfinite(w) + or abs(z) < 1e-12 or abs(w) < 1e-12 + or max(abs(z), abs(w)) > 1e12): + return z, w, np.inf, 0.0, False + if np.max(np.abs(step)) < 5e-14: + break + rp = abs(_eval_laurent(Dp, z, w)) / max(_laurent_scale(Dp, z, w), 1e-300) + ru = abs(_eval_laurent(Du, z, w)) / max(_laurent_scale(Du, z, w), 1e-300) + H = np.array([[_eval_laurent(Dpp, z, w), _eval_laurent(Dpu, z, w)], + [_eval_laurent(Dpu, z, w), _eval_laurent(Duu, z, w)]]) + cond = float(np.linalg.cond(H)) if np.all(np.isfinite(H)) else np.inf + return z, w, max(float(rp), float(ru)), 1.0 / cond, True + + +def _solution_distance(a, b): + return max(abs(a[0] - b[0]) / max(1.0, abs(a[0]), abs(b[0])), + abs(a[1] - b[1]) / max(1.0, abs(a[1]), abs(b[1]))) + + +def _one_projection(A, alpha, expected, root_tol, jacobian_rcond_min): + Dp, Du = _derivative_tables(A) + F = _projected_polynomial(Dp, alpha) + G = _projected_polynomial(Du, alpha) + report = {"alpha": alpha, "expected_roots": int(expected), + "pencil_size": 0, "finite_eigenvalues": 0, + "verified_complex_roots": 0, "min_jacobian_rcond": 0.0, + "decline": None} + try: + S = _sylvester_matrix_polynomial(F, G) + L0, L1 = _linearize_matrix_polynomial(S) + report["pencil_size"] = int(L0.shape[0]) + eig, left, right = linalg.eig( + L0, L1, left=True, right=True, homogeneous_eigvals=True, + check_finite=False) + except (ValueError, linalg.LinAlgError) as exc: + report["decline"] = "singular resultant construction: %s" % exc + return [], report + + aa, bb = eig + pair_scale = np.hypot(np.abs(aa), np.abs(bb)) + finite = ((np.abs(bb) > 100.0 * np.finfo(float).eps * pair_scale) + & np.isfinite(aa) & np.isfinite(bb)) + report["finite_eigenvalues"] = int(np.count_nonzero(finite)) + bnorm = max(float(np.linalg.norm(L1, ord="fro")), 1e-300) + anorm = max(float(np.linalg.norm(L0, ord="fro")), 1e-300) + solutions = [] + jac_rconds = [] + eig_rconds = [] + eig_backward = [] + for idx in np.nonzero(finite)[0]: + t0 = aa[idx] / bb[idx] + if not np.isfinite(t0): + continue + St = sum(S[j] * (t0 ** j) for j in range(S.shape[0])) + y, x = left[:, idx], right[:, idx] + eig_rc = abs(np.vdot(y, L1 @ x)) / max( + np.linalg.norm(y) * np.linalg.norm(x) * bnorm, 1e-300) + eig_be = np.linalg.norm(L0 @ x - t0 * (L1 @ x)) / max( + (anorm + abs(t0) * bnorm) * np.linalg.norm(x), + 1e-300) + # The right null vector is a geometric sequence in the eliminated + # variable for a simple fibre. A projection collision makes its null + # space multidimensional; that is not guessed through with extra seeds + # but exposed by the BKK count / second-projection checks below. + # In this companion linearization the first block of the generalized + # eigenvector is already a null vector of S(t). It is jointly computed + # with t by QZ and is materially more accurate than recomputing the + # smallest singular vector at a rounded eigenvalue. Retain SVD only as + # a fallback for a zero first block. + v = right[:S.shape[1], idx] + if np.linalg.norm(v) == 0.0: + try: + _, _, vh = np.linalg.svd(St) + v = vh[-1].conj() + except np.linalg.LinAlgError: + continue + denom = np.vdot(v[:-1], v[:-1]) + if abs(denom) == 0.0: + continue + w0 = np.vdot(v[:-1], v[1:]) / denom + z0 = t0 - alpha * w0 + z, w, residual, jac_rcond, converged = _laurent_newton(A, z0, w0) + if not converged or residual > root_tol: + continue + if (not np.isfinite(z) or not np.isfinite(w) + or abs(z) < 1e-10 or abs(w) < 1e-10): + continue + rp = abs(_eval_laurent(Dp, z, w)) / max( + _laurent_scale(Dp, z, w), 1e-300) + ru = abs(_eval_laurent(Du, z, w)) / max( + _laurent_scale(Du, z, w), 1e-300) + if max(rp, ru) > 10.0 * root_tol: + continue + candidate = (z, w, residual, jac_rcond, float(eig_rc)) + close = [_solution_distance(candidate, old) for old in solutions] + if not close or min(close) > 5e-8: + solutions.append(candidate) + jac_rconds.append(jac_rcond) + eig_rconds.append(float(eig_rc)) + eig_backward.append(float(eig_be)) + + report["verified_complex_roots"] = len(solutions) + report["min_jacobian_rcond"] = float(min(jac_rconds, default=0.0)) + report["min_pencil_eigen_rcond"] = float(min(eig_rconds, default=0.0)) + report["max_pencil_backward_error"] = float(max(eig_backward, default=np.inf)) + if min(jac_rconds, default=0.0) < jacobian_rcond_min: + report["decline"] = "singular or ill-conditioned stationary Jacobian" + return solutions, report + if len(solutions) != expected: + report["decline"] = "BKK root-count mismatch (%d != %d)" % ( + len(solutions), expected) + return solutions, report + if max(eig_backward, default=np.inf) > root_tol: + report["decline"] = "resultant eigenproblem failed its backward-error check" + return solutions, report + return solutions, report + + +def _angle_eval(A, points, order=(0, 0)): + K = (A.shape[0] - 1) // 2 + Q = (A.shape[1] - 1) // 2 + k = np.arange(-K, K + 1)[:, None] + q = np.arange(-Q, Q + 1)[None, :] + a, b = order + factor = (1j * k) ** a * (1j * q) ** b + phi = points[:, 0, None, None] + u = points[:, 1, None, None] + phase = np.exp(1j * (phi * k[None] + u * q[None])) + return np.real(np.sum(phase * factor[None] * A[None], axis=(1, 2))) + + +def _torus_points(solutions, torus_on_tol, torus_off_tol): + """Classify roots using the real-field reciprocal-conjugate involution. + + A torus root is a fixed point of ``(z,w)->(1/conj(z),1/conj(w))``. + A genuinely complex root has a distinct partner. This is stronger than an + ``abs(abs(z)-1) < tol`` filter: a close off-torus pair is declared ambiguous + and declines the whole solve instead of being rounded onto or away from the + torus. + """ + points = [] + ambiguous = 0 + roots = [(s[0], s[1]) for s in solutions] + for i, (z, w) in enumerate(roots): + involution = (1.0 / np.conj(z), 1.0 / np.conj(w)) + distance = np.asarray([_solution_distance(involution, other) + for other in roots]) + order = np.argsort(distance) + nearest = int(order[0]) + match_error = float(distance[nearest]) + self_error = float(distance[i]) + if nearest == i and match_error <= torus_on_tol: + points.append((np.mod(np.angle(z), 2.0 * np.pi), + np.mod(np.angle(w), 2.0 * np.pi))) + elif (nearest != i and match_error <= torus_on_tol + and self_error >= torus_off_tol): + # A resolved non-real reciprocal-conjugate pair: safely off torus. + continue + else: + ambiguous += 1 + return np.asarray(points, dtype=float).reshape((-1, 2)), ambiguous + + +def _periodic_assignment_distance(a, b): + if len(a) != len(b): + return np.inf + if len(a) == 0: + return 0.0 + delta = (a[:, None, :] - b[None, :, :] + np.pi) % (2.0 * np.pi) - np.pi + cost = np.linalg.norm(delta, axis=-1) + row, col = linear_sum_assignment(cost) + return float(np.max(cost[row, col])) + + +def _dedupe_periodic(points, tolerance=1e-7): + keep = [] + for point in np.asarray(points, dtype=float).reshape((-1, 2)): + if not keep: + keep.append(point) + continue + delta = (np.asarray(keep) - point + np.pi) % (2.0 * np.pi) - np.pi + if np.min(np.linalg.norm(delta, axis=1)) > tolerance: + keep.append(point) + return np.asarray(keep, dtype=float).reshape((-1, 2)) + + +def enumerate_torus_maxima( + C, *, projections=(0.371 + 0.193j, -0.227 + 0.419j), + root_tol=2e-9, jacobian_rcond_min=2e-10, + torus_on_tol=2e-7, torus_off_tol=2e-5, + projection_match_tol=2e-6): + """Enumerate every isolated local maximum of ``g(phi,u)`` algebraically. + + Certification is conditional on a regular zero-dimensional stationary + system. ``ok=False`` is the promised behavior for exact/near stationary + degeneracy, ill-conditioned resultants, ambiguous torus membership, or + disagreement between the independent projections. Such a result may carry + definite best-effort targets, but never claims them as complete. + """ + C = np.asarray(C, dtype=np.complex128) + empty_p = np.zeros((0, 2), dtype=float) + empty_h = np.zeros((0, 2, 2), dtype=float) + empty_v = np.zeros(0, dtype=float) + report = {"ok": False, "mixed_volume": 0, "n_stationary": 0, + "n_maxima": 0, "projections": [], "decline": None} + try: + A = canonical_laurent_table(C) + expected = stationary_mixed_volume(C) + except (ValueError, RuntimeError) as exc: + report["decline"] = str(exc) + return StationaryPointEnumeration( + empty_p, empty_h, empty_v, empty_p, False, report) + report["mixed_volume"] = int(expected) + if expected <= 0: + report["decline"] = "stationary system is not two-dimensional" + return StationaryPointEnumeration( + empty_p, empty_h, empty_v, empty_p, False, report) + scale = float(np.max(np.abs(A))) + if not scale > 0.0: + report["decline"] = "constant field has a positive-dimensional stationary set" + return StationaryPointEnumeration( + empty_p, empty_h, empty_v, empty_p, False, report) + A = A / scale + + torus_sets = [] + complete_sets = [] + for alpha in projections: + solutions, one = _one_projection( + A, complex(alpha), expected, float(root_tol), + float(jacobian_rcond_min)) + report["projections"].append(one) + points, ambiguous = _torus_points( + solutions, float(torus_on_tol), float(torus_off_tol)) + one["torus_roots"] = int(len(points)) + one["ambiguous_torus_roots"] = int(ambiguous) + if ambiguous and one["decline"] is None: + one["decline"] = "ambiguous unit-torus root" + if len(points): + torus_sets.append(points) + if one["decline"] is None: + complete_sets.append(points) + + certified = False + if len(complete_sets) >= 2: + mismatch = _periodic_assignment_distance(complete_sets[0], complete_sets[1]) + report["projection_match_error"] = mismatch + certified = bool(np.isfinite(mismatch) and mismatch <= projection_match_tol) + if not certified: + report["decline"] = "independent projections disagree on torus roots" + else: + report["decline"] = "fewer than two algebraically complete projections" + if not torus_sets: + return StationaryPointEnumeration( + empty_p, empty_h, empty_v, empty_p, False, report) + + # On an uncertified solve keep the UNION of every definitely-on-torus root. + # A downstream cover bound can safely validate this best-effort targeting + # set; returning no candidates would force a dense fallback unnecessarily. + stationary = _dedupe_periodic(np.concatenate(torus_sets, axis=0)) + # Refine in real angles. Algebra supplies all seeds; Newton only restores + # unit-modulus/roundoff accuracy and never supplies completeness. + real_ok = np.ones(len(stationary), dtype=bool) + for _ in range(8): + gp = _angle_eval(A, stationary, (1, 0)) + gu = _angle_eval(A, stationary, (0, 1)) + gpp = _angle_eval(A, stationary, (2, 0)) + gpu = _angle_eval(A, stationary, (1, 1)) + guu = _angle_eval(A, stationary, (0, 2)) + for i in range(len(stationary)): + H = np.array([[gpp[i], gpu[i]], [gpu[i], guu[i]]]) + try: + step = np.linalg.solve(H, -np.array([gp[i], gu[i]])) + except np.linalg.LinAlgError: + real_ok[i] = False + continue + if not np.all(np.isfinite(step)) or np.linalg.norm(step) > 1.0: + real_ok[i] = False + continue + stationary[i] = np.mod(stationary[i] + step, 2.0 * np.pi) + + stationary = _dedupe_periodic(stationary[real_ok]) + if len(stationary) == 0: + if report["decline"] is None: + report["decline"] = "no usable real stationary candidates" + return StationaryPointEnumeration( + empty_p, empty_h, empty_v, empty_p, False, report) + + gp = _angle_eval(A, stationary, (1, 0)) + gu = _angle_eval(A, stationary, (0, 1)) + gpp = _angle_eval(A, stationary, (2, 0)) + gpu = _angle_eval(A, stationary, (1, 1)) + guu = _angle_eval(A, stationary, (0, 2)) + hessian = np.stack((np.stack((gpp, gpu), axis=-1), + np.stack((gpu, guu), axis=-1)), axis=-2) + eig_h = np.linalg.eigvalsh(hessian) + hscale = max(float(np.max(np.abs(eig_h))), 1e-300) + grad_resid = np.hypot(gp, gu) + report["max_stationary_residual"] = float(np.max(grad_resid, initial=0.0)) + usable = ((np.min(np.abs(eig_h), axis=1) > jacobian_rcond_min * hscale) + & (grad_resid <= 5e-8)) + if not np.all(usable): + certified = False + report["decline"] = "degenerate or unconverged real stationary candidate" + stationary = stationary[usable] + hessian = hessian[usable] + eig_h = eig_h[usable] + + is_max = np.all(eig_h < 0.0, axis=1) + maxima = stationary[is_max] + max_h = hessian[is_max] * scale + values = _angle_eval(A, maxima, (0, 0)) * scale + report["n_stationary"] = int(len(stationary)) + report["n_maxima"] = int(len(maxima)) + report["ok"] = bool(certified) + return StationaryPointEnumeration( + maxima, max_h, values, stationary, bool(certified), report) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/cuda_Q_inner_product.cu b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/cuda_Q_inner_product.cu index 8e0c9982c..a4b395838 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/cuda_Q_inner_product.cu +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/cuda_Q_inner_product.cu @@ -63,6 +63,7 @@ extern "C" { const int * index_start, const double * fractional_offset, int window_size, + int time_stride, int num_time_points, int num_extrinsic_samples, int num_lms, @@ -81,7 +82,7 @@ extern "C" { for (size_t i_time = t_idx; i_time < window_size; i_time+=blockDim.y) { size_t i_output = sample_idx*window_size + i_time; - int q_time = i_first_time + (int)i_time; + int q_time = i_first_time + (int)i_time*time_stride; double out_re = 0.0; double out_im = 0.0; diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py index 8334227ef..8ae687389 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood.py @@ -76,6 +76,9 @@ TIME_QUADRATURE_DEFAULT = 'simpson' from .vectorized_lal_tools import ComputeDetAMResponse,TimeDelayFromEarthCenter +from .vectorized_lal_tools import (SourcePolarizationBasis, SourcePropagationDirection, + ComputeDetAMResponsePrecomputed, + TimeDelayFromEarthCenterPrecomputed) import os if 'PROFILE' not in os.environ: @@ -202,6 +205,104 @@ def marginalization_time_grid(integration_window_half, deltaT, xpy=np): useNR=False distMpcRef = 1000 # a fiducial distance for the template source. +# Nodes used ONLY to locate the maximum of the band-limited psi exponent, for the max +# subtraction in NetworkLogLikelihoodPolarizationMarginalized. The exponent carries only +# harmonics 2 and 4 in psi, so this grid needs to resolve a quarter period, not the peak. +N_PSI_MAX_GRID = 256 + +# --- per-detector constants, cached across likelihood calls ------------------- +# DetectorPrefixToLALDetector() plus two host->device transfers of a 3-vector and a +# 3x3 matrix were being redone on EVERY likelihood evaluation, once per detector. +# The values are fixed properties of the interferometer, so cache them keyed by +# (prefix, backend). Keyed on id(xpy) rather than the module object so numpy and +# cupy arrays never get mixed. +_DETECTOR_GEOMETRY_CACHE = {} + + +def _detector_geometry(det, xpy): + """(location, response) for detector prefix ``det`` as ``xpy`` arrays, cached.""" + key = (det, id(xpy)) + hit = _DETECTOR_GEOMETRY_CACHE.get(key) + if hit is None: + detector = lalsim.DetectorPrefixToLALDetector(det) + hit = (xpy.asarray(detector.location), xpy.asarray(detector.response)) + _DETECTOR_GEOMETRY_CACHE[key] = hit + return hit + + +# --- Simpson quadrature weights, cached across likelihood calls --------------- +# The time integral is a FIXED linear functional at fixed dx, so simps(y) == y . w with +# w = simps(I). Evaluating it as one matrix-vector product reads the (npts_extrinsic, +# npts) integrand once, instead of the several strided slices and full-size temporaries +# the composite-Simpson implementation builds; measured 1.765 ms -> 0.049 ms at +# production shapes. npts and deltaT are fixed for a run, so the weights are built once. +# +# NOT bitwise against simps(): a gemv reassociates the summation. The RULE is identical +# -- the weights come from the very same simps implementation the call site would have +# used, so the even='avg'-vs-Cartwright distinction that separates the vendored GPU copy +# from scipy's is preserved, and only the order of the additions changes. The measured +# discrepancy is at the floating-point noise floor; see +# DESIGN_noloop_per_detector_glue.md for the number. +_SIMPS_WEIGHTS_CACHE = {} + + +def _simps_weights(simps, npts, deltaT, xpy, block=256): + """Quadrature weight vector w with simps(y, dx=deltaT, axis=-1) == y . w. + + Built a BLOCK OF ROWS AT A TIME rather than from a full (npts, npts) identity. + npts is 2*data_integration_window_half*srate, and the batch driver's DEFAULT srate + is 16384, so npts is 2457 in the default configuration, not the 614 of an + srate-4096 run: a whole identity is then 48 MB, and cupy's pool holds ~97 MB + across the simps call -- a transient that lands inside the first likelihood + evaluation of every --vectorized --gpu run, in a function whose n_chunk is already + bounded by device memory. Blocking caps it at block*npts*8 bytes (5 MB at the + default). simps reduces along axis=-1, so rows are independent and the blocked + result is bitwise identical to the whole-identity one. + + Keyed on the quadrature FUNCTION as well as (npts, dx, backend): on GPU `simps` is + the vendored old-scipy copy with even='avg' and on CPU it is scipy's Cartwright + form, and those two disagree by 0.405 nats on an under-resolved peak. Today the + rule is a pure function of the backend so id(xpy) would suffice, but this helper is + module-level and nothing stops a future caller passing a different rule at the same + shape -- which would silently serve the other rule's weights. + """ + key = (int(npts), float(deltaT), id(xpy), id(simps)) + w = _SIMPS_WEIGHTS_CACHE.get(key) + if w is None: + n = int(npts) + parts = [] + for lo in range(0, n, int(block)): + hi = min(lo + int(block), n) + rows = xpy.zeros((hi - lo, n), dtype=np.float64) + idx = xpy.arange(hi - lo) + rows[idx, idx + lo] = 1.0 + parts.append(simps(rows, dx=deltaT, axis=-1)) + w = xpy.concatenate(parts) if len(parts) > 1 else parts[0] + _SIMPS_WEIGHTS_CACHE[key] = w + return w + + +# --- mode-list identity, cached across likelihood calls ----------------------- +# The Ylm array depends only on (modes, inclination, phiref) -- NOT on the detector -- +# but was recomputed once per detector per call. To share it we need to know which +# detectors carry the same mode list, and lookupNKDict[det] may be a DEVICE array, so +# comparing it per call would force a synchronization. Instead memoize a hashable +# host-side key per array OBJECT. The array itself is kept in the cache so its id() +# cannot be recycled onto a different object while the entry lives; the dicts are built +# once per event by ILE, so this stays a handful of entries. +_MODE_KEY_CACHE = {} + + +def _mode_list_key(lms): + """Hashable host-side key identifying a mode list, memoized on the array object.""" + hit = _MODE_KEY_CACHE.get(id(lms)) + if hit is not None and hit[0] is lms: + return hit[1] + host = lms.get() if hasattr(lms, "get") else lms + key = tuple(map(tuple, np.asarray(host).tolist())) + _MODE_KEY_CACHE[id(lms)] = (lms, key) + return key + tWindowExplore = [-0.15, 0.15] # Not used in main code. Provided for backward compatibility for ROS. Should be consistent with t_ref_wind in ILE. rosDebugMessages = True rosDebugMessagesDictionary = {} # Mutable after import (passed by reference). Not clear if it can be used by caling routines @@ -967,7 +1068,7 @@ def NetworkLogLikelihoodPolarizationMarginalized(epoch,rholmsDictionary,crossTer for pair1 in rholmsDictionary[det]: for pair2 in rholmsDictionary[det]: term2a += F[det] * np.conj(F[det]) * ( crossTerms[det][(pair1,pair2)])* np.conj(Ylms[pair1]) * Ylms[pair2] - term2b += F[det]*F[det]*Ylms[pair1]*Ylms[pair2]*crossTermsV[(pair1,pair2)] #((-1)**pair1[0])*crossTerms[det][((pair1[0],-pair1[1]),pair2)] + term2b += F[det]*F[det]*Ylms[pair1]*Ylms[pair2]*crossTermsV[det][(pair1,pair2)] #((-1)**pair1[0])*crossTerms[det][((pair1[0],-pair1[1]),pair2)] term2a = -np.real(term2a) / 4. /(distMpc/distMpcRef)**2 term2b = -term2b/4./(distMpc/distMpcRef)**2 # coefficient of exp(-4ipsi) @@ -981,12 +1082,26 @@ def NetworkLogLikelihoodPolarizationMarginalized(epoch,rholmsDictionary,crossTer if False: #xgterm2a+np.abs(term2b)+np.abs(term1)>100: return term2a+ np.log(special.iv(0,np.abs(term1))) # an approximation, ignoring term2b entirely! else: - # marginalize over phase. Ideally done analytically. Only works if the terms are not too large -- otherwise overflow can occur. - # Should probably implement a special solution if overflow occurs + # marginalize over psi. The exponent is band-limited in psi -- exactly the harmonics + # 2 and 4 -- so a coarse grid locates its maximum to far better than a nat, and that + # is all the subtraction below needs. WITHOUT the subtraction exp() overflows to inf + # (and log(inf) -> nan) as soon as max_psi lnL exceeds ln(DBL_MAX) = 709, a network + # SNR near 38: measured nan at max lnL 725 before this was added. + _psi_max_grid = np.arange(N_PSI_MAX_GRID)*np.pi/N_PSI_MAX_GRID + _expon_grid = term2a + np.real(term2b*np.exp(-4.j*_psi_max_grid) + + term1*np.exp(+2.j*_psi_max_grid)) + _i_max = int(np.argmax(_expon_grid)) + expon_max = float(_expon_grid[_i_max]) def fnIntegrand(x): - return np.exp( term2a+ np.real(term2b*np.exp(-4.j*x)+ term1*np.exp(+2.j*x)))/np.pi # remember how the two terms enter -- note signs! - LmargPsi = integrate.quad(fnIntegrand,0,np.pi,limit=100,epsrel=1e-4)[0] - return np.log(LmargPsi) + # remember how the two terms enter -- note signs! + return np.exp( term2a+ np.real(term2b*np.exp(-4.j*x)+ term1*np.exp(+2.j*x)) - expon_max)/np.pi + # epsabs=0 is required, not cosmetic: after the subtraction the integral is the peak + # WIDTH, ~1/rho, so quad's default epsabs=1.49e-8 is met by a coarse rule that has not + # resolved the peak at all. Measured at exponent max 4194: default epsabs returns an + # answer 20 nat low. 'points' pins the first subdivision at the located maximum. + LmargPsi = integrate.quad(fnIntegrand,0,np.pi,points=[float(_psi_max_grid[_i_max])], + limit=200,epsabs=0,epsrel=1e-10)[0] + return np.log(LmargPsi) + expon_max def SingleDetectorLogLikelihood(rholm_vals, crossTerms,crossTermsV, Ylms, F, dist): """ @@ -1169,8 +1284,22 @@ def InterpolateRholms(rholms, t,verbose=False): return rholm_intp +# --- opt-in batching of the mode cross terms --------------------------------------- +# OFF by default: the batched path reorders the frequency reduction and so does not +# reproduce the shipped rounding bit-for-bit (~1e-15 of max|U|; see +# DESIGN_precompute_crossterm_batching.md). Enable per run with +# RIFT_PRECOMPUTE_BATCHED_CROSSTERMS=1 +# or per call with the `batched=` keyword, which overrides the environment. +_CROSSTERM_BATCH_CALLS = [0] # observability: a run that never took the path reads 0 + + +def _crossterm_batched_default(): + return os.environ.get("RIFT_PRECOMPUTE_BATCHED_CROSSTERMS", "0").strip() in ("1", "true", "True") + + def ComputeModeCrossTermIP(hlmsA, hlmsB, psd, fmin, fMax, fNyq, deltaF, - analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0., verbose=True,prefix="U",same_waveform_Q=False): + analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0., verbose=True,prefix="U",same_waveform_Q=False, + batched=None): """ Compute the 'cross terms' between waveform modes, i.e. < h_lm | h_l'm' >. @@ -1189,6 +1318,37 @@ def ComputeModeCrossTermIP(hlmsA, hlmsB, psd, fmin, fMax, fNyq, deltaF, crossTerms = {} + if batched is None: + batched = _crossterm_batched_default() + if batched: + # One GEMM over the whole (mode x mode) block instead of Na*Nb separate full-length + # reductions. The symmetry bookkeeping below is kept identical to the loop, so the + # exact Hermitian/transpose relations the shipped path guarantees still hold. + _CROSSTERM_BATCH_CALLS[0] += 1 + modesA = list(hlmsA.keys()) + modesB = list(hlmsB.keys()) + M = IP.ip_matrix([hlmsA[m] for m in modesA], [hlmsB[m] for m in modesB]) + if same_waveform_Q: + assert modesA == modesB, "same_waveform_Q requires the same mode keys on both sides" + for i, mode in enumerate(modesA): + crossTerms[ (mode,mode) ] = M[i,i] + for i, mode1 in enumerate(modesA): + for j in range(i+1, len(modesA)): + mode2 = modesA[j] + crossTerms[ (mode1,mode2) ] = M[i,j] + if prefix == "V": + crossTerms[ (mode2,mode1) ] = crossTerms[(mode1,mode2)] + else: + crossTerms[ (mode2,mode1) ] = np.conj(crossTerms[(mode1,mode2)]) + else: + for i, mode1 in enumerate(modesA): + for j, mode2 in enumerate(modesB): + crossTerms[ (mode1,mode2) ] = M[i,j] + if verbose: + print(" : ", prefix, " populated ", (mode1, mode2), " = ",\ + crossTerms[(mode1,mode2) ]) + return crossTerms + if same_waveform_Q: mode_list = list(hlmsA.keys()) pairs = combinations(mode_list, 2) # all pairs, no diagonal terms, sorted order ! @@ -2198,7 +2358,113 @@ def _factored_lnL_helper(kappa_sq, rho_sq): return kappa_sq - 0.5 * rho_sq -def _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts): +def build_reflected_q_pregrid(rholms, factor=8, xpy=np): + """Build a one-time finer Q grid without changing the likelihood time grid. + + The finite cut Q window is reflected before FFT interpolation so its unlike + endpoints are never identified. Only the forward interval is retained; + consequently the epoch is unchanged and every ``factor``-th sample must + reproduce the input. This helper is intentionally opt-in at the driver. + """ + factor = int(factor) + if factor < 1: + raise ValueError("Q pregrid factor must be positive") + if factor == 1: + return rholms, dict(factor=1, input_bytes=int(rholms.nbytes), + output_bytes=int(rholms.nbytes), roundtrip_max=0.0) + retained_view = time_quadrature_module.reflected_bandlimited_upsample( + xpy.asarray(rholms), factor, xpy=xpy) + # reflected_bandlimited_upsample returns a short VIEW into the full 2*N*factor + # inverse FFT. Copy it so retaining the useful forward interval does not pin + # the much larger backing allocation for the whole ILE run. + dense = xpy.array(retained_view, copy=True) + del retained_view + scale = float(xpy.max(xpy.abs(rholms))) + mismatch = float(xpy.max(xpy.abs(dense[..., ::factor] - rholms))) + relative = mismatch / scale if scale else mismatch + if not np.isfinite(relative) or relative > 5e-12: + raise RuntimeError("Q pregrid round-trip failed: %.3g" % relative) + full_dense_bytes = int(rholms.nbytes)*2*factor + peak_bytes = (int(rholms.nbytes)*4 + 2*full_dense_bytes + int(dense.nbytes)) + return dense, dict(factor=factor, input_bytes=int(rholms.nbytes), + retained_bytes=int(dense.nbytes), output_bytes=int(dense.nbytes), + peak_allocation_bytes=peak_bytes, roundtrip_max=relative) + + +def prepare_reflected_q_pregrid(rholms_by_detector, factor=8, transfer=None, + cleanup=None): + """Transactionally build and optionally transfer a detector Q pregrid. + + A backend OOM after one detector transfer cannot leave a mixed host/device, + coarse/fine dictionary. Partial temporaries are dropped, ``cleanup`` is + invoked (normally CuPy's memory-pool release), and the original coarse Q + dictionary is transferred instead. The caller can then restore its prior + stencil and continue with an explicit fallback telemetry record. + """ + original = dict(rholms_by_detector) + transfer = (lambda value: value) if transfer is None else transfer + prepared = {} + reports = [] + try: + host_fine = {} + for det, values in original.items(): + host_fine[det], report = build_reflected_q_pregrid(values, factor=factor) + report['detector'] = det + reports.append(report) + for det, values in host_fine.items(): + prepared[det] = transfer(values) + return prepared, reports, None + except Exception as error: + allocation_failure = (isinstance(error, (MemoryError, RuntimeError)) or + error.__class__.__name__ == 'OutOfMemoryError') + if not allocation_failure: + raise + failure = dict(type=error.__class__.__name__, repr=repr(error)) + prepared.clear() + reports[:] = [] + try: + host_fine.clear() + except UnboundLocalError: + pass + if cleanup is not None: + cleanup() + fallback = {} + for det, values in original.items(): + fallback[det] = transfer(values) + # Never return ``error`` itself: its traceback retains this frame and + # therefore the last expanded host Q array that triggered backend OOM. + return fallback, reports, failure + + +def _q_sample_positions(t_det, tvals, integration_delta_t, q_delta_t, + time_interp, explicit_time_values, xpy=np): + """Map geocentric integration nodes onto an independently spaced Q grid.""" + q_delta_t = float(q_delta_t) + integration_delta_t = float(integration_delta_t) + if q_delta_t <= 0 or integration_delta_t <= 0: + raise ValueError("time-grid spacings must be positive") + separate_grid = not np.isclose(q_delta_t, integration_delta_t, + rtol=0.0, atol=1e-15*integration_delta_t) + ratio = integration_delta_t/q_delta_t + stride = int(round(ratio)) if separate_grid else 1 + regular_stride = (not separate_grid or + abs(ratio - stride) <= 1e-12*max(1.0, abs(ratio))) + per_time = bool(explicit_time_values or not regular_stride) + if per_time: + samples = ((t_det[:, None] + xpy.asarray(tvals)[None, :]) / q_delta_t) + else: + samples = (t_det + tvals[0]) / q_delta_t + if time_interp == 'nearest': + starts = (xpy.rint(samples) + 0.5).astype(np.int32) + fractions = None + else: + starts = xpy.floor(samples).astype(np.int32) + fractions = (samples - xpy.floor(samples)).astype(np.float64) + return starts, fractions, per_time, stride + + +def _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts, + time_stride=1): """Return cubic-interpolated Q windows with zero extension. Q_block has shape (n_time, n_lm). The returned array has shape @@ -2212,7 +2478,7 @@ def _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts): tgrid = np.arange(npts) n_time = Q_block.shape[0] for i in range(npts_extrinsic): - idxs = int(start_indices[i]) + tgrid + idxs = int(start_indices[i]) + tgrid*int(time_stride) u = float(fractional_offsets[i]) u2 = u*u u3 = u2*u @@ -2382,7 +2648,7 @@ def validate_time_interp(time_interp, on_gpu=False): def _q_window_numpy_interp(Q_block, start_indices, fractional_offsets, npts, time_interp, - xpy=np): + xpy=np, time_stride=1): """CPU Q-window dispatch. start_indices must already match the stencil: 'nearest' rounds, the interpolating stencils floor and carry the fractional part separately.""" if time_interp == 'nearest': @@ -2390,7 +2656,8 @@ def _q_window_numpy_interp(Q_block, start_indices, fractional_offsets, npts, tim if time_interp == 'sinc': return _sinc_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts) if time_interp == 'cubic': - return _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts) + return _cubic_Q_window_numpy(Q_block, start_indices, fractional_offsets, npts, + time_stride=time_stride) # Named explicitly rather than falling through to cubic. A bare `return cubic` here would # reinstate exactly the silent-wrong-stencil behaviour this work exists to remove: callers # reaching the dispatcher directly (the tests do) would get cubic for a typo and never find @@ -2399,7 +2666,8 @@ def _q_window_numpy_interp(Q_block, start_indices, fractional_offsets, npts, tim % (time_interp, TIME_INTERP_CHOICES)) -def _q_inner_product_gpu(Q, A, start_indices, fractional_offsets, npts, time_interp): +def _q_inner_product_gpu(Q, A, start_indices, fractional_offsets, npts, time_interp, + time_stride=1): """GPU Q-product dispatch: the device-side counterpart of _q_window_numpy_interp. Same stencil contract as the CPU dispatch, deliberately: the four GPU call sites (here x2, @@ -2414,7 +2682,8 @@ def _q_inner_product_gpu(Q, A, start_indices, fractional_offsets, npts, time_int Q, A, start_indices, fractional_offsets, npts) if time_interp == 'cubic': return Q_inner_product.Q_inner_product_cubic_cupy( - Q, A, start_indices, fractional_offsets, npts) + Q, A, start_indices, fractional_offsets, npts, + time_stride=time_stride) # Explicit, for the same reason as the CPU dispatcher above: no silent fallthrough to cubic. raise ValueError("unknown time_interp %r; expected one of %r" % (time_interp, TIME_INTERP_CHOICES)) @@ -2483,7 +2752,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,time_quadrature=None,explicit_time_values=False,return_time_draw=False,time_draw_uniforms=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,q_deltaT=None): """ DiscreteFactoredLogLikelihoodViaArray uses the array-ized data structures to compute the log likelihood, either as an array vs time *or* marginalized in time. @@ -2534,6 +2803,17 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic Distance-marginalization table+params for the fused distmarg kernel; see RIFT.likelihood.Q_fused_calmarg.Q_fused_calmarg_distmarg_cupy. + loglikelihood : callable(kappa_sq, rho_sq) -> lnL(t) + MUST NOT WRITE INTO ``rho_sq``. rho_sq is time-independent, so it is passed as + a stride-0 ``broadcast_to`` view over an ``(npts_extrinsic,)`` vector rather than + as a materialized ``(npts_extrinsic, npts)`` array. Every ``npts`` column + therefore aliases one address: on CPU an in-place write raises + ``ValueError: output array is read-only``, but on GPU cupy's broadcast is + WRITABLE and an in-place write races, giving a wrong and irreproducible answer + with no error. Every in-tree callback allocates (``_factored_lnL_helper`` and the + driver's ``distmarg_loglikelihood``), so nothing is broken today; a caller + supplying its own must allocate too, or call ``xpy.ascontiguousarray`` first. + time_interp : {'nearest', 'cubic', 'sinc'} Detector-time sampling convention for the data term. 'nearest' preserves the historical NoLoop integer-bin gather. 'cubic' evaluates @@ -2672,6 +2952,16 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic deltaT = float(P_vec.deltaT) # this is stored as a scalar + q_deltaT = (float(getattr(P_vec, 'q_deltaT', deltaT)) + if q_deltaT is None else float(q_deltaT)) + if q_deltaT <= 0: + raise ValueError("q_deltaT must be positive") + if q_deltaT != deltaT and n_cal != 1: + raise NotImplementedError("an independently spaced Q pregrid is not implemented for calibration marginalization") + if q_deltaT != deltaT and time_interp != 'cubic': + raise NotImplementedError( + "an independently spaced Q pregrid currently implements only the " + "strided cubic gather; nearest/sinc would silently use the wrong stride") # Convert tref to greenwich mean sidereal time @@ -2687,8 +2977,23 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # Used to accumulate kappa^2 and rho^2 over all detectors. They are just # the sum in quadrature of the individual detector contributions. - kappa_sq = xpy.zeros((npts_extrinsic, npts), dtype=np.complex128) - rho_sq = xpy.zeros((npts_extrinsic, npts), dtype=np.float64) + # kappa_sq is the (npts_extrinsic, npts) data term: 98 MB of complex128 at production + # shapes, and the single most expensive thing in this function. It used to be + # zero-filled and then read-modify-written once per detector, with the distance scaling + # allocating a further full-size temporary each time. Start from the first detector's + # own output buffer and scale it in place instead: same arithmetic, three fewer + # full-size passes over 98 MB for a three-detector network. + kappa_sq = None + # rho_sq is the term. It is TIME-INDEPENDENT: every detector contributes + # rho_sq_det of shape (npts_extrinsic,), which used to be broadcast into a dense + # (npts_extrinsic, npts) accumulator. At production shapes that is ~49 MB of float64 + # zero-filled once and read-modify-written once per detector, to store npts identical + # copies of each value. Accumulate the vector instead and expose the 2-D shape as a + # stride-0 view after the loop; downstream arithmetic is elementwise and sees no + # difference, and the additions happen in the same order on the same scalars, so the + # result is bitwise unchanged. (The calibration path already did exactly this with + # broadcast_to for rho_sq_cal; this brings the ordinary path in line.) + rho_sq_vec = xpy.zeros(npts_extrinsic, dtype=np.float64) # When marginalizing over calibration (n_cal>1), cache the per-detector data # term inputs here; the calibration-independent rho_sq is still accumulated @@ -2716,12 +3021,32 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic else: raise NotImplementedError("Backend not supported: {}".format(xpy)) + # ---- source-only geometry: built ONCE, shared by every detector ------------- + # None of this depends on the interferometer, only on the extrinsic samples, but it + # used to be rebuilt inside the detector loop. At production n_extrinsic these are + # small arrays, so the cost is launch-bound: ~30 kernels per detector, all but the + # response/location contractions redundant. The per-detector calls below consume + # these and perform exactly the same contractions as before, so results are bitwise + # unchanged. + XY_basis = SourcePolarizationBasis( + RA, DEC, psi, greenwich_mean_sidereal_time_tref, + xpy=xpy, + ) + ehat_src = SourcePropagationDirection( + RA, DEC, float(greenwich_mean_sidereal_time_tref), + xpy=xpy, + ) + + # Ylm depends on (modes, incl, phiref) only. Detectors that share a mode list -- + # in practice all of them, since the modes come from one waveform -- share the + # array. Keyed by mode list so a genuinely heterogeneous dict still gets a correct + # (merely unshared) result rather than a wrong shared one. + _ylm_by_modes = {} + # strings right now - need to change to make ufunc-able for det in detectors: - # Compute the detector's location and response matrix - detector = lalsim.DetectorPrefixToLALDetector(det) - detector_location = xpy.asarray(detector.location) - detector_response = xpy.asarray(detector.response) + # Compute the detector's location and response matrix (cached; fixed per IFO) + detector_location, detector_response = _detector_geometry(det, xpy) # These do not depend on extrinsic params. # Arrays of shape (n_lms, n_lms). @@ -2734,19 +3059,27 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # These do depend on extrinsic params # Array of shape (npts_extrinsic, n_lms,) - Ylms_vec = SphericalHarmonicsVectorized( - lms, incl, -phiref, - xpy=xpy, - l_max=Lmax, - ) + _mode_key = _mode_list_key(lms) + Ylms_vec = _ylm_by_modes.get(_mode_key) + if Ylms_vec is None: + Ylms_vec = SphericalHarmonicsVectorized( + lms, incl, -phiref, + xpy=xpy, + l_max=Lmax, + ) + _ylm_by_modes[_mode_key] = Ylms_vec + if phase_marginalization: + # The phase-marginalization branch below CONJUGATES Ylms_vec in place, and + # rho_sq_det above must see the un-conjugated array. Hand each detector its + # own copy so sharing cannot leak one detector's conjugation into the next + # detector's self-term. A copy of an (n_extrinsic, n_lms) array is still far + # cheaper than rebuilding the harmonics. + Ylms_vec = Ylms_vec.copy() # Array of shape (npts_extrinsic,) # F_vec_old = xpy.asarray(lalF(det, RA, DEC, psi, tref)) - F_vec = ComputeDetAMResponse( - detector_response, - RA, DEC, psi, - greenwich_mean_sidereal_time_tref, - xpy=xpy + F_vec = ComputeDetAMResponsePrecomputed( + detector_response, XY_basis[0], XY_basis[1], xpy=xpy, ) # Scalar -- is constant for each IFO @@ -2756,29 +3089,12 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # Note that to save on precision compared to ...NoLoopOrig, we CHANGE the t_det definition to be relative to the IFO statt time t_ref # ... this means we don't keep a 1e9 out in front, so we have more significant digits in the event time (and can if needed reduce precision in GPU ops) # an array of shape (npts_extrinsic,) - t_det = float(tref - float(t_ref)) + TimeDelayFromEarthCenter( - detector_location, RA, DEC, - float(greenwich_mean_sidereal_time_tref), - xpy=xpy + t_det = float(tref - float(t_ref)) + TimeDelayFromEarthCenterPrecomputed( + detector_location, ehat_src, xpy=xpy, ) - 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: - 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) + ifirst, frac_first, _q_per_time, _q_time_stride = _q_sample_positions( + t_det, tvals, deltaT, q_deltaT, time_interp, + explicit_time_values, xpy=xpy) # ilast = ifirst + npts @@ -2879,30 +3195,29 @@ 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) - if explicit_time_values: + if _q_per_time: 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) + Q, FY_conj, ifirst, frac_first, npts, time_interp, + time_stride=_q_time_stride) else: # Use old code completely unchanged ... very wasteful on memory management! - Q_block = rholmsArrayDict[det].T - if explicit_time_values: + Q_block = (Q if phase_marginalization and _q_per_time + else rholmsArrayDict[det].T) + if _q_per_time: 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") + xpy=xpy, time_stride=_q_time_stride) + if phase_marginalization and not _q_per_time: Qlms[:, :, 1] = xpy.conj(Qlms[:, :, 1]) - if not explicit_time_values: + if not _q_per_time: FY_dummy_t = np.broadcast_to( (F_vec_dummy_lm * Ylms_vec)[:, np.newaxis], Qlms.shape, @@ -2913,7 +3228,14 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic np.conj(FY_dummy_t), Qlms, ) - kappa_sq += Q_prod_result * (distMpcRef/distMpc)[..., np.newaxis] + # Scale in place into the buffer the Q kernel just handed us -- it is freshly + # allocated per detector and not aliased anywhere -- rather than allocating a + # full-size temporary for the product. + xpy.multiply(Q_prod_result, invDistMpc[..., np.newaxis], out=Q_prod_result) + if kappa_sq is None: + kappa_sq = Q_prod_result + else: + kappa_sq += Q_prod_result else: # ---- calibration-marginalization path (Option B): cache pieces ---- # The rholm timeseries hold n_cal contiguous realizations; realization c @@ -2934,7 +3256,7 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # Accumulate term2 into the time-dependent log likelihood. # Have to create a view with an extra axis so they broadcast. - rho_sq += rho_sq_det[..., np.newaxis] + rho_sq_vec += rho_sq_det # lnL_t_accum += term2[..., np.newaxis] # print lnL_t_accum.shape, lnL_t.shape @@ -2942,10 +3264,24 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # lnL_t_accum += lnL_t + # The (npts_extrinsic, npts) shape every consumer below expects, as a stride-0 view + # over the vector accumulated above. Consumers that need real backing memory -- the + # fused CUDA kernels, which index raw device pointers, and the non-Simpson quadrature + # helpers, which are free to write -- go through _dense_rho_sq() and pay exactly what + # they paid before. + rho_sq = xpy.broadcast_to(rho_sq_vec[:, np.newaxis], (npts_extrinsic, npts)) + + def _dense_rho_sq(a): + """A writable, contiguous copy of a possibly stride-0 rho_sq view.""" + return a if getattr(a, "flags", None) is not None and a.flags.c_contiguous \ + else xpy.ascontiguousarray(a) + if n_cal == 1: # Fused-calmarg self-term fix also applies to a SINGLE calibration draw: the data # carries C_0, so its self-term is rho_sq_c = = rho_sq_cal[0], not the # cal-independent . Falls back to rho_sq for the ordinary (no-cal) likelihood. + if kappa_sq is None: # no detectors: preserve the old all-zeros behaviour + kappa_sq = xpy.zeros((npts_extrinsic, npts), dtype=np.complex128) rho_sq_here = rho_sq if not _use_rho_sq_cal else xpy.broadcast_to(rho_sq_cal[0][:, np.newaxis], (npts_extrinsic, npts)) if phase_marginalization: lnL_t = loglikelihood(xpy.abs(kappa_sq), rho_sq_here) @@ -2983,7 +3319,7 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # 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, + kappa_sq, _dense_rho_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) @@ -3001,13 +3337,14 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # 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, + kappa_sq, _dense_rho_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) + # simps(L_t, dx, axis=-1) as a single matrix-vector product; see _simps_weights. + L = L_t.dot(_simps_weights(simps, npts, deltaT, xpy)) # 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. @@ -3056,7 +3393,7 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic N_window_block = cal_cache[dets[0]][3] # Simpson quadrature weight vector (incl. dx=deltaT), so time integration # matches the loop path's simps() exactly. simps is linear -> weights = simps(I). - w_t = simps(xpy.eye(npts, dtype=np.float64), dx=deltaT, axis=-1) + w_t = _simps_weights(simps, npts, deltaT, xpy) # invDistMpc is a scalar when distance is marginalized (P.dist fixed at the # fiducial) and a vector when distance is sampled; the kernel wants one value # per extrinsic sample, so broadcast to (npts_extrinsic,). @@ -3067,17 +3404,17 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic if xpy is np: # CPU: pure-numpy fused (no CUDA); independent cross-check of the kernel return Q_fused_calmarg.Q_fused_calmarg_numpy( - Q_stack, A_stack, ifirst_stack, invDist_vec, rho_sq, w_t, + Q_stack, A_stack, ifirst_stack, invDist_vec, _dense_rho_sq(rho_sq), w_t, n_cal, N_window_block, distmarg=cal_distmarg, cal_log_weights=cal_log_weights, phase_marginalization=phase_marginalization, rho_sq_cal=rho_sq_cal) if cal_distmarg is None: return Q_fused_calmarg.Q_fused_calmarg_cupy( - Q_stack, A_stack, ifirst_stack, invDist_vec, rho_sq, w_t, + Q_stack, A_stack, ifirst_stack, invDist_vec, _dense_rho_sq(rho_sq), w_t, n_cal, N_window_block, cal_log_weights=cal_log_weights, phase_marginalization=phase_marginalization, rho_sq_cal=rho_sq_cal) else: return Q_fused_calmarg.Q_fused_calmarg_distmarg_cupy( - Q_stack, A_stack, ifirst_stack, invDist_vec, rho_sq, w_t, + Q_stack, A_stack, ifirst_stack, invDist_vec, _dense_rho_sq(rho_sq), w_t, n_cal, N_window_block, cal_distmarg, cal_log_weights=cal_log_weights, phase_marginalization=phase_marginalization, rho_sq_cal=rho_sq_cal) @@ -3137,7 +3474,8 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # RAW per-realization time-integrated log L (no importance weight), stable: # log( simps_t exp(lnL_t,c) ) = m + log( simps_t exp(lnL_t,c - m) ) m_raw = xpy.max(lnL_t_c, axis=-1, keepdims=True) - cal_components[:, c] = m_raw[:, 0] + xpy.log(simps(xpy.exp(lnL_t_c - m_raw), dx=deltaT, axis=-1)) + cal_components[:, c] = m_raw[:, 0] + xpy.log( + xpy.exp(lnL_t_c - m_raw).dot(_simps_weights(simps, npts, deltaT, xpy))) # fold in this realization's importance log-weight lnL_t_c = lnL_t_c + cal_log_w[c] @@ -3183,7 +3521,7 @@ def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop(tvals, P_vec, lookupNKDic # (the time integral is NOT taken; downstream resamples this timeseries). return running_max + xpy.log(S) - cal_log_w_norm - L = simps(S, dx=deltaT, axis=-1) + L = S.dot(_simps_weights(simps, npts, deltaT, xpy)) # lnL = max + log( sum_c exp(log_w[c]) \int dt exp(lnL_t - max) ) - log(n_cal) # 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 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py index f9378ee45..b5e83360b 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_freqresponse.py @@ -215,6 +215,31 @@ def response_coefficients(det, RA, DEC, psi, tref, Qmax, L_arm=None): return b +def response_coefficients_vector(det, RA, DEC, psi, tref, Qmax, L_arm=None): + """``response_coefficients`` for a BLOCK of extrinsic samples: {p: (npts_ex,) complex}. + + Same algebra as the scalar routine, evaluated once for the whole block. The two things + that do NOT depend on the sample -- the detector geometry (LAL lookup, arm unit vectors, + response tensor) and gmst(tref) -- are computed once here rather than once per sample; + everything that does depend on the sample (the polarization triad, the long-wavelength + contraction, the arm projections, beta_q) is an array operation over the block. + + The rotation path's ``factored_likelihood_with_rotation.rotation_coefficients_vector`` is + the same idea for the sidereal-harmonic coefficients; this is its finite-size twin. + """ + import lal + RA = np.atleast_1d(np.asarray(RA, dtype=float)) + DEC = np.atleast_1d(np.asarray(DEC, dtype=float)) + psi = np.atleast_1d(np.asarray(psi, dtype=float)) + gmst = float(lal.GreenwichMeanSiderealTime(lal.LIGOTimeGPS(float(tref)))) + geom = sfr.finite_size_geometry_vector(det, RA, DEC, psi, gmst=gmst, L_arm=L_arm) + beta = sfr.finite_size_beta(geom, Qmax) # (Qmax+1, npts_ex) + b = {0: np.asarray(geom['F0'], dtype=complex)} + for q in range(Qmax + 1): + b[1 + q] = np.asarray(beta[q], dtype=complex) + return b + + def FactoredLogLikelihoodFreqResponse(extr_params, rholms_intp_fr, crossTerms_fr, crossTermsV_fr, meta, Lmax): """Finite-size analogue of factored_likelihood.FactoredLogLikelihood. @@ -382,14 +407,15 @@ def _L_of(det): for det in rho_by_p: n_lms = len(lookupNKDict[det]) Ylms = FL.ComputeYlmsArrayVector(lookupNKDict[det], incl, -phiref).T # (npts_ex,n_lms) - # per-sample response coefficients b_p (npts_ex,) - bvec = {} - for i in range(npts_ex): - bi = response_coefficients(det, float(RA[i]), float(DEC[i]), float(psi[i]), - P_vec.tref, Qmax, L_arm=_L_of(det)) - for p in p_list: - bvec.setdefault(p, np.zeros(npts_ex, dtype=complex)) - bvec[p][i] = bi[p] + # per-sample response coefficients b_p (npts_ex,), for the whole block at once. + # This was a Python loop over samples calling the scalar response_coefficients. It + # dominated the likelihood: in a profiled five-detector CE+ET+K run it re-did the + # detector geometry 1.4e6 times for five distinct answers, and the eager default of + # `bvec.setdefault(p, np.zeros(npts_ex))` allocated an npts_ex-long array on every + # (sample, p) iteration -- 8.5e6 allocations, a quarter of the integration. + # DESIGN_freqresponse_vectorized_coefficients.md carries the measurement. + bvec = response_coefficients_vector(det, RA, DEC, psi, P_vec.tref, Qmax, + L_arm=_L_of(det)) t_ref = epochDict[det] # Precision-preserving time reference (see rotation NoLoop): keep (tref - epoch) and the diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_rotating_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_rotating_freqresponse.py new file mode 100644 index 000000000..e1818fb84 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_rotating_freqresponse.py @@ -0,0 +1,243 @@ +"""Combined Earth-rotation and finite-arm detector response for conventional ILE. + +The finite-arm response is first factored as ``sum_b beta_b(sky,t) W_b(f)``. Each +``beta_b`` is a finite sidereal Fourier series: its half-width is 2 for the exact +long-wavelength term and ``q+2`` for the arm-projection term of order ``q``. The +slow-delay Taylor expansion then composes with each frequency basis element, giving + + chi_(b,p,n) = M_n d_t^p [W_b h_lm] + +with all sky dependence in a short coefficient vector. This is the direct composition +described in the generalized-response section of the RIFT scaling paper. It acts on the +full inertial-frame modes, so precession and higher modes require no special treatment. + +The compound bank is deliberately opt-in. Its U/V precompute scales as the square of +the number of ``(b,p,n)`` elements; this implementation targets long, loud BNS-like +signals, not short eccentric mergers for which finite-arm response alone is sufficient. +""" +from __future__ import division, print_function + +import math +import numpy as np + +from . import factored_likelihood_with_rotation as flwr +from . import slowrot_freqresponse as sfr + + +def response_harmonic_width(basis_index): + """Exact sidereal half-width of finite-response coefficient ``b``.""" + return 2 if basis_index == 0 else (basis_index - 1) + 2 + + +def compound_index_set(Qmax, p_max): + """Nonzero elementary indices ``a=(b,p,n)`` for the compound response.""" + out = [] + for b in range(Qmax + 2): + for p in range(p_max + 1): + width = response_harmonic_width(b) + p + out.extend((b, p, n) for n in range(-width, width + 1)) + return out + + +def _basis_harmonics_geom(response, x_arm, y_arm, DEC, psi, Qmax): + """Fourier coefficients ``A[b][n]`` of each finite-response sky coefficient. + + A small exact DFT is used instead of maintaining separate symbolic expressions for + every arm-projection power. The sampled functions are polynomials in sin(g),cos(g) + with known half-width ``q+2``, so ``2*(Qmax+2)+1`` samples recover them to roundoff. + ``DEC`` and ``psi`` are vectors of extrinsic samples. + """ + DEC = np.atleast_1d(np.asarray(DEC, dtype=float)) + psi = np.atleast_1d(np.asarray(psi, dtype=float)) + width = Qmax + 2 + ngrid = 2 * width + 1 + g = 2.0 * np.pi * np.arange(ngrid, dtype=float) / float(ngrid) + X, Y, nhat = sfr._triad(DEC[:, None], psi[:, None], g[None, :]) + + Fp, Fc = sfr._lwl_response(response, X, Y) + zx = np.einsum('...i,i->...', X, x_arm) + 1j * np.einsum('...i,i->...', Y, x_arm) + zy = np.einsum('...i,i->...', X, y_arm) + 1j * np.einsum('...i,i->...', Y, y_arm) + ax = np.einsum('...i,i->...', nhat, x_arm) + ay = np.einsum('...i,i->...', nhat, y_arm) + + values = {0: Fp + 1j * Fc} + for q in range(Qmax + 1): + values[1 + q] = 0.5 * (zx ** 2 * ax ** q - zy ** 2 * ay ** q) + + harmonics = {} + for b, vals in values.items(): + bw = response_harmonic_width(b) + harmonics[b] = { + n: np.mean(vals * np.exp(-1j * n * g)[None, :], axis=1) + for n in range(-bw, bw + 1) + } + return harmonics + + +def combined_response_coefficients_vector(det, RA, DEC, psi, tref, p_max, + Qmax=4, L_arm=None): + """Compound coefficients ``{(b,p,n): C}`` for a vector of extrinsic samples.""" + import lal + import lalsimulation as lalsim + from . import slowrot_response as srr + + RA = np.atleast_1d(np.asarray(RA, dtype=float)) + DEC = np.atleast_1d(np.asarray(DEC, dtype=float)) + psi = np.atleast_1d(np.asarray(psi, dtype=float)) + response, x_arm, y_arm, _ = sfr.detector_geometry(det, L_arm=L_arm) + A = _basis_harmonics_geom(response, x_arm, y_arm, DEC, psi, Qmax) + gmst = float(lal.GreenwichMeanSiderealTime(lal.LIGOTimeGPS(float(tref)))) + g_ev = gmst - RA + Atil = {b: {n: val * np.exp(1j * n * g_ev) for n, val in Ab.items()} + for b, Ab in A.items()} + + lald = lalsim.DetectorPrefixToLALDetector(det) + Bd = srr.delay_harmonics_vector(lald.location, DEC) + Btil = {m: val * np.exp(1j * m * g_ev) for m, val in Bd.items()} + tau0 = np.real(sum(Btil.values())) + drift = dict(Btil) + drift[0] = drift[0] - tau0 + neg_drift = {m: -val for m, val in drift.items()} + + C = {} + E = {0: np.ones_like(g_ev, dtype=complex)} + for p in range(p_max + 1): + if p: + E = flwr._convolve_harmonics(E, neg_drift) + inv_fact = 1.0 / math.factorial(p) + for b, Ab in Atil.items(): + for n, an in Ab.items(): + for m, em in E.items(): + key = (b, p, n + m) + C[key] = C.get(key, 0j) + inv_fact * an * em + return C + + +def _arm_for_detector(L_arm, det): + return L_arm.get(det, None) if isinstance(L_arm, dict) else L_arm + + +def PrecomputeLikelihoodTermsRotatingFreqResponse( + event_time_geo, t_window, P, data_dict, psd_dict, Lmax, fMax, + Qmax=4, L_arm=None, p_max=0, f_sidereal=flwr.F_SIDEREAL, + analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0., + verbose=True, quiet=False, skip_interpolation=False, **hlm_kwargs): + """Build the intrinsic compound bank indexed by ``(b,p,n)``.""" + from . import factored_likelihood as FL + from .. import lalsimutils as lsu + + if data_dict.keys() != psd_dict.keys(): + raise ValueError("data and PSD detector sets differ") + detectors = list(data_dict) + P.dist = FL.distMpcRef * 1e6 * lsu.lsu_PC + P.deltaF = data_dict[detectors[0]].deltaF + hlms, hlms_conj = FL.internal_hlm_generator(P, Lmax, verbose=verbose, quiet=quiet, + **hlm_kwargs) + modes = list(hlms) + nfreq = hlms[modes[0]].data.length + fvals = flwr.evaluate_fvals_from_length(nfreq, hlms[modes[0]].deltaF) + a_list = compound_index_set(Qmax, p_max) + + rholms = {}; rholms_intp = {}; cross = {}; cross_v = {}; lengths = {} + for det in detectors: + _, _, _, length = sfr.detector_geometry(det, L_arm=_arm_for_detector(L_arm, det)) + lengths[det] = float(length) + weights = sfr.finite_size_response_weights( + fvals, {'L': float(length), 'T': float(length) / sfr.C_SI}, Qmax) + + weighted = {b: {lm: sfr_weighted_mode(hlms[lm], weights[b]) for lm in modes} + for b in range(Qmax + 2)} + weighted_conj = { + b: {lm: sfr_weighted_mode(hlms_conj[lm], weights[b]) for lm in modes} + for b in range(Qmax + 2)} + deriv = {(b, p): {lm: flwr.fd_apply_time_derivative(weighted[b][lm], p) + for lm in modes} + for b in range(Qmax + 2) for p in range(p_max + 1)} + deriv_conj = { + (b, p): {lm: flwr.fd_apply_time_derivative(weighted_conj[b][lm], p) + for lm in modes} + for b in range(Qmax + 2) for p in range(p_max + 1)} + chi = {a: {lm: flwr.fd_apply_sidereal_modulation( + deriv[(a[0], a[1])][lm], a[2], f_sidereal, 0.0) for lm in modes} + for a in a_list} + chi_conj = {a: {lm: flwr.fd_apply_sidereal_modulation( + deriv_conj[(a[0], a[1])][lm], a[2], f_sidereal, 0.0) + for lm in modes} + for a in a_list} + + data = data_dict[det]; psd = psd_dict[det] + t_det = FL.ComputeArrivalTimeAtDetector(det, P.phi, P.theta, event_time_geo) + rho_epoch = data.epoch - hlms[modes[0]].epoch + t_shift = float(float(t_det) - float(t_window) - float(rho_epoch)) + n_shift = int(t_shift / P.deltaT + 0.5) + n_window = int(2 * t_window / P.deltaT) + tgrid = np.arange(n_window) * P.deltaT + float(rho_epoch + n_shift * P.deltaT) + + rholms[det] = {}; rholms_intp[det] = {} + for a in a_list: + rho = FL.ComputeModeIPTimeSeries( + chi[a], data, psd, P.fmin, fMax, 1. / 2. / P.deltaT, + n_shift, n_window, analyticPSD_Q, inv_spec_trunc_Q, T_spec) + rholms[det][a] = rho + rholms_intp[det][a] = (None if skip_interpolation else + FL.InterpolateRholms(rho, tgrid, verbose=verbose)) + + cross[det] = {}; cross_v[det] = {} + for a in a_list: + for ap in a_list: + cross[det][(a, ap)] = FL.ComputeModeCrossTermIP( + chi[a], chi[ap], psd, P.fmin, fMax, 1. / 2. / P.deltaT, + P.deltaF, analyticPSD_Q, inv_spec_trunc_Q, T_spec, + verbose=False, same_waveform_Q=False) + cross_v[det][(a, ap)] = FL.ComputeModeCrossTermIP( + chi_conj[a], chi[ap], psd, P.fmin, fMax, 1. / 2. / P.deltaT, + P.deltaF, analyticPSD_Q, inv_spec_trunc_Q, T_spec, prefix="V", + verbose=False, same_waveform_Q=False) + + meta = dict(feature='rotation_freqresponse', Qmax=Qmax, p_max=p_max, + f_sidereal=f_sidereal, a_list=a_list, modes=modes, + event_time_geo=float(event_time_geo), L=lengths, L_arm=L_arm, + post_phase_required=True) + return rholms_intp, cross, cross_v, rholms, meta + + +def sfr_weighted_mode(hf, weight): + """Apply a finite-response weight without importing the scalar likelihood module.""" + out = flwr._copy_freqseries(hf) + out.data.data[:] = hf.data.data * weight + return out + + +def pack_rotating_freqresponse_arrays(meta, rholms, cross, cross_v): + """Pack the compound bank, with dense U/V arrays to avoid O(A^2) kernel launches.""" + lookup, rho, u_dict, v_dict, epoch = flwr.pack_rotation_arrays( + meta, rholms, cross, cross_v) + a_list = list(meta['a_list']) + u_dense = {}; v_dense = {} + for det in u_dict: + u_dense[det] = np.stack([ + np.stack([u_dict[det][(a, ap)] for ap in a_list], axis=0) + for a in a_list], axis=0) + v_dense[det] = np.stack([ + np.stack([v_dict[det][(a, ap)] for ap in a_list], axis=0) + for a in a_list], axis=0) + return lookup, rho, u_dense, v_dense, epoch + + +def DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, P_vec, meta, lookupNKDict, rho_by_a, U_by_aa, V_by_aa, epochDict, + Lmax=2, array_output=False, time_interp='nearest', xpy=np): + """Vectorized CPU/GPU likelihood for simultaneous rotation and finite-arm response.""" + Qmax = int(meta['Qmax']) + L_arm = meta.get('L_arm') + + def coefficients(det, RA, DEC, psi, tref, p_max): + return combined_response_coefficients_vector( + det, RA, DEC, psi, tref, p_max, Qmax=Qmax, + L_arm=_arm_for_detector(L_arm, det)) + + return flwr.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( + tvals, P_vec, meta, lookupNKDict, rho_by_a, U_by_aa, V_by_aa, epochDict, + Lmax=Lmax, array_output=array_output, time_interp=time_interp, xpy=xpy, + coefficient_function=coefficients, + reflection_function=lambda a: (a[0], a[1], -a[2])) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py index c932e5b4e..88714f508 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/factored_likelihood_with_rotation.py @@ -769,7 +769,8 @@ def pack_rotation_arrays(meta, rholms_rot, crossTerms_rot, crossTermsV_rot): def DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( tvals, P_vec, meta, lookupNKDict, rho_by_a, U_by_aa, V_by_aa, epochDict, - Lmax=2, array_output=False, time_interp='nearest', xpy=np): + Lmax=2, array_output=False, time_interp='nearest', xpy=np, + coefficient_function=None, reflection_function=None): """Vectorized rotation-aware lnL (Path A). GPU: pass xpy=cupy with rho_by_a/U_by_aa/V_by_aa already on device (the ILE converts them @@ -785,6 +786,11 @@ 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,). + + ``coefficient_function`` and ``reflection_function`` are internal extension points for + response models with the same sidereal-banded contraction. The combined slow-rotation + + finite-arm response uses them to reuse this maintained CPU/GPU implementation rather than + carrying a second copy of the likelihood algebra. """ # 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 @@ -815,6 +821,10 @@ def _h(v): # host (numpy) copy -- the antenna-harmonic / Ylm / delay helpers a a_list = list(meta['a_list']) p_max = meta['p_max'] + if coefficient_function is None: + coefficient_function = rotation_coefficients_vector + if reflection_function is None: + reflection_function = lambda a: (a[0], -a[1]) RA = _h(P_vec.phi); DEC = _h(P_vec.theta) incl = _h(P_vec.incl); phiref = _h(P_vec.phiref); psi = _h(P_vec.psi) npts = len(tvals); npts_ex = len(RA) @@ -825,7 +835,7 @@ def _h(v): # host (numpy) copy -- the antenna-harmonic / Ylm / delay helpers a for det in rho_by_a: n_lms = len(lookupNKDict[det]) Ylms = FL.ComputeYlmsArrayVector(lookupNKDict[det], incl, -phiref).T # (npts_ex, n_lms) - C = rotation_coefficients_vector(det, RA, DEC, psi, P_vec.tref, p_max) # {(p,n): (npts_ex,)} + C = coefficient_function(det, RA, DEC, psi, P_vec.tref, p_max) zeroC = np.zeros(npts_ex, dtype=complex) def Cg(a): @@ -886,7 +896,7 @@ def _ph(m): def _apply_post_phase(a, coef_ex, res): """conj(C~_a) Q^a = conj(C_a) e^{-i n_a omega delta_ij} Q^a_ij.""" - pe, pt = _ph(-a[1]) + pe, pt = _ph(-a[-1]) if pe is None: return coef_ex[:, None] * res return (coef_ex * pe)[:, None] * (pt[None, :] * res) @@ -922,15 +932,40 @@ def _apply_post_phase(a, coef_ex, res): # distinct m (4*n_harmonics+1 of them, so 4*(2+p_max)+1 at the default width) # instead of one per pair. term2_by_m = {} - for a in a_list: - aR = (a[0], -a[1]) - for ap in a_list: - val = xpy.conj(Cg_d(a)) * Cg_d(ap) * xpy.einsum( - 'xi,xj,ij->x', conjY_d, Ylms_d, xpy.asarray(U_by_aa[det][(a, ap)])) - val = val + Cg_d(aR) * Cg_d(ap) * xpy.einsum( - 'xi,xj,ij->x', Ylms_d, Ylms_d, xpy.asarray(V_by_aa[det][(a, ap)])) - m = ap[1] - a[1] - term2_by_m[m] = term2_by_m[m] + val if m in term2_by_m else val + if isinstance(U_by_aa[det], dict): + for a in a_list: + aR = reflection_function(a) + for ap in a_list: + val = xpy.conj(Cg_d(a)) * Cg_d(ap) * xpy.einsum( + 'xi,xj,ij->x', conjY_d, Ylms_d, + xpy.asarray(U_by_aa[det][(a, ap)])) + val = val + Cg_d(aR) * Cg_d(ap) * xpy.einsum( + 'xi,xj,ij->x', Ylms_d, Ylms_d, + xpy.asarray(V_by_aa[det][(a, ap)])) + m = ap[-1] - a[-1] + term2_by_m[m] = term2_by_m[m] + val if m in term2_by_m else val + else: + # Compound response banks carry tens to hundreds of bands. A Python/einsum + # launch for every ordered pair is prohibitive on a GPU, so their packer emits + # dense (A,A,K,K) banks. Contract all pairs sharing one post-phase difference + # in a single launch; the number of launches is then O(harmonic width), not A^2. + U_dense = xpy.asarray(U_by_aa[det]) + V_dense = xpy.asarray(V_by_aa[det]) + coeff = xpy.stack([Cg_d(a) for a in a_list], axis=1) + coeff_r = xpy.stack([Cg_d(reflection_function(a)) for a in a_list], axis=1) + n_index = np.asarray([a[-1] for a in a_list], dtype=int) + differences = n_index[None, :] - n_index[:, None] + for m in np.unique(differences): + ia, iap = np.nonzero(differences == m) + ia_d = xpy.asarray(ia) + iap_d = xpy.asarray(iap) + u_mode = xpy.einsum('xi,xj,pij->xp', conjY_d, Ylms_d, + U_dense[ia_d, iap_d]) + v_mode = xpy.einsum('xi,xj,pij->xp', Ylms_d, Ylms_d, + V_dense[ia_d, iap_d]) + term2_by_m[int(m)] = xpy.sum( + xpy.conj(coeff[:, ia_d]) * coeff[:, iap_d] * u_mode + + coeff_r[:, ia_d] * coeff[:, iap_d] * v_mode, axis=1) # Re[] is linear, so accumulate the real part per m and keep the persistent array real. term2 = xpy.zeros((npts_ex, npts), dtype=np.float64) for m, val in term2_by_m.items(): diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_anglemarg_memory.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_anglemarg_memory.md new file mode 100644 index 000000000..85938683e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_anglemarg_memory.md @@ -0,0 +1,157 @@ +# JAX angle-marginalization memory model + +These are logical array-size and lifetime models for the JAX-only +angle-marginalization kernels. Except for the historical XLA allocation request +identified below, they are not measurements of CUDA allocator peak memory. +They must not be read as the footprint of conventional production ILE. + +The evaluation cap in `samplers.py` protects only callers using `eval_lnL*`. +Direct `log_likelihood` calls and scalar value/gradient/Hessian entry points +bypass it, and a fraction of reported device memory does not bound the sum of +live buffers, allocator reservations, or reverse-mode residuals. + +Let `S` be batch size, `T=data.npts`, `F` a phi chunk, `D` a distance block, +`Q=16` the Laplace u chunk, `E` the exact dense-angle chunk, `G` the exact +distance block, and `P` the rolled sample-time point block. Float64 and +complex128 occupy 8 and 16 bytes. + +## Common storage + +For source mode bound `m`, the coefficient tables have shapes +`(m+1,3,S,T)` and `(2m+1,5,S,T)` complex128. Together they contain + +``` +16 S T [3(m+1) + 5(2m+1)] bytes. +``` + +At `m=2` this is `544 S T` bytes: 2.42 GiB at `S=4000,T=1193`. +Their angle-sample loop is rolled, but coefficient construction is not yet +tiled over the evaluation sample/time axes. These tables persist across the +phi scan; the quoted number is their logical payload, not an allocator peak. + +## Exact + +The dense angle grid is scanned in `E=8` chunks and distance in `G=32` +blocks. The dominant exponent slab is `(E S,T,G)` float64, or +`8 E G S T = 2048 S T` bytes (9.10 GiB at `4000 x 1193`). Grid length is +bounded; sample and time still multiply the slab. Exact therefore remains +under the conservative outer cap pending point-axis tiling. + +## Laplace + +Before this patch, one step of the u scan formed the logical f64 result +`blk` with shape + +``` +(Q,D,F,S,T) float64 = 8 Q D F S T = 8192 S T bytes +``` + +at shipped `Q=16,D=4,F=16`. It is reduced over `Q` immediately; the distance +and phi scans do not keep all of their blocks simultaneously. The complex128 +products used to form `blk` have the same shape but are eligible for compiler +fusion. At `S=4000,T=1193`, the f64 `blk` alone is 36.407 GiB. Commit +`c5b81dd6` records that XLA requested this single allocation during a pre-cap +SNR-40 JAX acceptance run against a 25-GiB cgroup. This investigation does not +have the original allocator log, did not reproduce that run, and did not +measure a 36-GiB current-production footprint. + +The other source-visible live values include the persistent coefficient tables, +five phi fields (`64 F S T` bytes: A0/B0 real, A1/B1/B2 complex), distance-scan +carries and, for differentiated calls, residuals selected by XLA/AD. Their +simultaneous physical lifetime cannot be obtained by summing source-level +shapes and requires an allocator profile. + +Laplace now flattens the independent `(S,T)` axes, edge-pads only the last +tile, and maps distance/psi marginalization over fixed tiles. Its expensive +slab is bounded by + +``` +8 Q D F min(S T,P), P=LAPLACE_POINT_BLOCK=4096, +``` + +or 32 MiB with shipped inner blocks for a direct call whose only batched axes +are the explicit `S,T` axes. Padding repeats a finite edge point and is discarded +before the phi reduction. Every real bin retains the same distance nodes, psi +quadrature, per-bin reduction order, phi reduction, and Simpson time +marginalization. The map body is checkpointed for reverse AD. Coefficient tables +and phi fields remain `O(S T)`, so this is neither a claim that total memory is +32 MiB nor a bound on an arbitrary transformed caller. + +In particular, `flowMC` applies an outer `vmap` over its chains to the scalar AD +target. The scalar wrapper has explicit `S=1`, so its `pblk` calculation cannot +see that mapped chain axis. For the usual 20-chain driver call at `T=1193`, the +corresponding logical primal slab is at most about 186 MiB before accounting for +AD residuals, not 36.41 GiB, but it is also not covered by the 32-MiB statement. + +## Production call paths + +Conventional `integrate_likelihood_extrinsic_batchmode` does not call this JAX +kernel. Its maintained GPU NoLoop path samples distance, phi and psi and carries +primarily `(S,T)` arrays (`kappa_sq` complex128 and `rho_sq` float64); it has no +`Q*D*F` angle-quadrature multiplier. Operation on 4-GB cards therefore does not +contradict the JAX shape above. + +The separate `integrate_likelihood_extrinsic_jax` reaches this kernel only for +the distance+phi+psi-marginalized mode with a resolved Laplace scheme. Its host +pilot/reweight evaluations call `angle_marg_eval_chunk`; the sampler helpers do +the same. At `T=1193`, the 4-GiB fallback target caps the old model at `S=439`, +so the current production call path does not submit `S=4000`. Scalar +value/gradient/Hessian calls use explicit `S=1`; flowMC normally maps those over +20 chains. + +When a device limit is known, `_angle_marg_buffer_target()` now always applies +the configured fraction: a reported 4-GiB card therefore gets a 2-GiB target at +the default fraction. The historical 4-GiB value is reserved for the +unknown-device fallback. If the modeled payload for one sample exceeds the +target, the evaluation helper raises a resource preflight error instead of +returning a fictitious chunk size of one. This remains a source-level working-set +model, not a bound on total allocator use; direct `log_likelihood` calls bypass +the helper altogether. + +## Peak-local + +The u-node axis is already streamed with `U_live<=8`, and phi with `F=16`. +The node body per sample-time point is `8 F N_x 4 U_live` bytes: 1 MiB at +`N_x=256`. The phi scan also returns every step before reducing it, so its +stacked `(n_phi,N_x)` f64 result adds `8 n_phi N_x` bytes per sample-time point. +The outer evaluation cap budgets the sum and refuses a call when even one sample +does not fit. For example, at `T=1193,N_x=256,m_max=2`, `A=450` gives +`n_phi=352` and a 1.966-GiB one-sample model, while `A=12500` gives +`n_phi=1792` and a 5.242-GiB model. + +This does not fix hidden transformed axes. Nested `vmap(vmap(_one))` still +multiplies the body and scan result by explicit `S T`, and flowMC applies an +additional outer chain `vmap` to the scalar likelihood that this preflight +cannot see. A follow-up must roll those axes around `_one` and GPU-profile a +suitably smaller point tile before peak-local can claim a total-memory bound. + +## Validation boundary + +Checkpointing the exact/Laplace phi scans and peak-local phi/u scans bounds +saved loop residuals, but does not by itself shrink primal `S*T` +vectorization. Tests inspect the traced Laplace kernel-input shape and compare +tiled versus one-block values and gradients, including a padded tail. + +CPU tests cannot establish CUDA allocator peaks, GPU XLA fusion, or the +throughput-optimal `P`. Before relaxing `angle_marg_eval_chunk`, profile all +three schemes on a production CUDA host at `T≈1193`, batches spanning the +current cap and nominal 1000/4000, and exercise value, gradient, and +Fisher/Hessian calls while recording allocator peak statistics. Profile the +flowMC outer-vmap path separately: explicit point tiling does not bound that +hidden chain axis. + +## 2026-09-08: jax 0.9.2 never populates `largest_free_block_bytes` + +On jax 0.9.2 (ldas-pcdev11, idle 24 GiB card) `largest_free_block_bytes` and +`pool_bytes` both read 0 before the first allocation, so `_device_available_bytes` +now treats a bare 0 in either field as "not reported" and falls through, rather +than as a full device. + +Adversarial review of that fix found it reachable on a *busy* card too, before +this process's own first allocation, where 0/0 read the same as idle but the 4 GiB +fallback it falls through to is not safe. `_angle_marg_buffer_target` now forces +one tiny allocation with `_probe_allocate` before reading `memory_stats()`, so the +pool signal exists to read; if the resulting pool is small next to `bytes_limit` +(the on-demand-allocator shape, where the pool only grows to fit what has been +requested so far), availability is bounded by `bytes_limit - bytes_in_use` rather +than trusted. A missing `bytes_in_use` key is now read as unknown, not 0. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_planner.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_planner.md new file mode 100644 index 000000000..8ac622ea2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_planner.md @@ -0,0 +1,230 @@ +# Error- and resource-budgeted direct-marginalization planner + +## Status and verdict + +The policy engine in `direct_marginalization_planner.py` is implemented and +tested, but its RIFT scheme catalog is deliberately **not wired into the JAX +driver or wrapper**. It is an opt-in planning API, not a new production +`auto` mode. + +That boundary is load-bearing. The current angle, distance, and time schemes +do not yet expose comparable proof-carrying error bounds and measured costs on +a common unit. Wiring a selector before those adapters exist would require the +planner to invent numbers, or to call a calibrated grid "certified". The +framework note explicitly rules that out. The implemented layer can make the +decision once real adapters supply those records; until then, a strict +three-axis request declines honestly. + +No existing behavior changes: + +- `ANGLE_MARG_DEFAULT` remains `exact`; +- `choose_angle_marg_scheme` is unmodified, including its existing + amplitude crossover and GH compatibility behavior; +- `angle_marg="auto"`, the time default, and both distance-grid defaults keep + their old paths; +- no new CLI choice is registered. + +The only way to use this work is to import the new module, construct explicit +scheme offers, and call `plan_direct_marginalization` or +`plan_jax_direct_marginalization`. + +## Inputs and units + +A request has four independent inputs. + +1. A positive error ceiling in **absolute marginalized log-likelihood error, + nats**, for every requested axis. There is no implicit sharing of a total + budget: the caller must perform that allocation. +2. A compute ceiling and a peak-memory ceiling. Both are mandatory. Compute + estimates must use one common unit within the request. Memory is bytes. +3. One or more `SchemeOffer` objects per axis. Every offer carries its error + assessment, resource estimate, warrant, prerequisites, incompatibilities, + and provenance. +4. Concrete capabilities established for this dataset, such as + `gh-laplace-supported` after `gh_laplace_supported` has checked the actual + coefficient tables. Missing capabilities are refusals, not false values to + route around. + +By default the planner sums compute contributions and sums live-memory +contributions. Direct marginalization nests axes, so a production adapter +should pass a combination-aware `resource_model` when those interactions +matter. That callback returns the same provenance-carrying `ResourceEstimate` +type and is allowed to conservatively over-count buffers whose lifetimes do not +overlap; it may not assume reuse that it has not measured. The default is safe +for additive evidence packets and tests, not a claim that nested kernel costs +are separable. Either form is a hard resource guard, not a wall-time predictor. + +## Warrants are not accuracy labels + +The warrant union follows `DESIGN_peak_local_framework.md`: + +| warrant | can support a certificate? | current use | +|---|---:|---| +| `exact-band-limit` | yes | time exponent reconstruction | +| `exact-trig-degree` | yes | finite angular stationary set | +| `bounded-stationary-set` | yes | support-aware distance candidates | +| `effective-bandwidth-with-margin` | no | amplitude-sized dense angle grids | +| `empirical-calibration` | no | validation envelopes | +| `none` | no | fixed historical grids | + +"Can support" is still weaker than "implemented". `Warrant` therefore has a +separate `certificate_available` field. `AccuracyAssessment(CERTIFIED, ...)` +is rejected at construction unless both conditions hold. In particular, +calling the angle scheme `exact` refers to exact coefficient reconstruction; +the subsequent quadrature over `exp(lnL)` is sized from an effective bandwidth +and remains best-effort with a runtime label. The profile forbids relabeling it +as a proof. + +## Current JAX profiles + +The module records structural facts already enforced in the shipped call +sites. It does not attach error or cost numbers to them. + +| axis / scheme | recorded warrant | important compatibility fact | +|---|---|---| +| angle `grid` | none | cannot drive the amplitude-sized log-uniform distance grid | +| angle `exact` | effective bandwidth with margin | requires the data-derived amplitude estimate | +| angle `laplace` | effective bandwidth with margin for the complete angle result | GH additionally requires the measured `A0==0/B1==0` identity | +| angle `peak-local` | exact trig degree only on psi; effective bandwidth for the still-dense phi axis | requires an explicit feature warrant and refuses GH | +| distance `uniform` | none | historical fixed grid | +| distance `loguniform` | bounded stationary set, no implemented end-to-end certificate | requires full prior support, an interior peak, and a passing endpoint budget | +| distance `gh` | bounded stationary set, no implemented error certificate | currently the volumetric-prior kernel | +| time `simpson` | none | historical fixed grid | +| time `bandlimited` | exact band limit, no implemented per-request certificate | the nonlinear JAX ANGLE wrappers currently refuse this ordering; the distance wrapper applies its reduction on the refined nodes and accepts it | + +The last row carries both kinds of caveat at once, and is why a production +three-axis error-budgeted plan is not merely waiting for an angle cost table. +The band limit is genuine structure, so the warrant kind could support a +certificate; but the shipped rule derives its refinement factor from a measured +peak width and remeasures it, and reports measured reconstruction errors rather +than a proved bound on the marginalized log likelihood, so no certificate is +advertised and `CERTIFIED` is refused at offer construction. It is in any case +not compatible on the direct distance/angle-marginalized JAX path, while the +compatible Simpson rule has no per-request error bound either. No shipped +profile is therefore certificate-bearing today: `cheapest-certified` is +reachable only for a future scheme that implements and validates its bound. + +## Decision policy + +The planner enumerates the small Cartesian product of per-axis offers and +records, for every combination: + +- missing prerequisites and active conflicts; +- missing conditional warrants (for example GH plus Laplace); +- certification status and error-budget excess, per axis; +- compute and memory totals and any resource-budget excess. + +Among compatible, affordable combinations certified inside every axis budget, +it chooses the least compute, then least memory, then the smaller normalized +error. This is the `cheapest-certified` result. + +If none exists, it ranks compatible affordable combinations by the worst +per-axis normalized assessed error, then total normalized error and evidence +strength. Under the default policy this candidate is only `suggested` and the +decision action is `decline`. `require_selection()` raises +`MarginalizationPlanDeclined`, so code cannot accidentally execute the +suggestion as if it were a selection. + +Only `allow_best_effort=True` promotes that candidate to a runnable +`most-accurate-affordable` decision. Its record says `certified=False` and +separately says whether its numerical assessments meet the requested budgets. +This explicit authority is the only way for the *planner* to promote its own +suggestion. The production failure-resolution path below is separate and does +not rewrite the planner's claim. + +Every result is JSON-ready through `PlanDecision.as_dict()`. The record embeds +the complete input budgets, capabilities, offer provenance, warrant provenance, +resource provenance, selection basis, and combination decline ledger. + +## Production-safe method-decline resolution + +A marginalization method can fail its warrant at runtime even after selection. +The important example is an incomplete stationary-root enumeration. That +event says that the preferred quadrature cannot certify or complete its result; +it says nothing about whether the waveform or the underlying likelihood point +is valid. Converting it into the generic waveform-failure sentinel would +silently drop a scientifically valid sample. + +`resolve_plan_for_production` therefore keeps three outcomes distinct: + +- `use-preferred` when the selected method remains runnable; +- `use-conservative-fallback` after either a fail-closed planning decision or + an explicit runtime `MethodDecline`; +- `waveform-failure` only when the waveform/base-likelihood layer explicitly + supplies an independent `WaveformFailure` record. + +The resolver never invents or silently selects a fallback. Production setup +must supply a `ConservativeFallbackPolicy` with exactly one reserve offer for +each replaced axis, a separate hard reserve resource budget, provenance, and a +finite-output contract. For the shipped JAX catalog, +`make_jax_production_fallback_policy` restricts this role to the historical +support-covering, non-root-enumerating paths: angle `exact` (dense phi/psi), +distance `uniform`, and time `simpson`. This role does **not** relabel those +methods as error-certified. The resolution ledger reports their actual error +evidence and whether it meets the original request. + +A runtime decline on one axis replaces that axis and retains the other selected +axes. A runtime decline with no axis cannot identify which selected warrant was +lost, so it conservatively replaces the complete selected plan. A planning +decline likewise has no executable partial selection, so its fallback must cover +all requested axes. Missing coverage, incompatibility, or excess of the reserve +budget raises `FallbackConfigurationError` during resolution; none of those +configuration defects is returned as an invalid likelihood sample. +The ledger preserves the original warrant/resource refusal, the runtime root +postcondition when present, the chosen reserve, both budgets, and all +provenance. `ProductionResolution.require_selection()` returns either the +preferred or reserve plan and raises only for an explicit waveform failure. + +This is still an adapter contract rather than live wrapper wiring. A future +wrapper must call the resolver at the root-enumeration postcondition, evaluate +the selected dense reserve, verify that its returned value is finite, and only +then classify any independent non-finite waveform/base-likelihood condition. +It must not catch `MethodDecline` as a waveform exception. + +## Why amplitude alone is insufficient + +The old angle selector is intentionally retained as a compatibility API. Its +crossover is an accuracy crossover, while its own source records a different +and much higher measured cost crossover. A single amplitude threshold cannot +simultaneously express: + +- a caller's error tolerance; +- whether the dataset satisfies a scheme's warrant; +- distance/time compatibility; +- a device-memory ceiling; +- a measured execution-cost calibration. + +The focused amplitude-ladder test therefore supplies a synthetic evidence +packet in which the Laplace error and the dense-rule cost have different +crossings. The planner selects exact at low amplitude (Laplace misses the +error budget), exact at moderate amplitude (both are accurate but exact is +cheaper), and Laplace at high amplitude (both are accurate and Laplace is +cheaper). Those numbers test policy only and are explicitly not RIFT kernel +measurements. + +This follows the manuscript's Section IV policy at the structural level: no +single method is presumed to cover the whole amplitude range, cost and returned +quality are separate deliverables, and an approximation is not made correct by +being affordable. Section IV concerns samplers, so none of its performance +numbers are reused as quadrature calibration. + +## Production gate + +Before exposing a driver option, each live adapter must provide all of the +following from the concrete data and device: + +1. a per-axis quantitative accuracy assessment whose evidence class is honest; +2. an implemented certificate if the offer is to enter the strict pool; +3. compute on a common measured unit and a conservative live-memory estimate; +4. static and conditional compatibility tokens from the existing build-time + predicates; +5. a wrapper-level application test showing that a planner/runtime method + decline runs the configured finite reserve and cannot become either a + default scheme or a dropped waveform point; +6. low/moderate/high-amplitude campaign measurements, including the overlap + regions and device classes on which cost ordering changes. + +Until that evidence exists, the planner should remain an explicit prototype. +Its useful production contribution today is the typed contract: it makes the +missing evidence visible and prevents the next selector from encoding it as +another unexplained crossover. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_policy.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_policy.md new file mode 100644 index 000000000..a0acd7d51 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_direct_marginalization_policy.md @@ -0,0 +1,645 @@ +# The cross-axis direct-marginalization policy + +Module `direct_marginalization_policy.py`; driver flag +`--direct-marginalization-policy {off,auto}`, default `off`. +Companion to `DESIGN_direct_marginalization_planner.md` (the generic +error/resource planner, still unwired) and to PR #268, whose controller this +policy runs. + +## Status + +Opt-in, value-only. Wired on 2026-09-07 for `--mode flowmc-phipsimarg`. +Nothing selects it by default. Gradient parity is not validated; see the +gate below. + +## Why not widen `--angle-marg-scheme auto` + +That selector chooses between `exact` and `laplace` on one amplitude +crossover, controls angles only, and excludes both peak-local kernels by a +pinned test. The composite method owns four axes at once and decides per +evaluation from diagnostics. It is a different object and gets a different +flag. + +## What one evaluation does + +For a batch of extrinsic rows `(ra, dec, incl)`: + +1. `anglemarg.angle_coefficient_tables(..., guard=G)` builds the exact + coefficient tables once with `G` primitive-only support samples at each + end. The norm table is collapsed per row; its deviation over time is + recorded and a row whose norm moves with time is marked unusable. +2. Per row, under `vmap`: `rank_joint_starts_from_uvq_device` at angular + oversample 1 (base) and 2 (extra), then + `make_all_axis_mode_plan_pair_device` refines both in one shared pass and + freezes two nested plans. The plans are placed under `stop_gradient`. +3. `empirical_enrichment_with_exact_reserve_sequential_batch` applies the + local gate per row and, on a decline only, executes the reserve. +4. The reserve is the exact-angle coefficient integral on a time rule + refined `reserve_time_refine` times (default 4) over the native cadence. + It is warranted by the two-guard comparison (`G` against `G/2`) and by + the half-refined check rule evaluated in the same declined branch. The + warrant is a convergence statement about the refined rules. The native + Simpson rule is not the check rule: on a peak narrower than a sample it + is the unconverged one, and its error is what the refinement removes. + The refined rules are trapezoid, not Simpson: Simpson aliases at half the + node spacing on a sub-sample peak (0.03 to 1.8 nat at refine 4 for peaks + of 0.05 to 0.2 samples), the trapezoid rule converges exponentially. + The wiring test measures the native rule's error on its analytic fixture. +5. The selected value and the ledger come back per row. + +Acceptance diagnostics, all required for the local branch: + +| diagnostic | ledger key | +|---|---| +| norm table time-independent | `norm_time_invariant` | +| no capacity truncation, base and enriched | `base_capacity_ok`, `enriched_capacity_ok` | +| finite stationary modes | `base_and_enriched_values_finite`, `decline_no_modes` | +| no competitive start pinned to the time or distance boundary | `boundary_maximum_ok`, `decline_boundary_maximum` | +| valid, nested local geometry | `geometry_nesting_ok`, `decline_geometry` | +| base/enriched mode agreement | `mode_nesting_ok` | +| nested quadrature convergence | `decline_quadrature`, `decline_enrichment` | +| two-guard time agreement | `decline_time_reconstruction` | +| omitted-time mass under budget | `time_omitted_mass_ok` | +| total empirical error score under budget | `value_error_budget_ok` | + +Reserve warrant: `reserve_time_guard_validated`, +`reserve_time_resolution_validated`, `reserve_time_error_budget_ok`, combined +in `reserve_time_warranted`. A reserve that fails its warrant is escalated: +the rule is doubled and re-checked against its own half, up to +`reserve_time_refine_max` (default 32). A row still unwarranted, or whose +norm table varies with time, is `usable=False` and its likelihood is `nan`. +The finite diagnostic stays in the ledger under `selected_value` and never +reaches the sampler. The driver raises on the first `nan` it evaluates and +refuses to publish samples or evidence that contain one (external review of +PR #278, P1). A MALA step onto a `nan` target is rejected, so chains do not +carry such rows either. + +No SNR threshold appears anywhere. The transitions reported in the paper +(reserve at 40 and 80, local at 160 and 320) emerge from these diagnostics. + +## Operating point + +`PolicyConfig` defaults follow the configuration that accepted on production +tables at rho 163 and 326 (RIFT_roboto_paper +`analyses/va_sequence_20260902/RESULTS_20260907_aap268_ladder.md`): angular +oversample 2 and 4, 16 modes, radius 6, guard 128, 14 refine iterations. +PR #268's test values (oversample 1 and 2, 4 and 8 modes) overflowed capacity +on synthetic carrier tables and were never a measured production point. + +The guard is data, not a knob: the gather returns nonfinite samples past the +stored buffer with no error. The wrapper probes a coarse sky grid at +construction and refuses a guard the buffer cannot supply, and every row +carries `tables_finite`; a nonfinite row is `input_nonfinite`, `nan`, and +never a method decline. + +## Gradient memory + +Measured with XLA's compile-time memory analysis of `value_and_grad` on a +rho 163 production row (614-sample window, guard 128, 83200 dense angles, +16 GH distance nodes): + +| stage | temp memory | +|---|---| +| production exact scheme | 2.6 GiB | +| tables, planning, local gate (base and enriched) | under 0.02 GiB | +| one reserve tier at refine 4, inside `lax.cond`, dense chunk 8 | 10.5 GiB | +| full policy, tiers to refine 32, dense chunk 8 | 84.9 GiB (158 before the branches were rematerialized) | +| reserve at refine 2 / 4 / 8, dense chunk 8 | 5.3 / 10.5 / 21.0 GiB | +| reserve at refine 2 / 4 / 8, dense chunk 64 | 0.73 / 1.46 / 2.92 GiB | + +The local branch is essentially free to differentiate. The cost is the dense +reserve's reverse pass: one carry per dense-angle scan step, so it is +proportional to the refined time nodes and inversely to the dense chunk, and +`lax.cond` reserves memory for the largest tier whether or not it runs. The +policy's reserve chunk is therefore 64 (the kernel default is 8), which +brings the full policy at ceiling 32 to about 12 GiB per evaluated row. +The controller's branches and every tier are under `jax.checkpoint` and the +planning inputs are under `stop_gradient` (planning is control data). The +escalation ceiling is exposed as +`--direct-marginalization-reserve-time-refine-max` because it bounds gradient +memory; the value path does not depend on it below the ceiling. A cropped +reserve over the plans' certified time cover would cut both cost and memory +by the ratio of window to cover and is the follow-up. + +## Measures + +The local branch integrates `x**-4 dx dt_sample dphi du`. The reserve and +the exact scheme average both angles, weight distance by the normalized +`log_w_grid` (fixed grid) or the normalized volumetric measure +(`JAX_ILE_DISTMARG_GH`), and integrate time in seconds by Simpson. +`policy_log_normalization` derives the conversion (the wiring test checks +it end to end against an independent fine-time reference on both distance +paths): + +| term | constant | +|---|---| +| angles | `-2 log(2 pi)` | +| time | `log(deltaT) + log(sum(w_t) / ((npts-1) deltaT))` | +| distance, fixed grid | `3 log(Dref) - log(sum_i d_i^2 dd)` with `dd` read off the grid | +| distance, GH | `log 3 - log(x_min^-3 - x_max^-3)` | + +## Refusals + +The policy refuses, with a message, any of: a resolved angle scheme other +than `exact`; a time rule other than `simpson`; a distance prior other than +volumetric; a distance grid other than uniform-in-d; a `time_guard` below 2; +a reserve refinement that is not an even integer of at least 2; a request +for `lnL(t)`. The driver refuses at parse time a policy request in any mode +other than `flowmc-phipsimarg`, and any of the three policy knobs when the +policy is off (external review of PR #278, P1). Refusal rather than silence +is the standing rule on this arm. + +## Cost + + + + +Planning is vectorized. Accepted rows pay fixed local work per retained mode; +declined rows pay three exact reserve evaluations (refined rule at two guards, +plus the half-refined check). The sampler's angle-scheme chunk cap applies. + +`PolicyConfig.reserve_batch_rows` sets how many rows run under one `vmap`; +`--direct-marginalization-batch-rows` exposes it. At 1 the rows run one at a +time under `lax.map`, the graph PR #268 measured. Above 1 the tier-escalation +`lax.cond` becomes a `select`, so every reserve tier runs for every row. Nothing else changes: +`test_row_batch_size_changes_cost_not_values_decisions_or_gradients` requires +`lnL` bitwise equal, every ledger key and summary count equal, and the +gradient equal to one ulp, over the full-batch, whole-multiple and remainder +paths. + +### Device workspace + +XLA buffer assignment, ladder-2 tables, NVIDIA RTX PRO 4000 Blackwell, +jax 0.9.2, `--n-phi 32 --n-psi 8 --distance-grid-points 256`, +`reserve_time_refine_max` 32: + +| `reserve_batch_rows` | temp GiB at rho 40.8 | temp GiB at rho 652.3 | +|---|---|---| +| 1 | 0.425 | 0.432 | +| 2 | 0.812 | 0.822 | +| 4 | 1.588 | 1.592 | +| 8 | 3.132 | 3.132 | +| 16 | 6.212 | 6.212 | +| 32 | 12.371 | 12.371 | +| 64 | 24.692 | 24.692 | + +The fit is `0.046 + 0.385 B` GiB, maximum residual 6 MiB, at the grid sizes in +the caption. The rung does not enter. Measured device use at B=32 was 23.3 GiB +against the 12.4 GiB analysis figure, so buffer assignment understates the card +about twofold: a 24 GiB card holds B=32 and not B=64. + +The tier count multiplies it. At B=8, `reserve_time_refine_max` 32 costs +3.114 GiB and 138 s to compile; at 4 (one tier) the same batch costs +0.415 GiB and 29 s. A batched row pays every tier. + +### Throughput + +Batching is a REGRESSION, not a saving. It converts three `lax.cond`s to +selects, not the one the knob's first version named: + +| site | becomes, under vmap | +|---|---| +| tier escalation in `_row` | every reserve tier runs for every row | +| `all_axis_peaklocal.py:2504` accept/reserve | every ACCEPTED row also runs the dense reserve | +| `all_axis_peaklocal.py:1760` per-mode live | every dead mode slot runs a 4-D quadrature | + +So the penalty scales with the locally accepted fraction and the tier count. +Measured on the wiring fixture, 2 of 6 rows accepting, second timed call: + +| tiers | B=1 | B=2 | B=6 | +|---|---|---|---| +| 1 (`reserve_time_refine_max` 4) | 2.50 s/row | 3.28 (+31%) | 4.08 (+63%) | +| 2 (`reserve_time_refine_max` 8) | 6.49 s/row | | 7.81 (+20%) | + +An earlier production-table point read 96.3 s/row at B=1 against 99.9 at B=8 +and was reported as cost-neutral. It ran at `reserve_time_refine_max` 4, where +`tiers` has length one and the escalation cond is absent, with `accepted_local` +0 of 8, so the accept/reserve cond took its cheap branch everywhere. Both +penalties were inert, making it the best case. The docstring at +`all_axis_peaklocal.py:2344` is false above B=1. + +`reserve_batch_rows` defaults to 1, kept for the equivalence it pins rather +than for a saving. A single row takes the sequential path whatever is +requested: `_scalar` evaluates one row, so `value_and_grad` and `hessian` +would otherwise pay every select with nothing to amortize. + +## Peak-local time reserve + +The reserve schemes are pairs (angular kernel, time rule): + +| scheme | angular kernel | time rule | +|---|---|---| +| exact | exact angles | whole window, refine 4 escalating to 32 | +| laplace | psi-Laplace | whole window | +| peaklocal | psi-Laplace | peak-local, fixed count | +| peaklocal-exact (config only) | exact angles | peak-local, fixed count | + +The whole-window rule's node count grows with rho: the peak of exp(lnL(t)) +has width sigma_t = 1/(2 pi rho sigma_f), and every node pays the angular +kernel. At rung 41 the refine-4 rule already misses its own 1e-3 warrant on +two of three rows (0.0018, 0.0040, 0.0114 nat, ladder controller +2026-09-08), and the exact kernel costs 18.6 / 72.6 / 290 / 1278 s per +evaluation at rungs 41 / 82 / 163 / 326 on 2453 nodes. + +### The rule + +One lattice per row. The scan puts `scan_refine` nodes per native sample +across the window (default 2). The fine lattice divides each scan cell by an +integer `m`, chosen so that its spacing is at most `sigma_t / 3` for the +predicted width; every node of the rule is a point of that lattice. Around +each of the row's time maxima sits a block of `n_fine` consecutive fine nodes +(default 73, so the block spans 24 predicted sigma). Overlapping blocks and +scan nodes inside a block coincide and carry no weight; the only +spacing changes are at block ends, where the mass is negligible. Mixing two +lattices instead cost 0.02 nat on a 0.75-sample peak: the trapezoid rule is +spectrally accurate only on uniform spacing. + +The maxima come from the primitive, not from the local branch. The field +`max_x (x A(t, phi, u) - x^2 B(phi, u) / 2)` is evaluated on a search grid of +8 nodes per sample over an 8 x 8 angular lattice; the four highest local +maxima are polished by Newton steps on the fixed-angle field with the +reflected spectrum's exact derivatives. On a synthetic peak 0.16 samples wide +the local branch's plan carried no live mode at all, so a reserve leaning on it +inherits that decline. The plan's centre and width are reported beside the +locator's as a cross-check. + +The scan is support-limited too: `reserve_peaklocal_scan_nodes` (65) nodes +on the scan lattice across the hull of the live maxima widened by +`reserve_peaklocal_scan_margin_sigmas` (16) predicted sigma each side, never +the window. The mass outside the hull is bounded from the locator's own +search profile, summed over the outside search nodes, plus +`reserve_peaklocal_outside_slack_nats` (5) for what lies between nodes. The +angle- and distance-maximized exponent is an upper bound on the marginal at +each search node. That bound is charged through the kernel's cropped-cover +warrant (`reserve_time_cover`). The node +count is `65 + 4 x 73 = 357` whatever the window; a count that grows with the +window is the wrong design (RO, 2026-09-09). + +The check rule is the half-refined scan plus every other fine node, on the +same lattice with the same endpoints, so the kernel's structural resolution +warrant, its two-guard comparison and its tail bound apply unchanged. A +failed warrant doubles the fine lattice at fixed span, at most +`reserve_peaklocal_escalations` times (default 2), and the ledger counts it. +The rule and its check share the block, so a block off the maximum agrees +with itself: measured 0.22 nat, warranted, on a rung-160 row whose block sat +0.17 samples (2.7 sigma) from the maximum. The kernel therefore carries a +focus certificate (`reserve_time_focus_ok`): the node with the largest +evaluated `lnL(t)` must lie within a quarter of the block span of the block +centre, or the rule is unwarranted. + +Under jit, XLA fuses the difference of two bitwise-equal products into a +multiply-add that rounds to -1e-15, so the kernel's node checks compare +slices rather than take `diff >= 0`, and repeated positions are accepted as +non-decreasing. + +### Sizing from the prediction + +`sigma_t = 1 / (2 pi rho sigma_f)` per row. `rho` comes from the locator's +profile maximum, `rho^2 = 2 P_max`, which is the row's own `lnL` maximum; the +angular triangle bound over the norm's triangle lower bound is the fallback +and is reported beside it (`rho_bound`). The bound ran 2.8x over on a +rung-160 row (453 against 157), which narrowed the block span to 4 sigma of +the true peak; the profile value does not. `sigma_f` is the two-sided +rms frequency of the stored Q (`q_effective_bandwidth_hz`). The ladder record +measured the same quantity at 0.01557 cycles per sample on the ladder-2 tables +at every rung. The row's own table spectrum is the fallback and is printed +beside it. The raw moment is deliberate. The phi-marginalized field of a +(2, +-2) signal with two polarization weights is `|alpha kappa + beta kappa*|`, +which keeps carrier-scale sub-peaks unless the signal is circular. The raw +moment bounds the curvature of anything the band-limited primitive can make at +that amplitude. The central moment gives the face-on envelope, the widest +case. A factor sqrt(2) was proposed for the linear limit from the curvature +of `|zeta|^2` at its maximum; it is not there. The integrand is +`exp(lnL)` with `lnL = (rho^2 / 2) |zeta_hat|^2`, and for `cos^2(omega t)` +the curvature `2 omega^2` and the prefactor `rho^2 / 2` multiply to +`rho^2 omega^2`, so `sigma_t = 1 / (rho omega) = 1 / (2 pi rho f_c)`, the +raw moment, with no extra factor. The carrier fixture agrees to 5 percent +(located 0.0585 samples against predicted 0.0616 at rho 12.65), and the +reserve-scheme session's fit of the log-integrand's curvature on a linearly +polarized carrier gave measured / predicted of 1.0008, 1.0000, 0.9999 and +0.9999 at rho 12.65, 40.77, 163.08 and 652.31 (their 59b48e3e). The lattice +is sized from the narrower of the prediction and the located curvature width +at each maximum, so a row narrower than predicted sets its own spacing. The +located and plan widths are reported. `prediction_consistent` flags a +located width outside `[0.5, 4]` times the prediction. + +The ledger prints, per row: `rho_pred`, both `sigma_f`, the predicted, +located, used and plan widths, the fine spacing and block span, the live +block count and first centre, the escalation count, and the consistency flag. +`predict_time_rule` gives the pair selector the same numbers from `rho` and +`sigma_f` before any row runs. + +### Measured + +Rung 40.77: `likedata_snr40.pkl`, rho 40.8, 614-sample window, 8 rows in the +ladder's truth-centred sky box, one RTX PRO 4000 Blackwell. `peaklocal-exact` +is compared with the shipped `exact` scheme so only the time rule differs. +Policy defaults after #301, whole-window scan, `reserve_batch_rows` 1. The +exact column is the earlier run of the same rows. + +| row | exact reserve | peak-local minus exact, nat | nodes exact / peak-local | escalations exact / peak-local | wall s exact / peak-local | live blocks | +|---|---|---|---|---|---|---| +| 0 | 655.431925 | -2.3e-13 | 4905 / 1519 | 1 / 0 | 187 / 155 | 2 | +| 1 | accepted locally (575.772711) | | | | | | +| 2 | accepted locally (656.379930) | | | | | | +| 3 | 604.921535 | +7.5e-11 | 4905 / 1807 | 1 / 1 | 130 / 139 | 1 | +| 4 | accepted locally (433.744903) | | | | | | +| 5 | accepted locally (626.523450) | | | | | | +| 6 | accepted locally (591.821831) | | | | | | +| 7 | 647.095757 | +0.0e+00 | 2453 / 1807 | 0 / 1 | 44 / 92 | 2 | + +Row 0 of each scheme carries the compile. Rows 1, 2 and 6 accept locally +under the #301 capacities and show the local branch's value. The two +peak-local escalations (rows 3 and 7) are the prediction one doubling short +on those rows: the first tier's check missed 1e-3 and the second met it. + +Predicted against located width, same rows (samples; the locator's width is +the envelope's curvature at the maximum, the plan's is the local branch's +Newton width): + +| row | rho (located) | sigma_t predicted | sigma_t located | sigma_t plan | located / predicted | escalations | +|---|---|---|---|---|---|---| +| 0 | 36.7 | 0.279 | 0.332 | 0.309 | 1.19 | 0 | +| 3 | 35.4 | 0.289 | 0.181 | 0.294 | 0.63 | 1 | +| 7 | 36.6 | 0.280 | 0.305 | 0.306 | 1.09 | 1 | + +The stored Q's bandwidth is 0.01556 cycles per sample; the rows' own table +spectra give 0.0135 to 0.0150. The located width agrees with the plan's +within 40 percent and the prediction sits at or below the located width. + +Rung 163.08: `likedata_snr160.pkl`, rho 163.1, same window, rows 0 to 3 of the +seed-7 draw, one Blackwell. Under the #301 capacities every row accepts locally +(issue #308 has the values), so the decline is forced with `max_modes` 1. The +peak-local reserve has its support-limited scan here: 65 scan nodes plus +4 x 73 block nodes, 357 in all. It is compared with the whole-window exact +reserve on the rows where that exists. Row 1's is refine 32 with its refine-16 +check agreeing to 6e-10; the exact rows 0, 2 and 3 were stopped on 2026-09-09 (RO: +the exact kernel is not the one that drops out at 652) and are not run: + +| row | exact reserve | peak-local | difference, nat | nodes exact / peak-local | escalations exact / peak-local | wall s exact / peak-local | sigma_t predicted | located | plan | rho (located) | resolution error | +|---|---|---|---|---|---|---|---|---|---|---|---| +| 0 | not run | 13165.892582 | | / 357 | / 0 | / 186 | 0.0629 | 0.0619 | 0.0619 | 162.5 | 1.9e-10 | +| 1 | 12335.902574 | 12335.902574 | +3.2e-07 | 19617 / 357 | 3 / 0 | 9805 / 107 | 0.0650 | 0.0649 | 0.0649 | 157.4 | 9.1e-12 | +| 2 | not run | 10083.512460 | | / 357 | / 0 | / 107 | 0.0719 | 0.0700 | 0.0700 | 142.3 | 8.9e-07 | +| 3 | not run | 8474.729892 | | / 357 | / 0 | / 108 | 0.0786 | 0.0782 | 0.0782 | 130.2 | 0.0e+00 | + +Row 1 agrees with the converged whole-window value to every printed digit at +2 percent of its node count and 1 percent of its wall time. No row escalated: +the predicted widths sit within 3 percent of the located ones, and the located +`rho` within 20 percent of the network 163. The scan hull was 2.1 samples wide +(the maxima +-16 sigma) against the 614-sample window. + +### Locator search sizing (rung 652, 2026-09-09) + +The first `peaklocal` rows at rung 652 (`likedata_snr640.pkl`, S=8 draw, guard +128, 16 distance nodes, decline forced) refused 3 of 8 rows on the focus +certificate: the rule's argmax sat 0.18 to 0.33 samples from the block centre +(spans 0.22 to 0.37). On row 0 a 1/256-sample lattice of the psi-Laplace +kernel peaks at 307.109 (lnL 212147.7, half-maximum width 0.027 samples); the +locator had centred the block at 307.334, 107 nat lower. A sweep of the +locator on that row: + +| search refine | phi nodes | Newton steps / clip, rad | centre | value | note | +|---|---|---|---|---|---| +| 8 | 64 | 3 / 0.1 | 307.334 | 211532 | shipped #304 | +| 8 | 64 | 8 / 1.0 | 307.334 | 212067 | angles reach, time cell wrong | +| 64 | 64 | 3 / 0.1 | 306.559 | 211536 | | +| 64 | 64 | 8 / 1.0 | 307.464 | 211908 | | +| 128 | 64 | 8 / 1.0 | 307.482 | 211880 | | + +Refining the time search does not help: the search grid maximizes the angles on +a 64-node phi lattice, whose ripple is about `rho^2 (pi / n_phi)^2` = 1000 nat at +rho 652 against a 25-nat change of the profile across one search cell +(`(rho^2 / 2)(0.125 / tau)^2` with the envelope width tau about 7.6 samples), so +the search maximum lands on whichever cell the ripple favours, up to 0.4 +samples away, beyond the polish's reach of 1.33 cells. The Newton polish of +(phi, u) was clipped to 0.1 rad per step for 3 steps against a lattice offset +of up to 0.4 rad, 500 nat low on the same row. + +Sizing, from the amplitude: the phi count must hold the ripple under the +per-cell change, so `n_phi >= pi rho` for 1 nat; it is a static shape, so the +policy carries `reserve_peaklocal_search_phi_nodes` = 4096 (under 1 nat up to +rho 1300; the grid is `(t, phi, u)` = 4905 x 4096 x 8 doubles, 1.3 GB, a fraction +of a second on the GPU) and `reserve_peaklocal_newton_steps` = 8 within +`reserve_peaklocal_newton_step_max` = 1 rad. Test: a carrier at rho 632 with +tau 8 is located within tau / rho of its centre at the profile's true maximum. + +Rung 652 with the sizing, S=8 draw, `peaklocal` (psi-Laplace), decline forced: + +| row | reserve value | nodes | escalations | s per row | sigma_t predicted | located | rho located | focus offset, samples | resolution error, nat | warranted | +|---|---|---|---|---|---|---|---|---|---|---| +| 0 | 212136.145450 | 357 | 0 | 152 | 0.0157 | 0.0154 | 651.4 | 0.001 | 2.9e-11 | yes | +| 1 | 211785.826592 | 357 | 0 | 28 | 0.0157 | 0.0154 | 650.9 | 0.001 | 2.9e-11 | yes | +| 2 | 205203.742743 | 357 | 0 | 28 | 0.0160 | 0.0157 | 640.7 | 0.000 | 2.9e-11 | yes | +| 3 | 211361.706073 | 357 | 0 | 28 | 0.0157 | 0.0154 | 650.2 | 0.002 | 2.9e-11 | yes | +| 4 | 139434.603116 | 357 | 0 | 28 | 0.0194 | 0.0190 | 528.2 | 0.001 | 0.0e+00 | yes | +| 5 | 145870.675689 | 357 | 0 | 28 | 0.0189 | 0.0182 | 540.2 | 0.002 | 0.0e+00 | yes | +| 6 | 203151.295219 | 357 | 0 | 28 | 0.0160 | 0.0157 | 637.5 | 0.001 | 2.9e-11 | yes | +| 7 | 141631.499287 | 357 | 0 | 28 | 0.0192 | 0.0185 | 532.3 | 0.001 | 0.0e+00 | yes | + +Before the sizing rows 0, 1 and 3 were refused by the focus certificate +(offsets 0.334, 0.183, 0.183 samples after two escalations, 1221 nodes, +129 to 250 s); the other five rows keep their values to all printed digits. +Rung 163: all eight rows unchanged at 9 s per row. + +## Gate before this can be a default + +PR #268 warrants scalar values. Differentiating the composite differentiates +a truncated fixed-plan integral, and the JAX sampler's hill climbing and +MALA/HMC steps consume those gradients. Required before any default change: +value and gradient parity through the SNR and higher-mode ladder, on +production tables, recorded in the paper repository +(`development/OPEN_jax_direct_marginalization_policy.md`). + +## Known adversarial items + +- Acceptance was not completeness at a support boundary. On the 32-sample + synthetic window the exact lnL(t) peaks at the first sample; the planner's + boundary starts were rejected as non-stationary, the interior modes were + accepted with every diagnostic passing, and the value was 22.86 nat + against an exact 45.5. The plan now records a live start pinned to the + time or distance boundary within 30 nat of the best value, and the gate + declines on it (`decline_boundary_maximum`). A one-sided local region for + boundary maxima is the eventual fix; the decline is the fail-closed one. +- Sampler cost (wiring review): flowMC vmaps the scalar AD target over + chains, and under `vmap` a `lax.cond` with a batched predicate lowers to + `select_n`, so every chain executes the local branch and every reserve + tier whatever its own disposition. The `lax.map` batch is also sequential + in rows, so an 8000-row pilot is hours. Neither is a correctness problem; + both make the policy impractical as the sampler's target until a + host-compacted path exists (vmap the local gate, run the reserve on the + declined subset). +Items 1 and 2 below come from the adversarial review of PR #268 through this +wiring (2026-09-07) and are verified on synthetic tables only. They are the +first questions for the production-table ladder. + +- On synthetic 22-only carrier tables the base portfolio at angular + oversample 1 overflows `max_starts=32`: the 9-point phi lattice's + max-over-angles time profile ripples with the rotating carrier phase and + produces spurious time peaks, so every row declines on capacity and the + local branch never runs. Oversample 2 fits but the same tables then + decline on nested quadrature (13 vs 19 nodes, 8e-3 nat). PR #268's real + SNR-160 capture reports 4 base candidates with no overflow, so the + synthetic result does not transfer directly; the ladder must measure the + acceptance rate on production tables before the paper's "local at 160 and + 320" is quoted from this code. Suggested fix if it does transfer: rank + time peaks from the triangle envelope the time-cover step already + computes, not from the lattice profile. +- Sub-sample peaks: at 150 Hz and 4096 Hz the time peak is about 4.35/rho + samples wide, so above rho of a few tens the reserve needs refinement well + beyond 4. The escalation ceiling and the trapezoid rules address the + warrant; the cost (three dense evaluations per tier) is the ladder's to + measure. The norm lower bound also loosens with inclination (0.46 of the + norm edge-on), which can exhaust the 64 retained time nodes. +- Capacity at higher harmonic order: random m_max=4 tables show 5 to 12 + angular lattice maxima per time node against `max_modes=4`, and the u + lattice does not grow with oversample, so enrichment refines phi only. +- The error score double-charged common-mode terms. Base and enriched plans + measure the same quadrature, guard, and omitted-time discrepancies, and + the score summed both. On the analytic fixture at 10x amplitude a value + correct to 1.5e-4 nat was refused at a score of 1.06e-3 (two copies of one + 5.27e-4 guard discrepancy). Fixed in PR #278: each pair is charged once as + its maximum; the per-plan terms stay in the ledger. +- A distance peak below the prior's support pins every optimizer lane to the + boundary with an identical gradient norm growing like rho^2. PR #268 + declines such a row (`decline_no_modes`); the signature matches the one + the PR #270 ladder reported for its enriched tier, which points at a + distance-support problem in that harness rather than a refinement defect. +- The local box radius is a common-mode term. Base and enriched plans use + the same whitened radius, so mass outside the box is invisible to the + enrichment gate and to the error score. At radius 3 the accepted value on + the wiring test's analytic fixture was 0.015 nat below an independent + fine-time reference while every diagnostic passed. The default is the + library's 6; the wiring test pins the accepted value against the external + reference at 2e-3 nat, which the gate alone would not have caught. +- PR #270's host controller (`multipeak_planner`) declined on every rung of a + production ladder because its enriched tier never refined (identical + gradient norms across starts). This policy does not use that controller. + The same failure class, degenerate or unrefined enriched starts, must be + probed on this device pipeline with per-start gradient norms on production + tables before the ladder result is trusted. +- The audit ledger is evaluated on a subsample of exported rows after + sampling. It describes the exported cloud, not every evaluation the + sampler made. Unwarranted rows are not a labelling matter: they are `nan` + and stop the run. + +## Choosing the (local, reserve) pair from analysis + +RO, 2026-09-08: rely on analysis and the known physics to pick the pair, rather +than try-then-decline-then-refine. `predict_reserve_pair` computes, from the +precomputed inputs and before any row is evaluated: + +| quantity | source | decides | +|---|---|---| +| network SNR | the driver's own response-derived guess | amplitude | +| `sigma_f` | second moment of the stored Q's spectrum | peak width | +| `A = rho^2/2` vs `ANGLE_MARG_CROSSOVER_AMPLITUDE` (450) | measured crossover in `anglemarg` | exact or laplace angles | +| `sigma_t = 1/(2 pi rho sigma_f)` vs the local cover budget | `max_time_nodes` | can the LOCAL branch hold the peak | +| the same width vs the reserve's node budget | `reserve_time_refine_max` | can the whole-window reserve resolve it | + +The verdict and its reasons are printed before sampling. When no implemented +reserve is adequate the run REFUSES; it does not fall back, because a fallback +to whole-window refinement carries the rows in a method nobody chose. + +### Which bandwidth, and why it is physics not convention + +The reserve marginalizes phi exactly, so the field in time is `|zeta|` with +`zeta = alpha kappa + beta kappa*`. Face-on the carrier term vanishes and the +peak is the ENVELOPE; linearly polarized the envelope is modulated at the +carrier and each sub-peak is far narrower. Measured on a carrier fixture +(f_c = 200 Hz): + +| polarization | measured peak | matches | +|---|---|---| +| circular | 11.3 Hz equivalent | central moment, 5.6 Hz | +| linear | 309.6 Hz equivalent | raw moment, 200.1 Hz | + +So RAW is the narrowest peak the primitive can make and CENTRAL the widest. A +rule that must not under-resolve sizes on the raw one. Two errors were made +here and are recorded so they are not repeated: sizing on the central moment +(the Cramer-Rao bound is about an ESTIMATOR's variance, not how sharply the +INTEGRAND varies), and a claimed sqrt(2) correction that came from reading the +curvature of `|zeta|^2` without its `rho^2/2` prefactor. Against the actual +log-integrand, raw is exact: measured/predicted 1.0008, 1.0000, 0.9999, 0.9999 +at rho 12.65, 40.77, 163.08, 652.31. + +### The node budget is PROVISIONAL and known to be the wrong law + +The budget is points-per-sigma, i.e. an ALGEBRAIC convergence model. Measured at +rho 40.77 on 64 rows: + +| refine | nodes | warrant error | +|---|---|---| +| 4 | 2453 | 1.8e-03 .. 1.14e-02 | +| 8 | 4905 | 5e-11 .. 7.8e-09 | + +Read against a tolerance, which has since moved: refine 4 fails the 1e-3 that +was shipped when this was measured, and straddles the 1e-2 RIFT PR #301 adopted +(only the top of the range exceeds it). #301 measured the other side of the same +quantity at the same rung — escalations 2 -> 0 and 321 s -> 147 s going from +1e-3 to 1e-2, lnL moving 4.8e-12 — which is what this range predicts. Two +measurements of one effect, taken independently; neither confirms the other's +method, and together they say the escalation at this rung was being driven by +the tolerance rather than by the rule. + +Doubling improved the error by ~1e6 where an algebraic rule gives 4. That is the +trapezoid rule on a BAND-LIMITED reconstruction: spectrally accurate once the +band is resolved, `exp(-c R)` not `R^-2`. The 4905 nodes that succeeded are 67% +of what the budget demands at that rung and land five orders INSIDE tolerance. + +Consequences, and they are limits on what may be claimed: + +* A refusal produced by this budget means UNPROVEN, not shown inadequate. +* No statement about WHERE the whole-window reserve stops being adequate + follows from it. Such a claim was made and withdrawn twice, on two different + mechanisms; it is not restated here. +* The replacement is a band-resolution criterion fitted to a MEASURED + convergence law. The deciding test is rung 163.08 at refine 4, 8 and 16 -- + three points, because two fit either law. + +The warrant is what certifies a row. This budget only predicts which method to +reach for, and it must not be hardened into a threshold anyone tunes against. + +## The roster: which reserves this run may choose from + +`predict_reserve_pair` takes `available=`. That argument is not a preference +list. It says which schemes the **data and the distance quadrature** can support +at all, and it is computed before the analysis runs. + +| scheme | on the roster when | why | +|---|---|---| +| `exact` | always | what the composite dispatches (`empirical_enrichment_with_exact_reserve`) | +| `laplace` | per-sample adaptive distance quadrature is ON **and** `gh_laplace_supported_for_data` holds | the placement is derived from A0 == 0 / B1 == 0 | +| `peaklocal` | never, today | RIFT PR #304 | + +The laplace conditions are separate and both necessary. On a **static** distance +grid the laplace reserve is measured at 43.2 nats at rho 163 — that is the +grid's cost, not the scheme's, and it is why the reserve may not use it there. +The loguniform static grid is not admitted either: it is sized from the angle +amplitude and may well be adequate, but nothing has measured it. + +Two rules follow, both of which the code got wrong first: + +- The roster is checked on the **angular** branch, not only where the selector + chooses `peaklocal`. It was checked only on the branch that could never have + chosen `peaklocal` anyway, so a run with laplace off the roster still selected + laplace as soon as A cleared the crossover. +- An explicit `requested=` overrides the **analysis**, not the roster. Forcing a + scheme whose premise is absent is not an override. + +## Declared, executable, and the gap between them + +`RESERVE_SCHEME_CHOICES` is what may be named. `RESERVE_SCHEME_EXECUTABLE` is +what the composite dispatches, which is `("exact",)`. `validate_policy_config` +refuses the difference. + +Without that refusal, `PolicyConfig(reserve_scheme="laplace")` would be +accepted, printed in the policy line, and computed as exact — a field the +composite never reads is worse than a missing one, because it answers. + +The laplace table-level kernel exists (`coefficient_table_distphipsimarg_laplace`, +extracted from the fused laplace path so the two cannot drift). What is missing +is the dispatch: `empirical_enrichment_with_exact_reserve` names its kernel. +Wiring it is a change to `all_axis_peaklocal.py`, which is #304's file. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_bandlimited_distmarg.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_bandlimited_distmarg.md new file mode 100644 index 000000000..785c33bbb --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_jax_bandlimited_distmarg.md @@ -0,0 +1,300 @@ +# Band-limited time quadrature for the distance-marginalized JAX likelihood + +RECORD, 2026-09-08. Method and measured numbers for +`time_quadrature="bandlimited"` on `JAXDistanceMarginalizedLikelihood`. The code +carries the tested constants and a pointer here. It carries no numbers. + +## What changed + +`fused_log_likelihood_distmarg` used to reduce over distance on the data time +grid and hand the reduced `lnL(t)` to `_time_marginalize_terminal`, which +refuses `bandlimited`. Interpolating an already-reduced nonlinear field can +converge to the wrong function. The kernel now gathers the guarded complex +primitive, refines it, and applies the same distance quadrature at every refined +node. + +Order of operations under `bandlimited`: + +1. `_accumulate_unit(..., guard=g)` with `g = bandlimited_time_guard(npts)[1]`, + the certified guard. +2. Per row, inside `lax.map`: raised-cosine support pad, then even-extension FFT + refinement of `kappa(t)` at a curvature-derived power-of-two factor. +3. Crop to the integrated window. +4. `log sum_g exp(K x_g - R x_g^2 / 2 + log w_g)` at every remaining node, in + blocks sized from the refined row length. +5. Stable trapezoid over the original closed window. + +Step 3 comes before step 4. `reduce_fn` is pointwise in the node, so the value +is unchanged, and the distance reduction then runs on `(npts-1)*f+1` nodes +instead of the guard-padded `2*(npts+2g-1)*f`. Reducing first cost 6.50 GB for +one scalar `value_and_grad`; cropping first costs 2.56 GB. + +Certificates are three of the fixed-distance four, applied to the +distance-marginalized field: factor doubling to 1e-3 nat, remeasured peak-width +resolution, and agreement between the certified guard and half of it. The +fourth, the 15-nat endpoint gap, is off on this path; the section "The endpoint +certificate" below records why. A row that fails any certificate returns NaN and +the driver refuses the run. + +`bandlimited_time_guard` is new and is the single definition of the +`(initial, certified)` guard pair. The fixed-distance kernel, the +distance-marginalized kernel and the driver's storage-window sizing all read it. + +The curvature probe moved inside the row-local `lax.map`. It used to be computed +for the whole batch, which for the distance path would have carried an +`(S, 2*n_guarded, block)` temporary. + +`JAX_ILE_DISTMARG_GH` is not read by `fused_log_likelihood_distmarg` on either +quadrature. The per-sample Gauss-Hermite placement is implemented only in the +phi/psi-marginalized kernels, which refuse `bandlimited`. Both branches of the +distance-marginalized kernel call one helper, so a future GH branch there lands +on the refined grid as well as the coarse one. + +## Scope, and why the other wrappers still refuse + +- `JAXDistPhiMargLikelihood` and `fused_log_likelihood_distphimarg`: the phi_ref + grid sum streams one primitive per grid point through a `lax.scan` carrying + only `(S, npts)`. Primitive-first requires `(nphi, n_fine)` per row. At the + shipped `nphi = 32` and the certified factor cap that is about 200 times the + row-local budget. This is a cost limit, not a correctness gap. +- `JAXDistPsiMargLikelihood`, `JAXDistPhiPsiMargLikelihood`, the exact-angle and + Laplace schemes: these reach `anglemarg.py` coefficient-table kernels that + return an already-reduced `lnL(t)`. There is no primitive at that seam to + refine without an adapter. `time_first_peaklocal.py` prototypes one. + +## Measurement setup + +Zero-noise synthetic injection through the production `PrecomputeLikelihoodTerms`: +35 + 30 Msun, `IMRPhenomD`, H1L1, `fmin = fref = 40 Hz`, `fmax = 300 Hz`, +`deltaF = 0.25 Hz`, srate 1024 Hz, integration half-window 75 ms (npts 153, +guard 128 initial / 256 certified), distance grid 512 uniform nodes over +[50, 4000] Mpc with the Euclidean prior, evaluated at the injected angles +(1.2, -0.4, 0.7, 0.9, 2.1). Host `ldas-grid`, CPU, float64, +`~/.cache/jaxci_venv/bin/python`. + +Amplitude is set by the injected distance. The rho below is the network +optimal SNR of the zero-noise data, sqrt(sum ) over 40-300 Hz: 390 Mpc +gives rho 45.9 (H1 28.5, L1 36.0) and 48.5 Mpc gives rho 369 (H1 229, L1 289). +`extras["guess_snr"]` reads 19.92 and 160.18 on the same data, 2.305x lower: +it is the precompute's guess sqrt(sum max|Q_lm|^2 / U_lm) / 2.3, quoted as rho +in earlier versions of this record. + +## Agreement with an independent reference + +The reference reconstructs the primitive with a plain periodic zero-padded FFT, +reduces over distance with a numpy log-sum-exp, and integrates with a numpy +trapezoid. Its extension is periodic where the shipped one is an even +reflection, and its guard taper ramps on `(k+1)/(g+1)` where the shipped one +ramps on `k/g`. + +| rho | shipped bandlimited | reference (guard 512, factor 512) | difference (nat) | +|---|---|---|---| +| 19.92 | 589.877456830 | 589.877457058 | -2.3e-07 | +| 160.18 | 39193.075978776 | 39193.075993296 | -1.5e-05 | + +Reference convergence, tapered. Successive differences along each ladder: + +| ladder | rho 45.9 | rho 369 | +|---|---|---| +| factor 64→128→256→512→1024 at guard 512 | 0, 0, 0, 0 | -5.6e-02, +1.2e-03, +7e-12, 0 | +| guard 128→256→512→1024 at factor 512 | +7.4e-07, +1.8e-07, +4.5e-08 | +4.7e-05, +1.2e-05, +2.9e-06 | + +The taper is required. An untapered periodic reconstruction leaves a step at the +periodic seam, and its Gibbs ringing decays like 1/guard: + +| untapered reference, factor 128 | rho 45.9 | rho 369 | +|---|---|---| +| guard 128 | 589.929584 | 39196.403097 | +| guard 256 | 589.897169 | 39194.333390 | +| guard 512 | 589.883333 | 39193.449964 | +| guard 1024 | 589.878738 | 39193.156559 | + +Each doubling halves the residual instead of removing it. At guard 1024 the +untapered reference is still 1.3e-03 nat (rho 46) and 8.1e-02 nat (rho 369) from +the shipped value, so it certifies nothing at the tolerance this work asserts. + +The shipped value is not sitting on its own stopping tolerance. Forcing +`_TIME_ADAPTIVE_RTOL` to 1e-4 selects the same factor and returns the same +number, 39193.075978776. At 1e-5 and 1e-6 the row returns NaN, because the +doubling cannot be met within `_TIME_ADAPTIVE_FACTOR_MAX`. + +## The option changes the answer + +| rho | native Simpson | bandlimited | gap (nat) | +|---|---|---|---| +| 19.92 | 541.265029 | 589.877457 | 48.6 | +| 160.18 | 35923.702300 | 39193.075979 | 3269.4 | + +Reduce-then-refine, built explicitly in the test with the same numpy pieces, +lands 3.31 nat (rho 46) and 99.35 nat (rho 369) from the shipped value. + +Mutating `at_factor` to reduce on the coarse grid and refine the reduced field +makes the rho 46 agreement test fail by -3.366 nat, about 3400 times its +tolerance, and makes the rho 369 row return NaN. The resolution and doubling +certificates reject the wrong-order field on their own. + +## The endpoint certificate + +The fixed-distance kernel refuses a row whose refined endpoint is within 15 nat +of its peak. On the distance-marginalized field that rule rejects rows whose +integral is converged, and it rejects most blind draws. + +The reason is a floor. At every node the distance sum is at least the +far-distance prior mass, so the field never falls below about the value it takes +where the template is orthogonal to the data. A row's peak-to-endpoint contrast +is therefore bounded by its own peak height, and a fixed 15-nat gap cannot be met +by any row with a peak under about 15 nat, however well the trapezoid has +converged. The fixed-distance field has no floor: it falls to -rho^2/2 away from +the peak, so the same gap measures something there. + +Blind full-sky, isotropic-orientation draws are exactly the low-contrast rows. +Every prior-seeded driver mode evaluates them by the thousand (`--mode map` and +`nuts` pilot on 4000, `laplace-is` on `n_max/4`, `prior-mc` on `n_max`), and the +driver stops the run on one NaN. + +Measured, same injection as above at 900 Mpc (rho 19.9), 20 ms half-window, +seeds 0-3 of the driver's prior, 64 rows each: + +| certificate set | uncertified rows | +|---|---| +| all four (gap on) | 90 / 256 | +| gap off | 36 / 256 | +| fixed-distance 6-D kernel, blind distance, gap on (unchanged) | 92 / 256 | + +Every one of the 90 failed the endpoint gap and nothing else. Twenty-four of +them, six per seed, against the independent reference (periodic FFT, guard 128, +factor 512): worst disagreement 1.06e-04 nat, and all of them 130-180 nat below +the batch maximum. + +The 36 that remain with the gap off all have a detector arrival peak at or beyond +the window edge. A wrong sky moves a detector's arrival by up to 2 R_earth / c, +about 43 ms, past a 20 ms half-window. Two signatures: the trapezoid value drops +by ln 2 per doubling (the integrand is an edge sliver narrower than a sample), or +the two guards disagree at 1e-3 to 1e-2 nat with the doubling converged to 1e-14 +(structure under the taper, outside the window). Neither is a certificate the +path should drop; the window has to contain the shifts. + +With the window widened to contain the shifts, same seeds and rows, gap off: + +| half-window | gap on | gap off | 6-D kernel, gap on | +|---|---|---|---| +| 20 ms | 90 / 256 | 36 / 256 | 92 / 256 | +| 50 ms | 14 / 256 | 0 / 256 | 32 / 256 | +| 75 ms | 14 / 256 | 0 / 256 | 30 / 256 | + +The 14 rows the gap alone rejects at 50 and 75 ms agree with the reference to +1.0e-06 and 3.8e-07 nat at worst, and sit 137-169 nat below the batch maximum. + +The 6-D kernel's column is base-branch behaviour and is unchanged by this +branch. Its blind-draw stop is a separate follow-up. + +The driver's stop is unchanged. Its message now counts the failed rows in the +chunk, prints the first three, and names the window as the usual cause. + +Driver runs on the merged tree, same injection, 50 ms half-window, 32 distance +nodes, seed 3, `--n-max 400` unless stated: + +| mode | quadrature | result | +|---|---|---| +| prior-mc | bandlimited | row written, lnZ 167.4, neff 2.0, 11 s | +| map | bandlimited | row written, peak lnL 186.3, 47 s | +| laplace-is, n_max 4000 | bandlimited | refused: adapted lnZ 165.3 is 6.3 nat below the prior pilot's Markov floor, neff 6.1 | +| laplace-is, n_max 4000 | simpson | row written, lnZ 166.5, neff 16.4 | + +The laplace-is refusal is the estimator's own gate, reached after every +evaluation certified: the resolved peak is narrower than one moment-matched +Gaussian covers. It is not a certificate failure. The default mode therefore +needs a larger budget or another mode under `bandlimited` on a narrow posterior. + +The injected angles are not where this likelihood peaks. On the 900 Mpc +injection the fixed-distance lnL at the injected angles is 84.0 in the +conventional code and 84.3 in JAX. The 2-D (psi, phiref) maximum is 194.9 in +both codes, so the offset is a property of the fixture and not of the +quadrature. The fixture hands the precompute the same P as the injection. For +IMRPhenomD the template route (`SimInspiralTDModesFromPolarizations`) bakes +that P's phiref and psi into the (2,2) mode as exp(-2i phiref) exp(+4i psi) +(measured, lalsimulation 6.2.0). ILE then applies both again through Y_lm and +F. Production drivers zero P.phiref and P.psi before the precompute. +`lalsimutils` is unchanged. + +### The fixed-distance kernel (follow-up, 2026-09-08) + +The 6-D kernel kept the gap when the rows above were measured. Re-measured on +rift_O4d `d84597c2a` (92 / 256 at 20 ms, unchanged by this branch) and on this +branch with the gap switched per row: same injection, seeds 0-3 of the driver's +prior with distance, 64 rows each, ldas-grid, `~/.cache/jaxci_venv`. + +| half-window | gap on | gap off | gap alone | +|---|---|---|---| +| 20 ms | 92 / 256 | 35 / 256 | 57 | +| 50 ms | 32 / 256 | 0 / 256 | 32 | + +Twenty-four gap-only rows per window against the reference (periodic FFT, +guard twice the certified value, factor 512): worst 1.1e-04 nat at 20 ms and +6.1e-06 nat at 50 ms, all 54-128 nat below the batch maximum, at 1600-3900 Mpc. + +The floor argument has a fixed-distance twin. The field is Re kappa(t) - +rho^2/2, so its contrast is at most 2 max|kappa|, which scales with the row's +own amplitude. A far or wrong-sky draw has no 15 nat to give up, converged or +not. The gap certifies the row's amplitude, not the quadrature. + +Decision: the endpoint certificate is off on both fields. The kernel's +`endpoint_log_gap` defaults to `None`; the threshold constant stays for the +tests that pin what it rejected. The other three certificates are unchanged, +and the 35 edge-peak rows at 20 ms still stop the driver. + +The driver refuses `--mode prior-mc`, `laplace-is`, `map` and `nuts` with +`bandlimited` at parse time when `--data-integration-window-half` is below +2 R_earth / c = 42.6 ms. Those modes push full-sky draws through `eval_lnL`, +which stops on one uncertified row; the flowMC family and `multistart-nuts` +pilot through the samplers' own draw and are not refused. Before the change, +`--mode prior-mc` without distance marginalization at 50 ms stopped in its +first chunk with 30 of 400 rows failed, every one at the gap. + +Test: `test/jax/test_jax_bandlimited_6d_blind.py`, ten tests. The end-to-end +runs cover `prior-mc` and `map`; `laplace-is` at this injection walks off the +peak and fails `require_finite_evidence` at pilots of 100, 250 and 500, on the +distance-marginalized field of the unchanged base as well, so it is covered +by the parse-time test only. + +## Memory + +Peak RSS, one JAX process, same data. `vmap 8` is eight chains of +`value_and_grad`, which is the shape flowMC's MALA proposal builds. + +| path | scalar | vmap 8 | +|---|---|---| +| fixed-distance 6-D bandlimited (shipped) | 1.15 GB | 3.25 GB | +| distance-marginalized, reduce before crop | 6.50 GB | - | +| distance-marginalized, crop before reduce | 2.56 GB | 9.96 GB | + +The remaining factor over the fixed-distance path is the distance grid itself. +`_BANDLIMITED_GRID_ELEMENTS` sits near the minimum of a two-sided trade-off, so +lowering it makes things worse. At `1<<22` the `vmap 8` peak is 9.97 GB; at +`1<<18` it is 22.77 GB, because a smaller block multiplies the scan carries the +reverse pass retains, and smaller values exceed the 25 GiB per-user cgroup. + +Driver runs of `--mode flowmc --distance-marginalization`, integration +half-window 20 ms, 32 distance nodes, one training and one production loop: + +| steps (local = global) | quadrature | peak RSS | wall | result row | +|---|---|---|---|---| +| 20 | simpson | 1.43 GB | 0:21 | yes | +| 20 | bandlimited | 23.4 GB | 4:50 | yes | +| 4 | bandlimited | 6.36 GB | 2:03 | yes | +| 2 | bandlimited | 4.46 GB | 2:09 | yes | + +flowMC unrolls its per-step proposal, so the compiled graph is multiplied by the +step count. The eleven refinement branches make that graph large, and 20 steps +exceeds the 25 GiB cgroup on `ldas-grid`. The gated driver test uses 2 steps and +brackets the peak with `--n-prior-pilot` instead. + +## Files + +- `core.py`: `bandlimited_time_guard`, `_time_marginalize_reflected_primitive` + (`reduce_fn`, row-local probe, crop before reduce), + `_logsumexp_grid_scanned`, `fused_log_likelihood_distmarg`. +- `wrapper.py`: `JAXDistanceMarginalizedLikelihood` accepts the option and + publishes `time_guard_initial` and `time_guard_certified`. +- `bin/integrate_likelihood_extrinsic_jax`: `eval_lnL` failure message. +- `test/jax/test_jax_bandlimited_distmarg.py`: 20 tests, about 3.5 min on `ldas-grid` (the flowMC driver run is 100 s of it; CI deselects that one). diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md index c6a98d4d8..f29f3480e 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md @@ -42,12 +42,29 @@ response, geometric time delay, spin-(-2) spherical harmonics, the `kappa`/`rho^2` assembly, continuous time-shift interpolation, time marginalization, and **analytic distance marginalization**. +### Phase marginalization and the packed mode set + +`phase_marginalization=True` is implemented for the `(2,2)`/`(2,-2)` pair only: +the reduction conjugates the `m = -2` component of the harmonic, the antenna +response and the `rholm` timeseries, which is specific to a single +`m = +2`/`m = -2` pair. Any other mode set raises `NotImplementedError` rather +than silently dropping a mode. + +**Either packed ORDER is accepted.** The column order of `lms`, `Q`, `U` and `V` +follows the iteration order of the precompute's mode dictionary, not anything the +caller chooses, so both `[(2,2), (2,-2)]` and `[(2,-2), (2,2)]` arrive in practice; +the accumulator canonicalizes internally. Note that `U` and `V` carry the mode +index on BOTH axes -- any code reordering a packed bank by hand must permute both, +or it returns a wrong likelihood with no error. + ### 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 +`JAXExtrinsicLikelihood` (6-D, including analytic phase marginalization) and by +`JAXDistanceMarginalizedLikelihood` (5-D, the wrapper `--mode nuts`, +`multistart-nuts` and `flowmc` use under `--distance-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 @@ -59,12 +76,13 @@ 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. +Nyquist frequency. Likelihood mass at the ends of the short integration window +is covered by the guard-agreement certificate below, not by an endpoint gap. +The 15-nat endpoint gap of 2026-08-29 was switched off on 2026-09-08. A row's +peak-to-endpoint contrast is bounded by its own amplitude, so the gap rejected +every blind or far draw whatever the quadrature did. The rows it rejected +alone agree with an independent reference to 1e-4 nat (DESIGN record, "The +endpoint certificate"). 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 @@ -89,11 +107,40 @@ 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 distance reduction runs on the refined nodes, not on an interpolated +`lnL(t)`. `fused_log_likelihood_distmarg` gathers the guarded primitive, hands +the refiner the same blocked distance quadrature the Simpson path uses, and that +quadrature is evaluated inside the row-local `lax.map` at every fine node before +the trapezoid. The block count for the fine grid is derived from the refined row +length rather than inherited from `grid_block`, so a 2048x row does not scale the +working set with it. Three of the fixed-distance certificates -- factor +doubling, two-guard agreement, and remeasured resolution -- apply to the +distance-marginalized field, and the endpoint gap applies to neither. With a +full-sky prior the integration half-window must contain the detector arrival +shifts of a wrong-sky draw (up to 2 R_earth / c, 42.6 ms). A row whose arrival +peak sits at the window edge fails the doubling or guard certificate and stops +the driver. The driver therefore refuses the modes that push full-sky draws +through that stop (`prior-mc`, `laplace-is`, `map`, `nuts`) at parse time when +the half-window is below that bound. +`return_lnLt` is refused under `bandlimited`: there is no reduced field on the +data grid to return. + +The phi, psi, exact-angle, and Laplace-marginalized wrappers still refuse +`bandlimited`, and the reason is now specific rather than generic. The phi_ref +grid sum streams one primitive per grid point through a `lax.scan` that carries +only `(S, npts)`; refining first means carrying `(nphi, n_fine)` per row, which +at the shipped `nphi` and the certified factor is two orders of magnitude more +scratch than the row-local budget allows. The psi and exact-angle wrappers +reach the coefficient-table kernels in `anglemarg.py`, which return an +already-reduced `lnL(t)`: there is no primitive at that seam to refine without +an adapter. Both continue to use the unchanged Simpson default. +`time_first_peaklocal.py` contains an unwired, fixed-shape prototype of that +primitive-first composition: it reconstructs one raw complex correlation per +downstream distance/angle quadrature state, builds a certified time-cell cover, +and only then performs the nonlinear reduction on local nodes. It returns an +explicit validity ledger and changes no wrapper or CLI default. Production +wiring still needs a tighter Hermite certificate, two-guard convergence, and an +adapter from the coefficient-table angle kernels. 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 @@ -117,8 +164,33 @@ executables without dying during option parsing. `(ra, dec, psi, incl, phiref, distMpc)`. - `fused_log_likelihood_distmarg(...)` — **distance- and time-marginalized** lnL over the 5 angular parameters (regulates the amplitude degeneracy; see - below). + below). Honours `time_quadrature="bandlimited"` by refining the primitive + first. - `make_distance_grid(...)`, `JAXLikelihoodData`, `build_likelihood_data`. +- `time_first_peaklocal.py` — experimental primitive-first time-cover planner + and distance adapter; not selected by any production endpoint. +- `direct_marginalization_policy.py` — opt-in cross-axis policy + (`--direct-marginalization-policy auto`): per evaluation, the four-axis + peak-local controller of `all_axis_peaklocal.py` under its acceptance + ledger, with a reserve on decline. The reserve is a pair named by + `--direct-marginalization-reserve-scheme`: exact angles on the whole-window + refined rule (default), psi-Laplace angles on that rule, or `peaklocal`, + psi-Laplace angles on a fixed-count time rule sized from the predicted + peak width `1 / (2 pi rho sigma_f)` around maxima located on the primitive + (`peaklocal_time_reserve.py`), with the coarse scan limited to the + support of the located maxima, so the node count grows with neither rho + nor the window. + Value-only; see `DESIGN_direct_marginalization_policy.md`. +- `peaklocal_time_reserve.py` — the peak-local time rule: width prediction + from the row's table and the stored Q's bandwidth, primitive-based maxima + locator, one commensurate lattice, and the pair selector's node-count + prediction. +- `../bivariate_trig_stationary.py` — host reference for complete finite-order + `(phi_ref, 2 psi)` stationary enumeration by a Sylvester resultant and + generalized eigenproblem. It records BKK expected/found counts, + conditioning, cross-projection agreement, and supplies best-effort targets + only behind an outside-cover bound; no sampled phi grid is called + enumeration. A fixed-capacity JAX plan adapter remains future work. - `wrapper.py` — `build_data_from_precompute` (runs the production precompute + packing and returns a device-resident `JAXLikelihoodData`), and the convenience classes `JAXExtrinsicLikelihood` (6-D, value/grad/Fisher) and @@ -173,8 +245,54 @@ integral — exactly the ordering of the production `distmarg_loglikelihood`. T result is smooth, bounded, and peaks at the correct sky location, and is the right object for gradient-based exploration. +The fixed distance grid under-resolves the per-sample integrand's peak (width +`~d0/SNR`) at high SNR. The driver's `--distance-gh-nodes N` places `N` +Gauss-Hermite-style nodes centred on that peak, PER SAMPLE, resolving it to +machine precision at any SNR with a few dozen nodes; `N=0` (default) keeps the +legacy fixed grid. Equivalent to the environment variable +`JAX_ILE_DISTMARG_GH`, still honoured for compatibility (`core. +set_distmarg_gh_nodes`); the two are refused, not silently reconciled, if set +to different nonzero values. See `core.make_distance_gh` / +`core._distmarg_gh_logL` and `DESIGN_jax_distance_quadrature.md`. + ## Driver +### Value-only adaptive volume and portfolio + +The opt-in ``--sampler-method AV`` and ``--sampler-method portfolio`` paths use +the same JAX likelihood selected by ``--mode`` but do not differentiate it +during integration. Likelihood rows are evaluated in one fixed JAX shape; +``--jax-av-eval-chunk`` therefore controls accelerator memory independently of +the larger ``--n-chunk`` used to cover and contract the adaptive volume. + +Portfolio defaults to AV plus a defensive GMM member. An optional Fisher-sky +initializer pays an explicit, one-time AD cost for hill climbing and local +curvature; every integration evaluation remains value-only. A finite seed +cloud does not itself guarantee prior support. For blind/full-prior inference, +use the defensive portfolio rather than interpreting seeded standalone AV as a +global calculation. + +For deliberately local tests, AV/portfolio honor +``--limit-right-ascension``, ``--limit-declination``, ``--limit-psi``, and +``--limit-inclination`` as comma-separated sampling limits. These restrict the +domain sampled while the integrand retains the normalized full physical prior. +Consequently the evidence is the full-prior contribution from that domain; it +is not conditional on the box and must not receive an inverse-volume correction. +A widened-box repeat and a posterior edge-contact check are required before the +boxed contribution can be identified with the all-sky evidence. RA windows +that cross 0/2pi are refused because one AV hyperrectangle cannot represent the +wrapped union. + +``JAXFixedDistanceLikelihood`` provides a five-angular-coordinate view of the +six-dimensional likelihood for controlled validation problems. It can also +shift the periodic phase coordinate so a narrow mode at physical phase zero is +not split across the sampler's box boundary; exported points must be mapped +back with ``to_physical_coordinates``. +``JAXRotatedPhaseLikelihood`` supplies the conventional +``--internal-rotate-phase`` sum/difference coordinates on a redundant +``[0,4 pi)`` cover, making the leading phase--polarization ridge axis-aligned +for AV as well as for gradient samplers. + `bin/integrate_likelihood_extrinsic_jax` mirrors the ILE CLI/output conventions and uses the JAX likelihood. @@ -287,6 +405,112 @@ 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`. +### Persistent compilation cache + +The shipped driver enables JAX's cross-process compilation cache by default. +RIFT disables JAX's auxiliary per-fusion autotune cache while doing so. In JAX +0.9.2 that auxiliary cache places its absolute directory in the executable +cache key, so leaving it enabled makes an otherwise compatible exported bundle +miss after import at a different filesystem path. The persistent compiled- +executable cache remains enabled and is the transferable cache described here. +It selects a stable directory under `$RIFT_JAX_CACHE_ROOT` (or +`$XDG_CACHE_HOME/rift/jax`, normally `~/.cache/rift/jax`) and adds a +compatibility namespace derived from Python, JAX/JAXLIB, the CUDA plugin, +backend/platform version, GPU kind, and compute capability. JAX's own keys then +separate static argument shapes and compiler options inside that namespace. + +Use `--jax-cache-dir /shared/rift-jax-cache` to choose a shared root, or +`--no-jax-persistent-cache` for a diagnostic cold run. The standard +`JAX_COMPILATION_CACHE_DIR` variable remains an exact-directory expert +override. The selected directory contains its provenance manifest. +Runtime identity and durable imported-bundle profile/static-shape provenance +are stored separately. Each distinct contributing bundle gets an atomic record +keyed by its manifest digest, so neither a later ordinary startup nor a second +compatible bundle import can erase the earlier provenance. +On Condor, an unset root falls back to +`$_CONDOR_SCRATCH_DIR/.rift_cache/jax`; transfer that directory or set a shared +root to reuse it across jobs. An unwritable cache disables itself with a warning +rather than failing the ILE calculation. + +Condor scratch is job-local, so default enablement there avoids duplicate +compilation only within that job; it does not provide automatic cross-job +persistence. To reuse a survey/full-run cache, transfer the bundle as an input +and append `--jax-cache-bundle rift-o4-laplace.zip --jax-cache-profile +o4-laplace` to the ordinary ILE arguments. The driver validates and imports it +before importing modules that construct ILE JITs. Sites with a genuinely shared +writable filesystem can instead set `RIFT_JAX_CACHE_ROOT` in the submit +environment. + +Warm with the real production command, then package that active namespace and +record the important static shapes: + +```sh +integrate_likelihood_extrinsic_jax --jax-cache-dir /scratch/rift-cache \ + +rift_jax_cache --cache-root /scratch/rift-cache export rift-o4-laplace.zip \ + --profile o4-laplace --shape detectors=3 --shape l_max=2 \ + --shape n_chunk=8000 --shape distance_grid=256 --shape n_phi=8 +``` + +On a compatible target host/container, import and reuse it: + +```sh +rift_jax_cache --cache-root /shared/rift-cache import rift-o4-laplace.zip \ + --expect-profile o4-laplace +integrate_likelihood_extrinsic_jax --jax-cache-dir /shared/rift-cache \ + +``` + +Import rejects a different JAX/JAXLIB/CUDA backend, GPU kind/capability, +Python, requested profile, unexpected archive members, or checksum failure. +Different static shapes safely miss JAX's inner cache and compile normally; +the bundle's shape metadata makes those misses explainable. +Import also bounds member count, individual/total uncompressed size, and +compression ratio and streams entries through their checksum, so a corrupt or +hostile archive cannot expand without limit. Cache bundles contain compiler +artifacts and should still be accepted only from a trusted build workflow. + +The compatibility namespace does not cover `XLA_FLAGS`, and does not need to: +JAX covers it. `jax/_src/cache_key.py::_hash_xla_flags` reads the `XLA_FLAGS` +and `LIBTPU_INIT_ARGS` environment variables and every `--xla*` token in +`sys.argv`, and hashes each into the key, skipping only the dump/debug flags in +`xla_flags_to_exclude_from_cache_key`. Measured on jax 0.9.2: adding +`--xla_cpu_enable_fast_math=true` to `XLA_FLAGS` produced a second, distinct +`jit_work-*` entry rather than reusing the first, and rerunning with unchanged +flags reused it. So a numerics-affecting flag varies the key rather than +silently reusing a kernel compiled under a different one, and flags need not be +held fixed per cache root. + +Cache entries are written by JAX, not by RIFT, and `LRUCache.put` writes them +with a plain `write_bytes` rather than through a temporary file, so a reader can +observe a partial entry. Measured on jax 0.9.2 against a populated cache: an +entry truncated to 60%, an entry with one byte flipped mid-blob, and an entry +with 4 KiB of random bytes spliced in each produced the same lnL as the +uncorrupted run, via a `UserWarning: Error reading persistent compilation cache +entry` and a recompile. Corruption therefore fails CLOSED: a shared root +degrades to recompilation under contention and does not hand back a wrong +kernel. + +It does not self-heal, and that is the operational cost. `LRUCache.put` returns +early when the path already exists, so a process killed mid-write -- a node +reboot, an OOM, a thread-budget abort -- leaves a truncated entry that no later +process rewrites. In the measurement above the file stayed at its truncated +2943 bytes across reruns. That key then recompiles forever while only warning. +If a warm cache stops saving time, delete the compatibility-keyed directory +(`rift_jax_cache fingerprint` names it) and re-warm; there is no partial repair. + +The amplitude-adequacy diagnostic of the amp-sized schemes (exact, Laplace, +peak-local, phi-local) is deliberately data returned +by a pure JIT, not a `jax.debug.callback`: JAX does not persist graphs with host +callbacks. The driver synchronously accumulates the maximum over every pilot, +reweight, and final production/output-cloud batch and records that deterministic +scope in result provenance. Transient flow-training-only proposals are not +claimed; they do not enter the reported evidence or exported cloud. The +`direct-marginalization-policy` composite exposes no such metric and is +therefore unlabelled. A tripped +check still leaves likelihood values finite and labels the artifacts +`SUSPECT-ANGLE-GRID` rather than silently excising the affected region. + ## Status and next steps **Done & validated:** the AD likelihood core (1e-13 vs reference), gradients, diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py index ef02fe145..8ddb7ac09 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py @@ -47,9 +47,11 @@ build_data_from_precompute, build_rotation_data_from_precompute, build_freqresponse_data_from_precompute, + build_rotating_freqresponse_data_from_precompute, EXTRINSIC_PARAM_ORDER, ) -from .banded import build_rotation_data, build_freqresponse_data +from .banded import (build_rotation_data, build_freqresponse_data, + build_rotating_freqresponse_data) from .coordinates import ( build_network_frame, equatorial_to_network, @@ -70,8 +72,10 @@ "build_data_from_precompute", "build_rotation_data_from_precompute", "build_freqresponse_data_from_precompute", + "build_rotating_freqresponse_data_from_precompute", "build_rotation_data", "build_freqresponse_data", + "build_rotating_freqresponse_data", "EXTRINSIC_PARAM_ORDER", "build_network_frame", "equatorial_to_network", diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/all_axis_peaklocal.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/all_axis_peaklocal.py new file mode 100644 index 000000000..79e8aa56a --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/all_axis_peaklocal.py @@ -0,0 +1,2898 @@ +"""Fixed-shape multi-peak marginalization over time, polarization, phase and distance. + +This module is the device-evaluation half of the all-variable peak-local design. +It deliberately separates three jobs which must not be conflated: + +* the compact norm table produced by the upstream packed ``U,V`` contraction is + summarized once and used by the host planner to rank candidate basins; +* JAX gradients and Hessians refine supplied starts and size local boxes; +* a separate acceptance layer decides whether the local integral is usable, + either from a formal omitted-mass warrant or from an explicitly empirical + stronger discovery/quadrature enrichment with conservative reserve. Optimizer + convergence alone is never treated as proof that every mode was found. + +The integration kernel consumes a padded :class:`AllAxisModePlan`, so its live +workspace is ``O(local_order**4)`` and is independent of the dense time/angle/ +distance resolutions. Time is reconstructed from the reflected primitive only +at the local nodes. The two angular axes use the exact finite Fourier tables, and +distance uses the exact quadratic likelihood with the volumetric ``x**-4`` +Jacobian. Modes are streamed through ``lax.scan``; no ``mode x local-grid`` tensor +is retained under nested JIT/AD/vmap transformations. Values are unnormalized: +the caller-provided ``log_normalization`` must include the time-sample Jacobian +and normalized time, angle, and distance priors appropriate to its application. + +This is an explicit prototype seam, not a sampler policy. ``ok=False`` means the +caller must evaluate its dense/exact reserve and keep the sample. It never means +waveform failure and the diagnostic local value must not be substituted silently. +The primitive ``ok`` is only a scalar-value usability gate: the outside-mass +bound may be certified, but the nested quadrature comparison is validated rather +than a formal error bound. :func:`empirical_enrichment_marginalize` implements +the more practical one-step operational gate and labels its non-rigorous basis. +Although the fixed-shape kernel is compatible with outer JIT/AD, +differentiating it holds +the host plan and its regions fixed and therefore differentiates the truncated +local integral. A production gradient/Hessian consumer needs a separate omitted- +derivative warrant; this prototype does not claim one. +""" + +from typing import NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np + +from .time_first_peaklocal import (_evaluate_time_spectrum, + _time_primitive_spectrum, + spectral_time_derivative_bound) + + +__all__ = [ + "AllAxisModePlan", + "UVHarmonicSummary", + "JointStartPlan", + "DeviceJointStartPlan", + "summarize_uv_norm_table", + "rank_time_starts_from_uv", + "rank_joint_starts_from_uvq", + "rank_joint_starts_from_uvq_device", + "combine_device_start_plans", + "algebraic_angle_starts_from_uv", + "refine_all_axis_starts", + "select_refined_modes", + "mode_local_geometry", + "make_all_axis_mode_plan", + "make_all_axis_mode_plan_device", + "make_all_axis_mode_plan_pair_device", + "all_axis_peak_local_marginalize", + "empirical_enrichment_marginalize", + "empirical_enrichment_with_exact_reserve", + "empirical_enrichment_with_exact_reserve_sequential_batch", +] + + +class AllAxisModePlan(NamedTuple): + """Padded host plan consumed by the fixed-shape device kernel. + + Coordinates are ``(time_sample, phi_ref, u=2*psi, x=Dref/D)``. + ``outside_log_bound`` bounds the *unnormalized* integral outside the union + of the exact transformed regions + ``center + local_transform @ [-local_radius,local_radius]^4`` (with the two + angular coordinates interpreted periodically). A finite value is + correctness-bearing only when ``outside_bound_certified`` is true. The + axis-aligned ``half_widths`` are derived conservative enclosures used only + for support and disjointness checks; they never define the certified cover. + ``time_reconstruction_certified`` is a separate guard/seam warrant. Real + unguarded captures must leave it false; an outside-mass certificate cannot + certify the noninteger reflected-time reconstruction inside a region. + ``time_outside_log_bound`` has a deliberately narrower meaning: it bounds + the full angle/distance integral in time cells discarded before basin + localization. ``time_cover_min_sample`` and ``time_cover_max_sample`` + enclose every retained scout cell associated with that bound, permitting a + reserve to integrate the contiguous enclosing interval without silently + dropping a retained time basin. Nonfinite endpoints mean that no such + cropped reserve warrant is available. The time bound is not a bound on + missing angular modes inside retained + cells and can only augment, never replace, the global outside-cover warrant. + ``enumeration_complete`` is a separate + diagnostic statement about the supplied root set. It is deliberately not + an acceptance requirement: a missed algebraic root is scientifically + harmless when the independent bound proves that *all* mass outside the + integrated regions fits the error budget. Algebraic roots, optimizer + starts, and an outside bound have different failure modes and remain + separate here. ``discovery_capacity_ok`` freezes whether the upstream + bounded start portfolio fit without truncation; the empirical gate declines + rather than trusting a caller-supplied boolean at evaluation time. + ``boundary_maximum_pinned`` records that a live start with a competitive + value ended on the time or distance boundary of the support after + refinement and was rejected by the stationarity filter. Such a start is + a constrained maximum the local integral does not cover: on a synthetic + window whose exact lnL(t) peaks at the first sample, the plans that + ignored it accepted a value 22.7 nat below the exact reserve with every + other diagnostic passing. The gate declines on it. + """ + + centers: jax.Array + half_widths: jax.Array + local_transforms: jax.Array + local_radius: jax.Array + live: jax.Array + outside_log_bound: jax.Array + enumeration_complete: jax.Array + outside_bound_certified: jax.Array + time_reconstruction_certified: jax.Array + time_outside_log_bound: jax.Array + time_outside_bound_certified: jax.Array + time_cover_min_sample: jax.Array + time_cover_max_sample: jax.Array + boxes_disjoint: jax.Array + discovery_capacity_ok: jax.Array + boundary_maximum_pinned: jax.Array + + +class UVHarmonicSummary(NamedTuple): + """Compact structural summary of the norm table derived from ``U,V``. + + ``C_B`` is the exact harmonic table already produced upstream from the + packed self terms. This class does not claim to repeat or count that + contraction. The lower/upper and derivative entries are triangle- + inequality bounds, not fits. They are host-planning data and should be + cached outside JIT/vmap. + """ + + C_B: np.ndarray + b_lower: float + b_upper: float + phi_derivative_bound: float + u_derivative_bound: float + time_invariant: bool + time_max_deviation: float + summary_build_count: int + input_harmonic_coefficients: int + + +class JointStartPlan(NamedTuple): + """Bounded U,V/Q-informed starts for joint four-axis refinement. + + The angular lattice is sized from the exact harmonic orders and an explicit + oversampling factor, never from SNR. It is a targeting device rather than + a completeness proof. ``capacity_ok`` is false instead of silently + discarding excess candidates; the empirical controller must then enrich or + use its reserve. + """ + + starts: np.ndarray + scores: np.ndarray + time_starts: np.ndarray + time_profile: np.ndarray + n_phi_lattice: int + n_u_lattice: int + n_lattice_evaluations: int + n_exact_symmetry_shifts: int + n_candidates_before_cap: int + capacity_ok: bool + + +class DeviceJointStartPlan(NamedTuple): + """Fixed-capacity U,V/Q basin portfolio produced inside JAX. + + Unlike :class:`JointStartPlan`, every array has a static leading dimension + and can therefore cross ``jit`` and ``vmap`` boundaries. The angular + lattice is fixed by the finite harmonic degrees, not by SNR; it is a cheap + basin-placement device, not an integration grid or completeness proof. + ``capacity_ok`` is false whenever more local lattice maxima exist than fit + in ``starts``. Such a row must enrich or execute the exact reserve. + The time-cover endpoints enclose every retained scout cell and travel with + the corresponding discarded-cell integral bound. + """ + + starts: jax.Array + scores: jax.Array + live: jax.Array + n_lattice_candidates_before_symmetry: jax.Array + n_candidates_before_cap: jax.Array + capacity_ok: jax.Array + n_phi_lattice: jax.Array + n_u_lattice: jax.Array + n_time_lattice: jax.Array + n_retained_time_samples: jax.Array + n_lattice_evaluations: jax.Array + n_time_scout_evaluations: jax.Array + n_time_cells_retained: jax.Array + n_time_nodes_retained: jax.Array + time_outside_log_bound: jax.Array + time_cover_min_sample: jax.Array + time_cover_max_sample: jax.Array + time_scout_peak_lower: jax.Array + time_cover_certified: jax.Array + time_capacity_ok: jax.Array + norm_nonnegative: jax.Array + n_exact_symmetry_shifts: jax.Array + + +def _kp_weights_numpy(n): + out = np.ones(int(n), dtype=float) + out[1:] = 2.0 + return out + + +def summarize_uv_norm_table(C_B_t, *, invariance_atol=1.0e-10): + """Collapse the ``U,V``-derived norm table and form exact harmonic bounds. + + ``C_B_t`` may be ``(KP,2KS+1)`` or the historical + ``(KP,2KS+1,Ntime)`` table. Ordinary (non-rotation) ILE has a + time-independent norm; the latter representation repeats it at every time. + Arrival-time-dependent input is reported in ``time_invariant`` and must make + a peak-local plan decline rather than being averaged away. + """ + table = np.asarray(C_B_t, dtype=np.complex128) + if table.ndim == 2: + base = table + deviation = 0.0 + elif table.ndim == 3: + base = table[..., 0] + deviation = float(np.max(np.abs(table - base[..., None]))) + else: + raise ValueError("C_B_t must have shape (KP,2KS+1[,Ntime])") + scale = max(1.0, float(np.max(np.abs(base)))) + invariant = bool(np.isfinite(deviation) + and deviation <= float(invariance_atol) * scale) + + kp = np.arange(base.shape[0], dtype=float)[:, None] + ks_max = (base.shape[1] - 1) // 2 + ks = np.arange(-ks_max, ks_max + 1, dtype=float)[None, :] + weight = _kp_weights_numpy(base.shape[0])[:, None] + magnitude = weight * np.abs(base) + centre = float(base[0, ks_max].real) + remainder = float(np.sum(magnitude) - abs(base[0, ks_max])) + # B= is non-negative. Combining that identity with the harmonic + # triangle inequality makes the lower bound tighter but never optimistic. + b_lower = max(0.0, centre - remainder) + b_upper = abs(centre) + remainder + m_phi = float(np.sum(magnitude * np.abs(kp))) + m_u = float(np.sum(magnitude * np.abs(ks))) + return UVHarmonicSummary( + np.ascontiguousarray(base), b_lower, b_upper, m_phi, m_u, + invariant, deviation, 1, int(table.size)) + + +def rank_time_starts_from_uv(C_A_t, uv_summary, x_min, x_max, *, + max_starts=16, min_separation=2): + """Rank time basins with a true ``U,V``-informed likelihood envelope. + + For every retained time sample, ``A_upper=sum w_k |C_A|`` bounds the data + term over both angles. ``uv_summary.b_lower`` bounds the norm from below, + so maximizing ``x*A_upper - B_lower*x**2/2`` on the physical distance + interval gives an upper envelope. The volumetric ``-4 log(x)`` term is + separately bounded at ``x_min``. This ranks starts cheaply; it does *not* + certify that unselected time cells are negligible. That remains the + outside-mass warrant in :class:`AllAxisModePlan`. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + if C_A_t.ndim != 3: + raise ValueError("C_A_t must have shape (KP,2KS+1,Ntime)") + if not isinstance(uv_summary, UVHarmonicSummary): + raise TypeError("uv_summary must come from summarize_uv_norm_table") + if not uv_summary.time_invariant: + raise ValueError("arrival-time-dependent U,V norm cannot be collapsed") + x_min, x_max = float(x_min), float(x_max) + if not (0.0 < x_min < x_max): + raise ValueError("need 0 < x_min < x_max") + max_starts = int(max_starts) + min_separation = int(min_separation) + if max_starts < 1 or min_separation < 0: + raise ValueError("invalid start-count policy") + + weight = _kp_weights_numpy(C_A_t.shape[0])[:, None, None] + a_upper = np.sum(weight * np.abs(C_A_t), axis=(0, 1)) + if uv_summary.b_lower > 0.0: + x_star = np.clip(a_upper / uv_summary.b_lower, x_min, x_max) + else: + x_star = np.full_like(a_upper, x_max) + envelope = (x_star * a_upper + - 0.5 * uv_summary.b_lower * np.square(x_star) + - 4.0 * np.log(x_min)) + + # Endpoints are legitimate boundary basins. Interior starts are drawn only + # from local maxima, then ranked by the structural upper envelope. + is_peak = np.ones(envelope.size, dtype=bool) + if envelope.size > 2: + is_peak[1:-1] = ((envelope[1:-1] >= envelope[:-2]) + & (envelope[1:-1] >= envelope[2:])) + candidates = np.flatnonzero(is_peak) + candidates = candidates[np.argsort(envelope[candidates])[::-1]] + selected = [] + for index in candidates: + if all(abs(int(index) - old) > min_separation for old in selected): + selected.append(int(index)) + if len(selected) == max_starts: + break + if not selected: + selected = [int(np.argmax(envelope))] + return np.asarray(selected, dtype=np.int32), envelope + + +def _distance_profile_numpy(A, B, x_min, x_max): + """Maximize ``x*A-x**2*B/2-4log(x)`` on a finite interval.""" + A = np.asarray(A, dtype=float) + B = np.asarray(B, dtype=float) + scale = max(1.0, float(np.max(np.abs(B)))) + if np.min(B) < -1.0e-9 * scale: + raise ValueError("U,V norm table is negative on the planning lattice") + B = np.maximum(B, 0.0) + x0 = np.full_like(A, float(x_min)) + x1 = np.full_like(A, float(x_max)) + + def value(x): + return x * A - 0.5 * B * x * x - 4.0 * np.log(x) + + v0, v1 = value(x0), value(x1) + choose_hi = v1 > v0 + best_x = np.where(choose_hi, x1, x0) + best_v = np.where(choose_hi, v1, v0) + discriminant = A * A - 16.0 * B + valid = (B > 0.0) & (discriminant >= 0.0) + root = np.where( + valid, + (A + np.sqrt(np.maximum(discriminant, 0.0))) + / np.where(B > 0.0, 2.0 * B, 1.0), + x0) + valid &= (root >= float(x_min)) & (root <= float(x_max)) + root_safe = np.where(valid, root, x0) + root_value = value(root_safe) + improve = valid & (root_value > best_v) + return (np.where(improve, root_value, best_v), + np.where(improve, root_safe, best_x)) + + +def _harmonic_lattice(table, n_phi, n_u): + """Evaluate a stored real Fourier half-plane on a periodic lattice.""" + table = np.asarray(table, dtype=np.complex128) + kp = np.arange(table.shape[0], dtype=float) + ks = np.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, dtype=float) + weight = _kp_weights_numpy(table.shape[0]) + phi = 2.0 * np.pi * np.arange(int(n_phi), dtype=float) / int(n_phi) + u = 2.0 * np.pi * np.arange(int(n_u), dtype=float) / int(n_u) + ep = weight[None, :] * np.exp(1j * phi[:, None] * kp[None, :]) + eu = np.exp(1j * u[:, None] * ks[None, :]) + if table.ndim == 2: + result = np.einsum("pk,uq,kq->pu", ep, eu, table, + optimize=True).real + elif table.ndim == 3: + result = np.einsum("pk,uq,kqt->put", ep, eu, table, + optimize=True).real + else: + raise ValueError("harmonic table must have shape (KP,2KS+1[,Ntime])") + return phi, u, result + + +def _exact_angular_translation_symmetries(C_A_t, C_B, *, rtol=1.0e-10): + """Find common coefficient-certified translations on a degree grid.""" + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + C_B = np.asarray(C_B, dtype=np.complex128) + k_phi = max(C_A_t.shape[0] - 1, C_B.shape[0] - 1) + k_u = max((C_A_t.shape[1] - 1) // 2, (C_B.shape[1] - 1) // 2) + n_phi = max(1, 2 * k_phi) + n_u = max(1, 2 * k_u) + + def invariant(table, dphi, du): + kp = np.arange(table.shape[0], dtype=float)[:, None] + ks = np.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, dtype=float)[None, :] + phase = np.exp(1j * (kp * dphi + ks * du)) + if table.ndim == 3: + phase = phase[..., None] + scale = max(float(np.max(np.abs(table))), 1.0) + return (float(np.max(np.abs(table * (phase - 1.0)))) + <= float(rtol) * scale) + + shifts = [] + for i in range(n_phi): + dphi = 2.0 * np.pi * i / n_phi + for j in range(n_u): + du = 2.0 * np.pi * j / n_u + if invariant(C_A_t, dphi, du) and invariant(C_B, dphi, du): + shifts.append((dphi, du)) + return np.asarray(shifts, dtype=float).reshape((-1, 2)) + + +def rank_joint_starts_from_uvq( + C_A_t, uv_summary, x_min, x_max, *, time_guard=0, + max_time_starts=4, max_starts=64, min_time_separation=2, + angular_oversample=2): + """Build a bounded distance-following start set from U,V and Q tables. + + The exact U,V norm harmonics and Q data harmonics are evaluated on a lattice + sized by their finite polynomial degrees. Distance is profiled analytically + at each lattice point. Only angular local maxima at the highest-ranked time + basins become starts, followed by coefficient-certified translation orbits. + + This procedure performs no sampled ``delta lnL`` pruning: a legitimate peak + becomes arbitrarily narrow with SNR and may lie between coarse nodes. The + lattice is for basin placement, not likelihood integration. Increasing + ``angular_oversample`` and the bounded capacities defines the independent + enrichment used by the operational convergence gate. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + if not isinstance(uv_summary, UVHarmonicSummary): + raise TypeError("uv_summary must come from summarize_uv_norm_table") + _validate_tables(C_A_t, uv_summary.C_B) + if not uv_summary.time_invariant: + raise ValueError("arrival-time-dependent U,V norm cannot be collapsed") + if not (0.0 < float(x_min) < float(x_max)): + raise ValueError("need 0 < x_min < x_max") + time_guard = int(time_guard) + n_time = C_A_t.shape[-1] - 2 * time_guard + if time_guard < 0 or n_time < 2: + raise ValueError("time_guard must leave at least two target samples") + if min(int(max_time_starts), int(max_starts)) < 1: + raise ValueError("start capacities must be positive") + angular_oversample = int(angular_oversample) + if angular_oversample < 1: + raise ValueError("angular_oversample must be positive") + target = (C_A_t if time_guard == 0 + else C_A_t[..., time_guard:-time_guard]) + + k_phi = uv_summary.C_B.shape[0] - 1 + k_u = (uv_summary.C_B.shape[1] - 1) // 2 + n_phi = max(9, 2 * angular_oversample * k_phi + 1) + n_u = max(9, 2 * angular_oversample * k_u + 1) + phi, u, A = _harmonic_lattice(target, n_phi, n_u) + _, _, B = _harmonic_lattice(uv_summary.C_B, n_phi, n_u) + profile, x_best = _distance_profile_numpy( + A, B[..., None], float(x_min), float(x_max)) + time_profile = np.max(profile, axis=(0, 1)) + + peak_t = np.ones(time_profile.size, dtype=bool) + if time_profile.size > 2: + peak_t[1:-1] = ((time_profile[1:-1] >= time_profile[:-2]) + & (time_profile[1:-1] >= time_profile[2:])) + candidates_t = np.flatnonzero(peak_t) + candidates_t = candidates_t[np.argsort(time_profile[candidates_t])[::-1]] + time_starts = [] + for time_index in candidates_t: + if all(abs(int(time_index) - old) > int(min_time_separation) + for old in time_starts): + time_starts.append(int(time_index)) + if len(time_starts) == int(max_time_starts): + break + if not time_starts: + time_starts = [int(np.argmax(time_profile))] + + raw = [] + for time_index in time_starts: + surface = profile[..., time_index] + local = np.ones(surface.shape, dtype=bool) + for dphi in (-1, 0, 1): + for du in (-1, 0, 1): + if dphi or du: + local &= surface >= np.roll( + np.roll(surface, dphi, axis=0), du, axis=1) + angular_indices = np.argwhere(local) + if not len(angular_indices): + angular_indices = np.asarray([ + np.unravel_index(np.argmax(surface), surface.shape)]) + for iphi, iu in angular_indices: + raw.append(( + float(surface[iphi, iu]), + (float(time_index), float(phi[iphi]), float(u[iu]), + float(x_best[iphi, iu, time_index])))) + + shifts = _exact_angular_translation_symmetries( + target, uv_summary.C_B) + orbit = [] + for score, start in raw: + for dphi, du in shifts: + candidate = ( + start[0], (start[1] + dphi) % (2.0 * np.pi), + (start[2] + du) % (2.0 * np.pi), start[3]) + if not any( + abs(candidate[0] - old[1][0]) <= 1.0e-10 + and _periodic_distance(candidate[1], old[1][1]) <= 1.0e-10 + and _periodic_distance(candidate[2], old[1][2]) <= 1.0e-10 + and abs(candidate[3] - old[1][3]) <= 1.0e-10 + for old in orbit): + orbit.append((score, candidate)) + orbit.sort(key=lambda item: item[0], reverse=True) + n_candidates = len(orbit) + capacity_ok = n_candidates <= int(max_starts) + kept = orbit[:int(max_starts)] + return JointStartPlan( + np.asarray([item[1] for item in kept], dtype=float).reshape((-1, 4)), + np.asarray([item[0] for item in kept], dtype=float), + np.asarray(time_starts, dtype=np.int32), time_profile, + int(n_phi), int(n_u), int(n_phi * n_u * n_time), + int(len(shifts)), int(n_candidates), bool(capacity_ok)) + + +def _harmonic_lattice_device(table, n_phi, n_u): + """JAX counterpart of :func:`_harmonic_lattice` for fixed-shape plans.""" + table = jnp.asarray(table, dtype=jnp.complex128) + kp = jnp.arange(table.shape[0], dtype=jnp.float64) + ks = jnp.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, dtype=jnp.float64) + weight = jnp.where(kp == 0.0, 1.0, 2.0) + phi = (2.0 * jnp.pi / float(n_phi) + * jnp.arange(int(n_phi), dtype=jnp.float64)) + u = (2.0 * jnp.pi / float(n_u) + * jnp.arange(int(n_u), dtype=jnp.float64)) + ep = weight[None, :] * jnp.exp(1j * phi[:, None] * kp[None, :]) + eu = jnp.exp(1j * u[:, None] * ks[None, :]) + if table.ndim == 2: + field = jnp.einsum("pk,uq,kq->pu", ep, eu, table).real + elif table.ndim == 3: + field = jnp.einsum("pk,uq,kqt->put", ep, eu, table).real + else: + raise ValueError("harmonic table must have shape (KP,2KS+1[,Ntime])") + return phi, u, field + + +def _distance_profile_device(A, B, x_min, x_max): + """JAX support-aware distance maximum used only to rank basins.""" + A = jnp.asarray(A, dtype=jnp.float64) + B = jnp.asarray(B, dtype=jnp.float64) + B_safe = jnp.maximum(B, 0.0) + x0 = jnp.full_like(A, float(x_min)) + x1 = jnp.full_like(A, float(x_max)) + + def value(x): + return x * A - 0.5 * B_safe * x * x - 4.0 * jnp.log(x) + + v0, v1 = value(x0), value(x1) + choose_hi = v1 > v0 + best_x = jnp.where(choose_hi, x1, x0) + best_v = jnp.where(choose_hi, v1, v0) + discriminant = A * A - 16.0 * B_safe + valid = (B_safe > 0.0) & (discriminant >= 0.0) + root = ((A + jnp.sqrt(jnp.maximum(discriminant, 0.0))) + / jnp.where(B_safe > 0.0, 2.0 * B_safe, 1.0)) + valid &= (root >= float(x_min)) & (root <= float(x_max)) + root_safe = jnp.where(valid, root, x0) + root_value = value(root_safe) + improve = valid & (root_value > best_v) + return (jnp.where(improve, root_value, best_v), + jnp.where(improve, root_safe, best_x)) + + +def _distance_upper_profile_device(A_upper, B_lower, x_min, x_max): + """Maximize a likelihood upper envelope with a possibly negative B bound.""" + A_upper = jnp.asarray(A_upper, dtype=jnp.float64) + B_lower = jnp.asarray(B_lower, dtype=jnp.float64) + x0 = jnp.full_like(A_upper, float(x_min)) + x1 = jnp.full_like(A_upper, float(x_max)) + + def value(x): + return x * A_upper - 0.5 * B_lower * x * x - 4.0 * jnp.log(x) + + v0, v1 = value(x0), value(x1) + best = jnp.maximum(v0, v1) + discriminant = A_upper * A_upper - 16.0 * B_lower + valid = (B_lower > 0.0) & (discriminant >= 0.0) + root = ((A_upper + jnp.sqrt(jnp.maximum(discriminant, 0.0))) + / jnp.where(B_lower > 0.0, 2.0 * B_lower, 1.0)) + valid &= (root >= float(x_min)) & (root <= float(x_max)) + return jnp.where(valid, jnp.maximum(best, value(root)), best) + + +def _time_cell_cover_device( + C_A_t, C_B, x_min, x_max, *, time_guard, keep_nats, + scout_size): + """Certify omitted full-angle/distance mass for discarded time cells. + + A constant-size angular scout supplies only a lower reference used to choose + cells. Correctness comes instead from a spectral derivative bound on every + Q coefficient, the angular triangle inequality, and a triangle lower bound + on the U,V norm polynomial. Thus a poor scout can retain extra cells but + cannot make the omitted-time integral optimistic. + """ + time_guard = int(time_guard) + scout_size = int(scout_size) + target = (C_A_t if time_guard == 0 + else C_A_t[..., time_guard:-time_guard]) + n_time = target.shape[-1] + _, _, scout_A = _harmonic_lattice_device( + target, scout_size, scout_size) + _, _, scout_B = _harmonic_lattice_device( + C_B, scout_size, scout_size) + scout_profile, _ = _distance_profile_device( + scout_A, scout_B[..., None], float(x_min), float(x_max)) + scout_peak_lower = jnp.max(scout_profile) + + kp_weight = jnp.where( + jnp.arange(C_A_t.shape[0]) == 0, 1.0, 2.0) + lane_derivative = spectral_time_derivative_bound( + C_A_t.reshape((-1, C_A_t.shape[-1])), 1.0, + guard=time_guard, order=1) + coefficient_derivative_bound = jnp.sum( + kp_weight[:, None] + * lane_derivative.reshape(C_A_t.shape[:-1])) + a_upper_node = jnp.sum( + kp_weight[:, None, None] * jnp.abs(target), axis=(0, 1)) + a_cell_upper = jnp.minimum( + a_upper_node[:-1] + coefficient_derivative_bound, + a_upper_node[1:] + coefficient_derivative_bound) + + ks0 = (C_B.shape[1] - 1) // 2 + b_weight = jnp.where( + jnp.arange(C_B.shape[0]) == 0, 1.0, 2.0)[:, None] + b_magnitude = b_weight * jnp.abs(C_B) + b_centre = C_B[0, ks0].real + b_remainder = jnp.sum(b_magnitude) - jnp.abs(C_B[0, ks0]) + b_triangle_lower = b_centre - b_remainder + cell_peak_upper = _distance_upper_profile_device( + a_cell_upper, jnp.full_like(a_cell_upper, b_triangle_lower), + float(x_min), float(x_max)) + cell_mass_upper = (cell_peak_upper + + jnp.log((2.0 * jnp.pi) ** 2 + * (float(x_max) - float(x_min)))) + finite = (jnp.all(jnp.isfinite(C_A_t.real)) + & jnp.all(jnp.isfinite(C_A_t.imag)) + & jnp.all(jnp.isfinite(C_B.real)) + & jnp.all(jnp.isfinite(C_B.imag)) + & jnp.isfinite(coefficient_derivative_bound) + & jnp.isfinite(scout_peak_lower) + & jnp.all(jnp.isfinite(cell_mass_upper))) + live_cells = cell_mass_upper >= scout_peak_lower - float(keep_nats) + # Invalid arithmetic retains the full time support and then fails the + # explicit certificate/capacity gates downstream. + live_cells = jnp.where(finite, live_cells, jnp.ones_like(live_cells)) + outside_log_bound = jax.scipy.special.logsumexp(jnp.where( + live_cells, -jnp.inf, cell_mass_upper)) + live_nodes = jnp.concatenate(( + live_cells[:1], live_cells[:-1] | live_cells[1:], live_cells[-1:])) + return { + "live_cells": live_cells, + "live_nodes": live_nodes, + "cell_mass_upper": cell_mass_upper, + "outside_log_bound": outside_log_bound, + "scout_peak_lower": scout_peak_lower, + "coefficient_derivative_bound": coefficient_derivative_bound, + "b_triangle_lower": b_triangle_lower, + "certified": finite, + "n_scout_evaluations": jnp.asarray(scout_size * scout_size * n_time), + } + + +def _exact_angular_translation_symmetries_device( + C_A_t, C_B, *, rtol=1.0e-10): + """Return a fixed grid of coefficient-certified translations and a mask.""" + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + k_phi = max(C_A_t.shape[0] - 1, C_B.shape[0] - 1) + k_u = max((C_A_t.shape[1] - 1) // 2, + (C_B.shape[1] - 1) // 2) + n_phi = max(1, 2 * k_phi) + n_u = max(1, 2 * k_u) + dphi = (2.0 * jnp.pi / float(n_phi) + * jnp.arange(n_phi, dtype=jnp.float64)) + du = (2.0 * jnp.pi / float(n_u) + * jnp.arange(n_u, dtype=jnp.float64)) + DPHI, DU = jnp.meshgrid(dphi, du, indexing="ij") + shifts = jnp.stack((DPHI.reshape(-1), DU.reshape(-1)), axis=1) + + def invariant(table): + kp = jnp.arange(table.shape[0], dtype=jnp.float64) + ks = jnp.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, + dtype=jnp.float64) + phase = jnp.exp(1j * ( + shifts[:, 0, None, None] * kp[None, :, None] + + shifts[:, 1, None, None] * ks[None, None, :])) + if table.ndim == 3: + residual = jnp.max(jnp.abs( + table[None, ...] * (phase[..., None] - 1.0)), axis=(1, 2, 3)) + else: + residual = jnp.max(jnp.abs( + table[None, ...] * (phase - 1.0)), axis=(1, 2)) + scale = jnp.maximum(1.0, jnp.max(jnp.abs(table))) + return residual <= float(rtol) * scale + + live = invariant(C_A_t) & invariant(C_B) + return shifts, live + + +def rank_joint_starts_from_uvq_device( + C_A_t, C_B, x_min, x_max, *, time_guard=0, max_starts=32, + angular_oversample=2, norm_rtol=1.0e-9, + symmetry_rtol=1.0e-10, time_keep_nats=30.0, + max_time_nodes=64, time_scout_size=4): + """Rank a static U,V/Q basin portfolio inside ``jit``/``vmap``. + + The finite harmonic orders determine a small ``(phi_ref, 2*psi)`` lattice. + A constant-size angular scout and a spectral coefficient derivative bound + first retain complete time cells and certify an upper bound on all discarded + time-cell mass. The full harmonic lattice is then evaluated only at a + fixed-capacity set of retained time nodes; only joint angular maxima at + time-profile maxima become optimizer starts. + This is deliberately the device analogue of + :func:`rank_joint_starts_from_uvq`, with a fixed padded result rather than a + variable host list. It performs no likelihood-drop pruning and never claims + completeness. Capacity overflow or a negative reconstructed norm is + explicit and must force enrichment/reserve downstream. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + _validate_tables(C_A_t, C_B) + if not (0.0 < float(x_min) < float(x_max)): + raise ValueError("need 0 < x_min < x_max") + time_guard = int(time_guard) + n_time = C_A_t.shape[-1] - 2 * time_guard + if time_guard < 0 or n_time < 2: + raise ValueError("time_guard must leave at least two target samples") + max_starts = int(max_starts) + angular_oversample = int(angular_oversample) + max_time_nodes = int(max_time_nodes) + time_scout_size = int(time_scout_size) + if max_starts < 1 or angular_oversample < 1: + raise ValueError("start capacity and angular oversampling must be positive") + if max_time_nodes < 2 or time_scout_size < 1: + raise ValueError("time capacity and scout size must be positive") + if not np.isfinite(float(time_keep_nats)) or float(time_keep_nats) <= 0.0: + raise ValueError("time_keep_nats must be finite and positive") + if not np.isfinite(float(norm_rtol)) or float(norm_rtol) < 0.0: + raise ValueError("norm_rtol must be finite and non-negative") + if not np.isfinite(float(symmetry_rtol)) or float(symmetry_rtol) < 0.0: + raise ValueError("symmetry_rtol must be finite and non-negative") + target = (C_A_t if time_guard == 0 + else C_A_t[..., time_guard:-time_guard]) + k_phi = max(C_A_t.shape[0] - 1, C_B.shape[0] - 1) + k_u = max((C_A_t.shape[1] - 1) // 2, + (C_B.shape[1] - 1) // 2) + n_phi = max(9, 2 * angular_oversample * k_phi + 1) + n_u = max(9, 2 * angular_oversample * k_u + 1) + time_capacity = min(max_time_nodes, n_time) + n_lattice = n_phi * n_u * time_capacity + if max_starts > n_lattice: + raise ValueError("max_starts exceeds the structural lattice size") + + time_cover = _time_cell_cover_device( + C_A_t, C_B, float(x_min), float(x_max), time_guard=time_guard, + keep_nats=float(time_keep_nats), scout_size=time_scout_size) + n_live_time_nodes = jnp.count_nonzero(time_cover["live_nodes"]) + time_capacity_ok = n_live_time_nodes <= time_capacity + cell_upper = time_cover["cell_mass_upper"] + node_priority = jnp.maximum( + jnp.concatenate((jnp.asarray([-jnp.inf]), cell_upper)), + jnp.concatenate((cell_upper, jnp.asarray([-jnp.inf])))) + node_priority = jnp.where( + time_cover["live_nodes"], node_priority, -jnp.inf) + selected_priority, selected_time_index = jax.lax.top_k( + node_priority, time_capacity) + selected_time_live = jnp.isfinite(selected_priority) + # Restore chronological order so adjacency in the compact vector retains + # its time meaning. Inactive padding sorts after every physical sample. + sort_key = jnp.where( + selected_time_live, selected_time_index, + n_time + jnp.arange(time_capacity)) + chronological = jnp.argsort(sort_key) + selected_time_index = selected_time_index[chronological] + selected_time_live = selected_time_live[chronological] + selected_time_index_safe = jnp.where( + selected_time_live, selected_time_index, 0) + selected_target = jnp.take( + target, selected_time_index_safe, axis=-1) + + phi, u, A = _harmonic_lattice_device(selected_target, n_phi, n_u) + _, _, B = _harmonic_lattice_device(C_B, n_phi, n_u) + profile, x_best = _distance_profile_device( + A, B[..., None], float(x_min), float(x_max)) + b_scale = jnp.maximum(1.0, jnp.max(jnp.abs(B))) + norm_nonnegative = jnp.min(B) >= -float(norm_rtol) * b_scale + + # The structural time profile identifies basins without resolving their + # SNR-narrow interior. The continuous refiner performs that second job. + time_profile = jnp.max(profile, axis=(0, 1)) + compact_index = jnp.arange(time_capacity) + left_adjacent = ( + selected_time_live + & (compact_index > 0) + & jnp.roll(selected_time_live, 1) + & (selected_time_index == jnp.roll(selected_time_index, 1) + 1)) + right_adjacent = ( + selected_time_live + & (compact_index + 1 < time_capacity) + & jnp.roll(selected_time_live, -1) + & (jnp.roll(selected_time_index, -1) == selected_time_index + 1)) + time_left = jnp.where(left_adjacent, jnp.roll(time_profile, 1), -jnp.inf) + time_right = jnp.where( + right_adjacent, jnp.roll(time_profile, -1), -jnp.inf) + time_peak = (time_profile >= time_left) & (time_profile >= time_right) + angular_peak = jnp.ones(profile.shape, dtype=bool) + for dphi in (-1, 0, 1): + for du in (-1, 0, 1): + if dphi or du: + angular_peak &= profile >= jnp.roll( + jnp.roll(profile, dphi, axis=0), du, axis=1) + candidate = (angular_peak & time_peak[None, None, :] + & selected_time_live[None, None, :] & norm_nonnegative) + n_lattice_candidates = jnp.count_nonzero(candidate) + ranked = jnp.where(candidate, profile, -jnp.inf).reshape(-1) + scores, flat = jax.lax.top_k(ranked, max_starts) + live = jnp.isfinite(scores) + compact_time_index = flat % time_capacity + angular_flat = flat // time_capacity + u_index = angular_flat % n_u + phi_index = angular_flat // n_u + representative_starts = jnp.stack(( + selected_time_index_safe[compact_time_index].astype(jnp.float64), + phi[phi_index], u[u_index], + x_best.reshape(-1)[flat]), axis=1) + fallback = jnp.asarray([ + 0.5 * (n_time - 1.0), 0.0, 0.0, + 0.5 * (float(x_min) + float(x_max))]) + representative_starts = jnp.where( + live[:, None], representative_starts, fallback[None, :]) + + # Odd degree-sized targeting lattices need not contain an exact symmetry + # translate of their best representative. Complete its orbit from the + # coefficients themselves; otherwise base and enrichment can agree on the + # same one-quarter quadrupole cover. The fixed expansion is small + # (<=64 shifts through m_max=4), and overflow remains an explicit decline. + shifts, shift_live = _exact_angular_translation_symmetries_device( + target, C_B, rtol=float(symmetry_rtol)) + expanded = jnp.broadcast_to( + representative_starts[:, None, :], + (max_starts, shifts.shape[0], 4)) + expanded = expanded.at[..., 1:3].set(jnp.mod( + expanded[..., 1:3] + shifts[None, :, :], 2.0 * jnp.pi)) + expanded_live = live[:, None] & shift_live[None, :] + expanded_scores = jnp.where( + expanded_live, scores[:, None], -jnp.inf).reshape(-1) + final_scores, final_index = jax.lax.top_k( + expanded_scores, max_starts) + starts = expanded.reshape((-1, 4))[final_index] + live = jnp.isfinite(final_scores) + starts = jnp.where(live[:, None], starts, fallback[None, :]) + n_symmetry = jnp.count_nonzero(shift_live) + n_candidates = n_lattice_candidates * n_symmetry + live_cell_index = jnp.arange(n_time - 1, dtype=jnp.int32) + has_live_time_cell = jnp.any(time_cover["live_cells"]) + time_cover_min = jnp.where( + has_live_time_cell, + jnp.min(jnp.where( + time_cover["live_cells"], live_cell_index, n_time)), + jnp.nan) + time_cover_max = jnp.where( + has_live_time_cell, + jnp.max(jnp.where( + time_cover["live_cells"], live_cell_index, -1)) + 1, + jnp.nan) + return DeviceJointStartPlan( + starts, final_scores, live, n_lattice_candidates, n_candidates, + (norm_nonnegative & time_cover["certified"] & time_capacity_ok + & (n_candidates <= max_starts)), + jnp.asarray(n_phi), jnp.asarray(n_u), jnp.asarray(time_capacity), + jnp.asarray(n_time), jnp.asarray(n_lattice), + time_cover["n_scout_evaluations"], + jnp.count_nonzero(time_cover["live_cells"]), n_live_time_nodes, + time_cover["outside_log_bound"], time_cover_min, time_cover_max, + time_cover["scout_peak_lower"], + time_cover["certified"], time_capacity_ok, + norm_nonnegative, n_symmetry) + + +def combine_device_start_plans(base, extra): + """Form a stronger fixed portfolio that contains every base start. + + Duplicates are intentionally retained here and removed only after continuous + refinement, where basin identity is meaningful. This construction makes + the empirical controller's nesting premise structural: the stronger pass + cannot silently omit a base optimizer start. Either input overflow remains + a fail-closed capacity flag. + """ + if not isinstance(base, DeviceJointStartPlan): + raise TypeError("base must be DeviceJointStartPlan") + if not isinstance(extra, DeviceJointStartPlan): + raise TypeError("extra must be DeviceJointStartPlan") + for name, plan in (("base", base), ("extra", extra)): + if (plan.starts.ndim != 2 or plan.starts.shape[1] != 4 + or plan.scores.shape != (plan.starts.shape[0],) + or plan.live.shape != (plan.starts.shape[0],)): + raise ValueError("%s device start plan has inconsistent shapes" % name) + same_time_support = ( + base.n_retained_time_samples == extra.n_retained_time_samples) + starts = jnp.concatenate((base.starts, extra.starts), axis=0) + scores = jnp.concatenate((base.scores, extra.scores), axis=0) + live = jnp.concatenate((base.live, extra.live), axis=0) + return DeviceJointStartPlan( + starts, scores, live, + (base.n_lattice_candidates_before_symmetry + + extra.n_lattice_candidates_before_symmetry), + base.n_candidates_before_cap + extra.n_candidates_before_cap, + base.capacity_ok & extra.capacity_ok & same_time_support, + jnp.maximum(base.n_phi_lattice, extra.n_phi_lattice), + jnp.maximum(base.n_u_lattice, extra.n_u_lattice), + jnp.maximum(base.n_time_lattice, extra.n_time_lattice), + jnp.maximum(base.n_retained_time_samples, + extra.n_retained_time_samples), + base.n_lattice_evaluations + extra.n_lattice_evaluations, + base.n_time_scout_evaluations + extra.n_time_scout_evaluations, + base.n_time_cells_retained + extra.n_time_cells_retained, + base.n_time_nodes_retained + extra.n_time_nodes_retained, + jnp.minimum(base.time_outside_log_bound, + extra.time_outside_log_bound), + jnp.minimum(base.time_cover_min_sample, + extra.time_cover_min_sample), + jnp.maximum(base.time_cover_max_sample, + extra.time_cover_max_sample), + jnp.maximum(base.time_scout_peak_lower, + extra.time_scout_peak_lower), + base.time_cover_certified & extra.time_cover_certified, + base.time_capacity_ok & extra.time_capacity_ok, + base.norm_nonnegative & extra.norm_nonnegative, + jnp.maximum(base.n_exact_symmetry_shifts, + extra.n_exact_symmetry_shifts)) + + +def _numpy_angular_field(C, phi, u): + kp = np.arange(C.shape[0], dtype=float)[:, None] + ks_max = (C.shape[1] - 1) // 2 + ks = np.arange(-ks_max, ks_max + 1, dtype=float)[None, :] + weight = _kp_weights_numpy(C.shape[0])[:, None] + return float(np.sum(weight * C * np.exp(1j * (kp * phi + ks * u))).real) + + +def _distance_start(K, R, x_min, x_max): + """Best support-aware stationary/boundary candidate for ``x^-4 L``.""" + candidates = [float(x_min), float(x_max)] + R = max(float(R), 0.0) + K = float(K) + discriminant = K * K - 16.0 * R + if R > 0.0 and discriminant >= 0.0: + x_plus = (K + np.sqrt(discriminant)) / (2.0 * R) + if x_min <= x_plus <= x_max: + candidates.append(float(x_plus)) + values = [K * x - 0.5 * R * x * x - 4.0 * np.log(x) + for x in candidates] + return candidates[int(np.argmax(values))] + + +def algebraic_angle_starts_from_uv(C_A_t, uv_summary, time_starts, + x_min, x_max): + """Build sparse four-axis starts from U,V ranking and algebraic maxima. + + One U,V-informed distance probe is used per selected time basin. At that + probe the exact bivariate trigonometric stationary system is enumerated by + :func:`RIFT.likelihood.bivariate_trig_stationary.enumerate_torus_maxima`. + Every returned maximum is then assigned its support-aware analytic distance + candidate. No generic angle or distance seed lattice is constructed. + + The returned ``all_enumerations_ok`` covers only the angular solves at the + probed time/distance slices. It is deliberately *not* suitable for + ``AllAxisModePlan.enumeration_complete``: completeness of the joint 4-D + modes still needs the independent outside-cover warrant. + """ + from RIFT.likelihood.bivariate_trig_stationary import enumerate_torus_maxima + from RIFT.likelihood.joint_angle_peak_local import joint_table + + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + time_starts = np.asarray(time_starts, dtype=np.int32).ravel() + if not isinstance(uv_summary, UVHarmonicSummary): + raise TypeError("uv_summary must come from summarize_uv_norm_table") + _validate_tables(C_A_t, uv_summary.C_B) + starts = [] + reports = [] + all_ok = True + weight = _kp_weights_numpy(C_A_t.shape[0])[:, None] + b_scale = max(uv_summary.b_lower, + float(np.abs(uv_summary.C_B[0, + (uv_summary.C_B.shape[1] - 1) // 2])), 1.0e-30) + for time_index in time_starts: + if not (0 <= int(time_index) < C_A_t.shape[-1]): + raise ValueError("time start outside retained support") + C_A = C_A_t[..., int(time_index)] + a_upper = float(np.sum(weight * np.abs(C_A))) + x_probe = float(np.clip(a_upper / b_scale, x_min, x_max)) + result = enumerate_torus_maxima( + joint_table(C_A, uv_summary.C_B, x_probe)) + report = dict(result.report) + report.update(time_index=int(time_index), x_probe=x_probe) + reports.append(report) + all_ok = all_ok and bool(result.ok) + for phi, u in result.points: + K = _numpy_angular_field(C_A, phi, u) + R = _numpy_angular_field(uv_summary.C_B, phi, u) + x = _distance_start(K, R, float(x_min), float(x_max)) + starts.append((float(time_index), float(phi), float(u), x)) + return (np.asarray(starts, dtype=float).reshape((-1, 4)), + bool(all_ok), reports) + + +def _validate_tables(C_A_t, C_B): + if C_A_t.ndim != 3: + raise ValueError("C_A_t must have shape (KP,2KS+1,Ntime)") + if C_B.ndim != 2: + raise ValueError("C_B must be the collapsed (KP,2KS+1) norm table") + if C_A_t.shape[1] % 2 != 1 or C_B.shape[1] % 2 != 1: + raise ValueError("angular harmonic axes must have odd length") + if C_A_t.shape[0] > C_B.shape[0] or C_A_t.shape[1] > C_B.shape[1]: + raise ValueError("C_B must contain every harmonic represented by C_A") + + +def _angular_field(C, phi, u): + """Evaluate a real stored-half-plane angular Fourier table at one point.""" + kp = jnp.arange(C.shape[0], dtype=jnp.float64) + ks_max = (C.shape[1] - 1) // 2 + ks = jnp.arange(-ks_max, ks_max + 1, dtype=jnp.float64) + weight = jnp.where(kp == 0.0, 1.0, 2.0) + phase = jnp.exp(1j * (kp[:, None] * phi + ks[None, :] * u)) + return jnp.sum(weight[:, None] * C * phase).real + + +def _scalar_log_density(theta, coeff, frequency, offset, C_A_shape, C_B, + x_min, x_max): + """Unnormalized four-axis log density at one continuous coordinate.""" + t, phi, u, x = theta + flat = _evaluate_time_spectrum( + coeff, frequency, jnp.atleast_1d(t), offset)[:, 0] + C_A = flat.reshape(C_A_shape[:-1]) + A = _angular_field(C_A, phi, u) + B = _angular_field(C_B, phi, u) + inside = ((t >= 0.0) & (t <= C_A_shape[-1] - 1.0) + & (x >= x_min) & (x <= x_max) & (x > 0.0)) + value = x * A - 0.5 * x * x * B - 4.0 * jnp.log(jnp.maximum(x, 1e-300)) + return jnp.where(inside, value, -jnp.inf) + + +def refine_all_axis_starts(C_A_t, C_B, starts, x_min, x_max, *, + time_guard=0, + iterations=12, ridge=1.0e-8, + max_step=(2.0, 0.5, 0.5, 0.25), + time_localize_iterations=32, live=None): + """Refine four-axis starts with fixed-iteration JAX gradient/Hessian steps. + + This is local optimization only. The return values report stationarity and + local curvature; they do not assert completeness. Angular coordinates are + wrapped, while time and distance remain on their physical support. An + optional fixed-shape ``live`` mask skips optimizer work for padded starts + and returns finite geometry placeholders with value ``-inf`` for them. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + starts = jnp.asarray(starts, dtype=jnp.float64) + _validate_tables(C_A_t, C_B) + if starts.ndim != 2 or starts.shape[1] != 4: + raise ValueError("starts must have shape (N,4)") + if live is None: + live = jnp.ones(starts.shape[0], dtype=bool) + else: + live = jnp.asarray(live, dtype=bool) + if live.shape != (starts.shape[0],): + raise ValueError("live must match the start capacity") + if int(iterations) < 1: + raise ValueError("iterations must be positive") + if int(time_localize_iterations) < 1: + raise ValueError("time_localize_iterations must be positive") + time_guard = int(time_guard) + n_time = C_A_t.shape[-1] - 2 * time_guard + if time_guard < 0 or n_time < 2: + raise ValueError("time_guard must leave at least two integration samples") + coeff, frequency, offset = _time_primitive_spectrum( + C_A_t.reshape((-1, C_A_t.shape[-1])), time_guard) + model_shape = C_A_t.shape[:-1] + (n_time,) + fn = lambda th: _scalar_log_density( + th, coeff, frequency, offset, model_shape, C_B, + float(x_min), float(x_max)) + grad_fn = jax.grad(fn) + hess_fn = jax.hessian(fn) + max_step = jnp.asarray(max_step, dtype=jnp.float64) + + def _project(th): + return jnp.asarray([ + jnp.clip(th[0], 0.0, n_time - 1.0), + jnp.mod(th[1], 2.0 * jnp.pi), + jnp.mod(th[2], 2.0 * jnp.pi), + jnp.clip(th[3], float(x_min), float(x_max)), + ]) + + def _one(start): + # The U,V/Q portfolio is ranked on native time samples. Localize the + # continuous time maximum inside that basin before the coupled Newton + # solve; the basin narrows with SNR while this work remains fixed. + left = jnp.maximum(0.0, start[0] - 1.0) + right = jnp.minimum(float(n_time - 1), start[0] + 1.0) + golden_ratio = 0.5 * (jnp.sqrt(5.0) - 1.0) + c = right - golden_ratio * (right - left) + d = left + golden_ratio * (right - left) + + def _at_time(value): + return fn(start.at[0].set(value)) + + fc, fd = _at_time(c), _at_time(d) + + def _golden_step(_, state): + lo, hi, ca, da, fca, fda = state + choose_left = fca >= fda + new_hi = jnp.where(choose_left, da, hi) + new_lo = jnp.where(choose_left, lo, ca) + new_c = jnp.where( + choose_left, + new_hi - golden_ratio * (new_hi - new_lo), da) + new_d = jnp.where( + choose_left, ca, + new_lo + golden_ratio * (new_hi - new_lo)) + new_fc = jnp.where(choose_left, _at_time(new_c), fda) + new_fd = jnp.where(choose_left, fca, _at_time(new_d)) + return new_lo, new_hi, new_c, new_d, new_fc, new_fd + + left, right, c, d, fc, fd = jax.lax.fori_loop( + 0, int(time_localize_iterations), _golden_step, + (left, right, c, d, fc, fd)) + localized = start.at[0].set(jnp.where(fc >= fd, c, d)) + candidates = jnp.stack((start, localized), axis=0) + start = candidates[jnp.argmax(jax.vmap(fn)(candidates))] + + def _step(th, _): + g = grad_fn(th) + H = hess_fn(th) + eigenvalue, eigenvector = jnp.linalg.eigh(-H) + safe = jnp.maximum(eigenvalue, float(ridge)) + step = eigenvector @ ((eigenvector.T @ g) / safe) + # Keep the coupled Newton direction. Component-wise clipping can + # reverse its directional derivative for the narrow correlated + # time/angle basins seen in the matched SNR ladder. + ratio = max_step / jnp.maximum( + jnp.abs(step), jnp.finfo(jnp.float64).tiny) + step = step * jnp.minimum(1.0, jnp.min(ratio)) + proposals = jax.vmap( + lambda scale: _project(th + scale * step))( + jnp.concatenate(( + jnp.exp2(-jnp.arange(13, dtype=jnp.float64)), + jnp.zeros(1, dtype=jnp.float64)))) + values = jax.vmap(fn)(proposals) + return proposals[jnp.argmax(values)], None + + point, _ = jax.lax.scan(_step, _project(start), None, + length=int(iterations)) + value = fn(point) + gradient = grad_fn(point) + hessian = hess_fn(point) + curvature = jnp.linalg.eigvalsh(-hessian) + return point, value, gradient, hessian, curvature + + def _inactive(start): + point = _project(start) + return (point, jnp.asarray(-jnp.inf, dtype=start.dtype), + jnp.zeros(4, dtype=start.dtype), + -jnp.eye(4, dtype=start.dtype), + jnp.ones(4, dtype=start.dtype)) + + def _mapped(args): + start, is_live = args + return jax.lax.cond( + is_live, jax.checkpoint(_one), _inactive, start) + + return jax.lax.map(_mapped, (starts, live)) + + +def select_refined_modes(points, values, gradients, curvatures, *, + max_modes, gradient_tol=1.0e-6, + coordinate_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + scaled_step_tol=None): + """Host-side stationarity filter, rank and periodic deduplication. + + A rejected optimizer result is not a missed-mode decision: callers retain + the original basin in their completeness accounting and must add starts or + decline if its mass has not independently been bounded. In particular, + constrained maxima on the time or distance boundary need a future one-sided + local region; the full-gradient/positive-curvature filter here rejects them + and relies on the outside warrant or conservative reserve. + """ + points = np.asarray(points, dtype=float) + values = np.asarray(values, dtype=float).ravel() + gradients = np.asarray(gradients, dtype=float) + curvatures = np.asarray(curvatures, dtype=float) + if (points.ndim != 2 or points.shape[1] != 4 + or gradients.shape != points.shape + or curvatures.shape != points.shape + or values.shape[0] != points.shape[0]): + raise ValueError("inconsistent refined-mode arrays") + tolerance = np.asarray(coordinate_tol, dtype=float) + if tolerance.shape != (4,) or np.any(tolerance <= 0.0): + raise ValueError("coordinate_tol must contain four positive values") + if scaled_step_tol is None: + scaled_step_tol = float(np.min(tolerance)) + if float(scaled_step_tol) <= 0.0: + raise ValueError("scaled_step_tol must be positive") + gradient_norm = np.linalg.norm(gradients, axis=1) + min_curvature = np.min(curvatures, axis=1) + # Absolute gradients scale with lnL and hence with SNR. For a positive + # Hessian, ||g||/lambda_min bounds the Newton displacement, providing an + # SNR-stable stationarity criterion alongside the legacy absolute gate. + scaled_stationary = gradient_norm <= min_curvature * float(scaled_step_tol) + stationary = (np.all(np.isfinite(points), axis=1) + & np.isfinite(values) + & np.all(np.isfinite(gradients), axis=1) + & ((gradient_norm <= float(gradient_tol)) | scaled_stationary) + & np.all(curvatures > 0.0, axis=1)) + order = np.flatnonzero(stationary) + order = order[np.argsort(values[order])[::-1]] + max_modes = int(max_modes) + if max_modes < 1: + raise ValueError("max_modes must be positive") + selected = [] + for index in order: + point = points[index] + duplicate = False + for old in selected: + delta = np.abs(point - points[old]) + delta[1] = _periodic_distance(point[1], points[old, 1]) + delta[2] = _periodic_distance(point[2], points[old, 2]) + if np.all(delta <= tolerance): + duplicate = True + break + if not duplicate: + selected.append(int(index)) + if len(selected) > max_modes: + raise ValueError( + "unique stationary mode count exceeds fixed plan capacity") + selected = np.asarray(selected, dtype=np.int32) + return selected, stationary + + +def mode_local_geometry(hessians, *, w_sigma=8.0, + eigenvalue_floor=1.0e-12): + """Return Hessian-whitening transforms and conservative box enclosures. + + If ``L L.T = inv(-H)``, local coordinates are ``theta=center+L z`` with + ``z`` in a fixed ``[-w_sigma,w_sigma]^4`` cube. Cholesky's lower-triangular + form is computationally important: time depends only on ``z[0]``, so exact + selected-point time reconstruction still needs just ``local_order`` points, + not ``local_order**4``. The returned axis-aligned half-width encloses the + transformed cube and is used for conservative support/overlap checks. + """ + hessians = np.asarray(hessians, dtype=float) + if hessians.ndim != 3 or hessians.shape[1:] != (4, 4): + raise ValueError("hessians must have shape (N,4,4)") + transforms = np.full_like(hessians, np.nan) + half_widths = np.full((hessians.shape[0], 4), np.nan) + for i, H in enumerate(hessians): + eigenvalue = np.linalg.eigvalsh(-H) + if (not np.all(np.isfinite(eigenvalue)) + or np.min(eigenvalue) <= float(eigenvalue_floor)): + continue + covariance = np.linalg.inv(-H) + try: + transform = np.linalg.cholesky(covariance) + except np.linalg.LinAlgError: + continue + transforms[i] = transform + half_widths[i] = float(w_sigma) * np.sum(np.abs(transform), axis=1) + return transforms, half_widths + + +def _periodic_distance(a, b): + return abs((float(a) - float(b) + np.pi) % (2.0 * np.pi) - np.pi) + + +def _boxes_disjoint(centers, half_widths): + for i in range(len(centers)): + for j in range(i): + separated = ( + abs(centers[i, 0] - centers[j, 0]) + >= half_widths[i, 0] + half_widths[j, 0] + or _periodic_distance(centers[i, 1], centers[j, 1]) + >= half_widths[i, 1] + half_widths[j, 1] + or _periodic_distance(centers[i, 2], centers[j, 2]) + >= half_widths[i, 2] + half_widths[j, 2] + or abs(centers[i, 3] - centers[j, 3]) + >= half_widths[i, 3] + half_widths[j, 3]) + if not separated: + return False + return True + + +def make_all_axis_mode_plan(centers, *, max_modes, local_transforms, + local_radius=1.0, + outside_log_bound=np.inf, + enumeration_complete=False, + outside_bound_certified=False, + time_reconstruction_certified=False, + time_outside_log_bound=np.inf, + time_outside_bound_certified=False, + time_cover_min_sample=np.nan, + time_cover_max_sample=np.nan, + discovery_capacity_ok=True, + boundary_maximum_pinned=False): + """Pad a host mode set and freeze its independent acceptance warrants.""" + centers = np.asarray(centers, dtype=float) + if centers.ndim != 2 or centers.shape[1] != 4: + raise ValueError("centers must have shape (N,4)") + local_transforms = np.asarray(local_transforms, dtype=float) + if local_transforms.shape != (len(centers), 4, 4): + raise ValueError("local_transforms must have shape (N,4,4)") + if not (float(local_radius) > 0.0): + raise ValueError("local_radius must be positive") + # This enclosure is a theorem for an affine image of a cube, not caller + # policy: |L z|_i <= radius * sum_j |L_ij|. Deriving it here prevents an + # undersized supplied box from blessing overlapping or out-of-support + # quadrature regions. + half_widths = (float(local_radius) + * np.sum(np.abs(local_transforms), axis=2)) + max_modes = int(max_modes) + if max_modes < 1 or len(centers) > max_modes: + raise ValueError("mode count exceeds fixed plan capacity") + valid = (np.all(np.isfinite(centers), axis=1) + & np.all(np.isfinite(half_widths) & (half_widths > 0.0), axis=1) + & np.all(np.isfinite(local_transforms), axis=(1, 2))) + if not np.all(valid): + raise ValueError( + "every supplied mode needs finite, nondegenerate local geometry") + upper = np.triu(local_transforms, k=1) + if not np.all(upper == 0.0): + raise ValueError( + "local_transforms must be lower triangular; the selected-point " + "time topology does not evaluate upper-triangle entries") + if not np.all(np.diagonal(local_transforms, axis1=1, axis2=2) > 0.0): + raise ValueError("local_transforms must have positive diagonal") + kept_centers = centers + kept_widths = half_widths + kept_transforms = local_transforms + padded_centers = np.zeros((max_modes, 4), dtype=float) + padded_widths = np.ones((max_modes, 4), dtype=float) + padded_transforms = np.repeat(np.eye(4)[None, ...], max_modes, axis=0) + # ``lax.scan`` traces/evaluates the padded lanes even though their values + # are masked from the log-sum. Reuse one finite live geometry so inactive + # lanes cannot manufacture NaNs (notably log(x) at x<=0) which would leak + # into outer gradients or Hessians through the masked branch. + if len(kept_centers): + padded_centers[:] = kept_centers[0] + padded_widths[:] = kept_widths[0] + padded_transforms[:] = kept_transforms[0] + live = np.zeros(max_modes, dtype=bool) + padded_centers[:len(kept_centers)] = kept_centers + padded_widths[:len(kept_widths)] = kept_widths + padded_transforms[:len(kept_transforms)] = kept_transforms + live[:len(kept_centers)] = True + disjoint = _boxes_disjoint(kept_centers, kept_widths) + return AllAxisModePlan( + jnp.asarray(padded_centers), jnp.asarray(padded_widths), + jnp.asarray(padded_transforms), jnp.asarray(float(local_radius)), + jnp.asarray(live), jnp.asarray(float(outside_log_bound)), + jnp.asarray(bool(enumeration_complete)), + jnp.asarray(bool(outside_bound_certified)), + jnp.asarray(bool(time_reconstruction_certified)), + jnp.asarray(float(time_outside_log_bound)), + jnp.asarray(bool(time_outside_bound_certified)), + jnp.asarray(float(time_cover_min_sample)), + jnp.asarray(float(time_cover_max_sample)), + jnp.asarray(disjoint), + jnp.asarray(bool(discovery_capacity_ok)), + jnp.asarray(bool(boundary_maximum_pinned))) + + +def _boxes_disjoint_device(centers, half_widths, live): + """Fixed-shape periodic counterpart of :func:`_boxes_disjoint`.""" + delta = jnp.abs(centers[:, None, :] - centers[None, :, :]) + angular = jnp.abs(jnp.mod( + centers[:, None, 1:3] - centers[None, :, 1:3] + jnp.pi, + 2.0 * jnp.pi) - jnp.pi) + delta = delta.at[..., 1:3].set(angular) + separated = jnp.any( + delta >= half_widths[:, None, :] + half_widths[None, :, :], + axis=-1) + index = jnp.arange(centers.shape[0]) + pair = ((index[:, None] > index[None, :]) + & live[:, None] & live[None, :]) + return jnp.all((~pair) | separated) + + +def _validate_device_mode_plan_arguments( + start_plan, max_modes, local_radius, coordinate_tol, + scaled_step_tol, eigenvalue_floor): + if not isinstance(start_plan, DeviceJointStartPlan): + raise TypeError("start_plan must be DeviceJointStartPlan") + if (start_plan.starts.ndim != 2 or start_plan.starts.shape[1] != 4 + or start_plan.scores.shape != (start_plan.starts.shape[0],) + or start_plan.live.shape != (start_plan.starts.shape[0],)): + raise ValueError("device start plan has inconsistent shapes") + max_modes = int(max_modes) + if max_modes < 1 or max_modes > start_plan.starts.shape[0]: + raise ValueError("max_modes must fit inside the start capacity") + if not (float(local_radius) > 0.0): + raise ValueError("local_radius must be positive") + tolerance = jnp.asarray(coordinate_tol, dtype=jnp.float64) + if tolerance.shape != (4,) or np.any(np.asarray(coordinate_tol) <= 0.0): + raise ValueError("coordinate_tol must contain four positive values") + if scaled_step_tol is None: + scaled_step_tol = float(np.min(np.asarray(coordinate_tol))) + if float(scaled_step_tol) <= 0.0: + raise ValueError("scaled_step_tol must be positive") + if (not np.isfinite(float(eigenvalue_floor)) + or float(eigenvalue_floor) <= 0.0): + raise ValueError("eigenvalue_floor must be finite and positive") + return max_modes, tolerance, float(scaled_step_tol) + + +def _assemble_all_axis_mode_plan_device( + C_A_t, start_plan, refined, x_min, x_max, *, max_modes, + local_radius, time_guard, gradient_tol, tolerance, + scaled_step_tol, eigenvalue_floor, + time_reconstruction_certified, boundary_keep_nats=30.0): + """Select fixed-shape local geometry from an existing device refinement.""" + points, values, gradients, hessians, curvatures = refined + if (points.shape != start_plan.starts.shape + or values.shape != (start_plan.starts.shape[0],) + or gradients.shape != start_plan.starts.shape + or hessians.shape != (start_plan.starts.shape[0], 4, 4) + or curvatures.shape != start_plan.starts.shape): + raise ValueError("refinement arrays do not match the start plan") + + gradient_norm = jnp.linalg.norm(gradients, axis=1) + min_curvature = jnp.min(curvatures, axis=1) + scaled_stationary = ( + gradient_norm <= min_curvature * float(scaled_step_tol)) + stationary = ( + start_plan.live + & jnp.all(jnp.isfinite(points), axis=1) + & jnp.isfinite(values) + & jnp.all(jnp.isfinite(gradients), axis=1) + & ((gradient_norm <= float(gradient_tol)) | scaled_stationary) + & jnp.all(curvatures > float(eigenvalue_floor), axis=1)) + + fisher = -0.5 * (hessians + jnp.swapaxes(hessians, 1, 2)) + # Never factor an invalid lane. Padded/non-stationary starts still exist + # in the fixed shape and an indefinite inverse can emit NaNs that leak into + # outer AD even when that lane is later masked. Identity is a finite + # tracing placeholder only; such a lane remains non-stationary. + safe_fisher = jnp.where( + stationary[:, None, None], fisher, + jnp.eye(4, dtype=jnp.float64)[None, :, :]) + covariance = jnp.linalg.inv(safe_fisher) + transforms = jnp.linalg.cholesky(covariance) + geometry_finite = jnp.all(jnp.isfinite(transforms), axis=(1, 2)) + stationary &= geometry_finite + order = jnp.argsort(jnp.where(stationary, values, -jnp.inf))[::-1] + + n_time = C_A_t.shape[-1] - 2 * int(time_guard) + # A live start that refinement pinned to the support boundary and the + # stationarity filter then rejected is a constrained maximum no mode + # covers. If its value is within ``boundary_keep_nats`` of the best live + # value it can carry the integral, so the plan records it and the gate + # declines rather than integrating the interior modes alone. + x_span = max(1.0e-300, float(x_max) - float(x_min)) + at_time_bound = ((points[:, 0] <= 1.0e-9) + | (points[:, 0] >= n_time - 1.0 - 1.0e-9)) + at_x_bound = ((points[:, 3] <= float(x_min) + 1.0e-9 * x_span) + | (points[:, 3] >= float(x_max) - 1.0e-9 * x_span)) + live_finite = (start_plan.live & jnp.all(jnp.isfinite(points), axis=1) + & jnp.isfinite(values)) + best_live_value = jnp.max(jnp.where(live_finite, values, -jnp.inf)) + pinned = (live_finite & (~stationary) & (at_time_bound | at_x_bound) + & (values >= best_live_value - float(boundary_keep_nats))) + boundary_maximum_pinned = jnp.any(pinned) + fallback_center = jnp.asarray([ + 0.5 * (n_time - 1.0), 0.0, 0.0, + 0.5 * (float(x_min) + float(x_max))]) + fallback_transform = jnp.diag(jnp.asarray([ + 1.0, 0.1, 0.1, + max(1.0e-12, 0.1 * (float(x_max) - float(x_min)))], + dtype=jnp.float64)) + selected_centers = jnp.broadcast_to( + fallback_center, (max_modes, 4)).copy() + selected_transforms = jnp.broadcast_to( + fallback_transform, (max_modes, 4, 4)).copy() + selected_live = jnp.zeros(max_modes, dtype=bool) + + def _select(index, state): + centers, local_transforms, live, n_unique, overflow = state + candidate_index = order[index] + candidate = points[candidate_index] + candidate_transform = transforms[candidate_index] + delta = jnp.abs(centers - candidate[None, :]) + angular = jnp.abs(jnp.mod( + centers[:, 1:3] - candidate[None, 1:3] + jnp.pi, + 2.0 * jnp.pi) - jnp.pi) + delta = delta.at[:, 1:3].set(angular) + duplicate = jnp.any(live & jnp.all(delta <= tolerance[None, :], axis=1)) + unique = stationary[candidate_index] & (~duplicate) + has_room = n_unique < max_modes + add = unique & has_room + slot = jnp.minimum(n_unique, max_modes - 1) + + def _write(payload): + old_centers, old_transforms, old_live = payload + return (old_centers.at[slot].set(candidate), + old_transforms.at[slot].set(candidate_transform), + old_live.at[slot].set(True)) + + centers, local_transforms, live = jax.lax.cond( + add, _write, lambda payload: payload, + (centers, local_transforms, live)) + return (centers, local_transforms, live, n_unique + add, + overflow | (unique & (~has_room))) + + (selected_centers, selected_transforms, selected_live, + n_selected, selection_overflow) = jax.lax.fori_loop( + 0, start_plan.starts.shape[0], _select, + (selected_centers, selected_transforms, selected_live, + jnp.asarray(0, dtype=jnp.int32), jnp.asarray(False))) + half_widths = (float(local_radius) + * jnp.sum(jnp.abs(selected_transforms), axis=2)) + disjoint = _boxes_disjoint_device( + selected_centers, half_widths, selected_live) + discovery_capacity_ok = start_plan.capacity_ok & (~selection_overflow) + plan = AllAxisModePlan( + selected_centers, half_widths, selected_transforms, + jnp.asarray(float(local_radius)), selected_live, + jnp.asarray(jnp.inf), jnp.asarray(False), jnp.asarray(False), + jnp.asarray(bool(time_reconstruction_certified)), + start_plan.time_outside_log_bound, + start_plan.time_cover_certified, + start_plan.time_cover_min_sample, + start_plan.time_cover_max_sample, + disjoint, + discovery_capacity_ok, + boundary_maximum_pinned) + ledger = { + "boundary_maximum_pinned": boundary_maximum_pinned, + "n_boundary_pinned_starts": jnp.count_nonzero(pinned), + "n_optimizer_starts": jnp.count_nonzero(start_plan.live), + "n_refined_stationary": jnp.count_nonzero(stationary), + "n_selected_modes": n_selected, + "selection_overflow": selection_overflow, + "start_capacity_ok": start_plan.capacity_ok, + # start_capacity_ok is an AND of four terms (:862): the candidate count + # against max_starts, and these three. Exposing only the conjunction + # makes decline_capacity a single label for four different failures, + # three of which no cap value can repair. Additive; nothing reads a + # planning dict by position. + "time_cover_certified": start_plan.time_cover_certified, + "time_capacity_ok": start_plan.time_capacity_ok, + "discovery_capacity_ok": discovery_capacity_ok, + "norm_nonnegative": start_plan.norm_nonnegative, + "n_lattice_candidates_before_symmetry": + start_plan.n_lattice_candidates_before_symmetry, + "n_candidates_before_cap": start_plan.n_candidates_before_cap, + "n_exact_symmetry_shifts": start_plan.n_exact_symmetry_shifts, + "n_lattice_evaluations": start_plan.n_lattice_evaluations, + "n_time_scout_evaluations": start_plan.n_time_scout_evaluations, + "n_phi_lattice": start_plan.n_phi_lattice, + "n_u_lattice": start_plan.n_u_lattice, + "n_time_lattice": start_plan.n_time_lattice, + "n_retained_time_samples": start_plan.n_retained_time_samples, + "n_time_cells_retained": start_plan.n_time_cells_retained, + "n_time_nodes_retained": start_plan.n_time_nodes_retained, + "time_outside_log_bound": start_plan.time_outside_log_bound, + "time_cover_min_sample": start_plan.time_cover_min_sample, + "time_cover_max_sample": start_plan.time_cover_max_sample, + "time_scout_peak_lower": start_plan.time_scout_peak_lower, + "time_cover_certified": start_plan.time_cover_certified, + "time_capacity_ok": start_plan.time_capacity_ok, + "max_gradient_norm": jnp.max(jnp.where( + stationary, gradient_norm, -jnp.inf)), + "min_selected_curvature": jnp.min(jnp.where( + stationary, min_curvature, jnp.inf)), + "global_completeness_certified": jnp.asarray(False), + "derivative_warrant_certified": jnp.asarray(False), + } + return plan, ledger + + +def make_all_axis_mode_plan_device( + C_A_t, C_B, start_plan, x_min, x_max, *, max_modes, + local_radius=6.0, time_guard=0, iterations=12, + time_localize_iterations=32, ridge=1.0e-8, + max_step=(2.0, 0.5, 0.5, 0.25), gradient_tol=1.0e-6, + coordinate_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + scaled_step_tol=None, eigenvalue_floor=1.0e-12, + time_reconstruction_certified=False): + """Refine and deduplicate a fixed-shape device start portfolio. + + This is the per-row planning seam needed by nested JAX callers. It never + transfers a tracer to NumPy: starts are refined, ranked, deduplicated, and + converted to Hessian-whitened local regions entirely on device. Overflow + is recorded in ``discovery_capacity_ok`` rather than silently truncating a + mode set. No outside-mass or derivative certificate is manufactured here; + callers must use empirical enrichment plus exact reserve, or supply a + separately derived certificate through a future API. + """ + max_modes, tolerance, scaled_step_tol = ( + _validate_device_mode_plan_arguments( + start_plan, max_modes, local_radius, coordinate_tol, + scaled_step_tol, eigenvalue_floor)) + refined = refine_all_axis_starts( + C_A_t, C_B, start_plan.starts, x_min, x_max, + time_guard=int(time_guard), iterations=int(iterations), ridge=float(ridge), + max_step=max_step, + time_localize_iterations=int(time_localize_iterations), + live=start_plan.live) + return _assemble_all_axis_mode_plan_device( + C_A_t, start_plan, refined, x_min, x_max, + max_modes=max_modes, local_radius=float(local_radius), + time_guard=int(time_guard), gradient_tol=float(gradient_tol), + tolerance=tolerance, scaled_step_tol=scaled_step_tol, + eigenvalue_floor=float(eigenvalue_floor), + time_reconstruction_certified=time_reconstruction_certified) + + +def make_all_axis_mode_plan_pair_device( + C_A_t, C_B, base_starts, extra_starts, x_min, x_max, *, max_modes, + enriched_max_modes=None, + local_radius=6.0, time_guard=0, iterations=12, + time_localize_iterations=32, ridge=1.0e-8, + max_step=(2.0, 0.5, 0.5, 0.25), gradient_tol=1.0e-6, + coordinate_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + scaled_step_tol=None, eigenvalue_floor=1.0e-12, + time_reconstruction_certified=False): + """Build nested base/enriched plans with one shared optimizer pass. + + ``extra_starts`` is combined with ``base_starts`` structurally before the + refinement, so the enriched portfolio contains every base lane. The + independent per-lane optimizer means the prefix of the shared result is + exactly the refinement that a separate base call would have produced. + Reusing it removes the otherwise duplicated base hill-climbing work without + changing either mode selection or any completeness warrant. + ``enriched_max_modes`` may raise the stronger plan's output capacity; it + defaults to the base ``max_modes``. + + The fifth return value records actual shared optimizer work and the work + avoided relative to two separate base-plus-enriched refinements. This is a + cost transformation only: neither start nesting nor optimizer convergence + supplies a global omitted-mass or derivative certificate. + """ + base_max_modes, tolerance, scaled_step_tol = ( + _validate_device_mode_plan_arguments( + base_starts, max_modes, local_radius, coordinate_tol, + scaled_step_tol, eigenvalue_floor)) + combined = combine_device_start_plans(base_starts, extra_starts) + if enriched_max_modes is None: + enriched_max_modes = base_max_modes + enriched_max_modes, _, _ = _validate_device_mode_plan_arguments( + combined, enriched_max_modes, local_radius, coordinate_tol, + scaled_step_tol, eigenvalue_floor) + refined = refine_all_axis_starts( + C_A_t, C_B, combined.starts, x_min, x_max, + time_guard=int(time_guard), iterations=int(iterations), ridge=float(ridge), + max_step=max_step, + time_localize_iterations=int(time_localize_iterations), + live=combined.live) + base_capacity = base_starts.starts.shape[0] + base_refined = tuple(value[:base_capacity] for value in refined) + base_plan, base_ledger = _assemble_all_axis_mode_plan_device( + C_A_t, base_starts, base_refined, x_min, x_max, + max_modes=base_max_modes, local_radius=float(local_radius), + time_guard=int(time_guard), gradient_tol=float(gradient_tol), + tolerance=tolerance, scaled_step_tol=scaled_step_tol, + eigenvalue_floor=float(eigenvalue_floor), + time_reconstruction_certified=time_reconstruction_certified) + enriched_plan, enriched_ledger = _assemble_all_axis_mode_plan_device( + C_A_t, combined, refined, x_min, x_max, + max_modes=enriched_max_modes, local_radius=float(local_radius), + time_guard=int(time_guard), gradient_tol=float(gradient_tol), + tolerance=tolerance, scaled_step_tol=scaled_step_tol, + eigenvalue_floor=float(eigenvalue_floor), + time_reconstruction_certified=time_reconstruction_certified) + base_live = jnp.count_nonzero(base_starts.live) + extra_live = jnp.count_nonzero(extra_starts.live) + shared_ledger = { + "n_optimizer_starts_executed": base_live + extra_live, + "n_optimizer_starts_avoided": base_live, + "n_optimizer_starts_previous_two_pass": 2 * base_live + extra_live, + "start_nesting_structural": jnp.asarray(True), + } + return (base_plan, enriched_plan, base_ledger, enriched_ledger, + shared_ledger) + + +def _legendre_rule(order): + nodes, weights = np.polynomial.legendre.leggauss(int(order)) + return (jnp.asarray(nodes, dtype=jnp.float64), + jnp.asarray(weights, dtype=jnp.float64)) + + +def _mode_integral(coeff, frequency, offset, C_A_shape, C_B, center, + transform, radius, nodes, log_weights, concentration): + if float(concentration) > 0.0: + scaled = (jnp.sinh(float(concentration) * nodes) + / jnp.sinh(float(concentration))) + log_jacobian_shape = ( + jnp.log(float(concentration)) + + jnp.log(jnp.cosh(float(concentration) * nodes)) + - jnp.log(jnp.sinh(float(concentration)))) + else: + scaled = nodes + log_jacobian_shape = jnp.zeros_like(nodes) + z = radius * scaled + # Lower-triangular whitening preserves a separable reconstruction topology: + # t has n points, (t,phi) n^2, (t,phi,u) n^3, and only the final exponent + # has n^4. This is the central memory property of the all-axis kernel. + t = center[0] + transform[0, 0] * z + phi = jnp.mod( + center[1] + transform[1, 0] * z[:, None] + + transform[1, 1] * z[None, :], 2.0 * jnp.pi) + u = jnp.mod( + center[2] + transform[2, 0] * z[:, None, None] + + transform[2, 1] * z[None, :, None] + + transform[2, 2] * z[None, None, :], 2.0 * jnp.pi) + x = (center[3] + transform[3, 0] * z[:, None, None, None] + + transform[3, 1] * z[None, :, None, None] + + transform[3, 2] * z[None, None, :, None] + + transform[3, 3] * z[None, None, None, :]) + + flat = _evaluate_time_spectrum(coeff, frequency, t, offset) + C_A = flat.reshape(C_A_shape[:-1] + (nodes.size,)) + + kp_a = jnp.arange(C_A.shape[0], dtype=jnp.float64) + ks_a = jnp.arange(-(C_A.shape[1] - 1) // 2, + (C_A.shape[1] - 1) // 2 + 1, dtype=jnp.float64) + wa = jnp.where(kp_a == 0.0, 1.0, 2.0) + EA = jnp.exp(1j * (phi[:, :, None, None, None] * kp_a[None, None, None, :, None] + + u[:, :, :, None, None] * ks_a[None, None, None, None, :])) + EA = EA * wa[None, None, None, :, None] + A = jnp.einsum("tpukq,kqt->tpu", EA, C_A).real + + kp_b = jnp.arange(C_B.shape[0], dtype=jnp.float64) + ks_b = jnp.arange(-(C_B.shape[1] - 1) // 2, + (C_B.shape[1] - 1) // 2 + 1, dtype=jnp.float64) + wb = jnp.where(kp_b == 0.0, 1.0, 2.0) + EB = jnp.exp(1j * (phi[:, :, None, None, None] * kp_b[None, None, None, :, None] + + u[:, :, :, None, None] * ks_b[None, None, None, None, :])) + EB = EB * wb[None, None, None, :, None] + B = jnp.einsum("tpukq,kq->tpu", EB, C_B).real + + exponent = (A[..., None] * x + - 0.5 * B[..., None] * jnp.square(x) + - 4.0 * jnp.log(jnp.maximum(x, 1.0e-300))) + sign, log_det = jnp.linalg.slogdet(transform) + mapped_log_weights = (log_weights + log_jacobian_shape + + jnp.log(radius)) + logw = (mapped_log_weights[:, None, None, None] + + mapped_log_weights[None, :, None, None] + + mapped_log_weights[None, None, :, None] + + mapped_log_weights[None, None, None, :] + + log_det) + return jnp.where(sign != 0.0, + jax.scipy.special.logsumexp(exponent + logw), -jnp.inf) + + +def _evaluate_plan_at_order(C_A_t, C_B, plan, order, concentration, + time_guard): + nodes, weights = _legendre_rule(order) + log_weights = jnp.log(weights) + coeff, frequency, offset = _time_primitive_spectrum( + C_A_t.reshape((-1, C_A_t.shape[-1])), int(time_guard)) + model_shape = C_A_t.shape[:-1] + ( + C_A_t.shape[-1] - 2 * int(time_guard),) + + def _step(total, args): + center, transform, live = args + def _live_mode(local_args): + local_center, local_transform = local_args + return _mode_integral( + coeff, frequency, offset, model_shape, C_B, + local_center, local_transform, plan.local_radius, nodes, + log_weights, concentration) + + value = jax.lax.cond( + live, _live_mode, lambda _: jnp.asarray(-jnp.inf), + (center, transform)) + return jnp.logaddexp(total, value), None + + value, _ = jax.lax.scan( + jax.checkpoint(_step), jnp.asarray(-jnp.inf), + (plan.centers, plan.local_transforms, plan.live)) + n_live = jnp.count_nonzero(plan.live) + n_eval = n_live * int(order) ** 4 + # Conservative live-array accounting for one streamed mode. It is a + # deterministic shape counter, not a device allocator measurement. + nt = int(order) + lanes = int(np.prod(C_A_t.shape[:-1])) + n_frequency = 2 * C_A_t.shape[-1] - 2 + angle_terms_a = C_A_t.shape[0] * C_A_t.shape[1] + angle_terms_b = C_B.shape[0] * C_B.shape[1] + table_bytes = (((nt + lanes) * n_frequency + lanes * nt + + nt ** 3 * (angle_terms_a + angle_terms_b)) * 16 + + C_B.size * 16) + field_bytes = (n_frequency + 3 * nt ** 3 + nt ** 2 + + 3 * nt ** 4 + 10 * nt) * 8 + workspace_bytes = jnp.asarray(table_bytes + field_bytes) + n_selected_time_points = n_live * nt + n_time_frequency_terms = n_live * nt * n_frequency * lanes + n_angle_harmonic_terms = ( + n_live * nt ** 3 * (angle_terms_a + angle_terms_b)) + return (value, n_eval, workspace_bytes, n_selected_time_points, + n_time_frequency_terms, n_angle_harmonic_terms) + + +def all_axis_peak_local_marginalize( + C_A_t, C_B, plan, x_min, x_max, *, local_order=5, + check_order=9, quadrature_tol_nats=1.0e-5, + outside_tol_nats=-23.0, log_normalization=0.0, + node_concentration=1.0, time_guard=0, + time_guard_tol_nats=1.0e-3): + """Marginalize a padded multi-mode plan with explicit fail-closed ledger. + + ``C_A_t`` is the primitive data table ``(mmax+1,3,Ntime)`` and ``C_B`` is + the cached ``U,V`` norm table ``(2*mmax+1,5)``. The returned value is + diagnostic unless ``ok`` is true. Here ``ok`` is a validated value-only + disposition, not a certified quadrature bound or a gradient/Hessian + certificate. On any decline the caller must use the + dense/exact reserve; ``fallback_required`` is provided to make that branch + hard to omit accidentally. + + With ``time_guard >= 2``, ``C_A_t`` contains support on both sides of the + target window. The high-order local integral is repeated after trimming to + ``time_guard//2`` support and acceptance requires their difference to meet + ``time_guard_tol_nats``. A guarded input must pass that comparison and + cannot be rescued by the plan's external warrant; an unguarded input needs + the external warrant. This is an operational convergence validation, not + a rigorous interpolation-error bound, and the ledger labels it accordingly. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + _validate_tables(C_A_t, C_B) + if not isinstance(plan, AllAxisModePlan): + raise TypeError("plan must be AllAxisModePlan") + if not (0.0 < float(x_min) < float(x_max)): + raise ValueError("need 0 < x_min < x_max") + if int(local_order) < 2 or int(check_order) <= int(local_order): + raise ValueError("need 2 <= local_order < check_order") + if float(node_concentration) < 0.0: + raise ValueError("node_concentration must be non-negative") + time_guard = int(time_guard) + n_time = C_A_t.shape[-1] - 2 * time_guard + if time_guard < 0 or n_time < 2: + raise ValueError("time_guard must leave at least two integration samples") + if time_guard == 1: + raise ValueError("two-guard validation requires time_guard=0 or >=2") + if float(time_guard_tol_nats) <= 0.0: + raise ValueError("time_guard_tol_nats must be positive") + + (value_lo, eval_lo, bytes_lo, time_lo, + time_terms_lo, angle_terms_lo) = _evaluate_plan_at_order( + C_A_t, C_B, plan, int(local_order), float(node_concentration), + time_guard) + (value_hi, eval_hi, bytes_hi, time_hi, + time_terms_hi, angle_terms_hi) = _evaluate_plan_at_order( + C_A_t, C_B, plan, int(check_order), float(node_concentration), + time_guard) + if time_guard: + inner_guard = time_guard // 2 + trim = time_guard - inner_guard + inner_table = C_A_t[..., trim:-trim] + (value_guard_inner, guard_eval_hi, guard_bytes_hi, guard_time_hi, + guard_time_terms_hi, guard_angle_terms_hi) = _evaluate_plan_at_order( + inner_table, C_B, plan, int(check_order), + float(node_concentration), inner_guard) + guard_error = jnp.abs(value_hi - value_guard_inner) + guard_validated = (jnp.isfinite(value_guard_inner) + & (guard_error <= float(time_guard_tol_nats))) + else: + inner_guard = 0 + value_guard_inner = jnp.asarray(jnp.nan) + guard_error = jnp.asarray(jnp.inf) + guard_validated = jnp.asarray(False) + guard_eval_hi = jnp.asarray(0) + guard_bytes_hi = jnp.asarray(0) + guard_time_hi = jnp.asarray(0) + guard_time_terms_hi = jnp.asarray(0) + guard_angle_terms_hi = jnp.asarray(0) + value_lo = value_lo + float(log_normalization) + value_hi = value_hi + float(log_normalization) + outside = plan.outside_log_bound + float(log_normalization) + time_outside = plan.time_outside_log_bound + float(log_normalization) + + centers = plan.centers + widths = plan.half_widths + inside_time_distance = ( + (centers[:, 0] - widths[:, 0] >= 0.0) + & (centers[:, 0] + widths[:, 0] <= n_time - 1.0) + & (centers[:, 3] - widths[:, 3] >= float(x_min)) + & (centers[:, 3] + widths[:, 3] <= float(x_max))) + angular_single_cover = jnp.all(widths[:, 1:3] <= jnp.pi, axis=1) + support_ok = jnp.all(jnp.where( + plan.live, inside_time_distance & angular_single_cover, True)) + finite = (jnp.all(jnp.isfinite(C_A_t.real)) + & jnp.all(jnp.isfinite(C_A_t.imag)) + & jnp.all(jnp.isfinite(C_B.real)) + & jnp.all(jnp.isfinite(C_B.imag)) + & jnp.isfinite(value_hi) + & jnp.any(plan.live)) + quadrature_error = jnp.abs(value_hi - value_lo) + quadrature_ok = quadrature_error <= float(quadrature_tol_nats) + tail_margin = outside - value_hi + tail_ok = tail_margin < float(outside_tol_nats) + # The outside-cover certificate is the correctness-bearing completeness + # warrant. Requiring the algebraic root report as well would incorrectly + # reject an otherwise bounded missed root. Conversely, a perfect root + # report cannot replace an integral bound outside the local regions. + # A supplied guarded reconstruction owns its own convergence check. Do not + # let a stale external warrant mask a failed outer/inner comparison. + time_warranted = (guard_validated if time_guard + else plan.time_reconstruction_certified) + cover_warranted = plan.outside_bound_certified & time_warranted + + decline_nonfinite = ~finite + decline_incomplete = finite & (~plan.outside_bound_certified) + decline_time_reconstruction = ( + finite & plan.outside_bound_certified + & (~time_warranted)) + decline_overlap = finite & cover_warranted & (~plan.boxes_disjoint) + decline_support = (finite & cover_warranted & plan.boxes_disjoint + & (~support_ok)) + decline_quadrature = (finite & cover_warranted & plan.boxes_disjoint & support_ok + & (~quadrature_ok)) + decline_tail = (finite & cover_warranted & plan.boxes_disjoint & support_ok + & quadrature_ok & (~tail_ok)) + ok = (finite & cover_warranted & plan.boxes_disjoint & support_ok + & quadrature_ok & tail_ok) + reconciles = (ok.astype(jnp.int32) + + decline_nonfinite.astype(jnp.int32) + + decline_incomplete.astype(jnp.int32) + + decline_time_reconstruction.astype(jnp.int32) + + decline_overlap.astype(jnp.int32) + + decline_support.astype(jnp.int32) + + decline_quadrature.astype(jnp.int32) + + decline_tail.astype(jnp.int32)) == 1 + ledger = { + "accepted": ok, + "fallback_required": ~ok, + "decline_is_waveform_failure": jnp.asarray(False), + "fixed_plan_autodiff_only": jnp.asarray(True), + "derivative_warrant_certified": jnp.asarray(False), + "decline_nonfinite": decline_nonfinite, + "decline_incomplete": decline_incomplete, + "decline_time_reconstruction": decline_time_reconstruction, + "decline_overlap": decline_overlap, + "decline_support": decline_support, + "decline_quadrature": decline_quadrature, + "decline_tail": decline_tail, + "reconciles": reconciles, + "enumeration_complete": plan.enumeration_complete, + "outside_bound_certified": plan.outside_bound_certified, + "time_reconstruction_certified": plan.time_reconstruction_certified, + "time_outside_bound_certified": plan.time_outside_bound_certified, + "time_outside_log_bound": time_outside, + "time_outside_tail_margin": time_outside - value_hi, + "time_guard_validated": guard_validated, + "time_reconstruction_warranted": time_warranted, + "time_guard_error_certified": jnp.asarray(False), + "time_guard": jnp.asarray(time_guard), + "time_guard_inner": jnp.asarray(inner_guard), + "time_guard_error": guard_error, + "time_guard_tol_nats": jnp.asarray(float(time_guard_tol_nats)), + "time_guard_inner_value": value_guard_inner + float(log_normalization), + "boxes_disjoint": plan.boxes_disjoint, + "support_ok": support_ok, + "quadrature_ok": quadrature_ok, + "quadrature_error_certified": jnp.asarray(False), + "value_warrant_certified": jnp.asarray(False), + "tail_ok": tail_ok, + "quadrature_error": quadrature_error, + "tail_margin": tail_margin, + "outside_log_bound": outside, + "n_modes": jnp.count_nonzero(plan.live), + "n_mode_capacity": jnp.asarray(plan.live.size), + "n_local_evaluations_lo": eval_lo, + "n_local_evaluations_hi": eval_hi, + "n_guard_local_evaluations_hi": guard_eval_hi, + "n_total_local_evaluations_hi": eval_hi + guard_eval_hi, + "n_selected_time_points_lo": time_lo, + "n_selected_time_points_hi": time_hi, + "n_guard_selected_time_points_hi": guard_time_hi, + "n_total_selected_time_points_hi": time_hi + guard_time_hi, + "n_time_frequency_terms_lo": time_terms_lo, + "n_time_frequency_terms_hi": time_terms_hi, + "n_guard_time_frequency_terms_hi": guard_time_terms_hi, + "n_total_time_frequency_terms_hi": time_terms_hi + guard_time_terms_hi, + "n_angle_harmonic_terms_lo": angle_terms_lo, + "n_angle_harmonic_terms_hi": angle_terms_hi, + "n_guard_angle_harmonic_terms_hi": guard_angle_terms_hi, + "n_total_angle_harmonic_terms_hi": angle_terms_hi + guard_angle_terms_hi, + "workspace_bytes_lo": bytes_lo, + "workspace_bytes_hi": bytes_hi, + "workspace_bytes_guard_hi": guard_bytes_hi, + "workspace_bytes_peak_bound_hi": jnp.maximum(bytes_hi, guard_bytes_hi), + } + return value_hi, ok, ledger + + +def empirical_enrichment_marginalize( + C_A_t, C_B, base_plan, enriched_plan, x_min, x_max, *, + base_order=13, base_check_order=19, + enriched_order=19, enriched_check_order=25, + convergence_tol_nats=1.0e-3, time_guard=0, + time_guard_tol_nats=1.0e-3, log_normalization=0.0, + time_outside_tol_nats=-23.0, + total_value_error_budget_nats=1.0e-3, + node_concentration=1.0, + mode_match_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + geometry_match_rtol=1.0e-3, geometry_match_atol=1.0e-8): + """Apply the operational one-step enrichment gate to two fixed plans. + + ``enriched_plan`` must come from a strictly stronger discovery portfolio + that includes the base starts, and uses the stronger quadrature orders + supplied here. Acceptance requires finite values, explicit start capacity, + disjoint in-support regions, healthy nested quadrature and time-guard + diagnostics, certified omitted-time bounds for both plans to clear + ``time_outside_tol_nats``, and agreement within ``convergence_tol_nats``. + The empirical discovery, nested-quadrature, guarded-time, and certified + omitted-time contributions must also fit one shared + ``total_value_error_budget_nats``. Each of the three paired terms is + charged once, as the maximum over the base and enriched plans; their + cancellation-resistant sum is an operational error score, not a formal + global bound: enrichment and quadrature differences remain empirical + convergence diagnostics. + Every base + mode must recur with matching local geometry. Additional enriched basins + are probes, not automatically part of the accepted cover: if a probe has + broad/overlapping geometry but changes the positive local integral by less + than the same convergence budget, the valid base cover is retained and its + value returned. This prevents a manifestly negligible low-curvature HM + basin from forcing dense reserve while still making a changed relevant + basin or an invalid base cover decline. It does not claim formal global + completeness or derivative accuracy. + + Any decline returns the finite enriched diagnostic with + ``fallback_required=True``. The caller must execute and retain the + dense/exact reserve; a decline is never a waveform failure. + """ + if not (int(base_order) < int(base_check_order) + <= int(enriched_order) < int(enriched_check_order)): + raise ValueError( + "need base_order < base_check_order <= enriched_order " + "< enriched_check_order") + if float(convergence_tol_nats) <= 0.0: + raise ValueError("convergence_tol_nats must be positive") + if (not np.isfinite(float(total_value_error_budget_nats)) + or float(total_value_error_budget_nats) <= 0.0): + raise ValueError( + "total_value_error_budget_nats must be finite and positive") + if (not np.isfinite(float(time_outside_tol_nats)) + or float(time_outside_tol_nats) >= 0.0): + raise ValueError("time_outside_tol_nats must be finite and negative") + mode_match_tol = np.asarray(mode_match_tol, dtype=float) + if mode_match_tol.shape != (4,) or np.any(mode_match_tol <= 0.0): + raise ValueError("mode_match_tol must contain four positive values") + if float(geometry_match_rtol) < 0.0 or float(geometry_match_atol) < 0.0: + raise ValueError("geometry match tolerances must be non-negative") + + base_value, _, base = all_axis_peak_local_marginalize( + C_A_t, C_B, base_plan, x_min, x_max, + local_order=int(base_order), check_order=int(base_check_order), + quadrature_tol_nats=float(convergence_tol_nats), + log_normalization=float(log_normalization), + node_concentration=float(node_concentration), time_guard=int(time_guard), + time_guard_tol_nats=float(time_guard_tol_nats)) + enriched_value, _, enriched = all_axis_peak_local_marginalize( + C_A_t, C_B, enriched_plan, x_min, x_max, + local_order=int(enriched_order), + check_order=int(enriched_check_order), + quadrature_tol_nats=float(convergence_tol_nats), + log_normalization=float(log_normalization), + node_concentration=float(node_concentration), time_guard=int(time_guard), + time_guard_tol_nats=float(time_guard_tol_nats)) + + has_modes = (base["n_modes"] > 0) & (enriched["n_modes"] > 0) + values_finite = jnp.isfinite(base_value) & jnp.isfinite(enriched_value) + # An empty padded plan evaluates to -inf by construction. That is a + # discovery disposition, not numerical corruption: keep it eligible for + # the explicit ``decline_no_modes`` branch below. A nonfinite value from + # a plan that does contain modes remains a numerical decline. + finite = values_finite | (~has_modes) + capacity_ok = (base_plan.discovery_capacity_ok + & enriched_plan.discovery_capacity_ok) + boundary_ok = ~(base_plan.boundary_maximum_pinned + | enriched_plan.boundary_maximum_pinned) + time_ok = (base["time_reconstruction_warranted"] + & enriched["time_reconstruction_warranted"]) + any_time_cover = (base_plan.time_outside_bound_certified + | enriched_plan.time_outside_bound_certified) + time_cover_pair = (base_plan.time_outside_bound_certified + & enriched_plan.time_outside_bound_certified) + base_time_tail_margin = (base["time_outside_log_bound"] - base_value) + enriched_time_tail_margin = ( + enriched["time_outside_log_bound"] - enriched_value) + time_tail_bounds_ok = ( + (base_time_tail_margin < float(time_outside_tol_nats)) + & (enriched_time_tail_margin < float(time_outside_tol_nats))) + time_omitted_ok = time_cover_pair & time_tail_bounds_ok + base_geometry_ok = base["boxes_disjoint"] & base["support_ok"] + enriched_geometry_ok = (enriched["boxes_disjoint"] + & enriched["support_ok"]) + quadrature_ok = base["quadrature_ok"] & enriched["quadrature_ok"] + delta = jnp.abs(base_plan.centers[:, None, :] + - enriched_plan.centers[None, :, :]) + angular_delta = jnp.abs(jnp.mod( + delta[..., 1:3] + jnp.pi, 2.0 * jnp.pi) - jnp.pi) + delta = delta.at[..., 1:3].set(angular_delta) + matches = (jnp.all(delta <= jnp.asarray(mode_match_tol), axis=-1) + & enriched_plan.live[None, :]) + width_scale = (float(geometry_match_atol) + + float(geometry_match_rtol) + * jnp.abs(base_plan.half_widths[:, None, :])) + width_matches = jnp.all( + jnp.abs(base_plan.half_widths[:, None, :] + - enriched_plan.half_widths[None, :, :]) <= width_scale, + axis=-1) + transform_scale = (float(geometry_match_atol) + + float(geometry_match_rtol) + * jnp.abs(base_plan.local_transforms[:, None, :, :])) + transform_matches = jnp.all( + jnp.abs(base_plan.local_transforms[:, None, :, :] + - enriched_plan.local_transforms[None, :, :, :]) + <= transform_scale, axis=(-2, -1)) + geometry_matches = matches & width_matches & transform_matches + # Padded base rows are vacuously retained. A stronger plan may add modes, + # but it may not silently lose one that contributed to the base value. + mode_nesting_ok = jnp.all( + jnp.where(base_plan.live, jnp.any(matches, axis=1), True)) + geometry_nesting_ok = jnp.all(jnp.where( + base_plan.live, jnp.any(geometry_matches, axis=1), True)) + geometry_ok = base_geometry_ok & geometry_nesting_ok + convergence_error = jnp.abs(enriched_value - base_value) + converged = convergence_error <= float(convergence_tol_nats) + # A single operational allowance prevents several individually acceptable + # diagnostics from silently spending the full science tolerance apiece. + # Sums, rather than maxima, are deliberate: base/enriched quadrature or + # guard errors can cancel in their observed difference. Unguarded plans + # contribute zero here only after their external reconstruction warrants + # have passed ``time_ok`` above. The discarded-time terms are genuine + # integral bounds, converted to maximum log-integral corrections, while + # the other terms remain empirical diagnostics. + base_guard_score = jnp.where( + int(time_guard) > 0, base["time_guard_error"], 0.0) + enriched_guard_score = jnp.where( + int(time_guard) > 0, enriched["time_guard_error"], 0.0) + base_time_tail_correction = jnp.logaddexp( + 0.0, base_time_tail_margin) + enriched_time_tail_correction = jnp.logaddexp( + 0.0, enriched_time_tail_margin) + # Each paired term is the same physical quantity measured on the two + # nested plans (same table, same recurring modes, same G vs G/2 + # comparison). Charging both against the budget double-counted a + # common-mode error: on the analytic wiring fixture at ten times the + # unit amplitude a value correct to 1.5e-4 nat was refused at a score of + # 1.06e-3, of which 2 x 5.27e-4 was one guard discrepancy counted twice. + # The maximum over the pair bounds whichever plan's value is selected and + # counts it once. The per-plan terms stay in the ledger. + quadrature_score = jnp.maximum( + base["quadrature_error"], enriched["quadrature_error"]) + guard_score = jnp.maximum(base_guard_score, enriched_guard_score) + tail_score = jnp.maximum( + base_time_tail_correction, enriched_time_tail_correction) + empirical_value_error_score = ( + convergence_error + quadrature_score + guard_score + tail_score) + error_budget_complete = time_cover_pair & time_ok + error_budget_ok = ( + error_budget_complete + & jnp.isfinite(empirical_value_error_score) + & (empirical_value_error_score + <= float(total_value_error_budget_nats))) + + decline_nonfinite = ~finite + decline_capacity = finite & (~capacity_ok) + decline_no_modes = finite & capacity_ok & (~has_modes) + decline_boundary_maximum = (finite & capacity_ok & has_modes + & (~boundary_ok)) + decline_mode_nesting = (finite & capacity_ok & has_modes & boundary_ok + & (~mode_nesting_ok)) + decline_time = (finite & capacity_ok & has_modes & boundary_ok + & mode_nesting_ok & (~time_ok)) + decline_time_cover = ( + finite & capacity_ok & has_modes & boundary_ok & mode_nesting_ok + & time_ok & (~time_cover_pair)) + decline_time_omitted_bound = ( + finite & capacity_ok & has_modes & boundary_ok & mode_nesting_ok + & time_ok & time_cover_pair & (~time_tail_bounds_ok)) + # Umbrella science diagnostic retained for callers that need only the + # broad reason. The two exclusive fields above own reconciliation. + decline_time_omitted = decline_time_cover | decline_time_omitted_bound + decline_geometry = (finite & capacity_ok & has_modes & boundary_ok + & mode_nesting_ok & time_ok & time_omitted_ok + & (~geometry_ok)) + decline_quadrature = (finite & capacity_ok & has_modes & boundary_ok + & mode_nesting_ok & time_ok & time_omitted_ok + & geometry_ok & (~quadrature_ok)) + decline_enrichment = (finite & capacity_ok & has_modes & boundary_ok + & mode_nesting_ok & time_ok & time_omitted_ok + & geometry_ok & quadrature_ok & (~converged)) + decline_error_budget = ( + finite & capacity_ok & has_modes & boundary_ok & mode_nesting_ok + & time_ok & time_omitted_ok & geometry_ok & quadrature_ok & converged + & (~error_budget_ok)) + accepted = (finite & capacity_ok & has_modes & boundary_ok + & mode_nesting_ok & time_ok & time_omitted_ok & geometry_ok + & quadrature_ok & converged & error_budget_ok) + accepted_value_uses_base_geometry = accepted & (~enriched_geometry_ok) + accepted_value = jnp.where( + accepted_value_uses_base_geometry, base_value, enriched_value) + reconciles = ( + accepted.astype(jnp.int32) + + decline_nonfinite.astype(jnp.int32) + + decline_capacity.astype(jnp.int32) + + decline_no_modes.astype(jnp.int32) + + decline_boundary_maximum.astype(jnp.int32) + + decline_mode_nesting.astype(jnp.int32) + + decline_time.astype(jnp.int32) + + decline_time_cover.astype(jnp.int32) + + decline_time_omitted_bound.astype(jnp.int32) + + decline_geometry.astype(jnp.int32) + + decline_quadrature.astype(jnp.int32) + + decline_enrichment.astype(jnp.int32) + + decline_error_budget.astype(jnp.int32)) == 1 + ledger = { + "accepted": accepted, + "fallback_required": ~accepted, + "decline_is_waveform_failure": jnp.asarray(False), + "acceptance_is_empirical_enrichment": jnp.asarray(True), + "global_completeness_certified": jnp.asarray(False), + "empirical_value_error_certified": jnp.asarray(False), + "derivative_warrant_certified": jnp.asarray(False), + "decline_nonfinite": decline_nonfinite, + "decline_capacity": decline_capacity, + "decline_no_modes": decline_no_modes, + "decline_boundary_maximum": decline_boundary_maximum, + "boundary_maximum_ok": boundary_ok, + "decline_mode_nesting": decline_mode_nesting, + "decline_time_reconstruction": decline_time, + "decline_time_cover_incomplete": decline_time_cover, + "decline_time_omitted_mass_bound": decline_time_omitted_bound, + "decline_time_omitted_mass": decline_time_omitted, + "decline_geometry": decline_geometry, + "decline_quadrature": decline_quadrature, + "decline_enrichment": decline_enrichment, + "decline_error_budget": decline_error_budget, + "reconciles": reconciles, + "base_value": base_value, + "enriched_value": enriched_value, + "base_and_enriched_values_finite": values_finite, + "convergence_error": convergence_error, + "convergence_tol_nats": jnp.asarray(float(convergence_tol_nats)), + "empirical_value_error_score_nats": empirical_value_error_score, + "total_value_error_budget_nats": jnp.asarray( + float(total_value_error_budget_nats)), + "value_error_budget_complete": error_budget_complete, + "value_error_budget_ok": error_budget_ok, + "value_error_budget_is_empirical": jnp.asarray(True), + "value_error_budget_is_formal_bound": jnp.asarray(False), + "error_score_discovery_nats": convergence_error, + "error_score_quadrature_nats": quadrature_score, + "error_score_time_guard_nats": guard_score, + "error_score_omitted_time_nats": tail_score, + "error_score_pairs_charged_as_max": jnp.asarray(True), + "error_score_base_quadrature_nats": base["quadrature_error"], + "error_score_enriched_quadrature_nats": + enriched["quadrature_error"], + "error_score_base_time_guard_nats": base_guard_score, + "error_score_enriched_time_guard_nats": enriched_guard_score, + "error_score_base_omitted_time_nats": base_time_tail_correction, + "error_score_enriched_omitted_time_nats": + enriched_time_tail_correction, + "base_capacity_ok": base_plan.discovery_capacity_ok, + "enriched_capacity_ok": enriched_plan.discovery_capacity_ok, + "base_n_modes": base["n_modes"], + "enriched_n_modes": enriched["n_modes"], + "mode_nesting_ok": mode_nesting_ok, + "geometry_nesting_ok": geometry_nesting_ok, + "base_geometry_ok": base_geometry_ok, + "enriched_geometry_ok": enriched_geometry_ok, + "accepted_value_uses_base_geometry": + accepted_value_uses_base_geometry, + "base_quadrature_error": base["quadrature_error"], + "enriched_quadrature_error": enriched["quadrature_error"], + "base_time_guard_error": base["time_guard_error"], + "enriched_time_guard_error": enriched["time_guard_error"], + "time_outside_cover_used": time_cover_pair, + "time_outside_cover_any": any_time_cover, + "time_outside_cover_pair": time_cover_pair, + "time_omitted_mass_ok": time_omitted_ok, + "time_outside_tol_nats": jnp.asarray(float(time_outside_tol_nats)), + "base_time_outside_tail_margin": base_time_tail_margin, + "enriched_time_outside_tail_margin": enriched_time_tail_margin, + "base_total_local_evaluations_hi": + base["n_total_local_evaluations_hi"], + "enriched_total_local_evaluations_hi": + enriched["n_total_local_evaluations_hi"], + "total_local_evaluations_hi": ( + base["n_total_local_evaluations_hi"] + + enriched["n_total_local_evaluations_hi"]), + "workspace_bytes_peak_bound_hi": jnp.maximum( + base["workspace_bytes_peak_bound_hi"], + enriched["workspace_bytes_peak_bound_hi"]), + } + return accepted_value, accepted, ledger + + +def empirical_enrichment_with_exact_reserve( + C_A_t, C_B, base_plan, enriched_plan, x_min, x_max, *, + reserve_x_grid, reserve_log_weights, time_weights, + reserve_amp_sizing, reserve_m_max=None, + reserve_dense_chunk=8, reserve_grid_block=32, + reserve_time_nodes=None, reserve_time_resolution_warranted=False, + reserve_time_check_value=np.nan, + reserve_time_check_nodes=None, reserve_time_check_weights=None, + reserve_time_resolution_tol_nats=1.0e-3, + base_order=13, base_check_order=19, + enriched_order=19, enriched_check_order=25, + convergence_tol_nats=1.0e-3, time_guard=0, + time_guard_tol_nats=1.0e-3, local_log_normalization=0.0, + time_outside_tol_nats=-23.0, + total_value_error_budget_nats=1.0e-3, + reserve_log_offset=0.0, node_concentration=1.0, + mode_match_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + reserve_angular_kernel=None, reserve_time_focus=None, + reserve_time_cover=None): + """Select an accepted local value or execute the exact table reserve. + + This is the first operational fixed-point composition seam. The caller + supplies two immutable host- or device-built mode plans; this device + function evaluates the empirical gate and uses + :func:`anglemarg.coefficient_table_distphipsimarg_exact` only on a decline. + Thus accepted high-SNR rows pay fixed local work per retained mode after + bounded discovery; broad, + unresolved, capacity-limited, or otherwise unhealthy rows retain the sample + through the established dense/exact coefficient reserve. + + Measures remain explicit. ``reserve_log_weights`` owns the fixed-grid + distance quadrature measure and ``time_weights`` owns the reserve time + integral. If ``reserve_time_nodes`` is supplied, it names sub-sample + positions in the unguarded target window and the coefficient table is + reconstructed there only inside the declined branch. This band-limited + reserve must cover either the complete target window or the contiguous + interval enclosing every retained scout cell. A cropped interval is usable + only with the plan's discarded-time bound. In both cases the node spacing + must remain everywhere finer than one native sample and carry an external + resolution warrant through ``reserve_time_resolution_warranted`` and the + independently evaluated lower-resolution ``reserve_time_check_value``, as + well as the same two-guard convergence comparison used by the local path. + Resolution error, guard error, and the maximum omitted-time correction are + charged to one ``total_value_error_budget_nats`` allowance. Without + nodes, the legacy reserve uses native target samples only as a diagnostic + and is always fail-closed. Native-sample angular exactness does not certify + a time integral once its peak is narrower than a sample. The external + warrant is a scalar JAX boolean so callers may bind it to a per-row + convergence record; the caller remains responsible for proving that it + describes the exact nodes, weights, and full interval passed here. + Alternatively ``reserve_time_check_nodes``/``reserve_time_check_weights`` + name a strictly coarser rule on the same target window. The controller + then evaluates that rule itself, inside the declined branch only, and the + resolution warrant is structural: both rules come from the same reflected + primitive, the check rule is coarser, and the two values must agree within + ``reserve_time_resolution_tol_nats``. An accepted local row never pays + for either reserve evaluation. + + If ``JAX_ILE_DISTMARG_GH`` is active, the established reserve + instead reads the support from ``reserve_x_grid`` and uses its normalized + volumetric ``x**-4`` measure. The + local branch owns a continuous ``x**-4 dx dtime_sample dphi du`` integral, + so ``local_log_normalization`` must convert that measure to the reserve's + normalization. ``reserve_log_offset`` is the reserve's separately + recorded global log-measure conversion. It is applied to both the + included reserve value and a cropped reserve's discarded-time bound, so + their omitted/included ratio is invariant to that conversion. Neither + conversion is inferred from a distance-prior name. This prevents an + unnormalized prototype value from silently replacing a production result. + + ``reserve_angular_kernel(target_table, C_B) -> lnL_t`` is the angular + and distance kernel evaluated on the reserve's time nodes; ``None`` is + :func:`anglemarg.coefficient_table_distphipsimarg_exact` with the + arguments above. ``reserve_time_focus=(centre, half_width)`` (samples) + names where the rule is fine; the node carrying the largest evaluated + ``lnL_t`` must then lie within ``half_width`` of ``centre`` or the rule + is unwarranted (``reserve_time_focus_ok``). A peak-local rule and its + check share the block, so a misplaced block agrees with itself; this is + the certificate that sees it (measured: 0.22 nat, warranted, on a + rung-160 row with the block 0.17 samples off the maximum). + ``reserve_time_cover=(lo, hi, outside_log_bound)`` replaces the plans' + time cover in the cropped-cover warrant: the rule must span ``[lo, hi]`` + (samples) and ``outside_log_bound`` (already in the reserve's units) is + charged as the omitted mass. It is how a support-limited reserve rule + certifies what it left out without leaning on the local branch's plan. A peak-local time rule (``peaklocal_time_reserve``) + may repeat a position; a repeated position carries zero weight, so the + node checks below ask for a NON-DECREASING rule. + + Ledger field ``accepted_local`` is the empirical local disposition, while + the returned ``usable`` describes the selected result after reserve + execution. A local + decline is never a waveform failure. A nonfinite reserve is reported as an + integration failure and remains unusable; it is not relabeled as a missed + waveform evaluation. + """ + from . import anglemarg as _anglemarg + from .core import _time_marginalize + + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + time_guard = int(time_guard) + n_target = C_A_t.shape[-1] - 2 * time_guard + time_weights = jnp.asarray(time_weights, dtype=jnp.float64) + use_bandlimited_time = reserve_time_nodes is not None + if use_bandlimited_time: + if time_guard < 2: + raise ValueError( + "band-limited reserve requires time_guard >= 2") + reserve_time_nodes = jnp.asarray( + reserve_time_nodes, dtype=jnp.float64) + if reserve_time_nodes.ndim != 1 or reserve_time_nodes.size < 2: + raise ValueError( + "reserve_time_nodes must be a one-dimensional rule") + if (time_weights.ndim != 1 + or time_weights.shape[0] != reserve_time_nodes.shape[0]): + raise ValueError( + "time_weights must match reserve_time_nodes") + elif time_weights.ndim != 1 or time_weights.shape[0] != n_target: + raise ValueError("time_weights must match the unguarded target window") + reserve_time_resolution_warranted = jnp.asarray( + reserve_time_resolution_warranted, dtype=bool) + if reserve_time_resolution_warranted.ndim != 0: + raise ValueError("reserve_time_resolution_warranted must be scalar") + reserve_time_check_value = jnp.asarray( + reserve_time_check_value, dtype=jnp.float64) + if reserve_time_check_value.ndim != 0: + raise ValueError("reserve_time_check_value must be scalar") + use_internal_check = reserve_time_check_nodes is not None + if use_internal_check: + if not use_bandlimited_time: + raise ValueError( + "reserve_time_check_nodes requires reserve_time_nodes") + if reserve_time_check_weights is None: + raise ValueError( + "reserve_time_check_nodes requires reserve_time_check_weights") + reserve_time_check_nodes = jnp.asarray( + reserve_time_check_nodes, dtype=jnp.float64) + reserve_time_check_weights = jnp.asarray( + reserve_time_check_weights, dtype=jnp.float64) + if (reserve_time_check_nodes.ndim != 1 + or reserve_time_check_nodes.size < 2 + or reserve_time_check_weights.shape + != reserve_time_check_nodes.shape): + raise ValueError( + "reserve_time_check_nodes/weights must be one matching rule") + if (not np.isfinite(float(reserve_time_resolution_tol_nats)) + or not float(reserve_time_resolution_tol_nats) > 0.0): + raise ValueError( + "reserve_time_resolution_tol_nats must be finite and positive") + if reserve_m_max is None: + reserve_m_max = int(C_A_t.shape[0] - 1) + reserve_m_max = int(reserve_m_max) + if reserve_angular_kernel is None: + def reserve_angular_kernel(table, norm_table): + return _anglemarg.coefficient_table_distphipsimarg_exact( + table, norm_table, reserve_x_grid, reserve_log_weights, + amp_sizing=float(reserve_amp_sizing), m_max=reserve_m_max, + dense_chunk=int(reserve_dense_chunk), + grid_block=int(reserve_grid_block)) + + local_value, accepted_local, local_ledger = empirical_enrichment_marginalize( + C_A_t, C_B, base_plan, enriched_plan, x_min, x_max, + base_order=int(base_order), base_check_order=int(base_check_order), + enriched_order=int(enriched_order), + enriched_check_order=int(enriched_check_order), + convergence_tol_nats=float(convergence_tol_nats), + time_guard=time_guard, + time_guard_tol_nats=float(time_guard_tol_nats), + time_outside_tol_nats=float(time_outside_tol_nats), + total_value_error_budget_nats=float(total_value_error_budget_nats), + log_normalization=float(local_log_normalization), + node_concentration=float(node_concentration), + mode_match_tol=mode_match_tol) + + def _reserve(_): + if use_bandlimited_time: + flat_table = C_A_t.reshape((-1, C_A_t.shape[-1])) + coeff, frequency, offset = _time_primitive_spectrum( + flat_table, time_guard) + target_table = _evaluate_time_spectrum( + coeff, frequency, reserve_time_nodes, offset).reshape( + C_A_t.shape[:-1] + (reserve_time_nodes.size,)) + elif time_guard: + target_table = C_A_t[..., time_guard:-time_guard] + else: + target_table = C_A_t + lnL_t = reserve_angular_kernel(target_table, C_B) + reserve_value = (_time_marginalize(lnL_t, time_weights)[0] + + float(reserve_log_offset)) + if use_bandlimited_time: + inner_guard = time_guard // 2 + trim = time_guard - inner_guard + inner_table = C_A_t[..., trim:-trim] + inner_flat = inner_table.reshape((-1, inner_table.shape[-1])) + coeff_inner, frequency_inner, offset_inner = ( + _time_primitive_spectrum(inner_flat, inner_guard)) + target_inner = _evaluate_time_spectrum( + coeff_inner, frequency_inner, reserve_time_nodes, + offset_inner).reshape( + C_A_t.shape[:-1] + (reserve_time_nodes.size,)) + lnL_inner = reserve_angular_kernel(target_inner, C_B) + guard_value = (_time_marginalize(lnL_inner, time_weights)[0] + + float(reserve_log_offset)) + else: + guard_value = jnp.asarray(jnp.nan, dtype=jnp.float64) + if use_internal_check: + check_table = _evaluate_time_spectrum( + coeff, frequency, reserve_time_check_nodes, offset).reshape( + C_A_t.shape[:-1] + (reserve_time_check_nodes.size,)) + lnL_check = reserve_angular_kernel(check_table, C_B) + check_value = ( + _time_marginalize(lnL_check, reserve_time_check_weights)[0] + + float(reserve_log_offset)) + else: + check_value = reserve_time_check_value + if use_bandlimited_time: + peak_node = reserve_time_nodes[jnp.argmax(lnL_t[0])] + else: + peak_node = jnp.asarray(jnp.nan, dtype=jnp.float64) + return reserve_value, reserve_value, guard_value, check_value, peak_node + + def _accepted(_): + nan = jnp.asarray(jnp.nan, dtype=jnp.float64) + return local_value, nan, nan, reserve_time_check_value, nan + + # Both branches are rematerialized. Reverse-mode AD through ``lax.cond`` + # stores backward residuals for BOTH branches whatever the predicate, and + # the dense reserve's residuals at production amplitude are tens of GiB + # (a 158 GiB allocation was requested on a rho 163 production row, PR #278 + # follow-up). With checkpointing the backward pass recomputes the taken + # branch instead, so gradient memory is one forward evaluation. + (selected_value, reserve_value, reserve_guard_value, + reserve_time_check_value, reserve_time_peak_node) = jax.lax.cond( + accepted_local, jax.checkpoint(_accepted), jax.checkpoint(_reserve), + operand=None) + if reserve_time_focus is None: + reserve_time_focus_ok = jnp.asarray(True) + reserve_time_focus_offset = jnp.asarray(jnp.nan, dtype=jnp.float64) + else: + focus_centre = jnp.asarray(reserve_time_focus[0], dtype=jnp.float64) + focus_half = jnp.asarray(reserve_time_focus[1], dtype=jnp.float64) + reserve_time_focus_offset = jnp.abs(reserve_time_peak_node - focus_centre) + reserve_time_focus_ok = ( + (~(~accepted_local)) + | (jnp.isfinite(reserve_time_focus_offset) + & (reserve_time_focus_offset <= focus_half))) + if use_internal_check: + # Structural warrant: same primitive, same window, strictly coarser + # check rule with a valid measure. The value comparison itself is + # still applied below through reserve_time_resolution_validated. + check_rule_valid = ( + jnp.all(jnp.isfinite(reserve_time_check_nodes)) + & jnp.all(reserve_time_check_nodes[1:] >= reserve_time_check_nodes[:-1]) + & (reserve_time_check_nodes[0] == reserve_time_nodes[0]) + & (reserve_time_check_nodes[-1] == reserve_time_nodes[-1]) + & (jnp.max(jnp.diff(reserve_time_check_nodes)) + > jnp.max(jnp.diff(reserve_time_nodes))) + & jnp.all(jnp.isfinite(reserve_time_check_weights)) + & jnp.all(reserve_time_check_weights >= 0.0) + & (jnp.sum(reserve_time_check_weights) > 0.0)) + reserve_time_resolution_warranted = ( + reserve_time_resolution_warranted | check_rule_valid) + reserve_executed = ~accepted_local + reserve_finite = jnp.isfinite(reserve_value) + if use_bandlimited_time: + reserve_time_nodes_finite = jnp.all(jnp.isfinite(reserve_time_nodes)) + # Non-decreasing: a peak-local rule repeats a position where a mode + # slot is dead or a block is clipped, and the repeat carries no weight. + # Compared as slices, not as ``diff >= 0``: under jit XLA fuses the + # subtraction of two bitwise-equal products into a multiply-add that + # rounds to -1e-15 (measured), which would fail a rule numpy calls + # sorted. + reserve_time_nodes_increasing = jnp.all( + reserve_time_nodes[1:] >= reserve_time_nodes[:-1]) + reserve_time_subsampled = jnp.max( + jnp.diff(reserve_time_nodes)) < 1.0 + reserve_time_weights_valid = ( + jnp.all(jnp.isfinite(time_weights)) + & jnp.all(time_weights >= 0.0) + & (jnp.sum(time_weights) > 0.0)) + reserve_time_nodes_in_support = jnp.all( + (reserve_time_nodes >= 0.0) + & (reserve_time_nodes <= float(n_target - 1))) + reserve_time_nodes_cover_target = ( + (reserve_time_nodes[0] == 0.0) + & (reserve_time_nodes[-1] == float(n_target - 1))) + if reserve_time_cover is None: + reserve_required_time_min = jnp.minimum( + base_plan.time_cover_min_sample, + enriched_plan.time_cover_min_sample) + reserve_required_time_max = jnp.maximum( + base_plan.time_cover_max_sample, + enriched_plan.time_cover_max_sample) + reserve_time_plan_cover_certified = ( + base_plan.time_outside_bound_certified + & enriched_plan.time_outside_bound_certified) + cover_outside_log_bound = ( + jnp.minimum(base_plan.time_outside_log_bound, + enriched_plan.time_outside_log_bound) + + float(local_log_normalization) + + float(reserve_log_offset)) + else: + reserve_required_time_min = jnp.asarray(reserve_time_cover[0], dtype=jnp.float64) + reserve_required_time_max = jnp.asarray(reserve_time_cover[1], dtype=jnp.float64) + reserve_time_plan_cover_certified = jnp.asarray(True) + cover_outside_log_bound = jnp.asarray(reserve_time_cover[2], dtype=jnp.float64) + reserve_time_plan_cover_finite = ( + jnp.isfinite(reserve_required_time_min) + & jnp.isfinite(reserve_required_time_max) + & (reserve_required_time_min < reserve_required_time_max)) + reserve_time_nodes_cover_plans = ( + (reserve_time_nodes[0] <= reserve_required_time_min) + & (reserve_time_nodes[-1] >= reserve_required_time_max)) + reserve_time_cropped_cover_warranted = ( + reserve_time_plan_cover_finite + & reserve_time_nodes_cover_plans + & reserve_time_plan_cover_certified) + reserve_time_interval_warranted = ( + reserve_time_nodes_cover_target + | reserve_time_cropped_cover_warranted) + reserve_time_outside_log_bound = jnp.where( + reserve_time_nodes_cover_target, -jnp.inf, cover_outside_log_bound) + reserve_time_tail_margin = ( + reserve_time_outside_log_bound - reserve_value) + reserve_time_tail_correction = jnp.logaddexp( + 0.0, reserve_time_tail_margin) + reserve_time_tail_ok = ( + reserve_time_tail_margin < float(time_outside_tol_nats)) + reserve_time_guard_error = jnp.abs( + reserve_value - reserve_guard_value) + reserve_time_guard_validated = ( + reserve_executed & reserve_finite + & jnp.isfinite(reserve_guard_value) + & (reserve_time_guard_error <= float(time_guard_tol_nats))) + reserve_time_resolution_error = jnp.abs( + reserve_value - reserve_time_check_value) + reserve_time_resolution_validated = ( + reserve_executed & reserve_finite + & reserve_time_resolution_warranted + & jnp.isfinite(reserve_time_check_value) + & (reserve_time_resolution_error + <= float(reserve_time_resolution_tol_nats))) + reserve_time_error_score = ( + reserve_time_guard_error + reserve_time_resolution_error + + reserve_time_tail_correction) + reserve_time_error_budget_ok = ( + jnp.isfinite(reserve_time_error_score) + & (reserve_time_error_score + <= float(total_value_error_budget_nats))) + reserve_time_warranted = ( + reserve_time_focus_ok + & reserve_time_nodes_finite + & reserve_time_nodes_increasing + & reserve_time_subsampled + & reserve_time_weights_valid + & reserve_time_nodes_in_support + & reserve_time_interval_warranted + & reserve_time_tail_ok + & reserve_time_resolution_validated + & reserve_time_guard_validated + & reserve_time_error_budget_ok) + reserve_time_points = reserve_time_nodes.size + reserve_time_min = jnp.min(reserve_time_nodes) + reserve_time_max = jnp.max(reserve_time_nodes) + else: + reserve_time_nodes_finite = jnp.asarray(True) + reserve_time_nodes_increasing = jnp.asarray(True) + reserve_time_subsampled = jnp.asarray(False) + reserve_time_weights_valid = ( + jnp.all(jnp.isfinite(time_weights)) + & jnp.all(time_weights >= 0.0) + & (jnp.sum(time_weights) > 0.0)) + reserve_time_nodes_in_support = jnp.asarray(True) + reserve_time_nodes_cover_target = jnp.asarray(True) + reserve_required_time_min = jnp.asarray(0.0) + reserve_required_time_max = jnp.asarray(float(n_target - 1)) + reserve_time_plan_cover_finite = jnp.asarray(False) + reserve_time_nodes_cover_plans = jnp.asarray(False) + reserve_time_plan_cover_certified = jnp.asarray(False) + reserve_time_cropped_cover_warranted = jnp.asarray(False) + reserve_time_interval_warranted = jnp.asarray(False) + reserve_time_outside_log_bound = jnp.asarray(jnp.inf) + reserve_time_tail_margin = jnp.asarray(jnp.inf) + reserve_time_tail_correction = jnp.asarray(jnp.inf) + reserve_time_tail_ok = jnp.asarray(False) + reserve_time_guard_error = jnp.asarray(jnp.nan) + reserve_time_guard_validated = jnp.asarray(False) + reserve_time_resolution_error = jnp.asarray(jnp.nan) + reserve_time_resolution_validated = jnp.asarray(False) + reserve_time_error_score = jnp.asarray(jnp.inf) + reserve_time_error_budget_ok = jnp.asarray(False) + reserve_time_warranted = jnp.asarray(False) + reserve_time_points = n_target + reserve_time_min = jnp.asarray(0.0) + reserve_time_max = jnp.asarray(float(n_target - 1)) + reserve_time_failed = reserve_executed & (~reserve_time_warranted) + reserve_failed = reserve_executed & ( + (~reserve_finite) | reserve_time_failed) + usable = accepted_local | ( + reserve_executed & reserve_finite & reserve_time_warranted) + nphi_reserve, nu_reserve = _anglemarg._dense_grid_sizes( + float(reserve_amp_sizing), m_max=reserve_m_max) + reserve_gh_nodes = int(_anglemarg._core._DISTMARG_GH_N) + reserve_input_distance_points = int(jnp.asarray(reserve_x_grid).size) + ledger = dict(local_ledger) + ledger.update({ + "local_fallback_required": local_ledger["fallback_required"], + "local_reconciles": local_ledger["reconciles"], + }) + ledger.update({ + "accepted": usable, + "fallback_required": reserve_failed, + "reconciles": ( + usable.astype(jnp.int32) + + reserve_failed.astype(jnp.int32)) == 1, + "accepted_local": accepted_local, + "selected_value_is_local": accepted_local, + "selected_value_is_warranted_reserve": ( + reserve_executed & reserve_finite & reserve_time_warranted), + "reserve_executed": reserve_executed, + "reserve_value": reserve_value, + "reserve_finite": reserve_finite, + "reserve_failed": reserve_failed, + "reserve_time_failed": reserve_time_failed, + "reserve_uses_bandlimited_time": jnp.asarray(use_bandlimited_time), + "reserve_uses_native_time": jnp.asarray(not use_bandlimited_time), + "reserve_native_time_warranted": jnp.asarray(False), + "reserve_time_resolution_warranted": ( + reserve_time_resolution_warranted), + "reserve_time_check_rule_internal": jnp.asarray(use_internal_check), + "reserve_time_peak_node": reserve_time_peak_node, + "reserve_time_focus_offset_samples": reserve_time_focus_offset, + "reserve_time_focus_ok": reserve_time_focus_ok, + "reserve_time_check_value": reserve_time_check_value, + "reserve_time_resolution_error_nats": ( + reserve_time_resolution_error), + "reserve_time_resolution_tol_nats": jnp.asarray( + float(reserve_time_resolution_tol_nats)), + "reserve_time_resolution_validated": ( + reserve_time_resolution_validated), + "reserve_time_nodes_finite": reserve_time_nodes_finite, + "reserve_time_nodes_increasing": reserve_time_nodes_increasing, + "reserve_time_subsampled": reserve_time_subsampled, + "reserve_time_weights_valid": reserve_time_weights_valid, + "reserve_time_nodes_in_support": reserve_time_nodes_in_support, + "reserve_time_nodes_cover_target": reserve_time_nodes_cover_target, + "reserve_required_time_min_sample": reserve_required_time_min, + "reserve_required_time_max_sample": reserve_required_time_max, + "reserve_time_plan_cover_finite": reserve_time_plan_cover_finite, + "reserve_time_nodes_cover_plans": reserve_time_nodes_cover_plans, + "reserve_time_plan_cover_certified": ( + reserve_time_plan_cover_certified), + "reserve_time_cropped_cover_warranted": ( + reserve_time_cropped_cover_warranted), + "reserve_time_interval_warranted": reserve_time_interval_warranted, + "reserve_time_outside_log_bound": reserve_time_outside_log_bound, + "reserve_time_tail_margin": reserve_time_tail_margin, + "reserve_time_tail_correction_nats": ( + reserve_time_tail_correction), + "reserve_time_tail_ok": reserve_time_tail_ok, + "reserve_time_guard_validated": reserve_time_guard_validated, + "reserve_time_guard_error": reserve_time_guard_error, + "reserve_time_guard_value": reserve_guard_value, + "reserve_time_error_score_nats": reserve_time_error_score, + "reserve_time_error_budget_ok": reserve_time_error_budget_ok, + "reserve_time_warranted": reserve_time_warranted, + "reserve_time_min_sample": reserve_time_min, + "reserve_time_max_sample": reserve_time_max, + "usable": usable, + "sample_retained_after_local_decline": ( + reserve_executed & reserve_finite & reserve_time_warranted), + "decline_is_waveform_failure": jnp.asarray(False), + "selected_nonfinite_is_integration_failure": ( + reserve_executed & (~reserve_finite)), + "reserve_nphi": jnp.asarray(nphi_reserve), + "reserve_nu": jnp.asarray(nu_reserve), + "reserve_angle_points": jnp.asarray(nphi_reserve * nu_reserve), + "reserve_distance_support_points": jnp.asarray( + reserve_input_distance_points), + "reserve_distance_points": jnp.asarray( + reserve_gh_nodes if reserve_gh_nodes + else reserve_input_distance_points), + "reserve_uses_adaptive_distance": jnp.asarray( + reserve_gh_nodes > 0), + "reserve_distance_gh_nodes": jnp.asarray(reserve_gh_nodes), + "reserve_time_points": jnp.asarray(reserve_time_points), + "reserve_dense_chunk": jnp.asarray(int(reserve_dense_chunk)), + "reserve_grid_block": jnp.asarray(int(reserve_grid_block)), + "local_log_normalization": jnp.asarray( + float(local_log_normalization)), + "reserve_log_offset": jnp.asarray(float(reserve_log_offset)), + "disposition_reconciles": ( + accepted_local.astype(jnp.int32) + + reserve_executed.astype(jnp.int32)) == 1, + }) + return selected_value, usable, ledger + + +def empirical_enrichment_with_exact_reserve_sequential_batch( + C_A_t, C_B, base_plans, enriched_plans, x_min, x_max, *, + reserve_x_grid, reserve_log_weights, time_weights, + reserve_amp_sizing, reserve_m_max=None, + reserve_dense_chunk=8, reserve_grid_block=32, + reserve_time_nodes=None, reserve_time_resolution_warranted=False, + reserve_time_check_value=np.nan, + reserve_time_check_nodes=None, reserve_time_check_weights=None, + reserve_time_resolution_tol_nats=1.0e-3, + base_order=13, base_check_order=19, + enriched_order=19, enriched_check_order=25, + convergence_tol_nats=1.0e-3, time_guard=0, + time_guard_tol_nats=1.0e-3, local_log_normalization=0.0, + time_outside_tol_nats=-23.0, + total_value_error_budget_nats=1.0e-3, + reserve_log_offset=0.0, node_concentration=1.0, + mode_match_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5), + reserve_angular_kernel=None, reserve_time_focus=None, + reserve_time_cover=None): + """Apply the scalar controller sequentially to a fixed-size batch. + + A direct ``vmap`` of :func:`empirical_enrichment_with_exact_reserve` + rewrites its scalar conditional as a batched selection and can therefore + execute the dense reserve for accepted rows. This wrapper deliberately + uses ``lax.map`` so each row reaches the scalar conditional independently; + reserve workspace scales with one row rather than the batch size. For a + sparse, variable-size decline set, a host controller should instead vmap + only :func:`empirical_enrichment_marginalize`, compact the declined rows, + and invoke the reserve on that compact set. + + ``C_B`` and the time/distance rules are shared across the batch. Every + leaf of ``base_plans`` and ``enriched_plans`` must have a leading batch + dimension. ``reserve_time_resolution_warranted`` and + ``reserve_time_check_value`` may each be one scalar policy value or one + scalar per row. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + if C_A_t.ndim != 4: + raise ValueError("C_A_t must have shape (batch,KP,KS,Ntime)") + batch = C_A_t.shape[0] + if C_B.ndim == 2: + C_B = jnp.broadcast_to(C_B, (batch,) + C_B.shape) + elif C_B.ndim != 3 or C_B.shape[0] != batch: + raise ValueError("C_B must be shared or have a leading batch axis") + for name, plans in (("base_plans", base_plans), + ("enriched_plans", enriched_plans)): + for leaf in jax.tree.leaves(plans): + if jnp.asarray(leaf).ndim < 1 or jnp.asarray(leaf).shape[0] != batch: + raise ValueError("%s must have a leading batch axis" % name) + warrant = jnp.asarray(reserve_time_resolution_warranted, dtype=bool) + if warrant.ndim == 0: + warrant = jnp.broadcast_to(warrant, (batch,)) + elif warrant.shape != (batch,): + raise ValueError( + "reserve_time_resolution_warranted must be scalar or per-row") + check_value = jnp.asarray(reserve_time_check_value, dtype=jnp.float64) + if check_value.ndim == 0: + check_value = jnp.broadcast_to(check_value, (batch,)) + elif check_value.shape != (batch,): + raise ValueError("reserve_time_check_value must be scalar or per-row") + + def _row(args): + (table, norm, base_plan, enriched_plan, row_warrant, + row_check_value) = args + return empirical_enrichment_with_exact_reserve( + table, norm, base_plan, enriched_plan, x_min, x_max, + reserve_x_grid=reserve_x_grid, + reserve_log_weights=reserve_log_weights, + time_weights=time_weights, + reserve_amp_sizing=reserve_amp_sizing, + reserve_m_max=reserve_m_max, + reserve_dense_chunk=reserve_dense_chunk, + reserve_grid_block=reserve_grid_block, + reserve_time_nodes=reserve_time_nodes, + reserve_time_resolution_warranted=row_warrant, + reserve_time_check_value=row_check_value, + reserve_time_check_nodes=reserve_time_check_nodes, + reserve_time_check_weights=reserve_time_check_weights, + reserve_time_resolution_tol_nats=( + reserve_time_resolution_tol_nats), + base_order=base_order, base_check_order=base_check_order, + enriched_order=enriched_order, + enriched_check_order=enriched_check_order, + convergence_tol_nats=convergence_tol_nats, + time_guard=time_guard, time_guard_tol_nats=time_guard_tol_nats, + local_log_normalization=local_log_normalization, + time_outside_tol_nats=time_outside_tol_nats, + total_value_error_budget_nats=total_value_error_budget_nats, + reserve_log_offset=reserve_log_offset, + node_concentration=node_concentration, + mode_match_tol=mode_match_tol, + reserve_angular_kernel=reserve_angular_kernel, + reserve_time_focus=reserve_time_focus, + reserve_time_cover=reserve_time_cover) + + selected, usable, ledger = jax.lax.map( + _row, (C_A_t, C_B, base_plans, enriched_plans, warrant, + check_value)) + ledger = dict(ledger) + ledger["reserve_batch_execution_sequential"] = jnp.ones( + (batch,), dtype=bool) + return selected, usable, ledger diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py index c29a4bea3..323d93307 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py @@ -68,6 +68,8 @@ by grid size. """ +import sys + import numpy as np import jax import jax.numpy as jnp @@ -85,11 +87,15 @@ "angle_sample_grid_sizes", "angle_coefficient_tables", "estimate_angle_amplitude", + "coefficient_table_distphipsimarg_exact", + "coefficient_table_distphipsimarg_laplace", "fused_log_likelihood_distphipsimarg_exact", "fused_log_likelihood_distphipsimarg_laplace", "choose_angle_marg_scheme", "fused_log_likelihood_distphipsimarg_peaklocal", + "fused_log_likelihood_distphipsimarg_multipeak", "gh_laplace_supported", + "gh_laplace_supported_for_data", "ANGLE_MARG_CROSSOVER_AMPLITUDE", ] @@ -128,11 +134,17 @@ # 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. +# ANGLE_MARG_CROSSOVER_AMPLITUDE (rho ~26-30; see that constant's note for why +# 26 and not the 21 this line used to say), which is an ACCURACY crossover. +# But 'laplace' on the DEFAULT UNIFORM distance grid 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. (This +# comment used to say laplace 'cannot use the per-sample adaptive distance +# quadrature'. It can, and does, for m_max <= _GH_PSI_M_MAX via +# _gh_psi_node_offsets, gated by gh_laplace_supported; the 43.2 nats is what +# the UNIFORM grid costs, not what the scheme costs.) The log-uniform distance +# grid is opt-in, so a caller who selects laplace without moving the distance +# axis with it pays that 43.2 nats. # 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 @@ -149,7 +161,8 @@ # 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") +ANGLE_MARG_CHOICES = ("grid", "exact", "laplace", "peak-local", "phi-local", + "multipeak", "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 @@ -161,10 +174,64 @@ # --------------------------------------------------------------------------- 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 +# auto selector compares the MARGINED data-derived bound to this, so laplace +# engages below rho = 30. TWO DIFFERENT NUMBERS LIVE HERE and an earlier version +# of this comment ran them together: +# * the INTENDED margin is the `margin=2.0` argument of +# estimate_angle_amplitude -- a deliberate parameter, not an estimate; +# * the ratio the SWITCH actually keys on was measured AT THE LADDER INJECTION +# (rho = 40.77): bound 1109.17 against the nominal rho^2/2 = 831.1, i.e. +# 1.335, not 2. +# That gives 450 / 1.335 -> engagement at nominal A ~ 337, rho ~ 26.0, which is +# what the manuscript quotes; this comment previously said rho ~ 21 by assuming +# the factor equalled the margin. Keep the two distinct or the code and the +# paper quote different crossovers for the same switch. +# +# THREE LIMITS ON 1.335, so it is not read as more than it is: +# (a) the denominator is the NOMINAL rho^2/2, not a measured maximum of the +# (phi,psi) exponent. Reading "the raw estimator sits at 0.667x TRUE A" +# goes through the identification true A == rho^2/2, which is this file's +# own convention but is not a measurement. +# (b) the ratio is constant to 6e-5 across rungs rho = 40.77 ... 652.31. That +# is ARITHMETIC, NOT EVIDENCE: the ladder is one injection replayed at +# scaled amplitudes, so the exponent rescales uniformly and the ratio is +# forced. Four decades of agreement validate nothing about the margin. +# (c) it is ONE injection's sky-sample realization. The shortfall's size is +# set by how sharp the sky peak is relative to the sample -- the sample is +# random draws plus a coarse uniform grid and does not contain the +# injection's sky position, while the exponent is sharp enough that 1 of +# 9824 (sky, time) points sits within 23 nats of the peak. A shortfall is +# expected by design there, and nothing here bounds it for another event. +# So: at this injection the margin is load-bearing rather than decorative -- by +# about 1.5x -- and that is the whole claim. +# +# TWO OTHER RATIOS CIRCULATE FOR THIS LADDER AND NEITHER IS THE MARGIN. Measured +# on it: raw-estimator/(rho^2/2) = 0.1888 and margined-bound/guess = 7.069. Both +# are the SNR-GUESS DEFICIT SQUARED -- guess_amp == guess_snr^2/2 exactly, and +# rho/guess_snr = 2.3014 constant, so 0.1888 = 1/2.3014^2 and 7.069 = +# 2.3014^2 * 1.33465. guess_snr is the ABANDONED sizing route (external review +# removed it precisely because an underestimated SNR silently under-resolved the +# dense quadrature). If 7.069 lands here as "the margin" it inflates a ~1.5x +# effect to 7x and credits the live estimator with the dead route's deficit. +# They are easy to accept because they AGREE with the conclusion above -- for an +# unrelated reason -- so they read as corroboration and are not. +# The genuinely reportable fact in them: on this ladder guess_snr sits 2.30x +# below the true rho, so the abandoned route would have sized the dense grids +# from an amplitude too small BY A FACTOR THAT DEPENDS ON WHAT YOU DIVIDE BY -- +# 7.069x against the LIVE data-derived bound (the thing that sizes grids +# today, so this is the operative figure), and +# 5.296x against the nominal rho^2/2, +# the two differing by exactly the 1.335 above. A reader handed "7.069x" with no +# denominator cannot tell which, and will be off by 1.335 either way: that is the +# same unnamed-denominator defect this block exists to guard against, and I +# shipped it here one commit before fixing it. The docstring's stated failure +# mode, measured on a real configuration. One injection, one guess_snr. What actually protects the general case is +# _runtime_amp_failsafe, which recomputes the amplitude from the tables at the +# point of use and warns if it exceeds amp_sizing -- independent of whether the +# margin was well chosen. (Ratios measured by the paper-1 ladder session.) +# Early engagement is safe by measurement either way: laplace is at -1.8e-04 nats +# by A = 200 on the injection ladder (the table above, spelled to match it so a +# grep finds both) 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 @@ -202,7 +269,7 @@ def _data_m_max(data): def angle_coefficient_tables(data, ra, dec, incl, interp=JAX_INTERP_DEFAULT, - sample_chunk=None): + sample_chunk=None, guard=0): """Exact 2-D Fourier coefficient tables of A = Re kappa_unit, B = rho^2_unit. Samples :func:`core._accumulate_unit` on the Nyquist-sized @@ -218,15 +285,24 @@ def angle_coefficient_tables(data, ra, dec, incl, interp=JAX_INTERP_DEFAULT, 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) + ``guard`` requests primitive-only reconstruction support from the same + accumulation operation as the terminal band-limited path. The returned + time axis then has ``data.npts + 2*guard`` samples; callers must discard the + support after reconstruction and compare two guard widths before accepting. + + Memory: the tables are (m_max+1, 3, S, ntime) and (2*m_max+1, 5, S, ntime) 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)``. + Returns ``(C_A, C_B, meta)`` with grid sizes and the effective ``guard`` + and ``ntime`` support recorded in ``meta``. """ m_max = _data_m_max(data) + guard = int(guard) + if guard < 0: + raise ValueError("guard must be non-negative") nphi_s, npsi_s = angle_sample_grid_sizes(m_max) if sample_chunk is None: sample_chunk = npsi_s @@ -255,7 +331,7 @@ def _phase_table(kp_max, ks_max): dec = jnp.asarray(dec, dtype=jnp.float64) incl = jnp.asarray(incl, dtype=jnp.float64) S = ra.shape[0] - npts = data.npts + npts = data.npts + 2 * guard c = int(sample_chunk) nsteps = Ns // c @@ -273,7 +349,7 @@ def _step(carry, x): 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) + interp, False, guard=guard) A = ku.real.reshape(c, S, npts) B = rs.reshape(c, S, npts) CA = CA + jnp.einsum("ckq,cst->kqst", pA, A) @@ -283,7 +359,8 @@ def _step(carry, x): 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) + meta = dict(m_max=m_max, nphi_s=nphi_s, npsi_s=npsi_s, + guard=guard, ntime=npts) return C_A, C_B, meta @@ -750,45 +827,60 @@ def _draw(n, rng): 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. + No barrier is needed: the record is written synchronously at the Python + boundary by :func:`record_amp_failsafe`, never by a queued device effect. """ - 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. + """Host-side record of the deterministic output-cloud amplitude checks. - ``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. + ``barrier`` is accepted for API compatibility and ignored: there are no + queued device effects left to drain. The jitted kernels RETURN their + amplitude metric as ordinary data and the wrapper accumulates it + synchronously once each batch is ready, so a read here is already + ordered after every batch the caller has taken delivery of. + + Keeping host effects out of these graphs is load-bearing beyond tidiness: + ``jax/_src/compiler.py::_cache_write`` refuses to write a persistent cache + entry for any module carrying host callbacks ("because it uses host + callbacks"), so a ``jax.debug.print``/``jax.debug.callback`` anywhere in + the angle-marginalization graph makes the most expensive compile in RIFT + permanently uncacheable. 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.""" +def record_amp_failsafe(amp_call, amp_sizing, scheme_name): + """Accumulate one already-evaluated batch's amplitude maximum. + + Runs at the Python boundary after the device result is ready, never inside + a JIT. Maxima accumulate across chunks and calls for the whole event; the + likelihood values are neither altered nor filtered. + """ + amp_call = float(np.max(np.asarray(amp_call))) + amp_sizing = float(amp_sizing) + tripped = amp_call > AMP_FAILSAFE_TRIP_FACTOR * amp_sizing _AMP_FAILSAFE["n_calls"] += 1 - if bool(tripped): + _AMP_FAILSAFE["worst_amp"] = max(_AMP_FAILSAFE["worst_amp"], amp_call) + _AMP_FAILSAFE["amp_sizing"] = amp_sizing + _AMP_FAILSAFE["scheme"] = scheme_name + if 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 + sys.stderr.write( + "WARNING anglemarg/%s: this batch's coefficient tables reach an " + "amplitude scale ~%.4g (analytic over-reading expression), above " + "%gx the amp_sizing=%.4g 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.\n" + % (scheme_name, amp_call, AMP_FAILSAFE_TRIP_FACTOR, amp_sizing)) def _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, scheme_name): @@ -803,10 +895,11 @@ def _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, scheme_name): 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. + margin in N, i.e. 4x in amplitude). It RETURNS the metric as ordinary JAX + data; no value is altered and there are deliberately no host effects here, + because such effects make this graph ineligible for JAX's persistent + compilation cache. 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)) @@ -822,21 +915,9 @@ def _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, scheme_name): # 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) + # reasoning and the two rejected alternatives. This function alters no value; it + # RETURNS the metric as ordinary JAX data and the caller records it at the Python + # boundary. Everything is under stop_gradient so the check never enters the AD graph. # DELIBERATELY NOT FATAL, AND DELIBERATELY NOT A NaN. # # An earlier version returned NaN to "fail closed". That was worse than the @@ -851,29 +932,30 @@ def _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, scheme_name): # 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. + # So: the value is untouched, the run completes, and the metric is RETURNED + # for the host to record, 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. + # + # WHY THIS IS RETURNED RATHER THAN REPORTED FROM INSIDE THE GRAPH, which is + # what it used to be. Two reasons, and the second is why it changed. + # (1) Reliability: debug-callback effects may be dropped, duplicated or + # reordered under transformation, so a clean read never proved adequacy and + # every consumer had to say so. (2) Cacheability, which is load-bearing: + # jax/_src/compiler.py::_cache_write declines to write a persistent cache + # entry for any module that carries host callbacks. With one in this + # function the whole angle-marginalization graph -- the most expensive + # compile in RIFT, minutes for peak-local -- could never be persistently + # cached, on any scheme, in any run. Returning the metric fixes both: the + # record is now synchronous and exact for every batch the caller takes + # delivery of. # - # 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) + # The COVERAGE that buys is narrower than "every traced call", and is stated + # rather than implied: the wrapper records on the BATCHED path (pilot, + # reweight and final output-cloud evaluations, i.e. every point that reaches + # a published artifact) and deliberately not on the scalar AD/flow-training + # path, whose proposals do not enter those artifacts. + return amp_call def _require_amp_sizing(amp_sizing): @@ -919,46 +1001,72 @@ def _pad_chunks(values, 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. +def coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w_grid, *, amp_sizing=None, m_max=None, + dense_chunk=8, grid_block=32, return_amp=False): + """Stream the exact angle/distance reserve from coefficient tables. + + This is the common fixed-point seam between the dense/exact reserve and + all-axis peak-local work. ``C_A`` and ``C_B`` may be the batched + ``(KP,KS,S,Ntime)`` tables returned by :func:`angle_coefficient_tables`, or + an unbatched ``C_A`` of shape ``(KP,KS,Ntime)`` together with a collapsed, + time-independent ``C_B`` of shape ``(KP,KS)``. The latter is exactly the + compact representation used by ``all_axis_peaklocal``. + + The result has shape ``(S,Ntime)`` and is normalized over the two periodic + angles. With the fixed-grid distance path, normalization is entirely + determined by ``log_w_grid``. When ``JAX_ILE_DISTMARG_GH`` is enabled, the + existing adaptive-distance contract instead reads only the support from + ``x_grid`` and applies its built-in normalized volumetric ``x**-4`` measure. + Time is deliberately not integrated here. Keeping those measures explicit + prevents an empirical local result using continuous ``x**-4 dx`` from being + silently compared with a differently normalized production distance prior. + + Dense angle coordinates are generated procedurally from each scan index and + distance blocks are streamed, so peak workspace is controlled by + ``dense_chunk`` and ``grid_block`` rather than the complete dense angle + lattice. No waveform or packed U,V/Q contraction is repeated. """ + 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) 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 + if C_A.ndim == 3: + C_A = C_A[:, :, None, :] + if C_A.ndim != 4: + raise ValueError("C_A must have shape (KP,KS[,S],Ntime)") + if C_B.ndim == 2: + C_B = jnp.broadcast_to( + C_B[:, :, None, None], + C_B.shape + (C_A.shape[2], C_A.shape[3])) + if C_B.ndim != 4 or C_B.shape[2:] != C_A.shape[2:]: + raise ValueError( + "C_B must be collapsed (KP,KS) or match C_A sample/time axes") + if C_A.shape[1] % 2 != 1 or C_B.shape[1] % 2 != 1: + raise ValueError("angular harmonic axes must have odd length") + if x_grid.ndim != 1 or log_w_grid.shape != x_grid.shape or x_grid.size < 2: + raise ValueError("x_grid/log_w_grid must be matching one-dimensional grids") + if not (int(dense_chunk) > 0 and int(grid_block) > 0): + raise ValueError("dense_chunk and grid_block must be positive") + inferred_m_max = int(C_A.shape[0] - 1) + if m_max is None: + m_max = inferred_m_max + m_max = int(m_max) + if m_max != inferred_m_max: + raise ValueError("m_max does not match the C_A harmonic order") + if C_B.shape[0] < 2 * m_max + 1: + raise ValueError("C_B does not contain the required norm harmonics") + S = C_A.shape[2] + npts = C_A.shape[3] 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") + amp_call = _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, + "exact-tables") + nphi_d, nu_d = _dense_grid_sizes(amp_sizing, m_max=m_max) c = int(dense_chunk) - phi_x, u_x, lw_x = _pad_chunks([PH.ravel(), UU.ravel()], c) n_dense = nphi_d * nu_d + nsteps = (n_dense + c - 1) // c + lane = jnp.arange(c, dtype=jnp.int32) a_g = x_grid b_g = -0.5 * jnp.square(x_grid) @@ -968,30 +1076,74 @@ def fused_log_likelihood_distphipsimarg_exact( x_min = jnp.min(x_grid) x_max = jnp.max(x_grid) - def _step(carry, x): + def _step(carry, step): m, s = carry - phw, uw, lww = x - A = _reconstruct_field(C_A, phw, uw) # (c,S,npts) + flat = step * c + lane + live = flat < n_dense + safe = jnp.minimum(flat, n_dense - 1) + iphi = safe // nu_d + iu = safe - iphi * nu_d + phw = (2.0 * jnp.pi / float(nphi_d)) * iphi + uw = (2.0 * jnp.pi / float(nu_d)) * iu + lww = jnp.where(live, 0.0, -jnp.inf) + A = _reconstruct_field(C_A, phw, uw) 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 = _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)) + (m, s), _ = jax.lax.scan( + jax.checkpoint(_step), (m0, s0), + jnp.arange(nsteps, dtype=jnp.int32)) 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) + return (lnL_t, amp_call) if return_amp else lnL_t + + +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, + return_amp=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) + lnL_t, amp_call = coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w_grid, amp_sizing=amp_sizing, + m_max=meta["m_max"], dense_chunk=dense_chunk, + grid_block=grid_block, return_amp=True) + out = lnL_t if return_lnLt else _time_marginalize_terminal( + lnL_t, data, time_quadrature) + return (out, amp_call) if return_amp else out # --------------------------------------------------------------------------- @@ -1040,6 +1192,30 @@ def _step(carry, x): # global maximum and carries negligible weight. _LAPLACE_MAX_ROOTS = 4 +#: Maximum number of independent ``(sample, time)`` points presented to one +#: distance/psi kernel invocation. This is an execution-only tile: neither a +#: quadrature count nor an accuracy knob. At the shipped ``QCH=16``, +#: ``dist_block=4`` and ``phi_chunk=16``, the largest pure-quadrature slab is +#: +#: 16 * 4 * 16 * LAPLACE_POINT_BLOCK * sizeof(float64) = 32 MiB. +#: +#: This is a bound on the EXPLICIT sample/time axes of a direct kernel call, +#: not on arbitrary enclosing transformations: an outer ``vmap`` (flowMC maps +#: its scalar AD target over chains) adds another batch axis that this function +#: cannot see when it chooses ``pblk``. +#: +#: Before this point axis was rolled, that last factor was ``S * npts``. The +#: historical pre-cap JAX acceptance call at ``S=4000, npts=1193`` therefore +#: asked XLA for one 36.41-GiB buffer. That number is the repository-recorded +#: allocation request for this ONE logical f64 value, not a measurement of +#: current production ILE peak memory. The sampler-side device cap reduces S +#: for current host-batched callers, but direct ``log_likelihood`` calls do not, +#: and a device-memory fraction does not bound the total live graph or its AD +#: residuals. The coefficient tables and the output still scale as O(S*npts); +#: this constant removes only that multiplicative quadrature slab for explicit +#: batches. +LAPLACE_POINT_BLOCK = 4096 + def _psi_lnI_amplitudes(c1, c2): """(b, d, t_amp) for the kernel and the block dispatcher: harmonic @@ -1475,75 +1651,59 @@ def _gh_psi_node_offsets(n_nodes): 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. +def coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w_grid, *, amp_sizing=None, m_max=None, + phi_chunk=16, dist_block=4, point_block=LAPLACE_POINT_BLOCK, + return_amp=False): + """Stream the psi-Laplace angle/distance reserve from coefficient tables. + + Same seam and normalization contract as + :func:`coefficient_table_distphipsimarg_exact`: it consumes the tables + :func:`angle_coefficient_tables` returns, produces ``(S, Ntime)`` + normalized over the two periodic angles, and does NOT integrate time. + Extracted from :func:`fused_log_likelihood_distphipsimarg_laplace`, which + is now a thin wrapper over it, so the two paths cannot drift. + + It exists because the four-axis policy's reserve consumes already-built + tables at refined time nodes, and the only table-level reserve was the + exact one. That made the reserve method un-selectable: the composite could + only ever fall back to exact angles, whatever the amplitude. + + Costs ~sqrt(A) rather than ~A -- the lattice is dense in phi only, the psi + axis being removed analytically -- and its error SHRINKS with amplitude, so + it is the reserve to use above the selector crossover. The caller owns that + choice; this function does not select. + + The adaptive distance quadrature is honoured exactly as in the fused + kernel, and carries the same restriction: the psi-marginal node placement + rests on the A0 == 0 / B1 == 0 identity, established for mode content up to + ``_GH_PSI_M_MAX``. That identity is a property of the DATA and cannot be + measured here, where the tables are tracers; the caller must have gated it + with :func:`gh_laplace_supported` on concrete tables. """ - # 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,)) + C_A = jnp.asarray(C_A, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + if C_A.ndim == 3: + C_A = C_A[:, :, None, :] + if C_A.ndim != 4: + raise ValueError("C_A must have shape (KP,KS[,S],Ntime)") + if C_B.ndim == 2: + C_B = jnp.broadcast_to( + C_B[:, :, None, None], + C_B.shape + (C_A.shape[2], C_A.shape[3])) + if C_B.ndim != 4 or C_B.shape[2:] != C_A.shape[2:]: + raise ValueError( + "C_B must be collapsed (KP,KS) or match C_A sample/time axes") + inferred_m_max = int(C_A.shape[0] - 1) + if m_max is None: + m_max = inferred_m_max + m_max = int(m_max) + if m_max != inferred_m_max: + raise ValueError("m_max does not match the C_A harmonic order") + S = int(C_A.shape[2]) + npts = int(C_A.shape[3]) 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 @@ -1554,18 +1714,20 @@ def fused_log_likelihood_distphipsimarg_laplace( # 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 " + "distance-GH-nodes is set (--distance-gh-nodes / " + "JAX_ILE_DISTMARG_GH) 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." + "exact, or pass --distance-gh-nodes 0 (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") + amp_call = _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) @@ -1577,6 +1739,9 @@ def fused_log_likelihood_distphipsimarg_laplace( kpB = jnp.arange(2 * m_max + 1, dtype=jnp.float64) G = x_grid.shape[0] blk = int(dist_block) + pblk = min(int(point_block), S * npts) + if pblk < 1: + raise ValueError("point_block must be at least 1") # 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 @@ -1635,81 +1800,111 @@ def _step(carry, x): # 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): + # Roll the combined independent (sample,time) point axis BEFORE adding + # distance and quadrature axes. The old body formed + # (quad_chunk, dist_block, phi_chunk, S, npts) at once; the sampler cap + # only hid that from some callers. Edge padding is safe because each + # padded result is discarded before the phi reduction. Repeating the + # edge (rather than zero-padding the coefficients) also keeps every + # branch finite, which matters to reverse-mode AD even for dead outputs. + npoint = S * npts + n_pblk = (npoint + pblk - 1) // pblk + pad_p = n_pblk * pblk - npoint + + def _pack_points(v): + v = v.reshape(c, npoint) + if pad_p: + v = jnp.concatenate( + [v, jnp.broadcast_to(v[:, -1:], (c, pad_p))], axis=1) + return jnp.swapaxes(v.reshape(c, n_pblk, pblk), 0, 1) + + fields = tuple(_pack_points(v) for v in (A0, A1, B0, B1, B2)) + + def _point_step(field_block): + A0p, A1p, B0p, B1p, B2p = field_block # (c,pblk) + + # 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 - 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 + xgb, lwgb = xw # (blk,) + xg = xgb[:, None, None] # (g,1,1) + lwg = lwgb[:, None, None] + av = xg * A0p[None] - 0.5 * jnp.square(xg) * B0p[None] + c1 = xg * A1p[None] - 0.5 * jnp.square(xg) * B1p[None] + c2 = -0.5 * jnp.square(xg) * B2p[None] + e = _laplace_psi_lnI_block(av, c1, c2) + lwg # (g,c,pblk) 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) + 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)); see the derivation + # above _gh_psi_node_offsets. + w_st = B0p * A1p - jnp.conj(A1p) * B2p + ph1 = jnp.conj(w_st) / jnp.maximum(jnp.abs(w_st), 1e-300) + sgn = jnp.where((A1p * ph1).real >= 0, 1.0, -1.0) + ph1 = ph1 * sgn # e^{i u*} + A_st = A0p + (A1p * ph1).real # A(u*) + B_st = B0p + (B1p * ph1).real + (B2p * ph1 * ph1).real + R_lo = B0p - jnp.abs(B1p) - jnp.abs(B2p) # <= 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], + x_min, x_max) + + xg = _node(zb) # (g,c,pblk) + # Composite-trapezoid weight, index-clamped at both ends: + # identical to core._distmarg_gh_logL's convention. + 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] # -inf on pad slots + av = xg * A0p[None] - 0.5 * jnp.square(xg) * B0p[None] + c1 = xg * A1p[None] - 0.5 * jnp.square(xg) * B1p[None] + c2 = -0.5 * jnp.square(xg) * B2p[None] + e = _laplace_psi_lnI_block(av, c1, c2) + lwg + return _lse_update(mx, sx, e, axis=0), None + + mx0 = jnp.full((c, pblk), -jnp.inf, dtype=jnp.float64) + sx0 = jnp.zeros((c, pblk), dtype=jnp.float64) + (mx, sx), _ = jax.lax.scan( + _gh_dist_step, (mx0, sx0), + (zg_blk, zpg_blk, zng_blk, zpad_blk)) + return (mx + jnp.where( + sx > 0, jnp.log(jnp.maximum(sx, 1e-300)), -jnp.inf) + + gh_C0) + + mx0 = jnp.full((c, pblk), -jnp.inf, dtype=jnp.float64) + sx0 = jnp.zeros((c, pblk), 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) + _dist_step, (mx0, sx0), (xg_blk, lwg_blk)) + return mx + jnp.where(sx > 0, + jnp.log(jnp.maximum(sx, 1e-300)), -jnp.inf) + + # Avoid wrapping the overwhelmingly common scalar/small-test case in a + # one-trip map: it buys no memory and adds another control-flow region + # for XLA/AD to compile. Production batches cross the bound and take + # the rolled path below. + if n_pblk == 1: + lnI_blk = _point_step(tuple(v[0] for v in fields))[None] + else: + lnI_blk = jax.lax.map(jax.checkpoint(_point_step), fields) + lnI = (jnp.swapaxes(lnI_blk, 0, 1).reshape(c, n_pblk * pblk) + [:, :npoint].reshape(c, S, npts) + lww[:, None, None]) # (c,S,npts) m_new, s_new = _lse_update(m, s, lnI, axis=0) return (m_new, s_new), None @@ -1718,9 +1913,88 @@ def _node(zz): 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) + return (lnL_t, amp_call) if return_amp else lnL_t + + +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, point_block=LAPLACE_POINT_BLOCK, + time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False, + return_amp=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 of the multiplicative quadrature slab is bounded by ``phi_chunk`` x + ``dist_block`` x ``point_block``, never by the full sample x time product or + by grid sizes. ``point_block`` rolls independent ``(sample, time)`` bins and + changes no quadrature rule or reduction order within a bin. + """ + # 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( + "distance-GH-nodes is set (--distance-gh-nodes / " + "JAX_ILE_DISTMARG_GH), 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 pass --distance-gh-nodes 0 " + "(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) + lnL_t, amp_call = coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w_grid, amp_sizing=amp_sizing, + m_max=meta["m_max"], phi_chunk=phi_chunk, dist_block=dist_block, + point_block=point_block, return_amp=True) + out = lnL_t if return_lnLt else _time_marginalize_terminal( + lnL_t, data, time_quadrature) + return (out, amp_call) if return_amp else out # Relative size at which A0 / B1 count as nonzero. The identity the psi-marginal @@ -1818,6 +2092,84 @@ def psi_harmonics_at_phi(C_A, C_B, phi, m_max): MB(4) + jnp.conj(MB(0))) # B2 +# 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). +# +# These lived in wrapper.py, reached from its one call site. They are here +# because there are now TWO consumers -- the angle scheme and the policy's +# reserve roster -- and a second copy of a probe direction is a second +# definition of what "the identity holds on this data" means. +_GH_PROBE_RA = (1.0,) +_GH_PROBE_DEC = (0.3,) +_GH_PROBE_INCL = (1.0,) + + +def _gh_laplace_precondition(m_max, feature): + """The two conditions that need no tables. ``None`` means both hold. + + Split out because there are two entry points and the cheap half must run + FIRST at both of them. ``gh_laplace_supported_for_data`` builds tables, and + building them for a banded dataset takes the banded accumulation route -- + a different code path, which fails on its own terms before the response + model is ever looked at. So the caller that starts from a dataset has to + be able to refuse before it builds anything. + """ + if int(m_max) > _GH_PSI_M_MAX: + # The tables are SIZED by m_max, so a mismatched m_max is a shape error + # rather than a measurement. + return 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 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,)) + return None + + +def gh_laplace_supported_for_data(data, interp=JAX_INTERP_DEFAULT): + """:func:`gh_laplace_supported` on tables this builds at a probe direction. + + THE ONLY WAY to ask the question of a dataset rather than of a table. The + predicate itself cannot be called under jit/grad -- the tables are tracers + there -- so every caller needs concrete tables, and every caller that builds + its own would be choosing its own probe direction. + + Returns ``(ok, info)`` exactly as :func:`gh_laplace_supported` does. Costs + one O(1) table build. Says nothing about whether the distance quadrature in + use NEEDS the identity: that is the caller's condition (it is needed by the + per-sample adaptive node placement, not by a static grid). + """ + m_max = _data_m_max(data) + feature = getattr(data, "feature", None) + # Cheap half first: see _gh_laplace_precondition. Building the probe + # tables for a banded dataset would take the banded route and raise on its + # own missing fields, so a refusal here must not depend on tables. + refused = _gh_laplace_precondition(m_max, feature) + if refused is not None: + return False, refused + C_A, C_B = angle_coefficient_tables( + data, + jnp.asarray(_GH_PROBE_RA, dtype=jnp.float64), + jnp.asarray(_GH_PROBE_DEC, dtype=jnp.float64), + jnp.asarray(_GH_PROBE_INCL, dtype=jnp.float64), + interp)[:2] + return gh_laplace_supported(C_A, C_B, m_max, feature=feature) + + def gh_laplace_supported(C_A, C_B, m_max, feature=None): """May 'laplace' use the per-sample adaptive distance quadrature on THIS data? @@ -1844,28 +2196,9 @@ def gh_laplace_supported(C_A, C_B, m_max, feature=None): # 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,)) + refused = _gh_laplace_precondition(m_max, feature) + if refused is not None: + return False, refused 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) @@ -1903,7 +2236,7 @@ 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): + phi_chunk=None, return_amp=False): """Distance-, phi_ref- AND psi-marginalized lnL: PEAK-LOCAL scheme. Same contract and normalization as @@ -1911,9 +2244,10 @@ def fused_log_likelihood_distphipsimarg_peaklocal( 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. + integrated on a window set by its own curvature. Windowed cells need only 48 nodes, + but rejected Newton centres span whole cells, so production sizes the shared static + count from ``amp_sizing``. The node axis is streamed in fixed-size blocks; cost grows + as sqrt(amplitude), while its live memory does not. THE PHI AXIS IS STILL DENSE HERE and is sized by :func:`~RIFT.likelihood.jax_ile.joint_anglemarg_peaklocal.required_n_phi` from the @@ -1929,10 +2263,12 @@ def fused_log_likelihood_distphipsimarg_peaklocal( """ 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.") + "distance-GH-nodes is set (--distance-gh-nodes / " + "JAX_ILE_DISTMARG_GH), 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 pass --distance-gh-nodes 0 (or unset " + "JAX_ILE_DISTMARG_GH).") _require_amp_sizing(amp_sizing) from . import joint_anglemarg_peaklocal as _jp @@ -1946,10 +2282,19 @@ def fused_log_likelihood_distphipsimarg_peaklocal( # 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") + amp_call = _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)} + # Size the u axis through the SINGLE SOURCE OF TRUTH rather than letting the kernel + # fall back to its own constant: the batch-memory guard in samplers.py models this + # same number from the same amp_sizing, and the two live in different files. Passing + # it explicitly is what makes them provably the same value rather than two defaults + # that happen to agree. The derived count is streamed inside the kernel, so raising + # accuracy does not materialize that entire axis across the outer batches. + kw = {"n_nodes": _jp.u_nodes_in_use(amp_sizing)} + if phi_chunk is not None: + kw["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). @@ -1960,9 +2305,255 @@ 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) + out = lnL_t if return_lnLt else _time_marginalize_terminal( + lnL_t, data, time_quadrature) + return (out, amp_call) if return_amp else out + + +def fused_log_likelihood_distphipsimarg_phi_local( + data, ra, dec, incl, x_grid, log_w_grid, + interp=JAX_INTERP_DEFAULT, amp_sizing=None, + time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False, + x_chunk=None, pt_chunk=None, n_slots=None, return_ok=False, + return_amp=False): + """Distance-, phi_ref- AND psi-marginalized lnL with BOTH ANGLE AXES LOCALIZED. + + Same contract and normalization as the other ``fused_log_likelihood_distphipsimarg_*`` + entries. What changes against ``peak-local`` is the phi axis: rather than a dense grid + sized ``~sqrt(A)``, phi is localized around the maxima of the u-profile and the omitted + mass is BOUNDED, so cost stops growing with amplitude. + + IT DECLINES, AND THE CALLER GETS THE DENSE ANSWER WHEN IT DOES. ``phi_local_lnI`` is + fail-closed: a row whose omitted-mass bound, convergence probes or u sizing do not pass + returns ``ok = False``, and across a distance grid ``ok`` is the CONJUNCTION over nodes. + This entry evaluates the dense peak-local scheme as well and selects elementwise, so a + decline costs time and never accuracy. Pass ``return_ok=True`` to get the mask and + account for how often the localized path actually carried the row -- a scheme that + silently fell back on every sample would otherwise look like it worked. + + THAT MAKES THIS SLOWER THAN ``peak-local`` UNTIL THE FALLBACK CAN BE SKIPPED, which + needs an acceptance rate measured on production tables rather than assumed. It is + therefore reachable only by name and is not in ``auto``. + + ``JAX_ILE_DISTMARG_GH`` is REFUSED for the same reason the dense peak-local branch + refuses it: no psi-marginal node placement exists yet. The seam it will attach to, + :func:`~RIFT.likelihood.jax_ile.joint_anglemarg_peaklocal.phi_local_lnI_at_distance`, + is public and takes one distance node, so that work does not have to modify this + function. + """ + if _core._DISTMARG_GH_N > 0: + raise ValueError( + "distance-GH-nodes is set (--distance-gh-nodes / " + "JAX_ILE_DISTMARG_GH), but the 'phi-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 pass --distance-gh-nodes 0 (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) + amp_call = _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, + "phi-local") + + n_phi = _jp.required_n_phi(amp_sizing, m_max=_data_m_max(data)) + u_nodes = _jp.u_nodes_in_use(amp_sizing) + kw = {"u_nodes": u_nodes, + "n_bound": int(_jp.required_bound_grid(amp_sizing)), + "pt_chunk": int(pt_chunk if pt_chunk is not None else _jp.PT_CHUNK_DEFAULT)} + if n_slots is not None: + kw["n_slots"] = int(n_slots) + + A = jnp.moveaxis(jnp.asarray(C_A), (2, 3), (0, 1)) + B = jnp.moveaxis(jnp.asarray(C_B), (2, 3), (0, 1)) + xc = int(x_chunk if x_chunk is not None else _jp.X_CHUNK_DEFAULT) + + def _one(a, b): + loc, ok, _ = _jp.joint_lnL_phi_local(a, b, x_grid, log_w_grid, + x_chunk=xc, **kw) + dense = _jp.joint_lnL_phi_dense(a, b, x_grid, log_w_grid, n_phi=n_phi, + n_nodes=u_nodes) + return jnp.where(ok, loc, dense), ok + + lnL_t, ok_t = jax.vmap(jax.vmap(_one))(A, B) # (S, npts) each + # return_amp appends the amplitude metric as the LAST element whatever the + # other flags select, so the three optional returns compose unambiguously. + out = lnL_t if return_lnLt else _time_marginalize_terminal( + lnL_t, data, time_quadrature) + if return_ok: + return (out, ok_t, amp_call) if return_amp else (out, ok_t) + return (out, amp_call) if return_amp else out + + +def _multipeak_sigma_t(rows_A, guard): + """Predicted time-peak width in native samples, from the tables alone. + + sigma_t = 1 / (2 pi rho sigma_f). sigma_f is the RMS frequency of the time + primitive's own spectrum, in cycles per sample, so no rate conversion enters; + rho comes from the exponent's peak. Both are properties of the data, which is + the point: the reserve's node placement is PREDICTED before any evaluation + rather than discovered by evaluating everywhere. + """ + from .time_first_peaklocal import _time_primitive_spectrum + flat = jnp.asarray(rows_A[0]).reshape((-1, rows_A.shape[-1])) + coeff, freq, _ = _time_primitive_spectrum(flat, int(guard)) + p = np.abs(np.asarray(coeff)) ** 2 + f = np.asarray(freq) + w = p.sum(axis=0) + sigma_f = float(np.sqrt((w * f * f).sum() / max(w.sum(), 1.0e-300))) + env = np.abs(np.asarray(rows_A)).sum(axis=(1, 2)) + rho = float(np.sqrt(2.0 * np.max(env))) + return 1.0 / max(2.0 * np.pi * rho * sigma_f, 1.0e-300) + + +def _multipeak_reserve_rule(CA_full, guard, data, sigma_t, n_sigma, pts_per_sigma): + """Peak-local time nodes and trapezoid weights, in production time units. + + Node count does not grow with rho: the window is +-n_sigma sigma_t and the + spacing is sigma_t / pts_per_sigma, so the two scale together. + """ + from .all_axis_peaklocal import (_time_primitive_spectrum, + _evaluate_time_spectrum) + n_nat = CA_full.shape[-1] - 2 * int(guard) + enum = np.arange(0.0, n_nat - 1 + 1.0e-9, 0.25) + flat = jnp.asarray(CA_full).reshape((-1, CA_full.shape[-1])) + coeff, freq, off = _time_primitive_spectrum(flat, int(guard)) + env = np.abs(np.asarray(_evaluate_time_spectrum( + coeff, freq, jnp.asarray(enum), off))).sum(axis=0) + keep = env >= env.max() * np.exp(-0.5 * 40.0 / max(env.max(), 1.0e-300)) + peaks = [j for j in range(1, len(env) - 1) + if env[j] >= env[j - 1] and env[j] >= env[j + 1] and keep[j]] + if not peaks: + peaks = [int(np.argmax(env))] + half = float(n_sigma) * float(sigma_t) + step = float(sigma_t) / float(pts_per_sigma) + wins = sorted((max(0.0, enum[j] - half), min(float(n_nat - 1), enum[j] + half)) + for j in peaks) + merged = [list(wins[0])] + for w0, w1 in wins[1:]: + if w0 <= merged[-1][1]: + merged[-1][1] = max(merged[-1][1], w1) + else: + merged.append([w0, w1]) + nd, wt = [], [] + for w0, w1 in merged: + k = max(2, int(np.ceil((w1 - w0) / step)) + 1) + g = np.linspace(w0, w1, k) + h = (w1 - w0) / (k - 1) + ww = np.full(k, h) + ww[0] = ww[-1] = 0.5 * h + nd.append(g) + wt.append(ww) + nodes = np.concatenate(nd) + # BUCKET THE SHAPE so a batch compiles once, not once per row. Padding is + # zero-weight, so the value is identical. + weights = np.concatenate(wt) + tgt = next((b for b in (256, 512, 1024, 2048) if b >= len(nodes)), len(nodes)) + pad = tgt - len(nodes) + if pad > 0: + # compute `pad` BEFORE reassigning nodes: deriving it from len(nodes) + # afterwards gives zero and leaves weights shorter than nodes. + nodes = np.concatenate([nodes, np.full(pad, nodes[-1])]) + weights = np.concatenate([weights, np.zeros(pad)]) + assert len(nodes) == len(weights) == tgt + w_t = np.asarray(data.w_t, dtype=float) + scale = float(np.sum(w_t)) / ((int(data.npts) - 1) * float(data.deltaT)) + return jnp.asarray(nodes), jnp.asarray(weights * float(data.deltaT) * scale) + + +def fused_log_likelihood_distphipsimarg_multipeak( + data, ra, dec, incl, x_grid, log_w_grid, + interp=JAX_INTERP_DEFAULT, amp_sizing=None, guard=16, + tier0=(2, 3, 24), tier1=(3, 5, 48), log_integral_tol=1.0e-3, + cell_sigma=5.0, quadrature_order=7, refine_iterations=18, + reserve_sigma=12.0, reserve_pts_per_sigma=8.0, return_record=False): + """Distance-, phi_ref-, psi- AND time-marginalized lnL: MULTIPEAK scheme. + + The four-axis controller of + :func:`~RIFT.likelihood.jax_ile.multipeak_planner.multipeak_local_marginalize`, + reachable from ``--angle-marg-scheme multipeak``. Unlike every other entry + in this family it OWNS THE TIME INTEGRAL, so there is no ``lnL(t)`` and no + ``time_quadrature``: the caller gets one value per sample. Callers that + need ``lnL(t)`` must use another scheme. + + The default operating point is the one measured on the ladder-2 injection at + rho 40.77, 163.08 and 652.31 (64 rows per rung, inclination banded +-0.20 rad + about the injection): 64/64 accepted at every rung, 4.74-5.17 s per row and + ~218 MiB peak device memory, error against a peak-local reference of 2.2e-05 + nats median at rho 40.77 and 5.6e-04 at 163.08. See + analyses/va_sequence_20260902/RESULTS_20260909_multipeak_ladder.md in the + RIFT_roboto_paper record store. ``guard`` defaults to the value that + measurement used; the driver's production default is larger and the caller + passes it explicitly. + + The reserve is a PEAK-LOCAL time rule, not a refined whole window: nodes are + placed only in +-``reserve_sigma`` sigma_t windows about the peaks of the + coefficient envelope, at spacing sigma_t/``reserve_pts_per_sigma``. Window + and spacing both scale as sigma_t, so the node count does not grow with rho. + A refined whole-window reserve was measured at 4905 nodes for a peak 0.06 + native samples wide at rho 163, and its own half-refined warrant failed at + rho 40.77; this rule reproduced it to 0.0 on every row with 193 nodes. + """ + from . import multipeak_planner as _mp + from . import all_axis_peaklocal as _aap + from . import direct_marginalization_policy as _pol + from .core import _time_marginalize + + _require_amp_sizing(amp_sizing) + guard = int(guard) + if guard < 2: + raise ValueError("multipeak needs guard >= 2 for the time primitive") + 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, + guard=guard) + _runtime_amp_failsafe(C_A, C_B, x_grid, amp_sizing, "multipeak") + log_measure, _ = _pol.policy_log_normalization( + data, np.asarray(x_grid), np.asarray(log_w_grid)) + x_min = float(np.min(np.asarray(x_grid))) + x_max = float(np.max(np.asarray(x_grid))) + m_max = int(meta["m_max"]) + + rows_A = np.moveaxis(np.asarray(C_A), 2, 0) # (S,KP,KS,ntime+2g) + rows_B = np.moveaxis(np.asarray(C_B), 2, 0) + sigma_t = _multipeak_sigma_t(rows_A, guard) + + out, records = [], [] + for i in range(rows_A.shape[0]): + CA_full = rows_A[i] + CA_i = CA_full[..., guard:-guard] + CB_i = rows_B[i][..., 0][..., None] * np.ones(CA_i.shape[-1]) + nodes, weights = _multipeak_reserve_rule( + CA_full, guard, data, sigma_t, + float(reserve_sigma), float(reserve_pts_per_sigma)) + + CB_flat = rows_B[i][..., 0] # (KP,KS), time-independent + + def _reserve(CA_full=CA_full, CB_flat=CB_flat, nodes=nodes, + weights=weights): + flat = jnp.asarray(CA_full).reshape((-1, CA_full.shape[-1])) + coeff, freq, off = _aap._time_primitive_spectrum(flat, guard) + tgt = _aap._evaluate_time_spectrum( + coeff, freq, nodes, off).reshape( + CA_full.shape[:-1] + (nodes.size,)) + lnLt = coefficient_table_distphipsimarg_laplace( + tgt, jnp.asarray(CB_flat), x_grid, log_w_grid, + amp_sizing=amp_sizing, m_max=m_max) + return float(np.asarray(_time_marginalize(lnLt, weights)[0])) + + res = _mp.multipeak_local_marginalize( + CA_i, CB_i, x_min, x_max, _reserve, + log_integral_tol=float(log_integral_tol), + tier0=tuple(int(v) for v in tier0), + tier1=tuple(int(v) for v in tier1), + refine_iterations=int(refine_iterations), + cell_sigma=float(cell_sigma), + quadrature_order=int(quadrature_order), + log_measure=float(log_measure), label="multipeak_row%d" % i) + out.append(float(res.value)) + records.append(res) + values = jnp.asarray(np.asarray(out, dtype=float)) + return (values, records) if return_record else values def choose_angle_marg_scheme(amplitude, gh_enabled=None, diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py index 6bcc67d3e..2510dc1f7 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py @@ -24,6 +24,7 @@ from .core import build_likelihood_data, DIST_MPC_REF from . import response_slowrot as _rs from . import response_freqresponse as _rf +from . import response_rotating_freqresponse as _rrf def _stack_bank(by_key, keys, det): @@ -192,3 +193,59 @@ def build_freqresponse_data(meta, lookupNKDict, rho_by_p, U_by_pp, V_by_pp, refl_idx=np.asarray(_rf.reflection_index(p_list), dtype=np.int64), ) return data + + +def build_rotating_freqresponse_data(meta, lookupNKDict, rho_by_a, U_by_aa, + V_by_aa, epochDict, deltaT, tvals, det_geom, + distMpcRef=DIST_MPC_REF): + """Banded data for simultaneous slow rotation and finite-arm response.""" + if not bool(meta.get("post_phase_required", False)): + raise ValueError("compound response bank must require the arrival-time post-phase") + a_list = [tuple(int(v) for v in a) for a in meta["a_list"]] + tref = float(meta["event_time_geo"]) + detectors = list(rho_by_a.keys()) + a0 = a_list[0] + + def _pair(bank, det, i, j, a, ap): + return bank[det][(a, ap)] if isinstance(bank[det], dict) else bank[det][i, j] + + packed_scalar = {} + for det in detectors: + packed_scalar[det] = dict( + lms=np.asarray(lookupNKDict[det]), + rholmArray=np.asarray(rho_by_a[det][a0], dtype=np.complex128), + U=np.asarray(_pair(U_by_aa, det, 0, 0, a0, a0), dtype=np.complex128), + V=np.asarray(_pair(V_by_aa, det, 0, 0, a0, a0), dtype=np.complex128), + epoch=float(epochDict[det])) + data = _base_data(packed_scalar, deltaT, tref, tvals, distMpcRef) + + A = len(a_list) + for det in detectors: + dd = data.detectors[det] + Q_bank = _stack_bank(rho_by_a, a_list, det) + dd["Q_bank"] = jnp.asarray(np.ascontiguousarray(np.transpose(Q_bank, (0, 2, 1)))) + K = len(dd["lms"]) + U = np.empty((A, A, K, K), dtype=np.complex128) + V = np.empty((A, A, K, K), dtype=np.complex128) + for i, a in enumerate(a_list): + for j, ap in enumerate(a_list): + U[i, j] = np.asarray(_pair(U_by_aa, det, i, j, a, ap)) + V[i, j] = np.asarray(_pair(V_by_aa, det, i, j, a, ap)) + dd["U_bank"] = jnp.asarray(U) + dd["V_bank"] = jnp.asarray(V) + _resp, x_arm, y_arm, length = det_geom[det] + dd["x_arm"] = jnp.asarray(np.asarray(x_arm, dtype=np.float64)) + dd["y_arm"] = jnp.asarray(np.asarray(y_arm, dtype=np.float64)) + dd["L_arm"] = float(length) + + m_values, term1_idx, term2_idx = _rs.post_phase_bucketing(a_list) + data.feature = "rotation_freqresponse" + data.band = dict( + a_list=a_list, Qmax=int(meta["Qmax"]), p_max=int(meta["p_max"]), + refl_idx=np.asarray(_rrf.reflection_index(a_list), dtype=np.int64), + f_sidereal=float(meta["f_sidereal"]), post_phase_required=True, + pp_m_values=np.asarray(m_values, dtype=np.int64), + pp_term1_idx=np.asarray(term1_idx, dtype=np.int64), + pp_term2_idx=np.asarray(term2_idx, dtype=np.int64), + ) + return data diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py index 59cd4162f..0371c5f4a 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py @@ -82,10 +82,31 @@ # Nodes centred per-sample on x* with scale 1/sqrt(R) (trapezoid, gradient-stable # placement via stop_gradient) integrate it to machine precision at any SNR with # a few dozen nodes -- removing the evidence bias. Enable with env -# JAX_ILE_DISTMARG_GH= (e.g. 64); 0 keeps the legacy uniform grid. +# JAX_ILE_DISTMARG_GH= (e.g. 64), or the driver's --distance-gh-nodes, +# which calls set_distmarg_gh_nodes() below; 0 keeps the legacy uniform grid. # See make_distance_gh / _distmarg_gh_logL. _DISTMARG_GH_N = int(os.environ.get("JAX_ILE_DISTMARG_GH", "0")) + +def set_distmarg_gh_nodes(n): + """Set the per-sample Gauss-Hermite distance-quadrature node count. + + Every reader of ``_DISTMARG_GH_N`` in this module (a bare global lookup) + and in the sibling modules that hold a reference to this one (``_core. + _DISTMARG_GH_N``, an attribute lookup) resolves it dynamically at CALL + time, not at import time -- so calling this after those modules have + already been imported is sufficient; nothing needs to be re-imported. + This is what makes the driver's ``--distance-gh-nodes`` option reachable + without import-order fragility. ``n=0`` restores the legacy uniform grid. + """ + global _DISTMARG_GH_N + _DISTMARG_GH_N = int(n) + + +def get_distmarg_gh_nodes(): + """Return the currently active per-sample Gauss-Hermite node count.""" + return _DISTMARG_GH_N + import lal import lalsimulation as lalsim @@ -93,6 +114,7 @@ from .spherical import spherical_harmonics_vectorized from . import response_slowrot as _rs from . import response_freqresponse as _rf +from . import response_rotating_freqresponse as _rrf # Fiducial template distance (Mpc); identical to factored_likelihood.distMpcRef. DIST_MPC_REF = 1000.0 @@ -121,10 +143,21 @@ class JAXLikelihoodData: """ def __init__(self, detectors, deltaT, gmst, tvals, tref, - distMpcRef=DIST_MPC_REF): + distMpcRef=DIST_MPC_REF, q_time_pregrid_factor=1): self.detector_names = list(detectors.keys()) self.detectors = detectors # name -> dict (see build_likelihood_data) self.deltaT = float(deltaT) + # Integer refinement of the stored Q sampling. 1 == the historical + # behaviour, Q sampled at deltaT. With factor f the stored Q arrays are + # sampled at deltaT/f, so a position expressed in COARSE samples must be + # multiplied by f before it indexes them. deltaT itself, tvals, and the + # Simpson weights below are DELIBERATELY unchanged: the pregrid refines + # the interpolation of Q, not the cadence the likelihood integrates on + # (mirroring ``--q-time-pregrid-factor`` on the conventional arm). + self.q_time_pregrid_factor = int(q_time_pregrid_factor) + if self.q_time_pregrid_factor < 1: + raise ValueError("q_time_pregrid_factor must be >= 1, got %r" + % (q_time_pregrid_factor,)) self.gmst = float(gmst) self._tref = float(tref) self.tvals = jnp.asarray(tvals, dtype=jnp.float64) @@ -142,8 +175,71 @@ def lms(self): return self.detectors[self.detector_names[0]]["lms"] +def build_q_time_pregrid(rho, factor): + """Refine a packed ``(K, npts_full)`` rholm block onto a ``factor``x time grid. + + THIS IS A THIN WRAPPER, ON PURPOSE. The arithmetic is + ``factored_likelihood.build_reflected_q_pregrid`` -- the SAME host-side + builder the conventional NoLoop arm uses for ``--q-time-pregrid-factor`` + (RIFT PR #261). Calling it rather than restating it here is the whole point: + the two arms are meant to be answering with the same refined Q, and a second + implementation of a boundary convention is exactly how the two drivers came + to ship opposite stencil defaults (issue #233). It also inherits that + function's round-trip guard -- every ``factor``-th refined sample must + reproduce the input to 5e-12 relative -- and its host-side (numpy) execution, + so the transient ``2 n factor`` reflection never lands on the accelerator. + + WHICH REFLECTION, and why it is not this module's ``_reflected_fft_upsample``. + The two differ, and the difference is measured, not stylistic: + + * ``_reflected_fft_upsample`` periodizes ``[x0..x_{n-1}, x_{n-2}..x_1]`` + (period ``2(n-1)``). It is right for what IT is used for -- reconstructing + ``kappa`` on the *terminal* integration window, where a series sitting at + exactly Nyquist must keep reconstructing ``cos(pi t)``; duplicating the + turning samples inserts a flat pair and breaks that. + * ``reflected_bandlimited_upsample`` (what this uses, via #261) periodizes + ``[x0..x_{n-1}, x_{n-1}..x0]`` (period ``2n``). That is the right choice + HERE, for a different reason: the rholm buffer is a CROP of a longer series + (``ComputeModeIPTimeSeries`` ends in ``CutCOMPLEX16TimeSeries(rhoTS, 0, + N_window)``), and on crop-shaped fixtures the ``2n`` form measured 2e-8 to + 2.5e-6 nats against 2e-6 to 2.5e-4 for ``2(n-1)`` -- see + DESIGN_time_marginalization_quadrature.md, "Finite-window reconstruction". + + Both docstrings assert their own convention is the correct one; they are + describing different problems and both are right about theirs. Neither is a + substitute for the exact-period oracle -- see + ``test/jax/test_jax_q_time_pregrid.py``, which measures both against a Q + built as a genuine crop of an exactly periodic band-limited series. MEASURED + THERE, on this arm's own fixture: routed through ``_reflected_fft_upsample`` + the factor-8 pregrid saturates at 7.8e-4 relative and does not improve at + factor 16 (7.6e-4); through the ``2n`` form it reaches 4.6e-5 and keeps + converging. Getting this wrong costs a 17x floor and all of the convergence, + while every "every factor-th sample reproduces the input" check still passes. + + The ``factor == 1`` short-circuit returns the INPUT OBJECT and does not import + ``factored_likelihood`` at all, so the default path is untouched -- bit-identity + is a property of the code path, not of an agreement to 1e-15. + + Returns ``(rho_fine, report)`` with ``rho_fine`` shaped + ``(K, (npts_full-1)*factor + 1)`` and ``report`` the #261 telemetry dict. + """ + factor = int(factor) + if factor < 1: + raise ValueError("q_time_pregrid_factor must be >= 1, got %r" % (factor,)) + if factor == 1: + return rho, dict(factor=1) + # LOCAL import, deliberately. factored_likelihood pulls in lalsimutils and numba; + # this module is imported by lightweight consumers (the stencil-parity tests, the + # coordinate helpers) that never build data, and a module-level import would make + # them pay for it. Data building already imports it via wrapper.py anyway. + from RIFT.likelihood import factored_likelihood as _fl + dense, report = _fl.build_reflected_q_pregrid(np.asarray(rho), factor=factor, + xpy=np) + return np.asarray(dense), report + + def build_likelihood_data(packed_per_detector, deltaT, tref, tvals, - distMpcRef=DIST_MPC_REF): + distMpcRef=DIST_MPC_REF, q_time_pregrid_factor=1): """Assemble a :class:`JAXLikelihoodData` from packed numpy arrays. Parameters @@ -170,15 +266,24 @@ def build_likelihood_data(packed_per_detector, deltaT, tref, tvals, helper ``bin/integrate_likelihood_extrinsic_batchmode`` uses (issue #146). """ gmst = float(lal.GreenwichMeanSiderealTime(tref)) + q_time_pregrid_factor = int(q_time_pregrid_factor) detectors = {} for det, d in packed_per_detector.items(): lms = [(int(l), int(m)) for (l, m) in np.asarray(d["lms"])] rho = np.asarray(d["rholmArray"], dtype=np.complex128) # (K, npts_full) + npts_full_coarse = int(rho.shape[-1]) + # factor 1 returns ``rho`` itself, so the historical path is not merely + # numerically equal, it is the SAME array object -- see + # test_factor_one_is_bit_identical. + rho, q_report = build_q_time_pregrid(rho, q_time_pregrid_factor) Q = jnp.asarray(np.ascontiguousarray(rho.T)) # (npts_full, K) D = lalsim.DetectorPrefixToLALDetector(det) detectors[det] = { "lms": lms, "Q": Q, + "q_time_pregrid_factor": q_time_pregrid_factor, + "q_time_pregrid_report": q_report, + "npts_full_coarse": npts_full_coarse, "U": jnp.asarray(np.asarray(d["U"], dtype=np.complex128)), "V": jnp.asarray(np.asarray(d["V"], dtype=np.complex128)), "epoch": float(d["epoch"]), @@ -187,7 +292,8 @@ def build_likelihood_data(packed_per_detector, deltaT, tref, tvals, "npts_full": int(Q.shape[0]), "l_max": max(l for (l, m) in lms), } - return JAXLikelihoodData(detectors, deltaT, gmst, tvals, tref, distMpcRef) + return JAXLikelihoodData(detectors, deltaT, gmst, tvals, tref, distMpcRef, + q_time_pregrid_factor=q_time_pregrid_factor) def _gather_nearest(Q_col, pos, u=None): @@ -374,6 +480,128 @@ def _separable_u(p0): return (p0 - jnp.floor(p0))[:, None] +# Stencil footprint in STORED samples, i.e. how far a gather reaches either side of +# its base index. One definition, consumed by both accumulators' support checks. +_STENCIL_MARGIN = {"nearest": 1, "linear": 2, "cubic": 3, + "sinc": SINC_HALFWIDTH_DEFAULT + 1} + + +def _check_stored_q_length(dd, stored_npts, factor, what): + """Fail closed when a stored Q bank was not refined to its declared factor. + + ``_q_sample_positions`` scales every index by ``factor``. If the array it + indexes was never refined, that scaling silently reads the wrong samples -- + a factor-8 window would cover an eighth of the intended span and land on + whatever happens to be there. Nothing downstream can see it: the shapes + still broadcast, the likelihood still returns finite numbers, and they are + wrong. + + The shape it guards against is reachable. ``banded._base_data`` builds the + scaffold through :func:`build_likelihood_data`, then attaches an + INDEPENDENTLY packed ``Q_bank`` that :func:`build_q_time_pregrid` never sees, + and ``_accumulate_unit_banded`` indexes that bank. ``_base_data`` takes no + factor today, so the two cannot disagree yet; giving it one without also + refining the bank would produce exactly this. + + Cheap (a python int comparison at trace time), so there is no reason to make + it conditional. + """ + declared = dd.get("q_time_pregrid_factor") + if declared is not None and int(declared) != int(factor): + raise ValueError( + "%s was built at q_time_pregrid_factor=%d but is being indexed at " + "factor %d" % (what, int(declared), int(factor))) + coarse = dd.get("npts_full_coarse") + if coarse is None: + # A hand-built detector dict (tests, benchmark shims) carrying no + # refinement metadata. At factor 1 there is nothing to check: no index + # is scaled, so an unrefined buffer is the correct buffer. + # + # Above factor 1 the absence of the metadata is itself the fault, and + # returning here was a hole. `build_q_time_pregrid` sets both keys + # together, so a dict that declares a factor without `npts_full_coarse` + # was not built by it, and its Q is coarse. `_q_sample_positions` would + # still scale every index by the factor, reading an eighth of the + # intended span at factor 8. Shapes broadcast, the likelihood returns + # finite numbers, and they are wrong. Refuse instead. + if int(factor) != 1: + raise ValueError( + "%s is indexed at q_time_pregrid_factor=%d but carries no " + "'npts_full_coarse'; refinement metadata is required above " + "factor 1, because the stored Q cannot be shown to have been " + "refined and every index would be scaled regardless. Build it " + "with build_q_time_pregrid, or index at factor 1." + % (what, int(factor))) + return + expected = (int(coarse) - 1)*int(factor) + 1 if int(factor) != 1 else int(coarse) + if int(stored_npts) != expected: + raise ValueError( + "%s has %d samples but q_time_pregrid_factor=%d over a %d-sample " + "coarse buffer requires %d; the stored Q was not refined to the " + "declared factor" % (what, int(stored_npts), int(factor), + int(coarse), expected)) + + +def _q_sample_positions(data, p0, t_offsets, interp): + """Map a coarse-sample window onto the stored (possibly pre-refined) Q grid. + + ``p0`` (shape ``(S,)``) and ``t_offsets`` (shape ``(npts,)``) are in units of + ``data.deltaT``, the cadence the likelihood integrates on. The STORED Q may be + sampled ``f = data.q_time_pregrid_factor`` times finer, so an index into it is + ``f`` times larger. Returns ``(pos, u_sep)`` in stored-sample units. + + ``f == 1`` returns what the accumulators computed inline before the pregrid + existed, the same expressions in the same order, so that path is bit-identical. + ``test_factor_one_positions_are_bit_identical_to_the_pre_pregrid_expressions`` + pins these positions bitwise against those expressions; the whole-likelihood + identity against base ``bec19ad5`` is in the PR, over 52 toy arrays and 35 + from a rebuilt production likelihood. + + SEPARABILITY IS THE PRECONDITION. ``_separable_u`` computes ONE fractional + offset per sample and hands it to the gatherer for every time column; that is + only legitimate while the time offsets are exact integers in the units the + gather indexes, which ``t_offsets * f`` (integer ``t_offsets``, integer ``f``) + keeps them. The form written here is additive to match the factor-1 branch + line for line. That is a readability choice and carries no accuracy claim: + ``(p0 + t) * f`` gives bit-identical ``frac`` and ``floor`` at f = 8 for ``p0`` + from 5e2 to 5e5, 2000 samples and 742 columns per decade (re-measured + 2026-09-07). + + ``nearest`` is REFUSED with a pregrid rather than quietly allowed. It would + gather correctly (snapping to a finer sample is strictly better), but + :func:`_accumulate_unit_banded` reconstructs the arrival time its post-phase + applies as ``rint(p0)`` in COARSE samples, which is no longer the sample the + gather read; the data term and the model norm would drift apart by up to half a + coarse bin. Refusing costs nothing -- a pregrid exists to buy sub-sample + accuracy, which is precisely what 'nearest' declines to use. + """ + # getattr, not attribute access: benchmark shims and several existing tests + # duck-type ``data`` as a SimpleNamespace. The default is the historical + # behaviour, and it is SAFE rather than merely convenient -- the paired + # ``_check_stored_q_length`` refuses a detector dict whose declared factor + # disagrees with the one being indexed, so a refined Q reaching a namespace + # that forgot the attribute raises instead of being read at the wrong stride. + factor = int(getattr(data, "q_time_pregrid_factor", 1)) + if factor == 1: + pos = p0[:, None] + t_offsets[None, :] + # 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) + return pos, u_sep + if interp == "nearest": + raise NotImplementedError( + "interp='nearest' is not supported with q_time_pregrid_factor=%d: the " + "banded post-phase reconstructs the gathered arrival time in COARSE " + "samples, so it would no longer match the sample a refined-grid nearest " + "gather reads. Use interp='cubic' (the validated pregrid stencil)." + % (factor,)) + p0_q = p0 * factor + pos = p0_q[:, None] + (t_offsets * factor)[None, :] + return pos, _separable_u(p0_q) + + _GATHERERS = {"nearest": _gather_nearest, "linear": _gather_linear, "cubic": _gather_cubic, "sinc": _make_gather_sinc(SINC_HALFWIDTH_DEFAULT)} @@ -416,6 +644,61 @@ def _guarded_window(data, guard): jnp.arange(-guard, data.npts + guard, dtype=jnp.float64)) +# The mode order :func:`_accumulate_unit`'s phase-marginalized branch is written +# against. That branch is position-dependent -- it conjugates the m=-2 column of +# Y and Q and pairs it with conj(F) -- but the packed column order is NOT the +# caller's to choose: it comes from a python dict's iteration order in the +# precompute upstream, so a correctly-configured run can arrive with the pair the +# other way round. Permute to canonical rather than refuse. +_PHASE_MARG_MODES = ((2, 2), (2, -2)) + + +def _phase_marg_permutation(lms): + """Index permutation taking ``lms`` to ``[(2,2), (2,-2)]``, or ``None``. + + ``None`` means ``lms`` is ALREADY canonical. The caller must then skip the + permutation entirely rather than apply an identity one, because the ordering + that works today has to keep producing bit-for-bit the same numbers, and a + ``take`` with an identity index vector is not guaranteed to leave the XLA + graph -- and so the rounding -- untouched. + + Only the ORDER is free. Any other mode SET still raises: the conjugation + the accumulator applies is specific to one m=+2 / m=-2 pair, so a third mode + is a real gap in the method, not a relabelling. + + The set guard is what makes the DIRECTION of the permutation below safe. On + two modes the only non-canonical order is a transposition, which is its own + inverse, so ``order.index(...)`` and its inverse are the same map and no test + can tell them apart -- verified by enumerating the accepted inputs. They + diverge at K >= 3. So if this is ever widened past the pair, the widening + must come with a test that pins the direction; today's suite cannot. + """ + order = [(int(l), int(m)) for (l, m) in lms] + if sorted(order) != sorted(_PHASE_MARG_MODES): + raise NotImplementedError( + "phase marginalization currently requires modes " + "[(2,2),(2,-2)] (either order); got %r" % (order,)) + if order == list(_PHASE_MARG_MODES): + return None + return [order.index(lm) for lm in _PHASE_MARG_MODES] + + +def _permute_modes(lms, Q, U, V, perm): + """Reorder one detector's packed mode axis so column k becomes old ``perm[k]``. + + ``Q`` is (npts_full, K) -- mode on axis 1. ``U`` and ``V`` are (K, K) and + carry the mode index on BOTH axes: they are contracted as + ``sum_ij Ybar_i Y_j U_ij``, so permuting only one axis pairs each mode's + coefficient with the other mode's harmonic and returns a wrong likelihood + with no error. Both axes, or neither. + """ + p = np.asarray(perm, dtype=np.intp) + return ([lms[i] for i in perm], + jnp.take(Q, p, axis=1), + jnp.take(jnp.take(U, p, axis=0), p, axis=1), + jnp.take(jnp.take(V, p, axis=0), p, axis=1)) + + def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, phase_marginalization, guard=0): """Network kappa and rho^2 at the *fiducial* distance (invDist == 1). @@ -464,14 +747,18 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, V = dd["V"] K = len(lms) + if phase_marginalization: + # Canonicalize BEFORE Y is built, so the whole branch below stays the + # literal code it was: when the order is already canonical nothing is + # touched at all, and the working path is unchanged by construction. + perm = _phase_marg_permutation(lms) + if perm is not None: + lms, Q, U, V = _permute_modes(lms, Q, U, V, perm) + F = compute_detamresponse(dd["response"], ra, dec, psi, gmst) Y = spherical_harmonics_vectorized(lms, incl, -phiref, l_max=dd["l_max"]) if phase_marginalization: - if [tuple(x) for x in lms] != [(2, 2), (2, -2)]: - raise NotImplementedError( - "phase marginalization currently requires modes " - "[(2,2),(2,-2)]; got %r" % (lms,)) Y = Y.at[:, 1].set(jnp.conj(Y[:, 1])) F_lm = jnp.stack([F, jnp.conj(F)], axis=-1) Q = jnp.concatenate([Q[:, 0:1], jnp.conj(Q[:, 1:2])], axis=1) @@ -486,19 +773,16 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, FY_conj = jnp.conj(F_lm * Y) t_det = (data.tref_minus_epoch(det) + time_delay_from_earth_center(dd["location"], ra, dec, gmst)) + _check_stored_q_length(dd, Q.shape[0], + getattr(data, "q_time_pregrid_factor", 1), + "detector %s Q" % det) p0 = (t_det + data.tval0) * inv_deltaT - pos = p0[:, None] + t_offsets[None, :] + pos, u_sep = _q_sample_positions(data, p0, t_offsets, interp) if guard: - stencil_margin = {"nearest": 1, "linear": 2, "cubic": 3, - "sinc": SINC_HALFWIDTH_DEFAULT + 1}[interp] + stencil_margin = _STENCIL_MARGIN[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): @@ -524,17 +808,18 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, 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 + Banks carrying slow rotation have that dependence: their 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. + ``rho_sq`` there; the rotation+frequency-response bank shares this behavior. + The baseline accumulator (static ``F``) and the ``freqresponse``-only 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" + return getattr(data, "feature", None) in ("rotation", "rotation_freqresponse") def _banded_coefficients(data, det, ra, dec, psi): @@ -555,6 +840,10 @@ def _banded_coefficients(data, det, ra, dec, psi): return _rf.response_coefficients_packed( dd["response"], dd["x_arm"], dd["y_arm"], ra, dec, psi, data.gmst, b["Qmax"], b["p_list"]) + if data.feature == "rotation_freqresponse": + return _rrf.coefficients_packed( + dd["response"], dd["location"], dd["x_arm"], dd["y_arm"], + ra, dec, psi, data.gmst, b["Qmax"], b["p_max"], b["a_list"]) raise ValueError("unknown banded feature %r" % (data.feature,)) @@ -580,7 +869,7 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, ``term2``). ``aR`` is the V-term reflection (``(p,-n)`` for rotation, the identity for finite-size), supplied as ``data.band['refl_idx']``. - ARRIVAL-TIME POST-PHASE (``feature == "rotation"`` only). + ARRIVAL-TIME POST-PHASE (features carrying slow rotation). The bank's elementary templates ``chi_a(u) = e^{i n_a Omega u} h^{(p_a)}(u)`` live on the template's INTRINSIC time ``u``, while the physical response modulation lives on absolute time. Placing the template at arrival time ``t`` (``t' = u + t``) factorizes @@ -611,7 +900,7 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, == -1`` -- one bin off the FRONT of the rholm buffer -- where ``_gather_nearest``'s ``trunc(. + 0.5)`` index rounds to sample 0; see the note at the ``samp0`` assignment. - ``freqresponse`` (Path D) has NO post-phase -- its basis is not a sidereal modulation + ``freqresponse`` alone (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`. @@ -682,18 +971,18 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, t_det = (data.tref_minus_epoch(det) + 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) + # ``pos`` indexes the STORED Q (refined by q_time_pregrid_factor); ``p0`` + # stays in coarse samples, because the post_phase block below converts it to + # a physical arrival time via data.deltaT. + _check_stored_q_length(dd, Q_bank.shape[1], + getattr(data, "q_time_pregrid_factor", 1), + "detector %s Q_bank" % det) + pos, u_sep = _q_sample_positions(data, p0, t_offsets, interp) # (S, npts) if guard: - stencil_margin = {"nearest": 1, "linear": 2, "cubic": 3, - "sinc": SINC_HALFWIDTH_DEFAULT + 1}[interp] + stencil_margin = _STENCIL_MARGIN[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. @@ -736,21 +1025,42 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, kappa_unit = kappa_unit + kappa_det # --- term2: 0.5 Re[ sum_{a,a'} conj(C~_a)C~_a' YbarUY + C~_aR C~_a' YVY ] --- - # YUY[a,a'] = einsum(conjY, Y, U_bank[a,a']); YVY[a,a'] = einsum(Y, Y, V) - YUY = jnp.einsum("si,sj,abij->abs", conjY, Y, U_bank) # (A,A,S) - YVY = jnp.einsum("si,sj,abij->abs", Y, Y, V_bank) # (A,A,S) - # conj(C_a) C_a' and C_aR C_a' contracted over (a,a') -- the post-phase is - # applied below, since it depends only on m = n_a' - n_a for both contractions. - CC_U = jnp.einsum("as,bs->abs", jnp.conj(C), C) # (A,A,S) - CC_V = jnp.einsum("as,bs->abs", C_refl, C) # (A,A,S) - pair = CC_U * YUY + CC_V * YVY # (A,A,S) complex - if post_phase: + if post_phase and data.feature == "rotation_freqresponse": + # A compound bank can have A=50 at pmax=0 and A=112 at pmax=1. + # Materializing four (A,A,S) arrays makes the temporary footprint the + # practical limit. Contract each sidereal-difference bucket directly, + # matching the conventional dense-bank path; peak scratch is then + # max_m(number of pairs in m)*S instead of A*A*S. + bucket_values = [] + pp_t2_host = np.asarray(band["pp_term2_idx"], dtype=np.int64) + for im in range(M): + ia, iap = np.nonzero(pp_t2_host == im) + ia_d = jnp.asarray(ia, dtype=jnp.int32) + iap_d = jnp.asarray(iap, dtype=jnp.int32) + yuy = jnp.einsum("si,sj,pij->ps", conjY, Y, + U_bank[ia_d, iap_d]) + yvy = jnp.einsum("si,sj,pij->ps", Y, Y, + V_bank[ia_d, iap_d]) + bucket_values.append(jnp.sum( + jnp.conj(C[ia_d]) * C[iap_d] * yuy + + C_refl[ia_d] * C[iap_d] * yvy, axis=0)) + val_m = jnp.stack(bucket_values, axis=0) + rho_sq_det = 0.5 * jnp.einsum("ms,mt->st", val_m * pe, pt).real + else: + # The smaller individual-feature banks retain the original fused + # contraction, which minimizes compilation overhead for their usual A. + YUY = jnp.einsum("si,sj,abij->abs", conjY, Y, U_bank) + YVY = jnp.einsum("si,sj,abij->abs", Y, Y, V_bank) + CC_U = jnp.einsum("as,bs->abs", jnp.conj(C), C) + CC_V = jnp.einsum("as,bs->abs", C_refl, C) + pair = CC_U * YUY + CC_V * YVY + if post_phase and data.feature != "rotation_freqresponse": # BOTH contractions carry exp(i (n_a' - n_a) omega delta), so bucket the pairs # by m and pay one rank-1 phase per distinct m (M of them) instead of A^2. val_m = jnp.zeros((M, S), dtype=jnp.complex128).at[pp_t2].add(pair) # rho_sq becomes arrival-time dependent: (S, npts), not a broadcast scalar. rho_sq_det = 0.5 * jnp.einsum("ms,mt->st", val_m * pe, pt).real - else: + elif not post_phase: term2_c = jnp.sum(pair, axis=(0, 1)) # (S,) complex rho_sq_det = 0.5 * term2_c.real # (S,) rho_sq_unit = rho_sq_unit + (rho_sq_det if post_phase @@ -769,7 +1079,15 @@ def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, _TIME_ADAPTIVE_FACTOR_MAX = 1024 _TIME_ADAPTIVE_SAFETY = 2.0 _TIME_ADAPTIVE_RTOL = 1e-3 +# Threshold of the endpoint certificate. No production caller applies it since +# 2026-09-08 (kernel default None); kept for the tests that pin what it rejected. +# Evidence: DESIGN_jax_bandlimited_distmarg.md, "The endpoint certificate". _TIME_ENDPOINT_LOG_GAP_MIN = 15.0 +# Element budget for one distance-quadrature block on the refined time grid. +# The refined row is up to 2048x the data grid, so the block count is derived +# from the row length rather than inherited from the coarse ``grid_block``. +_BANDLIMITED_GRID_ELEMENTS = 1 << 22 +_LOG_ZERO = -1e300 def default_time_guard(npts): @@ -786,6 +1104,19 @@ def default_time_guard(npts): return max(32, int(npts) // 2) +def bandlimited_time_guard(npts): + """``(guard_initial, guard_certified)`` for the conventional bandlimited API. + + ONE definition, because three sites need the same pair: the fixed-distance + kernel, the distance-marginalized kernel, and the driver's storage-window + sizing. A second copy would let the gathered support and the guard the + quadrature actually uses drift apart, which is a wrong likelihood and not an + error. + """ + g_initial = 1 << int(np.ceil(np.log2(default_time_guard(int(npts))))) + return g_initial, 2 * g_initial + + def _upsample_bandlimited(x, factor, axis=-1): """Band-limited resampling of ``x`` by an integer ``factor``, ASSUMING PERIODICITY. @@ -1062,15 +1393,36 @@ def refine_one(args): def _time_marginalize_reflected_primitive(kappa_t, rho_sq, deltaT, phase_marginalization=False, - guard=0): + guard=0, reduce_fn=None, + endpoint_log_gap=None): """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. + Endpoint mass is covered by the guard-agreement certificate, not by an + endpoint gap. The 15-nat gap of 2026-08-29 was switched off on 2026-09-08: + a row's peak-to-endpoint contrast is bounded by its own amplitude (at most + 2 max|kappa| on the fixed-distance field; a prior-mass floor on the + distance-marginalized one), so a fixed gap rejected every blind or far draw + whatever the trapezoid did, and the rows it rejected alone agree with an + independent reference to 1e-4 nat on both fields. + + ``reduce_fn(kappa, rho_sq) -> lnL`` is the shape-preserving nonlinear + reduction applied AFTER refinement. ``None`` is the fixed-distance + reduction; :func:`fused_log_likelihood_distmarg` passes the distance + quadrature so the same certificates cover it. It is applied to the + curvature probe as well as the fine grid, because the starting factor has + to be derived from the field that is integrated. + + Every per-row quantity -- probe, factor selection, refinement -- is built + inside the ``lax.map`` body, so scratch is set by one row and not by the + sampler batch. + + ``endpoint_log_gap`` is the endpoint certificate's threshold in nats, or + ``None`` (the default) for no endpoint certificate; the resolution, + doubling and guard-agreement certificates are in force either way. No + production caller passes a threshold. """ guard = int(guard) if guard < 0 or 2 * guard >= kappa_t.shape[-1] - 1: @@ -1079,6 +1431,10 @@ def _time_marginalize_reflected_primitive(kappa_t, rho_sq, deltaT, inner_guard = guard // 2 if guard and inner_guard < 1: raise ValueError("guard convergence requires at least two samples per end") + if reduce_fn is None: + def reduce_fn(kappa, rho): + return ((jnp.abs(kappa) if phase_marginalization else kappa.real) + - 0.5 * rho) 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) @@ -1097,26 +1453,6 @@ def taper_support(x, support_guard): 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): @@ -1134,21 +1470,28 @@ def at_factor(kappa, rho, f, support_guard): # mass is negligible. kappa = taper_support(kappa, support_guard) dense_kappa = _reflected_fft_upsample(kappa, f) + # Crop to the integrated window BEFORE reducing. Refinement needs + # the guard columns; nothing downstream reads a reduced value + # outside the crop. ``reduce_fn`` is pointwise in the node, so this + # is the same number, but the distance reduction is the expensive + # one and the padded row is several times the window. + start = support_guard * f + dense_kappa = dense_kappa[start:start + (npts - 1) * f + 1] # 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] + dense = reduce_fn(dense_kappa, dense_rho) 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 + if endpoint_log_gap is None: + boundary_ok = True + else: + peak = jnp.max(dense) + endpoint = jnp.maximum(dense[0], dense[-1]) + boundary_ok = endpoint <= peak - float(endpoint_log_gap) return value, resolved, boundary_ok def branch(args): @@ -1167,17 +1510,37 @@ def branch(args): branches = tuple(make_branch(f) for f in powers) def refine_one(args): - kappa, rho, row_factor = args + kappa, rho = args + # 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(kappa, guard), 2) + probe_kappa = probe_kappa[2 * guard:2 * guard + (npts - 1) * 2 + 1] + probe_rho = jnp.broadcast_to(rho[:1], probe_kappa.shape) + probe = reduce_fn(probe_kappa, probe_rho) + 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) + row_factor = jnp.minimum( + factor_float, float(_TIME_ADAPTIVE_FACTOR_MAX)).astype(jnp.int32) 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)) + return jnp.where(too_sharp, jnp.nan, + 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) + jax.checkpoint(refine_one), (clean_kappa, clean_rho)) return jnp.where(finite_rows, refined, jnp.nan) @@ -1192,7 +1555,11 @@ def _time_marginalize_terminal(lnL_t, data, time_quadrature=TIME_QUAD_DEFAULT, if not bandlimited_safe: raise ValueError( "bandlimited terminal interpolation is invalid after nonlinear " - "distance/phase/polarization marginalization; use 'simpson'") + "distance/phase/polarization marginalization; use 'simpson'. " + "This refuses the ORDERING, not the option: a caller that can " + "refine the primitive first and reduce on the refined nodes does " + "not reach here. fused_log_likelihood_distmarg does exactly " + "that and supports 'bandlimited'.") return _time_marginalize_reflected_fft(lnL_t, data.deltaT, data.w_t) @@ -1273,9 +1640,7 @@ def fused_log_likelihood(data, ra, dec, psi, incl, phiref, distMpc, 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 + guard = bandlimited_time_guard(data.npts)[1] else: guard = 0 distMpc = jnp.asarray(distMpc, dtype=jnp.float64) @@ -1337,13 +1702,38 @@ def fused_log_likelihood_distmarg(data, ra, dec, psi, incl, phiref, Grid of ``x = distMpcRef / d`` values. log_w_grid : array_like, shape (G,) Log quadrature weights (including the distance prior) for each grid point. + time_quadrature : {"simpson", "bandlimited"} + ``"simpson"`` (default) applies the distance quadrature on the data time + grid and integrates it with the fixed Simpson weights. ``"bandlimited"`` + refines the complex primitive kappa(t) first, applies the SAME distance + quadrature on the refined nodes, and integrates by trapezoid under the + resolution, endpoint and guard certificates of + :func:`_time_marginalize_reflected_primitive`. Refining the reduced + lnL(t) instead is what that ordering exists to avoid. """ + if time_quadrature not in _TIME_QUAD_CHOICES: + raise ValueError("time_quadrature must be one of %r, got %r" + % (_TIME_QUAD_CHOICES, time_quadrature)) + bandlimited = time_quadrature == "bandlimited" + if bandlimited: + if return_lnLt: + raise ValueError( + "return_lnLt returns the reduced field on the DATA grid; the " + "band-limited path has no such field, because the distance " + "reduction is applied on the refined grid. Use " + "time_quadrature='simpson' to read lnL_t.") + if _norm_is_arrival_time_dependent(data): + raise ValueError( + "time_quadrature='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_quadrature='simpson' for rotation data.") x_grid = jnp.asarray(x_grid, dtype=jnp.float64) log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64) + guard = bandlimited_time_guard(data.npts)[1] if bandlimited else 0 kappa_unit, rho_sq_unit = _accumulate_unit( - data, ra, dec, psi, incl, phiref, interp, phase_marginalization) - K = jnp.abs(kappa_unit) if phase_marginalization else kappa_unit.real - R = rho_sq_unit + data, ra, dec, psi, incl, phiref, interp, phase_marginalization, + guard=guard) # log-sum-exp over the distance grid -> (S, npts), done in a few *vectorized* # blocks combined by a running log-sum-exp. The block loop is a plain Python @@ -1353,12 +1743,80 @@ def fused_log_likelihood_distmarg(data, ra, dec, psi, incl, phiref, # bounded. Mathematically identical to the previous scan. 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) + if bandlimited: + def _reduce(kappa, rho): + k = jnp.abs(kappa) if phase_marginalization else kappa.real + shape = k.shape + n = 1 + for dim in shape: + n *= int(dim) + block = min(max(1, _BANDLIMITED_GRID_ELEMENTS // max(n, 1)), + int(x_grid.shape[0])) + out = _logsumexp_grid_scanned( + k.reshape(n), rho.reshape(n), a, b, log_w_grid, block) + return out.reshape(shape) + + # The endpoint certificate is OFF here, and only here. The marginal + # field has a floor: at every node the distance sum is at least the + # far-distance prior mass, so a row's peak-to-endpoint contrast is + # bounded by its own peak height, and a fixed 15-nat gap rejects every + # low-contrast row outright, converged or not. Those rows are exactly + # the blind sky/orientation draws a prior-seeded mode evaluates by the + # thousand. The reconstruction at the window edge is still certified + # by the guard-agreement check, and the value by the doubling check. + # Numbers: DESIGN_jax_bandlimited_distmarg.md, "The endpoint + # certificate". The fixed-distance kernel keeps the gap. + return _time_marginalize_reflected_primitive( + kappa_unit, rho_sq_unit, data.deltaT, + phase_marginalization=phase_marginalization, guard=guard, + reduce_fn=_reduce, endpoint_log_gap=None) + K = jnp.abs(kappa_unit) if phase_marginalization else kappa_unit.real + lnL_t = _logsumexp_grid_blocked(K, rho_sq_unit, a, b, log_w_grid, grid_block) if return_lnLt: return lnL_t return _time_marginalize_terminal(lnL_t, data, time_quadrature) +def _logsumexp_grid_scanned(K, R, a, b, log_w, block): + """Same integral as :func:`_logsumexp_grid_blocked`, traced ONCE. + + That function's Python loop unrolls, which is what makes reverse mode fast + on the data grid. The refined time grid is up to 2048x longer, so a block + sized to bound the working set implies many blocks, and the unrolled copies + would then be multiplied by the certificate's eleven refinement branches. + Scanning keeps the graph one block wide at any refinement factor. + + Padding weights are a finite log-zero rather than -inf: a -inf entry can + become the block maximum and turn the shift into inf-inf. + """ + G = int(a.shape[0]) + block = max(1, min(int(block), G)) + n_blocks = (G + block - 1) // block + pad = n_blocks * block - G + if pad: + zeros = jnp.zeros((pad,), dtype=a.dtype) + a = jnp.concatenate([a, zeros]) + b = jnp.concatenate([b, zeros]) + log_w = jnp.concatenate( + [log_w, jnp.full((pad,), _LOG_ZERO, dtype=log_w.dtype)]) + a = a.reshape(n_blocks, block) + b = b.reshape(n_blocks, block) + log_w = log_w.reshape(n_blocks, block) + + def step(carry, node): + m, s = carry + a_b, b_b, w_b = node + e = K[..., None] * a_b + R[..., None] * b_b + w_b + m_blk = jnp.max(e, axis=-1) + s_blk = jnp.sum(jnp.exp(e - m_blk[..., None]), axis=-1) + m_new = jnp.maximum(m, m_blk) + return (m_new, s * jnp.exp(m - m_new) + s_blk * jnp.exp(m_blk - m_new)), None + + init = (jnp.full(K.shape, -jnp.inf), jnp.zeros(K.shape)) + (m, s), _ = jax.lax.scan(step, init, (a, b, log_w)) + return m + jnp.log(s) + + def _logsumexp_grid_blocked(K, R, a, b, log_w, block): """Stable log sum_g exp(K*a_g + R*b_g + log_w_g) over the grid axis. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_planner.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_planner.py new file mode 100644 index 000000000..036adbbb3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_planner.py @@ -0,0 +1,1206 @@ +"""Opt-in planner for error- and resource-budgeted direct marginalization. + +This module is deliberately separate from :func:`choose_angle_marg_scheme`. +Importing it changes no default and the existing ``angle_marg='auto'`` path +continues to use the measured amplitude crossover. A caller must construct +scheme offers, provide every requested per-axis error budget and both resource +budgets, and explicitly call :func:`plan_direct_marginalization`. + +The planner does not turn a calibration into a proof. Each offer carries an +accuracy assessment, a completeness warrant and its provenance. Only an +assessment marked ``CERTIFIED`` under a warrant with an implemented +certificate participates in the ``cheapest-certified`` choice. An empirical +or unknown offer can only be run when the caller explicitly sets +``allow_best_effort=True``; otherwise it is returned as a non-executable +suggestion on a structured decline. + +Production callers may separately pass that fail-closed decision, or a runtime +method-warrant refusal, to :func:`resolve_plan_for_production`. This API never +promotes the planner's suggestion. It requires an explicitly provisioned +support-covering fallback and records its real (possibly uncertified) accuracy +label. A method decline cannot become a waveform-failure/sample-drop result. +Only independent waveform/base-likelihood evidence can authorize that outcome. + +By default, resource estimates are conservative additive contributions on a +common unit: compute and peak-memory contributions are summed. A nested JAX +adapter can instead supply a combination-aware ``resource_model`` whose return +value carries its own provenance. It may over-count buffers whose lifetimes do +not overlap, but may not under-count them; an optimistic lifetime model would +be another silent OOM fallback. +""" + +from dataclasses import dataclass, field +from enum import Enum +from itertools import product +import math +from types import MappingProxyType + + +__all__ = [ + "AccuracyAssessment", + "ConditionalRequirement", + "ConservativeFallbackPolicy", + "EvidenceKind", + "FallbackConfigurationError", + "JAX_CONSERVATIVE_FALLBACK_SCHEMES", + "JAX_DIRECT_MARGINALIZATION_AXES", + "JAX_SCHEME_PROFILES", + "MarginalizationPlanDeclined", + "MethodDecline", + "PlanDecision", + "ProductionResolution", + "ResolutionAction", + "ResourceBudget", + "ResourceEstimate", + "SchemeOffer", + "SchemeProfile", + "Warrant", + "WarrantKind", + "WaveformFailure", + "WaveformLikelihoodFailure", + "make_jax_production_fallback_policy", + "make_jax_scheme_offer", + "plan_direct_marginalization", + "plan_jax_direct_marginalization", + "resolve_plan_for_production", +] + + +class WarrantKind(str, Enum): + """Finite structures which may warrant a completeness certificate. + + ``EFFECTIVE_BANDWIDTH_WITH_MARGIN`` is intentionally represented even + though it cannot certify completeness. Naming it lets the planner refuse + a proof claim instead of treating every amplitude-sized grid as exact. + """ + + EXACT_BAND_LIMIT = "exact-band-limit" + EXACT_TRIG_DEGREE = "exact-trig-degree" + BOUNDED_STATIONARY_SET = "bounded-stationary-set" + EFFECTIVE_BANDWIDTH_WITH_MARGIN = "effective-bandwidth-with-margin" + EMPIRICAL_CALIBRATION = "empirical-calibration" + NONE = "none" + + +class EvidenceKind(str, Enum): + """Strength of a quantitative per-axis error assessment.""" + + CERTIFIED = "certified" + VALIDATED = "validated" + ESTIMATED = "estimated" + UNKNOWN = "unknown" + + +class ResolutionAction(str, Enum): + """Production disposition, separate from the planner's proof claim.""" + + USE_PREFERRED = "use-preferred" + USE_CONSERVATIVE_FALLBACK = "use-conservative-fallback" + WAVEFORM_FAILURE = "waveform-failure" + + +_POTENTIALLY_CERTIFYING_WARRANTS = frozenset(( + WarrantKind.EXACT_BAND_LIMIT, + WarrantKind.EXACT_TRIG_DEGREE, + WarrantKind.BOUNDED_STATIONARY_SET, +)) + + +def _enum_value(value, enum_type, field_name): + try: + return value if isinstance(value, enum_type) else enum_type(value) + except ValueError: + raise ValueError("unknown %s %r" % (field_name, value)) + + +def _finite_nonnegative(value, field_name): + value = float(value) + if not math.isfinite(value) or value < 0.0: + raise ValueError("%s must be finite and non-negative; got %r" + % (field_name, value)) + return value + + +def _nonnegative_integer(value, field_name): + if isinstance(value, bool): + raise ValueError("%s must be a non-negative integer" % field_name) + try: + as_float = float(value) + as_int = int(value) + except (TypeError, ValueError, OverflowError): + raise ValueError("%s must be a non-negative integer" % field_name) + if (not math.isfinite(as_float) or as_float < 0.0 + or as_float != float(as_int)): + raise ValueError("%s must be a non-negative integer; got %r" + % (field_name, value)) + return as_int + + +@dataclass(frozen=True) +class Warrant: + """Completeness warrant carried by one implementation. + + ``certificate_available`` means the implementation actually discharges a + quantitative error inequality. A mathematical structure that could + support a future certificate is not sufficient. + """ + + kind: WarrantKind + scope: str + certificate_available: bool + provenance: str + + def __post_init__(self): + object.__setattr__(self, "kind", _enum_value( + self.kind, WarrantKind, "warrant kind")) + if not self.scope or not self.provenance: + raise ValueError("warrant scope and provenance must be non-empty") + if (self.certificate_available + and self.kind not in _POTENTIALLY_CERTIFYING_WARRANTS): + raise ValueError( + "warrant %s cannot advertise a completeness certificate" + % self.kind.value) + + def as_dict(self): + return dict(kind=self.kind.value, scope=self.scope, + certificate_available=bool(self.certificate_available), + provenance=self.provenance) + + +@dataclass(frozen=True) +class AccuracyAssessment: + """Quantitative error information for one axis and scheme. + + The unit is absolute error in the marginalized log likelihood (nats). + ``UNKNOWN`` must carry ``max_error_nats=None``. The other evidence kinds + need a finite non-negative value, but only ``CERTIFIED`` is a hard bound. + """ + + evidence: EvidenceKind + max_error_nats: object + provenance: str + + def __post_init__(self): + object.__setattr__(self, "evidence", _enum_value( + self.evidence, EvidenceKind, "evidence kind")) + if not self.provenance: + raise ValueError("accuracy provenance must be non-empty") + if self.evidence is EvidenceKind.UNKNOWN: + if self.max_error_nats is not None: + raise ValueError( + "UNKNOWN accuracy must not carry a numerical error") + else: + object.__setattr__(self, "max_error_nats", _finite_nonnegative( + self.max_error_nats, "max_error_nats")) + + def as_dict(self): + return dict(evidence=self.evidence.value, + max_error_nats=self.max_error_nats, + provenance=self.provenance) + + +@dataclass(frozen=True) +class ResourceEstimate: + """Conservative contribution to a plan's compute and peak memory.""" + + compute_units: float + memory_bytes: int + provenance: str + + def __post_init__(self): + object.__setattr__(self, "compute_units", _finite_nonnegative( + self.compute_units, "compute_units")) + object.__setattr__(self, "memory_bytes", _nonnegative_integer( + self.memory_bytes, "memory_bytes")) + if not self.provenance: + raise ValueError("resource provenance must be non-empty") + + def as_dict(self): + return dict(compute_units=self.compute_units, + memory_bytes=self.memory_bytes, + provenance=self.provenance) + + +@dataclass(frozen=True) +class ResourceBudget: + """Hard request-level ceilings. + + The fields may be ``None`` only so a missing budget can produce a + structured decline. A complete request needs both. + """ + + max_compute_units: object + max_memory_bytes: object + + def __post_init__(self): + if self.max_compute_units is not None: + object.__setattr__(self, "max_compute_units", _finite_nonnegative( + self.max_compute_units, "max_compute_units")) + if self.max_memory_bytes is not None: + object.__setattr__(self, "max_memory_bytes", _nonnegative_integer( + self.max_memory_bytes, "max_memory_bytes")) + + def validation_errors(self): + errors = [] + if self.max_compute_units is None: + errors.append("max_compute_units") + if self.max_memory_bytes is None: + errors.append("max_memory_bytes") + return tuple(errors) + + def as_dict(self): + return dict(max_compute_units=self.max_compute_units, + max_memory_bytes=self.max_memory_bytes) + + +@dataclass(frozen=True) +class ConditionalRequirement: + """Capability required only when another scheme/token is selected.""" + + trigger: str + capability: str + reason: str + + def __post_init__(self): + if not self.trigger or not self.capability or not self.reason: + raise ValueError("conditional requirement fields must be non-empty") + + def as_dict(self): + return dict(trigger=self.trigger, capability=self.capability, + reason=self.reason) + + +@dataclass(frozen=True) +class SchemeOffer: + """One runnable scheme offered for one marginalized axis.""" + + axis: str + scheme: str + accuracy: AccuracyAssessment + resources: ResourceEstimate + warrant: Warrant + provenance: str + requires: frozenset = field(default_factory=frozenset) + provides: frozenset = field(default_factory=frozenset) + conflicts: frozenset = field(default_factory=frozenset) + conditional_requirements: tuple = field(default_factory=tuple) + + def __post_init__(self): + if not self.axis or not self.scheme or not self.provenance: + raise ValueError("offer axis, scheme and provenance must be non-empty") + object.__setattr__(self, "requires", frozenset(self.requires)) + object.__setattr__(self, "provides", frozenset(self.provides)) + object.__setattr__(self, "conflicts", frozenset(self.conflicts)) + object.__setattr__(self, "conditional_requirements", + tuple(self.conditional_requirements)) + if (self.accuracy.evidence is EvidenceKind.CERTIFIED + and not self.warrant.certificate_available): + raise ValueError( + "%s cannot claim CERTIFIED accuracy: its %s warrant has no " + "implemented certificate" % (self.key, self.warrant.kind.value)) + + @property + def key(self): + return "%s:%s" % (self.axis, self.scheme) + + def as_dict(self): + return dict( + key=self.key, axis=self.axis, scheme=self.scheme, + accuracy=self.accuracy.as_dict(), + resources=self.resources.as_dict(), warrant=self.warrant.as_dict(), + provenance=self.provenance, requires=sorted(self.requires), + provides=sorted(self.provides), conflicts=sorted(self.conflicts), + conditional_requirements=[r.as_dict() + for r in self.conditional_requirements]) + + +class MarginalizationPlanDeclined(RuntimeError): + """Raised when a caller tries to execute a declined decision.""" + + +class FallbackConfigurationError(RuntimeError): + """Raised when a method decline has no runnable fail-safe policy.""" + + +class WaveformLikelihoodFailure(RuntimeError): + """Raised only for an explicitly reported waveform/likelihood failure.""" + + +@dataclass(frozen=True) +class PlanDecision: + """Structured planner result. ``action`` is either ``run`` or ``decline``.""" + + action: str + basis: str + reason_code: str + reason: str + selected: tuple + suggested: tuple + resource_use: object + suggested_resource_use: object + certified: bool + meets_error_budget: bool + ledger: dict + + def require_selection(self): + """Return the selected offers, refusing a declined recommendation.""" + if self.action != "run": + raise MarginalizationPlanDeclined( + "%s: %s" % (self.reason_code, self.reason)) + return self.selected + + def as_dict(self): + return dict( + action=self.action, basis=self.basis, + reason_code=self.reason_code, reason=self.reason, + selected=[o.as_dict() for o in self.selected], + suggested=[o.as_dict() for o in self.suggested], + resource_use=(None if self.resource_use is None + else self.resource_use.as_dict()), + suggested_resource_use=( + None if self.suggested_resource_use is None + else self.suggested_resource_use.as_dict()), + certified=bool(self.certified), + meets_error_budget=bool(self.meets_error_budget), + ledger=self.ledger) + + +@dataclass(frozen=True) +class MethodDecline: + """A planning or runtime marginalizer refusal, never a waveform failure. + + Runtime implementations should use this record for events such as an + incomplete stationary-root enumeration. Such an event invalidates the + preferred *method's* warrant, not the waveform or the likelihood point. + """ + + code: str + reason: str + provenance: str + axis: object = None + stage: str = "runtime" + ledger: dict = field(default_factory=dict) + + def __post_init__(self): + if not self.code or not self.reason or not self.provenance: + raise ValueError( + "method decline code, reason and provenance must be non-empty") + if not self.stage: + raise ValueError("method decline stage must be non-empty") + if self.axis is not None and not self.axis: + raise ValueError("method decline axis must be non-empty or None") + + def as_dict(self): + return dict(code=self.code, reason=self.reason, + provenance=self.provenance, axis=self.axis, + stage=self.stage, ledger=self.ledger) + + +@dataclass(frozen=True) +class WaveformFailure: + """Independent evidence that the waveform/base likelihood is unusable. + + The production resolver never constructs this object from a planner or + marginalizer decline. A caller must report it explicitly from the + waveform/base-likelihood layer. + """ + + code: str + reason: str + provenance: str + ledger: dict = field(default_factory=dict) + + def __post_init__(self): + if not self.code or not self.reason or not self.provenance: + raise ValueError( + "waveform failure code, reason and provenance must be non-empty") + + def as_dict(self): + return dict(code=self.code, reason=self.reason, + provenance=self.provenance, ledger=self.ledger) + + +@dataclass(frozen=True) +class ConservativeFallbackPolicy: + """Explicit reserve plan used after a marginalization-method decline. + + The fallback has its own hard resource budget because a resource-limited + preferred plan may need a finite, slower reserve path. The non-empty + ``finite_output_contract`` is an adapter assertion that these offers cover + the full finite domain without relying on the declined shortcut. It is + provenance, not an error certificate; accuracy labels remain unchanged. + """ + + offers: tuple + resource_budget: ResourceBudget + provenance: str + finite_output_contract: str + + def __post_init__(self): + object.__setattr__(self, "offers", tuple(self.offers)) + if not self.offers: + raise ValueError("a conservative fallback needs at least one offer") + if not all(isinstance(offer, SchemeOffer) for offer in self.offers): + raise TypeError("fallback offers must be SchemeOffer objects") + if not self.provenance or not self.finite_output_contract: + raise ValueError( + "fallback provenance and finite-output contract are required") + axes = [offer.axis for offer in self.offers] + if len(axes) != len(set(axes)): + raise ValueError( + "a conservative fallback may offer only one scheme per axis") + budget = self.resource_budget + if isinstance(budget, dict): + budget = ResourceBudget(budget.get("max_compute_units"), + budget.get("max_memory_bytes")) + object.__setattr__(self, "resource_budget", budget) + if not isinstance(budget, ResourceBudget): + raise TypeError("fallback resource_budget must be ResourceBudget") + missing = budget.validation_errors() + if missing: + raise ValueError("fallback resource budget is missing %r" + % (missing,)) + + def as_dict(self): + return dict(offers=[offer.as_dict() for offer in self.offers], + resource_budget=self.resource_budget.as_dict(), + provenance=self.provenance, + finite_output_contract=self.finite_output_contract) + + +@dataclass(frozen=True) +class ProductionResolution: + """Executable production disposition with complete failure provenance.""" + + action: ResolutionAction + selected: tuple + resource_use: object + certified: bool + meets_error_budget: bool + drops_sample: bool + method_decline: object + waveform_failure: object + ledger: dict + + def __post_init__(self): + object.__setattr__(self, "action", _enum_value( + self.action, ResolutionAction, "resolution action")) + object.__setattr__(self, "selected", tuple(self.selected)) + is_waveform_failure = self.action is ResolutionAction.WAVEFORM_FAILURE + if is_waveform_failure: + if self.waveform_failure is None or self.selected: + raise ValueError( + "waveform-failure resolution needs failure evidence and " + "no selection") + if self.method_decline is not None or not self.drops_sample: + raise ValueError( + "waveform failure cannot be conflated with a method decline") + else: + if not self.selected or self.waveform_failure is not None: + raise ValueError( + "runnable resolution needs a selection and no waveform " + "failure") + if self.drops_sample: + raise ValueError("a runnable resolution cannot drop the sample") + if (self.action is ResolutionAction.USE_CONSERVATIVE_FALLBACK + and self.method_decline is None): + raise ValueError("fallback resolution needs a method decline") + if (self.action is ResolutionAction.USE_PREFERRED + and self.method_decline is not None): + raise ValueError("preferred resolution cannot carry a decline") + + def require_selection(self): + """Return a runnable plan; raise only for explicit waveform failure.""" + if self.action is ResolutionAction.WAVEFORM_FAILURE: + raise WaveformLikelihoodFailure( + "%s: %s" % (self.waveform_failure.code, + self.waveform_failure.reason)) + return self.selected + + def as_dict(self): + return dict( + action=self.action.value, + selected=[offer.as_dict() for offer in self.selected], + resource_use=(None if self.resource_use is None + else self.resource_use.as_dict()), + certified=bool(self.certified), + meets_error_budget=bool(self.meets_error_budget), + drops_sample=bool(self.drops_sample), + method_decline=(None if self.method_decline is None + else self.method_decline.as_dict()), + waveform_failure=(None if self.waveform_failure is None + else self.waveform_failure.as_dict()), + ledger=self.ledger) + + +def _resource_use(offers, resource_model=None): + if resource_model is None: + return ResourceEstimate( + sum(o.resources.compute_units for o in offers), + sum(o.resources.memory_bytes for o in offers), + "conservative additive aggregation of selected offer estimates") + use = resource_model(tuple(offers)) + if not isinstance(use, ResourceEstimate): + raise TypeError("resource_model must return ResourceEstimate") + return use + + +def _resource_reasons(use, budget): + reasons = [] + if use.compute_units > float(budget.max_compute_units): + reasons.append("compute %.9g exceeds budget %.9g" + % (use.compute_units, + float(budget.max_compute_units))) + if use.memory_bytes > int(budget.max_memory_bytes): + reasons.append("memory %d exceeds budget %d" + % (use.memory_bytes, int(budget.max_memory_bytes))) + return reasons + + +def _compatibility_reasons(offers, capabilities): + capabilities = frozenset(capabilities) + tokens = set(capabilities) + for offer in offers: + tokens.add(offer.key) + tokens.update(offer.provides) + reasons = [] + for offer in offers: + missing = sorted(offer.requires.difference(tokens)) + if missing: + reasons.append("%s missing requirements %r" % (offer.key, missing)) + conflicts = sorted(offer.conflicts.intersection(tokens)) + if conflicts: + reasons.append("%s conflicts with %r" % (offer.key, conflicts)) + for requirement in offer.conditional_requirements: + if (requirement.trigger in tokens + and requirement.capability not in capabilities): + reasons.append( + "%s with %s requires capability %s: %s" + % (offer.key, requirement.trigger, + requirement.capability, requirement.reason)) + return reasons + + +def _error_reasons(offers, error_budget, certified_only): + reasons = [] + for offer in offers: + assessment = offer.accuracy + if certified_only and assessment.evidence is not EvidenceKind.CERTIFIED: + reasons.append("%s accuracy is %s, not certified" + % (offer.key, assessment.evidence.value)) + continue + if assessment.max_error_nats is None: + reasons.append("%s has no quantitative error assessment" % offer.key) + continue + limit = float(error_budget[offer.axis]) + if assessment.max_error_nats > limit: + reasons.append("%s error %.9g exceeds axis budget %.9g" + % (offer.key, assessment.max_error_nats, limit)) + return reasons + + +def _accuracy_rank(offers, error_budget, resource_model): + unknown = sum(o.accuracy.max_error_nats is None for o in offers) + ratios = [o.accuracy.max_error_nats / float(error_budget[o.axis]) + for o in offers if o.accuracy.max_error_nats is not None] + worst = max(ratios) if ratios else math.inf + total = sum(ratios) if ratios else math.inf + evidence_order = {EvidenceKind.CERTIFIED: 0, EvidenceKind.VALIDATED: 1, + EvidenceKind.ESTIMATED: 2, EvidenceKind.UNKNOWN: 3} + evidence = sum(evidence_order[o.accuracy.evidence] for o in offers) + use = _resource_use(offers, resource_model) + return (unknown, worst, total, evidence, use.compute_units, + use.memory_bytes, tuple(o.key for o in offers)) + + +def _cost_rank(offers, error_budget, resource_model): + use = _resource_use(offers, resource_model) + ratios = [o.accuracy.max_error_nats / float(error_budget[o.axis]) + for o in offers] + return (use.compute_units, use.memory_bytes, max(ratios), sum(ratios), + tuple(o.key for o in offers)) + + +def _preflight_decline(reason_code, reason, axes, error_budget, + resource_budget, capabilities, details): + if resource_budget is None: + resource_record = None + elif isinstance(resource_budget, dict): + resource_record = dict(resource_budget) + else: + resource_record = resource_budget.as_dict() + return PlanDecision( + action="decline", basis="decline", reason_code=reason_code, + reason=reason, selected=(), suggested=(), resource_use=None, + suggested_resource_use=None, certified=False, + meets_error_budget=False, + ledger=dict(required_axes=list(axes), + error_budget=None if error_budget is None + else dict(error_budget), + resource_budget=resource_record, + capabilities=sorted(capabilities), details=details, + combinations=[])) + + +def plan_direct_marginalization(offers, error_budget, resource_budget, *, + required_axes=None, capabilities=(), + allow_best_effort=False, resource_model=None): + """Choose a direct-marginalization plan without changing any RIFT default. + + The primary policy is the least-compute plan whose per-axis errors are + certified within budget and whose summed resource estimates fit. If none + exists, the most accurate affordable compatible plan is recorded as a + suggestion. It becomes executable only under the explicit + ``allow_best_effort=True`` policy. ``resource_model``, when supplied, is + called on each complete offer combination and must return a provenance- + carrying :class:`ResourceEstimate`; exceptions are never converted to a + decline. + """ + offers = tuple(offers) + capabilities = frozenset(capabilities) + keys = [offer.key for offer in offers] + if len(keys) != len(set(keys)): + raise ValueError("offer keys must be unique; got %r" % keys) + axes = tuple(required_axes) if required_axes is not None else tuple(sorted( + set(offer.axis for offer in offers))) + if not axes: + return _preflight_decline( + "missing-axis", "no marginalization axes were requested", axes, + error_budget, resource_budget, capabilities, {}) + + duplicate_axes = sorted(set(axis for axis in axes if axes.count(axis) > 1)) + if duplicate_axes: + # One scheme per axis is the planner's contract. A repeated axis would + # otherwise enter the Cartesian product twice, select the same offer + # twice and double-count its compute and memory. + return _preflight_decline( + "duplicate-axis", + "required axes repeat %r; each axis may be marginalized once" + % duplicate_axes, axes, error_budget, resource_budget, + capabilities, dict(duplicate_axes=duplicate_axes)) + + by_axis = {axis: tuple(o for o in offers if o.axis == axis) for axis in axes} + unsupported = [axis for axis in axes if not by_axis[axis]] + if unsupported: + return _preflight_decline( + "unsupported-axis", "no scheme offers for axes %r" % unsupported, + axes, error_budget, resource_budget, capabilities, + dict(unsupported_axes=unsupported)) + + if error_budget is None: + return _preflight_decline( + "missing-error-budget", "a per-axis error budget is required", + axes, error_budget, resource_budget, capabilities, + dict(missing_axes=list(axes))) + missing_axes = [axis for axis in axes if axis not in error_budget] + if missing_axes: + return _preflight_decline( + "missing-error-budget", + "error budget is missing axes %r" % missing_axes, + axes, error_budget, resource_budget, capabilities, + dict(missing_axes=missing_axes)) + clean_error_budget = {} + for axis in axes: + value = float(error_budget[axis]) + if not math.isfinite(value) or value <= 0.0: + raise ValueError("error budget for %s must be finite and positive" + % axis) + clean_error_budget[axis] = value + + if resource_budget is None: + return _preflight_decline( + "missing-resource-budget", + "both compute and memory budgets are required", axes, + clean_error_budget, resource_budget, capabilities, + dict(missing=("max_compute_units", "max_memory_bytes"))) + if isinstance(resource_budget, dict): + resource_budget = ResourceBudget( + resource_budget.get("max_compute_units"), + resource_budget.get("max_memory_bytes")) + missing_resources = resource_budget.validation_errors() + if missing_resources: + return _preflight_decline( + "missing-resource-budget", "resource budget is missing %r" + % (missing_resources,), axes, clean_error_budget, + resource_budget, capabilities, + dict(missing=missing_resources)) + + combinations = [] + compatible = [] + affordable = [] + certified = [] + certified_affordable = [] + for combination in product(*(by_axis[axis] for axis in axes)): + use = _resource_use(combination, resource_model) + compat_reasons = _compatibility_reasons(combination, capabilities) + resource_reasons = _resource_reasons(use, resource_budget) + certified_error_reasons = _error_reasons( + combination, clean_error_budget, certified_only=True) + numeric_error_reasons = _error_reasons( + combination, clean_error_budget, certified_only=False) + record = dict( + schemes=[o.key for o in combination], + compatibility_reasons=compat_reasons, + resource_reasons=resource_reasons, + certified_error_reasons=certified_error_reasons, + numeric_error_reasons=numeric_error_reasons, + resource_use=use.as_dict()) + combinations.append(record) + if compat_reasons: + continue + compatible.append(combination) + if not resource_reasons: + affordable.append(combination) + if not certified_error_reasons: + certified.append(combination) + if not resource_reasons: + certified_affordable.append(combination) + + ledger = dict( + required_axes=list(axes), error_budget=clean_error_budget, + resource_budget=resource_budget.as_dict(), + capabilities=sorted(capabilities), + allow_best_effort=bool(allow_best_effort), + offers=[offer.as_dict() for offer in offers], + combinations=combinations) + + if certified_affordable: + chosen = min(certified_affordable, + key=lambda c: _cost_rank( + c, clean_error_budget, resource_model)) + use = _resource_use(chosen, resource_model) + return PlanDecision( + action="run", basis="cheapest-certified", reason_code="selected", + reason="least-compute compatible plan certified within every " + "axis and resource budget", + selected=tuple(chosen), suggested=(), resource_use=use, + suggested_resource_use=None, certified=True, + meets_error_budget=True, ledger=ledger) + + best = (min(affordable, + key=lambda c: _accuracy_rank( + c, clean_error_budget, resource_model)) + if affordable else None) + best_use = (_resource_use(best, resource_model) + if best is not None else None) + best_numeric_ok = bool(best is not None and not _error_reasons( + best, clean_error_budget, certified_only=False)) + + if best is not None and allow_best_effort: + return PlanDecision( + action="run", basis="most-accurate-affordable", + reason_code="best-effort-authorized", + reason="no affordable fully certified plan; caller explicitly " + "authorized the most accurate affordable compatible plan", + selected=tuple(best), suggested=(), resource_use=best_use, + suggested_resource_use=None, certified=False, + meets_error_budget=best_numeric_ok, ledger=ledger) + + if not compatible: + code = "no-compatible-plan" + reason = "all scheme combinations violate declared compatibility" + elif certified and not certified_affordable: + code = "resource-budget-exceeded" + reason = "certified plans exist, but none fits both resource budgets" + elif not certified: + code = "no-certified-plan" + reason = "no compatible plan is certified within every axis budget" + else: + code = "no-affordable-plan" + reason = "no compatible plan fits both resource budgets" + return PlanDecision( + action="decline", basis="decline", reason_code=code, reason=reason, + selected=(), suggested=tuple(best) if best is not None else (), + resource_use=None, suggested_resource_use=best_use, + certified=False, meets_error_budget=False, ledger=ledger) + + +def _decision_error_reasons(offers, decision, certified_only): + """Assess a resolved plan without manufacturing a missing error budget.""" + error_budget = decision.ledger.get("error_budget") + required_axes = tuple(decision.ledger.get("required_axes", ())) + if not isinstance(error_budget, dict): + return ["the preferred request had no complete error budget"] + missing = [axis for axis in required_axes if axis not in error_budget] + if missing: + return ["the preferred request omitted error budgets for %r" % missing] + return _error_reasons(offers, error_budget, certified_only) + + +def _planner_method_decline(decision): + return MethodDecline( + code=decision.reason_code, + reason=decision.reason, + provenance="fail-closed PlanDecision from the preferred planner", + stage="planning", + ledger=dict(basis=decision.basis, + suggested=[offer.key for offer in decision.suggested])) + + +def resolve_plan_for_production(preferred_decision, fallback_policy=None, *, + method_decline=None, waveform_failure=None, + capabilities=(), resource_model=None): + """Resolve proof failure separately from waveform/likelihood failure. + + A runnable preferred decision passes through unchanged. A fail-closed + planning decision, or an explicit runtime :class:`MethodDecline`, selects + the explicitly configured conservative fallback. The fallback may use a + separate reserve resource budget but retains its real certification and + error labels. Missing, incompatible, or unaffordable fallback setup is a + configuration error; it is never returned as an invalid likelihood point. + + Only a separately constructed :class:`WaveformFailure` can produce a + ``drops_sample=True`` resolution. In particular, callers must report an + incomplete root enumeration as ``method_decline``, not as an exception to + be caught and converted into a waveform failure. + """ + if not isinstance(preferred_decision, PlanDecision): + raise TypeError("preferred_decision must be a PlanDecision") + if method_decline is not None and not isinstance( + method_decline, MethodDecline): + raise TypeError("method_decline must be a MethodDecline") + if waveform_failure is not None and not isinstance( + waveform_failure, WaveformFailure): + raise TypeError("waveform_failure must be a WaveformFailure") + if method_decline is not None and waveform_failure is not None: + raise ValueError( + "a marginalization-method decline is not a waveform failure") + + if waveform_failure is not None: + return ProductionResolution( + action=ResolutionAction.WAVEFORM_FAILURE, selected=(), + resource_use=None, certified=False, meets_error_budget=False, + drops_sample=True, method_decline=None, + waveform_failure=waveform_failure, + ledger=dict( + preferred_decision=preferred_decision.as_dict(), + resolution_policy=( + "sample invalidation requires independent waveform/base-" + "likelihood failure evidence"))) + + if preferred_decision.action == "run" and method_decline is None: + return ProductionResolution( + action=ResolutionAction.USE_PREFERRED, + selected=preferred_decision.selected, + resource_use=preferred_decision.resource_use, + certified=preferred_decision.certified, + meets_error_budget=preferred_decision.meets_error_budget, + drops_sample=False, method_decline=None, waveform_failure=None, + ledger=dict( + preferred_decision=preferred_decision.as_dict(), + resolution_policy="preferred plan remained runnable")) + + if preferred_decision.action == "decline" and method_decline is None: + method_decline = _planner_method_decline(preferred_decision) + elif preferred_decision.action not in ("run", "decline"): + raise ValueError("unknown PlanDecision action %r" + % preferred_decision.action) + + if fallback_policy is None: + raise FallbackConfigurationError( + "%s is a marginalization-method decline, not a waveform failure; " + "an explicit conservative fallback policy is required" + % method_decline.code) + if not isinstance(fallback_policy, ConservativeFallbackPolicy): + raise TypeError("fallback_policy must be ConservativeFallbackPolicy") + + required_axes = tuple(preferred_decision.ledger.get( + "required_axes", ())) + if not required_axes: + required_axes = tuple(offer.axis for offer in + preferred_decision.selected) + base = ({offer.axis: offer for offer in preferred_decision.selected} + if preferred_decision.action == "run" else {}) + fallback_by_axis = {offer.axis: offer + for offer in fallback_policy.offers} + extra = sorted(set(fallback_by_axis).difference(required_axes)) + if extra: + raise FallbackConfigurationError( + "fallback contains unrequested axes %r" % extra) + # A runtime decline without an axis cannot identify which selected method + # lost its warrant. Treat it conservatively as a decline of the complete + # selected plan: every requested axis must be supplied by the explicit + # fallback policy. Retaining the preferred plan and replacing only an + # unrelated axis would report a runnable resolution that still contains the + # method that may have declined. + declined_axes = (required_axes if method_decline.axis is None + else (method_decline.axis,)) + missing_replacements = [axis for axis in declined_axes + if axis not in fallback_by_axis] + if missing_replacements: + raise FallbackConfigurationError( + "fallback does not replace declined axes %r" + % missing_replacements) + for axis in declined_axes: + if axis in base and fallback_by_axis[axis].key == base[axis].key: + raise FallbackConfigurationError( + "fallback repeats declined method %s" + % base[axis].key) + base.update(fallback_by_axis) + missing = [axis for axis in required_axes if axis not in base] + if missing: + raise FallbackConfigurationError( + "fallback does not cover requested axes %r" % missing) + selected = tuple(base[axis] for axis in required_axes) + + active_capabilities = set(preferred_decision.ledger.get( + "capabilities", ())) + active_capabilities.update(capabilities) + compatibility_reasons = _compatibility_reasons( + selected, active_capabilities) + use = _resource_use(selected, resource_model) + resource_reasons = _resource_reasons( + use, fallback_policy.resource_budget) + if compatibility_reasons or resource_reasons: + details = compatibility_reasons + resource_reasons + raise FallbackConfigurationError( + "%s is a method decline, but its configured fallback is not " + "runnable: %s" % (method_decline.code, "; ".join(details))) + + certified_error_reasons = _decision_error_reasons( + selected, preferred_decision, certified_only=True) + numeric_error_reasons = _decision_error_reasons( + selected, preferred_decision, certified_only=False) + fallback_record = dict( + schemes=[offer.key for offer in selected], + compatibility_reasons=compatibility_reasons, + resource_reasons=resource_reasons, + certified_error_reasons=certified_error_reasons, + numeric_error_reasons=numeric_error_reasons, + resource_use=use.as_dict()) + return ProductionResolution( + action=ResolutionAction.USE_CONSERVATIVE_FALLBACK, + selected=selected, resource_use=use, + certified=not certified_error_reasons, + meets_error_budget=not numeric_error_reasons, + drops_sample=False, method_decline=method_decline, + waveform_failure=None, + ledger=dict( + preferred_decision=preferred_decision.as_dict(), + method_decline=method_decline.as_dict(), + fallback_policy=fallback_policy.as_dict(), + fallback_evaluation=fallback_record, + resolution_policy=( + "method/warrant failure selects the explicit finite fallback; " + "it does not invalidate the likelihood point"))) + + +@dataclass(frozen=True) +class SchemeProfile: + """Static compatibility and warrant facts for a shipped JAX scheme.""" + + axis: str + scheme: str + warrant: Warrant + provenance: str + requires: frozenset = field(default_factory=frozenset) + conflicts: frozenset = field(default_factory=frozenset) + conditional_requirements: tuple = field(default_factory=tuple) + + def __post_init__(self): + object.__setattr__(self, "requires", frozenset(self.requires)) + object.__setattr__(self, "conflicts", frozenset(self.conflicts)) + object.__setattr__(self, "conditional_requirements", + tuple(self.conditional_requirements)) + + @property + def key(self): + return "%s:%s" % (self.axis, self.scheme) + + +def _warrant(kind, scope, available, provenance): + return Warrant(kind, scope, available, provenance) + + +_FRAMEWORK = "RIFT/likelihood/DESIGN_peak_local_framework.md" +_ANGLE = "RIFT/likelihood/jax_ile/anglemarg.py" +_DISTANCE = "RIFT/likelihood/jax_ile/DESIGN_jax_distance_quadrature.md" +_TIME = "RIFT/likelihood/time_marginalization_quadrature.py" + + +def _profile(axis, scheme, warrant, provenance, requires=(), conflicts=(), + conditional_requirements=()): + return SchemeProfile(axis, scheme, warrant, provenance, + frozenset(requires), frozenset(conflicts), + tuple(conditional_requirements)) + + +# These profiles state structural facts only. In particular they intentionally +# do not invent error or wall-time envelopes for the current schemes. +_JAX_PROFILE_LIST = ( + _profile("angle", "grid", + _warrant(WarrantKind.NONE, "fixed legacy product grid", False, + _ANGLE), _ANGLE, + conflicts=("distance:loguniform",)), + _profile("angle", "exact", + _warrant(WarrantKind.EFFECTIVE_BANDWIDTH_WITH_MARGIN, + "exact angle coefficients, amplitude-sized exp grid", + False, _FRAMEWORK), _ANGLE, + requires=("angle-amplitude-estimate",)), + _profile("angle", "laplace", + _warrant(WarrantKind.EFFECTIVE_BANDWIDTH_WITH_MARGIN, + "dense phi plus enumerated psi Laplace rule", False, + _FRAMEWORK), _ANGLE, + requires=("angle-amplitude-estimate",), + conditional_requirements=(ConditionalRequirement( + "distance:gh", "gh-laplace-supported", + "the A0==0/B1==0 identity must hold on concrete tables"),)), + _profile("angle", "peak-local", + _warrant(WarrantKind.EFFECTIVE_BANDWIDTH_WITH_MARGIN, + "exact-trig-degree psi cells but amplitude-sized dense phi", + False, _FRAMEWORK), _ANGLE, + requires=("angle-amplitude-estimate", + "angle-peak-local-warranted"), + conflicts=("distance:gh",)), + _profile("distance", "uniform", + _warrant(WarrantKind.NONE, "fixed uniform-in-distance grid", False, + _DISTANCE), _DISTANCE), + _profile("distance", "loguniform", + _warrant(WarrantKind.BOUNDED_STATIONARY_SET, + "interior Gaussian peak on finite distance support", + False, _DISTANCE), _DISTANCE, + requires=("angle-amplitude-estimate", "distance-full-prior", + "distance-peak-interior", + "distance-endpoint-error-ok")), + _profile("distance", "gh", + _warrant(WarrantKind.BOUNDED_STATIONARY_SET, + "support-aware per-sample distance nodes", False, + _FRAMEWORK), + "RIFT/likelihood/jax_ile/core.py:_distmarg_gh_logL", + requires=("distance-volumetric-prior",)), + _profile("time", "simpson", + _warrant(WarrantKind.NONE, "fixed native time grid", False, + _TIME), _TIME), + # The band limit is a real structural fact, so this warrant kind COULD + # support a certificate. The shipped implementation does not discharge one: + # it derives the refinement factor from a curvature-measured peak width and + # remeasures it on the dense grid, and its accuracy record is a table of + # measured nonzero reconstruction errors, not a per-request inequality on + # the marginalized log likelihood. Advertising a certificate here would let + # any caller-supplied CERTIFIED assessment enter cheapest-certified with an + # arbitrarily tight budget and no executable proof, which is exactly the + # relabeling the warrant/certificate split exists to refuse. + _profile("time", "bandlimited", + _warrant(WarrantKind.EXACT_BAND_LIMIT, + "band-limited kappa with time-independent self term", + False, _TIME), _TIME, + requires=("time-exact-band-limit", "time-independent-rho-sq", + "n-cal-one"), + conflicts=("jax-direct-nonlinear-time",)), +) + +JAX_SCHEME_PROFILES = MappingProxyType( + {profile.key: profile for profile in _JAX_PROFILE_LIST}) +JAX_DIRECT_MARGINALIZATION_AXES = ("angle", "distance", "time") +JAX_CONSERVATIVE_FALLBACK_SCHEMES = MappingProxyType({ + # These are support-covering, non-root-enumerating historical paths. The + # designation is a finite-execution role, not an error certificate. + "angle": frozenset(("exact",)), + "distance": frozenset(("uniform",)), + "time": frozenset(("simpson",)), +}) + + +def make_jax_scheme_offer(axis, scheme, accuracy, resources, *, + provenance, requires=(), provides=(), conflicts=(), + conditional_requirements=()): + """Attach measured request-specific evidence to a shipped scheme profile. + + Static incompatibilities cannot be removed here; callers may only add more + restrictive request-specific facts. This prevents an adapter from making + an unsupported combination look runnable by omission. + """ + key = "%s:%s" % (axis, scheme) + try: + profile = JAX_SCHEME_PROFILES[key] + except KeyError: + raise ValueError("unknown JAX direct-marginalization scheme %r" % key) + return SchemeOffer( + axis=axis, scheme=scheme, accuracy=accuracy, resources=resources, + warrant=profile.warrant, + provenance="%s; request evidence: %s" % ( + profile.provenance, provenance), + requires=profile.requires.union(requires), provides=provides, + conflicts=profile.conflicts.union(conflicts), + conditional_requirements=(profile.conditional_requirements + + tuple(conditional_requirements))) + + +def _validate_jax_offer_profiles(offers): + for offer in offers: + try: + profile = JAX_SCHEME_PROFILES[offer.key] + except KeyError: + raise ValueError("unknown JAX direct-marginalization offer %r" + % offer.key) + if offer.warrant != profile.warrant: + raise ValueError("%s does not carry the shipped warrant profile" + % offer.key) + if not profile.requires.issubset(offer.requires): + raise ValueError("%s omits shipped requirements %r" + % (offer.key, sorted( + profile.requires.difference(offer.requires)))) + if not profile.conflicts.issubset(offer.conflicts): + raise ValueError("%s omits shipped conflicts %r" + % (offer.key, sorted( + profile.conflicts.difference(offer.conflicts)))) + missing_conditionals = [ + requirement for requirement in profile.conditional_requirements + if requirement not in offer.conditional_requirements] + if missing_conditionals: + raise ValueError("%s omits a shipped conditional requirement" + % offer.key) + + +def make_jax_production_fallback_policy( + offers, resource_budget, *, provenance, finite_output_contract): + """Build an explicit JAX fallback from support-covering dense schemes. + + This helper deliberately accepts no root-enumerating angle scheme. It + still requires request-specific error/resource evidence through normal + offers and does not relabel the fallback as certified. + """ + offers = tuple(offers) + _validate_jax_offer_profiles(offers) + unsupported = [offer.key for offer in offers + if offer.scheme not in + JAX_CONSERVATIVE_FALLBACK_SCHEMES.get( + offer.axis, frozenset())] + if unsupported: + raise ValueError( + "schemes %r are not registered JAX conservative fallbacks" + % unsupported) + return ConservativeFallbackPolicy( + offers, resource_budget, provenance, finite_output_contract) + + +def plan_jax_direct_marginalization(offers, error_budget, resource_budget, *, + capabilities=(), allow_best_effort=False, + required_axes=None, resource_model=None): + """RIFT-specific entry point; still entirely opt-in and side-effect free. + + The static profile is rechecked here rather than trusted to the offer + builder. A caller may use :func:`plan_direct_marginalization` for an + experimental catalog, but this entry point cannot be made to forget a + shipped incompatibility by manually constructing a weaker offer. + """ + axes = (JAX_DIRECT_MARGINALIZATION_AXES if required_axes is None + else tuple(required_axes)) + offers = tuple(offers) + _validate_jax_offer_profiles(offers) + + active_capabilities = set(capabilities) + if "time" in axes and "angle" in axes: + # The JAX angle wrappers reduce through coefficient-table kernels that + # return an already-reduced lnL(t), so bandlimited has no primitive to + # refine there and _validate_nonlinear_time_quadrature refuses it. The + # distance axis is NOT in this condition: its reduction consumes the + # refined (kappa, rho^2) directly and the distance-marginalized wrapper + # applies it on the refined nodes. This is an active execution-context + # fact, not a capability callers should have to remember to declare. + active_capabilities.add("jax-direct-nonlinear-time") + return plan_direct_marginalization( + offers, error_budget, resource_budget, required_axes=axes, + capabilities=active_capabilities, + allow_best_effort=allow_best_effort, + resource_model=resource_model) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_policy.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_policy.py new file mode 100644 index 000000000..0a0ca7bbc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_policy.py @@ -0,0 +1,1450 @@ +"""Opt-in cross-axis direct-marginalization policy for the JAX ILE arm. + +``--direct-marginalization-policy auto`` composes, per likelihood evaluation, +the four-axis peak-local controller of :mod:`all_axis_peaklocal` with the +established exact-angle reserve: + +1. build the guarded coefficient tables once (the same U,V/Q contraction the + exact scheme uses), collapse the time-independent norm table per row; +2. rank a base and an enriched U,V/Q start portfolio on device, refine both in + one shared optimizer pass, and freeze two nested fixed-shape plans; +3. attempt the four-axis local integral over ``(t, phi_ref, u=2 psi, x)`` + under the empirical enrichment gate; +4. on any decline, execute the band-limited exact-angle reserve on a refined + time rule, warranted by the two-guard comparison and a half-refined check + rule evaluated in the same declined branch; +5. return the selected value per row with the complete acceptance ledger. + +No SNR threshold is coded. Which branch runs is decided by the diagnostics +listed in :func:`policy_acceptance_diagnostics`. A local decline is never a +waveform failure. A reserve that fails its own warrant is escalated once per +doubling of the time rule up to ``reserve_time_refine_max``; a row that is +still unwarranted, or whose norm table varies with time, returns ``nan``. +The driver refuses to publish a run containing such rows. The ledger keeps +the finite diagnostic value under ``selected_value`` for the record; it is +never handed to the sampler. + +Measures. The local branch integrates ``x**-4 dx dt_sample dphi du``. +The reserve, and the production exact scheme it must agree with, average the +two angles (``dphi/2pi``, ``dpsi/pi``), weight distance by the normalized +``log_w_grid`` (or the normalized volumetric measure under +``JAX_ILE_DISTMARG_GH``), and integrate time in seconds with Simpson weights. +:func:`policy_log_normalization` derives the constant that converts the local +measure to that convention; it is a derivation from the prior's stated form, +never an inferred number, and it refuses any prior it cannot derive. + +Not claimed here: derivative accuracy. The plans are frozen under +``stop_gradient``; differentiating the composite differentiates the truncated +fixed-plan local integral or the reserve. Value and gradient parity through +an SNR/HM ladder is the gate before this policy can become a default. +""" + +from typing import NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np + +from . import all_axis_peaklocal as _aap +from . import anglemarg as _anglemarg +from . import core as _core +from . import peaklocal_time_reserve as _plr + +__all__ = [ + "POLICY_CHOICES", + "POLICY_DEFAULT", + "RESERVE_SCHEME_TEST_ONLY", + "reserve_pair", + "resolve_reserve_angular_kernel", + "q_bandwidth_cycles_per_sample", + "PolicyConfig", + "validate_policy_request", + "validate_policy_config", + "policy_log_normalization", + "policy_time_rules", + "probe_guarded_tables", + "policy_acceptance_diagnostics", + "fused_log_likelihood_four_axis_policy", + "summarize_policy_ledger", + "RESERVE_SCHEME_CHOICES", + "RESERVE_SCHEME_DEFAULT", + "RESERVE_SCHEME_EXECUTABLE", + "q_effective_bandwidth_hz", + "predict_reserve_pair", + "format_reserve_pair", +] + +POLICY_CHOICES = ("off", "auto") +POLICY_DEFAULT = "off" +_LAPLACE_TABLE_KERNEL = "coefficient_table_distphipsimarg_laplace" + +RESERVE_SCHEME_CHOICES = ("auto", "exact", "laplace", "peaklocal") +RESERVE_SCHEME_DEFAULT = "exact" +# A reserve scheme names a PAIR (angular kernel, time rule). "auto" is +# resolved by predict_reserve_pair BEFORE construction; the composite refuses +# it. "peaklocal-exact" is the peak-local time rule with the exact angular +# kernel: the accuracy reference the tests need, accepted on PolicyConfig +# and not offered on the command line. +_RESERVE_PAIRS = { + "exact": ("exact", "window"), + "laplace": ("laplace", "window"), + "peaklocal": ("laplace", "peaklocal"), + "peaklocal-exact": ("exact", "peaklocal"), +} +RESERVE_SCHEME_TEST_ONLY = ("peaklocal-exact",) + +# WHICH OF THOSE THE COMPOSITE CAN ACTUALLY EXECUTE TODAY. Kept separate from +# the choices tuple, and checked in validate_policy_config, because a config +# field the composite never reads is worse than a missing one: it accepts the +# value, reports it, and runs something else. The reserve is dispatched through +# empirical_enrichment_with_exact_reserve, which names its kernel -- so 'exact' +# is the whole of what is wired. +# +# 'laplace' has its table-level kernel (coefficient_table_distphipsimarg_laplace, +# extracted from the fused laplace path in this same PR) but no dispatch: the +# selector may CHOOSE it, and a run that needs it is refused with that reason +# rather than carried by exact. 'peaklocal' belongs to RIFT PR #304. +# Executable = dispatched by the composite through reserve_pair: the exact +# and psi-Laplace kernels on the whole-window rule, and the peak-local time +# rule (peaklocal_time_reserve) under the psi-Laplace kernel. +RESERVE_SCHEME_EXECUTABLE = ("exact", "laplace", "peaklocal") + + +# The controller's own decline reasons, in the order they gate acceptance. +# Every key is a boolean per row in the returned ledger. +_DECLINE_KEYS = ( + "decline_nonfinite", + "decline_capacity", + "decline_no_modes", + "decline_boundary_maximum", + "decline_mode_nesting", + "decline_time_reconstruction", + "decline_time_cover_incomplete", + "decline_time_omitted_mass_bound", + "decline_time_omitted_mass", + "decline_geometry", + "decline_quadrature", + "decline_enrichment", + "decline_error_budget", +) + + +class PolicyConfig(NamedTuple): + """Operating point of the composite. + + These are the values PR #268's device composition test and real captures + ran with, exposed so the validation ladder can move them. They are not a + measured production operating point yet. + """ + + # Guard: the ladder record (RIFT_roboto_paper analyses/va_sequence_20260902/ + # RESULTS_20260907_aap268_ladder.md) accepted identically at guard 128 and + # 1024 on production tables; 128 is the largest the driver's default + # 0.15 s storage window supports. A guard past the stored buffer is + # refused at construction (see the wrapper's probe), not read as a decline. + time_guard: int = 128 + reserve_time_refine: int = 4 + # Bounded escalation of the reserve rule on a failed warrant: the rule is + # doubled (and re-checked against its own half) until it is warranted or + # this factor is reached. Rows still unwarranted return nan. + reserve_time_refine_max: int = 32 + # Start capacity of the local plan. Sized with max_time_nodes, not + # independently: on the same 64 rows the two capacities decline almost the + # same rows, 36 failing on time nodes and 37 on starts, with only 7 failing + # on starts alone. Raising either one by itself buys nothing. Measured + # acceptance at rho 652.3, full-sky prior draws, as (max_time_nodes, + # base_max_starts): (64, 32) 28%, (256, 32) 31%, (256, 128) 75%, + # (512, 256) 77%. The pair moves acceptance from 28% to 75%; the second + # doubling buys 2 points. + # + # The resize is free. Compiled workspace is 0.546 GiB -- 5% of a 24 GiB + # card -- and is UNCHANGED across that whole 16x capacity range, because + # peak scratch is set by the reserve's amplitude-sized grids, which are + # compiled into the graph whether or not any row reaches them. So there is + # no memory argument for a small capacity here, only a compile-time one. + base_max_starts: int = 128 + # Time-node capacity of the local plan, SET IN ADVANCE rather than + # discovered and then declined. rank_joint_starts_from_uvq_device + # defaults it to 64 and the policy never passed it, so the value that + # gated the local path was unreachable from any config field or flag. + # + # Measured, rho 652.3, full-sky prior draws, 64 rows: the live-node count + # a row needs is min 18, median 72, p90 217, max 614. A cap of 64 holds + # 44% of rows, 128 holds 73%, 256 holds 91%, 1024 holds 100%. The median + # row misses the old cap by eight nodes. + # + # 256 with base_max_starts 128, approved by RO on 2026-09-08 (evening). + # What shipped was 64/32 -- the value rank_joint_starts_from_uvq_device + # used while the policy passed nothing, reachable from no field or flag. + # Measured acceptance at rho 652 as (max_time_nodes, base_max_starts): + # (64, 32) 28%, (256, 32) 31%, (256, 128) 75%, (512, 256) 77%. They move + # together: of 64 rows, 36 declines fail on time nodes and 37 on starts, + # only 7 on starts alone, so raising either alone plateaus near 16%. + # + # The price, and it is not zero: base_max_starts 32 -> 128 shifts + # already-accepted values by up to 3.5e-3 nats at rho 163, so an A/B across + # this change may not treat its difference as noise. Accepted because a + # declined row runs the exact reserve at hours per row at rho 652 while an + # accepted one costs ~2.3 s. Workspace is unaffected -- 0.546 GiB, flat + # from (64, 32) to (1024, 256) on a 24 GiB card -- so the cost of the + # resize is compile time, 13 s to ~70 s, not memory. + max_time_nodes: int = 256 + # Angular oversample 2/4 and 16 modes are the configuration that accepted + # on production tables at rho 163 and 326 (same record as above; 8 to 12 + # candidates against 32 starts). PR #268's test values 1/2 and 4/8 + # overflowed capacity on synthetic carrier tables. + base_oversample: int = 2 + enriched_oversample: int = 4 + max_modes: int = 16 + enriched_max_modes: int = 16 + # 6 whitened sigmas, the library default. PR #268's composition test used + # 3.0; on the wiring test's analytic fixture that truncated ~1% of the + # four-dimensional mass and read 0.015 nat LOW against an independent + # fine-time reference while the gate accepted, because base and enriched + # plans share the truncation. The gate cannot see this term; the wiring + # test pins it against the external reference instead. + local_radius: float = 6.0 + refine_iterations: int = 14 + base_order: int = 13 + base_check_order: int = 19 + enriched_order: int = 19 + enriched_check_order: int = 25 + convergence_tol_nats: float = 1.0e-3 + time_guard_tol_nats: float = 1.0e-3 + total_value_error_budget_nats: float = 1.0e-2 + time_outside_tol_nats: float = -23.0 + # 64, not the kernel's 8: the reserve's reverse pass keeps one carry per + # dense-angle scan step, so gradient memory falls ~7x from 8 to 64 + # (5.3 -> 0.73 GiB at refine 2, 21 -> 2.9 GiB at refine 8 on a rho 163 + # production row). The value is unchanged; per-step forward memory grows + # with the chunk. + reserve_dense_chunk: int = 64 + reserve_grid_block: int = 32 + # Rows the controller executes together under one ``vmap``. 1 is the + # row-at-a-time path (``lax.map`` with no ``batch_size``), whose reserve + # workspace is one row's. Above 1, ``B`` rows share a scan step, device + # workspace grows linearly in ``B``, and the tier-escalation ``lax.cond`` + # becomes a ``select`` that evaluates EVERY tier for EVERY row in the + # batch. Values, branch decisions and gradients are unchanged at every + # size. 0 means one full batch of all rows. + # + # The default is 1 because batching is a COST REGRESSION. The + # tier-escalation cond is not the only one: the accept/reserve cond in + # ``all_axis_peaklocal`` also becomes a select, so a locally ACCEPTED row + # executes the dense reserve it would otherwise skip. The penalty scales + # with the locally accepted fraction, and the docstring there claiming an + # accepted row never pays for the reserve is false above B=1. + # + # WITHDRAWN: this comment previously read "measured and does not pay", + # citing 96.3 s per row at B=1 against 99.9 at B=8 on ladder-2 tables at + # rho 40.8. Every row in those runs DECLINED (accepted_local 0 of 8 and 0 + # of 2), so the figures price the decline path, and 37-50% of the rows were + # nan under reserve_time_refine_max=4. They also ran at refine == refine + # _max, where the escalation cond is absent from the graph, so both penalty + # mechanisms were inert. Do not cite them. What stands from that work is + # the workspace law, 0.046 + 0.385 B GiB, and the value equivalence. + # See DESIGN_direct_marginalization_policy.md. + reserve_batch_rows: int = 1 + # WHICH reserve the composite falls back to. Default 'exact' is what + # shipped, so no existing command line changes. 'laplace' is the same + # accuracy crossover the angle selector already validates + # (ANGLE_MARG_CROSSOVER_AMPLITUDE): above it laplace is the MORE accurate + # scheme and costs ~sqrt(A) rather than ~A. 'auto' chooses from the + # precomputed inputs via predict_reserve_pair and REFUSES when the analysis + # says the signal needs a method that is not implemented, rather than + # falling back to whole-window refinement. A string so a third value plugs + # in without touching the composite. + reserve_scheme: str = RESERVE_SCHEME_DEFAULT + norm_invariance_rtol: float = 1.0e-10 + # Peak-local time rule (reserve schemes 'peaklocal', 'peaklocal-exact'): + # fine nodes per plan mode block (49 = 8 sigma either side at 3 nodes per + # predicted sigma), the coarse scan refinement of the window, the number + # of warrant escalations (each doubles the fine lattice at fixed span), + # and an override of the predicted width (nan = predict per row from + # rho and the Q bandwidth; an experiment knob, printed in the ledger). + # Measured operating point: DESIGN_direct_marginalization_policy.md, + # "Peak-local time reserve". + reserve_peaklocal_fine_nodes: int = 73 + reserve_peaklocal_scan_refine: int = 2 + reserve_peaklocal_escalations: int = 2 + # The row's time maxima come from the primitive (locate_time_maxima): + # this many candidates, on a search grid of this many nodes per native + # sample, over this angular lattice. The local branch's plan is only a + # cross-check in the ledger. + reserve_peaklocal_blocks: int = 4 + reserve_peaklocal_search_refine: int = 8 + reserve_peaklocal_angular_lattice: int = 8 + # The coarse scan is support-limited: this many nodes across the hull of + # the live maxima widened by this many predicted sigma each side; the + # mass outside is bounded from the locator's search profile plus this + # slack and charged to the warrant. Node count never grows with the + # window (RO, 2026-09-09). + reserve_peaklocal_scan_nodes: int = 65 + reserve_peaklocal_scan_margin_sigmas: float = 16.0 + # Locator search sizing (measured 2026-09-09, rung 652 row 0): the search + # grid's phi lattice leaves a ripple of about rho^2 (pi / n_phi)^2 nat in + # the profile, 1000 nat at rho 652 on 64 nodes against a 25-nat time + # structure per search cell, so the search maximum landed 0.4 samples off + # and the focus certificate refused the row. 4096 nodes hold the ripple + # under 1 nat up to rho 1300; the count is a static shape, so it is a + # config field rather than a per-row prediction. The Newton polish of + # (phi, u) at each polished node was clipped to 0.1 rad per step for 3 + # steps, a 0.3 rad reach against a 0.4 rad lattice offset: 500 nat low on + # the same row. Eight steps within a 1 rad trust region reach it. + reserve_peaklocal_search_phi_nodes: int = 4096 + reserve_peaklocal_newton_steps: int = 8 + reserve_peaklocal_newton_step_max: float = 1.0 + reserve_peaklocal_outside_slack_nats: float = 5.0 + reserve_peaklocal_sigma_t_override_samples: float = float("nan") + + +def validate_batch_rows(batch_rows): + """Refuse a row-batch size the controller cannot execute. + + Returns the integer. Negative sizes and non-integers are refused here so + the driver and the library agree on the same rule. + """ + try: + b = int(batch_rows) + except (TypeError, ValueError): + raise ValueError("PolicyConfig.reserve_batch_rows must be an integer, " + "got %r" % (batch_rows,)) + if b != batch_rows: + raise ValueError("PolicyConfig.reserve_batch_rows must be an integer, " + "got %r" % (batch_rows,)) + if b < 0: + raise ValueError("PolicyConfig.reserve_batch_rows must be >= 0 " + "(1 = row at a time, 0 = one full batch), got %d" % b) + return b + + +def reserve_pair(scheme): + """``(angular_kernel, time_rule)`` named by a reserve scheme. + + ``auto`` is not a pair: it is resolved by :func:`predict_reserve_pair` + before the likelihood is built, and the composite refuses it. + """ + if scheme not in _RESERVE_PAIRS: + if scheme == "auto": + raise ValueError( + "reserve_scheme='auto' must be resolved by predict_reserve_pair " + "before construction; the composite takes a concrete scheme") + raise ValueError("reserve scheme must be one of %r (plus %r for tests), " + "got %r" % (tuple(k for k in RESERVE_SCHEME_CHOICES + if k != "auto"), + RESERVE_SCHEME_TEST_ONLY, scheme)) + return _RESERVE_PAIRS[scheme] + + +def q_bandwidth_cycles_per_sample(data): + """:func:`q_effective_bandwidth_hz` in cycles per native sample, or nan + when the data carries no stored Q (synthetic tables).""" + try: + hz = float(q_effective_bandwidth_hz(data)) + except (AttributeError, KeyError, TypeError): + return float("nan") + return hz * float(data.deltaT) if np.isfinite(hz) else float("nan") + + +def resolve_reserve_angular_kernel(name, x_grid, log_w_grid, *, amp_sizing, + m_max, dense_chunk, grid_block): + """``None`` for the kernel's own exact default, else a table callable. + + The psi-Laplace table kernel is provided by ``anglemarg`` under the name + ``coefficient_table_distphipsimarg_laplace`` (the --direct-marginalization- + reserve-scheme laplace work); a tree without it refuses the request here, + at construction, rather than at trace time. + """ + if name not in ("exact", "laplace"): + raise ValueError("angular kernel must be exact or laplace, got %r" % (name,)) + if name == "exact": + return None + fn = getattr(_anglemarg, _LAPLACE_TABLE_KERNEL, None) + if fn is None: + raise ValueError( + "the laplace angular kernel needs anglemarg.%s, which this tree " + "does not provide; use reserve scheme exact or peaklocal-exact" + % _LAPLACE_TABLE_KERNEL) + x_grid = jnp.asarray(x_grid, dtype=jnp.float64) + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64) + + # ``dense_chunk`` and ``grid_block`` are the EXACT kernel's streaming knobs; + # the Laplace kernel streams by ``phi_chunk``/``dist_block``/``point_block`` + # and rejects them. They stay in this signature so the caller is uniform, + # and are deliberately not forwarded (measured 2026-09-09: the first + # psi-Laplace reserve evaluation at rung 652 died on the keyword). + del dense_chunk, grid_block + + def kernel(table, norm_table): + return fn(table, norm_table, x_grid, log_w_grid, + amp_sizing=float(amp_sizing), m_max=int(m_max)) + return kernel + + +def validate_policy_config(config): + """Refuse a PolicyConfig the composite would only reject at trace time.""" + if not isinstance(config, PolicyConfig): + raise TypeError("policy_config must be a PolicyConfig") + if int(config.time_guard) < 2: + raise ValueError("PolicyConfig.time_guard must be >= 2: the local " + "path and the reserve both need the two-guard " + "comparison") + f, fm = int(config.reserve_time_refine), int(config.reserve_time_refine_max) + if f < 2 or f % 2: + raise ValueError("reserve_time_refine must be an even integer >= 2 " + "so the check rule is the half-refined rule") + if fm < f or fm % 2: + raise ValueError("reserve_time_refine_max must be an even integer >= " + "reserve_time_refine") + if not (np.isfinite(float(config.total_value_error_budget_nats)) + and float(config.total_value_error_budget_nats) > 0.0): + raise ValueError("total_value_error_budget_nats must be finite and " + "positive") + if int(config.max_time_nodes) < 2: + raise ValueError("PolicyConfig.max_time_nodes must be >= 2: the time " + "cover needs at least one interior node pair") + if int(config.base_max_starts) < 1: + raise ValueError("PolicyConfig.base_max_starts must be >= 1") + if int(config.base_oversample) < 1 or int(config.enriched_oversample) <= int( + config.base_oversample): + raise ValueError("enriched_oversample must exceed base_oversample " + "(the enriched portfolio must be strictly stronger)") + if int(config.max_modes) < 1 or int(config.enriched_max_modes) < int( + config.max_modes): + raise ValueError("enriched_max_modes must be >= max_modes >= 1") + if not float(config.local_radius) > 0.0: + raise ValueError("local_radius must be positive") + validate_batch_rows(config.reserve_batch_rows) + scheme = config.reserve_scheme + if scheme not in RESERVE_SCHEME_CHOICES + RESERVE_SCHEME_TEST_ONLY: + raise ValueError("PolicyConfig.reserve_scheme must be one of %r, got %r" + % (RESERVE_SCHEME_CHOICES, scheme)) + # 'auto' is resolved by predict_reserve_pair against the precomputed inputs, + # before this config ever reaches the composite; it is not a value the + # composite executes, so it is admitted here and refused there if the + # analysis lands on something unimplemented. + if (scheme != "auto" and scheme not in RESERVE_SCHEME_EXECUTABLE + and scheme not in RESERVE_SCHEME_TEST_ONLY): + raise ValueError( + "PolicyConfig.reserve_scheme=%r is a declared choice but is NOT " + "WIRED into the composite: the reserve is dispatched through " + "empirical_enrichment_with_exact_reserve and only %r is executable " + "today. Refused rather than run as 'exact', which is what a " + "silently ignored field would do -- the run would report the " + "scheme you asked for and compute the other one." + % (scheme, RESERVE_SCHEME_EXECUTABLE)) + if int(config.max_time_nodes) < 2: + raise ValueError("max_time_nodes must be >= 2") + if scheme == "auto": + return config + angular, time_rule = reserve_pair(scheme) + if (angular == "laplace" + and getattr(_anglemarg, _LAPLACE_TABLE_KERNEL, None) is None): + raise ValueError( + "reserve scheme %r needs anglemarg.%s, which this tree does not " + "provide" % (config.reserve_scheme, _LAPLACE_TABLE_KERNEL)) + if time_rule == "peaklocal": + _plr.validate_peaklocal_rule_arguments( + 2, config.reserve_peaklocal_fine_nodes, + config.reserve_peaklocal_scan_refine) + if int(config.reserve_peaklocal_escalations) < 0: + raise ValueError("reserve_peaklocal_escalations must be >= 0") + if (int(config.reserve_peaklocal_blocks) < 1 + or int(config.reserve_peaklocal_search_refine) < 1 + or int(config.reserve_peaklocal_angular_lattice) < 2): + raise ValueError("reserve_peaklocal_blocks >= 1, search_refine >= 1 " + "and angular_lattice >= 2 are required") + ns = int(config.reserve_peaklocal_scan_nodes) + if ns < 3 or ns % 2 == 0: + raise ValueError("reserve_peaklocal_scan_nodes must be an odd integer >= 3") + if int(config.reserve_peaklocal_search_phi_nodes) < 4: + raise ValueError("reserve_peaklocal_search_phi_nodes must be >= 4") + if int(config.reserve_peaklocal_newton_steps) < 1: + raise ValueError("reserve_peaklocal_newton_steps must be >= 1") + if not (float(config.reserve_peaklocal_newton_step_max) > 0.0): + raise ValueError("reserve_peaklocal_newton_step_max must be positive") + if not (float(config.reserve_peaklocal_scan_margin_sigmas) > 0.0 + and float(config.reserve_peaklocal_outside_slack_nats) >= 0.0): + raise ValueError("reserve_peaklocal_scan_margin_sigmas must be positive " + "and reserve_peaklocal_outside_slack_nats >= 0") + ov = float(config.reserve_peaklocal_sigma_t_override_samples) + if np.isfinite(ov) and ov <= 0.0: + raise ValueError("reserve_peaklocal_sigma_t_override_samples must " + "be positive or nan") + return config + + +# --------------------------------------------------------------------------- +# Analysis-driven reserve selection +# --------------------------------------------------------------------------- +# RO, 2026-09-08: "we are learning a hard lesson about refinement and the +# 'reserve' not protecting us; we need to rely on ANALYSIS and the known physics +# ... have a hierarchy of methods and pick the expected bounding pairs as needed +# that apply to our signal." +# +# The failure mode being named is try-then-decline-then-refine: run the local +# branch, discover it declined, escalate a whole-window refinement, and discover +# at the END that most rows were carried by a method nobody chose. Everything +# below is computable from the PRECOMPUTED inputs before any row is evaluated, +# so the pair is chosen and PRINTED up front and the run can be read in its +# first line instead of its last. + +# PROVISIONAL, AND KNOWN TO BE THE WRONG MODEL FOR THE RESERVE. +# +# The points-per-sigma budget below treats the reserve's trapezoid rule as +# ALGEBRAICALLY convergent, so the node count it demands scales as +# window / sigma_t. A direct test at rho 40.77 on 64 rows says otherwise: +# +# refine=4 2453 nodes warrant 1.8e-03 .. 1.14e-02 +# refine=8 4905 nodes warrant 5e-11 .. 7.8e-09 +# +# The tolerance those are read against has MOVED (#301: 1e-3 -> 1e-2), so the +# numbers are recorded without a verdict: refine=4 fails 1e-3 and straddles +# 1e-2. See DESIGN_direct_marginalization_policy.md. +# +# Doubling the rule improved the quadrature error by ~1e6. An algebraic rule +# would give 4. That is the signature of the trapezoid rule on a BAND-LIMITED +# reconstruction, which is spectrally accurate once the band is resolved: +# error ~ exp(-c R), not R^-2. The measured 4905 nodes is 67% of what this +# budget demands at that rung and lands five orders INSIDE tolerance. +# +# So the correct criterion is band resolution -- node spacing against the +# integrand's highest frequency -- not points per sigma, and the replacement +# must be FITTED to a measured convergence law rather than assumed. Until a +# second rung is measured (163.08 at refine 4 and 8 is the deciding test), the +# time verdict below is provisional and MUST NOT be hardened into a threshold +# anyone tunes against. It is retained because refusing is the conservative +# direction, but a refusal it produces is "unproven", not "shown inadequate". +# +# Fraction of a peak sigma the local branch's time cover must resolve. The +# cover keeps cells above a mass threshold, so a peak narrower than the node +# spacing puts its mass in one cell and the cover cannot localize it. +_TIME_NODES_PER_SIGMA = 3.0 + + +def q_effective_bandwidth_hz(data, moment="raw"): + """Bandwidth of the stored Q(t), in Hz. TWO different quantities. + + ``moment="raw"`` (default, and THE one to size a time rule with) returns + sqrt() over the full two-sided spectrum. + + ``moment="central"`` returns the RMS bandwidth about the mean over positive + frequencies -- the ENVELOPE bandwidth. + + WHICH ONE, AND WHY IT IS PHYSICS RATHER THAN CONVENTION. The reserve + marginalizes phi exactly, so the field in time is |zeta| with + zeta = alpha kappa + beta kappa*, kappa = E(t) e^{i theta}, theta ~ 2 pi f_c t, + and alpha, beta the polarization weights. Then + + |zeta|^2 = (|alpha|^2 + |beta|^2) E^2 + 2 Re(alpha beta* E^2 e^{2 i theta}) + + Face-on (beta -> 0) the carrier term vanishes, the field is the envelope, and + the peak width is the envelope one -- the CENTRAL moment. Linearly polarized + (|alpha| = |beta|) the envelope is modulated at the carrier and each sub-peak + is far narrower -- the RAW moment. Every real row lies between, set by its + own psi and inclination. + + So raw is the NARROWEST peak the primitive can produce at that amplitude and + central the WIDEST. A rule that must not under-resolve, and a selector that + must not call a whole-window rule adequate when it is not, must size on the + NARROW one. Measured on a carrier fixture (f_c = 200 Hz, Gaussian envelope): + circular gives a peak of 11.3 Hz equivalent against a central moment of + 5.6 Hz; linear gives 309.6 Hz against a raw moment of 200.1 Hz. + + RAW IS EXACT FOR THE INTEGRAND, NOT sqrt(2) OPTIMISTIC. An earlier revision + of this docstring claimed the latter, from measuring the curvature of + |zeta|^2 itself. That is not the integrand. The quadrature integrates + exp(lnL) with lnL = (rho^2/2) |zeta_hat|^2, so in the linear limit + |zeta_hat|^2 = cos^2(omega t) ~ 1 - omega^2 t^2 gives + lnL ~ const - (rho^2/2) omega^2 t^2 and hence sigma_t = 1/(rho omega) + = 1/(2 pi rho f_c) exactly -- the raw moment, no factor. The sqrt(2) + appears only if the curvature of |zeta|^2 is read as a Gaussian width + WITHOUT the rho^2/2 prefactor; the two multiply. Verified against the + log-integrand on a carrier fixture over four decades of rho: + measured/predicted = 1.0008, 1.0000, 0.9999, 0.9999 at rho 12.65, 40.77, + 163.08, 652.31. + + Both are returned by name so the two can never be silently confused. An + earlier revision of this function made central the default and fed it to + sigma_t, which is the optimistic error: it predicts the envelope width for + rows whose likelihood is actually carrier-modulated. + + Q is stored at deltaT / q_time_pregrid_factor, so the frequency axis uses the + REFINED spacing; using deltaT would understate the bandwidth by that factor. + """ + import numpy as _np + if moment not in ("central", "raw"): + raise ValueError("moment must be 'central' or 'raw', got %r" % (moment,)) + num1 = num2 = den = 0.0 + for det in data.detector_names: + d = data.detectors[det] + Q = _np.asarray(d["Q"]) # (npts_full, K) + f_ref = int(d.get("q_time_pregrid_factor", 1)) or 1 + dt = float(data.deltaT) / float(f_ref) + n = Q.shape[0] + if n < 4: + continue + # Q (rholm) is COMPLEX, so this is the full two-sided transform: a + # real-input transform rejects it outright, and dropping the imaginary + # part would discard half the phase structure. + freqs = _np.fft.fftfreq(n, d=dt) + spec = _np.abs(_np.fft.fft(Q, axis=0)) ** 2 # (n, K) + w = spec.sum(axis=1) + if moment == "central": + keep = freqs > 0.0 + freqs, w = freqs[keep], w[keep] + num1 += float((freqs * w).sum()) + num2 += float((freqs ** 2 * w).sum()) + den += float(w.sum()) + if den <= 0.0: + return float("nan") + m1, m2 = num1 / den, num2 / den + if moment == "raw": + return float(_np.sqrt(max(m2, 0.0))) + return float(_np.sqrt(max(m2 - m1 * m1, 0.0))) + + +def predict_reserve_pair(data, guess_snr, *, reserve_time_refine_max, + crossover_amplitude, max_time_nodes, + requested="auto", available=("exact", "laplace")): + """Choose the (local, reserve) pair from the precomputed inputs. + + Returns ``(scheme, info)``. ``scheme`` is None when the analysis says the + signal needs a method that is not implemented; the caller must REFUSE rather + than fall back, which is the whole point of predicting. + + The four quantities, all available before any row is evaluated: + + * ``rho`` -- the network SNR guess the driver already computes from the + detector response. + * ``sigma_f`` -- the effective bandwidth of the stored Q (above). + * ``A = rho^2 / 2`` against ``ANGLE_MARG_CROSSOVER_AMPLITUDE``, the + selector's own VALIDATED accuracy crossover between exact and laplace + angles. Above it laplace is the more accurate scheme AND costs ~sqrt(A) + rather than ~A. + * ``sigma_t = 1 / (2 pi rho sigma_f)`` -- the expected width of the time + peak, in native samples, against what the whole-window refined reserve + can AFFORD: ``(npts - 1) * reserve_time_refine_max + 1`` nodes. A peak + the escalation ceiling cannot resolve is the regime where refinement + carries the rows without resolving them, which is what this function + exists to predict rather than discover at the end of a run. + + NOTE the comparison is deliberately against the RESERVE's node budget and + NOT against ``max_time_nodes``. The latter caps the LOCAL branch's time + CELL COVER, which is a different quantity: the cover exceeds its budget + when lnL(t) is BROAD (many live cells, far from truth), not when the peak + is narrow. Comparing a whole-window node count against a cover budget + mixes the two, which an earlier draft of this function did. + """ + import numpy as _np + rho = float(guess_snr) if guess_snr else float("nan") + # RAW: the narrowest peak the primitive can make, which is what a time + # rule must resolve. Central is reported as the envelope bandwidth. + sigma_f = q_effective_bandwidth_hz(data, moment='raw') + sigma_f_env = q_effective_bandwidth_hz(data, moment='central') + A = 0.5 * rho * rho if _np.isfinite(rho) else float("nan") + dt = float(data.deltaT) + if _np.isfinite(rho) and _np.isfinite(sigma_f) and rho > 0 and sigma_f > 0: + sigma_t = 1.0 / (2.0 * _np.pi * rho * sigma_f) + else: + sigma_t = float("nan") + width_samples = sigma_t / dt if _np.isfinite(sigma_t) else float("nan") + # Nodes the cover would need to put _TIME_NODES_PER_SIGMA across one sigma + # over the whole window, which is what a whole-window rule has to do. + if _np.isfinite(width_samples) and width_samples > 0: + nodes_needed = _TIME_NODES_PER_SIGMA * float(data.npts) / width_samples + else: + nodes_needed = float("inf") + nodes_available = (float(data.npts) - 1.0) * float( + reserve_time_refine_max) + 1.0 + time_reserve_ok = nodes_needed <= nodes_available + + # FIRST question, and the one about the branch actually under test: can the + # LOCAL cover hold the peak? Its budget is max_time_nodes, and a peak + # narrower than a native sample needs the cover to place its nodes inside + # one sample rather than across the window. MEASURED, not modelled: at + # rho 163 raising the cover 64 -> 256 took acceptance 31% -> 75%, which is + # why the start cap appeared to "plateau" -- the plateau was time capacity + # binding, not the start cap saturating. + if _np.isfinite(width_samples) and width_samples > 0: + cover_nodes_needed = _TIME_NODES_PER_SIGMA / width_samples + else: + cover_nodes_needed = float("inf") + local_cover_ok = cover_nodes_needed <= float(max_time_nodes) + time_local_ok = time_reserve_ok + + info = dict(rho=rho, sigma_f_hz=sigma_f, + sigma_f_envelope_hz=sigma_f_env, amplitude_A=A, + crossover_amplitude=float(crossover_amplitude), + sigma_t_s=sigma_t, peak_width_samples=width_samples, + cover_nodes_needed=cover_nodes_needed, + max_time_nodes=int(max_time_nodes), + local_cover_resolves_peak=bool(local_cover_ok), + whole_window_nodes_needed=nodes_needed, + whole_window_nodes_available=nodes_available, + reserve_time_refine_max=int(reserve_time_refine_max), + time_peak_resolvable_whole_window=bool(time_local_ok), + requested=requested, available=tuple(available)) + + if requested != "auto": + # An explicit request overrides the ANALYSIS. It does not override the + # ROSTER: `available` says which schemes this data and this distance + # quadrature can support at all, and forcing a scheme whose premise is + # absent is not an override, it is an unnoticed wrong answer. + if requested not in available: + info["reason"] = ("explicit request %r is not on the roster for " + "this run (available: %s); the roster is a " + "property of the data and the distance " + "quadrature, not of the analysis, so it is not " + "overridable" + % (requested, ", ".join(available))) + return None, info + info["reason"] = "explicit request, no analysis applied" + return requested, info + + if not _np.isfinite(A) or not _np.isfinite(sigma_f): + info["reason"] = ("cannot analyse: rho=%r sigma_f=%r; refusing rather " + "than guessing" % (rho, sigma_f)) + return None, info + + # Angles: the validated accuracy crossover. + angular = "laplace" if A > float(crossover_amplitude) else "exact" + + # Time: if a whole-window rule cannot resolve the peak within the local + # branch's node budget, the pair needs a peak-local time reserve. Saying so + # and refusing is the point; falling back to refinement is what RO is + # calling the hard lesson. + if not time_local_ok: + if "peaklocal" in available: + info["reason"] = ( + "local cover %s (needs %.1f nodes of %d); " + "A=%.4g > crossover %.4g selects laplace angles; peak is %.3g " + "native samples wide and a whole-window rule would need %.0f " + "nodes but the escalation ceiling affords %.0f, so the time " + "reserve must be peak-local" + % ("holds the peak" if local_cover_ok else "CANNOT hold the peak", + cover_nodes_needed, max_time_nodes, + A, crossover_amplitude, + width_samples, nodes_needed, nodes_available)) + return "peaklocal", info + info["reason"] = ( + "local cover %s (needs %.1f nodes of %d); " + "peak is %.3g native samples wide; a whole-window time rule would " + "need %.0f nodes but the escalation ceiling affords only %.0f, so this signal needs a " + "peak-local-in-time reserve, which is NOT IMPLEMENTED. Refusing " + "rather than falling back to whole-window refinement, which would " + "carry the rows without anyone choosing it." + % ("holds the peak" if local_cover_ok else "CANNOT hold the peak", + cover_nodes_needed, max_time_nodes, + width_samples, nodes_needed, nodes_available)) + return None, info + + # The roster is checked HERE too, not only on the peak-local branch. It was + # not, and the effect was that a run with laplace off the roster still + # selected laplace whenever A cleared the crossover: the caller's roster was + # honoured for the scheme it could not have chosen anyway and ignored for + # the one it could. + if angular not in available: + info["reason"] = ( + "the accuracy crossover selects %s angles (A=%.4g against %.4g), " + "and %s is NOT on the roster for this run (available: %s). " + "Refusing rather than running the other scheme under the selected " + "one's name." + % (angular, A, crossover_amplitude, angular, ", ".join(available))) + return None, info + + info["reason"] = ( + "local cover %s (needs %.1f nodes of %d at this width); " + "A=%.4g against crossover %.4g selects %s angles; peak is %.3g native " + "samples wide and a whole-window rule needs %.0f nodes within the %.0f " + "the ceiling affords, so the refined whole-window time reserve is adequate" + % ("holds the peak" if local_cover_ok else "CANNOT hold the peak", + cover_nodes_needed, max_time_nodes, + A, crossover_amplitude, angular, width_samples, nodes_needed, + nodes_available)) + return angular, info + + +def format_reserve_pair(scheme, info): + """One line for the run log, printed BEFORE any row is evaluated.""" + return ("RESERVE-PAIR local=four-axis reserve=%s rho=%.4g sigma_f=%.4gHz " + "A=%.4g crossover=%.4g peak=%.3gsamples cover_needs=%.1f/%d " + "reserve_needs=%.0f/%.0f :: %s" + % (scheme if scheme else "REFUSED", info.get("rho", float("nan")), + info.get("sigma_f_hz", float("nan")), + info.get("amplitude_A", float("nan")), + info.get("crossover_amplitude", float("nan")), + info.get("peak_width_samples", float("nan")), + info.get("cover_nodes_needed", float("nan")), + info.get("max_time_nodes", 0), + info.get("whole_window_nodes_needed", float("nan")), + info.get("whole_window_nodes_available", 0), + info.get("reason", ""))) + + +def validate_policy_request(policy, *, angle_marg_scheme, time_quadrature, + d_prior, dist_grid, reserve_scheme=None): + """Refuse every combination the composite cannot honour. + + Refusal is explicit because an ignored request on this arm has a history + of reading as a result (see the ``--angle-marg-scheme`` notes). + """ + if policy not in POLICY_CHOICES: + raise ValueError("direct_marginalization_policy must be one of %r, " + "got %r" % (POLICY_CHOICES, policy)) + if policy == "off": + return + if reserve_scheme is not None and reserve_scheme not in RESERVE_SCHEME_CHOICES: + raise ValueError( + "direct-marginalization reserve scheme must be one of %r, got %r" + % (RESERVE_SCHEME_CHOICES, reserve_scheme)) + if angle_marg_scheme != "exact": + raise ValueError( + "--direct-marginalization-policy auto needs the exact-angle " + "reserve: the resolved --angle-marg-scheme is %r, and this policy " + "does not compose with grid, laplace, peak-local or phi-local. " + "Use --angle-marg-scheme exact." % (angle_marg_scheme,)) + if time_quadrature != "simpson": + raise ValueError( + "--direct-marginalization-policy auto owns the time integral " + "(local four-axis or refined band-limited reserve) and only " + "composes with the simpson terminal rule as its check rule; " + "got %r." % (time_quadrature,)) + if d_prior not in ("euclidean", "volumetric"): + raise ValueError( + "--direct-marginalization-policy auto derives its local measure " + "from the volumetric distance prior p(d) ~ d^2 only; got %r. " + "The cosmological prior is out of scope for the composite." + % (d_prior,)) + if dist_grid != "uniform": + raise ValueError( + "--direct-marginalization-policy auto reads the reserve's distance " + "normalization off a uniform-in-d grid; --distance-grid-scheme %r " + "is not supported by the composite." % (dist_grid,)) + + +def policy_log_normalization(data, x_grid, log_w_grid, *, d_prior="euclidean", + gh_nodes=None): + """Constant converting the local ``x**-4 dx dt_sample dphi du`` integral to + the reserve convention. Returns ``(local_log_normalization, info)``. + + * angles: the reserve averages, so ``-2 log(2 pi)``; + * time: sample units to seconds, scaled by whatever constant the + production Simpson weights carry (``sum w_t == (npts-1) deltaT`` for the + plain rule; the ratio is measured rather than assumed); + * distance, fixed grid: ``p(d) dd = d^2 dd / N`` with + ``d = Dref/x`` gives ``Dref^3 x^-4 dx / N``; ``N`` is recovered from the + first weight, ``N = d_0^2 |d_1 - d_0| / w_0``, which is exact for the + uniform grid the request validator requires; + * distance, adaptive GH: the built-in normalized volumetric measure, + ``3 / (x_min^-3 - x_max^-3)``. + """ + if d_prior not in ("euclidean", "volumetric"): + raise ValueError("policy_log_normalization supports the volumetric " + "prior only, got %r" % (d_prior,)) + x = np.asarray(x_grid, dtype=float) + lw = np.asarray(log_w_grid, dtype=float) + if x.ndim != 1 or x.size < 2 or lw.shape != x.shape: + raise ValueError("x_grid/log_w_grid must be matching 1-D grids") + if gh_nodes is None: + gh_nodes = int(_core._DISTMARG_GH_N) + deltaT = float(data.deltaT) + npts = int(data.npts) + w_t = np.asarray(data.w_t, dtype=float) + plain = (npts - 1) * deltaT + time_scale = float(np.sum(w_t)) / plain + log_time = np.log(deltaT) + np.log(time_scale) + log_angles = -2.0 * np.log(2.0 * np.pi) + if int(gh_nodes) > 0: + x_min, x_max = float(np.min(x)), float(np.max(x)) + log_dist = np.log(3.0) - np.log(x_min ** -3 - x_max ** -3) + dist_mode = "gh-volumetric" + else: + dref = float(data.distMpcRef) + d = dref / x + dd = np.abs(d[1] - d[0]) + if not np.allclose(np.abs(np.diff(d)), dd, rtol=1.0e-8, atol=0.0): + raise ValueError("policy_log_normalization needs a uniform-in-d " + "distance grid") + norm = d[0] ** 2 * dd / np.exp(lw[0]) + log_dist = 3.0 * np.log(dref) - np.log(norm) + dist_mode = "fixed-grid-volumetric" + total = float(log_angles + log_time + log_dist) + info = dict(log_angles=float(log_angles), log_time=float(log_time), + log_distance=float(log_dist), distance_mode=dist_mode, + time_weight_scale=float(time_scale), + local_log_normalization=total) + return total, info + + +def _refined_rule(npts, deltaT, refine, scale): + """Trapezoid rule on the ``refine``-times finer grid, in seconds. + + Trapezoid, not Simpson: on a peak narrower than the node spacing Simpson's + alternating weights alias at half the spacing (review of PR #278 measured + 0.03 to 1.8 nat at refine 4 for peaks of 0.05 to 0.2 native samples), while + the trapezoid rule converges exponentially on a smooth peak as the spacing + shrinks, so a passed half-rule check means what it says. + """ + n_nodes = (npts - 1) * refine + 1 + h = deltaT / float(refine) + nodes = np.arange(n_nodes, dtype=float) / float(refine) + nodes[-1] = float(npts - 1) + weights = np.full(n_nodes, h) + weights[0] = weights[-1] = 0.5 * h + return nodes, weights * scale + + +def probe_guarded_tables(data, interp, guard, n_ra=6, decs=(-1.0, 0.0, 1.0)): + """Refuse a guard the stored data buffer cannot supply. + + ``core._guarded_window`` gathers from ``-guard`` to ``npts+guard-1``; + samples the build never stored come back nonfinite with no error, and a + nonfinite table empties the start plan and reads as a method decline + (ladder record, aap268_ladder README). The reachable guard depends on the + storage window and on the per-detector arrival offsets, which vary with the + sky position, so the probe sweeps a coarse sky grid at construction and + raises with the remedy if any table is nonfinite. This is a preflight, + not a certificate: the per-row ``tables_finite`` flag still gates every + evaluation. + """ + ra = jnp.asarray(np.tile(np.linspace(0.0, 2.0 * np.pi, int(n_ra), + endpoint=False), len(decs))) + dec = jnp.asarray(np.repeat(np.asarray(decs, dtype=float), int(n_ra))) + incl = jnp.full(ra.shape, 0.5 * np.pi) + C_A, C_B, _ = _anglemarg.angle_coefficient_tables( + data, ra, dec, incl, interp, guard=int(guard)) + finite = bool(jnp.all(jnp.isfinite(C_A)) and jnp.all(jnp.isfinite(C_B))) + if not finite: + raise ValueError( + "direct-marginalization policy: the guarded coefficient tables are " + "not finite at time_guard=%d for this build. The guard gathers " + "%d samples beyond each end of the %d-sample window, and the " + "stored data buffer (--internal-data-storage-window-half, minus " + "the per-detector arrival offsets) does not reach that far. " + "Lower --direct-marginalization-time-guard or widen the storage " + "window; a nonfinite table is not a likelihood decline." + % (int(guard), int(guard), int(data.npts))) + return True + + +def policy_time_rules(data, refine): + """Refined reserve rule and its coarser check rule on the target window. + + Positions are in native samples of the unguarded window, ``0 .. npts-1``. + Weights are trapezoid weights in seconds, carrying the same constant as + the production ``data.w_t`` (so the reserve lands in production units + without a separate offset). The reserve rule refines the native cadence ``refine`` + times; the check rule refines it ``refine/2`` times (the native production + rule itself when ``refine == 2``). Agreement between the two is the + resolution warrant, so the warrant is a convergence statement about the + refined rules and does not require the native rule to be converged. + """ + refine = int(refine) + if refine < 2 or refine % 2: + raise ValueError("reserve_time_refine must be an even integer >= 2 " + "so the check rule is the half-refined rule") + npts = int(data.npts) + deltaT = float(data.deltaT) + w_t = np.asarray(data.w_t, dtype=float) + scale = float(np.sum(w_t)) / ((npts - 1) * deltaT) + nodes, weights = _refined_rule(npts, deltaT, refine, scale) + if refine == 2: + check_nodes, check_weights = np.arange(npts, dtype=float), w_t + else: + check_nodes, check_weights = _refined_rule( + npts, deltaT, refine // 2, scale) + return (jnp.asarray(nodes), jnp.asarray(weights), + jnp.asarray(check_nodes), jnp.asarray(check_weights)) + + +def policy_acceptance_diagnostics(): + """Names of the per-row booleans that must all hold for local acceptance, + then the reserve warrant flags. Documentation and audit order only.""" + return dict( + local=("tables_finite", "norm_time_invariant", "base_capacity_ok", + "enriched_capacity_ok", "boundary_maximum_ok", + "base_and_enriched_values_finite", + "mode_nesting_ok", "geometry_nesting_ok", + "time_omitted_mass_ok", "value_error_budget_ok", + "accepted_local"), + reserve=("reserve_executed", "reserve_finite", + "reserve_time_guard_validated", + "reserve_time_resolution_validated", + "reserve_time_error_budget_ok", "reserve_time_warranted"), + declines=_DECLINE_KEYS) + + +def _strong(tree): + """Strip weak types so both branches of a ``lax.cond`` agree.""" + return jax.tree.map( + lambda x: jax.lax.convert_element_type(jnp.asarray(x), + jnp.asarray(x).dtype), tree) + + +def fused_log_likelihood_four_axis_policy( + data, ra, dec, incl, x_grid, log_w_grid, *, interp, amp_sizing, + config=None, local_log_normalization=None, return_ledger=False): + """Distance-, phi_ref-, psi- AND time-marginalized lnL under the policy. + + Same contract as :func:`anglemarg.fused_log_likelihood_distphipsimarg_exact` + without ``return_lnLt``: the composite owns the time integral, so there is + no ``lnL(t)`` to hand back. With ``return_ledger`` the per-row ledger of + the controller is returned alongside (every leaf shaped ``(S,)``). + """ + if config is None: + config = PolicyConfig() + validate_policy_config(config) + guard = int(config.time_guard) + batch_rows = validate_batch_rows(config.reserve_batch_rows) + if local_log_normalization is None: + local_log_normalization, _ = policy_log_normalization( + data, x_grid, log_w_grid) + x_grid = jnp.asarray(x_grid, dtype=jnp.float64) + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64) + x_min = float(np.min(np.asarray(x_grid))) + x_max = float(np.max(np.asarray(x_grid))) + nodes, weights, check_nodes, check_weights = policy_time_rules( + data, config.reserve_time_refine) + + C_A, C_B, meta = _anglemarg.angle_coefficient_tables( + data, ra, dec, incl, interp, guard=guard) + # (KP,KS,S,Ntime) -> (S,KP,KS,Ntime): one row per extrinsic sample. + rows_A = jnp.moveaxis(C_A, 2, 0) + rows_B = jnp.moveaxis(C_B, 2, 0) + # Ordinary ILE has an arrival-time-independent norm. Collapse it per row + # and record the deviation; a row whose norm moves with time cannot be + # planned by this composite and is reported, never averaged away. + norm0 = rows_B[..., 0] + norm_dev = jnp.max(jnp.abs(rows_B - norm0[..., None]), axis=(1, 2, 3)) + norm_scale = jnp.maximum(1.0, jnp.max(jnp.abs(norm0), axis=(1, 2))) + norm_time_invariant = norm_dev <= float(config.norm_invariance_rtol) * norm_scale + # A guard past the stored data buffer gathers samples the build never + # stored; they come back nonfinite with no error, empty the plan and would + # read as a method decline. Name it as the input error it is. + tables_finite = (jnp.all(jnp.isfinite(rows_A), axis=(1, 2, 3)) + & jnp.all(jnp.isfinite(rows_B), axis=(1, 2, 3))) + + def _plan_row(table, norm): + base = _aap.rank_joint_starts_from_uvq_device( + table, norm, x_min, x_max, time_guard=guard, + max_starts=int(config.base_max_starts), + max_time_nodes=int(config.max_time_nodes), + angular_oversample=int(config.base_oversample)) + extra = _aap.rank_joint_starts_from_uvq_device( + table, norm, x_min, x_max, time_guard=guard, + max_starts=int(config.base_max_starts), + max_time_nodes=int(config.max_time_nodes), + angular_oversample=int(config.enriched_oversample)) + (base_plan, enriched_plan, base_planning, enriched_planning, + shared_planning) = _aap.make_all_axis_mode_plan_pair_device( + table, norm, base, extra, x_min, x_max, + max_modes=int(config.max_modes), + enriched_max_modes=int(config.enriched_max_modes), + local_radius=float(config.local_radius), + time_guard=guard, iterations=int(config.refine_iterations), + time_reconstruction_certified=False) + # Row-local control data (inputs are already under stop_gradient; the + # output cut is kept so a direct caller of _plan_row gets the same + # contract). Until derivative parity is established the discrete + # rank/dedup decisions are not part of the differentiated graph. + base_plan = jax.tree.map(jax.lax.stop_gradient, base_plan) + enriched_plan = jax.tree.map(jax.lax.stop_gradient, enriched_plan) + planning = dict( + base_n_selected_modes=base_planning["n_selected_modes"], + enriched_n_selected_modes=enriched_planning["n_selected_modes"], + base_n_optimizer_starts=base_planning["n_optimizer_starts"], + enriched_n_optimizer_starts=enriched_planning["n_optimizer_starts"], + optimizer_starts_executed=shared_planning[ + "n_optimizer_starts_executed"], + base_n_lattice_evaluations=base_planning["n_lattice_evaluations"], + enriched_n_lattice_evaluations=enriched_planning[ + "n_lattice_evaluations"], + # How far over the cap a declining row actually was. decline_capacity + # says only that n_candidates exceeded base_max_starts; without the + # count there is no way to tell a row that missed by one from a row + # that would need ten times the cap, and therefore no way to judge + # whether raising the cap would recover anything. + # + # SCOPE, because the name would otherwise mislead exactly as the + # sibling `enriched_*` keys misled a reader on 2026-09-08: the + # second plan is built from `combine_device_start_plans(base, extra)` + # (all_axis_peaklocal.py:1602 onward), so its count is base PLUS + # extra (`:901`), while `capacity_ok` ANDs the two plans' own flags, + # each already compared against base_max_starts separately (`:902`). + # Comparing the combined count against the cap is therefore not a + # test of anything. Named `combined_` so the units travel with it. + base_n_candidates_before_cap=base_planning[ + "n_candidates_before_cap"], + combined_n_candidates_before_cap=enriched_planning[ + "n_candidates_before_cap"], + # decline_capacity is charged for THREE different causes and the + # ledger named only the union. `capacity_ok` at :2067 is + # discovery_capacity_ok on both plans; :1514 makes that + # start_capacity_ok & ~selection_overflow; and the combined plan's + # start_capacity_ok at :902 is itself + # base.capacity_ok & extra.capacity_ok & same_time_support. So a + # row can carry decline_capacity with every candidate count under + # the cap, and raising the cap cannot recover it. Without these + # two flags a count-based estimate of what a larger cap buys is an + # upper bound and reads as if it were the answer. + base_start_capacity_ok=base_planning["start_capacity_ok"], + combined_start_capacity_ok=enriched_planning["start_capacity_ok"], + base_selection_overflow=base_planning["selection_overflow"], + combined_selection_overflow=enriched_planning[ + "selection_overflow"], + base_norm_nonnegative=base_planning["norm_nonnegative"], + combined_norm_nonnegative=enriched_planning["norm_nonnegative"], + base_time_cover_certified=base_planning["time_cover_certified"], + combined_time_cover_certified=enriched_planning[ + "time_cover_certified"], + base_time_capacity_ok=base_planning["time_capacity_ok"], + combined_time_capacity_ok=enriched_planning["time_capacity_ok"]) + return base_plan, enriched_plan, planning + + # Planning is control data. Cutting the tangents at its INPUTS, not only + # at the plan outputs, keeps reverse-mode AD from tracing the 14-step + # Newton refinement over ~100 starts and stacking its residuals: with the + # cut at the outputs only, a rho 163 production row still asked for 85 GiB + # (158 GiB before the branches were rematerialized). + base_plans, enriched_plans, planning = jax.vmap(_plan_row)( + jax.lax.stop_gradient(rows_A), jax.lax.stop_gradient(norm0)) + + angular_name, time_rule = reserve_pair(config.reserve_scheme) + angular_kernel = resolve_reserve_angular_kernel( + angular_name, x_grid, log_w_grid, + amp_sizing=float(amp_sizing), m_max=int(meta["m_max"]), + dense_chunk=int(config.reserve_dense_chunk), + grid_block=int(config.reserve_grid_block)) + peaklocal = time_rule == "peaklocal" + if peaklocal: + # Tiers refine the fine lattice at fixed span (m -> 2m, n -> 2n-1); + # the scan is fixed. Node counts never depend on the row. The + # first tier is sized from the PREDICTED width; a tier beyond it is + # a report on the prediction, counted in reserve_escalations. + n0 = int(config.reserve_peaklocal_fine_nodes) + tiers = [] + n_fine, mult = n0, 1 + for _ in range(int(config.reserve_peaklocal_escalations) + 1): + tiers.append(("peaklocal", n_fine, mult)) + n_fine, mult = 2 * n_fine - 1, 2 * mult + # Seconds per native sample carrying the production w_t constant, as + # policy_time_rules does. + dt_scale = float(np.sum(np.asarray(data.w_t, dtype=float))) / ( + int(data.npts) - 1) + # Bandwidth of the stored Q, cycles per sample (nan without a Q; the + # row's own table then supplies it, and both are in the ledger). + sigma_f_q = q_bandwidth_cycles_per_sample(data) + sigma_t_override = float(config.reserve_peaklocal_sigma_t_override_samples) + else: + refine0 = int(config.reserve_time_refine) + refine_max = int(config.reserve_time_refine_max) + if refine_max < refine0: + raise ValueError("reserve_time_refine_max must be >= reserve_time_refine") + tiers = [] + f = refine0 + while f <= refine_max: + tiers.append((f,) + tuple(policy_time_rules(data, f))) + f *= 2 + + def _predict_row(table, norm, rho_located): + # rho from the located profile maximum (rho^2 = 2 P_max); the angular + # triangle bound is the fallback and is reported beside it (it ran + # 2.8x over on a rung-160 row, 453 against 157). + rho_bound = _plr.row_amplitude(table, norm, guard) + rho = jnp.where(jnp.isfinite(rho_located) & (rho_located > 0.0), + rho_located, rho_bound) + sigma_f_table = _plr.table_bandwidth_cycles(table, guard) + sigma_f = jnp.where(jnp.isfinite(sigma_f_q), sigma_f_q, sigma_f_table) + sigma_t = _plr.predicted_width_samples(rho, sigma_f) + if np.isfinite(sigma_t_override): + sigma_t = jnp.asarray(sigma_t_override, dtype=jnp.float64) + return dict(rho=rho, rho_bound=rho_bound, sigma_f_table=sigma_f_table, + sigma_f=sigma_f, sigma_t=sigma_t) + + def _rule_for_tier(tier, table, norm, base_plan, enriched_plan): + if tier[0] == "peaklocal": + found = _plr.locate_time_maxima( + table, norm, guard, int(data.npts), x_min, x_max, + n_candidates=int(config.reserve_peaklocal_blocks), + search_refine=int(config.reserve_peaklocal_search_refine), + angular_lattice=int(config.reserve_peaklocal_angular_lattice), + n_phi=int(config.reserve_peaklocal_search_phi_nodes), + newton_steps=int(config.reserve_peaklocal_newton_steps), + newton_step_max=float(config.reserve_peaklocal_newton_step_max)) + pred = _predict_row(table, norm, found["rho_located"]) + sigma_for_margin = jnp.minimum( + jnp.where(jnp.isfinite(pred["sigma_t"]), pred["sigma_t"], jnp.inf), + jnp.min(jnp.where(found["live"] & (found["widths"] > 0.0), + found["widths"], jnp.inf))) + sigma_for_margin = jnp.where(jnp.isfinite(sigma_for_margin), + sigma_for_margin, 1.0) + margin = float(config.reserve_peaklocal_scan_margin_sigmas) * sigma_for_margin + rule = _plr.peaklocal_time_rule( + found["centres"], found["widths"], found["live"], + int(data.npts), dt_scale, + sigma_t_samples=pred["sigma_t"], n_fine=int(tier[1]), + scan_refine=int(config.reserve_peaklocal_scan_refine), + fine_refine_multiplier=int(tier[2]), + n_scan=int(config.reserve_peaklocal_scan_nodes), + margin_samples=margin, + search_positions=found["search_positions"], + search_profile=found["search_profile"], + outside_slack_nats=float(config.reserve_peaklocal_outside_slack_nats)) + # The local branch's plan, as a cross-check only: its narrowest + # live Newton width and the distance from its first live centre + # to the locator's first block. + plan_live = jnp.concatenate((base_plan.live, enriched_plan.live)).astype(bool) + plan_c = jnp.concatenate((base_plan.centers[:, 0], enriched_plan.centers[:, 0])) + plan_w = jnp.abs(jnp.concatenate((base_plan.local_transforms[:, 0, 0], + enriched_plan.local_transforms[:, 0, 0]))) + plan_width = jnp.min(jnp.where(plan_live & (plan_w > 0.0), plan_w, jnp.inf)) + plan_centre = jnp.where(jnp.any(plan_live), + plan_c[jnp.argmax(plan_live)], jnp.nan) + ratio = rule["sigma_t_located_samples"] / rule["sigma_t_pred_samples"] + extra = dict( + reserve_time_refine_used=jnp.asarray(0), + reserve_time_rule_peaklocal=jnp.asarray(True), + reserve_peaklocal_fine_nodes_used=jnp.asarray(int(tier[1])), + reserve_peaklocal_fine_refine=rule["fine_refine"], + reserve_peaklocal_fine_spacing_samples=rule["fine_spacing_samples"], + reserve_peaklocal_block_span_samples=rule["block_span_samples"], + reserve_peaklocal_live_blocks=rule["n_live_blocks"], + reserve_peaklocal_first_block_centre_samples=rule[ + "first_block_centre_samples"], + reserve_peaklocal_rho_pred=pred["rho"], + reserve_peaklocal_rho_bound=pred["rho_bound"], + reserve_peaklocal_search_phi_nodes=jnp.asarray( + int(config.reserve_peaklocal_search_phi_nodes), dtype=jnp.int32), + # Where the rule is fine: the kernel's focus certificate. + reserve_peaklocal_focus_centre_samples=rule["first_block_centre_samples"], + reserve_peaklocal_focus_half_width_samples=0.25 * rule["block_span_samples"], + reserve_peaklocal_scan_lo_samples=rule["scan_lo_samples"], + reserve_peaklocal_scan_hi_samples=rule["scan_hi_samples"], + reserve_peaklocal_scan_margin_samples=margin, + reserve_peaklocal_outside_log_bound=rule["outside_log_bound"], + reserve_peaklocal_outside_search_nodes=rule["n_outside_search_nodes"], + reserve_peaklocal_sigma_f_q_cycles=jnp.asarray( + sigma_f_q, dtype=jnp.float64), + reserve_peaklocal_sigma_f_table_cycles=pred["sigma_f_table"], + reserve_peaklocal_sigma_t_pred_samples=rule["sigma_t_pred_samples"], + reserve_peaklocal_sigma_t_located_samples=rule["sigma_t_located_samples"], + reserve_peaklocal_sigma_t_used_samples=rule["sigma_t_used_samples"], + reserve_peaklocal_sigma_t_plan_samples=plan_width, + reserve_peaklocal_plan_centre_offset_samples=jnp.abs( + plan_centre - rule["first_block_centre_samples"]), + reserve_peaklocal_prediction_finite=rule["prediction_finite"], + # The located curvature width against the prediction. The + # prediction is the NARROWEST peak the primitive can make at + # this amplitude (raw rms frequency), so a face-on envelope + # is legitimately wider, by the raw-to-central moment ratio: + # consistent means located / predicted in [0.5, 4]. Outside + # it the prediction is reported as disagreeing; the warrant + # decides the row either way. + reserve_peaklocal_prediction_consistent=( + jnp.isfinite(ratio) & (ratio >= 0.5) & (ratio <= 4.0))) + return (rule["nodes"], rule["weights"], rule["check_nodes"], + rule["check_weights"], extra) + refine, nodes, weights, check_nodes, check_weights = tier + nan = jnp.asarray(jnp.nan, dtype=jnp.float64) + extra = dict( + reserve_time_refine_used=jnp.asarray(int(refine)), + reserve_time_rule_peaklocal=jnp.asarray(False), + reserve_peaklocal_fine_nodes_used=jnp.asarray(0), + reserve_peaklocal_fine_refine=jnp.asarray(0, dtype=jnp.int32), + reserve_peaklocal_fine_spacing_samples=nan, + reserve_peaklocal_block_span_samples=nan, + reserve_peaklocal_live_blocks=jnp.asarray(0, dtype=jnp.int32), + reserve_peaklocal_first_block_centre_samples=nan, + reserve_peaklocal_rho_pred=nan, + reserve_peaklocal_rho_bound=nan, + reserve_peaklocal_search_phi_nodes=jnp.asarray(0, dtype=jnp.int32), + reserve_peaklocal_focus_centre_samples=nan, + reserve_peaklocal_focus_half_width_samples=nan, + reserve_peaklocal_scan_lo_samples=nan, + reserve_peaklocal_scan_hi_samples=nan, + reserve_peaklocal_scan_margin_samples=nan, + reserve_peaklocal_outside_log_bound=nan, + reserve_peaklocal_outside_search_nodes=jnp.asarray(0, dtype=jnp.int32), + reserve_peaklocal_sigma_f_q_cycles=nan, + reserve_peaklocal_sigma_f_table_cycles=nan, + reserve_peaklocal_sigma_t_pred_samples=nan, + reserve_peaklocal_sigma_t_located_samples=nan, + reserve_peaklocal_sigma_t_used_samples=nan, + reserve_peaklocal_sigma_t_plan_samples=nan, + reserve_peaklocal_plan_centre_offset_samples=nan, + reserve_peaklocal_prediction_finite=jnp.asarray(False), + reserve_peaklocal_prediction_consistent=jnp.asarray(False)) + return nodes, weights, check_nodes, check_weights, extra + + def _controller(table, norm, base_plan, enriched_plan, tier): + nodes, weights, check_nodes, check_weights, extra = _rule_for_tier( + tier, table, norm, base_plan, enriched_plan) + sel, ok, led = _aap.empirical_enrichment_with_exact_reserve( + table, norm, base_plan, enriched_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w_grid, + time_weights=weights, + reserve_amp_sizing=float(amp_sizing), + reserve_m_max=int(meta["m_max"]), + reserve_dense_chunk=int(config.reserve_dense_chunk), + reserve_grid_block=int(config.reserve_grid_block), + reserve_time_nodes=nodes, + reserve_time_check_nodes=check_nodes, + reserve_time_check_weights=check_weights, + reserve_time_resolution_tol_nats=float( + config.total_value_error_budget_nats), + base_order=int(config.base_order), + base_check_order=int(config.base_check_order), + enriched_order=int(config.enriched_order), + enriched_check_order=int(config.enriched_check_order), + convergence_tol_nats=float(config.convergence_tol_nats), + time_guard=guard, + time_guard_tol_nats=float(config.time_guard_tol_nats), + local_log_normalization=float(local_log_normalization), + time_outside_tol_nats=float(config.time_outside_tol_nats), + total_value_error_budget_nats=float( + config.total_value_error_budget_nats), + reserve_log_offset=0.0, + reserve_angular_kernel=angular_kernel, + reserve_time_focus=( + None if not peaklocal else + (extra["reserve_peaklocal_focus_centre_samples"], + extra["reserve_peaklocal_focus_half_width_samples"])), + reserve_time_cover=( + None if not peaklocal else + (extra["reserve_peaklocal_scan_lo_samples"], + extra["reserve_peaklocal_scan_hi_samples"], + extra["reserve_peaklocal_outside_log_bound"]))) + led = dict(led) + led.update(extra) + return _strong((sel, ok, led)) + + def _row(args): + table, norm, base_plan, enriched_plan = args + # Each tier is rematerialized: reverse-mode AD otherwise keeps the + # residuals of every tier's dense reserve alive at once. + state = jax.checkpoint( + lambda t, nm, bp, ep: _controller(t, nm, bp, ep, tiers[0]))( + table, norm, base_plan, enriched_plan) + escalations = jnp.asarray(0) + for tier in tiers[1:]: + sel, ok, led = state + need = (led["reserve_executed"] & led["reserve_finite"] + & (~led["reserve_time_warranted"])) + run_tier = jax.checkpoint( + lambda _, tier=tier: _controller( + table, norm, base_plan, enriched_plan, tier)) + state = jax.lax.cond(need, run_tier, lambda st: st, state) + escalations = escalations + need.astype(escalations.dtype) + sel, ok, led = state + led = dict(led) + led["reserve_escalations"] = escalations + return sel, ok, led + + n_rows = int(rows_A.shape[0]) + xs = (rows_A, norm0, base_plans, enriched_plans) + if batch_rows == 1 or n_rows == 1: + # Row at a time. Kept as a distinct call rather than batch_size=1 so + # the graph is the one PR #268 measured: batch_size=1 would still wrap + # the body in a vmap, paying the cond-to-select cost for no occupancy. + # + # n_rows == 1 takes this path whatever was requested. The wrapper's + # _scalar evaluates ONE row, so value_and_grad and hessian always land + # here; a vmap over a single row would convert both conds to selects + # and pay every reserve tier and both accept/reserve branches with no + # second row to amortize them. Requesting a batch must not make the + # gradient path more expensive than not requesting one. + selected, usable, ledger = jax.lax.map(_row, xs) + elif batch_rows == 0 or batch_rows >= n_rows: + # One full batch. Spelled as an explicit vmap rather than delegated + # to batch_size: jax 0.9.2 documents batch_size=0 as a full vmap, but + # the IGWN environment's jax 0.7.1 computes n // batch_size first and + # raises ZeroDivisionError, and a batch_size above the row count is a + # zero-length scan plus a remainder in both. The explicit vmap is the + # same computation in every version. + selected, usable, ledger = jax.vmap(_row)(xs) + else: + selected, usable, ledger = jax.lax.map(_row, xs, + batch_size=batch_rows) + ledger = dict(ledger) + # Truthful, not decorative: this key read True unconditionally before the + # batch size was a knob. Record what EXECUTED, not what was asked for: + # a request of 8 against 6 rows runs a 6-row vmap, and a request of 8 on + # the single-row gradient path runs sequentially. Reporting the request + # on the one key whose purpose is truthfulness is how the old hardcoded + # True happened. lax.map's trailing remainder means the batch a given row + # landed in is still not recoverable per row, so this is the size of the + # scanned batch; the request is kept beside it. + if batch_rows == 1 or n_rows == 1: + batch_executed = 1 + elif batch_rows == 0 or batch_rows >= n_rows: + batch_executed = n_rows + else: + batch_executed = batch_rows + ledger["reserve_batch_execution_sequential"] = jnp.full( + (n_rows,), batch_executed == 1, dtype=bool) + ledger["reserve_batch_rows_executed"] = jnp.full( + (n_rows,), batch_executed, dtype=jnp.int32) + ledger["reserve_batch_rows_requested"] = jnp.full( + (n_rows,), batch_rows, dtype=jnp.int32) + usable = usable & norm_time_invariant & tables_finite + # Fail closed: a value the controller could not warrant is not a + # likelihood. nan, never the finite diagnostic, reaches the sampler; the + # driver refuses to publish a run that contains such rows. + lnL = jnp.where(usable, selected, jnp.nan) + ledger.update(planning) + ledger["norm_time_invariant"] = norm_time_invariant + ledger["norm_time_deviation"] = norm_dev + ledger["tables_finite"] = tables_finite + ledger["input_nonfinite"] = ~tables_finite + ledger["usable"] = usable + ledger["selected_value"] = selected + ledger["lnL"] = lnL + if return_ledger: + return lnL, ledger + return lnL + + +def summarize_policy_ledger(ledger): + """Host-side counts for the run record. ``ledger`` leaves are ``(S,)``.""" + def _count(key): + return int(np.sum(np.asarray(ledger[key], dtype=bool))) + n = int(np.asarray(ledger["usable"]).shape[0]) + out = dict( + rows=n, + accepted_local=_count("accepted_local"), + reserve_executed=_count("reserve_executed"), + reserve_warranted=_count("selected_value_is_warranted_reserve"), + usable=_count("usable"), + unusable=n - _count("usable"), + norm_time_invariant=_count("norm_time_invariant"), + tables_finite=_count("tables_finite"), + reconciles=_count("reconciles"), + disposition_reconciles=_count("disposition_reconciles"), + ) + declines = {} + for key in _DECLINE_KEYS: + if key in ledger: + c = _count(key) + if c: + declines[key] = c + out["declines"] = declines + if "reserve_escalations" in ledger: + out["reserve_escalations"] = int(np.sum( + np.asarray(ledger["reserve_escalations"]))) + if "reserve_batch_rows_executed" in ledger: + ex = np.asarray(ledger["reserve_batch_rows_executed"]) + req = np.asarray(ledger["reserve_batch_rows_requested"]) + out["reserve_batch_rows"] = int(ex[0]) if ex.size else 0 + out["reserve_batch_rows_requested"] = int(req[0]) if req.size else 0 + out["reserve_batch_execution_sequential"] = bool(np.all( + np.asarray(ledger["reserve_batch_execution_sequential"], + dtype=bool))) + if "lnL" in ledger: + out["nan_rows"] = int(np.sum(~np.isfinite( + np.asarray(ledger["lnL"], dtype=float)))) + score = np.asarray(ledger["empirical_value_error_score_nats"], dtype=float) + finite = score[np.isfinite(score)] + out["max_local_error_score_nats"] = ( + float(np.max(finite)) if finite.size else float("nan")) + return out diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py index 41ceda4aa..3bc30aaef 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py @@ -1,9 +1,24 @@ """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 NumPy reference ``RIFT.likelihood.joint_angle_peak_local`` obtains BOTH-angle targets +from the finite algebraic stationary set implemented in +``RIFT.likelihood.bivariate_trig_stationary``. This device kernel is NOT a transcription +of that rule and does not enumerate: it is exact in u on the cell partition, and on phi it +offers two paths, a dense scan (:func:`joint_lnL_phi_dense`) and a peak-local rule +(:func:`phi_local_lnI`) that Newton-iterates from a phi GRID of seeds. + +A SAMPLED PHI GRID IS NOT ALGEBRAIC ENUMERATION, and phi_local_lnI does not claim to be +one. What stands behind it is an omitted-mass bound -- a seed the grid misses raises +``area_outside`` and the row declines -- gated further by empirical convergence checks +that are labelled as estimates where they are used. See :func:`phi_local_lnI` for exactly +which part of ``ok`` is a bound and which parts are not. + +An in-kernel Sylvester-resultant seeder was built here and removed: it duplicated +``bivariate_trig_stationary`` without that module's BKK root count, Jacobian conditioning, +torus classification, cross-projection agreement or ``ok`` flag, and solving it inside a +traced function is the substitution that module says must not be made. Production wiring +stays unchanged until a host-built, fixed-capacity algebraic plan and its dense fallback +can cross the JAX boundary honestly. 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 @@ -11,8 +26,8 @@ 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. +``_merge_boxes`` on cannot arise. Everything is then static at trace time: 4 roots, +4 candidate cells, and an amplitude-derived quadrature count streamed in fixed blocks. 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 @@ -23,20 +38,30 @@ 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 +``A^-1/2``. A local window therefore needs a fixed count, but a rejected Newton centre +falls back to a whole cell and needs ``~sqrt(A)`` nodes. Production uses that conservative +count for every cell because fallback is data-dependent; streaming preserves the memory +economy even though the arithmetic cost is no longer claimed constant. + +SCOPE OF THIS KERNEL, AND IT IS TWO RULES RATHER THAN ONE. Both localize u exactly on +the cell partition; they differ on phi. + +:func:`joint_lnL_phi_dense` scans phi as a dense grid in chunks -- 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. This is the production path. + +:func:`phi_local_lnI` localizes phi as well -- the (phi localized, psi localized) cell of +the family. An earlier version of this paragraph said that was "not attempted here", +which was true when written and stopped being true in the same file: the profile +``F(phi)`` and its envelope derivatives are :func:`u_profile`, and phi_local_lnI is built +on them. Its cost stops growing with amplitude, and it DECLINES rather than returning an +unbounded number, on a certificate that is one genuine bound plus two empirical gates -- +see that function for exactly which is which. It is not wired into production. + +MEMORY. Bounded by ``phi_chunk`` and ``U_NODE_STREAM_CHUNK`` through rolled loops, never +by the full phi or u grids: the largest u transient is +``(phi_chunk, n_x, 4, U_NODE_STREAM_CHUNK)``. These are cost knobs and cannot change the result beyond floating-point reassociation. """ @@ -47,12 +72,30 @@ __all__ = [ "required_n_phi", + "required_u_nodes", + "u_nodes_in_use", "U_WINDOW_SIGMA", "U_NODES_PER_CELL", + "U_PTS_PER_SIGMA", + "U_NODE_STREAM_CHUNK", "PHI_CHUNK_DEFAULT", "u_stationary_roots", "log_inner_u_integral", "joint_lnL_phi_dense", + "u_profile", + "eval_g2", + "phi_local_lnI", + "phi_local_lnI_at_distance", + "joint_lnL_phi_local", + "X_CHUNK_DEFAULT", + "PT_CHUNK_DEFAULT", + "PHI_SEEDS", + "PHI_WINDOW_SIGMA", + "PHI_NODES_PER_REGION", + "PHI_BOUND_GRID", + "OUTSIDE_TOL_NATS", + "phi_derivative_bound", + "profile_derivative_bounds", ] #: Local u-window half-width in units of the local sigma, CLIPPED to the cell. The cell @@ -68,12 +111,90 @@ #: 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). +#: +#: THAT AMPLITUDE-INDEPENDENCE HOLDS FOR A WINDOWED CELL AND NOT FOR A FALLBACK ONE. +#: A cell whose Newton centre is rejected (stalled on a boundary, large stationary +#: residual) is integrated WHOLE, and 48 nodes then span the entire cell rather than +#: +-12 sigma. The numpy twin measured 1.7e-03 nats of inner-u error that way, so the +#: honest statement is: this default resolves WINDOWED cells at any amplitude. The +#: production caller may hit a fallback at any phi/distance point, so it uses the +#: amplitude-derived :func:`u_nodes_in_use` policy instead of relying on this floor. U_NODES_PER_CELL = 48 +#: Maximum number of u nodes materialized at once. The production count grows as +#: sqrt(amplitude), but the quadrature is accumulated through a rolled scan so that its +#: live node axis -- and therefore the batch-memory model -- stays bounded. +U_NODE_STREAM_CHUNK = 8 + #: phi points per scan step. PHI_CHUNK_DEFAULT = 16 +def u_nodes_in_use(amp_sizing=None): + """The u-node count the peak-local kernel WILL ACTUALLY REQUEST at this amplitude. + + SINGLE SOURCE OF TRUTH, and it exists because the batch-memory guard in + :mod:`~RIFT.likelihood.jax_ile.samplers` has to model the same number the kernel + requests, and the two are in different files. External review found the trap before + it fired: the guard hard-coded ``U_NODES_PER_CELL``, so anyone wiring + :func:`required_u_nodes` into the kernel would silently invalidate it -- at the + production floor ``amp_sizing = 450`` that is 896 nodes against a modeled 48, and the + documented live slab goes from 3.6 GiB to 67 GiB at chunk one. An automated agent + then did exactly that wiring, and left the guard untouched, which is the trap firing. + + Both the kernel (:func:`joint_lnL_phi_dense`, whose ``n_nodes`` defaults to ``None`` + and resolves here) and the guard call this, and the fused caller passes the same + ``amp_sizing`` to both. An earlier version of this docstring claimed that while only + the guard called it and the kernel still defaulted straight to ``U_NODES_PER_CELL`` -- + a single source of truth that only one side read, which is no single source of truth + at all and is exactly the divergence this helper exists to prevent. Caught in review. + + A direct low-level call without an amplitude retains the validated 48-node windowed + floor. Production always supplies ``amp_sizing`` and therefore gets the derived, + uncapped whole-cell requirement. The quadrature streams that count in + ``U_NODE_STREAM_CHUNK``-sized blocks, so accuracy grows with amplitude without making + the live node dimension grow with it. + """ + if amp_sizing is None: + return U_NODES_PER_CELL + return required_u_nodes(amp_sizing) + + +#: Trapezoid points per curvature length on the u axis. Shared by :func:`required_u_nodes`, +#: which sizes a fallback cell from an amplitude PROXY before the table is built, and by +#: :func:`u_profile`, which applies the same density to the EXACT per-cell curvature bound +#: once it has one. One constant so the static budget and the in-kernel adequacy test +#: cannot drift apart. +U_PTS_PER_SIGMA = 3.0 + + +def required_u_nodes(amplitude, pts_per_sigma=None, cap=None): + """u nodes per cell adequate for a FALLBACK (whole-cell) integration at ``amplitude``. + + Derived, not tuned. The u-spectrum has two terms, so ``|d2g/du2| <= M2u`` exactly, + and at exponent amplitude ``A`` the coefficients scale with ``A`` giving + ``M2u ~ 5 A``: 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. A fallback cell can span most of the circle, so the requirement + is ``2 pi * sqrt(M2u) * pts_per_sigma``. + + JAX NEEDS THIS STATICALLY, which is why it is a caller-side helper rather than an + adaptation inside the kernel: shapes cannot depend on traced values. The numpy twin + derives the same quantity per call because it can. + + ``cap`` is available only for explicit diagnostic callers. It is deliberately + ``None`` in production: truncating the requested count recreates the inside-cover + accuracy failure this policy exists to prevent. Memory is bounded independently by + streaming the node axis rather than by silently reducing the quadrature. + """ + a = max(float(amplitude), 1.0) + if pts_per_sigma is None: + pts_per_sigma = U_PTS_PER_SIGMA + need = int(np.ceil(2.0 * np.pi * np.sqrt(5.0 * a) * float(pts_per_sigma))) + 1 + need = max(need, U_NODES_PER_CELL) + return int(need if cap is None else min(need, int(cap))) + + def required_n_phi(amplitude, m_max=2): """phi-grid size for a given exponent amplitude. @@ -200,7 +321,14 @@ def _newton(uc, _): # 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. + # |g_u|/M_1 = 0.33. A cell failing it is integrated WHOLE -- which ADDS NO NODES, it + # spreads the same n_nodes over the whole cell, so the fallback is COARSER than the + # window it replaces. (An earlier comment here claimed "can only add nodes"; that was + # wrong, and the numpy twin measured the inner-u error recorded on + # U_NODES_PER_CELL from it.) JAX + # cannot adapt n_nodes -- shapes may not depend on traced values -- so the sizing is + # exposed to the caller as required_u_nodes() rather than fixed here; see its docstring + # for why raising it by default is the wrong trade. 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| @@ -216,13 +344,33 @@ def _newton(uc, _): 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,) + # STREAM THE NODE AXIS. Materializing (4, n_nodes) here is multiplied by the outer + # phi, distance, time and sample batches. At the production floor the accurate + # fallback policy asks for 896 nodes, which would turn the documented 48-node live + # slab into ~67 GiB even at sample chunk one. A rolled scan keeps only + # U_NODE_STREAM_CHUNK nodes live while accumulating the identical trapezoid sum. + n_nodes = int(n_nodes) + if n_nodes < 2: + raise ValueError("n_nodes must be at least 2") + n_blocks = int(np.ceil(n_nodes / U_NODE_STREAM_CHUNK)) + local_idx = jnp.arange(U_NODE_STREAM_CHUNK) + + def _node_block(block_i, log_sum): + idx = block_i * U_NODE_STREAM_CHUNK + local_idx + live = idx < n_nodes + s = idx / float(n_nodes - 1) + uu = lo[:, None] + width[:, None] * s[None, :] + gg = _g_u(a, c1, c2, uu, 0) + endpoint = (idx == 0) | (idx == n_nodes - 1) + log_trap = jnp.where(endpoint, -jnp.log(2.0), 0.0) + terms = jnp.where(live[None, :], gg + log_trap[None, :], -jnp.inf) + block = jax.scipy.special.logsumexp(terms, axis=-1) + return jnp.logaddexp(log_sum, block) + + cell_sum = lax.fori_loop(0, n_blocks, jax.checkpoint(_node_block), + jnp.full(4, -jnp.inf)) + log_scale = jnp.log(jnp.where(width > 0, width, 1.0)) - jnp.log(n_nodes - 1) + cell = cell_sum + log_scale cell = jnp.where(width > 0, cell, -jnp.inf) return jax.scipy.special.logsumexp(cell) @@ -237,7 +385,7 @@ def _joint_table(C_A, C_B, x): 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): + n_nodes=None): """Distance-, phi- and psi-marginalized value at one ``(sample, time)``. Same normalization as ``anglemarg.fused_log_likelihood_distphipsimarg_*``: uniform @@ -246,6 +394,8 @@ def joint_lnL_phi_dense(C_A, C_B, x_grid, log_w_grid, n_phi=256, ``phi`` is a dense grid scanned in chunks; ``u`` is exact per the cell partition. """ + if n_nodes is None: + n_nodes = u_nodes_in_use() 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() @@ -265,23 +415,1015 @@ def one_phi(phi): return jax.vmap(log_inner_u_integral, in_axes=(0, 0, 0, None))( a[:, 0], c1[:, 0], c2[:, 0], n_nodes) # (nx,) + # REDUCE INTO THE CARRY, do not stack. The phi reduction is a logsumexp, so the + # scan can carry the running (nx,) accumulator instead of returning a value per + # chunk. Returning one made `out` a live (n_chunk, phi_chunk, nx) array -- the + # whole phi axis, materialized before the reduction that immediately collapses it. + # Under the caller's sample/time vmaps that is `S * npts * n_phi * n_x * 8` bytes, + # which reached 19.97 GiB at rung 640 of the ladder-2 injection and OOMed a card + # with 18: measured 19.99 GiB requested against 19.97 predicted, and the same + # formula to three significant figures at three other sizes. Nothing about the u + # axis was involved -- pinning the per-cell node count to its 48 floor, 468x below + # production, left the peak identical to the tenth of a MiB. + # `anglemarg.coefficient_table_distphipsimarg_exact` has carried its logsumexp in + # the scan state from the start; this is the same shape of fix, and it makes the + # live phi footprint one chunk rather than the whole axis. + # See development/BLOCKER_peaklocal_scheme_oom_20260908.md in RIFT_roboto_paper. 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 + # jnp.logaddexp against the running total, the accumulation + # `log_inner_u_integral` already uses over its own cells. The padded lanes are + # -inf and are the identity here, which is what retires the [:n_phi] slice: the + # mask alone now excludes them. + return jnp.logaddexp(carry, jax.scipy.special.logsumexp(vals, axis=0)), None # 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, + acc, _ = lax.scan(jax.checkpoint(step), + jnp.full((x_grid.size,), -jnp.inf, dtype=jnp.float64), (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) + per_x = acc - 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) + + +#: Distance nodes evaluated at once by :func:`joint_lnL_phi_local`. The phi-local +#: quadrature grid is ``n_slots * n_nodes * 4 * u_nodes`` per distance node and NOTHING +#: about it is streamed, so the distance axis is the one that has to be rolled or the +#: live slab is multiplied by the whole grid. This is the analogue of +#: ``PHI_CHUNK_DEFAULT`` on the dense path, and it exists for the same reason. +#: ONE, and that is a memory verdict rather than a preference. The batch guard models +#: `x_chunk * pt_chunk * 4 * u_nodes * terms * 16` bytes live, and at the production u +#: count (896) even a chunk of 8 needs 43.8 GiB for a 64-point time window -- the guard +#: refuses it outright. Vectorizing the distance axis has to wait for the u axis to be +#: streamed inside u_profile, or for the u count itself to come down. +X_CHUNK_DEFAULT = 1 + +#: Quadrature points evaluated at once inside :func:`phi_local_lnI`. The grid is +#: ``n_slots * n_nodes`` phi points, each costing a u profile of ``4 * u_nodes`` -- and +#: ``eval_g2`` materializes ``(points, KP, 2KS+1)`` COMPLEX, so the live slab is +#: ``points * 4 * u_nodes * KP * (2KS+1) * 16`` bytes. At the production u count +#: (``u_nodes_in_use(450) = 896``) that is about a gigabyte per distance node unrolled, +#: which is why this axis is rolled and not merely counted. Same role as +#: ``PHI_CHUNK_DEFAULT`` on the dense path. +#: 32, chosen so the guard passes at the production u count rather than by taste: +#: 32 * 4 * 896 * 25 * 16 = 45.9 MB live per sample-time point, 2.9 GiB across a 64-point +#: window, against the 4 GiB default allowance. Raising it is the first thing to try on a +#: bigger card, and the guard will say so rather than let it OOM. +PT_CHUNK_DEFAULT = 32 + + +def _prof_scan(prof, pts, chunk): + """``vmap(prof)`` over ``pts``, rolled in fixed-size blocks. Values identical. + + Only the profile VALUE and the three fallback counts are kept: the phi derivatives + the Newton step needs are not wanted here, and carrying them would defeat the point + by keeping a second array of the same length alive. + """ + n = int(pts.shape[0]) + # CLAMPED, because a chunk larger than the data is not "unrolled", it is PADDED. A + # caller asking for pt_chunk = 1e6 on 388 points was evaluating a million, 999612 of + # them padding -- external review measured 6.56 GB of temporaries for one test. The + # natural way to ask for an unrolled reference is a huge chunk, so the function has to + # mean it rather than take it literally. + chunk = int(min(int(chunk), n)) + n_chunk = int(np.ceil(n / chunk)) + pad = n_chunk * chunk - n + pp = jnp.concatenate([pts, jnp.zeros(pad)]) + + def step(carry, blk): + F, _, _, fb, rk, st = jax.vmap(prof)(blk) + return carry, (F, fb, rk, st) + + _, (F, fb, rk, st) = lax.scan(jax.checkpoint(step), None, + pp.reshape(n_chunk, chunk)) + keep = lambda a: a.reshape(-1)[:n] + return keep(F), keep(fb), keep(rk), keep(st) + + +def phi_local_lnI_at_distance(C_A, C_B, x, **kw): + """The ``(phi, u)`` torus integral at ONE distance node. THE STACKABLE UNIT. + + ``log int dphi int du exp(x A - x^2/2 B)`` with both angle axes localized, returned + as ``(value, ok, info)`` exactly as :func:`phi_local_lnI` does. + + THIS IS THE SEAM, and it is public so that a distance quadrature does not have to be + built into this module to be used with it. The distance rule is entirely a matter of + WHICH ``x`` a caller evaluates and WHAT WEIGHTS it applies; nothing here assumes the + caller's nodes are a grid, are equally spaced, or come from any particular rule. A + Gauss-Hermite placement, an adaptive rule, or the plain grid + :func:`joint_lnL_phi_local` uses all sit on top of this same call. + + The normalization is the bare torus integral -- NOT the ``(2 pi)^-2`` prior factor, + which belongs to whoever closes the distance sum. Getting that split wrong is how a + stacked quadrature would silently double-apply or drop it, so it is stated here and + applied in exactly one place, :func:`joint_lnL_phi_local`. + """ + return phi_local_lnI(_joint_table(C_A, C_B, x), **kw) + + +def joint_lnL_phi_local(C_A, C_B, x_grid, log_w_grid, x_chunk=X_CHUNK_DEFAULT, **kw): + """Distance-, phi- and psi-marginalized value with BOTH angle axes localized. + + The default combiner over :func:`phi_local_lnI_at_distance`: it sums the caller's + distance grid, the same contract :func:`joint_lnL_phi_dense` has, and it is a THIN + and replaceable layer. A caller with its own distance quadrature should call the + per-node function directly rather than reach through this one. + + Returns ``(value, ok, info)``. + + ``ok`` IS THE CONJUNCTION OVER NODES, which is the fail-closed reading: the distance + sum is only as trustworthy as the least trustworthy node in it, and a declining node + still returns a finite number that would otherwise be summed in silently. A node + whose weight makes it negligible cannot currently exempt itself -- deciding that + would need the very value the decline says not to trust -- so this is conservative + and deliberately so. + + THE DISTANCE AXIS IS ROLLED. Unlike the dense path there is no streaming inside the + phi-local kernel, so its live slab is the whole ``n_slots * n_nodes * 4 * u_nodes`` + quadrature grid; multiplying that by an unrolled distance grid is what makes the + scheme unusable rather than merely expensive. ``x_chunk`` bounds it, and the batch + guard in :mod:`~RIFT.likelihood.jax_ile.samplers` must model the SAME number. + """ + x_grid = jnp.asarray(x_grid, dtype=jnp.float64).ravel() + log_w_grid = jnp.asarray(log_w_grid, dtype=jnp.float64).ravel() + n_x = int(x_grid.shape[0]) + n_chunk = int(np.ceil(n_x / x_chunk)) + pad = n_chunk * x_chunk - n_x + xs = jnp.concatenate([x_grid, jnp.zeros(pad)]) + lw = jnp.concatenate([log_w_grid, jnp.full(pad, -jnp.inf)]) + + def step(carry, args): + xc, lwc = args + v, ok, _ = jax.vmap( + lambda z: phi_local_lnI_at_distance(C_A, C_B, z, **kw))(xc) + # a padded node is switched off by its weight, and must not vote on `ok` + live = jnp.isfinite(lwc) + return carry, (jnp.where(live, v, -jnp.inf), jnp.logical_or(ok, ~live)) + + _, (vals, oks) = lax.scan(jax.checkpoint(step), None, + (xs.reshape(n_chunk, x_chunk), + lw.reshape(n_chunk, x_chunk))) + per_x = vals.reshape(-1)[:n_x] + value = jax.scipy.special.logsumexp(per_x + log_w_grid) - 2.0 * jnp.log(2.0 * jnp.pi) + return value, oks.reshape(-1)[:n_x].all(), {"per_x": per_x} + + +# ------------------------------------------------------- phi localization + +#: phi seeds. These are SEEDS, not a quadrature grid: Newton moves each to a maximum of +#: the profile and overlapping windows merge, so the count sets how many distinct modes +#: can be found, not the accuracy. It does not scale with amplitude -- the number of +#: maxima of F is set by the bidegree, which is mode content, not SNR. +PHI_SEEDS = 32 + +#: phi window half-width in units of the profile's local sigma, and nodes per region. +#: Same Poisson-summation argument as U_NODES_PER_CELL: at +-12 sigma with 96 nodes the +#: spacing is sigma/4 and the trapezoid error on a Gaussian is ~1e-137. +PHI_WINDOW_SIGMA = 12.0 +#: Odd, so the grid is NESTED and both probes are free: the even indices are a trapezoid +#: at half the density and the odd indices are exactly its midpoints. Every point +#: evaluated enters the answer; neither probe costs an evaluation. +#: +#: 193 rather than 97, and the count is set by what the probes can see rather than by the +#: accuracy of the answer. A rule's own aliases are invisible in its own samples, so the +#: probes always certify the COARSE rule -- here the 97-node one -- and the answer rides +#: a level finer. At 97 the probes drop to 49/48 and the measured consequence is not +#: subtle: on the aliasing counterexample both probes read 1.1e-13 while the answer is +#: 0.02017 nats wrong, i.e. it ACCEPTS. At 193 the same case returns the right answer to +#: 1e-6 and still declines, because the coarse rule was bad. This is the same evaluation +#: count the discarded second grid used to cost, spent on the answer instead of a +#: diagnostic. +PHI_NODES_PER_REGION = 193 + +#: Grid on which the phi omitted-mass bound is evaluated. Not a tuning knob: it sets the +#: half-spacing ``delta`` of the second-order lift, so a coarser grid gives a LOOSER +#: (still valid) bound and more declines, never a wrong accept. +PHI_BOUND_GRID = 256 + +#: Accept when the certified mass outside the covered phi regions is this many nats below +#: the value. Same number and same meaning as the numpy reference's OUTSIDE_TOL_NATS. +OUTSIDE_TOL_NATS = -23.0 + + +def phi_derivative_bound(C, order=0): + """TRUE bound on ``|d^order_phi g|`` by the triangle inequality on the table. + + The one construction here that cannot be a fit -- the 1-D phi analogue of the numpy + reference's :func:`~RIFT.likelihood.joint_angle_peak_local.derivative_bound`. + """ + KP = C.shape[0] + k = jnp.arange(KP)[:, None] + w = jnp.where(k > 0, 2.0, 1.0) # k>0 stored once, counted twice (real field) + return (w * jnp.abs(C) * (jnp.abs(k) ** order)).sum() + + +def sup_g_bound(C, phi): + """``log(2 pi) + max_u g(phi, u)``: an EXACT upper bound on ``F(phi)``, with no u + quadrature in it at all. + + F(phi) = log int_0^{2pi} exp(g) du <= log(2 pi) + max_u g(phi, u) + + and ``max_u g`` is exact: on a phi slice ``g`` is ``a + Re(c1 e^{iu}) + Re(c2 e^{2iu})``, + whose maximum over the circle is attained at one of the four stationary angles + :func:`u_stationary_roots` already returns. + + WHY THIS EXISTS RATHER THAN A FINER PROFILE GRID. The outside certificate needs an + upper bound on ``F``, not ``F`` itself, and buying accuracy through the profile costs a + full u quadrature at every bound-grid point -- ``n_bound * u_nodes``, and BOTH grow with + amplitude. Measured: at amplitude 3e4 that is 8192 * 4 * 7302 evaluations and the + process is killed outright. + + It is CHEAPER and BETTER-CONDITIONED, though not uniformly tighter: it discards the + Laplace width (see the slack note below), so on a grid fine enough for the profile + route's second-order term to vanish that route would win. What matters is that no + affordable grid is that fine. Its Lipschitz constant is + ``M10``, by the envelope inequality + ``|max_u g(phi1,.) - max_u g(phi2,.)| <= max_u |g(phi1,u) - g(phi2,u)| <= M10 |dphi|``, + so refining the grid buys a LINEAR reduction where the profile route is pinned at second + order by ``M2F ~ M10^2``. And it cannot be corrupted by the u-quadrature fallback that + adversarial review flagged for the profile route, because it never calls it. + + The slack is the Laplace width it discards: ``F ~ max_u g + log(sigma_u sqrt(2 pi))`` + with ``sigma_u = |g_uu|^-1/2``, so it over-estimates by ``log(sqrt(2 pi) / sigma_u)`` -- + about 6 nats at amplitude 3e4, against a threshold with tens of nats of headroom. + """ + 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 * k) * w + D = lambda q: (ph * C[:, KS + q]).sum() + a = D(0).real + c1 = D(1) + jnp.conj(D(-1)) + c2 = D(2) + jnp.conj(D(-2)) + # A MAX OVER ROOTS IS A LOWER BOUND UNLESS THE ROOTS ARE RIGHT, and this returned one + # as if it were the maximum. ``u_stationary_roots`` builds a companion matrix for + # ``c2 z^4 + ...`` and substitutes ``lead = 1`` when ``c2 = 0``, which solves a + # DIFFERENT polynomial, so for a table with no q = +-2 content the returned angles need + # not contain the maximizer. Measured over four such draws, this came back 0.024 to + # 0.092 nats BELOW ``log(2 pi) + max_u g``: not a loose bound, an invalid one, and the + # whole outside certificate rests on this inequality. Adversarial review, found by + # constructing the degenerate table rather than by reading the algebra. + # + # ``a + |c1| + |c2| >= max_u g`` holds for every table and needs no roots. It is used + # wherever the quartic cannot be trusted -- and in exactly that regime it is also TIGHT, + # since ``c2 -> 0`` makes ``max_u g -> a + |c1|``. Where c2 is healthy the roots are + # exact, and the argmax must still pass a stationarity test against the axis's own + # derivative bound before it is believed. + m1u = jnp.abs(c1) + 2.0 * jnp.abs(c2) + u = u_stationary_roots(c1, c2) + gv = _g_u(a, c1, c2, u, 0) + i = jnp.argmax(gv) + resid = jnp.abs(_g_u(a, c1, c2, u[i], 1)) + bad = ((jnp.abs(c2) <= 1e-8 * jnp.abs(c1)) + | (resid > 1e-6 * jnp.maximum(m1u, 1e-300))) + gmax = jnp.where(bad, a + jnp.abs(c1) + jnp.abs(c2), gv[i]) + return jnp.log(2.0 * jnp.pi) + gmax + + +def required_bound_grid(amplitude, tol_nats=5.0, m_max=2): + """Bound-grid points that keep the Lipschitz lift ``M10 * delta`` under ``tol_nats``. + + Derived, not tuned, and the same shape as :func:`required_u_nodes`. The lift on a grid + of half-spacing ``delta = pi / n`` is ``M10 * delta``, and ``M10 <= 2 * m_max * A`` for + a table of amplitude ``A``, so ``n >= pi * 2 * m_max * A / tol_nats``. Host side and + static, because JAX needs the shape before it sees the table. + """ + a = max(float(amplitude), 1.0) + return int(np.ceil(np.pi * 2.0 * float(m_max) * a / float(tol_nats))) + 1 + + +def profile_derivative_bounds(C): + """Exact bounds ``(M1F, M2F)`` on ``|F'|`` and ``|F''|`` for the u-profile ``F``. + + The envelope identities are ``F' = E[d_phi g]`` and ``F'' = E[d^2_phi g] + + Var(d_phi g)``, the expectation being under the normalized ``exp(g) du``. So + ``|F'| <= sup|d_phi g| <= M10`` and, since a variable confined to a range of width + ``2 M10`` has variance at most ``M10^2``, ``|F''| <= M20 + M10^2``. Both follow from + the coefficient table alone -- no sample, no fit, and in particular NOT the measured + ``F''`` at a point, which is what an estimate-promoted-to-bound would use here. + """ + m10 = phi_derivative_bound(C, 1) + m20 = phi_derivative_bound(C, 2) + return m10, m20 + m10 * m10 + + +def eval_g2(C, phi, u, order=(0, 0)): + """``d^a_phi d^b_u g`` at matching ``(phi, u)``, from the 2-D table.""" + KP = C.shape[0] + KS = (C.shape[1] - 1) // 2 + k = jnp.arange(KP)[None, :, None] + q = jnp.arange(-KS, KS + 1)[None, None, :] + w = jnp.where(jnp.arange(KP) > 0, 2.0, 1.0)[None, :, None] + a, b = order + phi = jnp.atleast_1d(phi) + u = jnp.atleast_1d(u) + E = jnp.exp(1j * (phi[:, None, None] * k + u[:, None, None] * q)) + return (E * ((1j * k) ** a) * ((1j * q) ** b) * (w * C[None])).sum((1, 2)).real + + +def u_profile(C, phi, n_nodes=U_NODES_PER_CELL, window_sigma=U_WINDOW_SIGMA): + """``F(phi) = log int du exp(g)``, its first two EXACT phi-derivatives, and TWO + THREE fallback counts: how many u cells were integrated whole, how many of those were + also under-sampled for the narrowest STATIONARY scale the coefficients admit, and how + many were under-sampled for the narrower non-stationary scale ``1/M1u``. Only the + second is gated; see the notes beside them for why the first declines every table + there is and why the third is a measure of the gap rather than a usable requirement. + + Differentiating under the integral gives them from the SAME nodes at no extra + evaluation cost: + + F' = E[d_phi g] F'' = E[d^2_phi g] + Var(d_phi g) + + the expectation being under the normalized ``exp(g) du`` on the u axis. That + variance term is why phi cannot inherit the u axis's economy: it grows with + amplitude, so ``F`` sharpens as the signal does even though ``g`` does not. + """ + 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 * k) * w + D = lambda q: (ph * C[:, KS + q]).sum() + a = D(0).real + c1 = D(1) + jnp.conj(D(-1)) + c2 = D(2) + jnp.conj(D(-2)) + + u = jnp.sort(u_stationary_roots(c1, c2)) + 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) + + 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) + return jnp.clip(uc + jnp.clip(step, -0.5, 0.5), lo_c, mid), None + + ustar, _ = lax.scan(_newton, u, None, length=8) + g1s = _g_u(a, c1, c2, ustar, 1) + g2s = _g_u(a, c1, c2, ustar, 2) + # A CLIPPED NEWTON POINT IS NOT A PEAK, however negative the curvature -- the SAME + # defect log_inner_u_integral already gates, reintroduced here because this function + # was written as a fresh copy of that iteration rather than as a call to it. 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 +-window_sigma window on a + # non-stationary point, sizes sigma from the wrong curvature, and can EXCLUDE the true + # maximum -- underestimating F while the docstring calls the derivatives exact. + # Measured in the numpy twin: 18% of cells that g'' < 0 accepts fail this gate, worst + # at |g_u|/M_1 = 0.33. Require stationarity against the axis's own exact derivative + # bound AND interior placement; a cell failing either is integrated WHOLE. + 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)) + sig = jnp.where(peaked, 1.0 / jnp.sqrt(jnp.where(peaked, -g2s, 1.0)), jnp.inf) + lo = jnp.where(peaked, jnp.maximum(ustar - window_sigma * sig, lo_c), lo_c) + hi = jnp.where(peaked, jnp.minimum(ustar + window_sigma * sig, mid), mid) + width = jnp.maximum(hi - lo, 0.0) + + s = jnp.linspace(0.0, 1.0, n_nodes) + uu = (lo[:, None] + width[:, None] * s[None, :]).ravel() # (4n,) + pp = jnp.full(uu.shape, phi) + gg = eval_g2(C, pp, uu, (0, 0)) + gp = eval_g2(C, pp, uu, (1, 0)) + gpp = eval_g2(C, pp, uu, (2, 0)) + wq = jnp.full(n_nodes, 1.0 / (n_nodes - 1)).at[0].mul(0.5).at[-1].mul(0.5) + lw = (jnp.log(jnp.where(width > 0, width, 1e-300))[:, None] + + jnp.log(wq)[None, :]).ravel() + lw = jnp.where(jnp.repeat(width > 0, n_nodes), lw, -jnp.inf) + + m = gg.max() + wt = jnp.exp(gg - m + lw) + Z = wt.sum() + e1 = (wt * gp).sum() / Z + F = m + jnp.log(Z) + ddF = (wt * (gpp + gp * gp)).sum() / Z - e1 * e1 + # how many of the four cells were integrated WHOLE rather than windowed. Reported + # because a fallback cell spreads the same static node count over a wider interval, so + # it is the one place F itself can be inaccurate -- and no bound on this axis can see + # that, since the omitted-mass certificate covers what is outside the regions. + n_fallback = (~peaked).sum() + # ...AND OF THOSE, HOW MANY COULD HAVE HIDDEN A MAXIMUM. The two are not the same + # count and only the second can invert a bound built on F. A cell whose stationary + # point is a MINIMUM has no peak to window; integrating it whole is the design, not a + # shortfall, and its contribution to F is exponentially subdominant to the maximum + # cells, so its quadrature error cannot move F at the scale a certificate cares about. + # A cell with g'' < 0 that failed the stationarity or interior test is the other case: + # a genuine maximum may sit inside it unresolved, F is then UNDERESTIMATED, and a + # Taylor lift applied to an underestimate bounds nothing. + # + # THIS DISTINCTION IS WHY THE OBVIOUS FIX IS WRONG. External review asked for a + # decline whenever any profile evaluation fell back. Every generic table has four + # u-stationary points, two of them minima, so n_fallback >= 2 ALWAYS and that gate + # declines every row unconditionally -- measured: 0 of 2 accepted on cases accurate to + # 1e-5. The finding is real; the remedy as stated is not implementable. + # + # AND "DID IT FALL BACK" IS STILL THE WRONG QUESTION -- measured, it fires on 127 of + # 256 bound-grid points for tables accurate to 1e-5, because an 8-step Newton misses + # the 1e-8 relative residual on plenty of perfectly ordinary maxima. The question the + # bound actually needs answered is whether the whole-cell quadrature was ADEQUATE for + # the sharpest feature the cell can hold, which is review's other remedy and is exact + # here: the u spectrum has two terms, so |d2g/du2| <= |c1| + 4|c2| everywhere, nothing + # is narrower than 1/sqrt(M2u), and a cell of `width` sampled at U_PTS_PER_SIGMA per + # curvature length needs width*sqrt(M2u)*U_PTS_PER_SIGMA nodes. Same derivation as + # required_u_nodes, against the true per-cell curvature instead of an amplitude proxy. + m2u = jnp.abs(c1) + 4.0 * jnp.abs(c2) # exact bound on |d2 g / du2| + need_u = width * jnp.sqrt(m2u) * U_PTS_PER_SIGMA + 1.0 + n_risky = ((g2s < 0.0) & (~peaked) & (need_u > n_nodes)).sum() + # ...and the STRICTER criterion, reported and never gated. Where g is steep but not + # turning, exp(g) varies on 1/M1u rather than 1/sqrt(M2u), and that boundary-layer + # scale -- not the stationary one -- is what the integrand actually has there. This + # is the count against that scale. It is the honest measure of how far the whole-cell + # quadrature is from something that could certify, and gating on it declines rows + # accurate to 1e-5, which is precisely why this axis is described as empirically + # gated rather than bounded. + need_strict = width * jnp.maximum(jnp.sqrt(m2u), m1u) * U_PTS_PER_SIGMA + 1.0 + n_strict = ((~peaked) & (need_strict > n_nodes)).sum() + return F, e1, ddF, n_fallback, n_risky, n_strict + + +def _merge_sorted_intervals(lo, hi, n): + """Merge overlapping 1-D intervals under jit, without data-dependent shapes. + + Sorting by ``lo`` makes merging a running maximum: a new group starts exactly where + an interval begins beyond the running max of the ``hi`` seen so far. Group ids are + then a cumsum, and the merged bounds are segment reductions over a FIXED number of + slots. Empty slots come back as an inverted interval and are dropped by the + ``width > 0`` mask downstream, so nothing needs compaction. + + This is the jittable form of the reference's ``_merge_boxes``; merging is not + tidiness but what stops the mass between two windows being counted twice. + """ + idx = jnp.argsort(lo) + lo, hi = lo[idx], hi[idx] + run = jax.lax.cummax(hi) + fresh = jnp.concatenate([jnp.array([True]), lo[1:] > run[:-1]]) + gid = jnp.cumsum(fresh) - 1 + seg_lo = jax.ops.segment_min(lo, gid, num_segments=n, indices_are_sorted=True) + seg_hi = jax.ops.segment_max(hi, gid, num_segments=n, indices_are_sorted=True) + return seg_lo, seg_hi + + +#: Trapezoid points per curvature length inside a phi region. Not a tolerance: it is the +#: sampling density at which the trapezoid resolves a feature of scale ``1/sqrt(M2F)``. +PHI_PTS_PER_SIGMA = 3.0 + +#: Halving the phi nodes must move the answer by less than this for the integration to be +#: called resolved. A CHOICE, and stated as one -- but not a knife-edge: measured, cases +#: accurate to ~1e-5 move by 1.4e-07 to 2.0e-05, and cases wrong by 0.16-0.78 nats move by +#: 0.64 to 4.6. Five decades separate them and this sits in the middle, so nothing turns +#: on where in the gap it is placed. +PHI_CONVERGENCE_NATS = 1.0e-3 + + +def required_phi_nodes(width, m2f, pts_per_sigma=PHI_PTS_PER_SIGMA): + """Nodes a phi region of ``width`` needs, from the EXACT bound ``|F''| <= m2f``. + + Nothing in the region is narrower than ``1/sqrt(m2f)``, so ``width * sqrt(m2f)`` counts + the curvature lengths it spans and the requirement is that times the sampling density. + Bound, not estimate: ``m2f`` comes from :func:`profile_derivative_bounds`, i.e. from the + coefficient table. + + This is what distinguishes a WINDOWED region from a COVERING one. A window spans a few + ``sigma`` and needs a few tens of nodes at any amplitude; a region spanning the whole + circle spans ``2 pi sqrt(m2f)`` curvature lengths and needs thousands. Both were given + the same fixed 96. + """ + return width * jnp.sqrt(jnp.maximum(m2f, 0.0)) * pts_per_sigma + + +def phi_local_lnI(C, n_seed=PHI_SEEDS, w_sigma=PHI_WINDOW_SIGMA, + n_nodes=PHI_NODES_PER_REGION, u_nodes=U_NODES_PER_CELL, + n_bound=PHI_BOUND_GRID, tol_nats=OUTSIDE_TOL_NATS, + n_slots=None, pt_chunk=PT_CHUNK_DEFAULT): + """``log int dphi int du exp(g)`` with BOTH axes localized, jittable. + + Returns ``(value, ok, info)``. ``ok`` is False when the omitted-mass bound on the phi + axis could not be made small enough; the value is returned either way for diagnosis, + but a value with ``ok=False`` is NOT to be used. + + u is exact on the cell partition; phi is localized around the maxima of the profile + ``F`` using its exact derivatives (see :func:`u_profile`). phi has no algebraic + completeness warrant -- ``F`` is a log-integral, not a trig polynomial -- so the seeds + are targeting only and correctness rests on the certificate below. + + WHAT ``ok`` ACTUALLY ASSERTS, because two reviews found this overstated in two + different places. It is ONE genuine bound and TWO empirical gates, and the chain is + only as strong as its weakest link: + + * ``margin`` IS a bound. Mass outside the covered regions is at most + ``area_outside * exp(sup_outside F)``, with ``sup_outside F`` lifted from a grid by + an exact second-order remainder built from the coefficient table. + * ``resolved`` is NOT. It compares nested quadrature rules, and no rule can see its + own aliases in its own samples, so it certifies the COARSE rule and infers the fine + one. The shifted companion adds the odd multiples of the coarse sampling frequency + and still shares the even ones. Estimates, used only to decline. + * ``u_sizing_ok`` is NOT. Three samples per curvature length is a sampling rule, not + an enclosure of the quadrature error, and it uses the stationary scale + ``1/sqrt(M2u)`` where a boundary layer has the narrower ``1/M1u``. + + And ``margin``'s own soundness runs through ``Fb``, which the empirical gates are what + stand behind -- so the chain is empirical END TO END. THIS PATH IS EMPIRICALLY GATED, + NOT FAIL-CLOSED, and it must not be described as certified. Bounds tight enough to + replace the gates were looked for and do not appear to exist at usable tightness: the + exact ``M2F`` requirement demands 3.8e3-2.3e4 phi nodes for cases right to 1e-4, and the + ``1/M1u`` requirement declines rows right to 1e-5. Both collapse to "always decline", + which is why the claims are narrowed instead. + + READ THIS BEFORE PROMOTING THIS PATH -- AND THE COST ARGUMENT BELOW IS WITHDRAWN. + + THIS FUNCTION LOCALIZES ON THE WRONG OBJECT. It Newton-iterates on the maxima of + ``F(phi) = log int du exp(g)`` from ``PHI_SEEDS`` arbitrary seeds. ``F`` is a + log-integral and has no completeness warrant -- but ``g`` ITSELF DOES, and it is the + same warrant psi has. The orbital phase enters the modes as ``e^{-i m phi}``, so ``A`` + carries phi-harmonics to ``m_max`` and ``B``, being quadratic, to ``2 m_max``. The + combined table's ``k_max = KP-1 = 2 m_max`` is therefore EXACT, and ``dg/dphi = 0`` + under ``z = e^{i phi}`` is a polynomial of degree ``2 k_max``. Knowing the mode + content fixes the stationary count; the 2-D system with ``dg/du = 0`` has a + mixed-volume bound of ``16 k_max``. The numpy reference already does this -- + :func:`~RIFT.likelihood.joint_angle_peak_local.enumerate_modes` solves the algebraic + system -- and it finds MORE maxima than the seeded search: 13 against 8 at ``KP=5``, + 30 against 20 at ``KP=13``. + + So the earlier conclusion here -- that certifying phi costs more than the dense grid + because ``n_bound ~ A`` -- was reasoning about a construction chosen in this file, not + about the phi axis. Seeded algebraically the region count is mode-order-bounded and + provable, and the cost comparison has to be redone on that basis. It is NOT restated + here in a corrected form, because two successive versions of it were wrong; the + measurements are on the PR and the argument needs rebuilding, not patching. + + MEASURED LIMITS OF THE SHIPPED CONSTANTS (Blackwell, jax 0.9.2, x64): + * ``PHI_SEEDS = 32`` is an undocumented assumption about mode content. At + ``m_max = 2`` the region count plateaus by 32 seeds (7-8 regions, unchanged at + 64 and 128). At ``m_max = 6`` it does NOT: 32 seeds find 14-19 regions where 64+ + find 19-21. FAIL-CLOSED -- every such case declines, none returns an accepted + wrong value -- and the missed regions are subdominant, changing the value by less + than 1e-5. At ``m_max = 6`` the rule declines universally, so high mode content + is outside its reach for reasons beyond the seed count. + * 94-97% of the phi work is on EMPTY slots: ``2 * PHI_SEEDS = 64`` static slots are + allocated and every one is evaluated in full, while production tables use 2-4. + Since the midpoint companion was added the region grid is evaluated TWICE per + slot -- ``n_nodes`` trapezoid nodes and ``n_nodes - 1`` midpoints -- so this waste + now costs twice what the figures below were measured at. That is the price of a + convergence check that can see its own leading error term; no subset of an + existing node set can see aliasing at its own sampling harmonic. + That is the price of static shapes without an enumeration; it is not recoverable + by shrinking the allocation, because shrinking starves the seeds as well and + converts silent waste into declines (measured: 2 regions accept at 8 seeds and + decline at 4). + * Per-evaluation device memory was 0.098 GiB against the dense path's 0.001 GiB -- + MEASURED BEFORE the midpoint companion, so the region-quadrature part of it has + since roughly doubled and the figure has not been re-measured on a GPU. It + it scales LINEARLY with the vmap product because nothing here chunks. + :func:`joint_lnL_phi_dense` bounds its own memory with ``lax.scan`` over + ``phi_chunk`` and is flat in ``n_phi`` (0.39 GiB at 256, 1024 and 4096 alike). + The (2,+-2) tables carry an EXACT ORDER-4 SYMMETRY, which reduces the bound grid to a + QUARTER domain -- worth 4x against a shortfall of 80x, real but not the answer. + Measured on the exponent itself, which is the object this code evaluates, and not on + the coefficient table it is built from: + + S : (phi, u) -> (phi + pi/2, u + pi) generator, order 4 + + rung 1 S^1..S^4 deviations 2.2e-15 2.6e-15 4.0e-15 3.8e-15 + rung 3 2.8e-15 2.8e-15 4.6e-15 4.3e-15 + + ``S^2 = (phi + pi, u)`` is therefore also exact, which is where the phi half-period + comes from; ``(phi, u + pi)`` and ``(phi + pi, u + pi)`` are NOT symmetries (relative + deviation 1.32 each), so there is no u half-period on its own. Every maximum carries + exactly FOUR copies and the enumeration confirms it: rung 3's four maxima are ONE orbit + of four (one distinct maximum, which is why they are exactly degenerate), and rung 1's + eight are TWO orbits of four, matching its two distinct exponent values. + + TWO EARLIER VERSIONS OF THIS NOTE WERE WRONG HERE, in opposite directions, and both + times the CONCLUSION that ``F`` is pi-periodic survived: first ``(phi+pi, u+pi)`` with + multiplicity four, taken from a coefficient parity measured in another convention; then + ``(phi+pi, u)`` with multiplicity two, from testing only the shifts I had thought to + list. The generator was never among them. Enumerate the group from the maxima's own + offsets rather than guessing which shifts to test. + """ + prof = lambda p: u_profile(C, p, n_nodes=u_nodes) + # SEEDS ARE TARGETING, NOT AN ENUMERATION, and this file no longer pretends otherwise. + # A uniform linspace has no completeness claim and measurably under-resolves: at + # m_max = 6 it finds 14-19 regions where 64+ seeds find 19-21. The omitted-mass bound + # is what stands behind that -- a missed region raises area_outside and the row + # declines -- and it is the only thing that does. + # + # AN IN-KERNEL ALGEBRAIC SEEDER WAS BUILT HERE AND HAS BEEN REMOVED. It solved the + # Sylvester resultant on the roots of unity inside the traced function, which bought + # completeness of the seed set but is exactly the substitution + # :mod:`RIFT.likelihood.bivariate_trig_stationary` says must not be made: that module + # is the authoritative enumerator, it is fail-closed on the BKK root count, Jacobian + # conditioning, torus classification and cross-projection agreement, and it is host + # side because none of those has an honest static-shape JAX transcription. A second, + # weaker enumerator with no ``ok`` flag living in this tree was the duplication both + # integration reviews asked not to carry. + # + # The honest way back to algebraic seeds is a host-built fixed-capacity plan passed IN + # as ``seeds``, built from ``enumerate_torus_maxima`` and declining when its ``ok`` is + # false. That is not implemented; until it is, these are a grid and are labelled one. + seeds = jnp.linspace(0.0, 2.0 * jnp.pi, n_seed, endpoint=False) + + def _newton(p, _): + _, d1, d2, _, _, _ = jax.vmap(prof)(p) + step = jnp.where(d2 < 0, -d1 / jnp.where(d2 < 0, d2, -1.0), 0.0) + return jnp.mod(p + jnp.clip(step, -0.3, 0.3), 2.0 * jnp.pi), None + + p, _ = lax.scan(_newton, seeds, None, length=24) + F, d1, d2, n_fb, _, _ = jax.vmap(prof)(p) + peaked = d2 < 0.0 + sig = jnp.where(peaked, 1.0 / jnp.sqrt(jnp.where(peaked, -d2, 1.0)), 0.0) + + # non-maxima are pushed past every real interval so they form empty groups; no + # tolerance decides membership, which is deliberate -- a threshold on |F'| would be + # exactly the estimate-promoted-to-bound this design refuses. + big = 1.0e6 + lo = jnp.where(peaked, p - w_sigma * sig, big) + hi = jnp.where(peaked, p + w_sigma * sig, big) + + # SPLIT AT THE SEAM BEFORE MERGING, for the reason the numpy reference had to: a + # linear merge never joins a window near 0 to one near 2 pi, yet every region is + # integrated at mod(., 2 pi), so both cover both peaks and the mass is counted twice + # (+log 2, accepted, because the error is inside the regions). Each interval yields + # AT MOST two pieces, so 2*n_seed slots is a static bound and nothing has to be + # compacted; a piece that does not exist is emitted empty and drops out downstream. + wdt = jnp.clip(hi - lo, 0.0, 2.0 * jnp.pi) + # A WINDOW THAT ALREADY SPANS A FULL CIRCUIT HAS NO SEAM TO SPLIT AT, and splitting + # one anyway made the region count -- and the ANSWER -- a one-ulp coin flip. After the + # clip `wdt` is EXACTLY 2 pi, so the two pieces are [a0, 2 pi] and [0, a0 + 2 pi - 2 pi] + # and they are adjacent by construction. In floating point they are adjacent only when + # (a0 + 2 pi) - 2 pi comes back >= a0. ONLY THE LOW SIDE BREAKS, and an earlier version + # of this note had the criterion wrong: the round-trip landing one ulp ABOVE a0 is fine, + # the pieces then overlap and _merge_sorted_intervals folds them. The loss is a0's low + # bits -- ulp(a0 + 2 pi) is 1.78e-15 whatever a0 is, against ulp(a0) from 6.9e-18 to + # 8.9e-16, so the sum is 2x to 64x coarser. Measured on pure float64, no jax involved: + # 34.4% of a0 drawn uniformly on the circle round LOW. + # + # When they do, the pieces are one ulp apart, the merge (which joins only exactly + # touching intervals, by design -- no tolerance decides membership here) leaves them + # separate, `total` comes out 8.9e-16 below 2 pi, and the `wrapped` clamp below does not + # fire. The rule then runs a two-region trapezoid with a seam instead of the PERIODIC + # trapezoid on the full circle, and a periodic trapezoid is spectrally accurate where a + # seamed one is O(h^2): on the harmonic-alias table of + # test_the_halving_check_is_blind_at_the_sampling_harmonic, 2.1e-3 nats wrong instead of + # 2.1e-8, a factor of 1e5, from a two-ulp difference in the Newton fixed point. + # + # jax 0.9.2 and 0.10.2 land on opposite sides of it -- same host, same python, same + # numpy -- which is how this arrived as an environment-dependent test failure rather + # than as a bug. + # + # THIS PATH IS NOT FAIL-CLOSED AGAINST THE DEFECT, and an earlier version of this note + # claimed it was ("both sides return ok=False, so nothing was ever accepted wrong"). + # That is true of the shipped fixture and false in general. Sweeping the peak location + # pre-fix on jax 0.9.2, the seam fires WITH ok=True: 5 of 96 at kappa=300 and 400, 12 of + # 96 at kappa=550 and 700. What bounds the accepted error is the convergence probe, not + # the decline -- the worst accepted case found was 6.0e-06 nats, so the CONCLUSION that + # nothing was accepted materially wrong survives, but not for the reason first given. + # + # The fix is exact and not a tolerance: a full circuit is anchored at 0 and emitted as + # the single piece [0, 2 pi]. This is the tree's existing idiom, not a new one -- + # multipeak_planner._periodic_segments anchors the same way before splitting. The numpy + # twin instead closes the seam AFTER the fact with a 1e-12 tolerance + # (_merge_boxes' caller in joint_angle_peak_local); this path had neither. + full_circuit = peaked & (wdt >= 2.0 * jnp.pi) + a0 = jnp.where(peaked, + jnp.where(full_circuit, 0.0, jnp.mod(lo, 2.0 * jnp.pi)), + big) + crosses = peaked & (~full_circuit) & (a0 + wdt > 2.0 * jnp.pi) + lo2 = jnp.concatenate([a0, + jnp.where(crosses, 0.0, big)]) + hi2 = jnp.concatenate([jnp.where(crosses, 2.0 * jnp.pi, a0 + wdt), + jnp.where(crosses, a0 + wdt - 2.0 * jnp.pi, big)]) + n_out = int(2 * PHI_SEEDS if n_slots is None else n_slots) + seg_lo, seg_hi = _merge_sorted_intervals(lo2, hi2, n_out) + n_seed = n_out + # There are always more slots than groups, and an EMPTY slot comes back from the + # segment reductions as (+inf, -inf). Masking its weight is not enough: the node + # positions are still built from it, jnp.mod(inf, 2 pi) is NaN, and NaN * 0 is NaN, + # so the poison reaches the sum through a term that was supposed to be switched off. + # Neutralize the POSITION, not just the weight. + seg_lo = jnp.where(jnp.isfinite(seg_lo), seg_lo, 0.0) + seg_hi = jnp.where(jnp.isfinite(seg_hi), seg_hi, 0.0) + width = jnp.clip(seg_hi - seg_lo, 0.0, 2.0 * jnp.pi) + + # CLAMP TO ONE CIRCUIT. At low amplitude sigma is huge and the windows span more + # than 2 pi; integrating that literally wraps the circle and counts the same mass + # repeatedly (measured +1.84 nats, a factor of 6.3, on real tables in the numpy + # reference -- and ACCEPTED, because a region covering everything leaves nothing + # outside for the certificate to object to). + # close the circle: if some piece ends at 2 pi and another starts at 0 they are one + # region. Left unjoined they are still DISJOINT, so nothing is double-counted -- the + # only cost is one extra region and a seam the quadrature treats as an edge. + total = width.sum() + wrapped = total >= 2.0 * jnp.pi + seg_lo = jnp.where(wrapped, jnp.where(jnp.arange(n_seed) == 0, 0.0, big), seg_lo) + width = jnp.where(wrapped, + jnp.where(jnp.arange(n_seed) == 0, 2.0 * jnp.pi, 0.0), width) + + s = jnp.linspace(0.0, 1.0, n_nodes) + pp = (seg_lo[:, None] + width[:, None] * s[None, :]).ravel() + # ROLLED, because this is the axis that decides whether the rule is usable. An + # unrolled (n_slots * n_nodes) grid costs `points * 4 * u_nodes * KP * (2KS+1) * 16` + # bytes live in eval_g2's intermediate -- about a gigabyte per distance node at the + # production u count -- and the distance axis then multiplies it. Chunking here is + # what lets `joint_lnL_phi_local` be wired at all; the values are identical, only the + # peak allocation changes. + Fv, nfb_v, nrisk_v, nstrict_v = _prof_scan(prof, jnp.mod(pp, 2.0 * jnp.pi), pt_chunk) + # EMPTY SLOTS MUST NOT VOTE. A slot with no region is neutralized for the VALUE by + # zeroing its position and masking its weight, but its nodes are still evaluated -- at + # the artificial point phi = 0 -- and their fallback counts were summed with the rest. + # A risky cell there could decline a row whose every contributing node was adequate, + # and the reported counters were contaminated the same way. Measured: 5 risky at + # n_slots=2 rising to 176 at n_slots=8 while the region count only went 2 -> 4, so the + # growth was entirely empty slots and raising the allocation alone could flip a row. + # Found independently by this session's review and by external review. + live_pt = jnp.repeat(width > 0, n_nodes) + nfb_v = jnp.where(live_pt, nfb_v, 0) + nrisk_v = jnp.where(live_pt, nrisk_v, 0) + nstrict_v = jnp.where(live_pt, nstrict_v, 0) + wq = jnp.full(n_nodes, 1.0 / (n_nodes - 1)).at[0].mul(0.5).at[-1].mul(0.5) + lw = (jnp.log(jnp.where(width > 0, width, 1e-300))[:, None] + + jnp.log(wq)[None, :]).ravel() + lw = jnp.where(jnp.repeat(width > 0, n_nodes), lw, -jnp.inf) + value = jax.scipy.special.logsumexp(Fv + lw) + + # NESTED, SO NOTHING IS EVALUATED THAT DOES NOT ENTER THE ANSWER. The first version + # of the companion evaluated a SECOND grid of n-1 midpoints used only for the probe + # and then thrown away -- 1.85x the cost for a diagnostic. With an odd n_nodes the + # one grid already contains both sub-rules: the even indices are a trapezoid at half + # the density, and the odd indices are exactly ITS midpoints. Same evaluation count, + # and the returned value is the FINE rule rather than the coarse one. + # + # That is not only cheaper, it is more accurate where it matters. On the aliasing + # counterexample the old arrangement returned the 97-node rule, which is 0.02017 nats + # wrong, and used 96 extra points to notice. The nested arrangement spends the same + # points on the answer and returns it correct to 1e-6, with the probes still firing + # because the COARSE rule was bad. + Fr = Fv.reshape(-1, n_nodes) + # CONVERGENCE, MEASURED, FROM THE NODES ALREADY EVALUATED. n_nodes is odd, so indices + # 0, 2, ... n-1 span the same interval at double the spacing: a half-resolution estimate + # for free, no second integration. This replaces two gates that did not work -- the + # exact M2F bound demands 3.8e3-2.3e4 nodes and declines cases right to 1e-4, and a + # local-curvature rule declines cases right to 1e-5, because the trapezoid on a periodic + # integrand converges spectrally and any real-space "points per sigma" is far too + # conservative for a region spanning the circle. + # + # It is an ESTIMATE of the discretization error, not a bound, and is used ONLY to + # decline -- the conservative direction. It cannot certify; it can only refuse. + n_h = (n_nodes + 1) // 2 + whq = jnp.full(n_h, 1.0 / (n_h - 1)).at[0].mul(0.5).at[-1].mul(0.5) + lwh = (jnp.log(jnp.where(width > 0, width, 1e-300))[:, None] + + jnp.log(whq)[None, :]).ravel() + lwh = jnp.where(jnp.repeat(width > 0, n_h), lwh, -jnp.inf) + value_half = jax.scipy.special.logsumexp(Fr[:, ::2].ravel() + lwh) + conv = jnp.abs(value - value_half) + + # THE HALVING CHECK CANNOT SEE THE ERROR THAT MATTERS, and no subset of the nodes + # already evaluated ever can. A periodic n-interval rule's error is the sum of the + # aliased harmonics at multiples of n; the n/2 rule aliases at multiples of n/2, which + # CONTAINS every multiple of n, so the two share the whole leading term and `conv` + # cancels it. Detecting content AT the sampling harmonic requires points the rule did + # not sample -- this is Nyquist, not an implementation shortfall. + # + # The companion is the composite MIDPOINT rule on the same regions: n-1 nodes at the + # interval midpoints, uniform weight. On a periodic region it is exactly the + # half-shifted trapezoid, whose error is sum (-1)^k c_{kn}, so the difference from + # `value` is 2 * sum_{k odd} c_{kn} -- the leading alias itself, the term halving + # cancels. On a window it is the classic O(h^2) companion with error -1/2 the + # trapezoid's, so the difference is 1.5x the true error: an estimator, not a bound, + # used only to decline. + # + # Adversarial review supplied the case this closes: F = 1000 cos(phi - pi/96) on the + # full circle at n = 97. The 96- and 48-interval rules agree to 1.1e-13 while both are + # 0.02017 nats wrong -- the phase makes the c_48 alias vanish exactly and leaves c_96. + # The midpoint companion reads 3.99e-2 and declines. On every accurate case measured + # (kappa 4.5-1e4, windows of 3-12 sigma, and the same table resolved at n = 385) it + # reads 0.0 to 1.3e-5, so it does not cost a single good row. + n_m = n_nodes // 2 + lwm = jnp.broadcast_to((jnp.log(jnp.where(width > 0, width, 1e-300)) + - jnp.log(float(n_m)))[:, None], + (width.shape[0], n_m)).ravel() + lwm = jnp.where(jnp.repeat(width > 0, n_m), lwm, -jnp.inf) + value_mid = jax.scipy.special.logsumexp(Fr[:, 1::2].ravel() + lwm) + # compared against the COARSE rule, which is the rule it is the midpoint companion + # OF. Comparing it to the fine value would conflate a shift with a refinement. + conv_shift = jnp.abs(value_half - value_mid) + + # ------------------------------------------------- the phi omitted-mass bound + # THE ONE PART OF `ok` THAT IS A BOUND. Everything else gating this return is an + # empirical convergence estimate; see the note in the docstring. + # WITHOUT THIS THE RETURN VALUE IS AN ESTIMATE WEARING A LIKELIHOOD'S CLOTHES. The + # seeds are targeting, not an enumeration -- phi has no algebraic completeness warrant + # because F is a log-integral, not a trig polynomial -- so a missed maximum or an + # unconverged seed is silently omitted and a finite number comes back regardless. + # External review found this exposed with no bound, no validity result and no fallback + # signal, and it is the house rule of this whole family violated in its own code. + # + # The bound: mass outside the covered regions is at most + # area_outside * exp(sup_outside F), + # and sup_outside F is obtained from a grid of F values LIFTED by a true remainder, + # never from the grid maximum itself -- a grid max is a LOWER bound on a supremum and + # the gap grows with amplitude. Both F and F' come back from u_profile at no extra + # cost, so the lift is second order: + # F(x) <= F(x_i) + |F'(x_i)| * delta + M2F * delta^2 / 2, delta = half spacing + # with M2F from profile_derivative_bounds, i.e. from the coefficient table alone. + # A first-order Lipschitz lift was tried first in the numpy twin and is USELESS at + # amplitude -- it put the bound above the integral by +1225 nats. + gb = jnp.linspace(0.0, 2.0 * jnp.pi, n_bound, endpoint=False) + delta = jnp.pi / n_bound # half of the grid spacing + m1f, m2f = profile_derivative_bounds(C) + # THE BOUND GRID NO LONGER RUNS THE U QUADRATURE, and that is what makes this + # certificate usable rather than merely correct. It used to call the full profile at + # every point and lift by ``|F'| delta + M2F delta^2 / 2``, which failed twice over. + # M2F is ~99.5% the ``M10^2`` variance term -- a worst-case bound on a variance that + # CONCENTRATES as amplitude rises, so it is loosest exactly where the physics is + # tightest. Measured at amplitude 3e4: that lift sat 2.7e5 nats above the integral + # while the TRUE margin was about -66, i.e. the row deserved to be accepted by a wide + # margin and was declined by five orders of magnitude. Refining the grid to fix it is + # what kills the process, because ``n_bound * u_nodes`` both grow with amplitude. + # + # :func:`sup_g_bound` replaces the whole construction: an exact upper bound on F for + # four quartic roots per point, Lipschitz in M10 rather than M10^2 so refinement is + # linear, and independent of the u quadrature -- which also retires the review finding + # that the lift could be applied to a profile the fallback had underestimated. There + # is no longer a profile value on this grid to underestimate. + hb = jax.vmap(lambda q: sup_g_bound(C, q))(gb) + ub = hb + m1f * delta + + # A GRID POINT COUNTS AS OUTSIDE UNLESS ITS WHOLE delta-BALL IS COVERED. Testing the + # point alone leaves a band of width delta beside every region boundary belonging to + # no test at all, and the bound would then be a bound on the wrong set. Regions are + # therefore ERODED by delta before the test, which over-estimates the outside -- the + # safe direction. A region already spanning the circle stays covering: that is the + # low-amplitude case where the rule has degenerated into the dense grid on purpose, + # and eroding it would report an uncovered band and decline every such row. + full = width >= 2.0 * jnp.pi - 1e-12 + eff_lo = jnp.where(full, -1.0, seg_lo + delta) + eff_hi = jnp.where(full, 2.0 * jnp.pi + 1.0, seg_lo + width - delta) + d = gb[None, :] - eff_lo[:, None] + covered = (((d >= 0.0) & (gb[None, :] <= eff_hi[:, None])) + | ((d + 2.0 * jnp.pi >= 0.0) + & (gb[None, :] + 2.0 * jnp.pi <= eff_hi[:, None]))).any(axis=0) + + area_outside = jnp.clip(2.0 * jnp.pi - width.sum(), 0.0, 2.0 * jnp.pi) + sup_outside = jnp.max(jnp.where(covered, -jnp.inf, ub)) + outside = jnp.where(area_outside > 0.0, + jnp.log(jnp.where(area_outside > 0.0, area_outside, 1.0)) + + sup_outside, + -jnp.inf) + # AN EMPTY OUTSIDE IS NOT A CORRECT ANSWER. area_outside = 0 says nothing was left + # OUT; it says nothing whatever about the quadrature INSIDE, and the two were being + # conflated -- a full cover gave margin = -inf and an unconditional accept. Measured + # at KP=13, amplitude 1e2: full cover, margin -inf, accepted, value 0.196 nats wrong. + # The same conflation cost the numpy reference 0.36 nats on production tables. + # + # So the accept now also requires that every non-empty region is RESOLVED at the node + # count actually used. The requirement is a bound, not an estimate: nothing in a + # region is narrower than 1/sqrt(M2F), so a region spanning `width` needs + # width*sqrt(M2F) curvature lengths sampled. A windowed region spans a few sigma and + # passes at any amplitude; a region spanning 2 pi does not, which is exactly the case + # that was being accepted wrongly. + # WHAT MAKES 96 NODES DEFENSIBLE IS THE WINDOW, NOT THE COUNT. A region of +-w_sigma + # spans 2*w_sigma curvature lengths whatever the amplitude, so 2*w_sigma*PTS_PER_SIGMA + # = 72 nodes resolve it and 96 has margin -- that is where the constant came from, and + # it holds for as long as a region IS a window. + # + # It stops holding when the rule stops localizing. The `wrapped` branch above fires + # when the windows already span the circle and replaces them with ONE region of width + # 2 pi: that is the rule degenerating into a dense grid on purpose, and 96 nodes across + # 2 pi is not the same claim as 96 nodes across 24 sigma. It is also exactly the branch + # that leaves area_outside = 0 and so would otherwise accept unconditionally. + # + # Sizing this from the exact bound M2F instead was tried and is useless: M2F is 99.5% + # the M10^2 variance term, so it demands 3.8e3 - 2.3e4 nodes for cases that are right to + # 1e-4 at 96 and would decline everything. A bound too loose to distinguish the good + # case from the bad one cannot be the gate, however true it is. + # THE GATE IS SHARPNESS, and it is the numpy reference's criterion: _log_box_integral + # sizes each box from the LOCAL curvature, so a region of `width` carrying a feature of + # scale 1/sqrt(|F''|) needs width*sqrt(|F''|)*PTS_PER_SIGMA nodes. + # + # For a WINDOW this is automatic and amplitude-free: width = 2*w_sigma/sqrt(|F''|), so + # the requirement is 2*w_sigma*PTS_PER_SIGMA = 72, which is where 96 came from. For a + # region that grew -- merged, or the whole circle after `wrapped` -- the width no longer + # tracks the curvature and the requirement can exceed 96. Measured: at amplitude 4.5 a + # full circle needs ~40 nodes and is right to 1e-5; at amplitude 1e2 with KP=13 it needs + # ~190 and is 0.196 nats wrong at 96. The gate separates exactly those. + # + # M2F was tried as the curvature and is useless here: 99.5% of it is the M10^2 variance + # term, so it demands 3.8e3-2.3e4 nodes for cases right to 1e-4 and declines everything. + # A bound too loose to tell the good case from the bad one cannot be the gate. It is + # still reported, because it IS a bound and the measured curvature is not. + # THE HALVING CHECK IS BLIND TO ITS OWN LEADING ERROR TERM, and that has to be closed + # by an assumption made explicit rather than left implicit. The n-node and n/2-node + # trapezoids share EVERY aliased harmonic at multiples of n, so `conv` measures the + # n/2 aliasing and infers the n aliasing from smoothness. Adversarial review built the + # counterexample: a table with phi-content at exactly harmonic n makes F periodic on the + # node spacing, both rules sample one phase, conv comes back at 1e-7 and the value is + # 0.02-0.066 nats wrong -- accepted. + # + # The assumption is enforceable here because the mode content is EXACT: g is a trig + # polynomial in phi of degree k_max = KP-1 = 2 m_max, so requiring the node count to + # Nyquist-resolve k_max rules out content at the sampling harmonic by construction. + # Production (k_max = 4) needs 8 and has 97; the counterexample (k_max = 96) needs 192, + # has 97, and now DECLINES instead of accepting. This is the phi warrant paying for + # itself a second time. + # NECESSARY, NOT SUFFICIENT -- AND THE EARLIER NOTE HERE CLAIMED OTHERWISE. It said + # that Nyquist-resolving k_max "rules out content at the sampling harmonic by + # construction", and that is false: the warrant is a statement about `g`, while the + # outer trapezoid integrates exp(F) with F = log int du exp(g). Neither F nor exp(F) + # is band-limited because g is. The counterexample above has k_max = 1, passes this + # guard trivially at 97 > 2, and is still 0.02 nats wrong. The guard is kept because + # a rule that cannot resolve g certainly cannot resolve exp(F), but what actually + # closes the aliasing family is `conv_shift`, which samples points this rule does not. + k_max = C.shape[0] - 1 + alias_safe = n_nodes > 2 * k_max + need_max = jnp.max(jnp.where(width > 0, required_phi_nodes(width, m2f), 0.0)) + resolved = ((conv < PHI_CONVERGENCE_NATS) + & (conv_shift < PHI_CONVERGENCE_NATS) + & alias_safe) + margin = outside - value + + # THE OUTSIDE BOUND MAY NOT LIFT A PROFILE THAT WAS ITSELF UNDERESTIMATED. `ub` is + # Fb + |d1b| delta + M2F delta^2 / 2, an upper bound on the true F outside the cover + # ONLY IF Fb and d1b are the true profile at the bound-grid points. When a u cell + # fails u_profile's stationarity gate it is integrated WHOLE at the same node count -- + # the branch that function documents as able to underestimate F -- and lifting an + # underestimate does not bound anything. The count was being discarded at this call + # entirely, so a row could be accepted on a non-conservative certificate with no + # signal that it had happened: info["n_u_fallback"] carried only the Newton-seed + # evaluation, not this one and not the quadrature grid. + # + # Fail closed on the bound grid, because that is where the certificate's soundness + # lives. The quadrature and seed grids are reported instead of gated: a fallback there + # perturbs the VALUE, which `conv`/`conv_shift` already measure, rather than inverting + # the direction of a bound. + # + # The gate is the RISKY count, not the fallback count, and the difference decides + # whether this function returns anything at all. Gating on every whole-cell + # integration declines universally -- two of the four u cells hold minima in any + # generic table -- so the count that matters is the cells with negative curvature that + # failed the stationarity or interior test, which are the ones that can hide a maximum + # and underestimate Fb. See u_profile for why the other two are safe. + # NOT AN ERROR BOUND, AND NO LONGER NAMED AS IF IT WERE. `need_u` is + # width*sqrt(M2u)*U_PTS_PER_SIGMA: bounding |d2g/du2| identifies the narrowest + # STATIONARY scale the coefficients admit, but choosing three samples per scale is a + # sampling rule and does not enclose the quadrature error. Review is right that + # calling the result `bound_exact` promoted an estimate into a certificate. + # + # It also misses the non-stationary case: where g is steep but not turning, exp(g) + # varies on 1/M1u, not 1/sqrt(M2u), and that is the scale the integrand actually has + # in a boundary layer. The numpy twin says the same thing at its own u integral and + # leaves a measured residual. `n_u_understood_bound` below reports the count against + # THAT criterion. It is deliberately reported and not gated: applying it declines the + # amplitude-19 case that is accurate to 1e-5, so it would be a wall rather than a + # requirement -- which is exactly the evidence that this axis is empirically gated and + # not certified, and it belongs in the info dict where a caller can see it. + u_sizing_ok = nrisk_v.sum() == 0 + ok = (margin < tol_nats) & resolved & u_sizing_ok + + info = {"margin": margin, + "area_outside": area_outside, + "sup_outside": sup_outside, + "n_phi_regions": (width > 0).sum(), + # the cover itself, so the outside bound can be tested against the set it is a + # bound ON. A soundness check that compares it to the GLOBAL sup of h instead + # reads ~w_sigma^2/2 = 72 nats low and condemns a correct bound -- which is + # exactly what happened here before these were exported. + "seg_lo": seg_lo, + "seg_width": width, + # INTERNAL accuracy, which the certificate above CANNOT see: it bounds the + # mass left OUTSIDE the regions and says nothing about the quadrature inside + # one. Reported separately and never folded into `margin`. + "n_u_fallback": n_fb.sum(), + # THE QUADRATURE GRID IS WHERE THIS BELONGS NOW. It used to be read off the + # bound grid, because that was where an underestimated profile could invert an + # upper bound. sup_g_bound removed that exposure entirely, so the remaining + # question is whether the quadrature that produced `value` was adequate -- and + # that is a property of the grid `value` came from. + "n_u_fallback_quad": nfb_v.sum(), + "n_u_risky_quad": nrisk_v.sum(), + # the stricter 1/M1u criterion: reported, never gated. See u_profile. + "n_u_understood_quad": nstrict_v.sum(), + "u_sizing_ok": u_sizing_ok, + # INTERNAL accuracy, reported beside the omitted-mass margin and never folded + # into it: they are independent failures and both are needed. + # the M2F-derived requirement is a TRUE bound and is reported; it is not the + # gate, because it is too loose to separate the good case from the bad one. + "phi_nodes_needed": need_max, + "phi_convergence": conv, + # the companion rule that samples points the trapezoid does not; this is the + # one that closes the aliasing family, conv alone cannot. + "phi_convergence_shift": conv_shift, + # separate from conv: conv can be small because the check is blind, and this + # says whether it was entitled to be believed at all. + "phi_alias_safe": jnp.asarray(alias_safe), + "phi_resolved": resolved} + return value, ok, info diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/multipeak_planner.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/multipeak_planner.py new file mode 100644 index 000000000..73e5d36ae --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/multipeak_planner.py @@ -0,0 +1,1525 @@ +"""Diagnostic planner for joint peak-local JAX marginalization. + +This module deliberately exposes an opt-in seam rather than changing the +production likelihood dispatch. The fixed-shape device controller in +:mod:`all_axis_peaklocal` is the canonical four-axis path; the shared host +primitives (norm-table summary, harmonic lattice, distance profile, angular +field) are imported from there. What remains here is the diagnostic variant: +symmetry-orbit start ranking, strict sequential Newton/polish refinement, +axis-aligned overlap partition, and the frozen hierarchical cover. Its primary path tests whether a small, +mode-order-sized start portfolio plus empirical enrichment can replace global +time/angle/distance work while retaining a finite exact/dense reserve. + +The planner keeps three statements separate: + +* U,V/Q structure proposes a small set of four-dimensional optimizer starts; +* JAX hill climbing refines those starts, without claiming mode completeness; +* a richer structural tier repeats placement and the overlap-partitioned local + integral; agreement and local diagnostics warrant the result, otherwise the + caller-supplied finite reserve is returned. + +The angular lattice resolves the finite coefficient polynomial, not +``exp(lnL)``. It is targeting only. The optional hierarchical cover below is +a frozen diagnostic: real 22 tables showed it remained thousands of nats too +loose after 5,000 boxes, so it is deliberately absent from the runtime +decision. It bounds only the finite reflected coefficient model; a real ILE +caller still owns the guard-sample warrant connecting that model to physical +support. Any local decline returns the reserve, never a waveform failure. + +Returning the reserve because a diagnostic failed its budget and returning it +because the planner raised are different events, and this module reports them +separately: see ``MultiPeakResult.decline_kind`` and ``fault``. Only the +second is logged, and only the second is made fatal by ``fail_on_fallback``. +""" + +import heapq +import logging +import math +from typing import NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np +from scipy.special import logsumexp as scipy_logsumexp + +# The host-side primitives shared with the device controller live in +# all_axis_peaklocal, which is canonical. This module keeps only the +# diagnostic-seam variants that differ in semantics (symmetry-orbit start +# ranking, strict sequential Newton/polish refinement, axis-aligned overlap +# partition, and the frozen hierarchical cover). +from .all_axis_peaklocal import ( # noqa: E402 + UVHarmonicSummary as UVQSummary, + _angular_field as _angular_field_jax, + _distance_profile_numpy as _distance_profile, + _harmonic_lattice, + _kp_weights_numpy as _kp_weights, + _periodic_distance, + _validate_tables, + summarize_uv_norm_table, +) + + +# A planner FAULT is reported here rather than through ``warnings.warn``. A +# warning's delivery and its fatality are both governed by process-global +# filters that this module does not own: under ``-W error::RuntimeWarning`` the +# warn call itself raises, which would make the DEFAULT path +# (``fail_on_fallback=False``) fatal, and under the default filters the warnings +# registry de-duplicates per call site, so a campaign of identical faults from +# one call site reports once instead of once per call. A logger has neither +# property. The primary channel is still the returned record +# (``decline_kind``/``fault``), which no filter can suppress. +logger = logging.getLogger(__name__) + + +__all__ = [ + "UVQSummary", + "HarmonicSymmetry", + "StartPortfolio", + "CoverReport", + "LocalIntegralReport", + "MultiPeakResult", + "FallbackFault", + "MultiPeakFallbackError", + "DECLINE_DIAGNOSTIC", + "DECLINE_FAULT", + "summarize_uv_norm_table", + "infer_harmonic_symmetry", + "rank_joint_starts_from_uvq", + "refine_joint_starts_jax", + "select_refined_modes", + "axis_local_geometry", + "integrate_refined_modes_tensor", + "multipeak_local_marginalize", + "hierarchical_union_cover", +] + + +class HarmonicSymmetry(NamedTuple): + """Finite angular symmetry group verified on the full coefficient tables.""" + + shifts: np.ndarray + group_order: int + harmonic_lattice_index: int + max_abs_residual: float + relative_residual: float + certified: bool + + +class StartPortfolio(NamedTuple): + """Small distance-following start set and its explicit work counters.""" + + starts: np.ndarray + scores: np.ndarray + raw_starts: np.ndarray + raw_scores: np.ndarray + group_action: np.ndarray + symmetry: HarmonicSymmetry + time_starts: np.ndarray + time_profile: np.ndarray + n_phi_lattice: int + n_u_lattice: int + n_lattice_evaluations: int + n_raw_candidates: int + capacity_truncated: bool + + +class CoverReport(NamedTuple): + """Ledger for a hierarchical bound on the local-union complement. + + ``outside_log_upper`` is an absolute integral upper bound for the finite + coefficient model when ``bound_certified`` is true. ``budget_met`` is a + separate comparison with the caller's diagnostic target value. + """ + + outside_log_upper: float + tail_margin: float + bound_certified: bool + budget_met: bool + cap_reached: bool + n_boxes_evaluated: int + n_subdivisions: int + n_outside_leaves: int + n_owned_leaves: int + n_overlap_owned: int + max_depth: np.ndarray + initial_tail_margin: float + best_tail_margin: float + stalled: bool + progress_checks: np.ndarray + owned_centers: np.ndarray + owned_half_widths: np.ndarray + owned_mode: np.ndarray + + +class LocalIntegralReport(NamedTuple): + """Empirical local-union integral and its bounded-work diagnostics.""" + + value: float + value_half: float + quadrature_delta: float + ok: bool + finite: bool + hessian_ok: bool + edge_ok: bool + local_geometry_ok: bool + overlap_ok: bool + contribution_ok: bool + cell_tail_ok: bool + n_input_modes: int + n_retained_modes: int + n_dropped_modes: int + n_evaluations: int + modeled_peak_bytes: int + retained_indices: np.ndarray + contribution_proxy: np.ndarray + dropped_proxy_relative: float + cell_tail_proxy_relative: float + cell_axis_extents: np.ndarray + edge_sigma: np.ndarray + min_core_separation: float + active_node_fraction: float + + +class _MultiPeakRecord(NamedTuple): + """The 13-element tuple record. Construct :class:`MultiPeakResult`.""" + + value: float + accepted: bool + used_reserve: bool + provenance: str + delta_log_integral: float + tier0: LocalIntegralReport + tier1: LocalIntegralReport + tier0_portfolio: StartPortfolio + tier1_portfolio: StartPortfolio + total_lattice_evaluations: int + total_refinement_steps: int + total_local_evaluations: int + modeled_peak_bytes: int + + +class MultiPeakResult(_MultiPeakRecord): + """Two-tier opt-in result with an explicit finite-reserve provenance. + + ``modeled_peak_bytes`` counts explicit planner/evaluator arrays. It is a + portable sizing model, not measured RSS or device high-water memory: JAX + compilation caches, allocator retention, host/device duplication, and AD + workspace must be measured separately on the production GPU. + + THE TUPLE IS 13 ELEMENTS AND STAYS 13. ``decline_kind`` and ``fault`` + annotate the record; they are attributes only, and are not tuple + elements. A ``NamedTuple`` is a tuple, so appending two fields would have + changed ``len()``, indexing, and iteration -- and any existing caller + writing ``a, b, ..., m = result`` would get ``ValueError: too many values + to unpack``. Defaulted fields prevent that break on CONSTRUCTION, not on + UNPACKING, which is the contract callers actually hold. This is a core + path; an existing caller must see exactly the previous behaviour. + + So ``len()``, iteration, indexing, ``_fields`` and ``_asdict()`` all cover + the same 13 elements they did before. The two annotations are reached by + attribute, and are preserved across ``_replace``, ``copy`` and ``pickle``. + """ + + # decline_kind (Optional[str]) and fault (Optional[FallbackFault]) are set + # per instance below. There are no class-level fallbacks: __new__ is the + # only door, since _make, _replace, copy and pickle all route through it, + # so a fallback would be code no test could distinguish. + def __new__(cls, *args, decline_kind=None, fault=None, **kwargs): + self = super().__new__(cls, *args, **kwargs) + # The tuple payload is immutable, and so are these: set them behind + # the __setattr__ guard below. + object.__setattr__(self, "decline_kind", decline_kind) + object.__setattr__(self, "fault", fault) + return self + + def __setattr__(self, name, value): + raise AttributeError( + "MultiPeakResult is immutable; use _replace(%s=...)" % name) + + def __delattr__(self, name): + raise AttributeError("MultiPeakResult is immutable") + + @classmethod + def _make(cls, iterable, *, decline_kind=None, fault=None): + # namedtuple._make is tuple.__new__ bound as a classmethod: it skips + # __new__, so it must be overridden or the annotations go missing. + return cls(*iterable, decline_kind=decline_kind, fault=fault) + + # No __reduce__ is needed: this subclass has a __dict__, so the default + # pickle/copy protocol carries the annotations as instance state on top of + # the namedtuple's __getnewargs__ payload. Removing an explicit __reduce__ + # changed no test, which is how it was found to be dead. The round trip is + # pinned by test_annotations_are_attributes_and_survive_replace_copy_pickle. + def _replace(self, **kwargs): + decline_kind = kwargs.pop("decline_kind", self.decline_kind) + fault = kwargs.pop("fault", self.fault) + values = dict(zip(_MultiPeakRecord._fields, self)) + unexpected = set(kwargs) - set(values) + if unexpected: + raise ValueError( + "Got unexpected field names: %r" % sorted(unexpected)) + values.update(kwargs) + return type(self)( + *(values[name] for name in _MultiPeakRecord._fields), + decline_kind=decline_kind, fault=fault) + + def __repr__(self): + return "%s(%s, decline_kind=%r, fault=%r)" % ( + type(self).__name__, + ", ".join("%s=%r" % (name, value) + for name, value in zip(_MultiPeakRecord._fields, self)), + self.decline_kind, self.fault) + + +class _DenseReserveError(Exception): + """A caller-supplied reserve failed; do not relabel it as planner decline.""" + + +class MultiPeakFallbackError(Exception): + """A planner fault was converted to a reserve under ``fail_on_fallback``. + + Deliberately outside the ``(RuntimeError, ValueError, LinAlgError)`` set the + planner catches, for the same reason as :class:`_DenseReserveError`: a + nested or repeated call must not swallow it back into a decline. + """ + + +class FallbackFault(NamedTuple): + """Why a fallback was a fault rather than a diagnostic decline. + + ``stage`` names the planner step that raised, so a tier-specific defect is + attributable without re-running. ``error_type``/``message`` carry the + exception itself. ``None`` in ``MultiPeakResult.fault`` means no fault. + """ + + stage: str + error_type: str + message: str + + +#: ``MultiPeakResult.decline_kind`` values. ``None`` means the local branch was +#: accepted. A diagnostic decline is a normal outcome -- a diagnostic +#: legitimately failed its budget. A fault means something raised, and the +#: reserve is standing in for a planner that could not run. These are separate +#: values so downstream analysis never has to parse ``provenance``. +DECLINE_DIAGNOSTIC = "diagnostic" +DECLINE_FAULT = "fault" + + +def infer_harmonic_symmetry(C_A_t, C_B, *, support_rtol=1.0e-12, + invariance_rtol=1.0e-12): + """Infer and verify the finite angular translation group of U,V and Q. + + Significant integer harmonics generate a rank-two lattice ``L``. The + finite symmetry group dual to ``Z^2/L`` has order equal to the gcd of the + two-by-two minors. Candidate translations are enumerated on that exact + denominator and then verified against *all* coefficients, including those + below the support threshold. A noisy near-zero can therefore remove the + ``certified`` label but can never silently invent a group action. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + C_B = np.asarray(C_B, dtype=np.complex128) + _validate_tables(C_A_t, C_B) + scale = max(1.0, float(np.max(np.abs(C_A_t))), float(np.max(np.abs(C_B)))) + support = [] + tables = (C_A_t, C_B) + for table in tables: + ks_max = (table.shape[1] - 1) // 2 + for kp in range(table.shape[0]): + for column, ks in enumerate(range(-ks_max, ks_max + 1)): + if float(np.max(np.abs(table[kp, column]))) > support_rtol * scale: + support.append((int(kp), int(ks))) + support = sorted(set(support)) + index = 0 + for i, first in enumerate(support): + for second in support[:i]: + determinant = abs(first[0] * second[1] - first[1] * second[0]) + index = math.gcd(index, int(determinant)) + if index == 0: + return HarmonicSymmetry( + np.zeros((1, 2)), 1, 0, np.inf, np.inf, False) + + actions = [] + for iphi in range(index): + for iu in range(index): + if all((kp * iphi + ks * iu) % index == 0 + for kp, ks in support): + actions.append((2.0 * np.pi * iphi / index, + 2.0 * np.pi * iu / index)) + shifts = np.asarray(actions, dtype=float).reshape((-1, 2)) + maximum = 0.0 + for shift_phi, shift_u in shifts: + for table in tables: + kp = np.arange(table.shape[0], dtype=float)[:, None] + ks = np.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, + dtype=float)[None, :] + factor = np.exp(1j * (kp * shift_phi + ks * shift_u)) - 1.0 + if table.ndim == 3: + factor = factor[..., None] + maximum = max(maximum, float(np.max(np.abs(table * factor)))) + relative = maximum / scale + certified = bool(len(shifts) == index + and np.isfinite(relative) + and relative <= float(invariance_rtol)) + if not certified: + shifts = np.zeros((1, 2)) + return HarmonicSymmetry( + shifts, int(len(shifts)), int(index), maximum, relative, certified) + + +def rank_joint_starts_from_uvq( + C_A_t, uv_summary, x_min, x_max, *, max_time_starts=4, + max_starts=16, min_time_separation=2, keep_nats=None, + angular_oversample=2): + """Build sparse four-axis starts from the exact U,V/Q harmonics. + + U,V supplies the full angle-dependent norm, not only a scalar bound. Q + supplies the data harmonics at every retained time. On a lattice sized by + their exact harmonic orders, distance is optimized analytically at every + point; the angular placement therefore follows distance rather than using a + single frozen distance slice. Only periodic angular maxima at a small + number of ranked time basins become JAX starts. + + This is a mode-order-sized targeting lattice. It is not an integration + grid and carries no completeness semantics. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + if not isinstance(uv_summary, UVQSummary): + raise TypeError("uv_summary must come from summarize_uv_norm_table") + _validate_tables(C_A_t, uv_summary.C_B) + if not uv_summary.time_invariant: + raise ValueError("arrival-time-dependent U,V norm cannot be collapsed") + if not (0.0 < float(x_min) < float(x_max)): + raise ValueError("need 0 < x_min < x_max") + if min(int(max_time_starts), int(max_starts)) < 1: + raise ValueError("start capacities must be positive") + angular_oversample = int(angular_oversample) + if angular_oversample < 1: + raise ValueError("angular_oversample must be positive") + symmetry = infer_harmonic_symmetry(C_A_t, uv_summary.C_B) + + k_phi = uv_summary.C_B.shape[0] - 1 + k_u = (uv_summary.C_B.shape[1] - 1) // 2 + n_phi = max(9, 2 * angular_oversample * k_phi + 1) + n_u = max(9, 2 * angular_oversample * k_u + 1) + phi, u, A = _harmonic_lattice(C_A_t, n_phi, n_u) + _, _, B = _harmonic_lattice(uv_summary.C_B, n_phi, n_u) + profile, x_best = _distance_profile( + A, B[..., None], float(x_min), float(x_max)) + time_profile = np.max(profile, axis=(0, 1)) + + time_peak = np.zeros(time_profile.size, dtype=bool) + if time_profile.size == 1: + time_peak[0] = True + elif time_profile.size >= 2: + # The reflected time support is not periodic. Endpoints must pass the + # available one-sided comparison before they can displace a real basin. + time_peak[0] = time_profile[0] >= time_profile[1] + time_peak[-1] = time_profile[-1] >= time_profile[-2] + if time_profile.size > 2: + time_peak[1:-1] = ((time_profile[1:-1] >= time_profile[:-2]) + & (time_profile[1:-1] >= time_profile[2:])) + candidate_time = np.flatnonzero(time_peak) + candidate_time = candidate_time[ + np.argsort(time_profile[candidate_time])[::-1]] + selected_time = [] + for time_index in candidate_time: + if all(abs(int(time_index) - old) > int(min_time_separation) + for old in selected_time): + selected_time.append(int(time_index)) + if len(selected_time) == int(max_time_starts): + break + if not selected_time: + selected_time = [int(np.argmax(time_profile))] + + raw = [] + for time_index in selected_time: + surface = profile[..., time_index] + is_peak = np.ones(surface.shape, dtype=bool) + for dphi in (-1, 0, 1): + for du in (-1, 0, 1): + if dphi or du: + is_peak &= surface >= np.roll( + np.roll(surface, dphi, axis=0), du, axis=1) + indices = np.argwhere(is_peak) + if not len(indices): + indices = np.asarray([ + np.unravel_index(np.argmax(surface), surface.shape)]) + for iphi, iu in indices: + raw.append(( + float(surface[iphi, iu]), + (float(time_index), float(phi[iphi]), float(u[iu]), + float(x_best[iphi, iu, time_index])))) + raw.sort(key=lambda item: item[0], reverse=True) + # Keep the odd, unshifted targeting grid: forcing its phase to align with + # the group can move every seed outside a narrow high-SNR basin. Instead, + # reduce sampled cells modulo the *exact* group before capacity. Copies + # can differ by one grid cell because the exact translation generally lies + # between lattice sites; two resolved cells are aliases only when the time + # index agrees and some proven action brings both angular coordinates + # within the corresponding lattice resolution. + orbit_representatives = [] + phi_resolution = 2.0 * np.pi / n_phi + u_resolution = 2.0 * np.pi / n_u + for candidate in raw: + start = np.asarray(candidate[1]) + duplicate = False + for _, representative in orbit_representatives: + representative = np.asarray(representative) + if int(start[0]) != int(representative[0]): + continue + for shift in symmetry.shifts: + shifted = representative[1:3] + shift + delta = (start[1:3] - shifted + np.pi) % ( + 2.0 * np.pi) - np.pi + if (abs(delta[0]) <= phi_resolution + 1.0e-13 + and abs(delta[1]) <= u_resolution + 1.0e-13): + duplicate = True + break + if duplicate: + break + if not duplicate: + orbit_representatives.append(candidate) + raw = orbit_representatives + # The default does not prune on sampled height. At high SNR a genuine + # narrow basin can land tens of nats below its true peak on this deliberately + # small lattice (measured -37 sampled versus -12 after refinement for the + # second lmax=4 mode). A fixed sampled-height cut would therefore become + # *less* complete as amplitude grows even when extrema locations do not + # change. Capacity is the default work bound; an explicit keep_nats remains + # available only as a diagnostic experiment. + if keep_nats is None: + kept_raw = raw + else: + best = raw[0][0] + kept_raw = [item for item in raw + if item[0] >= best - float(keep_nats)] + # Never truncate a proven orbit. The total capacity limits the number of + # representatives; every retained representative receives every verified + # group action, with the action index recorded for the placement audit. + n_representative = max(1, int(max_starts) // symmetry.group_order) + n_raw_candidates = len(kept_raw) + capacity_truncated = n_raw_candidates > n_representative + kept_raw = kept_raw[:n_representative] + expanded = [] + action = [] + for score, start in kept_raw: + for action_index, shift in enumerate(symmetry.shifts): + copied = np.asarray(start, dtype=float).copy() + copied[1:3] = np.mod(copied[1:3] + shift, 2.0 * np.pi) + expanded.append((score, copied)) + action.append(action_index) + return StartPortfolio( + np.asarray([item[1] for item in expanded], dtype=float).reshape((-1, 4)), + np.asarray([item[0] for item in expanded], dtype=float), + np.asarray([item[1] for item in kept_raw], dtype=float).reshape((-1, 4)), + np.asarray([item[0] for item in kept_raw], dtype=float), + np.asarray(action, dtype=np.int32), symmetry, + np.asarray(selected_time, dtype=np.int32), time_profile, + int(n_phi), int(n_u), int(n_phi * n_u * C_A_t.shape[-1]), + int(n_raw_candidates), bool(capacity_truncated)) + + +def _reflected_spectrum(C_A_t): + reflected = jnp.concatenate( + (C_A_t, jnp.flip(C_A_t[..., 1:-1], axis=-1)), axis=-1) + return (jnp.fft.fft(reflected, axis=-1) / reflected.shape[-1], + jnp.fft.fftfreq(reflected.shape[-1])) + + +def _evaluate_spectrum(coeff, frequency, time): + phase = jnp.exp(2j * jnp.pi * time * frequency) + if coeff.shape[-1] % 2 == 0: + phase = phase.at[coeff.shape[-1] // 2].set(jnp.cos(jnp.pi * time)) + return jnp.einsum("kqn,n->kq", coeff, phase) + + +def _bounded_ascent_direction(gradient, hessian, ridge, max_step=None): + """Modified-Newton direction, bounded WITHOUT losing the ascent property. + + ``eigenvector @ ((eigenvector.T @ g) / safe)`` always ascends: + ``g . d = sum_i (v_i . g)^2 / safe_i > 0`` because every ``safe_i`` is + positive. Bounding it by clipping each COORDINATE independently does not + preserve that, because the clip rescales the coordinates by different + factors. Measured on the row that declined the 2026-09-07 ladder + campaign, where ``eigh(-H)`` is indefinite and the raw step is therefore + about 1e8 gradients long, so every coordinate saturates: + + g = [-54.47, +51.02, +28.92, 0] + clipped = [ +2.00, -0.50, +0.50, +0.25] g . d = -119.9 + + A descent direction has no improving lane, the value-only search takes its + zero lane, and the iterate is a fixed point of the whole loop. Scaling by + one factor instead keeps every ratio, hence the sign of ``g . d``, while + respecting the same per-coordinate cap. + + Written as ``min(max_step / |d|)`` to match + ``all_axis_peaklocal.refine_all_axis_starts``, which bounds its own step + this way for the same reason. Measured, so that the next reader does not + have to re-derive it: the two spellings are equivalent here, and BOTH + return a zero step once ``|d|`` reaches about 1e308, because the scale + factor is then denormal and the product underflows. That needs + ``|g| >~ 1e299`` at the default ridge, which this likelihood cannot reach. + ``tiny`` is defensive, not load-bearing: ``max_step`` is validated positive + below, so ``max_step / 0`` is ``+inf`` rather than a NaN, and the step for + a zero direction is zero either way. + + Returns the eigenvalues of ``-H`` alongside the direction so a caller that + needs both does not decompose the same 4x4 twice. + """ + eigenvalue, eigenvector = jnp.linalg.eigh(-hessian) + safe = jnp.maximum(eigenvalue, float(ridge)) + direction = eigenvector @ ((eigenvector.T @ gradient) / safe) + if max_step is None: + return direction, eigenvalue + ratio = max_step / jnp.maximum(jnp.abs(direction), + jnp.finfo(jnp.float64).tiny) + return direction * jnp.minimum(1.0, jnp.min(ratio)), eigenvalue + + +def refine_joint_starts_jax( + C_A_t, C_B, starts, x_min, x_max, *, iterations=12, + ridge=1.0e-8, max_step=(2.0, 0.5, 0.5, 0.25)): + """Refine the small portfolio with bounded sequential JAX Newton steps. + + The reflected time primitive is evaluated only at each proposed time. A + ``lax.map`` over starts avoids a start-by-frequency-by-Hessian batch. This + is local hill climbing only; convergence never asserts completeness. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + starts = jnp.asarray(starts, dtype=jnp.float64) + _validate_tables(C_A_t, C_B) + if starts.ndim != 2 or starts.shape[1] != 4: + raise ValueError("starts must have shape (N,4)") + if int(iterations) < 1: + raise ValueError("iterations must be positive") + coeff, frequency = _reflected_spectrum(C_A_t) + max_step = jnp.asarray(max_step, dtype=jnp.float64) + # The rescale below is meaningless for a non-positive bound, and fails + # QUIETLY rather than loudly: a zero bound scales every step to exactly + # zero, which is the stall this function exists to prevent, and a negative + # bound is silently exceeded (bound -0.5 returns a component of -3.5). + # Neither is non-finite, so nothing downstream would notice. + if max_step.shape != (4,) or not bool(jnp.all(max_step > 0.0)): + raise ValueError("max_step must be four positive coordinate bounds") + + def log_density(theta): + time, phi, u, x = theta + C_A = _evaluate_spectrum(coeff, frequency, time) + A = _angular_field_jax(C_A, phi, u) + B = _angular_field_jax(C_B, phi, u) + inside = ((time >= 0.0) & (time <= C_A_t.shape[-1] - 1.0) + & (x >= float(x_min)) & (x <= float(x_max)) & (x > 0.0)) + value = x * A - 0.5 * x * x * B - 4.0 * jnp.log( + jnp.maximum(x, 1.0e-300)) + return jnp.where(inside, value, -jnp.inf) + + gradient_fn = jax.grad(log_density) + hessian_fn = jax.hessian(log_density) + + def project(theta): + return jnp.asarray([ + jnp.clip(theta[0], 0.0, C_A_t.shape[-1] - 1.0), + jnp.mod(theta[1], 2.0 * jnp.pi), + jnp.mod(theta[2], 2.0 * jnp.pi), + jnp.clip(theta[3], float(x_min), float(x_max)), + ]) + + def one(start): + def step(theta, _): + gradient = gradient_fn(theta) + hessian = hessian_fn(theta) + direction, _ = _bounded_ascent_direction( + gradient, hessian, ridge, max_step) + proposal = jax.vmap( + lambda scale: project(theta + scale * direction))( + jnp.asarray([1.0, 0.5, 0.25, 0.125, 0.0])) + values = jax.vmap(log_density)(proposal) + return proposal[jnp.argmax(values)], None + + point, _ = jax.lax.scan(step, project(start), None, + length=int(iterations)) + + # At loud SNR the Newton improvement can be below one ulp of lnL while + # the remaining absolute gradient is still visible. A value-only line + # search then selects its zero-step lane forever. Two final root-polish + # steps accept only a strict gradient-norm reduction, positive local + # curvature, and no value loss beyond roundoff. This tightens the + # stationarity result; it does not relax the downstream gate. + def polish(theta, _): + gradient = gradient_fn(theta) + hessian = hessian_fn(theta) + # Unbounded, as before: the polish's own guard below accepts a + # step only at a strict maximum with a smaller gradient. One + # decomposition serves both the step and that guard. + direction, eigenvalue = _bounded_ascent_direction( + gradient, hessian, ridge) + proposal = project(theta + direction) + proposal_gradient = gradient_fn(proposal) + value = log_density(theta) + proposal_value = log_density(proposal) + tolerance = 32.0 * jnp.finfo(jnp.float64).eps * jnp.maximum( + 1.0, jnp.abs(value)) + use = (jnp.all(eigenvalue > 0.0) + & (jnp.linalg.norm(proposal_gradient) + < jnp.linalg.norm(gradient)) + & (proposal_value >= value - tolerance)) + return jnp.where(use, proposal, theta), None + + point, _ = jax.lax.scan(polish, point, None, length=2) + value = log_density(point) + gradient = gradient_fn(point) + hessian = hessian_fn(point) + curvature = jnp.linalg.eigvalsh(-hessian) + return point, value, gradient, hessian, curvature + + return jax.lax.map(jax.checkpoint(one), starts) + + +def select_refined_modes(points, values, gradients, curvatures, *, + max_modes, gradient_tol=2.0e-6, + coordinate_tol=(0.25, 1.0e-4, 1.0e-4, 1.0e-5)): + """Filter, rank, and periodically deduplicate refined local maxima.""" + points = np.asarray(points, dtype=float) + values = np.asarray(values, dtype=float).ravel() + gradients = np.asarray(gradients, dtype=float) + curvatures = np.asarray(curvatures, dtype=float) + if (points.ndim != 2 or points.shape[1] != 4 + or gradients.shape != points.shape + or curvatures.shape != points.shape + or values.shape != (points.shape[0],)): + raise ValueError("inconsistent refined-mode arrays") + stationary = (np.all(np.isfinite(points), axis=1) + & np.isfinite(values) + & np.all(np.isfinite(gradients), axis=1) + & (np.linalg.norm(gradients, axis=1) <= float(gradient_tol)) + & np.all(curvatures > 0.0, axis=1)) + order = np.flatnonzero(stationary) + order = order[np.argsort(values[order])[::-1]] + tolerance = np.asarray(coordinate_tol, dtype=float) + selected = [] + for index in order: + duplicate = False + for old in selected: + delta = np.abs(points[index] - points[old]) + delta[1] = _periodic_distance(points[index, 1], points[old, 1]) + delta[2] = _periodic_distance(points[index, 2], points[old, 2]) + if np.all(delta <= tolerance): + duplicate = True + break + if not duplicate: + selected.append(int(index)) + if len(selected) > int(max_modes): + raise ValueError("unique stationary mode count exceeds capacity") + return np.asarray(selected, dtype=np.int32), stationary + + +def axis_local_geometry(hessians, *, w_sigma=6.0, + eigenvalue_floor=1.0e-12): + """Axis-aligned boxes enclosing ``w_sigma`` marginal Hessian widths.""" + hessians = np.asarray(hessians, dtype=float) + if hessians.ndim != 3 or hessians.shape[1:] != (4, 4): + raise ValueError("hessians must have shape (N,4,4)") + widths = np.full((len(hessians), 4), np.nan) + for index, hessian in enumerate(hessians): + eigenvalue = np.linalg.eigvalsh(-hessian) + if (np.all(np.isfinite(eigenvalue)) + and np.min(eigenvalue) > float(eigenvalue_floor)): + covariance = np.linalg.inv(-hessian) + widths[index] = float(w_sigma) * np.sqrt( + np.maximum(np.diag(covariance), 0.0)) + return widths + + +def _evaluate_points_jax(C_A_t, C_B, points, x_min, x_max, chunk_size): + """Stream the four-axis exponent without materializing point x frequency.""" + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + points = jnp.asarray(points, dtype=jnp.float64) + coeff, frequency = _reflected_spectrum(C_A_t) + chunk_size = int(chunk_size) + if chunk_size < 1: + raise ValueError("chunk_size must be positive") + n_point = points.shape[0] + n_chunk = (n_point + chunk_size - 1) // chunk_size + padding = n_chunk * chunk_size - n_point + padded = jnp.pad(points, ((0, padding), (0, 0))) + + def one(theta): + time, phi, u, x = theta + C_A = _evaluate_spectrum(coeff, frequency, time) + A = _angular_field_jax(C_A, phi, u) + B = _angular_field_jax(C_B, phi, u) + inside = ((time >= 0.0) & (time <= C_A_t.shape[-1] - 1.0) + & (x >= float(x_min)) & (x <= float(x_max)) & (x > 0.0)) + value = x * A - 0.5 * x * x * B - 4.0 * jnp.log( + jnp.maximum(x, 1.0e-300)) + return jnp.where(inside, value, -jnp.inf) + + def step(_, block): + return None, jax.vmap(one)(block) + + _, values = jax.lax.scan( + jax.checkpoint(step), None, + padded.reshape((n_chunk, chunk_size, 4))) + return values.reshape((-1,))[:n_point] + + +def _periodic_delta_rows(points, center): + delta = np.asarray(points, dtype=float) - np.asarray(center, dtype=float) + delta[..., 1] = (delta[..., 1] + np.pi) % (2.0 * np.pi) - np.pi + delta[..., 2] = (delta[..., 2] + np.pi) % (2.0 * np.pi) - np.pi + return delta + + +def _laplace_contribution_proxy(values, hessians): + result = np.full(len(values), -np.inf) + for index, (value, hessian) in enumerate(zip(values, hessians)): + sign, logdet = np.linalg.slogdet(-hessian) + if sign > 0 and np.isfinite(logdet): + result[index] = (float(value) + 2.0 * np.log(2.0 * np.pi) + - 0.5 * logdet) + return result + + +def integrate_refined_modes_tensor( + C_A_t, C_B, points, values, hessians, x_min, x_max, *, + log_integral_tol=1.0e-3, contribution_cutoff_nats=-18.0, + cell_sigma=5.0, quadrature_order=7, chunk_size=64, + max_condition=1.0e10, edge_guard_sigma=0.5, + core_overlap_sigma=4.0, log_measure=0.0): + """Integrate the union of finite full-Hessian local cells. + + Every retained maximum supplies the finite affine cell + ``theta = mu + L z``, ``|z_i| <= cell_sigma``, with + ``L L.T = (-H)^-1``. Tensor Gauss--Hermite rules sample the corresponding + Gaussian mixture, while an exact cell-union indicator zeros nodes outside + those finite regions. Dividing by the *full mixture density* partitions + overlaps automatically, rather than integrating shared tails once per + mode. Orders ``n`` and ``n-2`` provide the empirical convergence check + and are exact for the quadratic local limit. The tensor is streamed in + bounded chunks. ``core_overlap_sigma`` is retained as a separation + telemetry scale; overlap itself is not a rejection because the full + mixture density owns it exactly. + + This is an empirical warrant, not a deterministic omitted-mass proof. Its + ``ok`` flag is therefore allowed to choose a finite dense reserve, never to + delete the outer likelihood point. Coordinates here are time-sample index, + two radians, and inverse distance. ``log_measure`` carries any constant + physical time step and normalized-prior factors required by the caller; + the dense reserve must use the same convention. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + C_B = np.asarray(C_B, dtype=np.complex128) + points = np.asarray(points, dtype=float) + values = np.asarray(values, dtype=float).ravel() + hessians = np.asarray(hessians, dtype=float) + if (points.ndim != 2 or points.shape[1] != 4 + or values.shape != (len(points),) + or hessians.shape != (len(points), 4, 4) + or len(points) == 0): + raise ValueError("points, values, and hessians must describe N>0 modes") + if int(quadrature_order) < 4: + raise ValueError("quadrature_order must be at least 4") + if not (np.isfinite(float(log_integral_tol)) + and 0.0 < float(log_integral_tol)): + raise ValueError("log_integral_tol must be positive") + if not (np.isfinite(float(contribution_cutoff_nats)) + and float(contribution_cutoff_nats) <= 0.0): + raise ValueError("contribution_cutoff_nats must be finite and nonpositive") + if not (np.isfinite(float(cell_sigma)) and 0.0 < float(cell_sigma)): + raise ValueError("cell_sigma must be positive") + if not np.isfinite(float(log_measure)): + raise ValueError("log_measure must be finite") + + proxy = _laplace_contribution_proxy(values, hessians) + best_proxy = float(np.max(proxy)) + retained = np.flatnonzero( + proxy >= best_proxy + float(contribution_cutoff_nats)) + dropped = np.setdiff1d(np.arange(len(points)), retained) + dropped_proxy = (-np.inf if not len(dropped) else + float(scipy_logsumexp(proxy[dropped]) - best_proxy)) + contribution_ok = bool( + dropped_proxy < math.log(float(log_integral_tol)) - 2.0) + normal_mass = math.erf(float(cell_sigma) / math.sqrt(2.0)) + cell_tail_proxy = math.log(max( + 1.0 - normal_mass ** 4, np.finfo(float).tiny)) + cell_tail_ok = bool( + cell_tail_proxy < math.log(float(log_integral_tol)) - 2.0) + + modes = points[retained] + fishers = -hessians[retained] + cholesky = [] + hessian_ok = bool(len(retained)) + for fisher in fishers: + try: + eigenvalue = np.linalg.eigvalsh(fisher) + this_condition = float(np.max(eigenvalue) / np.min(eigenvalue)) + covariance = np.linalg.inv(fisher) + factor = np.linalg.cholesky(covariance) + good = (np.all(np.isfinite(eigenvalue)) and np.min(eigenvalue) > 0.0 + and np.isfinite(this_condition) + and this_condition <= float(max_condition)) + except np.linalg.LinAlgError: + factor = np.full((4, 4), np.nan) + this_condition = np.inf + good = False + cholesky.append(factor) + hessian_ok &= good + cholesky = np.asarray(cholesky) + + # Exact row-wise extent of the affine parallelepiped. Angular extents must + # stay below half a period so nearest-copy mixture densities are unambiguous. + extents = float(cell_sigma) * np.sum(np.abs(cholesky), axis=2) + marginal_sigma = np.sqrt(np.maximum( + np.diagonal(cholesky @ np.swapaxes(cholesky, 1, 2), axis1=1, axis2=2), + 0.0)) + edge_sigma = np.minimum( + (modes[:, 0] - 0.0) / np.maximum(marginal_sigma[:, 0], 1.0e-300), + (C_A_t.shape[-1] - 1.0 - modes[:, 0]) + / np.maximum(marginal_sigma[:, 0], 1.0e-300)) + edge_sigma = np.minimum( + edge_sigma, + np.minimum( + (modes[:, 3] - float(x_min)) + / np.maximum(marginal_sigma[:, 3], 1.0e-300), + (float(x_max) - modes[:, 3]) + / np.maximum(marginal_sigma[:, 3], 1.0e-300))) + edge_ok = bool(np.all(edge_sigma >= float(edge_guard_sigma))) + local_ok = bool(np.all(extents[:, 1:3] < np.pi) + and np.all(extents[:, 0] < 0.5 * (C_A_t.shape[-1] - 1.0)) + and np.all(extents[:, 3] < 0.5 * (x_max - x_min))) + + separation = [] + for first in range(len(modes)): + for second in range(first): + delta = _periodic_delta_rows(modes[first:first + 1], modes[second])[0] + # Symmetric local metric: use the smaller of the two Fisher lengths. + separation.append(min( + math.sqrt(max(0.0, float(delta @ fishers[first] @ delta))), + math.sqrt(max(0.0, float(delta @ fishers[second] @ delta))))) + min_separation = min(separation) if separation else np.inf + # The mixture denominator partitions ordinary affine-cell overlaps. The + # only ambiguous case is a cell wide enough to meet more than one periodic + # image of itself; the strict angular-locality gate rejects that geometry. + overlap_ok = bool(np.all(extents[:, 1:3] < np.pi)) + n_high_model = len(modes) * int(quadrature_order) ** 4 + n_rule = int(quadrature_order) ** 4 + # Conservative count of the explicit simultaneous arrays in the Python + # quadrature and streamed JAX evaluator: samples and their construction + # copies, rule weights, mixture matrix, union mask, periodic deltas/local-z, + # density/proposal/integrand vectors, and a device points/value payload. + # Backend compilation, allocator retention, AD, and host/device duplication + # outside these arrays are deliberately not represented (see result doc). + modeled_peak_bytes = int( + C_A_t.nbytes + C_B.nbytes + + n_high_model * (209 + len(modes) * 8) + + n_rule * (4 * 8 + 8) + + int(chunk_size) * 2 * (C_A_t.shape[-1] - 1) * 16) + + finite_structure = bool(hessian_ok and np.all(np.isfinite(cholesky)) + and np.all(np.isfinite(proxy[retained]))) + if not finite_structure: + return LocalIntegralReport( + np.nan, np.nan, np.inf, False, False, hessian_ok, edge_ok, + local_ok, overlap_ok, contribution_ok, cell_tail_ok, + len(points), len(retained), len(dropped), 0, + modeled_peak_bytes, + retained.astype(np.int32), proxy, dropped_proxy, cell_tail_proxy, + extents, edge_sigma, + float(min_separation), 0.0) + + log_normalization = 2.0 * np.log(2.0 * np.pi) + + def integrate_order(order): + node, weight = np.polynomial.hermite.hermgauss(int(order)) + z_axis = np.sqrt(2.0) * node + log_weight_axis = np.log(weight) - 0.5 * np.log(np.pi) + mesh = np.meshgrid(z_axis, z_axis, z_axis, z_axis, indexing="ij") + z = np.stack(mesh, axis=-1).reshape((-1, 4)) + weight_mesh = np.meshgrid( + log_weight_axis, log_weight_axis, log_weight_axis, + log_weight_axis, indexing="ij") + log_rule_weight = np.sum( + np.stack(weight_mesh, axis=-1), axis=-1).reshape(-1) + samples = [] + rule_weights = [] + for mode, factor in zip(modes, cholesky): + theta = mode[None, :] + z @ factor.T + theta[:, 1:3] = np.mod(theta[:, 1:3], 2.0 * np.pi) + samples.append(theta) + rule_weights.append(log_rule_weight) + samples = np.concatenate(samples, axis=0) + rule_weights = np.concatenate(rule_weights) - np.log(len(modes)) + + log_component = np.full((len(samples), len(modes)), -np.inf) + inside_union = np.zeros(len(samples), dtype=bool) + for mode_index, (mode, factor) in enumerate(zip(modes, cholesky)): + delta = _periodic_delta_rows(samples, mode) + local_z = np.linalg.solve(factor, delta.T).T + inside_union |= np.all( + np.abs(local_z) <= float(cell_sigma) + 1.0e-12, axis=1) + logdet = float(np.sum(np.log(np.diag(factor)))) + log_component[:, mode_index] = ( + -0.5 * np.sum(np.square(local_z), axis=1) + - log_normalization - logdet) + log_proposal = (scipy_logsumexp(log_component, axis=1) + - np.log(len(modes))) + log_density = np.asarray(_evaluate_points_jax( + C_A_t, C_B, samples, x_min, x_max, chunk_size), dtype=float) + log_density = np.where(inside_union, log_density, -np.inf) + log_integrand = log_density - log_proposal + rule_weights + return (float(scipy_logsumexp(log_integrand) + float(log_measure)), + log_density, + log_proposal, len(samples)) + + value, log_density, log_proposal, n_high = integrate_order( + int(quadrature_order)) + value_half, _, _, n_low = integrate_order(int(quadrature_order) - 2) + quadrature_delta = abs(value - value_half) + active_node_fraction = float(np.mean(np.isfinite(log_density))) + finite = bool(np.isfinite(value) and np.isfinite(value_half) + and np.all(np.isfinite(log_proposal))) + ok = bool(finite and hessian_ok and edge_ok and local_ok and overlap_ok + and contribution_ok and cell_tail_ok + and quadrature_delta <= float(log_integral_tol)) + return LocalIntegralReport( + value, value_half, quadrature_delta, ok, finite, hessian_ok, + edge_ok, local_ok, overlap_ok, contribution_ok, cell_tail_ok, + len(points), len(retained), + len(dropped), n_high + n_low, modeled_peak_bytes, + retained.astype(np.int32), proxy, + dropped_proxy, cell_tail_proxy, extents, edge_sigma, + float(min_separation), active_node_fraction) + + +def _run_structural_tier(C_A_t, uv_summary, x_min, x_max, *, + angular_oversample, max_time_starts, max_starts, + refine_iterations, integral_kwargs): + portfolio = rank_joint_starts_from_uvq( + C_A_t, uv_summary, x_min, x_max, + angular_oversample=angular_oversample, + max_time_starts=max_time_starts, max_starts=max_starts) + result = tuple(np.asarray(item) for item in refine_joint_starts_jax( + C_A_t, uv_summary.C_B, portfolio.starts, x_min, x_max, + iterations=refine_iterations)) + points, values, gradients, hessians, curvatures = result + selected, _ = select_refined_modes( + points, values, gradients, curvatures, max_modes=max_starts) + if not len(selected): + raise RuntimeError( + "structural tier found no strict stationary maximum") + integral = integrate_refined_modes_tensor( + C_A_t, uv_summary.C_B, points[selected], values[selected], + hessians[selected], x_min, x_max, **integral_kwargs) + return portfolio, integral + + +def multipeak_local_marginalize( + C_A_t, C_B_t, x_min, x_max, dense_reserve, *, + log_integral_tol=1.0e-3, tier0=(2, 3, 24), tier1=(3, 5, 48), + refine_iterations=18, contribution_cutoff_nats=-18.0, + cell_sigma=5.0, quadrature_order=7, chunk_size=64, + log_measure=0.0, fail_on_fallback=False, label=None): + """Two-tier empirical four-axis marginal with a finite reserve fallback. + + The tuple for each tier is ``(angular_oversample, time_starts, start_cap)``. + Acceptance requires both local-mixture quadratures to pass their internal + diagnostics and agree within ``log_integral_tol``. Otherwise the supplied + dense/exact reserve is evaluated and returned with explicit provenance. + ``dense_reserve`` should normally be a zero-argument callable, so an + accepted local row never pays for the fallback. A finite scalar is also + accepted when a caller already has the reserve value. ``log_measure`` is + the caller-owned constant measure/normalization for time, angles, and the + inverse-distance prior; both paths must use the same convention. + + Two different things return the reserve and they are reported separately. + A *diagnostic* decline (``decline_kind == DECLINE_DIAGNOSTIC``) is a normal + outcome: the tiers ran and one of their budgets was not met. A *fault* + (``decline_kind == DECLINE_FAULT``, ``fault`` populated) means the planner + raised, so the reserve is standing in for a step that did not run. The + fault is reported in two places: in the returned record, which is the + primary channel because no configuration can suppress it, and on the module + logger (``RIFT.likelihood.jax_ile.multipeak_planner``) at WARNING, once per + CALL -- not once per call site, which is what ``warnings.warn`` would give. + ``fail_on_fallback=True`` raises + :class:`MultiPeakFallbackError` instead of evaluating the reserve. It + defaults off: this is a core path and an existing caller must see exactly + the previous behaviour, under any logging or warnings configuration. + + ``label`` is an opaque caller tag (a row id, an event name) echoed in the + log line so a fault in a large campaign is attributable without re-running. + """ + def evaluate_reserve(): + try: + value = dense_reserve() if callable(dense_reserve) else dense_reserve + value = float(value) + if not np.isfinite(value): + raise ValueError("dense_reserve must produce a finite value") + return value + except Exception as error: + # This wrapper is intentionally outside the planner exception + # hierarchy. A failing reserve is invoked once and propagated, + # never retried or mislabeled as a local-planner exception. + raise _DenseReserveError("dense reserve evaluation failed") from error + + # Names the planner step in flight, so a fault reports WHERE it happened + # rather than only that something raised. Read only in the handler below. + stage = "uv-summary" + try: + uv_summary = summarize_uv_norm_table(C_B_t) + if not uv_summary.time_invariant: + raise ValueError( + "the four-axis prototype requires time-independent U,V") + integral_kwargs = dict( + log_integral_tol=log_integral_tol, + contribution_cutoff_nats=contribution_cutoff_nats, + cell_sigma=cell_sigma, quadrature_order=quadrature_order, + chunk_size=chunk_size, log_measure=log_measure) + stage = "tier0" + portfolio0, result0 = _run_structural_tier( + C_A_t, uv_summary, x_min, x_max, + angular_oversample=int(tier0[0]), + max_time_starts=int(tier0[1]), max_starts=int(tier0[2]), + refine_iterations=int(refine_iterations), + integral_kwargs=integral_kwargs) + stage = "tier1" + portfolio1, result1 = _run_structural_tier( + C_A_t, uv_summary, x_min, x_max, + angular_oversample=int(tier1[0]), + max_time_starts=int(tier1[1]), max_starts=int(tier1[2]), + refine_iterations=int(refine_iterations), + integral_kwargs=integral_kwargs) + stage = "acceptance" + delta = abs(result1.value - result0.value) + accepted = bool(result0.ok and result1.ok and np.isfinite(delta) + and delta <= float(log_integral_tol)) + if accepted: + value = result1.value + provenance = "uvq-multipeak-tier1" + decline_kind = None + else: + # Both tiers ran and reported. This is a budget outcome, not a + # fault: it is silent by design and must stay silent. + value = evaluate_reserve() + provenance = "dense-reserve:enrichment-or-local-diagnostic" + decline_kind = DECLINE_DIAGNOSTIC + input_bytes = (np.asarray(C_A_t).nbytes + + np.asarray(uv_summary.C_B).nbytes) + + def portfolio_bytes(portfolio): + return int(input_bytes + + portfolio.n_phi_lattice * portfolio.n_u_lattice + * (3 * np.asarray(C_A_t).shape[-1] + 1) * 8) + return MultiPeakResult( + float(value), accepted, not accepted, provenance, float(delta), + result0, result1, portfolio0, portfolio1, + int(portfolio0.n_lattice_evaluations + + portfolio1.n_lattice_evaluations), + (int(refine_iterations) + 2) * (len(portfolio0.starts) + + len(portfolio1.starts)), + int(result0.n_evaluations + result1.n_evaluations), + max(result0.modeled_peak_bytes, result1.modeled_peak_bytes, + portfolio_bytes(portfolio0), portfolio_bytes(portfolio1)), + decline_kind=decline_kind, fault=None) + except (RuntimeError, ValueError, np.linalg.LinAlgError) as error: + # Keep the return finite even when the local planner itself cannot form + # a trustworthy report. Re-run failures should be diagnosed upstream; + # they must never be reclassified as waveform failures. + # + # This is a FAULT, not a decline: the reserve stands in for a step that + # did not run. A whole campaign once read as a conservative controller + # because this path was silent, so it is reported once per call and can + # be made fatal. The report precedes the reserve so a fault is still + # visible when the reserve itself then fails. + # + # A LOGGER, not warnings.warn. Two properties of the warnings module + # are wrong for this: under ``-W error::RuntimeWarning`` the warn call + # RAISES, which would make this default (fail_on_fallback=False) path + # fatal on a filter the caller may have set for unrelated reasons; and + # under the default filters the registry de-duplicates per (message, + # category, call site), so N identical faults from one call site report + # ONCE. That is worst in the campaign case this change exists for, + # where label=None leaves every message identical. Neither the + # record below nor this logger has either property. + fault = FallbackFault( + stage, type(error).__name__, str(error)) + logger.warning( + "multipeak_local_marginalize: the local planner RAISED at stage " + "%r and fell back to the dense reserve. This is a fault, not a " + "budget decline: %s: %s. label=%r, C_A_t.shape=%s, " + "C_B_t.shape=%s, x=[%r, %r], tier0=%r, tier1=%r, " + "refine_iterations=%r. Diagnose it upstream; pass " + "fail_on_fallback=True to make it fatal.", + fault.stage, fault.error_type, fault.message, label, + np.shape(C_A_t), np.shape(C_B_t), x_min, x_max, + tuple(tier0), tuple(tier1), refine_iterations, + stacklevel=2) + if fail_on_fallback: + raise MultiPeakFallbackError( + "multipeak_local_marginalize declined by fault at stage %r " + "(%s: %s), label=%r; fail_on_fallback is set" + % (fault.stage, fault.error_type, fault.message, label) + ) from error + empty = LocalIntegralReport( + np.nan, np.nan, np.inf, False, False, False, False, False, False, + False, False, 0, 0, 0, 0, 0, np.empty(0, dtype=np.int32), + np.empty(0), -np.inf, np.inf, np.empty((0, 4)), np.empty(0), + np.nan, 0.0) + empty_symmetry = HarmonicSymmetry( + np.zeros((1, 2)), 1, 0, np.inf, np.inf, False) + empty_portfolio = StartPortfolio( + np.empty((0, 4)), np.empty(0), np.empty((0, 4)), np.empty(0), + np.empty(0, dtype=np.int32), empty_symmetry, + np.empty(0, dtype=np.int32), np.empty(0), 0, 0, 0, 0, False) + return MultiPeakResult( + evaluate_reserve(), False, True, + "dense-reserve:planner-exception:%s" % type(error).__name__, + np.inf, + empty, empty, empty_portfolio, empty_portfolio, 0, 0, 0, 0, + decline_kind=DECLINE_FAULT, fault=fault) + + +def _periodic_box_contains(box_center, box_half, mode_center, mode_half): + if mode_half >= np.pi: + return True + if box_half >= np.pi: + return False + return (_periodic_distance(box_center, mode_center) + box_half + <= mode_half + 1.0e-14) + + +def _box_owners(lo, hi, centers, half_widths): + midpoint = 0.5 * (lo + hi) + half = 0.5 * (hi - lo) + owners = [] + for index, (center, width) in enumerate(zip(centers, half_widths)): + linear = ((lo[0] >= center[0] - width[0]) + and (hi[0] <= center[0] + width[0]) + and (lo[3] >= center[3] - width[3]) + and (hi[3] <= center[3] + width[3])) + angular = (_periodic_box_contains( + midpoint[1], half[1], center[1], width[1]) + and _periodic_box_contains( + midpoint[2], half[2], center[2], width[2])) + if linear and angular: + owners.append(index) + if not owners: + return [], -1 + scores = [] + for index in owners: + delta = midpoint - centers[index] + delta[1] = _periodic_distance(midpoint[1], centers[index, 1]) + delta[2] = _periodic_distance(midpoint[2], centers[index, 2]) + score = float(np.sum(np.square( + delta / np.maximum(half_widths[index], 1.0e-300)))) + scores.append((score, index)) + return owners, min(scores)[1] + + +def _periodic_segments(center, half_width): + """Represent a periodic interval as closed segments on ``[0, 2 pi]``.""" + if half_width >= np.pi: + return [(0.0, 2.0 * np.pi)] + lower = (float(center) - float(half_width)) % (2.0 * np.pi) + upper = (float(center) + float(half_width)) % (2.0 * np.pi) + if lower <= upper: + return [(lower, upper)] + return [(0.0, upper), (lower, 2.0 * np.pi)] + + +def _local_boundary_splits(lo, hi, centers, half_widths, available): + """Return exact local-union boundaries cutting an intersecting box. + + Splitting at these coordinates is what allows the ledger to remove a leaf + only when the *same axis-aligned region* will be locally integrated. + """ + candidates = [[] for _ in range(4)] + epsilon = 64.0 * np.finfo(float).eps + for center, width in zip(centers, half_widths): + intervals = [ + [(center[0] - width[0], center[0] + width[0])], + _periodic_segments(center[1], width[1]), + _periodic_segments(center[2], width[2]), + [(center[3] - width[3], center[3] + width[3])], + ] + intersects = True + for axis in range(4): + if not any(max(lo[axis], left) < min(hi[axis], right) + for left, right in intervals[axis]): + intersects = False + break + if not intersects: + continue + for axis in range(4): + if not available[axis]: + continue + for left, right in intervals[axis]: + for boundary in (left, right): + tolerance = epsilon * max( + 1.0, abs(float(lo[axis])), abs(float(hi[axis]))) + if lo[axis] + tolerance < boundary < hi[axis] - tolerance: + candidates[axis].append(float(boundary)) + return candidates + + +def _time_fourier_enclosure(C_A_t, max_order=8): + """Return exact reflected coefficients and global derivative remainders.""" + reflected = np.concatenate( + (C_A_t, np.flip(C_A_t[..., 1:-1], axis=-1)), axis=-1) + coefficient = np.fft.fft(reflected, axis=-1) / reflected.shape[-1] + frequency = np.fft.fftfreq(reflected.shape[-1]) + omega = 2.0 * np.pi * frequency + magnitude = np.abs(coefficient) + derivative_bounds = np.stack([ + np.sum(magnitude * np.abs(omega) ** order, axis=-1) + for order in range(1, int(max_order) + 1) + ]) + return (coefficient, frequency, derivative_bounds, + np.sum(magnitude, axis=-1)) + + +def _evaluate_spectrum_numpy(coefficient, frequency, time, derivative=0): + """Evaluate the reflected finite Fourier polynomial or its derivative.""" + omega = 2.0 * np.pi * frequency + phase = np.exp(1j * omega * float(time)) + nyquist = None + if coefficient.shape[-1] % 2 == 0: + nyquist = coefficient.shape[-1] // 2 + factor = np.power(1j * omega, int(derivative)) * phase + if nyquist is not None: + factor[nyquist] = (np.pi ** int(derivative) + * np.cos(np.pi * float(time) + + 0.5 * np.pi * int(derivative))) + return np.einsum("kqn,n->kq", coefficient, factor, optimize=True) + + +def _field_variation(table, phi, u, half_phi, half_u): + table = np.asarray(table, dtype=np.complex128) + kp = np.arange(table.shape[0], dtype=float)[:, None] + ks = np.arange(-(table.shape[1] - 1) // 2, + (table.shape[1] - 1) // 2 + 1, dtype=float)[None, :] + weight = _kp_weights(table.shape[0])[:, None] + phase = np.exp(1j * (kp * float(phi) + ks * float(u))) + value = float(np.sum(weight * table * phase).real) + phase_span = np.minimum( + 2.0, np.abs(kp) * float(half_phi) + np.abs(ks) * float(half_u)) + variation = float(np.sum(weight * np.abs(table) * phase_span)) + return value, variation + + +def _box_log_upper(C_A_t, C_B, time_enclosure, lo, hi, + inherited_point_upper=np.inf): + center = 0.5 * (lo + hi) + half = 0.5 * (hi - lo) + coefficient, frequency, derivative_bounds, magnitude_bound = time_enclosure + dt = float(half[0]) + C_A_center = _evaluate_spectrum_numpy( + coefficient, frequency, center[0], derivative=0) + A0, A_angle = _field_variation( + C_A_center, center[1], center[2], half[1], half[2]) + # Each Taylor expression is independently rigorous for the reflected + # finite Fourier polynomial. Their minimum is rigorous too. Higher-order + # local cancellation matters for a narrow band-limited time peak: a global + # first-derivative lift did not contract fast enough on real 22 tables. + candidates = [magnitude_bound + np.abs(C_A_center)] + partial = np.zeros_like(magnitude_bound) + factorial = 1.0 + power = 1.0 + for order, bound in enumerate(derivative_bounds, start=1): + factorial *= order + power *= dt + if order > 1: + previous = _evaluate_spectrum_numpy( + coefficient, frequency, center[0], derivative=order - 1) + partial = partial + np.abs(previous) * ( + dt ** (order - 1)) / math.factorial(order - 1) + candidates.append(partial + bound * power / factorial) + time_remainder = np.minimum.reduce(candidates) + A_time = float(np.sum( + _kp_weights(C_A_t.shape[0])[:, None] * time_remainder)) + B0, B_angle = _field_variation( + C_B, center[1], center[2], half[1], half[2]) + A_upper = A0 + A_angle + A_time + B_lower = max(0.0, B0 - B_angle) # intersect with B= >= 0 + profile, _ = _distance_profile( + np.asarray(A_upper), np.asarray(B_lower), lo[3], hi[3]) + volume = float(np.prod(hi - lo)) + if not np.isfinite(volume) or volume <= 0.0: + return -np.inf, np.zeros(4) + magnitude = (abs(A0) + A_angle + A_time + abs(B0) + B_angle + + abs(float(profile)) + 1.0) + upper = float(profile) + 64.0 * np.finfo(float).eps * magnitude + # A child is a subset of its parent. Capping its pointwise enclosure by + # the inherited parent enclosure is exact and makes the ledger's integral + # upper bound monotone under subdivision. + upper = min(upper, float(inherited_point_upper)) + angle_total = max(half[1] + half[2], 1.0e-300) + score = np.asarray([ + float(hi[3]) * A_time, + (float(hi[3]) * A_angle + 0.5 * hi[3] ** 2 * B_angle) + * half[1] / angle_total, + (float(hi[3]) * A_angle + 0.5 * hi[3] ** 2 * B_angle) + * half[2] / angle_total, + (abs(A_upper) + hi[3] * max(B0 + B_angle, 0.0) + + 4.0 / max(lo[3], 1.0e-300)) * half[3], + ]) + return math.log(volume) + upper, score, upper + + +def _logsumexp(values): + values = np.asarray(list(values), dtype=float) + if values.size == 0: + return -np.inf + top = float(np.max(values)) + if not np.isfinite(top): + return top + return top + math.log(float(np.sum(np.exp(values - top)))) + + +def hierarchical_union_cover( + C_A_t, uv_summary, centers, half_widths, x_min, x_max, *, + target_log_value, outside_tol_nats=-23.0, max_boxes=50000, + max_depth=(14, 10, 10, 10), progress_interval=512, + stall_checks=3, min_progress_nats=0.5): + """Adaptively upper-bound mass outside a union of local four-axis boxes. + + The largest unresolved coefficient-space box is split first. A box wholly + inside multiple local regions is assigned to one canonical nearest owner; + overlaps therefore become disjoint tiles rather than a failure. The cap + limits work, not safety: every unresolved leaf retains a valid upper bound. + """ + C_A_t = np.asarray(C_A_t, dtype=np.complex128) + if not isinstance(uv_summary, UVQSummary): + raise TypeError("uv_summary must come from summarize_uv_norm_table") + _validate_tables(C_A_t, uv_summary.C_B) + if not uv_summary.time_invariant: + raise ValueError("arrival-time-dependent U,V norm cannot be collapsed") + centers = np.asarray(centers, dtype=float) + half_widths = np.asarray(half_widths, dtype=float) + if (centers.ndim != 2 or centers.shape[1] != 4 or len(centers) == 0 + or half_widths.shape != centers.shape): + raise ValueError("centers and half_widths must have shape (N,4), N>0") + if np.any(~np.isfinite(centers)) or np.any(~np.isfinite(half_widths)): + raise ValueError("local boxes must be finite") + if np.any(half_widths <= 0.0) or np.any(half_widths[:, 1:3] > np.pi): + raise ValueError("half-widths must be positive and angular widths <= pi") + if not (0.0 < float(x_min) < float(x_max)): + raise ValueError("need 0 < x_min < x_max") + if not np.isfinite(float(target_log_value)): + raise ValueError("target_log_value must be finite") + max_boxes = int(max_boxes) + max_depth = np.asarray(max_depth, dtype=np.int32) + if max_boxes < 1 or max_depth.shape != (4,) or np.any(max_depth < 0): + raise ValueError("invalid cover cap") + + time_enclosure = _time_fourier_enclosure(C_A_t) + domain_lo = np.asarray([0.0, 0.0, 0.0, float(x_min)]) + domain_hi = np.asarray([ + float(C_A_t.shape[-1] - 1), 2.0 * np.pi, 2.0 * np.pi, + float(x_max)]) + heap = [] + owned = [] + counter = 0 + evaluated = 0 + overlap_owned = 0 + max_seen = np.zeros(4, dtype=np.int32) + progress = [] + best_tail = np.inf + stalled = False + + def add_box(lo, hi, depth, inherited_point_upper=np.inf): + nonlocal counter, evaluated, overlap_owned + owners, owner = _box_owners(lo, hi, centers, half_widths) + evaluated += 1 + max_seen[:] = np.maximum(max_seen, depth) + if owners: + owned.append((0.5 * (lo + hi), 0.5 * (hi - lo), owner)) + overlap_owned += int(len(owners) > 1) + return + log_upper, score, point_upper = _box_log_upper( + C_A_t, uv_summary.C_B, time_enclosure, lo, hi, + inherited_point_upper) + counter += 1 + heapq.heappush( + heap, (-log_upper, counter, lo, hi, depth, score, point_upper)) + + add_box(domain_lo, domain_hi, np.zeros(4, dtype=np.int32)) + initial_outside = _logsumexp(-item[0] for item in heap) + initial_tail = initial_outside - float(target_log_value) + best_tail = initial_tail + subdivisions = 0 + cap_reached = False + while heap: + outside = _logsumexp(-item[0] for item in heap) + tail = outside - float(target_log_value) + best_tail = min(best_tail, tail) + if outside - float(target_log_value) < float(outside_tol_nats): + break + if evaluated + 2 > max_boxes: + cap_reached = True + break + item = heapq.heappop(heap) + _, _, lo, hi, depth, score, parent_point_upper = item + available = depth < max_depth + if not np.any(available): + heapq.heappush(heap, item) + cap_reached = True + break + boundary = _local_boundary_splits( + lo, hi, centers, half_widths, available) + boundary_axes = np.asarray([bool(values) for values in boundary]) + split_score = np.where(available, score, -np.inf) + if np.any(boundary_axes): + # Use the same physics sensitivity score to order exact local-union + # cuts. These cuts establish ownership; midpoint cuts then tighten + # the complement enclosure. + axis = int(np.argmax(np.where(boundary_axes, score, -np.inf))) + midpoint = 0.5 * (lo[axis] + hi[axis]) + middle = min(boundary[axis], key=lambda value: abs(value - midpoint)) + else: + axis = int(np.argmax(split_score)) + if not np.isfinite(split_score[axis]) or hi[axis] <= lo[axis]: + axis = int(np.flatnonzero(available)[0]) + middle = 0.5 * (lo[axis] + hi[axis]) + child_depth = depth.copy() + child_depth[axis] += 1 + left_hi = hi.copy() + left_hi[axis] = middle + right_lo = lo.copy() + right_lo[axis] = middle + add_box(lo.copy(), left_hi, child_depth.copy(), parent_point_upper) + add_box(right_lo, hi.copy(), child_depth.copy(), parent_point_upper) + subdivisions += 1 + + if int(progress_interval) > 0 and evaluated >= ( + len(progress) + 1) * int(progress_interval): + checkpoint = _logsumexp(-leaf[0] for leaf in heap) + checkpoint_tail = checkpoint - float(target_log_value) + progress.append(checkpoint_tail) + best_tail = min(best_tail, checkpoint_tail) + if (len(progress) > int(stall_checks) + and checkpoint_tail > float(outside_tol_nats) + and progress[-1 - int(stall_checks)] - checkpoint_tail + < float(min_progress_nats)): + stalled = True + break + + outside = _logsumexp(-item[0] for item in heap) + tail_margin = outside - float(target_log_value) + owned_centers = np.asarray( + [item[0] for item in owned], dtype=float).reshape((-1, 4)) + owned_widths = np.asarray( + [item[1] for item in owned], dtype=float).reshape((-1, 4)) + owned_mode = np.asarray([item[2] for item in owned], dtype=np.int32) + finite = bool(np.isfinite(outside) or outside == -np.inf) + return CoverReport( + float(outside), float(tail_margin), finite, + bool(finite and tail_margin < float(outside_tol_nats)), + bool(cap_reached), int(evaluated), int(subdivisions), int(len(heap)), + int(len(owned)), int(overlap_owned), max_seen, + float(initial_tail), float(best_tail), bool(stalled), + np.asarray(progress, dtype=float), + owned_centers, owned_widths, owned_mode) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/peaklocal_time_reserve.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/peaklocal_time_reserve.py new file mode 100644 index 000000000..c1c07c1fd --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/peaklocal_time_reserve.py @@ -0,0 +1,529 @@ +"""Peak-local time rule for the four-axis policy's reserve. + +The policy's reserve integrates time on a rule that refines the WHOLE window +(``policy_time_rules``: ``refine`` nodes per native sample, escalating by +doubling to ``reserve_time_refine_max``). The peak of ``exp(lnL(t))`` has +width ``sigma_t = 1 / (2 pi rho sigma_f)``, so that rule's node count grows +with rho and the angular kernel is paid at every node: 18.6 / 72.6 / 290 / +1278 s per evaluation at rho 41 / 82 / 163 / 326 with the exact kernel on +2453 nodes (ladder controller, 2026-09-08), and at rho 41 the refine-4 rule +already misses its own 1e-3 warrant on two of three rows. + +This rule is sized from the PREDICTED width and its node count does not +depend on the row: + +* the width is predicted per row, ``sigma_t = 1 / (2 pi rho sigma_f)``, with + ``rho`` from the row's own coefficient table (the angular triangle bound of + the field over the U,V norm) and ``sigma_f`` the two-sided rms frequency of + the stored Q (the policy passes it; the table's own spectrum is the per-row + fallback and cross-check). Both are bounds in the safe direction: the + narrowest peak the primitive can make at the row's amplitude; +* ONE fine lattice per row, commensurate with the coarse scan: fine spacing + ``h_f = h_scan / m`` with the integer ``m`` chosen so that + ``h_f <= sigma_t / nodes_per_sigma``. Every node of the rule sits on that + lattice, so overlapping blocks and scan nodes inside a block are EXACT + duplicates (zero trapezoid weight), and the only spacing changes are at + block ends, where the mass is negligible. A dead slot collapses onto the + first scan node. Mixing two incommensurate + lattices was measured to cost 0.02 nat on a 0.75-sample peak: the + trapezoid rule is spectrally accurate only on uniform spacing; +* a coarse scan of the whole window at ``scan_refine`` nodes per native + sample (the broad, low-lying part of the field and any peak the locator + missed), and one block of ``n_fine`` consecutive fine-lattice nodes + centred on each of the row's time maxima. The maxima come from the + PRIMITIVE, not from the local branch's plan: :func:`locate_time_maxima` + maximizes the field over a dense phi grid, a u lattice and the distance at + every node of a fixed search grid (the envelope the marginalized reserve + integrates), keeps the ``n_candidates`` highest local maxima, and polishes + each on that envelope with parabolic steps. The plan's Newton centres and widths are reported beside + them as a cross-check and never decide the rule, so the reserve does not + inherit a decline of the local branch. A dead slot sits on the scan + lattice and adds nothing. + +The check rule is the same construction at half the scan refinement and every +other fine node (spacing ``2 h_f``), so it is everywhere coarser with the same +endpoints and the kernel's structural resolution warrant applies unchanged. +The two-guard comparison and the tail bound are the kernel's. Nothing here +is a bound: a peak the plans missed and the scan cannot resolve fails the +resolution warrant and the row escalates (``m`` and ``n_fine`` double, same +span) or returns ``nan``. An escalation is a report on the prediction, not +a loop. + +Evidence: DESIGN_direct_marginalization_policy.md, "Peak-local time reserve". +""" + +import jax +import jax.numpy as jnp +import numpy as np + +__all__ = ["NODES_PER_SIGMA", "peaklocal_time_rule", "peaklocal_rule_size", + "validate_peaklocal_rule_arguments", "row_amplitude", + "table_bandwidth_cycles", "predicted_width_samples", + "predict_time_rule", "locate_time_maxima"] + +# Fine nodes per predicted sigma_t. Three matches the policy's whole-window +# node accounting (_TIME_NODES_PER_SIGMA); the trapezoid rule on a Gaussian at +# spacing sigma/3 is converged to ~1e-8, and its half-refined check at +# 2 sigma/3 to ~1e-4, so the warrant's 1e-3 has margin on both sides. +NODES_PER_SIGMA = 3.0 + + +def validate_peaklocal_rule_arguments(n_target, n_fine, scan_refine): + n_target = int(n_target) + n_fine = int(n_fine) + scan_refine = int(scan_refine) + if n_target < 2: + raise ValueError("n_target must be >= 2") + if n_fine < 3 or n_fine % 2 == 0: + raise ValueError("n_fine must be an odd integer >= 3 so the check " + "rule keeps the block endpoints") + if scan_refine < 2 or scan_refine % 2: + raise ValueError("scan_refine must be an even integer >= 2 so the " + "check scan is the half-refined scan") + return n_target, n_fine, scan_refine + + +def peaklocal_rule_size(n_target, n_blocks, n_fine, scan_refine, n_scan=None): + """Node counts ``(reserve, check)`` of the rule: a shape statement, so a + test can pin that the count is independent of the row. ``n_scan`` is + the cropped scan's node count; ``None`` is the whole-window scan.""" + n_target, n_fine, scan_refine = validate_peaklocal_rule_arguments( + n_target, n_fine, scan_refine) + n_blocks = int(n_blocks) + if n_scan is None: + n_s = (n_target - 1) * scan_refine + 1 + n_sc = (n_target - 1) * (scan_refine // 2) + 1 + else: + n_s = int(n_scan) + if n_s < 3 or n_s % 2 == 0: + raise ValueError("n_scan must be an odd integer >= 3") + n_sc = (n_s - 1) // 2 + 1 + reserve = n_s + n_blocks * n_fine + check = n_sc + n_blocks * ((n_fine - 1) // 2 + 1) + return reserve, check + + +# ------------------------------------------------------------ prediction +def _kp_weights(kp): + return jnp.where(jnp.arange(kp) == 0, 1.0, 2.0) + + +def row_amplitude(C_A_t, C_B, guard): + """Predicted rho of one row from its tables: ``rho^2 = A_max^2 / B_min``. + + ``A_max`` is the angular triangle bound of the field over the target + window (``sum_kp,ks w |C_A|``, maximized over time) and ``B_min`` the + triangle lower bound of the U,V norm polynomial, the same bounds + ``all_axis_peaklocal._time_cell_cover_device`` uses. An upper bound on + rho, so a lower bound on the width: the lattice it sizes is never + coarser than the peak needs. + """ + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + guard = int(guard) + target = C_A_t[..., guard:-guard] if guard else C_A_t + wa = _kp_weights(C_A_t.shape[0]) + a_upper = jnp.max(jnp.sum(wa[:, None, None] * jnp.abs(target), axis=(0, 1))) + ks0 = (C_B.shape[1] - 1) // 2 + wb = _kp_weights(C_B.shape[0])[:, None] + b_centre = C_B[0, ks0].real + b_lower = b_centre - (jnp.sum(wb * jnp.abs(C_B)) - jnp.abs(C_B[0, ks0])) + rho_sq = jnp.where(b_lower > 0.0, jnp.square(a_upper) / b_lower, jnp.inf) + return jnp.sqrt(rho_sq) + + +def table_bandwidth_cycles(C_A_t, guard): + """Effective bandwidth of the row's own primitive, cycles per sample. + + ``sqrt()`` over ALL bins of the reflected series' spectrum (the + spectrum ``_time_primitive_spectrum`` refines with), summed over lanes + with the ``kp`` weights: the same two-sided convention as + ``q_effective_bandwidth_hz`` and the ladder record (0.01557 cycles per + sample on the ladder-2 tables, identical at rho 41, 163 and 652). The + RAW moment, not the spread about a carrier: a primitive with content on + both sides of zero frequency (the (2, -2) mode, a single polarization) + keeps carrier-scale structure after the angle marginalization, and + ```` bounds the curvature of anything the band-limited primitive can + make, so ``1 / (2 pi rho sqrt)`` is the NARROWEST peak the row can + produce at its amplitude. A face-on envelope is wider (by the ratio of + the raw to the central moment), which the plan's Newton width reports. + """ + from .time_first_peaklocal import _time_primitive_spectrum + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + flat = C_A_t.reshape((-1, C_A_t.shape[-1])) + coeff, frequency, _ = _time_primitive_spectrum(flat, int(guard)) + wa = jnp.broadcast_to(_kp_weights(C_A_t.shape[0])[:, None], + C_A_t.shape[:-1]).reshape(-1) + power = jnp.sum(wa[:, None] * jnp.square(jnp.abs(coeff)), axis=0) + total = jnp.sum(power) + msq = jnp.sum(jnp.square(frequency) * power) / total + return jnp.where(total > 0.0, jnp.sqrt(jnp.maximum(msq, 0.0)), jnp.nan) + + +def predicted_width_samples(rho, sigma_f_cycles): + """``sigma_t = 1 / (2 pi rho sigma_f)`` in native samples.""" + rho = jnp.asarray(rho, dtype=jnp.float64) + sigma_f = jnp.asarray(sigma_f_cycles, dtype=jnp.float64) + ok = jnp.isfinite(rho) & jnp.isfinite(sigma_f) & (rho > 0.0) & (sigma_f > 0.0) + return jnp.where(ok, 1.0 / (2.0 * jnp.pi * jnp.where(ok, rho, 1.0) + * jnp.where(ok, sigma_f, 1.0)), jnp.nan) + + +def predict_time_rule(rho, sigma_f_cycles, n_target, *, n_blocks, n_fine, + scan_refine, n_scan=None, nodes_per_sigma=NODES_PER_SIGMA): + """Host-side prediction for the pair selector: what each rule needs. + + Pure numpy on scalars the selector already has. ``prefer_peaklocal`` is + true when a whole-window rule at ``nodes_per_sigma`` per predicted sigma + would need more than twice the peak-local rule's fixed node count. + ``whole_window_nodes_needed`` is a linear density: a COST comparison + between two rules, not a convergence model (the trapezoid converges + spectrally, so the warrant, not this count, certifies a value). + """ + rho = float(rho) + sigma_f = float(sigma_f_cycles) + n_target, n_fine, scan_refine = validate_peaklocal_rule_arguments( + n_target, n_fine, scan_refine) + if rho > 0.0 and sigma_f > 0.0 and np.isfinite(rho) and np.isfinite(sigma_f): + sigma_t = 1.0 / (2.0 * np.pi * rho * sigma_f) + else: + sigma_t = float("nan") + if np.isfinite(sigma_t) and sigma_t > 0.0: + window_needed = (n_target - 1) * nodes_per_sigma / sigma_t + 1.0 + fine_refine = max(1, int(np.ceil(nodes_per_sigma / (sigma_t * scan_refine)))) + else: + window_needed = float("inf") + fine_refine = 1 + peaklocal_nodes, _ = peaklocal_rule_size(n_target, n_blocks, n_fine, scan_refine, + n_scan=n_scan) + h_f = 1.0 / (scan_refine * fine_refine) + return dict(rho=rho, sigma_f_cycles=sigma_f, sigma_t_samples=sigma_t, + fine_refine=fine_refine, fine_spacing_samples=h_f, + block_span_samples=(n_fine - 1) * h_f, + whole_window_nodes_needed=window_needed, + peaklocal_nodes=int(peaklocal_nodes), + prefer_peaklocal=bool(window_needed > 2.0 * peaklocal_nodes)) + + + +# --------------------------------------------------------------- locator +def _angular_grid_field(table, C_B, n_phi, n_u): + """``A[t, phi, u]`` and ``B[phi, u]`` on a dense phi grid and a u lattice + from the coefficient tables (``table`` is ``(KP,KS,Nt)``).""" + phi = jnp.arange(int(n_phi), dtype=jnp.float64) * (2.0 * jnp.pi / int(n_phi)) + u = jnp.arange(int(n_u), dtype=jnp.float64) * (2.0 * jnp.pi / int(n_u)) + kp_a = jnp.arange(table.shape[0], dtype=jnp.float64) + ks_a = jnp.arange(-(table.shape[1] - 1) // 2, (table.shape[1] - 1) // 2 + 1, + dtype=jnp.float64) + EA = (jnp.exp(1j * (phi[:, None, None, None] * kp_a[None, None, :, None] + + u[None, :, None, None] * ks_a[None, None, None, :])) + * _kp_weights(table.shape[0])[None, None, :, None]) + A = jnp.einsum("pukq,kqt->tpu", EA, table).real + kp_b = jnp.arange(C_B.shape[0], dtype=jnp.float64) + ks_b = jnp.arange(-(C_B.shape[1] - 1) // 2, (C_B.shape[1] - 1) // 2 + 1, + dtype=jnp.float64) + EB = (jnp.exp(1j * (phi[:, None, None, None] * kp_b[None, None, :, None] + + u[None, :, None, None] * ks_b[None, None, None, :])) + * _kp_weights(C_B.shape[0])[None, None, :, None]) + B = jnp.einsum("pukq,kq->pu", EB, C_B).real + return A, B + + +def _profile(A, B, x_min, x_max): + """``max_x (x A - x^2 B / 2)`` on ``[x_min, x_max]``; ``-inf`` where the + norm is not positive.""" + safe = jnp.where(B > 0.0, B, 1.0) + x = jnp.clip(A / safe, float(x_min), float(x_max)) + val = x * A - 0.5 * jnp.square(x) * B + return jnp.where(B > 0.0, val, -jnp.inf) + + +def locate_time_maxima(C_A_t, C_B, guard, n_target, x_min, x_max, *, + n_candidates, search_refine=8, angular_lattice=8, + n_phi=64, polish_iterations=4, keep_nats=30.0, + newton_steps=3, newton_step_max=0.1): + """Time maxima of the angle- and distance-maximized field of one row. + + Fixed shape. The ENVELOPE profile ``P(t) = max_{phi,u,x} (x A - x^2 + B / 2)`` is what the angle-marginalized reserve integrates in time, so + the maxima are found on it: phi on a dense ``n_phi`` grid (the maximum + over phi of the carrier's phase is what removes the carrier), u on an + ``angular_lattice`` grid, x in closed form, on a search grid of + ``search_refine`` nodes per native sample. The ``n_candidates`` highest + local maxima are each polished by ``polish_iterations`` parabolic steps + on ``P(t)`` at the candidate's ``u`` (phi and x re-maximized at every + evaluation), the step shrinking by four each time. A fixed-angle Newton + polish was measured to land up to ``1 / (16 f_c)`` samples (four samples + on a 64 Hz carrier) from the envelope maximum, because the fixed-angle + field peaks where its carrier phase does, not where the envelope does. + + Returns a dict of ``(n_candidates,)`` arrays: ``centres`` (native samples + of the unguarded window), ``widths`` (``1 / sqrt(-P'')``, the peak's + sigma from the envelope's curvature at the maximum), ``values`` (``P`` + there), ``live`` (finite, interior, curved, within ``keep_nats`` of the + best), plus ``n_search_nodes``. + """ + from .time_first_peaklocal import _time_primitive_spectrum, _evaluate_time_spectrum + C_A_t = jnp.asarray(C_A_t, dtype=jnp.complex128) + C_B = jnp.asarray(C_B, dtype=jnp.complex128) + guard = int(guard) + n_target = int(n_target) + K = int(n_candidates) + R = int(search_refine) + L = int(angular_lattice) + if K < 1 or R < 1 or L < 2 or int(n_phi) < 4 or int(polish_iterations) < 1: + raise ValueError("need n_candidates >= 1, search_refine >= 1, " + "angular_lattice >= 2, n_phi >= 4, polish_iterations >= 1") + flat = C_A_t.reshape((-1, C_A_t.shape[-1])) + coeff, frequency, offset = _time_primitive_spectrum(flat, guard) + shape = C_A_t.shape[:-1] + t_s = jnp.arange((n_target - 1) * R + 1, dtype=jnp.float64) / float(R) + table = _evaluate_time_spectrum(coeff, frequency, t_s, offset).reshape( + shape + (t_s.size,)) + A, B = _angular_grid_field(table, C_B, n_phi, L) # (t,phi,u), (phi,u) + prof = _profile(A, B[None], x_min, x_max) + P_u = jnp.max(prof, axis=1) # (t, u): max over phi + P = jnp.max(P_u, axis=1) # (t,) + u_idx = jnp.argmax(P_u, axis=1) # (t,) + left = jnp.concatenate((jnp.array([-jnp.inf]), P[:-1])) + right = jnp.concatenate((P[1:], jnp.array([-jnp.inf]))) + is_max = (P >= left) & (P >= right) & jnp.isfinite(P) + scores, idx = jax.lax.top_k(jnp.where(is_max, P, -jnp.inf), K) + t0 = t_s[idx] + u_star = u_idx[idx] # (K,) + + kp_a = jnp.arange(C_A_t.shape[0], dtype=jnp.float64) + ks_a = jnp.arange(-(C_A_t.shape[1] - 1) // 2, (C_A_t.shape[1] - 1) // 2 + 1, + dtype=jnp.float64) + kp_b = jnp.arange(C_B.shape[0], dtype=jnp.float64) + ks_b = jnp.arange(-(C_B.shape[1] - 1) // 2, (C_B.shape[1] - 1) // 2 + 1, + dtype=jnp.float64) + wa = _kp_weights(C_A_t.shape[0]) + wb = _kp_weights(C_B.shape[0]) + ang = jnp.arange(L, dtype=jnp.float64) * (2.0 * jnp.pi / L) + u_val = ang[u_star] # (K,) + phi_grid = jnp.arange(int(n_phi), dtype=jnp.float64) * (2.0 * jnp.pi / int(n_phi)) + + def _profile_at(phi, u, a_lanes): + """Scalar profile at one (phi, u) from one time node's table.""" + EA = jnp.exp(1j * (phi * kp_a[:, None] + u * ks_a[None, :])) * wa[:, None] + A = jnp.sum(EA * a_lanes).real + EB = jnp.exp(1j * (phi * kp_b[:, None] + u * ks_b[None, :])) * wb[:, None] + Bv = jnp.sum(EB * C_B).real + return _profile(A, Bv, x_min, x_max) + + def _profile_ang(ang, a_lanes): + return _profile_at(ang[0], ang[1], a_lanes) + + grad_ang = jax.grad(_profile_ang) + hess_ang = jax.hessian(_profile_ang) + + def _envelope(t, return_angles=False): + """``P(t)`` at the candidates' own (phi, u), both re-maximized: the + best of a dense phi grid at the lattice u, then two Newton steps in + (phi, u) jointly on the trigonometric polynomial. A grid maximum + alone leaves a ripple of period ``1 / (n_phi f_c)`` in t that a + parabola reads as curvature, and a lattice u moves the t-maximum + (measured 0.17 samples, 2.7 sigma, on a rung-160 production row).""" + tab = _evaluate_time_spectrum(coeff, frequency, t, offset).reshape( + shape + (t.size,)) # (KP,KS,K) + lanes = jnp.moveaxis(tab, -1, 0) # (K,KP,KS) + grid = jax.vmap(lambda a, u: jax.vmap( + lambda p: _profile_at(p, u, a))(phi_grid))(lanes, u_val) # (K,n_phi) + ang = jnp.stack((phi_grid[jnp.argmax(grid, axis=1)], u_val), axis=1) # (K,2) + + def _newton(a, _): + g = jax.vmap(grad_ang)(a, lanes) # (K,2) + h = jax.vmap(hess_ang)(a, lanes) # (K,2,2) + # Newton on a maximum. A direction with no curvature (a table + # with no u content has h_uu = 0 exactly) must not stall the + # other: ridge the Hessian and zero the step in any direction + # whose own curvature is not negative. + ridge = 1.0e-8 * (1.0 + jnp.abs(h[:, 0, 0]) + jnp.abs(h[:, 1, 1])) + h_reg = h - ridge[:, None, None] * jnp.eye(2)[None] + step = -jnp.linalg.solve(h_reg, g[..., None])[..., 0] + concave = jnp.stack((h[:, 0, 0] < 0.0, h[:, 1, 1] < 0.0), axis=1) + step = jnp.where(concave, jnp.clip(step, -float(newton_step_max), + float(newton_step_max)), 0.0) + return a + step, None + + ang, _ = jax.lax.scan(_newton, ang, None, length=int(newton_steps)) + vals = jax.vmap(_profile_ang)(ang, lanes) + return (vals, ang) if return_angles else vals + + def _step(carry, _): + t, delta = carry + p0 = _envelope(t) + pp = _envelope(jnp.clip(t + delta, 0.0, float(n_target - 1))) + pm = _envelope(jnp.clip(t - delta, 0.0, float(n_target - 1))) + curv = (pp + pm - 2.0 * p0) / jnp.square(delta) + move = jnp.where(curv < 0.0, -0.5 * delta * (pp - pm) + / jnp.where(curv < 0.0, pp + pm - 2.0 * p0, -1.0), 0.0) + move = jnp.clip(move, -delta, delta) + t_new = jnp.clip(t + move, 0.0, float(n_target - 1)) + return (t_new, delta * 0.25), curv + + delta0 = jnp.full((K,), 1.0 / float(R)) + (t_ref, _), curvs = jax.lax.scan(_step, (t0, delta0), None, + length=int(polish_iterations)) + curv = curvs[-1] + values, angles = _envelope(t_ref, return_angles=True) + curved = curv < 0.0 + widths = jnp.where(curved, 1.0 / jnp.sqrt(jnp.where(curved, -curv, 1.0)), jnp.inf) + best = jnp.max(values) + # The profile maximum IS the row's lnL maximum, so rho^2 = 2 P_max: the + # amplitude the prediction should use. The angular triangle bound is an + # upper bound that ran 2.8x over on a rung-160 production row (453 + # against 157), which narrowed the predicted width and the block span by + # the same factor. + rho_located = jnp.where(jnp.isfinite(best) & (best > 0.0), + jnp.sqrt(2.0 * jnp.maximum(best, 0.0)), jnp.nan) + live = (jnp.isfinite(values) & curved & jnp.isfinite(scores) + & (values >= best - float(keep_nats)) + & (t_ref > 0.0) & (t_ref < float(n_target - 1))) + return dict(centres=t_ref, widths=widths, values=values, live=live, + phi=angles[:, 0], u=angles[:, 1], rho_located=rho_located, + search_centres=t0, n_search_nodes=jnp.asarray(t_s.size), + search_positions=t_s, search_profile=P) + + +# ------------------------------------------------------------------ rule +def _trapezoid_weights(sorted_nodes): + """Trapezoid weights on a sorted, possibly repeated, rule (sample units).""" + left = jnp.concatenate((sorted_nodes[:1], sorted_nodes[:-1])) + right = jnp.concatenate((sorted_nodes[1:], sorted_nodes[-1:])) + # Clamped: a repeated position differs by a rounding unit under XLA's + # fused subtraction, and a weight of -1e-16 would fail the kernel's + # non-negativity check. + return jnp.maximum(0.5 * (right - left), 0.0) + + +def peaklocal_time_rule(centres, widths, live, n_target, dt_scale, *, + sigma_t_samples, n_fine, scan_refine, + fine_refine_multiplier=1, + nodes_per_sigma=NODES_PER_SIGMA, + n_scan=None, margin_samples=None, + search_positions=None, search_profile=None, + outside_slack_nats=5.0): + """Per-row reserve rule and check rule on one commensurate lattice. + + ``centres``, ``widths`` (sigma at each maximum, samples) and ``live`` are + the row's time maxima (:func:`locate_time_maxima`). ``sigma_t_samples`` + is the predicted width (traced scalar); the fine spacing is the smaller + of the prediction and the narrowest live maximum, divided by + ``nodes_per_sigma``, so the lattice is never coarser than either says. + ``dt_scale`` is the seconds-per-sample constant of the production rule + (``sum(w_t) / (npts - 1)``) so the weights land in the units of + ``policy_time_rules``. ``fine_refine_multiplier`` is the escalation + tier (1, 2, 4, ...): the fine lattice is refined by it and ``n_fine`` + must have been widened to match (``n -> 2n - 1``) so the block span is + unchanged. + + With ``n_scan`` the scan is SUPPORT-LIMITED: ``n_scan`` nodes on the + scan lattice across the hull of the live maxima widened by + ``margin_samples`` each side (never spanning the window), and the mass + outside the hull is bounded from the locator's ``search_profile`` (the + angle- and distance-maximized exponent, an upper bound on the marginal at + each search node) summed over the outside search nodes, plus + ``outside_slack_nats`` for what lies between nodes. That bound goes to + the kernel's cropped-cover warrant, which charges it to the error budget. + A node count that grows with the window is the wrong design (RO, + 2026-09-09); the count here is ``n_scan + n_blocks n_fine`` whatever the + window. Returns a dict with the fixed-shape ``nodes``, ``weights``, + ``check_nodes``, ``check_weights`` and the numbers behind them. + """ + n_target, n_fine, scan_refine = validate_peaklocal_rule_arguments( + n_target, n_fine, scan_refine) + mult = int(fine_refine_multiplier) + if mult < 1: + raise ValueError("fine_refine_multiplier must be >= 1") + centres = jnp.asarray(centres, dtype=jnp.float64) + widths = jnp.asarray(widths, dtype=jnp.float64) + live = jnp.asarray(live).astype(bool) & jnp.isfinite(centres) + sigma_pred = jnp.asarray(sigma_t_samples, dtype=jnp.float64) + located = jnp.min(jnp.where(live & (widths > 0.0), widths, jnp.inf)) + sigma_t = jnp.minimum(jnp.where(jnp.isfinite(sigma_pred) & (sigma_pred > 0.0), + sigma_pred, jnp.inf), located) + finite = jnp.isfinite(sigma_t) & (sigma_t > 0.0) + need = float(nodes_per_sigma) / (jnp.where(finite, sigma_t, 1.0) * scan_refine) + m_pred = jnp.where(finite, jnp.maximum(1.0, jnp.ceil(need)), 1.0) + # Guard the lattice against an absurd prediction: the fine index space + # is (n_target - 1) * scan_refine * m long and must stay exact in + # float64 and cheap to build. + m_pred = jnp.minimum(m_pred, 2.0 ** 20).astype(jnp.int64) + m = m_pred * mult + h_f = 1.0 / (scan_refine * m).astype(jnp.float64) + n_index = (n_target - 1) * scan_refine * m # last fine index + + half_span = 0.5 * (n_fine - 1) # in fine indices + start = jnp.round(centres / h_f - half_span) + start = jnp.clip(start, 0.0, jnp.maximum(0.0, n_index - (n_fine - 1))) + start = jnp.where(live, start, 0.0).astype(jnp.int64) + k = jnp.arange(n_fine) + fine_live_idx = jnp.minimum(start[:, None] + k[None, :], n_index) # (K, n_fine) + cropped = n_scan is not None + if not cropped: + scan_idx = jnp.arange((n_target - 1) * scan_refine + 1) * m + scan_idx_c = jnp.arange((n_target - 1) * (scan_refine // 2) + 1) * (2 * m) + hull_lo = jnp.asarray(0.0) + hull_hi = jnp.asarray(float(n_target - 1)) + outside_log_bound = jnp.asarray(-jnp.inf) + n_outside = jnp.asarray(0, dtype=jnp.int32) + else: + n_s = int(n_scan) + if n_s < 3 or n_s % 2 == 0: + raise ValueError("n_scan must be an odd integer >= 3") + if margin_samples is None or search_positions is None or search_profile is None: + raise ValueError("a cropped scan needs margin_samples, " + "search_positions and search_profile") + margin = jnp.asarray(margin_samples, dtype=jnp.float64) + any_live = jnp.any(live) + c_lo = jnp.min(jnp.where(live, centres, jnp.inf)) + c_hi = jnp.max(jnp.where(live, centres, -jnp.inf)) + hull_lo = jnp.where(any_live, jnp.clip(c_lo - margin, 0.0, float(n_target - 1)), 0.0) + hull_hi = jnp.where(any_live, jnp.clip(c_hi + margin, 0.0, float(n_target - 1)), + float(n_target - 1)) + # Scan lattice indices (multiples of m) across the hull: a step of + # whole scan cells, at least one, so the scan stays on the lattice. + lo_idx = jnp.floor(hull_lo / h_f / m).astype(jnp.int64) * m + hi_idx = jnp.minimum(jnp.ceil(hull_hi / h_f / m).astype(jnp.int64) * m, n_index) + step = jnp.maximum(1, jnp.ceil((hi_idx - lo_idx) / (m * (n_s - 1))).astype(jnp.int64)) * m + scan_idx = jnp.minimum(lo_idx + jnp.arange(n_s) * step, n_index) + scan_idx_c = jnp.minimum(lo_idx + jnp.arange((n_s - 1) // 2 + 1) * (2 * step), n_index) + hull_lo = lo_idx.astype(jnp.float64) * h_f + hull_hi = jnp.minimum(lo_idx + (n_s - 1) * step, n_index).astype(jnp.float64) * h_f + t_s = jnp.asarray(search_positions, dtype=jnp.float64) + P_s = jnp.asarray(search_profile, dtype=jnp.float64) + outside = (t_s < hull_lo) | (t_s > hull_hi) + ds = (t_s[1] - t_s[0]) * float(dt_scale) + outside_log_bound = (jax.scipy.special.logsumexp(jnp.where(outside, P_s, -jnp.inf)) + + jnp.log(ds) + float(outside_slack_nats)) + n_outside = jnp.count_nonzero(outside).astype(jnp.int32) + # A dead slot collapses onto the scan's first node: exact duplicates, + # zero weight, and no gap that the check rule would share. + fine_idx = jnp.where(live[:, None], fine_live_idx, scan_idx[0]) + idx = jnp.sort(jnp.concatenate((scan_idx, fine_idx.ravel()))) + nodes = idx.astype(jnp.float64) * h_f + # Check rule: the half-refined scan and every other fine node, on the + # same lattice, same endpoints, everywhere coarser. + fine_idx_c = fine_idx[:, ::2] + idx_c = jnp.sort(jnp.concatenate((scan_idx_c, fine_idx_c.ravel()))) + check_nodes = idx_c.astype(jnp.float64) * h_f + weights = _trapezoid_weights(nodes) * float(dt_scale) + check_weights = _trapezoid_weights(check_nodes) * float(dt_scale) + n_live = jnp.count_nonzero(live).astype(jnp.int32) + first_centre = jnp.where(n_live > 0, centres[jnp.argmax(live)], jnp.nan) + return dict(nodes=nodes, weights=weights, check_nodes=check_nodes, + check_weights=check_weights, + scan_lo_samples=hull_lo, scan_hi_samples=hull_hi, + outside_log_bound=outside_log_bound, + n_outside_search_nodes=n_outside, + sigma_t_pred_samples=sigma_pred, + sigma_t_located_samples=located, + sigma_t_used_samples=sigma_t, + fine_refine=m.astype(jnp.int32), + fine_spacing_samples=h_f, + block_span_samples=(n_fine - 1) * h_f, + n_live_blocks=n_live, + first_block_centre_samples=first_centre, + prediction_finite=finite) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_rotating_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_rotating_freqresponse.py new file mode 100644 index 000000000..bb4f12326 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_rotating_freqresponse.py @@ -0,0 +1,115 @@ +"""JAX coefficients for simultaneous slow rotation and finite-arm response. + +This is the JAX analogue of +``factored_likelihood_rotating_freqresponse.combined_response_coefficients_vector``. +The packed elementary-template index is ``a=(b,p,n)``: finite-frequency basis, +delay-derivative order, and sidereal harmonic. +""" + +import math + +import numpy as np +import jax.numpy as jnp + +from . import response_freqresponse as _rf +from . import response_slowrot as _rs + + +def response_harmonic_width(basis_index): + return 2 if int(basis_index) == 0 else int(basis_index) + 1 + + +def _basis_harmonics(response, x_arm, y_arm, dec, psi, Qmax): + """Exact small-DFT coefficients of each finite-response sky factor.""" + dec = jnp.asarray(dec, dtype=jnp.float64) + psi = jnp.asarray(psi, dtype=jnp.float64) + D = jnp.asarray(response, dtype=jnp.float64) + xa = jnp.asarray(np.asarray(x_arm, dtype=float)) + ya = jnp.asarray(np.asarray(y_arm, dtype=float)) + + width = int(Qmax) + 2 + ngrid = 2 * width + 1 + g = 2.0 * jnp.pi * jnp.arange(ngrid, dtype=jnp.float64) / float(ngrid) + X, Y, nhat = _rf._triad_jax(dec[:, None], psi[:, None], g[None, :]) + Fp, Fc = _rf._lwl_response_jax(D, X, Y) + zx = jnp.einsum('...i,i->...', X, xa) + 1j * jnp.einsum('...i,i->...', Y, xa) + zy = jnp.einsum('...i,i->...', X, ya) + 1j * jnp.einsum('...i,i->...', Y, ya) + ax = jnp.einsum('...i,i->...', nhat, xa) + ay = jnp.einsum('...i,i->...', nhat, ya) + + values = {0: Fp + 1j * Fc} + for q in range(int(Qmax) + 1): + values[1 + q] = 0.5 * (zx ** 2 * ax ** q - zy ** 2 * ay ** q) + + out = {} + for b, vals in values.items(): + bw = response_harmonic_width(b) + out[b] = { + n: jnp.mean(vals * jnp.exp(-1j * n * g)[None, :], axis=1) + for n in range(-bw, bw + 1) + } + return out + + +def coefficients_dict(response, location, x_arm, y_arm, RA, DEC, psi, + gmst_tref, Qmax, p_max): + """Return ``{(b,p,n): (S,)}`` compound response coefficients.""" + RA = jnp.asarray(RA, dtype=jnp.float64) + DEC = jnp.asarray(DEC, dtype=jnp.float64) + psi = jnp.asarray(psi, dtype=jnp.float64) + g_ev = float(gmst_tref) - RA + + basis = _basis_harmonics(response, x_arm, y_arm, DEC, psi, Qmax) + basis_tref = { + b: {n: value * jnp.exp(1j * n * g_ev) for n, value in harmonics.items()} + for b, harmonics in basis.items() + } + + delay = _rs._delay_harmonics_jax(location, DEC) + delay_tref = {m: value * jnp.exp(1j * m * g_ev) + for m, value in delay.items()} + tau0 = jnp.real(sum(delay_tref.values())) + drift = dict(delay_tref) + drift[0] = drift[0] - tau0 + neg_drift = {m: -value for m, value in drift.items()} + + out = {} + expansion = {0: jnp.ones_like(g_ev, dtype=jnp.complex128)} + for p in range(int(p_max) + 1): + if p: + expansion = _rs._convolve_harmonics(expansion, neg_drift) + inv_fact = 1.0 / math.factorial(p) + for b, harmonics in basis_tref.items(): + for n, amplitude in harmonics.items(): + for m, delay_amplitude in expansion.items(): + key = (b, p, n + m) + out[key] = (out.get(key, 0.0) + + inv_fact * amplitude * delay_amplitude) + return out + + +def coefficients_packed(response, location, x_arm, y_arm, RA, DEC, psi, + gmst_tref, Qmax, p_max, a_list): + """Return compound coefficients as an ``(A,S)`` array aligned to ``a_list``.""" + S = int(jnp.asarray(RA).shape[0]) + coeff = coefficients_dict(response, location, x_arm, y_arm, RA, DEC, psi, + gmst_tref, Qmax, p_max) + rows = [] + for a in a_list: + key = tuple(int(v) for v in a) + rows.append(jnp.broadcast_to(coeff.get( + key, jnp.zeros((S,), dtype=jnp.complex128)), (S,))) + return jnp.stack(rows, axis=0).astype(jnp.complex128) + + +def reflection_index(a_list): + """Map ``(b,p,n)`` to ``(b,p,-n)`` for the V contraction.""" + keys = [tuple(int(v) for v in a) for a in a_list] + pos = {a: i for i, a in enumerate(keys)} + reflected = [] + for b, p, n in keys: + key = (b, p, -n) + if key not in pos: + raise ValueError("reflection partner %r absent from compound a_list" % (key,)) + reflected.append(pos[key]) + return np.asarray(reflected, dtype=np.int64) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py index 31b984bb6..bcb06ec75 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py @@ -223,11 +223,11 @@ def rotation_coefficients_packed(response, location, RA, DEC, psi, gmst_tref, # factored_likelihood_with_rotation.rotation_post_phase). # --------------------------------------------------------------------------- def harmonic_indices(a_list): - """Sidereal harmonic ``n_a`` of each elementary template ``a = (p, n)``. + """Sidereal harmonic ``n_a`` (the final field of an elementary index). Returns an ``(A,)`` int numpy array (static; used to index the post-phase table). """ - return np.asarray([int(n) for (_p, n) in a_list], dtype=np.int64) + return np.asarray([int(a[-1]) for a in a_list], dtype=np.int64) def post_phase_bucketing(a_list): diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py index b53dc39c6..5bfcca84c 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py @@ -19,6 +19,9 @@ estimate. * :func:`flowmc_sample` -- a normalizing-flow sampler (flowMC) that trains a flow to the full multimodal geometry in a single run. +* :func:`adaptive_volume_sample` -- the production RIFT AV or AV+GMM portfolio + control logic around fixed-shape, value-only JAX likelihood batches. Its + optional AD work is confined to a short Fisher-sky initializer. The priors are the standard physical ones (uniform sky/orientation): ``ra ~ U(0, 2pi)``, ``sin(dec) ~ U(-1, 1)``, ``psi ~ U(0, pi)``, @@ -38,13 +41,37 @@ import jax import jax.numpy as jnp +# ``core`` pulls in lal/lalsimulation (see core.py's own module-level import). +# test/jax/test_nuts_phimarg.py deliberately loads *this* file standalone, by +# path, with no lal on hand -- so the relative import below must not run in +# that context. A module loaded via importlib.util.spec_from_file_location() +# with no parent package gets __package__ == "" (falsy); a normal package +# import (``import RIFT.likelihood.jax_ile.samplers``) gets the real dotted +# package name (truthy). Gate on that instead of a bare ``from . import +# core``, which raises ImportError unconditionally outside a package. +if __package__: + from . import core as _core +else: + _core = None + # Default chunk for the batched lnL evals. The per-sample distance quadrature -# (JAX_ILE_DISTMARG_GH=G) materialises a (chunk, npts, G) array, ~G/ (grid_block) -# more device memory than the legacy grid, so a 4000-row chunk OOMs the 11GB -# 2080Ti. Shrink the chunk when per-sample is active so the per-sample path fits -# small-VRAM GPUs too (don't force every high-SNR job onto a 24GB node). -_GH_NODES = int(os.environ.get("JAX_ILE_DISTMARG_GH", "0")) -_EVAL_CHUNK = max(500, 4000 * 16 // max(16, _GH_NODES)) if _GH_NODES > 0 else 4000 +# (JAX_ILE_DISTMARG_GH=G, or the driver's --distance-gh-nodes) materialises a +# (chunk, npts, G) array, ~G/(grid_block) more device memory than the legacy +# grid, so a 4000-row chunk OOMs the 11GB 2080Ti. Shrink the chunk when +# per-sample is active so the per-sample path fits small-VRAM GPUs too (don't +# force every high-SNR job onto a 24GB node). +# +# Resolved through _core.get_distmarg_gh_nodes() at CALL time (see +# _default_eval_chunk below), not baked in here at import time: the CLI path +# (core.set_distmarg_gh_nodes(), called from the driver's option parsing) runs +# AFTER this module is first imported, so a module-level constant read from +# os.environ here would miss a node count set only via --distance-gh-nodes. +# Standalone-loaded (_core is None): fall back to the env var directly, same +# as the pre-CLI behaviour, since there is no driver to call set_ from. +def _default_eval_chunk(): + n = (_core.get_distmarg_gh_nodes() if _core is not None + else int(os.environ.get("JAX_ILE_DISTMARG_GH", "0"))) + return max(500, 4000 * 16 // max(16, n)) if n > 0 else 4000 # Parameter order used everywhere in this module. ANG_NAMES = ("ra", "dec", "psi", "incl", "phiref") @@ -231,7 +258,7 @@ def _log_prior_jax(theta5): # --------------------------------------------------------------------------- # Batched lnL evaluation (chunked to bound memory) # --------------------------------------------------------------------------- -# Largest single XLA buffer of the anglemarg laplace path, per sample per +# Historical 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 @@ -240,8 +267,372 @@ def _log_prior_jax(theta5): # 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. +# +# The laplace kernel now tiles the combined sample-time point axis internally, +# so this is no longer its literal largest-buffer model. Keep the outer cap as +# a conservative bound on still-live coefficient tables and phi fields, and for +# exact, which does not share that tiler. +# +# "BOTH" MEANS EXACT AND LAPLACE, AND NOTHING ELSE. A reviewer read it as covering +# every scheme and concluded this constant understates peak-local by ~128x. It does -- +# peak-local's live slab is about 1 MiB per sample-point, not 8 KiB -- but peak-local +# does not USE this number as its model: angle_marg_eval_chunk raises `bytes_per` to a +# scheme-specific peak-local model with max(), so 8192 acts only as a floor there. The +# two figures are both right, for different schemes. Do not "reconcile" them. _ANGLE_MARG_BYTES_PER_SAMPLE_PT = 8192 -_ANGLE_MARG_BUFFER_TARGET = 4 << 30 # ~4 GiB largest single buffer + +#: Largest single buffer we will let the anglemarg eval request. 4 GiB was chosen on +#: 2026-08-28 against the 25 GiB per-UID cgroup of the machine the OOM was reproduced on, +#: with a deliberate ~6x margin. On a card with more memory it throttles the accurate +#: schemes for no reason -- at npts=1230 it caps the eval chunk at 426 where the nominal +#: chunk is 1000, so `exact`/`laplace`/`peak-local` run at under half the batch `grid` +#: gets, and small batches are exactly where their per-sample cost is worst. +#: So DERIVE it from the device when we can see one, and keep 4 GiB as the fallback for the +#: machine we cannot measure. Deliberately a fraction of free VRAM rather than all of it: +#: this bounds ONE buffer, and the rest of the graph has to live alongside it. +#: +#: THIS IS NOT A FLOOR, and an earlier revision of this file wrongly said it was. 4 GiB is +#: what we use when we cannot SEE the device; it carries no guarantee about a device we can. +#: It was measured safe against one 25 GiB cgroup and says nothing about a 6 GiB card. +_ANGLE_MARG_BUFFER_TARGET_FALLBACK = 4 << 30 + +#: Fraction of the device's AVAILABLE memory to allow for this ONE buffer. +#: NOT a fraction of the reported limit, and review caught that it was: `bytes_limit` and +#: `bytes_reservable_limit` are capacity CEILINGS, not free memory. These cards are +#: SHARED -- a contemporaneous survey of the interactive hosts found all four GPUs at 100% +#: utilisation with 18-22 GiB of 24 GiB already held by other users -- so half of a 24 GiB +#: ceiling is 12 GiB on a card with 2 GiB left, i.e. exactly the RESOURCE_EXHAUSTED this +#: cap exists to prevent, wearing device awareness as a costume. The fraction is HEADROOM +#: ON WHAT IS FREE; the ceiling never licenses an allowance by itself. +#: WHY 0.5 RATHER THAN A MEASURED NUMBER: the remaining margin has to cover the rest of the +#: graph alongside this buffer, and that has NOT been measured -- an attempt was defeated by +#: the interactive hosts' thread cap. 0.5 is therefore a JUDGEMENT, not a result: it is +#: twice the first guess and still leaves half the reported limit. Override it when you +#: know your card is yours: +#: RIFT_ANGLEMARG_BUFFER_FRACTION=0.8 +#: and if you measure the true overhead, replace this constant with the measurement and say +#: so here. +_ANGLE_MARG_BUFFER_FRACTION_DEFAULT = 0.5 + + +def _read_buffer_fraction(env=None): + """Parse RIFT_ANGLEMARG_BUFFER_FRACTION, refusing a value that cannot bound anything. + + Refuses LOUDLY rather than quietly substituting the default. An override that is + silently ignored is worse than no override at all: the caller goes on believing a + bound is in force that is not, which is precisely how the buffer gets sized wrong. + Not being set is not an error -- only a value we were handed and cannot use. + + Above 1.0 is rejected rather than clamped because it asks for a buffer larger than + the device reports FREE, i.e. it asks this function to cause the OOM it exists to + prevent. A caller who really wants everything currently free writes 1.0. + """ + if env is None: + env = os.environ + raw = env.get("RIFT_ANGLEMARG_BUFFER_FRACTION") + if raw is None: + return _ANGLE_MARG_BUFFER_FRACTION_DEFAULT + try: + val = float(raw) + except (TypeError, ValueError): + raise ValueError( + "RIFT_ANGLEMARG_BUFFER_FRACTION=%r is not a number; give a fraction in " + "(0, 1], e.g. 0.8" % (raw,)) + # NaN fails this comparison too, which is the intent. + if not (0.0 < val <= 1.0): + raise ValueError( + "RIFT_ANGLEMARG_BUFFER_FRACTION=%r is outside (0, 1]; above 1 would size this " + "buffer larger than the device reports, and at or below 0 it bounds nothing" + % (raw,)) + return val + + +_ANGLE_MARG_BUFFER_FRACTION = _read_buffer_fraction() + + +def _read_buffer_bytes(env=None): + """Parse RIFT_ANGLEMARG_BUFFER_BYTES, an ABSOLUTE allowance in bytes, or None. + + WHY A SECOND KNOB EXISTS. ``angle_marg_eval_chunk`` now REFUSES a configuration + whose single sample already exceeds the allowance, because returning a chunk of 1 + there breaks the bound it advertises. On a machine whose device we cannot read, the + allowance being refused against is ``_ANGLE_MARG_BUFFER_TARGET_FALLBACK`` -- a + documented guess that the comment above is explicit carries no guarantee. Failing + closed against a guess with no way to override it turns "we could not see your + device" into "you may not run", which is an outage, not a bound. + + RIFT_ANGLEMARG_BUFFER_FRACTION cannot serve this purpose: it is a fraction OF a + reported FREE figure, and the paths that need the escape -- no readable device, or a + device that reports a ceiling but never says how much of it is free -- are exactly + the ones with no such figure to take a fraction of. + + Read per call rather than once at import so a caller can set it before the eval + without re-importing the module. Refused loudly on garbage, for the same reason the + fraction is: an override that is silently dropped leaves the caller believing a + bound is in force that is not. + """ + if env is None: + env = os.environ + raw = env.get("RIFT_ANGLEMARG_BUFFER_BYTES") + if raw is None: + return None + try: + val = int(float(raw)) + except (TypeError, ValueError, OverflowError): + # OverflowError is in the list because int(float('inf')) raises it and not + # ValueError, so 'inf' would otherwise escape as an unhandled OverflowError + # instead of the actionable message. 'nan' goes the ValueError route. + raise ValueError( + "RIFT_ANGLEMARG_BUFFER_BYTES=%r is not a usable number of bytes; give a " + "positive integer, e.g. %d for 12 GiB" % (raw, 12 << 30)) + if val <= 0: + raise ValueError( + "RIFT_ANGLEMARG_BUFFER_BYTES=%r is not positive; a non-positive allowance " + "bounds nothing and refuses every chunk" % (raw,)) + return val + + +def _device_available_bytes(stats): + """Bytes we can actually expect to get from the device NOW, or None if unknowable. + + THE CEILING IS NOT THE ANSWER, which was a review finding on this file. Neither + ``bytes_limit`` nor ``bytes_reservable_limit`` says anything about what is free: they + are what the allocator may grow to, on a card another process may already be sitting + on. Sizing off either one returns a 12 GiB allowance on a shared 24 GiB GPU with + 2 GiB left, which is the failure this cap exists to prevent. + + Only keys that mean "free" are read: + + * ``largest_free_block_bytes`` -- the largest contiguous block the allocator can + serve right now. It answers the question actually being asked, because the thing + being bounded is ONE allocation, not a total. But on jax 0.9.2 (measured + 2026-09-08) this key is simply never populated: it reads 0 both before any + allocation and after one, on an otherwise-idle card. A bare 0 here is therefore + read as "not reported", not "the device is full", and falls through to the pool + signal below. + * failing that, the reserved pool minus what we hold in it. Memory already + reserved for this process cannot be taken by another one, so ``pool - in_use`` is + genuinely ours in a way the ceiling is not. ``pool_bytes`` is also 0 on jax 0.9.2 + until the first allocation grows the pool, which is likewise "not yet known", not + "nothing is free" -- a 0 pool falls through the same way a 0 block does. + + Returns None when neither is reported (or both report 0). The caller must read that + as "we could not see how much of this device is free" -- NOT as zero, and emphatically + not as the ceiling that is sitting right there in the same dict. + + A device that genuinely has nothing left -- a nonzero pool fully consumed + (``pool_bytes > 0`` and ``pool_bytes - bytes_in_use == 0``) -- still returns 0, not + None, and that is deliberate: it is a reading, not a failure to read. Falling back to + the 4 GiB guess there would hand out memory we have just been told does not exist. + + A missing ``bytes_in_use`` key (review MINOR, 2026-09-08) is read as unknown + occupancy, not as 0: ``stats.get(...) or 0`` could not tell "reported zero" from + "never reported", so a pool with no occupancy figure at all was read as entirely + free. ``stats.get("bytes_in_use")`` is None in both cases; only the missing-key + case must fall through to "unknown" here. + """ + block = stats.get("largest_free_block_bytes") + if block: + return max(0, int(block)) + pool = stats.get("pool_bytes") or stats.get("bytes_reserved") + if pool: + in_use = stats.get("bytes_in_use") + if in_use is None: + return None + return max(0, int(pool) - int(in_use)) + return None + + +def _probe_allocate(jax_module, dev): + """Force one tiny allocation on `dev` so the allocator's pool actually exists + before `_angle_marg_buffer_target` reads `memory_stats()`. + + Review MAJOR, 2026-09-08: `_angle_marg_buffer_target` is called before this + process's first device allocation, so on jax 0.9.2 `pool_bytes` and + `largest_free_block_bytes` both read 0 -- not because the device is busy, but + because the allocator has not been asked to reserve anything yet. `#285` made + that unread state fall through to the blind 4 GiB fallback, which is safe on an + idle card but is 8x too generous on a busy shared one (~0.5 GiB truly free, + simulated in the review). A tiny real allocation is the only way to make the + pool signal exist: under JAX's default preallocating allocator it only succeeds + because the memory it reserves genuinely was free, so a successful probe means + `pool_bytes` afterward is memory this process actually holds, not a ceiling. + + Exceptions are the caller's problem on purpose: a device with no room even for + this allocation is exactly the "we could not read this device" case the + fallback already exists for. + """ + jax_module.device_put(jax_module.numpy.zeros(1), dev).block_until_ready() + + +def _angle_marg_buffer_target(): + """Bytes to allow for the largest single anglemarg buffer. + + Derived from the device's FREE memory rather than assumed, because the constant this + replaces was sized on the smallest machine anyone had run on. Any failure to read the + device -- no jax, no GPU, an API that moved, or stats that report a ceiling but no + availability -- returns the historical 4 GiB, so a machine we cannot interrogate + behaves exactly as before rather than getting a larger number by accident. + + An explicit RIFT_ANGLEMARG_BUFFER_BYTES wins over both, and is read OUTSIDE the + try below on purpose: inside it, the blanket `except Exception` would swallow the + ValueError from a malformed override and hand back the fallback -- silently ignoring + the one number in this function a human asserted about the machine in front of them. + + Before reading stats, `_probe_allocate` forces one tiny allocation so the pool + signal exists at all on jax 0.9.2 (review MAJOR, 2026-09-08). A small pool next + to a much larger `bytes_limit` (below half of it) means the on-demand allocator + is in play, where the pool grows only to fit what has actually been requested so + far and `pool - bytes_in_use` is not a free-memory reading; that case is bounded + by `bytes_limit - bytes_in_use` instead of trusted, same fallback ceiling as an + unreadable device. + """ + explicit = _read_buffer_bytes() + if explicit is not None: + return explicit + try: + import jax + devs = [d for d in jax.devices() if getattr(d, "platform", "") == "gpu"] + if not devs: + return _ANGLE_MARG_BUFFER_TARGET_FALLBACK + dev = devs[0] + try: + _probe_allocate(jax, dev) + except Exception: + # No room even for this allocation, or a fake/incomplete jax in tests -- + # either way, still try to read whatever memory_stats() reports below + # rather than giving up immediately. + pass + stats = dev.memory_stats() or {} + limit = stats.get("bytes_limit") or stats.get("bytes_reservable_limit") + pool = stats.get("pool_bytes") or stats.get("bytes_reserved") + block = stats.get("largest_free_block_bytes") + if not block and pool and limit and int(pool) < 0.5 * int(limit): + # On-demand allocator (XLA_PYTHON_CLIENT_PREALLOCATE=false): the pool + # grows incrementally with each request instead of claiming a big + # fraction up front, so a small pool relative to the limit does not mean + # little is free -- it means pool - bytes_in_use cannot be trusted here. + # Bound the blind guess by what the device could still give this + # process rather than reach for a figure this allocator shape cannot + # support. + in_use = stats.get("bytes_in_use") + in_use = int(in_use) if in_use is not None else 0 + return max(0, min(_ANGLE_MARG_BUFFER_TARGET_FALLBACK, + int(limit) - in_use)) + avail = _device_available_bytes(stats) + if avail is None: + # We can see a device but not how much of it is free. The conservative + # fallback stands; an operator who knows their card asserts otherwise with + # RIFT_ANGLEMARG_BUFFER_BYTES. Reaching for `bytes_limit` here instead is + # the exact regression review flagged -- see _device_available_bytes. + return _ANGLE_MARG_BUFFER_TARGET_FALLBACK + # The ceiling is still worth reading, but only DOWNWARD: availability cannot + # legitimately exceed what the allocator may hold, so a runtime reporting a free + # block bigger than its own limit is misreporting and must not inflate this. + if limit: + avail = min(avail, int(limit)) + # NO max() WITH THE FALLBACK HERE. Flooring at 4 GiB would defeat the whole + # point in the one direction that matters for safety: a card with 6 GiB free + # would be handed a 4 GiB single buffer, and one with under 4 GiB free would be + # handed more than it has. That is the failure this function exists to prevent, + # wearing device awareness as a costume. A busy device gets a small allowance. + # + # AN EARLIER VERSION OF THIS COMMENT SAID a device too small for the model merely + # "goes slow, not wrong", because angle_marg_eval_chunk floored the chunk at 1. + # That was false and review caught it: a chunk of one still requests + # bytes_per * npts, so once ONE sample exceeds the allowance the floor returns a + # chunk that BREAKS the bound rather than a chunk that is slow. There is no + # kernel-level tiling of that buffer -- the sample axis is the only axis this cap + # can divide -- so angle_marg_eval_chunk now refuses instead of pretending. + # + # max(0, ...), not max(1, ...): a device with nothing free must produce an + # allowance of nothing, and let angle_marg_eval_chunk refuse with the message + # that names the knobs. A one-byte floor would be the same lie in miniature. + return max(0, int(avail * _ANGLE_MARG_BUFFER_FRACTION)) + except Exception: + return _ANGLE_MARG_BUFFER_TARGET_FALLBACK + + +#: Kept as a module attribute so existing readers (and tests) still see a number. +_ANGLE_MARG_BUFFER_TARGET = _ANGLE_MARG_BUFFER_TARGET_FALLBACK + + +def _peaklocal_bytes_per_sample_pt(like): + """Conservative source-level payload for one peak-local sample/time point. + + The streamed nonlinear body, one phi chunk's values, the phi accumulator and the + per-distance-node joint tables have distinct shapes, and all have to be budgeted. This is still not a CUDA allocator + measurement and cannot see an outer transformation such as flowMC's chain + ``vmap``; callers of the scalar AD target require separate profiling. + """ + from . import anglemarg as _am + from . import joint_anglemarg_peaklocal as _jp + + n_x = int(np.size(getattr(like, "x_grid", ())) or 1) + info = getattr(like, "angle_marg_info", None) or {} + # Production wrappers record the floored sizing amplitude. Preserve the + # same floor for small test doubles and legacy readers that omit the ledger. + amp_sizing = info.get("amp_sizing") + if amp_sizing is None: + amp_sizing = _am.ANGLE_MARG_CROSSOVER_AMPLITUDE + n_u_live = min(_jp.u_nodes_in_use(amp_sizing), _jp.U_NODE_STREAM_CHUNK) + + data = getattr(like, "data", None) + lms = getattr(data, "lms", None) + m_max = (int(np.max(np.abs(np.asarray(lms)[:, 1]))) + if lms is not None else 2) + + streamed_body = _jp.PHI_CHUNK_DEFAULT * n_x * 4 * n_u_live * 8 + # The phi scan REDUCES into its carry rather than returning a value per chunk, so + # what is live is one chunk's (phi_chunk, n_x) block and the (n_x,) accumulator -- + # not the whole phi axis. This term used to read `n_phi * n_x * 8`, the stacked + # output, and that was the real 19.97 GiB at ladder-2 rung 640. The model moves in + # the same commit as the kernel: this file's own history is a guard that kept an old + # model while the kernel's sizing moved. + phi_chunk_values = _jp.PHI_CHUNK_DEFAULT * n_x * 8 + accumulator = n_x * 8 + # `tables` in joint_lnL_phi_dense: one complex (KP, 2KS+1) per distance node, live + # for the whole call. Previously unmodelled, and small against the streamed body + # (~10% at n_x 256), but it is live and cheap to state. + joint_tables = n_x * (2 * m_max + 1) * 5 * 16 + return int(streamed_body + phi_chunk_values + accumulator + joint_tables) + + +def _philocal_bytes_per_sample_pt(like): + """Conservative source-level payload for one phi-local sample/time point. + + THE PHI-LOCAL KERNEL ROLLS TWO AXES AND THE GUARD MUST MODEL BOTH, because unlike the + dense path its quadrature grid is not streamed on the u axis: `pt_chunk` phi points + are evaluated at once, each costing a u profile of `4 * u_nodes`, and `eval_g2` + materializes `(points, KP, 2KS+1)` COMPLEX -- 16 bytes a term, not 8. `x_chunk` + distance nodes are in flight simultaneously. Miss the complex width or the table + terms and this undercounts by ~90x. + + The entry ALSO evaluates the dense peak-local scheme for the fallback, so the dense + model is added rather than maxed: both are live in the same trace. + """ + from . import joint_anglemarg_peaklocal as _jp + from . import anglemarg as _am + + info = getattr(like, "angle_marg_info", None) or {} + amp_sizing = info.get("amp_sizing") + if amp_sizing is None: + amp_sizing = _am.ANGLE_MARG_CROSSOVER_AMPLITUDE + u_nodes = _jp.u_nodes_in_use(amp_sizing) + + data = getattr(like, "data", None) + lms = getattr(data, "lms", None) + m_max = (int(np.max(np.abs(np.asarray(lms)[:, 1]))) + if lms is not None else 2) + KP = 2 * m_max + 1 + terms = KP * 5 # (KP, 2KS+1) with the u degree pinned at 2 + + live_pts = _jp.PT_CHUNK_DEFAULT * 4 * u_nodes + body = _jp.X_CHUNK_DEFAULT * live_pts * terms * 16 + # the per-node values the distance scan stacks before its reduction + n_x = int(np.size(getattr(like, "x_grid", ())) or 1) + stacked = n_x * 8 + return int(body + stacked + _peaklocal_bytes_per_sample_pt(like)) def angle_marg_eval_chunk(like, chunk): @@ -249,7 +640,7 @@ def angle_marg_eval_chunk(like, chunk): 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 + pattern as the _default_eval_chunk() shrink above. Grid-scheme and 4/5-param likelihoods pass through unchanged. """ # NOT the scheme default. "grid" here is a SENTINEL meaning "this object @@ -261,42 +652,85 @@ def angle_marg_eval_chunk(like, chunk): # (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. + # phi chunks, four cells and a streamed u-node block, so the batch multiplies the + # same way the dense schemes do. 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"): + "peak-local", "phi-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)) + if getattr(like, "angle_marg_scheme", None) == "phi-local": + # Modelled in the SAME change that added the kernel, because the trap this guard + # exists for is a kernel whose sizing moved while the guard kept its old model. + bytes_per = max(bytes_per, _philocal_bytes_per_sample_pt(like)) + elif getattr(like, "angle_marg_scheme", None) == "peak-local": + # The streamed (phi_chunk,n_x,4,u_live) body, one chunk's (phi_chunk,n_x) + # values, the (n_x,) phi accumulator and the per-distance-node joint tables. + # None of these grows with n_phi now that the scan reduces into its carry. + bytes_per = max(bytes_per, _peaklocal_bytes_per_sample_pt(like)) + target = _angle_marg_buffer_target() + per_sample = bytes_per * npts + if per_sample > target: + # FAIL CLOSED. This branch used to be `cap = max(1, target // per_sample)`, + # which returns 1 here and therefore hands back a chunk whose buffer is + # `per_sample` bytes -- larger than the target this function exists to enforce. + # The floor made the bound silently untrue on any device small enough, which is + # not the same failure as being slow. peak-local reaches it at production + # dimensions: phi_chunk 16, n_x 256, four cells, an 8-node stream block, the + # stacked phi scan and npts 1230 is 2.03 GiB for ONE sample, so a 2 GiB card + # (1 GiB allowance at the + # default fraction) cannot honour the bound at any chunk size. + # + # The alternative repair is kernel-level tiling of the buffer itself. That is a + # real option and a much larger change; until someone does it, the honest thing + # is to say the bound cannot be met rather than to report a chunk that breaks it. + # + # MemoryError, matching RIFT.likelihood.time_posterior's + # validate_time_posterior_working_set: same shape (a preflight refusal of a + # dense working set, with the estimate, the dimensions, the limit and the knobs + # in the message), so it gets the same type. It also lets a caller that wants + # to fall back to a cheaper scheme catch this narrowly instead of every + # RuntimeError the eval path can raise. + raise MemoryError( + "angle-marginalization resource preflight: scheme %r cannot honour the " + "buffer bound at ANY chunk size: one " + "sample needs %d bytes (%.2f GiB) -- %d bytes per sample per time point x " + "npts=%d -- against an allowance of %d bytes (%.2f GiB). Returning a chunk " + "of 1 would ask the device for the full %.2f GiB and OOM, so this refuses " + "instead. Act on one of: raise the allowance with " + "RIFT_ANGLEMARG_BUFFER_FRACTION (a fraction, at most 1.0, of the memory the " + "device reports FREE -- it has no effect when that could not be read, and " + "note that the free figure moves with whoever else is on the card) or " + "RIFT_ANGLEMARG_BUFFER_BYTES (an absolute byte allowance, which wins over " + "both the device probe and the %d-byte fallback); shorten the time window " + "(npts); shrink the distance grid (n_x), which drives the peak-local model; " + "or run a cheaper angle_marg_scheme. The sample axis is the only axis this " + "cap can divide, so no chunk size is a fix; reducing the outer " + "evaluation chunk cannot make this call fit." + % (getattr(like, "angle_marg_scheme", None), per_sample, + per_sample / float(1 << 30), bytes_per, npts, target, + target / float(1 << 30), per_sample / float(1 << 30), + _ANGLE_MARG_BUFFER_TARGET_FALLBACK)) + # No max(..., 1) here, deliberately: the refusal above is what guarantees + # `per_sample <= target`, so the floor division is already at least 1. Restoring the + # floor would restore the defect -- it is the floor, not the division, that broke the + # bound. And a floor LARGER than one breaks it in the other direction for long but + # valid time windows (npts=65537 with a floor of 64 requested ~32 GiB). + cap = target // per_sample 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): +def eval_lnL(like, theta, chunk=None): """Evaluate the distance-marginalized lnL on an ``(N, 5)`` array in chunks. Chunking bounds peak device memory (the distance grid multiplies the batch dimension inside the likelihood). """ + if chunk is None: + chunk = _default_eval_chunk() theta = np.atleast_2d(theta) chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] @@ -445,6 +879,35 @@ def _moment_match(theta, logL): return mu, cov +def regularize_cov(cov, rel=1e-12): + """The covariance a Gaussian proposal must use for BOTH its Cholesky draw + and its density. + + Two properties, and each one alone was a live defect (issue #227): + + RELATIVE, not absolute. An ``+ eps*I`` regularizer with a fixed ``eps`` + is only negligible if the covariance is O(1). A production extrinsic + posterior is not: on the real S250114ax point in #227 (rho ~ 49) the driver's + own Fisher gives angular scales ~1e-3 rad, and an ADAPTING proposal contracts + far below that -- to 3e-21 there. ``1e-12`` then stops being a conditioning + nudge and becomes the proposal. Scaling by ``trace(cov)/dim`` makes the nudge + a fixed fraction of the covariance at every scale. + + ONE matrix. Callers must pass this return value to the Cholesky *and* to + ``_gaussian_logq``/``_mixture_logq``. Drawing from ``cov + eps*I`` while + scoring under bare ``cov`` computes importance weights against a + distribution that was never sampled; with ``cov ~ 3e-21`` and ``eps=1e-12`` + the Mahalanobis term is ``(1e-6/sqrt(3e-21))**2 ~ 3e8`` per dimension, which + is how ``--mode laplace-is`` returned ``lnZ = 5.8e9`` and exited 0. + """ + cov = np.asarray(cov, dtype=float) + d = cov.shape[-1] + scale = float(np.trace(cov)) / d + if not np.isfinite(scale) or scale <= 0.0: + scale = 1.0 # degenerate/zero covariance: fall back to absolute + return cov + (rel * scale) * np.eye(d) + + def _finalize_evidence(logZ, sigma_over_Z, neff, max_lnL): """Flag an importance-evidence estimate as unreliable (nan) when it cannot be trusted: log Z must satisfy ``log Z <= lnL_max`` for a normalized prior @@ -777,6 +1240,8 @@ def extract(s): mus, covs = [mu], [cov * proposal_inflate] n_comp = len(mus) weights = np.full(n_comp, 1.0 / n_comp) + # ONE matrix per component for both the draw and _mixture_logq below (#227). + covs = [regularize_cov(cv) for cv in covs] # draw from the mixture counts = rng.multinomial(n_is, weights) @@ -784,7 +1249,7 @@ def extract(s): for c in range(n_comp): if counts[c] == 0: continue - Lc = np.linalg.cholesky(covs[c] + 1e-12 * np.eye(5)) + Lc = np.linalg.cholesky(covs[c]) z = rng.standard_normal((counts[c], 5)) draws.append(mus[c][None, :] + z @ Lc.T) th_is = np.concatenate(draws, axis=0) @@ -946,8 +1411,8 @@ def logpdf(theta5, data): logZ = sigma_over_Z = neff = np.nan if len(theta) >= 6: mu, cov = _moment_match(theta, np.zeros(len(theta))) - cov = cov * 2.0 - Lc = np.linalg.cholesky(cov + 1e-12 * np.eye(n_dim)) + cov = regularize_cov(cov * 2.0) # ONE matrix: draw and density (#227) + Lc = np.linalg.cholesky(cov) n_is = 40000 z = rng.standard_normal((n_is, n_dim)) th_is = mu[None, :] + z @ Lc.T @@ -1015,8 +1480,10 @@ def _log_prior_4_jax(theta4): return jnp.where(inb, logp, -1e30) -def eval_lnL_4(like, theta, chunk=_EVAL_CHUNK, desc="lnL"): +def eval_lnL_4(like, theta, chunk=None, desc="lnL"): """Evaluate the 4-param (phi-marginalised) lnL on an ``(N, 4)`` array.""" + if chunk is None: + chunk = _default_eval_chunk() theta = np.atleast_2d(theta) chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] @@ -1112,8 +1579,10 @@ def _log_prior_3_jax(theta3): return jnp.where(inb, logp, -1e30) -def eval_lnL_3(like, theta, chunk=_EVAL_CHUNK, desc="lnL"): +def eval_lnL_3(like, theta, chunk=None, desc="lnL"): """Evaluate the 3-param (phi+psi-marginalised) lnL on an ``(N, 3)`` array.""" + if chunk is None: + chunk = _default_eval_chunk() theta = np.atleast_2d(theta) chunk = angle_marg_eval_chunk(like, chunk) N = theta.shape[0] @@ -1594,8 +2063,8 @@ def _ess(next_invT): print(" [evidence] Laplace diag failed: %r" % e) elif len(theta) >= 6: mu, cov = _moment_match(theta, np.zeros(len(theta))) - cov = cov * 2.0 - Lc = np.linalg.cholesky(cov + 1e-12 * np.eye(n_dim)) + cov = regularize_cov(cov * 2.0) # ONE matrix: draw and density (#227) + Lc = np.linalg.cholesky(cov) n_is = 40000 z = rng.standard_normal((n_is, n_dim)) th_is = mu[None, :] + z @ Lc.T @@ -1628,8 +2097,10 @@ def _ess(next_invT): if mapT is not None: cov_is = (float(fisher_is_inflate) ** 2) * (A_is @ A_is.T) cov_is = 0.5 * (cov_is + cov_is.T) + # ONE matrix: this is what _gaussian_logq is given below (#227). + cov_is = regularize_cov(cov_is) try: - Lc = np.linalg.cholesky(cov_is + 1e-12 * np.eye(n_dim)) + Lc = np.linalg.cholesky(cov_is) N = int(fisher_is_samples) z = rng.standard_normal((N, n_dim)) th_is = mapT[None, :] + z @ Lc.T @@ -1876,6 +2347,415 @@ def _ess(db): lnL_map=float(np.max(lnL)) if len(lnL) else np.nan) +# --------------------------------------------------------------------------- +# 2c. Adaptive-volume / portfolio backends (value-only JAX likelihood) +# --------------------------------------------------------------------------- + +_FULL_EXTRINSIC_ORDER = ("ra", "dec", "psi", "incl", "phiref", "distMpc") + + +def _av_param_order(like): + """Parameter order exposed by a JAX likelihood, including the bare 6-D case.""" + order = tuple(getattr(like, "ANGULAR_PARAM_ORDER", ())) + if order: + return order + return _FULL_EXTRINSIC_ORDER + + +def _av_sample_bounds(order, d_min, d_max, sample_d_min=None, + sample_d_max=None, sample_bounds=None): + """Validate and resolve AV sampling limits without renormalizing the prior.""" + requested = dict(sample_bounds or {}) + if sample_d_min is not None or sample_d_max is not None: + requested["distMpc"] = ( + d_min if sample_d_min is None else sample_d_min, + d_max if sample_d_max is None else sample_d_max) + unknown = set(requested) - set(order) + if unknown: + raise ValueError("sampling bounds name coordinates absent from likelihood: %s" + % ", ".join(sorted(unknown))) + resolved = {} + for name in order: + physical_lo, physical_hi, _ = _av_prior_spec(name, d_min, d_max) + lo, hi = requested.get(name, (physical_lo, physical_hi)) + lo, hi = float(lo), float(hi) + if not np.isfinite(lo) or not np.isfinite(hi) or lo >= hi: + raise ValueError("invalid sampling bounds for %s: (%r, %r)" + % (name, lo, hi)) + if lo < physical_lo or hi > physical_hi: + raise ValueError("sampling bounds for %s must lie within [%g, %g]" + % (name, physical_lo, physical_hi)) + resolved[name] = (lo, hi) + return resolved + + +def _av_prior_draw(order, n, rng, d_min, d_max, sample_bounds=None): + """Draw the physical prior conditioned only on the AV sampling window.""" + bounds = _av_sample_bounds(order, d_min, d_max, + sample_bounds=sample_bounds) + def interval(name): + return bounds.get(name, _av_prior_spec(name, d_min, d_max)[:2]) + ra_lo, ra_hi = interval("ra") if "ra" in order else (0.0, _TWO_PI) + dec_lo, dec_hi = interval("dec") if "dec" in order else (-_PI / 2, _PI / 2) + psi_lo, psi_hi = interval("psi") if "psi" in order else (0.0, _PI) + inc_lo, inc_hi = interval("incl") if "incl" in order else (0.0, _PI) + phase_name = "phiref_shifted" if "phiref_shifted" in order else "phiref" + phase_lo, phase_hi = interval(phase_name) if phase_name in order else (0.0, _TWO_PI) + pp_lo, pp_hi = interval("phase_p") if "phase_p" in order else (0.0, 2.0 * _TWO_PI) + pm_lo, pm_hi = interval("phase_m") if "phase_m" in order else (0.0, 2.0 * _TWO_PI) + dist_lo, dist_hi = interval("distMpc") if "distMpc" in order else (d_min, d_max) + draws = { + "ra": rng.uniform(ra_lo, ra_hi, n), + "dec": np.arcsin(rng.uniform(np.sin(dec_lo), np.sin(dec_hi), n)), + "psi": rng.uniform(psi_lo, psi_hi, n), + "incl": np.arccos(rng.uniform(np.cos(inc_hi), np.cos(inc_lo), n)), + "phiref": rng.uniform(phase_lo, phase_hi, n), + "phiref_shifted": rng.uniform(phase_lo, phase_hi, n), + "phase_p": rng.uniform(pp_lo, pp_hi, n), + "phase_m": rng.uniform(pm_lo, pm_hi, n), + "distMpc": np.cbrt(rng.uniform(dist_lo ** 3, dist_hi ** 3, n)), + } + return np.column_stack([draws[name] for name in order]) + + +def _av_prior_spec(name, d_min, d_max, sample_d_min=None, sample_d_max=None, + sample_bounds=None): + """Return ``(lo, hi, physical_density)`` for one wrapper coordinate.""" + if name == "ra": + spec = (0.0, _TWO_PI, lambda x: np.ones(np.shape(x)) / _TWO_PI) + elif name == "dec": + spec = (-_PI / 2, _PI / 2, + lambda x: 0.5 * np.maximum(np.cos(np.asarray(x)), 0.0)) + elif name == "psi": + spec = (0.0, _PI, lambda x: np.ones(np.shape(x)) / _PI) + elif name == "incl": + spec = (0.0, _PI, + lambda x: 0.5 * np.maximum(np.sin(np.asarray(x)), 0.0)) + elif name in ("phiref", "phiref_shifted"): + spec = (0.0, _TWO_PI, lambda x: np.ones(np.shape(x)) / _TWO_PI) + elif name in ("phase_p", "phase_m"): + spec = (0.0, 2.0 * _TWO_PI, + lambda x: np.ones(np.shape(x)) / (2.0 * _TWO_PI)) + elif name == "distMpc": + norm = 3.0 / (float(d_max) ** 3 - float(d_min) ** 3) + spec = (float(d_min if sample_d_min is None else sample_d_min), + float(d_max if sample_d_max is None else sample_d_max), + lambda x, _norm=norm: _norm * np.asarray(x) ** 2) + else: + raise ValueError("unsupported JAX-ILE adaptive-volume parameter %r" % (name,)) + if sample_bounds and name in sample_bounds: + return float(sample_bounds[name][0]), float(sample_bounds[name][1]), spec[2] + return spec + + +def _fixed_shape_value_callback(like, n_dim, eval_chunk): + """Build an AV callback that compiles only one JAX batch shape. + + AV's number of occupied bins changes as it contracts, so its raw draw length + changes slightly from cycle to cycle. Passing that length straight to a jitted + likelihood recompiles the likelihood every cycle. Split into fixed blocks and + pad only the last block; padding changes neither returned rows nor the integral. + """ + eval_chunk = int(angle_marg_eval_chunk(like, max(1, int(eval_chunk)))) + + def lnL(*cols): + # AV can be configured with cupy even though the JAX likelihood owns the + # accelerator. ``np.asarray(cupy_array)`` intentionally raises; ``get`` + # is the explicit host handoff needed before JAX transfers each fixed + # block to its device. This also covers portfolio's selfish AV update, + # which does not have integrate_log's host-fallback wrapper. + host_cols = [(c.get() if hasattr(c, "get") else np.asarray(c)) + for c in cols] + x = np.column_stack([np.asarray(c, dtype=float).reshape(-1) + for c in host_cols]) + if x.shape[1] != n_dim: + raise ValueError("JAX-AV callback expected %d columns, got %d" + % (n_dim, x.shape[1])) + out = np.empty(len(x), dtype=float) + for start in range(0, len(x), eval_chunk): + stop = min(start + eval_chunk, len(x)) + block = x[start:stop] + if len(block) < eval_chunk: + block = np.concatenate( + [block, np.repeat(block[-1:], eval_chunk - len(block), axis=0)], + axis=0) + val = like.log_likelihood(*[block[:, j] for j in range(n_dim)]) + out[start:stop] = np.asarray(val, dtype=float)[:stop - start] + return out + + lnL.eval_chunk = eval_chunk + return lnL + + +def _sky_distance(a, b): + return float(np.arccos(np.clip( + np.sin(a[1]) * np.sin(b[1]) + + np.cos(a[1]) * np.cos(b[1]) * np.cos(a[0] - b[0]), -1.0, 1.0))) + + +def _fisher_sky_seed(like, order, lnL, rng, d_min, d_max, n_seed, + n_pilot, n_modes, sky_inflate, prior_frac, + initial_points=None, sample_bounds=None, verbose=False): + """Hill-climb modes, then draw Fisher sky / physical-prior other coordinates. + + This is deliberately a proposal initializer, not part of the estimator. The + likelihood calls in the AV integration remain value-only. A fraction of the + cloud is left as a full-prior draw; in portfolio mode the GMM's defensive + component supplies the actual full-support guarantee (a finite point cloud + alone cannot provide one). + """ + from scipy.optimize import minimize + + n_pilot = max(int(n_pilot), int(n_modes), 1) + pilot = _av_prior_draw(order, n_pilot, rng, d_min, d_max, sample_bounds) + pilot_lnL = lnL(*pilot.T) + ranked = np.argsort(np.where(np.isfinite(pilot_lnL), pilot_lnL, -np.inf))[::-1] + seeds = [] + if initial_points is not None: + supplied = np.atleast_2d(np.asarray(initial_points, dtype=float)) + if supplied.shape[1] != len(order): + raise ValueError("initial_points must have %d columns, got %d" + % (len(order), supplied.shape[1])) + seeds.extend(supplied) + for idx in ranked: + if len(seeds) >= int(n_modes): + break + candidate = pilot[idx] + if not seeds or all(_sky_distance(candidate, old) >= 0.25 for old in seeds): + seeds.append(candidate) + if len(seeds) >= int(n_modes): + break + if not seeds: + raise RuntimeError("the Fisher-sky prior pilot found no finite likelihood") + + bounds = [_av_prior_spec(name, d_min, d_max, + sample_bounds=sample_bounds)[:2] for name in order] + # Avoid evaluating the Jacobian-singular orientation endpoints during AD. + bounds = [(lo + 1e-6 if name in ("dec", "incl") else lo, + hi - 1e-6 if name in ("dec", "incl") else hi) + for name, (lo, hi) in zip(order, bounds)] + + def objective(x): + value, grad = like.value_and_grad(x) + return -float(value), -np.asarray(grad, dtype=float) + + modes = [] + for seed_theta in seeds: + try: + result = minimize(objective, seed_theta, jac=True, method="L-BFGS-B", + bounds=bounds, options={"maxiter": 120}) + theta = np.asarray(result.x if np.isfinite(result.fun) else seed_theta) + value = -float(result.fun) if np.isfinite(result.fun) else float( + lnL(*seed_theta[:, None])[0]) + fisher = np.asarray(like.fisher(theta), dtype=float) + fisher = 0.5 * (fisher + fisher.T) + # Marginalize over the nuisance angles. Inverting only F[:2,:2] + # would give the conditional sky covariance and is too narrow when + # sky is correlated with polarization, inclination, or phase. + full_cov = np.linalg.pinv(fisher, rcond=1e-12) + sky_cov = 0.5 * (full_cov[:2, :2] + full_cov[:2, :2].T) + eig, vec = np.linalg.eigh(sky_cov) + # Cap the one-sigma sky width at one radian: weak curvature should + # remain broad, while a sharp high-SNR sky peak gets its 1/rho scale. + floor = max(float(np.max(eig)) * 1e-10, 1e-16) + var = float(sky_inflate) ** 2 * np.clip(eig, floor, None) + cov = (vec * np.minimum(var, 1.0)) @ vec.T + modes.append((theta, cov, value)) + except Exception as exc: # one failed mode must not discard the good ones + if verbose: + print(" [JAX-AV seed] hill climb skipped: %s" % exc) + if not modes: + raise RuntimeError("all Fisher-sky hill climbs failed") + modes.sort(key=lambda item: item[2], reverse=True) + + n_seed = max(int(n_seed), len(order) + 2) + n_prior = int(np.clip(float(prior_frac), 0.0, 1.0) * n_seed) + n_focus = n_seed - n_prior + focused = _av_prior_draw(order, n_focus, rng, d_min, d_max, sample_bounds) + counts = np.full(len(modes), n_focus // len(modes), dtype=int) + counts[:n_focus % len(modes)] += 1 + cursor = 0 + for (mode, cov, _), count in zip(modes, counts): + if count == 0: + continue + sky = rng.multivariate_normal(mode[:2], cov, size=count) + ra_lo, ra_hi = bounds[0] + dec_lo, dec_hi = bounds[1] + sky[:, 0] = np.mod(sky[:, 0], _TWO_PI) + # A restricted, non-wrapping RA window cannot use periodic wrap as a + # boundary condition. Clip the seed proposal to the declared window; + # integration weights remain governed by the physical prior. + focused[cursor:cursor + count, 0] = np.clip( + sky[:, 0], ra_lo + 1e-9, ra_hi - 1e-9) + focused[cursor:cursor + count, 1] = np.clip( + sky[:, 1], dec_lo + 1e-9, dec_hi - 1e-9) + cursor += count + cloud = focused + if n_prior: + cloud = np.vstack([cloud, _av_prior_draw( + order, n_prior, rng, d_min, d_max, sample_bounds)]) + if verbose: + print(" [JAX-AV seed] %d hill-climbed sky mode(s), %d seed points " + "(%d full-prior)" % (len(modes), len(cloud), n_prior)) + return cloud, np.asarray([m[0] for m in modes]), np.asarray([m[2] for m in modes]) + + +def adaptive_volume_sample(like, d_min, d_max, sampler_method="AV", + portfolio_members=("AV", "GMM"), nmax=300000, + neff=1000, n_chunk=8000, eval_chunk=None, seed=0, + seed_method="none", seed_pilot=4000, seed_modes=4, + seed_points=None, seed_initial_points=None, + initial_samples=None, + sky_inflate=2.0, + seed_prior_frac=0.1, anisotropic_bins=True, + gmm_components=2, + verbose=False, sample_d_min=None, sample_d_max=None, + sample_bounds=None): + """Run production AV/portfolio control logic on a value-only JAX likelihood. + + ``sampler_method`` is ``AV`` or ``portfolio``. The optional ``fisher-sky`` + seed pays a small, explicit AD startup cost (multi-start hill climb plus local + Hessians); all integration calls use only ``like.log_likelihood``. Portfolio + defaults to AV+GMM so its defensive mixture retains full support even when the + seeded AV live volume covers only selected sky modes. + """ + from RIFT.integrators import mcsamplerAdaptiveVolume as AV + + method = str(sampler_method) + if method not in ("AV", "portfolio"): + raise ValueError("sampler_method must be 'AV' or 'portfolio', got %r" % method) + order = _av_param_order(like) + n_dim = len(order) + resolved_bounds = _av_sample_bounds( + order, d_min, d_max, sample_d_min, sample_d_max, sample_bounds) + # Decouple AV's coverage cloud from the accelerator batch. AV often needs + # a large n_chunk to hit a narrow sky mode, while the marginalized JAX + # kernel has a much smaller memory-efficient batch. The callback loops over + # this fixed shape, so increasing coverage does not increase device memory. + eval_chunk = int(eval_chunk or min(int(n_chunk), _default_eval_chunk())) + lnL = _fixed_shape_value_callback(like, n_dim, eval_chunk) + + av_member = AV.MCSampler(n_chunk=int(n_chunk)) + if method == "AV": + sampler = av_member + else: + from RIFT.integrators import mcsamplerEnsemble as GMM + from RIFT.integrators import mcsamplerPortfolio as Portfolio + members = [] + member_setup_args = [] + for name in portfolio_members: + key = str(name).strip().upper() + if key == "AV": + members.append(av_member if not any(m is av_member for m in members) + else AV.MCSampler(n_chunk=int(n_chunk))) + member_setup_args.append({}) + elif key == "GMM": + members.append(GMM.MCSampler()) + # At least two components are required whenever a narrow mode + # crosses a periodic box boundary (Event B has phiref=0). One + # Euclidean Gaussian spans the whole [0,2pi] box and destroys + # the sky/phase correlations in an otherwise excellent seed. + member_setup_args.append({"n_comp": max(2, int(gmm_components))}) + else: + raise ValueError("JAX portfolio member %r is unsupported; use AV or GMM" + % name) + if not members: + raise ValueError("JAX portfolio needs at least one member") + sampler = Portfolio.MCSampler(portfolio=members, n_chunk=int(n_chunk)) + + for name in order: + lo, hi = resolved_bounds[name] + prior = _av_prior_spec(name, d_min, d_max)[2] + sampler.add_parameter(name, pdf=None, left_limit=lo, right_limit=hi, + prior_pdf=prior, adaptive_sampling=True) + setup_kwargs = {"anisotropic_bins": bool(anisotropic_bins)} + if method == "portfolio": + setup_kwargs["portfolio_args"] = member_setup_args + sampler.setup(**setup_kwargs) + + seed_cloud = seed_modes_theta = seed_modes_lnL = None + if initial_samples is not None: + if seed_method not in (None, "none"): + raise ValueError("initial_samples and seed_method are mutually exclusive") + seed_cloud = np.atleast_2d(np.asarray(initial_samples, dtype=float)) + if seed_cloud.shape[1] != n_dim: + raise ValueError("initial_samples must have %d columns, got %d" + % (n_dim, seed_cloud.shape[1])) + for j, name in enumerate(order): + lo, hi = resolved_bounds[name] + if np.any(seed_cloud[:, j] < lo) or np.any(seed_cloud[:, j] > hi): + raise ValueError("initial_samples fall outside sampling bounds for %s" + % name) + sampler.bootstrap_from_samples(seed_cloud, params=order, seed=seed) + if verbose: + print(" [JAX-AV seed] bootstrapped from %d caller/oracle samples" + % len(seed_cloud)) + elif seed_method not in (None, "none"): + if seed_method != "fisher-sky": + raise ValueError("unknown JAX-AV seed method %r" % seed_method) + seed_cloud, seed_modes_theta, seed_modes_lnL = _fisher_sky_seed( + like, order, lnL, np.random.default_rng(seed), d_min, d_max, + n_seed=(seed_points or n_chunk), n_pilot=seed_pilot, + n_modes=seed_modes, sky_inflate=sky_inflate, + prior_frac=seed_prior_frac, initial_points=seed_initial_points, + sample_bounds=resolved_bounds, verbose=verbose) + if method == "portfolio": + sampler.bootstrap_from_samples(seed_cloud, params=order, seed=seed) + else: + print(" [JAX-AV seed] WARNING: standalone seeded AV has compact " + "support; the prior seed fraction is a diagnostic safety net, " + "not a full-support guarantee. Prefer --sampler-method portfolio.") + sampler.bootstrap_from_samples(seed_cloud, params=order, seed=seed) + + # The production integrators use numpy's legacy module RNG internally. + # Isolate that stream so --seed controls this run without perturbing a + # caller's RNG (notably later events in the same JAX-ILE batch). + numpy_rng_state = np.random.get_state() + np.random.seed(int(seed)) + try: + result = sampler.integrate_log( + lnL, *order, nmax=int(nmax), neff=float(neff), n=int(n_chunk), + no_protect_names=True, verbose=bool(verbose), save_intg=True, + tempering_exp=1.0, anisotropic_bins=bool(anisotropic_bins), + # Standalone AV can keep device-typed internal arrays when cupy is + # importable even though this adapter evaluates on the host. Its + # legacy in-integrator fair draw mixes those backends. Return the + # retained weighted population instead; callers have the exact + # log_weight below and can resample without changing the integral. + igrand_fairdraw_samples=(method != "AV"), + igrand_fairdraw_samples_max=max(int(1.5 * float(neff)), 1)) + finally: + np.random.set_state(numpy_rng_state) + logZ, log_var, eff_samp, diagnostics = result + if logZ is None: + raise RuntimeError("JAX-%s terminated without an evidence estimate" % method) + if method == "AV" and diagnostics.get("live_volume_collapsed", False): + raise RuntimeError("JAX-AV live-volume collapse: %s" % + diagnostics.get("collapse_reason", "unspecified")) + + theta = np.column_stack([np.asarray(sampler._rvs[name], dtype=float) + for name in order]) + out_lnL = np.asarray(sampler._rvs["log_integrand"], dtype=float) + already_fair = bool(getattr(sampler, "_rvs_is_fairdraw", False)) + log_weight = None + if not already_fair: + log_weight = (out_lnL + + np.asarray(sampler._rvs["log_joint_prior"], dtype=float) + - np.asarray(sampler._rvs["log_joint_s_prior"], dtype=float)) + sigma_over_Z = float(np.exp(0.5 * float(log_var) - float(logZ))) + peak = float(np.max(out_lnL)) if len(out_lnL) else np.nan + logZ, sigma_over_Z, eff_samp = _finalize_evidence( + float(logZ), sigma_over_Z, float(eff_samp), peak) + return dict(theta=theta, lnL=out_lnL, logZ=logZ, + sigma_over_Z=sigma_over_Z, neff=eff_samp, + n_eval=int(getattr(sampler, "ntotal", nmax)), + log_weight=log_weight, sampler=sampler, + diagnostics=diagnostics, eval_chunk=lnL.eval_chunk, + seed_cloud=seed_cloud, seed_modes=seed_modes_theta, + seed_mode_lnL=seed_modes_lnL, + sample_bounds=resolved_bounds) + + # --------------------------------------------------------------------------- # 3. Fisher-preconditioned importance sampling (high-SNR) # --------------------------------------------------------------------------- @@ -1969,11 +2849,15 @@ def fisher_is_sample(like, n_samples=20000, n_starts=16, n_prior_pilot=20000, var = inflate / np.clip(w, inflate / max_std ** 2, None) # cap variance cov = (V * var) @ V.T cov = 0.5 * (cov + cov.T) - Lc = np.linalg.cholesky(cov + 1e-12 * np.eye(5) * np.trace(cov) / 5) + # Already relative before #227 -- routed through the helper so that the + # draw and the density are literally the same object at every site, and so + # the grep for the defect pattern returns nothing. + cov_q = regularize_cov(cov) + Lc = np.linalg.cholesky(cov_q) z = rng.standard_normal((n_samples, 5)) theta = _wrap_angles(th0[None, :] + z @ Lc.T) - logq = _gaussian_logq(th0[None, :] + z @ Lc.T, th0, cov) # q on the raw draw + logq = _gaussian_logq(th0[None, :] + z @ Lc.T, th0, cov_q) # q on the raw draw logp = log_prior(theta) valid = np.isfinite(logp) lnL = np.full(n_samples, -np.inf) @@ -2225,13 +3109,15 @@ def model(): mu, cov = _moment_match(theta, np.zeros(len(theta))) mus, covs = [mu], [cov * 2.0] weights = np.full(len(mus), 1.0 / len(mus)) + # ONE matrix per component for both the draw and _mixture_logq below (#227). + covs = [regularize_cov(cv) for cv in covs] counts = rng.multinomial(n_is, weights) draws, comp_of_draw = [], [] for c in range(len(mus)): if counts[c] == 0: continue - Lc = np.linalg.cholesky(covs[c] + 1e-12 * np.eye(4)) + Lc = np.linalg.cholesky(covs[c]) z = rng.standard_normal((counts[c], 4)) draws.append(mus[c][None, :] + z @ Lc.T) comp_of_draw.append(np.full(counts[c], c)) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/time_first_peaklocal.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/time_first_peaklocal.py new file mode 100644 index 000000000..d7807fec9 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/time_first_peaklocal.py @@ -0,0 +1,498 @@ +"""Time-first peak-local marginalization of band-limited JAX primitives. + +This module is the deliberately small composition seam missing from the JAX +likelihood. A caller supplies one *primitive correlation* row for every fixed +state of the axes that will subsequently be marginalized (distance, angle, or +their Cartesian product). The rows are reconstructed in time before the +nonlinear log-sum-exp over those axes is formed. There is intentionally no API +that accepts a sampled, already-marginalized ``lnL(t)``: that object is not +band-limited and interpolating it is mathematically the wrong operation. + +The implementation is a fixed-shape prototype rather than production wiring. +It provides the two pieces needed to make that wiring safe: + +* :func:`plan_time_cover` builds a finite cell cover and an omitted-mass bound + from reconstructed primitive values plus a true spectral derivative bound; +* :func:`time_first_peak_local_marginalize` evaluates the nonlinear downstream + marginal only at nodes in that cover and returns ``(value, ok, ledger)``. + +``ok`` owns no fallback policy. A production time adapter should fail closed +to the existing dense primitive reconstruction when it is false. Keeping that +choice at the call site follows ``DESIGN_peak_local_framework.md`` and prevents +one axis's policy from leaking into another. + +Scope +----- +The model norm must be time-independent. ``kappa_t`` has shape +``(n_lanes, n_support)`` and ``rho_sq`` has shape ``(n_lanes,)``; the latter +shape makes the precondition explicit. A lane is a fixed downstream +quadrature state with exponent + + q_l(t) = Re[kappa_l(t)] - rho_sq_l / 2. + +The marginal integrand is ``sum_l exp(log_weight_l + q_l(t))``. Consequently +one lane can represent a distance node, an angle node, or one point of their +product. :func:`time_first_distance_peak_local_marginalize` is a convenience +adapter for the RIFT distance form ``x Re(kappa_unit) - x^2 rho_unit^2 / 2``. + +The current reconstruction topology matches the existing JAX terminal path: +an endpoint-nonduplicating even extension, with optional raised-cosine support +guards. Guard convergence is not certified here; production wiring must apply +the same two-guard comparison as ``core._time_marginalize_reflected_primitive``. + +Why the cover bound is valid +---------------------------- +For the finite Fourier series defining each reconstructed primitive, + + |kappa'_l(t)| <= M1_l = sum_k |K_lk| |omega_k|. + +On an enumeration cell of width ``h``, either endpoint therefore bounds the +whole lane by ``q_l(endpoint) + M1_l h``. Taking the smaller of the two +endpoint-derived log-sum-exp bounds gives a true upper bound on the downstream +marginal over that cell. The omitted integral is then bounded by the sum of +``h * exp(cell_upper)`` over cells outside the cover. The first-derivative +bound is intentionally conservative; a future production adapter can replace +it with the shared Hermite/M4 certificate without changing the plan contract. +""" + +from typing import NamedTuple + +import jax +import jax.numpy as jnp + +from .core import _upsample_bandlimited + + +__all__ = [ + "TimeCoverPlan", + "reconstruct_time_primitive", + "evaluate_time_primitive_points", + "spectral_time_derivative_bound", + "plan_time_cover", + "time_first_peak_local_marginalize", + "time_first_distance_peak_local_marginalize", +] + + +class TimeCoverPlan(NamedTuple): + """Fixed-shape result of the time-axis planner. + + ``live_cells`` identifies complete enumeration cells included in the local + quadrature. ``cell_log_upper`` is a certified supremum bound for every + cell, not a sampled maximum. ``outside_log_bound`` bounds the integral over + all cells not in the cover. ``peak_lower`` is the largest reconstructed + nodal value and is used only for targeting; correctness does not depend on + it being the continuous maximum. + """ + + live_cells: jax.Array + cell_log_upper: jax.Array + outside_log_bound: jax.Array + peak_lower: jax.Array + enum_step: jax.Array + + +def _validate_primitive_shapes(kappa_t, rho_sq, log_lane_weight, guard): + if kappa_t.ndim != 2: + raise ValueError( + "kappa_t must have shape (n_lanes, n_support); an already-" + "marginalized lnL(t) is deliberately not accepted") + if rho_sq.ndim != 1 or rho_sq.shape[0] != kappa_t.shape[0]: + raise ValueError( + "rho_sq must have shape (n_lanes,), making the time-independent " + "norm precondition explicit") + if log_lane_weight.ndim != 1 or log_lane_weight.shape[0] != kappa_t.shape[0]: + raise ValueError("log_lane_weight must have shape (n_lanes,)") + if kappa_t.shape[-1] - 2 * guard < 2: + raise ValueError("guard must leave at least two integration samples") + + +def _tapered_support(kappa_t, guard): + """Move the artificial reflection seam through support-only tapering.""" + guard = int(guard) + if guard == 0: + return kappa_t + n_keep = kappa_t.shape[-1] - 2 * guard + u = jnp.arange(guard + 1, dtype=jnp.float64) / float(guard) + ramp = 0.5 * (1.0 - jnp.cos(jnp.pi * u)) + taper = jnp.concatenate( + (ramp[:-1], jnp.ones((n_keep,), dtype=jnp.float64), + jnp.flip(ramp[:-1]))) + return kappa_t * taper[None, :] + + +def _reflected_series(kappa_t, guard): + supported = _tapered_support(kappa_t, guard) + return jnp.concatenate( + (supported, jnp.flip(supported[..., 1:-1], axis=-1)), axis=-1) + + +def reconstruct_time_primitive(kappa_t, factor, guard=0): + """Reconstruct raw complex correlations on a uniformly refined time grid. + + The returned interval contains the original unguarded closed window only; + guard samples influence the Fourier reconstruction but are never integrated. + This is the primitive operation that must precede every distance/angle + reduction in this module. + """ + factor = int(factor) + guard = int(guard) + if factor < 1: + raise ValueError("factor must be >= 1") + kappa_t = jnp.asarray(kappa_t, dtype=jnp.complex128) + if kappa_t.ndim != 2: + raise ValueError("kappa_t must have shape (n_lanes, n_support)") + n_keep = kappa_t.shape[-1] - 2 * guard + if n_keep < 2: + raise ValueError("guard must leave at least two integration samples") + + reflected = _reflected_series(kappa_t, guard) + dense = _upsample_bandlimited(reflected, factor, axis=-1) + # The forward half of the endpoint-nonduplicating reflection has + # (n_support - 1) * factor + 1 points. Crop support after refinement so + # both integration endpoints remain exact input samples. + forward = dense[..., :(kappa_t.shape[-1] - 1) * factor + 1] + start = guard * factor + return forward[..., start:start + (n_keep - 1) * factor + 1] + + +def _time_primitive_spectrum(kappa_t, guard): + """Coefficients of the exact reflected series and the unguarded offset.""" + series = _reflected_series(kappa_t, guard) + n = series.shape[-1] + coeff = jnp.fft.fft(series, axis=-1) / float(n) + frequency = jnp.fft.fftfreq(n) + return coeff, frequency, int(guard) + + +def _evaluate_time_spectrum(coeff, frequency, positions, offset): + """Evaluate a reflected finite Fourier series at selected sample positions.""" + positions = jnp.asarray(positions, dtype=jnp.float64).ravel() + shifted = positions + float(offset) + phase = jnp.exp( + 2j * jnp.pi * shifted[:, None] * frequency[None, :]) + if coeff.shape[-1] % 2 == 0: + # ``core._upsample_bandlimited`` splits the even-length Nyquist bin + # evenly between +Nyquist and -Nyquist. Its continuous contribution + # is therefore c_N cos(pi x), rather than the one-sided + # c_N exp(-i pi x) represented by fftfreq's Nyquist entry. + phase = phase.at[:, coeff.shape[-1] // 2].set( + jnp.cos(jnp.pi * shifted)) + return jnp.einsum("lk,pk->lp", coeff, phase) + + +def evaluate_time_primitive_points(kappa_t, positions, guard=0): + """Reconstruct raw correlations only at selected unguarded positions. + + ``positions`` is measured in input-sample units from the first unguarded + sample, so integers reproduce the original row. Unlike + :func:`reconstruct_time_primitive`, the returned shape depends only on the + requested point count, never on a global refinement factor. This is the + memory-bounded primitive used by the local-cover evaluator. + """ + guard = int(guard) + kappa_t = jnp.asarray(kappa_t, dtype=jnp.complex128) + if kappa_t.ndim != 2: + raise ValueError("kappa_t must have shape (n_lanes, n_support)") + n_keep = kappa_t.shape[-1] - 2 * guard + if n_keep < 2: + raise ValueError("guard must leave at least two integration samples") + positions = jnp.asarray(positions, dtype=jnp.float64).ravel() + coeff, frequency, offset = _time_primitive_spectrum(kappa_t, guard) + return _evaluate_time_spectrum(coeff, frequency, positions, offset) + + +def spectral_time_derivative_bound(kappa_t, delta_t, guard=0, order=1): + """True per-lane bound on ``|d^order kappa/dt^order|``. + + The coefficients are those of the exact reflected finite Fourier series + used by :func:`reconstruct_time_primitive`. This is a triangle-inequality + bound, never a fit to samples. + """ + guard = int(guard) + order = int(order) + if order < 0: + raise ValueError("order must be non-negative") + if not (float(delta_t) > 0.0): + raise ValueError("delta_t must be positive") + kappa_t = jnp.asarray(kappa_t, dtype=jnp.complex128) + if kappa_t.ndim != 2: + raise ValueError("kappa_t must have shape (n_lanes, n_support)") + if kappa_t.shape[-1] - 2 * guard < 2: + raise ValueError("guard must leave at least two integration samples") + series = _reflected_series(kappa_t, guard) + n = series.shape[-1] + coeff = jnp.fft.fft(series, axis=-1) / float(n) + omega = 2.0 * jnp.pi * jnp.fft.fftfreq(n, d=float(delta_t)) + return jnp.sum(jnp.abs(coeff) * (jnp.abs(omega)[None, :] ** order), axis=-1) + + +def _lane_log_integrand(kappa, rho_sq, log_lane_weight): + """Nonlinear downstream marginal, evaluated only after reconstruction.""" + exponent = kappa.real - 0.5 * rho_sq[:, None] + return jax.scipy.special.logsumexp( + exponent + log_lane_weight[:, None], axis=0) + + +def plan_time_cover(kappa_enum, rho_sq, log_lane_weight, derivative_bound, + enum_step, keep_nats=40.0): + """Plan complete time cells and certify the mass outside their union. + + ``kappa_enum`` must already be a reconstruction of the primitive. The API + accepts no marginalized time series. ``derivative_bound[l]`` must be a true + bound on ``|kappa'_l|``; use :func:`spectral_time_derivative_bound`. + """ + kappa_enum = jnp.asarray(kappa_enum, dtype=jnp.complex128) + rho_sq = jnp.asarray(rho_sq, dtype=jnp.float64) + log_lane_weight = jnp.asarray(log_lane_weight, dtype=jnp.float64) + derivative_bound = jnp.asarray(derivative_bound, dtype=jnp.float64) + if kappa_enum.ndim != 2 or kappa_enum.shape[-1] < 2: + raise ValueError("kappa_enum must have shape (n_lanes, n_enum >= 2)") + n_lane = kappa_enum.shape[0] + for name, value in (("rho_sq", rho_sq), + ("log_lane_weight", log_lane_weight), + ("derivative_bound", derivative_bound)): + if value.ndim != 1 or value.shape[0] != n_lane: + raise ValueError("%s must have shape (n_lanes,)" % name) + if not (float(enum_step) > 0.0): + raise ValueError("enum_step must be positive") + if not (float(keep_nats) > 0.0): + raise ValueError("keep_nats must be positive") + + node_log = _lane_log_integrand(kappa_enum, rho_sq, log_lane_weight) + peak_lower = jnp.max(node_log) + q = kappa_enum.real - 0.5 * rho_sq[:, None] + lift = derivative_bound[:, None] * float(enum_step) + + # Each endpoint-derived expression bounds the ENTIRE cell. The minimum + # of two upper bounds is still an upper bound and is often much tighter. + left_upper = jax.scipy.special.logsumexp( + q[:, :-1] + lift + log_lane_weight[:, None], axis=0) + right_upper = jax.scipy.special.logsumexp( + q[:, 1:] + lift + log_lane_weight[:, None], axis=0) + cell_upper = jnp.minimum(left_upper, right_upper) + + # Target from the reconstructed nodes, certify from cell_upper. Selection + # is intentionally stopped: changing which cells belong to a cover is a + # discrete planner decision, not a differentiable likelihood operation. + live = jax.lax.stop_gradient(cell_upper >= peak_lower - float(keep_nats)) + omitted = jnp.where( + live, -jnp.inf, cell_upper + jnp.log(float(enum_step))) + outside = jax.scipy.special.logsumexp(omitted) + return TimeCoverPlan(live, cell_upper, outside, peak_lower, + jnp.asarray(enum_step, dtype=jnp.float64)) + + +def _node_weights(live_cells, fine_factor, enum_factor, delta_t): + """Composite-trapezoid weights for a union of complete enum cells.""" + sub = int(fine_factor) // int(enum_factor) + fine_cells = jnp.repeat(live_cells, sub) + h = float(delta_t) / float(fine_factor) + # Every live fine cell contributes h/2 at each end. Adjacent cells + # therefore give their shared point weight h, without double counting. + middle = 0.5 * h * (fine_cells[:-1].astype(jnp.float64) + + fine_cells[1:].astype(jnp.float64)) + return jnp.concatenate( + (jnp.asarray([0.5 * h * fine_cells[0]], dtype=jnp.float64), + middle, + jnp.asarray([0.5 * h * fine_cells[-1]], dtype=jnp.float64))) + + +def _evaluate_cover_at_factor(kappa_t, rho_sq, log_lane_weight, plan, + delta_t, enum_factor, factor, guard, max_nodes): + sub = int(factor) // int(enum_factor) + nodes_per_cell = sub + 1 + live = jax.lax.stop_gradient(jnp.asarray(plan.live_cells, dtype=bool)) + n_live = jnp.count_nonzero(live) + # Cells are integrated independently. Shared endpoints therefore appear + # twice with half weight from each neighbour, exactly reproducing composite + # trapezoid weights without constructing the global refined grid. + n_local = n_live * nodes_per_cell + capacity_ok = n_local <= int(max_nodes) + dense_nodes = (kappa_t.shape[-1] - 2 * int(guard) - 1) * int(factor) + 1 + # Static shape refusal: do not construct even one local node vector when a + # single cell is larger than the declared memory capacity. Reporting the + # decline only after evaluating that cell would make ``max_nodes`` advisory + # rather than a hard bound. + if nodes_per_cell > int(max_nodes): + return (jnp.asarray(-jnp.inf), n_local, jnp.asarray(False), + dense_nodes) + # Compact the selected cells into a fixed-capacity vector. The mask and + # count remain discrete planner outputs, while every compiled scan body has + # the same small local-node shape. Avoiding a conditional around the + # Fourier evaluator is important on accelerators: the conditional/FFT + # combination produces a disproportionately large compiled program. + cell_capacity = min(live.size, + max(1, int(max_nodes) // nodes_per_cell)) + cell_index = jnp.nonzero(live, size=cell_capacity, fill_value=0)[0] + active = jnp.arange(cell_capacity) < n_live + + coeff, frequency, offset = _time_primitive_spectrum(kappa_t, guard) + local_index = jnp.arange(nodes_per_cell, dtype=jnp.float64) + log_trap = jnp.where( + (local_index == 0) | (local_index == sub), + -jnp.log(2.0), 0.0) + log_h = jnp.log(float(delta_t) / float(factor)) + + def _cell_value(cell_index): + positions = ((cell_index * sub + local_index) / float(factor)) + # Reconstruct primitive values FIRST and apply the nonlinear reduction + # over lanes only at these local nodes. No globally refined primitive + # or marginalized time row exists on this path. + primitive_local = _evaluate_time_spectrum( + coeff, frequency, positions, offset) + log_t = _lane_log_integrand( + primitive_local, rho_sq, log_lane_weight) + return jax.scipy.special.logsumexp(log_t + log_trap) + log_h + + def _step(total, args): + selected_cell, cell_live = args + contribution = jax.lax.cond( + cell_live, _cell_value, lambda _: jnp.asarray(-jnp.inf), + selected_cell) + return jnp.logaddexp(total, contribution), None + + value, _ = jax.lax.scan( + _step, jnp.asarray(-jnp.inf), + (cell_index.astype(jnp.float64), active)) + return value, n_local, capacity_ok, dense_nodes + + +def time_first_peak_local_marginalize( + kappa_t, rho_sq, log_lane_weight, delta_t, *, guard=0, + enum_factor=8, fine_factor=32, max_nodes=8192, + keep_nats=40.0, tail_tol_nats=-23.0, quadrature_tol_nats=1.0e-5): + """Peak-local joint marginal with time applied to primitives first. + + Parameters other than the three lane arrays are planner policy and are + expected to be static under :func:`jax.jit`. ``fine_factor`` is checked + against ``2*fine_factor``; the latter value is returned. The cover is + planned once on ``enum_factor`` and reused by both quadratures. + + Returns + ------- + value : scalar + Local-cover integral at ``2*fine_factor``. It is diagnostic only when + ``ok`` is false. + ok : bool scalar + True iff the node capacity, local quadrature convergence, finite-input + check, and certified omitted-mass threshold all pass. + ledger : dict of JAX scalars + Named diagnostics. A caller owns the fail-closed fallback. + """ + guard = int(guard) + enum_factor = int(enum_factor) + fine_factor = int(fine_factor) + max_nodes = int(max_nodes) + if enum_factor < 1: + raise ValueError("enum_factor must be >= 1") + if fine_factor < enum_factor or fine_factor % enum_factor: + raise ValueError("fine_factor must be a multiple of enum_factor") + if max_nodes < 2: + raise ValueError("max_nodes must be at least 2") + if not (float(delta_t) > 0.0): + raise ValueError("delta_t must be positive") + + kappa_t = jnp.asarray(kappa_t, dtype=jnp.complex128) + rho_sq = jnp.asarray(rho_sq, dtype=jnp.float64) + log_lane_weight = jnp.asarray(log_lane_weight, dtype=jnp.float64) + _validate_primitive_shapes(kappa_t, rho_sq, log_lane_weight, guard) + + derivative_bound = spectral_time_derivative_bound( + kappa_t, delta_t, guard=guard, order=1) + kappa_enum = reconstruct_time_primitive( + kappa_t, enum_factor, guard=guard) + plan = plan_time_cover( + kappa_enum, rho_sq, log_lane_weight, derivative_bound, + float(delta_t) / enum_factor, keep_nats=keep_nats) + + value_lo, n_lo, cap_lo, dense_lo = _evaluate_cover_at_factor( + kappa_t, rho_sq, log_lane_weight, plan, delta_t, enum_factor, + fine_factor, guard, max_nodes) + value_hi, n_hi, cap_hi, dense_hi = _evaluate_cover_at_factor( + kappa_t, rho_sq, log_lane_weight, plan, delta_t, enum_factor, + 2 * fine_factor, guard, max_nodes) + + quadrature_error = jnp.abs(value_hi - value_lo) + # A capacity refusal must never masquerade as zero likelihood. Preserve a + # finite diagnostic lower-resolution value when available; if neither rule + # could be evaluated, the finite nodal peak is explicitly diagnostic only. + diagnostic_value = jnp.where( + jnp.isfinite(value_hi), value_hi, + jnp.where(jnp.isfinite(value_lo), value_lo, plan.peak_lower)) + tail_margin = plan.outside_log_bound - diagnostic_value + finite_inputs = (jnp.all(jnp.isfinite(kappa_t.real)) + & jnp.all(jnp.isfinite(kappa_t.imag)) + & jnp.all(jnp.isfinite(rho_sq)) + & jnp.all(jnp.isfinite(derivative_bound)) + & jnp.all(jnp.isfinite(log_lane_weight) + | jnp.isneginf(log_lane_weight))) + capacity_ok = cap_lo & cap_hi + quadrature_ok = quadrature_error <= float(quadrature_tol_nats) + tail_ok = tail_margin < float(tail_tol_nats) + # Priority makes the decline reasons disjoint. A caller can therefore + # reconcile one and only one terminal state without interpreting a set of + # overlapping diagnostic predicates. + decline_nonfinite = ~finite_inputs + decline_capacity = finite_inputs & (~capacity_ok) + decline_quadrature = finite_inputs & capacity_ok & (~quadrature_ok) + decline_tail = finite_inputs & capacity_ok & quadrature_ok & (~tail_ok) + ok = finite_inputs & capacity_ok & quadrature_ok & tail_ok + reconciles = (ok.astype(jnp.int32) + + decline_nonfinite.astype(jnp.int32) + + decline_capacity.astype(jnp.int32) + + decline_quadrature.astype(jnp.int32) + + decline_tail.astype(jnp.int32)) == 1 + + ledger = { + "accepted": ok, + "decline_nonfinite": decline_nonfinite, + "decline_capacity": decline_capacity, + "decline_quadrature": decline_quadrature, + "decline_tail": decline_tail, + "reconciles": reconciles, + "capacity_ok": capacity_ok, + "returned_high_order": jnp.isfinite(value_hi), + "quadrature_ok": quadrature_ok, + "tail_ok": tail_ok, + "finite_inputs": finite_inputs, + "quadrature_error": quadrature_error, + "tail_margin": tail_margin, + "outside_log_bound": plan.outside_log_bound, + "peak_lower": plan.peak_lower, + "n_live_cells": jnp.count_nonzero(plan.live_cells), + "n_cells": jnp.asarray(plan.live_cells.size), + "n_local_lo": n_lo, + "n_local_hi": n_hi, + "n_dense_lo": jnp.asarray(dense_lo), + "n_dense_hi": jnp.asarray(dense_hi), + } + return diagnostic_value, ok, ledger + + +def time_first_distance_peak_local_marginalize( + kappa_unit_t, rho_sq_unit, x_grid, log_weight, delta_t, **kwargs): + """Distance adapter for :func:`time_first_peak_local_marginalize`. + + ``kappa_unit_t`` is the raw unit-distance complex correlation, including + optional support guards. Distance scaling is applied lane-by-lane *before* + reconstruction; linearity then makes reconstructing the scaled lanes + identical to scaling the reconstructed primitive. The nonlinear distance + log-sum-exp is formed only after reconstruction at each requested time. + + This helper handles one outer sample. Batch it with :func:`jax.vmap`. + """ + kappa_unit_t = jnp.asarray(kappa_unit_t, dtype=jnp.complex128) + if kappa_unit_t.ndim != 1: + raise ValueError("kappa_unit_t must have shape (n_support,); use vmap for batches") + x_grid = jnp.asarray(x_grid, dtype=jnp.float64).ravel() + log_weight = jnp.asarray(log_weight, dtype=jnp.float64).ravel() + if x_grid.shape != log_weight.shape: + raise ValueError("x_grid and log_weight must have identical shape") + rho_sq_unit = jnp.asarray(rho_sq_unit, dtype=jnp.float64) + if rho_sq_unit.ndim != 0: + raise ValueError("rho_sq_unit must be a scalar (time-independent norm)") + kappa_lanes = x_grid[:, None] * kappa_unit_t[None, :] + rho_lanes = jnp.square(x_grid) * rho_sq_unit + return time_first_peak_local_marginalize( + kappa_lanes, rho_lanes, log_weight, delta_t, **kwargs) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py index 43d02c9cd..4f37d315e 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py @@ -35,14 +35,12 @@ DIST_GRID_TOL_DEFAULT, DIST_GRID_SCHEMES, estimate_distance_peak, phi_ref_grid, psi_grid, 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] + TIME_QUAD_DEFAULT, _TIME_QUAD_CHOICES, + bandlimited_time_guard) +# The probe direction for the build-time identity check moved to +# anglemarg.gh_laplace_supported_for_data: the policy's reserve roster asks the +# same question, and a second probe direction here would be a second definition +# of it. from . import core as _core from .anglemarg import (ANGLE_MARG_DEFAULT, ANGLE_MARG_LEGACY, # noqa: F401 ANGLE_MARG_CHOICES) @@ -56,9 +54,7 @@ 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 + g0, g_certificate = bandlimited_time_guard(len(tvals)) # 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) @@ -67,6 +63,15 @@ def bandlimited_storage_requirement(deltaT, integration_window_half): def _validate_nonlinear_time_quadrature(time_quadrature, endpoint): + """Refusal for the endpoints whose reduction has no refinable primitive here. + + The pure distance reduction is not one of them: it consumes the same + ``(kappa, rho^2)`` the refinement produces, so + :class:`JAXDistanceMarginalizedLikelihood` applies it on the refined nodes + instead of calling this. The phi/psi/exact-angle endpoints either stream a + per-phi primitive the refined grid cannot hold or receive an already-reduced + lnL(t) from the coefficient-table kernels. + """ if time_quadrature not in _TIME_QUAD_CHOICES: raise ValueError("time_quadrature must be one of %r" % (_TIME_QUAD_CHOICES,)) if time_quadrature == "bandlimited": @@ -180,12 +185,48 @@ def _L_of(det): return data, extras +def build_rotating_freqresponse_data_from_precompute( + P, data_dict, psd_dict, fiducial_epoch, integration_window_half, + Lmax, fMax, t_window=0.1, Qmax=4, L_arm=None, p_max=0, + analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0.0, + tvals=None, verbose=False, **precompute_kwargs): + """One-call builder for the compound rotation + finite-response likelihood.""" + import RIFT.likelihood.factored_likelihood_rotating_freqresponse as flrr + import RIFT.likelihood.slowrot_freqresponse as sfr + from .banded import build_rotating_freqresponse_data + + bk = flrr.PrecomputeLikelihoodTermsRotatingFreqResponse( + fiducial_epoch, t_window, P, data_dict, psd_dict, Lmax, fMax, + Qmax=Qmax, L_arm=L_arm, p_max=p_max, + analyticPSD_Q=analyticPSD_Q, inv_spec_trunc_Q=inv_spec_trunc_Q, + T_spec=T_spec, verbose=verbose, quiet=not verbose, + skip_interpolation=True, **precompute_kwargs) + meta = bk[4] + lk, rba, uba, vba, ep = flrr.pack_rotating_freqresponse_arrays( + meta, bk[3], bk[1], bk[2]) + + def _L_of(det): + return L_arm.get(det, None) if isinstance(L_arm, dict) else L_arm + det_geom = {det: sfr.detector_geometry(det, L_arm=_L_of(det)) + for det in data_dict} + deltaT = float(P.deltaT) + if tvals is None: + tvals = factored_likelihood.marginalization_time_grid( + integration_window_half, deltaT, xpy=np) + data = build_rotating_freqresponse_data( + meta, lk, rba, uba, vba, ep, deltaT, tvals, det_geom) + extras = dict(meta=meta, rho_by_a=rba, U_by_aa=uba, V_by_aa=vba, + epochDict=ep, lookupNKDict=lk, det_geom=det_geom) + return data, extras + + def build_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, storage_window_half, integration_window_half, Lmax, fMax, analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0.0, tvals=None, verbose=False, skip_interpolation=True, + q_time_pregrid_factor=1, **precompute_kwargs): """Run the production precompute + packing, return a JAXLikelihoodData. @@ -210,6 +251,12 @@ def build_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, against the numpy reference should still pass ``data.tvals`` to the reference rather than rebuild a grid. + ``q_time_pregrid_factor`` (``--q-time-pregrid-factor``) refines the sampling + of the STORED rholm buffers by that integer factor before they reach the + device, leaving ``deltaT``, ``tvals`` and the Simpson weights alone. 1 is the + default and the historical behaviour; see + :func:`RIFT.likelihood.jax_ile.core.build_q_time_pregrid`. + Returns ------- (data, extras) where ``data`` is a :class:`JAXLikelihoodData` and @@ -244,7 +291,8 @@ def build_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, tvals = factored_likelihood.marginalization_time_grid( integration_window_half, deltaT, xpy=np) - data = build_likelihood_data(packed, deltaT, float(fiducial_epoch), tvals) + data = build_likelihood_data(packed, deltaT, float(fiducial_epoch), tvals, + q_time_pregrid_factor=q_time_pregrid_factor) extras = dict(rholms=rholms, cross_terms=cross_terms, cross_terms_V=cross_terms_V, guess_snr=guess_snr, rholms_intp=rholms_intp) @@ -270,9 +318,8 @@ def __init__(self, data, interp=JAX_INTERP_DEFAULT, phase_marginalization=False, 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 + (self.time_guard_initial, + self.time_guard_certified) = bandlimited_time_guard(data.npts) def _batched(ra, dec, psi, incl, phiref, distMpc): return fused_log_likelihood( @@ -317,6 +364,139 @@ def fisher(self, theta6): return -H +class JAXFixedDistanceLikelihood: + """Five-angle view of :class:`JAXExtrinsicLikelihood` at fixed distance. + + This is the fixed-distance geometry used by the high-SNR Event-B validation: + ``theta5 = (ra, dec, psi, incl, phiref)`` is sampled while ``distMpc`` is a + constant. Keeping this as a likelihood view (instead of an extremely narrow + distance prior) removes a numerically artificial sixth direction from both + hill climbing and the observed Fisher matrix. + """ + + ANGULAR_PARAM_ORDER = ("ra", "dec", "psi", "incl", "phiref") + + def __init__(self, likelihood, dist_mpc, phase_shift=0.0): + if not isinstance(likelihood, JAXExtrinsicLikelihood): + raise TypeError("likelihood must be a JAXExtrinsicLikelihood") + self.likelihood = likelihood + self.data = likelihood.data + self.dist_mpc = float(dist_mpc) + self.phase_shift = float(phase_shift) + if self.phase_shift: + self.ANGULAR_PARAM_ORDER = ( + "ra", "dec", "psi", "incl", "phiref_shifted") + self.interp = likelihood.interp + self.phase_marginalization = likelihood.phase_marginalization + self.time_quadrature = likelihood.time_quadrature + if hasattr(likelihood, "time_guard_initial"): + self.time_guard_initial = likelihood.time_guard_initial + self.time_guard_certified = likelihood.time_guard_certified + + distance = jnp.asarray([self.dist_mpc], dtype=jnp.float64) + + def _scalar(theta5): + physical5 = theta5.at[4].set( + jnp.mod(theta5[4] + self.phase_shift, 2.0 * jnp.pi)) + return likelihood._scalar(jnp.concatenate([physical5, distance])) + + self._scalar = _scalar + self._value_and_grad = jax.jit(jax.value_and_grad(_scalar)) + self._hessian = jax.jit(jax.hessian(_scalar)) + + def log_likelihood(self, ra, dec, psi, incl, phiref): + ra = jnp.asarray(ra) + distance = jnp.full_like(ra, self.dist_mpc) + return self.likelihood.log_likelihood( + ra, dec, psi, incl, + jnp.mod(jnp.asarray(phiref) + self.phase_shift, 2.0 * jnp.pi), + distance) + + def to_sampler_coordinates(self, theta5): + """Map physical ``phiref`` to this view's shifted sampler coordinate.""" + theta5 = np.array(theta5, dtype=float, copy=True) + theta5[..., 4] = np.mod(theta5[..., 4] - self.phase_shift, 2.0 * np.pi) + return theta5 + + def to_physical_coordinates(self, theta5): + """Map this view's sampler coordinate back to physical ``phiref``.""" + theta5 = np.array(theta5, dtype=float, copy=True) + theta5[..., 4] = np.mod(theta5[..., 4] + self.phase_shift, 2.0 * np.pi) + return theta5 + + def value(self, theta5): + return float(self._scalar(jnp.asarray(theta5, dtype=jnp.float64))) + + def value_and_grad(self, theta5): + value, grad = self._value_and_grad( + jnp.asarray(theta5, dtype=jnp.float64)) + return float(value), np.asarray(grad) + + def fisher(self, theta5): + hessian = np.asarray( + self._hessian(jnp.asarray(theta5, dtype=jnp.float64))) + return -hessian + + +class JAXRotatedPhaseLikelihood: + """Rotate ``(psi, phiref)`` into AV-friendly sum/difference coordinates. + + This mirrors conventional ILE's ``--internal-rotate-phase``. Both rotated + coordinates live on ``[0, 4 pi)``; the redundant cover preserves the flat + physical angle prior and straightens the leading quadrupole degeneracy. + """ + + ANGULAR_PARAM_ORDER = ("ra", "dec", "phase_p", "incl", "phase_m") + + def __init__(self, likelihood): + if tuple(getattr(likelihood, "ANGULAR_PARAM_ORDER", ())) not in ( + ("ra", "dec", "psi", "incl", "phiref"), + ("ra", "dec", "psi", "incl", "phiref_shifted")): + raise TypeError("likelihood must expose ra,dec,psi,incl,phase coordinates") + self.likelihood = likelihood + for name in ("data", "interp", "phase_marginalization", "time_quadrature"): + setattr(self, name, getattr(likelihood, name)) + + def _scalar(theta): + psi = jnp.mod(0.5 * (theta[2] - theta[4]), jnp.pi) + phase = jnp.mod(0.5 * (theta[2] + theta[4]), 2.0 * jnp.pi) + physical = jnp.stack([theta[0], theta[1], psi, theta[3], phase]) + return likelihood._scalar(physical) + + self._scalar = _scalar + self._value_and_grad = jax.jit(jax.value_and_grad(_scalar)) + self._hessian = jax.jit(jax.hessian(_scalar)) + + def log_likelihood(self, ra, dec, phase_p, incl, phase_m): + psi = jnp.mod(0.5 * (phase_p - phase_m), jnp.pi) + phase = jnp.mod(0.5 * (phase_p + phase_m), 2.0 * jnp.pi) + return self.likelihood.log_likelihood(ra, dec, psi, incl, phase) + + def to_sampler_coordinates(self, theta5): + base = self.likelihood.to_sampler_coordinates(theta5) + out = np.array(base, dtype=float, copy=True) + out[..., 2] = np.mod(base[..., 4] + base[..., 2], 4.0 * np.pi) + out[..., 4] = np.mod(base[..., 4] - base[..., 2], 4.0 * np.pi) + return out + + def to_physical_coordinates(self, theta5): + theta5 = np.asarray(theta5, dtype=float) + base = np.array(theta5, copy=True) + base[..., 2] = np.mod(0.5 * (theta5[..., 2] - theta5[..., 4]), np.pi) + base[..., 4] = np.mod(0.5 * (theta5[..., 2] + theta5[..., 4]), 2.0 * np.pi) + return self.likelihood.to_physical_coordinates(base) + + def value(self, theta5): + return float(self._scalar(jnp.asarray(theta5, dtype=jnp.float64))) + + def value_and_grad(self, theta5): + value, grad = self._value_and_grad(jnp.asarray(theta5, dtype=jnp.float64)) + return float(value), np.asarray(grad) + + def fisher(self, theta5): + return -np.asarray(self._hessian(jnp.asarray(theta5, dtype=jnp.float64))) + + class JAXDistanceMarginalizedLikelihood: """Distance- and time-marginalized lnL over the 5 angular parameters. @@ -336,9 +516,12 @@ def __init__(self, data, d_min, d_max, n_grid=256, d_prior="euclidean", 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") + 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": + (self.time_guard_initial, + self.time_guard_certified) = bandlimited_time_guard(data.npts) 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) @@ -547,9 +730,18 @@ def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, 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): + dist_grid="uniform", dist_grid_tol=DIST_GRID_TOL_DEFAULT, + direct_marginalization_policy=None, policy_config=None, + multipeak_guard=16): self.data = data self.interp = interp # the instance's stencil; sample_phi_ref defaults to it + from . import direct_marginalization_policy as _policy + if direct_marginalization_policy is None: + direct_marginalization_policy = _policy.POLICY_DEFAULT + if direct_marginalization_policy not in _policy.POLICY_CHOICES: + raise ValueError("direct_marginalization_policy must be one of %r, " + "got %r" % (_policy.POLICY_CHOICES, + direct_marginalization_policy)) _validate_nonlinear_time_quadrature( time_quadrature, "distance/phase/polarization marginalization") self.time_quadrature = time_quadrature @@ -591,12 +783,13 @@ def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, # 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)) + "dist_grid=%r cannot be combined with distance-GH-nodes=%d " + "(--distance-gh-nodes / JAX_ILE_DISTMARG_GH): 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. " + "Pass --distance-gh-nodes 0 (or 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)): @@ -906,15 +1099,8 @@ def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, # 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)) + gh_ok, gh_info = _anglemarg.gh_laplace_supported_for_data( + data, interp) if angle_marg == "auto": scheme, sel_info = _anglemarg.choose_angle_marg_scheme( amp_data, gh_laplace_ok=gh_ok) @@ -923,13 +1109,14 @@ def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, 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-GH-nodes set (--distance-gh-nodes / " + "JAX_ILE_DISTMARG_GH), 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." + "do not hold. Use --angle-marg-scheme exact, or pass " + "--distance-gh-nodes 0 (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, @@ -942,43 +1129,190 @@ def __init__(self, data, d_min, d_max, nphi=32, npsi=16, n_grid=256, # 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"): + if scheme in ("exact", "laplace", "peak-local", "phi-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): + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=False): + if return_amp: + raise ValueError( + "the 'grid' scheme has no amp_sizing and no runtime " + "amplitude failsafe, so there is no metric to return") 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): + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=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) + return_lnLt=return_lnLt, return_amp=return_amp) 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): + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=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) + return_lnLt=return_lnLt, return_amp=return_amp) + elif scheme == "multipeak": + # The four-axis controller: it OWNS the time integral, so there is no + # lnL(t) and time_quadrature does not reach it. Reachable only by + # name; not in 'auto'. + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=False): + if return_lnLt: + raise ValueError( + "--angle-marg-scheme multipeak marginalizes time inside " + "the controller; there is no lnL(t) to return. Use " + "another scheme if you need the time series.") + v = _anglemarg.fused_log_likelihood_distphipsimarg_multipeak( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, guard=int(multipeak_guard)) + return (v, jnp.asarray(amp_sizing)) if return_amp else v + + elif scheme == "phi-local": + # BOTH angle axes localized, with a dense fallback wherever the certificate + # declines. By name only, and deliberately not in 'auto': it is slower than + # 'peak-local' until the fallback can be skipped, which needs a measured + # acceptance rate on production tables. + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=False): + return _anglemarg.fused_log_likelihood_distphipsimarg_phi_local( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, time_quadrature=time_quadrature, + return_lnLt=return_lnLt, return_amp=return_amp) else: # laplace - def _fused(data_, ra, dec, incl, return_lnLt=False): + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=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) - + return_lnLt=return_lnLt, return_amp=return_amp) + + # Cross-axis policy (opt-in). It REPLACES the per-scheme _fused above: + # the composite owns time, distance and both angles per row, and uses + # the exact scheme only as its reserve. Every combination it cannot + # honour is refused here, not ignored. + self.direct_marginalization_policy = direct_marginalization_policy + self.policy_info = None + self.policy_config = None + self._batched_ledger = None + if direct_marginalization_policy != "off": + _policy.validate_policy_request( + direct_marginalization_policy, angle_marg_scheme=scheme, + time_quadrature=time_quadrature, d_prior=d_prior, + dist_grid=dist_grid) + cfg = policy_config if policy_config is not None else ( + _policy.PolicyConfig()) + _policy.validate_policy_config(cfg) + lln, norm_info = _policy.policy_log_normalization( + data, xg, lwg, d_prior=d_prior) + _policy.probe_guarded_tables(data, interp, int(cfg.time_guard)) + self.policy_config = cfg + self.policy_info = dict( + norm_info, policy=direct_marginalization_policy, + # The reserve's ANGLE scheme. Not PolicyConfig.reserve_scheme, + # which names WHICH reserve runs -- two different quantities + # that shared this key while 'exact' was the only reserve. The + # driver overwrites "reserve_scheme" with the resolved pair and + # keeps this one under its own name. + reserve_angle_scheme=scheme, + reserve_scheme=scheme, + time_guard=int(cfg.time_guard), + reserve_time_refine=int(cfg.reserve_time_refine), + reserve_distance_gh_nodes=int(_core._DISTMARG_GH_N), + reserve_batch_rows=int(cfg.reserve_batch_rows), + reserve_pair=str(cfg.reserve_scheme), + max_time_nodes=int(cfg.max_time_nodes), + reserve_peaklocal_fine_nodes=int(cfg.reserve_peaklocal_fine_nodes), + total_value_error_budget_nats=float( + cfg.total_value_error_budget_nats), + # The plan-sizing and tolerance knobs are reported for the same + # reason the resolved angle scheme is: they decide whether the + # controller can accept at all, and a caller that passed one and + # got the default back had no way to see it from the log. + max_modes=int(cfg.max_modes), + enriched_max_modes=int(cfg.enriched_max_modes), + base_oversample=int(cfg.base_oversample), + enriched_oversample=int(cfg.enriched_oversample), + base_max_starts=int(cfg.base_max_starts), + convergence_tol_nats=float(cfg.convergence_tol_nats), + time_guard_tol_nats=float(cfg.time_guard_tol_nats)) + self.angle_marg_info["direct_marginalization_policy"] = ( + direct_marginalization_policy) + + def _fused(data_, ra, dec, incl, return_lnLt=False, + return_amp=False): + if return_amp: + raise ValueError( + "direct_marginalization_policy=%r does not expose the " + "angle-grid amplitude metric" + % (direct_marginalization_policy,)) + if return_lnLt: + raise ValueError( + "direct_marginalization_policy=%r marginalizes time " + "inside the composite; there is no lnL(t) to return" + % (direct_marginalization_policy,)) + return _policy.fused_log_likelihood_four_axis_policy( + data_, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, config=cfg, + local_log_normalization=lln) + + def _batched_ledger(ra, dec, incl): + return _policy.fused_log_likelihood_four_axis_policy( + data, ra, dec, incl, xg, lwg, interp=interp, + amp_sizing=amp_sizing, config=cfg, + local_log_normalization=lln, return_ledger=True) + self._batched_ledger = jax.jit(_batched_ledger) + + self._fused = _fused + + # WHICH CALLS ARE COVERED BY THE AMPLITUDE FAILSAFE, AND WHY IT IS THE + # BATCHED PATH. The kernels return their amplitude metric instead of + # reporting it from inside the graph, because a host callback anywhere + # in a jitted module makes that module ineligible for JAX's persistent + # compilation cache (jax/_src/compiler.py::_cache_write) -- and the + # angle-marginalization graph is the most expensive compile in RIFT. + # + # The batched path is every pilot, reweight and final output-cloud + # evaluation, i.e. every point that reaches a published artifact, so + # the recorded coverage is exactly the set of points the label speaks + # for. The scalar AD/flow-training path deliberately does not report: + # its proposals do not enter those artifacts, and asking for the metric + # there would put a second output on the differentiated graph. + self._amp_record = None + if scheme in ("exact", "laplace", "peak-local", "phi-local") and ( + direct_marginalization_policy == "off"): + self._amp_record = lambda amp: _anglemarg.record_amp_failsafe( + amp, amp_sizing, scheme) + + # _batched KEEPS its lnL-only contract, and the metric-bearing graph is + # a SEPARATE jit. Folding the amplitude into _batched made its arity + # depend on the construction options, so `np.asarray(like._batched(...))` + # -- which test_angle_marg_peaklocal_wiring.py and + # test_direct_marginalization_policy.py both do, on the SAME line as a + # policy-enabled sibling whose _batched still returned one array -- + # raised "inhomogeneous shape" for the amp-sized schemes only. Both jits + # are lazy, and production reaches only the one log_likelihood calls, so + # nothing is compiled or cached twice. def _batched(ra, dec, incl): return _fused(data, ra, dec, incl) self._batched = jax.jit(_batched) + self._batched_amp = None + if self._amp_record is not None: + def _batched_amp(ra, dec, incl): + return _fused(data, ra, dec, incl, return_amp=True) + self._batched_amp = jax.jit(_batched_amp) + def _scalar(theta3): v = _fused(data, theta3[0:1], theta3[1:2], theta3[2:3]) return v[0] @@ -988,7 +1322,16 @@ def _scalar(theta3): def log_likelihood(self, ra, dec, incl): """lnL for arrays of 3 angular parameters (ra, dec, incl), shape (S,).""" - return self._batched(jnp.asarray(ra), jnp.asarray(dec), jnp.asarray(incl)) + if self._batched_amp is None: + return self._batched(jnp.asarray(ra), jnp.asarray(dec), + jnp.asarray(incl)) + # One deliberate device->host read per batch, after the values are + # already required on the host anyway. The maximum accumulates across + # calls for the whole event; the values are returned unchanged. + values, amp_call = self._batched_amp( + jnp.asarray(ra), jnp.asarray(dec), jnp.asarray(incl)) + self._amp_record(amp_call) + return values def value(self, theta3): return float(self._scalar(jnp.asarray(theta3, dtype=jnp.float64))) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py index fb20f348e..408b60b36 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py @@ -22,42 +22,44 @@ 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. +HOW THE MODES ARE FOUND. Both angular derivatives are finite Laurent polynomials. +``bivariate_trig_stationary.enumerate_torus_maxima`` clears their Laurent powers, +eliminates one variable with a Sylvester resultant, and solves that resultant as a +generalized polynomial eigenproblem. A generic affine hidden variable separates +stationary points that share exactly the same phi or u. The solve must recover the +mixed-volume (BKK) root count and agree under two independent projections; otherwise +this path declines. Enumeration cost is fixed by bidegree and never by amplitude. + +NO ON-CIRCLE FILTER, deliberately. Roots are classified with the real polynomial's +reciprocal-conjugate involution: a torus root is a fixed point, while a complex root has +a distinct partner. A close pair whose status is numerically ambiguous declines the +whole solve. Thus a tolerance can never silently remove a possible real mode. 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. + * Both angular axes are certified at enumeration time for a regular, + zero-dimensional stationary system: the finite algebraic solve recovers its BKK + root count and two projections agree. + * Degenerate or ill-conditioned systems are not certified. Any definite + candidates they retain are explicitly partial and require the outside-cover gate. * 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. + * Every retained cover, including one built from a complete root set, must pass + an independent doubled-rule check on its inside-box quadrature. Root + completeness and an omitted-mass bound say nothing about that error. - 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. + An incomplete algebraic set is used only when that omitted-mass bound passes. + Otherwise this reference executes its dense-phi/exact-u fallback and returns a + finite answer with the algebraic ledger and fallback reason attached. """ import numpy as np +from .bivariate_trig_stationary import enumerate_torus_maxima + __all__ = [ "W_SIGMA", "MERGE_MAX_PASSES", @@ -68,6 +70,7 @@ "enumerate_modes", "derivative_bound", "outside_bound", + "dense_phi_exact_u_marginalize", "joint_marginalize_peak_local", "joint_marginalize_over_distance", "u_profile", @@ -99,6 +102,14 @@ #: exp(-23) ~ 1e-10 of the mass. OUTSIDE_TOL_NATS = -23.0 +# Keep the host fallback independent of the optional JAX stack. These are the +# phi-axis pieces of jax_ile.anglemarg._dense_grid_sizes: the calibration point +# is m_max=2 and the count is rounded up to a multiple of 16. Importing that +# private helper here made the advertised NumPy fallback fail before producing +# a value whenever JAX was not installed. +_DENSE_K_PHI = 16.0 +_DENSE_FLOOR_PHI = 128 + def joint_table(C_A, C_B, x=1.0): """Coefficient table of ``g = x*A - x**2/2 * B`` from the anglemarg tables. @@ -135,6 +146,18 @@ def _kq(C): #: here because each point is summed independently. _POINT_CHUNK = 200_000 +#: Per-axis ceiling on a box's trapezoid. NOT a free tuning knob: it is the point at +#: which the local integration stops honouring the curvature it derived, and the +#: certificate cannot report that -- the omitted-mass bound covers what is outside the +#: boxes, so a capped box can carry ``margin = -inf`` and still be wrong. Measured on +#: the rho=163.08 production tables (amplitude ~2.7e4, ``area_outside == 0``) against a +#: torus reference self-converged to 2e-12: at 256 the value was off by up to 0.36 nats, +#: at 512 by 3e-4, at 1024 exact to 1e-4. Cost went 0.07 s -> 0.21 s -> 0.83 s. 512 +#: buys three orders of magnitude for 3x, and 1024 buys almost nothing more for 12x. +#: Raising this does not widen the certificate's REACH -- declines are omitted-mass +#: declines and this is internal accuracy; the two are independent, and both are needed. +_BOX_MAX_PTS = 512 + def eval_g(C, phi, u, order=(0, 0)): """``d^a_phi d^b_u g`` at points ``(phi, u)``; ``order=(a, b)``. @@ -197,58 +220,18 @@ def _wrap(d): 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)``. +def enumerate_modes(C, n_phi=None, newton_iters=None, _return_report=False): + """All isolated torus maxima from the finite bivariate polynomial. - 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`. + ``n_phi`` and ``newton_iters`` remain accepted for source compatibility with the + former grid-seeded reference, but no sampled grid enters enumeration. When + ``_return_report`` is true the internal caller also receives the fail-closed + algebraic ledger. """ - 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] + result = enumerate_torus_maxima(C) + if _return_report: + return result.points, result.hessians, result.ok, result.report + return result.points, result.hessians def _merge_boxes(cen, half): @@ -378,15 +361,40 @@ def outside_bound(C, cen, half, n_grid=256): #: 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.""" +# A local-cover value is accepted only after this independent doubled-rule +# comparison. The outside bound cannot diagnose quadrature error inside a box. +_LOCAL_QUADRATURE_TOL_NATS = 1.0e-6 + + +def _log_box_integral(C, c, h, pts_per_sigma=_PTS_PER_SIGMA, max_pts=_BOX_MAX_PTS): + """``log int_box exp(g)`` by a tensor trapezoid sized from the LOCAL curvature. + + Returns ``(value, n_points, capped)``. ``capped`` is True when ``max_pts`` bound the + curvature-derived count on either axis -- i.e. the sizing rule ASKED FOR MORE NODES + THAN IT GOT. That is a truncated request, NOT a verdict that the value is wrong: + measured on the ladder, rung 1 (rho=40.77, amplitude ~2.5e3) caps on every + mass-carrying point and is still exact to 0.00000 nats against a converged reference, + while rung 3 (rho=163.08, amplitude ~2.8e4) caps and is wrong by the amount recorded + on _BOX_MAX_PTS -- stated there once rather than repeated here. The trapezoid + on a periodic integrand converges fast enough that the derived count is conservative + at low amplitude and binding at high. So treat the flag as "look here", not "this is + broken" -- it is the only signal available, because the certificate cannot see inside + a box at all. It has to be reported, + because the certificate cannot see it: the omitted-mass bound covers what is OUTSIDE + the boxes and says nothing about the quadrature inside one, so a capped box is exactly + the case where ``margin`` can read ``-inf`` (nothing omitted at all) while the value is + still wrong. Measured on the rho=163 production tables: at the shipped cap of 256 the + value sat that far from a converged torus reference with ``area_outside == 0``. + """ n = [] + capped = False 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 + if want > max_pts: + capped = True 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]) @@ -396,27 +404,90 @@ def _log_box_integral(C, c, h, pts_per_sigma=_PTS_PER_SIGMA, max_pts=256): 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] + return m + np.log(np.sum(np.exp(g - m + W))), n[0] * n[1], capped + + +def dense_phi_exact_u_marginalize(C, n_phi=None, n_u_nodes=64): + """Finite dense-phi/exact-u fallback for one coefficient table. + + This is the host reference analogue of the shipped JAX ``laplace`` member: + phi uses the amplitude- and mode-order-derived dense sizing rule, while + :func:`u_profile` integrates the finite degree-two u polynomial by its + algebraic cell partition. The returned value is always finite for finite + input. A doubled-phi comparison is reported rather than silently treating + the requested floor as proof of convergence. + """ + C = np.asarray(C, dtype=np.complex128) + m_max = max(1, int(np.ceil((C.shape[0] - 1) / 2.0))) + amplitude_bound = max(derivative_bound(C, (0, 0)), 25.0) + m_scale = max(1.0, float(m_max) / 2.0) + derived = max(int(np.ceil(_DENSE_FLOOR_PHI * m_scale)), + int(np.ceil(_DENSE_K_PHI * m_scale + * np.sqrt(amplitude_bound)))) + derived = ((derived + 15) // 16) * 16 + base = max(int(derived), int(n_phi) if n_phi is not None else 0) + + def one(count): + phi = np.linspace(0.0, 2.0 * np.pi, count, endpoint=False) + F, _, _ = u_profile(C, phi, n_nodes=int(n_u_nodes)) + peak = float(np.max(F)) + return (peak + np.log(np.exp(F - peak).sum()) + - np.log(float(count)) - np.log(2.0 * np.pi)) + + lo = one(base) + hi = one(2 * base) + return float(hi), { + 'n_phi': int(2 * base), + 'n_phi_coarse': int(base), + 'n_u_nodes_floor': int(n_u_nodes), + 'amplitude_bound': float(amplitude_bound), + 'doubling_error': float(abs(hi - lo)), + } 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. + Returns ``(value, ok, report)`` with an explicit three-level hierarchy: + + 1. use the BKK-complete algebraic maxima when enumeration is certified; + 2. use either a complete or partial candidate union only when + :func:`outside_bound` proves omitted impact below ``tol_nats`` and a doubled + local rule verifies inside-cover quadrature; + 3. otherwise return :func:`dense_phi_exact_u_marginalize`. + + Thus incomplete root accounting can cost speed but never silently deletes a + likelihood sample. A dense-fallback implementation failure is raised rather than + converted to ``-inf``. """ C = np.asarray(C) rep = {'n_modes': 0, 'n_regions': 0, 'n_local_points': 0, + 'n_boxes_pts_capped': 0, 'margin': np.inf, 'area_outside': np.nan, 'sup_outside': np.nan, - 'decline': None} + 'enumeration_certified': False, 'result_path': None, + 'fallback_reason': None, 'decline': None} + + def dense_fallback(reason): + rep['fallback_reason'] = str(reason) + value, dense_report = dense_phi_exact_u_marginalize(C, n_phi=n_phi) + if not np.isfinite(value): + # Do not turn an implementation failure into a zero-likelihood + # sample. A raised error is visible; -inf would be silent deletion. + raise FloatingPointError("dense fallback returned a non-finite value") + rep['result_path'] = 'dense-phi/exact-u' + rep['dense_fallback'] = dense_report + rep['decline'] = None + return float(value), True, rep - P, H = enumerate_modes(C, n_phi=n_phi) + P, H, enum_ok, enum_report = enumerate_modes( + C, n_phi=n_phi, _return_report=True) + rep['enumeration'] = enum_report + rep['enumeration_certified'] = bool(enum_ok) rep['n_modes'] = int(P.shape[0]) if P.shape[0] == 0: - rep['decline'] = 'no modes enumerated' - return -np.inf, False, rep + return dense_fallback('algebraic enumeration produced no usable maxima: ' + + str(enum_report['decline'])) # marginal sigmas of the local Gaussian: sqrt of the diagonal of (-H)^-1 half = np.empty_like(P) @@ -429,15 +500,20 @@ def joint_marginalize_peak_local(C, n_phi=64, n_bound_grid=256, 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 + return dense_fallback('regions still overlap after MERGE_MAX_PASSES') - parts, npts = [], 0 + parts, npts, n_capped = [], 0, 0 for c, h in zip(cen, half): - v, k = _log_box_integral(C, c, h) + v, k, capped = _log_box_integral(C, c, h) parts.append(v) npts += k + n_capped += int(capped) rep['n_local_points'] = int(npts) + # a capped box had its node request truncated and the certificate CANNOT see inside a + # box at all, so surface it: it is the only available signal that 'nothing omitted' + # might be sitting on a quadrature error. Capped does NOT mean wrong -- rung 1 caps + # everywhere and is exact -- it means this is where to look if a value is doubted. + rep['n_boxes_pts_capped'] = int(n_capped) parts = np.array(parts) m = parts.max() log_inside = m + np.log(np.exp(parts - m).sum()) @@ -450,10 +526,48 @@ def joint_marginalize_peak_local(C, n_phi=64, n_bound_grid=256, 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 + local_value = float(log_inside - 2.0 * np.log(2.0 * np.pi)) + bound_ok = rep['margin'] < tol_nats + if bound_ok: + # The outside bound certifies MISSED modes, not quadrature inside the + # retained regions. Enumeration completeness cannot change that: a + # complete cover can have area_outside == 0 while a narrow diagonal + # ridge is badly under-resolved by an axis-aligned tensor rule. Always + # perform a doubled local rule before accepting the cover. If that + # independent error budget fails, level three of the hierarchy is the + # finite dense fallback -- never a sample deletion. + parts_hi = [] + capped_hi = False + for c, h in zip(cen, half): + v_hi, _, cap_hi = _log_box_integral( + C, c, h, pts_per_sigma=2 * _PTS_PER_SIGMA, + max_pts=2 * _BOX_MAX_PTS) + parts_hi.append(v_hi) + capped_hi |= bool(cap_hi) + parts_hi = np.asarray(parts_hi) + mh = float(np.max(parts_hi)) + log_inside_hi = mh + np.log(np.exp(parts_hi - mh).sum()) + quadrature_error = float(abs(log_inside_hi - log_inside)) + rep['local_quadrature_error'] = quadrature_error + rep['local_quadrature_capped'] = bool(capped_hi) + # Preserve the existing best-effort ledger names for consumers of an + # incomplete algebraic solve; the common names above cover both paths. + if not enum_ok: + rep['best_effort_quadrature_error'] = quadrature_error + rep['best_effort_quadrature_capped'] = bool(capped_hi) + if capped_hi or quadrature_error > _LOCAL_QUADRATURE_TOL_NATS: + return dense_fallback( + 'inside-cover quadrature did not converge ' + '(doubled_error=%.6g, doubled_capped=%s)' + % (quadrature_error, bool(capped_hi))) + local_value = float(log_inside_hi - 2.0 * np.log(2.0 * np.pi)) + rep['result_path'] = ('algebraic-certified' if enum_ok + else 'algebraic-best-effort/bound-certified') + if not enum_ok: + rep['fallback_reason'] = str(enum_report['decline']) + return local_value, True, rep + return dense_fallback('omitted-mass bound too large (margin %.6g >= %.6g)' + % (rep['margin'], tol_nats)) def joint_marginalize_over_distance(C_A_st, C_B_st, x_grid, log_w_grid, @@ -649,8 +763,11 @@ def u_profile(C, phi, n_nodes=64, window_sigma=12.0): # 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. + # A cell failing either is integrated WHOLE rather than windowed. That is the + # conservative branch for the CENTRE -- it never moves onto a non-stationary + # point -- but it is NOT conservative for the resolution: see the node-count + # derivation below, which exists because the whole-cell branch spreads the same + # count over a wider interval. g1c = eval_g(C, pv, ustar, (0, 1)) g2c = _g_uu_at(C, p, ustar) _m1u = max(derivative_bound(C, (0, 1)), 1e-300) @@ -663,9 +780,12 @@ def u_profile(C, phi, n_nodes=64, window_sigma=12.0): 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. + # FELL BACK spans the whole cell with the same nodes. An earlier version of the + # comment above called that branch conservative because it "can only add nodes", + # which was simply false: it adds none and spreads them wider, so rejecting a + # peak made the resolution WORSE. (That false sentence outlived its own + # retraction here by 14 lines until a grep for the NUMBER, not the paragraph, + # turned it up -- correcting a claim means finding every copy of it.) # Measured on a searched counterexample: 1.7e-03 nats at 64 nodes, converging # only by n = 1024. # @@ -713,7 +833,9 @@ 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 + LEGACY PROFILE EXPERIMENT, not the bivariate algebraic enumerator used by + :func:`joint_marginalize_peak_local`. 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 @@ -760,7 +882,20 @@ def phi_local_marginalize(C, n_seed=64, w_sigma=12.0, n_nodes=64, 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: + if wdt >= 2.0 * np.pi: + # A WINDOW THAT ALREADY SPANS THE CIRCLE IS NOT SPLIT. Splitting it emits + # (a, 2 pi) and (0, a + wdt - 2 pi), and that second endpoint is the round + # trip fl(fl(a + 2 pi) - 2 pi), which misses `a` by one ulp in a direction + # nothing here controls. The seam-close below then joins the halves into one + # region of width 2 pi MINUS ONE ULP, the clamp `sum >= 2 pi` does not fire, + # and the certificate sees area_outside = 8.9e-16 instead of 0. Measured on + # F = 1000 cos(phi - pi/96) at w_sigma = 200: one region [-0.00864509, + # 6.27454022], margin -0.657, DECLINED on omitted mass that does not exist. + # `wdt` is a min AT 2 pi, so this test is exact and needs no tolerance. The + # jax port carries the same fix in phi_local_lnI, where the halves could also + # fail to merge at all and cost 20x in the value. + pieces.append((0.0, 2.0 * np.pi)) + elif a + wdt <= 2.0 * np.pi: pieces.append((a, a + wdt)) else: pieces.append((a, 2.0 * np.pi)) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/q_time_pregrid.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/q_time_pregrid.py new file mode 100644 index 000000000..ace29fa3d --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/q_time_pregrid.py @@ -0,0 +1,376 @@ +"""Pipeline-side validation and emission-refusal for --q-time-pregrid-factor. + +WHAT THIS IS. bin/integrate_likelihood_extrinsic_batchmode carries a certified, opt-in +Q_lm pregrid (PR #261): factor 8 reflects each finite Q window, FFT-interpolates it onto +an 8x finer grid once after packing, and evaluates detector arrival times off that dense +grid with four-tap cubic interpolation, while leaving the geocentric time-integration grid +at the data deltaT. Factor 1 (the default) is the historical, unchanged path. The driver +enforces this itself, at first-job time: + + if opts.q_time_pregrid_factor not in Q_TIME_PREGRID_CHOICES: # imported from this module + raise ValueError(...) + if opts.q_time_pregrid_factor == 8: + if not opts.vectorized or opts.rotation_slow or opts.freqresponse or opts.calibration_envelope_directory: + raise NotImplementedError(...) + if not opts._interp_time_from_default and opts._noloop_time_interp != "cubic": + raise ValueError(...) + +This module is the pipeline-side MIRROR of that guard -- the same discipline +RIFT.likelihood.time_marginalization_quadrature applies to --time-marginalization-quadrature +-- so a workflow build refuses an unhonourable request before a whole queue-slot cycle is +spent discovering it at the driver. The choice tuple is defined HERE and the driver imports +it (rather than retyping its own), so the pipeline-side and driver-side legal-factor sets +cannot silently drift apart -- this used to be a docstring claim the driver did not honour +(PR #291 review, MAJOR #2: the driver hard-coded ``(1, 8)`` independently and a parity test +widening it to accept factor 4 passed 9/9); it is now enforced by the import itself, not by +convention. The exact conflict wording (STENCIL_CONFLICT_MESSAGE below) is reproduced from +the driver's raise rather than imported, since it is a string inside a conditional the driver +never delegates. + +THE ONE ENTANGLEMENT. Factor 8 forces the arrival-time stencil to cubic. That is silent +and harmless when the caller never named a stencil (the driver's own default path), but it +is refused when the caller passed an EXPLICIT --interpolate-time that is not itself cubic -- +"remove the explicit --interpolate-time option or set it to cubic" -- because silently +overriding a stencil the user asked for by name is exactly the kind of inert-flag failure +this whole family of checks exists to rule out. + +ABBREVIATION RESOLUTION. ``--manual-extra-ile-args`` can carry any ILE flag at all +(util_RIFT_pseudo_pipe.py), including an optparse-abbreviated one. Resolving a token like +``--vec`` needs the driver's REAL option namespace, not a length floor on this module's own +few flags: a token that is a prefix of more than one driver option (e.g. ``--rotation``, +which prefixes --rotation-n-harmonics, --rotation-p-max, AND --rotation-slow) is one optparse +itself refuses as ambiguous, and treating it as a match to just one of them here mis-attributed +that refusal (PR #291 review, MAJOR #1). _resolve_driver_option below reads the driver's +add_option calls with ``ast`` (no import: the driver is a top-level script with no __main__ +guard) to resolve abbreviations, and reject exactly the tokens optparse itself would reject. +""" + +import ast +import os + +Q_TIME_PREGRID_CHOICES = (1, 8) + +ILE_Q_TIME_PREGRID_FLAG = '--q-time-pregrid-factor' +ILE_INTERPOLATE_TIME_FLAG = '--interpolate-time' +ILE_CALIBRATION_ENVELOPE_DIRECTORY_FLAG = '--calibration-envelope-directory' + +# Legacy --interpolate-time spellings the ILE driver itself still accepts (see +# bin/integrate_likelihood_extrinsic_batchmode's _TI_LEGACY_BOOLEAN): a truthy value meant +# 'cubic', a falsy one meant 'nearest'. Reproduced here only so a hand-passed +# --manual-extra-ile-args using the legacy spelling is not misread as "no stencil requested". +_LEGACY_TRUTHY = ("1", "true", "t", "yes", "y", "on") +_LEGACY_FALSY = ("0", "false", "f", "no", "n", "off", "none") + +# The driver's own wording (bin/integrate_likelihood_extrinsic_batchmode, q_time_pregrid_factor +# == 8 branch), reproduced VERBATIM so a workflow-build-time refusal and the driver's own +# first-job refusal read identically. +STENCIL_CONFLICT_MESSAGE = ( + "--q-time-pregrid-factor 8 uses four-tap cubic interpolation; remove the " + "explicit --interpolate-time option or set it to cubic") + +_PIPELINE_REQUIRED_ILE_FLAGS = ( + ('--vectorized', + 'the driver restricts --q-time-pregrid-factor 8 to ordinary vectorized NoLoop'), +) +# Pure boolean flags: the driver reads these as opts.rotation_slow / opts.freqresponse +# directly, so presence on the command line is exactly the excluded condition. +_PIPELINE_EXCLUDING_ILE_FLAGS = ( + ('--rotation-slow', + 'the driver restricts --q-time-pregrid-factor 8 to ordinary vectorized NoLoop without rotation'), + ('--freqresponse', + 'the driver restricts --q-time-pregrid-factor 8 to ordinary vectorized NoLoop without ' + 'frequency-dependent response'), +) +# --calibration-envelope-directory takes a VALUE, and the driver's own guard +# (bin/integrate_likelihood_extrinsic_batchmode:509, and the same idiom at lines 774/880/920) +# reads it as `opts.calibration_envelope_directory` truthiness, not presence: an empty string +# is falsy, so the driver treats `--calibration-envelope-directory ""` as "not set" and does +# NOT refuse. Matching that (rather than refusing on token presence alone, as this module did +# before PR #281's follow-up review) keeps this module an exact mirror of the driver instead of +# a stricter one. +_PIPELINE_EXCLUDING_VALUE_ILE_FLAGS = ( + (ILE_CALIBRATION_ENVELOPE_DIRECTORY_FLAG, + 'the driver restricts --q-time-pregrid-factor 8 to ordinary vectorized NoLoop without ' + 'calibration marginalization'), +) + + +def validate_q_time_pregrid_factor(value): + """Return the canonical int factor, or raise ValueError. + + Mirrors the driver's own ``if opts.q_time_pregrid_factor not in (1, 8): raise`` exactly, + so this module and the driver can never disagree about the legal set. + """ + try: + factor = int(value) + except (TypeError, ValueError): + raise ValueError( + "--q-time-pregrid-factor must be an integer, got %r" % (value,)) + if factor not in Q_TIME_PREGRID_CHOICES: + raise ValueError( + "--q-time-pregrid-factor currently accepts only %s, got %r (same restriction as " + "bin/integrate_likelihood_extrinsic_batchmode)." + % ("|".join(str(c) for c in Q_TIME_PREGRID_CHOICES), factor)) + return factor + + +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. Copied from + RIFT.likelihood.time_marginalization_quadrature rather than imported, so this leaf module + has no dependency on that one's numpy/scipy-facing internals. + """ + 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 + + +_HERE = os.path.dirname(os.path.abspath(__file__)) +_DRIVER_PATH = os.path.normpath(os.path.join( + _HERE, '..', '..', 'bin', 'integrate_likelihood_extrinsic_batchmode')) + +_driver_long_option_names_cache = None + + +def _driver_long_option_names(): + """Every long option string ('--foo-bar') the ILE driver's optparse parser registers, + read from source with ``ast`` (the driver is a top-level script with no ``__main__`` + guard, so importing it would run the whole thing). Cached: the driver source does not + change within a process, and a DAG build calls this module's guards many times. + """ + global _driver_long_option_names_cache + if _driver_long_option_names_cache is None: + with open(_DRIVER_PATH) as handle: + tree = ast.parse(handle.read(), filename=_DRIVER_PATH) + names = set() + 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 + for arg in node.args: + if (isinstance(arg, ast.Constant) and isinstance(arg.value, str) + and arg.value.startswith('--')): + names.add(arg.value) + _driver_long_option_names_cache = frozenset(names) + return _driver_long_option_names_cache + + +def _resolve_driver_option(token): + """Resolve ``token`` against the driver's real long-option namespace the way optparse's + own ``_match_abbrev`` would (cpython Lib/optparse.py): an exact match wins outright; a + token that is a prefix of exactly one registered long option resolves to it; a token that + is a prefix of MORE than one is ambiguous -- the driver refuses it with + 'ambiguous option: ...' before any of this module's own guards ever run. + + Returns ``(resolved_name, candidates)``. ``resolved_name`` is the flag ``token`` stands + for, or ``None`` if it stands for nothing recognisable. ``candidates`` is a non-empty + sorted tuple -- matching optparse's own ``possibilities.sort()`` -- only when ``token`` is + ambiguous. + """ + names = _driver_long_option_names() + if token in names: + return token, () + if not (token.startswith('--') and len(token) > 2): + return None, () + candidates = sorted(name for name in names if name.startswith(token)) + if len(candidates) == 1: + return candidates[0], () + if len(candidates) > 1: + return None, tuple(candidates) + return None, () + + +def _matches(flag, token): + """True if ``token`` is ``flag`` exactly, or an UNAMBIGUOUS optparse abbreviation of it. + + Resolves ``token`` against the driver's real option namespace (``_resolve_driver_option``), + not against ``flag`` alone: the previous version checked only whether ``token`` was a + prefix of THIS ONE flag, so an ambiguous token matched here regardless of how many driver + options it actually prefixes, misattributing the driver's own 'ambiguous option' refusal to + whichever flag this module happened to be checking (PR #291 review, MAJOR #1). An + ambiguous token now resolves to nothing here; see ``_ambiguous_abbreviation_refusals``, + which surfaces the ambiguity itself instead of silently absorbing it into one guard. + """ + if token == flag: + return True + resolved, _candidates = _resolve_driver_option(token) + return resolved == flag + + +def _ambiguous_abbreviation_refusals(toks): + """Every 'ambiguous option' refusal ``toks`` would trigger from the driver's own optparse, + in first-occurrence order. Wording matches ``str(optparse.AmbiguousOptionError)`` exactly: + 'ambiguous option: TOKEN (CANDIDATE, CANDIDATE?)', candidates sorted. A token this + ambiguous makes the driver refuse before any of the prerequisite guards below even run, so + reporting it here keeps this module's classification of "why the driver would refuse" + accurate instead of blaming an unrelated guard (PR #291 review, MAJOR #1). + """ + out = [] + seen = set() + for t in toks: + if t in seen or not (t.startswith('--') and len(t) > 2): + continue + seen.add(t) + _resolved, candidates = _resolve_driver_option(t) + if candidates: + out.append("ambiguous option: {} ({}?)".format(t, ", ".join(candidates))) + return out + + +def find_q_time_pregrid_in_ile_args(ile_args): + """Every value given to ``--q-time-pregrid-factor`` in ``ile_args``, in order.""" + toks = _ile_tokens(ile_args) + out = [] + for n, t in enumerate(toks): + if _matches(ILE_Q_TIME_PREGRID_FLAG, t): + out.append(toks[n + 1] if n + 1 < len(toks) else None) + return out + + +def find_interpolate_time_in_ile_args(ile_args): + """Every value given to ``--interpolate-time`` in ``ile_args``, in order.""" + toks = _ile_tokens(ile_args) + out = [] + for n, t in enumerate(toks): + if _matches(ILE_INTERPOLATE_TIME_FLAG, t): + out.append(toks[n + 1] if n + 1 < len(toks) else None) + return out + + +def find_calibration_envelope_directory_in_ile_args(ile_args): + """Every value given to ``--calibration-envelope-directory`` in ``ile_args``, in order.""" + toks = _ile_tokens(ile_args) + out = [] + for n, t in enumerate(toks): + if _matches(ILE_CALIBRATION_ENVELOPE_DIRECTORY_FLAG, t): + out.append(toks[n + 1] if n + 1 < len(toks) else None) + return out + + +# Dispatch for _PIPELINE_EXCLUDING_VALUE_ILE_FLAGS: routes each value-taking excluded flag to +# its own find_*_in_ile_args helper, rather than reimplementing the same last-occurrence scan +# inline. find_calibration_envelope_directory_in_ile_args was added (mirroring the two find_* +# helpers above it) but never wired in here -- dead code (PR #291 review, MINOR #3). +_VALUE_FLAG_FINDERS = { + ILE_CALIBRATION_ENVELOPE_DIRECTORY_FLAG: find_calibration_envelope_directory_in_ile_args, +} + + +def _resolve_stencil_token(value): + """Canonical stencil name for an --interpolate-time VALUE, or None if unrecognised. + + Mirrors the driver's own resolution (nearest|cubic|sinc verbatim, or a legacy boolean). + An unrecognised spelling is left for the driver's own parser to reject; it is not this + module's job to duplicate that error. + """ + v = str(value).strip().lower() + if v in ("nearest", "cubic", "sinc"): + return v + if v in _LEGACY_TRUTHY: + return "cubic" + if v in _LEGACY_FALSY: + return "nearest" + return None + + +def q_time_pregrid_pipeline_prereqs(factor, ile_args): + """Missing/violated prerequisites for ``factor`` in an ILE argument string. + + ``ile_args`` is the assembled ILE command line the workflow is about to write + (``args_ile.txt`` / ``helper_ile_args.txt``). Returns a list of human-readable reasons; + empty means the configuration can honour the request. Factor 1 -- the default -- always + returns an empty list, since it is what ILE does anyway. + """ + factor = validate_q_time_pregrid_factor(factor) + if factor == 1: + return [] + toks = _ile_tokens(ile_args) + missing = [] + # A token this ambiguous makes the driver refuse via optparse itself, before it ever + # reaches the prerequisite checks below -- report that refusal AS ambiguity rather than + # letting one of the checks below misattribute it (PR #291 review, MAJOR #1). + missing.extend(_ambiguous_abbreviation_refusals(toks)) + for flag, why in _PIPELINE_REQUIRED_ILE_FLAGS: + 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)) + for flag, why in _PIPELINE_EXCLUDING_VALUE_ILE_FLAGS: + values = _VALUE_FLAG_FINDERS[flag](ile_args) + # optparse takes the LAST occurrence, matching find_calibration_envelope_directory_in_ + # ile_args / find_interpolate_time_in_ile_args elsewhere in this module. + if values and values[-1]: + missing.append("incompatible {} (value {!r}) ({})".format(flag, values[-1], why)) + interp_values = find_interpolate_time_in_ile_args(ile_args) + if interp_values: + # optparse takes the LAST occurrence, so that is the one the driver will actually see. + resolved = _resolve_stencil_token(interp_values[-1]) + if resolved is not None and resolved != "cubic": + missing.append(STENCIL_CONFLICT_MESSAGE) + return missing + + +def refuse_unhonourable_q_time_pregrid(factor, ile_args, where): + """Raise unless ``ile_args`` can honour ``factor``. + + The raise lives HERE, not at the call sites, so it is executable in a unit test: both + pipeline scripts are top-level scripts that need real data before they reach their guard. + """ + missing = q_time_pregrid_pipeline_prereqs(factor, ile_args) + if missing: + raise ValueError( + "--q-time-pregrid-factor {!r} was requested, but {} cannot honour it: {}. " + "Refusing rather than running the historical factor=1 grid while reporting that " + "you asked for something else.".format(factor, where, "; ".join(missing))) + + +def refuse_unless_q_time_pregrid_emitted(factor, ile_args, where): + """Raise unless the REQUESTED factor is the one the bytes actually carry. + + ``factor`` of ``None`` or ``1`` means "nothing forced": the flag may legitimately be + absent (the pipeline option was never set) or may equal the historical default. If + something is on the line anyway -- --manual-extra-ile-args, or an ini -- it is validated + and prerequisite-checked exactly like a pipeline-driven request, which is the same + "hold a hand-passed value to the same standard" discipline + RIFT.likelihood.time_marginalization_quadrature.refuse_unless_time_quadrature_emitted uses. + """ + found = find_q_time_pregrid_in_ile_args(ile_args) + if len(found) > 1: + raise ValueError( + "{} carries {} occurrences of {} ({!r}). optparse takes the LAST, so the factor " + "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_Q_TIME_PREGRID_FLAG, found)) + if factor is None or int(factor) == 1: + if found: + refuse_unhonourable_q_time_pregrid( + validate_q_time_pregrid_factor(found[0]), ile_args, where) + return + factor = validate_q_time_pregrid_factor(factor) + if not found: + raise ValueError( + "--q-time-pregrid-factor {!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 the historical factor=1 grid.".format( + factor, where, ILE_Q_TIME_PREGRID_FLAG)) + found_val = validate_q_time_pregrid_factor(found[0]) + if found_val != factor: + raise ValueError( + "--q-time-pregrid-factor {!r} was requested but {} carries {!r}. " + "Refusing.".format(factor, where, found[0])) + refuse_unhonourable_q_time_pregrid(factor, ile_args, where) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/slowrot_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/slowrot_freqresponse.py index cd0fb567d..1cf3f1932 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/slowrot_freqresponse.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/slowrot_freqresponse.py @@ -392,9 +392,86 @@ def finite_size_geometry(det, ra, dec, psi, gmst=0.0, L_arm=None): F0=complex(Fp_lwl) + 1j * complex(Fc_lwl)) +_DETECTOR_GEOMETRY_CACHE = {} + + +def detector_geometry_cached(det, L_arm=None): + """``detector_geometry`` memoized on (det, L_arm). ARRAYS ARE READ-ONLY. + + The geometry depends on the detector and its arm length only. The likelihood + evaluates the response coefficients once per Monte Carlo sample, so calling + ``detector_geometry`` from there re-did the LAL detector lookup and the arm + trigonometry for every sample: 1.4e6 times in one profiled ILE run, for five + distinct answers. + + The cached arrays are shared by every caller and are marked non-writeable so a + caller that mutates one fails loudly instead of corrupting the next. Callers that + need to write should use ``detector_geometry``, which is unchanged and uncached. + """ + key = (str(det), None if L_arm is None else float(L_arm)) + got = _DETECTOR_GEOMETRY_CACHE.get(key) + if got is None: + response, x_arm, y_arm, L = detector_geometry(det, L_arm=L_arm) + response = np.array(response, dtype=float) + x_arm = np.array(x_arm, dtype=float) + y_arm = np.array(y_arm, dtype=float) + for a in (response, x_arm, y_arm): + a.setflags(write=False) + got = (response, x_arm, y_arm, float(L)) + _DETECTOR_GEOMETRY_CACHE[key] = got + return got + + +def finite_size_geometry_vector(det, ra, dec, psi, gmst=0.0, L_arm=None): + """``finite_size_geometry`` for a BLOCK of sky samples. + + ra, dec, psi are broadcastable arrays of the same shape; the returned ax, ay, zx, zy + and F0 carry that shape, while T and L stay scalar (they depend on the detector only). + Same algebra, same order of operations, as the scalar routine -- the only difference is + that the triad, the long-wavelength contraction and the arm projections are evaluated + for the whole block at once. + + ``finite_size_beta`` accepts the result unchanged and returns (Qmax+1,)+shape. + """ + response, x_arm, y_arm, L = detector_geometry_cached(det, L_arm=L_arm) + ra = np.asarray(ra, dtype=float) + dec = np.asarray(dec, dtype=float) + psi = np.asarray(psi, dtype=float) + g = gmst - ra + X, Y, nhat = _triad(dec, psi, g) # (..., 3) + Fp_lwl, Fc_lwl = _lwl_response(response, X, Y) + Xx = np.einsum('...i,i->...', X, x_arm) + Yx = np.einsum('...i,i->...', Y, x_arm) + Xy = np.einsum('...i,i->...', X, y_arm) + Yy = np.einsum('...i,i->...', Y, y_arm) + return dict(T=L / C_SI, L=L, + ax=np.einsum('...i,i->...', nhat, x_arm), + ay=np.einsum('...i,i->...', nhat, y_arm), + zx=Xx + 1j * Yx, zy=Xy + 1j * Yy, + F0=Fp_lwl + 1j * Fc_lwl) + + +def _csquare(z): + """z*z written out in real arithmetic: (zr^2 - zi^2) + i(zr zi + zi zr). + + This is CPython's own complex multiply, term for term, so it reproduces the Python + scalar ``z ** 2`` BIT FOR BIT -- which numpy's `complex128 ** 2` does not, because its + SIMD complex loop rounds differently (measured: ~29% of random samples differ by 1 ulp). + Written out here so the vectorized and scalar geometry paths give the same beta_0, + rather than two answers that differ in the last bit. + """ + zr, zi = z.real, z.imag + return (zr * zr - zi * zi) + 1j * (zr * zi + zi * zr) + + def finite_size_beta(geom, Qmax): - """Analytic sky/pol coefficients beta_q = (1/2)[zx^2 a_x^q - zy^2 a_y^q], q=0..Qmax.""" - zx2, zy2, ax, ay = geom['zx'] ** 2, geom['zy'] ** 2, geom['ax'], geom['ay'] + """Analytic sky/pol coefficients beta_q = (1/2)[zx^2 a_x^q - zy^2 a_y^q], q=0..Qmax. + + ``geom`` may come from ``finite_size_geometry`` (scalars) or from + ``finite_size_geometry_vector`` (arrays over a block of sky samples); the returned + array has shape (Qmax+1,) + the sample shape. + """ + zx2, zy2, ax, ay = _csquare(geom['zx']), _csquare(geom['zy']), geom['ax'], geom['ay'] return np.array([0.5 * (zx2 * ax ** q - zy2 * ay ** q) for q in range(Qmax + 1)], dtype=complex) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_precompute_crossterm_batching.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_precompute_crossterm_batching.py new file mode 100644 index 000000000..3537bc34b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_precompute_crossterm_batching.py @@ -0,0 +1,219 @@ +"""Precompute cross-term batching: the bit-identical fixes stay bit-identical, and the +opt-in batched path stays off by default and agrees with the loop it replaces. + +Companion to DESIGN_precompute_crossterm_batching.md. Three of the four changes under +test claim BIT identity, not approximate agreement, so each is checked against a replay +of the exact code it replaced rather than against a tolerance. +""" +import os + +import numpy as np +import pytest + +lal = pytest.importorskip("lal") + +import RIFT.lalsimutils as lsu +from RIFT.likelihood import factored_likelihood as FL + +FNYQ = 512.0 +DELTAF = 1.0 / 8.0 +FMIN, FMAX = 20.0, 256.0 +LEN1 = int(FNYQ / DELTAF) + 1 +LEN2 = 2 * (LEN1 - 1) +TSPEC = 1.0 # N_spec = 1024 < LEN2/2 = 4096 + + +def _psd_array(dead_bin=True, negative_bin=False): + f = np.arange(LEN1) * DELTAF + psd = np.zeros(LEN1) + b = (f >= FMIN) & (f <= FMAX) + psd[b] = 1e-46 * ((f[b] / 100.0) ** -4.14 + 2.0 + 0.5 * (f[b] / 100.0) ** 2) + if dead_bin: + psd[300] = 0.0 # a dead bin inside the band + if negative_bin: + psd[400] = -1e-46 # pins `!= 0` (the shipped loop) against `> 0` + return psd + + +def _psd_series(arr): + s = lal.CreateREAL8FrequencySeries("psd", lal.LIGOTimeGPS(0.0), 0.0, DELTAF, + lsu.lsu_HertzUnit, LEN1) + s.data.data[:] = arr + return s + + +def _series(n, seed): + rng = np.random.default_rng(seed) + out = [] + for _ in range(n): + s = lal.CreateCOMPLEX16FrequencySeries("h", lal.LIGOTimeGPS(0.0), 0.0, DELTAF, + lsu.lsu_DimensionlessUnit, LEN2) + s.data.data[:] = (rng.standard_normal(LEN2) + 1j * rng.standard_normal(LEN2)) * 1e-23 + out.append(s) + return out + + +# -------------------------------------------------------------------------------------- +# bit-identity of the three loop -> vector rewrites +# -------------------------------------------------------------------------------------- +def test_array_psd_weights_bit_identical(): + """The vectorised array-PSD fill reproduces the per-bin loop exactly, INCLUDING its + `!= 0` mask -- the REAL8FrequencySeries branch uses `> 0`, which would silently drop a + negative bin, and the two branches must not be conflated.""" + psd = _psd_array(negative_bin=True) + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, psd, False, False, 0.0) + ref = np.zeros(LEN1) + for i in range(ip.minIdx, ip.maxIdx): + if psd[i] != 0.0: + ref[i] = 1.0 / psd[i] * 1.0 + assert np.array_equal(ip.weights, ref) + assert ip.weights[400] < 0, "negative PSD bin must survive the `!= 0` mask" + + +def test_array_psd_weights_bit_identical_psi4(): + psd = _psd_array() + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, psd, False, False, 0.0, + waveform_is_psi4=True) + ref = np.zeros(LEN1) + for i in range(ip.minIdx, ip.maxIdx): + if psd[i] != 0.0: + ew = 1.0 / (2 * np.pi * i * DELTAF) / (2 * np.pi * i * DELTAF) + ref[i] = 1.0 / psd[i] * ew + assert np.array_equal(ip.weights, ref) + + +def test_inv_spec_trunc_weights_bit_identical(): + """The slice-assignment zeroing reproduces the per-element SWIG loop exactly.""" + ser = _psd_series(_psd_array()) + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, ser, False, True, TSPEC) + + w = np.zeros(LEN1) + iv = np.arange(ip.minIdx, ip.maxIdx) + ok = iv[ser.data.data[iv] > 0] + w[ok] = 1.0 / ser.data.data[ok] * np.ones(LEN1)[ok] + n_spec = int(TSPEC / ip.deltaT) + wfd = lal.CreateCOMPLEX16FrequencySeries("w", lal.LIGOTimeGPS(0.0), 0.0, DELTAF, + lsu.lsu_DimensionlessUnit, ip.len1side) + wtd = lal.CreateREAL8TimeSeries("w", lal.LIGOTimeGPS(0.0), 0.0, ip.deltaT, + lsu.lsu_DimensionlessUnit, ip.len2side) + fwd = lal.CreateForwardREAL8FFTPlan(ip.len2side, 0) + rev = lal.CreateReverseREAL8FFTPlan(ip.len2side, 0) + wfd.data.data[:] = np.sqrt(w) + wfd.data.data[0] = wfd.data.data[-1] = 0.0 + lal.REAL8FreqTimeFFT(wtd, wfd, rev) + for i in range(int(n_spec / 2), ip.len2side - int(n_spec / 2)): + wtd.data.data[i] = 0.0 + lal.REAL8TimeFreqFFT(wfd, wtd, fwd) + wfd.data.data[0] = wfd.data.data[-1] = 0.0 + assert np.array_equal(ip.weights, np.abs(wfd.data.data * wfd.data.data)) + + +def test_ip_bit_identical_without_epoch_differences(): + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, _psd_array(), False, False, 0.0) + h1, h2 = _series(2, 11) + shipped = np.sum(np.conj(h1.data.data) * h2.data.data + * np.ones(LEN2) * ip.weights2side) * 2.0 * ip.deltaF + assert ip.ip(h1, h2) == shipped + + +def test_ip_epoch_difference_path_still_reached(): + """The `include_epoch_differences` branch must still apply a phase, not silently no-op.""" + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, _psd_array(), False, False, 0.0) + h1, h2 = _series(2, 12) + h1.epoch = lal.LIGOTimeGPS(0.25) + plain = ip.ip(h1, h2) + shifted = ip.ip(h1, h2, include_epoch_differences=True) + assert plain != shifted + + +# -------------------------------------------------------------------------------------- +# band support +# -------------------------------------------------------------------------------------- +def test_band_support_matches_actual_nonzeros(): + """Recorded support must be read off the weights, never assumed from (fmin, fMax): + inverse spectrum truncation -- ON by default in the driver -- smears the band to full + support, and a batched path that skipped to [fmin, fMax] there would drop ~1e-6 of the + weight.""" + narrow = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, _psd_array(), False, False, 0.0) + nz = np.nonzero(narrow.weights2side)[0] + assert (narrow.band_lo2side, narrow.band_hi2side) == (int(nz[0]), int(nz[-1]) + 1) + assert narrow.band_hi2side - narrow.band_lo2side < narrow.len2side + + trunc = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, _psd_series(_psd_array()), + False, True, TSPEC) + nz = np.nonzero(trunc.weights2side)[0] + assert (trunc.band_lo2side, trunc.band_hi2side) == (int(nz[0]), int(nz[-1]) + 1) + + +# -------------------------------------------------------------------------------------- +# the batched path +# -------------------------------------------------------------------------------------- +@pytest.mark.parametrize("trunc", [False, True]) +def test_ip_matrix_matches_ip_loop(trunc): + psd = _psd_series(_psd_array()) if trunc else _psd_array() + ip = lsu.ComplexIP(FMIN, FMAX, FNYQ, DELTAF, psd, False, trunc, TSPEC if trunc else 0.0) + A, B = _series(4, 21), _series(3, 22) + loop = np.array([[ip.ip(a, b) for b in B] for a in A]) + mat = ip.ip_matrix(A, B) + assert mat.shape == loop.shape + assert np.max(np.abs(mat - loop)) <= 1e-13 * np.max(np.abs(loop)) + + +def _hlms(keys, seed): + return dict(zip(keys, _series(len(keys), seed))) + + +MODES = [(2, -2), (2, 0), (2, 2), (3, -3)] + + +def test_crossterm_batched_matches_loop_general(): + psd = _psd_array() + a, b = _hlms(MODES, 31), _hlms(MODES, 32) + kw = dict(analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0.0, verbose=False) + ref = FL.ComputeModeCrossTermIP(a, b, psd, FMIN, FMAX, FNYQ, DELTAF, batched=False, **kw) + got = FL.ComputeModeCrossTermIP(a, b, psd, FMIN, FMAX, FNYQ, DELTAF, batched=True, **kw) + assert set(ref) == set(got) + scale = max(abs(v) for v in ref.values()) + assert max(abs(got[k] - ref[k]) for k in ref) <= 1e-13 * scale + + +@pytest.mark.parametrize("prefix", ["U", "V"]) +def test_crossterm_batched_preserves_same_waveform_symmetry(prefix): + """same_waveform_Q fills the lower triangle by mirroring the upper one. The batched + path computes the full matrix, so it must re-impose that mirror rather than keep the + independently-computed element -- otherwise an exact symmetry becomes approximate.""" + psd = _psd_array() + a = _hlms(MODES, 41) + kw = dict(analyticPSD_Q=False, inv_spec_trunc_Q=False, T_spec=0.0, verbose=False, + prefix=prefix, same_waveform_Q=True) + ref = FL.ComputeModeCrossTermIP(a, a, psd, FMIN, FMAX, FNYQ, DELTAF, batched=False, **kw) + got = FL.ComputeModeCrossTermIP(a, a, psd, FMIN, FMAX, FNYQ, DELTAF, batched=True, **kw) + assert set(ref) == set(got) + scale = max(abs(v) for v in ref.values()) + assert max(abs(got[k] - ref[k]) for k in ref) <= 1e-13 * scale + # Off-diagonal only: the shipped `pairs = combinations(mode_list, 2)` loop never + # touches the diagonal, so no mirror is claimed there. The loop result is asserted + # first, as a control on the criterion itself -- if it did not hold for the shipped + # path, holding for the batched one would mean nothing. + for i, m1 in enumerate(MODES): + for m2 in MODES[i + 1:]: + for d in (ref, got): + mirror = d[(m2, m1)] if prefix == "V" else np.conj(d[(m2, m1)]) + assert d[(m1, m2)] == mirror + + +def test_batched_is_off_unless_asked(monkeypatch): + """Default must be the shipped path. A flag nobody can see fire is a silent no-op, so + the batched path bumps a counter and this pins that it stays put.""" + monkeypatch.delenv("RIFT_PRECOMPUTE_BATCHED_CROSSTERMS", raising=False) + assert FL._crossterm_batched_default() is False + psd = _psd_array() + a, b = _hlms(MODES[:2], 51), _hlms(MODES[:2], 52) + before = FL._CROSSTERM_BATCH_CALLS[0] + FL.ComputeModeCrossTermIP(a, b, psd, FMIN, FMAX, FNYQ, DELTAF, verbose=False) + assert FL._CROSSTERM_BATCH_CALLS[0] == before + + monkeypatch.setenv("RIFT_PRECOMPUTE_BATCHED_CROSSTERMS", "1") + assert FL._crossterm_batched_default() is True + FL.ComputeModeCrossTermIP(a, b, psd, FMIN, FMAX, FNYQ, DELTAF, verbose=False) + assert FL._CROSSTERM_BATCH_CALLS[0] == before + 1 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_time_pregrid.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_time_pregrid.py new file mode 100644 index 000000000..365a0fbf3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_q_time_pregrid.py @@ -0,0 +1,216 @@ +#!/usr/bin/env python3 +# RIFT-CI-GATE: q-window-stencil +"""Focused tests for the opt-in reflected Q pregrid.""" + +import numpy as np +import gc +import weakref +from types import SimpleNamespace +from unittest.mock import patch + +from RIFT.likelihood.factored_likelihood import ( + _cubic_Q_window_numpy, + _q_inner_product_explicit_times, + _q_sample_positions, + build_reflected_q_pregrid, + prepare_reflected_q_pregrid, +) +from RIFT.likelihood import factored_likelihood as fl +from RIFT.likelihood import time_marginalization_quadrature as tmq + +try: + import cupy + HAVE_GPU = cupy.cuda.runtime.getDeviceCount() > 0 +except Exception: + cupy = None + HAVE_GPU = False + + +def test_reflected_pregrid_roundtrip_odd_even_and_size(): + rng = np.random.RandomState(811) + for n_time in (31, 32): + coarse = rng.normal(size=(3, n_time)) + 1j*rng.normal(size=(3, n_time)) + fine, report = build_reflected_q_pregrid(coarse, factor=8) + assert fine.shape == (3, (n_time - 1)*8 + 1) + np.testing.assert_allclose(fine[..., ::8], coarse, rtol=5e-13, atol=5e-13) + assert report['factor'] == 8 + assert report['output_bytes'] == fine.nbytes + assert report['retained_bytes'] == fine.nbytes + assert report['peak_allocation_bytes'] > fine.nbytes + assert fine.flags.owndata + assert fine.base is None + + +def test_backend_oom_rolls_back_whole_dictionary_and_cleans_up(): + original = {'H1': np.ones((2, 9)), 'L1': np.ones((2, 9))*2} + calls = [] + cleaned = [] + expanded_refs = [] + + def transfer(value): + calls.append(value.shape[-1]) + if value.shape[-1] == 65: + expanded_refs.append(weakref.ref(value)) + if calls == [65, 65]: + raise MemoryError('forced device OOM') + return np.array(value, copy=True) + + got, reports, error = prepare_reflected_q_pregrid( + original, factor=8, transfer=transfer, cleanup=lambda: cleaned.append(True)) + assert error == {'type': 'MemoryError', 'repr': "MemoryError('forced device OOM')"} + assert reports == [] + assert cleaned == [True] + assert calls == [65, 65, 9, 9] + for det in original: + np.testing.assert_array_equal(got[det], original[det]) + gc.collect() + assert all(reference() is None for reference in expanded_refs) + + +def test_reflection_is_load_bearing_at_both_nonperiodic_edges(): + # A smooth finite-window ramp has deliberately unlike endpoints. Direct + # periodic interpolation joins them and rings; even reflection preserves + # the local continuation at both edges. This test fails if reflection is + # mutated to direct periodic upsampling. + n = 64 + factor = 8 + x = np.linspace(-1.0, 1.0, n) + coarse = (x + 0.15*x**2)[None, :] + direct = tmq.bandlimited_upsample(coarse, factor)[0] + reflected, _ = build_reflected_q_pregrid(coarse, factor=factor) + dense_x = np.linspace(-1.0, 1.0, (n - 1)*factor + 1) + truth = dense_x + 0.15*dense_x**2 + edge = np.r_[1:factor, len(truth)-factor:len(truth)-1] + reflected_error = np.max(np.abs(reflected[0, edge] - truth[edge])) + direct_error = np.max(np.abs(direct[edge] - truth[edge])) + assert reflected_error < 0.2*direct_error, (reflected_error, direct_error) + + +def test_separate_grid_refuses_unimplemented_stencils(): + p = SimpleNamespace(deltaT=1.0, q_deltaT=0.125, phi=np.array([0.0]), + theta=np.array([0.0]), phiref=np.array([0.0]), + incl=np.array([0.0]), psi=np.array([0.0]), + dist=np.array([fl.distMpcRef*1e6*fl.lal.PC_SI]), tref=0.0) + args = (np.arange(2.0), p, {}, {}, {}, {}, {}) + for stencil in ('nearest', 'sinc'): + try: + fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + *args, time_interp=stencil, return_lnLt=True) + except NotImplementedError: + pass + else: + raise AssertionError('%s silently accepted a separate Q spacing' % stencil) + + +def test_separate_q_spacing_preserves_coarse_integration_nodes(): + t_det = np.array([10.25, 11.5]) + tvals = np.arange(7)*0.25 - 0.5 + starts, fractions, per_time, stride = _q_sample_positions( + t_det, tvals, 0.25, 0.25/8, 'cubic', False) + assert not per_time + assert stride == 8 + target = (t_det + tvals[0])/(0.25/8) + np.testing.assert_array_equal(starts, np.floor(target).astype(np.int32)) + np.testing.assert_allclose(fractions, target - np.floor(target)) + # The geocentric nodes are still separated by the original 0.25 seconds; + # only their coordinates on Q advance by eight samples. + grid = np.arange(200, dtype=float) + q = (grid**3 - 2*grid + 1).astype(complex)[:, None] + got = _cubic_Q_window_numpy(q, np.array([20]), np.array([0.25]), 7, + time_stride=stride)[0, :, 0] + x = 20.25 + np.arange(7)*8 + np.testing.assert_allclose(got, x**3 - 2*x + 1, rtol=2e-13) + + +def test_factor_one_keeps_historical_scalar_window_gather(): + starts, fractions, per_time, stride = _q_sample_positions( + np.array([4.25, 8.75]), np.arange(5)*0.5 - 1.0, + 0.5, 0.5, 'cubic', False) + assert not per_time + assert stride == 1 + assert starts.shape == (2,) + expected_samples = (np.array([4.25, 8.75]) - 1.0)/0.5 + np.testing.assert_allclose(fractions, expected_samples - np.floor(expected_samples)) + + +def test_cubic_explicit_gather_matches_cubic_truth_and_zero_extends_edges(): + # A cubic polynomial is reproduced exactly by the four-tap stencil. + grid = np.arange(20, dtype=float) + q = (grid**3 - 2*grid**2 + 0.5*grid + 3).astype(complex)[:, None] + starts = np.array([[4, 8, 12]], dtype=np.int32) + fractions = np.array([[0.125, 0.5, 0.875]]) + amplitude = np.array([[2.0 - 0.25j]]) + got = _q_inner_product_explicit_times( + q, amplitude, starts, fractions, 'cubic', xpy=np) + x = starts + fractions + truth = amplitude[0, 0]*(x**3 - 2*x**2 + 0.5*x + 3) + np.testing.assert_allclose(got, truth, rtol=2e-13, atol=2e-12) + + # Far outside the captured Q interval every tap is unavailable: fail closed + # to zero rather than wrapping reflected-pregrid samples across an edge. + outside = _q_inner_product_explicit_times( + q, amplitude, np.array([[-10, 30]], dtype=np.int32), + np.array([[0.5, 0.5]]), 'cubic', xpy=np) + np.testing.assert_array_equal(outside, 0.0) + + +def _phase_noloop(q_rows, q_delta_t, stride, fractional): + n_time = q_rows.shape[-1] + start = 8 + integration_dt = q_delta_t*stride + t_det = (start + fractional)*q_delta_t + p = SimpleNamespace( + deltaT=integration_dt, q_deltaT=q_delta_t, + phi=np.array([0.1]), theta=np.array([0.2]), + phiref=np.array([0.3]), incl=np.array([0.4]), psi=np.array([0.5]), + dist=np.array([fl.distMpcRef*1e6*fl.lal.PC_SI]), tref=0.0) + y = np.array([[1.2 + 0.4j, -0.7 + 0.2j]]) + response = np.array([0.8 - 0.3j]) + tvals = np.arange(3)*integration_dt + lookup = {'H1': np.array([[2, 2], [2, -2]])} + rho = {'H1': q_rows} + zeros = {'H1': np.zeros((2, 2), dtype=complex)} + epochs = {'H1': 0.0} + with patch.object(fl, '_detector_geometry', return_value=(None, None)), \ + patch.object(fl, 'SourcePolarizationBasis', return_value=(None, None)), \ + patch.object(fl, 'SourcePropagationDirection', return_value=None), \ + patch.object(fl, 'ComputeDetAMResponsePrecomputed', return_value=response), \ + patch.object(fl, 'TimeDelayFromEarthCenterPrecomputed', + return_value=np.array([t_det])), \ + patch.object(fl, 'SphericalHarmonicsVectorized', return_value=y.copy()): + got = fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + tvals, p, lookup, rho, zeros, zeros, epochs, Lmax=2, xpy=np, + return_lnLt=True, phase_marginalization=True, time_interp='cubic') + q_block = np.column_stack((q_rows[0], np.conj(q_rows[1]))) + sampled = _cubic_Q_window_numpy( + q_block, np.array([start]), np.array([fractional]), 3, + time_stride=stride)[0] + y_phase = y.copy(); y_phase[:, 1] = np.conj(y_phase[:, 1]) + factors = np.array([[response[0], np.conj(response[0])]])*y_phase + expected = np.abs(np.einsum('ti,i->t', sampled, np.conj(factors[0]))) + np.testing.assert_allclose(got[0], expected, rtol=2e-13, atol=2e-13) + + +def test_cpu_phase_marginalization_scalar_and_pregrid_match_reference(): + grid = np.arange(40.0) + coarse = np.vstack((np.exp(0.08j*grid), (1 + 0.01*grid)*np.exp(-0.05j*grid))) + _phase_noloop(coarse, 1.0, 1, 0.25) + fine, _ = build_reflected_q_pregrid(coarse, factor=8) + _phase_noloop(fine, 1.0/8, 8, 0.25) + + +def test_gpu_stride8_cubic_matches_cpu_at_fractional_and_edge_starts(): + if not HAVE_GPU: + import pytest + pytest.skip('CUDA device unavailable; stride-8 kernel parity is GPU-gated') + rng = np.random.RandomState(91) + q = rng.normal(size=(70, 3)) + 1j*rng.normal(size=(70, 3)) + amplitude = rng.normal(size=(4, 3)) + 1j*rng.normal(size=(4, 3)) + starts = np.array([-2, 3, 58, 68], dtype=np.int32) + fractions = np.array([0.2, 0.75, 0.4, 0.9]) + cpu_q = _cubic_Q_window_numpy(q, starts, fractions, 5, time_stride=8) + expected = np.einsum('eti,ei->et', cpu_q, amplitude) + got = fl.Q_inner_product.Q_inner_product_cubic_cupy( + cupy.asarray(q), cupy.asarray(amplitude), cupy.asarray(starts), + cupy.asarray(fractions), 5, time_stride=8) + np.testing.assert_allclose(cupy.asnumpy(got), expected, rtol=2e-12, atol=2e-12) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse.py index 86c0ecb58..0fe6e319e 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse.py @@ -397,6 +397,83 @@ def test_ce_is_100x_longer_effect(): assert 30 < ratio < 250, "unexpected CE/LIGO scaling: %g" % ratio +# --------------------------------------------------------------------------- +# (E) BLOCK EVALUATION. The likelihood evaluates the response geometry once per Monte +# Carlo sample; finite_size_geometry_vector does the whole block at once, and +# detector_geometry_cached stops the sample-independent part being redone per sample. +# These pin the two properties that make that substitution safe: the answer does not +# change, and the detector lookup does not scale with the block. +# --------------------------------------------------------------------------- +def test_detector_geometry_cached_matches_and_is_readonly(): + """The memo returns the same geometry, in arrays no caller can mutate.""" + for det in DETECTORS: + for L_arm in (None, 40000.0): + r0, x0, y0, L0 = fr.detector_geometry(det, L_arm=L_arm) + r1, x1, y1, L1 = fr.detector_geometry_cached(det, L_arm=L_arm) + assert np.array_equal(r0, r1) and np.array_equal(x0, x1) and np.array_equal(y0, y1) + assert L0 == L1 + # a second call returns the SAME objects (it is a memo, not a rebuild) + r2, x2, y2, _ = fr.detector_geometry_cached(det, L_arm=L_arm) + assert r2 is r1 and x2 is x1 and y2 is y1 + for a in (r1, x1, y1): + assert not a.flags.writeable, "cached geometry must be read-only" + + +def test_geometry_vector_matches_scalar_loop(): + """finite_size_geometry_vector matches finite_size_geometry sample by sample. + + Every field is a rotation or a contraction of (ra, dec, psi) with sample-independent + detector vectors. NumPy may evaluate the vector and scalar paths in a different order, + so require agreement at machine precision rather than bit-for-bit identity. + """ + rng = np.random.RandomState(20260909) + n = 200 + ra = rng.uniform(0, 2 * np.pi, n) + dec = np.arcsin(rng.uniform(-1, 1, n)) + psi = rng.uniform(0, np.pi, n) + gmst = 1.234 + for det, L_arm in (("H1", None), ("K1", 40000.0), ("V1", 10000.0)): + gv = fr.finite_size_geometry_vector(det, ra, dec, psi, gmst=gmst, L_arm=L_arm) + assert gv['T'] == fr.finite_size_geometry(det, ra[0], dec[0], psi[0], + gmst=gmst, L_arm=L_arm)['T'] + for i in range(n): + gs = fr.finite_size_geometry(det, ra[i], dec[i], psi[i], gmst=gmst, L_arm=L_arm) + for key in ('ax', 'ay', 'zx', 'zy', 'F0'): + np.testing.assert_allclose( + gv[key][i], gs[key], rtol=8 * np.finfo(float).eps, + atol=8 * np.finfo(float).eps, + err_msg="%s sample %d field %s" % (det, i, key)) + + +def test_beta_block_matches_scalar_beta(): + """finite_size_beta on a block reproduces the per-sample beta_q. + + The scalar and vector paths can differ by a few ulps because NumPy and CPython may use + different evaluation orders and power implementations. This is a last-bit rounding + difference, not a change of formula, so all orders are bounded at machine precision. + """ + rng = np.random.RandomState(11) + n, Qmax = 300, 4 + ra = rng.uniform(0, 2 * np.pi, n) + dec = np.arcsin(rng.uniform(-1, 1, n)) + psi = rng.uniform(0, np.pi, n) + det, L_arm = "K1", 40000.0 + gv = fr.finite_size_geometry_vector(det, ra, dec, psi, gmst=0.7, L_arm=L_arm) + bv = fr.finite_size_beta(gv, Qmax) + worst = 0.0 + for i in range(n): + gs = fr.finite_size_geometry(det, ra[i], dec[i], psi[i], gmst=0.7, L_arm=L_arm) + bs = fr.finite_size_beta(gs, Qmax) + for q in range(Qmax + 1): + d = abs(bv[q][i] - bs[q]) + np.testing.assert_allclose( + bv[q][i], bs[q], rtol=1e-14, + atol=8 * np.finfo(float).eps, + err_msg="sample %d beta_%d" % (i, q)) + worst = max(worst, d / max(abs(bs[q]), 1e-300)) + print("(E) beta block-vs-scalar: worst relative difference %.3e" % worst) + + if __name__ == "__main__": test_unpaired_extreme_bin_predicate() test_weights_hermitian_on_the_grid() @@ -413,5 +490,9 @@ def test_ce_is_100x_longer_effect(): print("-" * 78) test_in_band_magnitude_ligo_vs_ce() test_ce_is_100x_longer_effect() + print("-" * 78) + test_detector_geometry_cached_matches_and_is_readonly() + test_geometry_vector_matches_scalar_loop() + test_beta_block_matches_scalar_beta() print("=" * 78) print("ALL SLOWROT FREQ-RESPONSE CHECKS PASSED (worst f=0 residual %.3e)" % wA) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse_likelihood.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse_likelihood.py index d83001403..6806912f5 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse_likelihood.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_freqresponse_likelihood.py @@ -289,7 +289,70 @@ def run_strong(fmax=2000., seglen=8., SCALE_strong=40., m1=15., m2=13., Qmax=6): return HALF, base, fin, gain +def run_response_coefficients_block(): + """(V5) BLOCK COEFFICIENTS: response_coefficients_vector == the per-sample loop. + + The NoLoop likelihood used to build b_p one Monte Carlo sample at a time; it now builds + the whole block at once. Nothing else about the likelihood changed, so this is the + single place the substitution can go wrong, and the reference is the scalar routine + itself -- still shipped, still the definition of b_p. + + b_0..b_3 must be bit-identical. b_4 and b_5 carry a_x**3 and a_x**4, where numpy's + power loop and CPython's libm pow round differently in the last bit; that is bounded + here, and the end-to-end consequence was measured at zero (see + DESIGN_freqresponse_vectorized_coefficients.md). + """ + rng = np.random.RandomState(20260909) + n, Qmax, tref = 400, 4, 1e9 + RA = rng.uniform(0, 2 * np.pi, n) + DEC = np.arcsin(rng.uniform(-1, 1, n)) + PSI = rng.uniform(0, np.pi, n) + worst = 0.0 + for det, L_arm in (("H1", None), ("K1", 40000.0), ("V1", 10000.0)): + bv = flfr.response_coefficients_vector(det, RA, DEC, PSI, tref, Qmax, L_arm=L_arm) + assert sorted(bv.keys()) == list(range(Qmax + 2)) + for i in range(n): + bs = flfr.response_coefficients(det, float(RA[i]), float(DEC[i]), float(PSI[i]), + tref, Qmax, L_arm=L_arm) + for p in range(Qmax + 2): + d = abs(bv[p][i] - bs[p]) + if p <= 3: + assert d == 0.0, ("b_%d must be bit-identical to the scalar routine " + "(%s sample %d: |d|=%g)" % (p, det, i, d)) + worst = max(worst, d / max(abs(bs[p]), 1e-300)) + print("\n(V5) BLOCK COEFFICIENTS: b_0..b_3 bit-identical; worst relative " + "difference over all p = %.3e" % worst) + assert worst < 1e-14, "block coefficients drifted past one ulp: %g" % worst + + # STRUCTURAL: the detector lookup must not scale with the block. This is what the + # change is for -- a regression to the per-sample loop passes every value check above + # and only shows up here. + calls = [0] + real = sfr.detector_geometry + + def counting(*a, **k): + calls[0] += 1 + return real(*a, **k) + sfr.detector_geometry = counting + try: + sfr._DETECTOR_GEOMETRY_CACHE.clear() + flfr.response_coefficients_vector("H1", RA[:1], DEC[:1], PSI[:1], tref, Qmax) + one = calls[0] + sfr._DETECTOR_GEOMETRY_CACHE.clear() + calls[0] = 0 + flfr.response_coefficients_vector("H1", RA, DEC, PSI, tref, Qmax) + many = calls[0] + finally: + sfr.detector_geometry = real + sfr._DETECTOR_GEOMETRY_CACHE.clear() + print(" detector_geometry calls: %d for 1 sample, %d for %d samples" % (one, many, n)) + assert many == one, ("the detector geometry is sample-independent; %d samples cost %d " + "lookups, 1 sample costs %d" % (n, many, one)) + return worst + + if __name__ == "__main__": run() run_strong() + run_response_coefficients_block() print("\nALL SLOWROT FREQRESPONSE LIKELIHOOD CHECKS PASSED") diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_rotating_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_rotating_freqresponse.py new file mode 100644 index 000000000..8beba36cc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/test_slowrot_rotating_freqresponse.py @@ -0,0 +1,135 @@ +"""Unit tests for the simultaneous slow-rotation + finite-arm response.""" +from __future__ import division, print_function + +import numpy as np +import lal +import lalsimulation as lalsim + +from RIFT.likelihood import factored_likelihood_rotating_freqresponse as combined +from RIFT.likelihood import factored_likelihood as fl +from RIFT.likelihood import factored_likelihood_freqresponse as flfr +from RIFT.likelihood import slowrot_freqresponse as sfr +from RIFT import lalsimutils as lsu + +if not getattr(fl, "numba_on", True): + fl.lalylm = np.vectorize(lal.SpinWeightedSphericalHarmonic, otypes=[complex]) + + +def test_compound_index_set_carries_exact_nonzero_harmonics(): + a0 = combined.compound_index_set(Qmax=4, p_max=0) + a1 = combined.compound_index_set(Qmax=4, p_max=1) + assert len(a0) == 50 + assert len(a1) == 112 + assert len(a0) == len(set(a0)) + assert len(a1) == len(set(a1)) + for b, p, n in a1: + assert abs(n) <= combined.response_harmonic_width(b) + p + + +def test_basis_harmonics_reconstruct_frequency_response_coefficients(): + rng = np.random.RandomState(20260909) + nsamp = 5 + ra = rng.uniform(0.0, 2.0 * np.pi, nsamp) + dec = np.arcsin(rng.uniform(-1.0, 1.0, nsamp)) + psi = rng.uniform(0.0, np.pi, nsamp) + tref = 1126259462.4 + qmax = 4 + C = combined.combined_response_coefficients_vector( + 'H1', ra, dec, psi, tref, p_max=0, Qmax=qmax, L_arm=40000.0) + gmst0 = float(lal.GreenwichMeanSiderealTime(lal.LIGOTimeGPS(tref))) + + worst = 0.0 + for dt in (0.0, 1234.0, 7200.0): + post = lambda n: np.exp(1j * n * 2.0 * np.pi * combined.flwr.F_SIDEREAL * dt) + for i in range(nsamp): + geom = sfr.finite_size_geometry( + 'H1', ra[i], dec[i], psi[i], + gmst=gmst0 + 2.0 * np.pi * combined.flwr.F_SIDEREAL * dt, + L_arm=40000.0) + want = {0: geom['F0']} + want.update({1 + q: beta for q, beta in enumerate(sfr.finite_size_beta(geom, qmax))}) + for b in range(qmax + 2): + got = sum(C.get((b, 0, n), np.zeros(nsamp))[i] * post(n) + for n in range(-qmax - 2, qmax + 3)) + worst = max(worst, abs(got - want[b])) + assert worst < 1e-11, "compound sidereal reconstruction error %g" % worst + + +def test_delay_order_extends_each_frequency_basis_width_without_leakage(): + ra = np.array([0.7, 2.4]) + dec = np.array([-0.3, 0.8]) + psi = np.array([0.2, 1.1]) + qmax = 3 + pmax = 2 + C = combined.combined_response_coefficients_vector( + 'L1', ra, dec, psi, 1126259462.4, pmax, Qmax=qmax, L_arm=40000.0) + allowed = set(combined.compound_index_set(qmax, pmax)) + assert set(C).issubset(allowed) + for b in range(qmax + 2): + for p in range(pmax + 1): + edge = combined.response_harmonic_width(b) + p + assert all(abs(n) <= edge for bb, pp, n in C if bb == b and pp == p) + + +def test_compound_likelihood_reduces_to_freqresponse_at_zero_rotation_rate(): + """Exercise waveform generation, compound packing, and the shared NoLoop contraction.""" + fsample = 1024.0 + event_time = 1000000000.0 + t_window = 0.08 + fmax = 400.0 + psig = lsu.ChooseWaveformParams( + fmin=30.0, radec=True, incl=0.3, phiref=0.0, theta=0.2, phi=1.0, + psi=0.4, m1=30 * lal.MSUN_SI, m2=25 * lal.MSUN_SI, detector='H1', + dist=200e6 * lal.PC_SI, deltaT=1.0 / fsample, tref=event_time, + deltaF=0.5) + data = {'H1': lsu.non_herm_hoff(psig)} + psd = {'H1': lalsim.SimNoisePSDaLIGOZeroDetHighPower} + + _, ct_f, ctv_f, rho_f, meta_f = flfr.PrecomputeLikelihoodTermsFreqResponse( + event_time, t_window, psig, data, psd, 2, fmax, Qmax=0, L_arm=40000.0, + analyticPSD_Q=True, verbose=False, quiet=True, skip_interpolation=True) + lk_f, rho_arr_f, u_f, v_f, ep_f = flfr.pack_freqresponse_arrays( + meta_f, rho_f, ct_f, ctv_f) + + _, ct_c, ctv_c, rho_c, meta = combined.PrecomputeLikelihoodTermsRotatingFreqResponse( + event_time, t_window, psig, data, psd, 2, fmax, Qmax=0, L_arm=40000.0, + p_max=1, f_sidereal=0.0, analyticPSD_Q=True, verbose=False, quiet=True, + skip_interpolation=True) + lk_c, rho_arr_c, u_c, v_c, ep_c = combined.pack_rotating_freqresponse_arrays( + meta, rho_c, ct_c, ctv_c) + + pvec = psig.manual_copy() + for name, val in [('phi', 1.0), ('theta', 0.2), ('incl', 0.7), + ('phiref', 0.9), ('psi', 0.5)]: + setattr(pvec, name, np.full(2, val)) + pvec.dist = np.full(2, 300e6 * lal.PC_SI) + pvec.tref = event_time + pvec.deltaT = 1.0 / fsample + tvals = np.linspace(-0.04, 0.04, 64) + + want = flfr.DiscreteFactoredLogLikelihoodFreqResponseNoLoop( + tvals, pvec, meta_f, lk_f, rho_arr_f, u_f, v_f, ep_f, Lmax=2, + array_output=True, time_interp='cubic') + got = combined.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, pvec, meta, lk_c, rho_arr_c, u_c, v_c, ep_c, Lmax=2, + array_output=True, time_interp='cubic') + worst = float(np.max(np.abs(got - want))) + assert got.shape == want.shape + assert worst < 1e-8, "compound zero-rotation limit mismatch %g" % worst + + ip = lsu.ComplexIP(30.0, fmax, fsample / 2.0, data['H1'].deltaF, + psd['H1'], True, False, 0.0) + half_dd = 0.5 * ip.ip(data['H1'], data['H1']).real + assert np.max(got) <= half_dd + 1e-7, "zero-rate compound likelihood violates bound" + + _, ct_w, ctv_w, rho_w, meta_w = combined.PrecomputeLikelihoodTermsRotatingFreqResponse( + event_time, t_window, psig, data, psd, 2, fmax, Qmax=0, L_arm=40000.0, + p_max=1, f_sidereal=combined.flwr.F_SIDEREAL, analyticPSD_Q=True, + verbose=False, quiet=True, skip_interpolation=True) + lk_w, rho_arr_w, u_w, v_w, ep_w = combined.pack_rotating_freqresponse_arrays( + meta_w, rho_w, ct_w, ctv_w) + got_w = combined.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, pvec, meta_w, lk_w, rho_arr_w, u_w, v_w, ep_w, Lmax=2, + array_output=True, time_interp='cubic') + assert np.all(np.isfinite(got_w)) + assert np.max(got_w) <= half_dd + 1e-7, "rotating finite-arm likelihood violates bound" diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py index f9e437f93..9a2efb6ec 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/time_marginalization_quadrature.py @@ -161,6 +161,9 @@ reconciled across realizations, which is untested here. """ +import os +import warnings + import numpy as np __all__ = [ @@ -260,6 +263,245 @@ #: is not floating-point noise. _DENSE_CHUNK_BYTES = 128 * 1024 * 1024 + +def _cpu_fft_workers(): + """Bounded CPU FFT parallelism, respecting scheduler CPU affinity. + + The reflected transforms have awkward production lengths (for example + ``2*307``), and dominate the AV band-limited path. SciPy's pocketfft can + parallelize the independent row transforms, while NumPy's public FFT API + cannot. Never request more CPUs than the process affinity mask exposes; + ``RIFT_TIME_FFT_WORKERS`` can lower the cap or raise the default cap of four. + """ + try: + available = len(os.sched_getaffinity(0)) + except (AttributeError, OSError): + available = os.cpu_count() or 1 + requested = int(os.environ.get("RIFT_TIME_FFT_WORKERS", "4")) + return max(1, min(requested, available)) + + +def _fft_rows(x, inverse=False, n=None, xpy=np): + if xpy is np: + from scipy import fft as scipy_fft + fn = scipy_fft.ifft if inverse else scipy_fft.fft + return fn(x, n=n, axis=-1, workers=_cpu_fft_workers()) + fn = xpy.fft.ifft if inverse else xpy.fft.fft + return fn(x, n=n, axis=-1) + + +class _RetainedFFTUnsupported(RuntimeError): + """The optional retained-grid transform cannot honour this input.""" + + +def _retained_fft_backend(xpy): + """Return the supported backend name without moving an array to the host.""" + if xpy is np: + return "numpy" + if getattr(xpy, "__name__", None) == "cupy": + return "cupy" + raise _RetainedFFTUnsupported( + "retained-grid FFT supports only the numpy and cupy backends") + + +def _retained_fft_plan(period, factor, dtype, xpy=np): + """Build stable Bluestein chirps for the forward half of a reflected row. + + The plan evaluates the same Fourier polynomial as zero padding to + ``period * factor``, but only at the ``(period/2 - 1)*factor + 1`` samples + consumed by the finite-window integral. Integer modular phases avoid the + unit-circle drift of forming a high power of one approximate complex root. + Plans live only for one marginalization call, so large GPU chirps cannot + become an unbounded process-wide cache. + """ + _retained_fft_backend(xpy) + period = int(period) + factor = int(factor) + dtype = np.dtype(dtype) + if period < 4 or period % 2: + raise _RetainedFFTUnsupported( + "reflected FFT period must be even and at least four") + if factor <= 1 or factor & (factor - 1): + raise _RetainedFFTUnsupported( + "retained-grid FFT requires a power-of-two factor above one") + if dtype != np.dtype(np.complex128): + raise _RetainedFFTUnsupported( + "retained-grid FFT is certified only for complex128 spectra, got %s" + % dtype) + + # The Nyquist coefficient is represented at both signed endpoints, hence + # period+1 input coefficients. Linear Bluestein convolution needs the sum + # of input and output lengths minus one. next_fast_len is a host-side + # integer calculation only; all arrays and FFTs stay on xpy's device. + n_coeff = period + 1 + n_out = (period // 2 - 1) * factor + 1 + n_chirp = max(n_coeff, n_out) + if n_chirp > 3037000499 or period * factor > np.iinfo(np.int64).max // 2: + raise _RetainedFFTUnsupported( + "retained-grid dimensions exceed the exact int64 chirp-phase range") + from scipy.fft import next_fast_len + n_fft = int(next_fast_len(n_coeff + n_out - 1)) + + k = xpy.arange(n_chirp, dtype=np.int64) + denominator = period * factor + # exp(+i*pi*k**2/denominator), reduced exactly modulo 2*denominator + # before conversion to float. The largest supported production grid is + # safely within int64 (roughly 1e14 at npts=2457, factor=4096). + phase_index = (k * k) % (2 * denominator) + wk2 = xpy.exp((1j * np.pi / denominator) * phase_index) + wk2 = xpy.asarray(wk2, dtype=np.complex128) + kernel = 1.0 / xpy.concatenate( + (wk2[n_coeff - 1:0:-1], wk2[:n_out])) + kernel_fft = _fft_rows(kernel, n=n_fft, xpy=xpy) + + j = xpy.arange(n_out, dtype=np.int64) + shift_index = j % (2 * factor) + signed_frequency_shift = xpy.exp( + (-1j * np.pi / factor) * shift_index) + post = (wk2[:n_out] * signed_frequency_shift) / float(period) + return { + "input_chirp": wk2[:n_coeff], + "kernel_fft": kernel_fft, + "post_chirp": post, + "n_fft": n_fft, + "n_out": n_out, + "period": period, + "factor": factor, + } + + +def _reflected_bandlimited_upsample_retained(x, factor, plan_cache=None, + xpy=np): + """Evaluate exactly the retained forward grid of the reflected interpolant. + + This is a pruned *evaluation* of :func:`reflected_bandlimited_upsample`, not + a different interpolant. It preserves the literal ``[x, flip(x)]`` + boundary condition and the half-weight split of the even-period Nyquist bin. + """ + x = xpy.asarray(x) + factor = int(factor) + if factor == 1: + return x + n = int(x.shape[-1]) + period = 2 * n + reflected = xpy.concatenate((x, xpy.flip(x, axis=-1)), axis=-1) + spectrum = _fft_rows(reflected, xpy=xpy) + dtype = np.dtype(spectrum.dtype) + cache_key = (period, factor, dtype.str) + if plan_cache is None: + plan_cache = {} + plan = plan_cache.get(cache_key) + if plan is None: + plan = _retained_fft_plan(period, factor, dtype, xpy=xpy) + plan_cache[cache_key] = plan + + half = period // 2 + # Consecutive signed-frequency coefficients k=-half,...,+half. Splitting + # the Nyquist bin across the two endpoints is exactly what the full padded + # inverse FFT does in bandlimited_upsample for an even-length row. + coeff = xpy.empty(spectrum.shape[:-1] + (period + 1,), dtype=spectrum.dtype) + coeff[..., 0] = 0.5 * spectrum[..., half] + coeff[..., 1:half] = spectrum[..., half + 1:] + coeff[..., half] = spectrum[..., 0] + coeff[..., half + 1:period] = spectrum[..., 1:half] + coeff[..., period] = 0.5 * spectrum[..., half] + + transformed = _fft_rows( + coeff * plan["input_chirp"], n=plan["n_fft"], xpy=xpy) + transformed *= plan["kernel_fft"] + convolved = _fft_rows(transformed, inverse=True, xpy=xpy) + retained = convolved[..., period:period + plan["n_out"]] + retained *= plan["post_chirp"] + if retained.shape[-1] != (n - 1) * factor + 1: + raise RuntimeError("retained-grid FFT returned an inconsistent shape") + return retained + + +def _record_transform(report, key, n_rows, period, factor, plan=None): + report[key + "_batches"] += 1 + report[key + "_rows"] += int(n_rows) + report["max_reflected_period"] = max(report["max_reflected_period"], + int(period)) + report["max_dense_factor"] = max(report["max_dense_factor"], int(factor)) + report["max_reference_full_fft_length"] = max( + report["max_reference_full_fft_length"], int(period) * int(factor)) + if plan is not None: + report["max_retained_fft_length"] = max( + report["max_retained_fft_length"], int(plan["n_fft"])) + report["max_retained_grid_length"] = max( + report["max_retained_grid_length"], int(plan["n_out"])) + + +def _new_transform_report(): + return dict( + retained_fft_batches=0, + retained_fft_rows=0, + full_fft_selected_batches=0, + full_fft_selected_rows=0, + full_fft_selected_reasons={}, + full_fft_fallback_batches=0, + full_fft_fallback_rows=0, + full_fft_fallback_reasons={}, + warned_fallback_reasons=set(), + max_reflected_period=0, + max_dense_factor=1, + max_reference_full_fft_length=0, + max_retained_fft_length=0, + max_retained_grid_length=0, + ) + + +def _reflected_upsample_for_integration(x, factor, plan_cache, + transform_report, xpy=np): + """Use the retained-grid transform, visibly falling back to the reference. + + An optimization failure is not a waveform or likelihood failure. Any + unsupported input or transform exception therefore retries the established + full-padding implementation and records why. The likelihood callback is + deliberately outside this function, so its failures are never mislabeled or + swallowed as FFT fallbacks. + """ + period = 2 * int(x.shape[-1]) + # Pocketfft measurements across all production npts found the retained + # convolution neutral-to-slower at factors 2 and 4; that small dense grid is + # not the bottleneck. Preserve the cheaper reference algorithm there. On + # CuPy the retained path won at every tested factor 2--64. + if xpy is np and int(factor) in (2, 4): + reason = "numpy factor %d is below the measured retained-FFT crossover" % factor + reasons = transform_report["full_fft_selected_reasons"] + reasons[reason] = reasons.get(reason, 0) + int(x.shape[0]) + _record_transform(transform_report, "full_fft_selected", x.shape[0], + period, factor) + return reflected_bandlimited_upsample(x, factor, xpy=xpy) + try: + out = _reflected_bandlimited_upsample_retained( + x, factor, plan_cache=plan_cache, xpy=xpy) + plan = next((value for (plan_period, plan_factor, _), value + in plan_cache.items() + if plan_period == period and plan_factor == int(factor)), None) + _record_transform(transform_report, "retained_fft", x.shape[0], + period, factor, plan) + return out + except Exception as exc: + reason = "%s: %s" % (type(exc).__name__, str(exc)) + reasons = transform_report["full_fft_fallback_reasons"] + reasons[reason] = reasons.get(reason, 0) + int(x.shape[0]) + _record_transform(transform_report, "full_fft_fallback", x.shape[0], + period, factor) + if reason not in transform_report["warned_fallback_reasons"]: + # Warning filters are allowed to promote RuntimeWarning to an + # exception. Diagnostics must not turn a successful reference-path + # retry into a dropped likelihood point, so contain that policy here. + try: + warnings.warn( + "retained-grid band-limited FFT unavailable ({}); using the " + "established full-padding sinc reconstruction for these rows" + .format(reason), RuntimeWarning, stacklevel=2) + except Exception: + pass + transform_report["warned_fallback_reasons"].add(reason) + return reflected_bandlimited_upsample(x, factor, xpy=xpy) + _LAST_REPORT = {} @@ -269,7 +511,16 @@ def last_report(): 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``. + ``n_flat_rows``, ``n_refined_rows``, and retained-transform provenance. + + ``bandlimited_fft_strategy`` says whether the production-only optimization + used the retained-grid ZoomFFT, intentionally selected the established full + transform below a measured CPU crossover, fell back to it after a transform + decline, used a mixture, or needed no dense transform. The corresponding + ``*_batches`` and ``*_rows`` fields distinguish these cases; the reason maps + make a cost selection or declined optimization auditable without converting + either into a failed waveform point. The reported reference, retained-grid, + and convolution lengths expose the padding mismatch for performance records. The diagnostic row counts are deliberately kept apart because they mean different things: @@ -527,7 +778,7 @@ def bandlimited_upsample(x, factor, xpy=np): return x n = x.shape[-1] lead = x.shape[:-1] - X = xpy.fft.fft(x, axis=-1) + X = _fft_rows(x, xpy=xpy) 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] @@ -538,7 +789,7 @@ def bandlimited_upsample(x, factor, xpy=np): Xup[..., -n_pos:] = X[..., n // 2 + 1:] else: Xup[..., -n_pos:] = X[..., n_pos + 1:] - return xpy.fft.ifft(Xup, axis=-1) * factor + return _fft_rows(Xup, inverse=True, xpy=xpy) * factor def reflected_bandlimited_upsample(x, factor, xpy=np): @@ -972,7 +1223,18 @@ def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, # auditable claim.) refined = has_peak & (factors > 1) - out = _log_simps_rows(lnL_coarse, deltaT, simps, xpy=xpy) + # Do not pay for the historical coarse-grid integral on rows that we already + # know will be overwritten by the dense reconstruction below. In ordinary + # AV ILE the coarse likelihood has already been evaluated for classification; + # the old unconditional call added another exp/reduction over every + # extrinsic×time point even when every row required refinement. Allocate the + # result once and run Simpson only on the rows for which it is the answer. + out = xpy.empty((n_rows,), dtype=xpy.asarray(lnL_coarse).dtype) + unrefined = ~refined + if bool(xpy.any(unrefined)): + idx_unrefined = xpy.where(unrefined)[0] + out[idx_unrefined] = _log_simps_rows( + lnL_coarse[idx_unrefined], deltaT, simps, xpy=xpy) time_draw = None lnL_at_draw = None if return_time_draw: @@ -991,6 +1253,11 @@ def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, hist = {} n_refine_total = 0 sigma_seen = np.inf + # Reuse chirps across every batch at a given factor, but only for this + # marginalization call. In particular, do not pin successively larger GPU + # plans in a process-global cache after a high-SNR cell has finished. + retained_plan_cache = {} + transform_report = _new_transform_report() for f in xpy.unique(xpy.where(refined, factors, 1)): f = int(f) if f == 1: @@ -1000,18 +1267,36 @@ def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, if not n_sel: continue idx = xpy.where(sel)[0] - vals, f_used, n_ref, s_min, drawn_t, drawn_lnL = _integrate_group( + vals, group_hist, 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) + t0=t0, retained_plan_cache=retained_plan_cache, + transform_report=transform_report, 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 + for f_used, n_used in group_hist.items(): + hist[int(f_used)] = hist.get(int(f_used), 0) + int(n_used) n_refine_total += n_ref sigma_seen = min(sigma_seen, s_min) + strategies = [] + if transform_report["retained_fft_batches"]: + strategies.append("retained-grid-zoomfft") + if transform_report["full_fft_selected_batches"]: + strategies.append("full-padding-selected") + if transform_report["full_fft_fallback_batches"]: + strategies.append("full-padding-fallback") + transform_strategy = (strategies[0] if len(strategies) == 1 else + ("mixed:" + ",".join(strategies) if strategies + else "not-used")) + transform_report.pop("warned_fallback_reasons") + transform_report.update( + bandlimited_fft_strategy=transform_strategy, + n_retained_fft_plans=len(retained_plan_cache), + ) + _LAST_REPORT.clear() _LAST_REPORT.update( upsample_factor=max(hist) if hist else 1, @@ -1023,6 +1308,8 @@ def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, n_unmeasurable_rows=int(xpy.sum(unmeasurable)), n_flat_rows=int(xpy.sum(flat)), n_refined_rows=int(xpy.sum(refined)), + cpu_fft_workers=(_cpu_fft_workers() if xpy is np else None), + **transform_report ) if return_time_draw: return out, time_draw, lnL_at_draw @@ -1031,16 +1318,28 @@ def time_marginalize_bandlimited(kappa, rho_sq, deltaT, loglikelihood, def _integrate_group(kappa_rows, rho_col_rows, npts, deltaT, factor, loglikelihood, _term, draw_uniforms_rows=None, t0=0.0, - xpy=np): + retained_plan_cache=None, transform_report=None, xpy=np): """Refine and integrate one group of rows that share a derived factor. - Returns ``(values, factor_used, n_refinements, sigma_dense_min, + Returns ``(values, factor_histogram, 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] + if retained_plan_cache is None: + retained_plan_cache = {} + if transform_report is None: + transform_report = _new_transform_report() n_refine = 0 - while True: + remaining = xpy.arange(n_rows) + values = xpy.empty((n_rows,), dtype=np.float64) + time_values = (xpy.empty((n_rows,), dtype=np.float64) + if draw_uniforms_rows is not None else None) + draw_lnL_values = (xpy.empty((n_rows,), dtype=np.float64) + if draw_uniforms_rows is not None else None) + factor_hist = {} + sigma_seen = np.inf + while int(remaining.size): if factor > UPSAMPLE_FACTOR_MAX: raise RuntimeError( "band-limited time marginalization needs an upsampling factor above " @@ -1054,25 +1353,28 @@ def _integrate_group(kappa_rows, rho_col_rows, npts, deltaT, factor, # 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)))) + n_remaining = int(remaining.size) + chunk = max(1, min(n_remaining, 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) + sigma_pieces = [] + for start in range(0, n_remaining, chunk): + take = remaining[start:start + chunk] + k_up = _reflected_upsample_for_integration( + kappa_rows[take], factor, retained_plan_cache, + transform_report, xpy=xpy) + rho_up = xpy.broadcast_to(rho_col_rows[take], 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))) + sigma_pieces.append(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) + uniforms=draw_uniforms_rows[take], xpy=xpy) draw_time_pieces.append(drawn_t) draw_lnL_pieces.append(drawn_lnL) @@ -1081,13 +1383,28 @@ def _integrate_group(kappa_rows, rho_col_rows, npts, deltaT, factor, # 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 - + current_values = xpy.concatenate(pieces) if len(pieces) > 1 else pieces[0] + current_sigma = (xpy.concatenate(sigma_pieces) + if len(sigma_pieces) > 1 else sigma_pieces[0]) + finite_sigma = xpy.isfinite(current_sigma) + if bool(xpy.any(finite_sigma)): + sigma_seen = min(sigma_seen, float(xpy.min(current_sigma[finite_sigma]))) + resolved = (~finite_sigma) | (dx_dense <= current_sigma / UPSAMPLE_SAFETY) + accepted = remaining[resolved] + values[accepted] = current_values[resolved] + n_accepted = int(xpy.sum(resolved)) + if n_accepted: + factor_hist[int(factor)] = factor_hist.get(int(factor), 0) + n_accepted + if draw_uniforms_rows is not None: + current_t = (xpy.concatenate(draw_time_pieces) if len(draw_time_pieces) > 1 + else draw_time_pieces[0]) + current_draw_lnL = (xpy.concatenate(draw_lnL_pieces) + if len(draw_lnL_pieces) > 1 else draw_lnL_pieces[0]) + time_values[accepted] = current_t[resolved] + draw_lnL_values[accepted] = current_draw_lnL[resolved] + remaining = remaining[~resolved] + if not int(remaining.size): + return (values, factor_hist, n_refine, sigma_seen, + time_values, draw_lnL_values) factor *= 2 n_refine += 1 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/vectorized_lal_tools.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/vectorized_lal_tools.py index 1a5278e62..bd01ff2d5 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/vectorized_lal_tools.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/vectorized_lal_tools.py @@ -35,9 +35,32 @@ def TimeDelayFromEarthCenter( ------- time_delay_from_earth_center : array_like, shape = det_shape + sample_shape """ - negative_speed_of_light = xpy.asarray(-299792458.0) + ehat_src = SourcePropagationDirection( + source_right_ascension_radians, source_declination_radians, + greenwich_mean_sidereal_time, xpy=xpy, dtype=dtype, + ) + return TimeDelayFromEarthCenterPrecomputed( + detector_earthfixed_xyz_metres, ehat_src, xpy=xpy, + ) + + +def SourcePropagationDirection( + source_right_ascension_radians, + source_declination_radians, + greenwich_mean_sidereal_time, + xpy=xpy_default, dtype=numpy.float64, + ): + """Unit vector from Earth's center towards the source, in Earth-fixed frame. - det_shape = detector_earthfixed_xyz_metres.shape[:-1] + This depends only on the SOURCE, not on the detector, so a caller looping over + detectors with a fixed set of extrinsic samples can build it once and hand it to + ``TimeDelayFromEarthCenterPrecomputed`` for each detector instead of recomputing + three trig evaluations per detector. + + Returns + ------- + ehat_src : array_like, shape = sample_shape + (3,) + """ sample_shape = source_right_ascension_radians.shape cos_dec = xpy.cos(source_declination_radians) @@ -52,6 +75,20 @@ def TimeDelayFromEarthCenter( ehat_src[...,1] = -cos_dec * xpy.sin(greenwich_hour_angle) ehat_src[...,2] = xpy.sin(source_declination_radians) + return ehat_src + + +def TimeDelayFromEarthCenterPrecomputed( + detector_earthfixed_xyz_metres, ehat_src, xpy=xpy_default, + ): + """Per-detector half of :func:`TimeDelayFromEarthCenter`. + + ``ehat_src`` comes from :func:`SourcePropagationDirection`. The arithmetic is the + same ``inner`` contraction the combined function performs, so results are bitwise + identical to calling ``TimeDelayFromEarthCenter`` directly. + """ + negative_speed_of_light = xpy.asarray(-299792458.0) + neg_separation = xpy.inner(detector_earthfixed_xyz_metres, ehat_src) return xpy.divide( neg_separation, negative_speed_of_light, @@ -89,9 +126,34 @@ def ComputeDetAMResponse( ------- F : array_like, shape = det_shape + sample_shape """ - det_shape = detector_response_matrix.shape[:-1] + X, Y = SourcePolarizationBasis( + source_right_ascension_radians, source_declination_radians, + source_polarization_radians, greenwich_mean_sidereal_time, + xpy=xpy, dtype_real=dtype_real, + ) + return ComputeDetAMResponsePrecomputed( + detector_response_matrix, X, Y, xpy=xpy, + ) + + +def SourcePolarizationBasis( + source_right_ascension_radians, + source_declination_radians, + source_polarization_radians, + greenwich_mean_sidereal_time, + xpy=xpy_default, dtype_real=numpy.float64, + ): + """The (X, Y) polarization basis vectors in the Earth-fixed frame. + + Six trig evaluations and twelve elementwise combinations, none of which depend on + the DETECTOR -- only the contraction with the response matrix does. A caller + looping over detectors at fixed extrinsic samples should build this once. + + Returns + ------- + X, Y : array_like, shape = sample_shape + (3,) + """ sample_shape = source_right_ascension_radians.shape - matrix_shape = 3, 3 # Initialize trig matrices. X = xpy.empty(sample_shape+(3,), dtype=dtype_real) @@ -119,6 +181,20 @@ def ComputeDetAMResponse( Y[...,1] = sin_psi*cos_gha + cos_psi*sin_gha*sin_dec Y[...,2] = cos_psi*cos_dec + return X, Y + + +def ComputeDetAMResponsePrecomputed( + detector_response_matrix, X, Y, xpy=xpy_default, + ): + """Per-detector half of :func:`ComputeDetAMResponse`. + + ``X, Y`` come from :func:`SourcePolarizationBasis`. The contractions are the same + ``inner`` calls in the same order as the combined function, so results are bitwise + identical to calling ``ComputeDetAMResponse`` directly. (A single batched einsum + over stacked detectors would be fewer launches still, but reassociates the + contraction and is only equal to ~4e-16; that is deliberately not done here.) + """ # Compute F for each polarization state. F_plus = ( X*xpy.inner(X, detector_response_matrix) - diff --git a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py index 8f03a5685..c7191e33a 100755 --- a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py +++ b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py @@ -38,6 +38,11 @@ from RIFT.likelihood.time_marginalization_quadrature import ( TIME_QUADRATURE_CHOICES, validate_time_quadrature, refuse_unless_time_quadrature_emitted) +# Same leaf-module reasoning: the choice tuple and the stencil-conflict wording are IMPORTED, +# not re-typed, so this helper and the ILE driver's own guard cannot silently disagree. +from RIFT.likelihood.q_time_pregrid import ( + Q_TIME_PREGRID_CHOICES, validate_q_time_pregrid_factor, + refuse_unless_q_time_pregrid_emitted) lalapps_path2cache = which('lal_path2cache') ligolw_add = 'igwn_ligolw_add' if not(which(ligolw_add)): @@ -238,6 +243,7 @@ def get_observing_run(t): 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: 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-ile-q-time-pregrid-factor",default=None,type=int,choices=list(Q_TIME_PREGRID_CHOICES),help="OPT-IN certified Q_lm pregrid (PR #261): %s. Default None = emit nothing, so ILE keeps its own default (factor 1, the historical unchanged path) and args_ile.txt is byte-identical to today. Factor 8 reflects each finite Q window, FFT-interpolates it onto an 8x finer grid once after packing, and evaluates detector arrival times off that grid with four-tap CUBIC interpolation -- the geocentric time-integration grid is left at the data deltaT. Emitted as --q-time-pregrid-factor on the ILE command line, so a completed run's pregrid setting is readable off the .sub file. Requires --vectorized and excludes --rotation-slow / --freqresponse / calibration marginalization; also FORCES the cubic stencil and REFUSES a conflicting explicit --internal-ile-interpolate-time (i.e. one naming a stencil other than cubic) rather than silently overriding it. 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 in an ini that a Makefile also sets." % ("|".join(str(c) for c in Q_TIME_PREGRID_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!") @@ -306,6 +312,13 @@ def get_observing_run(t): if time_quadrature_choice is not None: validate_time_quadrature(time_quadrature_choice) +# Same, for the Q_lm pregrid factor: 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. +q_time_pregrid_factor = opts.internal_ile_q_time_pregrid_factor +if q_time_pregrid_factor is not None: + validate_q_time_pregrid_factor(q_time_pregrid_factor) + # 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] @@ -1254,6 +1267,24 @@ def crit_m2(delta): # then have to catch it. Make it structurally impossible instead. helper_ile_args = helper_ile_args.rstrip() + " --time-marginalization-quadrature " + time_quadrature_choice + " " +if q_time_pregrid_factor is not None: + # Validated at parse time, so by here it is one of Q_TIME_PREGRID_CHOICES. The value goes + # on the ILE command line verbatim, so a completed run's pregrid setting is readable off + # the .sub file. Prerequisites (--vectorized, the exclusions, and the forced-cubic-stencil + # conflict) are checked on the FULLY ASSEMBLED command line below, by + # refuse_unless_q_time_pregrid_emitted -- not here, because --vectorized/--gpu are added by + # the strategy branches further down this 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 the historical factor=1 grid. + print(" ==> Q_lm pregrid factor: {} (emitted as --q-time-pregrid-factor; the ILE driver " + "refuses rather than ignores if its configuration cannot honour it)".format( + q_time_pregrid_factor)) + # rstrip(), for the same reason as the quadrature emission just above: the flag gluing onto + # its neighbour would make it invisible to the emission guard. + helper_ile_args = helper_ile_args.rstrip() + " --q-time-pregrid-factor " + str(q_time_pregrid_factor) + " " + if opts.internal_ile_auto_logarithm_offset and not opts.internal_ile_use_lnL: helper_ile_args += " --auto-logarithm-offset " rescaled_base_ile = True @@ -1963,6 +1994,9 @@ def lambda_m_estimate(m): # 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") +# Same discipline, same reason, for the Q_lm pregrid factor. +refuse_unless_q_time_pregrid_emitted( + q_time_pregrid_factor, 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: diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic index afa3dfe86..d937519eb 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic @@ -1044,7 +1044,7 @@ if not opts.time_marginalization: param_limits["t_ref"][0], param_limits["t_ref"][1]) sampler.add_parameter("t_ref", pdf = tref_sampler, - cdf_inv = None, + cdf_inv = tref_sampler_cdf_inv, left_limit = param_limits["t_ref"][0], right_limit = param_limits["t_ref"][1], prior_pdf = mcsampler.ret_uniform_samp_vector_alt(param_limits["t_ref"][0], param_limits["t_ref"][1])) diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index c6db29b7f..8ec7a59b2 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -51,6 +51,13 @@ 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 +# Q_TIME_PREGRID_CHOICES (RIFT/likelihood/q_time_pregrid.py) is the SINGLE definition of the +# legal --q-time-pregrid-factor set; this driver calls the shared validator below rather than +# retyping its own tuple+message, so the pipeline-side mirror (q_time_pregrid.py) and this +# driver cannot silently drift apart (PR #291 review, MAJOR #2: this used to be an independent +# literal here, and a pipeline-side docstring claimed drift was impossible while nothing +# enforced it). +from RIFT.likelihood.q_time_pregrid import validate_q_time_pregrid_factor 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 @@ -249,6 +256,7 @@ optp.add_option("--resample-time-marginalization",action='store_true', help="If 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("--psi-marginalization", action="store_true", default=False, help="Opt-in: analytically marginalize the polarization angle psi over its uniform [0,pi) prior with factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized, instead of sampling it. Only reachable on the legacy scalar (non-vectorized, non-GPU, non-time-marginalized) likelihood path, which is the only call site this function fits; REFUSED (not ignored) with --time-marginalization, --vectorized, --gpu, --distance-marginalization, --rotation-slow, --freqresponse, calibration marginalization, an explicit --interpolate-time, --sampler-method GMM/portfolio, --internal-rotate-phase, --limit-psi, and --zero-likelihood. Default is false (unreachable previously; see issue tracking marginalization-audit). The exported 'psi'/'polarization' column is NaN, NOT a sample: psi is integrated out, so there is no per-sample draw to report, and downstream PE-sample converters copy that column verbatim. The reported lnL carries the SAME psi prior mass the sampled-psi path carries (1/pi over (0,2 pi), i.e. 2), so lnZ is comparable with ordinary rows in the same all.net; note that mass is not 1, and the RIFT JAX driver uses psi in [0,pi] instead.") 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") @@ -277,10 +285,10 @@ optp.add_option("--extrinsic-proposal-adapt",action='store_true',default=False, 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)") optp.add_option("--gpu", action="store_true", help="Perform manipulations of lm and timeseries using numpy arrays, CONVERTING TO GPU when available. You MUST use this option with --vectorized (otherwise it is a no-op). You MUST have a suitable version of cupy installed, your cuda operational, etc") optp.add_option("--force-gpu-only", action="store_true", help="Hard fail if no GPU present (assessed by cupy not loading)") -optp.add_option("--rotation-slow", action="store_true", help="[Path A] Slow-rotation likelihood: account for the sidereal time-dependence of the antenna pattern F(t) over the signal (harmonic modulation). Requires --vectorized; supports --gpu (n_cal=1, no glitch/cal marg). Uses factored_likelihood_with_rotation.") +optp.add_option("--rotation-slow", action="store_true", help="[Path A] Slow-rotation likelihood: account for the sidereal time-dependence of the antenna pattern F(t) over the signal (harmonic modulation). Requires --vectorized; supports --gpu (n_cal=1, no glitch/cal marg). May be combined with --freqresponse for long, loud 3G/BNS-like signals.") optp.add_option("--rotation-n-harmonics", type=int, default=2, help="Number of sidereal harmonics for --rotation-slow (antenna pattern needs 2, i.e. n=-2..2).") optp.add_option("--rotation-p-max", type=int, default=0, help="[Path B] Max delay-derivative order for --rotation-slow (0 = amplitude drift only; >=1 adds propagation-delay drift).") -optp.add_option("--freqresponse", action="store_true", help="[Path D] Finite-size (frequency-dependent) detector-response likelihood: account for the finite light-travel-time transfer across the arms (matters for 3G/CE-ET). Requires --vectorized; supports --gpu (n_cal=1, no glitch/cal marg). Uses factored_likelihood_freqresponse (sky-harmonic route b, sky stays extrinsic).") +optp.add_option("--freqresponse", action="store_true", help="[Path D] Finite-size (frequency-dependent) detector-response likelihood: account for the finite light-travel-time transfer across the arms (matters for 3G/CE-ET). Requires --vectorized; supports --gpu (n_cal=1, no glitch/cal marg). May be combined with --rotation-slow; the compound bank is substantially more expensive and intended for long, loud sources.") optp.add_option("--freqresponse-qmax", type=int, default=4, help="Highest power of the arm projection retained for --freqresponse (basis size Qmax+2). Higher for larger fL/c (heavier systems / higher fmax).") optp.add_option("--freqresponse-arm-length", default=None, help="Arm-length override [m] for --freqresponse. Either a single float applied to ALL detectors (e.g. 40000 for 40-km CE), or per-detector 'C1=40000,E1=10000,...' (needed for mixed CE+ET networks, since LAL's cached C1 arm is a placeholder). Default: each detector's native LAL arm length.") optp.add_option("--force-xpy", action="store_true", help="Use the xpy code path. Use with --vectorized --gpu to use the fallback CPU-based code path. Useful for debugging.") @@ -334,6 +342,8 @@ integration_params.add_option("--internal-gmm-max-components",type=int,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=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("--q-time-pregrid-factor", default=1, type=int, + help="OPT-IN ordinary-NoLoop Q pregrid. Value 8 reflects each finite cut Q window, FFT-interpolates it onto an 8x finer grid once after packing, and uses four-tap cubic interpolation for detector arrival times while leaving the geocentric time-integration grid at the data deltaT. Default 1 preserves current behavior and memory. Other factors are refused until separately validated.") 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") @@ -500,6 +510,20 @@ else: "--interpolate-time: unrecognised value %r. Use a stencil name (nearest|cubic|sinc) " "or a legacy boolean (%s)." % (opts.interpolate_time, "|".join(_TI_LEGACY_BOOLEAN))) +opts.q_time_pregrid_factor = validate_q_time_pregrid_factor(opts.q_time_pregrid_factor) +if opts.q_time_pregrid_factor == 8: + if not opts.vectorized or opts.rotation_slow or opts.freqresponse or opts.calibration_envelope_directory: + raise NotImplementedError( + "--q-time-pregrid-factor 8 is currently restricted to ordinary vectorized " + "NoLoop without rotation, frequency-dependent response, or calibration marginalization") + if not opts._interp_time_from_default and opts._noloop_time_interp != "cubic": + raise ValueError( + "--q-time-pregrid-factor 8 uses four-tap cubic interpolation; remove the " + "explicit --interpolate-time option or set it to cubic") + opts._q_pregrid_fallback_interp = opts._noloop_time_interp + opts._noloop_time_interp = "cubic" + print(" Q_lm pregrid: ENABLED factor=8 boundary=even-reflection arrival_stencil=cubic " + "integration_grid=unchanged") # 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 @@ -545,8 +569,6 @@ if opts.freqresponse: # Path D finite-size response: wired into BOTH the CPU-vectorized and GPU (xpy) branches; # the NoLoop reuses the baseline fused Q_inner_product kernel per basis weight p on GPU # (same memory footprint as the baseline), mirroring --rotation-slow. - if opts.rotation_slow: - raise ValueError("--freqresponse and --rotation-slow both replace the likelihood; use at most one") if not opts.vectorized: raise ValueError("--freqresponse requires --vectorized") # Same never-firing calibration guard as --rotation-slow above, with the same history; @@ -602,6 +624,7 @@ if not(opts.use_gwsignal): import RIFT.likelihood.factored_likelihood as factored_likelihood import RIFT.likelihood.factored_likelihood_with_rotation as factored_likelihood_with_rotation import RIFT.likelihood.factored_likelihood_freqresponse as factored_likelihood_freqresponse +import RIFT.likelihood.factored_likelihood_rotating_freqresponse as factored_likelihood_rotating_freqresponse if opts.use_gwsignal and not(factored_likelihood.has_GWS): print(" HARD FAILURE: this node could not import gwsignal ! ") @@ -880,6 +903,52 @@ if opts._time_quadrature != 'simpson' and _tq_missing: "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))) +# --psi-marginalization: analytic polarization-angle marginalization, reachable only on the +# legacy SCALAR likelihood path (factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized +# has no vectorized/GPU/NoLoop counterpart) and only alone -- it marginalizes psi ALONE, with no +# joint sum over time, so it cannot stand in for --time-marginalization's likelihood. Refuse +# rather than silently ignore, same discipline as --time-marginalization-quadrature above. +if opts.psi_marginalization: + _psi_marg_prereqs = ( + ('not --time-marginalization (the analytic psi marginal has no joint time-sum form)', + not bool(opts.time_marginalization)), + ('not --vectorized (the analytic marginal exists only on the legacy scalar path)', + not bool(opts.vectorized)), + ('not --gpu (same restriction: scalar path only)', not bool(opts.gpu)), + ('not --distance-marginalization (the scalar path this needs does not implement the ' + 'distance lookup table)', not bool(opts.distance_marginalization)), + ('not --rotation-slow (separate likelihood, not audited for this)', + 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)), + ('default --interpolate-time (the scalar path always evaluates rholm(t) through its ' + 'own interpolator, independent of the sinc/cubic/nearest NoLoop stencil choice, so an ' + 'explicit request here would be silently ignored)', opts.interpolate_time is None), + ('not --sampler-method GMM/portfolio (their adaptive phi/psi pairing indexes "psi" ' + 'directly in sampler.params, which this option removes -- ValueError at sampler setup, ' + 'not audited as a combination)', + opts.sampler_method not in ('GMM', 'portfolio')), + ('not --internal-rotate-phase (its phi_orb/psi joint reparam reads a per-sample psi ' + 'that no longer exists; the exported "psi"/"polarization" columns would be a silently ' + 'wrong reconstruction rather than the fiducial placeholder this option writes)', + not bool(opts.internal_rotate_phase)), + ('not --limit-psi (there is no psi sampler left to apply the box to; the box would be ' + 'silently ignored)', not bool(opts.limit_psi)), + ('not --zero-likelihood (its debug likelihood keys results off a "psi" kwarg that this ' + 'option removes from the call signature)', not bool(opts.zero_likelihood)), + ) + _psi_marg_missing = [name for name, ok in _psi_marg_prereqs if not ok] + if _psi_marg_missing: + raise ValueError( + "--psi-marginalization was requested, but this configuration cannot honour it: %s. " + "Refusing rather than silently running the ordinary sampled-psi likelihood while " + "reporting that psi was analytically marginalized." % "; ".join(_psi_marg_missing)) + print(" Polarization angle: ANALYTICALLY MARGINALIZED (--psi-marginalization). psi is " + "removed from the sampled extrinsic parameters; " + "NetworkLogLikelihoodPolarizationMarginalized replaces it with an exact quadrature " + "over its uniform [0, pi) prior at each extrinsic point.") if opts._time_quadrature == 'bandlimited': if opts.time_posterior_export == 'grid': raise ValueError( @@ -1773,17 +1842,49 @@ else: # Psi -- polarization angle # sampler: uniform in [0, pi) # -psi_sampler = mcsampler.ret_uniform_samp_vector_alt( - param_limits["psi"][0], param_limits["psi"][1]) -psi_sampler_cdf_inv = functools.partial(mcsampler.uniform_samp_cdf_inv_vector, - param_limits["psi"][0], param_limits["psi"][1]) -sampler.add_parameter("psi", - pdf = psi_sampler, - cdf_inv = psi_sampler_cdf_inv, - left_limit = param_limits["psi"][0], - right_limit = param_limits["psi"][1], - prior_pdf = mcsampler.uniform_samp_psi, - adaptive_sampling=opts.internal_rotate_phase or opts.force_adapt_all) +if not opts.psi_marginalization: + psi_sampler = mcsampler.ret_uniform_samp_vector_alt( + param_limits["psi"][0], param_limits["psi"][1]) + psi_sampler_cdf_inv = functools.partial(mcsampler.uniform_samp_cdf_inv_vector, + param_limits["psi"][0], param_limits["psi"][1]) + sampler.add_parameter("psi", + pdf = psi_sampler, + cdf_inv = psi_sampler_cdf_inv, + left_limit = param_limits["psi"][0], + right_limit = param_limits["psi"][1], + prior_pdf = mcsampler.uniform_samp_psi, + adaptive_sampling=opts.internal_rotate_phase or opts.force_adapt_all) +# else: psi is not a sampled dimension at all -- --psi-marginalization integrates it out +# analytically inside the likelihood (see the likelihood_function branch below), matching the +# --distance-marginalization precedent of skipping add_parameter entirely for a parameter that +# is marginalized rather than sampled. +# +# EVIDENCE NEUTRALITY. The psi prior this driver USES is not normalized: the sampled path puts +# uniform_samp_psi = 1/pi over param_limits["psi"] = (0, 2 pi), so its psi prior integrates to 2, +# and every ordinary ILE lnZ on this path carries that +ln 2. NetworkLogLikelihoodPolarizationMarginalized +# instead returns the NORMALIZED marginal, (1/pi) int_0^pi exp(lnL) dpsi, which integrates to 1. +# Without the correction below the flag would report lnL exactly ln 2 = 0.693 nat BELOW an +# otherwise identical sampled-psi run, so marginalized and sampled rows could not share an +# all.net. Measured on two fixtures, AV, 3 seeds each: 0.7027 and 0.6602 +/- 0.036. +# Derived from the sampler's OWN prior and limits rather than written as a literal ln 2, so that +# changing either (e.g. --internal-rotate-phase's (0, 4 pi) psi range) cannot silently desync it. +# NOTE FOR REVIEW: this matches the incumbent, it does not decide the convention. The RIFT JAX +# driver uses psi in [0, pi]; until one convention is chosen, evidences compared ACROSS the two +# drivers differ by ln 2 regardless of this flag. +psi_marginalization_ln_prior_mass = 0.0 +if opts.psi_marginalization: + _psi_prior_density = float(numpy.atleast_1d( + mcsampler.uniform_samp_psi(numpy.atleast_1d( + 0.5*(param_limits["psi"][0]+param_limits["psi"][1]))))[0]) + _psi_prior_mass = _psi_prior_density*(param_limits["psi"][1]-param_limits["psi"][0]) + if not (_psi_prior_mass > 0): + raise ValueError("--psi-marginalization: psi prior mass {} is not positive".format(_psi_prior_mass)) + psi_marginalization_ln_prior_mass = float(numpy.log(_psi_prior_mass)) + print(" --psi-marginalization: sampled-path psi prior mass is {:.6f} over ({:.4f}, {:.4f}) at " + "density {:.6f}; adding ln(mass) = {:+.6f} nat to the normalized analytic marginal so " + "lnZ matches an otherwise identical sampled-psi run.".format( + _psi_prior_mass, param_limits["psi"][0], param_limits["psi"][1], + _psi_prior_density, psi_marginalization_ln_prior_mass)) # # Phi - orbital phase @@ -2088,7 +2189,7 @@ if not opts.time_marginalization: param_limits["t_ref"][0], param_limits["t_ref"][1]) sampler.add_parameter("t_ref", pdf = tref_sampler, - cdf_inv = None, + cdf_inv = tref_sampler_cdf_inv, left_limit = param_limits["t_ref"][0], right_limit = param_limits["t_ref"][1], # Reuse the backend-portable closure above. AV exposes @@ -2124,6 +2225,20 @@ if opts.skymap_file: # pinned_params = get_pinned_params(opts) unpinned_params = get_unpinned_params(opts, sampler.params) +if opts.psi_marginalization: + # psi is integrated out analytically, so it must not ALSO be an integration dimension. + # It is absent because add_parameter("psi", ...) above is skipped; this is the live check + # of that contract, at the one place unpinned_params is actually built for this path. + # (A `unpinned_params.remove('psi')` used to sit in the GMM-only branch below, which this + # flag refuses outright, so it could never run.) + _psi_still_sampled = [p for p in sampler.params + if p == 'psi' or (isinstance(p, tuple) and 'psi' in p)] + if _psi_still_sampled or 'psi' in unpinned_params: + raise ValueError( + "--psi-marginalization: 'psi' is still a sampled dimension ({}); it would be " + "integrated twice.".format(_psi_still_sampled or sorted(unpinned_params))) + print(" --psi-marginalization: psi is NOT a sampled dimension; integrating over {}".format( + sorted(str(p) for p in unpinned_params))) print( "{0:<25s} {1:>5s} {2:>5s} {3:>20s} {4:<10s}".format("parameter", "lower limit", "upper limit", "pinned?", "pin value")) plen = len(sorted(sampler.params, key=lambda p: len(p))[-1]) for p in sampler.params: @@ -3486,30 +3601,101 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t rholmArrayDict={} rholms_intpArrayDict={} epochDict={} + q_deltaT = float(P.deltaT) + _q_pregrid_reports = [] for det in rholms_intp.keys(): print( " Packing ", det) lookupNKDict[det],lookupKNDict[det], lookupKNconjDict[det], ctUArrayDict[det], ctVArrayDict[det], rholmArrayDict[det], rholms_intpArrayDict[det], epochDict[det] = factored_likelihood.PackLikelihoodDataStructuresAsArrays( rholms[det].keys(), rholms_intp[det], rholms[det], cross_terms[det],cross_terms_V[det]) if _have_cal_crossterms: ctUArrayDict_cal[det], ctVArrayDict_cal[det] = factored_likelihood.PackCalCrossTermsAsArrays( list(rholms[det].keys()), lookupKNDict[det], cross_terms_cal[det], cross_terms_cal_V[det]) - if opts.gpu and (not xpy_default is np): + if opts.q_time_pregrid_factor == 8: + _q_transfer = cupy.asarray if opts.gpu and (not xpy_default is np) else None + _q_cleanup = (lambda: cupy.get_default_memory_pool().free_all_blocks()) \ + if _q_transfer is not None else None + rholmArrayDict, _q_pregrid_reports, _q_pregrid_error = \ + factored_likelihood.prepare_reflected_q_pregrid( + rholmArrayDict, factor=8, transfer=_q_transfer, cleanup=_q_cleanup) + if _q_pregrid_error is None: + q_deltaT = float(P.deltaT) / 8.0 + print(" Q_lm pregrid telemetry: status=active q_deltaT={:.12g} input_bytes={} " + "retained_bytes={} peak_allocation_bytes={} max_roundtrip={:.3g}".format( + q_deltaT, + sum(item['input_bytes'] for item in _q_pregrid_reports), + sum(item['retained_bytes'] for item in _q_pregrid_reports), + max(item['peak_allocation_bytes'] for item in _q_pregrid_reports), + max(item['roundtrip_max'] for item in _q_pregrid_reports))) + else: + q_deltaT = float(P.deltaT) + opts.q_time_pregrid_factor = 1 + opts._noloop_time_interp = opts._q_pregrid_fallback_interp + print(" Q_lm pregrid telemetry: status=fallback reason={!r} q_deltaT={:.12g} " + "arrival_stencil={}".format( + _q_pregrid_error, q_deltaT, opts._noloop_time_interp)) + if opts.gpu and (not xpy_default is np): + for det in rholmArrayDict: lookupNKDict[det] = cupy.asarray(lookupNKDict[det]) - rholmArrayDict[det] = cupy.asarray(rholmArrayDict[det]) + # Q was transferred inside the pregrid transaction. The + # default/fallback path still needs its ordinary transfer. + if opts.q_time_pregrid_factor != 8 and not isinstance(rholmArrayDict[det], cupy.ndarray): + rholmArrayDict[det] = cupy.asarray(rholmArrayDict[det]) ctUArrayDict[det] = cupy.asarray(ctUArrayDict[det]) ctVArrayDict[det] = cupy.asarray(ctVArrayDict[det]) epochDict[det] = cupy.asarray(epochDict[det]) if _have_cal_crossterms: ctUArrayDict_cal[det] = cupy.asarray(ctUArrayDict_cal[det]) ctVArrayDict_cal[det] = cupy.asarray(ctVArrayDict_cal[det]) + # NoLoop keeps P.deltaT as the geocentric integration spacing and reads + # this independent spacing only for Q-grid coordinates. + P.q_deltaT = q_deltaT # Pass None (not empty dicts) downstream when the fix is inactive, so the # likelihood keeps its exact cal-independent behavior. if not _have_cal_crossterms: ctUArrayDict_cal = None ctVArrayDict_cal = None + # Combined Path A/B+D: compose finite-frequency basis weights with the sidereal + # modulation and delay-derivative operators. This branch must precede the two + # individual response branches: both flags now request one compound physical model. + rotating_freqresponse_data = None + if opts.rotation_slow and opts.freqresponse: + _pmax = int(opts.rotation_p_max) + _qmax = int(opts.freqresponse_qmax) + _arm = opts.freqresponse_arm_length + if _arm is not None: + if '=' in str(_arm): + _arm = {kv.split('=')[0]: float(kv.split('=')[1]) + for kv in str(_arm).split(',')} + else: + _arm = float(_arm) + _rint_rf, _ct_rf, _ctV_rf, _rho_rf, _meta_rf = \ + factored_likelihood_rotating_freqresponse.PrecomputeLikelihoodTermsRotatingFreqResponse( + fiducial_epoch, t_window, P, data_dict, psd_dict, opts.l_max, fmax, + Qmax=_qmax, L_arm=_arm, p_max=_pmax, analyticPSD_Q=False, + inv_spec_trunc_Q=inv_spec_trunc_Q, T_spec=T_spec, + verbose=opts.verbose, quiet=not opts.verbose, skip_interpolation=True) + _lkRF, _rhoA, _uAA, _vAA, _epRF = \ + factored_likelihood_rotating_freqresponse.pack_rotating_freqresponse_arrays( + _meta_rf, _rho_rf, _ct_rf, _ctV_rf) + if opts.gpu and (not xpy_default is np): + for _det in _rhoA: + for _a in _rhoA[_det]: + _rhoA[_det][_a] = cupy.asarray(_rhoA[_det][_a]) + _uAA[_det] = cupy.asarray(_uAA[_det]) + _vAA[_det] = cupy.asarray(_vAA[_det]) + rotating_freqresponse_data = dict( + meta=_meta_rf, lookupNKDict=_lkRF, rho_by_a=_rhoA, + U_by_aa=_uAA, V_by_aa=_vAA, epochDict=_epRF) + _nbasis = len(_meta_rf['a_list']) + print(" [rotation-slow+freqresponse] compound precompute complete; " + "p_max", _pmax, "Qmax", _qmax, "basis elements", _nbasis, + "ordered U/V pairs", _nbasis * _nbasis, "arm-length", + opts.freqresponse_arm_length, + "(GPU)" if (opts.gpu and not xpy_default is np) else "(CPU)") + # [Path A] slow-rotation precompute: build the harmonic-indexed bank and pack it. rotation_slow_data = None - if opts.rotation_slow: + if opts.rotation_slow and not opts.freqresponse: _pmax = int(opts.rotation_p_max) # --rotation-n-harmonics is a FLOOR, not the literal width: the response # coefficients C_{(p,ntilde)} reach |ntilde| <= 2 + p_max (issue #142), and the @@ -3547,7 +3733,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t # [Path D] finite-size (frequency-dependent) response precompute: fold each W_p(f) # into the modes once and pack the response-basis overlap bank. freqresponse_data = None - if opts.freqresponse: + if opts.freqresponse and not opts.rotation_slow: _qmax = int(opts.freqresponse_qmax) _arm = opts.freqresponse_arm_length if _arm is not None: @@ -3748,6 +3934,58 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t # Likelihood if not opts.time_marginalization: + if opts.psi_marginalization: + + # psi is not one of the sampled parameters (see the sampler.add_parameter("psi", ...) + # guard above) -- it is integrated out analytically at every extrinsic point by + # factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized, which sums the + # closed-form quadrature over its uniform [0, pi) prior. The value handed to the + # 'psi' argument below is a placeholder: the antenna pattern is periodic with period + # pi (F(psi) = F(0)*exp(-2i*psi)) and the function integrates over one full period, so + # its return value does not depend on which reference psi the placeholder names -- + # verified in test/test_psi_marginalization.py. + _psi_marg_detectors = list(rholms_intp.keys()) + _psi_marg_ln_prior_mass = psi_marginalization_ln_prior_mass + + def likelihood_function(right_ascension, declination, t_ref, phi_orb, + inclination, distance): + + dec = numpy.copy(declination).astype(numpy.float64) + if opts.declination_cosine_sampler: + dec = numpy.pi/2 - numpy.arccos(dec) + incl = numpy.copy(inclination).astype(numpy.float64) + if opts.inclination_cosine_sampler: + incl = numpy.arccos(incl) + if opts.d_prior_redshift: + distance = redshift_to_distance(distance) + + # use EXTREMELY many bits + lnL = numpy.zeros(right_ascension.shape,dtype=RiftFloat) + i = 0 + for ph, th, tr, phr, ic, di in zip(right_ascension, dec, + t_ref, phi_orb, incl, distance): + P.phi = ph # right ascension + P.theta = th # declination + P.tref = fiducial_epoch + tr # ref. time (rel to epoch for data taking) + P.phiref = phr # ref. orbital phase + P.incl = ic # inclination + P.dist = di* 1.e6 * lalsimutils.lsu_PC # luminosity distance + + lnL[i] = factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized( + fiducial_epoch, rholms_intp, cross_terms, cross_terms_V, + P.tref, P.phi, P.theta, P.incl, P.phiref, 0.0, P.dist, + opts.l_max, _psi_marg_detectors) + i+=1 + # restore the prior mass the SAMPLED psi path carries (see the derivation at the + # skipped sampler.add_parameter("psi") above): the analytic marginal is normalized, + # this driver's psi prior is not. Without it the flag reports lnZ - ln 2. + lnL += _psi_marg_ln_prior_mass + if return_lnL: + return lnL - manual_avoid_overflow_logarithm + return numpy.exp(lnL - manual_avoid_overflow_logarithm) + + else: + def likelihood_function(right_ascension, declination, t_ref, phi_orb, inclination, psi, distance): @@ -3778,7 +4016,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t opts.l_max) i+=1 if return_lnL: - return lnL - manual_avoid_overflow_logarithm + return lnL - manual_avoid_overflow_logarithm return numpy.exp(lnL - manual_avoid_overflow_logarithm) else: # Sum over time for every point in other extrinsic params @@ -3862,7 +4100,13 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t P.psi= psi_true P.phiref = phi_orb_true - if opts.rotation_slow: + if opts.rotation_slow and opts.freqresponse: + lnL = factored_likelihood_rotating_freqresponse.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, P, rotating_freqresponse_data['meta'], rotating_freqresponse_data['lookupNKDict'], + rotating_freqresponse_data['rho_by_a'], rotating_freqresponse_data['U_by_aa'], + rotating_freqresponse_data['V_by_aa'], rotating_freqresponse_data['epochDict'], + Lmax=opts.l_max, time_interp=opts._noloop_time_interp) + elif opts.rotation_slow: lnL = factored_likelihood_with_rotation.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( tvals, P, rotation_slow_data['meta'], rotation_slow_data['lookupNKDict'], rotation_slow_data['rho_by_n'], rotation_slow_data['U_by_nn'], @@ -3930,7 +4174,13 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t P.phiref = phi_orb_true - if opts.rotation_slow: + if opts.rotation_slow and opts.freqresponse: + lnL = factored_likelihood_rotating_freqresponse.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, P, rotating_freqresponse_data['meta'], rotating_freqresponse_data['lookupNKDict'], + rotating_freqresponse_data['rho_by_a'], rotating_freqresponse_data['U_by_aa'], + rotating_freqresponse_data['V_by_aa'], rotating_freqresponse_data['epochDict'], + Lmax=opts.l_max, time_interp=opts._noloop_time_interp, xpy=xpy_default) + elif opts.rotation_slow: lnL = factored_likelihood_with_rotation.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( tvals, P, rotation_slow_data['meta'], rotation_slow_data['lookupNKDict'], rotation_slow_data['rho_by_n'], rotation_slow_data['U_by_nn'], @@ -4037,10 +4287,17 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t print( " Using direct phase marginalization ") for det in lookupNKDict: - if set((lm[0], lm[1]) for lm in lookupNKDict[det]) != {(2, 2), (2, -2)}: + # ``lookupNKDict`` is moved to CuPy above. Iterating its rows + # yields device arrays (and, with current CuPy, unhashable + # zero-dimensional array elements). The inverse lookup stays + # on the host and already has canonical ``(l,m)`` tuple keys, + # so use it for this structural identity check. This avoids a + # device round trip and is representation-independent. + modes_here = set(lookupKNDict[det]) + if modes_here != {(2, 2), (2, -2)}: raise Exception( " Phase marginalization is implemented only for 2-2 modes, " - f"while the modes consired here are {lookupNKDict[det]}." + f"while the modes considered here are {sorted(modes_here)}." ) def likelihood_function(right_ascension, declination, inclination, psi): @@ -5285,6 +5542,31 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t samples, sampler, min(opts.fairdraw_extrinsic_output_n_max, opts.n_eff), convert=identity_convert, use_lnL=rvs_integrand_is_lnL) + # Polarization column under --psi-marginalization. There is no per-sample psi draw: + # NetworkLogLikelihoodPolarizationMarginalized integrated psi out, it never sampled one. + # Write NaN, NOT a fiducial 0.0. bin/convert_output_format_ile2inference copies this + # column verbatim into the PE-samples 'psi' column under a header that does not mark it, + # and the phi_orb/distance precedents for a fiducial value are always followed by a + # resampling step that psi has none of -- so a fiducial 0.0 reaches a consumer as a delta + # function at 0 that reads like a polarization MEASUREMENT. NaN cannot. + # Shaped off 'right_ascension' (always sampled on this path) rather than off 'psi', since + # 'psi' is exactly the key that is missing. + if opts.psi_marginalization: + if "psi" in samples: + raise ValueError( + "--psi-marginalization: the sampler returned a 'psi' column, but psi was supposed " + "to be integrated out analytically and never sampled. Refusing to export a " + "polarization column whose provenance is unknown.") + samples["psi"] = np.full_like( + np.asarray(samples["right_ascension"], dtype=np.float64), np.nan) + print(" --psi-marginalization: exported 'psi'/'polarization' column is NaN -- psi was " + "marginalized, not sampled; there is nothing to report per sample.") + elif "psi" not in samples: + # Never silently invent one: every other route to this block samples psi, so a missing + # column here is a wiring bug, and the old fallback would have written a fake 0.0. + raise KeyError( + "no 'psi' column in the retained samples, and --psi-marginalization was not requested; " + "refusing to write a fabricated polarization column") # Insert reference distance if it was marginalized over if "distance" not in samples: # Not distance output is the same as internal calculations: in *Mpc* diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa index 8e3c929b5..889628840 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode_lisa @@ -1163,7 +1163,7 @@ if not opts.time_marginalization: param_limits["t_ref"][0], param_limits["t_ref"][1]) sampler.add_parameter("t_ref", pdf = tref_sampler, - cdf_inv = None, + cdf_inv = tref_sampler_cdf_inv, 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])) diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax index 00be1a6f8..1c5248ba7 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax @@ -64,19 +64,52 @@ from optparse import OptionParser, OptionGroup import numpy as np import jax +import jax.numpy as jnp jax.config.update("jax_enable_x64", True) +# Configure the persistent cache before importing modules that construct ILE +# JITs. RIFT selects a compatibility-keyed child directory, so a shared cache +# root is safe across heterogeneous GPU/JAX installations. +from RIFT.jax_cache import (argv_option, configure_persistent_cache, + import_bundle, runtime_compatibility) +_JAX_CACHE_DIR = configure_persistent_cache(jax, sys.argv[1:]) +_JAX_CACHE_BUNDLE = (argv_option(sys.argv[1:], "--jax-cache-bundle") + or os.environ.get("RIFT_JAX_CACHE_BUNDLE")) +_JAX_CACHE_PROFILE = argv_option(sys.argv[1:], "--jax-cache-profile") +if _JAX_CACHE_BUNDLE: + if _JAX_CACHE_DIR is None: + raise RuntimeError("a cache bundle was requested but no writable JAX cache is available") + import_bundle(_JAX_CACHE_BUNDLE, _JAX_CACHE_DIR.parent, + runtime_compatibility(jax), _JAX_CACHE_PROFILE, + destination=_JAX_CACHE_DIR) + import lal import lalsimulation as lalsim import RIFT.lalsimutils as lalsimutils -from RIFT.likelihood.jax_ile import build_data_from_precompute +from RIFT.likelihood.jax_ile import ( + build_data_from_precompute, build_rotation_data_from_precompute, + build_freqresponse_data_from_precompute, + build_rotating_freqresponse_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.samplers import regularize_cov as _regularize_cov +# ONE definition, imported rather than re-typed: a second copy of the +# "is this evidence trustworthy" rule is a copy that drifts. +from RIFT.likelihood.jax_ile.samplers import _finalize_evidence 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) + ANGLE_MARG_CHOICES, + ANGLE_MARG_CROSSOVER_AMPLITUDE) +from RIFT.likelihood.jax_ile.direct_marginalization_policy import ( + POLICY_CHOICES as DIRECT_MARG_POLICY_CHOICES, + POLICY_DEFAULT as DIRECT_MARG_POLICY_DEFAULT, + PolicyConfig as DirectMargPolicyConfig, + RESERVE_SCHEME_CHOICES as DIRECT_MARG_RESERVE_SCHEME_CHOICES, + reserve_pair as _direct_marg_reserve_pair, + summarize_policy_ledger as _summarize_policy_ledger, + predict_reserve_pair, format_reserve_pair, RESERVE_SCHEME_EXECUTABLE) _JAX_GATHERER_NAMES = tuple(_JAX_GATHERERS) from RIFT.likelihood.jax_ile.wrapper import ( JAXExtrinsicLikelihood, JAXDistanceMarginalizedLikelihood, @@ -125,6 +158,22 @@ _TEMPERED_MODES = frozenset(( # 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")) +# Modes whose blind full-sky prior draws go through eval_lnL, which stops the +# run on one uncertified bandlimited row. The flowMC family and multistart-nuts +# pilot through the samplers' own draw and are not on this list. +_EVAL_LNL_STOP_MODES = frozenset(("prior-mc", "laplace-is", "map", "nuts")) +# Largest shift a wrong-sky draw makes to one detector's arrival time: the +# window must contain it or a blind row's arrival peak sits beyond the edge, +# which no certificate can converge on (measured 35/256 rows at 20 ms, 0/256 at +# 50 ms; DESIGN_jax_bandlimited_distmarg.md, "The endpoint certificate"). +_BANDLIMITED_FULLSKY_HALF_WINDOW_MIN = 2.0 * lal.REARTH_SI / lal.C_SI +# Modes whose posterior is already phi_ref-marginalised ANALYTICALLY (a grid +# sum baked into the likelihood, not a sampled axis): --phase-marginalization +# has nothing left to do on these and is reported IGNORED rather than refused +# (RO'S 2026-09-08: missing knobs are fine to no-op for compatibility, but +# must warn). +_PHASE_ANALYTIC_MODES = frozenset(( + "flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg", "nuts-phimarg")) # Boolean (zero-argument) ILE options (action=store_true/false). _ILE_BOOL_OPTS = { @@ -161,6 +210,7 @@ _ILE_BOOL_OPTS = { "--internal-reparam-dl-incl", "--internal-use-lnL", "--distance-slice-all-fresh", "--distance-slice-randomize", + "--psi-marginalization", } # ILE options taking repeated values (action=append). _ILE_APPEND_OPTS = { @@ -234,7 +284,7 @@ _ILE_ALL_OPTS = { "--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", + "--psd-window-shape", "--psi-marginalization", "--random-event", "--reference-freq", "--resample-time-marginalization", "--restricted-mode-list-file", "--rom-group", "--rom-integrate-intrinsic", "--rom-limit-basis-size-to", "--reject-collapsed-live-volume", "--rom-param", "--rom-use-basis", @@ -295,6 +345,39 @@ def check_critical_and_report(opts, optp): return bool(v) if name in _ILE_BOOL_OPTS else (v is not None and v != []) fatal = [] + # --distance-gh-nodes / JAX_ILE_DISTMARG_GH resolution (RO'S 2026-09-08, + # revised same day per adversarial review: the option must default to + # None, not 0, or an explicit ``--distance-gh-nodes 0`` is indistinguish- + # able from "not passed" and a nonzero JAX_ILE_DISTMARG_GH silently wins + # -- the exact opposite of the documented "CLI wins"). The per-sample + # Gauss-Hermite distance quadrature must be reachable by an ILE argument, + # not only the environment variable core.py reads at import. Resolution: + # the CLI flag, when GIVEN (including an explicit 0), always wins over + # the environment variable; the environment variable is consulted only + # when the CLI flag was not given at all. The two are REFUSED, not + # silently reconciled, whenever BOTH are given and differ -- including an + # explicit CLI 0 against a nonzero env, which is a conflict like any + # other: silently keeping either value would hide the disagreement. + # Only the RESOLUTION happens here (needed below by the + # distance-grid-scheme combination check); the actual global mutation and + # banner are deferred until after the fatal gate below, so a refused + # command line never mutates core's module-level GH state. + _gh_cli = getattr(opts, "distance_gh_nodes", None) + _gh_env_raw = os.environ.get("JAX_ILE_DISTMARG_GH") + _gh_env = int(_gh_env_raw) if _gh_env_raw not in (None, "") else None + if _gh_cli is not None and _gh_cli < 0: + fatal.append("--distance-gh-nodes must be >= 0, got %d" % _gh_cli) + elif _gh_cli is not None and _gh_env is not None and _gh_cli != _gh_env: + fatal.append( + "--distance-gh-nodes %d conflicts with JAX_ILE_DISTMARG_GH=%d in " + "the environment; unset the environment variable or pass the " + "same value on the command line" % (_gh_cli, _gh_env)) + if _gh_cli is not None: + _gh_resolved = _gh_cli + elif _gh_env is not None: + _gh_resolved = _gh_env + else: + _gh_resolved = 0 if getattr(opts, "calibration_n_realizations", None) not in (None, 1) \ or is_set("--calibration-export-posterior") \ or is_set("--calibration-envelope-directory"): @@ -311,6 +394,189 @@ 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)") + _is_n = int(getattr(opts, "smc_is_samples", 60000)) + if _is_n < 0: + fatal.append("--smc-is-samples must be >= 0; a negative count reaches " + "the SMC proposal draw, whose exception handler would " + "silently fall back to the raw SMC evidence") + _sampler_method = getattr(opts, "sampler_method", None) + _jax_av_active = _sampler_method in ("AV", "portfolio") + if _jax_av_active: + try: + resolve_av_angular_limits(opts) + except SystemExit as exc: + fatal.append(str(exc).strip()) + if getattr(opts, "mode", None) in ("map", "nuts", "multistart-nuts", + "nuts-phimarg"): + fatal.append("--sampler-method %s cannot override the chain/optimizer " + "mode --mode %s; use laplace-is or a flowmc-* mode to " + "select the desired JAX likelihood geometry" + % (_sampler_method, getattr(opts, "mode", None))) + if int(getattr(opts, "n_eff", None) or 1000) <= 0: + fatal.append("--n-eff must be positive for --sampler-method %s" + % _sampler_method) + if int(getattr(opts, "jax_av_seed_pilot", 4000)) <= 0: + fatal.append("--jax-av-seed-pilot must be positive") + if int(getattr(opts, "jax_av_seed_modes", 4)) <= 0: + fatal.append("--jax-av-seed-modes must be positive") + _seed_points = getattr(opts, "jax_av_seed_points", None) + if _seed_points is not None and int(_seed_points) <= 0: + fatal.append("--jax-av-seed-points must be positive") + _eval_chunk = getattr(opts, "jax_av_eval_chunk", None) + if _eval_chunk is not None and int(_eval_chunk) <= 0: + fatal.append("--jax-av-eval-chunk must be positive") + if float(getattr(opts, "jax_av_sky_inflate", 2.0)) <= 0: + fatal.append("--jax-av-sky-inflate must be positive") + _prior_frac = float(getattr(opts, "jax_av_seed_prior_frac", 0.1)) + if not (0.0 <= _prior_frac <= 1.0): + fatal.append("--jax-av-seed-prior-frac must lie in [0,1]") + _members = [s.strip().upper() + for item in (getattr(opts, "sampler_portfolio", None) or []) + for s in str(item).split(",") if s.strip()] + if _sampler_method == "AV" and _members: + fatal.append("--sampler-portfolio is inert with --sampler-method AV") + if any(name not in ("AV", "GMM") for name in _members): + fatal.append("the JAX portfolio supports only AV and GMM members; got %s" + % ",".join(_members)) + if getattr(opts, "jax_av_seed", "none") == "none": + for _name in ("--jax-av-seed-pilot", "--jax-av-seed-modes", + "--jax-av-seed-points", "--jax-av-sky-inflate", + "--jax-av-seed-prior-frac"): + if was_supplied(opts, _name): + fatal.append("%s is inert with --jax-av-seed none" % _name) + for _name in ("--auto-adapt-weight-exponent", "--target-export-ess-frac", + "--allow-degenerate-tempering", "--adapt-adapt", + "--temper-init", "--temper-ess-frac", + "--temper-max-stages", "--temper-max-dbeta", + "--fisher-precondition", "--fisher-is-samples", + "--smc-puffball", "--smc-walkers", "--smc-move-steps", + "--smc-is-samples", "--smc-puff-scale"): + if was_supplied(opts, _name): + fatal.append("%s configures the mode's gradient/SMC sampler and " + "is inert under --sampler-method %s" + % (_name, _sampler_method)) + else: + for _name in ("--sampler-portfolio", "--sampler-anisotropic-bins", + "--n-eff"): + if was_supplied(opts, _name): + fatal.append("%s is inert unless --sampler-method AV/portfolio" + % _name) + for _name in ("--jax-av-seed", "--jax-av-seed-pilot", + "--jax-av-seed-modes", "--jax-av-seed-points", + "--jax-av-sky-inflate", "--jax-av-seed-prior-frac", + "--jax-av-eval-chunk"): + if was_supplied(opts, _name): + fatal.append("%s requires --sampler-method AV/portfolio" % _name) + # ``was_supplied`` deliberately fails open when an in-process caller + # has no argv provenance. The non-default value itself is conclusive. + if getattr(opts, "jax_av_seed", "none") != "none" \ + and not any("--jax-av-seed requires" in item for item in fatal): + fatal.append("--jax-av-seed requires --sampler-method AV/portfolio") + if is_set("--psi-marginalization"): + fatal.append("--psi-marginalization (analytic polarization-angle " + "marginalization) is not implemented for psi-sampling modes " + "on this driver; use --mode flowmc-phipsimarg or " + "flowmc-dpsimarg, which marginalize psi by construction") + # Cross-axis policy scope, at parse time (external review of #278, P1): + # only --mode flowmc-phipsimarg reads the policy, so a request anywhere + # else would otherwise complete on the ordinary likelihood, and the + # policy's knobs are inert unless the policy is on. + _policy = getattr(opts, "direct_marginalization_policy", + DIRECT_MARG_POLICY_DEFAULT) + _policy_knobs = ("--direct-marginalization-time-guard", + "--direct-marginalization-reserve-time-refine", + "--direct-marginalization-reserve-time-refine-max", + "--direct-marginalization-error-budget-nats", + "--direct-marginalization-batch-rows", + "--direct-marginalization-policy-probe-rows", + "--direct-marginalization-policy-probe-only", + "--direct-marginalization-max-modes", + "--direct-marginalization-enriched-max-modes", + "--direct-marginalization-base-oversample", + "--direct-marginalization-enriched-oversample", + "--direct-marginalization-max-starts", + "--direct-marginalization-max-time-nodes", + "--direct-marginalization-convergence-tol-nats", + "--direct-marginalization-time-guard-tol-nats") + if _policy != "off": + if getattr(opts, "mode", None) != "flowmc-phipsimarg": + fatal.append("--direct-marginalization-policy %s applies only to " + "--mode flowmc-phipsimarg; --mode %s would run the " + "ordinary likelihood and silently ignore the request" + % (_policy, getattr(opts, "mode", None))) + _g = int(getattr(opts, "direct_marginalization_time_guard", 16)) + if _g < 2: + fatal.append("--direct-marginalization-time-guard must be >= 2 " + "(the two-guard comparison needs a half guard)") + _f = int(getattr(opts, "direct_marginalization_reserve_time_refine", 4)) + if _f < 2 or _f % 2: + fatal.append("--direct-marginalization-reserve-time-refine must be " + "an even integer >= 2 (the check rule is the " + "half-refined rule)") + _fm = int(getattr(opts, "direct_marginalization_reserve_time_refine_max", + 32)) + if _fm < _f or _fm % 2: + fatal.append("--direct-marginalization-reserve-time-refine-max must " + "be an even integer >= the refine factor") + _b = float(getattr(opts, "direct_marginalization_error_budget_nats", + 1.0e-3)) + if not (np.isfinite(_b) and _b > 0.0): + fatal.append("--direct-marginalization-error-budget-nats must be " + "finite and positive") + # --d-prior is accepted by the parser but not forwarded to the JAX + # wrapper (which is always volumetric); refuse a non-volumetric request + # under the policy here rather than let the wrapper's refusal be + # unreachable (wiring review, item 6). + _dp = getattr(opts, "d_prior", None) + if _dp not in (None, "", "euclidean", "volumetric"): + fatal.append("--direct-marginalization-policy %s derives its " + "measure from the volumetric distance prior; " + "--d-prior %s is not supported by the composite" + % (_policy, _dp)) + _br = int(getattr(opts, "direct_marginalization_batch_rows", 1)) + if _br < 0: + fatal.append("--direct-marginalization-batch-rows must be >= 0 " + "(1 = row at a time, 0 = one full batch)") + _pr = int(getattr(opts, "direct_marginalization_policy_probe_rows", 0)) + if _pr < 0: + fatal.append("--direct-marginalization-policy-probe-rows must be " + ">= 0 (0 disables the pre-sampling probe)") + if getattr(opts, "direct_marginalization_policy_probe_only", False) \ + and _pr <= 0: + fatal.append("--direct-marginalization-policy-probe-only needs " + "--direct-marginalization-policy-probe-rows > 0; " + "otherwise it would exit having measured nothing") + for _k, _d in (("--direct-marginalization-max-modes", + "direct_marginalization_max_modes"), + ("--direct-marginalization-enriched-max-modes", + "direct_marginalization_enriched_max_modes"), + ("--direct-marginalization-base-oversample", + "direct_marginalization_base_oversample"), + ("--direct-marginalization-enriched-oversample", + "direct_marginalization_enriched_oversample"), + ("--direct-marginalization-max-starts", + "direct_marginalization_max_starts")): + if int(getattr(opts, _d, 1)) < 1: + fatal.append("%s must be >= 1" % _k) + for _k, _d in (("--direct-marginalization-convergence-tol-nats", + "direct_marginalization_convergence_tol_nats"), + ("--direct-marginalization-time-guard-tol-nats", + "direct_marginalization_time_guard_tol_nats")): + _v = float(getattr(opts, _d, 1.0e-3)) + if not (np.isfinite(_v) and _v > 0.0): + fatal.append("%s must be finite and positive" % _k) + if (int(getattr(opts, "direct_marginalization_enriched_max_modes", 8)) + < int(getattr(opts, "direct_marginalization_max_modes", 4))): + fatal.append("--direct-marginalization-enriched-max-modes must be " + ">= --direct-marginalization-max-modes: the enriched " + "plan has to be able to nest the base plan, and a " + "narrower one declines on mode nesting every row") + else: + for _k in _policy_knobs: + if was_supplied(opts, _k): + fatal.append("%s is inert without --direct-marginalization-" + "policy auto; pass the policy or drop the option" + % _k) # 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). @@ -324,8 +590,8 @@ def check_critical_and_report(opts, optp): 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) + "data-derived angle amplitude, which the 'grid' scheme you " + "asked for 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 " @@ -344,21 +610,128 @@ def check_critical_and_report(opts, optp): "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: + if _gh_resolved > 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) + "--distance-gh-nodes/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") + # NOT ``... or 1``: 0 is falsy, so that idiom would silently promote an invalid + # ``--q-time-pregrid-factor 0`` to the default and report nothing. (It did, until + # test_jax_dropin_manifest... caught it.) + _qf = getattr(opts, "q_time_pregrid_factor", 1) + _qf = 1 if _qf is None else int(_qf) + if _qf < 1: + fatal.append("--q-time-pregrid-factor must be >= 1, got %d" % _qf) + elif _qf != 1: + if (getattr(opts, "rotation_slow", False) + or getattr(opts, "freqresponse", False)): + fatal.append("--q-time-pregrid-factor is not implemented for banded " + "rotation/frequency-response likelihoods") + # NOT a blanket refusal any more (it was, before the JAX Q path grew a + # pregrid). What survives is the narrower, still-true refusal, and it + # mirrors conventional ILE's (#261): the pregrid buys sub-sample accuracy + # with a FOUR-tap cubic, so a stencil that cannot use a sub-sample position + # ('nearest') or a different stencil chosen behind the user's back would + # both make the two arms answer differently for the same command line -- + # which is exactly how they came to ship opposite stencil defaults (#233). + # + # 'nearest' is separately UNIMPLEMENTED, not merely suboptimal: see + # core._q_sample_positions -- a refined-grid nearest gather no longer reads + # the sample the banded post-phase reconstructs from the coarse index, so + # the data term and the model norm would drift apart by up to half a coarse + # bin. + # The ``!= JAX_INTERP_DEFAULT`` clause is not redundant with was_supplied: + # was_supplied() FAILS OPEN by design ("no record -> assume not supplied"), + # which is the safe direction for the conflict checks it was written for and + # the WRONG one here, where the consequence of guessing "not supplied" is + # silently replacing a stencil the caller chose. A caller with no supplied + # record whose interp is not the module default has demonstrably chosen it, + # so treat that as explicit too. + _explicit_interp = (was_supplied(opts, "--interp") + or was_supplied(opts, "--interpolate-time") + or getattr(opts, "interp", None) != JAX_INTERP_DEFAULT) + if _explicit_interp and getattr(opts, "interp", None) != "cubic": + fatal.append( + "--q-time-pregrid-factor %d uses four-tap cubic interpolation " + "(--interp %r was requested); remove the explicit stencil option " + "or set it to cubic" % (_qf, getattr(opts, "interp", None))) + else: + opts._q_pregrid_fallback_interp = getattr(opts, "interp", None) + opts.interp = "cubic" + print(" Q_lm pregrid: ENABLED factor=%d boundary=even-reflection " + "arrival_stencil=cubic integration_grid=unchanged" % _qf) + _tq = getattr(opts, "time_marginalization_quadrature", "simpson") + if getattr(opts, "rotation_slow", False) and _tq == "bandlimited": + fatal.append("bandlimited time quadrature is not valid when slow rotation " + "makes arrival-time dependent; use simpson") + if ((getattr(opts, "rotation_slow", False) or getattr(opts, "freqresponse", False)) + and getattr(opts, "phase_marginalization", False)): + fatal.append("phase marginalization is not implemented for banded " + "rotation/frequency-response likelihoods") + if int(getattr(opts, "rotation_p_max", 0) or 0) < 0: + fatal.append("--rotation-p-max must be >= 0") + if int(getattr(opts, "rotation_n_harmonics", 2) or 2) < 2: + fatal.append("--rotation-n-harmonics must be >= 2") + if int(getattr(opts, "freqresponse_qmax", 4) or 4) < 0: + fatal.append("--freqresponse-qmax must be >= 0") + _rotation = bool(getattr(opts, "rotation_slow", False)) + _freqresponse = bool(getattr(opts, "freqresponse", False)) + if not _rotation: + for _name in ("--rotation-n-harmonics", "--rotation-p-max"): + if was_supplied(opts, _name): + fatal.append("%s is inert without --rotation-slow" % _name) + elif _freqresponse and was_supplied(opts, "--rotation-n-harmonics"): + fatal.append("--rotation-n-harmonics does not set the compound bank width; " + "its exact support is fixed by Qmax and pmax") + if not _freqresponse: + for _name in ("--freqresponse-qmax", "--freqresponse-arm-length"): + if was_supplied(opts, _name): + fatal.append("%s is inert without --freqresponse" % _name) + elif getattr(opts, "freqresponse_arm_length", None) not in (None, ""): + try: + _parse_freqresponse_arm_length(opts.freqresponse_arm_length) + except (TypeError, ValueError) as exc: + fatal.append("invalid --freqresponse-arm-length: %s" % exc) + # Refused HERE, before precompute: with a full-sky prior the stop below is + # not a chance event but the first chunk (see _EVAL_LNL_STOP_MODES). + _hw = getattr(opts, "data_integration_window_half", None) + if (_tq == "bandlimited" and (getattr(opts, "mode", None) in _EVAL_LNL_STOP_MODES + or _jax_av_active) + and _hw is not None + and float(_hw) < _BANDLIMITED_FULLSKY_HALF_WINDOW_MIN): + optp.error( + "--time-marginalization-quadrature bandlimited with --mode %s " + "evaluates full-sky prior draws and stops on one uncertified row, and " + "--data-integration-window-half %.4g s is below 2 R_earth / c = %.4f s, " + "the largest arrival shift a wrong-sky draw makes at one detector: a " + "row whose arrival peak lies beyond the window edge cannot be " + "certified, and the first chunk contains such rows (measured 35 of 256 " + "blind rows at 0.02 s, 0 of 256 at 0.05 s). Raise " + "--data-integration-window-half to at least %.4f s plus the " + "event-time uncertainty (0.05 s is the measured clean value), or use " + "--time-marginalization-quadrature simpson." + % (opts.mode, float(_hw), _BANDLIMITED_FULLSKY_HALF_WINDOW_MIN, + _BANDLIMITED_FULLSKY_HALF_WINDOW_MIN)) if fatal: optp.error("Cannot run as a faithful drop-in: " + "; ".join(fatal) + ". (These would silently change the result if ignored.)") + # Nothing above raised: safe to APPLY the resolved distance-GH-nodes count + # (deferred from the resolution comment above) and report it, once, here. + from RIFT.likelihood.jax_ile.core import set_distmarg_gh_nodes + set_distmarg_gh_nodes(_gh_resolved) + opts._distance_gh_nodes_resolved = _gh_resolved + print(" distance quadrature: per-sample Gauss-Hermite nodes=%d%s" + % (_gh_resolved, " (legacy uniform grid)" if _gh_resolved == 0 else + " (source: %s)" % ("--distance-gh-nodes" if _gh_cli is not None + else "JAX_ILE_DISTMARG_GH"))) + # Report accepted-but-ignored options that were actually passed. ignored = [] implemented = {"--cache-file", "--channel-name", "--psd-file", @@ -369,34 +742,50 @@ 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", "--limit-distance", + "--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", "--time-marginalization-quadrature", "--interpolate-time", - "--vectorized", "--use-gwsignal"} + "--vectorized", "--use-gwsignal", "--rotation-slow", + "--rotation-n-harmonics", "--rotation-p-max", + "--freqresponse", "--freqresponse-qmax", + "--freqresponse-arm-length"} # 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: + if getattr(opts, "sampler_method", None) in ("AV", "portfolio"): + implemented |= {"--sampler-method", "--n-eff", + "--sampler-anisotropic-bins", + "--limit-right-ascension", "--limit-declination", + "--limit-psi", "--limit-inclination"} + if getattr(opts, "sampler_method", None) == "portfolio": + implemented.add("--sampler-portfolio") + _tempered_mode_active = (mode in _TEMPERED_MODES + and getattr(opts, "sampler_method", None) + not in ("AV", "portfolio")) + if _tempered_mode_active: # 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): + # --d-prior is handled separately below with a substantive message (the + # JAX driver's distance prior, not merely "not yet implemented"), so it is + # excluded from the generic bag here rather than lumped into it. + for name in sorted(_ILE_ALL_OPTS - implemented - {"--d-prior"}): 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: + if not _tempered_mode_active: inert = [n for n in ("--auto-adapt-weight-exponent", "--allow-degenerate-tempering") if getattr(opts, _dest(n), False)] @@ -406,6 +795,33 @@ def check_critical_and_report(opts, optp): print("Note: %s only act on the tempered modes (%s); --mode %s ignores " "them." % (" ".join(sorted(set(inert))), " ".join(sorted(_TEMPERED_MODES)), mode)) + # --d-prior (RO'S 2026-09-08: missing knobs no-op for compatibility, but + # must warn). --d-prior is accepted (ILE compatibility) but never + # forwarded: the JAX driver always integrates against the volumetric d^2 + # prior on [--d-min, --d-max]. 'Euclidean'/'volumetric' (any case) IS + # that prior, so passing it explicitly is not a deviation and gets no note. + _dp = getattr(opts, "d_prior", None) + if _dp not in (None, "") and str(_dp).strip().lower() not in ( + "euclidean", "volumetric"): + print("Note: --d-prior %r is accepted but IGNORED by the JAX driver: " + "it always integrates distance against the volumetric d^2 " + "prior on [--d-min, --d-max]." % (_dp,)) + # --phase-marginalization: forwarded on the 6-D/5-D likelihoods, but the + # four phi_ref-marginalised modes already integrate phase out analytically + # (a grid sum baked into the likelihood, not a sampled axis) -- outside + # _ILE_ALL_OPTS, so the generic loop above cannot see it either. + if getattr(opts, "phase_marginalization", False) and mode in _PHASE_ANALYTIC_MODES: + print("Note: --phase-marginalization is accepted but IGNORED for " + "--mode %s (phase is already marginalized analytically)." % mode) + # --sky-coordinates network: wired only for --mode multistart-nuts (the + # other modes' samplers assume the equatorial (ra, sin(dec)) parameterization + # end to end -- reparameterizing them is more than plumbing, see the PR). + _skyc = getattr(opts, "sky_coordinates", "equatorial") + if _skyc != "equatorial" and mode != "multistart-nuts": + print("Note: --sky-coordinates %s is accepted but IGNORED for --mode " + "%s (only --mode multistart-nuts samples the sky in network-" + "frame coordinates; other modes use the default equatorial " + "parameterization)." % (_skyc, 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):") @@ -427,6 +843,20 @@ def _jax_core_dist_tol_default(): def build_parser(): optp = OptionParser(usage="%prog [options]", description=__doc__) + g = OptionGroup(optp, "JAX compilation cache") + g.add_option("--jax-cache-dir", default=None, + help="Persistent cache root. RIFT adds a JAX/JAXLIB/backend/" + "GPU compatibility namespace (default: $RIFT_JAX_CACHE_ROOT " + "or $XDG_CACHE_HOME/rift/jax).") + g.add_option("--no-jax-persistent-cache", action="store_true", default=False, + help="Disable cross-process JAX compilation caching for this run.") + g.add_option("--jax-cache-bundle", default=None, + help="Validate and import a warmed rift_jax_cache bundle before " + "constructing any ILE JIT (also $RIFT_JAX_CACHE_BUNDLE).") + g.add_option("--jax-cache-profile", default=None, + help="Require --jax-cache-bundle to declare this warmup profile.") + optp.add_option_group(g) + g = OptionGroup(optp, "Data input (frame mode)") g.add_option("--cache-file", default=None) g.add_option("--channel-name", action="append", default=[], @@ -487,6 +917,23 @@ def build_parser(): g.add_option("--internal-data-storage-window-half", type=float, default=0.15) optp.add_option_group(g) + g = OptionGroup(optp, "Time-dependent detector response") + g.add_option("--rotation-slow", action="store_true", default=False, + help="Use the sidereal slow-rotation likelihood. May be combined " + "with --freqresponse.") + g.add_option("--rotation-n-harmonics", type=int, default=2, + help="Minimum sidereal harmonic half-width (default 2).") + g.add_option("--rotation-p-max", type=int, default=0, + help="Maximum propagation-delay derivative order (default 0).") + g.add_option("--freqresponse", action="store_true", default=False, + help="Use finite-arm frequency response. Combining this with " + "--rotation-slow constructs a compound banded likelihood.") + g.add_option("--freqresponse-qmax", type=int, default=4, + help="Highest finite-arm projection power (default 4).") + g.add_option("--freqresponse-arm-length", default=None, + help="Arm length in metres, globally or DET=value comma list.") + optp.add_option_group(g) + g = OptionGroup(optp, "Extrinsic exploration / sampling") g.add_option("--mode", default="laplace-is", choices=["prior-mc", "laplace-is", "map", "nuts", @@ -539,6 +986,37 @@ def build_parser(): "factored-likelihood amplitude degeneracy; required for nuts).") g.add_option("--n-max", type=int, default=300000) g.add_option("--n-chunk", type=int, default=8000) + g.add_option("--n-eff", type=int, default=None, + help="Target effective sample count for --sampler-method AV/portfolio " + "(default 1000 on those backends).") + g.add_option("--sampler-method", default=None, + help="Optional non-AD integration backend: AV or portfolio. The " + "existing --mode still selects the JAX likelihood geometry; " + "other legacy values remain accepted as compatibility no-ops.") + g.add_option("--sampler-portfolio", action="append", default=[], + help="Portfolio members (repeat or comma-separate; JAX path supports " + "AV and GMM, default AV,GMM).") + g.add_option("--sampler-anisotropic-bins", action="store_true", default=False, + help="Allocate AV bins preferentially along compressed coordinates.") + g.add_option("--jax-av-seed", type="choice", choices=("none", "fisher-sky"), + default="none", + help="Optional JAX initializer for AV/portfolio: hill-climb several " + "modes, use their Fisher curvature on sky, and draw all other " + "coordinates from the physical prior (default none).") + g.add_option("--jax-av-eval-chunk", type=int, default=None, + help="Fixed JAX likelihood batch inside each AV coverage chunk " + "(default min(--n-chunk, the JAX memory-aware default)).") + g.add_option("--jax-av-seed-pilot", type=int, default=4000, + help="Prior likelihood evaluations used to find hill-climb starts.") + g.add_option("--jax-av-seed-modes", type=int, default=4, + help="Maximum separated sky modes retained by the Fisher-sky seed.") + g.add_option("--jax-av-seed-points", type=int, default=None, + help="Seed-cloud size (default --n-chunk).") + g.add_option("--jax-av-sky-inflate", type=float, default=2.0, + help="Standard-deviation inflation of the local Fisher sky proposal.") + g.add_option("--jax-av-seed-prior-frac", type=float, default=0.1, + help="Fraction of the seed cloud drawn from the full physical prior. " + "The portfolio GMM, not this finite cloud, guarantees full support.") 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).") @@ -567,6 +1045,31 @@ def build_parser(): "the node count follows from it in closed form. " "Default %g. Only valid with that scheme." % _jax_core_dist_tol_default()) + g.add_option("--distance-gh-nodes", type=int, default=None, + help="Node count for the per-sample Gauss-Hermite-style " + "distance quadrature (RIFT.likelihood.jax_ile.core." + "make_distance_gh / _distmarg_gh_logL): nodes centred " + "PER SAMPLE on the Gaussian peak of the distance " + "integrand, resolving it to machine precision at any " + "SNR with a few dozen nodes, instead of the legacy " + "fixed grid (see --distance-grid-points) which " + "under-resolves that peak at high SNR. Not passing " + "this flag keeps the legacy uniform grid (current " + "behaviour, no change). Equivalent to the environment " + "variable JAX_ILE_DISTMARG_GH, which is still honoured " + "for compatibility, but ONLY when this flag is not " + "given: when the flag IS given -- including an " + "explicit 0 -- it wins over the environment variable, " + "and a DIFFERENT nonzero JAX_ILE_DISTMARG_GH is " + "REFUSED rather than silently picked (an explicit " + "--distance-gh-nodes 0 against a nonzero " + "JAX_ILE_DISTMARG_GH is a conflict too, and is refused " + "the same way: silently choosing 0 or the environment " + "value would both hide the disagreement). " + "Not compatible with --distance-grid-scheme loguniform " + "(the per-sample quadrature reads only the SUPPORT of " + "that grid, so combining them would be silently inert " + "while reporting as active -- see --distance-grid-tol).") g.add_option("--limit-distance", default=None, help="Restrict distance SAMPLING (and, when distance is " "marginalized, the distance QUADRATURE) to 'LO,HI' in Mpc, " @@ -591,15 +1094,41 @@ def build_parser(): 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 " + help="Rule for the time integral: historical fixed-grid " + "Simpson or adaptive reflected-FFT refinement followed " "by a converged trapezoid. No resolution knob is exposed; " - "the factor is derived and rechecked from lnL(t).") + "the factor is derived and rechecked per row. " + "'bandlimited' is honoured by the 6-D modes and, under " + "--distance-marginalization, by every mode except the " + "phi/psi-marginalized ones, which refuse it and say so. " + "One uncertified row stops the run. The modes that push " + "full-sky prior draws through that stop (prior-mc, " + "laplace-is, map, nuts) are refused at parse time unless " + "--data-integration-window-half is at least 2 R_earth / c " + "(0.0426 s), the largest arrival shift of a wrong-sky " + "draw; 0.05 s is the measured clean value.") 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("--q-time-pregrid-factor", type="int", default=1, + help="OPT-IN. Refine the STORED rholm buffers onto a factor-x finer " + "time grid ONCE, before sampling, by even reflection plus a " + "band-limited FFT interpolation, and evaluate detector arrival " + "times on it with the four-tap cubic stencil. The integration " + "cadence -- deltaT, tvals and the Simpson weights -- is " + "UNCHANGED: this refines how Q is INTERPOLATED, not what the " + "likelihood integrates over. Default 1 is the historical " + "behaviour and is bit-identical. 8 is the validated value on " + "both arms (RIFT PR #261 for conventional ILE); larger factors " + "are accepted but buy little, because past ~8 the residual is " + "the reflection boundary condition rather than the stencil " + "step. Selects --interp cubic and REFUSES an explicit " + "different stencil, exactly as conventional ILE does. Costs " + "factor-x device memory for Q (only Q -- not the per-sample " + "gather scratch, which is the large term). Measured accuracy " + "and cost: RIFT/likelihood/DESIGN_q_window_stencil.md 9.7.") 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).") @@ -640,6 +1169,246 @@ def build_parser(): "ran is printed. See RIFT.likelihood.jax_ile.anglemarg." % (ANGLE_MARG_DEFAULT, ANGLE_MARG_LEGACY, ANGLE_MARG_LEGACY, ANGLE_MARG_LEGACY)) + g.add_option("--direct-marginalization-policy", + default=DIRECT_MARG_POLICY_DEFAULT, + choices=sorted(DIRECT_MARG_POLICY_CHOICES), + help="OPT-IN cross-axis policy for --mode flowmc-phipsimarg " + "(default '%s'). 'auto': per likelihood evaluation, " + "build the bounded U,V,Q start portfolios, attempt the " + "four-axis peak-local integral over (t, phi_ref, psi, D), " + "and keep it only if every acceptance diagnostic passes " + "(capacity, finite stationary modes, nested geometry, " + "base/enriched agreement, nested quadrature, two-guard " + "time, omitted-time mass, total error score under " + "--direct-marginalization-error-budget-nats); otherwise " + "run the band-limited exact-angle reserve. No SNR " + "threshold is coded. Requires --angle-marg-scheme exact " + "(its reserve), the simpson time rule (its check rule), " + "the volumetric distance prior and a uniform distance " + "grid; anything else is refused, not ignored. Value-only: " + "gradient parity is not yet validated (see " + "DESIGN_direct_marginalization_policy.md), so keep this " + "off for production until that ladder is recorded. The " + "branch statistics of the exported samples are printed and " + "labelled in the output headers. A row the controller " + "cannot warrant (after doubling the reserve rule up to " + "16x) is nan and the run is NOT published. Applies only " + "to --mode flowmc-phipsimarg; any other mode is refused." + % DIRECT_MARG_POLICY_DEFAULT) + g.add_option("--direct-marginalization-time-guard", type=int, + default=DirectMargPolicyConfig().time_guard, + help="Primitive-only support samples added at each end of the " + "time window for the policy's guarded reconstruction " + f"(default {DirectMargPolicyConfig().time_guard}; must " + "be >= 2). Both the local integral and " + "the reserve are re-evaluated at half this guard and must " + "agree within the error budget; a too-small guard declines " + "to the reserve rather than biasing.") + g.add_option("--direct-marginalization-reserve-time-refine", type=int, + default=DirectMargPolicyConfig().reserve_time_refine, + help="Refinement factor of the reserve's band-limited time rule " + "over the native cadence (default " + f"{DirectMargPolicyConfig().reserve_time_refine}; even, " + ">= 2). The " + "half-refined rule is the check rule the reserve must " + "agree with, so the warrant is a convergence statement " + "about the refined rules; the native Simpson rule's own " + "error is what the refinement removes.") + g.add_option("--direct-marginalization-reserve-scheme", type="choice", + choices=list(DIRECT_MARG_RESERVE_SCHEME_CHOICES), + default="exact", + help="The policy's reserve, a PAIR (angular kernel, time rule): " + "exact = exact angles on the whole-window refined rule " + "(default, unchanged); laplace = psi-Laplace angles on " + "that rule; peaklocal = psi-Laplace angles on a fixed-" + "count time rule sized per row from the PREDICTED peak " + "width 1/(2 pi rho sigma_f) (rho from the row's table, " + "sigma_f the rms frequency of the stored Q), with fine " + "blocks on one commensurate lattice around the time " + "maxima located on the primitive itself plus a coarse " + "scan of the window; its node count and " + "cost do not grow with rho, and the ledger prints the " + "prediction behind every row. auto = choose the pair " + "from the precomputed inputs before any row runs " + "(predict_reserve_pair). The whole-window refine " + "options are inert under peaklocal; a failed warrant " + "there doubles the fine lattice at fixed span, at most " + "--direct-marginalization-peaklocal-escalations times, " + "and is reported. Measured: " + "DESIGN_direct_marginalization_policy.md, " + "\"Peak-local time reserve\".") + g.add_option("--direct-marginalization-peaklocal-escalations", type=int, + default=int(DirectMargPolicyConfig().reserve_peaklocal_escalations), + help="Under the peaklocal reserve schemes, how many times a " + "failed resolution warrant may double the fine lattice at " + "fixed span before the row is unusable (default: the " + "PolicyConfig value). The " + "first tier is sized from the predicted width, so an " + "escalation is a finding about the prediction, not a loop; " + "the ledger counts them.") + g.add_option("--direct-marginalization-reserve-time-refine-max", type=int, + default=DirectMargPolicyConfig().reserve_time_refine_max, + help="Ceiling of the reserve rule's escalation on a failed " + "warrant (default " + f"{DirectMargPolicyConfig().reserve_time_refine_max}; " + "even, >= the refine factor). It " + "also bounds REVERSE-MODE GRADIENT MEMORY: the dense " + "reserve's backward pass costs ~2.6 GiB per unit of " + "refinement per evaluated row on a 614-sample window " + "(measured, rho 163), and lax.cond reserves memory for the " + "largest tier whether or not it runs; 32 asked for 85 GiB. " + "LOWERING IT IS NOT FREE. It does not change any value " + "that gets computed, but it changes which rows get a " + "value at all: a declined row gets one reserve attempt " + "per tier, and a row still unwarranted at the ceiling " + "returns nan, which this driver refuses to publish. " + "Measured on full-sky prior draws at rho 163, refine-max " + "4 left 6 of 8 rows nan and the run could not complete; " + "at the default 32 the same rows escalated and all were " + "usable. So on a small card the run fails rather than " + "returning a slightly worse number. Lower it only with " + "the nan count in the ledger in front of you.") + g.add_option("--direct-marginalization-error-budget-nats", type=float, + default=DirectMargPolicyConfig().total_value_error_budget_nats, + help="Value-error allowance, in nats, for the RESERVE's " + "time-resolution warrant (default " + f"{DirectMargPolicyConfig().total_value_error_budget_nats:g}" + "; 1e-3 shipped). " + "It does NOT reach the local accept/decline decision, " + "despite what this help used to say: that decision is the " + "four-term boolean at all_axis_peaklocal.py:862 and carries " + "no tolerance, so acceptance measured identical (48/64 at " + "rho 652) across 1e-3, 1e-2 and 1e-1. What it buys is " + "reserve cost: at rho 41 on identical rows, 1e-3 -> 1e-2 " + "took escalations from 2 to 0 and the run from 321.1 s to " + "146.9 s, a 2.19x speedup, moving lnL by 4.8e-12 nats. " + "Raised to 1e-2 by RO on 2026-09-08 (evening). Measured at " + "rho 41; the speedup is not yet confirmed at high SNR.") + g.add_option("--direct-marginalization-max-time-nodes", type=int, + default=DirectMargPolicyConfig().max_time_nodes, + help="Time-node capacity of the policy's LOCAL plan (default " + f"{DirectMargPolicyConfig().max_time_nodes}" + "; 64 shipped, and was reachable from no flag). A row " + "whose live time-node " + "count exceeds this declines to the exact reserve, which " + "costs ~4.7 h at rho 652 against 2.3 s for an accepted " + "row, so the capacity is sized to avoid the fallback " + "rather than to bound the plan. Measured need at " + "rho 652 on full-sky prior draws: median 72, p90 217, " + "max 614; caps of 64/128/256/1024 hold 44/73/91/100% of " + "rows. The library default was 64 and was unreachable " + "from any flag. NOT value-neutral: it changes which " + "rows take the local branch.") + g.add_option("--direct-marginalization-batch-rows", type=int, default=1, + help="Rows the policy's controller executes together under " + "one vmap (default 1: row at a time; 0: one full batch). " + "VALUES are identical at every size -- lnL, every branch " + "decision and the gradients -- and a test pins that. " + "COST IS A REGRESSION, not neutral: under vmap the " + "accept/reserve cond becomes a select, so a locally " + "ACCEPTED row also executes the dense reserve it would " + "otherwise skip, and every escalation tier runs for " + "every row. The penalty therefore scales with the " + "locally accepted fraction. Device workspace grows as " + "0.046 + 0.385 B GiB, so a 24 GiB card holds 32 rows and " + "not 64. Leave it at 1 unless a measurement on YOUR " + "rows says otherwise; see " + "DESIGN_direct_marginalization_policy.md. " + "WITHDRAWN: an earlier version of this text said COST " + "ONLY and quoted 96.3 s per row at 1 against 99.9 at 8. " + "Those rows all declined (accepted_local 0), so they " + "priced the decline path, and 37-50% of them were nan " + "under a reserve-refine ceiling of 4. Do not cite them.") + g.add_option("--direct-marginalization-policy-probe-rows", type=int, + default=0, + help="Evaluate the policy ledger on this many prior draws " + "BEFORE sampling and print the accept/decline counts " + "(default 0 = off). The end-of-run note reports the same " + "counts on the exported rows, but only a run that reaches " + "the export can be classified that way, and a declining " + "run is exactly the one that does not get there: each " + "declined row pays three exact reserve evaluations and " + "rows execute one at a time under jax.lax.map. The probe " + "answers 'local or reserve' in one compile plus this many " + "rows. The compile is not extra work for a policy run -- " + "the end-of-run note jits the identical function.") + g.add_option("--direct-marginalization-policy-probe-only", + action="store_true", default=False, + help="Exit after the probe above, before any sampling. For " + "surveying operating points: it reports which acceptance " + "diagnostic is declining without paying for a posterior. " + "Writes no samples and no evidence.") + # Start-portfolio and mode-plan sizing. Every default below is the value + # PolicyConfig already carried, so an unchanged command line is unchanged. + # They are exposed because the four-axis controller's own accepting + # operating point (analyses/va_sequence_20260902/RESULTS_20260907_aap268_ladder.md, + # 7 of 8 accepts at rho 326.15) reaches the SAME function through a richer + # portfolio than PolicyConfig's defaults build, and the difference was + # unreachable from the command line -- so the policy could only ever be + # observed declining. + g.add_option("--direct-marginalization-max-modes", type=int, + default=DirectMargPolicyConfig().max_modes, + help="Modes retained in the policy's BASE plan (default " + f"{DirectMargPolicyConfig().max_modes}).") + g.add_option("--direct-marginalization-enriched-max-modes", type=int, + default=DirectMargPolicyConfig().enriched_max_modes, + help="Modes retained in the policy's ENRICHED plan (default " + f"{DirectMargPolicyConfig().enriched_max_modes}). " + "Acceptance compares the two plans, so an enriched " + "plan that cannot resolve the extra structure declines.") + g.add_option("--direct-marginalization-base-oversample", type=int, + default=DirectMargPolicyConfig().base_oversample, + help="Angular oversampling of the base start portfolio " + f"(default {DirectMargPolicyConfig().base_oversample}).") + g.add_option("--direct-marginalization-enriched-oversample", type=int, + default=DirectMargPolicyConfig().enriched_oversample, + help="Angular oversampling of the enriched start portfolio " + f"(default {DirectMargPolicyConfig().enriched_oversample}).") + g.add_option("--direct-marginalization-max-starts", type=int, + default=DirectMargPolicyConfig().base_max_starts, + help="Cap on ranked optimizer starts per plan (default " + f"{DirectMargPolicyConfig().base_max_starts}). " + "TWO THINGS THIS IS NOT. (1) It is not the whole of " + "decline_capacity. That ledger key reports the " + "conjunction at all_axis_peaklocal.py:862 -- " + "norm_nonnegative, a certified time cover, the time-node " + "capacity, AND this cap -- so a row is charged to " + "'capacity' for four different reasons. Raising this " + "ALONE recovers ~16% and then plateaus -- at rho 652 on " + "43 such rows, 36 failed the TIME-NODE capacity and only " + "7 were count-driven. That plateau describes the PAIR, " + "not this knob: it held while the time-node capacity was " + "pinned at the library default of 64. Size this with " + "--direct-marginalization-max-time-nodes, which now " + "exists. Measured acceptance at rho 652 as (time nodes, " + "starts): (64, 32) 28%, (256, 32) 31%, (256, 128) 75%, " + "(512, 256) 77%. (2) It is not a " + "cost knob: it selects which modes are retained, so it " + "moves the integral. Measured on identical rows at rho " + "163, 32 -> 128 shifted already-accepted values by 4.8e-4 " + "and 3.5e-3 nats, at and above the 1e-3 error budget. " + "An A/B on this may not treat its value difference as " + "noise. 32 -> 128 was approved by RO on 2026-09-08 " + "(evening) with that shift on the record, because a " + "declined row runs the exact reserve at hours per row at " + "rho 652 against ~2.3 s for an accepted one. It is " + "sized WITH --direct-marginalization-max-time-nodes; " + "either alone plateaus near 16%.") + g.add_option("--direct-marginalization-convergence-tol-nats", type=float, + default=DirectMargPolicyConfig().convergence_tol_nats, + help="Agreement required between the base and enriched plans, " + "and between each plan's own two quadrature orders, in " + "nats (default 1e-3). This is the tolerance the " + "'enrichment' and 'quadrature' declines are measured " + "against; it is SEPARATE from " + "--direct-marginalization-error-budget-nats, which is the " + "total value allowance.") + g.add_option("--direct-marginalization-time-guard-tol-nats", type=float, + default=DirectMargPolicyConfig().time_guard_tol_nats, + help="Agreement required between the full-guard and half-guard " + "evaluations, in nats (default 1e-3). This is the " + "tolerance the two-guard time warrant is measured " + "against, for the local branch and the reserve alike.") # 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, @@ -714,6 +1483,16 @@ def build_parser(): help="SMC puffball: number of walkers in the cloud (default 2000).") g.add_option("--smc-move-steps", type=int, default=10, help="SMC puffball: random-walk Metropolis moves per temperature rung.") + g.add_option("--smc-is-samples", type=int, default=60000, + help="Draws in the SMC cloud-Gaussian importance-sampling " + "evidence stage (default 60000, the sampler's own value). " + "This stage is a FIXED cost paid after the ladder " + "finishes, and on a likelihood whose rows execute one at " + "a time -- the direct-marginalization policy -- it " + "dominates the whole run. Lowering it trades the IS " + "evidence for the raw SMC evidence: the ESS gate then " + "declines the IS value and the reported logZ falls back " + "to the SMC estimator, which the run already prints.") g.add_option("--smc-puff-scale", type=float, default=1.0, help="SMC puffball: proposal scale x sqrt(cloud covariance) (~1.0).") g.add_option("--fisher-is-samples", type=int, default=0, @@ -925,15 +1704,67 @@ def resolve_distance_limit(opts): raise SystemExit(" --limit-distance: %s" % exc) +def resolve_av_angular_limits(opts): + """Return validated AV-only angular sampling windows. + + The likelihood prior remains normalized on the full physical domain. These + limits therefore restrict the region integrated; they do not define a new, + box-normalized prior. Wraparound RA windows are intentionally refused until + AV can represent their union as more than one hyperrectangle. + """ + specs = (("limit_right_ascension", "ra", 0.0, 2.0 * np.pi), + ("limit_declination", "dec", -0.5 * np.pi, 0.5 * np.pi), + ("limit_psi", "psi", 0.0, np.pi), + ("limit_inclination", "incl", 0.0, np.pi)) + out = {} + for attr, name, physical_lo, physical_hi in specs: + raw = getattr(opts, attr, None) + if raw in (None, ""): + continue + try: + parts = [float(item.strip()) for item in str(raw).split(",")] + except ValueError: + parts = [] + if len(parts) != 2 or not np.all(np.isfinite(parts)): + raise SystemExit(" --limit-%s must be a finite LO,HI pair" % + attr[len("limit_"):].replace("_", "-")) + lo, hi = parts + if lo >= hi: + raise SystemExit(" --limit-%s requires LO < HI (RA wraparound is " + "not supported)" % + attr[len("limit_"):].replace("_", "-")) + if lo < physical_lo or hi > physical_hi: + raise SystemExit(" --limit-%s must lie within [%g,%g]" % + (attr[len("limit_"):].replace("_", "-"), + physical_lo, physical_hi)) + out[name] = (lo, hi) + return out + + +def av_distance_sampling_kwargs(like, d_lo, d_hi): + """Apply distance limits only when distance is an sampled coordinate.""" + from RIFT.likelihood.jax_ile import samplers as _samplers + if "distMpc" not in _samplers._av_param_order(like): + return {} + return {"sample_d_min": d_lo, "sample_d_max": d_hi} + + 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. + density is the prior divided by that mass. The rule is per ESTIMATOR, not + per function: an estimator that assumes "proposal == prior" must subtract + this, and one that forms ln w = lnL + ln p - ln q explicitly must NOT, + because log_prior() is already normalized over the physical [d_min,d_max]. + + run_laplace_is contains ONE OF EACH, which is why this docstring no longer + names functions. Its adapted loop forms ln w explicitly and does not + subtract; its prior pilot (kept as the collapse reference, #227) draws from + sample_prior and therefore does. Both then estimate the SAME integral on + the SAME full-range normalization, which is the only reason the two are + comparable and the 5-nat guard means anything. """ if not with_distance: return 0.0 @@ -995,19 +1826,124 @@ def eval_lnL(like, theta, opts, with_distance): 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, "direct_marginalization_policy", "off") != "off" + and np.any(np.isnan(out[sl]))): + raise RuntimeError( + "--direct-marginalization-policy %s: %d of %d evaluated rows " + "could not be warranted (local gate declined AND the escalated " + "band-limited reserve failed its guard/resolution check, or the " + "norm table varies with time). No coarse likelihood is " + "substituted and the run is not published. Raise " + "--direct-marginalization-time-guard, raise the reserve " + "refinement, or run with the policy off." + % (like.direct_marginalization_policy, + int(np.sum(np.isnan(out[sl]))), int(sl.stop - sl.start))) if (getattr(like, "time_quadrature", "simpson") == "bandlimited" and np.any(np.isnan(out[sl]))): + bad = np.where(np.isnan(out[sl]))[0] + names = FULL_NAMES if with_distance else ANG_NAMES + rows = "; ".join( + ", ".join("%s=%.4f" % (nm, v) for nm, v in zip(names, theta[sl][j])) + for j in bad[:3]) + # One uncertified row fails the run: no coarse likelihood is + # substituted, and a row's contribution cannot be bounded without a + # value. With a full-sky prior the usual cause is a blind draw whose + # detector arrival peak lies at or beyond the window edge (a wrong + # sky shifts arrival by up to 2 R_earth/c ~ 43 ms), which the + # trapezoid and guard certificates cannot converge on; the + # integration window has to contain those shifts. Measured rates: + # DESIGN_jax_bandlimited_distmarg.md. 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.") + "--time-marginalization-quadrature bandlimited: %d of %d rows in " + "this chunk failed a certificate (peak-width resolution, factor " + "doubling, or guard agreement); no coarse likelihood is " + "substituted and the run stops. First failing rows: %s. A " + "peak at the window edge is the common cause under a full-sky " + "prior; raise --data-integration-window-half to cover the " + "detector arrival shifts (>= 2 R_earth / c = %.4f s plus the " + "event-time uncertainty; 0.05 s measured clean), or raise the " + "input/rholm sample rate." + % (len(bad), int(sl.stop - sl.start), rows, + _BANDLIMITED_FULLSKY_HALF_WINDOW_MIN)) return out # --------------------------------------------------------------------------- # Evidence helpers # --------------------------------------------------------------------------- +def require_finite_evidence(logZ, neff, mode): + """A NON-FINITE evidence is a FAILED event, not a result. + + The estimators return nan when the proposal never bracketed the peak (see + ``_finalize_evidence``), and publishing that as an ILE row hands a + downstream CIP fit a nan ``lnL`` for a template that simply was not + integrated -- the same "leave no artifact behind" rule ``write_samples`` + already enforces for a cloud that admits no fair draw. Raising here means + ``--soft-fail-event-range`` skips to the next event, and without it the run + exits nonzero: on the #227 configuration the shipped code exited 0. + """ + if np.isfinite(logZ): + return + raise RuntimeError( + "extrinsic integration produced a non-finite evidence (logZ=%r, " + "neff=%.3g, mode=%s): the proposal did not bracket the likelihood peak, " + "so no result row is written for this event. Every --mode whose evidence " + "comes from a single moment-matched Gaussian fitted to its own draws can " + "fail this way on a narrow, high-SNR extrinsic posterior; try another " + "--mode, and check its reported neff rather than only its lnZ." + % (logZ, neff, mode)) + + +# The prior pilot is an UNBIASED estimator of the same Z, not a bound on it, and +# the difference is load-bearing -- see prior_pilot_floor below. +PILOT_FLOOR_FP_RATE = 3.4e-4 # = exp(-8); see the operating curve below + + +def prior_pilot_floor(logZ_pilot, false_positive_rate=PILOT_FLOOR_FP_RATE): + """A lower confidence bound on ln Z built from the prior pilot's estimate. + + WHY A BOUND AND NOT THE ESTIMATE. The pilot is prior Monte Carlo: unbiased + for Z, but with a heavy RIGHT tail once the target occupies a tiny fraction + of the prior. For a mode of prior mass ``m``, a single draw landing near the + peak makes the estimate ~``L_max / n_pilot`` while the truth is ~``L_max m``, + so it overshoots by ``1/(n_pilot m)`` -- more than ``T`` nats whenever + ``n_pilot m < exp(-T)``. Using the raw estimate as a floor therefore rejects + a CORRECT adapted answer at exactly the rate that tail occurs. That is not + hypothetical here: at a synthetic width of 0.05 rad, ``n_pilot m = 1.6e-3`` + and the pilot was measured above the truth by up to +5.46 nats, P = 1.1e-3 + over 900 seeds. + + THE BOUND. Markov is enough and needs nothing but unbiasedness and + non-negativity, both of which hold: ``P(Zhat >= Z / rate) <= rate``. So + ``Zhat * rate`` is a lower confidence bound on Z at level ``1 - rate``, and + the floor is ``ln Zhat + ln rate``. The threshold is therefore a CHOSEN + false-positive rate rather than a tuned constant, and it is distribution-free + -- in particular it does NOT assume the pilot resolved anything. It cannot: + the pilot's own ESS is ~1 in every regime where this guard matters (measured + median 1.0-1.3 for widths 0.03-0.08 rad), so "the pilot rests on one draw" is + the normal state here, not an exceptional one. + + WHY exp(-8) AND NOT exp(-5). Measured over 5400 synthetic runs (widths + 0.03-0.15 rad), sweeping the threshold T: + + T Markov FP <= FP measured power on inaccurate runs + 5 6.7e-3 0/1180 0.651 + 6 2.5e-3 0/1180 0.636 + 8 3.4e-4 0/1180 0.606 + 10 5.0e-5 0/1180 0.578 + + Going from 5 to 8 costs 4.5 points of power and buys a 20x smaller worst-case + false-positive rate; a false positive here FAILS THE EVENT and writes no row, + so it is worth paying for. No false positive was observed at any threshold: + the largest pilot-minus-adapted gap on an accurate run was +1.654 nats, so 8 + clears the measured margin by 6.3 nats. The bound, not the measurement, is + what the guarantee rests on. + """ + if not np.isfinite(logZ_pilot): + return -np.inf + return float(logZ_pilot) + float(np.log(false_positive_rate)) + + def evidence_from_logweights(logw): """(logZ, sigma/Z, neff) for Z = E[w] from log importance weights.""" fin = np.isfinite(logw) @@ -1066,12 +2002,27 @@ def run_laplace_is(like, opts, rng, dim, with_distance, n_adapt=2): lnL_p = eval_lnL(like, theta_p, opts, with_distance) logL_p = lnL_p + log_prior(theta_p, opts, with_distance) mu, cov = _moment_match(theta_p, logL_p) + # KEEP the pilot's own (prior-proposal) evidence estimate. It is crude -- one + # prior scan, often with an ESS of order 1 -- but it estimates the SAME integral + # from a proposal that is guaranteed to cover the prior, and importance sampling + # from a proposal that MISSES mass is biased low. So an adapted estimate coming + # out far BELOW this one is evidence that the adaptation walked away from the + # peak: the failure mode that remains once #227's draw/density mismatch is fixed. + # It is a REFERENCE, not a floor -- prior_pilot_floor() turns it into one, and + # the distinction is the whole of external review's P1 on this PR. + logZ_pilot, _, _ = evidence_from_logweights( + lnL_p - log_distance_box_correction(opts, with_distance)) per_round = max(opts.n_max // (n_adapt + 1), 1) all_theta, all_logw, all_lnL = [], [], [] for r in range(n_adapt + 1): - cov_use = cov * opts.proposal_inflate - Lc = np.linalg.cholesky(cov_use + 1e-12 * np.eye(dim)) + # ONE matrix for the draw AND the density. Drawing from + # ``cov_use + 1e-12*I`` and scoring under bare ``cov_use`` is issue #227: + # once the adapted covariance falls below the absolute jitter the weights + # are computed against a distribution that was never sampled, and this + # mode returned lnZ = 5.8e9 on real O4 data while exiting 0. + cov_use = _regularize_cov(cov * opts.proposal_inflate) + Lc = np.linalg.cholesky(cov_use) z = rng.standard_normal((per_round, dim)) theta = mu[None, :] + z @ Lc.T logq = _gaussian_logq(theta, mu, cov_use) @@ -1089,6 +2040,20 @@ 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) + # log Z <= max lnL for a normalized prior, and a low neff means the proposal + # never bracketed the peak. Same rule the library samplers already apply; + # this driver applied none, which is why #227 exited 0 on lnZ = 5.8e9. + logZ, sig, neff = _finalize_evidence( + logZ, sig, neff, float(np.max(lnL)) if np.isfinite(lnL).any() else np.nan) + pilot_floor = prior_pilot_floor(logZ_pilot) + if np.isfinite(logZ) and np.isfinite(pilot_floor) and logZ < pilot_floor: + print(" [laplace-is] adapted proposal gives lnZ = %.3f, %.1f nats BELOW the " + "prior pilot's Markov floor (%.3f, from a pilot estimate of %.3f at a " + "%.1e false-positive rate): the adaptation moved off the peak, so this " + "evidence is reported as unreliable." + % (logZ, pilot_floor - logZ, pilot_floor, logZ_pilot, + PILOT_FLOOR_FP_RATE)) + logZ, sig = np.nan, np.nan # 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 @@ -1161,10 +2126,11 @@ def run_nuts(like, opts, rng, with_distance): # moment-matched to the posterior draws -> high neff vs prior seeding). mu, cov = _moment_match(theta, np.zeros(len(theta))) # posterior moments n_is = min(opts.n_max, 40000) - Lc = np.linalg.cholesky(cov * opts.proposal_inflate + 1e-12 * np.eye(5)) + cov_is = _regularize_cov(cov * opts.proposal_inflate) # one matrix (#227) + Lc = np.linalg.cholesky(cov_is) z = rng.standard_normal((n_is, 5)) th_is = mu[None, :] + z @ Lc.T - logq = _gaussian_logq(th_is, mu, cov * opts.proposal_inflate) + logq = _gaussian_logq(th_is, mu, cov_is) logp = log_prior(th_is, opts, with_distance=False) valid = np.isfinite(logp) lnL_is = np.full(n_is, -np.inf) @@ -1172,6 +2138,12 @@ 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) + # Bound with the NUTS CHAIN's peak, not the IS cloud's. Both are valid upper + # bounds on ln Z, the chain's is the larger (it sits ON the peak), and using the + # smaller one would fail a correct run whose IS cloud happened to be broad. This + # is also the argument samplers.py passes at its own _finalize_evidence sites. + logZ, sig, neff = _finalize_evidence( + logZ, sig, neff, float(np.max(lnL)) if np.isfinite(lnL).any() else np.nan) # 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 @@ -1213,28 +2185,71 @@ def samples_path(opts, out_index): return opts.output_file + "_" + str(out_index) + "_samples.dat" +def direct_marginalization_policy_note(like, theta, n_max=256, chunk=32, + return_values=False): + """Label the branch statistics of the exported rows under the policy. + + Returns "" when the policy is off. Otherwise evaluates the ledger on up to + ``n_max`` evenly spaced exported rows and reports how many took the local + branch, the warranted reserve, or came back unusable (reserve failed its + own warrant; the value is retained and the run is labelled). The note is + written into the sample and evidence headers next to the angle-grid label. + """ + policy = getattr(like, "direct_marginalization_policy", "off") + ledger_fn = getattr(like, "_batched_ledger", None) + if policy == "off" or ledger_fn is None: + return ("", None, None) if return_values else "" + theta = np.asarray(theta) + if theta.ndim != 2 or theta.shape[0] == 0: + _empty = "DIRECT-MARG-POLICY=%s rows=0" % policy + return (_empty, None, None) if return_values else _empty + n = int(theta.shape[0]) + idx = np.unique(np.linspace(0, n - 1, min(n, int(n_max))).astype(int)) + parts = [] + for start in range(0, idx.size, int(chunk)): + sub = theta[idx[start:start + int(chunk)]] + cols = [jnp.asarray(sub[:, j]) for j in range(sub.shape[1])] + _, ledger = ledger_fn(*cols) + parts.append({k: np.asarray(v) for k, v in ledger.items()}) + ledger = {k: np.concatenate([p[k] for p in parts]) for k in parts[0]} + summary = _summarize_policy_ledger(ledger) + declines = ",".join("%s:%d" % (k.replace("decline_", ""), v) + for k, v in sorted(summary["declines"].items())) + note = ("DIRECT-MARG-POLICY=%s rows=%d local=%d reserve=%d " + "warranted-reserve=%d escalations=%d unusable=%d reconciles=%d " + "declines=[%s] max-local-error-score-nats=%.3g" + % (policy, summary["rows"], summary["accepted_local"], + summary["reserve_executed"], summary["reserve_warranted"], + summary.get("reserve_escalations", 0), + summary["unusable"], summary["reconciles"], declines, + summary["max_local_error_score_nats"])) + if return_values: + # The per-row selected value, so two probes at ONE seed -- one whose + # config accepts, one forced onto the reserve -- can be differenced on + # exactly the same prior draws. An accepted row never evaluates the + # reserve, so this is the only way to price the local branch against it. + return (note, np.asarray(ledger["lnL"], dtype=float), + np.asarray(ledger["accepted_local"], dtype=bool)) + return note + + 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: + Returns one of three things, and the THIRD states the exact coverage: "" -- 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. + "ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS ..." -- checked cloud was adequate + "ANGLE-GRID-CHECK=NOT-PERFORMED ..." -- amp-sized scheme, NO check ran + + The pure JIT returns its amplitude metric as ordinary data, and the Python + boundary synchronously accumulates the maximum over every pilot, reweight, + and final production/output-cloud batch. That deterministic coverage is + enough to label the points used for the published flow evidence and sample + export. Transient flow-training-only proposals (which do not enter those + artifacts) are deliberately NOT claimed. This split also keeps host + callbacks out of the expensive graph so JAX can persistently cache it. """ st = _anglemarg.amp_failsafe_state() if st.get("tripped"): @@ -1246,10 +2261,35 @@ def angle_grid_suspect_note(scheme=None): # 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)") + # 'phi-local' belongs here too, and its omission was found by external review rather + # than by this comment being read. It runs _runtime_amp_failsafe and evaluates an + # amp-sized dense fallback on every row, so it is amplitude-sized in exactly the sense + # this label is about; leaving it out published its artifacts with an EMPTY note -- + # the silence the paragraph above says a reader would take for verification. + if scheme in ("exact", "laplace", "peak-local", "phi-local"): + # A PASS may only be claimed when a batch was actually inspected. The + # recorder is wired ONLY for direct_marginalization_policy="off" + # (wrapper.py sets self._amp_record under that condition), while + # self.angle_marg_scheme still names the amp-sized reserve scheme. So + # `--angle-marg-scheme exact --direct-marginalization-policy auto` + # reaches here with n_calls == 0 and amp_sizing None. Without this + # branch that formatted "%.6g" % None and raised TypeError, killing the + # event AFTER integration and BEFORE either writer -- and had amp_sizing + # merely been 0.0 it would instead have published OUTPUT-CLOUD-PASS + # worst_amp=0, an affirmative adequacy claim backed by zero checks, + # which is the exact false negative this label exists to prevent. + if not st.get("n_calls"): + return ("ANGLE-GRID-CHECK=NOT-PERFORMED scheme=%s (no batched " + "amp-sized evaluation was recorded for this event, so the " + "adequacy of the dense (phi,psi) grids is UNKNOWN -- this " + "is NOT a pass; rebuild at a larger amp_sizing if it " + "matters)" % (scheme,)) + return ("ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS worst_amp=%.6g " + "amp_sizing=%.6g scheme=%s (deterministic over pilot/reweight/" + "final output-cloud evaluations; transient training-only " + "proposals not inspected)" + % (st.get("worst_amp", float("nan")), + st.get("amp_sizing", float("nan")), st.get("scheme"))) return "" @@ -1785,6 +2825,21 @@ def write_samples(opts, out_index, theta, lnL, with_distance, logw=None, print("Wrote %s (%d samples)" % (sname, len(cols))) +def _parse_freqresponse_arm_length(value): + if value in (None, ""): + return None + text = str(value) + if "=" not in text: + return float(text) + result = {} + for item in text.split(","): + fields = item.split("=", 1) + if len(fields) != 2 or not fields[0].strip(): + raise ValueError("invalid --freqresponse-arm-length item %r" % item) + result[fields[0].strip()] = float(fields[1]) + return result + + def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, rng, out_index, event_id, flow_state=None): """Build the JAX likelihood for one intrinsic template and run --mode. @@ -1807,15 +2862,68 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, 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, - opts.internal_data_storage_window_half, opts.data_integration_window_half, - opts.l_max, opts.fmax, - analyticPSD_Q=analyticPSD_Q, verbose=opts.verbose, + print("Building JAX likelihood (production precompute + banded pack)...") + waveform_kw = dict( use_gwsignal=bool(getattr(opts, "use_gwsignal", False)), - use_gwsignal_approx=(opts.approximant if getattr(opts, "use_gwsignal", False) else None)) - print(" modes:", like_data.lms, " guessed SNR:", extras["guess_snr"]) + use_gwsignal_approx=(opts.approximant + if getattr(opts, "use_gwsignal", False) else None)) + if opts.rotation_slow and opts.freqresponse: + arm_length = _parse_freqresponse_arm_length(opts.freqresponse_arm_length) + like_data, extras = build_rotating_freqresponse_data_from_precompute( + P.copy(), data_dict, psd_dict, fiducial_epoch, + opts.data_integration_window_half, opts.l_max, opts.fmax, + t_window=opts.internal_data_storage_window_half, + Qmax=opts.freqresponse_qmax, L_arm=arm_length, + p_max=opts.rotation_p_max, analyticPSD_Q=analyticPSD_Q, + verbose=opts.verbose, **waveform_kw) + elif opts.rotation_slow: + nh = int(opts.rotation_n_harmonics) + like_data, extras = build_rotation_data_from_precompute( + P.copy(), data_dict, psd_dict, fiducial_epoch, + opts.data_integration_window_half, opts.l_max, opts.fmax, + t_window=opts.internal_data_storage_window_half, + harmonics=tuple(range(-nh, nh + 1)), p_max=opts.rotation_p_max, + analyticPSD_Q=analyticPSD_Q, verbose=opts.verbose, **waveform_kw) + elif opts.freqresponse: + arm_length = _parse_freqresponse_arm_length(opts.freqresponse_arm_length) + like_data, extras = build_freqresponse_data_from_precompute( + P.copy(), data_dict, psd_dict, fiducial_epoch, + opts.data_integration_window_half, opts.l_max, opts.fmax, + t_window=opts.internal_data_storage_window_half, + Qmax=opts.freqresponse_qmax, L_arm=arm_length, + analyticPSD_Q=analyticPSD_Q, verbose=opts.verbose, **waveform_kw) + else: + like_data, extras = build_data_from_precompute( + P.copy(), data_dict, psd_dict, fiducial_epoch, + opts.internal_data_storage_window_half, opts.data_integration_window_half, + opts.l_max, opts.fmax, + analyticPSD_Q=analyticPSD_Q, verbose=opts.verbose, + q_time_pregrid_factor=( + 1 if getattr(opts, "q_time_pregrid_factor", 1) is None + else int(opts.q_time_pregrid_factor)), **waveform_kw) + print(" feature:", getattr(like_data, "feature", None), + " modes:", like_data.lms, + " guessed SNR:", extras.get("guess_snr", "not estimated")) + if getattr(like_data, "feature", None) is not None: + _d0 = like_data.detectors[like_data.detector_names[0]] + _A = int(_d0["Q_bank"].shape[0]) + _bank_bytes = sum(int(like_data.detectors[d]["Q_bank"].nbytes + + like_data.detectors[d]["U_bank"].nbytes + + like_data.detectors[d]["V_bank"].nbytes) + for d in like_data.detector_names) + print(" [banded] basis=%d ordered-pairs=%d device-bank=%.1f MB " + "(%d detectors)" % (_A, _A * _A, _bank_bytes / 2.0**20, + len(like_data.detector_names))) + if like_data.q_time_pregrid_factor != 1: + _d0 = like_data.detectors[like_data.detector_names[0]] + _bytes = sum(int(like_data.detectors[d]["Q"].nbytes) + for d in like_data.detector_names) + print(" [q-pregrid] factor %d: Q sampled at deltaT/%d, %d -> %d samples, " + "%.1f MB of Q on device (%d detectors); integration cadence " + "unchanged at deltaT=%.6g s" + % (like_data.q_time_pregrid_factor, like_data.q_time_pregrid_factor, + _d0["npts_full_coarse"], _d0["Q"].shape[0], _bytes/2.0**20, + len(like_data.detector_names), like_data.deltaT)) with_distance = not opts.distance_marginalization # --limit-distance: (d_lo,d_hi) is what is SAMPLED / quadratured; the prior is @@ -1889,12 +2997,57 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, raise SystemExit( "--distance-grid-tol only applies to --distance-grid-scheme " "loguniform; it would be silently inert here.") - like = JAXDistPhiPsiMargLikelihood( - 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) + policy = getattr(opts, "direct_marginalization_policy", + DIRECT_MARG_POLICY_DEFAULT) + policy_config = None + if policy != "off": + _f0 = int(opts.direct_marginalization_reserve_time_refine) + _scheme = getattr(opts, "direct_marginalization_reserve_scheme", + "exact") + if _scheme != "auto": + _angular, _time_rule = _direct_marg_reserve_pair(_scheme) + print(" direct-marginalization reserve scheme: %s (angular " + "kernel %s, time rule %s)" % (_scheme, _angular, _time_rule)) + policy_config = DirectMargPolicyConfig( + time_guard=int(opts.direct_marginalization_time_guard), + reserve_scheme=_scheme, + reserve_peaklocal_escalations=int( + opts.direct_marginalization_peaklocal_escalations), + reserve_time_refine=_f0, + reserve_time_refine_max=int( + opts.direct_marginalization_reserve_time_refine_max), + total_value_error_budget_nats=float( + opts.direct_marginalization_error_budget_nats), + reserve_batch_rows=int( + opts.direct_marginalization_batch_rows), + max_modes=int(opts.direct_marginalization_max_modes), + enriched_max_modes=int( + opts.direct_marginalization_enriched_max_modes), + base_oversample=int( + opts.direct_marginalization_base_oversample), + enriched_oversample=int( + opts.direct_marginalization_enriched_oversample), + base_max_starts=int(opts.direct_marginalization_max_starts), + max_time_nodes=int( + opts.direct_marginalization_max_time_nodes), + convergence_tol_nats=float( + opts.direct_marginalization_convergence_tol_nats), + time_guard_tol_nats=float( + opts.direct_marginalization_time_guard_tol_nats)) + try: + like = JAXDistPhiPsiMargLikelihood( + 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, + direct_marginalization_policy=policy, + policy_config=policy_config) + except ValueError as e: + if policy == "off": + raise + raise SystemExit("--direct-marginalization-policy %s refused: %s" + % (policy, e)) # 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" @@ -1902,9 +3055,132 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, "; ".join("%s=%s" % kv for kv in sorted(like.angle_marg_info.items()) if kv[0] not in ("scheme", "requested")))) + print(" direct-marginalization policy: %s%s" + % (like.direct_marginalization_policy, + "" if like.policy_info is None else ": " + "; ".join( + "%s=%s" % kv for kv in sorted(like.policy_info.items())))) print(" distance grid: %s" % "; ".join("%s=%s" % kv for kv in sorted(getattr(like, "dist_grid_info", {}).items()))) + + # THE PAIR IS CHOSEN AND ANNOUNCED BEFORE ANY ROW IS EVALUATED. + # RO, 2026-09-08: rely on analysis and the known physics to pick the + # (local, reserve) pair, rather than try-then-decline-then-refine. A + # run whose branch can only be read from a ledger at the end is a run + # that spent its time in a method nobody chose. + if like.direct_marginalization_policy != "off": + _requested = getattr( + opts, "direct_marginalization_reserve_scheme", + DirectMargPolicyConfig().reserve_scheme) + # WHAT IS ON THE MENU IS A PROPERTY OF THIS DATA, not a constant. + # Offering laplace unconditionally would let the analysis choose a + # reserve whose premise is absent here, and the run would then be + # refused (or worse, carried) for a reason that has nothing to do + # with the signal. Two conditions, both about the DISTANCE + # integral the reserve's table is contracted against: + # + # 1. The per-sample adaptive Gauss-Hermite quadrature must be on. + # On the legacy static uniform grid laplace is measured to cost + # 43.2 nats at rho 163 -- that is what the GRID costs, not the + # scheme, and it is why the reserve may not use it there. The + # loguniform static grid is NOT admitted either: it is sized + # from the angle amplitude and may well be adequate, but no + # measurement exists, and "probably fine" is the thing this + # selector was built to stop saying. + # 2. The A0 == 0 / B1 == 0 identity the adaptive node placement is + # DERIVED from must hold on this data, measured on concrete + # tables. Same predicate, same probe direction, and the same + # one the angle scheme is gated by -- one definition, in + # anglemarg.gh_laplace_supported_for_data. + _available = ["exact"] + _lap_why = None + if int(getattr(opts, "_distance_gh_nodes_resolved", 0)) <= 0: + _lap_why = ("the per-sample adaptive distance quadrature is " + "off (--distance-gh-nodes 0 / JAX_ILE_DISTMARG_GH " + "unset), and on a static distance grid the laplace " + "reserve is measured at 43.2 nats of error at " + "rho 163") + else: + # The wrapper has already measured this whenever the ANGLE + # scheme could use the placement, and stored it. Reuse that + # answer rather than measuring a second time: two measurements + # of one property can disagree, and the one the run reports + # should be the one the run is gated on. + _am_info = getattr(like, "angle_marg_info", None) or {} + if "gh_laplace_ok" in _am_info: + _lap_ok, _lap_info = _am_info["gh_laplace_ok"], _am_info + else: + _lap_ok, _lap_info = ( + _anglemarg.gh_laplace_supported_for_data( + like_data, + getattr(like, "interp", None) + or JAX_INTERP_DEFAULT)) + if not _lap_ok: + _lap_why = _lap_info.get("gh_laplace_reason", + "identity absent") + if _lap_why is None: + _available.append("laplace") + else: + print(" reserve roster: laplace NOT offered -- %s" % _lap_why) + _pair, _pair_info = predict_reserve_pair( + like_data, extras.get("guess_snr"), + reserve_time_refine_max=int( + like.policy_config.reserve_time_refine_max), + crossover_amplitude=ANGLE_MARG_CROSSOVER_AMPLITUDE, + max_time_nodes=int(like.policy_config.max_time_nodes), + requested=_requested, available=tuple(_available)) + _line = format_reserve_pair(_pair, _pair_info) + print(" " + _line) + sys.stderr.write("NOTE integrate_likelihood_extrinsic_jax: %s\n" + % _line) + if _pair is None: + # REFUSE, do not fall back. Falling back to whole-window + # refinement is the failure this analysis exists to prevent: + # it would carry the rows without anyone choosing it. + raise SystemExit( + "--direct-marginalization-reserve-scheme auto: %s No " + "implemented reserve is adequate for this signal, so the " + "run is refused rather than silently falling back. Pass " + "an explicit --direct-marginalization-reserve-scheme to " + "override, and record that you did." + % _pair_info.get("reason", "")) + # THE PAIR THE ANALYSIS CHOSE MUST BE THE PAIR THAT RUNS. + # Announcing it and leaving policy_config.reserve_scheme == "auto" + # would make the printed line a claim about nothing: the composite + # reads the config, not this print. validate_policy_config admits + # "auto" ONLY because it is resolved here, so this write is what + # makes that admission true (found by the #304 session, whose + # composite dispatches on the string). + if _pair not in RESERVE_SCHEME_EXECUTABLE: + raise SystemExit( + "the analysis selected the %r reserve (%s), and the " + "composite dispatches only %r. Refused rather than " + "running a different reserve under the selected one's " + "name. Wiring %r is RIFT PR #304." + % (_pair, _pair_info.get("reason", ""), + RESERVE_SCHEME_EXECUTABLE, _pair)) + # The wrapper's jitted closures capture the config at + # construction, so writing the attribute alone would leave the + # composite running the string 'auto' (which reserve_pair refuses): + # rebuild the likelihood with the resolved pair. + if _pair != like.policy_config.reserve_scheme: + like = JAXDistPhiPsiMargLikelihood( + 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, + direct_marginalization_policy=policy, + policy_config=like.policy_config._replace( + reserve_scheme=_pair)) + # policy_info's "reserve_scheme" reported the resolved ANGLE + # scheme, which is a different quantity that happened to share the + # name while 'exact' was the only reserve. Both are reported now, + # under names that say which is which. + like.policy_info["reserve_angle_scheme"] = \ + like.policy_info.get("reserve_scheme") + like.policy_info["reserve_scheme"] = _pair + print(" reserve pair RESOLVED: reserve_scheme=%s" % _pair) with_distance = False dim = 3 elif opts.mode == "flowmc-dpsimarg": @@ -1952,6 +3228,68 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, "fine factor is curvature-derived per row with one doubling certificate" % (like.time_guard_initial, like.time_guard_certified)) + # In-flight branch classification, BEFORE any sampling. + # + # direct_marginalization_policy_note() below reports the same counts, but it + # runs on the EXPORTED rows, so only a run that reaches the export can be + # classified -- and a declining run is precisely the one that does not get + # there. A declined row executes the exact reserve (three evaluations: the + # refined rule, its half-refined check, and the guard comparison) and rows + # run one at a time under jax.lax.map, so declining costs enough that the + # 32-row prior pilot itself does not finish. "Did the controller accept" + # was therefore unanswerable until the end of a run that never ends. + # + # The probe draws from the same prior the pilot draws from and calls the + # same ledger, so it predicts the pilot rather than describing something + # else. Its compile is not extra work for a policy run: the end-of-run note + # jits the identical _batched_ledger. + _probe_rows = int(getattr(opts, "direct_marginalization_policy_probe_rows", 0)) + if (getattr(like, "direct_marginalization_policy", "off") != "off" + and _probe_rows > 0): + import time as _time + from RIFT.likelihood.jax_ile.samplers import sample_prior_3 as _sp3 + # sample_prior_3 is correct here ONLY because the policy is refused + # outside --mode flowmc-phipsimarg, whose parameter order is + # (ra, dec, incl). Check rather than assume: another order would hand + # the ledger rows of the wrong width. + _order = tuple(getattr(like, "ANGULAR_PARAM_ORDER", ())) + if len(_order) != 3: + raise RuntimeError( + "policy probe expects the 3-parameter (ra, dec, incl) " + "likelihood the policy is restricted to; this one is %r" + % (_order,)) + _theta_probe = _sp3(_probe_rows, np.random.default_rng(opts.seed)) + _t0 = _time.time() + _probe_note, _probe_vals, _probe_acc = ( + direct_marginalization_policy_note( + like, _theta_probe, n_max=_probe_rows, chunk=_probe_rows, + return_values=True)) + _dt = _time.time() - _t0 + # stderr AND stdout: the surveys read one or the other, and an + # unclassifiable run is the failure this whole probe exists to remove. + _msg = ("PROBE integrate_likelihood_extrinsic_jax: prior-draw policy " + "ledger on %d rows in %.1f s (%.2f s/row incl. compile): %s" + % (_probe_rows, _dt, _dt / max(_probe_rows, 1), _probe_note)) + if _probe_vals is not None: + # lnL PAIRED WITH THE DISPOSITION, per row. The counts alone + # cannot say whether the rows that decline are rows that matter: + # capacity declines concentrate far from truth, where the + # likelihood is negligible, and a decline fraction quoted without + # the likelihood of the declining rows overstates what it costs + # the posterior. + _msg += ("\nPROBE-VALUES seed=%d rows=%d lnL=[%s] local=[%s]" + % (int(opts.seed), _probe_rows, + ",".join("%.10g" % v for v in _probe_vals), + ",".join("1" if a else "0" for a in _probe_acc))) + print(_msg) + sys.stderr.write(_msg + "\n") + sys.stdout.flush() + sys.stderr.flush() + if getattr(opts, "direct_marginalization_policy_probe_only", False): + print("--direct-marginalization-policy-probe-only: exiting before " + "sampling; no samples and no evidence were written.") + return None, None + if opts.mode == "map": theta_map, lnL_map = run_map(like, opts, rng, dim, with_distance) fish = like.fisher(theta_map) @@ -1964,7 +3302,43 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, return lnL_map, None out_flow_state = None - if opts.mode in ("multistart-nuts", "flowmc", "flowmc-phimarg", + if getattr(opts, "sampler_method", None) in ("AV", "portfolio"): + from RIFT.likelihood.jax_ile import samplers as _samplers + _portfolio_members = [ + name.strip() for item in (opts.sampler_portfolio or []) + for name in str(item).split(",") if name.strip()] + if not _portfolio_members: + _portfolio_members = ["AV", "GMM"] + _sample_bounds = resolve_av_angular_limits(opts) + print(" JAX-%s: value-only integration, coverage chunk=%d, eval chunk=%s%s" + % (opts.sampler_method, opts.n_chunk, + (str(opts.jax_av_eval_chunk) + if opts.jax_av_eval_chunk is not None else "auto"), + ("; members=" + ",".join(_portfolio_members)) + if opts.sampler_method == "portfolio" else "")) + _distance_sampling = av_distance_sampling_kwargs(like, d_lo, d_hi) + res = _samplers.adaptive_volume_sample( + like, opts.d_min, opts.d_max, + sampler_method=opts.sampler_method, + portfolio_members=_portfolio_members, + nmax=opts.n_max, neff=(opts.n_eff or 1000), + n_chunk=opts.n_chunk, eval_chunk=opts.jax_av_eval_chunk, + seed=opts.seed, seed_method=opts.jax_av_seed, + seed_pilot=opts.jax_av_seed_pilot, + seed_modes=opts.jax_av_seed_modes, + seed_points=opts.jax_av_seed_points, + sky_inflate=opts.jax_av_sky_inflate, + seed_prior_frac=opts.jax_av_seed_prior_frac, + anisotropic_bins=opts.sampler_anisotropic_bins, + verbose=opts.verbose, sample_bounds=_sample_bounds, + **_distance_sampling) + theta, lnL = res["theta"], res["lnL"] + logZ, sig, neff = res["logZ"], res["sigma_over_Z"], res["neff"] + ntot = res["n_eval"] + logw_export = res.get("log_weight") + print(" JAX-%s resolved fixed eval chunk=%d" + % (opts.sampler_method, res["eval_chunk"])) + elif opts.mode in ("multistart-nuts", "flowmc", "flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg", "nuts-phimarg"): if opts.mode not in ("flowmc-phimarg", "flowmc-phipsimarg", "flowmc-dpsimarg", "nuts-phimarg") and with_distance: @@ -2000,6 +3374,7 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, res = _samplers.smc_puffball_sample( like, opts.d_min, opts.d_max, n_walkers=opts.smc_walkers, n_move=opts.smc_move_steps, + is_samples=opts.smc_is_samples, ess_frac=opts.temper_ess_frac, max_dbeta=opts.temper_max_dbeta, max_stages=max(opts.temper_max_stages, 80), puff_scale=opts.smc_puff_scale, @@ -2105,6 +3480,7 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, 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)) + require_finite_evidence(logZ, neff, opts.mode) # 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 @@ -2122,11 +3498,37 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, # 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 + # None, so the standing output-cloud 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) + # Record the resolved distance-GH-nodes count in the same artifact header + # line as the angle-marg/policy notes and (via write_samples' provenance + # line) the mode/ESS record, so a reader of either artifact can see + # whether the per-sample quadrature ran without re-deriving it from the + # command line or the environment. + _gh_n = int(getattr(opts, "_distance_gh_nodes_resolved", 0) or 0) + _ev_note = (_ev_note + " gh_nodes=%d" % _gh_n).strip() + # Cross-axis policy audit: which branch the exported rows actually took. + # Evaluated on a subsample AFTER sampling because the sampler consumes the + # value-only path; the ledger is the same computation with its record kept. + if (getattr(like, "direct_marginalization_policy", "off") != "off" + and not np.all(np.isfinite(np.asarray(lnL)))): + raise SystemExit( + "--direct-marginalization-policy %s: %d of %d exported rows have a " + "non-finite likelihood because the controller could not warrant " + "them; refusing to publish samples or evidence. See the ledger " + "note: %s" % (like.direct_marginalization_policy, + int(np.sum(~np.isfinite(np.asarray(lnL)))), + int(np.size(lnL)), + direct_marginalization_policy_note( + like, theta, n_max=32))) + _policy_note = direct_marginalization_policy_note(like, theta) + if _policy_note: + sys.stderr.write("NOTE integrate_likelihood_extrinsic_jax: %s\n" + % _policy_note) + _ev_note = (_ev_note + " " + _policy_note).strip() if _ev_note.startswith("SUSPECT-ANGLE-GRID"): sys.stderr.write( "WARNING integrate_likelihood_extrinsic_jax: angle-marginalization " @@ -2136,9 +3538,8 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, "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. + # Deterministic pass over the artifact-producing cloud, with scope + # stated in the label (training-only proposals are not claimed). 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, @@ -2152,6 +3553,8 @@ def analyze_one(opts, P, data_dict, psd_dict, analyticPSD_Q, fiducial_epoch, # Main # --------------------------------------------------------------------------- def main(argv=None): + if _JAX_CACHE_DIR is not None: + print("JAX persistent compilation cache:", _JAX_CACHE_DIR) optp = build_parser() argv = _normalize_interpolate_time_argv(argv) opts, _ = optp.parse_args(argv) diff --git a/MonteCarloMarginalizeCode/Code/bin/rift_jax_cache b/MonteCarloMarginalizeCode/Code/bin/rift_jax_cache new file mode 100755 index 000000000..e8abfc841 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/bin/rift_jax_cache @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +"""Inspect, export, and safely import RIFT JAX compilation caches.""" + +import argparse +import json +import os +import sys + +from RIFT.jax_cache import ( + compatibility_key, + configure_persistent_cache, + export_bundle, + import_bundle, + runtime_compatibility, +) + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cache-root", help="cache root (default: RIFT/XDG cache root)") + commands = parser.add_subparsers(dest="command", required=True) + commands.add_parser("fingerprint", help="print this runtime/device cache identity") + export = commands.add_parser("export", help="export the active warmed cache") + export.add_argument("output") + export.add_argument("--profile") + export.add_argument("--shape", action="append", default=[], metavar="NAME=VALUE") + ingest = commands.add_parser("import", help="validate and import a cache bundle") + ingest.add_argument("bundle") + ingest.add_argument("--expect-profile") + args = parser.parse_args(argv) + + # Keep ``--help`` usable in RIFT's base installation, where JAX is an + # optional dependency. Operational subcommands require JAX, but argparse + # exits for help before this import is reached. + import jax + + compatibility = runtime_compatibility(jax) + if args.command == "fingerprint": + print(json.dumps({"compatibility_key": compatibility_key(compatibility), + "compatibility": compatibility}, indent=2, sort_keys=True)) + return 0 + + configure_args = (["--jax-cache-dir", args.cache_root] if args.cache_root else []) + active = configure_persistent_cache(jax, configure_args) + if active is None: + parser.error("the JAX cache directory is unavailable; choose a writable --cache-root") + if args.command == "export": + shapes = {} + for item in args.shape: + if "=" not in item: + parser.error("--shape must be NAME=VALUE") + key, value = item.split("=", 1) + shapes[key] = value + manifest = export_bundle(active, args.output, compatibility, args.profile, shapes) + print(json.dumps(manifest, indent=2, sort_keys=True)) + else: + root = args.cache_root or str(active.parent) + # ``active`` is authoritative even when the standard JAX environment + # variable names an exact (non-namespaced) directory. Importing into a + # derived sibling would succeed but the next ILE would never read it. + destination = import_bundle(args.bundle, root, compatibility, + args.expect_profile, destination=active) + print(destination) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py index 656c072ad..e599829b5 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py @@ -65,6 +65,10 @@ from RIFT.likelihood.time_marginalization_quadrature import ( TIME_QUADRATURE_CHOICES, validate_time_quadrature, refuse_unhonourable_time_quadrature, refuse_unless_time_quadrature_emitted) +# Same reason, same leaf-module discipline, for the Q_lm pregrid factor. +from RIFT.likelihood.q_time_pregrid import ( + Q_TIME_PREGRID_CHOICES, validate_q_time_pregrid_factor, + refuse_unhonourable_q_time_pregrid, refuse_unless_q_time_pregrid_emitted) ligolw_prefix = 'igwn_' if not(which(ligolw_prefix + "ligolw_add")): ligolw_prefix = '' @@ -210,6 +214,17 @@ def run_lisa_known_sky_surface(opts): if opts.approx is None: print(" --lisa-known-sky requires --approx ") sys.exit(1) + if opts.use_jax_ile: + # This path hardcodes integrate_likelihood_extrinsic_batchmode_lisa below + # (a separate, LISA-specific driver) and never reads opts.use_jax_ile; it is + # not "the same code" as the LDG/OSG ILE selection, so silently ignoring the + # flag would leave a user believing they got the JAX driver when they did + # not. Refuse rather than run the wrong driver silently. + print(" --use-jax-ile has no effect on --lisa-known-sky: this path always " + "runs integrate_likelihood_extrinsic_batchmode_lisa, a separate " + "LISA-specific driver with no JAX equivalent in this repository. " + "Drop --use-jax-ile for LISA runs.") + sys.exit(1) if opts.use_ini is not None: # LISA production-ini path: scalars/algorithm options come from the # generic [rift-pseudo-pipe] parser; fill the per-channel data products @@ -495,6 +510,7 @@ def run_lisa_known_sky_surface(opts): 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-q-time-pregrid-factor",default=None,type=int,choices=list(Q_TIME_PREGRID_CHOICES),help="OPT-IN certified Q_lm pregrid in ILE (PR #261): %s. Default None = pass nothing, so the ILE default (factor 1, the historical unchanged path) is unchanged and the emitted args_ile.txt is byte-identical to today. Factor 8 reflects each finite Q window, FFT-interpolates it onto an 8x finer grid once after packing, and evaluates detector arrival times off that grid with four-tap CUBIC interpolation -- the geocentric time-integration grid is left at the data deltaT. Forwarded verbatim to helper_LDG_Events.py, which validates it and puts --q-time-pregrid-factor on the ILE command line. REFUSED, not ignored, at DAG-BUILD TIME if this workflow cannot honour it (calibration marginalization, --rotation-slow, --freqresponse, a configuration without --vectorized, or an explicit --internal-ile-interpolate-time naming a stencil other than cubic -- factor 8 forces cubic and will not silently override a different explicit request). IMPORTANT -- INI OVERRIDE: the RIFT ini parser OVERRIDES the command line for non-boolean options, so NEVER set this in a --use-ini that a Makefile or wrapper also sets on the command line. Rationale: bin/integrate_likelihood_extrinsic_batchmode (grep q_time_pregrid)." % ("|".join(str(c) for c in Q_TIME_PREGRID_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 @@ -517,6 +533,9 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--ile-xpu",action='store_true',help='Request ILE run on both GPU and CPU. Disables ile_force_gpu, if provided!') parser.add_argument("--ile-force-gpu",action='store_true') parser.add_argument("--ile-gpu-fanout",default=None,help="Multi-GPU ILE fan-out: split each ILE batch's intrinsic-grid range across N GPUs on the node (one shard per GPU). Integer N (also requests N GPUs+CPUs) or 'auto' (split across whatever GPUs are visible at runtime). Baked into the generated ile_pre.sh, so it needs no runtime environment. Equivalent to setting RIFT_ILE_GPU_FANOUT. Requires --ile-force-gpu.") +parser.add_argument("--ile-exe",default=None,type=str,help="Path to the ILE executable used for this workflow's ILE/ILE_puff/ILE_fetch/ILE_extr jobs (forwarded to create_event_parameter_pipeline_* as --ile-exe). Default: `which integrate_likelihood_extrinsic_batchmode`, or `which integrate_likelihood_extrinsic_jax` if --use-jax-ile is set. Mutually exclusive with --use-jax-ile.") +parser.add_argument("--use-jax-ile",action='store_true',help="Use `which integrate_likelihood_extrinsic_jax` as the ILE executable in place of the default batchmode driver, for every ILE/ILE_puff/ILE_fetch/ILE_extr job. The JAX driver does not implement calibration marginalization, ROM/NR-lookup templates, supplementary likelihood factors, --zero-likelihood, or --maximize-only (see check_critical_and_report in bin/integrate_likelihood_extrinsic_jax); combining it with --calmarg-envelope-directory is REFUSED at DAG-build time rather than left to fail at the first ILE job. Mode-specific JAX options (--mode, --angle-marg-scheme, and the rest of that driver's surface) are not separate pseudo_pipe options -- pass them through --manual-extra-ile-args. OSG/SINGULARITY CAVEAT: --use-osg unconditionally adds --use-singularity, and write_ILE_sub_simple then rewrites the condor executable to /integrate_likelihood_extrinsic_jax (basename preserved) -- but no container this repository builds (Dockerfile, containers/rift_container.def.in, rift_container_family.yaml) installs JAX or that driver, so every ILE job would fail at runtime. --use-jax-ile with --use-osg is therefore REFUSED at DAG-build time unless --jax-ile-container-ok is also given, which asserts the named SINGULARITY_RIFT_IMAGE/SINGULARITY_BASE_EXE_DIR image actually provides integrate_likelihood_extrinsic_jax and JAX. Also REFUSED with --lisa-known-sky, which always runs the separate integrate_likelihood_extrinsic_batchmode_lisa driver and has no JAX equivalent.") +parser.add_argument("--jax-ile-container-ok",action='store_true',help="Override the --use-jax-ile + --use-osg refusal (see --use-jax-ile help). Pass this ONLY when the container named by SINGULARITY_RIFT_IMAGE / resolved via SINGULARITY_BASE_EXE_DIR actually provides both a JAX installation and the integrate_likelihood_extrinsic_jax executable -- no container built by this repository does. Has no effect without --use-jax-ile, and does not affect the separate --lisa-known-sky refusal.") parser.add_argument("--fake-data-cache",type=str) parser.add_argument("--spin-magnitude-prior",default='default',type=str,help="options are default [uniform mag for precessing, zprior for aligned], volumetric, uniform_mag_prec, uniform_mag_aligned, zprior_aligned") parser.add_argument("--eccentricity-prior",default='uniform',type=str,choices=['uniform','log_uniform'],help="options are uniform in e ('uniform') and uniform in log(e) ('log_uniform')") # constrained: the value is forwarded verbatim to CIP, which only branches on the exact string 'log_uniform', so an unrecognized value here would silently run the uniform prior instead of failing @@ -731,6 +750,21 @@ 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) +# JAX ILE selection. Placed AFTER the --use-ini block above: the ini can set both +# use-jax-ile and calmarg-envelope-directory via [rift-pseudo-pipe], and a check +# above that block would validate values the ini is about to replace. +if opts.use_jax_ile and opts.ile_exe: + raise ValueError( + "--use-jax-ile and --ile-exe are mutually exclusive: --use-jax-ile already " + "resolves to `which integrate_likelihood_extrinsic_jax`. Pass that " + "executable's path via --ile-exe directly instead of setting both.") +if opts.use_jax_ile and opts.calmarg_envelope_directory: + raise ValueError( + "--use-jax-ile is incompatible with in-loop calibration marginalization " + "(--calmarg-envelope-directory): bin/integrate_likelihood_extrinsic_jax does " + "not implement --calibration-* options (see check_critical_and_report in " + "that driver). Drop --calmarg-envelope-directory or drop --use-jax-ile.") + # 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 @@ -765,6 +799,38 @@ def run_lisa_known_sky_surface(opts): "this pipeline's own options (checked before the helper runs, so the failure is " "immediate rather than after a workflow has been built)") +# Q_lm PREGRID FACTOR, part 1 of 2: everything refusable WITHOUT running the helper. Same +# placement reasoning as the quadrature block just above (after --use-ini, so a bad ini value +# is not checked before the ini is about to replace it). +if opts.internal_ile_q_time_pregrid_factor 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_q_time_pregrid_factor(opts.internal_ile_q_time_pregrid_factor) + if opts.lisa_known_sky: + # Same reasoning as the quadrature check: --lisa-known-sky builds args_ile.txt through + # helper_LISA_Events.py, which does not carry this option either. + raise ValueError( + "--internal-ile-q-time-pregrid-factor 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 (including a + # conflicting explicit --interpolate-time). + _qp_early = "" + if opts.calmarg_envelope_directory: + _qp_early += " --calibration-envelope-directory " + str(opts.calmarg_envelope_directory) + if opts.manual_extra_ile_args: + _qp_early += " " + str(opts.manual_extra_ile_args) + if _qp_early: + # Only the EXCLUSIONS (and a manually-passed stencil conflict) are checkable this early + # -- the required --vectorized is added by the helper -- so append it to keep the + # message about what is actually wrong. + refuse_unhonourable_q_time_pregrid( + opts.internal_ile_q_time_pregrid_factor, + "--vectorized " + _qp_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) @@ -780,6 +846,26 @@ def run_lisa_known_sky_surface(opts): opts.condor_local_nonworker_igwn_prefix=False opts.condor_nogrid_nonworker=False +# --use-jax-ile + --use-osg: --use-osg unconditionally pairs with --use-singularity +# (see the cmd built below), and write_ILE_sub_simple then rewrites the condor +# executable to /, discarding the +# `which integrate_likelihood_extrinsic_jax` resolution above. No container this +# repository builds carries JAX or that driver, so every ILE job would fail at +# runtime, silently at build time. Placed here (after opts.use_osg_public is +# folded into opts.use_osg) so it sees the final value of opts.use_osg regardless +# of which flag or ini key set it. +if opts.use_jax_ile and opts.use_osg and not opts.jax_ile_container_ok: + raise ValueError( + "--use-jax-ile with --use-osg is REFUSED at DAG-build time: --use-osg " + "always adds --use-singularity, and write_ILE_sub_simple rewrites the " + "condor executable to /, " + "but no container this repository builds (Dockerfile, " + "containers/rift_container.def.in, rift_container_family.yaml) installs " + "JAX or integrate_likelihood_extrinsic_jax -- every ILE job would fail at " + "runtime. If SINGULARITY_RIFT_IMAGE/SINGULARITY_BASE_EXE_DIR names an " + "image you have verified provides both, pass --jax-ile-container-ok to " + "proceed. Otherwise drop --use-jax-ile or --use-osg.") + if opts.ile_copies <=0: raise Exception(" Must have 1 or more ILE instances per intrinsic point") @@ -1430,6 +1516,9 @@ def approx_supports_precession(approx_name): # 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 opts.internal_ile_q_time_pregrid_factor is not None: + # HELPER passthrough, exactly like the two options above and for the same reason. + cmd += " --internal-ile-q-time-pregrid-factor " + str(opts.internal_ile_q_time_pregrid_factor) + " " 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 @@ -1712,6 +1801,11 @@ def approx_supports_precession(approx_name): # 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") +# Same discipline, same reason, for the Q_lm pregrid factor -- including the forced-cubic-stencil +# conflict, which can only be checked here: --interpolate-time is resolved and emitted by the +# helper, so it is not visible to this script until `line` is fully assembled. +refuse_unless_q_time_pregrid_emitted( + opts.internal_ile_q_time_pregrid_factor, line, "args_ile.txt") with open('args_ile.txt','w') as f: f.write(line) @@ -2235,7 +2329,15 @@ def approx_supports_precession(approx_name): 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) + ("" 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) +# Resolve the ILE executable: --use-jax-ile wins (mutual exclusion with --ile-exe was +# already enforced above), then an explicit --ile-exe, then the historical default. +if opts.use_jax_ile: + resolved_ile_exe = "`which integrate_likelihood_extrinsic_jax`" +elif opts.ile_exe: + resolved_ile_exe = "'{}'".format(opts.ile_exe) +else: + resolved_ile_exe = "`which integrate_likelihood_extrinsic_batchmode`" +cmd =cepp+ " --ile-n-events-to-analyze {} --input-grid proposed-grid.{} --ile-exe {} --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,resolved_ile_exe,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 diff --git a/MonteCarloMarginalizeCode/Code/test/asimov_integration/test_asimov_rift_template_contract.py b/MonteCarloMarginalizeCode/Code/test/asimov_integration/test_asimov_rift_template_contract.py index 2128ff9ff..d9be5ad92 100644 --- a/MonteCarloMarginalizeCode/Code/test/asimov_integration/test_asimov_rift_template_contract.py +++ b/MonteCarloMarginalizeCode/Code/test/asimov_integration/test_asimov_rift_template_contract.py @@ -253,6 +253,37 @@ def test_rift_liquid_template_option_blocks_land_safely(): assert "--zero-likelihood" in parser.get("rift-pseudo-pipe", "manual-extra-ile-args") +def test_rift_liquid_template_time_stencil_keys_default_absent(): + """RO'S directive 2026-09-08: the three time-stencil/quadrature ledger keys must default + to ABSENT, so an existing production's rendered ini is unchanged unless the ledger sets + one of them.""" + meta = _base_meta() + rendered, parser = _render(meta) + for key in ("internal-ile-interpolate-time", + "internal-ile-time-marginalization-quadrature", + "internal-ile-q-time-pregrid-factor"): + assert not parser.has_option("rift-pseudo-pipe", key), key + assert key not in rendered + + +def test_rift_liquid_template_time_stencil_keys_render_when_set(): + meta = _base_meta() + meta["sampler"]["ile"]["interpolate time"] = "cubic" + meta["sampler"]["ile"]["time marginalization quadrature"] = "bandlimited" + meta["sampler"]["ile"]["q time pregrid factor"] = 8 + + _rendered, parser = _render(meta) + + # String-valued options must be quoted before the generic pseudo_pipe ini-override loop + # eval()s them, exactly like ile-sampler-method/ile-distance-prior above -- otherwise + # eval("cubic") raises NameError rather than yielding the Python string "cubic". + assert parser.get("rift-pseudo-pipe", "internal-ile-interpolate-time").strip("'\"") == "cubic" + assert parser.get("rift-pseudo-pipe", + "internal-ile-time-marginalization-quadrature").strip("'\"") == "bandlimited" + # The pregrid factor is an int pipeline option, so it must render UNQUOTED. + assert parser.get("rift-pseudo-pipe", "internal-ile-q-time-pregrid-factor").strip() == "8" + + def test_rift_liquid_template_randomized_ledger_sanity(): rng = random.Random(190426) approximants = ["SEOBNRv5PHM", "IMRPhenomXPHM", "TaylorF2"] @@ -297,3 +328,13 @@ def test_rift_liquid_template_randomized_ledger_sanity(): for ifo in ifos: assert f'"{ifo}":"{ifo}_TEST_FRAME"' in parser.get("datafind", "types") assert f'"{ifo}":"{ifo}:TEST-STRAIN"' in parser.get("data", "channels") + + +def test_rift_liquid_template_use_jax_ile_defaults_false_and_follows_ledger(): + meta = _base_meta() + _rendered, parser = _render(meta) + assert parser.get("rift-pseudo-pipe", "use-jax-ile").strip() == "False" + + meta["sampler"]["ile"]["use jax ile"] = True + _rendered, parser = _render(meta) + assert parser.get("rift-pseudo-pipe", "use-jax-ile").strip() == "True" diff --git a/MonteCarloMarginalizeCode/Code/test/benchmark_bandlimited_retained_fft.py b/MonteCarloMarginalizeCode/Code/test/benchmark_bandlimited_retained_fft.py new file mode 100644 index 000000000..04153d2d3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/benchmark_bandlimited_retained_fft.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +"""Reproduce the full-padding versus retained-grid FFT microbenchmark. + +Run each timed arm in a fresh process so ``ru_maxrss`` and the CuPy memory pool +belong to that arm. ``parity`` evaluates both arms on one deterministic batch +and reports differences after a nonlinear likelihood-like map and time +integration. This is a transform/kernel benchmark, not an ILE evidence run. + +Examples (inside a RIFT environment):: + + python benchmark_bandlimited_retained_fft.py --backend cupy --arm full \ + --npts 614 --factor 64 + python benchmark_bandlimited_retained_fft.py --backend cupy --arm retained \ + --npts 614 --factor 64 + python benchmark_bandlimited_retained_fft.py --backend cupy --arm parity \ + --npts 614 --factor 64 +""" +import argparse +import json +import os +import resource +import time + +# A benchmark must not inherit a many-thread BLAS default and then measure +# thread creation or exceed a batch system's process limit during imports. +for _thread_env in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", + "MKL_NUM_THREADS", "NUMEXPR_NUM_THREADS"): + os.environ.setdefault(_thread_env, "1") + +import numpy as np + +from RIFT.likelihood import time_marginalization_quadrature as tmq + + +def _backend(name): + if name == "numpy": + from scipy.special import logsumexp + return np, logsumexp + import cupy + from cupyx.scipy.special import logsumexp + if cupy.cuda.runtime.getDeviceCount() < 1: + raise RuntimeError("--backend cupy requested but no CUDA device is visible") + return cupy, logsumexp + + +def _synchronize(xpy): + if xpy is not np: + xpy.cuda.Stream.null.synchronize() + + +def _memory_start(xpy): + if xpy is np: + return None + try: + xpy.fft.config.get_plan_cache().clear() + except Exception: + pass + xpy.get_default_memory_pool().free_all_blocks() + xpy.get_default_pinned_memory_pool().free_all_blocks() + _synchronize(xpy) + free, total = xpy.cuda.runtime.memGetInfo() + return free, total + + +def _memory_finish(xpy, start): + out = { + "host_maxrss_mib": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0, + "cupy_pool_total_mib": None, + "device_resident_delta_mib": None, + "device_total_mib": None, + } + if xpy is not np: + free, _ = xpy.cuda.runtime.memGetInfo() + out.update( + cupy_pool_total_mib=xpy.get_default_memory_pool().total_bytes() / 2**20, + device_resident_delta_mib=(start[0] - free) / 2**20, + device_total_mib=start[1] / 2**20, + ) + return out + + +def _inputs(nrows, npts, xpy): + rng = np.random.default_rng(20260905 + npts) + host = rng.normal(size=(nrows, npts)) + 1j * rng.normal( + size=(nrows, npts)) + host *= np.exp(0.013j * np.arange(nrows)[:, None]) + return xpy.asarray(host, dtype=np.complex128) + + +def _transform(arm, rows, factor, cache, xpy): + if arm == "full": + return tmq.reflected_bandlimited_upsample(rows, factor, xpy=xpy) + return tmq._reflected_bandlimited_upsample_retained( + rows, factor, plan_cache=cache, xpy=xpy) + + +def _timed(args, xpy): + batch = args.batch or max(1, int( + tmq._DENSE_CHUNK_BYTES // (args.npts * args.factor * 16 * 8))) + rows = _inputs(batch, args.npts, xpy) + if xpy is not np: + xpy.fft.fft(xpy.ones((1, 32), dtype=np.complex128)).sum().get() + start_memory = _memory_start(xpy) + cache = {} + checksum = xpy.zeros((), dtype=np.float64) + _synchronize(xpy) + start = time.perf_counter() + done = 0 + while done < args.rows: + take = min(batch, args.rows - done) + dense = _transform(args.arm, rows[:take], args.factor, cache, xpy) + checksum += xpy.sum(dense[..., ::args.factor].real) + del dense + done += take + _synchronize(xpy) + wall = time.perf_counter() - start + checksum = float(checksum if xpy is np else checksum.get()) + record = { + "arm": args.arm, + "backend": args.backend, + "npts": args.npts, + "factor": args.factor, + "rows": args.rows, + "batch": batch, + "dense_points_evaluated": args.rows * ((args.npts - 1) * args.factor + 1), + "wall_s": wall, + "rows_per_s": args.rows / wall, + "checksum": checksum, + "full_fft_length": 2 * args.npts * args.factor, + "retained_grid_length": (args.npts - 1) * args.factor + 1, + "retained_plan_fft_length": max( + (p["n_fft"] for p in cache.values()), default=None), + } + record.update(_memory_finish(xpy, start_memory)) + return record + + +def _parity(args, xpy, logsumexp): + batch = args.batch or 32 + rows = _inputs(batch, args.npts, xpy) + full = tmq.reflected_bandlimited_upsample(rows, args.factor, xpy=xpy) + retained = tmq._reflected_bandlimited_upsample_retained( + rows, args.factor, xpy=xpy) + # Smooth and nonlinear, as distance/phase marginalization is. The factor + # 100 makes transform-level roundoff visible instead of rounding to zero. + lnlt_full = 100.0 * xpy.logaddexp(0.0, full.real) + lnlt_retained = 100.0 * xpy.logaddexp(0.0, retained.real) + + def integrate(lnlt): + offset = xpy.max(lnlt, axis=-1) + density = xpy.exp(lnlt - offset[:, None]) + density[:, 0] *= 0.5 + density[:, -1] *= 0.5 + return offset + xpy.log(xpy.sum(density, axis=-1) / args.factor) + + il_full = integrate(lnlt_full) + il_retained = integrate(lnlt_retained) + lnz_full = logsumexp(il_full) - np.log(batch) + lnz_retained = logsumexp(il_retained) - np.log(batch) + _synchronize(xpy) + + def scalar(value): + return float(value if xpy is np else value.get()) + + return { + "arm": "parity", + "backend": args.backend, + "npts": args.npts, + "factor": args.factor, + "rows": batch, + "max_abs_delta_kappa": scalar(xpy.max(xpy.abs(retained - full))), + "max_abs_delta_lnL": scalar(xpy.max(xpy.abs(il_retained - il_full))), + "delta_lnZ": scalar(lnz_retained - lnz_full), + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--backend", choices=("numpy", "cupy"), default="numpy") + parser.add_argument("--arm", choices=("full", "retained", "parity"), required=True) + parser.add_argument("--npts", type=int, required=True) + parser.add_argument("--factor", type=int, required=True) + parser.add_argument("--rows", type=int, default=40000) + parser.add_argument("--batch", type=int) + args = parser.parse_args() + xpy, logsumexp = _backend(args.backend) + record = (_parity(args, xpy, logsumexp) if args.arm == "parity" + else _timed(args, xpy)) + print(json.dumps(record, sort_keys=True)) + + +if __name__ == "__main__": + main() 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 0470652a0..eedc631f5 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 @@ -309,6 +309,14 @@ "decision": "PORT", "reason": "Normalizing-flow persistence is detector-agnostic, but the LISA portfolio factory currently constructs only AV, GMM, and adaptive_cartesian_gpu members. Port the NF member construction and route load/save to that member before exposing these flags; hooks on the portfolio aggregate are a silent no-op because it has no flow API." }, + "OPTION:--psi-marginalization": { + "decision": "NA", + "reason": "Analytic polarization-angle marginalization via factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized. Per its own help text, this is only reachable on the legacy scalar (non-vectorized, non-GPU, non-time-marginalized) likelihood path -- the main driver itself REFUSES it with --vectorized, --gpu, and --time-marginalization, among others. The LISA driver has no such scalar path: analyze_event there is vectorized/time-marginalized by construction, so this option would be refused there too if it existed. Not a gap to close -- porting it would add an option that is dead on arrival." + }, + "OPTION:--q-time-pregrid-factor": { + "decision": "PORT", + "reason": "Refines the finite Q time grid before detector-arrival interpolation. LISA uses the same NoLoop Q gather and can carry the same interpolation bias, so this is not ground-detector-specific. Port only after separate LISA accuracy and memory validation: its long observation windows make an unconditional 8x retained grid potentially much more expensive than in the ground-based driver." + }, "OPTION:--random-event": { "decision": "PORT", "reason": "Pick a random event from the input file. Detector-agnostic; flagged dangerous in its own help text for oversampling reasons that apply equally to LISA." 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 dd8444c0a..d73d134e2 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 @@ -399,6 +399,12 @@ (r"^FUNC:_normalize_interpolate_time_argv$", "PORT", "Normalizes --interpolate-time argv forms. LISA exposes --interpolate-time, so " "the same normalization applies."), + (r"^OPTION:--q-time-pregrid-factor$", "PORT", + "Refines the finite Q time grid before detector-arrival interpolation. LISA uses " + "the same NoLoop Q gather and can carry the same interpolation bias, so this is " + "not ground-detector-specific. Port only after separate LISA accuracy and memory " + "validation: its long observation windows make an unconditional 8x retained grid " + "potentially much more expensive than in the ground-based driver."), (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 " @@ -438,6 +444,15 @@ "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"^OPTION:--psi-marginalization$", "NA", + "Analytic polarization-angle marginalization via " + "factored_likelihood.NetworkLogLikelihoodPolarizationMarginalized. Per its own help " + "text, this is only reachable on the legacy scalar (non-vectorized, non-GPU, " + "non-time-marginalized) likelihood path -- the main driver itself REFUSES it with " + "--vectorized, --gpu, and --time-marginalization, among others. The LISA driver has " + "no such scalar path: analyze_event there is vectorized/time-marginalized by " + "construction, so this option would be refused there too if it existed. Not a gap " + "to close -- porting it would add an option that is dead on arrival."), (r"^CONST:_SEQ_WS_PENDING$", "PORT", "Sentinel for the deferred sequential warm-start capture; ports with " "--sampler-sequential-warmstart."), diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_integrator_studies.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_integrator_studies.py new file mode 100644 index 000000000..d6b91ebe2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_integrator_studies.py @@ -0,0 +1,64 @@ +"""Run the five integrator quantitative studies as real subprocesses and require exit 0. + +WHY A WRAPPER, AND WHY IN ORDINARY CI. These five scripts carry the only assertions anyone has +written about AV warm-starting and portfolio allocation -- a 4-sigma bias gate, an anti-bias +ordering under a mis-placed proposal, a draw-allocation comparison against standalone AV, safety +under a decoy member, and an oracle finding a needle. Each ends in `raise SystemExit(1)` on +failure UNDER `--as-test`, so the pass/fail signal is real and machine-readable. That flag is not +optional here: every one of these scripts keeps its scientific comparisons and its `SystemExit(1)` +behind `if args.as_test`, and without it a biased or otherwise invalid result still prints and exits +0 -- the wrapper would then detect only crashes, not the behaviour it claims to gate. None of them +ran in CI at all before this: they have +a __main__ and argparse and no test functions, so pytest collects ZERO items and exits 5 -- "no +tests ran", which reads as a pass -- and .travis/ci_roster.txt carried them as HANDRUN. + +That roster entry called them "expensive", which is why the suggested fix was an opt-in wrapper +behind RIFT_RUN_EXPENSIVE. MEASURED, and the premise was wrong: on CIT with the IGWN python +(OMP_NUM_THREADS=1) they take 5, 2, 14, 4 and 4 seconds -- 29 s for all five. Nothing here needs +to be opt-in. + +FLAKE RISK, since these are Monte Carlo studies with tolerance-based gates: all five seed +explicitly (numpy RandomState(0/1/3) and np.random.seed), so they are deterministic rather than +merely lucky, and three consecutive runs of each exited 0. Three runs is not a flake proof; if +one does prove marginal in CI, tighten ITS seed or widen ITS stated tolerance, and do not +delete the gate. + +Subprocess rather than import: each is a __main__ script with argparse, and running it the way a +human runs it is the point -- it is what keeps the wrapper honest about the entry point. +""" + +import os +import subprocess +import sys + +import pytest + +HERE = os.path.dirname(os.path.abspath(__file__)) +CODE = os.path.normpath(os.path.join(HERE, "..", "..")) + +# name -> measured wall seconds on CIT, for whoever wonders what this costs +STUDIES = [ + ("test_AV_bootstrap.py", 5), + ("test_AV_warmstart_safety.py", 2), + ("test_portfolio_adaptive_alloc.py", 14), + ("test_portfolio_balance_heuristic.py", 4), + ("test_portfolio_oracle.py", 4), +] + + +@pytest.mark.parametrize("script,_secs", STUDIES) +def test_study_exits_clean(script, _secs): + path = os.path.join(HERE, script) + assert os.path.exists(path), ( + "%s is gone. It carried the only assertions on this behaviour; restore it or remove " + "this entry deliberately." % script) + env = dict(os.environ) + env["PYTHONPATH"] = CODE + os.pathsep + env.get("PYTHONPATH", "") + env.setdefault("OMP_NUM_THREADS", "1") + env.setdefault("MPLBACKEND", "Agg") + # --as-test is what turns each study from a printout into a gate; see module docstring. + pr = subprocess.run([sys.executable, path, "--as-test"], env=env, timeout=900, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + out = pr.stdout.decode("utf-8", "replace") + assert pr.returncode == 0, "%s --as-test exited %d; its own gate failed.\n%s" % ( + script, pr.returncode, out[-3000:]) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py index d49f31ec4..c6a36f921 100644 --- a/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py @@ -14,22 +14,70 @@ def _load_driver_helpers(): - """Import the helpers out of the driver script without executing it.""" + """Import the helpers out of the driver script without executing it. + + THE FAILURE MODE THIS GUARDS. Slicing functions out by regex means the copy here goes + stale silently whenever the driver grows a helper. It did: _lnZ_of_rvs and + _kish_neff_of_rvs were refactored to delegate to _lw_of, which was not on this list, so + inside the exec'd module _lw_of was undefined -- and the driver's own + `except Exception: return None` swallowed the NameError and returned None. Ten of the + fifteen tests then died on `None - float`, a diagnosis three steps from the cause, and the + file was reachable from no CI job so nobody saw it for weeks. + + So the name list is no longer the only defence. After exec, every global each sliced + function references must resolve, and the error names the missing helper. That turns "the + driver grew a helper" from a puzzle into a one-line fix. + """ here = os.path.dirname(os.path.abspath(__file__)) path = os.path.join(here, "..", "..", "bin", "integrate_likelihood_extrinsic_batchmode") src = open(os.path.normpath(path)).read() mod = types.ModuleType("drv") mod.numpy = numpy # ln_weights_from_rvs first: the others now delegate to it (one canonical definition of the - # importance weight, see the driver docstring). - for fn in ("_rvs_lnL_convention", "ln_weights_from_rvs", "_rvs_len", "_pool_replica_rvs", - "_lnZ_of_rvs", "_kish_neff_of_rvs"): + # importance weight, see the driver docstring). _lw_of is the shared weight reconstruction + # that _lnZ_of_rvs and _kish_neff_of_rvs both call. + names = ("_rvs_lnL_convention", "ln_weights_from_rvs", "_rvs_len", "_lw_of", + "_pool_replica_rvs", "_lnZ_of_rvs", "_kish_neff_of_rvs") + for fn in names: m = re.search(r"^def %s\(.*?(?=\n\ndef |\n\nclass )" % fn, src, re.S | re.M) assert m, "helper %s not found in the driver" % fn exec(compile(m.group(0), "", "exec"), mod.__dict__) + _assert_globals_resolve(mod, names) return mod +def _assert_globals_resolve(mod, names): + """Every global name the sliced functions reference must exist in the sliced module. + + Without this the next helper the driver factors out reaches these tests as a None return + (the driver catches Exception broadly) rather than as a missing name. + """ + import builtins + + def _referenced(code, seen): + for n in code.co_names: + seen.add(n) + for c in code.co_consts: + if isinstance(c, types.CodeType): + _referenced(c, seen) + return seen + + missing = set() + for fn in names: + for n in _referenced(getattr(mod, fn).__code__, set()): + if n in mod.__dict__ or hasattr(builtins, n): + continue + # Attribute names appear in co_names too (numpy.log -> "log"); only flag names + # that look like the driver's own module-level helpers. + if n.startswith("_") or n.endswith("_of_rvs") or n.startswith("ln_weights"): + missing.add((fn, n)) + assert not missing, ( + "sliced helpers reference names that were not sliced out of the driver: %s.\n" + "The driver factored out a helper these delegate to; add it to `names` above. " + "Without this check it arrives as a None return and fails as `None - float`." + % sorted(missing)) + + DRV = _load_driver_helpers() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/profile_rotating_freqresponse.py b/MonteCarloMarginalizeCode/Code/test/jax/profile_rotating_freqresponse.py new file mode 100644 index 000000000..00d9d3db3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/profile_rotating_freqresponse.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python +"""CPU usability profile for the compound conventional and JAX likelihoods. + +The default waveform is IMRPhenomD with only the (2,+/-2) mode pair. This is a +basis-scaling benchmark, not a science-accuracy benchmark: start cheap, identify +the usable Qmax range, then repeat with the intended long-duration source. +""" + +import argparse +import time + +import numpy as np +import jax +jax.config.update("jax_enable_x64", True) +import lal +import lalsimulation as lalsim + +import RIFT.lalsimutils as lsu +import RIFT.likelihood.factored_likelihood as fl +import RIFT.likelihood.factored_likelihood_rotating_freqresponse as flrr +import RIFT.likelihood.slowrot_freqresponse as sfr +from RIFT.likelihood.jax_ile.banded import build_rotating_freqresponse_data +from RIFT.likelihood.jax_ile.core import fused_log_likelihood + + +def timed(call): + start = time.perf_counter() + value = call() + if hasattr(value, "block_until_ready"): + value.block_until_ready() + return value, time.perf_counter() - start + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--qmax", default="0,1,2,3,4") + parser.add_argument("--pmax", type=int, default=0) + parser.add_argument("--samples", type=int, default=64) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--detectors", default="H1") + parser.add_argument("--mass1", type=float, default=30.0) + parser.add_argument("--mass2", type=float, default=25.0) + parser.add_argument("--fmin", type=float, default=30.0) + parser.add_argument("--fmax", type=float, default=512.0) + parser.add_argument("--srate", type=float, default=1024.0) + parser.add_argument("--delta-f", type=float, default=0.5) + parser.add_argument("--arm-length", type=float, default=40000.0) + args = parser.parse_args() + + event_time = 1.0e9 + deltaT = 1.0 / args.srate + detectors = tuple(x.strip() for x in args.detectors.split(",") if x.strip()) + P = lsu.ChooseWaveformParams( + fmin=args.fmin, radec=True, incl=0.3, phiref=0.0, theta=0.2, + phi=1.0, psi=0.4, m1=args.mass1 * lal.MSUN_SI, + m2=args.mass2 * lal.MSUN_SI, detector=detectors[0], + dist=200e6 * lal.PC_SI, deltaT=deltaT, tref=event_time, + deltaF=args.delta_f) + P.approx = lalsim.IMRPhenomD + data_dict = {} + for det in detectors: + Pd = P.manual_copy() + Pd.detector = det + data_dict[det] = lsu.non_herm_hoff(Pd) + psd_dict = {det: lalsim.SimNoisePSDaLIGOZeroDetHighPower for det in detectors} + tvals = fl.marginalization_time_grid(0.03, deltaT, xpy=np) + + rng = np.random.default_rng(20260909) + S = args.samples + Pv = P.manual_copy() + Pv.phi = rng.uniform(0, 2 * np.pi, S) + Pv.theta = np.arcsin(rng.uniform(-1, 1, S)) + Pv.psi = rng.uniform(0, np.pi, S) + Pv.incl = np.arccos(rng.uniform(-1, 1, S)) + Pv.phiref = rng.uniform(0, 2 * np.pi, S) + Pv.dist = rng.uniform(100, 800, S) * 1e6 * lal.PC_SI + Pv.tref = event_time + Pv.deltaT = deltaT + dist_mpc = np.asarray(Pv.dist) / (1e6 * lal.PC_SI) + + print("model=IMRPhenomD modes=(2,+/-2) detectors=%s samples=%d pmax=%d" % + (",".join(detectors), S, args.pmax)) + print("Qmax basis pairs bank_MB precompute_s numpy_s jax_compile_s jax_warm_s samples/s") + for qmax in [int(x) for x in args.qmax.split(",")]: + def precompute(): + bank = flrr.PrecomputeLikelihoodTermsRotatingFreqResponse( + event_time, 0.1, P.manual_copy(), data_dict, psd_dict, 2, + args.fmax, Qmax=qmax, L_arm=args.arm_length, + p_max=args.pmax, analyticPSD_Q=True, verbose=False, quiet=True, + skip_interpolation=True) + packed = flrr.pack_rotating_freqresponse_arrays( + bank[4], bank[3], bank[1], bank[2]) + return bank, packed + (bank, packed), pre_s = timed(precompute) + meta = bank[4] + lk, rba, uba, vba, ep = packed + det_geom = {d: sfr.detector_geometry(d, L_arm=args.arm_length) + for d in detectors} + jdata = build_rotating_freqresponse_data( + meta, lk, rba, uba, vba, ep, deltaT, tvals, det_geom) + A = len(meta["a_list"]) + bank_bytes = sum(int(jdata.detectors[d][name].nbytes) + for d in detectors for name in ("Q_bank", "U_bank", "V_bank")) + + _, numpy_s = timed(lambda: flrr.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + tvals, Pv, meta, lk, rba, uba, vba, ep, Lmax=2, + time_interp="nearest", xpy=np)) + fn = jax.jit(lambda ra, dec, psi, incl, phiref, dist: fused_log_likelihood( + jdata, ra, dec, psi, incl, phiref, dist, interp="nearest")) + values = (Pv.phi, Pv.theta, Pv.psi, Pv.incl, Pv.phiref, dist_mpc) + _, compile_s = timed(lambda: fn(*values)) + warm = [] + for _ in range(args.repeats): + _, dt = timed(lambda: fn(*values)) + warm.append(dt) + warm_s = min(warm) + print("%d %d %d %.2f %.3f %.3f %.3f %.6f %.1f" % + (qmax, A, A * A, bank_bytes / 2.0**20, pre_s, numpy_s, + compile_s, warm_s, S / warm_s)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_all_axis_peaklocal.py b/MonteCarloMarginalizeCode/Code/test/jax/test_all_axis_peaklocal.py new file mode 100644 index 000000000..5ce7b2950 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_all_axis_peaklocal.py @@ -0,0 +1,1480 @@ +"""Tests for fixed-shape all-variable multi-peak marginalization.""" + +import types + +import numpy as np +import pytest +from scipy import integrate, special + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile import all_axis_peaklocal as AAP +from RIFT.likelihood.jax_ile import anglemarg as AM +from RIFT.likelihood.jax_ile import core as JCORE + + +def _problem(n=129): + span = n - 1.0 + t = np.arange(n, dtype=float) + k0, kt, kp, ku, B = 5.0, 15.0, 8.0, 6.0, 10.0 + C_A = np.zeros((3, 3, n), dtype=np.complex128) + C_A[0, 1] = k0 - kt * np.cos(2.0 * np.pi * t / span) + C_A[2, 1] = 0.5 * kp + C_A[0, 0] = 0.5 * ku + C_A[0, 2] = 0.5 * ku + C_B = np.zeros((5, 5), dtype=np.complex128) + C_B[0, 2] = B + constants = dict(span=span, k0=k0, kt=kt, kp=kp, ku=ku, B=B) + return C_A, C_B, constants + + +def _joint_peak(constants): + A = sum(constants[k] for k in ("k0", "kt", "kp", "ku")) + B = constants["B"] + x = (A + np.sqrt(A * A - 16.0 * B)) / (2.0 * B) + centers = np.asarray([ + [constants["span"] / 2.0, 0.0, 0.0, x], + [constants["span"] / 2.0, np.pi, 0.0, x], + ]) + hessian = np.zeros((2, 4, 4)) + diagonal = np.asarray([ + -x * constants["kt"] * (2.0 * np.pi / constants["span"]) ** 2, + -4.0 * x * constants["kp"], + -x * constants["ku"], + -B + 4.0 / (x * x), + ]) + hessian[:, np.arange(4), np.arange(4)] = diagonal + return centers, hessian + + +def _analytic_log_integral(constants, x_min, x_max): + c = constants + + def log_i0(z): + return np.log(special.i0e(z)) + abs(z) + + def log_integrand(x): + return (-4.0 * np.log(x) - 0.5 * c["B"] * x * x + c["k0"] * x + + log_i0(c["kt"] * x) + log_i0(c["kp"] * x) + + log_i0(c["ku"] * x)) + + probe = np.linspace(x_min, x_max, 2001) + shift = max(log_integrand(x) for x in probe) + value = integrate.quad( + lambda x: np.exp(log_integrand(x) - shift), x_min, x_max, + epsabs=1.0e-13, epsrel=1.0e-13, limit=500)[0] + return (shift + np.log(value) + np.log(c["span"]) + + 2.0 * np.log(2.0 * np.pi)) + + +def test_angle_tables_forward_primitive_guard_and_report_support(monkeypatch): + data = types.SimpleNamespace(lms=np.asarray([[2, 2]]), npts=5) + seen = [] + + def fake_accumulate(data, ra, dec, psi, incl, phi, interp, + phase_marginalization, guard=0): + seen.append(int(guard)) + shape = (ra.shape[0], data.npts + 2 * int(guard)) + return jnp.ones(shape, dtype=jnp.complex128), jnp.ones(shape) + + monkeypatch.setattr(AM, "_accumulate_unit", fake_accumulate) + C_A, C_B, meta = AM.angle_coefficient_tables( + data, jnp.asarray([0.1]), jnp.asarray([0.2]), jnp.asarray([0.3]), + guard=3) + assert C_A.shape == (3, 3, 1, 11) + assert C_B.shape == (5, 5, 1, 11) + assert meta["guard"] == 3 and meta["ntime"] == 11 + assert seen and set(seen) == {3} + + +def test_uv_summary_bounds_the_exact_norm_and_collapses_time(): + rng = np.random.default_rng(813) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_B[0, 2] = 20.0 + perturbation = (rng.normal(size=(5, 5)) + + 1j * rng.normal(size=(5, 5))) * 0.02 + perturbation[0, 2] = 0.0 + C_B += perturbation + repeated = np.repeat(C_B[..., None], 17, axis=-1) + summary = AAP.summarize_uv_norm_table(repeated) + + phi = rng.uniform(0.0, 2.0 * np.pi, 1000) + u = rng.uniform(0.0, 2.0 * np.pi, 1000) + values = np.asarray(jax.vmap( + lambda p, q: AAP._angular_field(jnp.asarray(C_B), p, q))( + jnp.asarray(phi), jnp.asarray(u))) + assert summary.time_invariant + assert values.min() >= summary.b_lower - 1.0e-12 + assert values.max() <= summary.b_upper + 1.0e-12 + assert summary.summary_build_count == 1 + assert summary.input_harmonic_coefficients == repeated.size + + +def test_uv_envelope_ranks_the_interior_time_mode_without_dense_starts(): + C_A, C_B, constants = _problem() + summary = AAP.summarize_uv_norm_table(C_B) + starts, envelope = AAP.rank_time_starts_from_uv( + C_A, summary, 0.2, 7.0, max_starts=3, min_separation=4) + assert starts[0] == int(constants["span"] // 2) + assert len(starts) <= 3 + assert envelope[starts[0]] == pytest.approx(envelope.max()) + + +def test_loud_off_lattice_refiner_enters_narrow_coupled_basin(): + """A fixed structural lattice remains useful as the peak narrows.""" + C_A, C_B, _ = _problem(33) + scale = 64.0 + starts = np.asarray([[15.3, 0.2, 0.2, 5.0 / scale], + [16.7, 3.0, 0.1, 5.0 / scale]]) + result = tuple(np.asarray(item) for item in AAP.refine_all_axis_starts( + scale * C_A, scale * scale * C_B, starts, + 0.2 / scale, 7.0 / scale, iterations=18)) + points, values, gradients, _, curvatures = result + selected, stationary = AAP.select_refined_modes( + points, values, gradients, curvatures, max_modes=2, + gradient_tol=2.0e-6) + assert stationary.any() + assert len(selected) >= 1 + assert np.max(np.linalg.norm(gradients[selected], axis=1)) < 2.0e-6 + assert np.all(curvatures[selected] > 0.0) + + +def test_refiner_mask_skips_padding_without_changing_live_modes(): + C_A, C_B, constants = _problem(33) + centers, _ = _joint_peak(constants) + starts = np.vstack(( + centers + np.asarray([[0.4, 0.1, 0.1, -0.05], + [-0.4, -0.1, 0.1, 0.05]]), + np.repeat([[0.0, 0.0, 0.0, 0.2]], 6, axis=0))) + live = np.asarray([True, True] + [False] * 6) + masked = jax.jit(lambda seed, mask: AAP.refine_all_axis_starts( + C_A, C_B, seed, 0.2, 7.0, iterations=14, live=mask))( + jnp.asarray(starts), jnp.asarray(live)) + direct = AAP.refine_all_axis_starts( + C_A, C_B, starts[:2], 0.2, 7.0, iterations=14) + for masked_value, direct_value in zip(masked, direct): + np.testing.assert_allclose( + np.asarray(masked_value)[:2], np.asarray(direct_value), + rtol=1.0e-12, atol=1.0e-12) + assert np.all(np.isneginf(np.asarray(masked[1])[2:])) + assert np.all(np.asarray(masked[2])[2:] == 0.0) + assert np.all(np.asarray(masked[4])[2:] == 1.0) + + +def test_uv_ranked_time_start_feeds_algebraic_angles_and_analytic_distance(): + C_A, C_B, constants = _problem(65) + summary = AAP.summarize_uv_norm_table(C_B) + time_starts, _ = AAP.rank_time_starts_from_uv( + C_A, summary, 0.2, 7.0, max_starts=1, min_separation=4) + # This deliberately sparse symmetric polynomial includes zero/infinite + # generalized roots; the authoritative report, not NumPy's intermediate + # negative-power warning, carries their classification. + with np.errstate(invalid="ignore", divide="ignore"): + starts, algebraic_ok, reports = AAP.algebraic_angle_starts_from_uv( + C_A, summary, time_starts, 0.2, 7.0) + + # The exact enumerator supplies the two maxima without a dense seed lattice. + # Its independent completeness result is carried separately; definite + # maxima remain useful targeting data if a numerically marginal projection + # makes that result false on another LAPACK implementation. + assert starts.shape == (2, 4) + assert algebraic_ok == bool(reports[0]["ok"]) + assert len(reports) == 1 and reports[0]["n_maxima"] == 2 + assert np.allclose(starts[:, 0], constants["span"] / 2.0) + assert np.all((starts[:, 3] > 0.2) & (starts[:, 3] < 7.0)) + + +def test_joint_start_lattice_is_harmonic_order_sized_not_snr_sized(): + C_A, C_B, constants = _problem(65) + summary = AAP.summarize_uv_norm_table(C_B) + low = AAP.rank_joint_starts_from_uvq( + C_A, summary, 0.2, 7.0, max_time_starts=1, max_starts=8) + high = AAP.rank_joint_starts_from_uvq( + 32.0 * C_A, summary, 0.2, 7.0, + max_time_starts=1, max_starts=8) + + assert low.starts.shape[1] == 4 + assert low.time_starts[0] == int(constants["span"] // 2) + assert low.n_phi_lattice == high.n_phi_lattice == 17 + assert low.n_u_lattice == high.n_u_lattice == 9 + assert low.n_lattice_evaluations == high.n_lattice_evaluations == 17 * 9 * 65 + assert low.n_exact_symmetry_shifts == high.n_exact_symmetry_shifts == 2 + assert low.capacity_ok and high.capacity_ok + assert np.all((low.starts[:, 3] >= 0.2) & (low.starts[:, 3] <= 7.0)) + + +def test_joint_start_guard_discards_support_before_ranking(): + C_A, C_B, constants = _problem(65) + guard = 8 + guarded = np.full(C_A.shape[:-1] + (65 + 2 * guard,), + 1.0e8 + 2.0e8j, dtype=np.complex128) + guarded[..., guard:-guard] = C_A + summary = AAP.summarize_uv_norm_table(C_B) + plan = AAP.rank_joint_starts_from_uvq( + guarded, summary, 0.2, 7.0, time_guard=guard, + max_time_starts=1, max_starts=8) + + assert plan.time_starts.tolist() == [int(constants["span"] // 2)] + assert plan.n_lattice_evaluations == 17 * 9 * 65 + + +def test_joint_start_capacity_declines_instead_of_silent_truncation(): + C_A, C_B, _ = _problem(65) + summary = AAP.summarize_uv_norm_table(C_B) + plan = AAP.rank_joint_starts_from_uvq( + C_A, summary, 0.2, 7.0, max_time_starts=1, max_starts=1) + assert plan.starts.shape == (1, 4) + assert plan.n_candidates_before_cap > 1 + assert not plan.capacity_ok + + +def test_device_joint_start_portfolio_is_fixed_shape_jittable_and_vmappable(): + C_A, C_B, _ = _problem(33) + + def rank(table): + return AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=32, + angular_oversample=2) + + low = jax.jit(rank)(jnp.asarray(C_A)) + high = jax.jit(rank)(jnp.asarray(64.0 * C_A)) + assert low.starts.shape == high.starts.shape == (32, 4) + assert int(low.n_phi_lattice) == int(high.n_phi_lattice) == 17 + assert int(low.n_u_lattice) == int(high.n_u_lattice) == 9 + assert int(low.n_lattice_evaluations) == 17 * 9 * 33 + assert int(high.n_lattice_evaluations) == 17 * 9 * 33 + assert int(low.n_time_scout_evaluations) == 4 * 4 * 33 + assert int(low.n_retained_time_samples) == 33 + assert bool(low.time_cover_certified) and bool(high.time_cover_certified) + assert bool(low.time_capacity_ok) and bool(high.time_capacity_ok) + assert bool(low.norm_nonnegative) and bool(high.norm_nonnegative) + assert bool(low.capacity_ok) and bool(high.capacity_ok) + assert np.all(np.asarray(low.starts)[np.asarray(low.live), 3] >= 0.2) + assert np.all(np.asarray(high.starts)[np.asarray(high.live), 3] <= 7.0) + + batch = jax.jit(jax.vmap(rank))(jnp.asarray( + np.stack((C_A, 1.01 * C_A)))) + assert batch.starts.shape == (2, 32, 4) + np.testing.assert_array_equal( + np.asarray(batch.n_lattice_evaluations), [17 * 9 * 33] * 2) + + +def test_device_time_cover_bounds_discarded_cells_and_limits_full_lattice(): + n_time = 129 + t = np.arange(n_time, dtype=float) + C_A = np.zeros((1, 1, n_time), dtype=np.complex128) + C_A[0, 0] = 5.0 - 15.0 * np.cos(2.0 * np.pi * t / (n_time - 1.0)) + C_B = np.asarray([[10.0 + 0.0j]]) + cover = jax.jit(lambda table: AAP._time_cell_cover_device( + table, C_B, 0.5, 2.0, time_guard=0, keep_nats=5.0, + scout_size=4))(jnp.asarray(C_A)) + live = np.asarray(cover["live_cells"]) + cell_upper = np.asarray(cover["cell_mass_upper"]) + assert bool(cover["certified"]) + assert 0 < np.count_nonzero(live) < live.size + assert float(cover["outside_log_bound"]) == pytest.approx( + special.logsumexp(cell_upper[~live])) + + coeff, frequency, offset = AAP._time_primitive_spectrum( + jnp.asarray(C_A.reshape((1, n_time))), 0) + x = np.linspace(0.5, 2.0, 129) + log_volume = np.log((2.0 * np.pi) ** 2 * (2.0 - 0.5)) + for cell in np.flatnonzero(~live): + position = np.linspace(cell, cell + 1.0, 33) + amplitude = np.asarray(AAP._evaluate_time_spectrum( + coeff, frequency, jnp.asarray(position), offset))[0].real + sampled = (amplitude[:, None] * x[None, :] + - 5.0 * x[None, :] ** 2 - 4.0 * np.log(x[None, :])) + assert sampled.max() + log_volume <= cell_upper[cell] + 1.0e-10 + + plan = jax.jit(lambda table: AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.5, 2.0, max_starts=32, max_time_nodes=64, + time_keep_nats=5.0))(jnp.asarray(C_A)) + assert int(plan.n_retained_time_samples) == n_time + assert int(plan.n_time_lattice) == 64 + assert int(plan.n_lattice_evaluations) == 9 * 9 * 64 + assert int(plan.n_time_scout_evaluations) == 4 * 4 * n_time + assert int(plan.n_time_nodes_retained) <= 64 + assert bool(plan.time_capacity_ok) + assert float(plan.time_cover_min_sample) == float( + np.flatnonzero(live)[0]) + assert float(plan.time_cover_max_sample) == float( + np.flatnonzero(live)[-1] + 1) + # The angle-constant fixture is deliberately degenerate: every angular + # lattice point is a maximum, so the independent start-capacity gate still + # declines even though the time cover itself fits and is certified. + assert int(plan.n_candidates_before_cap) > 32 + assert not bool(plan.capacity_ok) + + +def test_device_joint_start_portfolio_fails_closed_on_capacity_and_norm(): + C_A, C_B, _ = _problem(33) + truncated = jax.jit(lambda table: AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=1))(jnp.asarray(C_A)) + assert int(truncated.n_candidates_before_cap) > 1 + assert not bool(truncated.capacity_ok) + + invalid_norm = C_B.copy() + invalid_norm[0, 2] = -10.0 + rejected = jax.jit(lambda table: AAP.rank_joint_starts_from_uvq_device( + C_A, table, 0.2, 7.0, max_starts=8))(jnp.asarray(invalid_norm)) + assert not bool(rejected.norm_nonnegative) + assert not bool(rejected.capacity_ok) + assert not np.any(np.asarray(rejected.live)) + + time_overflow = jax.jit( + lambda table: AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=8, max_time_nodes=2))( + jnp.asarray(C_A)) + assert not bool(time_overflow.time_capacity_ok) + assert not bool(time_overflow.capacity_ok) + + truncated = truncated._replace( + time_cover_min_sample=jnp.asarray(2.0), + time_cover_max_sample=jnp.asarray(5.0), + time_outside_log_bound=jnp.asarray(-11.0)) + rejected = rejected._replace( + time_cover_min_sample=jnp.asarray(3.0), + time_cover_max_sample=jnp.asarray(6.0), + time_outside_log_bound=jnp.asarray(-13.0)) + combined = AAP.combine_device_start_plans(truncated, rejected) + assert combined.starts.shape == (9, 4) + assert not bool(combined.capacity_ok) + assert not bool(combined.norm_nonnegative) + assert float(combined.time_cover_min_sample) == 2.0 + assert float(combined.time_cover_max_sample) == 6.0 + assert float(combined.time_outside_log_bound) == -13.0 + + +def test_device_mode_plan_refines_and_deduplicates_without_host_transfer(): + C_A, C_B, _ = _problem(33) + + @jax.jit + def build(table): + starts = AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=32) + return AAP.make_all_axis_mode_plan_device( + table, C_B, starts, 0.2, 7.0, max_modes=4, + local_radius=3.0, iterations=14, + time_reconstruction_certified=True) + + plan, ledger = build(jnp.asarray(C_A)) + assert plan.centers.shape == (4, 4) + assert plan.local_transforms.shape == (4, 4, 4) + assert int(ledger["n_optimizer_starts"]) >= 2 + assert int(ledger["n_selected_modes"]) == 2 + assert int(jnp.count_nonzero(plan.live)) == 2 + assert bool(ledger["discovery_capacity_ok"]) + assert not bool(ledger["selection_overflow"]) + assert not bool(ledger["global_completeness_certified"]) + assert not bool(ledger["derivative_warrant_certified"]) + assert bool(plan.time_reconstruction_certified) + live_transforms = np.asarray(plan.local_transforms)[np.asarray(plan.live)] + assert np.all(np.diagonal(live_transforms, axis1=1, axis2=2) > 0.0) + + plans, ledgers = jax.jit(jax.vmap(build))(jnp.asarray( + np.stack((C_A, 1.01 * C_A)))) + assert plans.centers.shape == (2, 4, 4) + assert ledgers["n_selected_modes"].shape == (2,) + assert np.all(np.asarray(ledgers["discovery_capacity_ok"])) + + @jax.jit + def overflow(table): + starts = AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=32) + return AAP.make_all_axis_mode_plan_device( + table, C_B, starts, 0.2, 7.0, max_modes=1, + local_radius=3.0, iterations=14, + time_reconstruction_certified=True) + + overflow_plan, overflow_ledger = overflow(jnp.asarray(C_A)) + assert int(jnp.count_nonzero(overflow_plan.live)) == 1 + assert bool(overflow_ledger["selection_overflow"]) + assert not bool(overflow_ledger["discovery_capacity_ok"]) + assert not bool(overflow_plan.discovery_capacity_ok) + + +def test_device_plans_compose_with_empirical_local_controller_under_vmap(): + C_A, C_B, _ = _problem(33) + + def evaluate(table): + base_starts = AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=32, + angular_oversample=1) + extra_starts = AAP.rank_joint_starts_from_uvq_device( + table, C_B, 0.2, 7.0, max_starts=32, + angular_oversample=2) + separate_base_plan, separate_base_planning = ( + AAP.make_all_axis_mode_plan_device( + table, C_B, base_starts, 0.2, 7.0, max_modes=4, + local_radius=3.0, iterations=14, + time_reconstruction_certified=True)) + (base_plan, enriched_plan, base_planning, enriched_planning, + shared_planning) = AAP.make_all_axis_mode_plan_pair_device( + table, C_B, base_starts, extra_starts, 0.2, 7.0, + max_modes=4, enriched_max_modes=8, + local_radius=3.0, iterations=14, + time_reconstruction_certified=True) + # Plans are row-local control data. Until derivative parity is + # established, outer differentiation must not interpret the discrete + # rank/dedup decisions as a full-marginal derivative certificate. + base_plan = jax.tree.map(jax.lax.stop_gradient, base_plan) + enriched_plan = jax.tree.map(jax.lax.stop_gradient, enriched_plan) + value, accepted, ledger = AAP.empirical_enrichment_marginalize( + table, C_B, base_plan, enriched_plan, 0.2, 7.0, + base_order=7, base_check_order=9, + enriched_order=9, enriched_check_order=11, + convergence_tol_nats=1.0e-3) + return (value, accepted, ledger, base_planning, enriched_planning, + shared_planning, separate_base_plan, + separate_base_planning, base_plan) + + (value, accepted, ledger, base_planning, enriched_planning, + shared_planning, separate_base_plan, separate_base_planning, + paired_base_plan) = jax.jit(evaluate)(jnp.asarray(C_A)) + assert np.isfinite(float(value)) + assert bool(accepted) + assert not bool(ledger["fallback_required"]) + assert bool(ledger["mode_nesting_ok"]) + assert int(base_planning["n_selected_modes"]) == 2 + assert int(enriched_planning["n_selected_modes"]) == 2 + assert int(enriched_planning["n_lattice_evaluations"]) == ( + 9 * 9 * 33 + 17 * 9 * 33) + n_base = int(base_planning["n_optimizer_starts"]) + n_enriched = int(enriched_planning["n_optimizer_starts"]) + assert int(shared_planning["n_optimizer_starts_executed"]) == n_enriched + assert int(shared_planning["n_optimizer_starts_avoided"]) == n_base + assert int(shared_planning["n_optimizer_starts_previous_two_pass"]) == ( + n_base + n_enriched) + assert bool(shared_planning["start_nesting_structural"]) + assert int(separate_base_planning["n_selected_modes"]) == 2 + for separate, paired in zip(jax.tree.leaves(separate_base_plan), + jax.tree.leaves(paired_base_plan)): + np.testing.assert_allclose(np.asarray(separate), np.asarray(paired), + rtol=0.0, atol=1.0e-12) + + batch = jax.jit(jax.vmap(evaluate))(jnp.asarray( + np.stack((C_A, 1.01 * C_A)))) + assert batch[0].shape == batch[1].shape == (2,) + assert np.all(np.isfinite(np.asarray(batch[0]))) + assert np.all(np.asarray(batch[1])) + + +def test_exact_coefficient_symmetry_completes_quadrupole_orbit(): + C_A = np.zeros((3, 3, 9), dtype=np.complex128) + C_A[2, 0] = 1.0 + C_A[2, 2] = 0.7 + C_B = np.zeros((5, 5), dtype=np.complex128) + C_B[0, 2] = 2.0 + shifts = AAP._exact_angular_translation_symmetries(C_A, C_B) + want = np.asarray([ + [0.0, 0.0], [np.pi, 0.0], + [0.5 * np.pi, np.pi], [1.5 * np.pi, np.pi]]) + assert shifts.shape == (4, 2) + for shift in want: + assert np.min(np.linalg.norm(shifts - shift, axis=1)) < 1.0e-12 + + device_shifts, device_live = jax.jit( + AAP._exact_angular_translation_symmetries_device)(C_A, C_B) + device_shifts = np.asarray(device_shifts)[np.asarray(device_live)] + assert device_shifts.shape == (4, 2) + for shift in want: + assert np.min(np.linalg.norm(device_shifts - shift, axis=1)) < 1.0e-12 + + +def test_mode_stationarity_uses_curvature_scaled_displacement_at_high_snr(): + point = np.asarray([[10.0, 1.0, 2.0, 1.5]]) + value = np.asarray([10000.0]) + gradient = np.asarray([[2.0e-4, 0.0, 0.0, 0.0]]) + curvature = np.asarray([[35.0, 200.0, 1000.0, 100000.0]]) + selected, stationary = AAP.select_refined_modes( + point, value, gradient, curvature, max_modes=1, + gradient_tol=2.0e-6) + assert stationary.tolist() == [True] + assert selected.tolist() == [0] + + curvature[0, 0] = 1.0 + selected, stationary = AAP.select_refined_modes( + point, value, gradient, curvature, max_modes=1, + gradient_tol=2.0e-6) + assert stationary.tolist() == [False] + assert selected.size == 0 + + +def test_jax_gradient_hessian_refinement_finds_both_angular_modes(): + C_A, C_B, constants = _problem(65) + centers, _ = _joint_peak(constants) + starts = centers + np.asarray([ + [1.3, 0.17, -0.11, -0.13], + [-1.1, -0.15, 0.09, 0.16], + ]) + refined, value, gradient, hessian, curvature = AAP.refine_all_axis_starts( + C_A, C_B, starts, 0.2, 7.0, iterations=14) + refined = np.asarray(refined) + angular_error = np.abs( + (refined[:, 1:3] - centers[:, 1:3] + np.pi) % (2.0 * np.pi) - np.pi) + assert np.max(np.abs(refined[:, [0, 3]] - centers[:, [0, 3]])) < 2.0e-7 + assert np.max(angular_error) < 2.0e-7 + assert np.max(np.linalg.norm(np.asarray(gradient), axis=1)) < 2.0e-7 + assert np.all(np.asarray(curvature) > 0.0) + assert np.all(np.isfinite(np.asarray(value))) + assert np.all(np.isfinite(np.asarray(hessian))) + selected, stationary = AAP.select_refined_modes( + refined, value, gradient, curvature, max_modes=4) + assert np.all(stationary) + assert selected.shape == (2,) + with pytest.raises(ValueError, match="exceeds fixed plan capacity"): + AAP.select_refined_modes( + refined, value, gradient, curvature, max_modes=1) + with pytest.raises(ValueError, match="must be positive"): + AAP.select_refined_modes( + refined, value, gradient, curvature, max_modes=0) + + +def test_multimode_local_primitive_matches_oracle_but_stays_uncertified(): + C_A, C_B, constants = _problem() + centers, hessian = _joint_peak(constants) + transforms, half_widths = AAP.mode_local_geometry(hessian, w_sigma=5.0) + x_min, x_max = 0.2, 7.0 + truth = _analytic_log_integral(constants, x_min, x_max) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=4, + local_transforms=transforms, local_radius=5.0, + outside_log_bound=np.inf, + enumeration_complete=True, outside_bound_certified=False) + value, ok, ledger = AAP.all_axis_peak_local_marginalize( + C_A, C_B, plan, x_min, x_max, local_order=13, check_order=19, + quadrature_tol_nats=1.0e-4) + + assert not bool(ok) + assert bool(ledger["decline_incomplete"]) + assert abs(float(value) - truth) < 2.0e-4 + assert int(ledger["n_modes"]) == 2 + assert int(ledger["n_local_evaluations_hi"]) == 2 * 19 ** 4 + assert int(ledger["n_selected_time_points_hi"]) == 2 * 19 + assert int(ledger["n_time_frequency_terms_hi"]) == ( + 2 * 19 * (2 * C_A.shape[-1] - 2) * np.prod(C_A.shape[:-1])) + assert int(ledger["n_angle_harmonic_terms_hi"]) == ( + 2 * 19 ** 3 * (np.prod(C_A.shape[:-1]) + C_B.size)) + assert int(ledger["workspace_bytes_hi"]) < 8_000_000 + assert bool(ledger["reconciles"]) + + +def test_empirical_enrichment_accepts_without_claiming_global_proof(): + C_A, C_B, constants = _problem(33) + C_A *= 0.1 + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + value, accepted, ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + + assert np.isfinite(float(value)) + assert bool(accepted) + assert bool(ledger["acceptance_is_empirical_enrichment"]) + assert not bool(ledger["global_completeness_certified"]) + assert not bool(ledger["empirical_value_error_certified"]) + assert float(ledger["convergence_error"]) <= 1.0e-3 + assert bool(ledger["value_error_budget_complete"]) + assert bool(ledger["value_error_budget_ok"]) + assert bool(ledger["value_error_budget_is_empirical"]) + assert not bool(ledger["value_error_budget_is_formal_bound"]) + assert (float(ledger["empirical_value_error_score_nats"]) + <= float(ledger["total_value_error_budget_nats"])) + # Paired terms are charged ONCE, as the max over the two plans: they are + # the same quantity measured on nested plans, and summing them double + # counted a common-mode error (refused a correct value at 10x amplitude). + score_components = [ + "error_score_discovery_nats", + "error_score_quadrature_nats", + "error_score_time_guard_nats", + "error_score_omitted_time_nats", + ] + components = np.asarray([float(ledger[key]) for key in score_components]) + score = float(ledger["empirical_value_error_score_nats"]) + assert score == pytest.approx(float(np.sum(components)), abs=1.0e-15) + assert bool(ledger["error_score_pairs_charged_as_max"]) + for pair, charged in ( + (("error_score_base_quadrature_nats", + "error_score_enriched_quadrature_nats"), + "error_score_quadrature_nats"), + (("error_score_base_time_guard_nats", + "error_score_enriched_time_guard_nats"), + "error_score_time_guard_nats"), + (("error_score_base_omitted_time_nats", + "error_score_enriched_omitted_time_nats"), + "error_score_omitted_time_nats")): + assert float(ledger[charged]) == pytest.approx( + max(float(ledger[pair[0]]), float(ledger[pair[1]])), abs=1e-15) + assert (float(ledger["error_score_base_time_guard_nats"]) + == float(ledger["error_score_enriched_time_guard_nats"]) == 0.0) + assert bool(ledger["mode_nesting_ok"]) + assert not bool(ledger["fallback_required"]) + assert bool(ledger["reconciles"]) + + loose_time_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=1.0e6, + time_outside_bound_certified=True) + _, accepted, time_tail_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, loose_time_plan, loose_time_plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert bool(time_tail_ledger["time_outside_cover_used"]) + assert not bool(time_tail_ledger["time_omitted_mass_ok"]) + assert bool(time_tail_ledger["decline_time_omitted_mass"]) + assert bool(time_tail_ledger["decline_time_omitted_mass_bound"]) + assert not bool(time_tail_ledger["decline_time_cover_incomplete"]) + assert bool(time_tail_ledger["reconciles"]) + + missing_time_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True) + _, accepted, missing_time_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, missing_time_plan, missing_time_plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert not bool(missing_time_ledger["time_outside_cover_used"]) + assert not bool(missing_time_ledger["value_error_budget_complete"]) + assert bool(missing_time_ledger["decline_time_omitted_mass"]) + assert bool(missing_time_ledger["decline_time_cover_incomplete"]) + assert not bool(missing_time_ledger["decline_time_omitted_mass_bound"]) + assert not bool(missing_time_ledger["decline_error_budget"]) + assert bool(missing_time_ledger["reconciles"]) + + _, accepted, one_time_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, missing_time_plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert bool(one_time_ledger["time_outside_cover_any"]) + assert not bool(one_time_ledger["time_outside_cover_used"]) + assert bool(one_time_ledger["decline_time_cover_incomplete"]) + assert bool(one_time_ledger["reconciles"]) + + # Each legacy diagnostic can clear its individual gate while their + # cancellation-resistant sum exceeds one shared allowance. + aggregate_budget = 0.5 * (float(np.max(components)) + score) + assert float(np.max(components)) < aggregate_budget < score + _, accepted, aggregate_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-3, + total_value_error_budget_nats=aggregate_budget) + assert not bool(accepted) + assert bool(aggregate_ledger["base_quadrature_error"] <= 1.0e-3) + assert bool(aggregate_ledger["enriched_quadrature_error"] <= 1.0e-3) + assert bool(aggregate_ledger["convergence_error"] <= 1.0e-3) + assert not bool(aggregate_ledger["value_error_budget_ok"]) + assert bool(aggregate_ledger["decline_error_budget"]) + assert bool(aggregate_ledger["reconciles"]) + + _, accepted_at_budget, boundary_ledger = ( + AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-3, + total_value_error_budget_nats=score)) + assert bool(accepted_at_budget) + assert bool(boundary_ledger["value_error_budget_ok"]) + _, accepted_below_budget, below_ledger = ( + AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-3, + total_value_error_budget_nats=np.nextafter(score, -np.inf))) + assert not bool(accepted_below_budget) + assert bool(below_ledger["decline_error_budget"]) + assert bool(below_ledger["reconciles"]) + + _, accepted, ordered_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-12, + total_value_error_budget_nats=1.0e-15) + assert not bool(accepted) + assert bool(ordered_ledger["decline_quadrature"]) + assert not bool(ordered_ledger["decline_error_budget"]) + assert bool(ordered_ledger["reconciles"]) + + for invalid_budget in (0.0, -1.0, np.nan, np.inf): + with pytest.raises(ValueError, match="total_value_error_budget_nats"): + AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, plan, x_min, x_max, + total_value_error_budget_nats=invalid_budget) + + # A stronger discovery pass may expose a broad diagnostic basin whose + # positive integral is negligible. It must not invalidate the unchanged, + # valid base cover when the enrichment delta remains inside the same budget. + extra_centers = np.vstack((centers, [[ + centers[0, 0], 0.1, 0.2, centers[0, 3]]])) + extra_transforms = np.concatenate(( + transforms, + np.asarray([np.diag([1.0e-8, 4.0, 4.0, 1.0e-8])])), axis=0) + diagnostic_plan = AAP.make_all_axis_mode_plan( + extra_centers, max_modes=3, local_transforms=extra_transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + retained, accepted, diagnostic_ledger = ( + AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, diagnostic_plan, x_min, x_max, + convergence_tol_nats=1.0e-3)) + assert bool(accepted) + assert bool(diagnostic_ledger["base_geometry_ok"]) + assert not bool(diagnostic_ledger["enriched_geometry_ok"]) + assert bool(diagnostic_ledger["geometry_nesting_ok"]) + assert bool(diagnostic_ledger["accepted_value_uses_base_geometry"]) + assert float(retained) == pytest.approx(float(ledger["base_value"]), abs=1e-12) + + shifted = centers.copy() + shifted[:, 1] += 0.1 + shifted_plan = AAP.make_all_axis_mode_plan( + shifted, max_modes=2, local_transforms=transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + _, accepted, nesting_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, plan, shifted_plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert bool(nesting_ledger["decline_mode_nesting"]) + assert bool(nesting_ledger["reconciles"]) + + truncated_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True, + discovery_capacity_ok=False) + _, accepted, capacity_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, truncated_plan, plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert bool(capacity_ledger["decline_capacity"]) + assert bool(capacity_ledger["fallback_required"]) + assert not bool(capacity_ledger["decline_is_waveform_failure"]) + assert bool(capacity_ledger["reconciles"]) + + empty_plan = AAP.make_all_axis_mode_plan( + np.empty((0, 4)), max_modes=2, + local_transforms=np.empty((0, 4, 4)), local_radius=1.0, + outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + _, accepted, empty_ledger = AAP.empirical_enrichment_marginalize( + C_A, C_B, empty_plan, empty_plan, x_min, x_max, + convergence_tol_nats=1.0e-3) + assert not bool(accepted) + assert not bool(empty_ledger["base_and_enriched_values_finite"]) + assert bool(empty_ledger["decline_no_modes"]) + assert not bool(empty_ledger["decline_nonfinite"]) + assert bool(empty_ledger["fallback_required"]) + assert bool(empty_ledger["reconciles"]) + + +def test_native_time_reserve_is_diagnostic_on_local_decline(): + C_A, C_B, constants = _problem(33) + C_A *= 0.1 + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + declined_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=True, + discovery_capacity_ok=False) + + x_grid = np.linspace(x_min, x_max, 129) + dx = np.empty_like(x_grid) + dx[1:-1] = 0.5 * (x_grid[2:] - x_grid[:-2]) + dx[0] = x_grid[1] - x_grid[0] + dx[-1] = x_grid[-1] - x_grid[-2] + log_w = np.log(dx * x_grid ** -4) + time_weights = np.ones(C_A.shape[-1]) + selected, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + C_A, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16) + + lnL_t = AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, amp_sizing=30.0, + dense_chunk=8, grid_block=16) + m = np.max(np.asarray(lnL_t)[0]) + expected = m + np.log(np.sum( + time_weights * np.exp(np.asarray(lnL_t)[0] - m))) + assert float(selected) == pytest.approx(expected, abs=2.0e-12) + assert not bool(usable) + assert not bool(ledger["accepted_local"]) + assert bool(ledger["decline_capacity"]) + assert bool(ledger["reserve_executed"]) + assert bool(ledger["reserve_finite"]) + assert not bool(ledger["selected_value_is_warranted_reserve"]) + assert not bool(ledger["sample_retained_after_local_decline"]) + assert not bool(ledger["decline_is_waveform_failure"]) + assert bool(ledger["local_fallback_required"]) + assert bool(ledger["fallback_required"]) + assert not bool(ledger["accepted"]) + assert bool(ledger["reconciles"]) + assert bool(ledger["disposition_reconciles"]) + assert int(ledger["reserve_distance_points"]) == x_grid.size + assert bool(ledger["reserve_uses_native_time"]) + assert not bool(ledger["reserve_native_time_warranted"]) + assert not bool(ledger["reserve_time_warranted"]) + + +def test_declined_controller_can_execute_guarded_bandlimited_time_reserve(): + C_A, C_B, constants = _problem(33) + C_A *= 0.1 + guard = 8 + support_time = np.arange(-guard, C_A.shape[-1] + guard, dtype=float) + guarded = np.zeros(C_A.shape[:-1] + (support_time.size,), + dtype=np.complex128) + guarded[0, 1] = 0.1 * ( + constants["k0"] - constants["kt"] + * np.cos(2.0 * np.pi * support_time / constants["span"])) + guarded[2, 1] = 0.05 * constants["kp"] + guarded[0, 0] = 0.05 * constants["ku"] + guarded[0, 2] = 0.05 * constants["ku"] + np.testing.assert_allclose(guarded[..., guard:-guard], C_A, atol=1e-14) + + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + declined_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=False, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True, + discovery_capacity_ok=False) + x_grid = np.linspace(x_min, x_max, 65) + dx = np.empty_like(x_grid) + dx[1:-1] = 0.5 * (x_grid[2:] - x_grid[:-2]) + dx[0] = x_grid[1] - x_grid[0] + dx[-1] = x_grid[-1] - x_grid[-2] + log_w = np.log(dx * x_grid ** -4) + time_nodes = np.linspace(0.0, constants["span"], 65) + time_weights = np.full(time_nodes.size, 0.5) + time_weights[[0, -1]] *= 0.5 + + coeff, frequency, offset = AAP._time_primitive_spectrum( + guarded.reshape((-1, guarded.shape[-1])), guard) + fine = AAP._evaluate_time_spectrum( + coeff, frequency, time_nodes, offset).reshape( + C_A.shape[:-1] + (time_nodes.size,)) + lnL_t = AM.coefficient_table_distphipsimarg_exact( + fine, C_B, x_grid, log_w, amp_sizing=30.0, + dense_chunk=8, grid_block=16) + m = np.max(np.asarray(lnL_t)[0]) + expected = m + np.log(np.sum( + time_weights * np.exp(np.asarray(lnL_t)[0] - m))) + selected, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16, + reserve_time_nodes=time_nodes, + reserve_time_resolution_warranted=True, + reserve_time_check_value=expected, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert float(selected) == pytest.approx(expected, abs=2.0e-12) + assert bool(usable) + assert not bool(ledger["accepted_local"]) + assert bool(ledger["reserve_executed"]) + assert bool(ledger["reserve_uses_bandlimited_time"]) + assert not bool(ledger["reserve_uses_native_time"]) + assert bool(ledger["reserve_time_resolution_warranted"]) + assert float(ledger["reserve_time_resolution_error_nats"]) == pytest.approx( + 0.0, abs=2.0e-12) + assert bool(ledger["reserve_time_resolution_validated"]) + assert bool(ledger["reserve_time_nodes_finite"]) + assert bool(ledger["reserve_time_nodes_increasing"]) + assert bool(ledger["reserve_time_weights_valid"]) + assert bool(ledger["reserve_time_nodes_in_support"]) + assert bool(ledger["reserve_time_nodes_cover_target"]) + assert bool(ledger["reserve_time_guard_validated"]) + assert float(ledger["reserve_time_guard_error"]) <= 1.0e-3 + assert bool(ledger["reserve_time_error_budget_ok"]) + assert bool(ledger["reserve_time_warranted"]) + assert not bool(ledger["reserve_time_failed"]) + assert bool(ledger["sample_retained_after_local_decline"]) + assert bool(ledger["reconciles"]) + + _, uncertified_usable, uncertified = ( + AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16, + reserve_time_nodes=time_nodes, + reserve_time_resolution_warranted=False, time_guard=guard, + time_guard_tol_nats=1.0e-3)) + assert not bool(uncertified_usable) + assert bool(uncertified["reserve_time_failed"]) + assert not bool(uncertified["sample_retained_after_local_decline"]) + assert bool(uncertified["fallback_required"]) + assert bool(uncertified["reconciles"]) + + clipped_nodes = np.linspace(1.0, constants["span"] - 1.0, 65) + _, clipped_usable, clipped = ( + AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16, + reserve_time_nodes=clipped_nodes, + reserve_time_resolution_warranted=True, time_guard=guard, + time_guard_tol_nats=1.0e-3)) + assert not bool(clipped_usable) + assert not bool(clipped["reserve_time_nodes_cover_target"]) + assert bool(clipped["reserve_time_failed"]) + assert bool(clipped["fallback_required"]) + assert bool(clipped["reconciles"]) + + accepted_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=False, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + batched_plans = jax.tree.map( + lambda accepted, declined: jnp.stack((accepted, declined)), + accepted_plan, declined_plan) + batched = jax.jit( + lambda tables, warrants, checks: + AAP.empirical_enrichment_with_exact_reserve_sequential_batch( + tables, C_B, batched_plans, batched_plans, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16, + reserve_time_nodes=time_nodes, + reserve_time_resolution_warranted=warrants, + reserve_time_check_value=checks, + time_guard=guard, time_guard_tol_nats=1.0e-3)) + batch_selected, batch_usable, batch_ledger = batched( + jnp.stack((guarded, guarded)), jnp.asarray([False, True]), + jnp.asarray([np.nan, expected])) + assert np.all(np.asarray(batch_usable)) + assert bool(batch_ledger["accepted_local"][0]) + assert not bool(batch_ledger["reserve_executed"][0]) + assert not bool(batch_ledger["accepted_local"][1]) + assert bool(batch_ledger["reserve_executed"][1]) + assert float(batch_selected[1]) == pytest.approx(expected, abs=2.0e-12) + assert np.all(np.asarray( + batch_ledger["reserve_batch_execution_sequential"])) + assert np.all(np.asarray(batch_ledger["reconciles"])) + + +def test_native_time_reserve_cannot_be_warranted(): + C_A, C_B, constants = _problem(17) + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + declined_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=False, + discovery_capacity_ok=False) + x_grid = np.linspace(x_min, x_max, 33) + dx = np.empty_like(x_grid) + dx[1:-1] = 0.5 * (x_grid[2:] - x_grid[:-2]) + dx[0] = x_grid[1] - x_grid[0] + dx[-1] = x_grid[-1] - x_grid[-2] + log_w = np.log(dx * x_grid ** -4) + time_weights = np.ones(C_A.shape[-1]) + + _, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + C_A, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=x_grid, reserve_log_weights=log_w, + time_weights=time_weights, reserve_amp_sizing=30.0, + reserve_dense_chunk=8, reserve_grid_block=16) + + assert not bool(usable) + assert not bool(ledger["reserve_native_time_warranted"]) + assert not bool(ledger["reserve_time_warranted"]) + assert bool(ledger["reserve_time_failed"]) + assert bool(ledger["fallback_required"]) + assert not bool(ledger["sample_retained_after_local_decline"]) + assert not bool(ledger["selected_nonfinite_is_integration_failure"]) + assert bool(ledger["reconciles"]) + + +def test_cropped_bandlimited_reserve_requires_certified_scout_union(monkeypatch): + C_A, C_B, constants = _problem(9) + guard = 2 + guarded = np.pad(C_A, ((0, 0), (0, 0), (guard, guard)), mode="edge") + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + raw_outside_bound = np.log(4.0) - 24.0 + + def plan(time_min, time_max, outside_bound=raw_outside_bound, + certified=True): + return AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=False, + time_outside_log_bound=outside_bound, + time_outside_bound_certified=certified, + time_cover_min_sample=time_min, + time_cover_max_sample=time_max, + discovery_capacity_ok=False) + + # This contract-level fixture has unit marginalized density on the named + # compact support. The external bound is therefore a conservative warrant + # for its deliberately negligible discarded cells; the test exercises how + # that warrant is propagated, not how a production scout derives it. + def fake_exact(table, _norm, _x_grid, _log_w_grid, **_kwargs): + return jnp.zeros((1, table.shape[-1]), dtype=jnp.float64) + + monkeypatch.setattr(AM, "coefficient_table_distphipsimarg_exact", + fake_exact) + base = plan(2.0, 5.0) + enriched = plan(3.0, 6.0, raw_outside_bound + 1.0) + nodes = np.linspace(2.0, 6.0, 9) + weights = JCORE._simpson_weights(nodes.size, 0.5) + offset = 3.0 + expected = np.log(4.0) + offset + + def evaluate(base_plan=base, enriched_plan=enriched, + time_nodes=nodes, check_value=expected): + return AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, base_plan, enriched_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=time_nodes, + reserve_time_resolution_warranted=True, + reserve_time_check_value=check_value, + reserve_log_offset=offset, time_guard=guard, + time_guard_tol_nats=1.0e-3) + + value, usable, ledger = evaluate() + assert bool(usable) + assert float(value) == pytest.approx(expected, abs=2.0e-12) + assert not bool(ledger["reserve_time_nodes_cover_target"]) + assert bool(ledger["reserve_time_plan_cover_finite"]) + assert bool(ledger["reserve_time_nodes_cover_plans"]) + assert bool(ledger["reserve_time_plan_cover_certified"]) + assert bool(ledger["reserve_time_cropped_cover_warranted"]) + assert bool(ledger["reserve_time_interval_warranted"]) + assert float(ledger["reserve_required_time_min_sample"]) == 2.0 + assert float(ledger["reserve_required_time_max_sample"]) == 6.0 + assert float(ledger["reserve_time_outside_log_bound"]) == pytest.approx( + raw_outside_bound + offset) + assert float(ledger["reserve_time_tail_margin"]) == pytest.approx(-24.0) + assert bool(ledger["reserve_time_tail_ok"]) + assert bool(ledger["reserve_time_error_budget_ok"]) + assert bool(ledger["selected_value_is_warranted_reserve"]) + + short_nodes = np.linspace(2.0, 5.5, 8) + short_weights = JCORE._simpson_weights(short_nodes.size, + short_nodes[1] - short_nodes[0]) + _, short_usable, short = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, base, enriched, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=short_weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=short_nodes, + reserve_time_resolution_warranted=True, + reserve_time_check_value=np.log(3.5) + offset, + reserve_log_offset=offset, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert not bool(short_usable) + assert not bool(short["reserve_time_nodes_cover_plans"]) + assert not bool(short["reserve_time_interval_warranted"]) + + _, uncertified_usable, uncertified = evaluate( + enriched_plan=plan(3.0, 6.0, raw_outside_bound + 1.0, False)) + assert not bool(uncertified_usable) + assert not bool(uncertified["reserve_time_plan_cover_certified"]) + assert not bool(uncertified["reserve_time_interval_warranted"]) + + loose = raw_outside_bound + 2.0 + _, loose_usable, loose_ledger = evaluate( + base_plan=plan(2.0, 5.0, loose), + enriched_plan=plan(3.0, 6.0, loose + 1.0)) + assert not bool(loose_usable) + assert not bool(loose_ledger["reserve_time_tail_ok"]) + assert bool(loose_ledger["reserve_time_failed"]) + assert bool(loose_ledger["fallback_required"]) + assert bool(loose_ledger["reconciles"]) + + +def test_bandlimited_reserve_rejects_invalid_time_rules(monkeypatch): + C_A, C_B, constants = _problem(9) + guard = 2 + guarded = np.pad(C_A, ((0, 0), (0, 0), (guard, guard)), mode="edge") + n_target = C_A.shape[-1] + x_min, x_max = 0.5, 2.0 + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + declined_plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, local_transforms=transforms, + local_radius=1.0, time_reconstruction_certified=False, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True, + discovery_capacity_ok=False) + + def fake_exact(table, _norm, _x_grid, _log_w_grid, **_kwargs): + return jnp.zeros((1, table.shape[-1]), dtype=jnp.float64) + + monkeypatch.setattr(AM, "coefficient_table_distphipsimarg_exact", + fake_exact) + fine = np.linspace(0.0, n_target - 1.0, 2 * n_target - 1) + valid_weights = np.ones(fine.size) + cases = [] + cases.append((fine[::-1], valid_weights, + "reserve_time_nodes_increasing")) + nan_node = fine.copy() + nan_node[4] = np.nan + cases.append((nan_node, valid_weights, "reserve_time_nodes_finite")) + outside = fine.copy() + outside[0] = -0.25 + cases.append((outside, valid_weights, "reserve_time_nodes_in_support")) + negative = valid_weights.copy() + negative[3] = -1.0 + cases.append((fine, negative, "reserve_time_weights_valid")) + not_finite = valid_weights.copy() + not_finite[3] = np.nan + cases.append((fine, not_finite, "reserve_time_weights_valid")) + cases.append((fine, np.zeros_like(valid_weights), + "reserve_time_weights_valid")) + cases.append((np.linspace(0.25, n_target - 1.25, fine.size), + valid_weights, "reserve_time_nodes_cover_target")) + cases.append((np.arange(n_target, dtype=float), np.ones(n_target), + "reserve_time_subsampled")) + + for nodes, weights, failed_field in cases: + _, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=nodes, + reserve_time_resolution_warranted=True, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert not bool(usable) + assert not bool(ledger[failed_field]) + assert bool(ledger["reserve_time_failed"]) + assert bool(ledger["fallback_required"]) + assert bool(ledger["reconciles"]) + + # A REPEATED position is a non-decreasing rule and is accepted: the + # peak-local time rule (peaklocal_time_reserve) repeats a position where + # a mode slot is dead or a block is clipped, and the repeat carries no + # trapezoid weight. Only a decreasing rule fails the node check. + duplicate = fine.copy() + duplicate[4] = duplicate[3] + _, _, ledger = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=valid_weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=duplicate, + reserve_time_resolution_warranted=True, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert bool(ledger["reserve_time_nodes_increasing"]) + + delta_t = 0.25 + physical_weights = JCORE._simpson_weights( + fine.size, delta_t / 2.0) + physical_expected = np.log((n_target - 1) * delta_t) + value, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=physical_weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=fine, + reserve_time_resolution_warranted=True, + reserve_time_check_value=physical_expected, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert bool(usable) + assert float(value) == pytest.approx(physical_expected, abs=2.0e-12) + assert bool(ledger["reserve_time_warranted"]) + + _, usable, aggregate = AAP.empirical_enrichment_with_exact_reserve( + guarded, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=physical_weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=fine, + reserve_time_resolution_warranted=True, + reserve_time_check_value=physical_expected - 6.0e-4, + reserve_time_resolution_tol_nats=1.0e-3, + total_value_error_budget_nats=5.0e-4, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert not bool(usable) + assert bool(aggregate["reserve_time_resolution_validated"]) + assert bool(aggregate["reserve_time_guard_validated"]) + assert not bool(aggregate["reserve_time_error_budget_ok"]) + assert bool(aggregate["reserve_time_failed"]) + assert bool(aggregate["fallback_required"]) + assert bool(aggregate["reconciles"]) + + def table_dependent_exact(table, _norm, _x_grid, _log_w_grid, + **_kwargs): + return jnp.real(table[0, 1])[None, :] + + monkeypatch.setattr(AM, "coefficient_table_distphipsimarg_exact", + table_dependent_exact) + corrupt_guard = guarded.copy() + corrupt_guard[0, 1, 0] += 100.0 + _, usable, ledger = AAP.empirical_enrichment_with_exact_reserve( + corrupt_guard, C_B, declined_plan, declined_plan, x_min, x_max, + reserve_x_grid=np.asarray([x_min, x_max]), + reserve_log_weights=np.zeros(2), time_weights=physical_weights, + reserve_amp_sizing=30.0, reserve_dense_chunk=8, + reserve_grid_block=16, reserve_time_nodes=fine, + reserve_time_resolution_warranted=True, + reserve_time_check_value=physical_expected, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert not bool(usable) + assert float(ledger["reserve_time_guard_error"]) > 1.0e-3 + assert not bool(ledger["reserve_time_guard_validated"]) + assert bool(ledger["reserve_time_failed"]) + assert bool(ledger["fallback_required"]) + assert bool(ledger["reconciles"]) + + +def test_missing_completeness_declines_to_reserve_not_waveform_failure(): + C_A, C_B, constants = _problem(33) + centers, hessian = _joint_peak(constants) + transforms, half_widths = AAP.mode_local_geometry(hessian, w_sigma=3.0) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=4, local_transforms=transforms, + local_radius=3.0, enumeration_complete=False, + outside_bound_certified=False) + value, ok, ledger = AAP.all_axis_peak_local_marginalize( + C_A, C_B, plan, 0.2, 7.0, local_order=3, check_order=5) + + assert np.isfinite(float(value)) + assert not bool(ok) + assert bool(ledger["fallback_required"]) + assert bool(ledger["decline_incomplete"]) + assert not bool(ledger["decline_is_waveform_failure"]) + assert bool(ledger["reconciles"]) + + +def test_two_guard_local_integral_validates_same_target_window(): + C_A, C_B, constants = _problem(33) + guard = 32 + support_time = np.arange(-guard, C_A.shape[-1] + guard, dtype=float) + guarded = np.zeros(C_A.shape[:-1] + (support_time.size,), dtype=np.complex128) + guarded[0, 1] = (constants["k0"] - constants["kt"] + * np.cos(2.0 * np.pi * support_time / constants["span"])) + guarded[2, 1] = 0.5 * constants["kp"] + guarded[0, 0] = 0.5 * constants["ku"] + guarded[0, 2] = 0.5 * constants["ku"] + np.testing.assert_allclose(guarded[..., guard:-guard], C_A, atol=1e-14) + + centers, hessian = _joint_peak(constants) + transforms, _ = AAP.mode_local_geometry(hessian, w_sigma=3.0) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=4, local_transforms=transforms, + local_radius=3.0, outside_bound_certified=False, + time_reconstruction_certified=True, + time_outside_log_bound=-np.inf, + time_outside_bound_certified=True) + value, ok, ledger = AAP.all_axis_peak_local_marginalize( + guarded, C_B, plan, 0.2, 7.0, + local_order=7, check_order=11, time_guard=guard, + time_guard_tol_nats=1.0e-3) + + assert np.isfinite(float(value)) + assert not bool(ok) # outside mass is deliberately still unwarranted + assert bool(ledger["time_guard_validated"]) + assert bool(ledger["time_reconstruction_warranted"]) + assert int(ledger["time_guard"]) == 32 + assert int(ledger["time_guard_inner"]) == 16 + assert float(ledger["time_guard_error"]) <= 1.0e-3 + assert int(ledger["n_guard_local_evaluations_hi"]) == 2 * 11 ** 4 + assert int(ledger["n_total_local_evaluations_hi"]) == 4 * 11 ** 4 + assert int(ledger["n_total_time_frequency_terms_hi"]) == ( + int(ledger["n_time_frequency_terms_hi"]) + + int(ledger["n_guard_time_frequency_terms_hi"])) + assert int(ledger["workspace_bytes_peak_bound_hi"]) == max( + int(ledger["workspace_bytes_hi"]), + int(ledger["workspace_bytes_guard_hi"])) + assert bool(ledger["decline_incomplete"]) + assert bool(ledger["reconciles"]) + + _, empirical_ok, empirical_ledger = ( + AAP.empirical_enrichment_marginalize( + guarded, C_B, plan, plan, 0.2, 7.0, + base_order=7, base_check_order=9, + enriched_order=9, enriched_check_order=11, + convergence_tol_nats=1.0e-3, time_guard=guard, + time_guard_tol_nats=1.0e-3, + total_value_error_budget_nats=1.0e-2)) + assert bool(empirical_ok) + assert float(empirical_ledger["error_score_base_time_guard_nats"]) == ( + pytest.approx(float(empirical_ledger["base_time_guard_error"]))) + assert float(empirical_ledger[ + "error_score_enriched_time_guard_nats"]) == pytest.approx( + float(empirical_ledger["enriched_time_guard_error"])) + assert float(empirical_ledger["error_score_base_time_guard_nats"]) > 0.0 + assert bool(empirical_ledger["value_error_budget_ok"]) + assert bool(empirical_ledger["reconciles"]) + + # Corrupt only support discarded by the inner guard. Integer target + # samples remain unchanged, but the outer Fourier seam rings into the local + # nodes; the two-guard comparison must see it rather than blessing exact + # retained-sample parity or trusting the plan's stale external warrant. + bad = guarded.copy() + bad[0, 1, :guard // 2] += 1.0e4 + _, _, bad_ledger = AAP.all_axis_peak_local_marginalize( + bad, C_B, plan, 0.2, 7.0, + local_order=7, check_order=11, time_guard=guard, + time_guard_tol_nats=1.0e-3) + assert not bool(bad_ledger["time_guard_validated"]) + assert not bool(bad_ledger["time_reconstruction_warranted"]) + assert float(bad_ledger["time_guard_error"]) > 1.0e-3 + _, bad_empirical_ok, bad_empirical_ledger = ( + AAP.empirical_enrichment_marginalize( + bad, C_B, plan, plan, 0.2, 7.0, + base_order=7, base_check_order=9, + enriched_order=9, enriched_check_order=11, + convergence_tol_nats=1.0e-3, time_guard=guard, + time_guard_tol_nats=1.0e-3, + total_value_error_budget_nats=1.0e-15)) + assert not bool(bad_empirical_ok) + assert bool(bad_empirical_ledger["decline_time_reconstruction"]) + assert not bool(bad_empirical_ledger["decline_error_budget"]) + assert bool(bad_empirical_ledger["reconciles"]) + + +def test_certified_omitted_mass_can_cover_an_incomplete_root_report(): + C_A, C_B, constants = _problem(33) + scale = 0.1 + C_A *= scale + constants = dict(constants) + for name in ("k0", "kt", "kp", "ku"): + constants[name] *= scale + x_min, x_max = 0.5, 2.0 + truth = _analytic_log_integral(constants, x_min, x_max) + # One diagonal affine region is exactly the complete t x phi x u x x + # support. Its complement has zero measure, so -inf is an actual outside + # integral bound rather than a fabricated tail assertion. + centers = np.asarray([[ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_min + x_max)]]) + transforms = np.asarray([np.diag([ + constants["span"] / 2.0, np.pi, np.pi, + 0.5 * (x_max - x_min)])]) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=2, + local_transforms=transforms, local_radius=1.0, + outside_log_bound=-np.inf, + # This models an incomplete/degenerate root report. The exact full- + # support cover, not the root count, owns the science error budget. + enumeration_complete=False, outside_bound_certified=True, + # This fixture is itself an exact finite reflected cosine series. Real + # packets need the independent two-guard convergence warrant. + time_reconstruction_certified=True) + value, ok, ledger = AAP.all_axis_peak_local_marginalize( + C_A, C_B, plan, x_min, x_max, local_order=13, check_order=19, + quadrature_tol_nats=2.0e-5) + + assert bool(ok) + assert float(value) == pytest.approx(truth, abs=2.0e-5) + assert not bool(ledger["enumeration_complete"]) + assert bool(ledger["outside_bound_certified"]) + assert not bool(ledger["fallback_required"]) + assert bool(ledger["reconciles"]) + + +@pytest.mark.parametrize("mutation", ("nan", "upper", "negative_diagonal")) +def test_plan_rejects_geometry_that_does_not_match_the_integrated_region(mutation): + _, _, constants = _problem(33) + centers, hessian = _joint_peak(constants) + transforms, _ = AAP.mode_local_geometry(hessian, w_sigma=3.0) + if mutation == "nan": + transforms[0, 0, 0] = np.nan + elif mutation == "upper": + transforms[0, 0, 1] = 1.0e-4 + else: + transforms[0, 0, 0] *= -1.0 + with pytest.raises(ValueError): + AAP.make_all_axis_mode_plan( + centers, max_modes=4, local_transforms=transforms, + local_radius=3.0) + + +def test_fixed_plan_is_transform_compatible_without_claiming_derivative_accuracy(): + C_A, C_B, constants = _problem(33) + centers, hessian = _joint_peak(constants) + transforms, half_widths = AAP.mode_local_geometry(hessian, w_sigma=3.0) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=4, local_transforms=transforms, + local_radius=3.0) + C_A = jnp.asarray(C_A) + C_B = jnp.asarray(C_B) + + @jax.jit + def value(scale): + answer, _, _ = AAP.all_axis_peak_local_marginalize( + scale * C_A, C_B, plan, 0.2, 7.0, + local_order=3, check_order=5) + return answer + + got = value(1.0) + gradient = jax.grad(value)(1.0) + hessian_value = jax.hessian(value)(1.0) + assert np.all(np.isfinite(np.asarray([got, gradient, hessian_value]))) + _, _, ledger = AAP.all_axis_peak_local_marginalize( + C_A, C_B, plan, 0.2, 7.0, local_order=3, check_order=5) + assert bool(ledger["fixed_plan_autodiff_only"]) + assert not bool(ledger["derivative_warrant_certified"]) + + +def test_padded_fixed_plan_survives_outer_vmap(): + C_A, C_B, constants = _problem(33) + centers, hessian = _joint_peak(constants) + transforms, _ = AAP.mode_local_geometry(hessian, w_sigma=3.0) + plan = AAP.make_all_axis_mode_plan( + centers, max_modes=4, local_transforms=transforms, + local_radius=3.0) + + def value(table): + return AAP.all_axis_peak_local_marginalize( + table, C_B, plan, 0.2, 7.0, + local_order=3, check_order=5)[0] + + batch = jnp.asarray(np.stack((C_A, 1.01 * C_A))) + result = jax.jit(jax.vmap(value))(batch) + assert result.shape == (2,) + assert np.all(np.isfinite(np.asarray(result))) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py index 13bf48a0a..d5c024dc7 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_compile_cost.py @@ -249,6 +249,71 @@ def shifted(dc): rtol=0, atol=1e-12) +def test_laplace_point_axis_is_really_tiled(monkeypatch): + """The fused driver must present at most ``point_block`` sample-time bins + to the expensive psi kernel. + + This is a trace-level allocation test, not an estimate from the public + sampler cap. With S=2 and npts=5, the old call handed the kernel all ten + points (as separate S,T axes); the tiled call below must hand it only three + at a time. Deleting the point map or moving it below the psi kernel makes + this fail while all value-only tests remain green. + """ + data = make_synth(npts=5) + xg, lwg = make_distance_grid(30.0, 3000.0, 4, + distMpcRef=data.distMpcRef) + seen = [] + real = AM._laplace_psi_lnI_block + + def spy(a, c1, c2): + seen.append(tuple(a.shape)) + return real(a, c1, c2) + + monkeypatch.setattr(AM, "_laplace_psi_lnI_block", spy) + + def f(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, amp_sizing=450.0, + phi_chunk=4, dist_block=2, point_block=3) + + jax.make_jaxpr(f)(jnp.asarray([0.9, 1.2]), + jnp.asarray([0.4, -0.2]), + jnp.asarray([1.1, 2.0])) + assert seen, "the fused path never called the block-dispatched psi kernel" + assert all(sh == (2, 4, 3) for sh in seen), seen + + +def test_laplace_point_tiling_preserves_value_and_gradient(): + """Tail padding/reassembly and the rolled map preserve values and AD. + + ``point_block=10`` is the one-block reference for S*npts=10; + ``point_block=3`` exercises three full blocks and a one-point tail. A + zero-padded tail, wrong transpose, dropped block, or stop_gradient around + the map fails this test. The tolerance covers only the dispatcher's + documented sub-roundoff choice of a cheaper quadrature rung per tile. + """ + data = make_synth(npts=5, kappa_boost=2.0) + xg, lwg = make_distance_grid(30.0, 3000.0, 4, + distMpcRef=data.distMpcRef) + theta = jnp.asarray([[0.9, 0.4, 1.1], [1.2, -0.2, 2.0]]) + + def call(th, point_block): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, th[:, 0], th[:, 1], th[:, 2], xg, lwg, + amp_sizing=450.0, phi_chunk=4, dist_block=2, + point_block=point_block) + + ref = call(theta, 10) + got = call(theta, 3) + np.testing.assert_allclose(np.asarray(got), np.asarray(ref), + rtol=0.0, atol=2e-12) + + g_ref = jax.grad(lambda th: jnp.sum(call(th, 10)))(theta) + g_got = jax.grad(lambda th: jnp.sum(call(th, 3)))(theta) + np.testing.assert_allclose(np.asarray(g_got), np.asarray(g_ref), + rtol=2e-11, atol=2e-11) + + # --------------------------------------------------------------------------- # Execution-side memory: the batched-eval chunk cap. # diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py index 5cdbe9eaa..8b82f7c0a 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_exact.py @@ -202,6 +202,31 @@ def test_time_marginalization_destroys_the_invariant(): assert C[3:].max() > 1e-9 * C.max() +def test_exact_reserve_reuses_batched_and_collapsed_coefficient_tables(): + """Peak-local declines can reach the exact reserve without recomputing Q/U/V.""" + data = make_synth(scale=0.1, npts=9) + x_grid, log_w = _dist_grid(data, n=16) + C_A, C_B, meta = AM.angle_coefficient_tables( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), INTERP) + assert np.max(np.abs(np.asarray(C_B[..., 0, :]) + - np.asarray(C_B[..., 0, :1]))) < 1.0e-12 + + from_tables = AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, amp_sizing=30.0, + m_max=meta["m_max"], dense_chunk=8, grid_block=8) + collapsed = AM.coefficient_table_distphipsimarg_exact( + C_A[:, :, 0, :], C_B[:, :, 0, 0], x_grid, log_w, + amp_sizing=30.0, m_max=meta["m_max"], + dense_chunk=8, grid_block=8) + wrapped = AM.fused_log_likelihood_distphipsimarg_exact( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=30.0, + dense_chunk=8, grid_block=8, return_lnLt=True) + + np.testing.assert_allclose(from_tables, wrapped, rtol=0.0, atol=2.0e-12) + np.testing.assert_allclose(collapsed, from_tables, rtol=0.0, atol=2.0e-12) + + # --------------------------------------------------------------------------- # 2. sample-grid sizing is derived and asserted, not settable # --------------------------------------------------------------------------- @@ -1031,8 +1056,7 @@ def test_higher_mode_dense_sizing_self_convergence(): 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. + """The pure-JIT metric must be recorded synchronously without excision. (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 @@ -1055,9 +1079,12 @@ def test_runtime_amp_failsafe_warns_and_records_without_excising(capfd): 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) + v, amp_call = AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, + return_amp=True) jax.block_until_ready(v) + AM.record_amp_failsafe( + amp_call, AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, "exact") out = capfd.readouterr() assert "WARNING anglemarg/exact" in out.out + out.err, "guard must warn" @@ -1074,9 +1101,10 @@ def test_runtime_amp_failsafe_warns_and_records_without_excising(capfd): # 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) + v2, amp_call2 = AM.fused_log_likelihood_distphipsimarg_exact( + *args, interp=INTERP, amp_sizing=amp, return_amp=True) jax.block_until_ready(v2) + AM.record_amp_failsafe(amp_call2, amp, "exact") out = capfd.readouterr() assert "WARNING anglemarg" not in out.out + out.err assert AM.amp_failsafe_state()["tripped"] is False @@ -1104,14 +1132,14 @@ def test_driver_labels_a_suspect_angle_grid_in_provenance(): # 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 + assert "def angle_grid_suspect_note(scheme=None)" in src wd = src[src.index("def write_dat("):] wd = wd[:wd.index("\ndef ")] - assert "angle_grid_suspect_note()" in wd, ( + assert "angle_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, ( + assert "_ev_note = angle_grid_suspect_note(_scheme)" 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, ( @@ -1120,34 +1148,26 @@ def test_driver_labels_a_suspect_angle_grid_in_provenance(): -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. - """ +def test_failsafe_jit_is_pure_and_host_record_accumulates_maximum(): + """The persisted graph has no host effects; host state spans all chunks.""" 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 + assert "debug.callback" not in src + assert "debug.print" not in src + assert "return amp_call" in src + + _AMmod.reset_amp_failsafe() + _AMmod.record_amp_failsafe(10.0, 100.0, "exact") + _AMmod.record_amp_failsafe(250.0, 100.0, "exact") + _AMmod.record_amp_failsafe(50.0, 100.0, "exact") + st = _AMmod.amp_failsafe_state() + assert st["tripped"] is True + assert st["n_calls"] == 3 + assert st["worst_amp"] == 250.0 + assert st["amp_sizing"] == 100.0 + assert st["scheme"] == "exact" + _AMmod.reset_amp_failsafe() assert _AMmod.amp_failsafe_state()["tripped"] is False diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py index 398dceef8..097597361 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_gh_selection.py @@ -247,10 +247,52 @@ def test_response_model_is_an_angle_independent_precondition(): assert "factorization" in info["gh_laplace_reason"] -def test_wrapper_passes_the_response_feature_through(): +def test_the_dataset_level_gate_forwards_the_response_feature(): + """The angle-independent half of the gate must run on the DATA's feature. + + Asserted on behaviour, not on source text. This test used to grep + ``JAXDistPhiPsiMargLikelihood.__init__`` for the literal + ``feature=getattr(data, "feature", None)``, which made it fail the moment + the build-time gate was extracted into + ``gh_laplace_supported_for_data`` -- an extraction that REMOVED a duplicate + probe direction, i.e. exactly the change the test should have been + indifferent to. A test pinned to a spelling forbids refactors without + checking anything the spelling was for. + """ + class _FeatureData: + """Wraps a real dataset and relabels only its response model.""" + def __init__(self, inner, feature): + self._inner, self.feature = inner, feature + + def __getattr__(self, name): + return getattr(self._inner, name) + + from test_angle_marg_exact import make_synth + data = make_synth(scale=1.0, npts=64) + assert getattr(data, "feature", None) is None + ok, info = AM.gh_laplace_supported_for_data(data) + assert ok is True, info.get("gh_laplace_reason") + assert info["feature"] is None + + for bad in ("rotation", "freqresponse", "something_added_later"): + ok, info = AM.gh_laplace_supported_for_data(_FeatureData(data, bad)) + assert ok is False, "feature %r was admitted by the dataset gate" % bad + assert info["feature"] == bad + assert "factorization" in info["gh_laplace_reason"] + + +def test_the_wrapper_and_the_policy_ask_the_SAME_gate(): + """One definition, and therefore one probe direction. + + The wrapper had its own ``_ANGLE_MARG_PROBE_*`` constants and built its own + tables. With the policy's reserve roster asking the same question, a second + copy would be a second definition of "the identity holds on this data" -- + and the two could answer differently while both looked right. + """ 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") + assert "gh_laplace_supported_for_data" in src, ( + "the wrapper no longer routes through the shared dataset-level gate") + assert "_ANGLE_MARG_PROBE" not in inspect.getsource(WR), ( + "a second probe direction is back in wrapper.py") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_laplace_table.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_laplace_table.py new file mode 100644 index 000000000..9c2b50d69 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_laplace_table.py @@ -0,0 +1,129 @@ +"""The table-level psi-Laplace reserve: same seam, same normalization. + +`coefficient_table_distphipsimarg_laplace` exists so the four-axis policy can +select a reserve METHOD. Its reserve consumes coefficient tables at refined +time nodes, and until now the only table-level kernel was the exact one, so the +composite could only ever fall back to exact angles whatever the amplitude -- +even where the selector's own calibration says laplace is the more accurate and +far cheaper choice. + +Two things are pinned here, and they are different claims: + +1. EXTRACTION FIDELITY. The fused laplace kernel is now a thin wrapper over + this function, so the two cannot drift. The body moved verbatim; if it did + not, this fails. +2. CROSS-SCHEME AGREEMENT. The table laplace and table exact kernels compute + the same quantity by different routes. The tolerance is MEASURED on these + fixtures and pinned, not copied from the selector's calibration table, which + compares the FUSED paths on real data at stated SNRs and is a different + comparison. +""" + +import numpy as np +import pytest + +import jax +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, _dist_grid + + +def _tables(data): + return AM.angle_coefficient_tables( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), INTERP) + + +def test_the_table_kernel_reproduces_the_fused_laplace_kernel(): + """Extraction fidelity: the fused kernel delegates, so they must agree to + round-off. A larger gap means the body changed in the move, which is the + one thing a refactor may not do.""" + data = make_synth(scale=0.1, npts=9) + x_grid, log_w = _dist_grid(data, n=16) + C_A, C_B, meta = _tables(data) + + from_tables = AM.coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w, amp_sizing=30.0, m_max=meta["m_max"]) + wrapped = AM.fused_log_likelihood_distphipsimarg_laplace( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + x_grid, log_w, interp=INTERP, amp_sizing=30.0, return_lnLt=True) + + np.testing.assert_allclose(np.asarray(from_tables), np.asarray(wrapped), + rtol=0.0, atol=2.0e-12) + + +def test_the_table_kernel_takes_the_collapsed_table_form(): + """The four-axis path holds an unbatched C_A and a time-independent C_B. + The reserve seam must accept that form, or the policy cannot call it.""" + data = make_synth(scale=0.1, npts=9) + x_grid, log_w = _dist_grid(data, n=16) + C_A, C_B, meta = _tables(data) + + batched = AM.coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w, amp_sizing=30.0, m_max=meta["m_max"]) + collapsed = AM.coefficient_table_distphipsimarg_laplace( + C_A[:, :, 0, :], C_B[:, :, 0, 0], x_grid, log_w, + amp_sizing=30.0, m_max=meta["m_max"]) + + np.testing.assert_allclose(np.asarray(collapsed), np.asarray(batched), + rtol=0.0, atol=2.0e-12) + + +@pytest.mark.parametrize("scale,kappa", [(1.0, 1.0), (3.0, 6.0), (6.0, 12.0)]) +def test_the_two_table_reserves_agree_at_machine_level(scale, kappa): + """Cross-scheme agreement on IDENTICAL tables, with the dense grid sized to + the data. + + MEASURED, not copied: 1.2e-14, 1.4e-14 and 2.8e-14 on these three fixtures. + The tolerance below is not the selector's calibrated laplace-vs-exact figure + (-1.6e-05 at rho 40), which compares the FUSED paths on real data at a + stated SNR and is a different comparison; at these synthetic amplitudes both + kernels are essentially exact, so the Laplace error regime is NOT exercised + here and this test must not be read as evidence about it. + """ + data = make_synth(scale=scale, kappa_boost=kappa, npts=9) + x_grid, log_w = _dist_grid(data, n=32) + C_A, C_B, meta = _tables(data) + + kw = dict(amp_sizing=3000.0, m_max=meta["m_max"]) + lap = np.asarray(AM.coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w, **kw)) + exa = np.asarray(AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, dense_chunk=8, grid_block=8, **kw)) + + assert np.all(np.isfinite(lap)), lap + np.testing.assert_allclose(lap, exa, rtol=0.0, atol=1.0e-12) + + +def test_an_undersized_dense_grid_degrades_the_exact_side_not_laplace(): + """Which kernel is fragile when the dense grid is too small for the data. + + This pins a correction to my own first version of this file, which held + amp_sizing fixed while raising the amplitude and read the growing + disagreement as the Laplace error growing. It is the opposite: laplace + removes the psi axis analytically and has no dense u grid to undersize, so + the degradation is entirely on the exact side. Measured on the loud + fixture: 5.7e-07 at amp_sizing 30 against 2.8e-14 at 3000, same tables. + + That is also the practical argument for the reserve hierarchy -- above the + crossover the exact reserve needs a grid that grows with amplitude, and the + laplace reserve does not. + """ + data = make_synth(scale=6.0, kappa_boost=12.0, npts=9) + x_grid, log_w = _dist_grid(data, n=32) + C_A, C_B, meta = _tables(data) + + def gap(amp): + lap = np.asarray(AM.coefficient_table_distphipsimarg_laplace( + C_A, C_B, x_grid, log_w, amp_sizing=amp, m_max=meta["m_max"])) + exa = np.asarray(AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, amp_sizing=amp, m_max=meta["m_max"], + dense_chunk=8, grid_block=8)) + return float(np.max(np.abs(lap - exa))) + + starved, sized = gap(30.0), gap(3000.0) + assert sized < 1.0e-12, sized + assert starved > 100.0 * sized, (starved, sized) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_multipeak_wiring.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_multipeak_wiring.py new file mode 100644 index 000000000..c1d8c1b6e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_multipeak_wiring.py @@ -0,0 +1,110 @@ +"""`--angle-marg-scheme multipeak`, as wired into the ILE likelihood. + +The four-axis controller validated on the ladder-2 injection +(analyses/va_sequence_20260902/RESULTS_20260909_multipeak_ladder.md: 64/64 +accepted at rho 40.77, 163.08 and 652.31, 4.74-5.17 s/row) was reachable only +from `multipeak_planner`; no driver path selected it. These pin the WIRING. +""" +import subprocess +import sys + +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_multipeak_is_an_offered_choice_and_is_exported(): + """optparse builds --angle-marg-scheme's choices from ANGLE_MARG_CHOICES, so + membership IS the CLI wiring; a scheme absent from it is unreachable.""" + assert "multipeak" in AM.ANGLE_MARG_CHOICES + assert "fused_log_likelihood_distphipsimarg_multipeak" in AM.__all__ + + +def test_multipeak_is_NOT_reachable_from_auto(): + """A scheme that changes the likelihood must not become a default on the + strength of one injection's ladder.""" + for amp in (1.0, 50.0, 500.0, 5.0e4, 5.0e6): + scheme, _ = AM.choose_angle_marg_scheme(amp) + assert scheme != "multipeak", (amp, scheme) + + +def test_multipeak_refuses_lnLt_because_it_owns_the_time_integral(): + """Every other scheme in this family can hand back lnL(t). This one cannot, + and must SAY so rather than return a wrong-shaped array.""" + data = make_synth(scale=2.0) + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="multipeak") + with pytest.raises(ValueError, match="no lnL"): + like._fused(data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), + return_lnLt=True) + + +def test_multipeak_returns_one_finite_value_per_sample(): + """The contract the sampler depends on: shape (S,), finite, no time axis.""" + data = make_synth(scale=2.0) + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="multipeak") + v = np.asarray(like._fused(data, jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL))) + assert v.shape == np.shape(RA), v.shape + assert np.all(np.isfinite(v)), v + + +def test_multipeak_records_its_provenance(): + """This pipeline has a history of silently-inert flags: the scheme actually + used must be visible in the 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="multipeak") + assert like.angle_marg_scheme == "multipeak" + assert like.angle_marg_info["requested"] == "multipeak" + + +def test_the_reserve_callable_runs_the_REAL_laplace_kernel(): + """The reserve is exercised against the real table kernel, not a stand-in. + + A fake here would pass while the real call raises: the policy layer's + resolve_reserve_angular_kernel forwards dense_chunk/grid_block, which + coefficient_table_distphipsimarg_laplace does not accept, and that TypeError + only appears on a genuine reserve evaluation. This calls the shipped + function with the shapes the multipeak reserve builds. + """ + data = make_synth(scale=2.0) + xg = jnp.linspace(0.4, 2.0, 16) + lwg = jnp.zeros(16) - np.log(16.0) + C_A, C_B, meta = AM.angle_coefficient_tables( + data, jnp.asarray(RA[:1]), jnp.asarray(DEC[:1]), jnp.asarray(INCL[:1]), + INTERP, guard=16) + CA = np.moveaxis(np.asarray(C_A), 2, 0)[0][..., 16:-16] + CB = np.moveaxis(np.asarray(C_B), 2, 0)[0][..., 0] + out = AM.coefficient_table_distphipsimarg_laplace( + jnp.asarray(CA), jnp.asarray(CB), xg, lwg, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, m_max=int(meta["m_max"])) + assert np.all(np.isfinite(np.asarray(out))) + + +def test_the_driver_CLI_accepts_it_and_rejects_a_typo(): + """A SUBPROCESS on purpose: optparse builds its choices at import time, so a + test that imports the module cannot see the CLI wiring.""" + import RIFT + drv = str(__import__("pathlib").Path(RIFT.__file__).parents[1] + / "bin" / "integrate_likelihood_extrinsic_jax") + # Assert on optparse's CHOICE VALIDATION, not on --help text: the option is + # built with choices=sorted(ANGLE_MARG_CHOICES) and optparse does not print + # the choice list, so a help-text grep tests the help string rather than the + # wiring and would pass or fail for the wrong reason. + bad = subprocess.run([sys.executable, drv, "--angle-marg-scheme", "multipeaks"], + capture_output=True, text=True, timeout=900) + assert bad.returncode != 0 + assert "multipeaks" in bad.stderr, bad.stderr[-800:] + good = subprocess.run([sys.executable, drv, "--angle-marg-scheme", "multipeak"], + capture_output=True, text=True, timeout=900) + # It must fail for a MISSING-INPUT reason, never because the scheme is invalid. + assert "angle-marg-scheme" not in good.stderr or "choice" not in good.stderr, ( + good.stderr[-800:]) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py index a23525eb1..7104efdbb 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py @@ -32,19 +32,49 @@ def test_peak_local_is_NOT_reachable_from_auto(): assert scheme != "peak-local", (amp, scheme) -@pytest.mark.parametrize("boost", [1.0, 30.0]) +@pytest.mark.parametrize("boost", [1.0, 10.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.""" + scheme it parallels. + + BOOST 30 WAS REPLACED BY 10, AND IT COST 807 SECONDS FOR NO WIRING COVERAGE. Since the + u node count became amplitude-derived and streamed, boost 30 puts amp_sizing at 2691 + and asks for 2188 nodes in 274 sequential stream blocks; that one parametrisation was + 807 s of a 1006 s file, and the jax gate went from ~24 min to 37-46 min. + + It bought nothing this file is for. `amp_sizing` FLOORS AT 450, so boost 1.0 already + requests 896 nodes -- 18.7x the U_NODES_PER_CELL floor -- and therefore already + exercises the amplitude-derived path end to end. What boost 30 added was numerical + stress at production amplitude, and this module's own docstring delegates that: "The + kernel's own numerics are tested in test_joint_anglemarg_peaklocal.py." + + 10.0 is kept rather than dropping to a second floored value because it is the first + boost whose amp_sizing (624) CLEARS the 450 floor -- so the pair still demonstrates that + the sizing tracks amplitude rather than being pinned to the crossover, which is the one + wiring property the second point exists to show. + """ 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" + # _batched returns lnL ALONE for every scheme. The amplitude metric the + # persistent-cache work introduced rides on a separate _batched_amp jit; when + # it rode on _batched the two np.asarray calls below raised "inhomogeneous + # shape" for the amp-sized schemes, because this line compares two wrappers + # whose _batched then had different arities under one attribute name. 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 a.shape == b.shape == np.shape(RA), (a.shape, b.shape) assert np.abs(a - b).max() < 1e-4, (boost, a, b) + # and the metric-bearing sibling returns the SAME likelihood: a pure caching + # change must be numerically inert, so pin it rather than assert it. + values, amp = ex._batched_amp(jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL)) + np.testing.assert_array_equal(np.asarray(values), a) + assert np.asarray(amp).shape == () and np.isfinite(float(amp)) + def test_peak_local_records_its_provenance(): """This pipeline has a documented history of silently-inert flags, so the scheme @@ -129,21 +159,26 @@ def test_peak_local_runs_the_runtime_amplitude_failsafe(): 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( + # size for a much quieter target than the data actually is: the check must notice. + # The kernel now RETURNS its metric rather than reporting it through a host + # callback -- a callback anywhere in the graph makes it ineligible for JAX's + # persistent compilation cache -- so the caller records, exactly as the + # wrapper's batched path does in production. + _, amp_call = AM.fused_log_likelihood_distphipsimarg_peaklocal( data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), x, lw, - interp=INTERP, amp_sizing=1.0) + interp=INTERP, amp_sizing=1.0, return_amp=True) + AM.record_amp_failsafe(amp_call, 1.0, "peak-local") 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(): +def test_peak_local_is_capped_by_the_batch_memory_rule(monkeypatch): """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.""" + phi chunks, four cells and the streamed u-node block, and its scan returns every + ``(phi,distance)`` value, so the batch multiplies the same way the dense schemes do. + Leaving it out of the cap kept an uncapped 8000-sample batch.""" from RIFT.likelihood.jax_ile import samplers as S class _Data(object): @@ -164,19 +199,170 @@ class _NoScheme(object): 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. + # grid inside every phi chunk, so its production-floor model is many times the dense + # model's 8192 bytes/sample/time-point even now that the phi scan reduces into its + # carry rather than stacking. 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 + # than a constant. Asserted on the model directly: every term in it is linear in + # n_x, so a 4x grid is a 4x model. class _Wide(_Like): x_grid = np.zeros(1024) - assert S.angle_marg_eval_chunk(_Wide(), 8000) <= capped + assert (S._peaklocal_bytes_per_sample_pt(_Wide()) + == 4 * S._peaklocal_bytes_per_sample_pt(_Like())) + # Fail-closed is pinned against an EXPLICIT target rather than the device probe. + # It used to ride on _Wide exceeding whatever the fallback guess was; the phi-scan + # fix shrank the model by ~4x and that incidental refusal stopped firing, which is + # a test measuring a magnitude while claiming to measure a behaviour. + one_sample = S._peaklocal_bytes_per_sample_pt(_Wide()) * _Data.npts + monkeypatch.setattr(S, "_angle_marg_buffer_target", lambda: one_sample - 1) + with pytest.raises(MemoryError, match="resource preflight"): + S.angle_marg_eval_chunk(_Wide(), 8000) + monkeypatch.setattr(S, "_angle_marg_buffer_target", lambda: one_sample) + assert S.angle_marg_eval_chunk(_Wide(), 8000) == 1 # the "grid" sentinel means "runs no dense angle scheme" and must stay uncapped assert S.angle_marg_eval_chunk(_NoScheme(), 8000) == 8000 +def test_known_four_gib_device_uses_configured_fraction(monkeypatch): + """The unknown-device 4-GiB reserve must never become a known-device floor.""" + from RIFT.likelihood.jax_ile import samplers as S + + class _Device(object): + platform = "gpu" + + def memory_stats(self): + return {"bytes_limit": 4 << 30, + "largest_free_block_bytes": 4 << 30} + + monkeypatch.setattr(S.jax, "devices", lambda: [_Device()]) + monkeypatch.setattr(S, "_ANGLE_MARG_BUFFER_FRACTION", 0.5) + assert S._angle_marg_buffer_target() == (2 << 30) + + +@pytest.mark.parametrize("amplitude,n_phi", [(450.0, 352), (12500.0, 1792)]) +def test_peak_local_model_is_flat_in_n_phi_because_the_scan_reduces( + amplitude, n_phi): + """The cap must account for every source-visible peak-local payload, and must NOT + carry an `n_phi * n_x` term any more. + + That term was the stacked phi-scan output, and it was real: 19.97 GiB predicted + against 19.99 requested at ladder-2 rung 640. `joint_lnL_phi_dense` now reduces + into its scan carry, so the live phi footprint is one chunk plus an (n_x,) + accumulator. The assertion is written as an EQUALITY against the enumerated terms + and, separately, as independence from `n_phi`: a model that silently regrew an + n_phi term would pass a loose inequality. + """ + from RIFT.likelihood.jax_ile import samplers as S + + class _Data(object): + npts = 1193 + lms = ((2, 2), (2, -2)) + + class _Like(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(256) + angle_marg_info = {"amp_sizing": amplitude} + + from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as _jp + live = min(_jp.u_nodes_in_use(amplitude), _jp.U_NODE_STREAM_CHUNK) + expected = (16 * 256 * 4 * live * 8 # streamed u body, one phi chunk + + 16 * 256 * 8 # one chunk of (phi, distance) values + + 256 * 8 # the (n_x,) phi accumulator + + 256 * 5 * 5 * 16) # the per-distance-node joint tables + per_point = S._peaklocal_bytes_per_sample_pt(_Like()) + assert per_point == expected, (per_point, expected, n_phi) + + +def test_peak_local_model_does_not_grow_with_the_phi_axis(): + """The defect this guards: a model that tracks n_phi means a kernel that + materializes the phi axis. Two amplitudes 256x apart put n_phi 16x apart and must + leave the per-point model unchanged.""" + from RIFT.likelihood.jax_ile import samplers as S + from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as _jp + + class _Data(object): + npts = 1193 + lms = ((2, 2), (2, -2)) + + def _like(amp): + class _L(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(256) + angle_marg_info = {"amp_sizing": amp} + return _L() + + lo, hi = 450.0, 450.0 * 256 + # NOT `== 16 *`: _dense_grid_sizes rounds n_phi up to a multiple of 16, so a 256x + # amplitude gives 352 -> 5440, a factor 15.45. The premise only needs n_phi to + # move a lot. + assert _jp.required_n_phi(hi) > 10 * _jp.required_n_phi(lo), "premise" + # the streamed u block is min(u_nodes, U_NODE_STREAM_CHUNK) and is already at the + # 8-node stream cap at both amplitudes, so the whole model must be identical + assert (min(_jp.u_nodes_in_use(lo), _jp.U_NODE_STREAM_CHUNK) + == min(_jp.u_nodes_in_use(hi), _jp.U_NODE_STREAM_CHUNK)), "premise" + assert (S._peaklocal_bytes_per_sample_pt(_like(lo)) + == S._peaklocal_bytes_per_sample_pt(_like(hi))) + + +def test_peak_local_resource_preflight_refuses_an_unfit_single_sample( + monkeypatch): + """One sample over the allowance must REFUSE, not floor the chunk at one. + + The size no longer comes from the amplitude. This used to read "A=12500 needs + 5.242 GiB/sample", which was true only while the model carried the stacked + `n_phi * n_x` term; with the phi scan reducing into its carry the model is flat in + amplitude, so the unfit sample is built from the dimensions that do still drive it. + """ + from RIFT.likelihood.jax_ile import samplers as S + + class _Data(object): + npts = 8192 + lms = ((2, 2), (2, -2)) + + class _Like(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(1024) + angle_marg_info = {"amp_sizing": 12500.0} + + one_sample = S._peaklocal_bytes_per_sample_pt(_Like()) * _Data.npts + assert one_sample > (4 << 30), "premise: one sample must not fit" + monkeypatch.setattr(S, "_angle_marg_buffer_target", lambda: 4 << 30) + with pytest.raises(MemoryError, match="reducing the outer evaluation chunk"): + S.angle_marg_eval_chunk(_Like(), 8000) + + +def test_peak_local_floor_amplitude_fits_one_sample_at_two_gib(monkeypatch): + """At the floor amplitude a 2 GiB target admits exactly one sample. + + The docstring used to say "A=450 needs 1.966 GiB/sample". That number was the + stacked-phi-scan model; it is now 1.32 GiB and does not depend on the amplitude at + all. The test still passed at the new size, which is how a stale measured number + survives in a green suite -- so the size is asserted here rather than narrated. + """ + from RIFT.likelihood.jax_ile import samplers as S + + class _Data(object): + npts = 1193 + lms = ((2, 2), (2, -2)) + + class _Like(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(256) + angle_marg_info = {"amp_sizing": 450.0} + + one_sample = S._peaklocal_bytes_per_sample_pt(_Like()) * _Data.npts + assert (2 << 30) // one_sample == 1, "premise: exactly one sample fits" + monkeypatch.setattr(S, "_angle_marg_buffer_target", lambda: 2 << 30) + assert S.angle_marg_eval_chunk(_Like(), 8000) == 1 + + 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 @@ -192,6 +378,90 @@ def test_peak_local_artifacts_carry_the_standing_best_effort_label(): mod = importlib.util.module_from_spec(spec) spec.loader.exec_module(mod) AM.reset_amp_failsafe() + # Nothing recorded: the label must STAND and must not read as a pass. The + # former wording was BEST-EFFORT; the metric is now returned data rather + # than a droppable callback, so an unchecked event says NOT-PERFORMED and a + # checked one says OUTPUT-CLOUD-PASS. What the P1 finding requires is + # unchanged -- peak-local never publishes silence. note = mod.angle_grid_suspect_note("peak-local") - assert note.startswith("ANGLE-GRID-CHECK=BEST-EFFORT"), note + assert note.startswith("ANGLE-GRID-CHECK=NOT-PERFORMED"), note + assert "NOT a pass" in note and "UNKNOWN" in note, note + AM.record_amp_failsafe(1.0, 100.0, "peak-local") + checked = mod.angle_grid_suspect_note("peak-local") + assert checked.startswith("ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS"), checked + AM.reset_amp_failsafe() assert mod.angle_grid_suspect_note("grid") == "" + + +def test_kernel_and_memory_guard_read_the_same_node_count(): + """Review P2. ``u_nodes_in_use`` was introduced as the single source of truth for the + u-node count, and its docstring said both the kernel and the batch-memory guard call + it -- but only the guard did. ``joint_lnL_phi_dense`` still defaulted straight to + ``U_NODES_PER_CELL`` and the fused caller passed no ``n_nodes``, so an + amplitude-dependent change would have moved the guard and left the kernel behind. A + single source of truth that only one side reads is not one. + + The invariant is NOT "both currently equal 48" -- production uses the uncapped + derived count. Both sides must read the same amplitude, while the guard models only + the streamed live block rather than the total quadrature work. + + The shape is deliberately NOT the production one. At npts=614 with 256 distance nodes + the cap is already pinned at its floor of 1 -- the measured "peak-local batches one + sample" result -- so quadrupling the node count cannot move it, and the guard assertion + would read ``1 < 1`` and fail while the wiring was correct. A saturated observable + cannot test the thing it saturates on. npts=64 with 32 distance nodes stays clear of + both the floor and the 8000 ceiling. + """ + from RIFT.likelihood.jax_ile import samplers as S + from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as JP + + class _Data(object): + npts = 64 + + class _Like(object): + data = _Data() + angle_marg_scheme = "peak-local" + x_grid = np.zeros(32) + angle_marg_info = {"amp_sizing": 450.0} + + seen = [] + real_helper = JP.u_nodes_in_use + real_inner = JP.log_inner_u_integral + + def _spy_inner(a, c1, c2, n_nodes=JP.U_NODES_PER_CELL, **kw): + seen.append(int(n_nodes)) + return real_inner(a, c1, c2, n_nodes, **kw) + + baseline_cap = S.angle_marg_eval_chunk(_Like(), 8000) + assert 1 < baseline_cap < 8000, baseline_cap # the observable is not saturated + + helper_args = [] + def _raised_policy(amp_sizing=None): + helper_args.append(amp_sizing) + return 4 * real_helper(amp_sizing) + + JP.u_nodes_in_use = _raised_policy + JP.log_inner_u_integral = _spy_inner + try: + # The guard must consult the helper with the production amplitude. Its cap does + # not shrink because the extra total work is streamed through the same live block. + raised_cap = S.angle_marg_eval_chunk(_Like(), 8000) + assert raised_cap == baseline_cap, (baseline_cap, raised_cap) + assert 450.0 in helper_args, helper_args + + # the KERNEL must follow it too, via n_nodes=None resolving through the helper + rng = np.random.default_rng(0) + C_A = rng.normal(size=(3, 3)) + 1j * rng.normal(size=(3, 3)) + C_B = rng.normal(size=(5, 5)) + 1j * rng.normal(size=(5, 5)) + C_B[0, 2] = abs(C_B[0, 2].real) + 3.0 + x_grid = jnp.asarray(np.linspace(0.5, 2.0, 8)) + lw = jnp.zeros(8) + JP.joint_lnL_phi_dense(jnp.asarray(C_A), jnp.asarray(C_B), x_grid, lw, n_phi=8) + assert seen, "kernel never reached log_inner_u_integral" + assert set(seen) == {4 * real_helper(None)}, (seen, real_helper(None)) + finally: + JP.u_nodes_in_use = real_helper + JP.log_inner_u_integral = real_inner + + # restoring the helper restores the cap exactly -- no hidden state + assert S.angle_marg_eval_chunk(_Like(), 8000) == baseline_cap diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py index ac43aa90b..6f6f090ee 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_smoke.py @@ -15,7 +15,9 @@ import ast import pathlib +import pytest import numpy as np +import jax import jax.numpy as jnp from RIFT.likelihood.jax_ile import anglemarg as AM @@ -24,6 +26,7 @@ _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 RA, DEC, INCL = 1.1, -0.35, 0.9 INTERP = "sinc" @@ -62,19 +65,24 @@ def test_amp_sizing_is_required_not_defaulted(): 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.""" +def test_failsafe_record_roundtrips_without_host_effects(): + """The host record spans calls while the persistable JIT stays pure.""" import inspect AM.reset_amp_failsafe() + AM.record_amp_failsafe(10.0, 100.0, "exact") + AM.record_amp_failsafe(250.0, 100.0, "exact") + AM.record_amp_failsafe(50.0, 100.0, "exact") 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) + assert st["tripped"] is True + assert st["n_calls"] == 3 + assert st["worst_amp"] == 250.0 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") + assert "debug.callback" not in src and "debug.print" not in src + assert "return amp_call" in src + AM.reset_amp_failsafe() + assert AM.amp_failsafe_state() == { + "tripped": False, "n_calls": 0, "worst_amp": 0.0, + "amp_sizing": None, "scheme": None} def _driver_src(): @@ -99,7 +107,7 @@ def test_driver_actually_passes_the_scheme_through(): "angle_marg must be forwarded as the parsed option, not a constant") -def test_driver_labels_both_artifacts_and_never_implies_verification(): +def test_driver_labels_both_artifacts_with_exact_checked_scope(): src = _driver_src() assert "def angle_grid_suspect_note(" in src # The note is computed ONCE in analyze_one from the RESOLVED scheme and @@ -117,13 +125,167 @@ def test_driver_labels_both_artifacts_and_never_implies_verification(): "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 "ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS" in src + assert "deterministic over pilot/reweight/" in src + assert "transient training-only" in src, ( + "the artifact must not claim coverage of proposals that do not enter " + "the reported evidence or exported output cloud") assert "SUSPECT-ANGLE-GRID" in src +def test_public_wrapper_records_output_calls_not_training_and_drives_note(): + """Pin the production wire from public batches to artifact provenance.""" + like = JAXDistPhiPsiMargLikelihood( + make_synth(), 30.0, 3000.0, n_grid=16, nphi=8, npsi=4, + interp=INTERP, guess_snr=5.0, angle_marg="exact") + assert like._amp_record is not None + assert like._batched_amp is not None + amplitudes = iter((10.0, 1000.0)) + like._batched_amp = lambda ra, dec, incl: ( + jnp.zeros_like(jnp.atleast_1d(ra)), jnp.asarray(next(amplitudes))) + + AM.reset_amp_failsafe() + like.log_likelihood([RA], [DEC], [INCL]) + like.log_likelihood([RA], [DEC], [INCL]) + state = AM.amp_failsafe_state() + assert state["n_calls"] == 2 + assert state["worst_amp"] == 1000.0 + assert state["tripped"] is True + + # Scalar/gradient calls model transient flow-training proposals and must + # not mutate the artifact-producing output-cloud record. Exercise the REAL + # graph before any stub. A mutation sweep found that stubbing _scalar and + # _value_and_grad first, then asserting the record is unchanged, passes even + # when the scalar path IS wired to request the amplitude -- the stub + # replaces exactly the code the assertion is about. Requesting it there + # would also make the scalar path return a tuple, so value_and_grad would be + # differentiating the wrong object; both are caught below. + # theta3 is three SCALARS; the module-level RA/DEC/INCL are 1-element + # arrays for the batched entry points, and passing them here would build a + # (3,1) theta and fail inside the coefficient scan. + # The scalar/AD path must stay a pure value graph. Traced with eval_shape + # rather than executed: this is a shape contract, and tracing costs no + # compile and no device memory. Asking that path for the amplitude makes + # _fused return a TUPLE, so v[0] becomes the (1,)-shaped lnL instead of the + # scalar, and value_and_grad would differentiate the wrong object. A + # mutation sweep found the stub-then-assert check below cannot see that: + # it replaces exactly the code it is meant to be testing. + theta3 = jnp.zeros(3, dtype=jnp.float64) + real_state = AM.amp_failsafe_state() + out_shape = jax.eval_shape(like._scalar, theta3) + assert out_shape.shape == (), ( + "the scalar/AD path must return a scalar, not a (value, amplitude) " + "pair; got %r" % (out_shape,)) + assert AM.amp_failsafe_state() == real_state, ( + "tracing the AD path must not enter the output-cloud record") + + like._scalar = lambda theta: jnp.asarray(1.0) + like._value_and_grad = lambda theta: (jnp.asarray(1.0), jnp.zeros(3)) + like.value([RA, DEC, INCL]) + like.value_and_grad([RA, DEC, INCL]) + assert AM.amp_failsafe_state() == state + + tree = ast.parse(_driver_src()) + note_fn = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "angle_grid_suspect_note") + namespace = {"_anglemarg": AM} + exec(compile(ast.Module(body=[note_fn], type_ignores=[]), + "", "exec"), namespace) + note = namespace["angle_grid_suspect_note"]("exact") + assert note.startswith("SUSPECT-ANGLE-GRID") + + AM.reset_amp_failsafe() + like._amp_record(10.0) + note = namespace["angle_grid_suspect_note"]("exact") + assert note.startswith("ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS") + assert "deterministic over pilot/reweight/final output-cloud" in note + assert "transient training-only proposals not inspected" in note + + + +def test_an_unchecked_amp_sized_scheme_is_labelled_not_performed(): + """A PASS label may only be published when a batch was actually checked. + + The recorder is wired ONLY for direct_marginalization_policy="off", while + ``angle_marg_scheme`` still names the amp-sized reserve scheme, so + ``--angle-marg-scheme exact --direct-marginalization-policy auto`` reaches + angle_grid_suspect_note("exact") having recorded nothing. As merged, that + formatted "%.6g" % None and raised TypeError -- killing the event after the + integration and before either writer. Had amp_sizing merely defaulted to + 0.0 it would instead have published OUTPUT-CLOUD-PASS worst_amp=0: an + adequacy claim backed by zero checks, which is the false negative the whole + label exists to prevent. + + Behavioural on purpose. The sibling checks in this file grep the driver + source; a source grep cannot see either failure, because the string it + matches is present in both the broken and the fixed driver. + """ + tree = ast.parse(_driver_src()) + note_fn = next(node for node in tree.body + if isinstance(node, ast.FunctionDef) + and node.name == "angle_grid_suspect_note") + namespace = {"_anglemarg": AM} + exec(compile(ast.Module(body=[note_fn], type_ignores=[]), + "", "exec"), namespace) + note = namespace["angle_grid_suspect_note"] + + AM.reset_amp_failsafe() + assert AM.amp_failsafe_state()["n_calls"] == 0 + for scheme in ("exact", "laplace", "peak-local", "phi-local"): + label = note(scheme) + assert label.startswith("ANGLE-GRID-CHECK=NOT-PERFORMED"), label + assert "OUTPUT-CLOUD-PASS" not in label + assert "UNKNOWN" in label and "NOT a pass" in label + + # one recorded batch, and the same call is entitled to claim the pass + AM.record_amp_failsafe(1.0, 100.0, "exact") + assert note("exact").startswith("ANGLE-GRID-CHECK=OUTPUT-CLOUD-PASS") + AM.reset_amp_failsafe() + + + + +def test_schemes_without_amp_sizing_refuse_to_invent_a_metric(): + """Asking an unsized scheme for the amplitude metric must raise. + + 'grid' has no amp_sizing and runs no failsafe, and the composite policy + exposes no metric either. Neither is reachable from the wrapper now that + the metric rides on a separate _batched_amp built only for the amp-sized + schemes, so a mutation sweep found both refusals survived: they could be + deleted and nothing failed. They still guard a direct caller, and the + failure they prevent is a silent one -- returning a metric that stands for + no check, which is what the NOT-PERFORMED label exists to keep out of + artifacts. + """ + data = make_synth() + kw = dict(n_grid=16, nphi=8, npsi=4, interp=INTERP, guess_snr=5.0) + ra = jnp.asarray([RA]); dec = jnp.asarray([DEC]); incl = jnp.asarray([INCL]) + + grid = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, + angle_marg="grid", **kw) + with pytest.raises(ValueError, match="no amp_sizing"): + grid._fused(data, ra, dec, incl, return_amp=True) + + pol = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, angle_marg="exact", + direct_marginalization_policy="auto", + **kw) + with pytest.raises(ValueError, match="does not expose"): + pol._fused(data, ra, dec, incl, return_amp=True) + +def test_the_policy_composite_leaves_the_amp_record_unwired(): + """The wiring fact that makes the case above reachable in production.""" + data = make_synth() + like = JAXDistPhiPsiMargLikelihood( + data, 30.0, 3000.0, n_grid=16, nphi=8, npsi=4, interp=INTERP, + guess_snr=5.0, angle_marg="exact", + direct_marginalization_policy="auto") + assert like.angle_marg_scheme == "exact", ( + "the composite still names an amp-sized reserve scheme, which is what " + "sends angle_grid_suspect_note down the labelled branch") + assert like._amp_record is None, ( + "the composite exposes no amplitude metric, so nothing records") + 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). diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_anglemarg_buffer_cap.py b/MonteCarloMarginalizeCode/Code/test/jax/test_anglemarg_buffer_cap.py new file mode 100644 index 000000000..5dea145cc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_anglemarg_buffer_cap.py @@ -0,0 +1,727 @@ +#!/usr/bin/env python3 +# Registered by NAME in .travis/test-jax.sh's FILES array -- that job selects by an +# explicit list, not by a marker. A '# RIFT-CI-GATE:' line here would name a gate +# that does not exist and the roster census refuses it, correctly. +"""The anglemarg eval-chunk cap: still bounds the buffer, no longer assumes 4 GiB. + +The cap exists because on 2026-08-28 the laplace path asked XLA for a single 36.41 GiB +buffer at chunk 4000 / npts 1193 and died RESOURCE_EXHAUSTED against a 25 GiB cgroup. +Making the target device-aware must not weaken that: these tests pin the bound itself, +not the constant that used to express it. Device-aware means AVAILABLE memory, not the +allocator's capacity ceiling -- these GPUs are shared, and a ceiling-sized allowance on a +card someone else is already holding is the same OOM with a nicer derivation. +""" +from __future__ import print_function +import pytest + +sam = pytest.importorskip("RIFT.likelihood.jax_ile.samplers") + + +class _Data(object): + def __init__(self, npts): self.npts = npts + + +class _Like(object): + def __init__(self, scheme, npts): + self.angle_marg_scheme = scheme + self.data = _Data(npts) + + +def _target(monkeypatch, byts): + monkeypatch.setattr(sam, "_angle_marg_buffer_target", lambda: byts) + + +def test_the_original_blowup_is_still_refused(monkeypatch): + """chunk 4000 at npts 1193 must not survive at the historical 4 GiB target.""" + _target(monkeypatch, 4 << 30) + got = sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) + assert got < 4000 + # the buffer the returned chunk implies must fit the target + assert got * sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT * 1193 <= (4 << 30) + + +@pytest.mark.parametrize("target", [4 << 30, 12 << 30, 24 << 30]) +def test_the_bound_holds_at_every_target(monkeypatch, target): + """Whatever the device reports, the implied buffer never exceeds it.""" + _target(monkeypatch, target) + for npts in (614, 1193, 4915, 32769): + got = sam.angle_marg_eval_chunk(_Like("laplace", npts), 4000) + assert got >= 1 + assert got * sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT * npts <= target + + +def test_a_bigger_device_lifts_the_throttle(monkeypatch): + """The point of the change: 4 GiB caps production npts below the nominal chunk.""" + npts = 1230 + _target(monkeypatch, 4 << 30) + small = sam.angle_marg_eval_chunk(_Like("laplace", npts), 1000) + _target(monkeypatch, 16 << 30) + big = sam.angle_marg_eval_chunk(_Like("laplace", npts), 1000) + assert small < 1000, "4 GiB should still throttle at production npts" + assert big == 1000, "a 16 GiB device should not throttle at all" + + +def test_grid_is_never_capped(monkeypatch): + """`grid` is a sentinel for 'no dense angle scheme' and must pass through.""" + _target(monkeypatch, 4 << 30) + assert sam.angle_marg_eval_chunk(_Like("grid", 32769), 4000) == 4000 + + +# --------------------------------------------------------------------------- +# Everything above stubs `_angle_marg_buffer_target` via `_target()`, which is right +# for testing the BOUND but means none of it touches the probe itself. An earlier +# revision of this file "covered" the probe with +# monkeypatch.setattr(s, "_angle_marg_buffer_target", lambda: FALLBACK) +# assert s._angle_marg_buffer_target() == (4 << 30) +# which replaces the function under test with a lambda and then asserts the lambda +# returns what it was written to return. It passes against ANY implementation, +# including no implementation. What follows drives the real function by faking the +# device, so the probe fails when the probe is wrong. +# --------------------------------------------------------------------------- + + +class _Dev(object): + """Minimal stand-in for a jax Device. + + `limit` is the allocator's CAPACITY CEILING and, deliberately, is not enough on its + own for the probe to size anything. The earlier version of this class modelled only + a limit, which is why it could not see the review finding below: every fake device it + built was an idle one, so a ceiling and free memory were the same number and treating + one as the other looked correct. `free` (largest servable block) and `pool`/`in_use` + are what say how much of the ceiling is actually obtainable. + """ + def __init__(self, platform, limit=None, key="bytes_limit", + free=None, pool=None, in_use=None): + self.platform = platform + self._limit = limit + self._key = key + self._free = free + self._pool = pool + self._in_use = in_use + + def memory_stats(self): + stats = {} + if self._limit is not None: + stats[self._key] = self._limit + if self._free is not None: + stats["largest_free_block_bytes"] = self._free + if self._pool is not None: + stats["pool_bytes"] = self._pool + if self._in_use is not None: + stats["bytes_in_use"] = self._in_use + return stats + + +def _idle_gpu(total): + """A card of `total` bytes with nobody else on it: ceiling AND free both `total`.""" + return _Dev("gpu", total, free=total) + + +def _fake_jax(monkeypatch, devices=None, raises=None): + """Install a fake `jax` module that the probe's local `import jax` will find.""" + import sys + import types + mod = types.ModuleType("jax") + if raises is not None: + def devs(): + raise raises + else: + def devs(): + return list(devices) + mod.devices = devs + monkeypatch.setitem(sys.modules, "jax", mod) + return mod + + +GIB = 1 << 30 + + +def test_probe_failure_falls_back_to_four_gib(monkeypatch): + """A device we cannot interrogate must behave exactly as before -- never larger.""" + _fake_jax(monkeypatch, raises=RuntimeError("no device")) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_no_gpu_falls_back_to_four_gib(monkeypatch): + """CPU-only: nothing to be device-aware about.""" + _fake_jax(monkeypatch, devices=[_Dev("cpu", 999 * GIB, free=999 * GIB)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_empty_memory_stats_falls_back_to_four_gib(monkeypatch): + """A GPU whose runtime reports nothing is a probe failure, not a zero limit.""" + _fake_jax(monkeypatch, devices=[_Dev("gpu", None)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_the_gpu_is_picked_out_of_a_mixed_device_list(monkeypatch): + """The platform filter must actually select, not just happen to be index 0.""" + _fake_jax(monkeypatch, devices=[_Dev("cpu", 999 * GIB, free=999 * GIB), + _idle_gpu(24 * GIB)]) + assert sam._angle_marg_buffer_target() == 12 * GIB + + +def test_the_reservable_limit_still_clamps_when_bytes_limit_is_absent(monkeypatch): + """The alternate ceiling spelling is still read -- but only downward. + + A runtime reporting a free block larger than its own allocator limit is misreporting; + the ceiling may shrink the allowance, never license one. + """ + _fake_jax(monkeypatch, + devices=[_Dev("gpu", 8 * GIB, key="bytes_reservable_limit", + free=999 * GIB)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_the_fraction_is_applied_to_available_memory(monkeypatch): + monkeypatch.setattr(sam, "_ANGLE_MARG_BUFFER_FRACTION", 0.25) + _fake_jax(monkeypatch, devices=[_idle_gpu(24 * GIB)]) + assert sam._angle_marg_buffer_target() == 6 * GIB + + +@pytest.mark.parametrize("free_gib", [1, 2, 4, 6, 8, 16, 24, 80]) +def test_the_allowance_never_exceeds_what_is_actually_free(monkeypatch, free_gib): + """THE regression this file exists for after review. + + An earlier revision returned max(4 GiB, limit * fraction). On a 6 GiB card that is + 4 GiB -- two thirds of the whole device for ONE buffer -- and on anything under 4 GiB + it hands out more memory than exists. 4 GiB is the answer for a device we cannot + SEE; it is not a safe minimum for a device we can. + """ + _fake_jax(monkeypatch, devices=[_idle_gpu(free_gib * GIB)]) + got = sam._angle_marg_buffer_target() + assert got <= free_gib * GIB, "allowance exceeds what the device reports free" + assert got == int(free_gib * GIB * sam._ANGLE_MARG_BUFFER_FRACTION) + + +def test_a_small_device_is_not_floored_at_four_gib(monkeypatch): + """Stated separately from the sweep so the failure names the defect.""" + _fake_jax(monkeypatch, devices=[_idle_gpu(6 * GIB)]) + assert sam._angle_marg_buffer_target() == 3 * GIB + + +# --- the ceiling is not free memory (review P1, second round) ---------------- +# Every fake device above this line was IDLE, so its ceiling and its free memory were +# the same number and a probe that read either looked correct. The cards this runs on +# are shared: a survey of the interactive hosts found 24 GiB GPUs with 18-22 GiB already +# held by other processes. `bytes_limit` does not move when that happens. + + +def test_a_busy_shared_card_is_not_sized_from_its_ceiling(monkeypatch): + """24 GiB ceiling, 22 GiB held by someone else, 2 GiB actually free. + + Sizing off the ceiling returns a 12 GiB allowance here -- six times what the card + has left -- and walks straight back into the RESOURCE_EXHAUSTED this cap exists to + prevent. The allowance must come from the 2 GiB, not the 24. + """ + _fake_jax(monkeypatch, devices=[_Dev("gpu", 24 * GIB, free=2 * GIB)]) + got = sam._angle_marg_buffer_target() + assert got < 12 * GIB, "allowance still derived from the capacity ceiling" + assert got <= 2 * GIB, "allowance exceeds the memory that is actually free" + assert got == 1 * GIB + + +def test_a_zero_largest_free_block_alone_is_not_known_full(monkeypatch): + """A zero largest-free-block reading by itself is NOT proof the device is full. + + Regression (2026-09-08, jax 0.9.2, ldas-pcdev11, idle 24 GiB RTX PRO 4000): + ``largest_free_block_bytes`` is simply never populated on this jax version -- it + reads 0 both before any allocation and after one. Treating a bare 0 as "full" + made the exact scheme's preflight refuse on an idle card with + "allowance of 0 bytes". With no pool information either, this is an unknown, not + a full device: fall back to the conservative 4 GiB, same as an unreadable device. + """ + _fake_jax(monkeypatch, devices=[_Dev("gpu", 24 * GIB, free=0)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + # cap = target // (8192 bytes/sample-pt * 1193) = 4 GiB // 9773056 = 439; the + # requested 4000 does not fit whole, so the point proven here is "no MemoryError", + # not "chunk unchanged". + assert sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) == 439 + + +def test_a_zero_largest_free_block_with_a_full_pool_is_still_known_full(monkeypatch): + """The pool signal can still report a genuinely full device even when the + free-block signal is uninformative (as on jax 0.9.2): the "full" case in the + docstring is real and must survive the fix to the "unreported" case above.""" + _fake_jax(monkeypatch, + devices=[_Dev("gpu", 24 * GIB, free=0, pool=24 * GIB, in_use=24 * GIB)]) + assert sam._angle_marg_buffer_target() == 0 + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) + + +# --- jax 0.9.2 default-preallocation dicts, measured 2026-09-08 on ldas-pcdev11 ----- +# largest_free_block_bytes is 0 in all three (never populated on this jax version); +# pool_bytes is 0 until the first allocation. This is the actual reproduction: PR #250 +# made 0 mean "full" for both signals, so `integrate_likelihood_extrinsic_jax +# --mode flowmc-phipsimarg --angle-marg-scheme exact` (exact is that mode's default) +# refused at the preflight on an idle, otherwise-healthy GPU. + +JAX092_BEFORE_ALLOC = { + "bytes_limit": 18895355904, + "bytes_in_use": 0, + "pool_bytes": 0, + "largest_free_block_bytes": 0, + "bytes_reservable_limit": None, +} + +JAX092_AFTER_8MB = { + "bytes_limit": 18895355904, + "bytes_in_use": 8388608, + "pool_bytes": 1259690240, + "largest_free_block_bytes": 0, +} + +JAX092_AFTER_512MB = { + "bytes_limit": 18895355904, + "bytes_in_use": 545259520, + "pool_bytes": 1259690240, + "largest_free_block_bytes": 0, +} + + +class _StatsDev(object): + """A device whose memory_stats() returns one exact dict, verbatim.""" + def __init__(self, platform, stats): + self.platform = platform + self._stats = stats + + def memory_stats(self): + return dict(self._stats) + + +def test_device_available_bytes_before_any_allocation_is_unknown(): + """Both signals report 0 before the pool exists; neither is a real reading.""" + assert sam._device_available_bytes(JAX092_BEFORE_ALLOC) is None + + +def test_device_available_bytes_after_8mb_is_positive(): + assert sam._device_available_bytes(JAX092_AFTER_8MB) == 1259690240 - 8388608 + + +def test_device_available_bytes_after_512mb_is_positive(): + assert sam._device_available_bytes(JAX092_AFTER_512MB) == 1259690240 - 545259520 + + +def test_jax092_before_allocation_falls_back_not_refuses(monkeypatch): + """THE reproduction. 614 points at the module's 8192 bytes/sample-point is + 5029888 bytes -- the exact figure the preflight quoted in the failure this PR + fixes -- and now fits under the 4 GiB fallback once 0 stops meaning "full", + capping the requested chunk of 4000 down to 853 (= 4 GiB // 5029888) instead + of refusing outright.""" + _fake_jax(monkeypatch, devices=[_StatsDev("gpu", JAX092_BEFORE_ALLOC)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + assert sam.angle_marg_eval_chunk(_Like("exact", 614), 4000) == 853 + + +def test_a_ceiling_with_no_free_report_falls_back_rather_than_guessing_up(monkeypatch): + """The device is visible but says nothing about occupancy. + + This is the shape the old fake device had, and the answer is NOT half the ceiling: + a limit alone cannot distinguish an idle card from a full one. Fall back to the + conservative 4 GiB and let the operator assert otherwise with the absolute override. + """ + _fake_jax(monkeypatch, devices=[_Dev("gpu", 24 * GIB)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_the_reserved_pool_minus_what_we_hold_is_used_when_no_block_is_reported( + monkeypatch): + """Second-choice availability signal: memory already reserved to us is genuinely + ours, unlike the ceiling, so pool - in_use is a real free figure.""" + _fake_jax(monkeypatch, + devices=[_Dev("gpu", 24 * GIB, pool=16 * GIB, in_use=4 * GIB)]) + assert sam._angle_marg_buffer_target() == 6 * GIB + + +def test_a_full_pool_yields_no_allowance_and_the_eval_refuses(monkeypatch): + """Nothing free is a READING, not a failure to read. + + Falling back to the 4 GiB guess here would hand out memory the runtime has just said + does not exist, so the target goes to zero and the eval refuses with the message that + names the knobs -- an outage the operator can act on, not a silent OOM later. + """ + _fake_jax(monkeypatch, + devices=[_Dev("gpu", 24 * GIB, pool=24 * GIB, in_use=24 * GIB)]) + assert sam._angle_marg_buffer_target() == 0 + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) + # and the sentinel still short-circuits before any of it + assert sam.angle_marg_eval_chunk(_Like("grid", 1193), 4000) == 4000 + + +# --- the advertised override ------------------------------------------------ + +def test_the_default_fraction_applies_when_unset(): + assert sam._read_buffer_fraction({}) == 0.5 + + +@pytest.mark.parametrize("raw,expect", [("0.8", 0.8), ("1.0", 1.0), ("0.25", 0.25)]) +def test_a_usable_override_is_honoured(raw, expect): + assert sam._read_buffer_fraction( + {"RIFT_ANGLEMARG_BUFFER_FRACTION": raw}) == expect + + +@pytest.mark.parametrize("raw", ["", "half", "0.5x", "1.5", "2", "0", "-0.5", "nan"]) +def test_an_unusable_override_is_refused_loudly(raw): + """Refused, NOT silently replaced by the default. + + A value above 1 asks for a buffer bigger than the device reports, i.e. asks this + code to cause the OOM it exists to prevent. A value at or below 0 bounds nothing. + Either way the caller believes a bound is in force, so failing quietly is worse + than failing. + """ + with pytest.raises(ValueError): + sam._read_buffer_fraction({"RIFT_ANGLEMARG_BUFFER_FRACTION": raw}) + + +# --- the constant the whole bound rests on ---------------------------------- + +def test_bytes_per_sample_point_still_reproduces_the_observed_allocation(): + """Pin _ANGLE_MARG_BYTES_PER_SAMPLE_PT against a number the code does not own. + + FOUND BY MUTATION, and it is why this test exists: halving the constant + 8192 -> 4096 left all 33 other tests in this file passing. Every one of them + computes the expected buffer as `got * sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT * npts` + -- reading the same constant the production code reads -- so the assertion is + self-consistent for ANY value of it. The bound would silently permit a buffer + twice the intended size and the suite would stay green. + + The independent reference is XLA's own report from 2026-08-28: at chunk 4000 + and npts 1193 the laplace path asked for a single buffer of 36.41 GiB. 8192 + reproduces that to 0.01%. This is an EXTERNAL measurement, not a restatement + of the constant, so it fails when the constant moves. + """ + observed_gib = 36.41 # from the RESOURCE_EXHAUSTED message itself + chunk, npts = 4000, 1193 # the configuration that produced it + implied = chunk * npts * sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT / float(GIB) + assert abs(implied / observed_gib - 1.0) < 0.01, ( + "%d bytes/sample-point implies a %.2f GiB buffer at chunk %d / npts %d, but " + "the allocation this cap was built from was %.2f GiB. If the per-point size " + "genuinely changed, re-measure it and update BOTH the constant and this " + "reference." % (sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT, implied, chunk, npts, + observed_gib)) + + +# --------------------------------------------------------------------------- +# THE BOUND WAS NOT ACTUALLY A BOUND: max(1, target // per_sample) +# +# Review P1 on #250. Every assertion above stubs a target that is comfortably larger +# than one sample, so none of them can reach the floor. Once ONE sample costs more than +# the target, `max(1, ...)` returns a chunk of 1 and the buffer that chunk implies is +# `bytes_per * npts` -- over the target, by construction. The floor turned "we cannot +# meet the bound" into "here is a chunk", silently. +# +# Two rules for the tests below, both learned on this file: +# * do NOT express the expected buffer as `got * sam._ANGLE_MARG_BYTES_PER_SAMPLE_PT`. +# That reads the same constant production reads and is self-consistent for any value +# of it -- the mistake the last section of this file documents. Targets here are +# explicit literals and the peak-local slab is written out as an explicit literal. +# * the peak-local dimensions are the REVIEWER'S worked example, checked against the +# kernel rather than taken on faith: PHI_CHUNK_DEFAULT=16, n_x=256, 4 cells, +# U_NODE_STREAM_CHUNK=8 live nodes, 8 bytes -> 1048576 bytes per sample-time-point. +# --------------------------------------------------------------------------- + +import numpy as np + +#: Streamed body plus the stacked phi-scan output at the production floor: +#: 16*256*4*8*8 + 352*256*8. +PEAKLOCAL_BYTES_PER_PT = 1185792 +PEAKLOCAL_ONE_SAMPLE = 1458524160 + + +class _PeakLocalLike(object): + """peak-local at the production dimensions of the review's worked example.""" + def __init__(self, npts=1230, n_x=256, amp_sizing=None): + self.angle_marg_scheme = "peak-local" + self.data = _Data(npts) + self.x_grid = np.zeros(n_x) + self.angle_marg_info = {"amp_sizing": amp_sizing} + + +def test_the_peak_local_slab_really_is_that_big(): + """Pin the reviewer's dimension model against the kernel's own constants. + + This is the number the P1 finding rests on, and it is NOT the module's + 8192 bytes/sample-point: that constant models the DENSE (exact/laplace) path and + peak-local overrides it upward with max(). Both are right, for different schemes; + the tension in the review was between a peak-local figure and a laplace constant. + """ + from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as jp + from RIFT.likelihood.jax_ile import anglemarg + streamed = jp.PHI_CHUNK_DEFAULT * 256 * 4 * jp.U_NODE_STREAM_CHUNK * 8 + modeled = (streamed + + jp.PHI_CHUNK_DEFAULT * 256 * 8 # one chunk of (phi, distance) values + + 256 * 8 # the (n_x,) phi accumulator + + 256 * 5 * 5 * 16) # per-distance-node joint tables + # NO n_phi TERM. It used to read `streamed + n_phi * 256 * 8`, the stacked phi + # scan, and that term was the 19.97 GiB that OOMed ladder-2 rung 640. The kernel + # now reduces into its scan carry, so the whole phi axis is never live. Keeping + # this as an explicit literal makes a regrown n_phi term a failing test rather than + # a silently larger buffer. + # The retired term GROWS as sqrt(amplitude) while everything left is flat in it, so + # the saving is small at the crossover and large in production. Both ends are + # pinned, because "it got smaller" at one amplitude would not show that. + n_phi_floor = jp.required_n_phi(anglemarg.ANGLE_MARG_CROSSOVER_AMPLITUDE, m_max=2) + n_phi_prod = jp.required_n_phi(283980.0, m_max=2) # ladder-2 rung 640 + assert n_phi_floor * 256 * 8 < modeled, ( + "at the crossover the retired stacked term was the smaller half") + assert n_phi_prod * 256 * 8 > 10 * modeled, ( + "at rung 640 the retired stacked term dominated by more than 10x") + assert modeled == PEAKLOCAL_BYTES_PER_PT, ( + "the peak-local live-slab model moved: kernel constants now imply %d bytes per " + "sample-time-point, the P1 review example assumed %d" % (modeled, + PEAKLOCAL_BYTES_PER_PT)) + assert PEAKLOCAL_BYTES_PER_PT * 1230 == PEAKLOCAL_ONE_SAMPLE + + +def test_one_sample_over_the_target_is_refused_not_floored(monkeypatch): + """THE P1 REGRESSION. A 2 GiB card at the default fraction 0.5 gives a 1 GiB + allowance; one peak-local sample at production dimensions is 1.20 GiB. The old + code returned chunk 1 and therefore a 1.20 GiB buffer -- over a bound it claimed to + enforce. It must refuse.""" + _target(monkeypatch, 1 << 30) # explicit literal, not a code constant + assert PEAKLOCAL_ONE_SAMPLE > (1 << 30) # the premise, stated in literals + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_PeakLocalLike(), 4000) + + +def test_the_refusal_names_what_the_user_can_change(monkeypatch): + """A bound that fails closed with a bare assertion is a different outage from one + that says which knob to turn. Pin the actionable content, not the wording.""" + _target(monkeypatch, 1 << 30) + with pytest.raises(MemoryError) as ei: + sam.angle_marg_eval_chunk(_PeakLocalLike(), 4000) + msg = str(ei.value) + for token in ("peak-local", "npts=1230", + str(PEAKLOCAL_ONE_SAMPLE), str(PEAKLOCAL_BYTES_PER_PT), + str(1 << 30), + "RIFT_ANGLEMARG_BUFFER_FRACTION", "RIFT_ANGLEMARG_BUFFER_BYTES"): + assert token in msg, "refusal does not mention %r:\n%s" % (token, msg) + + +@pytest.mark.parametrize("npts", [1193, 1230, 4915, 32769]) +def test_no_returned_chunk_ever_exceeds_the_target(monkeypatch, npts): + """The invariant, measured against a per-point size the module does NOT own. + + 36.41 GiB at chunk 4000 / npts 1193 is XLA's own report from 2026-08-28, so this + checks the returned chunk against an EXTERNAL measurement rather than against + _ANGLE_MARG_BYTES_PER_SAMPLE_PT. Either the call refuses, or the chunk it returns + implies a buffer inside the target -- there is no third outcome, and the old floor + produced exactly that third outcome. + """ + xla_bytes_per_pt = 36.41 * GIB / (4000 * 1193) + for target in (1 << 20, 8 << 20, 1 << 30, 4 << 30, 24 << 30): + _target(monkeypatch, target) + try: + got = sam.angle_marg_eval_chunk(_Like("laplace", npts), 4000) + except MemoryError: + # refusing is allowed ONLY when one sample genuinely does not fit + assert xla_bytes_per_pt * npts > target * 1.01, ( + "refused at target %d although one sample is only ~%.0f bytes" + % (target, xla_bytes_per_pt * npts)) + continue + assert got >= 1 + implied = got * xla_bytes_per_pt * npts + assert implied <= target * 1.01, ( + "chunk %d at npts %d implies ~%.2f GiB against a %.2f GiB target" + % (got, npts, implied / GIB, target / float(GIB))) + + +def test_a_sample_that_exactly_fills_the_target_is_allowed(monkeypatch): + """The boundary, so `>` cannot quietly become `>=`. + + Exactly at the allowance the bound IS met, at a chunk of one. A refusal here would + be over-tight and would take out a configuration that fits. + """ + _target(monkeypatch, PEAKLOCAL_ONE_SAMPLE) + assert sam.angle_marg_eval_chunk(_PeakLocalLike(), 4000) == 1 + _target(monkeypatch, PEAKLOCAL_ONE_SAMPLE - 1) + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_PeakLocalLike(), 4000) + + +def test_the_dense_schemes_reach_the_refusal_too(monkeypatch): + """Not a peak-local special case: any scheme whose sample outgrows the allowance.""" + _target(monkeypatch, 1 << 20) + for scheme in ("exact", "laplace"): + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_Like(scheme, 32769), 4000) + # and the sentinel still short-circuits before any of this + assert sam.angle_marg_eval_chunk(_Like("grid", 32769), 4000) == 4000 + + +# --- the absolute allowance override, which is what makes the refusal actionable ---- +# Failing closed against _ANGLE_MARG_BUFFER_TARGET_FALLBACK would be failing closed +# against a number the file itself calls a guess with no guarantee, on exactly the +# machines whose device we could not read. RIFT_ANGLEMARG_BUFFER_FRACTION cannot help +# there -- it is a fraction of a limit that path never obtained. + +def test_no_bytes_override_means_none(): + assert sam._read_buffer_bytes({}) is None + + +@pytest.mark.parametrize("raw,expect", [("1073741824", 1 << 30), + ("2e9", 2000000000), + ("12884901888", 12 << 30)]) +def test_a_usable_bytes_override_is_honoured(raw, expect): + assert sam._read_buffer_bytes({"RIFT_ANGLEMARG_BUFFER_BYTES": raw}) == expect + + +@pytest.mark.parametrize("raw", ["", "lots", "4GiB", "0", "-1", "nan", "inf"]) +def test_an_unusable_bytes_override_is_refused_loudly(raw): + with pytest.raises(ValueError): + sam._read_buffer_bytes({"RIFT_ANGLEMARG_BUFFER_BYTES": raw}) + + +def test_the_bytes_override_beats_the_device_probe(monkeypatch): + """It has to win over the probe, or it cannot rescue a machine the probe misreads. + + The fake card is idle, so the probe would otherwise answer 12 GiB: the 3 GiB below + is the override winning, not the fallback coinciding with it. + """ + _fake_jax(monkeypatch, devices=[_idle_gpu(24 * GIB)]) + monkeypatch.setenv("RIFT_ANGLEMARG_BUFFER_BYTES", str(3 * GIB)) + assert sam._angle_marg_buffer_target() == 3 * GIB + + +def test_the_bytes_override_beats_the_fallback_and_lifts_a_refusal(monkeypatch): + """The case the knob exists for: no readable device, and the 4 GiB guess refuses a + configuration the operator knows their machine can hold.""" + _fake_jax(monkeypatch, raises=RuntimeError("no device")) + big = _PeakLocalLike(npts=8192) # 9.05 GiB per sample, over the 4 GiB guess + one_sample = PEAKLOCAL_BYTES_PER_PT * 8192 + assert one_sample > 4 * GIB, "premise: the guess must refuse this" + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(big, 4000) + monkeypatch.setenv("RIFT_ANGLEMARG_BUFFER_BYTES", str(32 * GIB)) + assert sam.angle_marg_eval_chunk(big, 4000) == (32 * GIB) // one_sample == 3 + + +# --------------------------------------------------------------------------- +# MUTATION SWEEP of the section above (2026-09-05, 57 collected). Each mutation was +# applied to a pristine copy of samplers.py, verified present in the FILE ON DISK +# before running -- a replacement that changes no bytes reports as a surviving guard +# and is a harness bug, not a result -- and reverted afterwards. +# +# restore the pre-fix `cap = max(1, target // per_sample)` 9 failed KILLED +# `>` -> `>=` in the refusal 1 failed KILLED +# drop the peak-local slab model (use the 8192 constant) 4 failed KILLED +# make RIFT_ANGLEMARG_BUFFER_BYTES inert 13 failed KILLED +# read that override inside the probe's blanket except 1 failed KILLED +# strip the override names out of the refusal message 1 failed KILLED +# accept a zero/negative absolute allowance 2 failed KILLED +# put max(1, ...) back AROUND the surviving division 0 failed SURVIVED +# +# The survivor is an EQUIVALENT mutant, and it is recorded rather than chased: the +# refusal above guarantees `per_sample <= target` on every path that reaches the +# division, so `target // per_sample` is already >= 1 and the floor cannot change any +# value. It is the floor REPLACING the refusal (the first row) that was the defect, +# not the floor as such. No test can distinguish an unreachable branch, and writing +# one that appeared to would mean the refusal had a hole. +# --------------------------------------------------------------------------- + + +def test_a_malformed_bytes_override_is_not_swallowed_by_the_probe(monkeypatch): + """It is read OUTSIDE the probe's blanket `except Exception` on purpose: inside it, + a typo would be silently replaced by the 4 GiB fallback and the operator would never + learn their override did nothing.""" + _fake_jax(monkeypatch, devices=[_idle_gpu(24 * GIB)]) + monkeypatch.setenv("RIFT_ANGLEMARG_BUFFER_BYTES", "24GiB") + with pytest.raises(ValueError): + sam._angle_marg_buffer_target() + + +# --- forced probe allocation, review MAJOR (2026-09-08) --------------------- +# The tests above all fake `jax.devices()` but never touch `jax.device_put` or +# `jax.numpy`, so `_probe_allocate` fails with AttributeError against every fake +# device above and is silently swallowed -- which is exactly why none of those +# tests needed to change for this fix. The tests below drive the post-probe LOGIC +# directly: `_probe_allocate` is monkeypatched to a no-op and the fake device's +# `memory_stats()` is set to the dict a real probe would have produced, so a real +# GPU is never touched. +# +# Without a forced allocation, `_angle_marg_buffer_target()` on jax 0.9.2 read +# `pool_bytes=0` before this process's first allocation and returned the blind +# 4 GiB fallback even on a busy shared card with ~0.5 GiB truly free -- 8x too +# much (review MAJOR). The fix forces one tiny allocation first so the pool +# signal exists, then trusts `pool - bytes_in_use` only when the resulting pool is +# not small next to `bytes_limit` (the on-demand-allocator ambiguity, handled +# separately below). + +def _mock_probe(monkeypatch): + """Replace the real device-allocation probe with a no-op, so these tests + exercise the post-probe branch in `_angle_marg_buffer_target` without ever + calling into a real jax runtime.""" + monkeypatch.setattr(sam, "_probe_allocate", lambda jax_module, dev: None) + + +def test_probe_allocation_is_attempted_before_reading_stats(monkeypatch): + """The probe call itself must run: it is what makes the pool signal legible on + jax 0.9.2, where both `pool_bytes` and `largest_free_block_bytes` read 0 before + this process's first allocation (see JAX092_BEFORE_ALLOC above).""" + calls = [] + monkeypatch.setattr(sam, "_probe_allocate", + lambda jax_module, dev: calls.append(dev)) + dev = _Dev("gpu", 24 * GIB, free=0) + _fake_jax(monkeypatch, devices=[dev]) + sam._angle_marg_buffer_target() + assert calls == [dev] + + +def test_a_successful_probe_on_a_busy_card_is_trusted_not_blind(monkeypatch): + """THE MAJOR REVIEW FINDING, fixed. A PLAUSIBLE (not independently measured -- + see PR #285 reply) post-probe dict for a busy shared card under jax 0.9.2's + default preallocating allocator: the forced allocation only succeeds because + ~384 MiB genuinely was free, so `pool_bytes` reports memory this process + actually holds and `bytes_limit` (reported once the pool exists) tracks it, so + the pool is not "small" by the half-of-limit rule below. Before this fix, the + same before-probe state (`pool_bytes=0`) returned the blind 4 GiB fallback -- + 8x the true free memory.""" + _mock_probe(monkeypatch) + stats = {"bytes_limit": 402653184, "bytes_in_use": 4096, + "pool_bytes": 402653184, "largest_free_block_bytes": 0} + _fake_jax(monkeypatch, devices=[_StatsDev("gpu", stats)]) + got = sam._angle_marg_buffer_target() + assert got == int((402653184 - 4096) * sam._ANGLE_MARG_BUFFER_FRACTION) + assert got < 4 * GIB, "must not fall back to the blind guess once probed" + + +def test_a_small_pool_next_to_a_large_limit_is_unknown_not_free(monkeypatch): + """The on-demand allocator (`XLA_PYTHON_CLIENT_PREALLOCATE=false`): the pool + grows only to fit what has actually been requested, so after the forced probe + it can stay tiny next to a `bytes_limit` that reports the device's full + capacity regardless. `pool - bytes_in_use` would read as ~fully free here, + which is not a real measurement -- bound the blind guess by `bytes_limit - + bytes_in_use` instead of trusting it. Could not measure ldas-pcdev11 GPU 3 for + this PR (thread count 431 >= the 380 dispatch ceiling both times checked); this + dict is PLAUSIBLE, labelled as such, not measured.""" + _mock_probe(monkeypatch) + stats = {"bytes_limit": 18895355904, "bytes_in_use": 4096, + "pool_bytes": 2097152, "largest_free_block_bytes": 0} + _fake_jax(monkeypatch, devices=[_StatsDev("gpu", stats)]) + got = sam._angle_marg_buffer_target() + assert got == min(4 * GIB, 18895355904 - 4096) + + +def test_missing_bytes_in_use_is_unknown_not_free(monkeypatch): + """Review MINOR: a pool reported with no occupancy figure at all must not read + as fully free, through either entry point.""" + assert sam._device_available_bytes({"pool_bytes": 16 * GIB}) is None + _mock_probe(monkeypatch) + stats = {"bytes_limit": 24 * GIB, "pool_bytes": 16 * GIB} + _fake_jax(monkeypatch, devices=[_StatsDev("gpu", stats)]) + assert sam._angle_marg_buffer_target() == 4 * GIB + + +def test_a_full_pool_still_refuses_after_the_probe_fix(monkeypatch): + """Design item 4: the genuinely-full-pool refusal must survive the probe and + the new small-pool/large-limit branch must not swallow it -- `pool == limit` + here, so the half-of-limit rule does not divert it.""" + _mock_probe(monkeypatch) + stats = {"bytes_limit": 24 * GIB, "bytes_in_use": 24 * GIB, + "pool_bytes": 24 * GIB, "largest_free_block_bytes": 0} + _fake_jax(monkeypatch, devices=[_StatsDev("gpu", stats)]) + assert sam._angle_marg_buffer_target() == 0 + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_planner.py b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_planner.py new file mode 100644 index 000000000..edc57319f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_planner.py @@ -0,0 +1,462 @@ +"""Focused policy tests for the opt-in direct-marginalization planner. + +The amplitude ladder below is a synthetic calibration packet. The planner is +being tested, not a new accuracy claim for the shipped angle kernels: production +offers must bring their own measured resource and error provenance. +""" + +import json +import math + +import pytest + +from RIFT.likelihood.jax_ile import direct_marginalization_planner as P + + +def _certified_warrant(scope="synthetic finite spectrum"): + return P.Warrant(P.WarrantKind.EXACT_TRIG_DEGREE, scope, True, + "test fixture: analytic finite-spectrum bound") + + +def _offer(axis, scheme, error, compute, memory=64, *, + evidence=P.EvidenceKind.CERTIFIED, warrant=None, + requires=(), conflicts=(), conditional_requirements=()): + if warrant is None: + warrant = _certified_warrant() + accuracy = P.AccuracyAssessment( + evidence, error, + "test fixture: error envelope for %s:%s" % (axis, scheme)) + resources = P.ResourceEstimate( + compute, memory, + "test fixture: common-unit cost model for %s:%s" % (axis, scheme)) + return P.SchemeOffer( + axis, scheme, accuracy, resources, warrant, + "test fixture offer", requires=frozenset(requires), + conflicts=frozenset(conflicts), + conditional_requirements=tuple(conditional_requirements)) + + +def _amplitude_offers(amplitude): + """Synthetic measured envelopes with distinct accuracy and cost crossings.""" + amplitude = float(amplitude) + return ( + _offer("angle", "exact", error=1e-8, + compute=5.0 + amplitude / 50.0), + _offer("angle", "laplace", error=30.0 / amplitude ** 2, + compute=40.0 + math.sqrt(amplitude)), + ) + + +@pytest.mark.parametrize( + "amplitude, expected", + [(25.0, "exact"), (400.0, "exact"), (40000.0, "laplace")]) +def test_low_moderate_high_amplitude_choose_cheapest_certified( + amplitude, expected): + """Accuracy gates low A; measured cost, not one crossover, orders the rest.""" + decision = P.plan_direct_marginalization( + _amplitude_offers(amplitude), {"angle": 1e-2}, + P.ResourceBudget(2000.0, 1024), required_axes=("angle",)) + assert decision.action == "run" + assert decision.basis == "cheapest-certified" + assert decision.certified is True + assert decision.require_selection()[0].scheme == expected + + +def test_combination_resource_model_controls_nested_kernel_cost(): + """A measured whole-kernel model can override the additive safe default.""" + offers = ( + _offer("angle", "exact", error=1e-5, compute=1), + _offer("angle", "laplace", error=1e-5, compute=100), + _offer("distance", "uniform", error=1e-5, compute=1), + ) + + def nested_cost(combination): + angle = next(o.scheme for o in combination if o.axis == "angle") + return P.ResourceEstimate( + 10 if angle == "laplace" else 100, 50, + "fixture: measured complete nested-kernel cost") + + decision = P.plan_direct_marginalization( + offers, {"angle": 1e-3, "distance": 1e-3}, + P.ResourceBudget(200, 100), + required_axes=("angle", "distance"), + resource_model=nested_cost) + selected = {offer.axis: offer.scheme + for offer in decision.require_selection()} + assert selected == {"angle": "laplace", "distance": "uniform"} + assert "complete nested-kernel" in decision.resource_use.provenance + + +@pytest.mark.parametrize( + "error_budget, resource_budget, reason_code", + [ + (None, {"max_compute_units": 100, "max_memory_bytes": 100}, + "missing-error-budget"), + ({}, {"max_compute_units": 100, "max_memory_bytes": 100}, + "missing-error-budget"), + ({"angle": 0.1}, None, "missing-resource-budget"), + ({"angle": 0.1}, {"max_compute_units": 100}, + "missing-resource-budget"), + ]) +def test_missing_budget_declines_with_no_selection( + error_budget, resource_budget, reason_code): + decision = P.plan_direct_marginalization( + (_offer("angle", "exact", 1e-3, 10),), + error_budget, resource_budget, required_axes=("angle",)) + assert decision.action == "decline" + assert decision.reason_code == reason_code + assert decision.selected == () + with pytest.raises(P.MarginalizationPlanDeclined, match=reason_code): + decision.require_selection() + + +def test_repeated_required_axis_declines_instead_of_planning_it_twice(): + """One scheme per axis: a repeated axis is a malformed request, not a plan.""" + decision = P.plan_direct_marginalization( + (_offer("angle", "exact", 1e-5, 10),), {"angle": 1e-2}, + P.ResourceBudget(1000.0, 1024), required_axes=("angle", "angle")) + assert decision.action == "decline" + assert decision.reason_code == "duplicate-axis" + assert decision.selected == () + assert decision.resource_use is None + assert decision.ledger["details"]["duplicate_axes"] == ["angle"] + assert decision.ledger["combinations"] == [] + with pytest.raises(P.MarginalizationPlanDeclined, match="duplicate-axis"): + decision.require_selection() + json.dumps(decision.as_dict()) + + +def test_shipped_peak_local_plus_gh_is_an_unsupported_combination(): + """The real JAX profile declares this once; the planner refuses the pair.""" + def validated(label): + return P.AccuracyAssessment( + P.EvidenceKind.VALIDATED, 1e-4, + "fixture validation: " + label) + + def resources(label): + return P.ResourceEstimate(10.0, 10, "fixture cost: " + label) + + offers = ( + P.make_jax_scheme_offer( + "angle", "peak-local", validated("angle"), resources("angle"), + provenance="fixture request"), + P.make_jax_scheme_offer( + "distance", "gh", validated("distance"), resources("distance"), + provenance="fixture request"), + ) + decision = P.plan_jax_direct_marginalization( + offers, {"angle": 1e-3, "distance": 1e-3}, + P.ResourceBudget(100.0, 100), + required_axes=("angle", "distance"), + capabilities=("angle-amplitude-estimate", + "angle-peak-local-warranted", + "distance-volumetric-prior"), + allow_best_effort=True) + assert decision.action == "decline" + assert decision.reason_code == "no-compatible-plan" + records = decision.ledger["combinations"] + assert len(records) == 1 + assert any("angle:peak-local conflicts" in reason + for reason in records[0]["compatibility_reasons"]) + + +def test_conditional_gh_laplace_warrant_must_be_supplied(): + """GH+Laplace is supported only after the concrete identity predicate passes.""" + validated = P.AccuracyAssessment( + P.EvidenceKind.VALIDATED, 1e-4, "fixture validation") + resources = P.ResourceEstimate(10.0, 10, "fixture cost") + offers = ( + P.make_jax_scheme_offer("angle", "laplace", validated, resources, + provenance="fixture request"), + P.make_jax_scheme_offer("distance", "gh", validated, resources, + provenance="fixture request"), + ) + base_capabilities = ("angle-amplitude-estimate", + "distance-volumetric-prior") + refused = P.plan_jax_direct_marginalization( + offers, {"angle": 1e-3, "distance": 1e-3}, + P.ResourceBudget(100.0, 100), + required_axes=("angle", "distance"), + capabilities=base_capabilities, allow_best_effort=True) + assert refused.action == "decline" + assert "gh-laplace-supported" in str(refused.ledger["combinations"]) + + allowed = P.plan_jax_direct_marginalization( + offers, {"angle": 1e-3, "distance": 1e-3}, + P.ResourceBudget(100.0, 100), + required_axes=("angle", "distance"), + capabilities=base_capabilities + ("gh-laplace-supported",), + allow_best_effort=True) + assert allowed.action == "run" + assert allowed.basis == "most-accurate-affordable" + + +def test_jax_direct_path_injects_the_nonlinear_time_incompatibility(): + """Callers cannot omit the wrapper fact that currently excludes bandlimited.""" + validated = P.AccuracyAssessment( + P.EvidenceKind.VALIDATED, 1e-4, "fixture validation") + resources = P.ResourceEstimate(10.0, 10, "fixture cost") + offers = ( + P.make_jax_scheme_offer("angle", "exact", validated, resources, + provenance="fixture request"), + P.make_jax_scheme_offer("distance", "uniform", validated, resources, + provenance="fixture request"), + P.make_jax_scheme_offer("time", "bandlimited", validated, + resources, provenance="fixture request"), + ) + decision = P.plan_jax_direct_marginalization( + offers, {"angle": 1e-3, "distance": 1e-3, "time": 1e-3}, + P.ResourceBudget(100.0, 100), + capabilities=("angle-amplitude-estimate", "time-exact-band-limit", + "time-independent-rho-sq", "n-cal-one"), + allow_best_effort=True) + assert decision.action == "decline" + assert decision.reason_code == "no-compatible-plan" + assert "jax-direct-nonlinear-time" in decision.ledger["capabilities"] + + +def test_no_silent_fallback_and_best_effort_requires_explicit_authority(): + """An affordable estimate is a suggestion, never an implicit replacement.""" + exact = _offer("angle", "exact", error=1e-5, compute=200, memory=20) + empirical_warrant = P.Warrant( + P.WarrantKind.EMPIRICAL_CALIBRATION, "measured envelope", False, + "test fixture: empirical campaign") + approximate = _offer( + "angle", "approximate", error=2e-2, compute=10, memory=10, + evidence=P.EvidenceKind.VALIDATED, warrant=empirical_warrant) + budget = P.ResourceBudget(100, 100) + + strict = P.plan_direct_marginalization( + (exact, approximate), {"angle": 1e-2}, budget, + required_axes=("angle",)) + assert strict.action == "decline" + assert strict.reason_code == "resource-budget-exceeded" + assert strict.selected == () + assert [offer.scheme for offer in strict.suggested] == ["approximate"] + assert strict.meets_error_budget is False + with pytest.raises(P.MarginalizationPlanDeclined): + strict.require_selection() + + explicit = P.plan_direct_marginalization( + (exact, approximate), {"angle": 1e-2}, budget, + required_axes=("angle",), allow_best_effort=True) + assert explicit.action == "run" + assert explicit.basis == "most-accurate-affordable" + assert explicit.certified is False + assert explicit.meets_error_budget is False + assert explicit.require_selection()[0].scheme == "approximate" + record = explicit.as_dict() + assert record["selected"][0]["accuracy"]["provenance"] + assert record["selected"][0]["warrant"]["provenance"] + assert record["selected"][0]["resources"]["provenance"] + json.dumps(record) + + +def test_resource_decline_uses_explicit_reserve_without_dropping_sample(): + """A primary resource refusal remains recorded when dense exact is used.""" + dense_exact = _offer( + "angle", "dense-exact", error=1e-6, compute=200, memory=20) + approximate_warrant = P.Warrant( + P.WarrantKind.EMPIRICAL_CALIBRATION, "measured approximation", False, + "test fixture: empirical envelope") + preferred = _offer( + "angle", "shortcut", error=1e-3, compute=5, memory=5, + evidence=P.EvidenceKind.VALIDATED, + warrant=approximate_warrant) + decision = P.plan_direct_marginalization( + (dense_exact, preferred), {"angle": 1e-2}, + P.ResourceBudget(20, 100), required_axes=("angle",)) + assert decision.action == "decline" + assert decision.reason_code == "resource-budget-exceeded" + + fallback = P.ConservativeFallbackPolicy( + (dense_exact,), P.ResourceBudget(250, 100), + provenance="fixture: reserve-budget policy", + finite_output_contract="fixture: full finite angle grid") + resolution = P.resolve_plan_for_production(decision, fallback) + + assert resolution.action is P.ResolutionAction.USE_CONSERVATIVE_FALLBACK + assert resolution.require_selection()[0].scheme == "dense-exact" + assert resolution.drops_sample is False + assert resolution.waveform_failure is None + assert resolution.method_decline.code == "resource-budget-exceeded" + assert resolution.certified is True + assert resolution.ledger["fallback_policy"]["provenance"] + assert (resolution.ledger["preferred_decision"]["reason_code"] + == "resource-budget-exceeded") + json.dumps(resolution.as_dict()) + + +def test_uncertified_jax_plan_resolves_to_registered_dense_fallback(): + """Cannot certify preferred is a method result, not an invalid waveform.""" + validated = P.AccuracyAssessment( + P.EvidenceKind.VALIDATED, 1e-4, "fixture validation") + resources = P.ResourceEstimate(10.0, 10, "fixture cost") + peak_local = P.make_jax_scheme_offer( + "angle", "peak-local", validated, resources, + provenance="fixture preferred request") + dense_exact = P.make_jax_scheme_offer( + "angle", "exact", validated, + P.ResourceEstimate(50.0, 20, "fixture dense fallback cost"), + provenance="fixture fallback request") + decision = P.plan_jax_direct_marginalization( + (peak_local,), {"angle": 1e-3}, P.ResourceBudget(100, 100), + required_axes=("angle",), + capabilities=("angle-amplitude-estimate", + "angle-peak-local-warranted")) + assert decision.action == "decline" + assert decision.reason_code == "no-certified-plan" + fallback = P.make_jax_production_fallback_policy( + (dense_exact,), P.ResourceBudget(100, 100), + provenance="fixture: dense JAX reserve", + finite_output_contract="fixture: dense phi and psi cover full support") + + resolution = P.resolve_plan_for_production(decision, fallback) + + assert resolution.action is P.ResolutionAction.USE_CONSERVATIVE_FALLBACK + assert resolution.require_selection()[0].scheme == "exact" + assert resolution.certified is False + assert resolution.meets_error_budget is True + assert resolution.drops_sample is False + assert resolution.method_decline.code == "no-certified-plan" + assert resolution.waveform_failure is None + + +def test_incomplete_root_enumeration_replaces_method_not_likelihood_point(): + """Runtime root refusal switches to dense exact and retains the sample.""" + validated = P.AccuracyAssessment( + P.EvidenceKind.VALIDATED, 1e-4, "fixture validation") + peak_local = P.make_jax_scheme_offer( + "angle", "peak-local", validated, + P.ResourceEstimate(10.0, 10, "fixture shortcut cost"), + provenance="fixture preferred request") + dense_exact = P.make_jax_scheme_offer( + "angle", "exact", validated, + P.ResourceEstimate(50.0, 20, "fixture dense fallback cost"), + provenance="fixture fallback request") + decision = P.plan_jax_direct_marginalization( + (peak_local,), {"angle": 1e-3}, P.ResourceBudget(100, 100), + required_axes=("angle",), + capabilities=("angle-amplitude-estimate", + "angle-peak-local-warranted"), + allow_best_effort=True) + assert decision.action == "run" + fallback = P.make_jax_production_fallback_policy( + (dense_exact,), P.ResourceBudget(100, 100), + provenance="fixture: dense JAX reserve", + finite_output_contract="fixture: dense phi and psi cover full support") + root_decline = P.MethodDecline( + "incomplete-root-enumeration", + "stationary-root completeness check did not close", + "fixture: root enumeration postcondition", axis="angle", + stage="runtime-enumeration", ledger={"roots_found": 3}) + + resolution = P.resolve_plan_for_production( + decision, fallback, method_decline=root_decline) + + assert resolution.action is P.ResolutionAction.USE_CONSERVATIVE_FALLBACK + assert resolution.require_selection()[0].scheme == "exact" + assert resolution.method_decline.ledger == {"roots_found": 3} + assert resolution.waveform_failure is None + assert resolution.drops_sample is False + assert "incomplete-root-enumeration" in str(resolution.as_dict()) + + +def test_axisless_runtime_decline_requires_full_plan_replacement(): + """An unknown declined axis cannot leave any preferred method in service.""" + angle_fast = _offer("angle", "root-shortcut", 1e-5, 5) + time_fast = _offer("time", "time-shortcut", 1e-5, 5) + decision = P.plan_direct_marginalization( + (angle_fast, time_fast), {"angle": 1e-3, "time": 1e-3}, + P.ResourceBudget(100, 256), required_axes=("angle", "time")) + assert decision.action == "run" + decline = P.MethodDecline( + "runtime-warrant-lost", "runtime check did not identify its axis", + "fixture: axis-less runtime callback") + time_only = P.ConservativeFallbackPolicy( + (_offer("time", "simpson", 1e-6, 10),), + P.ResourceBudget(100, 256), "fixture: partial reserve", + "fixture: finite time support") + + with pytest.raises(P.FallbackConfigurationError, + match="does not replace declined axes.*angle"): + P.resolve_plan_for_production( + decision, time_only, method_decline=decline) + + complete = P.ConservativeFallbackPolicy( + (_offer("angle", "dense", 1e-6, 20), + _offer("time", "simpson", 1e-6, 10)), + P.ResourceBudget(100, 256), "fixture: complete reserve", + "fixture: finite full-axis support") + resolution = P.resolve_plan_for_production( + decision, complete, method_decline=decline) + + assert [offer.key for offer in resolution.require_selection()] == [ + "angle:dense", "time:simpson"] + assert resolution.drops_sample is False + assert resolution.method_decline is decline + assert resolution.waveform_failure is None + + +def test_method_decline_without_fallback_is_configuration_error_not_drop(): + preferred = _offer("angle", "shortcut", 1e-5, 5) + decision = P.plan_direct_marginalization( + (preferred,), {"angle": 1e-3}, P.ResourceBudget(100, 100), + required_axes=("angle",)) + decline = P.MethodDecline( + "incomplete-root-enumeration", "root postcondition failed", + "fixture: runtime postcondition", axis="angle") + with pytest.raises(P.FallbackConfigurationError, + match="not a waveform failure"): + P.resolve_plan_for_production(decision, method_decline=decline) + + +def test_only_explicit_waveform_failure_can_drop_sample(): + preferred = _offer("angle", "dense", 1e-5, 5) + decision = P.plan_direct_marginalization( + (preferred,), {"angle": 1e-3}, P.ResourceBudget(100, 100), + required_axes=("angle",)) + failure = P.WaveformFailure( + "waveform-generation-failed", "base waveform contains non-finite data", + "fixture: waveform validation", ledger={"finite": False}) + + resolution = P.resolve_plan_for_production( + decision, waveform_failure=failure) + + assert resolution.action is P.ResolutionAction.WAVEFORM_FAILURE + assert resolution.drops_sample is True + assert resolution.selected == () + assert resolution.method_decline is None + assert resolution.waveform_failure is failure + with pytest.raises(P.WaveformLikelihoodFailure, + match="waveform-generation-failed"): + resolution.require_selection() + + +def test_current_angle_profiles_cannot_be_mislabeled_certified(): + """Exact coefficients do not certify the amplitude-sized exp quadrature.""" + accuracy = P.AccuracyAssessment( + P.EvidenceKind.CERTIFIED, 1e-8, "invalid fixture claim") + resources = P.ResourceEstimate(1.0, 1, "fixture cost") + with pytest.raises(ValueError, match="no implemented certificate"): + P.make_jax_scheme_offer( + "angle", "exact", accuracy, resources, + provenance="attempted invalid offer") + + +def test_bandlimited_time_profile_cannot_be_mislabeled_certified(): + """A derived-and-remeasured refinement factor is not a per-request bound.""" + accuracy = P.AccuracyAssessment( + P.EvidenceKind.CERTIFIED, 1e-12, "invalid fixture claim") + resources = P.ResourceEstimate(1.0, 1, "fixture cost") + with pytest.raises(ValueError, match="no implemented certificate"): + P.make_jax_scheme_offer( + "time", "bandlimited", accuracy, resources, + provenance="attempted invalid offer") + + +def test_no_shipped_profile_advertises_a_certificate_yet(): + """cheapest-certified stays unreachable until some rule implements a bound.""" + advertised = sorted(key for key, profile in P.JAX_SCHEME_PROFILES.items() + if profile.warrant.certificate_available) + assert advertised == [] diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy.py b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy.py new file mode 100644 index 000000000..ae187262c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy.py @@ -0,0 +1,810 @@ +"""The opt-in cross-axis policy, as WIRED into the phi/psi-marginalized likelihood. + +These test the wiring: that the policy reaches the wrapper and the CLI, refuses +what it cannot honour, converts the local measure to the production convention, +executes a warranted band-limited reserve on decline, keeps every sample, and +records its ledger. The controller's own numerics are tested in +test_all_axis_peaklocal.py. Nothing here validates derivatives. +""" +import os +import subprocess +import sys +import types + +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 import core as _core +from RIFT.likelihood.jax_ile import direct_marginalization_policy as DP +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiPsiMargLikelihood +from test_all_axis_peaklocal import _problem +from test_angle_marg_exact import make_synth, RA, DEC, INCL, INTERP + + +# ------------------------------------------------------------------ fixtures + +def _guarded_problem(n, guard, scale=1.0): + """The analytic three-harmonic table of test_all_axis_peaklocal with + ``guard`` primitive-only support samples at each end (the cosine is + continued analytically, so the guarded reconstruction has a true target).""" + C_A, C_B, constants = _problem(n) + support = np.arange(-guard, n + guard, dtype=float) + guarded = np.zeros(C_A.shape[:-1] + (support.size,), dtype=np.complex128) + guarded[0, 1] = (constants["k0"] + - constants["kt"] * np.cos(2.0 * np.pi * support + / constants["span"])) + guarded[2, 1] = 0.5 * constants["kp"] + guarded[0, 0] = 0.5 * constants["ku"] + guarded[0, 2] = 0.5 * constants["ku"] + np.testing.assert_allclose(guarded[..., guard:-guard], C_A, atol=1e-14) + return scale * guarded, scale * scale * C_B, constants + + +def _fake_data(n, deltaT=1.0 / 4096.0): + return types.SimpleNamespace( + npts=n, deltaT=deltaT, distMpcRef=_core.DIST_MPC_REF, + w_t=jnp.asarray(_core._simpson_weights(n, deltaT)), + lms=np.asarray([[2, 2], [2, -2]])) + + +def _install_tables(monkeypatch, tables, guard): + """Route angle_coefficient_tables to fixed analytic tables, one per row.""" + C_A_rows, C_B = tables + seen = [] + + def fake(data, ra, dec, incl, interp=None, sample_chunk=None, guard=0): + seen.append(int(guard)) + S = int(np.asarray(ra).shape[0]) + assert S == len(C_A_rows) + C_A = jnp.asarray(np.stack(C_A_rows, axis=2)) # (KP,KS,S,Nt) + C_Bb = jnp.broadcast_to( + jnp.asarray(C_B)[:, :, None, None], + C_B.shape + (S, C_A.shape[-1])) + return C_A, C_Bb, dict(m_max=2, nphi_s=9, npsi_s=5, guard=guard, + ntime=C_A.shape[-1]) + + monkeypatch.setattr(AM, "angle_coefficient_tables", fake) + return seen + + +def _production_reference(C_A_target, C_B, data, x_grid, log_w, amp_sizing): + """What the exact scheme returns for these tables: exact angles on the + native window, production distance weights, production Simpson time.""" + lnL_t = AM.coefficient_table_distphipsimarg_exact( + C_A_target, C_B, x_grid, log_w, amp_sizing=amp_sizing, + dense_chunk=8, grid_block=32) + return float(_core._time_marginalize(lnL_t, data.w_t)[0]) + + +def _fine_reference(constants, C_B, data, x_grid, log_w, amp_sizing, + refine=16, scale=1.0): + """Independent time reference: the ANALYTIC table evaluated on a + ``refine``-times finer grid (no reflected primitive involved), exact + angles, production distance weights, Simpson in seconds.""" + n = int(data.npts) + t = np.arange((n - 1) * refine + 1, dtype=float) / float(refine) + C_A = np.zeros((3, 3, t.size), dtype=np.complex128) + C_A[0, 1] = (constants["k0"] + - constants["kt"] * np.cos(2.0 * np.pi * t / constants["span"])) + C_A[2, 1] = 0.5 * constants["kp"] + C_A[0, 0] = 0.5 * constants["ku"] + C_A[0, 2] = 0.5 * constants["ku"] + lnL_t = AM.coefficient_table_distphipsimarg_exact( + scale * C_A, C_B, x_grid, log_w, amp_sizing=amp_sizing, + dense_chunk=8, grid_block=32) + w = jnp.asarray(_core._simpson_weights(t.size, data.deltaT / refine)) + return float(_core._time_marginalize(lnL_t, w)[0]) + + +_GUARD = 8 +_N = 33 +_X_RANGE = (0.2, 7.0) + + +def _grid(n=1024): + d_min = _core.DIST_MPC_REF / _X_RANGE[1] + d_max = _core.DIST_MPC_REF / _X_RANGE[0] + return _core.make_distance_grid(d_min, d_max, n, distMpcRef=_core.DIST_MPC_REF) + + +# ------------------------------------------------------- choices and refusal + +def test_policy_is_opt_in_and_reaches_the_wrapper(): + assert DP.POLICY_DEFAULT == "off" + assert "auto" in DP.POLICY_CHOICES + data = make_synth(scale=2.0) + like = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="exact") + assert like.direct_marginalization_policy == "off" + assert like.policy_info is None + assert like._batched_ledger is None + with pytest.raises(ValueError, match="direct_marginalization_policy"): + JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="exact", + direct_marginalization_policy="autp") + + +@pytest.mark.parametrize("kw, needle", [ + (dict(angle_marg_scheme="laplace", time_quadrature="simpson", + d_prior="euclidean", dist_grid="uniform"), "exact-angle reserve"), + (dict(angle_marg_scheme="peak-local", time_quadrature="simpson", + d_prior="euclidean", dist_grid="uniform"), "exact-angle reserve"), + (dict(angle_marg_scheme="exact", time_quadrature="bandlimited", + d_prior="euclidean", dist_grid="uniform"), "time integral"), + (dict(angle_marg_scheme="exact", time_quadrature="simpson", + d_prior="uniform", dist_grid="uniform"), "volumetric"), + (dict(angle_marg_scheme="exact", time_quadrature="simpson", + d_prior="euclidean", dist_grid="loguniform"), "distance-grid-scheme"), +]) +def test_policy_refuses_what_it_cannot_compose(kw, needle): + with pytest.raises(ValueError, match=needle): + DP.validate_policy_request("auto", **kw) + DP.validate_policy_request("off", **kw) # off composes with anything + + +def test_wrapper_refuses_laplace_reserve_and_lnLt(): + data = make_synth(scale=2.0) + with pytest.raises(ValueError, match="exact-angle reserve"): + JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="laplace", + direct_marginalization_policy="auto") + with pytest.raises(ValueError, match="volumetric"): + JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, nphi=32, npsi=8, + interp=INTERP, angle_marg="exact", + d_prior="uniform", + direct_marginalization_policy="auto") + + +# ------------------------------------------------ measure conversion (accept) + +@pytest.mark.parametrize("gh_nodes", [0, 64]) +def test_accepted_local_value_lands_in_the_production_convention( + monkeypatch, gh_nodes): + """The whole point of the normalization constant: an ACCEPTED local row + must equal what the exact scheme returns for the same tables, in the + reserve's units (angles averaged, distance prior normalized, time in + seconds), on both distance paths.""" + monkeypatch.setattr(_core, "_DISTMARG_GH_N", gh_nodes) + guarded, C_B, constants = _guarded_problem(_N, _GUARD) + rows = [guarded, 1.01 * guarded] + seen = _install_tables(monkeypatch, (rows, C_B), _GUARD) + data = _fake_data(_N) + x_grid, log_w = _grid() + cfg = DP.PolicyConfig(time_guard=_GUARD, reserve_time_refine=4) + lln, info = DP.policy_log_normalization(data, x_grid, log_w) + assert info["distance_mode"] == ( + "gh-volumetric" if gh_nodes else "fixed-grid-volumetric") + assert info["time_weight_scale"] == pytest.approx(1.0) + + value, ledger = DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(2), jnp.zeros(2), jnp.zeros(2), x_grid, log_w, + interp=INTERP, amp_sizing=40.0, config=cfg, return_ledger=True) + assert seen == [_GUARD] # the guard reached the tables + assert np.all(np.asarray(ledger["accepted_local"])), { + k: np.asarray(v) for k, v in ledger.items() if k.startswith("decline")} + assert np.all(np.asarray(ledger["usable"])) + assert not np.any(np.asarray(ledger["reserve_executed"])) + assert np.all(np.asarray(ledger["norm_time_invariant"])) + assert np.all(np.asarray(ledger["reconciles"])) + for i, (table, row_scale) in enumerate(zip(rows, (1.0, 1.01))): + fine = _fine_reference(constants, C_B, data, x_grid, log_w, 40.0, + scale=row_scale) + native = _production_reference( + table[..., _GUARD:-_GUARD], C_B, data, x_grid, log_w, 40.0) + # 1e-3 nat is the controller's own budget; the rest is the fixed + # distance grid's quadrature error, which the 1024-point grid keeps + # below 1e-4 on this broad fixture. The native Simpson rule is NOT + # the reference: this fixture's time peak is ~0.7 samples wide and + # the production rule misses it by ~0.03 nat, which is the error the + # composite exists to remove. + assert abs(float(value[i]) - fine) < 2.0e-3, ( + gh_nodes, i, float(value[i]), fine, native) + assert abs(native - fine) > 5.0 * abs(float(value[i]) - fine), ( + "the fixture no longer distinguishes the composite from the " + "native rule", native, fine, float(value[i])) + + +def test_normalization_constant_is_derived_not_guessed(monkeypatch): + data = _fake_data(_N) + x_grid, log_w = _grid() + monkeypatch.setattr(_core, "_DISTMARG_GH_N", 0) + total, info = DP.policy_log_normalization(data, x_grid, log_w) + d = _core.DIST_MPC_REF / np.asarray(x_grid) + norm = np.sum(d ** 2) * abs(d[1] - d[0]) + expected = (-2.0 * np.log(2.0 * np.pi) + np.log(data.deltaT) + + 3.0 * np.log(_core.DIST_MPC_REF) - np.log(norm)) + assert total == pytest.approx(expected, rel=1e-12) + with pytest.raises(ValueError, match="volumetric"): + DP.policy_log_normalization(data, x_grid, log_w, d_prior="uniform") + # a non-uniform-in-d grid cannot supply the constant and must say so + x_lu = jnp.asarray(np.geomspace(0.2, 7.0, 64)) + with pytest.raises(ValueError, match="uniform-in-d"): + DP.policy_log_normalization(data, x_lu, jnp.zeros(64)) + + +# ---------------------------------------------- decline -> warranted reserve + +def test_declined_row_executes_warranted_bandlimited_reserve_and_keeps_sample( + monkeypatch): + """Capacity of one for a two-mode table forces the decline; the reserve + then runs on the refined time rule, its coarser check rule is the native + production rule, and the selected value is the production value to within + the resolution tolerance -- so no sample is lost and none is silently + substituted.""" + monkeypatch.setattr(_core, "_DISTMARG_GH_N", 0) + guarded, C_B, constants = _guarded_problem(_N, _GUARD) + _install_tables(monkeypatch, ([guarded], C_B), _GUARD) + data = _fake_data(_N) + x_grid, log_w = _grid() + cfg = DP.PolicyConfig(time_guard=_GUARD, reserve_time_refine=4, + max_modes=1, enriched_max_modes=1) + value, ledger = DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(1), jnp.zeros(1), jnp.zeros(1), x_grid, log_w, + interp=INTERP, amp_sizing=40.0, config=cfg, return_ledger=True) + L = {k: np.asarray(v)[0] for k, v in ledger.items()} + assert not L["accepted_local"] + assert L["reserve_executed"] + assert L["reserve_uses_bandlimited_time"] + assert L["reserve_time_check_rule_internal"] + assert L["reserve_time_resolution_warranted"] + assert L["reserve_time_resolution_validated"] + assert L["reserve_time_guard_validated"] + assert L["reserve_time_warranted"] + assert L["selected_value_is_warranted_reserve"] + assert L["usable"] and L["reconciles"] and L["disposition_reconciles"] + assert int(L["reserve_escalations"]) == 0 + assert int(L["reserve_time_refine_used"]) == cfg.reserve_time_refine + assert np.isfinite(float(value[0])) + assert not L["decline_is_waveform_failure"] + fine = _fine_reference(constants, C_B, data, x_grid, log_w, 40.0) + assert abs(float(value[0]) - fine) <= 2.0e-3, (float(value[0]), fine) + assert abs(float(L["reserve_time_check_value"]) - float(value[0])) <= float( + cfg.total_value_error_budget_nats) + 1e-9 + assert float(L["reserve_time_resolution_error_nats"]) <= float( + cfg.total_value_error_budget_nats) + summary = DP.summarize_policy_ledger(ledger) + assert summary["rows"] == 1 and summary["reserve_executed"] == 1 + assert summary["reserve_warranted"] == 1 and summary["unusable"] == 0 + assert "decline_capacity" in summary["declines"] + + +def test_ledger_carries_every_named_acceptance_diagnostic(monkeypatch): + monkeypatch.setattr(_core, "_DISTMARG_GH_N", 0) + guarded, C_B, _ = _guarded_problem(_N, _GUARD) + _install_tables(monkeypatch, ([guarded], C_B), _GUARD) + data = _fake_data(_N) + x_grid, log_w = _grid(256) + _, ledger = DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(1), jnp.zeros(1), jnp.zeros(1), x_grid, log_w, + interp=INTERP, amp_sizing=40.0, + config=DP.PolicyConfig(time_guard=_GUARD), return_ledger=True) + names = DP.policy_acceptance_diagnostics() + for group in ("local", "reserve", "declines"): + for key in names[group]: + assert key in ledger, (group, key) + assert np.asarray(ledger[key]).shape == (1,), key + for key in ("lnL", "selected_value", "reserve_escalations", + "reserve_time_refine_used"): + assert key in ledger and np.asarray(ledger[key]).shape == (1,), key + L = {k: np.asarray(v)[0] for k, v in ledger.items()} + assert (np.isfinite(L["lnL"]) == bool(L["usable"])) + with pytest.raises(ValueError, match="time_guard must be >= 2"): + DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(1), jnp.zeros(1), jnp.zeros(1), x_grid, log_w, + interp=INTERP, amp_sizing=40.0, + config=DP.PolicyConfig(time_guard=1)) + with pytest.raises(ValueError, match="reserve_time_refine"): + DP.policy_time_rules(data, 1) + with pytest.raises(ValueError, match="even"): + DP.policy_time_rules(data, 3) + nodes, weights, check_nodes, check_weights = DP.policy_time_rules(data, 2) + np.testing.assert_allclose(np.asarray(check_weights), np.asarray(data.w_t)) + assert np.max(np.diff(np.asarray(nodes))) == pytest.approx(0.5) + nodes, weights, check_nodes, check_weights = DP.policy_time_rules(data, 4) + assert np.max(np.diff(np.asarray(check_nodes))) == pytest.approx(0.5) + assert float(np.sum(weights)) == pytest.approx(float(np.sum(data.w_t))) + + +# ---------------------------------------------------- end to end, real tables + +@pytest.mark.parametrize("force_decline", [False, True]) +def test_wrapper_policy_on_real_synthetic_tables_fails_closed_and_labels( + force_decline): + """Real coefficient tables from the accumulate path with a guard, the + wrapper's own distance grid and amplitude sizing, on the 32-sample + synthetic window. The exact lnL(t) on this window peaks at the FIRST + sample and falls monotonically, so the mass sits on the time boundary; + the planner's boundary starts are non-stationary and the interior modes + it keeps integrate to 22.86 nat against an exact 45.5. Under the + production operating point the row must therefore decline on + ``decline_boundary_maximum`` (before this flag existed it ACCEPTED that + value with every diagnostic passing). With capacity forced to one mode + it declines on capacity first. Either way the window is too short for + the reflected primitive (the reserve at guard 16 and 8 disagree by + ~0.06 nat, refine 4 and 2 by ~0.2 nat), so the composite must (a) fail + the reserve warrant even after escalating to the ceiling, (b) return nan + while keeping the finite diagnostic in the ledger, (c) say why, and (d) + count the row as unusable for the run label.""" + from RIFT.likelihood.jax_ile.time_first_peaklocal import ( + _evaluate_time_spectrum, _time_primitive_spectrum) + data = make_synth(scale=2.0, kappa_boost=10.0) + kw = dict(nphi=32, npsi=8, interp=INTERP) + exact = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, + angle_marg="exact", **kw) + guard = 16 + cfg = DP.PolicyConfig(time_guard=guard, reserve_time_refine=4, + reserve_time_refine_max=8) + if force_decline: + cfg = cfg._replace(max_modes=1, enriched_max_modes=1) + pol = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, angle_marg="exact", + direct_marginalization_policy="auto", + policy_config=cfg, **kw) + assert pol.direct_marginalization_policy == "auto" + assert pol.angle_marg_scheme == "exact" + assert pol.angle_marg_info["direct_marginalization_policy"] == "auto" + assert pol.policy_info["time_guard"] == guard + assert "local_log_normalization" in pol.policy_info + + a = np.asarray(exact._batched(jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL))) + b, ledger = pol._batched_ledger(jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL)) + b = np.asarray(b) + c = np.asarray(pol._batched(jnp.asarray(RA), jnp.asarray(DEC), + jnp.asarray(INCL))) + L = {k: np.asarray(v) for k, v in ledger.items()} + summary = DP.summarize_policy_ledger(ledger) + np.testing.assert_allclose(b, c, rtol=0.0, atol=1e-12, equal_nan=True) + assert np.all(L["reconciles"]) and np.all(L["disposition_reconciles"]) + assert np.all(L["norm_time_invariant"]) + assert not np.any(L["decline_is_waveform_failure"]) + tol = float(cfg.total_value_error_budget_nats) + unusable = ~L["usable"] + assert summary["unusable"] == int(np.sum(unusable)) + assert summary["nan_rows"] == int(np.sum(unusable)) + # (a)-(c): an unwarranted reserve was escalated to the maximum rule, keeps + # its finite diagnostic in the ledger, returns nan, and names the failed + # check; on this window that is the time reconstruction. + for i in np.flatnonzero(unusable): + assert L["reserve_executed"][i] and L["reserve_time_failed"][i] + assert np.isnan(b[i]) + assert np.isfinite(L["selected_value"][i]) + assert np.isfinite(L["reserve_value"][i]) + assert int(L["reserve_time_refine_used"][i]) == cfg.reserve_time_refine_max + assert int(L["reserve_escalations"][i]) == 1 + assert (float(L["reserve_time_guard_error"][i]) > tol + or float(L["reserve_time_resolution_error_nats"][i]) > tol), ( + {k: L[k][i] for k in L if k.startswith("reserve_time_")}) + assert np.all(np.isfinite(b[L["usable"]])) + assert np.any(unusable), ("the 32-sample window became warrantable; " + "move this test's fail-closed claim", summary) + if force_decline: + assert np.all(L["decline_capacity"]), summary + else: + assert np.all(L["decline_boundary_maximum"]), summary + assert not np.any(L["accepted_local"]), summary + # the interior-only local diagnostic is the 22.7 nat miss + assert np.all(L["enriched_value"] < L["reserve_value"] - 10.0), ( + L["enriched_value"], L["reserve_value"]) + # (d) a warranted row, if any, against the 8x refined exact reference. + usable = L["usable"] + if np.any(usable): + C_A, C_B, meta = AM.angle_coefficient_tables( + data, jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL), INTERP, + guard=guard) + refine = 8 + t = np.arange((data.npts - 1) * refine + 1, dtype=float) / refine + coeff, freq, off = _time_primitive_spectrum( + jnp.asarray(C_A).reshape((-1, C_A.shape[-1])), guard) + fine = _evaluate_time_spectrum(coeff, freq, jnp.asarray(t), off).reshape( + C_A.shape[:-1] + (t.size,)) + C_Bf = jnp.broadcast_to(jnp.asarray(C_B)[..., :1], + tuple(C_B.shape[:-1]) + (t.size,)) + lnL_t = AM.coefficient_table_distphipsimarg_exact( + fine, C_Bf, pol.x_grid, pol.log_w_grid, + amp_sizing=pol.angle_marg_info["amp_sizing"], m_max=meta["m_max"]) + w = jnp.asarray(_core._simpson_weights(t.size, data.deltaT / refine)) + ref = np.asarray(_core._time_marginalize(lnL_t, w)) + assert np.max(np.abs(b - ref)[usable]) < 3.0e-3, (summary, b, ref, a) + + theta = jnp.asarray([RA[0], DEC[0], INCL[0]]) + v, g = pol._value_and_grad(theta) + if np.isfinite(b[0]): + assert np.isfinite(float(v)) and np.all(np.isfinite(np.asarray(g))) + assert float(v) == pytest.approx(float(b[0]), abs=1e-9) + else: + # nan is the fail-closed value on the AD path too: a MALA step on it + # is rejected rather than accepted on a number nobody stands behind. + assert np.isnan(float(v)) + + +def test_wrapper_policy_has_no_lnLt_path(): + """A consumer asking for lnL(t) must be refused, not handed the + time-marginalized value under that name. The exact scheme still serves + it, so the refusal is the policy's, not the wrapper's.""" + data = make_synth(scale=2.0) + kw = dict(nphi=32, npsi=8, interp=INTERP, angle_marg="exact") + exact = JAXDistPhiPsiMargLikelihood(data, 30.0, 3000.0, **kw) + lnLt = exact._fused(data, jnp.asarray(RA[:1]), jnp.asarray(DEC[:1]), + jnp.asarray(INCL[:1]), return_lnLt=True) + assert np.asarray(lnLt).shape == (1, data.npts) + pol = JAXDistPhiPsiMargLikelihood( + data, 30.0, 3000.0, direct_marginalization_policy="auto", + policy_config=DP.PolicyConfig(time_guard=4), **kw) + with pytest.raises(ValueError, match="no lnL"): + pol._fused(data, jnp.asarray(RA[:1]), jnp.asarray(DEC[:1]), + jnp.asarray(INCL[:1]), return_lnLt=True) + + +# ------------------------------------------------------- guard vs stored buffer + +def test_guard_past_the_stored_buffer_is_refused_at_construction(monkeypatch): + """A guard the data buffer cannot supply yields a nonfinite table with no + error from the gather; the wrapper must refuse it with the remedy, not let + it read as a method decline. Per row, the same condition is a distinct + input flag.""" + data = make_synth(scale=2.0) + real = AM.angle_coefficient_tables + + def poisoned(d, ra, dec, incl, interp=None, sample_chunk=None, guard=0): + C_A, C_B, meta = real(d, ra, dec, incl, interp, sample_chunk=sample_chunk, + guard=guard) + if guard >= 8: + C_A = C_A.at[..., 0].set(jnp.nan) + return C_A, C_B, meta + + monkeypatch.setattr(AM, "angle_coefficient_tables", poisoned) + kw = dict(nphi=32, npsi=8, interp=INTERP, angle_marg="exact") + with pytest.raises(ValueError, match="not finite at time_guard=8"): + JAXDistPhiPsiMargLikelihood( + data, 30.0, 3000.0, direct_marginalization_policy="auto", + policy_config=DP.PolicyConfig(time_guard=8), **kw) + # guard 4 passes the probe; the row-level flag is true + pol = JAXDistPhiPsiMargLikelihood( + data, 30.0, 3000.0, direct_marginalization_policy="auto", + policy_config=DP.PolicyConfig(time_guard=4), **kw) + lnL, ledger = pol._batched_ledger(jnp.asarray(RA[:1]), jnp.asarray(DEC[:1]), + jnp.asarray(INCL[:1])) + assert bool(np.asarray(ledger["tables_finite"])[0]) + assert not bool(np.asarray(ledger["input_nonfinite"])[0]) + # a nonfinite row at evaluation time is nan with the input flag set (the + # policy function is called directly so the poisoned table reaches it) + def poison_all(d, ra, dec, incl, interp=None, sample_chunk=None, guard=0): + C_A, C_B, meta = real(d, ra, dec, incl, interp, sample_chunk=sample_chunk, + guard=guard) + return C_A.at[..., 0].set(jnp.nan), C_B, meta + + monkeypatch.setattr(AM, "angle_coefficient_tables", poison_all) + lnL2, ledger2 = DP.fused_log_likelihood_four_axis_policy( + data, jnp.asarray(RA[:1]), jnp.asarray(DEC[:1]), jnp.asarray(INCL[:1]), + pol.x_grid, pol.log_w_grid, interp=INTERP, + amp_sizing=pol.angle_marg_info["amp_sizing"], + config=pol.policy_config, return_ledger=True) + assert bool(np.asarray(ledger2["input_nonfinite"])[0]) + assert not bool(np.asarray(ledger2["usable"])[0]) + assert np.isnan(float(lnL2[0])) + + +def test_defaults_are_the_production_measured_operating_point(): + """The defaults follow the ladder record that accepted on production + tables, not PR #268's test fixture values.""" + cfg = DP.PolicyConfig() + assert (cfg.base_oversample, cfg.enriched_oversample) == (2, 4) + assert (cfg.max_modes, cfg.enriched_max_modes) == (16, 16) + assert cfg.local_radius == 6.0 + assert cfg.time_guard == 128 + assert cfg.reserve_time_refine_max >= cfg.reserve_time_refine + + +# ------------------------------------------------------------------- the CLI + +def test_the_driver_CLI_offers_the_policy_and_rejects_a_typo(): + """Deliberately a SUBPROCESS, like the peak-local wiring test: optparse + builds its choices from POLICY_CHOICES at import time.""" + 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(*args): + p = subprocess.run([sys.executable, driver] + list(args), env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=600) + return p.returncode, p.stdout.decode("utf-8", "replace") + + rc, out = run("--direct-marginalization-policy", "autp") + assert rc != 0 and "invalid choice" in out, out[-1500:] + assert "auto" in out, out[-1500:] + # Scope (external review P1): the policy outside its one mode, and its + # knobs without the policy, are refused at parse time, not ignored. + rc, out = run("--mode", "laplace-is", "--direct-marginalization-policy", + "auto") + assert rc != 0 and "flowmc-phipsimarg" in out, out[-1500:] + rc, out = run("--mode", "flowmc-phipsimarg", + "--direct-marginalization-time-guard", "8") + assert rc != 0 and "inert" in out, out[-1500:] + rc, out = run("--mode", "flowmc-phipsimarg", + "--direct-marginalization-policy", "auto", + "--direct-marginalization-reserve-time-refine", "3") + assert rc != 0 and "even" in out, out[-1500:] + rc, out = run("--mode", "flowmc-phipsimarg", + "--direct-marginalization-policy", "auto", + "--direct-marginalization-reserve-time-refine-max", "2") + assert rc != 0 and "refine-max" in out, out[-1500:] + rc, out = run("--help") + assert "--direct-marginalization-policy" in out + assert "--direct-marginalization-time-guard" in out + assert "--direct-marginalization-reserve-time-refine" in out + assert "--direct-marginalization-reserve-time-refine-max" in out + assert "--direct-marginalization-error-budget-nats" in out + + +# ------------------------------------------- row batching is a COST knob only + +def _policy_like(batch_rows, data, **kw): + cfg = DP.PolicyConfig(time_guard=16, reserve_time_refine=4, + reserve_time_refine_max=8, + reserve_batch_rows=batch_rows) + return JAXDistPhiPsiMargLikelihood( + data, 30.0, 3000.0, angle_marg="exact", + direct_marginalization_policy="auto", policy_config=cfg, **kw) + + +def test_row_batch_size_changes_cost_not_values_decisions_or_gradients(): + """``reserve_batch_rows`` may only move VALUES, never cost. + + Cost it does move, and against the change: see the design note. This test + is about the values, decisions and gradients only. + + Executing B rows together turns the controller's tier-escalation + ``lax.cond`` into a ``select`` under ``vmap``, so EVERY reserve tier runs + for EVERY row in the batch instead of only for the rows that failed their + warrant. That is a real change to the executed graph, and it must leave + the selected value, every branch decision in the ledger, the summary + counts, and the reverse-mode gradient exactly where the row-at-a-time path + put them. The gradient is checked separately from the value because a + ``select`` propagates the untaken branch's nan into the cotangent even + when it discards the untaken branch's value -- the classic ``where`` + nan-gradient trap, which a value-only comparison cannot see. + """ + S = 6 + ra = np.linspace(0.55, 1.35, S) + dec = np.linspace(0.05, 0.75, S) + incl = np.linspace(0.35, 2.45, S) + data = make_synth(scale=2.0, kappa_boost=10.0) + kw = dict(nphi=32, npsi=8, interp=INTERP) + args = (jnp.asarray(ra), jnp.asarray(dec), jnp.asarray(incl)) + + ref_like = _policy_like(1, data, **kw) + ref_lnL, ref_led = ref_like._batched_ledger(*args) + ref_lnL = np.asarray(ref_lnL) + ref_led = {k: np.asarray(v) for k, v in ref_led.items()} + ref_sum = DP.summarize_policy_ledger(ref_led) + assert ref_sum["reserve_batch_execution_sequential"] is True + assert ref_sum["reserve_batch_rows"] == 1 + + # The comparison is only meaningful if the escalation predicate is + # actually mixed across the batch: with a uniform predicate the select + # would agree with the cond for trivial reasons. Assert the fixture + # still produces both, so a future retune cannot silently blind this. + esc = ref_led["reserve_escalations"] + assert esc.min() != esc.max(), ( + "fixture no longer mixes escalating and non-escalating rows " + "(escalations=%r); the batched/sequential comparison would be blind" + % (esc,)) + + mid = np.array([ra[S // 2], dec[S // 2], incl[S // 2]]) + ref_v, ref_g = ref_like.value_and_grad(mid) + assert np.all(np.isfinite(ref_g)), ref_g + + for B in (2, 4, 0): # 0 == one full batch of all S rows + like = _policy_like(B, data, **kw) + lnL, led = like._batched_ledger(*args) + lnL = np.asarray(lnL) + led = {k: np.asarray(v) for k, v in led.items()} + + np.testing.assert_array_equal(np.isnan(lnL), np.isnan(ref_lnL)) + np.testing.assert_allclose(lnL, ref_lnL, rtol=0.0, atol=1e-13, + equal_nan=True) + + assert set(led) == set(ref_led) + for key, want in sorted(ref_led.items()): + if key in ("reserve_batch_execution_sequential", + "reserve_batch_rows_requested", + "reserve_batch_rows_executed"): + continue + got = led[key] + if want.dtype == bool or np.issubdtype(want.dtype, np.integer): + np.testing.assert_array_equal( + got, want, err_msg="batch_rows=%d moved ledger key %r" + % (B, key)) + else: + np.testing.assert_array_equal( + np.isfinite(got), np.isfinite(want), + err_msg="batch_rows=%d moved finiteness of %r" % (B, key)) + np.testing.assert_allclose( + got, want, rtol=0.0, atol=1e-13, equal_nan=True, + err_msg="batch_rows=%d moved ledger key %r" % (B, key)) + + summ = DP.summarize_policy_ledger(led) + for key in ref_sum: + # The three batch keys are the ones that MUST differ; they are + # asserted explicitly below. + if key in ("reserve_batch_rows", + "reserve_batch_rows_requested", + "reserve_batch_execution_sequential"): + continue + want, got = ref_sum[key], summ[key] + # Counts must be exact. max_local_error_score_nats is not a count: + # it is a max over a per-row float diagnostic that is itself a + # reduction, so the select reassociates it. Measured 3.553e-15 on + # 2 of 6 rows (2.4e-13 relative) while lnL stayed bitwise equal and + # every branch decision held, so it is compared as a float. + if isinstance(want, float): + assert np.isclose(got, want, rtol=1e-11, atol=0.0, + equal_nan=True), (B, key, got, want) + else: + assert got == want, (B, key, got, want) + + # The ledger key must state what actually happened. It read True + # unconditionally before the batch size was a knob. + # The ledger reports what EXECUTED, not what was requested: a request + # above the row count runs a vmap of S, not of B. + assert summ["reserve_batch_execution_sequential"] is False + assert summ["reserve_batch_rows"] == (S if (B == 0 or B >= S) else B) + assert summ["reserve_batch_rows_requested"] == B + assert not np.any(led["reserve_batch_execution_sequential"]) + + # The gradient is compared at EVERY batch size. Dropping it to one + # size, and halving nphi, were both tried and neither moved the test's + # cost (640.8 s against 641.6 s on the ldas-grid CPU runner), so the + # cheaper variants bought nothing and this keeps the coverage. + # A single row has nothing to batch, and value_and_grad evaluates + # exactly one. If a batch request reached it, both the accept/reserve + # cond and every escalation tier would become selects and the gradient + # would cost more than at batch 1 for no amortization. Pin that the + # scalar path stays sequential whatever was asked for. + _, sled = like._batched_ledger(jnp.asarray(mid[0:1]), + jnp.asarray(mid[1:2]), + jnp.asarray(mid[2:3])) + ssum = DP.summarize_policy_ledger({k: np.asarray(v_) + for k, v_ in sled.items()}) + assert ssum["reserve_batch_rows"] == 1, (B, ssum["reserve_batch_rows"]) + assert ssum["reserve_batch_execution_sequential"] is True + + v, g = like.value_and_grad(mid) + assert np.all(np.isfinite(g)), (B, g) + assert abs(v - ref_v) <= 1e-13, (B, v, ref_v) # measured exactly 0 + # Relative, not absolute: the select re-associates the cotangent sum, + # so the gradient is equal to ~1 ulp rather than bitwise (measured + # 2.8e-14 on a Blackwell GPU and 5.6e-14 on CPU against |dlnL/dincl| + # ~ 141). A branch decision that actually moved would be O(1) nats. + np.testing.assert_allclose(g, ref_g, rtol=1e-12, atol=1e-12, + err_msg="batch_rows=%d moved the gradient" % B) + + +@pytest.mark.parametrize("bad", [-1, -8, 2.5, "4", None]) +def test_row_batch_size_refuses_what_it_cannot_execute(bad): + with pytest.raises(ValueError): + DP.validate_batch_rows(bad) + + +# ------------------------------------------- preset local-plan capacities + +def test_local_plan_capacities_reach_the_planner(monkeypatch): + """Both capacities are PRESET, and -- the point of the test -- reachable. + + The regression guarded here is not a wrong number but an unreachable one: + max_time_nodes defaulted to 64 inside rank_joint_starts_from_uvq_device and + the policy never passed it, so the value that gated the local path could + not be moved from any config field or flag. Asserting the default alone + would have passed against that bug, so spy on the call and read the kwargs + the planner is actually handed.""" + seen = [] + + class _Stop(Exception): + pass + + def spy(*a, **kw): + seen.append(kw) + raise _Stop + + monkeypatch.setattr(DP._aap, "rank_joint_starts_from_uvq_device", spy) + guarded, C_B, _ = _guarded_problem(_N, _GUARD) + _install_tables(monkeypatch, ([guarded], C_B), _GUARD) + data = _fake_data(_N) + x_grid, log_w = _grid(256) + with pytest.raises(_Stop): + DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(1), jnp.zeros(1), jnp.zeros(1), x_grid, log_w, + interp=INTERP, amp_sizing=40.0, + config=DP.PolicyConfig(time_guard=_GUARD, max_time_nodes=333, + base_max_starts=77)) + assert seen, "the planner was never called" + assert seen[0].get("max_time_nodes") == 333, seen[0] + assert seen[0].get("max_starts") == 77, seen[0] + + +def test_capacity_and_budget_defaults_are_the_approved_operating_point(): + """Pins the operating point RO approved on 2026-09-08 (evening). + + Shipped was (64, 32) with a 1e-3 budget. The capacities moved because + acceptance at rho 652 goes (64, 32) 28%, (256, 32) 31%, (256, 128) 75%; + they are pinned as a PAIR because of 64 rows, 36 declines fail on time + nodes and 37 on starts and only 7 on starts alone, so a later change that + moves one alone is a mistake this test should catch. The budget moved + because 1e-3 -> 1e-2 took reserve escalations from 2 to 0 at rho 41, a + 2.19x speedup, while moving lnL by 4.8e-12 nats. + + None of these is value-neutral, which is why they are pinned rather than + left to drift: base_max_starts 32 -> 128 alone shifts already-accepted + values by up to 3.5e-3 nats at rho 163.""" + cfg = DP.PolicyConfig() + assert (cfg.max_time_nodes, cfg.base_max_starts) == (256, 128) + assert cfg.total_value_error_budget_nats == 1.0e-2 + + +@pytest.mark.parametrize("kw", [{"max_time_nodes": 1}, {"max_time_nodes": 0}, + {"base_max_starts": 0}, + {"base_max_starts": -4}]) +def test_capacities_refuse_what_cannot_plan(kw): + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig(**kw)) + + +def test_driver_offers_the_capacity_knobs_and_scopes_them(): + 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(*args): + p = subprocess.run([sys.executable, driver] + list(args), env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=600) + return p.returncode, p.stdout.decode("utf-8", "replace") + + rc, out = run("--help") + assert "--direct-marginalization-max-time-nodes" in out + assert "--direct-marginalization-max-starts" in out + # --help exits before validation, so it proves only that the flags parse. + # Both must also be scoped: inert without the policy is an error, not a + # silently ignored flag. + for flag in ("--direct-marginalization-max-time-nodes", + "--direct-marginalization-max-starts"): + rc, out = run("--mode", "flowmc-phipsimarg", flag, "64") + assert rc != 0 and "inert" in out, (flag, out[-1500:]) + + +def test_driver_offers_the_batch_rows_knob_and_scopes_it(): + 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(*args): + p = subprocess.run([sys.executable, driver] + list(args), env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=600) + return p.returncode, p.stdout.decode("utf-8", "replace") + + rc, out = run("--help") + assert "--direct-marginalization-batch-rows" in out + # Inert without the policy, and refused rather than clamped when negative. + rc, out = run("--mode", "flowmc-phipsimarg", + "--direct-marginalization-batch-rows", "8") + assert rc != 0 and "inert" in out, out[-1500:] + rc, out = run("--mode", "flowmc-phipsimarg", + "--direct-marginalization-policy", "auto", + "--direct-marginalization-batch-rows", "-1") + assert rc != 0 and "batch-rows" in out, out[-1500:] diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy_cli.py b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy_cli.py new file mode 100644 index 000000000..fdcfb0f81 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_direct_marginalization_policy_cli.py @@ -0,0 +1,305 @@ +"""Driver-seam tests for the four-axis policy's observability and plan knobs. + +These are SUBPROCESS tests on purpose. Every option here is read in +``build_parser`` and validated in ``check_critical_and_report``, both of which +run before any likelihood is built, so a library-level test of the policy +module cannot see them at all: it would exercise ``PolicyConfig`` directly and +pass no matter what the command line does. The failure this guards against is +the one this pipeline keeps hitting -- a flag that is accepted and then +silently inert -- so each case asserts that a misuse is REFUSED rather than +ignored, and that the accepting combination is not refused. + +A separate file from test_direct_marginalization_policy.py because these need +no JAX device, no synthetic data and no fixture: they are parser and validator +behaviour only, and they run in well under a second each. +""" + +import os +import subprocess +import sys + +import pytest + + +_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +_DRIVER = os.path.join(_ROOT, "bin", "integrate_likelihood_extrinsic_jax") + +# Enough of a policy request to reach the policy's own validation. --sim-xml is +# a path that exists so argument parsing does not fail for an unrelated reason; +# the run never gets as far as reading it, because every case below is decided +# at parse time. +_POLICY = ("--mode", "flowmc-phipsimarg", "--distance-marginalization", + "--direct-marginalization-policy", "auto", + "--angle-marg-scheme", "exact", "--sim-xml", os.devnull) + + +def _run(*args): + env = dict(os.environ) + env["PYTHONPATH"] = _ROOT + os.pathsep + env.get("PYTHONPATH", "") + env["JAX_PLATFORMS"] = "cpu" + p = subprocess.run([sys.executable, _DRIVER] + list(args), env=env, + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, + timeout=600) + return p.returncode, p.stdout.decode("utf-8", "replace") + + +def test_the_help_lists_every_new_policy_and_smc_knob(): + rc, out = _run("--help") + assert rc == 0, out[-1500:] + for flag in ("--direct-marginalization-policy-probe-rows", + "--direct-marginalization-policy-probe-only", + "--direct-marginalization-max-modes", + "--direct-marginalization-enriched-max-modes", + "--direct-marginalization-base-oversample", + "--direct-marginalization-enriched-oversample", + "--direct-marginalization-max-starts", + "--direct-marginalization-convergence-tol-nats", + "--direct-marginalization-time-guard-tol-nats", + "--direct-marginalization-reserve-time-refine-max", + "--smc-is-samples"): + assert flag in out, flag + + +@pytest.mark.parametrize("flag,value", [ + ("--direct-marginalization-policy-probe-rows", "8"), + ("--direct-marginalization-max-modes", "16"), + ("--direct-marginalization-enriched-max-modes", "16"), + ("--direct-marginalization-base-oversample", "2"), + ("--direct-marginalization-enriched-oversample", "4"), + ("--direct-marginalization-max-starts", "64"), + ("--direct-marginalization-convergence-tol-nats", "0.5"), + ("--direct-marginalization-time-guard-tol-nats", "0.5"), + ("--direct-marginalization-reserve-time-refine-max", "8"), +]) +def test_a_policy_knob_without_the_policy_is_refused_not_ignored(flag, value): + """Each knob is inert unless the policy is on, so passing one without it is + a fatal mistake rather than a silently dropped request.""" + rc, out = _run("--mode", "flowmc-phipsimarg", flag, value) + assert rc != 0, out[-1500:] + assert "inert" in out, out[-1500:] + + +def test_a_negative_smc_is_sample_count_is_refused(): + """A negative count reaches the SMC proposal draw, whose bare exception + handler would swallow it and publish the raw SMC evidence instead of the IS + evidence, with nothing in the output saying the estimator changed.""" + rc, out = _run("--mode", "flowmc-phipsimarg", "--smc-is-samples", "-1") + assert rc != 0, out[-1500:] + assert "smc-is-samples" in out, out[-1500:] + + +def test_probe_only_without_probe_rows_is_refused(): + """--probe-only with zero rows would exit having measured nothing.""" + rc, out = _run(*(_POLICY + ("--direct-marginalization-policy-probe-only",))) + assert rc != 0, out[-1500:] + assert "probe-rows" in out, out[-1500:] + + +def test_an_enriched_plan_narrower_than_the_base_plan_is_refused(): + """Acceptance compares a base plan against an enriched one that must be + able to nest it; a narrower enriched cap declines on mode nesting every + row, which would read as a property of the data.""" + rc, out = _run(*(_POLICY + ("--direct-marginalization-max-modes", "16", + "--direct-marginalization-enriched-max-modes", + "8"))) + assert rc != 0, out[-1500:] + assert "nest" in out, out[-1500:] + + +def test_an_escalation_ceiling_below_its_floor_is_refused(): + rc, out = _run(*(_POLICY + ( + "--direct-marginalization-reserve-time-refine", "8", + "--direct-marginalization-reserve-time-refine-max", "4"))) + assert rc != 0, out[-1500:] + assert "reserve-time-refine-max" in out, out[-1500:] + + +@pytest.mark.parametrize("flag", [ + "--direct-marginalization-convergence-tol-nats", + "--direct-marginalization-time-guard-tol-nats", +]) +def test_a_nonpositive_tolerance_is_refused(flag): + rc, out = _run(*(_POLICY + (flag, "-1"))) + assert rc != 0, out[-1500:] + assert "finite and positive" in out, out[-1500:] + + +def test_the_recorded_accepting_operating_point_is_not_refused(): + """The counterpart to every case above: the combination that reaches the + four-axis branch must pass validation. It still exits nonzero, on a + missing --event-time, which is the point -- the policy's own validation is + behind it, so a future tightening that rejected this configuration would + be caught here rather than in a run.""" + rc, out = _run(*(_POLICY + ( + "--direct-marginalization-time-guard", "128", + "--direct-marginalization-max-modes", "16", + "--direct-marginalization-enriched-max-modes", "16", + "--direct-marginalization-base-oversample", "2", + "--direct-marginalization-enriched-oversample", "4", + "--direct-marginalization-reserve-time-refine", "4", + "--direct-marginalization-reserve-time-refine-max", "4", + "--direct-marginalization-policy-probe-rows", "8", + "--direct-marginalization-policy-probe-only"))) + assert rc != 0, out[-1500:] + assert "event-time" in out, out[-1500:] + assert "direct-marginalization" not in out.split("error:")[-1], out[-1500:] + + +# --------------------------------------------------- the note's return arity + +def _load_driver(): + """Import the driver script as a module (it has no .py extension).""" + import importlib.util + spec = importlib.util.spec_from_loader( + "ile_jax_driver", + importlib.machinery.SourceFileLoader("ile_jax_driver", _DRIVER)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +@pytest.mark.parametrize("policy,ledger,theta_rows", [ + ("off", object(), 4), # policy disabled + ("auto", None, 4), # no ledger built + ("auto", object(), 0), # no rows to evaluate +]) +def test_the_note_returns_three_values_on_every_early_path(policy, ledger, + theta_rows): + """``return_values=True`` must return the same NUMBER of values on every + path, including the ones that give up early. + + The caller unpacks three. Two early returns handed back two, so any caller + reaching them died on an unpacking error rather than on the condition the + early return was written to handle. None of the three is reachable from + the probe today, which is exactly why it needs pinning: the guard is + unexercised, so nothing else would notice it drifting. + """ + import numpy as np + mod = _load_driver() + + class _Like(object): + pass + + like = _Like() + like.direct_marginalization_policy = policy + like._batched_ledger = ledger + theta = np.zeros((theta_rows, 3)) + + out = mod.direct_marginalization_policy_note(like, theta, + return_values=True) + assert isinstance(out, tuple) and len(out) == 3, out + assert isinstance(out[0], str) + + plain = mod.direct_marginalization_policy_note(like, theta) + assert isinstance(plain, str), plain + + +def test_every_driver_option_string_is_registered_exactly_once(): + """optparse SILENTLY keeps the last registration of a duplicated option. + + A parallel branch merge produced two add_option calls for one flag; --help + rendered correctly from the first while the parser used the second, so the + driver ran at a different default from the one documented and nothing + raised. That is invisible to every other test in this file, which all go + through the same parser and would agree with each other. + + Read from the SOURCE rather than the parser for the same reason: the parser + only knows the winner. + """ + import collections + import re + + src = open(_DRIVER, encoding="utf-8").read() + names = re.findall(r'add_option\(\s*"(--[A-Za-z0-9-]+)"', src) + dupes = {n: c for n, c in collections.Counter(names).items() if c > 1} + assert not dupes, "option strings registered more than once: %r" % (dupes,) + + +def test_each_policy_flag_default_matches_its_PolicyConfig_field(): + """The driver default and the library default must be the SAME number. + + The symptom that started this guard: a run used max_starts 32 from the + driver while PolicyConfig carried 128, so --help, the config and the + executed run disagreed and nothing raised. Any knob whose two defaults + drift silently changes what a bare command line computes, and a default + change landing on one side only is the easiest way to produce that. + + Deliberately compares VALUES, not a hardcoded expectation: when a default is + legitimately changed (RO approved max_starts 32 -> 128), this test keeps + passing as long as BOTH sides move, and fails the moment only one does. + """ + from RIFT.likelihood.jax_ile.direct_marginalization_policy import PolicyConfig + + mod = _load_driver() + parser = mod.build_parser() + cfg = PolicyConfig() + + pairs = { + "direct_marginalization_time_guard": "time_guard", + "direct_marginalization_reserve_time_refine": "reserve_time_refine", + "direct_marginalization_reserve_time_refine_max": "reserve_time_refine_max", + "direct_marginalization_error_budget_nats": "total_value_error_budget_nats", + "direct_marginalization_max_modes": "max_modes", + "direct_marginalization_enriched_max_modes": "enriched_max_modes", + "direct_marginalization_base_oversample": "base_oversample", + "direct_marginalization_enriched_oversample": "enriched_oversample", + "direct_marginalization_max_starts": "base_max_starts", + "direct_marginalization_max_time_nodes": "max_time_nodes", + "direct_marginalization_convergence_tol_nats": "convergence_tol_nats", + "direct_marginalization_time_guard_tol_nats": "time_guard_tol_nats", + } + mismatched = {} + for dest, field in pairs.items(): + if not hasattr(cfg, field): + continue + drv = parser.defaults.get(dest, "") + if drv == "": + continue + lib = getattr(cfg, field) + if drv != lib: + mismatched[dest] = (drv, lib) + assert not mismatched, ( + "driver default != PolicyConfig default for %r " + "(driver, library)" % (mismatched,)) + + +def test_every_policy_flag_help_states_its_real_default(): + """``--help`` must not quote a number the flag no longer uses. + + The value guard above was passing while four help strings still read + "(default 4)", "(default 8)", "(default 1)", "(default 2)" -- the pre-#280 + portfolio -- and --direct-marginalization-time-guard read "(default 16)" + against a default of 128. The values had been repointed at PolicyConfig + and the prose had not, so the two defaults agreed with each other and + disagreed with what --help told the operator. A knob's documented default + is what someone reads before deciding whether to pass it, so a stale one + misconfigures a run exactly as a stale value does. + + Reads the parser's rendered help, not the source, so an interpolation that + silently fails to interpolate is caught too. A stated default may carry a + trailing constraint ("16; must be >= 2") -- only the leading token is the + number, and the rest of the parenthetical may explain the value + ("1: row at a time", "0 = off"). + """ + import re + + mod = _load_driver() + parser = mod.build_parser() + + stale = {} + for opt in parser._get_all_options(): + dest = opt.dest + if not dest or not str(dest).startswith("direct_marginalization"): + continue + for match in re.finditer(r"\(default ([^\s);,:=]+)", opt.help or ""): + stated = match.group(1).strip().strip("'\"") + actual = parser.defaults.get(dest) + try: + ok = float(stated) == float(actual) + except (TypeError, ValueError): + ok = stated == str(actual) + if not ok: + stale[opt.get_opt_string()] = (stated, actual) + assert not stale, ( + "help text states a default the flag does not use %r " + "(stated, actual)" % (stale,)) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_distance_gh_nodes_cli.py b/MonteCarloMarginalizeCode/Code/test/jax/test_distance_gh_nodes_cli.py new file mode 100644 index 000000000..9d2b4d5ab --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_distance_gh_nodes_cli.py @@ -0,0 +1,472 @@ +"""``--distance-gh-nodes``: make the per-sample Gauss-Hermite distance +quadrature reachable by an ILE argument, and the driver's warn-not-ignore +compatibility notes for three previously-silent knobs. + +RO'S dispositions (2026-09-08): (a) unreachable code (the per-sample distance +quadrature) gets an ILE-style CLI argument; (b) missing knobs stay no-op for +compatibility but must WARN. This file gates both, plus the CLI/env +compatibility contract the new option must honour without import-order +fragility (``core.set_distmarg_gh_nodes`` / ``get_distmarg_gh_nodes``, read at +CALL time by every consumer -- see core.py and the module docstring there). + +Layout: + * CLI/env resolution and refusal (parse-time, no subprocess: the real + ``check_critical_and_report`` is called in-process, exactly as + test_distance_grid_loguniform.py's parse-time tests do). + * The three "accepted but IGNORED, and SAYS SO" notes: --phase-marginalization + on the phi_ref-analytic modes, --sky-coordinates network on the modes that + do not implement it, --d-prior on any non-volumetric value. + * A numeric liveness check: the resolved node count actually changes the + constructed likelihood's VALUE (not just an echoed CLI flag) on a cheap + synthetic packed-data fixture -- no lal, no frames. + * One subprocess check against the real driver with --inj-mode (tiny + injection, stopped at a known post-construction validation error, exactly + as test_distance_grid_loguniform.py's own subprocess test does) that the + resolved count reaches the run log end to end. + +Run: + PYTHONPATH=<...>/Code python -m pytest -q test/jax/test_distance_gh_nodes_cli.py +""" +import contextlib +import importlib.machinery +import importlib.util +import io +import os +import subprocess +import sys +import tempfile + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import core as core_mod +from RIFT.likelihood.jax_ile import build_likelihood_data +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiMargLikelihood + +_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 _driver_module(): + """Import the driver BY PATH (no .py suffix, so plain import cannot see + it). Gives in-process access to the real build_parser/ + check_critical_and_report, so the checks below are executable coverage of + the shipping functions, not a grep of the source.""" + assert os.path.exists(_JAXDRIVER), "driver missing: %s" % _JAXDRIVER + loader = importlib.machinery.SourceFileLoader("_ghnodes_driver", _JAXDRIVER) + spec = importlib.util.spec_from_loader(loader.name, loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +@pytest.fixture +def saved_gh_env(): + """Every test that touches --distance-gh-nodes mutates process-global + state (core._DISTMARG_GH_N and/or os.environ), exactly like the existing + JAX_ILE_DISTMARG_GH tests in test_distance_grid_loguniform.py -- restore + both, unconditionally, so one test's mutation cannot leak into the next.""" + saved_env = os.environ.get("JAX_ILE_DISTMARG_GH") + saved_core = core_mod._DISTMARG_GH_N + try: + yield + finally: + if saved_env is None: + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + else: + os.environ["JAX_ILE_DISTMARG_GH"] = saved_env + core_mod._DISTMARG_GH_N = saved_core + + +def _parse(mod, args): + optp = mod.build_parser() + opts, _ = optp.parse_args(list(args)) + return optp, opts + + +def _run_checked(mod, args): + """check_critical_and_report(), capturing stdout+stderr as one string.""" + optp, opts = _parse(mod, args) + out, err = io.StringIO(), io.StringIO() + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + mod.check_critical_and_report(opts, optp) + return opts, out.getvalue() + err.getvalue() + + +def _run_refused(mod, args): + """Like _run_checked, but the call is expected to SystemExit; returns the + captured text. Raises AssertionError if it does not exit.""" + optp, opts = _parse(mod, args) + out, err = io.StringIO(), io.StringIO() + try: + with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): + mod.check_critical_and_report(opts, optp) + except SystemExit: + return out.getvalue() + err.getvalue() + raise AssertionError("expected a refusal (SystemExit); accepted %r" % (args,)) + + +# --------------------------------------------------------------------------- +# CLI/env resolution and refusal +# --------------------------------------------------------------------------- + +def test_option_is_registered_with_default_none(): + """Default must be None, not 0: 0 is a legal explicit value (BLOCKER, + external review of this PR) and a 0-default makes an explicit + ``--distance-gh-nodes 0`` indistinguishable from "not passed", so a + nonzero JAX_ILE_DISTMARG_GH would silently win instead of losing to the + explicit CLI 0.""" + mod = _driver_module() + optp = mod.build_parser() + opt = next(o for g in ([optp] + optp.option_groups) + for o in getattr(g, "option_list", []) + if "--distance-gh-nodes" in (o._long_opts or [])) + assert opt.default is None, ( + "unreachable feature (a) must default to None, not 0, so an " + "explicit 0 is distinguishable from not-passed: got %r" % (opt.default,)) + assert opt.type == "int" + + +def test_cli_alone_resolves_and_mutates_core_state(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + opts, text = _run_checked(mod, ["--distance-gh-nodes", "48"]) + assert opts._distance_gh_nodes_resolved == 48 + assert core_mod.get_distmarg_gh_nodes() == 48, ( + "set_distmarg_gh_nodes() must reach core's live module state") + assert "nodes=48" in text and "--distance-gh-nodes" in text + + +def test_env_alone_is_still_honoured(saved_gh_env): + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "17" + opts, text = _run_checked(mod, []) + assert opts._distance_gh_nodes_resolved == 17 + assert core_mod.get_distmarg_gh_nodes() == 17 + assert "nodes=17" in text and "JAX_ILE_DISTMARG_GH" in text + + +def test_default_is_zero_and_legacy_grid_is_named(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + opts, text = _run_checked(mod, []) + assert opts._distance_gh_nodes_resolved == 0 + assert core_mod.get_distmarg_gh_nodes() == 0 + assert "nodes=0" in text and "legacy uniform grid" in text + + +def test_cli_and_env_agreeing_is_accepted(saved_gh_env): + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "9" + opts, text = _run_checked(mod, ["--distance-gh-nodes", "9"]) + assert opts._distance_gh_nodes_resolved == 9 + assert core_mod.get_distmarg_gh_nodes() == 9 + + +def test_cli_and_env_conflict_is_REFUSED_not_reconciled(saved_gh_env): + """(a): the CLI must WIN, but silently picking one of two disagreeing + values is exactly the drift a compatibility path exists to prevent -- so + two different nonzero values must refuse, not resolve.""" + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "12" + text = _run_refused(mod, ["--distance-gh-nodes", "40"]) + assert "--distance-gh-nodes" in text + assert "JAX_ILE_DISTMARG_GH" in text + assert "conflicts" in text + + +def test_refusal_does_not_mutate_core_state(saved_gh_env): + """A refused command line must never apply its (disputed) resolution -- + mirrors test_distance_grid_loguniform.py's own precondition discipline for + this exact global.""" + mod = _driver_module() + core_mod._DISTMARG_GH_N = 0 + os.environ["JAX_ILE_DISTMARG_GH"] = "12" + try: + _run_refused(mod, ["--distance-gh-nodes", "40"]) + finally: + pass + assert core_mod._DISTMARG_GH_N == 0, ( + "a refused --distance-gh-nodes/JAX_ILE_DISTMARG_GH conflict must not " + "mutate core._DISTMARG_GH_N; got %r" % (core_mod._DISTMARG_GH_N,)) + + +def test_cli_route_still_trips_the_loguniform_incompatibility(saved_gh_env): + """F2's driver-side refusal (originally env-only) must also fire when the + node count arrives via --distance-gh-nodes, not just JAX_ILE_DISTMARG_GH -- + the parse-time check must read the RESOLVED value, not re-read the + environment directly (that would silently accept the CLI route).""" + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + text = _run_refused(mod, [ + "--mode", "flowmc-phipsimarg", "--distance-grid-scheme", "loguniform", + "--angle-marg-scheme", "auto", "--distance-gh-nodes", "32"]) + assert "distance-gh-nodes" in text or "JAX_ILE_DISTMARG_GH" in text + assert "inert" in text + + +def test_explicit_cli_zero_with_nonzero_env_is_REFUSED(saved_gh_env): + """The BLOCKER this file exists to close: with the old ``0``-default, + ``getattr(opts, "distance_gh_nodes", 0) or 0`` could not tell an explicit + ``--distance-gh-nodes 0`` apart from "not passed", so JAX_ILE_DISTMARG_GH + silently won. An explicit 0 against a nonzero env is a conflict like any + other -- it must refuse, not resolve to 0 and not resolve to 64.""" + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "64" + text = _run_refused(mod, ["--distance-gh-nodes", "0"]) + assert "--distance-gh-nodes" in text + assert "JAX_ILE_DISTMARG_GH" in text + assert "conflicts" in text + + +def test_explicit_cli_16_with_agreeing_env_16_is_accepted(saved_gh_env): + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "16" + opts, text = _run_checked(mod, ["--distance-gh-nodes", "16"]) + assert opts._distance_gh_nodes_resolved == 16 + assert core_mod.get_distmarg_gh_nodes() == 16 + assert "nodes=16" in text + + +def test_env_16_with_no_cli_resolves_to_16_and_banner_says_so(saved_gh_env): + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "16" + opts, text = _run_checked(mod, []) + assert opts._distance_gh_nodes_resolved == 16 + assert core_mod.get_distmarg_gh_nodes() == 16 + assert "nodes=16" in text and "JAX_ILE_DISTMARG_GH" in text + + +def test_cli_wins_regression_env_16_cli_32_must_refuse(saved_gh_env): + """Mutation-test target: a priority reversal (env wins over CLI) passes + every OTHER test in this file just as readily as "CLI wins" does, because + most cases here use only one of the two knobs. This is the one case that + tells them apart: CLI 32 disagrees with env 16, so the correct policy + (CLI wins when given; conflicts are refused, never silently reconciled) + must refuse. A reversed-priority implementation that silently resolves + to the env value (16) instead of refusing passes _run_refused's + SystemExit check trivially only if it also refuses -- if it does not, + _run_refused raises AssertionError itself, so this test fails loudly + rather than passing on an unresolved value.""" + mod = _driver_module() + os.environ["JAX_ILE_DISTMARG_GH"] = "16" + text = _run_refused(mod, ["--distance-gh-nodes", "32"]) + assert "--distance-gh-nodes" in text + assert "32" in text and "16" in text + assert "JAX_ILE_DISTMARG_GH" in text + assert "conflicts" in text + + +# --------------------------------------------------------------------------- +# (b) missing/scoped knobs: no-op for compatibility, but must WARN +# --------------------------------------------------------------------------- + +def test_phase_marginalization_ignored_note_on_phase_analytic_modes(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + for mode in sorted(mod._PHASE_ANALYTIC_MODES): + _, text = _run_checked(mod, ["--mode", mode, "--phase-marginalization"]) + assert "--phase-marginalization" in text and "IGNORED" in text, ( + "mode %s: expected an IGNORED note, got %r" % (mode, text)) + assert mode in text + + +def test_phase_marginalization_note_absent_where_the_flag_is_honoured(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + for mode in ("laplace-is", "prior-mc", "nuts"): + _, text = _run_checked(mod, [ + "--mode", mode, "--distance-marginalization", + "--phase-marginalization"]) + assert "--phase-marginalization" not in text, ( + "mode %s honours --phase-marginalization; it must not be reported " + "IGNORED: %r" % (mode, text)) + + +def test_sky_coordinates_ignored_note_off_multistart(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + for mode in ("nuts", "laplace-is", "flowmc-phimarg"): + _, text = _run_checked(mod, ["--mode", mode, + "--sky-coordinates", "network"]) + assert "--sky-coordinates" in text and "IGNORED" in text, ( + "mode %s: expected an IGNORED note, got %r" % (mode, text)) + + +def test_sky_coordinates_note_absent_on_multistart_nuts(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + _, text = _run_checked(mod, ["--mode", "multistart-nuts", + "--sky-coordinates", "network"]) + assert "--sky-coordinates" not in text, ( + "multistart-nuts implements --sky-coordinates network; got %r" % (text,)) + + +def test_sky_coordinates_default_equatorial_never_notes(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + for mode in ("nuts", "multistart-nuts", "laplace-is"): + _, text = _run_checked(mod, ["--mode", mode]) + assert "--sky-coordinates" not in text + + +def test_d_prior_ignored_note_for_a_real_alternative_prior(saved_gh_env): + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + _, text = _run_checked(mod, ["--d-prior", "cosmo"]) + assert "--d-prior" in text and "IGNORED" in text and "volumetric" in text + + +@pytest.mark.parametrize("value", ["Euclidean", "euclidean", "Volumetric", "volumetric"]) +def test_d_prior_no_note_for_the_volumetric_prior_itself(saved_gh_env, value): + """Euclidean/volumetric (any case) IS the JAX driver's prior, so stating + it explicitly is not a deviation and must not print a note.""" + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + _, text = _run_checked(mod, ["--d-prior", value]) + assert "--d-prior" not in text, "no note expected for --d-prior %s: %r" % (value, text) + + +def test_d_prior_excluded_from_the_generic_ignored_bag(saved_gh_env): + """--d-prior must not ALSO appear in the blanket 'accepted but IGNORED' + list (it would be reported twice, once generically and once with the + substantive message) -- moved out of `implemented` per the task, and + explicitly excluded from the generic bag alongside that.""" + mod = _driver_module() + os.environ.pop("JAX_ILE_DISTMARG_GH", None) + _, text = _run_checked(mod, ["--d-prior", "cosmo"]) + assert text.count("--d-prior") == 1, ( + "expected exactly one --d-prior mention (the substantive note); got %r" % (text,)) + + +# --------------------------------------------------------------------------- +# Numeric liveness: the resolved node count must change the VALUE, not just +# an echoed CLI flag. Cheap synthetic packed data -- no lal, no frames. +# --------------------------------------------------------------------------- + +def _synth(scale=1.0, seed=3, modes=((2, 2), (2, -2)), npts=32, + deltaT=1.0 / 1024, kappa_boost=6.0): + """Structurally-faithful packed data (Hermitian PD U, complex-symmetric V); + the same construction test_distance_grid_loguniform.py uses, duplicated + here so this file has no cross-file import-order dependency. kappa_boost + is large: the GH quadrature only differs measurably from the fixed grid + once the distance peak is narrow (high effective SNR).""" + 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 test_gh_nodes_changes_the_constructed_likelihood_value(saved_gh_env): + """The core claim of (a): setting the resolved node count through + core.set_distmarg_gh_nodes -- exactly what the driver's CLI path now does + -- must change what JAXDistPhiMargLikelihood.value() returns, at a FIXED + theta and FIXED data. A fresh likelihood object per setting (each gets + its own jax.jit closure), so no stale JIT trace masks the difference.""" + data = _synth() + theta4 = np.array([1.1, -0.3, 0.6, 1.0]) + + core_mod.set_distmarg_gh_nodes(0) + like_grid = JAXDistPhiMargLikelihood(data, 1.0, 10000.0, nphi=8, n_grid=64, + interp="sinc") + v_grid = like_grid.value(theta4) + + core_mod.set_distmarg_gh_nodes(48) + like_gh = JAXDistPhiMargLikelihood(data, 1.0, 10000.0, nphi=8, n_grid=64, + interp="sinc") + v_gh = like_gh.value(theta4) + + assert np.isfinite(v_grid) and np.isfinite(v_gh) + assert abs(v_gh - v_grid) > 1e-6, ( + "distance-gh-nodes=48 vs 0 must give a DIFFERENT lnL on this fixture " + "(grid=%.6f, gh=%.6f) -- if these agree, the resolved count is not " + "reaching the kernel" % (v_grid, v_gh)) + + +def test_gh_nodes_zero_reproduces_the_legacy_grid_exactly(saved_gh_env): + """The companion safety property: --distance-gh-nodes 0 (the default) + must be BIT IDENTICAL to never having called the setter at all.""" + data = _synth() + theta4 = np.array([1.1, -0.3, 0.6, 1.0]) + + core_mod._DISTMARG_GH_N = 0 + like_untouched = JAXDistPhiMargLikelihood(data, 1.0, 10000.0, nphi=8, + n_grid=64, interp="sinc") + v_untouched = like_untouched.value(theta4) + + core_mod.set_distmarg_gh_nodes(48) + core_mod.set_distmarg_gh_nodes(0) # explicit round trip back to off + like_reset = JAXDistPhiMargLikelihood(data, 1.0, 10000.0, nphi=8, n_grid=64, + interp="sinc") + v_reset = like_reset.value(theta4) + + assert v_untouched == v_reset, (v_untouched, v_reset) + + +# --------------------------------------------------------------------------- +# One subprocess check against the real driver (--inj-mode, tiny budget) +# --------------------------------------------------------------------------- + +def _run_driver(args, timeout=240): + env = dict(os.environ, PYTHONPATH=_CODE, OMP_NUM_THREADS="1", + JAX_PLATFORMS="cpu", JAX_ENABLE_X64="1") + env.pop("JAX_ILE_DISTMARG_GH", None) + return subprocess.run([sys.executable, _JAXDRIVER] + args, + capture_output=True, text=True, env=env, + cwd=tempfile.mkdtemp(), timeout=timeout) + + +_INJ_ARGS = [ + "--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", + # Deliberately invalid so the run stops right after construction -- + # cheap, and reaches the same point test_distance_grid_loguniform.py's + # own "reaches and uses" subprocess test relies on. + "--time-marginalization-quadrature", "bandlimited", + "--n-max", "1", "--n-chunk", "1", +] + + +def test_driver_reaches_and_reports_the_resolved_gh_nodes_end_to_end(): + """--distance-gh-nodes must reach the run log through the real subprocess + entry point, exactly as the equivalent test does for + --distance-grid-points in test_distance_grid_loguniform.py.""" + p = _run_driver(_INJ_ARGS + ["--distance-gh-nodes", "48"]) + out = p.stdout + p.stderr + assert "TypeError" not in out, out[-500:] + assert "nodes=48" in out and "--distance-gh-nodes" in out, out[-800:] + assert "time_quadrature='bandlimited' is not valid" in out, out[-500:] + + +def test_driver_reports_zero_nodes_without_the_flag(): + p = _run_driver(_INJ_ARGS) + out = p.stdout + p.stderr + assert "nodes=0" in out and "legacy uniform grid" in out, out[-500:] diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_is_proposal_jitter.py b/MonteCarloMarginalizeCode/Code/test/jax/test_is_proposal_jitter.py new file mode 100644 index 000000000..fd61c1a9e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_is_proposal_jitter.py @@ -0,0 +1,563 @@ +#!/usr/bin/env python +"""Issue #227: a Gaussian importance proposal must be SCORED under the matrix it +was DRAWN from. + +WHAT WENT WRONG. Seven sites drew ``theta ~ N(mu, cov + 1e-12 I)`` by Cholesky +and then evaluated ``logq`` under bare ``cov``. The regularizer is ABSOLUTE, so +it is negligible only while ``cov`` is O(1). MEASURED on the real S250114ax +H1/L1 point this was found on (rho ~ 49): ``--mode map`` reports a Fisher +diagonal of [5.5e5 9.8e4 5.3e5 5.3e5 1.6e4 1.6], i.e. angular scales ~1e-3 rad +(the issue quotes ~1e-5; 1e-3 is what the Fisher there actually gives, and the +argument does not need the smaller number). ``run_laplace_is``'s adaptation then +contracts far below that: ``cov`` reached 3e-21, the Mahalanobis term became +``(1e-6/sqrt(3e-21))**2 ~ 3e8`` per dimension, and ``--mode laplace-is`` -- the +driver's DEFAULT -- returned ``lnZ = 5.85e9`` with ``neff = 1.0`` and exited 0. + +WHY THESE TESTS LOOK LIKE THIS. A unit test of the jitter helper cannot see the +defect: the defect is not in either matrix, it is in the two of them being +different at the call site. So the behavioural tests below drive the REAL +``run_laplace_is`` on a synthetic likelihood (pure numpy -- no frames, no PSDs, +no jax evaluation), and a separate AST test gates the CLASS across both files, +because a fix at one of seven sites is not a fix. + +FLOATING POINT. Nothing in this file is precision-sensitive: the synthetic +likelihood, the proposal algebra and the reference integral are all numpy +float64 regardless of jax's x64 flag, and the 0.1-nat tolerance on the reference +comparison is far above float32 resolution. x64 is still requested below so that running +this file FIRST in a session cannot change what any later file sees. +""" + +import ast +import importlib.machinery +import importlib.util +import io +import contextlib +import os + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + +_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") +_SAMPLERS = os.path.join(_CODE, "RIFT", "likelihood", "jax_ile", "samplers.py") + + +def _driver(): + loader = importlib.machinery.SourceFileLoader("_isj_drv", _DRIVER) + spec = importlib.util.spec_from_loader("_isj_drv", loader) + mod = importlib.util.module_from_spec(spec) + mod.__name__ = "_isj_drv" # keep the __main__ guard from firing + loader.exec_module(mod) + return mod + + +# -------------------------------------------------------------------------- +# Synthetic likelihood: an isotropic Gaussian in the five angles. +# +# ``sig`` is the whole experiment. sig ~ 1e-5 is the production regime this +# issue was found in (rho ~ 49) and the regime no prior pilot can resolve; +# sig ~ 0.15 is a posterior a prior pilot CAN find, and is where the estimator +# is supposed to work and must keep working. +# -------------------------------------------------------------------------- +_MU = np.array([1.3, 0.2, 1.0, 1.1, 3.0]) + + +class _GaussianAngles(object): + def __init__(self, sig, peak): + self.sig = float(sig) + self.peak = float(peak) + + def log_likelihood(self, *cols): + th = np.stack([np.asarray(c) for c in cols], axis=-1) + d = (th[..., :5] - _MU[None, :]) / self.sig + return self.peak - 0.5 * np.sum(d * d, axis=-1) + + +def _run(sig, peak, n_max=120000, seed=3): + """Drive the real ``run_laplace_is``; return its outputs and its log.""" + mod = _driver() + optp = mod.build_parser() + opts, _ = optp.parse_args(["--inj-mode", "--n-max", str(n_max), + "--seed", str(seed)]) + rng = np.random.default_rng(seed) + buf = io.StringIO() + with contextlib.redirect_stdout(buf): + out = mod.run_laplace_is(_GaussianAngles(sig, peak), opts, rng, 5, False) + return mod, opts, out, buf.getvalue() + + +def _reference_logZ(mod, opts, sig, peak, n=2000000, seed=99): + """Independent estimate of ln Z = ln int L(theta) p(theta) dtheta. + + Deliberately shares NO machinery with the estimator under test: the + proposal is written out here, drawn directly, and centred on the known peak + at twice its width, which makes this a near-perfect importance proposal + (ESS ~ 0.25 n). It uses the driver's ``log_prior`` only because both + estimators must integrate against the SAME prior to be comparable. + """ + rng = np.random.default_rng(seed) + s = 2.0 * sig + th = _MU[None, :] + s * rng.standard_normal((n, 5)) + lnL = _GaussianAngles(sig, peak).log_likelihood(*[th[:, i] for i in range(5)]) + logp = mod.log_prior(th, opts, False) + logq = (-0.5 * np.sum(((th - _MU[None, :]) / s) ** 2, axis=1) + - 0.5 * 5 * np.log(2 * np.pi * s * s)) + lw = lnL + logp - logq + lw = lw[np.isfinite(lw)] + m = lw.max() + w = np.exp(lw - m) + return float(m + np.log(w.mean())) + + +### +### 1. The helper's contract +### + +def test_regularize_cov_regularizer_is_relative_to_the_covariance(): + """The whole defect is an absolute epsilon meeting a 1e-21 covariance.""" + from RIFT.likelihood.jax_ile.samplers import regularize_cov + base = np.diag([1.0, 2.0, 3.0, 4.0, 5.0]) + for scale in (1.0, 1e-10, 1e-20, 1e-30): + cov = scale * base + out = regularize_cov(cov) + added = np.diag(out - cov) + assert np.allclose(added, added[0]) # isotropic + # the nudge is a fixed FRACTION of the covariance scale, never a floor + assert added[0] == pytest.approx(1e-12 * np.trace(cov) / 5, rel=1e-12) + assert 0 < added[0] < 1e-11 * float(np.max(np.diag(cov))) + + +def test_regularize_cov_is_scale_equivariant(): + """f(a C) == a f(C): the property an absolute jitter does not have, and the + reason the proposal can no longer be swamped by its own regularizer.""" + from RIFT.likelihood.jax_ile.samplers import regularize_cov + rng = np.random.default_rng(0) + A = rng.standard_normal((5, 5)) + C = A @ A.T + for a in (1e-8, 1.0, 1e8): + assert np.allclose(regularize_cov(a * C), a * regularize_cov(C), + rtol=1e-12, atol=0.0) + + +def test_regularize_cov_still_conditions_a_degenerate_covariance(): + """trace == 0 has no scale to be relative to; it must still come back + Cholesky-able rather than raising inside a sampler.""" + from RIFT.likelihood.jax_ile.samplers import regularize_cov + out = regularize_cov(np.zeros((4, 4))) + np.linalg.cholesky(out) # must not raise + assert np.all(np.diag(out) > 0) + + +### +### 2. The CLASS gate. Seven sites shared the pattern; a fix at one is not a fix. +### + +def _cholesky_calls_with_an_inline_identity(path): + """Every ``np.linalg.cholesky(X)`` in ``path`` whose argument builds an + identity inline -- i.e. regularizes a matrix at the DRAW while some other + expression is what gets scored. Returns (lineno, source) pairs.""" + with open(path) as f: + src = f.read() + tree = ast.parse(src) + lines = src.splitlines() + bad = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + fn = node.func + if not (isinstance(fn, ast.Attribute) and fn.attr == "cholesky"): + continue + arg = node.args[0] if node.args else None + if arg is None: + continue + for sub in ast.walk(arg): + if (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Attribute) + and sub.func.attr in ("eye", "identity")): + bad.append((node.lineno, lines[node.lineno - 1].strip())) + break + return bad + + +@pytest.mark.parametrize("path", [_DRIVER, _SAMPLERS]) +def test_no_cholesky_regularizes_a_matrix_inline(path): + """Structural gate on the #227 pattern. + + HONEST SCOPE: this is a shape check, not a correctness proof -- it cannot + see a caller that passes two different *named* matrices. It exists because + the behavioural test below reaches only ONE of the seven sites (the default + mode); the other six are inside flowMC / NUTS / SMC paths that need numpyro, + flowMC and a real likelihood to run. Reverting ANY of the seven to + ``cholesky(cov + 1e-12*np.eye(d))`` fails this test. + + The fix is to call ``regularize_cov(cov)`` once and hand the SAME object to + the Cholesky and to _gaussian_logq/_mixture_logq. + """ + bad = _cholesky_calls_with_an_inline_identity(path) + assert not bad, ( + "%s regularizes a covariance inside np.linalg.cholesky(...):\n%s\n" + "Call regularize_cov(cov) once and score under the SAME matrix (#227)." + % (os.path.basename(path), + "\n".join(" line %d: %s" % b for b in bad))) + + +def test_the_class_gate_can_actually_fail(): + """A structural test that never fails on anything is not coverage. Prove + the detector fires on the exact pre-fix source line.""" + import tempfile + src = ("import numpy as np\n" + "def f(cov, dim):\n" + " Lc = np.linalg.cholesky(cov + 1e-12 * np.eye(dim))\n" + " return Lc\n") + with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f: + f.write(src) + tmp = f.name + try: + found = _cholesky_calls_with_an_inline_identity(tmp) + assert len(found) == 1 and found[0][0] == 3 + finally: + os.unlink(tmp) + + +def test_regularize_cov_is_the_single_definition_both_files_use(): + """The driver must not grow its own copy: the bug was one rule written + twice (drawn one way, scored another), and a second helper is the same + mistake one level up.""" + with open(_DRIVER) as f: + drv = f.read() + assert "from RIFT.likelihood.jax_ile.samplers import regularize_cov" in drv + assert "def regularize_cov" not in drv + + +### +### 3. Behaviour: the default mode on a posterior no prior pilot can resolve. +### This is the configuration that returned 5.85e9 on real data. +### + +_NARROW = dict(sig=2e-5, peak=1539.0) + + +def test_laplace_is_never_reports_evidence_above_the_peak_likelihood(): + """The #227 regression assertion, stated as the issue asks for it. + + For a NORMALIZED prior, Z = E_prior[L] <= max L, so ln Z <= max lnL always. + The shipped code returned ln Z = 5.85e9 against a peak lnL of 1808 on real + S250114ax data, and 4.9e9 - 5.6e9 here on every seed tried. Reproduces in + seconds with no frames, because the defect is in the proposal algebra rather + than in the physics. + """ + _mod, _opts, out, _log = _run(n_max=200000, seed=11, **_NARROW) + logZ, _sig, _neff, _n, _theta, lnL, _logw = out + max_lnL = float(np.nanmax(lnL[np.isfinite(lnL)])) + assert np.isnan(logZ) or logZ <= max_lnL + 5.0, ( + "ln Z = %r exceeds max lnL = %r; the weights were computed against a " + "distribution that was never sampled (#227)" % (logZ, max_lnL)) + + +@pytest.mark.parametrize("seed", [3, 5, 11, 12]) +def test_a_proposal_that_walked_off_the_peak_is_never_published(seed): + """Fixing the jitter is NOT sufficient, and this is the test that says so. + + With the draw and the density matched, the same configuration returns a + SELF-CONSISTENT number computed from a proposal that never found the peak: + a plausible wrong answer in place of an implausible one. + + THE ASSERTION IS A PROPERTY, NOT AN OUTCOME, and that is a correction from + external review. This test used to assert ``isnan`` on every seed, which + quietly made a SUCCESS into a test failure: had the adaptation ever recovered + the peak here, the suite would have reported a regression. A test that can + only pass while the code fails cannot witness the guard being too aggressive, + which is precisely the risk under review. So what is pinned is the property + that matters -- *no inaccurate number is ever published* -- with a recovered, + accurate answer explicitly allowed through. + """ + mod, opts, out, log = _run(n_max=120000, seed=seed, **_NARROW) + logZ = out[0] + if np.isnan(logZ): + assert "Markov floor" in log + return + ref = _reference_logZ(mod, opts, **_NARROW) + assert abs(logZ - ref) < 0.5, ( + "published lnZ = %r from a collapsed proposal (reference %r)" % (logZ, ref)) + + +def test_the_pilot_floor_is_a_markov_bound_not_the_raw_estimate(): + """External review, P1: the pilot is UNBIASED for Z, not a bound on it. + + For a mode of prior mass m a single lucky draw gives ~L_max/n_pilot against a + truth of ~L_max*m, overshooting by 1/(n_pilot*m). MEASURED, not argued: at a + synthetic width of 0.05 rad (n_pilot*m = 1.6e-3) the pilot ran up to +5.46 + nats ABOVE the truth, P = 1.1e-3 over 900 seeds -- so a raw-estimate floor + set at 5 nats rejects correct answers at about that rate. + + Markov needs only non-negativity and unbiasedness: P(Zhat >= Z/rate) <= rate, + so ln Zhat + ln rate is a lower confidence bound at level 1 - rate. The + threshold is a chosen false-positive rate, and it is distribution-free -- it + does not assume the pilot resolved anything, which matters because the + pilot's ESS is ~1 in every regime where this guard does any work. + """ + mod = _driver() + for rate in (3.4e-4, 1e-2, 0.5): + for lz in (-3.0, 0.0, 1234.5): + assert mod.prior_pilot_floor(lz, rate) == pytest.approx( + lz + np.log(rate), rel=0, abs=1e-12) + assert mod.prior_pilot_floor(lz, rate) < lz # always a DISCOUNT + # a smaller admitted false-positive rate must push the floor DOWN, never up + assert mod.prior_pilot_floor(0.0, 1e-6) < mod.prior_pilot_floor(0.0, 1e-2) + # the shipped rate is the one the operating curve was read at + assert mod.PILOT_FLOOR_FP_RATE == pytest.approx(np.exp(-8.0), rel=0.02) + # a pilot that estimated nothing must not manufacture a floor + assert mod.prior_pilot_floor(np.nan) == -np.inf + assert mod.prior_pilot_floor(-np.inf) == -np.inf + + +def test_an_inflated_pilot_does_not_reject_an_accurate_high_ess_answer(): + """The regression case external review asked for, and it is a REAL run. + + Reviewer's scenario: a sparse pilot hit both inflates the pilot's estimate + AND seeds a good proposal, so the correct adapted answer is compared against + a reference that its own lucky draw pushed up. Found by sweeping 5400 runs + for the shape -- pilot ABOVE truth, adapted accurate, high ESS. At + sig = 0.15, seed = 885 the pilot lands +1.35 nats above the reference while + the adapted estimate is right to 0.001 nats at neff ~ 1.4e4. + + The exhaustive sweep is the honest part: over those 5400 runs the largest + pilot-minus-adapted gap on an accurate run was +1.654 nats, so this case does + NOT reach the shipped threshold and no false positive was ever observed. + What this pins is the margin -- lower the threshold under ~1.7 nats, or go + back to comparing against the raw pilot at a 5-nat cut without the Markov + discount, and a correct high-ESS answer starts being thrown away. + """ + mod, opts, out, log = _run(sig=0.15, peak=100.0, n_max=120000, seed=885) + logZ, _s, neff = out[0], out[1], out[2] + ref = _reference_logZ(mod, opts, sig=0.15, peak=100.0) + assert neff > 1000.0 + assert abs(logZ - ref) < 0.1, "the case no longer has the reviewer's shape" + assert not np.isnan(logZ), "an accurate, high-ESS answer was rejected" + assert "Markov floor" not in log + + +### +### 4. The regime the estimator DOES work in must be untouched. +### + +_HEALTHY = [(0.30, 20.0), (0.20, 20.0), (0.15, 20.0)] + + +@pytest.mark.parametrize("sig,peak", _HEALTHY) +@pytest.mark.parametrize("seed", [3, 5, 12]) +def test_laplace_is_matches_an_independent_reference_where_it_works(sig, peak, seed): + """A posterior a prior pilot CAN resolve. Two jobs: + + 1. THE FIX CHANGED NOTHING HERE. The proposal covariance is ~1e-2, ten + orders above the old absolute jitter, so old and new differ in the + eleventh significant figure: on (sig=0.15, seed=3) the parent commit + gives 8.7451703051592702 and this one 8.7451703051683118. All nine + cases below print identically to 5 dp on both sides. + 2. THE GUARD IS NOT TRIGGER-HAPPY. An earlier version of this guard keyed + on the pilot's ESS, and ESS turned out to be a poor predictor: it fired + on (sig=0.15, seed=4) -- a run whose final neff was 17843 and whose + answer was right to 0.003 nats -- and made the answer 1.5 nats WORSE. + Several seeds and widths are swept because a single seed hid that. + """ + mod, opts, out, log = _run(sig=sig, peak=peak, n_max=120000, seed=seed) + logZ, _sig, neff, _n, _theta, _lnL, _logw = out + ref = _reference_logZ(mod, opts, sig=sig, peak=peak) + assert neff > 1000.0 + assert abs(logZ - ref) < 0.1, "ln Z = %.5f vs reference %.5f" % (logZ, ref) + assert "laplace-is]" not in log, "the guard fired on a healthy run: %s" % log + + +def test_the_evidence_sanity_rule_is_wired_into_both_driver_estimators(): + """WIRING, honestly labelled: a call-site check, not a behavioural one. + + ``_finalize_evidence`` (ln Z <= max lnL, neff >= 1.5) is the library + samplers' rule; these two driver estimators applied NO rule at all, which is + why #227's 5.8e9 was reported as a success. With the jitter fixed there is + no longer a synthetic that reaches it -- the pilot comparison above catches + the collapse first -- so what is testable is that the belt is still attached + to the braces. + """ + with open(_DRIVER) as f: + tree = ast.parse(f.read()) + for name in ("run_laplace_is", "run_nuts"): + fn = next(n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == name) + calls = {n.func.id for n in ast.walk(fn) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name)} + assert "_finalize_evidence" in calls, "%s does not finalize its evidence" % name + + +### +### 5. A non-finite evidence must FAIL the event, not be published as one. +### + +def test_require_finite_evidence_passes_a_number_and_refuses_a_nan(): + mod = _driver() + mod.require_finite_evidence(1462.36, 4.5, "laplace-is") # must not raise + for bad in (np.nan, np.inf, -np.inf): + with pytest.raises(RuntimeError) as ei: + mod.require_finite_evidence(bad, 1.0, "laplace-is") + assert "laplace-is" in str(ei.value) + + +def test_analyze_one_refuses_before_it_writes_either_artifact(): + """WIRING, not presence. On real S250114ax data the pre-fix driver wrote + ``lnL = 5.848741051595e+09`` into ``out_0_.dat`` and exited 0; the estimator + now returns nan there, and a nan row published to a CIP fit is no better. + So the refusal has to come BEFORE write_samples/write_dat -- checked by + statement order inside ``analyze_one``, because a call that runs after the + files exist is the same defect with a tidier log. + """ + with open(_DRIVER) as f: + tree = ast.parse(f.read()) + fn = next(n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "analyze_one") + seen = {} + for node in ast.walk(fn): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + seen.setdefault(node.func.id, node.lineno) + for name in ("require_finite_evidence", "write_samples", "write_dat"): + assert name in seen, "analyze_one never calls %s" % name + assert seen["require_finite_evidence"] < seen["write_samples"] + assert seen["require_finite_evidence"] < seen["write_dat"] + + +### +### 6. The pilot and the adapted estimator must integrate the SAME quantity. +### + +def test_pilot_and_explicit_weights_share_one_normalization_under_limit_distance(): + """THE MATH, not the wiring -- the wiring is the AST test below, and this + test is deliberately unable to see a caller that drops the term. (It was + written as a behavioural test first, and a mutation that deleted the term + from run_laplace_is left it green: it reimplements the two weight + expressions rather than reaching them, which is the same "tests the helper, + not the call site" mistake the #227 defect itself is an instance of.) + + What it does establish is why the term has to be there at all: the 5-nat + collapse guard compares two DIFFERENT estimators, and it means nothing + unless both are on the same normalization. + + They are built differently on purpose. The adapted loop forms + ``ln w = lnL + ln p - ln q`` explicitly, and ``log_prior`` is normalized over + the physical ``[d_min,d_max]``; the prior pilot draws from ``sample_prior``, + which samples the prior RESTRICTED to the ``--limit-distance`` box, so its + weights need ``-log_distance_box_correction`` to land on that same scale. + Drop that term -- which the docstring of ``log_distance_box_correction`` used + to tell a reader to do, because it named ``run_laplace_is`` as an estimator + that must not subtract -- and the two differ by a CONSTANT. A constant + offset does not make the guard noisy, it makes it wrong in one direction: + the box here is a factor of 17.86 in prior mass, i.e. 2.88 nats, and a wider + limit silently walks past the 5-nat threshold and disables the guard + entirely while every existing test still passes. + + Taking L == 1 makes both estimators exactly computable: Z = int_box p_full + = the box's prior mass fraction, so the assertion is against a number + derived on paper rather than against the other estimator. + """ + mod = _driver() + optp = mod.build_parser() + opts, _ = optp.parse_args(["--inj-mode", "--d-min", "10.0", "--d-max", "1000.0", + "--limit-distance", "200.0,400.0"]) + lo, hi = mod.resolve_distance_limit(opts) + assert (lo, hi) == (200.0, 400.0), "the box did not narrow; test is vacuous" + + mass = (hi ** 3 - lo ** 3) / (opts.d_max ** 3 - opts.d_min ** 3) + assert 0.0 < mass < 0.1 # a box small enough to matter + + rng = np.random.default_rng(7) + n = 400000 + theta_p, _ = mod.sample_prior(n, opts, rng, True) + + # (a) the PILOT form: proposal == restricted prior, lnL == 0 + lw_pilot = np.zeros(n) - mod.log_distance_box_correction(opts, True) + Z_pilot = np.exp(lw_pilot).mean() + + # (b) the ADAPTED form: ln w = lnL + ln p - ln q, written out against a + # uniform-in-d proposal over the same box (q is known exactly here). + # The angles keep their prior draw, so q's angular factor IS the driver's + # own angles-only log_prior -- taken from the code rather than re-derived + # here, since a hand-written copy of a density is the mistake this whole + # file is about. Only the distance proposal differs (uniform on the box). + th = np.array(theta_p, dtype=float, copy=True) + th[:, 5] = rng.uniform(lo, hi, size=n) + logq = mod.log_prior(th, opts, False) - np.log(hi - lo) + lw_adapt = mod.log_prior(th, opts, True) - logq + Z_adapt = np.exp(lw_adapt[np.isfinite(lw_adapt)]).mean() + + assert Z_pilot == pytest.approx(mass, rel=1e-12), ( + "pilot weights do not carry the full-range normalization: %r vs %r" + % (Z_pilot, mass)) + assert Z_adapt == pytest.approx(mass, rel=0.02), ( + "explicit lnL+lnp-lnq weights are on a different scale: %r vs %r" + % (Z_adapt, mass)) + # and therefore on the same scale as each other, which is the guard's premise + assert abs(np.log(Z_pilot) - np.log(Z_adapt)) < 0.05 + + +def test_the_pilot_normalization_test_would_catch_the_dropped_correction(): + """The mutation the test above exists to stop, stated so it is visible: with + the box correction dropped the pilot is low by ln(1/mass) -- 2.7 nats here, + and larger for a wider box, which is how it would slip past a 5-nat guard.""" + mod = _driver() + optp = mod.build_parser() + opts, _ = optp.parse_args(["--inj-mode", "--d-min", "10.0", "--d-max", "1000.0", + "--limit-distance", "200.0,400.0"]) + corr = mod.log_distance_box_correction(opts, True) + assert corr > 0.0 + assert corr == pytest.approx(2.8824, abs=0.01) + + +def test_the_pilot_reference_is_wired_to_the_box_correction(): + """WIRING, and the test that actually holds run_laplace_is to the maths above. + + The numerical test cannot: it writes the two weight expressions out itself, + so deleting ``- log_distance_box_correction(...)`` from run_laplace_is left + it green. This one reads the call site. ``evidence_from_logweights`` is + called twice in the file -- once in run_prior_mc, once for this pilot -- and + BOTH are "proposal == prior" estimators, so both arguments must carry the + correction. The adapted loop must not, and that is asserted too, because + the failure mode is symmetric: adding the term where ln w is already + explicit breaks the comparison just as thoroughly as dropping it here. + """ + with open(_DRIVER) as f: + src = f.read() + tree = ast.parse(src) + fn = next(n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "run_laplace_is") + + calls = [n for n in ast.walk(fn) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + and n.func.id == "evidence_from_logweights"] + # TWO, and they are the two halves of the comparison: the pilot (raw prior + # draws, needs the correction) and the adapted estimate (ln w already + # explicit, must not have it). Pinning the count is what makes "the pilot + # is the one with the correction" a statement rather than a search. + assert len(calls) == 2, ( + "expected exactly two evidence estimates in run_laplace_is (prior pilot " + "and adapted), found %d" % len(calls)) + pilot = [c for c in calls if "log_distance_box_correction" in ast.dump(c.args[0])] + assert len(pilot) == 1, ( + "the prior pilot's weights do not subtract log_distance_box_correction, so " + "under --limit-distance it sits ln(1/box mass) BELOW the adapted estimate it " + "is compared against -- 2.88 nats for the box in the test above, and a wider " + "limit walks straight past the 5-nat guard and disables it (#227).") + arg = pilot[0].args[0] + assert isinstance(arg, ast.BinOp) and isinstance(arg.op, ast.Sub), ( + "the correction must be SUBTRACTED from the pilot weights, not added") + other = [c for c in calls if c is not pilot[0]] + assert isinstance(other[0].args[0], ast.Name) and other[0].args[0].id == "logw", ( + "the adapted estimate should be taken on the explicit ln w array") + + # ... and the adapted loop, which forms ln w explicitly, must NOT correct again. + for node in ast.walk(fn): + if (isinstance(node, ast.Assign) and len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id == "logw"): + assert "log_distance_box_correction" not in ast.dump(node.value), ( + "the adapted loop already forms lnL + ln p - ln q against a " + "full-range-normalized log_prior; correcting again double-counts") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_av.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_av.py new file mode 100644 index 000000000..58b8be1a4 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_av.py @@ -0,0 +1,367 @@ +"""Contract tests for the value-only JAX -> AV/portfolio adapter.""" + +import importlib.machinery +import importlib.util +import os +import numpy as np +import pytest +import jax +import jax.numpy as jnp + +from RIFT.likelihood.jax_ile import samplers + +_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 _trapezoid(y, x): + """Integrate on NumPy versions before and after ``trapz`` was removed.""" + trapezoid = getattr(np, "trapezoid", None) + return trapezoid(y, x) if trapezoid is not None else np.trapz(y, x) + + +def _driver_module(): + loader = importlib.machinery.SourceFileLoader("_jax_av_driver", _DRIVER) + spec = importlib.util.spec_from_loader(loader.name, loader) + module = importlib.util.module_from_spec(spec) + loader.exec_module(module) + return module + + +class _ToySkyLikelihood: + ANGULAR_PARAM_ORDER = ("ra", "dec", "incl") + + def __init__(self): + centre = jnp.array([2.0, 0.2, 1.0]) + scale = jnp.array([0.5, 0.35, 0.4]) + + def scalar(theta): + return -0.5 * jnp.sum(((theta - centre) / scale) ** 2) + + self._scalar = scalar + self._value_and_grad = jax.jit(jax.value_and_grad(scalar)) + self._hessian = jax.jit(jax.hessian(scalar)) + self.batch_shapes = [] + + def log_likelihood(self, *cols): + theta = jnp.stack(cols, axis=-1) + self.batch_shapes.append(tuple(theta.shape)) + return jax.vmap(self._scalar)(theta) + + def value_and_grad(self, theta): + value, grad = self._value_and_grad(jnp.asarray(theta)) + return float(value), np.asarray(grad) + + def fisher(self, theta): + return -np.asarray(self._hessian(jnp.asarray(theta))) + + +def test_fixed_shape_callback_pads_only_the_hidden_tail(): + like = _ToySkyLikelihood() + callback = samplers._fixed_shape_value_callback(like, 3, 8) + theta = np.arange(57, dtype=float).reshape(19, 3) / 20.0 + + got = callback(*theta.T) + want = np.asarray(jax.vmap(like._scalar)(jnp.asarray(theta))) + + np.testing.assert_allclose(got, want) + assert like.batch_shapes[:3] == [(8, 3), (8, 3), (8, 3)] + assert got.shape == (19,) # padded rows never escape the adapter + + +def test_physical_coordinate_priors_are_normalized(): + for name in ("ra", "dec", "psi", "incl", "phiref", "distMpc"): + lo, hi, density = samplers._av_prior_spec(name, 10.0, 100.0) + x = np.linspace(lo, hi, 20001) + np.testing.assert_allclose(_trapezoid(density(x), x), 1.0, + rtol=2e-6, atol=2e-6) + + +def test_sampling_window_does_not_renormalize_physical_prior(): + bounds = {"ra": (1.17, 1.23), "dec": (0.27, 0.33)} + ra_lo, ra_hi, ra_pdf = samplers._av_prior_spec( + "ra", 10.0, 100.0, sample_bounds=bounds) + dec_lo, dec_hi, dec_pdf = samplers._av_prior_spec( + "dec", 10.0, 100.0, sample_bounds=bounds) + + assert (ra_lo, ra_hi) == bounds["ra"] + assert (dec_lo, dec_hi) == bounds["dec"] + np.testing.assert_allclose( + _trapezoid(ra_pdf(np.linspace(ra_lo, ra_hi, 10001)), + np.linspace(ra_lo, ra_hi, 10001)), + (ra_hi - ra_lo) / (2 * np.pi), rtol=1e-10) + np.testing.assert_allclose( + _trapezoid(dec_pdf(np.linspace(dec_lo, dec_hi, 10001)), + np.linspace(dec_lo, dec_hi, 10001)), + 0.5 * (np.sin(dec_hi) - np.sin(dec_lo)), rtol=1e-9) + + +def test_prior_draw_respects_sky_sampling_window(): + bounds = {"ra": (1.17, 1.23), "dec": (0.27, 0.33)} + draw = samplers._av_prior_draw( + ("ra", "dec", "incl"), 10000, np.random.default_rng(8), + 1.0, 100.0, bounds) + assert np.all((draw[:, 0] >= 1.17) & (draw[:, 0] <= 1.23)) + assert np.all((draw[:, 1] >= 0.27) & (draw[:, 1] <= 0.33)) + assert np.std(draw[:, 2]) > 0.6 + + +@pytest.mark.parametrize("bounds,match", [ + ({"ra": (6.1, 0.1)}, "invalid sampling bounds"), + ({"dec": (-2.0, 0.1)}, "must lie within"), + ({"bogus": (0.0, 1.0)}, "absent from likelihood"), +]) +def test_sampling_window_validation(bounds, match): + with pytest.raises(ValueError, match=match): + samplers._av_sample_bounds( + _ToySkyLikelihood.ANGULAR_PARAM_ORDER, 1.0, 100.0, + sample_bounds=bounds) + + +def test_fisher_sky_seed_targets_sky_but_randomizes_other_coordinates(): + like = _ToySkyLikelihood() + callback = samplers._fixed_shape_value_callback(like, 3, 64) + cloud, modes, mode_lnL = samplers._fisher_sky_seed( + like, like.ANGULAR_PARAM_ORDER, callback, np.random.default_rng(7), + 1.0, 100.0, n_seed=600, n_pilot=200, n_modes=2, + sky_inflate=1.5, prior_frac=0.1) + + assert cloud.shape == (600, 3) + assert modes.shape[1] == 3 and np.all(np.isfinite(mode_lnL)) + assert np.std(cloud[:, 2]) > 0.5 # inclination came from its broad prior + assert np.min(np.abs(np.angle(np.exp(1j * (cloud[:, 0] - 2.0))))) < 0.1 + assert np.all((-np.pi / 2 <= cloud[:, 1]) & (cloud[:, 1] <= np.pi / 2)) + + +def test_fisher_sky_seed_uses_marginal_not_conditional_covariance(): + class CorrelatedSky(_ToySkyLikelihood): + def __init__(self): + # F[:2,:2]^{-1} is diag(0.01), but marginalizing the correlated + # nuisance coordinate makes var(ra)=1.0. + fisher = jnp.array([[100.0, 0.0, 9.94987437], + [0.0, 100.0, 0.0], + [9.94987437, 0.0, 1.0]]) + centre = jnp.array([2.0, 0.2, 1.0]) + + def scalar(theta): + delta = theta - centre + return -0.5 * delta @ fisher @ delta + + self._scalar = scalar + self._value_and_grad = jax.jit(jax.value_and_grad(scalar)) + self._hessian = jax.jit(jax.hessian(scalar)) + self.batch_shapes = [] + + like = CorrelatedSky() + callback = samplers._fixed_shape_value_callback(like, 3, 64) + cloud, modes, _ = samplers._fisher_sky_seed( + like, like.ANGULAR_PARAM_ORDER, callback, np.random.default_rng(9), + 1.0, 100.0, n_seed=4000, n_pilot=10, n_modes=1, + sky_inflate=1.0, prior_frac=0.0, + initial_points=np.array([[2.0, 0.2, 1.0]])) + + dra = np.angle(np.exp(1j * (cloud[:, 0] - modes[0, 0]))) + assert np.std(dra) > 0.5 + + +def test_fixed_distance_likelihood_is_five_dimensional(monkeypatch): + from RIFT.likelihood.jax_ile import wrapper + + # Avoid constructing physical data: instantiate the class shell with the + # same public/JAX scalar contract used by the fixed-distance view. + raw = object.__new__(wrapper.JAXExtrinsicLikelihood) + raw.data = object() + raw.interp = "linear" + raw.phase_marginalization = False + raw.time_quadrature = "simpson" + raw._scalar = lambda theta: jnp.sum(theta ** 2) + raw.log_likelihood = lambda *cols: jnp.sum(jnp.stack(cols) ** 2, axis=0) + + fixed = wrapper.JAXFixedDistanceLikelihood(raw, 17.0) + theta = np.arange(15, dtype=float).reshape(3, 5) / 10.0 + got = fixed.log_likelihood(*theta.T) + want = np.sum(theta ** 2, axis=1) + 17.0 ** 2 + + assert fixed.ANGULAR_PARAM_ORDER == ("ra", "dec", "psi", "incl", "phiref") + np.testing.assert_allclose(got, want) + assert fixed.fisher(theta[0]).shape == (5, 5) + + +def test_fixed_distance_likelihood_can_center_periodic_phase(): + from RIFT.likelihood.jax_ile import wrapper + + raw = object.__new__(wrapper.JAXExtrinsicLikelihood) + raw.data = object() + raw.interp = "linear" + raw.phase_marginalization = False + raw.time_quadrature = "simpson" + raw._scalar = lambda theta: theta[4] + raw.log_likelihood = lambda *cols: cols[4] + + fixed = wrapper.JAXFixedDistanceLikelihood(raw, 17.0, phase_shift=np.pi) + physical = np.array([2.0, 0.2, 0.5, 1.0, 0.0]) + sampler_theta = fixed.to_sampler_coordinates(physical) + + assert fixed.ANGULAR_PARAM_ORDER[-1] == "phiref_shifted" + np.testing.assert_allclose(sampler_theta[-1], np.pi) + np.testing.assert_allclose(fixed.to_physical_coordinates(sampler_theta), physical) + np.testing.assert_allclose(fixed.value(sampler_theta), 0.0, atol=1e-12) + + +def test_rotated_phase_wrapper_roundtrips_and_preserves_likelihood(): + from RIFT.likelihood.jax_ile import wrapper + + raw = object.__new__(wrapper.JAXExtrinsicLikelihood) + raw.data = object(); raw.interp = "linear" + raw.phase_marginalization = False; raw.time_quadrature = "simpson" + raw._scalar = lambda theta: theta[2] + 2.0 * theta[4] + raw.log_likelihood = lambda ra, dec, psi, incl, phase, dist: psi + 2 * phase + fixed = wrapper.JAXFixedDistanceLikelihood(raw, 17.0, phase_shift=np.pi) + rotated = wrapper.JAXRotatedPhaseLikelihood(fixed) + physical = np.array([1.2, 0.3, 0.5, 1.05, 0.2]) + theta = rotated.to_sampler_coordinates(physical) + + assert rotated.ANGULAR_PARAM_ORDER == ( + "ra", "dec", "phase_p", "incl", "phase_m") + np.testing.assert_allclose(rotated.to_physical_coordinates(theta), physical) + np.testing.assert_allclose(rotated.value(theta), 0.9) + for name in ("phase_p", "phase_m"): + lo, hi, density = samplers._av_prior_spec(name, 1.0, 100.0) + assert (lo, hi) == (0.0, 4.0 * np.pi) + np.testing.assert_allclose(density(np.array([1.0])), 1.0 / (4.0 * np.pi)) + + +def test_av_and_seeded_portfolio_return_driver_contract(): + common = dict(d_min=1.0, d_max=100.0, nmax=20000, neff=25, + n_chunk=2000, eval_chunk=512, seed=11, verbose=False) + av = samplers.adaptive_volume_sample(_ToySkyLikelihood(), + sampler_method="AV", **common) + portfolio = samplers.adaptive_volume_sample( + _ToySkyLikelihood(), sampler_method="portfolio", + seed_method="fisher-sky", seed_pilot=100, seed_modes=2, + seed_points=500, **common) + + for result in (av, portfolio): + assert result["theta"].ndim == 2 and result["theta"].shape[1] == 3 + assert len(result["theta"]) == len(result["lnL"]) + assert np.isfinite(result["logZ"]) + assert result["neff"] >= 25 + assert result["n_eval"] <= 20000 + assert result["eval_chunk"] == 512 + assert av["log_weight"] is not None # weighted retained AV population + assert portfolio["log_weight"] is None # portfolio performed its fair draw + + +def test_pure_av_runs_inside_a_narrow_sky_sampling_window(): + bounds = {"ra": (1.7, 2.3), "dec": (-0.1, 0.5)} + result = samplers.adaptive_volume_sample( + _ToySkyLikelihood(), 1.0, 100.0, sampler_method="AV", + sample_bounds=bounds, nmax=20000, neff=25, n_chunk=2000, + eval_chunk=512, seed=14) + assert result["sample_bounds"]["ra"] == bounds["ra"] + assert np.all((result["theta"][:, 0] >= 1.7) & + (result["theta"][:, 0] <= 2.3)) + assert np.all((result["theta"][:, 1] >= -0.1) & + (result["theta"][:, 1] <= 0.5)) + + +def test_caller_supplied_oracle_cloud_bootstraps_portfolio(): + rng = np.random.default_rng(31) + centre = np.array([2.0, 0.2, 1.0]) + cloud = centre + rng.normal(size=(500, 3)) * np.array([0.2, 0.15, 0.2]) + result = samplers.adaptive_volume_sample( + _ToySkyLikelihood(), 1.0, 100.0, sampler_method="portfolio", + initial_samples=cloud, nmax=20000, neff=25, n_chunk=2000, + eval_chunk=512, seed=31) + + np.testing.assert_array_equal(result["seed_cloud"], cloud) + assert result["neff"] >= 25 + + +def test_portfolio_gmm_uses_two_components_for_periodic_boundary_modes(): + rng = np.random.default_rng(32) + cloud = np.column_stack([ + rng.normal(2.0, 0.1, 500), rng.normal(0.2, 0.08, 500), + rng.normal(1.0, 0.1, 500)]) + result = samplers.adaptive_volume_sample( + _ToySkyLikelihood(), 1.0, 100.0, sampler_method="portfolio", + portfolio_members=("GMM",), initial_samples=cloud, + nmax=10000, neff=20, n_chunk=1000, eval_chunk=256, seed=32) + + gmm = result["sampler"].portfolio_realizations[0] + assert gmm.integrator.n_comp == 2 + + +def test_oracle_cloud_and_internal_seed_are_mutually_exclusive(): + with pytest.raises(ValueError, match="mutually exclusive"): + samplers.adaptive_volume_sample( + _ToySkyLikelihood(), 1.0, 100.0, sampler_method="portfolio", + initial_samples=np.ones((10, 3)), seed_method="fisher-sky", + nmax=100, neff=2, n_chunk=20) + + +def test_driver_exposes_sampler_as_an_orthogonal_backend(monkeypatch): + monkeypatch.delenv("JAX_ILE_DISTMARG_GH", raising=False) + driver = _driver_module() + parser = driver.build_parser() + opts, _ = parser.parse_args([ + "--mode", "flowmc-phipsimarg", "--distance-marginalization", + "--sampler-method", "portfolio", "--sampler-portfolio", "AV,GMM", + "--jax-av-seed", "fisher-sky", "--n-eff", "321"]) + driver.check_critical_and_report(opts, parser) + + assert opts.mode == "flowmc-phipsimarg" # still chooses likelihood geometry + assert opts.sampler_method == "portfolio" + assert opts.sampler_portfolio == ["AV,GMM"] + assert opts.jax_av_seed == "fisher-sky" + assert opts.n_eff == 321 + + +def test_driver_validates_and_activates_av_sky_limits(monkeypatch): + monkeypatch.delenv("JAX_ILE_DISTMARG_GH", raising=False) + driver = _driver_module() + parser = driver.build_parser() + opts, _ = parser.parse_args([ + "--sampler-method", "AV", "--limit-right-ascension", "1.17,1.23", + "--limit-declination", "0.27,0.33"]) + driver.check_critical_and_report(opts, parser) + assert driver.resolve_av_angular_limits(opts) == { + "ra": (1.17, 1.23), "dec": (0.27, 0.33)} + + opts_bad, _ = parser.parse_args([ + "--sampler-method", "AV", "--limit-right-ascension", "6.1,0.1"]) + with pytest.raises(SystemExit): + driver.check_critical_and_report(opts_bad, parser) + + +def test_driver_refuses_seed_knobs_without_jax_av_backend(monkeypatch): + monkeypatch.delenv("JAX_ILE_DISTMARG_GH", raising=False) + driver = _driver_module() + parser = driver.build_parser() + opts, _ = parser.parse_args(["--jax-av-seed", "fisher-sky"]) + with pytest.raises(SystemExit): + driver.check_critical_and_report(opts, parser) + + +def test_driver_distance_limits_follow_likelihood_dimension(monkeypatch): + monkeypatch.delenv("JAX_ILE_DISTMARG_GH", raising=False) + driver = _driver_module() + assert driver.av_distance_sampling_kwargs(_ToySkyLikelihood(), 10.0, 90.0) == {} + + class WithDistance: + ANGULAR_PARAM_ORDER = ("ra", "dec", "distMpc") + + assert driver.av_distance_sampling_kwargs(WithDistance(), 10.0, 90.0) == { + "sample_d_min": 10.0, "sample_d_max": 90.0} + + +def test_fisher_sky_seed_respects_restricted_sky_window(): + like = _ToySkyLikelihood() + callback = samplers._fixed_shape_value_callback(like, 3, 64) + bounds = {"ra": (1.8, 2.2), "dec": (0.0, 0.4)} + cloud, _, _ = samplers._fisher_sky_seed( + like, like.ANGULAR_PARAM_ORDER, callback, np.random.default_rng(71), + 1.0, 100.0, n_seed=400, n_pilot=100, n_modes=1, + sky_inflate=2.0, prior_frac=0.1, sample_bounds=bounds) + assert np.all((cloud[:, 0] >= 1.8) & (cloud[:, 0] <= 2.2)) + assert np.all((cloud[:, 1] >= 0.0) & (cloud[:, 1] <= 0.4)) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_6d_blind.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_6d_blind.py new file mode 100644 index 000000000..172b5054e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_6d_blind.py @@ -0,0 +1,310 @@ +"""Blind full-sky prior draws through the fixed-distance (6-D) bandlimited kernel, +and the driver modes that stop the run on one uncertified row. + +Before 2026-09-08 the 6-D kernel applied a 15-nat endpoint gap that returned +NaN on 92 of 256 blind driver-prior draws at a 20 ms half-window and 32 of 256 +at 50 ms (rift_O4d d84597c2a), so ``--mode prior-mc / laplace-is / map`` with +``--time-marginalization-quadrature bandlimited`` and no distance +marginalization stopped in the first chunk. Evidence and decision: +DESIGN_jax_bandlimited_distmarg.md, "The fixed-distance kernel". + +TEST-TIMING: the two kernel fixtures build one injection each and refine +4 x 64 rows; the driver runs take ~10 s each on an 8-core ldas-grid slice. +""" +import importlib.util +import os +import pathlib +import subprocess +import sys + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) +lal = pytest.importorskip("lal") +lalsim = pytest.importorskip("lalsimulation") + +import jax.numpy as jnp # noqa: E402 +import RIFT.lalsimutils as lalsimutils # noqa: E402 +from RIFT.likelihood.jax_ile import core # noqa: E402 +from RIFT.likelihood.jax_ile.wrapper import ( # noqa: E402 + JAXExtrinsicLikelihood, build_data_from_precompute, + bandlimited_storage_requirement) + +MSUN, PC = lal.MSUN_SI, lal.PC_SI +_HERE = pathlib.Path(__file__).parent +_CODE = pathlib.Path(__file__).parents[2] +_DRIVER = _CODE / "bin" / "integrate_likelihood_extrinsic_jax" + +EPOCH = 1126259462.0 +DETECTORS = ("H1", "L1") +FMIN, FREF, FMAX = 40.0, 40.0, 300.0 +DELTAF = 0.25 +SRATE = 1024.0 +DIST_INJ = 900.0 # rho ~ 8.7 in H1L1: the driver test's injection +ANGLES = (1.2, -0.4, 0.7, 0.9, 2.1) +D_MIN, D_MAX = 50.0, 4000.0 +# 20 ms cannot contain a wrong-sky arrival shift (2 R_earth / c = 42.6 ms); +# 50 ms is the measured clean value. +HALF_WINDOW_NARROW, HALF_WINDOW_CLEAN = 0.02, 0.05 +SEEDS, N_ROWS, N_CHECK = 4, 64, 6 +# Measured worst gap-only disagreement: 1.1e-04 nat (20 ms), 6.1e-06 nat (50 ms). +TOL_REFERENCE = 1e-3 +# Measured: every gap-only row sat 54-128 nat below its batch maximum. +DEPTH_BELOW_MAX_MIN = 30.0 + + +def _load_module(path, name): + spec = importlib.util.spec_from_file_location(name, str(path)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +# The independent numpy reference pieces (periodic FFT upsample, guard taper, +# log-trapezoid) live with the distance-marginalized tests; loaded by path so +# this file does not depend on how pytest names that module. +_ref_pieces = _load_module(_HERE / "test_jax_bandlimited_distmarg.py", + "_jax_bl_distmarg_reference_pieces") + + +def _params(dist_mpc): + P = lalsimutils.ChooseWaveformParams() + P.m1, P.m2 = 35.0 * MSUN, 30.0 * MSUN + P.fmin, P.fref = FMIN, FREF + P.deltaT, P.deltaF = 1.0 / SRATE, DELTAF + P.dist = dist_mpc * 1e6 * PC + P.fmax = 0.0 + P.approx = lalsim.IMRPhenomD + P.radec = True + P.tref = EPOCH + P.phi, P.theta, P.psi, P.incl, P.phiref = ANGLES + return P + + +def _build(half_window): + _, _, gcert = bandlimited_storage_requirement(1.0 / SRATE, half_window) + ref_guard = 2 * gcert + storage = half_window + ref_guard / SRATE + 0.05 + 16.0 / SRATE + P = _params(DIST_INJ) + data_dict, psd_dict = {}, {} + for det in DETECTORS: + Pd = P.copy() + Pd.detector = det + data_dict[det] = lalsimutils.non_herm_hoff(Pd) + psd_dict[det] = lalsim.SimNoisePSDaLIGOZeroDetHighPower + data, _ = build_data_from_precompute( + P.copy(), data_dict, psd_dict, EPOCH, storage, half_window, 2, FMAX, + analyticPSD_Q=True, verbose=False) + return data, gcert, ref_guard + + +def _driver_prior_draws(n, seed): + """The driver's ``sample_prior`` with distance, re-typed on purpose: the + point is a full-sky, isotropic-orientation, volumetric-distance draw.""" + rng = np.random.default_rng(seed) + angles = [rng.uniform(0.0, 2 * np.pi, n), np.arcsin(rng.uniform(-1.0, 1.0, n)), + rng.uniform(0.0, np.pi, n), np.arccos(rng.uniform(-1.0, 1.0, n)), + rng.uniform(0.0, 2 * np.pi, n)] + u = rng.uniform(0.0, 1.0, n) + dist = (D_MIN ** 3 + u * (D_MAX ** 3 - D_MIN ** 3)) ** (1.0 / 3.0) + return angles, dist + + +def _kernel(data, gcert, angles, dist, endpoint_log_gap): + """The shipped refinement on the fixed-distance field with a CHOSEN gap.""" + kappa_u, rho_u = core._accumulate_unit( + data, *[jnp.asarray(a) for a in angles], core.JAX_INTERP_DEFAULT, False, + guard=gcert) + inv = jnp.asarray(float(data.distMpcRef) / np.asarray(dist)) + return np.asarray(core._time_marginalize_reflected_primitive( + kappa_u * inv[:, None], rho_u * jnp.square(inv)[:, None], data.deltaT, + False, guard=gcert, endpoint_log_gap=endpoint_log_gap)) + + +def _reference(data, ref_guard, row_angles, dist): + """Periodic FFT on a twice-wider guard at factor 512, fixed-distance + reduction, log-trapezoid: independent of the reflected primitive.""" + ang = [jnp.atleast_1d(jnp.asarray(v)) for v in row_angles] + kappa, rho = core._accumulate_unit(data, *ang, core.JAX_INTERP_DEFAULT, False, + guard=ref_guard) + inv = float(data.distMpcRef) / float(dist) + kappa = np.asarray(kappa)[0] * inv + rho = np.asarray(rho)[0] * inv * inv + npts, factor = int(data.npts), _ref_pieces.REF_FACTOR + fine = _ref_pieces._periodic_fft_upsample( + _ref_pieces._taper_guard(kappa, ref_guard), factor) + keep = slice(ref_guard * factor, ref_guard * factor + (npts - 1) * factor + 1) + lnL = fine[keep].real - 0.5 * rho[ref_guard] + return _ref_pieces._numpy_log_trapezoid(lnL, float(data.deltaT) / factor) + + +@pytest.fixture(scope="module", params=[HALF_WINDOW_NARROW, HALF_WINDOW_CLEAN], + ids=["20ms", "50ms"]) +def blind(request): + half_window = request.param + data, gcert, ref_guard = _build(half_window) + like = JAXExtrinsicLikelihood(data, time_quadrature="bandlimited") + seeds = [] + for seed in range(SEEDS): + angles, dist = _driver_prior_draws(N_ROWS, seed) + seeds.append(dict( + angles=angles, dist=dist, + shipped=np.asarray(like.log_likelihood(*angles, dist)), + gap_on=_kernel(data, gcert, angles, dist, core._TIME_ENDPOINT_LOG_GAP_MIN), + gap_off=_kernel(data, gcert, angles, dist, None))) + return dict(half_window=half_window, data=data, gcert=gcert, + ref_guard=ref_guard, seeds=seeds) + + +def _gap_only(row): + return np.where(np.isnan(row["gap_on"]) & np.isfinite(row["gap_off"]))[0] + + +# -------------------------------------------------------------------------- +# (a) the kernel +# -------------------------------------------------------------------------- +def test_shipped_6d_value_is_the_gap_off_refinement(blind): + """The wrapper's number is the ``endpoint_log_gap=None`` refinement, and + the gap parameter is still live: passing the old threshold rejects rows + the shipped path certifies (measured 57 of 256 at 20 ms, 32 at 50 ms).""" + n_gap_only = 0 + for row in blind["seeds"]: + assert np.allclose(row["shipped"], row["gap_off"], equal_nan=True) + n_gap_only += len(_gap_only(row)) + assert n_gap_only > 0, "the endpoint gap no longer rejects anything: revisit the pin" + + +def test_remaining_failures_are_the_window_not_the_certificates(blind): + """With the gap off, every row is certified once the half-window contains + a wrong-sky arrival shift (0 of 256 at 50 ms), and rows still fail at + 20 ms (measured 35 of 256). The driver's parse-time refusal rests on + both halves of this.""" + n_nan = sum(int(np.isnan(row["gap_off"]).sum()) for row in blind["seeds"]) + if blind["half_window"] >= HALF_WINDOW_CLEAN: + assert n_nan == 0, n_nan + else: + assert n_nan > 0, "20 ms is now clean: revisit the parse-time window refusal" + + +def test_gap_only_rows_agree_with_the_independent_reference(blind): + """The rows the gap rejected alone are converged: they match the periodic + reference to TOL_REFERENCE and lie far below the batch maximum, so the + certificate was measuring the row's amplitude, not the quadrature.""" + checked = 0 + for row in blind["seeds"]: + vmax = np.nanmax(row["gap_off"]) + for i in _gap_only(row)[:N_CHECK]: + ref = _reference(blind["data"], blind["ref_guard"], + [a[i] for a in row["angles"]], row["dist"][i]) + assert abs(row["gap_off"][i] - ref) <= TOL_REFERENCE, ( + blind["half_window"], i, row["gap_off"][i], ref) + assert row["gap_off"][i] <= vmax - DEPTH_BELOW_MAX_MIN, (i, row["gap_off"][i], vmax) + checked += 1 + assert checked >= SEEDS + + +# -------------------------------------------------------------------------- +# (b) the driver's parse-time refusal +# -------------------------------------------------------------------------- +def _load_driver(): + import importlib.machinery + loader = importlib.machinery.SourceFileLoader("_jax_bl_6d_blind_driver", str(_DRIVER)) + spec = importlib.util.spec_from_loader(loader.name, loader) + mod = importlib.util.module_from_spec(spec) + loader.exec_module(mod) + return mod + + +def _parse_and_check(drv, extra): + argv = ["--inj-mode", "--mass1", "35", "--mass2", "30"] + list(extra) + optp = drv.build_parser() + argv = drv._normalize_interpolate_time_argv(argv) + opts, _ = optp.parse_args(argv) + drv.record_supplied_options(opts, argv, optp) + drv.resolve_ile_interface_aliases(opts, optp) + drv.check_critical_and_report(opts, optp) + return opts + + +def test_driver_refuses_the_stop_modes_below_the_arrival_bound(capsys): + """The window bound is 2 R_earth / c, applied only to the modes whose + full-sky draws go through ``eval_lnL``'s stop; Simpson, a wide window + and the flowMC pilot (which the 20 ms flowMC test runs) are not refused.""" + drv = _load_driver() + bound = 2.0 * lal.REARTH_SI / lal.C_SI + assert drv._BANDLIMITED_FULLSKY_HALF_WINDOW_MIN == bound + assert 0.042 < bound < 0.043 + for mode in ("prior-mc", "laplace-is", "map", "nuts"): + with pytest.raises(SystemExit): + _parse_and_check(drv, ["--mode", mode, + "--time-marginalization-quadrature", "bandlimited", + "--data-integration-window-half", "0.02"]) + err = capsys.readouterr().err + assert "2 R_earth / c" in err and "--data-integration-window-half" in err, err + assert "0.05 s" in err, err + _parse_and_check(drv, ["--mode", "prior-mc", + "--time-marginalization-quadrature", "bandlimited", + "--data-integration-window-half", "0.05"]) + _parse_and_check(drv, ["--mode", "prior-mc", + "--time-marginalization-quadrature", "simpson", + "--data-integration-window-half", "0.02"]) + _parse_and_check(drv, ["--mode", "flowmc", "--distance-marginalization", + "--time-marginalization-quadrature", "bandlimited", + "--data-integration-window-half", "0.02"]) + + +# -------------------------------------------------------------------------- +# (c) the driver end to end +# -------------------------------------------------------------------------- +def _driver_argv(mode, half_window, out): + return [sys.executable, str(_DRIVER), + "--inj-mode", "--mass1", "35", "--mass2", "30", + "--inj-deltaF", "0.25", "--inj-detectors", "H1,L1", + "--inj-distance", str(DIST_INJ), + "--fmin-template", "40", "--reference-freq", "40", "--fmax", "300", + "--l-max", "2", "--approximant", "IMRPhenomD", "--srate", "1024", + "--data-integration-window-half", str(half_window), + "--internal-data-storage-window-half", "0.08", + "--d-min", str(D_MIN), "--d-max", str(D_MAX), + "--mode", mode, "--time-marginalization-quadrature", "bandlimited", + "--n-max", "400", "--seed", "3", "--output-file", str(out)] + + +def _run_driver(tmp_path, mode, half_window): + out = tmp_path / "ile" + env = dict(os.environ, PYTHONPATH=str(_CODE), OMP_NUM_THREADS="1", + JAX_PLATFORMS="cpu", JAX_ENABLE_X64="1") + proc = subprocess.run(_driver_argv(mode, half_window, out), capture_output=True, + text=True, env=env, cwd=str(tmp_path), timeout=3600) + return proc, proc.stdout + proc.stderr, out + + +@pytest.mark.parametrize("mode", ["prior-mc", "map"]) +def test_driver_6d_prior_seeded_modes_run_bandlimited(tmp_path, mode): + """No distance marginalization: the 6-D kernel under the prior-seeded + modes, at the clean half-window. Before the endpoint change this command + stopped in its first chunk (prior-mc: 30 of 400 rows failed, all at the + gap). + + ``laplace-is`` is not run end to end: at this injection and seed its + moment-matched proposal walks off the peak and ``require_finite_evidence`` + raises, at ``--n-max`` 400, 1000 and 2000 alike, and the distance- + marginalized variant does the same on the unchanged base (ldas-grid, + 2026-09-08). That is an evidence-quality failure, not a certificate one; + its blind draws take the same ``eval_lnL`` path as ``prior-mc``, and the + parse-time refusal above covers it.""" + proc, log, out = _run_driver(tmp_path, mode, HALF_WINDOW_CLEAN) + assert proc.returncode == 0, log[-3000:] + assert "failed a certificate" not in log, log[-3000:] + assert "time-marginalization quadrature: bandlimited" in log, log[-3000:] + row = np.atleast_2d(np.loadtxt(str(out) + "_0_.dat")) + assert row.shape[1] == 13 and np.isfinite(row[0, 9]), row + + +def test_driver_refuses_the_narrow_window_before_precompute(tmp_path): + proc, log, _ = _run_driver(tmp_path, "prior-mc", HALF_WINDOW_NARROW) + assert proc.returncode != 0 + assert "2 R_earth / c" in log, log[-3000:] + assert "Building JAX likelihood" not in log, log[-3000:] diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_distmarg.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_distmarg.py new file mode 100644 index 000000000..2c1152548 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_bandlimited_distmarg.py @@ -0,0 +1,687 @@ +"""Band-limited time quadrature on the DISTANCE-marginalized JAX likelihood. + +`fused_log_likelihood_distmarg` used to reduce over distance on the data time +grid and hand the reduced field to the terminal selector, which refuses +`bandlimited`: interpolating an already-reduced nonlinear field can converge to +the wrong function. It now refines the complex primitive kappa(t) first and +applies the SAME distance quadrature at every refined node. + +WHAT THE REFERENCE IS, AND WHY IT IS NOT THE CODE UNDER TEST. The fine-grid +reference below reconstructs the primitive with a PLAIN PERIODIC zero-padded FFT +on a guarded window, reduces over distance with a numpy log-sum-exp, and +integrates with a numpy trapezoid. Its extension is periodic where the shipped +one is an even reflection, and its guard taper is a half-cosine on the offset +grid (k+1)/(g+1) where the shipped one is on k/g, so agreement is not two +spellings of one routine. It is checked for convergence in BOTH the guard and +the refinement factor, because a reference that has not converged is not one. + +The taper is not optional and is not decoration. An UNTAPERED periodic +reconstruction leaves a step at the periodic seam whose Gibbs ringing decays +only like 1/guard: doubling the guard halves the error instead of removing it, +and the sequence never converges to the tolerance this file asserts. That is +what the reference's own convergence test would have caught, and it is why the +reference tapers. The DESIGN record carries the untapered ladder. + +A sample-rate ladder cannot serve as that reference here and the numbers say +why: the integrand's width is sigma_t ~ 1/(2 pi rho sigma_f), so Simpson at 8x +the native rate is still coarse at production amplitude. What the ladder DOES +show, and is asserted below, is that native Simpson moves monotonically toward +the band-limited value as the rate rises -- which is the claim the option makes. + +Numbers and method: RIFT/likelihood/jax_ile/DESIGN_jax_bandlimited_distmarg.md. +""" +import inspect +import pathlib +import subprocess +import sys + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + +lal = pytest.importorskip("lal") +lalsim = pytest.importorskip("lalsimulation") + +import RIFT.lalsimutils as lalsimutils # noqa: E402 +from RIFT.likelihood.jax_ile import core, wrapper # noqa: E402 +from RIFT.likelihood.jax_ile.wrapper import ( # noqa: E402 + JAXDistanceMarginalizedLikelihood, build_data_from_precompute, + bandlimited_storage_requirement) + +MSUN, PC = lal.MSUN_SI, lal.PC_SI +_CODE = pathlib.Path(__file__).parents[2] + +EPOCH = 1126259462.0 +DETECTORS = ("H1", "L1") +IWH = 0.075 # marginalization half-window, seconds +FMIN, FREF, FMAX = 40.0, 40.0, 300.0 +DELTAF = 0.25 +SRATE = 1024.0 +ANGLES = (1.2, -0.4, 0.7, 0.9, 2.1) +D_MIN, D_MAX, N_GRID = 50.0, 4000.0, 512 +# The injected angles are NOT where this likelihood peaks, and every test here +# evaluates at a fixed point, so none needs them to be. The precompute below +# receives the same P as the injection, and for IMRPhenomD the template route +# (hlmoft -> SimInspiralTDModesFromPolarizations) bakes that P's phiref and psi +# into the modes; ILE then applies both again through Y_lm and F. Production +# zeroes P.phiref/P.psi before the precompute (batchmode driver, and the JAX +# driver's make_template). Measured on the 900 Mpc case: lnL 84.0 at the +# injection against a 2-D (psi, phiref) maximum of 194.9; the conventional +# factored likelihood gives the same numbers, so this is the waveform interface, +# not the quadrature. + +# Injected distances. Named by the amplitude they produce, and the fixture +# asserts the amplitude rather than trusting the label. +DIST_QUIET, DIST_LOUD = 390.0, 48.5 + +# Agreement between the shipped quadrature and the independently reconstructed +# fine-grid reference, in nats. Measured; see the DESIGN record. +TOL_REFERENCE = 1e-3 + + +def _params(dist_mpc, srate): + P = lalsimutils.ChooseWaveformParams() + P.m1, P.m2 = 35.0 * MSUN, 30.0 * MSUN + P.s1z, P.s2z = 0.1, -0.2 + P.fmin, P.fref = FMIN, FREF + P.deltaT = 1.0 / srate + P.deltaF = DELTAF + P.dist = dist_mpc * 1e6 * PC + P.fmax = 0.0 + P.approx = lalsim.IMRPhenomD + P.radec = True + P.tref = EPOCH + P.phi, P.theta, P.psi, P.incl, P.phiref = ANGLES + return P + + +def _build(dist_mpc, srate, storage_half): + P = _params(dist_mpc, srate) + data_dict, psd_dict = {}, {} + for det in DETECTORS: + Pdet = P.copy() + Pdet.detector = det + data_dict[det] = lalsimutils.non_herm_hoff(Pdet) + psd_dict[det] = lalsim.SimNoisePSDaLIGOZeroDetHighPower + data, extras = build_data_from_precompute( + P.copy(), data_dict, psd_dict, EPOCH, storage_half, IWH, 2, FMAX, + analyticPSD_Q=True, verbose=False) + extras["network_snr"] = _network_snr(P, data_dict) + return data, extras + + +def _network_snr(P, data_dict): + """Optimal network SNR of the zero-noise data, sqrt(sum_det ), on the + band the likelihood integrates. Not ``extras["guess_snr"]``: that is the + precompute's deliberately deflated estimate, sqrt(sum max|Q_lm|^2 / U_lm) + divided by 2.3, and sits 2.305x below this on every fixture here.""" + IP = lalsimutils.ComplexIP( + fLow=FMIN, fNyq=0.5 / P.deltaT, deltaF=P.deltaF, fMax=FMAX, + psd=lalsim.SimNoisePSDaLIGOZeroDetHighPower, analyticPSD_Q=True) + return float(np.sqrt(sum(IP.norm(d) ** 2 for d in data_dict.values()))) + + +#: Guard and refinement factor at which the numpy reference is converged; both +#: are certified below by halving them. See the DESIGN record for the ladders. +REF_FACTOR = 512 + + +def _reference_guard(srate=SRATE): + """Guard for the numpy reference gathers: one doubling past the certified one.""" + return 2 * bandlimited_storage_requirement(1.0 / srate, IWH)[2] + + +def _storage_half(guard, srate=SRATE): + return IWH + guard / srate + 0.05 + 16.0 / srate + + +def _angles(): + return [np.atleast_1d(np.asarray(v, dtype=float)) for v in ANGLES] + + +def _like(data, quadrature): + return JAXDistanceMarginalizedLikelihood( + data, D_MIN, D_MAX, n_grid=N_GRID, time_quadrature=quadrature) + + +def _value(like): + return float(np.asarray(like.log_likelihood(*_angles()))[0]) + + +# -------------------------------------------------------------------------- +# Independent numpy reference +# -------------------------------------------------------------------------- +def _periodic_fft_upsample(x, factor): + """Plain periodic zero-padded FFT interpolation, Nyquist bin split evenly. + + Deliberately NOT the shipped even extension: the reference must not be able + to inherit a boundary-handling mistake from the code it certifies. Its own + periodic seam is pushed far outside the integrated window by the guard. + """ + n = x.shape[-1] + X = np.fft.fft(x) + half, n_out = n // 2, n * factor + if n % 2 == 0: + nyq = X[half:half + 1] * 0.5 + Y = np.concatenate([X[:half], nyq, + np.zeros(n_out - n - 1, dtype=complex), nyq, + X[half + 1:]]) + else: + Y = np.concatenate([X[:half + 1], np.zeros(n_out - n, dtype=complex), + X[half + 1:]]) + return np.fft.ifft(Y) * factor + + +def _numpy_distance_reduction(K, R, x_grid, log_w, block=16): + """log sum_g exp(K x_g - R x_g^2 / 2 + log w_g), blocked, in numpy.""" + m = np.full(K.shape, -np.inf) + s = np.zeros(K.shape) + for start in range(0, len(x_grid), block): + xs = x_grid[start:start + block] + e = (K[:, None] * xs[None, :] + - 0.5 * R[:, None] * np.square(xs)[None, :] + + log_w[start:start + block][None, :]) + m_b = np.max(e, axis=-1) + s_b = np.sum(np.exp(e - m_b[:, None]), axis=-1) + m_new = np.maximum(m, m_b) + s = s * np.exp(m - m_new) + s_b * np.exp(m_b - m_new) + m = m_new + return m + np.log(s) + + +def _numpy_log_trapezoid(lnL, dx): + peak = np.max(lnL) + y = np.exp(lnL - peak) + return peak + np.log(dx * (0.5 * y[0] + np.sum(y[1:-1]) + 0.5 * y[-1])) + + +def _taper_guard(kappa, guard): + """Half-cosine ramp over the guard samples, on the OFFSET grid (k+1)/(g+1). + + The shipped taper ramps on k/g. Both drive the seam to zero; neither is the + other, so a mistake in the shipped window shape cannot be inherited here. + """ + if not guard: + return kappa + ramp = 0.5 * (1.0 - np.cos(np.pi * np.arange(1, guard + 1) / (guard + 1))) + w = np.concatenate([ramp, np.ones(kappa.shape[-1] - 2 * guard), ramp[::-1]]) + return kappa * w + + +def _fine_grid_reference(data, x_grid, log_w, guard, factor, angles=None): + angles = _angles() if angles is None else angles + kappa, rho = core._accumulate_unit( + data, *angles, core.JAX_INTERP_DEFAULT, False, guard=guard) + kappa = np.asarray(kappa)[0] + rho = np.asarray(rho)[0] + npts = int(data.npts) + fine = _periodic_fft_upsample(_taper_guard(kappa, guard), factor) + keep = slice(guard * factor, guard * factor + (npts - 1) * factor + 1) + K = fine[keep].real + R = np.full(K.shape, rho[guard]) + lnL = _numpy_distance_reduction(K, R, x_grid, log_w) + return _numpy_log_trapezoid(lnL, float(data.deltaT) / factor) + + +def _reduce_then_refine(data, x_grid, log_w, factor): + """The WRONG order, built explicitly: reduce on the data grid, then refine. + + This is what `_time_marginalize_terminal` refuses, and what the code would + do if the primitive gather were dropped. Nothing asserts it is close to + anything; it exists so the ordering assertions below are load-bearing. + """ + kappa, rho = core._accumulate_unit( + data, *_angles(), core.JAX_INTERP_DEFAULT, False, guard=0) + K = np.asarray(kappa)[0].real + R = np.asarray(rho)[0] + lnL_t = _numpy_distance_reduction(K, R, x_grid, log_w) + fine = _periodic_fft_upsample(lnL_t.astype(complex), factor).real + fine = fine[:(len(lnL_t) - 1) * factor + 1] + return _numpy_log_trapezoid(fine, float(data.deltaT) / factor) + + +# -------------------------------------------------------------------------- +# Fixtures +# -------------------------------------------------------------------------- +@pytest.fixture(scope="module") +def cases(): + guard = _reference_guard() + out = {} + for tag, dist in (("quiet", DIST_QUIET), ("loud", DIST_LOUD)): + data, extras = _build(dist, SRATE, _storage_half(guard)) + like_s = _like(data, "simpson") + out[tag] = dict( + data=data, snr=float(extras["network_snr"]), + x_grid=np.asarray(like_s.x_grid), + log_w=np.asarray(like_s.log_w_grid), + simpson=_value(like_s), + bandlimited=_value(_like(data, "bandlimited")), + guard=guard) + return out + + +def test_fixture_amplitudes_are_the_two_regimes_claimed(cases): + """Guard on the fixture. Both assertions below are amplitude claims, and a + distance that quietly stopped producing the intended SNR would make the + 'differs at high amplitude' test pass or fail for the wrong reason. + The bounds are on the network optimal SNR of the data (measured 45.9 and + 369.2), not on ``guess_snr``, which is 2.305x lower by construction.""" + assert 40.0 < cases["quiet"]["snr"] < 55.0, cases["quiet"]["snr"] + assert cases["loud"]["snr"] > 300.0, cases["loud"]["snr"] + + +# -------------------------------------------------------------------------- +# (a) agreement with a converged, independently built reference +# -------------------------------------------------------------------------- +@pytest.mark.parametrize("tag", ["quiet", "loud"]) +def test_bandlimited_distmarg_matches_the_independent_fine_grid_reference( + cases, tag): + case = cases[tag] + ref = _fine_grid_reference(case["data"], case["x_grid"], case["log_w"], + case["guard"], REF_FACTOR) + assert abs(case["bandlimited"] - ref) < TOL_REFERENCE, ( + "%s: bandlimited %.6f vs independent reference %.6f (%+.3e nats)" + % (tag, case["bandlimited"], ref, case["bandlimited"] - ref)) + + +@pytest.mark.parametrize("tag", ["quiet", "loud"]) +def test_the_reference_is_converged_in_both_of_its_own_knobs(cases, tag): + """An unconverged reference certifies nothing. Halving the factor and + halving the guard must each leave it alone at the tolerance the agreement + test uses.""" + case = cases[tag] + args = (case["data"], case["x_grid"], case["log_w"]) + full = _fine_grid_reference(*args, guard=case["guard"], factor=REF_FACTOR) + coarse_factor = _fine_grid_reference(*args, guard=case["guard"], + factor=REF_FACTOR // 2) + narrow_guard = _fine_grid_reference(*args, guard=case["guard"] // 2, + factor=REF_FACTOR) + assert abs(full - coarse_factor) < TOL_REFERENCE, ( + "%s: reference not converged in the factor (%+.3e nats)" + % (tag, full - coarse_factor)) + assert abs(full - narrow_guard) < TOL_REFERENCE, ( + "%s: reference not converged in the guard (%+.3e nats)" + % (tag, full - narrow_guard)) + + +# -------------------------------------------------------------------------- +# (b) the knob is live +# -------------------------------------------------------------------------- +def test_bandlimited_differs_from_native_simpson_at_high_amplitude(cases): + """If this stops biting, the option is inert and everything above is + measuring the Simpson value twice.""" + loud = cases["loud"] + delta = loud["bandlimited"] - loud["simpson"] + assert abs(delta) > 100 * TOL_REFERENCE, ( + "bandlimited and native Simpson agree to %+.3e nats at rho %.1f; the " + "option is not changing the answer" % (delta, loud["snr"])) + + +def test_native_simpson_moves_toward_the_bandlimited_value_as_the_rate_rises( + cases): + """The physical claim, on regenerated data: Simpson's error is a resolution + error, so refining the DATA rate must close the gap the option closes at the + native rate. Not a converged reference -- 8x is still coarse against + sigma_t ~ 1/(2 pi rho sigma_f) -- which is why the assertion is on the + ordering of the residuals and not on their size.""" + loud = cases["loud"] + target = loud["bandlimited"] + gaps = [abs(loud["simpson"] - target)] + guard = bandlimited_storage_requirement(1.0 / SRATE, IWH)[2] + for mult in (4, 8): + data, _ = _build(DIST_LOUD, SRATE * mult, _storage_half(guard)) + gaps.append(abs(_value(_like(data, "simpson")) - target)) + assert gaps[1] < gaps[0], "4x Simpson did not improve on native: %r" % (gaps,) + assert gaps[2] < gaps[1], "8x Simpson did not improve on 4x: %r" % (gaps,) + + +def test_the_reduction_order_is_what_carries_the_agreement(cases): + """Reduce-then-refine is a different number, not a rounding difference. + + This is the assertion that was mutation-tested: swapping the shipped order + makes the agreement test above fail by orders of magnitude, and this records + by how much rather than leaving it to the reviewer to imagine.""" + loud = cases["loud"] + wrong = _reduce_then_refine(loud["data"], loud["x_grid"], loud["log_w"], + REF_FACTOR) + assert abs(loud["bandlimited"] - wrong) > 100 * TOL_REFERENCE, ( + "refining the already-reduced field lands within %.3e nats of the " + "shipped answer, so the ordering assertions prove nothing" + % abs(loud["bandlimited"] - wrong)) + + +# -------------------------------------------------------------------------- +# (c) the refusal set still refuses, and the enabled one no longer does +# -------------------------------------------------------------------------- +def test_the_refusal_set_still_refuses_and_names_the_endpoint(cases): + data = cases["quiet"]["data"] + refusing = [ + (wrapper.JAXDistPhiMargLikelihood, dict(nphi=4)), + (wrapper.JAXDistPsiMargLikelihood, dict(npsi=4)), + (wrapper.JAXDistPhiPsiMargLikelihood, dict(nphi=4, npsi=4, + angle_marg="grid")), + ] + for cls, kwargs in refusing: + with pytest.raises(ValueError, match="primitive time fields"): + cls(data, D_MIN, D_MAX, n_grid=32, + time_quadrature="bandlimited", **kwargs) + # ... and the same construction is fine on the default quadrature, so + # the refusal is about the COMBINATION and not about the wrapper. + cls(data, D_MIN, D_MAX, n_grid=32, **kwargs) + + +def test_only_the_angle_wrappers_call_the_nonlinear_refusal(cases): + """The distance wrapper must not merely stop raising -- it must stop calling + the validator, or a later edit re-enabling it would be invisible here.""" + src = inspect.getsource(wrapper.JAXDistanceMarginalizedLikelihood.__init__) + assert "_validate_nonlinear_time_quadrature" not in src + for cls in (wrapper.JAXDistPhiMargLikelihood, + wrapper.JAXDistPsiMargLikelihood, + wrapper.JAXDistPhiPsiMargLikelihood): + assert "_validate_nonlinear_time_quadrature" in inspect.getsource( + cls.__init__), cls.__name__ + + +def test_bandlimited_distmarg_refuses_what_it_cannot_honour(cases): + """Two fail-closed doors. `return_lnLt` has no meaning once the reduction + moves to the refined grid, and rotation data's norm depends on arrival time, + which the refinement holds fixed.""" + import types + data = cases["quiet"]["data"] + like = _like(data, "simpson") + with pytest.raises(ValueError, match="return_lnLt"): + core.fused_log_likelihood_distmarg( + data, *_angles(), like.x_grid, like.log_w_grid, + time_quadrature="bandlimited", return_lnLt=True) + with pytest.raises(ValueError, match="arrival time"): + core.fused_log_likelihood_distmarg( + types.SimpleNamespace(feature="rotation"), *_angles(), + like.x_grid, like.log_w_grid, time_quadrature="bandlimited") + + +def test_the_distance_reduction_has_one_definition_for_both_grids(): + """Coarse and refined paths must reach the same quadrature helper. A second + copy is how a future adaptive/GH distance branch would land on one grid and + not the other -- silently, since both would still return a number.""" + src = inspect.getsource(core.fused_log_likelihood_distmarg) + assert "_logsumexp_grid_blocked(" in src and "_logsumexp_grid_scanned(" in src + assert "reduce_fn=_reduce" in src + + +def test_the_guard_pair_has_one_definition(cases): + """Gathered support and integrated guard must be the same number. Three + sites need it; a re-typed copy is a wrong likelihood, not an error.""" + for fn in (core.fused_log_likelihood, core.fused_log_likelihood_distmarg, + wrapper.bandlimited_storage_requirement): + assert "bandlimited_time_guard(" in inspect.getsource(fn), fn.__name__ + npts = int(cases["quiet"]["data"].npts) + g0, gcert = core.bandlimited_time_guard(npts) + assert gcert == 2 * g0 and g0 >= core.default_time_guard(npts) + like = _like(cases["quiet"]["data"], "bandlimited") + assert (like.time_guard_initial, like.time_guard_certified) == (g0, gcert) + + +# -------------------------------------------------------------------------- +# (d) the endpoint certificate on a floored field +# -------------------------------------------------------------------------- +# The distance-marginalized field has a floor: at every node the distance sum +# is at least the far-distance prior mass, so a row's peak-to-endpoint contrast +# is bounded by its own peak height and the fixed-distance kernel's 15-nat +# endpoint gap rejects every low-contrast row -- converged or not. Blind +# full-sky draws, which every prior-seeded driver mode evaluates by the +# thousand, are mostly low-contrast rows. Measured on 256 such rows at 20 ms +# half-window: 35% rejected by the gap alone, all of them agreeing with the +# independent reference to 1e-4 nat. The distance path therefore runs with +# that certificate off and the guard-agreement and doubling certificates on. +# Numbers: the DESIGN record, "The endpoint certificate". +N_BLIND = 128 +#: The blind census runs at this distance, not on the ``quiet`` fixture: the +#: floor argument needs low-contrast rows, and at 390 Mpc the gap rejects none +#: of 256 (measured); at 900 Mpc it rejects 14 of 256 at this half-window. +DIST_BLIND = 900.0 + + +def _blind_draws(n, seed): + """The driver's ``sample_prior`` for the five angles, re-typed on purpose: + the point is a full-sky, isotropic-orientation draw, not a driver import.""" + rng = np.random.default_rng(seed) + return [rng.uniform(0.0, 2 * np.pi, n), np.arcsin(rng.uniform(-1.0, 1.0, n)), + rng.uniform(0.0, np.pi, n), np.arccos(rng.uniform(-1.0, 1.0, n)), + rng.uniform(0.0, 2 * np.pi, n)] + + +def _primitive_with_gap(like, angles, endpoint_log_gap): + """The shipped refinement on the distance field with a CHOSEN endpoint gap. + + Rebuilds ``fused_log_likelihood_distmarg``'s reduction so the certificate + can be switched without touching the module; the wrapper's own value is + asserted equal to the ``None`` setting below, so this helper cannot drift + from the code silently.""" + import jax.numpy as jnp + data = like.data + guard = like.time_guard_certified + kappa, rho = core._accumulate_unit( + data, *[jnp.asarray(v) for v in angles], like.interp, False, guard=guard) + a = jnp.asarray(like.x_grid) + b = -0.5 * jnp.square(a) + log_w = jnp.asarray(like.log_w_grid) + + def reduce_fn(k, r): + kk = k.real + n = int(np.prod(kk.shape)) + block = min(max(1, core._BANDLIMITED_GRID_ELEMENTS // max(n, 1)), + int(a.shape[0])) + return core._logsumexp_grid_scanned( + kk.reshape(n), r.reshape(n), a, b, log_w, block).reshape(kk.shape) + + return np.asarray(core._time_marginalize_reflected_primitive( + kappa, rho, data.deltaT, False, guard=guard, reduce_fn=reduce_fn, + endpoint_log_gap=endpoint_log_gap)) + + +@pytest.fixture(scope="module") +def blind(): + data, _ = _build(DIST_BLIND, SRATE, _storage_half(_reference_guard())) + like = _like(data, "bandlimited") + angles = _blind_draws(N_BLIND, seed=0) + return dict( + data=data, like=like, angles=angles, + shipped=np.asarray(like.log_likelihood(*angles)), + gap_off=_primitive_with_gap(like, angles, None), + gap_on=_primitive_with_gap(like, angles, core._TIME_ENDPOINT_LOG_GAP_MIN)) + + +def test_blind_draws_are_certified_without_the_endpoint_gap(blind): + """Every blind row gets a number from the shipped wrapper, and that number + is the ``endpoint_log_gap=None`` refinement and nothing else.""" + assert np.all(np.isfinite(blind["shipped"])), ( + "%d of %d blind rows uncertified with the endpoint gap off" + % (int(np.sum(~np.isfinite(blind["shipped"]))), N_BLIND)) + # The wrapper is jitted and the direct call is not; XLA fusion moves the + # last bits, so this is a tolerance and not an equality. + np.testing.assert_allclose(blind["shipped"], blind["gap_off"], rtol=0, atol=1e-8) + + +def test_the_endpoint_gap_was_rejecting_converged_rows(blind): + """The certificate is live on this field (it rejects a material fraction), + and what it rejects agrees with the independent reference. Without the + first assertion the second would be vacuous; without the second the first + would only show the gate is loud.""" + rejected = np.where(np.isnan(blind["gap_on"]) & np.isfinite(blind["gap_off"]))[0] + assert len(rejected) >= 4, ( + "the 15-nat gap rejected only %d of %d blind rows; the floor argument " + "is not exercised by this draw" % (len(rejected), N_BLIND)) + guard = _reference_guard() + like = blind["like"] + worst = 0.0 + for i in rejected[:8]: + angles = [np.atleast_1d(v[i]) for v in blind["angles"]] + ref = _fine_grid_reference(blind["data"], np.asarray(like.x_grid), + np.asarray(like.log_w_grid), guard, + REF_FACTOR, angles=angles) + worst = max(worst, abs(blind["gap_off"][i] - ref)) + assert abs(blind["gap_off"][i] - ref) < TOL_REFERENCE, ( + "rejected row %d: refinement %.6f vs reference %.6f" + % (i, blind["gap_off"][i], ref)) + + +def test_no_production_caller_applies_the_endpoint_gap(): + """Both fields run without the endpoint gap. The fixed-distance kernel + kept it when this file was written; the follow-up of 2026-09-08 measured + the same signature there (test_jax_bandlimited_6d_blind.py; DESIGN record, + "The fixed-distance kernel"). The kernel default is ``None``, the + threshold constant stays for the tests that pin what the gap rejected, and + no production caller passes one.""" + sig = inspect.signature(core._time_marginalize_reflected_primitive) + assert sig.parameters["endpoint_log_gap"].default is None + assert core._TIME_ENDPOINT_LOG_GAP_MIN == 15.0 + assert "endpoint_log_gap" not in inspect.getsource(core.fused_log_likelihood) + assert "endpoint_log_gap=None" in inspect.getsource(core.fused_log_likelihood_distmarg) + + +# -------------------------------------------------------------------------- +# (e) the driver +# -------------------------------------------------------------------------- +def _load_driver(): + import importlib.machinery + import importlib.util + path = _CODE / "bin" / "integrate_likelihood_extrinsic_jax" + loader = importlib.machinery.SourceFileLoader("_jax_bl_distmarg_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_eval_lnL_still_fails_closed_and_names_the_rows(): + """The hard stop on an uncertified row is kept (no coarse likelihood is + substituted); what changed is that the message counts the rows and prints + their parameters, so a failed run says where it failed.""" + import types + drv = _load_driver() + like = types.SimpleNamespace( + time_quadrature="bandlimited", + log_likelihood=lambda *cols: np.array([1.0, np.nan, 3.0])) + opts = types.SimpleNamespace(n_chunk=8000) + theta = np.array([[0.1, 0.2, 0.3, 0.4, 0.5]] * 3) + theta[1, 0] = 2.5 + with pytest.raises(RuntimeError) as err: + drv.eval_lnL(like, theta, opts, with_distance=False) + msg = str(err.value) + assert "1 of 3 rows" in msg and "ra=2.5000" in msg + assert "no coarse likelihood is substituted" in msg + + +@pytest.mark.parametrize("mode", ["prior-mc", "map"]) +def test_driver_prior_seeded_modes_run_distance_marginalized_bandlimited( + tmp_path, mode): + """The modes that evaluate blind prior draws through the driver's own + ``eval_lnL`` -- which stops the run on one uncertified row. Before the + endpoint change, every seed of this command failed inside the first chunk + (measured: 2-7 uncertified rows per 16 draws, seeds 0-29). ``map`` is + the one whose pilot is hard-coded to 4000 blind draws. No flowMC + dependency, so this is the executable coverage the CI ``jax-ile-check`` + job actually runs. + + ``laplace-is`` is deliberately not here. On this injection it gets + through every evaluation and then refuses its own evidence (adapted lnZ + below the prior pilot's Markov floor, neff 6 at n_max 4000) while the + Simpson control at the same budget passes with neff 16: the resolved peak + is narrower than its single moment-matched Gaussian covers. That is the + estimator's documented failure, not a certificate, and a test of it would + be a test of luck. + + The half-window is 50 ms, not the 20 ms of the flowMC test above: a + wrong-sky draw shifts a detector's arrival by up to 2 R_earth / c, about + 43 ms, and a row whose arrival peak sits at the window edge is one the + trapezoid and guard certificates legitimately cannot converge on.""" + import os + out = tmp_path / "ile" + env = dict(os.environ, PYTHONPATH=str(_CODE), OMP_NUM_THREADS="1", + JAX_PLATFORMS="cpu", JAX_ENABLE_X64="1") + proc = subprocess.run( + [sys.executable, str(_CODE / "bin" / "integrate_likelihood_extrinsic_jax"), + "--inj-mode", "--mass1", "35", "--mass2", "30", + "--inj-deltaF", "0.25", "--inj-detectors", "H1,L1", + "--inj-distance", "900", + "--fmin-template", "40", "--reference-freq", "40", "--fmax", "300", + "--l-max", "2", "--approximant", "IMRPhenomD", "--srate", "1024", + "--data-integration-window-half", "0.05", + "--internal-data-storage-window-half", "0.08", + "--d-min", "50", "--d-max", "4000", "--distance-grid-points", "32", + "--distance-marginalization", "--mode", mode, + "--time-marginalization-quadrature", "bandlimited", + "--n-max", "400", "--seed", "3", "--output-file", str(out)], + capture_output=True, text=True, env=env, cwd=str(tmp_path), timeout=3600) + log = proc.stdout + proc.stderr + assert proc.returncode == 0, log[-3000:] + assert "failed a certificate" not in log, log[-3000:] + assert "time-marginalization quadrature: bandlimited" in log, log[-3000:] + row = np.atleast_2d(np.loadtxt(str(out) + "_0_.dat")) + assert row.shape[1] == 13 and np.isfinite(row[0, 9]), row + if mode == "prior-mc": + assert np.isfinite(row[0, 12]) and row[0, 11] == 400, row + + +# -------------------------------------------------------------------------- +# (f) the flowMC driver path +# -------------------------------------------------------------------------- +def test_driver_runs_flowmc_distance_marginalized_bandlimited(tmp_path): + """End to end through the shipped executable, small budget. A library test + cannot see the storage-window widening, the option plumbing, or the fact + that the run publishes a row rather than a NaN. + + flowMC is an optional dependency the CI ``jax-ile-check`` job does not + install, so ``pytest.importorskip`` below is a real skip there, not a + pass; this is still full executable coverage on any host that has flowMC + (e.g. local development). A non-flowMC ``--mode`` (laplace-is, prior-mc, + map, nuts) was tried here first and dropped: all four reach the SAME + ``JAXDistanceMarginalizedLikelihood(..., time_quadrature="bandlimited")`` + construction (driver ``analyze_one``, the ``elif + opts.distance_marginalization:`` branch) via a blind full-sky prior draw, + and MEASURED against this exact injection that draw's reflected-FFT + convergence check hard-fails (no coarse likelihood substituted, by + design) on roughly a quarter of blind draws regardless of seed -- so a + small ``--n-max`` does not make the combination reliable, only rarer to + catch in one CI run. See the flagged follow-up on this fragility. + + The budget is set by memory, not by wall clock. flowMC unrolls its + per-step proposal, so the compiled graph -- already eleven refinement + branches wide -- is multiplied by ``--n-local-steps``; at twenty steps this + same run needs 23 GB, and at two it needs 4.5 GB. ``--n-prior-pilot`` is + what actually brackets the peak here, so the steps are what gets cut. + ``--internal-data-storage-window-half`` is set BELOW the band-limited + requirement on purpose, so the auto-widening has something to do and the + assertion on it is not vacuous.""" + pytest.importorskip("flowMC") + import os + out = tmp_path / "ile" + env = dict(os.environ, PYTHONPATH=str(_CODE), OMP_NUM_THREADS="1", + JAX_PLATFORMS="cpu", JAX_ENABLE_X64="1") + proc = subprocess.run( + [sys.executable, str(_CODE / "bin" / "integrate_likelihood_extrinsic_jax"), + "--inj-mode", "--mass1", "35", "--mass2", "30", + "--inj-deltaF", "0.25", "--inj-detectors", "H1,L1", + "--inj-distance", "900", + "--fmin-template", "40", "--reference-freq", "40", "--fmax", "300", + "--l-max", "2", "--approximant", "IMRPhenomD", "--srate", "1024", + "--data-integration-window-half", "0.02", + "--internal-data-storage-window-half", "0.08", + "--d-min", "50", "--d-max", "4000", "--distance-grid-points", "32", + "--distance-marginalization", "--mode", "flowmc", + "--time-marginalization-quadrature", "bandlimited", + "--n-training-loops", "1", "--n-production-loops", "1", + "--n-epochs", "2", "--n-local-steps", "2", "--n-global-steps", "2", + "--n-prior-pilot", "2000", "--output-file", str(out)], + capture_output=True, text=True, env=env, cwd=str(tmp_path), timeout=3600) + log = proc.stdout + proc.stderr + assert proc.returncode == 0, log[-3000:] + assert "is not valid for" not in log, log[-3000:] + assert "widening rholm storage for bandlimited time support" in log, ( + "the storage window was not widened for the distance-marginalized " + "bandlimited run: %s" % log[-3000:]) + assert "time-marginalization quadrature: bandlimited" in log, log[-3000:] + row = np.atleast_2d(np.loadtxt(str(out) + "_0_.dat")) + assert row.shape[1] == 13 and np.isfinite(row[0, 9]), row diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_cache.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_cache.py new file mode 100644 index 000000000..16413b203 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_cache.py @@ -0,0 +1,786 @@ +import concurrent.futures +import hashlib +import json +import os +import subprocess +import sys +import textwrap +import threading +import zipfile +from pathlib import Path + +import pytest + +from RIFT import jax_cache as cache + + +COMPAT = { + "python": "3.11.9", "jax": "0.4.35", "jaxlib": "0.4.35", + "accelerator_plugins": {"jax-cuda12-plugin": "0.4.35"}, + "backend": "gpu", "platform_version": "CUDA 12.4", + "device_kind": "NVIDIA A30", "compute_capability": "8.0", +} + + +class _Config: + jax_persistent_cache_enable_xla_caches = None + + def __init__(self): + self.updates = [] + + def update(self, name, value): + self.updates.append((name, value)) + + +class _Jax: + config = _Config() + + +def test_configure_uses_compatibility_namespace(tmp_path, monkeypatch): + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + monkeypatch.setattr(cache, "runtime_compatibility", lambda unused: COMPAT) + fake = _Jax() + selected = cache.configure_persistent_cache(fake, ["--jax-cache-dir", str(tmp_path)]) + assert selected == (tmp_path / cache.compatibility_key(COMPAT)).resolve() + assert os.environ["JAX_COMPILATION_CACHE_DIR"] == str(selected) + manifest = json.loads((selected / cache.MANIFEST_NAME).read_text()) + assert manifest["compatibility"] == COMPAT + assert ("jax_persistent_cache_enable_xla_caches", "") in fake.config.updates + assert ("jax_enable_compilation_cache", True) in fake.config.updates + + +def test_configure_supports_jax_before_auxiliary_xla_caches(tmp_path, monkeypatch): + """JAX 0.4.24 has executable caching but not the path-valued XLA option.""" + class LegacyConfig: + def __init__(self): + self.updates = [] + + def update(self, name, value): + if name == "jax_persistent_cache_enable_xla_caches": + raise AttributeError("Unrecognized config option") + self.updates.append((name, value)) + + class LegacyJax: + config = LegacyConfig() + + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + monkeypatch.setattr(cache, "runtime_compatibility", lambda unused: COMPAT) + selected = cache.configure_persistent_cache( + LegacyJax(), ["--jax-cache-dir", str(tmp_path)]) + assert selected == (tmp_path / cache.compatibility_key(COMPAT)).resolve() + assert ("jax_enable_compilation_cache", True) in LegacyJax.config.updates + + +@pytest.mark.parametrize("argv,env", [ + (["--no-jax-persistent-cache"], None), + ([], "RIFT_DISABLE_JAX_CACHE"), +]) +def test_disable_does_not_create_cache(tmp_path, monkeypatch, argv, env): + """Each opt-out separately, against a fake that COULD have succeeded. + + A mutation sweep replaced the whole disable condition with ``False`` and + this test still passed. The reason is that the bare ``_Jax()`` fake has no + default_backend/devices, so with the early return gone the run instead hit + the device-probe fail-open handler -- which ALSO returns None, ALSO records + jax_enable_compilation_cache=False, and ALSO leaves tmp_path empty. Every + observable the test checked was reproduced by a different code path, so it + was pinning nothing. Stubbing runtime_compatibility gives the fake a + working probe, and naming a cache root means a non-disabled run must create + a directory. Both halves of the condition get their own case, because the + sweep flipped them together. + """ + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + monkeypatch.setattr(cache, "runtime_compatibility", lambda unused: COMPAT) + if env: + monkeypatch.setenv(env, "1") + fake = _Jax() + assert cache.configure_persistent_cache( + fake, argv + ["--jax-cache-dir", str(tmp_path)]) is None + assert ("jax_enable_compilation_cache", False) in fake.config.updates + assert not list(tmp_path.iterdir()), ( + "a disabled cache must not create its namespace directory") + + +@pytest.mark.parametrize("spelling", ["separate", "equals"]) +def test_cache_root_is_read_from_either_cli_spelling(tmp_path, monkeypatch, + spelling): + """optparse accepts --jax-cache-dir X and --jax-cache-dir=X; so must this. + + configure_persistent_cache scans sys.argv itself, before the option parser + exists, so the two spellings are two separate branches. Only the separate + form was covered: a mutation sweep deleted the "=" branch and the whole + file still passed. With it gone, `--jax-cache-dir=/shared/cache` silently + selects the DEFAULT root instead -- no error, and no reuse of the shared + cache the operator asked for. + """ + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + monkeypatch.delenv("RIFT_JAX_CACHE_ROOT", raising=False) + monkeypatch.setattr(cache, "runtime_compatibility", lambda unused: COMPAT) + # a default root that is NOT tmp_path, so falling back is visible + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path / "default")) + argv = ([str(tmp_path / "asked")] if spelling == "separate" else []) + argv = (["--jax-cache-dir"] + argv if spelling == "separate" + else ["--jax-cache-dir=" + str(tmp_path / "asked")]) + selected = cache.configure_persistent_cache(_Jax(), argv) + assert selected == ( + tmp_path / "asked" / cache.compatibility_key(COMPAT)).resolve(), selected + assert not (tmp_path / "default").exists(), ( + "the requested root was ignored and the default was used") + + +def test_condor_scratch_is_the_default_root(tmp_path, monkeypatch): + monkeypatch.delenv("RIFT_JAX_CACHE_ROOT", raising=False) + monkeypatch.delenv("XDG_CACHE_HOME", raising=False) + monkeypatch.setenv("_CONDOR_SCRATCH_DIR", str(tmp_path)) + assert cache.default_cache_root() == tmp_path / ".rift_cache" / "jax" + + +def test_bundle_option_scan_uses_last_cli_value(): + assert cache.argv_option(["--jax-cache-bundle", "old.zip", + "--jax-cache-bundle=new.zip"], + "--jax-cache-bundle") == "new.zip" + + +def test_cache_cli_help_does_not_require_optional_jax(): + script = Path(__file__).resolve().parents[2] / "bin" / "rift_jax_cache" + code = textwrap.dedent(""" + import runpy + import sys + sys.modules["jax"] = None + sys.argv = ["rift_jax_cache", "--help"] + runpy.run_path(%r, run_name="__main__") + """ % str(script)) + completed = subprocess.run([sys.executable, "-c", code], check=False, + capture_output=True, text=True, timeout=30) + assert completed.returncode == 0, completed.stderr + assert "Inspect, export, and safely import" in completed.stdout + + +def test_unwritable_cache_disables_without_failing(monkeypatch, capsys): + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + monkeypatch.setattr(cache, "runtime_compatibility", lambda unused: COMPAT) + monkeypatch.setattr(Path, "mkdir", lambda *args, **kwargs: (_ for _ in ()).throw(OSError("read only"))) + fake = _Jax() + assert cache.configure_persistent_cache(fake, ["--jax-cache-dir", "/unwritable"]) is None + assert "disabling JAX persistent cache" in capsys.readouterr().err + assert ("jax_enable_compilation_cache", False) in fake.config.updates + + +def test_an_unusable_backend_disables_the_cache_instead_of_raising(tmp_path, + monkeypatch, + capsys): + """A device probe that fails must not kill the driver. + + runtime_compatibility() calls default_backend()/devices(), which force + backend init; unloadable CUDA libraries, or a card busy for every tenant, + raise there. configure_persistent_cache runs at driver IMPORT, before the + option parser exists, so an escaping exception makes the ILE unable even to + print --help. Verified against the real jaxlib: with JAX_PLATFORMS=cuda on + a CPU-only host, jax.devices() raised and this function propagated it. + """ + monkeypatch.delenv("JAX_COMPILATION_CACHE_DIR", raising=False) + + def _explode(unused): + raise RuntimeError("cuInit(0) failed: Unknown CUDA error 303") + + monkeypatch.setattr(cache, "runtime_compatibility", _explode) + fake = _Jax() + assert cache.configure_persistent_cache( + fake, ["--jax-cache-dir", str(tmp_path)]) is None + assert "disabling JAX persistent cache" in capsys.readouterr().err + assert ("jax_enable_compilation_cache", False) in fake.config.updates + assert not list(tmp_path.iterdir()), "a failed probe must not create a cache" + + +def test_manifest_updates_use_unique_atomic_temporary_files(tmp_path, monkeypatch): + sources = [] + lock = threading.Lock() + real_replace = os.replace + + def recording_replace(source, target): + with lock: + sources.append(str(source)) + real_replace(source, target) + + monkeypatch.setattr(cache.os, "replace", recording_replace) + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool: + list(pool.map(lambda i: cache._write_manifest(tmp_path, COMPAT, {"writer": i}), + range(24))) + assert len(sources) == 24 + assert len(set(sources)) == 24 + assert json.loads((tmp_path / cache.MANIFEST_NAME).read_text())["writer"] in range(24) + + +def test_runtime_fingerprint_records_current_accelerator_plugins(monkeypatch): + class Client: + platform_version = "PJRT CUDA 13" + + class Device: + client = Client() + device_kind = "Future GPU" + compute_capability = (13, 0) + + class Jax: + __version__ = "1.0" + + @staticmethod + def default_backend(): + return "gpu" + + @staticmethod + def devices(backend): + assert backend == "gpu" + return [Device()] + + versions = {"jaxlib": "1.0", "jax-cuda13-plugin": "1.0"} + monkeypatch.setattr(cache, "_package_version", versions.get) + identity = cache.runtime_compatibility(Jax) + assert identity["accelerator_plugins"] == {"jax-cuda13-plugin": "1.0"} + assert identity["compute_capability"] == "13.0" + + +def test_bundle_round_trip_and_profile_guard(tmp_path): + source = tmp_path / "source" + (source / "nested").mkdir(parents=True) + (source / "nested" / "compiled-entry").write_bytes(b"compiled") + bundle = tmp_path / "warm.zip" + cache.export_bundle(source, bundle, COMPAT, "o4-laplace", {"n_chunk": 8000}) + destination = cache.import_bundle(bundle, tmp_path / "target", COMPAT, "o4-laplace") + assert (destination / "nested" / "compiled-entry").read_bytes() == b"compiled" + records = sorted(destination.glob(cache.IMPORT_MANIFEST_PREFIX + "*.json")) + assert len(records) == 1 + manifest = json.loads(records[0].read_text()) + assert manifest["static_shapes"] == {"n_chunk": 8000} + cache._write_manifest(destination, COMPAT) + assert json.loads(records[0].read_text()) == manifest + + # Compatible bundles merge compiler entries, so provenance must retain + # every contributor rather than silently replacing the previous profile. + (source / "nested" / "second-entry").write_bytes(b"second") + second_bundle = tmp_path / "second.zip" + cache.export_bundle(source, second_bundle, COMPAT, "o4-exact", + {"n_chunk": 1000}) + cache.import_bundle(second_bundle, tmp_path / "target", COMPAT, + "o4-exact") + records = sorted(destination.glob(cache.IMPORT_MANIFEST_PREFIX + "*.json")) + assert len(records) == 2 + imported_profiles = { + json.loads(path.read_text())["imported_profile"] for path in records} + assert imported_profiles == {"o4-laplace", "o4-exact"} + cache.import_bundle(bundle, tmp_path / "target", COMPAT, "o4-laplace") + assert len(list(destination.glob( + cache.IMPORT_MANIFEST_PREFIX + "*.json"))) == 2 + + # A cache warmed by an older PR may still contain the former singular + # import record. Neither legacy nor current provenance is compiler data. + (destination / cache.IMPORT_MANIFEST_NAME).write_text("{}\n") + reexport = tmp_path / "reexport.zip" + reexport_manifest = cache.export_bundle(destination, reexport, COMPAT) + exported_names = {Path(rel).name for rel in reexport_manifest["files"]} + assert cache.MANIFEST_NAME not in exported_names + assert cache.IMPORT_MANIFEST_NAME not in exported_names + assert not any(name.startswith(cache.IMPORT_MANIFEST_PREFIX) + for name in exported_names) + with pytest.raises(ValueError, match="profile"): + cache.import_bundle(bundle, tmp_path / "wrong-profile", COMPAT, "other") + + exact = tmp_path / "standard-jax-exact-dir" + imported = cache.import_bundle(bundle, tmp_path / "ignored-root", COMPAT, + destination=exact) + assert imported == exact + assert (exact / "nested" / "compiled-entry").read_bytes() == b"compiled" + + +def test_import_publishes_cache_entries_atomically(tmp_path, monkeypatch): + source = tmp_path / "source" + source.mkdir() + (source / "entry").write_bytes(b"compiled") + bundle = tmp_path / "warm.zip" + cache.export_bundle(source, bundle, COMPAT) + replacements = [] + real_replace = os.replace + + def recording_replace(temporary, target): + replacements.append((Path(temporary), Path(target))) + real_replace(temporary, target) + + monkeypatch.setattr(cache.os, "replace", recording_replace) + destination = cache.import_bundle(bundle, tmp_path / "target", COMPAT) + entry_publications = [(temporary, target) for temporary, target in replacements + if target == destination / "entry"] + assert len(entry_publications) == 1 + temporary, target = entry_publications[0] + assert temporary.parent == target.parent + assert temporary != target + + +def test_bundle_rejects_runtime_mismatch_and_tampering(tmp_path): + source = tmp_path / "source" + source.mkdir() + (source / "entry").write_bytes(b"one") + bundle = tmp_path / "warm.zip" + cache.export_bundle(source, bundle, COMPAT) + mismatch = dict(COMPAT, jaxlib="0.5.0") + with pytest.raises(ValueError, match="incompatible"): + cache.import_bundle(bundle, tmp_path / "mismatch", mismatch) + + tampered = tmp_path / "tampered.zip" + with zipfile.ZipFile(bundle) as old, zipfile.ZipFile(tampered, "w") as new: + for name in old.namelist(): + new.writestr(name, b"two" if name == "cache/entry" else old.read(name)) + with pytest.raises(ValueError, match="checksum"): + cache.import_bundle(tampered, tmp_path / "tampered", COMPAT) + + unexpected = tmp_path / "unexpected.zip" + with zipfile.ZipFile(bundle) as old, zipfile.ZipFile(unexpected, "w") as new: + for name in old.namelist(): + new.writestr(name, old.read(name)) + new.writestr("unrelated", b"surprise") + with pytest.raises(ValueError, match="unexpected"): + cache.import_bundle(unexpected, tmp_path / "unexpected", COMPAT) + + +def test_bundle_rejects_oversized_or_overcompressed_members(tmp_path, monkeypatch): + source = tmp_path / "source" + source.mkdir() + (source / "entry").write_bytes(b"0" * 10_000) + bundle = tmp_path / "warm.zip" + cache.export_bundle(source, bundle, COMPAT) + + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 100) + with pytest.raises(ValueError, match="size limit"): + cache.import_bundle(bundle, tmp_path / "oversized", COMPAT) + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 20_000) + monkeypatch.setattr(cache, "MAX_BUNDLE_COMPRESSION_RATIO", 2) + with pytest.raises(ValueError, match="compression-ratio"): + cache.import_bundle(bundle, tmp_path / "overcompressed", COMPAT) + + +def test_real_jax_cache_reused_across_fresh_processes(tmp_path): + pytest.importorskip("jax") + code = textwrap.dedent(""" + import jax + import jax.numpy as jnp + from RIFT.jax_cache import configure_persistent_cache + configure_persistent_cache(jax, ["--jax-cache-dir", r"%s"]) + @jax.jit + def work(x): + for _ in range(8): + x = jnp.sin(x @ x + 0.01) + return x.sum() + print(float(work(jnp.eye(64)).block_until_ready())) + """ % tmp_path) + env = os.environ.copy() + env.update({ + "JAX_PLATFORMS": "cpu", + "JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS": "0", + "JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES": "0", + "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", + "TF_NUM_INTRAOP_THREADS": "1", "TF_NUM_INTEROP_THREADS": "1", + "XLA_FLAGS": "--xla_cpu_multi_thread_eigen=false --xla_force_host_platform_device_count=1", + }) + + def run(): + subprocess.run([sys.executable, "-c", code], env=env, check=True, + capture_output=True, text=True, timeout=120) + + def entries(): + return { + str(path.relative_to(tmp_path)): (hashlib.sha256(path.read_bytes()).hexdigest(), + path.stat().st_mtime_ns) + for path in tmp_path.rglob("*") + if path.is_file() and path.name != cache.MANIFEST_NAME + and not path.name.endswith(".tmp") + } + + run() + first = entries() + assert first, "the first fresh process did not populate JAX's persistent cache" + run() + assert entries() == first, "the second fresh process recompiled or rewrote cache entries" + + +def test_real_jax_cache_bundle_reused_from_different_absolute_root(tmp_path): + """A transferred executable must not be keyed by its original cache path.""" + jax = pytest.importorskip("jax") + source_root = tmp_path / "producer" / "cache" + target_root = tmp_path / "consumer-at-a-different-path" / "cache" + code = textwrap.dedent(""" + import jax + import jax.numpy as jnp + from RIFT.jax_cache import configure_persistent_cache + configure_persistent_cache(jax, ["--jax-cache-dir", r"%s"]) + @jax.jit + def transferred_work(x): + for _ in range(8): + x = jnp.sin(x @ x + 0.01) + return x.sum() + print(float(transferred_work(jnp.eye(64)).block_until_ready())) + """) + env = os.environ.copy() + env.pop("JAX_COMPILATION_CACHE_DIR", None) + env.update({ + "JAX_PLATFORMS": "cpu", + "JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS": "0", + "JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES": "0", + "JAX_DEBUG_LOG_MODULES": "jax._src.compiler,jax._src.compilation_cache", + "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", + "TF_NUM_INTRAOP_THREADS": "1", "TF_NUM_INTEROP_THREADS": "1", + "XLA_FLAGS": "--xla_cpu_multi_thread_eigen=false --xla_force_host_platform_device_count=1", + }) + + producer = subprocess.run( + [sys.executable, "-c", code % source_root], env=env, check=True, + capture_output=True, text=True, timeout=120) + compatibility = cache.runtime_compatibility(jax) + source = source_root / cache.compatibility_key(compatibility) + bundle = tmp_path / "portable.zip" + cache.export_bundle(source, bundle, compatibility, "different-root-test") + target = cache.import_bundle(bundle, target_root, compatibility, + "different-root-test") + + def entries(): + return { + str(path.relative_to(target)): (hashlib.sha256(path.read_bytes()).hexdigest(), + path.stat().st_mtime_ns) + for path in target.rglob("*") + if path.is_file() and not cache._is_provenance_file(path) + and not path.name.endswith(".tmp") + } + + imported = entries() + assert imported, "producer did not create any persistent executable entry" + consumer = subprocess.run( + [sys.executable, "-c", code % target_root], env=env, check=True, + capture_output=True, text=True, timeout=120) + assert consumer.stdout == producer.stdout + # JAX publishes an executable cache entry atomically after compilation. A + # fresh compile would therefore replace it and change its mtime; preserving + # every imported byte and mtime pins an actual persistent-cache load without + # depending on JAX's version-specific debug-log formatting. + assert entries() == imported, "consumer recompiled after cache-root transfer" + + +@pytest.mark.parametrize("scheme", ["exact", "laplace"]) +def test_angle_batched_kernel_persists_without_host_effects(tmp_path, scheme): + """Pin both real anglemarg graphs, not a toy matmul cache entry. + + JAX refuses to persist any graph containing debug callbacks. This test + executes the shipped exact coefficient/reconstruction/scan kernel in two + fresh processes and requires its named cache entry to survive unchanged; + reintroducing the former amplitude callback therefore fails behaviorally. + """ + pytest.importorskip("jax") + test_dir = Path(__file__).resolve().parent + code = textwrap.dedent(""" + import sys + sys.path.insert(0, r"%s") + import jax + import jax.numpy as jnp + from RIFT.jax_cache import configure_persistent_cache + configure_persistent_cache(jax, ["--jax-cache-dir", r"%s"]) + from test_angle_marg_exact import make_synth, _dist_grid, RA, DEC, INCL, INTERP + from RIFT.likelihood.jax_ile import anglemarg as AM + data = make_synth(npts=16) + xg, lwg = _dist_grid(data, n=16) + if %r == "exact": + @jax.jit + def persisted_work(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_exact( + data, ra, dec, incl, xg, lwg, interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, + dense_chunk=8, grid_block=8, return_amp=True) + else: + @jax.jit + def persisted_work(ra, dec, incl): + return AM.fused_log_likelihood_distphipsimarg_laplace( + data, ra, dec, incl, xg, lwg, interp=INTERP, + amp_sizing=AM.ANGLE_MARG_CROSSOVER_AMPLITUDE, + phi_chunk=8, dist_block=8, return_amp=True) + value, amp = persisted_work(jnp.asarray(RA), jnp.asarray(DEC), jnp.asarray(INCL)) + print(float(value.block_until_ready()[0]), float(amp.block_until_ready())) + """ % (test_dir, tmp_path, scheme)) + env = os.environ.copy() + env.update({ + "JAX_PLATFORMS": "cpu", + "JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS": "0", + "JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES": "0", + "JAX_DEBUG_LOG_MODULES": "jax._src.compiler,jax._src.compilation_cache", + "OMP_NUM_THREADS": "1", "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", "NUMEXPR_NUM_THREADS": "1", + "TF_NUM_INTRAOP_THREADS": "1", "TF_NUM_INTEROP_THREADS": "1", + "XLA_FLAGS": "--xla_cpu_multi_thread_eigen=false --xla_force_host_platform_device_count=1", + }) + + def run(): + return subprocess.run([sys.executable, "-c", code], env=env, check=True, + capture_output=True, text=True, timeout=180) + + def persisted_entries(): + return { + str(path.relative_to(tmp_path)): (hashlib.sha256(path.read_bytes()).hexdigest(), + path.stat().st_mtime_ns) + for path in tmp_path.rglob("*") + if path.is_file() and "jit_persisted_work-" in path.name + } + + first_run = run() + assert "because it uses host callbacks" not in first_run.stderr + first = persisted_entries() + assert first, "the shipped %s-angle batch graph was not persisted" % scheme + second_run = run() + assert "because it uses host callbacks" not in second_run.stderr + assert persisted_entries() == first, ( + "fresh-process %s kernel cache entry changed" % scheme) + + +# --------------------------------------------------------------------------- +# Bundle-validation guards. Every test below was added because the guard it +# covers SURVIVED a mutation sweep: its condition could be replaced with +# ``False`` and the whole file still passed. A limit nothing reaches is not a +# limit, and this module's whole job is refusing a bundle it should not trust. +# --------------------------------------------------------------------------- + + +def _one_entry_bundle(tmp_path, name="warm.zip", payload=b"compiled"): + source = tmp_path / "source" + source.mkdir(exist_ok=True) + (source / "entry").write_bytes(payload) + bundle = tmp_path / name + cache.export_bundle(source, bundle, COMPAT) + return source, bundle + + +def _rebuild(bundle, target, *, manifest=None, members=None, drop=()): + """Write a modified copy of *bundle*: patched manifest, extra/dropped members.""" + with zipfile.ZipFile(bundle) as old: + original = json.loads(old.read(cache.MANIFEST_NAME)) + data = {n: old.read(n) for n in old.namelist()} + if manifest is not None: + original = manifest(original) + with zipfile.ZipFile(target, "w") as new: + new.writestr(cache.MANIFEST_NAME, + json.dumps(original, indent=2, sort_keys=True) + "\n") + for name, blob in data.items(): + if name == cache.MANIFEST_NAME or name in drop: + continue + new.writestr(name, blob) + for name, blob in (members or {}).items(): + new.writestr(name, blob) + return target + + +def test_import_refuses_a_member_path_escaping_the_cache(tmp_path): + """A '..' member must be refused, not written outside the destination. + + The manifest is what names the files, so a hostile bundle controls those + strings. Without the traversal guard ``temp_root / rel`` resolves above + the extraction directory and the write lands wherever the relative path + points -- while the publication walk, which only rglobs INSIDE temp_root, + never sees the file and reports nothing. + """ + _, bundle = _one_entry_bundle(tmp_path) + escaped = tmp_path / "escape.zip" + with zipfile.ZipFile(bundle) as old: + manifest = json.loads(old.read(cache.MANIFEST_NAME)) + blob = old.read("cache/entry") + digest = manifest["files"].pop("entry") + manifest["files"]["../../escaped-entry"] = digest + with zipfile.ZipFile(escaped, "w") as new: + new.writestr(cache.MANIFEST_NAME, + json.dumps(manifest, indent=2, sort_keys=True) + "\n") + new.writestr("cache/../../escaped-entry", blob) + with pytest.raises(ValueError, match="unsafe"): + cache.import_bundle(escaped, tmp_path / "target", COMPAT) + assert not (tmp_path.parent / "escaped-entry").exists() + + absolute = tmp_path / "absolute.zip" + manifest["files"] = {"/etc/escaped": digest} + with zipfile.ZipFile(absolute, "w") as new: + new.writestr(cache.MANIFEST_NAME, + json.dumps(manifest, indent=2, sort_keys=True) + "\n") + new.writestr("cache//etc/escaped", blob) + with pytest.raises(ValueError, match="unsafe"): + cache.import_bundle(absolute, tmp_path / "target-abs", COMPAT) + + +def test_a_member_declaring_zero_compressed_size_is_refused(tmp_path): + """A member claiming N uncompressed bytes in 0 compressed bytes. + + Reachability, stated because it decides how much this guard is worth: the + ratio check below it short-circuits on compress_size == 0, so removing this + one raises no ZeroDivisionError, and such a member then fails the checksum + during extraction instead. It is defence in depth -- refuse at the header, + with a message naming the reason, rather than reporting a checksum failure + for an archive whose real defect is a lying header. A mutation sweep found + it survived: nothing reached it at all. + + Unit-level on purpose. zipfile writes consistent sizes, so producing this + member end to end means forging central-directory fields; the guard is a + pure predicate on ZipInfo, so drive it directly and say so. + """ + info = zipfile.ZipInfo("cache/entry") + info.file_size = 4096 + info.compress_size = 0 + with pytest.raises(ValueError, match="invalid compressed size"): + cache._validate_member(info) + + # and the honest member with the same declared size passes + info.compress_size = 4096 + cache._validate_member(info) + + +def test_import_refuses_an_unknown_format_version(tmp_path): + """A future bundle layout must fail closed, not be read with today's rules.""" + _, bundle = _one_entry_bundle(tmp_path) + future = _rebuild(bundle, tmp_path / "future.zip", + manifest=lambda m: dict(m, format_version= + cache.FORMAT_VERSION + 1)) + with pytest.raises(ValueError, match="unsupported cache bundle format"): + cache.import_bundle(future, tmp_path / "target", COMPAT) + + +def test_import_refuses_a_bundle_with_no_manifest(tmp_path): + """Without the manifest there is nothing to check compatibility against.""" + _, bundle = _one_entry_bundle(tmp_path) + headless = tmp_path / "headless.zip" + with zipfile.ZipFile(bundle) as old, zipfile.ZipFile(headless, "w") as new: + for name in old.namelist(): + if name != cache.MANIFEST_NAME: + new.writestr(name, old.read(name)) + with pytest.raises(ValueError, match=cache.MANIFEST_NAME): + cache.import_bundle(headless, tmp_path / "target", COMPAT) + + +def test_import_refuses_duplicate_archive_members(tmp_path): + """Two members with one name: readers disagree about which is the content. + + zipfile resolves a duplicate name to the LAST entry, while the checksum + walk and the name-set comparison both see the name only once, so a + duplicate is exactly how a validated bundle and an extracted bundle come + apart. + """ + _, bundle = _one_entry_bundle(tmp_path) + duplicated = tmp_path / "duplicated.zip" + with zipfile.ZipFile(bundle) as old: + data = {n: old.read(n) for n in old.namelist()} + with zipfile.ZipFile(duplicated, "w") as new: + for name, blob in data.items(): + new.writestr(name, blob) + new.writestr("cache/entry", b"second copy") + with pytest.raises(ValueError, match="duplicate"): + cache.import_bundle(duplicated, tmp_path / "target", COMPAT) + + +def test_import_refuses_a_manifest_naming_absent_members(tmp_path): + """Declared-but-missing is the mirror of the extra-member case. + + The suite already covered an UNEXPECTED member; a manifest that promises a + file the archive does not carry took the same branch's other side, and + nothing exercised it. + """ + source = tmp_path / "source" + source.mkdir() + (source / "one").write_bytes(b"a") + (source / "two").write_bytes(b"b") + bundle = tmp_path / "pair.zip" + cache.export_bundle(source, bundle, COMPAT) + truncated = _rebuild(bundle, tmp_path / "truncated.zip", + drop=("cache/two",)) + with pytest.raises(ValueError, match="do not match its manifest"): + cache.import_bundle(truncated, tmp_path / "target", COMPAT) + + +def test_import_bounds_member_count_and_total_size(tmp_path, monkeypatch): + """Both import-side aggregate limits, each on its own. + + They are separate guards with separate messages, and neither was reached: + the file-count ceiling is 100k and the total-size ceiling 16 GiB, so no + honest fixture gets near either. Lower the constants instead of building + a hostile archive. + """ + source = tmp_path / "source" + source.mkdir() + for i in range(4): + (source / ("entry%d" % i)).write_bytes(b"0" * 64) + bundle = tmp_path / "many.zip" + cache.export_bundle(source, bundle, COMPAT) + + monkeypatch.setattr(cache, "MAX_BUNDLE_FILES", 2) + with pytest.raises(ValueError, match="too many archive members"): + cache.import_bundle(bundle, tmp_path / "count", COMPAT) + + monkeypatch.setattr(cache, "MAX_BUNDLE_FILES", 100_000) + monkeypatch.setattr(cache, "MAX_BUNDLE_TOTAL_BYTES", 100) + with pytest.raises(ValueError, match="total size limit"): + cache.import_bundle(bundle, tmp_path / "total", COMPAT) + + +def test_import_member_size_is_checked_before_and_during_extraction(tmp_path, + monkeypatch): + """The declared size and the streamed size are two guards, not one. + + Each masked the other in the sweep: defeating either alone still raised + "size limit" from its partner, so both read as covered while neither was. + The header check refuses a member whose DECLARED size is too large; the + streaming check refuses one that lies about it and keeps producing bytes. + """ + _, bundle = _one_entry_bundle(tmp_path, payload=b"0" * 4_000) + + # Header guard alone. Matching "size limit" is NOT enough to pin it: the + # streaming guard raises the same message, so that assertion passes with + # this guard deleted. What only the header guard can do is refuse BEFORE + # any extraction begins, so make reaching extraction an error. + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 100) + + class _ExtractionReached(Exception): + pass + + def _no_extraction(*args, **kwargs): + raise _ExtractionReached("import began extracting an oversized member") + + monkeypatch.setattr(cache.tempfile, "TemporaryDirectory", _no_extraction) + with pytest.raises(ValueError, match="size limit"): + cache.import_bundle(bundle, tmp_path / "declared", COMPAT) + monkeypatch.undo() + + # Streaming guard alone: headers pass, extraction must still stop. A + # ZipInfo reporting a small file_size passes _validate_member, so only the + # byte counter in the extraction loop can catch the real length. + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 1_000) + real_validate = cache._validate_member + monkeypatch.setattr(cache, "_validate_member", + lambda info, **kw: None if not kw else + real_validate(info, **kw)) + with pytest.raises(ValueError, match="size limit"): + cache.import_bundle(bundle, tmp_path / "streamed", COMPAT) + + +def test_export_refuses_a_cache_too_large_or_too_numerous_to_bundle(tmp_path, + monkeypatch): + """The export-side ceilings, which no fixture came near either. + + Export builds the bundle from a directory this process already trusts, so + these are resource guards rather than security ones -- but an unbounded + export is how a 16 GiB cache becomes an OOM on a submit node. + """ + source = tmp_path / "source" + source.mkdir() + for i in range(3): + (source / ("entry%d" % i)).write_bytes(b"0" * 512) + + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 100) + with pytest.raises(ValueError, match="member exceeds the bundle size"): + cache.export_bundle(source, tmp_path / "a.zip", COMPAT) + + monkeypatch.setattr(cache, "MAX_BUNDLE_MEMBER_BYTES", 4 * 1024**3) + monkeypatch.setattr(cache, "MAX_BUNDLE_TOTAL_BYTES", 600) + with pytest.raises(ValueError, match="total bundle size limit"): + cache.export_bundle(source, tmp_path / "b.zip", COMPAT) + + monkeypatch.setattr(cache, "MAX_BUNDLE_TOTAL_BYTES", 16 * 1024**3) + monkeypatch.setattr(cache, "MAX_BUNDLE_FILES", 2) + with pytest.raises(ValueError, match="too many files"): + cache.export_bundle(source, tmp_path / "c.zip", COMPAT) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_phase_marg_mode_order.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_phase_marg_mode_order.py new file mode 100644 index 000000000..41c10d7d5 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_phase_marg_mode_order.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python +"""Phase marginalization must accept either packed order of the (2,+-2) pair. + +WHAT WENT WRONG. ``_accumulate_unit``'s phase-marginalized branch is +position-dependent: it conjugates column 1 of ``Y`` and of ``Q`` and pairs +column 1 with ``conj(F)``. It enforced that position by REFUSING any ``lms`` +other than the literal list ``[(2,2), (2,-2)]``:: + + NotImplementedError: phase marginalization currently requires modes + [(2,2),(2,-2)]; got [(2, -2), (2, 2)] + +The packed column order is not the caller's to choose -- it comes from a python +dict's iteration order in the precompute upstream -- so a correctly configured +``--phase-marginalization`` run could arrive with complete, valid data and simply +die. Found 2026-09-06 driving PR #266 configurations; the campaign worked around +it by permuting ``lms``, the ``Q`` columns and ``U``/``V`` on BOTH indices itself, +and measured that permutation neutral to 4.5e-13 nats. The fix moves that +permutation into the library. + +WHY THESE TESTS LOOK LIKE THIS. + + * The equality test drives BOTH orders through the same synthetic likelihood. + It is not a test of the permutation helper: a helper-level assertion cannot + see a call site that stops calling the helper, and this module's recurring + defect class is guards that look like coverage and are not. + + * ``U`` and ``V`` are (K,K) with the mode index on BOTH axes. Permuting one + axis returns a WRONG likelihood with no error, so + ``test_one_axis_relabelling_is_detectable`` pins that the fixture can see + that mistake -- without it the equality test would pass under a half-fixed + implementation. + + * The bitwise tests protect the ordering that already works. Nothing about + the numbers a working run produces may change, so the canonical order must + not merely agree to tolerance: it must take the untouched code path. + +FLOATING POINT. x64 is requested below; several assertions here are bitwise and +would be meaningless (or spuriously loose) in float32. +""" + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp # noqa: E402 + +from RIFT.likelihood.jax_ile import build_likelihood_data # noqa: E402 +from RIFT.likelihood.jax_ile import core as _core # noqa: E402 +from RIFT.likelihood.jax_ile.core import ( # noqa: E402 + _accumulate_unit, _permute_modes, _phase_marg_permutation, + make_log_likelihood) + +TREF = 1126259462.413 +CANONICAL = ((2, 2), (2, -2)) +SWAPPED = ((2, -2), (2, 2)) +INTERPS = ("nearest", "linear", "cubic", "sinc") + + +# --------------------------------------------------------------------------- +# Fixture. Structurally faithful packed data (U Hermitian PD, V complex +# symmetric), the construction test_angle_marg_smoke / test_distance_grid use. +# --------------------------------------------------------------------------- + +def _packed(seed=3, npts=32, deltaT=1.0 / 1024, modes=CANONICAL): + rng = np.random.default_rng(seed) + K = len(modes) + out = {} + 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)) + M = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + U = M @ M.conj().T + 3 * np.eye(K) + B = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = (B @ B.T) * 0.3 + out[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=TREF - 0.5) + return out + + +def _relabel(pk, perm, u_axes=(0, 1), v_axes=(0, 1)): + """Repack the SAME physics with the mode axis reordered by ``perm``. + + ``u_axes`` / ``v_axes`` exist only so a test can build a deliberately + HALF-relabelled bank; the honest relabelling permutes both axes of each. + """ + p = np.asarray(perm, dtype=int) + out = {} + for det, d in pk.items(): + U = np.asarray(d["U"]) + V = np.asarray(d["V"]) + for ax in u_axes: + U = np.take(U, p, axis=ax) + for ax in v_axes: + V = np.take(V, p, axis=ax) + out[det] = dict(lms=np.asarray(d["lms"])[p], + rholmArray=np.asarray(d["rholmArray"])[p], + U=U, V=V, epoch=d["epoch"]) + return out + + +def _data(pk, npts=32, deltaT=1.0 / 1024): + tw = npts * deltaT / 2.0 + return build_likelihood_data(pk, deltaT, TREF, np.linspace(-tw, tw, npts)) + + +def _angles(S=5, seed=11): + rng = np.random.default_rng(seed) + return [jnp.asarray(x) for x in (rng.uniform(0.0, 2 * np.pi, S), + rng.uniform(-1.2, 1.2, S), + rng.uniform(0.0, np.pi, S), + rng.uniform(0.0, np.pi, S), + rng.uniform(0.0, 2 * np.pi, S))] + + +def _acc(pk, interp="cubic", guard=0, th=None): + k, r = _accumulate_unit(_data(pk), *(th or _angles()), interp, True, + guard=guard) + return np.asarray(k), np.asarray(r) + + +# --------------------------------------------------------------------------- +# 1. The defect: both orders must be accepted, and must agree. +# --------------------------------------------------------------------------- + +def test_swapped_order_is_accepted_at_all(): + """The bug was an outright refusal, so pin the refusal's absence first -- + an equality test alone would report an ERROR, not a diagnosis.""" + _accumulate_unit(_data(_relabel(_packed(), [1, 0])), *_angles(), + "cubic", True) + + +def test_both_orders_give_the_same_accumulation(): + """Same physics, two packings: the accumulation must not know the + difference. Every stencil, guarded and unguarded.""" + pk = _packed() + th = _angles() + for interp in INTERPS: + for guard in (0, 3): + k0, r0 = _acc(pk, interp, guard, th) + k1, r1 = _acc(_relabel(pk, [1, 0]), interp, guard, th) + scale = max(np.abs(k0).max(), 1.0) + assert np.abs(k0 - k1).max() <= 1e-12 * scale, ( + "interp=%s guard=%d: kappa differs by %.3g between mode orders" + % (interp, guard, np.abs(k0 - k1).max())) + assert np.abs(r0 - r1).max() <= 1e-12 * max(np.abs(r0).max(), 1.0), ( + "interp=%s guard=%d: rho^2 differs by %.3g between mode orders" + % (interp, guard, np.abs(r0 - r1).max())) + + +def test_both_orders_give_the_same_lnL_through_the_public_seam(): + """Through ``make_log_likelihood`` -- the seam a driver actually calls -- + not only through the private accumulator.""" + pk = _packed() + ra, dec, psi, incl, phiref = _angles() + dist = jnp.full(ra.shape, 400.0) + f0 = make_log_likelihood(_data(pk), interp="cubic", + phase_marginalization=True) + f1 = make_log_likelihood(_data(_relabel(pk, [1, 0])), interp="cubic", + phase_marginalization=True) + l0 = np.asarray(f0(ra, dec, psi, incl, phiref, dist)) + l1 = np.asarray(f1(ra, dec, psi, incl, phiref, dist)) + assert np.isfinite(l0).all(), "fixture produced a non-finite lnL" + assert np.abs(l0 - l1).max() <= 1e-10, ( + "lnL differs by %.3g nats between packed mode orders" % np.abs(l0 - l1).max()) + + +# --------------------------------------------------------------------------- +# 2. U and V carry the mode index on BOTH axes. +# --------------------------------------------------------------------------- + +def test_one_axis_relabelling_is_detectable(): + """The equality test above is only a test of "permute BOTH axes" if the + fixture can tell a half-permutation apart. Assert that it can, per matrix + and per axis, in the units the assertion is made in. + + Without this, an implementation that permuted only ``U[perm]`` would pass + every other test in this file on a fixture whose U happened to be + symmetric, and would return a silently wrong likelihood in production. + """ + pk = _packed() + th = _angles() + _, r0 = _acc(pk, th=th) + tol = 1e-12 * max(np.abs(r0).max(), 1.0) + for name, kw in (("U axis 0 only", dict(u_axes=(0,))), + ("U axis 1 only", dict(u_axes=(1,))), + ("V axis 0 only", dict(v_axes=(0,))), + ("V axis 1 only", dict(v_axes=(1,)))): + _, rb = _acc(_relabel(pk, [1, 0], **kw), th=th) + assert np.abs(r0 - rb).max() > 1e6 * tol, ( + "%s is invisible in rho^2 (max diff %.3g <= %.3g): this fixture " + "cannot detect a half-permuted U/V, so the equality tests do not " + "gate it" % (name, np.abs(r0 - rb).max(), 1e6 * tol)) + + +def test_permute_modes_permutes_both_axes_of_U_and_V(): + """Element-level contract of the helper, independent of any contraction: + ``out[i, j] == in[perm[i], perm[j]]``.""" + rng = np.random.default_rng(5) + K = 2 + U = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + Q = rng.standard_normal((7, K)) + 1j * rng.standard_normal((7, K)) + perm = [1, 0] + lms2, Q2, U2, V2 = _permute_modes(list(SWAPPED), jnp.asarray(Q), + jnp.asarray(U), jnp.asarray(V), perm) + assert lms2 == list(CANONICAL) + for i in range(K): + assert np.array_equal(np.asarray(Q2)[:, i], Q[:, perm[i]]) + for j in range(K): + assert np.asarray(U2)[i, j] == U[perm[i], perm[j]], "U axis pair" + assert np.asarray(V2)[i, j] == V[perm[i], perm[j]], "V axis pair" + + +# --------------------------------------------------------------------------- +# 3. The ordering that already works must be untouched -- bitwise. +# --------------------------------------------------------------------------- + +def test_canonical_order_never_enters_the_permutation_path(): + """The strongest available statement of "no numerical change for data that + already works": the canonical order does not merely agree to tolerance, it + executes the SAME operations it executed before this fix, because the + permutation is skipped entirely rather than applied as an identity. + + Asserted by making the permutation helper fatal for the duration -- a + counter that is merely checked at the end can be defeated by a call that + happens somewhere the counter does not look. + """ + assert _phase_marg_permutation(list(CANONICAL)) is None, ( + "canonical order must yield None (skip), not an identity permutation: " + "an identity `take` is still a new node in the XLA graph") + + def _fatal(*a, **kw): + raise AssertionError( + "_permute_modes was called for the canonical mode order; the " + "already-working path is no longer bit-for-bit what it was") + + saved = _core._permute_modes + _core._permute_modes = _fatal + try: + for interp in INTERPS: + _accumulate_unit(_data(_packed()), *_angles(), interp, True) + finally: + _core._permute_modes = saved + + +def test_relabelling_is_bitwise_exact_not_merely_close(): + """A permutation moves bytes; it does not arithmetic on them. After + canonicalization the swapped bank is bit-identical to the canonical one, so + every downstream op sees identical inputs and the outputs must match to the + last bit. + + Pinned bitwise on purpose. A drift to ~1e-16 here would mean the + canonicalization stopped being exact data movement (a cast, a reassociating + fusion) -- worth a red build and a look, not a widened tolerance. + """ + pk = _packed() + th = _angles() + for interp in INTERPS: + k0, r0 = _acc(pk, interp, th=th) + k1, r1 = _acc(_relabel(pk, [1, 0]), interp, th=th) + assert k0.tobytes() == k1.tobytes(), ( + "interp=%s: kappa not bitwise identical across mode orders " + "(max diff %.3g)" % (interp, np.abs(k0 - k1).max())) + assert r0.tobytes() == r1.tobytes(), ( + "interp=%s: rho^2 not bitwise identical across mode orders " + "(max diff %.3g)" % (interp, np.abs(r0 - r1).max())) + + +def test_non_phase_marginalized_path_is_untouched(): + """``phase_marginalization=False`` must not canonicalize anything: it never + refused an order, and its contractions are order-symmetric, so touching it + would change working numbers for no benefit.""" + def _fatal(*a, **kw): + raise AssertionError("_permute_modes called with phase_marginalization=False") + + saved = _core._permute_modes + _core._permute_modes = _fatal + try: + for pk in (_packed(), _relabel(_packed(), [1, 0])): + _accumulate_unit(_data(pk), *_angles(), "cubic", False) + finally: + _core._permute_modes = saved + + +# --------------------------------------------------------------------------- +# 4. Only the ORDER is free. Every other mode set is still a real gap. +# --------------------------------------------------------------------------- + +@pytest.mark.parametrize("modes", [ + ((2, 2),), # one mode + ((2, 2), (2, -2), (3, 3)), # a third mode + ((2, 2), (2, 1)), # right count, wrong pair + ((2, 1), (2, -1)), # a different m pair entirely + ((2, 2), (2, 2)), # duplicated +]) +def test_other_mode_sets_still_raise(modes): + """Widening the ORDER must not have widened the SET. The conjugation the + accumulator applies is specific to one m=+2/m=-2 pair; silently accepting a + third mode would drop it from the likelihood instead of failing.""" + with pytest.raises(NotImplementedError): + _phase_marg_permutation([tuple(m) for m in modes]) + + +def test_the_refusal_is_reachable_from_the_accumulator(): + """...and the accumulator still surfaces it, rather than the helper being + correct in isolation while nothing calls it.""" + with pytest.raises(NotImplementedError): + _accumulate_unit(_data(_packed(modes=((2, 2), (2, 1)))), *_angles(), + "cubic", True) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_q_time_pregrid.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_q_time_pregrid.py new file mode 100644 index 000000000..9fbb5a355 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_q_time_pregrid.py @@ -0,0 +1,720 @@ +"""Q time pregrid on the JAX arm: correctness, the factor-1 identity, and its guards. + +WHAT IS BEING TESTED, AND AGAINST WHAT + +``--q-time-pregrid-factor`` refines the STORED rholm buffers onto a finer time grid +once, at build time, so the per-sample gather interpolates over a step ``factor`` +times shorter. The integration cadence does not change. The claim is therefore +purely about interpolation error, and the only honest reference for interpolation +error is a Q whose exact value at arbitrary times is known independently. + +THE ORACLE. ``ComputeModeIPTimeSeries`` ends in +``CutCOMPLEX16TimeSeries(rhoTS, 0, N_window)``: the rholm buffer is a CROP of an +inverse-FFT series that is exactly periodic over the whole data segment. So the +fixture here builds precisely that -- a band-limited series with a known finite +Fourier sum over a long period, cropped -- and evaluates the truth at arbitrary +real positions by summing that Fourier series directly. Nothing in the fixture +reuses the stencil, the reflection, or the FFT-upsampler under test. + +That matters more than it might look, because the alternative references are both +CIRCULAR here. Converging the Lanczos half-width ``a`` converges onto the +truncated-sinc/zero-extension limit; converging the pregrid factor converges onto +the reflected limit. They are different limits, they differ at the buffer ends, +and neither is the truth. The exact Fourier sum is. + +WHICH REFLECTION. The two reflected upsamplers in this codebase disagree, and +each docstring asserts its own convention is the right one: +``jax_ile.core._reflected_fft_upsample`` omits the duplicate turning samples +(period ``2(n-1)``), ``time_marginalization_quadrature.reflected_bandlimited_upsample`` +duplicates them (period ``2n``). They are describing different problems -- see +``core.build_q_time_pregrid`` -- and ``test_duplicated_reflection_is_the_right_one_for_a_crop`` +below measures them against the oracle so the choice is evidence rather than +inheritance. +""" +import os +import sys + +import numpy as np +import pytest + +import jax +import jax.numpy as jnp + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + +from RIFT.likelihood.jax_ile import core as C +from RIFT.likelihood.jax_ile import build_likelihood_data, fused_log_likelihood +from RIFT.likelihood.time_interp_choice import SINC_HALFWIDTH_DEFAULT + +# GEOMETRY MIRRORS PRODUCTION, and that turned out to matter more than any +# threshold in this file. The rholm buffer is 2*0.15 s at 4096 Hz = 1229 samples +# and the marginalization window is +-0.075 s = 614 samples centred in it, so a +# gathered position sits ~300 samples clear of the buffer ends. +# +# The first version of this fixture used a 257-sample crop and evaluated 10 samples +# from its ends. Every refined-grid number it produced was the REFLECTION BOUNDARY +# error, not the stencil error: factor 8 and factor 32 agreed to 1%, the measured +# convergence rate was 0.96x per doubling instead of ~16x, and the pregrid looked +# only 9x better than production instead of 408x. A fixture can be wrong about the +# thing it is measuring while every assertion in it still passes. +N_LONG = 4096 # the oracle's period; long relative to the crop, as the real segment is +# Band edge as a fraction of the long series' Nyquist. The rholm timeseries is +# close to critically sampled -- that is exactly why its stencil error is large -- +# so a fixture that band-limits gently would make every stencil look good and +# would not distinguish them. +BAND_FRACTION = 0.92 +CROP_START = 700 +CROP_N = 1229 # 2*0.15 s at 4096 Hz, the production buffer +CLEARANCE = 300 # production distance from a gathered position to the buffer end + + +def _oracle_series(seed=11, n_long=N_LONG, band_fraction=BAND_FRACTION): + """Coefficients of an exactly periodic, band-limited complex series. + + Returns ``(k, c)`` such that ``f(t) = sum_j c[j] exp(2i pi k[j] t / n_long)`` + is exact for any real ``t``, and ``f`` at integer ``t`` is a legitimate stand-in + for an inverse-FFT rholm series over its full period. + + The amplitude envelope decays with |k| so the series looks like a filtered + matched-filter output rather than white noise at the band edge, but the band + edge itself is hard, so the sampling theorem applies exactly. + """ + rng = np.random.default_rng(seed) + kmax = int(band_fraction * (n_long // 2)) + k = np.arange(-kmax, kmax + 1) + env = np.exp(-0.5 * (k / (0.55 * kmax)) ** 2) + 0.15 + c = (rng.standard_normal(k.size) + 1j * rng.standard_normal(k.size)) * env + return k, c + + +def _oracle_eval(k, c, t, n_long=N_LONG): + """Exact ``f(t)`` at arbitrary real ``t`` (array), by direct Fourier sum.""" + t = np.asarray(t, dtype=np.float64) + phase = np.exp(2j * np.pi * np.outer(t.ravel(), k) / float(n_long)) + return (phase @ c).reshape(t.shape) + + +def _cropped_Q(k, c, start=CROP_START, n=CROP_N, n_long=N_LONG): + """The buffer a detector actually gets: ``n`` samples cropped out of the period.""" + return _oracle_eval(k, c, np.arange(start, start + n), n_long) + + +def _gather_at(Q_col, positions, interp, factor=1): + """Evaluate one stencil at crop-local COARSE positions, via the shipped gatherers. + + ``factor`` selects the refined grid: the Q column is pre-refined and the + positions scaled, exactly as :func:`core._q_sample_positions` does. + """ + Q_col = np.asarray(Q_col) + if factor != 1: + fine, _ = C.build_q_time_pregrid(Q_col[None, :], factor) + Q_col = fine[0] + pos = np.asarray(positions, dtype=np.float64) * factor + gather = C._GATHERERS[interp] + return np.asarray(gather(jnp.asarray(Q_col), jnp.asarray(pos[None, :]), + None))[0] + + +def _interior_positions(n=CROP_N, seed=5, count=400, margin=None): + """Fractional crop-local positions at the production distance from the ends. + + ``margin`` defaults to :data:`CLEARANCE`, NOT to the stencil footprint. A + footprint-sized margin is legal for every stencil and still wrong, because it + measures the buffer's boundary condition rather than the interpolation -- see + the geometry note at the top of this file. + """ + if margin is None: + margin = CLEARANCE + rng = np.random.default_rng(seed) + return rng.uniform(margin, n - 1 - margin, size=count) + + +# -------------------------------------------------------------------------- +# 1. The factor-1 path is the historical path, bit for bit. +# -------------------------------------------------------------------------- + +def test_factor_one_returns_the_same_array_object(): + """No copy, no round trip, no reflection: factor 1 must not touch the data. + + Asserting identity rather than equality is deliberate. ``==`` would still pass + if factor 1 quietly went through an FFT and came back within an ulp, and an ulp + of Q is not nothing at rho 652. + """ + rho = np.arange(12, dtype=np.complex128).reshape(2, 6) + out, report = C.build_q_time_pregrid(rho, 1) + assert out is rho + assert report["factor"] == 1 + + +def test_factor_one_positions_are_bit_identical_to_the_pre_pregrid_expressions(): + """``_q_sample_positions`` at factor 1 reproduces the inline code it replaced. + + The two expressions below are verbatim what ``_accumulate_unit`` and + ``_accumulate_unit_banded`` computed before the pregrid landed. Bitwise, not + ``allclose``: the whole factor-1 claim is that nothing moved at all, and the + additive/multiplicative reassociation this change introduces on the factor>1 + branch is exactly the kind of thing that shifts a last bit. + """ + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals) + assert data.q_time_pregrid_factor == 1 + rng = np.random.default_rng(3) + p0 = jnp.asarray(rng.uniform(1000.0, 1200.0, 17)) + t_offsets = jnp.arange(-4, data.npts + 4, dtype=jnp.float64) + for interp in ("nearest", "linear", "cubic", "sinc"): + pos, u = C._q_sample_positions(data, p0, t_offsets, interp) + want_pos = p0[:, None] + t_offsets[None, :] + assert np.array_equal(np.asarray(pos), np.asarray(want_pos)) + if interp == "nearest": + assert u is None + else: + want_u = C._separable_u(p0) + assert np.array_equal(np.asarray(u), np.asarray(want_u)) + + +def test_factor_one_stencil_margins_match_the_table_they_replaced(): + """The margin table was inlined twice; hoisting it must not have changed a value.""" + assert C._STENCIL_MARGIN == {"nearest": 1, "linear": 2, "cubic": 3, + "sinc": SINC_HALFWIDTH_DEFAULT + 1} + + +def _toy_packed(detectors=("H1", "L1"), deltaT=1.0 / 4096, tw=0.02, seed=7): + """Small packed dict built on the oracle series, so the whole file shares one Q.""" + k, c = _oracle_series(seed=seed) + npts = int(2 * tw / deltaT) + tvals = (np.arange(npts) - npts // 2) * deltaT + tref = 1126259462.413 + modes = ((2, 2), (2, -2)) + K = len(modes) + rng = np.random.default_rng(seed + 1) + packed = {} + for j, det in enumerate(detectors): + rho = np.stack([_cropped_Q(*_oracle_series(seed=seed + 10 * j + kk), + n=1024) for kk in range(K)]) + U = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + V = rng.standard_normal((K, K)) + 1j * rng.standard_normal((K, K)) + packed[det] = dict(lms=np.array(modes, dtype=int), rholmArray=rho, + U=U, V=V, epoch=tref - 512 * deltaT) + return packed, tvals, deltaT, tref + + +# -------------------------------------------------------------------------- +# 2. The refined grid is a refinement: it reproduces what it refines. +# -------------------------------------------------------------------------- + +@pytest.mark.parametrize("factor", [2, 4, 8]) +def test_every_factor_th_refined_sample_reproduces_the_input(factor): + k, c = _oracle_series() + Q = _cropped_Q(k, c)[None, :] + fine, report = C.build_q_time_pregrid(Q, factor) + assert fine.shape == (1, (Q.shape[-1] - 1) * factor + 1) + assert np.max(np.abs(fine[..., ::factor] - Q)) < 1e-12 * np.max(np.abs(Q)) + assert report["factor"] == factor + assert report["roundtrip_max"] < 5e-12 + + +def test_refinement_commutes_with_conjugation(): + """Required by the phase-marginalized path, and easy to break unnoticed. + + ``_accumulate_unit`` conjugates the (2,-2) column of the ALREADY-REFINED Q:: + + Q = jnp.concatenate([Q[:, 0:1], jnp.conj(Q[:, 1:2])], axis=1) + + so it uses ``conj(refine(x))`` where the physics wants ``refine(conj(x))``. + Those agree only while the refinement is a real-linear operator -- true of the + reflect/zero-pad/inverse-FFT construction, and NOT something a reader can see + from the call site, which is why it is pinned here rather than left to a + comment. A boundary convention that treated the two halves of the spectrum + asymmetrically (e.g. dumping the whole Nyquist bin into one side instead of + splitting it) would break this and would bias only the phase-marginalized runs. + """ + k, c = _oracle_series() + Q = _cropped_Q(k, c)[None, :] + scale = np.max(np.abs(Q)) + for factor in (2, 8, 16): + a = np.conj(C.build_q_time_pregrid(Q, factor)[0]) + b = C.build_q_time_pregrid(np.conj(Q), factor)[0] + rel = float(np.max(np.abs(a - b)) / scale) + print(" factor %2d max rel |conj(refine) - refine(conj)| %.3e" % (factor, rel)) + assert rel < 1e-13 + + +def test_positions_stay_separable_on_the_refined_grid(): + """frac(pos) must not vary along the time axis, or ``_separable_u`` lies. + + ``_separable_u`` computes ONE fractional offset per sample and hands it to the + gatherer for every time column. That is only legitimate while the time offsets + are exact integers in the units the gather indexes. Scaling as ``pos * factor`` + instead of ``p0*factor + t_offsets*factor`` reassociates the product and breaks + it by up to an ulp per column -- silently, since the result stays finite and + close. This pins the property that makes the memory optimisation sound. + """ + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=8) + rng = np.random.default_rng(4) + p0 = jnp.asarray(rng.uniform(300.0, 700.0, 11)) + t_offsets = jnp.arange(0, data.npts, dtype=jnp.float64) + pos, u = C._q_sample_positions(data, p0, t_offsets, "cubic") + pos = np.asarray(pos) + base = np.floor(pos) + frac = pos - base + # The base index must advance by EXACTLY the factor per column. This is the + # part that is not a rounding nicety: if it did not, a refined window would + # cover the wrong span of time and the likelihood would still look finite. + assert np.array_equal(np.diff(base, axis=1), + np.full((pos.shape[0], pos.shape[1] - 1), 8.0)) + # and the single offset handed to the gatherer must be the fractional part + # every column actually has, to within a rounding of the position itself + # (adding an integer can cross a binade and drop a low mantissa bit). + tol = 8.0 * np.spacing(np.max(np.abs(pos))) + assert np.max(np.abs(frac - np.asarray(u))) <= tol + assert np.max(np.abs(frac - frac[:, :1])) <= tol + + +# -------------------------------------------------------------------------- +# 3. Accuracy against the exact oracle. This is the point of the change. +# -------------------------------------------------------------------------- + +def _stencil_errors(margin=None, count=400): + k, c = _oracle_series() + Q = _cropped_Q(k, c) + pos = _interior_positions(count=count, margin=margin) + exact = _oracle_eval(k, c, CROP_START + pos) + scale = np.max(np.abs(Q)) + out = {} + for name, interp, factor in (("cubic/1", "cubic", 1), + ("sinc/1", "sinc", 1), + ("cubic/8", "cubic", 8), + ("cubic/16", "cubic", 16), + ("sinc/8", "sinc", 8)): + got = _gather_at(Q, pos, interp, factor=factor) + out[name] = float(np.max(np.abs(got - exact)) / scale) + return out + + +def test_pregrid_cubic_beats_the_production_stencil_against_the_oracle(): + """The factor-8 pregrid must cut the interpolation error by orders of magnitude. + + Thresholds are set well inside the measured margins so this is a REGRESSION + gate, not a re-measurement; the measured values are printed for the record. + """ + err = _stencil_errors() + for name in sorted(err): + print(" oracle relative max error %-9s %.3e" % (name, err[name])) + # Measured on this fixture: sinc/1 1.88e-2, cubic/1 1.25e-1, cubic/8 4.62e-5, + # i.e. 408x and 2706x. The gate is 100x, well inside that. + assert err["cubic/8"] < err["sinc/1"] / 100.0 + assert err["cubic/8"] < err["cubic/1"] / 100.0 + # cubic on the refined grid must also beat the 16-tap Lanczos on the SAME + # refined grid (measured 4.62e-5 against 4.86e-4). This is why the pregrid + # ships with cubic rather than inheriting the arm's 'sinc' default: a fixed + # 2a-tap window does not gain from a finer grid the way a 4th-order stencil does. + assert err["cubic/8"] < err["sinc/8"] / 5.0 + + +def test_refined_cubic_error_falls_like_the_fourth_power_of_the_step(): + """Doubling the factor must cut the cubic-Lagrange error by ~16x. + + This is the property that makes the factor a CONVERGENCE knob rather than a + tuning constant: if the observed rate were ~1 the residual would be dominated + by something other than the stencil step and the factor would not be buying + what it claims. The band is wide (6x-40x) because the fixture's error is a max + over random sub-sample phases, not a smooth asymptotic. + """ + k, c = _oracle_series() + Q = _cropped_Q(k, c) + pos = _interior_positions(count=400) + exact = _oracle_eval(k, c, CROP_START + pos) + scale = np.max(np.abs(Q)) + errs = {} + for factor in (2, 4, 8, 16, 32): + got = _gather_at(Q, pos, "cubic", factor=factor) + errs[factor] = float(np.max(np.abs(got - exact)) / scale) + print(" factor %3d relative max error %.3e" % (factor, errs[factor])) + for lo, hi in ((2, 4), (4, 8)): + ratio = errs[lo] / errs[hi] + print(" ratio %d->%d: %.1f" % (lo, hi, ratio)) + assert 8.0 < ratio < 30.0 + # SATURATION, and it is the operational point of the whole measurement: past + # factor ~8 the residual is the reflection boundary condition, which no factor + # can reduce. Measured 8->16 9.5x but 16->32 only 1.4x. That is why the + # shipped factor is 8 and not 32: 32 costs 4x the Q memory for ~10x less error + # than the boundary floor already permits. + assert errs[16] / errs[32] < 4.0 + assert errs[4] / errs[8] > errs[16] / errs[32] + + +def test_duplicated_reflection_is_the_right_one_for_a_crop(): + """Settle 2n vs 2(n-1) against the oracle rather than by inheritance. + + ``core.build_q_time_pregrid`` uses the ``2n`` (duplicated turning samples) form + that the conventional arm shipped, NOT this module's own + ``_reflected_fft_upsample`` (``2(n-1)``). Both are boundary heuristics for a + crop; the choice has to be measured, and it is measured HERE, near the buffer + end where the two actually differ -- in the deep interior both are exact and + the test would be blind by construction. + + If a future change reroutes the pregrid through ``_reflected_fft_upsample`` + "for consistency", this fails. + """ + k, c = _oracle_series() + Q = _cropped_Q(k, c) + n = Q.size + factor = 8 + scale = np.max(np.abs(Q)) + dup, _ = C.build_q_time_pregrid(Q[None, :], factor) + half = np.asarray(C._reflected_fft_upsample(jnp.asarray(Q[None, :]), factor)) + gather = C._GATHERERS["cubic"] + + rng = np.random.default_rng(9) + bands = { + # where production actually gathers + "interior (clearance %d)" % CLEARANCE: _interior_positions(seed=9), + # and hard against the ends, where the two conventions differ most + "near-end": np.concatenate([rng.uniform(4.0, 24.0, 200), + rng.uniform(n - 25.0, n - 5.0, 200)]), + } + for band, pos in bands.items(): + exact = _oracle_eval(k, c, CROP_START + pos) + got = {} + for name, fine in (("2n (shipped)", dup), ("2(n-1)", half)): + v = np.asarray(gather(jnp.asarray(fine[0]), + jnp.asarray((pos * factor)[None, :]), None))[0] + got[name] = float(np.max(np.abs(v - exact)) / scale) + print(" %-26s 2n %.3e 2(n-1) %.3e ratio %.1f" + % (band, got["2n (shipped)"], got["2(n-1)"], + got["2(n-1)"] / got["2n (shipped)"])) + # Measured 16.8x in the interior and 6.9x near the ends. Gate at 2x so + # this is a direction check, not a re-measurement. + assert got["2n (shipped)"] * 2.0 < got["2(n-1)"] + + +# -------------------------------------------------------------------------- +# 4. Guards. Each of these is mutation-tested in the PR; see the description. +# -------------------------------------------------------------------------- + +def test_nearest_is_refused_on_a_refined_grid(): + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=8) + p0 = jnp.asarray([100.0, 200.0]) + t_offsets = jnp.arange(0, 4, dtype=jnp.float64) + with pytest.raises(NotImplementedError, match="nearest"): + C._q_sample_positions(data, p0, t_offsets, "nearest") + # and the same call is fine at factor 1, so the refusal is about the pair + data1 = build_likelihood_data(packed, deltaT, tref, tvals) + C._q_sample_positions(data1, p0, t_offsets, "nearest") + + +def test_unrefined_bank_with_a_declared_factor_fails_closed(): + """A stored Q that was never refined must not be silently mis-indexed. + + This is the failure ``banded._base_data`` would produce if the factor were ever + forwarded to it without refining ``Q_bank``: shapes still broadcast, the + likelihood still returns finite numbers, and a factor-8 window silently covers + an eighth of the intended span. + """ + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=8) + det = data.detector_names[0] + dd = data.detectors[det] + coarse = dd["npts_full_coarse"] + with pytest.raises(ValueError, match="not refined"): + C._check_stored_q_length(dd, coarse, 8, "detector %s Q" % det) + # the real, refined length passes + C._check_stored_q_length(dd, dd["Q"].shape[0], 8, "detector %s Q" % det) + # and the guard is live on the accumulator, not just callable directly + bad = dict(dd) + bad["Q"] = dd["Q"][:coarse] + data.detectors[det] = bad + with pytest.raises(ValueError, match="not refined"): + C._accumulate_unit(data, jnp.asarray([1.0]), jnp.asarray([0.3]), + jnp.asarray([0.5]), jnp.asarray([1.05]), + jnp.asarray([0.7]), "cubic", True) + + +def test_a_declared_factor_without_refinement_metadata_is_refused(): + """The metadata-free escape hatch is for factor 1 only. + + ``_check_stored_q_length`` used to return early whenever ``npts_full_coarse`` + was absent, on the grounds that such a dict cannot have been refined by + ``build_q_time_pregrid``. That is true and it is not the hazard. A caller + can hand-build a detector dict that DECLARES a factor above 1 and omits the + metadata; the early return then skipped every check, and + ``_q_sample_positions`` scaled each index by the factor over a coarse Q. A + factor-8 window covers an eighth of the intended span, shapes broadcast, and + the likelihood returns finite wrong numbers. + + Factor 1 keeps the hatch: no index is scaled, so an unrefined buffer is the + correct buffer. + """ + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals) + det = data.detector_names[0] + bare = {k: v for k, v in data.detectors[det].items() + if k not in ("npts_full_coarse", "q_time_pregrid_factor")} + assert "npts_full_coarse" not in bare + + # factor 1 is still allowed through with no metadata at all + C._check_stored_q_length(bare, bare["Q"].shape[0], 1, "detector Q") + + # above factor 1 the missing metadata is itself the fault + for factor in (2, 8): + with pytest.raises(ValueError, match="npts_full_coarse"): + C._check_stored_q_length(bare, bare["Q"].shape[0], factor, "detector Q") + + # and it is refused even when the stored length would satisfy the arithmetic + # a refined bank of that factor requires, so the guard is not accidentally + # passing on a length coincidence + n = bare["Q"].shape[0] + with pytest.raises(ValueError, match="npts_full_coarse"): + C._check_stored_q_length(bare, (n - 1)*8 + 1, 8, "detector Q") + + +def test_a_refined_bank_reaching_a_factorless_namespace_is_refused(): + """The paired half of the ``getattr(..., 1)`` default in ``_q_sample_positions``. + + Duck-typed ``data`` objects (benchmark shims, several tests in this directory) + do not carry ``q_time_pregrid_factor``, so the lookup defaults to 1. That is + only safe because the detector dict carries its OWN declaration and this check + refuses the mismatch: without it, a refined Q handed to such a namespace would + be indexed at the coarse stride and quietly evaluate the wrong samples. + """ + import types + packed, tvals, deltaT, tref = _toy_packed() + real = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=8) + det = real.detector_names[0] + dd = real.detectors[det] + shim = types.SimpleNamespace( + feature=None, detectors={det: dd}, detector_names=[det], + gmst=real.gmst, deltaT=real.deltaT, npts=real.npts, + tval0=real.tval0, tref_minus_epoch=real.tref_minus_epoch) + assert not hasattr(shim, "q_time_pregrid_factor") + with pytest.raises(ValueError, match="being indexed at factor 1"): + C._accumulate_unit(shim, jnp.asarray([1.0]), jnp.asarray([0.3]), + jnp.asarray([0.5]), jnp.asarray([1.05]), + jnp.asarray([0.7]), "cubic", True) + + +def test_bad_factors_are_rejected(): + """Each rejection at ITS OWN entry point, not just through the builder. + + There are two: ``build_q_time_pregrid`` and ``JAXLikelihoodData.__init__``. + Going through ``build_likelihood_data`` exercises neither in isolation -- + it calls the first, so the second is unreachable that way, and the second + would have caught a hole in the first. Mutation-testing showed BOTH + survived a test written that way: two redundant guards, each masking the + other, and the pair reads as coverage. ``build_q_time_pregrid`` is also + public (`banded` and offline analysis call it directly), so its own + rejection is not a formality. + """ + packed, tvals, deltaT, tref = _toy_packed() + rho = packed["H1"]["rholmArray"] + # ``match=`` IS THE ASSERTION. A bare ``pytest.raises(ValueError)`` cannot see + # this guard at all: delete it and #261's own "Q pregrid factor must be + # positive" raises ValueError one frame down, so the test passes on a + # different rejection. Measured 2026-09-07 by mutation -- the bare form + # survives deleting BOTH guards, because the value is refused deeper still. + # Defense in depth is fine; a test that cannot tell which layer refused is not. + for bad in (0, -3): + with pytest.raises(ValueError, match="q_time_pregrid_factor must be"): + C.build_q_time_pregrid(rho, bad) + with pytest.raises(ValueError, match="q_time_pregrid_factor must be"): + C.JAXLikelihoodData({}, deltaT, 0.0, tvals, tref, + q_time_pregrid_factor=bad) + with pytest.raises(ValueError, match="q_time_pregrid_factor must be"): + build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=bad) + + +# -------------------------------------------------------------------------- +# 5. End to end: the whole likelihood, and the gradient it exists to provide. +# -------------------------------------------------------------------------- + +def test_whole_likelihood_moves_toward_the_refined_answer(): + """lnL at factor 1 vs 8 vs 32: factor 8 must sit far closer to the converged value. + + The pregrid is a numerical-accuracy change, so "runs without error" is not the + assertion. ``factor 32`` stands in for the converged interpolant here (the + per-position convergence rate is pinned above); the claim is that the shipped + factor 8 removes most of the gap that the production stencil leaves. + """ + packed, tvals, deltaT, tref = _toy_packed() + rng = np.random.default_rng(12) + S = 24 + args = (rng.uniform(0, 2 * np.pi, S), rng.uniform(-1.2, 1.2, S), + rng.uniform(0, np.pi, S), rng.uniform(0, np.pi, S), + rng.uniform(0, 2 * np.pi, S), np.full(S, 400.0)) + vals = {} + for factor in (1, 8, 32): + data = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=factor) + vals[factor] = np.asarray(fused_log_likelihood(data, *args, + interp="sinc" if factor == 1 + else "cubic")) + gap1 = np.max(np.abs(vals[1] - vals[32])) + gap8 = np.max(np.abs(vals[8] - vals[32])) + print(" max|lnL(a=8 coarse) - lnL(f=32)| = %.4e" % gap1) + print(" max|lnL(f=8) - lnL(f=32)| = %.4e" % gap8) + assert gap8 < gap1 / 10.0 + + +def test_gradient_still_flows_through_the_refined_gather(): + """The whole reason this arm exists is ``jax.grad``; a pregrid must not break it. + + A gather that lost its dependence on ``pos`` -- e.g. by rounding the scaled + position to the refined grid -- would still return sensible values and a + silently ZERO sky gradient. + """ + packed, tvals, deltaT, tref = _toy_packed() + data = build_likelihood_data(packed, deltaT, tref, tvals, + q_time_pregrid_factor=8) + + def f(ra): + return fused_log_likelihood( + data, ra, jnp.asarray([0.3]), jnp.asarray([0.5]), + jnp.asarray([1.05]), jnp.asarray([0.7]), jnp.asarray([400.0]), + interp="cubic")[0] + + g = jax.grad(f)(jnp.asarray([1.2])) + assert np.all(np.isfinite(np.asarray(g))) + assert np.max(np.abs(np.asarray(g))) > 0.0 + + +# -------------------------------------------------------------------------- +# 6. The SEAMS. Everything above tests the library; a flag that never reaches +# it is still a no-op, and a --help that parses is not a wired option. +# -------------------------------------------------------------------------- + +def test_wrapper_forwards_the_factor_to_the_data(): + """``build_data_from_precompute`` must carry the factor into the built data. + + Exercised through the REAL function with the two expensive production calls + stubbed, rather than by reading the source: a keyword that is accepted, + documented and then dropped on the floor is exactly the silent-no-op shape + this option could most easily take. + """ + from RIFT.likelihood.jax_ile import wrapper as W + + packed, tvals, deltaT, tref = _toy_packed(detectors=("H1", "L1")) + + class _P: + deltaT = None + + P = _P() + P.deltaT = deltaT + dets = list(packed) + + def _fake_precompute(*a, **k): + empty = {d: {} for d in dets} + return empty, empty, empty, {d: {} for d in dets}, 1.0, None + + def _fake_pack(keys, intp, rho, ct, ctV, _packed=packed, _dets=iter(dets)): + det = next(_dets) + d = _packed[det] + return (d["lms"], None, None, d["U"], d["V"], d["rholmArray"], None, + d["epoch"]) + + old_pre = W.factored_likelihood.PrecomputeLikelihoodTerms + old_pack = W.factored_likelihood.PackLikelihoodDataStructuresAsArrays + try: + W.factored_likelihood.PrecomputeLikelihoodTerms = _fake_precompute + W.factored_likelihood.PackLikelihoodDataStructuresAsArrays = _fake_pack + data, _extras = W.build_data_from_precompute( + P, {d: None for d in dets}, {d: None for d in dets}, 1126259462.0, + 0.15, 0.075, 2, 1700.0, tvals=tvals, q_time_pregrid_factor=8) + finally: + W.factored_likelihood.PrecomputeLikelihoodTerms = old_pre + W.factored_likelihood.PackLikelihoodDataStructuresAsArrays = old_pack + + n_coarse = packed[dets[0]]["rholmArray"].shape[-1] + assert data.q_time_pregrid_factor == 8 + for det in dets: + dd = data.detectors[det] + assert dd["q_time_pregrid_factor"] == 8 + assert dd["npts_full_coarse"] == n_coarse + assert dd["Q"].shape[0] == (n_coarse - 1) * 8 + 1 + + +def test_driver_passes_the_factor_at_its_call_site(): + """The driver's own ``build_data_from_precompute`` call must name the option. + + Parsed with ``ast`` rather than grepped, so a mention in a comment, a help + string or a dead branch does not satisfy it. This is the one seam a library + test cannot reach: the driver is a script with no ``.py`` extension and its + ``analyze_one`` needs real frames to run. + """ + import ast + + here = os.path.dirname(os.path.abspath(__file__)) + driver = os.path.join(here, "..", "..", "bin", + "integrate_likelihood_extrinsic_jax") + tree = ast.parse(open(driver).read()) + sites = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + f = node.func + name = f.attr if isinstance(f, ast.Attribute) else getattr(f, "id", None) + if name != "build_data_from_precompute": + continue + sites.append({k.arg for k in node.keywords if k.arg}) + assert sites, "driver no longer calls build_data_from_precompute" + for kwargs in sites: + assert "q_time_pregrid_factor" in kwargs, ( + "a build_data_from_precompute call site does not forward " + "--q-time-pregrid-factor; the flag would parse and do nothing") + + +def test_pregrid_and_the_phase_marg_mode_permutation_compose(): + """Both packings of the (2,+-2) pair must agree ON A REFINED GRID. + + This path is born at the merge and neither side covers it. #272 made + ``_accumulate_unit`` permute ``lms``, ``Q``, ``U`` and ``V`` to canonical + order under phase marginalization; its fixtures never set + ``q_time_pregrid_factor``. This file exercises the pregrid; its fixtures + never pack the pair the other way round. The permutation takes ``Q`` on + axis 1 while every pregrid index acts on axis 0, so the two are expected to + be independent -- but "expected to be independent" is the claim, and the + merge is where it first has to hold. + + The last assertion is what stops this being vacuous. Permuting a mode axis + would agree at every factor even if the pregrid were doing nothing at all, + so the refined answer must first be shown to DIFFER from the coarse one. + """ + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + import test_jax_phase_marg_mode_order as M + + pk = M._packed() + swapped = M._relabel(pk, [1, 0]) + th = M._angles() + + def acc(packed, factor): + tw = 32 * (1.0 / 1024) / 2.0 + data = build_likelihood_data(packed, 1.0 / 1024, M.TREF, + np.linspace(-tw, tw, 32), + q_time_pregrid_factor=factor) + k, r = C._accumulate_unit(data, *th, "cubic", True, guard=0) + return np.asarray(k), np.asarray(r) + + for factor in (1, 2, 8): + k0, r0 = acc(pk, factor) + k1, r1 = acc(swapped, factor) + scale = max(np.abs(k0).max(), 1.0) + assert np.abs(k0 - k1).max() <= 1e-12 * scale, ( + "packing order changes kappa at q_time_pregrid_factor=%d" % factor) + assert np.abs(r0 - r1).max() <= 1e-12 * max(np.abs(r0).max(), 1.0), ( + "packing order changes rho^2 at q_time_pregrid_factor=%d" % factor) + + k1c, _ = acc(pk, 1) + k8c, _ = acc(pk, 8) + assert np.abs(k1c - k8c).max() / max(np.abs(k1c).max(), 1.0) > 1e-9, ( + "factor 8 reproduces factor 1 on this fixture, so the agreement above " + "says nothing about the refined grid") + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-s", "-q"])) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py index fcc6d379b..a9e67bfad 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py @@ -38,11 +38,13 @@ import RIFT.likelihood.factored_likelihood as fl import RIFT.likelihood.factored_likelihood_with_rotation as flwr import RIFT.likelihood.factored_likelihood_freqresponse as flfr +import RIFT.likelihood.factored_likelihood_rotating_freqresponse as flrr import RIFT.likelihood.slowrot_freqresponse as sfr from RIFT.likelihood.jax_ile.core import fused_log_likelihood from RIFT.likelihood.jax_ile.banded import (build_rotation_data, - build_freqresponse_data) + build_freqresponse_data, + build_rotating_freqresponse_data) from RIFT.likelihood.jax_ile.wrapper import JAXDistanceMarginalizedLikelihood if not getattr(fl, "numba_on", True): @@ -183,6 +185,42 @@ def test_freqresponse(): check_ad(check_freqresponse(), "freqresponse") +def check_rotating_freqresponse(): + """Manual real-precompute/JIT gate; deliberately excluded from per-PR CI.""" + qmax = 0 + bk = flrr.PrecomputeLikelihoodTermsRotatingFreqResponse( + event_time, t_window, Psig, data_dict, psd_dict, Lmax, fmax, + Qmax=qmax, L_arm=L_CE, p_max=0, analyticPSD_Q=True, + verbose=False, quiet=True, skip_interpolation=True) + meta = bk[4] + lk, rba, uba, vba, ep = flrr.pack_rotating_freqresponse_arrays( + meta, bk[3], bk[1], bk[2]) + Pv = _P_vec(K=16) + want = flrr.DiscreteFactoredLogLikelihoodRotatingFreqResponseNoLoop( + TVALS, Pv, meta, lk, rba, uba, vba, ep, Lmax=Lmax, + time_interp="nearest", xpy=np) + det_geom = {d: sfr.detector_geometry(d, L_arm=L_CE) for d in DETS} + data = build_rotating_freqresponse_data( + meta, lk, rba, uba, vba, ep, deltaT, TVALS, det_geom) + got = np.asarray(fused_log_likelihood( + data, Pv.phi, Pv.theta, Pv.psi, Pv.incl, Pv.phiref, _distMpc(Pv), + interp="nearest")) + finite = np.isfinite(want) & np.isfinite(got) + rel = np.max(np.abs(want[finite] - got[finite]) / (1 + np.abs(want[finite]))) + print("[combined] nearest vs numpy: max|rel|=%.3e A=%d" % + (rel, len(meta["a_list"]))) + assert rel < 1e-10 + point = tuple(jnp.asarray([x]) for x in + (1.0, 0.2, 0.4, 0.9, 1.1, 300.0)) + fixed = jax.jit(lambda *x: fused_log_likelihood( + data, *x, interp="linear")) + assert np.all(np.isfinite(np.asarray(fixed(*point)))) + scalar = lambda x: fused_log_likelihood( + data, x[None], point[1], point[2], point[3], point[4], point[5], + interp="linear")[0] + assert np.isfinite(float(jax.grad(scalar)(jnp.asarray(1.0)))) + + def check_ad(data, tag): print("--- AD checks (%s) ---" % tag) # (c) jit + vmap of the fixed-distance likelihood @@ -218,6 +256,7 @@ def check_ad(data, tag): test_rotation_path_a() test_rotation_path_b() test_freqresponse() + check_rotating_freqresponse() print("\nSLOWROT + FREQRESPONSE JAX VALIDATION PASSED") print(" (agreement with the NoLoop is necessary, not sufficient: the rotation VALUE is") print(" pinned by test/jax/test_jax_slowrot_cauchy_schwarz.py.)") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py index c2512e6cd..cb9195aea 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py @@ -22,9 +22,11 @@ import RIFT.likelihood.factored_likelihood_with_rotation as flwr import RIFT.likelihood.factored_likelihood_freqresponse as ffr +import RIFT.likelihood.factored_likelihood_rotating_freqresponse as frr import RIFT.likelihood.slowrot_freqresponse as sfr from RIFT.likelihood.jax_ile import response_slowrot as rs from RIFT.likelihood.jax_ile import response_freqresponse as rf +from RIFT.likelihood.jax_ile import response_rotating_freqresponse as rrf DETS = ["H1", "L1", "V1"] TREF = 1126259462.0 @@ -87,6 +89,34 @@ def test_freqresponse_coefficients(): worst = max(worst, d) print("[freqresponse coeff] max|jax-np| over dets/L = %.3e" % worst) assert worst < 1e-11, "freqresponse coefficient mismatch %g" % worst + # Keep compound CI coverage inside this already-collected algebraic test. It + # adds no waveform precompute, likelihood compilation, or new test shard. + check_rotating_freqresponse_coefficients() + + +def check_rotating_freqresponse_coefficients(): + rng = np.random.default_rng(11) + S = 24 + RA = rng.uniform(0, 2 * np.pi, S) + DEC = np.arcsin(rng.uniform(-1, 1, S)) + psi = rng.uniform(0, np.pi, S) + gmst = _gmst(TREF) + worst = 0.0 + for Qmax, p_max in ((0, 0), (2, 1)): + for det in DETS: + response, x_arm, y_arm, _ = sfr.detector_geometry(det, L_arm=40000.0) + lald = lalsim.DetectorPrefixToLALDetector(det) + got = rrf.coefficients_dict( + response, lald.location, x_arm, y_arm, RA, DEC, psi, + gmst, Qmax, p_max) + want = frr.combined_response_coefficients_vector( + det, RA, DEC, psi, TREF, p_max, Qmax=Qmax, L_arm=40000.0) + for key in set(got) | set(want): + err = np.max(np.abs(np.asarray(got.get(key, np.zeros(S))) + - np.asarray(want.get(key, np.zeros(S))))) + worst = max(worst, err) + print("[combined coeff] max|jax-np| = %.3e" % worst) + assert worst < 1e-11, "compound coefficient mismatch %g" % worst if __name__ == "__main__": diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py index 3c9f9dad2..d5c1cab59 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py @@ -21,6 +21,7 @@ from RIFT.likelihood.jax_ile import ( build_rotation_data_from_precompute, build_freqresponse_data_from_precompute, + build_rotating_freqresponse_data_from_precompute, ) from RIFT.likelihood.jax_ile.wrapper import JAXDistanceMarginalizedLikelihood from RIFT.likelihood.jax_ile.core import fused_log_likelihood @@ -54,18 +55,20 @@ distMpc = rng.uniform(100, 800, S) -def _run(builder, tag, **kw): +def _run(builder, tag, check_ad=True, **kw): data, extras = builder(P.manual_copy(), data_dict, psd_dict, event_time, IWH, Lmax, fmax, analyticPSD_Q=True, verbose=False, **kw) lnL = np.asarray(fused_log_likelihood(data, ra, dec, psi, incl, phiref, distMpc, interp="nearest")) assert np.all(np.isfinite(lnL)), "%s produced non-finite lnL" % tag - # differentiable distmarg path - dlike = JAXDistanceMarginalizedLikelihood(data, 5.0, 3000.0, n_grid=64) - v, g = dlike.value_and_grad([ra[0], dec[0], psi[0], incl[0], phiref[0]]) - assert np.isfinite(v) and np.all(np.isfinite(g)), "%s distmarg AD non-finite" % tag - print("[%s] one-call build OK: lnL[0]=%.3f distmarg lnL=%.3f |grad|=%.2f" - % (tag, lnL[0], v, np.linalg.norm(g))) + if check_ad: + dlike = JAXDistanceMarginalizedLikelihood(data, 5.0, 3000.0, n_grid=64) + v, g = dlike.value_and_grad([ra[0], dec[0], psi[0], incl[0], phiref[0]]) + assert np.isfinite(v) and np.all(np.isfinite(g)), "%s distmarg AD non-finite" % tag + print("[%s] one-call build OK: lnL[0]=%.3f distmarg lnL=%.3f |grad|=%.2f" + % (tag, lnL[0], v, np.linalg.norm(g))) + else: + print("[%s] one-call build OK: lnL[0]=%.3f" % (tag, lnL[0])) return data @@ -77,6 +80,16 @@ def test_one_call_builders(): L_arm=40000.0) +def check_combined_one_call_builder(): + """Manual production-wrapper gate; deliberately excluded from per-PR CI.""" + # The individual features above own the expensive distance-AD wrapper gate. + # Here the compound path needs to prove production precompute selection and a + # finite fixed-distance contraction without adding minutes of duplicate CPU CI. + _run(build_rotating_freqresponse_data_from_precompute, "combined", + check_ad=False, Qmax=0, p_max=0, L_arm=40000.0) + + if __name__ == "__main__": test_one_call_builders() + check_combined_one_call_builder() print("ONE-CALL BUILDER SMOKE TEST PASSED") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py index da83d882f..d631b593c 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_stencil_parity.py @@ -368,16 +368,50 @@ def test_accumulators_pass_separable_u(): # ... 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) + # + # RETARGETED 2026-09-07 (Q time pregrid). The construction used to be inlined in each + # accumulator, and this counted `u_sep = ...` assignments, expecting one per accumulator. + # It is now hoisted into ``_q_sample_positions``, which both accumulators call, so that + # count stopped describing the code and the guard failed on a refactor it should not have + # objected to. The invariant is unchanged and is checked against the structure that now + # carries it -- and more of it, since the accumulators must take the offset from the shared + # owner rather than fabricating one. + owner = [n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == "_q_sample_positions"] + assert len(owner) == 1, "expected exactly one _q_sample_positions, found %d" % len(owner) + owner = owner[0] + + built = [n for n in ast.walk(owner) if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) and n.func.id == "_separable_u"] + assert len(built) >= 2, ( + "_q_sample_positions builds the separable offset on %d of its paths; every path that " + "returns one must build it, or the memory fix is disabled on that path" % len(built)) + + gated = [n for n in ast.walk(owner) if isinstance(n, ast.IfExp) + and "_separable_u" in ast.unparse(n)] + assert gated, ("the separable offset is unconditional; it must stay gated off the stencils " + "that ignore u -- feeding it to 'nearest' cost >60% wall on the banded path") + + for r in [n for n in ast.walk(owner) if isinstance(n, ast.Return) and n.value is not None]: + if isinstance(r.value, ast.Tuple) and len(r.value.elts) == 2: + second = r.value.elts[1] + assert not (isinstance(second, ast.Constant) and second.value is None), \ + "_q_sample_positions returns a hard-coded None offset: %s" % ast.unparse(r.value) + + for fname in ("_accumulate_unit", "_accumulate_unit_banded"): + fn = [n for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == fname] + assert len(fn) == 1, "expected exactly one %s, found %d" % (fname, len(fn)) + binds = [n for n in ast.walk(fn[0]) if isinstance(n, ast.Assign) + and any(isinstance(t, ast.Tuple) + and any(isinstance(e, ast.Name) and e.id == "u_sep" for e in t.elts) + for t in n.targets)] + assert binds, "%s no longer binds u_sep from _q_sample_positions" % fname + for b in binds: + expr = ast.unparse(b.value) + assert "_q_sample_positions" in expr, ( + "%s builds its own offset (%s) instead of taking the shared one; the two " + "accumulators would then be free to drift" % (fname, expr)) def test_every_entry_point_defaults_to_the_same_stencil(): diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py index faa5d4ab8..46aad3354 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_terminal_time_marginalization.py @@ -303,6 +303,65 @@ def test_jax_dropin_manifest_covers_every_batchmode_option_with_same_arity(): assert not missing assert not mismatched + # The conventional factor-8 path is not an inert tuning flag: it changes how Q + # is represented and evaluated. It is now IMPLEMENTED on this arm (the stored + # rholm buffers are refined once, at build time), so the driver must accept it. + # + # CHANGED from "refuses factor 8": before the JAX Q pregrid landed, this asserted + # the refusal. Leaving that assertion in place would have been a test that + # FORBIDS the fix -- the acceptance below is the point of the change, and what + # survives is the narrower, still-true refusal. + # + # record_supplied_options is called for the same reason main() calls it: without + # it was_supplied() reports False for everything, so the "explicit stencil" + # branch below would never be exercised and this would be a test that cannot + # fail. Asserting the promotion (opts.interp becomes 'cubic') as well as the + # absence of an exit is what makes the acceptance leg load-bearing. + def _check(argv): + opts, _ = parser.parse_args(list(argv)) + drv.record_supplied_options(opts, list(argv), parser) + drv.resolve_ile_interface_aliases(opts, parser) + drv.check_critical_and_report(opts, parser) + return opts + + for factor in ("1", "2", "8"): + opts = _check(["--q-time-pregrid-factor", factor]) + assert opts.interp == ("cubic" if factor != "1" + else drv.JAX_INTERP_DEFAULT), factor + # Asking for the pregrid AND cubic explicitly is the documented combination. + opts = _check(["--q-time-pregrid-factor", "8", "--interp", "cubic"]) + assert opts.interp == "cubic" + # Any OTHER explicit stencil is refused rather than silently replaced -- both + # spellings, because --interpolate-time is an alias that rewrites opts.interp + # before this check runs and would otherwise look like "not supplied". + for argv in (["--q-time-pregrid-factor", "8", "--interp", "nearest"], + ["--q-time-pregrid-factor", "8", "--interp", "sinc"], + ["--q-time-pregrid-factor", "8", "--interp", "linear"], + ["--q-time-pregrid-factor", "8", "--interpolate-time", "sinc"]): + with pytest.raises(SystemExit): + _check(argv) + # ... and every one of those stencils is still perfectly legal at the default + # factor, so the refusal is about the COMBINATION and not about the stencil. + for stencil in ("nearest", "linear", "cubic", "sinc"): + opts = _check(["--interp", stencil]) + assert opts.interp == stencil + # 0 is FALSY: a `getattr(...) or 1` idiom would promote it to the default and + # report nothing. + with pytest.raises(SystemExit): + _check(["--q-time-pregrid-factor", "0"]) + + # WITHOUT a supplied-option record, which is how every caller that builds an + # options object directly reaches this code. was_supplied() FAILS OPEN there + # ("no record -> assume not supplied"), so relying on it alone would silently + # replace a stencil the caller chose; the check also treats a non-default + # interp as explicit. Every case above records, so without this one that + # clause is unexercised -- found by mutation-testing the guard, not by review. + argv = ["--q-time-pregrid-factor", "8", "--interp", "linear"] + opts, _ = parser.parse_args(argv) + assert not hasattr(opts, "_supplied_options") + with pytest.raises(SystemExit): + drv.check_critical_and_report(opts, parser) + def _load_driver(): import importlib.machinery diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py index 27109ab95..4468c99cf 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_time_quadrature.py @@ -364,6 +364,7 @@ def __init__(self, feature): self.feature = feature assert _norm_is_arrival_time_dependent(_StubData("rotation")) + assert _norm_is_arrival_time_dependent(_StubData("rotation_freqresponse")) assert not _norm_is_arrival_time_dependent(_StubData("freqresponse")) assert not _norm_is_arrival_time_dependent(_StubData(None)) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py index be770644b..e6a0f3550 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py @@ -9,6 +9,30 @@ from RIFT.likelihood.jax_ile import joint_anglemarg_peaklocal as JP +@pytest.fixture(autouse=True) +def _drop_jax_caches(): + """Release JAX's compiled-executable cache after every test in this file. + + THE PEAK HERE IS RETENTION, NOT ALLOCATION, and that distinction is the whole fix. + Measured: the largest single test peaks at 1408 MB against a 590 MB import baseline, + so no test costs more than ~800 MB -- yet the file as a whole peaked at 4268 MB. The + gap is JAX holding a compiled executable per distinct shape, and this file + deliberately sweeps many combinations of (n_slots, n_nodes, u_nodes, n_bound) because + the properties under test are ABOUT those parameters. Nothing is freed between tests. + + Trimming individual fixtures therefore does almost nothing: six reductions across the + heaviest tests moved the file peak from 4318 MB to 4268 MB, about one percent, and one + of them silently broke an acceptance assertion by starving u_sizing_ok. Dropping the + caches attacks the actual mechanism. + + The cost is recompilation, which is why this is scoped to this file rather than to the + session: it is the one that sweeps shapes. + """ + yield + if hasattr(jax, "clear_caches"): + jax.clear_caches() + + 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 @@ -26,7 +50,7 @@ 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) + u = np.linspace(0.0, 2 * np.pi, 60000, endpoint=False) for _ in range(4): c1 = scale * (rng.normal() + 1j * rng.normal()) c2 = scale * (rng.normal() + 1j * rng.normal()) @@ -42,7 +66,7 @@ def test_both_signs_of_q_enter_the_u_coefficients(): 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) + u = np.linspace(0.0, 2 * np.pi, 60000, 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)) @@ -71,7 +95,7 @@ def test_spurious_off_circle_roots_do_not_orphan_an_arc(): 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) + u = np.linspace(0.0, 2 * np.pi, 60000, 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) @@ -145,3 +169,763 @@ def test_gradient_is_finite_as_the_quartic_leading_coefficient_vanishes(): 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 + + +def test_required_u_nodes_is_derived_and_grows_like_sqrt_amplitude(): + """P1 from review: the fallback (whole-cell) branch integrates with the SAME fixed + node count spread over the entire cell, so rejecting a stalled Newton centre makes + the resolution worse rather than safer. JAX cannot adapt the count -- shapes may not + depend on traced values -- so the sizing is exposed as a caller-side helper, derived + from the exact bound |d2g/du2| <= M2u ~ 5A. + + Production uses this count because fallback is data-dependent. It is intentionally + uncapped: memory is bounded by streaming the node axis, not by truncating an accuracy + request inside a region the omitted-mass certificate cannot inspect. + """ + lo = JP.required_u_nodes(1.0) + mid = JP.required_u_nodes(100.0) + hi = JP.required_u_nodes(1.0e4) + assert lo == JP.U_NODES_PER_CELL # never below the windowed default + assert lo < mid < hi # grows with amplitude + assert JP.u_nodes_in_use(450.0) == JP.required_u_nodes(450.0) + assert hi > 2048 # production does not silently cap accuracy + # the growth is the sqrt law, not something steeper + assert 5.0 < mid / np.sqrt(100.0) < 60.0, mid + + +def test_a_fallback_cell_is_resolved_when_the_caller_sizes_it(): + """The helper must actually buy resolution: a whole-cell integration at a raised node + count must agree with a much finer one.""" + rng = np.random.default_rng(0) + worst = 0.0 + for _ in range(6): + sc = 10.0 ** rng.uniform(0.5, 2.0) + c1 = sc * (rng.normal() + 1j * rng.normal()) + c2 = sc * (rng.normal() + 1j * rng.normal()) + amp = abs(c1) + 2 * abs(c2) + n = JP.required_u_nodes(amp) + a = float(JP.log_inner_u_integral(0.0, c1, c2, n_nodes=n)) + b = float(JP.log_inner_u_integral(0.0, c1, c2, n_nodes=min(4 * n, 4096))) + worst = max(worst, abs(a - b)) + assert worst < 1e-4, worst + + +def test_large_fallback_policy_streams_a_fixed_live_node_block(): + """The accurate production count must not reappear as a materialized node axis. + + At the sizing floor the policy requests hundreds of nodes. Observe the shape handed + to the exponent evaluator while tracing the rolled loop: its live last axis must stay + at the stream chunk, independent of the total quadrature count. + """ + n = JP.u_nodes_in_use(450.0) + assert n > JP.U_NODE_STREAM_CHUNK + shapes = [] + real_g = JP._g_u + + def _spy_g(a, c1, c2, u, order=0): + if order == 0 and getattr(u, "ndim", 0) == 2: + shapes.append(tuple(u.shape)) + return real_g(a, c1, c2, u, order) + + JP._g_u = _spy_g + try: + out = JP.log_inner_u_integral(0.0, 2.0 + 1j, 0.7 - 0.3j, n_nodes=n) + assert np.isfinite(float(out)) + finally: + JP._g_u = real_g + + assert shapes, "stream body never reached the exponent evaluator" + assert max(shape[-1] for shape in shapes) <= JP.U_NODE_STREAM_CHUNK, shapes + +# --------------------------------------------- phi localization (both axes local) + +def _tables_scaled(seed, scale): + 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 _joint(A, B, x=1.0): + from RIFT.likelihood import joint_angle_peak_local as JN + return JN.joint_table(A, B, x=x) + + +def _torus_ref(C, n=2048): + from RIFT.likelihood import joint_angle_peak_local as JN + t = np.linspace(0.0, 2 * np.pi, n, endpoint=False) + P, U = np.meshgrid(t, t, indexing='ij') + g = JN.eval_g(C, P.ravel(), U.ravel()) + m = g.max() + return m + np.log(np.exp(g - m).mean()) + 2 * np.log(2 * np.pi) + + +def test_u_profile_derivatives_match_the_numpy_reference(): + """F' and F'' come from differentiating under the integral, so they are exact and + cost no extra evaluation. Two independent implementations must agree.""" + from RIFT.likelihood import joint_angle_peak_local as JN + A, B = _tables_scaled(3, 3.0) + C = _joint(A, B) + f = jax.jit(JP.u_profile) + for phi in np.linspace(0.4, 5.6, 5): + F, d1, d2, _, _, _ = f(jnp.asarray(C), float(phi)) + Fn, d1n, d2n = JN.u_profile(C, np.array([phi])) + assert abs(float(F) - Fn[0]) < 1e-4, (phi, F, Fn[0]) + scale = max(1.0, abs(d1n[0])) + assert abs(float(d1) - d1n[0]) < 1e-3 * scale, (phi, d1, d1n[0]) + + +@pytest.mark.parametrize("scale", [1.0, 10.0, 100.0]) +def test_phi_local_matches_a_dense_torus_reference(scale): + A, B = _tables_scaled(3, 1.0) + C = _joint(A * scale, B * scale) + got, ok, info = jax.jit(JP.phi_local_lnI)(jnp.asarray(C)) + assert abs(float(got) - _torus_ref(C)) < 1e-4, (scale, float(got)) + + +def test_empty_merge_slots_do_not_poison_the_sum_with_nan(): + """Regression. There are always more slots than groups, and an empty slot comes + back from the segment reductions as (+inf, -inf). Masking its WEIGHT is not enough: + the node positions are still built from it, jnp.mod(inf, 2pi) is NaN, and NaN * 0 is + NaN -- so the poison reached the sum through a term that was supposed to be switched + off. Every amplitude above ~400 returned NaN before the position was neutralized.""" + for scale in (10.0, 30.0, 100.0, 300.0): + A, B = _tables_scaled(3, 1.0) + got, _ok, _info = jax.jit(JP.phi_local_lnI)(jnp.asarray(_joint(A * scale, B * scale))) + assert np.isfinite(float(got)), (scale, float(got)) + + +def test_phi_local_cost_is_flat_in_amplitude(): + """The point of localizing BOTH axes. Measured wall time is ~0.19 s at every + amplitude from 42 to 12650; here we assert the structural property that makes that + true -- the work is set by static shapes, so the SAME jitted callable serves every + amplitude without recompiling.""" + f = jax.jit(JP.phi_local_lnI) + A, B = _tables_scaled(3, 1.0) + shapes = set() + for scale in (1.0, 10.0, 100.0): + C = jnp.asarray(_joint(A * scale, B * scale)) + shapes.add(C.shape) + assert np.isfinite(float(f(C)[0])) + assert len(shapes) == 1, shapes # one shape => one compilation + + +def test_u_profile_rejects_a_clipped_newton_point_as_a_peak(): + """External-review P1 on the phi-localization branch. ``u_profile`` classified a cell + as peaked from ``g'' < 0`` ALONE -- the same defect ``log_inner_u_integral`` already + gates, reintroduced because this function was written as a fresh copy of that Newton + iteration rather than as a call to it. The iteration is clamped to ``[lo_c, mid]``, so + it can come to rest ON a boundary carrying a large stationary residual; curvature then + centres a +-window on a non-stationary point and can EXCLUDE the true maximum, which + underestimates ``F`` while the docstring calls its derivatives exact. + + Non-vacuity is the point of this test: measured over 200 random coefficient draws, the + gate rejects 7.3% of the cells the curvature-only test accepted, the worst at + ``|g_u|/M_1 = 0.512``. A gate that rejected nothing would pass this file's other tests + just as happily. + """ + from jax import lax + rng = np.random.default_rng(3) + total = rejected = 0 + worst = 0.0 + for _ in range(120): + sc = 10.0 ** rng.uniform(0.5, 2.0) + c1 = complex(sc * rng.normal(), sc * rng.normal()) + c2 = complex(sc * rng.normal(), sc * rng.normal()) + u = jnp.sort(JP.u_stationary_roots(c1, c2)) + 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) + + def _step(uc, _): + g1 = JP._g_u(0.0, c1, c2, uc, 1) + g2 = JP._g_u(0.0, c1, c2, uc, 2) + st = jnp.where(jnp.abs(g2) > 0, -g1 / jnp.where(jnp.abs(g2) > 0, g2, 1.0), 0.0) + return jnp.clip(uc + jnp.clip(st, -0.5, 0.5), lo_c, mid), None + + ustar, _ = lax.scan(_step, u, None, length=8) + g1s = JP._g_u(0.0, c1, c2, ustar, 1) + g2s = JP._g_u(0.0, c1, c2, ustar, 2) + m1u = abs(c1) + 2.0 * abs(c2) + edge = 1e-9 * float(jnp.max(mid - lo_c)) + curvature_only = np.asarray(g2s < 0.0) + gated = np.asarray((g2s < 0.0) + & (jnp.abs(g1s) <= 1e-8 * max(m1u, 1e-300)) + & (ustar > lo_c + edge) & (ustar < mid - edge)) + assert not (gated & ~curvature_only).any(), "gate must only ever REMOVE cells" + dropped = curvature_only & ~gated + total += int(curvature_only.sum()) + rejected += int(dropped.sum()) + if dropped.any(): + r = np.asarray(jnp.abs(g1s)) / max(m1u, 1e-300) + worst = max(worst, float(r[dropped].max())) + assert total > 0 + assert rejected > 0, "gate rejected nothing -- it is decoration, not a check" + assert worst > 1e-3, "worst rejected residual %.3g is within tolerance of stationary" % worst + + +def test_phi_local_returns_a_certificate_that_actually_declines(): + """External-review P1: ``phi_local_lnI`` returned a bare float -- no bound, no validity + result, no fallback signal -- while its docstring claimed correctness rested on "the + caller's cover bound", a contract no caller implemented. Fixed seeds are targeting, + not an enumeration, so a missed maximum came back as a finite likelihood. + + It now returns ``(value, ok, info)`` with an omitted-mass bound on the phi axis: + ``area_outside * exp(sup_outside F)``, the supremum obtained by LIFTING grid values of + ``F`` with a true remainder from ``profile_derivative_bounds`` -- never the grid + maximum, which is a lower bound on a supremum. + + The assertion that matters is that it DECLINES: a certificate that always accepts is + decoration, and would have passed every other test in this file. + """ + rng = np.random.default_rng(0) + verdicts = [] + for scale in (0.3, 3.0, 40.0, 200.0): + C = (rng.normal(size=(3, 5)) + 1j * rng.normal(size=(3, 5))) * scale + val, ok, info = JP.phi_local_lnI(jnp.asarray(C)) + assert np.isfinite(float(val)) + for key in ("margin", "area_outside", "sup_outside", "n_phi_regions", + "n_u_fallback", "n_u_risky_quad"): + assert key in info, key + # THE CONTRACT CHANGED AND THIS TEST USED TO PIN THE DEFECT. It asserted that ok + # was exactly the margin test and that a full cover MUST be accepted -- which is + # precisely the conflation test_a_full_cover_no_longer_accepts_unconditionally + # exists to remove. Both assertions passed only because this test's four fixtures + # all happen to converge; adversarial review found them contradicting each other + # across files. ok is the margin test AND the resolution test AND the u sizing + # test -- three independent ways to be wrong, and the contract is their conjunction. + # + # u_sizing_ok used to be read off the BOUND grid, where it was near-vacuous: that + # grid's job was to bound F from outside, not to produce `value`. It now reads the + # QUADRATURE grid, so it reports whether the integration that produced the returned + # number was adequately sampled -- and it does fire here, on a fixture whose margin + # (-29.3) and resolution both pass at the default 48 u nodes. Omitting it from + # this identity is what made the test fail when the gate moved to the right grid. + assert bool(ok) == (float(info["margin"]) < JP.OUTSIDE_TOL_NATS + and bool(info["phi_resolved"]) + and bool(info["u_sizing_ok"])) + if float(info["area_outside"]) == 0.0: + # nothing omitted, so the margin is -inf; whether that ACCEPTS now depends on + # the integration having converged, which is the whole point of the change. + assert float(info["margin"]) == -np.inf + assert bool(ok) == bool(info["phi_resolved"]) + verdicts.append(bool(ok)) + assert any(verdicts), "certificate declined everything -- it is unusable, not strict" + assert not all(verdicts), "certificate accepted everything -- it is decoration" + + +def test_a_full_cover_no_longer_accepts_unconditionally(): + """The covering path used to conflate two different statements. ``area_outside = 0`` + says nothing was left OUT; it says nothing about the quadrature INSIDE, yet it gave + ``margin = -inf`` and an unconditional accept. Measured before the fix at KP=13, + amplitude 1e2 with algebraic seeds: full cover, accepted, value 0.196 nats wrong -- the + same conflation that cost the numpy reference 0.36 nats on production tables. + + ``ok`` now also requires the integration to have CONVERGED, measured from the nested + grid -- free, because ``PHI_NODES_PER_REGION`` is odd so the even indices are a + trapezoid at half the density and the odd ones are exactly its midpoints. Two gates + were tried first and rejected on evidence: the exact ``M2F`` bound demands 3.8e3-2.3e4 + nodes and declines cases right to 1e-4, and a local-curvature rule declines cases right + to 1e-5, because a periodic trapezoid converges spectrally and any points-per-sigma + rule is far too conservative. + + The full cover is forced with ``w_sigma`` rather than with the algebraic seeder that + used to produce one here. That seeder has been removed -- it duplicated + ``bivariate_trig_stationary`` more weakly -- and the ``wrapped`` branch is the other + way a cover comes to span the circle, so the defect is still reachable. + """ + KS = 2 + rng = np.random.default_rng(101) + C = rng.normal(size=(13, 2 * KS + 1)) + 1j * rng.normal(size=(13, 2 * KS + 1)) + C = jnp.asarray(C * (1e2 / np.sum(np.abs(C)))) + v, ok, info = JP.phi_local_lnI(C, w_sigma=400.0) + assert float(info["area_outside"]) == 0.0 # the cover IS full + assert not bool(info["phi_resolved"]) # but the integration is not converged + assert not bool(ok), "a full cover must not accept an unconverged integration" + assert float(info["phi_convergence"]) > JP.PHI_CONVERGENCE_NATS + + +def test_the_convergence_gate_does_not_decline_accurate_results(): + """A gate that refuses correct answers is as useless as one that accepts wrong ones, and + the two gates tried before this one both did. These cases are accurate to ~1e-5 against + a converged torus reference and MUST still accept.""" + KS = 2 + accepted = 0 + for amp in (4.5, 19.0): + rng = np.random.default_rng(101) + C = rng.normal(size=(3, 2 * KS + 1)) + 1j * rng.normal(size=(3, 2 * KS + 1)) + C = jnp.asarray(C * (amp / np.sum(np.abs(C)))) + v, ok, info = JP.phi_local_lnI(C) + assert abs(float(v) - _torus_ref(np.asarray(C))) < 1e-3, (amp, float(v)) + assert float(info["phi_convergence"]) < JP.PHI_CONVERGENCE_NATS, (amp,) + accepted += bool(ok) + assert accepted == 2, accepted + + +def test_the_convergence_check_is_guarded_against_its_own_blind_spot(): + """Adversarial review F3. ``conv`` halves the nodes and compares -- but the n and n/2 + trapezoids share EVERY aliased harmonic at multiples of n, so it measures the n/2 + aliasing and infers the rest from smoothness. Content at exactly harmonic n is + invisible to it: review built a table with a phi ripple at n and got values 0.83-0.99 + nats wrong with ``conv`` as low as 1.3e-04 -- BELOW the 1e-3 gate, so ``conv`` alone + accepted them. + + The guard tested here is ``n_nodes > 2 k_max``. IT IS NECESSARY AND NOT SUFFICIENT, + and this docstring used to claim otherwise -- that Nyquist-resolving ``k_max`` "rules + out content at the sampling harmonic by construction". That is a statement about + ``g``; the outer trapezoid integrates ``exp(F)`` with ``F = log int du exp(g)``, and + neither is band-limited because ``g`` is. A later review supplied a ``k_max = 1`` + table that passes this guard trivially and is still 0.02 nats wrong -- see + :func:`test_the_halving_check_is_blind_at_the_sampling_harmonic`, which covers the + part of the family this guard does not. + + Tested through ``n_nodes`` rather than by building the degree-1552 counterexample, + which is correct-but-unaffordable in CI. + """ + KS = 2 + rng = np.random.default_rng(101) + C = rng.normal(size=(9, 2 * KS + 1)) + 1j * rng.normal(size=(9, 2 * KS + 1)) + C = jnp.asarray(C * (1e2 / np.sum(np.abs(C)))) + k_max = 8 # KP - 1 + + # under-resolved: the check cannot see harmonic n, so it must not be believed + _, ok_bad, info_bad = JP.phi_local_lnI(C, n_nodes=2 * k_max - 1) + assert not bool(info_bad["phi_alias_safe"]) + assert not bool(ok_bad), "an unresolvable node count must never accept" + + # comfortably resolved: the guard must not be what blocks an otherwise good case + _, _, info_ok = JP.phi_local_lnI(C, n_nodes=JP.PHI_NODES_PER_REGION) + assert bool(info_ok["phi_alias_safe"]), (JP.PHI_NODES_PER_REGION, k_max) + + # and the guard is load-bearing, not decoration: it must be able to veto a case whose + # conv is below the threshold, which is exactly what the counterexample showed. + assert JP.PHI_NODES_PER_REGION > 2 * k_max + + +def _separable_phi_table(kappa, shift, r=6.0, KS=2): + """A table whose profile is EXACTLY ``F(phi) = kappa cos(phi - shift) + const``. + + Only ``C[1, q=0]`` and ``C[0, q=+2]`` are set, so ``c1 = 0`` and ``c2 = r`` are both + phi-independent: the u integral contributes a constant and the phi dependence is the + single harmonic. ``k_max = KP - 1 = 1``, and the double integral is closed form, + ``2 pi I_0(kappa) * 2 pi I_0(r)``, so the error is known rather than estimated. + """ + C = np.zeros((2, 2 * KS + 1), dtype=complex) + C[1, KS + 0] = 0.5 * kappa * np.exp(-1j * shift) + C[0, KS + 2] = r + from scipy.special import ive + exact = (np.log(2 * np.pi) + kappa + np.log(ive(0, kappa)) + + np.log(2 * np.pi) + r + np.log(ive(0, r))) + return jnp.asarray(C), exact + + +def test_the_halving_check_is_blind_at_the_sampling_harmonic(): + """Adversarial review. ``conv`` halves the nodes -- but the n and n/2 periodic rules + alias at multiples of n and n/2, and the second set CONTAINS the first, so the leading + error term cancels out of the difference. No subset of the nodes already evaluated can + ever see it; that is Nyquist, not an implementation shortfall. + + Review's case: ``F = 1000 cos(phi - pi/96)`` on the full circle at 96 intervals. The + phase makes the c_48 alias vanish exactly and leaves c_96, so the 96- and 48-interval + rules agree to 1e-13 while both are 0.02017 nats wrong. ``k_max = 1`` here, so the + ``n_nodes > 2 k_max`` guard reports it safe at 97 > 2 and cannot help. + + THE FIX IS THE NODE COUNT, NOT A SECOND GRID. Because a rule's own aliases are + invisible in its own samples, the probes can only ever certify the COARSE rule, so the + answer has to ride a level finer than the probes. With the nested grid at 193 the + answer IS the fine rule and comes back right, while the probes still fire because the + 97-node rule they measure was bad -- fail-closed, and correct as well. + + Both halves are asserted, including the blind one: at 97 the probes read ~1e-13 on a + 0.02-nat error. That is the measurement the default rests on, and it is a statement + about Nyquist, so it will not stop being true. + """ + C, exact = _separable_phi_table(1000.0, np.pi / 96) + + # w_sigma forces the wrapped branch: one region spanning 2 pi, which is where a + # periodic aliasing family can exist at all. + v, ok, info = JP.phi_local_lnI(C, w_sigma=200.0) + assert int(info["n_phi_regions"]) == 1, int(info["n_phi_regions"]) + assert abs(float(v) - exact) < 1e-4, float(v) - exact # the ANSWER is now right + assert float(info["phi_convergence_shift"]) > JP.PHI_CONVERGENCE_NATS + assert not bool(ok), "the coarse rule was bad; declining is the conservative direction" + + # why the default is 193 and not 97: at 97 BOTH probes are blind to the error, so the + # same table would come back wrong and unflagged. + v9, _, info9 = JP.phi_local_lnI(C, w_sigma=200.0, n_nodes=97) + assert abs(float(v9) - exact) > 1e-2, float(v9) - exact + assert float(info9["phi_convergence"]) < 1e-9 + assert float(info9["phi_convergence_shift"]) < 1e-9 + assert bool(info9["phi_alias_safe"]) # and the k_max guard says "safe" + + # ...and the companion is not merely a decline switch: resolved, the table accepts. + v2, ok2, info2 = JP.phi_local_lnI(C, w_sigma=200.0, n_nodes=769) + assert abs(float(v2) - exact) < 1e-6, float(v2) - exact + assert float(info2["phi_convergence_shift"]) < JP.PHI_CONVERGENCE_NATS + assert bool(ok2), dict(info2) + + +def test_a_full_circuit_phi_window_is_one_region_at_every_peak_location(): + """This is the pin the test above could not be, and the reason it could not is the + finding: at ``w_sigma = 200`` the window spans a full circuit, so the seam split + emits ``[a0, 2 pi]`` and ``[0, a0 + 2 pi - 2 pi]``, adjacent BY CONSTRUCTION -- and + adjacent in floating point only when ``(a0 + 2 pi) - 2 pi`` rounds back to ``a0``. + ONLY THE LOW SIDE BREAKS: a round-trip landing one ulp ABOVE ``a0`` is fine, the + pieces overlap and the merge folds them. What is lost is ``a0``'s low bits -- + ``ulp(a0 + 2 pi)`` is 1.78e-15 whatever ``a0`` is, against ``ulp(a0)`` from 6.9e-18 to + 8.9e-16. In pure float64, no jax involved, 34.4% of ``a0`` uniform on the circle + round low. Then the pieces sit one ulp apart, the merge (exact-touching, no + tolerance, by design) keeps them separate, ``total`` comes out 8.9e-16 under 2 pi and + the ``wrapped`` clamp misses. The rule runs a SEAMED two-region trapezoid where the + periodic one is spectrally accurate: 2.1e-3 nats wrong instead of 2.1e-8. + + So the single fixture in + :func:`test_the_halving_check_is_blind_at_the_sampling_harmonic` pinned the property + BY LUCK. jax 0.9.2 and 0.10.2 put the Newton fixed point two ulp apart on the same + host, same python 3.13, same numpy 2.4.6; 0.10.2 landed on the good side, which is + why CI was green while the local runs failed. + + THE GRID BELOW IS PART OF THE GUARD, not incidental to it. Pre-fix it fails on 10 of + these 64 shifts under jax 0.9.2 (max error 1.85e-02, 10 of them past the 1e-4 + assertion) and on 10 under 0.10.2 -- so it catches the bug in the environment where + CI was green. That is a property of THIS grid, though: an equally natural + golden-ratio sweep of the same length was measured at 0/64 pre-fix. Sweeping is + necessary and not sufficient; the count above is the evidence, and it should be + re-measured rather than assumed if the grid is ever changed. + + ``vmap`` is what makes it affordable: the per-call cost is host-side tracing, ~1.35 s, + flat in the node counts, so 64 separate calls run ~85 s against ~20-30 s batched + (the spread is host load, not the method). + """ + shifts = np.linspace(0.0, 2 * np.pi, 64, endpoint=False) + C = jnp.stack([_separable_phi_table(1000.0, s)[0] for s in shifts]) + exact = _separable_phi_table(1000.0, 0.0)[1] # shift-independent + v, _, info = jax.vmap(lambda c: JP.phi_local_lnI(c, w_sigma=200.0))(C) + + regions = np.asarray(info["n_phi_regions"]) + bad = shifts[regions != 1] + assert bad.size == 0, (regions[regions != 1][:4], bad[:4]) + + # EXACTLY 2 pi, not approximately: one ulp short IS the failure, and a tolerance here + # would pass the very state this test exists to forbid. + total = np.asarray(info["seg_width"]).sum(axis=1) + assert (total == 2 * np.pi).all(), total[total != 2 * np.pi][:4] - 2 * np.pi + + # and the seam costs ACCURACY, which is why the count is worth pinning at all + err = np.abs(np.asarray(v) - exact) + assert err.max() < 1e-4, (float(err.max()), float(shifts[err.argmax()])) + + # EVERYTHING ABOVE CONSTRAINS WHAT HAPPENS ONCE THE FULL-CIRCUIT BRANCH FIRES, AND + # NOTHING CONSTRAINS WHEN. Every table above is a full circuit by construction + # (w_sigma=200, sigma=1/sqrt(1000), so the window is 12.65 rad), so a kernel taking the + # branch for EVERY window satisfies all three assertions. Mutation-tested: + # `full_circuit = peaked` leaves this file's two seam tests green. + # + # WIDTH IS THE WRONG THING TO ASSERT HERE, and a first version of this block asserted + # it and was inert. That mutation preserves the width -- a narrow window stays 0.1265 + # rad wide either way -- and moves only the LOCATION, anchoring every region at 0. The + # invariant that separates them is that a narrow region sits ON its peak: measured, + # 64/64 covered with the fix against 2/64 under the mutation. + narrow = jax.vmap(lambda c: JP.phi_local_lnI(c, w_sigma=2.0))(C)[2] + nlo = np.asarray(narrow["seg_lo"]) + nw = np.asarray(narrow["seg_width"]) + assert (nw.sum(axis=1) < 2 * np.pi).all(), float(nw.sum(axis=1).max()) + peak = (shifts % (2 * np.pi))[:, None] + live = nw > 0 + on_peak = ((nlo - 1e-9 <= peak) & (peak <= nlo + nw + 1e-9)) + wrapped_hit = ((nlo - 1e-9 <= peak + 2 * np.pi) + & (peak + 2 * np.pi <= nlo + nw + 1e-9)) + covered = ((on_peak | wrapped_hit) & live).any(axis=1) + assert covered.all(), shifts[~covered][:4] + + +def test_the_phi_grid_is_nested_so_no_evaluation_is_spent_on_a_probe_alone(): + """The first version of the companion evaluated a SECOND grid of n-1 midpoints, used + only for the probe and then discarded: 1.85x the cost for a diagnostic. With an odd + node count one grid already contains both sub-rules -- even indices are a trapezoid at + half the density, odd indices are exactly its midpoints -- so both probes are free and + the returned value is the fine rule. + + Counted at the GRID level, which is the level that costs: under ``jax.vmap`` the + profile is traced once per grid, so the number of ``u_profile`` invocations is the + number of distinct grids the kernel builds. There are four -- the Newton step, the + seed evaluation, the quadrature grid and the bound grid -- and a separate midpoint + grid would make five. The probes must come out of the quadrature grid by striding, + not out of a grid of their own. + """ + calls = [] + real = JP.u_profile + + def counting(*a, **kw): + calls.append(1) + return real(*a, **kw) + + C, _ = _separable_phi_table(30.0, 0.3) + JP.u_profile = counting + try: + _, _, info = JP.phi_local_lnI(C, n_slots=4, n_seed=4) + finally: + JP.u_profile = real + # THREE, not four: the Newton step, the seed evaluation and the quadrature grid. The + # bound grid used to be a fourth, and no longer calls the profile at all -- sup_g_bound + # needs four quartic roots per point and no u quadrature. A fifth would mean a probe + # is paying its own way; a fourth would mean the bound grid is back on the profile. + assert len(calls) == 3, (len(calls), "the bound grid must not run the u quadrature") + assert "phi_convergence_shift" in info + + # and the striding is exact only for an odd count: the even indices must span the same + # interval and the odd ones must be their midpoints. + assert JP.PHI_NODES_PER_REGION % 2 == 1 + + +def test_the_outside_bound_does_not_depend_on_the_u_quadrature_at_all(): + """The review finding this replaces is retired BY CONSTRUCTION, not by a gate. + + ``Fb`` and ``d1b`` used to come from ``u_profile`` on the bound grid, so a whole-cell + fallback there could underestimate ``F`` and a lift applied to an underestimate bounds + nothing. The fix was a gate on that fallback. :func:`sup_g_bound` removes the + exposure instead: the outside bound is ``log(2 pi) + max_u g``, four quartic roots per + point, and never touches the quadrature. + + So the property to assert is not "the gate fires" but "the bound cannot move": vary + ``u_nodes`` over a factor of 8 and ``sup_outside`` must be bit-identical. That is a + much stronger statement than the gate ever made, and it cannot pass by accident. + """ + KS = 2 + rng = np.random.default_rng(101) + C = rng.normal(size=(3, 2 * KS + 1)) + 1j * rng.normal(size=(3, 2 * KS + 1)) + C = jnp.asarray(C * (1e3 / np.sum(np.abs(C)))) + sups = [float(JP.phi_local_lnI(C, u_nodes=un)[2]["sup_outside"]) + for un in (48, 96, 384)] + assert sups[0] == sups[1] == sups[2], sups + + +def test_sup_g_bound_is_actually_an_upper_bound_on_the_profile(): + """The whole certificate now rests on ``F(phi) <= log(2 pi) + max_u g(phi,u)``. If that + is ever violated the outside bound is not a bound and every accepted row is suspect, so + it is checked directly against the profile rather than assumed from the algebra. + + Measured slack is 1.2-5.5 nats across the range -- small enough that the bound is + usable, and the reason the certificate stopped declining rows whose true margin was + already tens of nats clear. + """ + KS = 2 + worst = 1e9 + for KP, amp in ((3, 30.0), (3, 3e3), (5, 1e3), (9, 1e4)): + rng = np.random.default_rng(7) + C = rng.normal(size=(KP, 2 * KS + 1)) + 1j * rng.normal(size=(KP, 2 * KS + 1)) + C = jnp.asarray(C * (amp / np.sum(np.abs(C)))) + # 128 u nodes and 17 phi: the claim is an INEQUALITY that was violated by + # 0.024-0.092 nats, so the reference does not need to be fine, only correct. + un = min(JP.required_u_nodes(amp), 128) + for phi in np.linspace(0.0, 2 * np.pi, 17, endpoint=False): + F = float(JP.u_profile(C, float(phi), n_nodes=un)[0]) + H = float(JP.sup_g_bound(C, float(phi))) + worst = min(worst, H - F) + assert worst >= 0.0, ("sup_g_bound is NOT an upper bound", worst) + assert worst < 20.0, ("bound is sound but so loose it cannot certify", worst) + + +def test_the_localized_regime_now_accepts_and_is_right(): + """What the whole exercise was for. Before the bound was rebuilt, a sweep over + KP x amplitude found exactly ONE case in 36 that both localized (area_outside > 0) and + accepted, at amplitude 100 -- the cost win of localizing phi was real and the + certificate refused every row that realized it. The Taylor lift sat 2.7e5 nats above + the integral at amplitude 3e4 while the true margin was about -66. + + These cases localize into several regions AND accept AND are right to machine + precision. If this test starts declining, the certificate has regressed to refusing + the regime it exists to serve. + """ + KS = 2 + # u_nodes 256 and 8 slots, not required_u_nodes(amp) = 2310 and 16. The sizing helper + # is a conservative UPPER bound derived from amplitude; the gate that actually decides, + # u_sizing_ok, measures risky cells and passes here at 9x fewer nodes. Sizing this + # test from the helper costs 5.9 GB in eval_g2's intermediate and is killed when the + # file runs as a whole -- a guard that cannot run in CI guards nothing. + for KP, amp in ((3, 3e3), (9, 1e3)): + rng = np.random.default_rng(7) + C = rng.normal(size=(KP, 2 * KS + 1)) + 1j * rng.normal(size=(KP, 2 * KS + 1)) + C = C * (amp / np.sum(np.abs(C))) + v, ok, info = JP.phi_local_lnI(jnp.asarray(C), + n_bound=int(JP.required_bound_grid(amp)), + u_nodes=256, n_slots=8, n_nodes=97) + assert float(info["area_outside"]) > 0.0, "not localized -- fixture is degenerate" + assert int(info["n_phi_regions"]) >= 4, int(info["n_phi_regions"]) + assert bool(ok), (KP, amp, float(info["margin"])) + assert abs(float(v) - _torus_ref(np.asarray(C))) < 1e-3, float(v) + + +def test_the_bound_grid_adequacy_gate_fires_and_is_cleared_by_sizing(): + """Non-vacuity, at the source rather than through the kernel so it stays affordable. + + A gate that never fires is decoration. This one must fire on a table sharp enough + that 48 nodes cannot resolve a whole cell, and must CLEAR when the node count is + raised to what the curvature bound asks for -- that is what makes it a sizing + requirement the caller can act on rather than a wall. ``required_u_nodes`` is the + static helper that computes the same quantity from an amplitude proxy, and both now + read ``U_PTS_PER_SIGMA`` so the budget and the check cannot drift apart. + """ + KS = 2 + rng = np.random.default_rng(101) + C = rng.normal(size=(3, 2 * KS + 1)) + 1j * rng.normal(size=(3, 2 * KS + 1)) + C = jnp.asarray(C * (1.0e4 / np.sum(np.abs(C)))) + + fired = cleared = 0 + for phi in np.linspace(0.0, 2 * np.pi, 12, endpoint=False): + _, _, _, fb_lo, risk_lo, _ = JP.u_profile(C, float(phi), n_nodes=48) + _, _, _, fb_hi, risk_hi, _ = JP.u_profile(C, float(phi), n_nodes=384) + assert int(fb_lo) > 0 # minima always fall back; that is fine + fired += int(risk_lo) > 0 + cleared += int(risk_hi) == 0 + assert fired > 0, "an adequacy gate that never fires cannot protect the bound" + assert cleared == 12, "sizing the quadrature must clear it, or it is not a requirement" + assert JP.required_u_nodes(1.0e4) > 48 + + +def test_sup_g_bound_survives_a_degenerate_u_quartic(): + """Adversarial review, and the worst defect in this branch. + + ``sup_g_bound`` took ``max`` over ``u_stationary_roots`` as if that set contained the + maximizer. A max over a candidate set is a LOWER bound unless it provably does, and + ``u_stationary_roots`` substitutes ``lead = 1`` when ``c2 == 0`` -- solving a different + polynomial -- so for a table with no ``q = +-2`` content it need not. The whole outside + certificate rests on ``bound >= F``, so this made margins understated rather than loose. + + Measured before the fix: 0.024 to 0.092 nats BELOW ``log(2 pi) + max_u g``. The + fixture is the degenerate table, because every table in the rest of this file carries + full mode content and none of them can see it. + """ + KS = 2 + for trial in range(4): + rng = np.random.default_rng(trial) + C = np.zeros((3, 2 * KS + 1), dtype=complex) + C[:, KS + 0] = rng.normal(size=3) + 1j * rng.normal(size=3) + C[:, KS + 1] = rng.normal(size=3) + 1j * rng.normal(size=3) # no q = +-2 + C = jnp.asarray(C * (50.0 / np.sum(np.abs(C)))) + for phi in np.linspace(0.0, 2 * np.pi, 13, endpoint=False): + H = float(JP.sup_g_bound(C, float(phi))) + u = np.linspace(0.0, 2 * np.pi, 8000, endpoint=False) + g = np.asarray(JP.eval_g2(C, jnp.full(u.shape, float(phi)), + jnp.asarray(u), (0, 0))) + assert H >= float(np.log(2 * np.pi) + g.max()) - 1e-9, (trial, phi) + + +def test_empty_slots_do_not_vote_on_the_u_sizing_gate(): + """Adversarial review, found by this session and by external review independently. + + Empty merged-region slots are neutralized for the VALUE -- position zeroed, weight + masked -- but their nodes are still evaluated at phi = 0, and their fallback counts + were summed into the gate and the reported counters. A risky cell at that artificial + point could decline a row whose every contributing node was adequate. + + The tell is that the counts tracked the SLOT ALLOCATION rather than the regions: + 5 risky at n_slots=2 and 176 at n_slots=8 while the region count only went 2 -> 4. + So the invariant to pin is that the counters do not move once the slots exceed the + regions, which no amount of real structure could cause. + """ + KS = 2 + rng = np.random.default_rng(101) + C = rng.normal(size=(3, 2 * KS + 1)) + 1j * rng.normal(size=(3, 2 * KS + 1)) + C = jnp.asarray(C * (1000.0 / np.sum(np.abs(C)))) + # SLOTS 4/8/16 AT 48 NODES, not 8/32/64 at 96. The invariant is that the counters + # stop moving once the slots exceed the regions, and with 4 regions that is shown just + # as well at 16 as at 64 -- for an eighth of the memory. The 64-slot version cost + # 954 MB in one call and helped kill the CI gate at exit 137. + seen = {} + for ns in (4, 8, 16): + _, _, i = JP.phi_local_lnI(C, u_nodes=48, n_slots=ns, n_nodes=97) + seen[ns] = (int(i["n_phi_regions"]), int(i["n_u_risky_quad"]), + int(i["n_u_fallback_quad"])) + assert len({v[0] for v in seen.values()}) == 1, ("regions moved", seen) + assert len({v[1] for v in seen.values()}) == 1, ("risky tracked slots", seen) + assert len({v[2] for v in seen.values()}) == 1, ("fallback tracked slots", seen) + + +def _AB(scale, seed=3, KP=3, KS=2): + rng = np.random.default_rng(seed) + A = (rng.normal(size=(KP, 2 * KS + 1)) + 1j * rng.normal(size=(KP, 2 * KS + 1))) * scale + B = (rng.normal(size=(KP + 2, 2 * KS + 1)) + + 1j * rng.normal(size=(KP + 2, 2 * KS + 1))) * scale + B[0, KS] = abs(B[0, KS].real) + 3.0 * scale + return jnp.asarray(A), jnp.asarray(B) + + +def test_phi_local_distance_combiner_matches_the_dense_scheme(): + """The wiring's correctness condition. ``joint_lnL_phi_local`` must reproduce the + shipping ``joint_lnL_phi_dense`` on the same inputs, or the normalization split + between the per-node seam (bare torus integral) and the combiner (the ``(2 pi)^-2`` + prior factor) is wrong -- and that error is invisible in the per-node value. + """ + # deliberately small: the claim is about NORMALIZATION, which a 6-node grid tests as + # well as a 32-node one, and the full-size version cannot run beside the rest of this + # file -- 640 MB of eval_g2 intermediate per chunk kills the process. + for scale in (1.0, 4.0, 12.0): + A, B = _AB(scale) + x = jnp.linspace(0.5, 2.0, 6) + lw = jnp.full(6, -np.log(6.0)) + d = JP.joint_lnL_phi_dense(A, B, x, lw, n_phi=512, n_nodes=48) + v, _, _ = JP.joint_lnL_phi_local(A, B, x, lw, u_nodes=48, n_slots=8, + n_nodes=97, x_chunk=2) + assert abs(float(d) - float(v)) < 1e-5, (scale, float(d), float(v)) + + +def test_an_arbitrary_distance_rule_stacks_on_the_seam(): + """RO asked for this to be stackable, so it is asserted rather than described. + + ``phi_local_lnI_at_distance`` is the unit; the distance rule is entirely a matter of + which ``x`` a caller evaluates and what weights it applies. A 5-node Gauss-Legendre + rule -- nothing grid-shaped about it -- driven by hand through the seam must equal the + same nodes routed through the default combiner. If those ever diverge, the combiner + has grown an assumption about the rule and stacking is broken. + """ + A, B = _AB(4.0) + gx, gw = np.polynomial.legendre.leggauss(5) + xs = 0.5 * (gx + 1.0) * 1.5 + 0.5 + lws = np.log(gw * 0.75) + kw = dict(u_nodes=48, n_slots=8, n_nodes=97) + vals = np.array([float(JP.phi_local_lnI_at_distance(A, B, float(z), **kw)[0]) + for z in xs]) + from scipy.special import logsumexp + stacked = logsumexp(vals + lws) - 2.0 * np.log(2.0 * np.pi) + through, _, _ = JP.joint_lnL_phi_local(A, B, jnp.asarray(xs), jnp.asarray(lws), + x_chunk=1, **kw) + assert abs(stacked - float(through)) < 1e-9, (stacked, float(through)) + + +def test_rolling_the_quadrature_axis_does_not_move_the_value(): + """``pt_chunk`` exists to bound memory and must be invisible in the answer -- the same + contract ``phi_chunk`` has on the dense path. Bit-identical, not merely close: a scan + that reassociated the reduction would show up here as a last-digit drift and would + mean the chunking is not a pure refactor.""" + KS = 2 + rng = np.random.default_rng(7) + C = rng.normal(size=(3, 2 * KS + 1)) + 1j * rng.normal(size=(3, 2 * KS + 1)) + C = jnp.asarray(C * (3e3 / np.sum(np.abs(C)))) + kw = dict(n_bound=int(JP.required_bound_grid(3e3)), u_nodes=96, + n_slots=4, n_nodes=97) + # The unrolled reference asks for a chunk of exactly the grid, not 1e6: `_prof_scan` + # now clamps, but a test should not depend on that and external review measured + # 6.56 GB of temporaries when it padded 388 points to a million. + n_pts = 4 * 97 + ref = float(JP.phi_local_lnI(C, pt_chunk=n_pts, **kw)[0]) + for pc in (16, 32, 256): + got = float(JP.phi_local_lnI(C, pt_chunk=pc, **kw)[0]) + # NOT exact equality. This asserted bit-identity and passed here, but external + # review measured 4.55e-13 nats of spread across chunk sizes on another CPU: the + # scan's reduction IS reassociated on some platforms, so bit-identity was a claim + # about this machine rather than about the code. The contract that matters is + # that chunking is invisible at the scale anything downstream cares about. + assert abs(got - ref) < 1e-9, (pc, got, ref) + + +def test_the_distance_combiner_is_fail_closed_across_nodes(): + """``ok`` is the CONJUNCTION over distance nodes. A declining node still returns a + finite number that would otherwise be summed in silently, so one bad node must sink + the row. Asserted by starving a single node's slot budget.""" + A, B = _AB(4.0) + x = jnp.linspace(0.5, 2.0, 4) + lw = jnp.full(4, -np.log(4.0)) + _, ok_starved, _ = JP.joint_lnL_phi_local(A, B, x, lw, u_nodes=48, n_slots=1, + n_nodes=97, x_chunk=2) + assert not bool(ok_starved), "a starved node must sink the distance sum" diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py b/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py index 62479cf7e..4ab2558c6 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_limit_distance_jax.py @@ -294,9 +294,16 @@ def test_driver_resolves_and_forwards_the_box(): 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 + # Baseline, slow rotation, finite response, combined response, and the + # direct-marginalization wrapper must all preserve the physical prior box. + assert src.count('d_prior_range=(opts.d_min, opts.d_max)') == 5 assert 'like_data, d_lo, d_hi' in src - assert '"--d-prior", "--limit-distance",' in src # in the `implemented` set + # in the `implemented` set (PR 286 moved --d-prior OUT of it, so pin the + # block's contents rather than a neighbouring token) + start = src.index('implemented = {') + block = src[start:src.index('}', start)] + assert '"--limit-distance"' in block + assert '"--d-prior"' not in block @pytest.mark.skipif(not os.path.exists(_DRIVER), reason='JAX driver not in this tree') diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_fallback_visibility.py b/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_fallback_visibility.py new file mode 100644 index 000000000..f87020ed2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_fallback_visibility.py @@ -0,0 +1,485 @@ +"""A planner fault must not be readable as a conservative policy decline. + +``multipeak_local_marginalize`` returns the caller's reserve for two unrelated +reasons. One is a budget outcome: both tiers ran and a diagnostic failed. The +other is a fault: something raised and the reserve is standing in for a step +that never ran. Before the visibility change both looked the same in the log +(silence) and differed in the record only by a substring of ``provenance``, so +a ladder campaign in which every row declined by ``RuntimeError`` was read as a +conservative controller rather than as a defect. + +These tests pin the separation itself, not the tier1 defect that exposed it: +the fault is reported and the budget decline is not, ``fail_on_fallback`` is +fatal on the first and inert on the second, the record carries a +machine-readable ``decline_kind``/``fault``, and the default path is +byte-for-byte what it was. + +Two properties of the REPORTING are pinned here because a warning cannot supply +them. The report must not be suppressible or made fatal by a process-global +filter the caller set for unrelated reasons -- ``-W error::RuntimeWarning`` +would otherwise turn the default ``fail_on_fallback=False`` path into a raise -- +and it must arrive once per CALL, not once per call site, because the campaign +case that motivated the change runs one call site with ``label=None`` and so +emits an identical message every time. +""" + +import copy +import logging +import pickle +import warnings + +import jax +import numpy as np +import pytest + +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile import multipeak_planner as planner # noqa: E402 + + +# The tuple contract MultiPeakResult has, and had before decline_kind/fault +# were added. Defaulted trailing NamedTuple fields keep a 13-argument +# CONSTRUCTION working, but they do NOT keep 13-name UNPACKING working: a +# NamedTuple is a tuple, so a 15th field makes `a, ..., m = result` raise +# ValueError. Unpacking is the contract callers hold, so decline_kind and +# fault are attributes and the tuple stays 13 elements. Frozen here so an +# insertion in the middle -- which would silently reorder a caller's tuple -- +# fails instead of passing. +_LEGACY_FIELDS = ( + "value", "accepted", "used_reserve", "provenance", "delta_log_integral", + "tier0", "tier1", "tier0_portfolio", "tier1_portfolio", + "total_lattice_evaluations", "total_refinement_steps", + "total_local_evaluations", "modeled_peak_bytes", +) + +# The two provenance strings the pre-change module produced. Downstream +# analysis that greps them must keep working; the new record fields are an +# addition, not a replacement. +_PROVENANCE_ACCEPTED = "uvq-multipeak-tier1" +_PROVENANCE_BUDGET = "dense-reserve:enrichment-or-local-diagnostic" +_PROVENANCE_FAULT_PREFIX = "dense-reserve:planner-exception:" + + +def _synthetic_tables(n_time=9): + """Small reflected-polynomial problem with an interior four-axis peak.""" + time = np.arange(n_time, dtype=float) + C_A = np.zeros((3, 3, n_time), dtype=np.complex128) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_A[0, 1] = 20.0 - 2.0 * np.cos(2.0 * np.pi * time / (n_time - 1)) + C_A[2, 0] = 0.25 + C_A[2, 2] = 0.25 + C_B[0, 2] = 4.0 + C_B[2, 1] = 0.02 + C_B[2, 3] = 0.02 + return C_A, C_B + + +# Settings that accept the local branch, and settings whose only difference is +# an unreachable agreement budget, so the same tables decline by budget. Both +# rows run both tiers to completion; neither raises. +_ACCEPT_KWARGS = dict( + log_integral_tol=0.1, tier0=(2, 2, 24), tier1=(3, 3, 48), + quadrature_order=7, cell_sigma=4.0, chunk_size=32) +_BUDGET_KWARGS = dict( + log_integral_tol=1.0e-12, tier0=(2, 2, 24), tier1=(3, 3, 48), + quadrature_order=5, cell_sigma=4.0, chunk_size=32) + + +# Captured once, before any test patches it. +_REAL_RUN_STRUCTURAL_TIER = planner._run_structural_tier + + +def _fault_logs(caplog): + """The planner's own WARNING records, ignoring anything else logging.""" + return [r for r in caplog.records + if r.name == planner.__name__ and r.levelno >= logging.WARNING] + + +class _CountingReserve(object): + """A finite reserve that records whether the planner actually paid for it.""" + + def __init__(self, value=123.456): + self.value = float(value) + self.calls = 0 + + def __call__(self): + self.calls += 1 + return self.value + + +def _fault_at(monkeypatch, which, message="deliberate tier fault"): + """Make the ``which``-th structural tier raise, and the others run for real. + + ``which=2`` reproduces the campaign's shape exactly: tier0 converges, tier1 + raises, and the row falls back. The failure is injected at the tier seam + rather than by feeding bad tables, so tier0's report is real and the stage + the planner names can be checked against a known answer. + """ + state = {"n": 0} + + def wrapper(*args, **kwargs): + state["n"] += 1 + if state["n"] == int(which): + raise RuntimeError(message) + # Always the pristine function, so patching twice in one test does not + # stack wrappers and fire on the wrong call. + return _REAL_RUN_STRUCTURAL_TIER(*args, **kwargs) + + monkeypatch.setattr(planner, "_run_structural_tier", wrapper) + return state + + +def test_budget_decline_and_accepted_row_stay_silent(caplog): + C_A, C_B = _synthetic_tables() + caplog.set_level(logging.DEBUG, logger=planner.__name__) + + # A budget decline is a normal outcome. It must not report, or a campaign + # that declines legitimately drowns the faults it is supposed to surface. + reserve = _CountingReserve() + caplog.clear() + declined = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, **_BUDGET_KWARGS) + assert declined.used_reserve and not declined.accepted + assert declined.decline_kind == planner.DECLINE_DIAGNOSTIC + assert _fault_logs(caplog) == [] + + # An accepted row must not report either. + caplog.clear() + accepted = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + assert accepted.accepted and accepted.decline_kind is None + assert _fault_logs(caplog) == [] + + +def test_fault_log_names_stage_exception_and_label(monkeypatch, caplog): + C_A, C_B = _synthetic_tables() + caplog.set_level(logging.DEBUG, logger=planner.__name__) + _fault_at(monkeypatch, 2, message="tier1 refinement is degenerate") + reserve = _CountingReserve() + caplog.clear() + result = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, label="ladder-row-7", **_ACCEPT_KWARGS) + records = _fault_logs(caplog) + + # Exactly one record per CALL, not per Newton step: a 40-row campaign gets + # 40 lines, which is readable; per-step would not be. + assert len(records) == 1 + assert records[0].levelno == logging.WARNING + text = records[0].getMessage() + assert "tier1" in text + assert "RuntimeError" in text + assert "tier1 refinement is degenerate" in text + assert "ladder-row-7" in text + assert "fault" in text.lower() + assert result.used_reserve and result.value == reserve.value + + +def test_record_separates_fault_from_budget_without_parsing_provenance( + monkeypatch): + C_A, C_B = _synthetic_tables() + + budget = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_BUDGET_KWARGS) + _fault_at(monkeypatch, 2, message="tier1 refinement is degenerate") + fault = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + + # Both used the reserve; only the second is a defect. The distinction is a + # field comparison, not a substring search on provenance. + assert budget.used_reserve and fault.used_reserve + assert budget.decline_kind == planner.DECLINE_DIAGNOSTIC + assert fault.decline_kind == planner.DECLINE_FAULT + assert planner.DECLINE_DIAGNOSTIC != planner.DECLINE_FAULT + assert budget.fault is None + assert isinstance(fault.fault, planner.FallbackFault) + assert fault.fault.stage == "tier1" + assert fault.fault.error_type == "RuntimeError" + assert fault.fault.message == "tier1 refinement is degenerate" + + +def test_fault_stage_names_the_step_that_raised(monkeypatch): + """The stage is measured, not assumed: tier0 and tier1 report differently.""" + C_A, C_B = _synthetic_tables() + _fault_at(monkeypatch, 1) + first = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + _fault_at(monkeypatch, 2) + second = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + assert first.fault.stage == "tier0" + assert second.fault.stage == "tier1" + + # A fault before either tier runs is attributed to its own stage, so a bad + # table is never reported as a tier defect. + repeated = np.repeat(C_B[..., None], 7, axis=-1) + repeated[1, 1, 3] += 1.0e-3 + early = planner.multipeak_local_marginalize( + C_A, repeated, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + assert early.decline_kind == planner.DECLINE_FAULT + assert early.fault.stage == "uv-summary" + assert early.fault.error_type == "ValueError" + + +def test_fail_on_fallback_raises_on_fault_and_is_inert_on_budget_decline( + monkeypatch): + C_A, C_B = _synthetic_tables() + + # Inert on the budget decline: same value, same record, no exception. + reserve = _CountingReserve() + declined = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, fail_on_fallback=True, **_BUDGET_KWARGS) + assert declined.used_reserve and declined.value == reserve.value + assert declined.decline_kind == planner.DECLINE_DIAGNOSTIC + assert reserve.calls == 1 + + # Inert on an accepted row. + accepted = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), fail_on_fallback=True, + **_ACCEPT_KWARGS) + assert accepted.accepted and accepted.provenance == _PROVENANCE_ACCEPTED + + # Fatal on the fault, and it does not pay for the reserve first: the point + # is to stop, not to produce a value nobody should trust. + _fault_at(monkeypatch, 2, message="tier1 refinement is degenerate") + fatal_reserve = _CountingReserve() + with pytest.raises(planner.MultiPeakFallbackError) as excinfo: + planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, fatal_reserve, fail_on_fallback=True, + label="ladder-row-7", **_ACCEPT_KWARGS) + assert fatal_reserve.calls == 0 + assert "tier1" in str(excinfo.value) + assert isinstance(excinfo.value.__cause__, RuntimeError) + + # MultiPeakFallbackError sits outside the set the planner catches, so a + # nested or repeated call cannot swallow it back into a decline. + assert not isinstance( + excinfo.value, (RuntimeError, ValueError, np.linalg.LinAlgError)) + + +def test_default_path_is_unchanged(monkeypatch): + """Same values, same provenance, same legacy record shape, no exception.""" + C_A, C_B = _synthetic_tables() + assert planner.MultiPeakResult._fields == _LEGACY_FIELDS + + # A caller built on the pre-change arity still constructs the record, and + # the record it gets back is still exactly that many elements. + legacy = planner.MultiPeakResult(*range(len(_LEGACY_FIELDS))) + assert len(legacy) == len(_LEGACY_FIELDS) + assert tuple(legacy) == tuple(range(len(_LEGACY_FIELDS))) + assert legacy.decline_kind is None and legacy.fault is None + + accepted = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + assert accepted.accepted and not accepted.used_reserve + assert accepted.provenance == _PROVENANCE_ACCEPTED + assert np.isfinite(accepted.value) + + reserve = _CountingReserve() + budget = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, **_BUDGET_KWARGS) + assert not budget.accepted and budget.used_reserve + assert budget.value == reserve.value + assert budget.provenance == _PROVENANCE_BUDGET + assert reserve.calls == 1 + + _fault_at(monkeypatch, 2) + fault_reserve = _CountingReserve(77.25) + fault = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, fault_reserve, **_ACCEPT_KWARGS) + # Default is off, so the fault still RETURNS the reserve exactly as before. + assert not fault.accepted and fault.used_reserve + assert fault.value == 77.25 + assert fault.provenance == _PROVENANCE_FAULT_PREFIX + "RuntimeError" + assert not np.isfinite(fault.delta_log_integral) + assert fault_reserve.calls == 1 + + # A failing reserve still surfaces as _DenseReserveError, not as a planner + # decline and not as the new fallback error. + def failing_reserve(): + raise ValueError("deliberate reserve failure") + + with pytest.raises(planner._DenseReserveError): + planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, failing_reserve, **_BUDGET_KWARGS) + + +def test_legacy_thirteen_name_unpacking_still_works(): + """The contract a caller holds is UNPACKING, and it is 13 names wide. + + Defaulted trailing NamedTuple fields keep 13-argument CONSTRUCTION working, + so a test that only constructs cannot see this break. A NamedTuple is a + tuple: two appended fields make ``len()`` 15 and every existing + ``a, ..., m = result`` raise ``ValueError: too many values to unpack``. + That is why decline_kind and fault are attributes and not tuple elements. + """ + record = planner.MultiPeakResult(*range(13)) + + # The failure this pins is a ValueError from the unpacking statement + # itself; there is no way to write it that a construction test also covers. + (value, accepted, used_reserve, provenance, delta_log_integral, + tier0, tier1, tier0_portfolio, tier1_portfolio, + total_lattice_evaluations, total_refinement_steps, + total_local_evaluations, modeled_peak_bytes) = record + assert (value, modeled_peak_bytes) == (0, 12) + + # Everything else that reads the tuple as a sequence agrees on 13. + assert len(record) == 13 + assert len(tuple(record)) == 13 + assert len(list(record)) == 13 + assert len(planner.MultiPeakResult._fields) == 13 + assert len(record._asdict()) == 13 + assert record[-1] == 12 + assert record + () == tuple(range(13)) + + # And a real result, not only a hand-built one. + C_A, C_B = _synthetic_tables() + live = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), **_ACCEPT_KWARGS) + assert len(live) == 13 + unpacked_value = list(live)[0] + assert unpacked_value == live.value + + +def test_annotations_are_attributes_and_survive_replace_copy_pickle(): + """They are reachable, immutable, and not smuggled into the tuple.""" + fault = planner.FallbackFault("tier1", "RuntimeError", "degenerate") + record = planner.MultiPeakResult( + *range(13), decline_kind=planner.DECLINE_FAULT, fault=fault) + + assert record.decline_kind == planner.DECLINE_FAULT + assert record.fault is fault + assert tuple(record) == tuple(range(13)) + assert planner.DECLINE_FAULT not in tuple(record) + assert fault not in tuple(record) + assert "decline_kind" not in record._fields + assert "fault" not in record._fields + + # Immutable like the tuple part, so a consumer cannot annotate a record + # after the fact and have it read as the planner's own finding. + with pytest.raises(AttributeError): + record.decline_kind = planner.DECLINE_DIAGNOSTIC + with pytest.raises(AttributeError): + del record.fault + + # _replace, _make, copy and pickle all bypass __new__ in some way; each + # must still carry the annotation, or a round-tripped fault reads as clean. + replaced = record._replace(value=99.0) + assert replaced.value == 99.0 + assert replaced.decline_kind == planner.DECLINE_FAULT + assert replaced.fault == fault + assert len(replaced) == 13 + assert record._replace(decline_kind=None).decline_kind is None + with pytest.raises(ValueError): + record._replace(no_such_field=1) + + remade = planner.MultiPeakResult._make(range(13), fault=fault) + assert remade.fault == fault and len(remade) == 13 + + assert copy.copy(record).fault == fault + assert copy.deepcopy(record).decline_kind == planner.DECLINE_FAULT + assert pickle.loads(pickle.dumps(record)).fault == fault + + # A record built the plain way answers None rather than raising, whichever + # construction route produced it. + assert planner.MultiPeakResult._make(range(13)).decline_kind is None + assert repr(record).count("decline_kind") == 1 + + +@pytest.mark.parametrize("filter_action", ["default", "always", "error"]) +def test_fault_reporting_does_not_depend_on_warning_filters( + monkeypatch, caplog, filter_action): + """The default path stays non-fatal, and the fault stays observable. + + ``-W error::RuntimeWarning`` is a filter a caller sets for unrelated + reasons. While the fault was announced with ``warnings.warn`` it raised at + the warn call, BEFORE the ``fail_on_fallback`` check and before the reserve + was evaluated, so that filter alone turned the default path fatal. + """ + C_A, C_B = _synthetic_tables() + caplog.set_level(logging.DEBUG, logger=planner.__name__) + _fault_at(monkeypatch, 2, message="tier1 refinement is degenerate") + reserve = _CountingReserve(77.25) + caplog.clear() + + with warnings.catch_warnings(record=True) as caught: + warnings.resetwarnings() + warnings.simplefilter(filter_action, RuntimeWarning) + # No pytest.raises: the point is that this RETURNS under every filter. + result = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, label="ladder-row-7", + **_ACCEPT_KWARGS) + + assert result.value == 77.25 + assert result.used_reserve and not result.accepted + assert reserve.calls == 1 + + # Observable in the record, which is the primary channel precisely because + # no filter reaches it... + assert result.decline_kind == planner.DECLINE_FAULT + assert result.fault.stage == "tier1" + assert result.fault.error_type == "RuntimeError" + + # ...and on the log, which is the secondary one. + records = _fault_logs(caplog) + assert len(records) == 1 + assert "ladder-row-7" in records[0].getMessage() + + # The module must not route this through the warnings machinery at all: + # under "always" a warn would show up here, and under "error" it would have + # raised above instead of returning. + assert [w for w in caught if issubclass(w.category, RuntimeWarning)] == [] + + +def test_fail_on_fallback_is_the_only_thing_that_makes_a_fault_fatal( + monkeypatch): + """Under warnings-as-errors, both settings keep their documented meaning.""" + C_A, C_B = _synthetic_tables() + + with warnings.catch_warnings(): + warnings.resetwarnings() + warnings.simplefilter("error") + + _fault_at(monkeypatch, 2) + quiet_reserve = _CountingReserve(5.5) + result = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, quiet_reserve, **_ACCEPT_KWARGS) + assert result.value == 5.5 and quiet_reserve.calls == 1 + + _fault_at(monkeypatch, 2) + fatal_reserve = _CountingReserve() + with pytest.raises(planner.MultiPeakFallbackError): + planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, fatal_reserve, fail_on_fallback=True, + **_ACCEPT_KWARGS) + # Still the fallback error, not a RuntimeWarning promoted to an + # exception, and still without paying for the reserve. + assert fatal_reserve.calls == 0 + + +def test_identical_faults_from_one_call_site_report_once_each( + monkeypatch, caplog): + """Not once in total. + + ``warnings.warn`` de-duplicates on (message, category, module, lineno) + under the default filters. The campaign this change exists for calls one + call site with ``label=None``, so every message is identical and the + warning would be shown for the first row only -- exactly the silence the + change is meant to remove. + """ + C_A, C_B = _synthetic_tables() + caplog.set_level(logging.DEBUG, logger=planner.__name__) + n_calls = 3 + caplog.clear() + + with warnings.catch_warnings(): + warnings.resetwarnings() # the DEFAULT filters, where dedup applies + for _ in range(n_calls): + _fault_at(monkeypatch, 2) + result = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, _CountingReserve(), label=None, + **_ACCEPT_KWARGS) + assert result.decline_kind == planner.DECLINE_FAULT + + records = _fault_logs(caplog) + assert len(records) == n_calls + assert len({r.getMessage() for r in records}) == 1 # identical text diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_planner.py b/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_planner.py new file mode 100644 index 000000000..f00d10f23 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_multipeak_planner.py @@ -0,0 +1,588 @@ +"""Load-bearing tests for the U,V/Q multi-peak diagnostic planner.""" + +import os + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile import multipeak_planner as planner # noqa: E402 + + +def _synthetic_tables(n_time=9): + """Small reflected-polynomial problem with an interior four-axis peak.""" + time = np.arange(n_time, dtype=float) + C_A = np.zeros((3, 3, n_time), dtype=np.complex128) + C_B = np.zeros((5, 5), dtype=np.complex128) + # DCT-compatible time dependence, with its interior maximum at t=4. + C_A[0, 1] = 20.0 - 2.0 * np.cos(2.0 * np.pi * time / (n_time - 1)) + C_A[2, 0] = 0.25 + C_A[2, 2] = 0.25 + C_B[0, 2] = 4.0 + C_B[2, 1] = 0.02 + C_B[2, 3] = 0.02 + return C_A, C_B + + +def test_uv_summary_rejects_time_dependent_norm(): + _, C_B = _synthetic_tables() + repeated = np.repeat(C_B[..., None], 7, axis=-1) + summary = planner.summarize_uv_norm_table(repeated) + assert summary.time_invariant + repeated[1, 1, 3] += 1.0e-3 + changed = planner.summarize_uv_norm_table(repeated) + assert not changed.time_invariant + C_A, _ = _synthetic_tables() + with pytest.raises(ValueError, match="arrival-time-dependent"): + planner.rank_joint_starts_from_uvq(C_A, changed, 1.0, 8.0) + + +def test_exact_symmetry_expansion_is_scale_invariant(): + C_A, C_B = _synthetic_tables() + summary = planner.summarize_uv_norm_table(C_B) + first = planner.rank_joint_starts_from_uvq( + C_A, summary, 1.0, 8.0, max_time_starts=2, max_starts=24) + + assert first.symmetry.certified + assert first.symmetry.group_order == 4 + np.testing.assert_allclose( + first.symmetry.shifts, + [[0.0, 0.0], [0.5 * np.pi, np.pi], + [np.pi, 0.0], [1.5 * np.pi, np.pi]], atol=1.0e-13) + assert len(first.starts) == first.symmetry.group_order * len( + first.raw_starts) + for raw_index in range(len(first.raw_starts)): + actions = first.group_action[ + raw_index * first.symmetry.group_order: + (raw_index + 1) * first.symmetry.group_order] + np.testing.assert_array_equal(actions, np.arange(4)) + + # A -> s A, B -> s^2 B, x -> x/s preserves every extrema location in + # (time, phi, u) and changes the exponent only by the constant 4 log(s). + scale = 4.0 + scaled = planner.rank_joint_starts_from_uvq( + scale * C_A, planner.summarize_uv_norm_table(scale * scale * C_B), + 1.0 / scale, 8.0 / scale, max_time_starts=2, max_starts=24) + assert len(scaled.raw_starts) == len(first.raw_starts) + assert len(scaled.starts) == len(first.starts) + np.testing.assert_allclose(scaled.starts[:, :3], first.starts[:, :3], + atol=1.0e-13) + np.testing.assert_allclose(scaled.starts[:, 3], first.starts[:, 3] / scale, + rtol=1.0e-13, atol=1.0e-13) + np.testing.assert_allclose( + scaled.scores - first.scores, 4.0 * np.log(scale), atol=1.0e-12) + + +def test_symmetry_orbits_are_reduced_before_representative_capacity(): + """A louder orbit's four copies must not evict a lower distinct orbit.""" + n_time = 9 + time = np.arange(n_time, dtype=float) + C_A = np.zeros((5, 3, n_time), dtype=np.complex128) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_A[0, 1] = 20.0 - 2.0 * np.cos( + 2.0 * np.pi * time / (n_time - 1)) + C_A[2, 0] = 0.25 + C_A[2, 2] = 0.25 + C_A[4, 1] = -0.25 + 0.05j + C_B[0, 2] = 4.0 + C_B[2, 1] = 0.02 + C_B[2, 3] = 0.02 + portfolio = planner.rank_joint_starts_from_uvq( + C_A, planner.summarize_uv_norm_table(C_B), 1.0, 8.0, + max_time_starts=1, max_starts=8) + assert portfolio.symmetry.group_order == 4 + assert len(portfolio.raw_starts) == 2 + assert len(portfolio.starts) == 8 + assert not portfolio.capacity_truncated + assert portfolio.raw_scores[1] < portfolio.raw_scores[0] + np.testing.assert_array_equal( + portfolio.group_action, np.tile(np.arange(4), 2)) + + +def test_ranked_starts_optimize_distance_at_each_angular_candidate(): + C_A, C_B = _synthetic_tables() + # A strong norm harmonic makes the analytic distance optimum follow angle. + C_B[2, 1] = 0.35 + C_B[2, 3] = 0.35 + portfolio = planner.rank_joint_starts_from_uvq( + C_A, planner.summarize_uv_norm_table(C_B), 1.0, 8.0, + max_time_starts=2, max_starts=24) + phi, u, A = planner._harmonic_lattice( + C_A, portfolio.n_phi_lattice, portfolio.n_u_lattice) + _, _, B = planner._harmonic_lattice( + C_B, portfolio.n_phi_lattice, portfolio.n_u_lattice) + for start in portfolio.raw_starts: + it = int(start[0]) + iphi = int(np.argmin(np.abs(phi - start[1]))) + iu = int(np.argmin(np.abs(u - start[2]))) + _, expected_x = planner._distance_profile( + np.asarray(A[iphi, iu, it]), np.asarray(B[iphi, iu]), 1.0, 8.0) + assert start[3] == pytest.approx(float(expected_x), abs=1.0e-13) + + +def test_time_endpoints_require_one_sided_maximum(): + C_A, C_B = _synthetic_tables() + C_A[0, 1] = np.linspace(12.0, 20.0, C_A.shape[-1]) + portfolio = planner.rank_joint_starts_from_uvq( + C_A, planner.summarize_uv_norm_table(C_B), 1.0, 8.0, + max_time_starts=3, max_starts=24) + assert 0 not in portfolio.time_starts + assert portfolio.time_starts.tolist() == [C_A.shape[-1] - 1] + + +def test_jax_refiner_reaches_strict_stationary_maximum(): + C_A, C_B = _synthetic_tables() + starts = np.asarray([[3.3, 0.2, 0.2, 5.0], + [4.7, 3.0, 0.1, 5.0]]) + result = tuple(np.asarray(item) for item in planner.refine_joint_starts_jax( + C_A, C_B, starts, 1.0, 8.0, iterations=18)) + points, values, gradients, hessians, curvatures = result + selected, stationary = planner.select_refined_modes( + points, values, gradients, curvatures, max_modes=2) + assert stationary.any() + assert len(selected) >= 1 + assert np.max(np.linalg.norm(gradients[selected], axis=1)) < 2.0e-6 + assert np.all(curvatures[selected] > 0.0) + assert np.all(np.isfinite(hessians[selected])) + + +def test_two_tier_local_integral_accepts_or_returns_finite_reserve(): + C_A, C_B = _synthetic_tables() + calls = [] + + def reserve(): + calls.append("called") + return 123.456 + + accepted = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, log_integral_tol=0.1, + tier0=(2, 2, 24), tier1=(3, 3, 48), quadrature_order=7, + cell_sigma=4.0, chunk_size=32) + assert accepted.accepted + assert not accepted.used_reserve + assert accepted.provenance == "uvq-multipeak-tier1" + assert np.isfinite(accepted.value) + assert accepted.delta_log_integral < 0.1 + assert calls == [] + + declined = planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, reserve, log_integral_tol=1.0e-12, + tier0=(2, 2, 24), tier1=(3, 3, 48), quadrature_order=5, + cell_sigma=4.0, chunk_size=32) + assert not declined.accepted + assert declined.used_reserve + assert declined.value == 123.456 + assert declined.provenance.startswith("dense-reserve:") + assert calls == ["called"] + + bad_calls = [] + + def failing_reserve(): + bad_calls.append("called") + raise ValueError("deliberate reserve failure") + + with pytest.raises(planner._DenseReserveError): + planner.multipeak_local_marginalize( + C_A, C_B, 1.0, 8.0, failing_reserve, + log_integral_tol=1.0e-12, tier0=(2, 2, 24), + tier1=(3, 3, 48), quadrature_order=5, + cell_sigma=4.0, chunk_size=32) + assert bad_calls == ["called"] + + +def test_affine_cell_overlap_is_partitioned_not_rejected(): + C_A, C_B = _synthetic_tables() + refined = tuple(np.asarray(item) for item in + planner.refine_joint_starts_jax( + C_A, C_B, np.asarray([[4.0, 0.0, 0.0, 5.0]]), + 1.0, 8.0, iterations=18)) + points, values, _, hessians, _ = refined + one = planner.integrate_refined_modes_tensor( + C_A, C_B, points, values, hessians, 1.0, 8.0, + log_integral_tol=0.1, cell_sigma=4.0, quadrature_order=7, + chunk_size=32) + duplicate = planner.integrate_refined_modes_tensor( + C_A, C_B, np.repeat(points, 2, axis=0), + np.repeat(values, 2), np.repeat(hessians, 2, axis=0), 1.0, 8.0, + log_integral_tol=0.1, cell_sigma=4.0, quadrature_order=7, + chunk_size=32) + assert duplicate.min_core_separation == pytest.approx(0.0) + assert duplicate.overlap_ok + assert duplicate.ok == one.ok + assert duplicate.value == pytest.approx(one.value, abs=2.0e-12) + + +def _log_density_dense(C_A, C_B, theta): + enclosure = planner._time_fourier_enclosure(C_A) + C_t = planner._evaluate_spectrum_numpy( + enclosure[0], enclosure[1], theta[0]) + A, _ = planner._field_variation(C_t, theta[1], theta[2], 0.0, 0.0) + B, _ = planner._field_variation(C_B, theta[1], theta[2], 0.0, 0.0) + x = theta[3] + return x * A - 0.5 * x * x * B - 4.0 * np.log(x) + + +def test_local_fourier_box_bound_dominates_dense_points(): + C_A, C_B = _synthetic_tables() + enclosure = planner._time_fourier_enclosure(C_A) + lo = np.asarray([3.25, 0.0, 0.0, 3.0]) + hi = np.asarray([4.75, 0.6, 0.7, 6.0]) + log_integral_upper, _, point_upper = planner._box_log_upper( + C_A, C_B, enclosure, lo, hi) + rng = np.random.default_rng(1729) + points = rng.uniform(lo, hi, size=(20000, 4)) + dense = np.asarray([_log_density_dense(C_A, C_B, p) for p in points]) + assert np.max(dense) <= point_upper + 2.0e-11 + assert log_integral_upper == pytest.approx( + point_upper + np.log(np.prod(hi - lo)), abs=1.0e-13) + + +def test_overlap_is_owned_once_in_exact_axis_box_geometry(): + C_A, C_B = _synthetic_tables() + summary = planner.summarize_uv_norm_table(C_B) + center = np.asarray([[4.0, np.pi, np.pi, 4.5], + [4.0, np.pi, np.pi, 4.5]]) + half = np.asarray([[4.0, np.pi, np.pi, 3.5], + [4.0, np.pi, np.pi, 3.5]]) + report = planner.hierarchical_union_cover( + C_A, summary, center, half, 1.0, 8.0, + target_log_value=0.0, max_boxes=10) + assert report.bound_certified + assert report.budget_met + assert report.n_owned_leaves == 1 + assert report.n_overlap_owned == 1 + assert report.n_outside_leaves == 0 + assert report.owned_mode.tolist() == [0] + np.testing.assert_allclose(report.owned_centers, center[:1]) + np.testing.assert_allclose(report.owned_half_widths, half[:1]) + + +def test_cover_cap_is_a_decline_not_a_failure_or_false_certificate(): + C_A, C_B = _synthetic_tables() + summary = planner.summarize_uv_norm_table(C_B) + report = planner.hierarchical_union_cover( + C_A, summary, np.asarray([[4.0, 0.0, 0.0, 5.0]]), + np.asarray([[0.1, 0.1, 0.1, 0.1]]), 1.0, 8.0, + target_log_value=1.0e6, outside_tol_nats=-23.0, max_boxes=1) + assert report.bound_certified + assert report.budget_met # A huge supplied target makes the comparison pass. + assert not report.cap_reached # No refinement was needed. + + declined = planner.hierarchical_union_cover( + C_A, summary, np.asarray([[4.0, 0.0, 0.0, 5.0]]), + np.asarray([[0.1, 0.1, 0.1, 0.1]]), 1.0, 8.0, + target_log_value=-1.0e6, outside_tol_nats=-23.0, max_boxes=1) + assert declined.bound_certified + assert not declined.budget_met + assert declined.cap_reached + assert declined.n_outside_leaves == 1 + + +# --------------------------------------------------------------------------- +# Refinement stall. The 2026-09-07 ladder campaign declined EVERY row of every +# rung with "planner-exception: RuntimeError". The cause was not degenerate +# starts and not the tier1 configuration: the bounded Newton loop had a hard +# FIXED POINT, and the single element responsible was the per-coordinate +# ``clip`` that bounded the step. +# +# The modified-Newton direction always ascends, because every ``safe_i`` in +# ``(v_i . g) / safe_i`` is positive. Clipping each coordinate independently +# rescales them by DIFFERENT factors and does not preserve that. At an +# indefinite Hessian the raw step is about 1e8 gradients long, so every +# coordinate saturates and only the signs survive. Measured on the campaign's +# own rung-160 table, row 0 tier1: +# +# eig(-H) = [-1.157e+01, 2.629e+02, 1.924e+03, 1.325e+05] +# g = [-54.47, +51.02, +28.92, 0] +# clipped = [ +2.00, -0.50, +0.50, +0.25] g . d = -119.9 +# lanes = [4878.19, 11062.76, 12679.19, 13091.06, 13237.70] <- zero wins +# +# A descent direction has no improving lane, so the value-only search took its +# zero lane, the iterate was bit-identical next step, and all eighteen steps +# ran without moving. Scaling by ONE factor keeps every ratio and so keeps the +# sign of g . d. all_axis_peaklocal.py, the independent four-axis +# implementation merged in #268, bounds its own Newton step the same way and +# for the same stated reason. +# +# Reporting the resulting decline under its own name is PR #277's subject, not +# this one's. +# +# Both tiers stalled, on different rows of the campaign, so the four gradient +# norms agreeing to ten digits was the certified order-4 symmetry orbit doing +# its job, not a degeneracy. Zero spread WITHIN an orbit is correct; a test +# asserting non-zero spread would fail on healthy rows. +# +# Tried and rejected, both measured against the real rung-160 table and neither +# shipped: taking |lambda| instead of flooring at ridge (the ablation shows the +# tier converges identically with and without it), and widening the +# backtracking ladder from 1/8 to 2**-15 (converged 16/64 with, 18/64 without, +# from the same random starts). +# --------------------------------------------------------------------------- + +_STALLED_SPECTRUM = (-1.157203234765e+01, 2.629300619639e+02, + 1.924365666846e+03, 1.324806698996e+05) +_STALLED_GRADIENT = (-5.446512310805e+01, 5.102215811225e+01, + 2.892359149328e+01, -1.818989403546e-12) +_MAX_STEP = (2.0, 0.5, 0.5, 0.25) + + +def _narrow_time_peak_tables(amp, n_time=65, width=0.35, centre=32.37): + """Fixture with the production GEOMETRY, not just its amplitude. + + ``_synthetic_tables`` puts its maximum exactly on the targeting lattice + and is nearly isotropic, so the Newton loop is barely exercised there and + the stall is invisible. What makes the production problem hard is that the + arrival-time peak is much narrower than one sample -- the campaign measured + deltaT/sigma_t between 4 and 65 -- which is what makes a fixed absolute + step bound span many peak widths. The centre is deliberately off-grid. + """ + time = np.arange(n_time, dtype=float) + bump = np.exp(-0.5 * ((time - centre) / width) ** 2) + C_A = np.zeros((3, 3, n_time), dtype=np.complex128) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_A[0, 1] = amp * (2.0 + 18.0 * bump) + C_A[2, 0] = 0.25 * amp * bump + C_A[2, 2] = 0.25 * amp * bump + # Scaling C_B with C_A holds the distance optimum inside [x_min, x_max]: + # x* ~ A/B. Scaling C_A alone would drive x* through x_max and the residual + # gradient would then be a boundary artefact rather than a stall. + C_B[0, 2] = 4.0 * amp + C_B[2, 1] = 0.02 * amp + C_B[2, 3] = 0.02 * amp + return C_A, C_B + + +def test_bounded_ascent_direction_ascends_and_respects_its_bound(): + """The BOUNDED step must still increase lnL to first order. + + ``g . d > 0`` is what makes a backtracking search able to succeed at all. + The unbounded direction has it by construction; the old per-coordinate clip + destroyed it, giving ``g . d = -119.9`` on the campaign's own stalled row. + """ + max_step = jnp.asarray(_MAX_STEP, dtype=jnp.float64) + gradient = jnp.asarray(_STALLED_GRADIENT, dtype=jnp.float64) + rng = np.random.default_rng(20260907) + checked = 0 + for trial in range(40): + # A random orthonormal frame carrying the measured spectrum, plus + # spectra with more than one negative eigenvalue. The invariant is a + # property of the construction, not of one matrix. + basis, _ = np.linalg.qr(rng.normal(size=(4, 4))) + if trial < 20: + spectrum = np.asarray(_STALLED_SPECTRUM) + else: + spectrum = rng.normal(size=4) * 10.0 ** rng.uniform(-1, 4, size=4) + hessian = jnp.asarray(-(basis * spectrum) @ basis.T, + dtype=jnp.float64) + this_gradient = (gradient if trial < 20 else + jnp.asarray(rng.normal(size=4) * 50.0)) + direction, _ = planner._bounded_ascent_direction( + this_gradient, hessian, 1.0e-8, max_step) + assert np.all(np.isfinite(np.asarray(direction))) + assert np.all(np.abs(np.asarray(direction)) + <= np.asarray(_MAX_STEP) * (1.0 + 1.0e-12)) + assert float(jnp.dot(this_gradient, direction)) > 0.0 + checked += 1 + assert checked == 40 + # A zero gradient gives a zero step, so a converged point stays put. This + # is not a NaN guard: max_step is validated positive, so the rescale is + # finite for a zero direction with or without the tiny floor. + at_rest, _ = planner._bounded_ascent_direction( + jnp.zeros(4), jnp.asarray(-np.eye(4)), 1.0e-8, max_step) + assert np.allclose(np.asarray(at_rest), 0.0) + + +def test_refinement_outcome_does_not_collapse_as_the_step_bound_widens(): + """``max_step`` is a safeguard; it must not decide whether the loop works. + + The bound is fixed in absolute coordinates while the peak width scales as + 1/rho, so widening it is the CI-affordable stand-in for raising amplitude: + production runs at 8 to 130 arrival-sample widths per unit of the time + bound, and a synthetic table narrow enough to reach that at the default + bound has no interior maximum left to find. Measured across this sweep, + frozen interior starts and converged modes: + + old 0, 3, 13 frozen; 4, 2, 0 converged + this commit 0, 0, 5 frozen; 3, 4, 3 converged + + The five at the widest bound were not converging under either version, so + the assertion below is on the two bounds where standing still is a defect, + plus the convergence count across all three. + """ + C_A, C_B = _narrow_time_peak_tables(256.0) + rng = np.random.default_rng(31) + starts = np.column_stack([ + rng.uniform(28.0, 37.0, 24), rng.uniform(0.0, 2.0 * np.pi, 24), + rng.uniform(0.0, 2.0 * np.pi, 24), rng.uniform(1.0, 8.0, 24)]) + converged, frozen_counts = [], [] + for bound in (_MAX_STEP, (8.0, 2.0, 2.0, 1.0), (32.0, 8.0, 8.0, 4.0)): + refined = tuple(np.asarray(item) for item in + planner.refine_joint_starts_jax( + C_A, C_B, starts, 1.0, 8.0, iterations=18, + max_step=bound)) + points, values, gradients, hessians, curvatures = refined + interior = ((points[:, 0] > 0.0) + & (points[:, 0] < C_A.shape[-1] - 1.0) + & (points[:, 3] > 1.0) & (points[:, 3] < 8.0)) + frozen_counts.append( + int((np.all(points == starts, axis=1) & interior).sum())) + selected, _ = planner.select_refined_modes( + points, values, gradients, curvatures, max_modes=len(starts)) + converged.append(len(selected)) + assert frozen_counts[0] == 0 and frozen_counts[1] == 0, ( + "interior starts frozen in place at the production and 4x bounds: %s" + % frozen_counts) + # An ABSOLUTE floor. Normalizing against converged[0] would compare the + # code under test with itself: [0, 0, 0] satisfies any ratio, and a real + # loss at the production bound would be absorbed into the baseline rather + # than caught. What the old construction did was reduce the loop to + # finding NOTHING as the bound widened (4, 2, 0); no bound may do that. + assert min(converged) >= 1, ( + "some step bound left the loop unable to converge any mode at all: %s" + % converged) + + +def test_refine_joint_starts_rejects_a_degenerate_step_bound(): + """The rescale divides by max_step, so a zero bound must be refused. + + Under the old per-coordinate clip a zero bound merely pinned that + coordinate; it now produces a non-finite step, so the check is required by + the change rather than decorative. + """ + C_A, C_B = _narrow_time_peak_tables(16.0) + start = np.asarray([[32.0, 0.5, 0.5, 4.0]]) + for bad in ((2.0, 0.5, 0.0, 0.25), (2.0, 0.5, -0.5, 0.25), + (2.0, 0.5, 0.5)): + with pytest.raises(ValueError): + planner.refine_joint_starts_jax(C_A, C_B, start, 1.0, 8.0, + max_step=bad) + + +def test_tier_starts_are_a_distinct_orbit_with_one_shared_gradient_norm(): + """Distinct starts; equal refined gradient norms WITHIN a symmetry orbit. + + The campaign's write-up read "min equals median to ten digits" as evidence + of degenerate starts. It is not: every retained representative receives + every certified group action, the log density is exactly invariant under + them, so an orbit's members must agree to roundoff. Pin both halves, so + that neither the distinctness nor the invariance can regress unnoticed. + """ + C_A, C_B = _narrow_time_peak_tables(256.0) + summary = planner.summarize_uv_norm_table(C_B) + portfolio = planner.rank_joint_starts_from_uvq( + C_A, summary, 1.0, 8.0, angular_oversample=3, max_time_starts=5, + max_starts=48) + assert portfolio.symmetry.certified + assert portfolio.symmetry.group_order > 1 + starts = portfolio.starts + assert len(np.unique(np.round(starts, 12), axis=0)) == len(starts) + refined = tuple(np.asarray(item) for item in + planner.refine_joint_starts_jax( + C_A, summary.C_B, starts, 1.0, 8.0, iterations=18)) + points, values, gradients, hessians, curvatures = refined + order = int(portfolio.symmetry.group_order) + norms = np.linalg.norm(gradients, axis=1) + # group_action is emitted representative-major, so members of one orbit are + # consecutive: [rep0 act0, rep0 act1, ..., rep1 act0, ...] + for base in range(0, len(starts), order): + block = values[base:base + order] + assert np.allclose(block, block[0], rtol=0.0, atol=1.0e-6), ( + "one symmetry orbit disagreed on its log density: %s" % block) + # The half this test is named for. The campaign write-up read "min + # equals median to ten digits" as evidence of degenerate starts; it is + # the exact group invariance, and it is pinned here rather than + # described. Measured spread within a block is 0 to 3.1e-12. + gradient_block = norms[base:base + order] + assert np.max(gradient_block) - np.min(gradient_block) <= 1.0e-9, ( + "one symmetry orbit disagreed on its gradient norm: %s" + % gradient_block) + assert np.all(curvatures[base:base + order] > 0.0) + +_HM_PACKET = "/tmp/hm51_Ctables_incl0.6.npz" +_SNR40_PACKET = ("/tmp/rift-paper-av-ladder/analyses/va_sequence_20260902/" + "records/angle_coeffs_rung40_n256.npz") +_SNR160_PACKET = ("/tmp/rift-paper-av-ladder/analyses/va_sequence_20260902/" + "records/angle_coeffs_rung160_n256.npz") + + +@pytest.mark.skipif(not os.path.exists(_HM_PACKET), + reason="external real-table validation packet is absent") +def test_hm_second_mode_survives_unsafe_proxy_gap(): + """Regression for the real Lmax=4 mode proxy that defeated PR267's line.""" + packet = np.load(_HM_PACKET) + C_A = packet["C_A"] + summary = planner.summarize_uv_norm_table(packet["C_B"]) + portfolio = planner.rank_joint_starts_from_uvq( + C_A, summary, 1000.0 / 720.0, 1000.0 / 240.0, + max_time_starts=3, max_starts=24) + # This is load-bearing: proxy pruning at the nominal -23 nat error budget, + # or even at -32, discards a mode whose refined contribution is relevant. + assert len(portfolio.raw_scores) >= 15 + assert portfolio.raw_scores[14] - portfolio.raw_scores[0] == pytest.approx( + -37.13988596, abs=2.0e-6) + + result = tuple(np.asarray(item) for item in planner.refine_joint_starts_jax( + C_A, summary.C_B, portfolio.starts, 1000.0 / 720.0, + 1000.0 / 240.0, iterations=18)) + points, values, gradients, _, curvatures = result + selected, _ = planner.select_refined_modes( + points, values, gradients, curvatures, max_modes=24) + assert len(selected) == 2 + delta = np.sort(values[selected] - np.max(values[selected])) + np.testing.assert_allclose(delta, [-11.671141934982415, 0.0], + rtol=0.0, atol=2.0e-8) + assert np.max(np.linalg.norm(gradients[selected], axis=1)) < 2.0e-6 + + +@pytest.mark.skipif(not os.path.exists(_HM_PACKET), + reason="external real-table validation packet is absent") +def test_hm_two_tier_integral_matches_overcomplete_oracle(): + packet = np.load(_HM_PACKET) + oracle = 1305.8219235157544 + result = planner.multipeak_local_marginalize( + packet["C_A"], packet["C_B"], 1000.0 / 720.0, 1000.0 / 240.0, + oracle, log_integral_tol=1.0e-3, quadrature_order=7, + cell_sigma=5.0, chunk_size=64) + assert result.accepted + assert not result.used_reserve + assert result.tier0.n_retained_modes == 2 + assert result.tier1.n_retained_modes == 2 + assert abs(result.value - oracle) < 1.0e-3 + assert result.modeled_peak_bytes < 32 * 1024 ** 2 + + +@pytest.mark.skipif(not os.path.exists(_SNR40_PACKET), + reason="external real-table validation packet is absent") +def test_real_low_snr_declines_to_finite_reserve(): + packet = np.load(_SNR40_PACKET) + C_A = packet["C_A"][:, :, 148, :] + C_B = packet["C_B"][:, :, 148, :] + oracle = 814.7510954543737 + result = planner.multipeak_local_marginalize( + C_A, C_B, 0.2, 7.0, oracle, log_integral_tol=1.0e-3, + quadrature_order=7, cell_sigma=5.0, chunk_size=64) + assert not result.accepted + assert result.used_reserve + assert result.value == oracle + assert result.provenance.startswith("dense-reserve:") + + +@pytest.mark.skipif(not os.path.exists(_SNR160_PACKET), + reason="external real-table validation packet is absent") +def test_real_high_snr_two_tier_path_matches_overcomplete_oracle(): + packet = np.load(_SNR160_PACKET) + C_A = packet["C_A"][:, :, 148, :] + C_B = packet["C_B"][:, :, 148, :] + oracle = 13255.018541583624 + result = planner.multipeak_local_marginalize( + C_A, C_B, 0.2, 7.0, oracle, log_integral_tol=1.0e-3, + quadrature_order=7, cell_sigma=5.0, chunk_size=64) + assert result.accepted + assert not result.used_reserve + assert result.tier0.n_retained_modes == 4 + assert result.tier1.n_retained_modes == 4 + assert result.delta_log_integral < 1.0e-3 + assert abs(result.value - oracle) < 1.0e-3 diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_policy_peaklocal_reserve.py b/MonteCarloMarginalizeCode/Code/test/jax/test_policy_peaklocal_reserve.py new file mode 100644 index 000000000..cba0317a2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_policy_peaklocal_reserve.py @@ -0,0 +1,567 @@ +"""The four-axis policy's peak-local time reserve, as wired. + +The rule (one commensurate lattice, fixed shape, coarser check, zero-weight +repeats), the width prediction from the row's tables, the scheme pair table +and its refusals, the angular-kernel seam, agreement of the peak-local reserve +with the whole-window refined reserve and with an analytic fine reference on +two synthetic tables (the policy tests' three-harmonic table, and a +Gaussian-envelope carrier whose peak is 0.16 samples wide), the escalation +report when the width cannot be predicted, fail-closed behaviour when the +escalations are exhausted, and the driver options. Evidence: +DESIGN_direct_marginalization_policy.md, "Peak-local time reserve". + +TEST-TIMING: five policy evaluations on 33-sample tables; the carrier +fixture's whole-window reserve escalates to refine 32. +""" +import types + +import numpy as np +import pytest + +jax = pytest.importorskip("jax") +import jax.numpy as jnp # noqa: E402 + +from RIFT.likelihood.jax_ile import anglemarg as AM # noqa: E402 +from RIFT.likelihood.jax_ile import core as _core # noqa: E402 +from RIFT.likelihood.jax_ile import direct_marginalization_policy as DP # noqa: E402 +from RIFT.likelihood.jax_ile import peaklocal_time_reserve as PLR # noqa: E402 +from test_direct_marginalization_policy import ( # noqa: E402 + _guarded_problem, _fake_data, _install_tables, _fine_reference, _grid, + _GUARD, _N) +from test_angle_marg_exact import INTERP # noqa: E402 + +AMP_SIZING = 40.0 +# The carrier drives the phi exponent to x A = 160, so its exact kernel is +# sized for that: an undersized phi grid leaves a ripple of period +# 1/(n_phi f_c) in lnL(t) that a coarser check rule aliases (measured +# 1.85e-3 nat at amp_sizing 40). +AMP_SIZING_CARRIER = 200.0 + + +# ------------------------------------------------------------------ the rule +def _rule(centres, widths, live, sigma_t, n_target=40, n_fine=9, scan_refine=2, + dt_scale=0.5, mult=1): + return PLR.peaklocal_time_rule(jnp.asarray(centres, dtype=jnp.float64), + jnp.asarray(widths, dtype=jnp.float64), + jnp.asarray(live, dtype=bool), n_target, + dt_scale, sigma_t_samples=sigma_t, + n_fine=n_fine, scan_refine=scan_refine, + fine_refine_multiplier=mult) + + +def test_rule_is_one_commensurate_lattice_of_fixed_shape(): + n_target, n_fine, scan = 40, 9, 2 + expected = PLR.peaklocal_rule_size(n_target, 3, n_fine, scan) + narrow = _rule([10.3, 25.0, 3.0], [0.1, 0.4, 2.0], [True, True, False], + sigma_t=0.1, n_target=n_target, n_fine=n_fine, scan_refine=scan) + wide = _rule([2.0, 30.0, 36.0], [3.0, 0.05, 1.0], [True, False, True], + sigma_t=4.0, n_target=n_target, n_fine=n_fine, scan_refine=scan) + # sigma 0.1 at 3 nodes per sigma against a 0.5 scan needs m = 15; the + # wide prediction (sigma 4) is capped by its narrowest live maximum + # (width 1.0), which needs m = 2. + assert int(narrow["fine_refine"]) == 15 and int(wide["fine_refine"]) == 2 + for r in (narrow, wide): + assert (int(r["nodes"].size), int(r["check_nodes"].size)) == expected + nodes, check = np.asarray(r["nodes"]), np.asarray(r["check_nodes"]) + h = float(r["fine_spacing_samples"]) + assert np.allclose(nodes / h, np.round(nodes / h), atol=1e-9) + assert np.allclose(check / h, np.round(check / h), atol=1e-9) + assert np.all(np.diff(nodes) >= 0.0) and np.all(np.diff(check) >= 0.0) + assert nodes[0] == 0.0 and nodes[-1] == n_target - 1 + assert check[0] == nodes[0] and check[-1] == nodes[-1] + assert np.max(np.diff(check)) > np.max(np.diff(nodes)) + assert np.max(np.diff(nodes)) < 1.0 + w, wc = np.asarray(r["weights"]), np.asarray(r["check_weights"]) + assert np.all(w >= 0.0) and np.all(wc >= 0.0) + # Trapezoid weights telescope: the rule integrates the whole window + # whatever the block positions. + assert np.isclose(w.sum(), (n_target - 1) * 0.5) + assert np.isclose(wc.sum(), (n_target - 1) * 0.5) + assert int(r["n_live_blocks"]) == 2 + # The narrow rule's live blocks span (n_fine - 1) h around their centres + # and no more; the dead block's nodes are scan nodes (zero extra weight). + nodes = np.asarray(narrow["nodes"]) + h = float(narrow["fine_spacing_samples"]) + assert float(narrow["block_span_samples"]) == pytest.approx((n_fine - 1) * h) + inside = nodes[(nodes > 10.3 - 0.5) & (nodes < 10.3 + 0.5)] + assert len(inside) >= n_fine + assert float(narrow["sigma_t_located_samples"]) == pytest.approx(0.1) + assert float(narrow["sigma_t_used_samples"]) == pytest.approx(0.1) + assert float(narrow["first_block_centre_samples"]) == pytest.approx(10.3) + # The lattice is never coarser than the narrowest located maximum. + coarse_pred = _rule([10.3, 25.0, 3.0], [0.02, 0.4, 2.0], [True, True, False], + sigma_t=0.1, n_target=n_target, n_fine=n_fine, scan_refine=scan) + assert int(coarse_pred["fine_refine"]) == 75 + # A tier doubles the lattice at fixed span. + tier = _rule([10.3, 25.0, 3.0], [0.1, 0.4, 2.0], [True, True, False], + sigma_t=0.1, n_target=n_target, n_fine=2 * n_fine - 1, + scan_refine=scan, mult=2) + assert int(tier["fine_refine"]) == 30 + assert float(tier["block_span_samples"]) == pytest.approx( + float(narrow["block_span_samples"])) + + +def test_rule_refuses_what_its_check_rule_cannot_express(): + with pytest.raises(ValueError): + PLR.validate_peaklocal_rule_arguments(40, 8, 2) # even fine count + with pytest.raises(ValueError): + PLR.validate_peaklocal_rule_arguments(40, 9, 3) # odd scan + with pytest.raises(ValueError): + _rule([1.0, 2.0], [1.0, 1.0], [True, True], sigma_t=1.0, mult=0) + + +def test_an_unpredictable_width_degenerates_to_the_scan(): + r = _rule([10.0, 20.0], [np.inf, np.inf], [False, False], sigma_t=float("nan")) + assert int(r["fine_refine"]) == 1 and not bool(r["prediction_finite"]) + nodes = np.asarray(r["nodes"]) + assert np.allclose(nodes * 2.0, np.round(nodes * 2.0)) + + +# --------------------------------------------------------- the prediction +def test_predict_time_rule_compares_the_two_node_counts(): + narrow = PLR.predict_time_rule(163.0, 0.01557, 614, n_blocks=8, n_fine=49, + scan_refine=2) + assert narrow["sigma_t_samples"] == pytest.approx( + 1.0 / (2 * np.pi * 163.0 * 0.01557)) + assert narrow["whole_window_nodes_needed"] > 2 * narrow["peaklocal_nodes"] + assert narrow["prefer_peaklocal"] + broad = PLR.predict_time_rule(8.0, 0.01557, 614, n_blocks=8, n_fine=49, + scan_refine=2) + assert not broad["prefer_peaklocal"] + assert PLR.predict_time_rule(float("nan"), 0.01, 100, n_blocks=2, n_fine=9, + scan_refine=2)["fine_refine"] == 1 + + +# ---------------------------------------------------------- the carrier +# f_c 0.1, not 0.2: the Gaussian envelope (tau 2) must be band-limited on the +# stored grid, and at 0.2 its amplitude 8e-4 past Nyquist aliased the +# reflected primitive 0.018 nat away from the analytic table (measured); at +# 0.1 the leak is 3e-6. +_CARRIER = dict(amp=20.0, B=10.0, tau=2.0, f_c=0.1, t0=16.37) + + +def _carrier_table(n, guard, t): + c = _CARRIER + C_A = np.zeros((3, 3, t.size), dtype=np.complex128) + C_A[2, 1] = (c["amp"] * np.exp(-0.5 * ((t - c["t0"]) / c["tau"]) ** 2) + * np.exp(2j * np.pi * c["f_c"] * (t - c["t0"]))) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_B[0, 2] = c["B"] + return C_A, C_B + + +def _carrier_problem(n, guard): + """A (2, +-2)-like carrier with a Gaussian envelope: the angular maximum + is 2 amp g(t), so lnL(t) peaks at 2 amp^2 / B with rho^2 = 4 amp^2 / B + and exp(lnL) has width tau / rho, here 0.159 samples.""" + support = np.arange(-guard, n + guard, dtype=float) + return _carrier_table(n, guard, support) + + +def _carrier_fine_reference(data, x_grid, log_w, refine=32): + n = int(data.npts) + t = np.arange((n - 1) * refine + 1, dtype=float) / float(refine) + C_A, C_B = _carrier_table(n, 0, t) + lnL_t = AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, amp_sizing=AMP_SIZING_CARRIER, dense_chunk=8, + grid_block=32) + w = jnp.asarray(_core._simpson_weights(t.size, data.deltaT / refine)) + return float(_core._time_marginalize(lnL_t, w)[0]) + + +def _carrier_sigma_t(): + """Width of the phi-marginalized peak (the envelope), and rho.""" + c = _CARRIER + rho = 2.0 * c["amp"] / np.sqrt(c["B"]) + return c["tau"] / rho, rho + + +def _carrier_fixed_angle_sigma(): + """Width from the curvature of the FIXED-angle field at its maximum, + which carries the carrier: ``1 / (rho sqrt(omega^2 + 1/tau^2))``. This + is what the locator measures and the narrowest structure the primitive + has; the phi-marginalized peak is wider by the same ratio the raw and + central moments differ by.""" + c = _CARRIER + _, rho = _carrier_sigma_t() + omega = 2.0 * np.pi * c["f_c"] + return 1.0 / (rho * np.sqrt(omega ** 2 + 1.0 / c["tau"] ** 2)) + + +def test_the_prediction_reads_rho_and_the_bandwidth_off_the_table(): + C_A, C_B = _carrier_problem(_N, _GUARD) + sigma_t_true, rho_true = _carrier_sigma_t() + rho = float(PLR.row_amplitude(C_A, C_B, _GUARD)) + sigma_f = float(PLR.table_bandwidth_cycles(C_A, _GUARD)) + # rho reads the envelope at the stored samples; the peak sits 0.37 of a + # sample off the lattice, so the triangle bound is 2% low here. + assert rho == pytest.approx(rho_true, rel=0.05) + # Raw two-sided rms frequency: the carrier plus the Gaussian envelope's + # power spread 1 / (2 sqrt2 pi tau), in quadrature. + expected_f = np.hypot(_CARRIER["f_c"], 1.0 / (2 * np.sqrt(2) * np.pi * _CARRIER["tau"])) + assert sigma_f == pytest.approx(expected_f, rel=0.05) + sigma_t = float(PLR.predicted_width_samples(rho, sigma_f)) + # The prediction is the narrowest peak the primitive can make: below the + # envelope's true width here, and within the factor-four band. + assert sigma_t_true / 4.0 < sigma_t < sigma_t_true + # The three-harmonic table's only frequency is its cosine, 1/32 cycles + # per sample diluted by the constant lanes; the raw moment reads it and + # predicts a width within the consistency band of the measured 0.75 + # samples (the prediction is the narrowest peak, so it sits below). + C_A3, C_B3, _ = _guarded_problem(_N, _GUARD) + sigma_f3 = float(PLR.table_bandwidth_cycles(C_A3, _GUARD)) + assert 0.005 < sigma_f3 < 1.0 / 32.0 + sigma_t3 = float(PLR.predicted_width_samples( + float(PLR.row_amplitude(C_A3, C_B3, _GUARD)), sigma_f3)) + assert 0.75 / 4.0 < sigma_t3 <= 0.75 * 1.05 + assert not np.isfinite(float(PLR.predicted_width_samples(10.0, 0.0))) + + +# ------------------------------------------------------- config and the seam +def test_scheme_pairs_and_defaults_leave_the_shipped_reserve_alone(): + assert DP.PolicyConfig().reserve_scheme == "exact" + assert DP.PolicyConfig().max_time_nodes >= 2 + assert DP.reserve_pair("exact") == ("exact", "window") + assert DP.reserve_pair("laplace") == ("laplace", "window") + assert DP.reserve_pair("peaklocal") == ("laplace", "peaklocal") + assert DP.reserve_pair("peaklocal-exact") == ("exact", "peaklocal") + assert "peaklocal-exact" not in DP.RESERVE_SCHEME_CHOICES + with pytest.raises(ValueError) as err: + DP.reserve_pair("auto") + assert "predict_reserve_pair" in str(err.value) + with pytest.raises(ValueError): + DP.reserve_pair("peak-local") + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="cover")) + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig( + reserve_scheme="peaklocal-exact", reserve_peaklocal_fine_nodes=48)) + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig( + reserve_scheme="peaklocal-exact", + reserve_peaklocal_sigma_t_override_samples=0.0)) + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig(max_time_nodes=1)) + + +def test_laplace_kernel_is_resolved_from_anglemarg_or_refused(monkeypatch): + x_grid, log_w = _grid(16) + assert DP.resolve_reserve_angular_kernel( + "exact", x_grid, log_w, amp_sizing=AMP_SIZING, m_max=2, dense_chunk=8, + grid_block=32) is None + monkeypatch.delattr(AM, DP._LAPLACE_TABLE_KERNEL, raising=False) + with pytest.raises(ValueError) as err: + DP.resolve_reserve_angular_kernel( + "laplace", x_grid, log_w, amp_sizing=AMP_SIZING, m_max=2, + dense_chunk=8, grid_block=32) + assert DP._LAPLACE_TABLE_KERNEL in str(err.value) + with pytest.raises(ValueError): + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="peaklocal")) + seen = {} + + def fake(table, norm_table, xg, lw, *, amp_sizing, m_max): + # The Laplace kernel's signature: no dense_chunk/grid_block (those are + # the exact kernel's); a fake that accepted them hid the seam until + # the first real rung-652 evaluation (2026-09-09). + seen.update(amp_sizing=amp_sizing, m_max=m_max, n=int(xg.shape[0])) + return jnp.zeros(table.shape[-1:]) + monkeypatch.setattr(AM, DP._LAPLACE_TABLE_KERNEL, fake, raising=False) + kernel = DP.resolve_reserve_angular_kernel( + "laplace", x_grid, log_w, amp_sizing=AMP_SIZING, m_max=2, + dense_chunk=8, grid_block=32) + out = kernel(jnp.zeros((3, 3, 7), dtype=jnp.complex128), + jnp.zeros((5, 5), dtype=jnp.complex128)) + assert out.shape == (7,) + assert seen == dict(amp_sizing=AMP_SIZING, m_max=2, n=16) + + +def test_laplace_kernel_runs_through_the_real_anglemarg_function(): + """The resolved Laplace kernel must call anglemarg's REAL function with + keywords it accepts: shape (npts,), finite, and within the psi-Laplace + approximation of the exact kernel on the carrier tables.""" + x_grid, log_w = _grid(16) + n, guard = 8, 4 + C_A, C_B = _carrier_problem(n, guard) + kernel = DP.resolve_reserve_angular_kernel( + "laplace", x_grid, log_w, amp_sizing=AMP_SIZING_CARRIER, m_max=2, + dense_chunk=8, grid_block=32) + lap = np.asarray(kernel(jnp.asarray(C_A), jnp.asarray(C_B))) + ex = np.asarray(AM.coefficient_table_distphipsimarg_exact( + C_A, C_B, x_grid, log_w, amp_sizing=AMP_SIZING_CARRIER, m_max=2, + dense_chunk=8, grid_block=32)) + # Both kernels return the batched (B, npts) contract, B = 1 here. + assert lap.shape == ex.shape == (1, C_A.shape[-1]) + assert np.all(np.isfinite(lap)) + lap, ex = lap[0], ex[0] + peak = int(np.argmax(ex)) + assert abs(float(lap[peak] - ex[peak])) < 0.5 * abs(float(ex[peak])) + assert abs(int(np.argmax(lap)) - peak) <= 1 + + +def test_locator_search_phi_grid_is_sized_for_the_amplitude(): + """At rho ~ 630 the 64-node phi grid's ripple, about rho^2 (pi/64)^2 = + 975 nat, exceeds the profile's change across one search cell, here + (rho^2/2)(0.125/tau)^2 = 49 nat with tau = 8, so the search maximum lands + on the wrong cell and the polish cannot reach the peak (measured on the + rung-652 production row, 2026-09-09). With the policy's 4096 nodes the + ripple is 0.24 nat and the locator lands within the marginal peak's + width tau / rho of the carrier's centre, at the profile's true maximum + rho^2 / 2.""" + n, guard = 40, 8 + c = dict(_CARRIER, amp=1000.0, tau=8.0, t0=19.37) + support = np.arange(-guard, n + guard, dtype=float) + C_A = np.zeros((3, 3, support.size), dtype=np.complex128) + C_A[2, 1] = (c["amp"] * np.exp(-0.5 * ((support - c["t0"]) / c["tau"]) ** 2) + * np.exp(2j * np.pi * c["f_c"] * (support - c["t0"]))) + C_B = np.zeros((5, 5), dtype=np.complex128) + C_B[0, 2] = c["B"] + rho = 2.0 * c["amp"] / np.sqrt(c["B"]) + width = c["tau"] / rho + x_lo, x_hi = 0.0, 1.0e6 + kw = dict(n_candidates=2, search_refine=8, angular_lattice=8, + newton_steps=8, newton_step_max=1.0) + sized = PLR.locate_time_maxima(C_A, C_B, guard, n, x_lo, x_hi, n_phi=4096, **kw) + k = int(np.argmax(np.where(np.asarray(sized["live"]), np.asarray(sized["values"]), -np.inf))) + assert abs(float(sized["centres"][k]) - c["t0"]) < width + assert abs(float(sized["values"][k]) - 0.5 * rho ** 2) < 1.0 + # The 64-node grid's failure is not pinned here: on this one-harmonic + # carrier the ripple's phase can favour the right cell by chance. The + # production row is the evidence (DESIGN, rung 652 row 0). + cfg = DP.PolicyConfig(reserve_scheme="peaklocal-exact") + assert int(cfg.reserve_peaklocal_search_phi_nodes) == 4096 + assert rho ** 2 * (np.pi / cfg.reserve_peaklocal_search_phi_nodes) ** 2 < 1.0 + + +# ------------------------------------------------ agreement on the reserve +def _policy(monkeypatch, tables, cfg, amp_sizing=AMP_SIZING): + monkeypatch.setattr(_core, "_DISTMARG_GH_N", 0) + _install_tables(monkeypatch, tables, _GUARD) + data = _fake_data(_N) + x_grid, log_w = _grid() + value, ledger = DP.fused_log_likelihood_four_axis_policy( + data, jnp.zeros(1), jnp.zeros(1), jnp.zeros(1), x_grid, log_w, + interp=INTERP, amp_sizing=amp_sizing, config=cfg, return_ledger=True) + L = {k: np.asarray(v)[0] for k, v in ledger.items()} + return float(value[0]), L, data, x_grid, log_w + + +_WINDOW = DP.PolicyConfig(time_guard=_GUARD, reserve_time_refine=4, + max_modes=1, enriched_max_modes=1) +_PEAK = _WINDOW._replace(reserve_scheme="peaklocal-exact") + + +def _assert_warranted_peaklocal(L, cfg): + assert not L["accepted_local"] + assert L["reserve_executed"] and L["reserve_time_rule_peaklocal"] + assert L["reserve_time_check_rule_internal"] + assert L["reserve_time_resolution_validated"] + assert L["reserve_time_guard_validated"] + assert L["reserve_time_warranted"] and L["usable"] and L["reconciles"] + assert int(L["reserve_time_refine_used"]) == 0 + assert int(L["reserve_peaklocal_live_blocks"]) >= 1 + expected, _ = PLR.peaklocal_rule_size( + _N, cfg.reserve_peaklocal_blocks, + int(L["reserve_peaklocal_fine_nodes_used"]), + cfg.reserve_peaklocal_scan_refine, n_scan=cfg.reserve_peaklocal_scan_nodes) + assert int(L["reserve_time_points"]) == expected + # Support-limited: the scan spans the hull of the live maxima plus the + # printed margin, not the window, and the omitted mass is bounded and + # charged through the cropped-cover warrant. + assert L["reserve_time_cropped_cover_warranted"] + assert L["reserve_time_tail_ok"] + assert np.isfinite(float(L["reserve_peaklocal_outside_log_bound"])) + assert float(L["reserve_peaklocal_scan_lo_samples"]) >= 0.0 + assert float(L["reserve_peaklocal_scan_hi_samples"]) <= _N - 1 + + +def test_carrier_peak_is_predicted_sized_and_matches_the_references(monkeypatch): + """The narrow carrier: the whole-window rule needs refine 32 to resolve a + 0.16-sample peak, the peak-local rule resolves it at its FIRST tier + (no escalation), the prediction agrees with the plan's Newton width, and + both reserves match the analytic fine reference.""" + tables = ([_carrier_problem(_N, _GUARD)[0]], _carrier_problem(_N, _GUARD)[1]) + v_w, L_w, data, x_grid, log_w = _policy(monkeypatch, tables, _WINDOW, + amp_sizing=AMP_SIZING_CARRIER) + v_p, L_p, _, _, _ = _policy(monkeypatch, tables, _PEAK, + amp_sizing=AMP_SIZING_CARRIER) + fine = _carrier_fine_reference(data, x_grid, log_w) + sigma_t_true, rho_true = _carrier_sigma_t() + assert L_w["reserve_time_warranted"] and not L_w["reserve_time_rule_peaklocal"] + assert int(L_w["reserve_time_refine_used"]) >= 8 + _assert_warranted_peaklocal(L_p, _PEAK) + assert int(L_p["reserve_escalations"]) == 0 + assert int(L_p["reserve_peaklocal_fine_nodes_used"]) == _PEAK.reserve_peaklocal_fine_nodes + assert L_p["reserve_peaklocal_prediction_finite"] + assert L_p["reserve_peaklocal_prediction_consistent"] + assert float(L_p["reserve_peaklocal_rho_pred"]) == pytest.approx(rho_true, rel=0.02) + assert float(L_p["reserve_peaklocal_rho_bound"]) >= 0.95 * rho_true + assert L_p["reserve_time_focus_ok"] + assert float(L_p["reserve_time_focus_offset_samples"]) < 0.05 + pred = float(L_p["reserve_peaklocal_sigma_t_pred_samples"]) + assert sigma_t_true / 4.0 < pred < sigma_t_true + # The locator measures the ENVELOPE's curvature (what the marginalized + # reserve integrates), so its width is the phi-marginalized peak's, wider + # than the raw-moment prediction by the carrier-to-envelope ratio. + located = float(L_p["reserve_peaklocal_sigma_t_located_samples"]) + assert located == pytest.approx(sigma_t_true, rel=0.05) + assert float(L_p["reserve_peaklocal_sigma_t_used_samples"]) == pytest.approx( + min(pred, located)) + assert not np.isfinite(float(L_p["reserve_peaklocal_sigma_f_q_cycles"])) + assert int(L_p["reserve_peaklocal_fine_refine"]) >= 8 + assert abs(float(L_p["reserve_peaklocal_first_block_centre_samples"]) + - _CARRIER["t0"]) < 1e-3 + # On this 33-sample toy window the whole-window rule at refine 8 is + # cheaper than the fixed 357; the advantage is a production-window + # statement, made through the same predictor the selector reads. + expected_pl, _ = PLR.peaklocal_rule_size(_N, _PEAK.reserve_peaklocal_blocks, + _PEAK.reserve_peaklocal_fine_nodes, + _PEAK.reserve_peaklocal_scan_refine, + n_scan=_PEAK.reserve_peaklocal_scan_nodes) + assert int(L_p["reserve_time_points"]) == expected_pl + production = PLR.predict_time_rule( + rho_true, float(L_p["reserve_peaklocal_sigma_f_table_cycles"]), 614, + n_blocks=_PEAK.reserve_peaklocal_blocks, + n_fine=_PEAK.reserve_peaklocal_fine_nodes, + scan_refine=_PEAK.reserve_peaklocal_scan_refine, + n_scan=_PEAK.reserve_peaklocal_scan_nodes) + assert production["prefer_peaklocal"] + # The count does not grow with the window: 33 samples or 614. + assert production["peaklocal_nodes"] == expected_pl + # The scan hull contains the carrier's maximum and is narrower than the + # window by the margin rule (16 sigma each side of one maximum). + assert float(L_p["reserve_peaklocal_scan_lo_samples"]) < _CARRIER["t0"] < float( + L_p["reserve_peaklocal_scan_hi_samples"]) + assert (float(L_p["reserve_peaklocal_scan_hi_samples"]) + - float(L_p["reserve_peaklocal_scan_lo_samples"])) < _N - 1 + assert abs(v_p - v_w) <= float(_PEAK.total_value_error_budget_nats), (v_p, v_w) + assert abs(v_p - fine) <= 2.0e-3, (v_p, fine) + assert abs(v_w - fine) <= 2.0e-3, (v_w, fine) + + +def test_the_locator_finds_the_carrier_maximum_from_the_primitive(): + """No plan involved: the search grid, the dense phi maximization and the + parabolic polish put the first block on the analytic maximum with the + analytic envelope width, and the far candidates are dead.""" + C_A, C_B = _carrier_problem(_N, _GUARD) + x_grid, _ = _grid(64) + found = PLR.locate_time_maxima( + C_A, C_B, _GUARD, _N, float(np.min(x_grid)), float(np.max(x_grid)), + n_candidates=4, search_refine=8, angular_lattice=8) + sigma_t_true, rho_true = _carrier_sigma_t() + live = np.asarray(found["live"]) + centres = np.asarray(found["centres"]) + widths = np.asarray(found["widths"]) + assert live[0] + assert abs(centres[0] - _CARRIER["t0"]) < 2e-3 + assert widths[0] == pytest.approx(sigma_t_true, rel=0.05) + # The fixed-angle field's own curvature width is narrower by the carrier + # ratio; the envelope is what the marginalized reserve integrates. + assert widths[0] > 1.5 * _carrier_fixed_angle_sigma() + assert float(np.asarray(found["values"])[0]) == pytest.approx( + 0.5 * rho_true ** 2, rel=0.05) + assert int(found["n_search_nodes"]) == (_N - 1) * 8 + 1 + # rho from the profile maximum, and the angles it was found at. + assert float(found["rho_located"]) == pytest.approx(rho_true, rel=0.02) + assert np.isfinite(float(found["phi"][0])) and np.isfinite(float(found["u"][0])) + + +def test_a_misplaced_block_fails_the_focus_certificate(monkeypatch): + """The rule and its check share the block, so a block off the maximum + agrees with itself (measured 0.22 nat, warranted, on a rung-160 row). + The kernel's focus certificate compares the evaluated argmax with the + block centre; here the locator is forced 1.2 samples off.""" + real = PLR.locate_time_maxima + + def shifted(*args, **kwargs): + found = real(*args, **kwargs) + found = dict(found) + found["centres"] = found["centres"] + 1.2 + return found + monkeypatch.setattr(PLR, "locate_time_maxima", shifted) + tables = ([_carrier_problem(_N, _GUARD)[0]], _carrier_problem(_N, _GUARD)[1]) + v, L, _, _, _ = _policy(monkeypatch, tables, + _PEAK._replace(reserve_peaklocal_escalations=0), + amp_sizing=AMP_SIZING_CARRIER) + assert L["reserve_executed"] and not L["reserve_time_focus_ok"] + assert float(L["reserve_time_focus_offset_samples"]) > float( + L["reserve_peaklocal_focus_half_width_samples"]) + assert not L["reserve_time_warranted"] and not L["usable"] + assert not np.isfinite(v) + + +@pytest.mark.parametrize("scale", [1.0, 3.0]) +def test_three_harmonic_table_is_sized_from_its_cosine_and_needs_no_escalation( + monkeypatch, scale): + """The policy tests' table at two amplitudes (peak 0.75 and 0.24 samples + wide). The raw moment of its cosine predicts a width inside the + consistency band, the first tier resolves the peak (no escalation), and + the value matches the window reserve and the analytic reference.""" + guarded, C_B, constants = _guarded_problem(_N, _GUARD, scale=scale) + v_w, L_w, data, x_grid, log_w = _policy(monkeypatch, ([guarded], C_B), _WINDOW) + v_p, L_p, _, _, _ = _policy(monkeypatch, ([guarded], C_B), _PEAK) + fine = _fine_reference(constants, C_B, data, x_grid, log_w, AMP_SIZING, + scale=scale) + _assert_warranted_peaklocal(L_p, _PEAK) + assert L_p["reserve_peaklocal_prediction_finite"] + assert L_p["reserve_peaklocal_prediction_consistent"] + assert int(L_p["reserve_escalations"]) == 0 + assert int(L_p["reserve_peaklocal_fine_refine"]) >= (2 if scale > 1 else 1) + assert abs(v_p - v_w) <= float(_PEAK.total_value_error_budget_nats), (v_p, v_w) + assert abs(v_p - fine) <= 2.0e-3, (v_p, fine) + + +def test_exhausted_escalations_fail_closed_and_keep_the_diagnostic(monkeypatch): + """The carrier with the width overridden 20x too wide: the first tier + cannot resolve the peak and with no escalation allowed the row is + unusable, its finite diagnostic kept in the ledger and never selected.""" + tables = ([_carrier_problem(_N, _GUARD)[0]], _carrier_problem(_N, _GUARD)[1]) + sigma_t_true, _ = _carrier_sigma_t() + # Override the prediction 40x too wide AND blind the locator (one + # candidate on a native-sample search grid, no Newton polish would still + # find it, so its width is also overridden by the coarse lattice): the + # first tier is then the scan alone. + cfg = _PEAK._replace(reserve_peaklocal_escalations=0, + reserve_peaklocal_sigma_t_override_samples=40.0 * sigma_t_true, + reserve_peaklocal_blocks=1, reserve_peaklocal_fine_nodes=3) + v, L, _, _, _ = _policy(monkeypatch, tables, cfg, amp_sizing=AMP_SIZING_CARRIER) + assert L["reserve_executed"] and not L["reserve_time_warranted"] + assert not L["usable"] and not np.isfinite(v) + assert int(L["reserve_escalations"]) == 0 + assert not L["reserve_peaklocal_prediction_consistent"] + assert np.isfinite(float(L["selected_value"])) + + +# ------------------------------------------------------------------ driver +def test_driver_offers_the_reserve_scheme_and_the_time_node_knobs(): + import importlib.machinery + import importlib.util + import pathlib + path = pathlib.Path(__file__).parents[2] / "bin" / "integrate_likelihood_extrinsic_jax" + loader = importlib.machinery.SourceFileLoader("_plr_driver", str(path)) + spec = importlib.util.spec_from_loader(loader.name, loader) + drv = importlib.util.module_from_spec(spec) + loader.exec_module(drv) + optp = drv.build_parser() + opts, _ = optp.parse_args(["--direct-marginalization-reserve-scheme", "peaklocal", + "--direct-marginalization-max-time-nodes", "256", + "--direct-marginalization-peaklocal-escalations", "1"]) + assert opts.direct_marginalization_reserve_scheme == "peaklocal" + assert opts.direct_marginalization_max_time_nodes == 256 + assert opts.direct_marginalization_peaklocal_escalations == 1 + opts, _ = optp.parse_args([]) + assert opts.direct_marginalization_reserve_scheme == "exact" + # One definition: the flag's default is the config's, whatever it is. + assert opts.direct_marginalization_max_time_nodes == DP.PolicyConfig().max_time_nodes + assert (opts.direct_marginalization_peaklocal_escalations + == DP.PolicyConfig().reserve_peaklocal_escalations) + with pytest.raises(SystemExit): + optp.parse_args(["--direct-marginalization-reserve-scheme", "peak-local"]) + with pytest.raises(SystemExit): + optp.parse_args(["--direct-marginalization-reserve-scheme", "peaklocal-exact"]) + # Each flag is added once (a merge once left two add_option calls and + # optparse silently took the last). + src = path.read_text() + for flag in ("--direct-marginalization-reserve-scheme", + "--direct-marginalization-max-time-nodes", + "--direct-marginalization-peaklocal-escalations"): + assert src.count('add_option("%s"' % flag) == 1, flag diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_reserve_pair_selection.py b/MonteCarloMarginalizeCode/Code/test/jax/test_reserve_pair_selection.py new file mode 100644 index 000000000..4c470bb61 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_reserve_pair_selection.py @@ -0,0 +1,378 @@ +"""The (local, reserve) pair is CHOSEN from analysis, before any row is evaluated. + +RO, 2026-09-08: rely on analysis and the known physics to pick the pair, rather +than try-then-decline-then-refine. Everything the selector uses is computable +from the precomputed inputs, so the choice and its reasons are printed in the +run's first lines instead of being discovered from a ledger at the end. + +The tests below pin the RULE, not a run. The most important one is +`test_the_predictor_reproduces_the_measured_refine4_failure`: an independent +session MEASURED the whole-window reserve failing its own convergence warrant at +the lowest rung, and the predictor says the same thing from the physics alone. +""" + +import os + +import numpy as np +import pytest + +from RIFT.likelihood.jax_ile import direct_marginalization_policy as DP +from RIFT.likelihood.jax_ile.anglemarg import ANGLE_MARG_CROSSOVER_AMPLITUDE + +from test_angle_marg_exact import make_synth + +_DRIVER = os.path.join( + os.path.dirname(os.path.dirname(os.path.dirname( + os.path.abspath(__file__)))), "bin", + "integrate_likelihood_extrinsic_jax") + +# Ladder-2 network SNRs. +RHO_40, RHO_160, RHO_320, RHO_640 = 40.7691, 163.0766, 326.1531, 652.3062 + + +@pytest.fixture(scope="module") +def data(): + return make_synth(scale=1.0, npts=614) + + +def _pick(data, rho, refine_max=32, **kw): + kw.setdefault("max_time_nodes", 64) + return DP.predict_reserve_pair( + data, rho, reserve_time_refine_max=refine_max, + crossover_amplitude=ANGLE_MARG_CROSSOVER_AMPLITUDE, **kw) + + +def test_the_bandwidth_comes_from_the_complex_Q_spectrum(data): + """Q (rholm) is complex. A real-input transform rejects it outright, and + taking only the real part would discard half the phase structure, so this + pins that the second moment is finite, positive and Nyquist-bounded.""" + sigma_f = DP.q_effective_bandwidth_hz(data) + assert np.isfinite(sigma_f) and sigma_f > 0.0 + f_nyq = 0.5 / (data.deltaT / data.q_time_pregrid_factor) + assert sigma_f < f_nyq + + +def test_the_angle_scheme_follows_the_validated_crossover(data): + """A = rho^2/2 against ANGLE_MARG_CROSSOVER_AMPLITUDE, which is a measured + accuracy crossover in anglemarg, not a tuning constant introduced here.""" + lo, _ = _pick(data, 1.0) # A = 0.5, far below 450 + assert lo == "exact" + hi, info = _pick(data, RHO_40) # A = 831, above 450 + assert hi == "laplace" + assert info["amplitude_A"] > info["crossover_amplitude"] + + +def test_a_peak_the_ceiling_cannot_resolve_is_refused_not_refined(data): + """The whole point. When the escalation ceiling cannot resolve the peak and + no peak-local time reserve is available, the selector returns None so the + caller REFUSES. Silently falling back to whole-window refinement is the + failure mode this exists to prevent.""" + scheme, info = _pick(data, RHO_640) + assert scheme is None + assert not info["time_peak_resolvable_whole_window"] + assert info["whole_window_nodes_needed"] > info["whole_window_nodes_available"] + assert "NOT IMPLEMENTED" in info["reason"] + + +def test_the_same_signal_selects_peaklocal_once_it_exists(data): + """The refusal is about availability, not about the signal. With a + peak-local time reserve on the menu the identical inputs select it, so the + day that kernel lands the selector starts choosing it with no rule change.""" + scheme, info = _pick(data, RHO_640, available=("exact", "laplace", "peaklocal")) + assert scheme == "peaklocal" + assert "peak-local" in info["reason"] + + +def test_the_predictor_reproduces_the_measured_refine4_failure(data): + """Independent confirmation, and the reason to trust the rule at all. + + The reserve-scheme session MEASURED the whole-window refined reserve at + refine=4 failing its OWN half-refined convergence warrant at the LOWEST + ladder rung, rho 40.77: 0.0018, 0.0040 and 0.0114 nats on three rows + against a 1e-3 target. The predictor is told nothing about that. From the + physics alone it says the refine-4 rule affords 2453 nodes on this window + while the peak needs more, i.e. under-resolved at the bottom of the ladder, + which is what they saw. At the ceiling of 32 the same rung is comfortable. + """ + _, tight = _pick(data, RHO_40, refine_max=4) + assert not tight["time_peak_resolvable_whole_window"] + assert tight["whole_window_nodes_available"] == pytest.approx( + (data.npts - 1) * 4 + 1) + + scheme, loose = _pick(data, RHO_40, refine_max=32) + assert loose["time_peak_resolvable_whole_window"] + assert scheme == "laplace" + + +def test_an_explicit_request_bypasses_the_analysis(data): + """Overrides stay overrides: 'auto' analyses, anything else is obeyed and + labelled as such so a run log cannot be misread as an analysed choice.""" + scheme, info = _pick(data, RHO_640, requested="exact") + assert scheme == "exact" + assert "explicit request" in info["reason"] + + +def test_the_pair_line_is_printable_and_names_the_refusal(data): + scheme, info = _pick(data, RHO_640) + line = DP.format_reserve_pair(scheme, info) + assert line.startswith("RESERVE-PAIR local=four-axis reserve=REFUSED") + for token in ("rho=", "sigma_f=", "A=", "peak=", "cover_needs=", + "reserve_needs="): + assert token in line + + +def test_the_local_cover_verdict_is_reported_and_leads(data): + """The FIRST quantity: can the local branch's time cover hold the peak? + + Measured elsewhere at rho 163: raising the cover 64 -> 256 took acceptance + 31% -> 75%. That is why the start cap appeared to plateau -- time capacity + was binding, not the start cap saturating -- so a selector that reports only + the reserve's budget would predict the wrong lever. The local branch is the + thing under test; the reserve is only what it falls back to. + """ + _, tight = _pick(data, RHO_640, max_time_nodes=64) + _, loose = _pick(data, RHO_640, max_time_nodes=4096) + assert tight["cover_nodes_needed"] == pytest.approx(loose["cover_nodes_needed"]) + assert not tight["local_cover_resolves_peak"] + assert loose["local_cover_resolves_peak"] + assert "local cover" in loose["reason"] + + +def test_the_bandwidth_definition_is_pinned_not_just_its_plausibility(data): + """A synthetic signal with a KNOWN analytic bandwidth, so a definition swap + fails here rather than shifting every threshold quietly. + + The earlier tests in this file only asserted sigma_f was finite, positive and + Nyquist-bounded. Every one of them passes with the RAW two-sided moment, + which is the wrong quantity for a timing width and was what this function + originally returned. Plausibility is not a definition. + + Construction: a complex tone at f0 with a Gaussian AMPLITUDE of width bw in + frequency. The moments weight |Qtilde|^2, so the POWER spectrum has width + bw/sqrt(2), and that -- not bw -- is what the central moment must return. + Getting this wrong was my first version of this test: the expectation, not + the code, was off by sqrt(2). The two-sided RAW moment must instead land + near f0, which is the error the test exists to catch. + """ + import numpy as np + + n, dt = 4096, 1.0 / 4096 + f0, bw = 200.0, 20.0 + freqs = np.fft.fftfreq(n, d=dt) + amp = np.exp(-0.5 * ((freqs - f0) / bw) ** 2) # positive-side only + x = np.fft.ifft(amp) + + class _D: + deltaT = dt + q_time_pregrid_factor = 1 + detector_names = ["H1"] + detectors = {"H1": {"Q": np.asarray(x)[:, None], + "q_time_pregrid_factor": 1}} + + central = DP.q_effective_bandwidth_hz(_D, moment="central") + raw = DP.q_effective_bandwidth_hz(_D, moment="raw") + default = DP.q_effective_bandwidth_hz(_D) + + bw_power = bw / np.sqrt(2.0) + assert central == pytest.approx(bw_power, rel=0.05), central + assert raw == pytest.approx(np.hypot(f0, bw_power), rel=0.05), raw + assert raw > 5.0 * central # the two are not interchangeable + # The DEFAULT must be the narrow one: sizing a time rule on the envelope + # bandwidth under-resolves every row whose likelihood is carrier-modulated. + assert default == pytest.approx(raw) + + +def test_the_selector_sizes_on_the_narrow_bandwidth(data): + """sigma_t must be built from the RAW moment -- the narrowest peak the + primitive can produce. If the envelope moment is ever wired back in, the + predicted width widens by the ratio of the two and the selector starts + calling a whole-window rule adequate when it is not.""" + _, info = _pick(data, RHO_160) + sigma_f = info["sigma_f_hz"] + assert sigma_f == pytest.approx( + DP.q_effective_bandwidth_hz(data, moment="raw")) + assert info["sigma_f_envelope_hz"] == pytest.approx( + DP.q_effective_bandwidth_hz(data, moment="central")) + expect = 1.0 / (2.0 * np.pi * RHO_160 * sigma_f) + assert info["sigma_t_s"] == pytest.approx(expect, rel=1e-9) + + +def test_the_gate_roster_lists_every_file_once_and_covers_this_one(): + """The CI roster is an explicit list, so a new test file is unrun until it + is added -- and a DUPLICATE entry runs the file twice and inflates the + collection floor, which then hides a later removal. + + Both halves are things I got wrong on this branch within one hour: a rebase + resolved the roster conflict by taking upstream's side and silently dropped + my entry, and the fix then added a second copy of an entry that was already + there because my check used a broken grep pattern. + """ + import collections + import os + import re + + # test/jax/ -> test/jax -> test -> Code -> MonteCarloMarginalizeCode + # -> repo root: FIVE levels, not four. + root = os.path.abspath(__file__) + for _ in range(5): + root = os.path.dirname(root) + gate = os.path.join(root, ".travis", "test-jax.sh") + src = open(gate, encoding="utf-8").read() + block = src[src.index("FILES=("):src.index("\n)", src.index("FILES=("))] + listed = re.findall(r'\$\{JAXDIR\}/(test_[A-Za-z0-9_]+\.py)', block) + + dupes = {n: c for n, c in collections.Counter(listed).items() if c > 1} + assert not dupes, "roster lists a file more than once: %r" % (dupes,) + assert os.path.basename(__file__) in listed, ( + "this file is not in the gate roster, so CI would not run it") + + +def test_the_gate_sets_its_floor_exactly_once(): + """Bash keeps the LAST assignment, so a duplicated constant leaves earlier + ones dead while they still read as authoritative in review. + + This file carried FIVE consecutive unconditional EXPECTED_TESTS= lines + (648, 629, 657, 648, 705) accumulated by parallel merges. Only 705 was + live. A reviewer checking "is the floor right?" would most likely read the + first, which had been dead for three merges. + """ + import os + import re + + root = os.path.abspath(__file__) + for _ in range(5): + root = os.path.dirname(root) + src = open(os.path.join(root, ".travis", "test-jax.sh"), encoding="utf-8").read() + assigns = re.findall(r"(?m)^EXPECTED_TESTS=(\d+)", src) + assert len(assigns) == 1, ( + "EXPECTED_TESTS assigned %d times (%s); bash keeps the last and the " + "rest are dead" % (len(assigns), ", ".join(assigns))) + + +def test_the_roster_is_honoured_on_the_ANGULAR_branch_too(data): + """A roster without laplace must not yield laplace. + + The roster was checked only where the selector chooses ``peaklocal`` -- the + branch that could not have chosen it anyway, since 'peaklocal' is never on + the roster today. On the branch that CAN choose laplace, the caller's + roster was ignored, so a run whose data cannot support the laplace reserve + still selected it the moment A cleared the crossover. + + The roster is not a preference. It says which schemes this data and this + distance quadrature can support at all -- for laplace, that the adaptive + node placement's A0 == 0 / B1 == 0 premise holds -- so ignoring it means + running exact under laplace's name, or worse. + """ + with_lap, _ = _pick(data, RHO_40, available=("exact", "laplace")) + assert with_lap == "laplace" + + scheme, info = _pick(data, RHO_40, available=("exact",)) + assert scheme is None, ( + "laplace selected off a roster that does not offer it: %r" + % (info["reason"],)) + assert "on the roster" in info["reason"] + assert "laplace" in info["reason"] + + +def test_a_roster_absence_is_not_overridable_by_an_explicit_request(data): + """An explicit request overrides the ANALYSIS, not the ROSTER. + + Forcing a scheme whose premise is absent is not an override; it is an + unnoticed wrong answer. The distinction matters because the driver's + refusal message invites the user to pass an explicit scheme -- that must + let them overrule the crossover, and must not let them overrule a measured + identity failure. + """ + ok, _ = _pick(data, RHO_40, requested="exact", available=("exact",)) + assert ok == "exact" + + scheme, info = _pick(data, RHO_40, requested="laplace", + available=("exact",)) + assert scheme is None + assert "not overridable" in info["reason"] + + +def test_an_unwired_reserve_scheme_is_refused_not_run_as_exact(): + """A config field the composite never reads is worse than a missing one. + + ``PolicyConfig.reserve_scheme`` is validated against the CHOICES tuple, and + the composite dispatches through ``reserve_pair``. A scheme declared in + CHOICES but absent from the pair table must be refused, not accepted, + reported in the policy line, and computed as exact. + """ + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="exact")) + # 'auto' is resolved before the composite sees it, so it is admitted here. + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="auto")) + + # laplace and peaklocal are wired through the pair table (#304), so every + # declared scheme executes today; the refusal is exercised by declaring a + # scheme the pair table does not carry. + for scheme in ("laplace", "peaklocal"): + assert scheme in DP.RESERVE_SCHEME_CHOICES + assert scheme in DP.RESERVE_SCHEME_EXECUTABLE + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme=scheme)) + saved = DP.RESERVE_SCHEME_CHOICES + DP.RESERVE_SCHEME_CHOICES = saved + ("nonesuch-declared",) + try: + assert "nonesuch-declared" not in DP.RESERVE_SCHEME_EXECUTABLE + with pytest.raises(ValueError, match="NOT WIRED"): + DP.validate_policy_config( + DP.PolicyConfig(reserve_scheme="nonesuch-declared")) + finally: + DP.RESERVE_SCHEME_CHOICES = saved + + with pytest.raises(ValueError, match="must be one of"): + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="nonesuch")) + + +def test_the_executable_roster_is_a_subset_of_the_choices(): + """Guards the pair as a pair: a scheme becomes executable by being wired, + and this fails if RESERVE_SCHEME_EXECUTABLE ever names something the + choices tuple does not, which would mean the two lists were edited apart.""" + assert set(DP.RESERVE_SCHEME_EXECUTABLE) <= set(DP.RESERVE_SCHEME_CHOICES) + assert "auto" not in DP.RESERVE_SCHEME_EXECUTABLE, ( + "'auto' is a resolution mode, not something the composite executes") + assert DP.RESERVE_SCHEME_DEFAULT in DP.RESERVE_SCHEME_EXECUTABLE, ( + "the default must be executable or every bare run refuses") + + +def test_auto_is_admitted_only_because_something_resolves_it(): + """``validate_policy_config`` admits 'auto' on a PREMISE: that the driver + writes the resolved pair back onto the config before the composite runs. + + Found by the #304 session -- the driver announced a pair and left + ``policy_config.reserve_scheme == "auto"``. Nothing dispatched on the + string yet, so nothing broke; the printed line was a claim about a value + the composite never read, and the admission's stated reason was false. + + A COUPLING GUARD, not a behaviour test, and labelled as one. The behaviour + -- 'auto' reaching the composite and being executed as something -- cannot + be observed until #304 dispatches on the string, because today the + composite ignores the field entirely. So this checks the two tokens that + have to co-exist, tolerant of formatting, and says what to do if it fires: + if the driver stops resolving, 'auto' must stop being admitted. + """ + import re + with open(_DRIVER) as fh: + src = fh.read() + assert re.search(r"_replace\(\s*reserve_scheme=_pair\s*\)", src), ( + "the driver no longer writes the resolved pair onto policy_config, so " + "validate_policy_config must stop admitting 'auto' -- an unresolved " + "'auto' reaching the composite is the silently-inert field this branch " + "refuses everywhere else") + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme="auto")) + + +def test_a_selected_reserve_the_composite_cannot_run_is_refused(): + """Selecting is not running. The roster says what the DATA supports; the + executable tuple says what the composite DISPATCHES, and the two are not + the same list. A pair that clears the first and fails the second must + refuse, not run exact under the selected scheme's name.""" + for scheme in DP.RESERVE_SCHEME_CHOICES: + if scheme == "auto": + continue + if scheme in DP.RESERVE_SCHEME_EXECUTABLE: + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme=scheme)) + else: + with pytest.raises(ValueError, match="NOT WIRED"): + DP.validate_policy_config(DP.PolicyConfig(reserve_scheme=scheme)) diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_time_first_peaklocal.py b/MonteCarloMarginalizeCode/Code/test/jax/test_time_first_peaklocal.py new file mode 100644 index 000000000..bbb3e21a3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_time_first_peaklocal.py @@ -0,0 +1,192 @@ +"""Tests for primitive-first time peak-local composition.""" + +import inspect + +import numpy as np +import pytest +from scipy import special + +jax = pytest.importorskip("jax") +import jax.numpy as jnp + +jax.config.update("jax_enable_x64", True) + +from RIFT.likelihood.jax_ile import time_first_peaklocal as TFP + + +def _log_i0(x): + x = np.asarray(x, dtype=float) + return np.log(special.i0e(x)) + np.abs(x) + + +def _cosine_samples(n, amplitude, harmonic): + span = n - 1.0 + t = np.arange(n, dtype=float) + return amplitude * np.cos(harmonic * np.pi * t / span) + + +def test_distance_and_time_known_integral_uses_fewer_nonlinear_nodes(): + """Distance nodes are lanes; each time integral is exactly an I0 integral.""" + n = 65 + span = n - 1.0 + K = 36.0 + harmonic = 3 + kappa = _cosine_samples(n, K, harmonic).astype(complex) + rho = 4.0 + x = np.array([0.45, 0.7, 1.0, 1.25]) + logw = np.log(np.array([0.1, 0.25, 0.4, 0.25])) + + got, ok, info = TFP.time_first_distance_peak_local_marginalize( + jnp.asarray(kappa), rho, jnp.asarray(x), jnp.asarray(logw), 1.0, + enum_factor=8, fine_factor=32, max_nodes=8192, + keep_nats=36.0, quadrature_tol_nats=2.0e-6) + want = special.logsumexp( + logw - 0.5 * rho * x * x + np.log(span) + _log_i0(K * x)) + + assert bool(ok), {k: np.asarray(v) for k, v in info.items()} + assert abs(float(got) - float(want)) < 2.0e-6 + assert int(info["n_local_hi"]) < int(info["n_dense_hi"]) + assert float(info["tail_margin"]) < -23.0 + + +def test_symmetric_angle_reduction_adversary_reconstructs_before_logsumexp(): + """A nonlinear marginal can be constant on samples and structured between them. + + The two lanes represent symmetry-related angle states with primitive + correlations ``+A cos(pi t)`` and ``-A cos(pi t)``. At integer input + samples their marginalized log integrand is the constant ``log cosh(A)``. + Interpolating that already-marginalized row therefore converges to the wrong + constant function. Reconstructing both primitive lanes first recovers + ``log cosh(A cos(pi t))`` and the known ``T I0(A)`` integral. + """ + n, amplitude = 17, 8.0 + base = amplitude * (-1.0) ** np.arange(n) + lanes = np.stack((base, -base)).astype(complex) + logw = np.full(2, -np.log(2.0)) + rho = np.zeros(2) + + got, ok, info = TFP.time_first_peak_local_marginalize( + jnp.asarray(lanes), jnp.asarray(rho), jnp.asarray(logw), 1.0, + enum_factor=8, fine_factor=32, max_nodes=8192, + keep_nats=20.0, quadrature_tol_nats=1.0e-7) + want = np.log(n - 1.0) + float(_log_i0(amplitude)) + wrong = np.log(n - 1.0) + np.log(np.cosh(amplitude)) + + assert bool(ok), {k: np.asarray(v) for k, v in info.items()} + assert abs(float(got) - want) < 1.0e-7 + assert abs(wrong - want) > 1.0 + + # Pin the ordering structurally as well as numerically: the evaluator has + # explicit selected-point primitive reconstruction -> downstream-reduction + # stages and never creates the globally refined primitive. + source = inspect.getsource(TFP._evaluate_cover_at_factor) + assert source.index("_evaluate_time_spectrum") < source.index( + "_lane_log_integrand") + assert "reconstruct_time_primitive" not in source + + +def test_selected_point_reconstruction_matches_the_dense_fft_grid(): + """The local DFT is the same reflected interpolant, evaluated sparsely.""" + rng = np.random.default_rng(390) + lanes = (rng.normal(size=(3, 41)) + + 1j * rng.normal(size=(3, 41))) + factor = 16 + dense = np.asarray(TFP.reconstruct_time_primitive( + jnp.asarray(lanes), factor)) + index = np.array([0, 1, 7, 31, 117, dense.shape[-1] - 1]) + sparse = np.asarray(TFP.evaluate_time_primitive_points( + jnp.asarray(lanes), jnp.asarray(index / factor))) + np.testing.assert_allclose(sparse, dense[:, index], atol=3e-12, rtol=0) + + +def test_local_evaluator_live_shape_does_not_contain_the_dense_factor(): + """The fine factor changes positions, not a materialized array dimension.""" + n = 33 + lanes = jnp.asarray(_cosine_samples(n, 14.0, 3)[None, :], + dtype=jnp.complex128) + rho = jnp.zeros(1) + logw = jnp.zeros(1) + enum_factor = 4 + k_enum = TFP.reconstruct_time_primitive(lanes, enum_factor) + m1 = TFP.spectral_time_derivative_bound(lanes, 1.0) + plan = TFP.plan_time_cover( + k_enum, rho, logw, m1, 1.0 / enum_factor, keep_nats=12.0) + + # The output's conceptual dense count grows, while every explicit phasor + # evaluated by the implementation has only sub+1 selected positions. + source = inspect.getsource(TFP._evaluate_cover_at_factor) + assert "positions[:, None]" not in source # phasor is isolated in the helper + for factor in (16, 128): + value, n_local, ok, n_dense = TFP._evaluate_cover_at_factor( + lanes, rho, logw, plan, 1.0, enum_factor, factor, 0, 65536) + assert np.isfinite(float(value)) and bool(ok) + assert int(n_local) < int(n_dense) + + +def test_cell_upper_bound_dominates_a_much_finer_reconstruction(): + """The planner's correctness-bearing output is an upper bound, not a grid max.""" + n = 49 + a = _cosine_samples(n, 13.0, 5) + b = (_cosine_samples(n, 7.0, 2) + + _cosine_samples(n, 3.0, 7)) + lanes = np.stack((a, b)).astype(complex) + rho = jnp.asarray([1.3, 0.7]) + logw = jnp.log(jnp.asarray([0.35, 0.65])) + enum_factor, truth_factor = 4, 128 + + k_enum = TFP.reconstruct_time_primitive(jnp.asarray(lanes), enum_factor) + m1 = TFP.spectral_time_derivative_bound(jnp.asarray(lanes), 1.0) + plan = TFP.plan_time_cover( + k_enum, rho, logw, m1, 1.0 / enum_factor, keep_nats=12.0) + k_truth = TFP.reconstruct_time_primitive(jnp.asarray(lanes), truth_factor) + g_truth = np.asarray(TFP._lane_log_integrand(k_truth, rho, logw)) + + sub = truth_factor // enum_factor + upper = np.asarray(plan.cell_log_upper) + observed = np.array([ + g_truth[i * sub:(i + 1) * sub + 1].max() + for i in range(upper.size) + ]) + assert np.all(observed <= upper + 2.0e-11), np.max(observed - upper) + + +def test_capacity_decline_is_ledgered_and_does_not_silently_widen(): + n = 33 + kappa = _cosine_samples(n, 20.0, 1)[None, :].astype(complex) + got, ok, info = TFP.time_first_peak_local_marginalize( + jnp.asarray(kappa), jnp.zeros(1), jnp.zeros(1), 1.0, + enum_factor=8, fine_factor=32, max_nodes=8, + keep_nats=30.0) + assert np.isfinite(float(got)) + assert not bool(ok) + assert not bool(info["capacity_ok"]) + assert bool(info["decline_capacity"]) + assert bool(info["reconciles"]) + assert sum(bool(info[k]) for k in ( + "decline_nonfinite", "decline_capacity", "decline_quadrature", + "decline_tail")) == 1 + assert int(info["n_local_hi"]) > 8 + + +def test_fixed_shape_kernel_jits_and_has_finite_gradient(): + n = 33 + shape = _cosine_samples(n, 1.0, 3) + + @jax.jit + def f(amplitude): + lanes = (amplitude * jnp.asarray(shape))[None, :].astype(jnp.complex128) + value, ok, _ = TFP.time_first_peak_local_marginalize( + lanes, jnp.zeros(1), jnp.zeros(1), 1.0, + enum_factor=4, fine_factor=16, max_nodes=4096, + keep_nats=30.0, quadrature_tol_nats=1.0e-5) + return jnp.where(ok, value, jnp.nan) + + value = f(12.0) + grad = jax.grad(f)(12.0) + assert np.all(np.isfinite(np.asarray([value, grad]))) + + +def test_api_rejects_an_already_marginalized_time_row(): + with pytest.raises(ValueError, match="already-marginalized"): + TFP.time_first_peak_local_marginalize( + jnp.ones(17), jnp.zeros(1), jnp.zeros(1), 1.0) diff --git a/MonteCarloMarginalizeCode/Code/test/run_bandlimited_retained_ile_validation.py b/MonteCarloMarginalizeCode/Code/test/run_bandlimited_retained_ile_validation.py new file mode 100644 index 000000000..3d4bb40f0 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/run_bandlimited_retained_ile_validation.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 +"""Run one accepted-reference ILE argv against a frozen RIFT commit.""" + +import argparse +import hashlib +import json +import math +import os +import re +import shlex +import statistics +import subprocess +import sys +import time +from datetime import datetime, timezone +from pathlib import Path + + +def _git(tree, *args): + return subprocess.check_output( + ["git", "-C", str(tree), *args], universal_newlines=True).strip() + + +def _option(argv, name): + return argv[argv.index(name) + 1] + + +def _set_option(argv, name, value): + where = argv.index(name) + 1 + argv[where] = str(value) + + +def _elapsed_seconds(resource_text): + match = re.search(r"Elapsed \(wall clock\) time.*?:\s*([0-9:.]+)$", resource_text, re.M) + if not match: + return None + fields = [float(item) for item in match.group(1).split(":")] + if len(fields) == 2: + return 60 * fields[0] + fields[1] + if len(fields) == 3: + return 3600 * fields[0] + 60 * fields[1] + fields[2] + return None + + +def _resource_value(pattern, resource_text, cast=int): + match = re.search(pattern, resource_text, re.M) + return cast(match.group(1)) if match else None + + +def _sha256(path): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--baseline-record", required=True, type=Path) + parser.add_argument("--rift-tree", required=True, type=Path) + parser.add_argument("--expected-commit", required=True) + parser.add_argument("--container", required=True, type=Path) + parser.add_argument("--output-root", required=True, type=Path) + parser.add_argument("--gpu", type=int, default=2) + parser.add_argument("--cpuset", default="0-7") + parser.add_argument("--control", choices=("retained", "full"), default="retained") + args = parser.parse_args() + + baseline = json.loads(args.baseline_record.read_text()) + if not baseline.get("accepted", False): + raise SystemExit("baseline record is not accepted") + argv = list(baseline["argv"]) + required = { + "--sampler-method": "AV", + "--time-marginalization-quadrature": "bandlimited", + "--interpolate-time": "sinc", + "--n-max": "4000000", + "--n-eff": "100", + "--n-chunk": "40000", + } + mismatch = {} + for key, expected in required.items(): + observed = _option(argv, key) if key in argv else None + if observed != expected: + mismatch[key] = (observed, expected) + if mismatch: + raise SystemExit("baseline argv does not meet production contract: %r" % mismatch) + + commit = _git(args.rift_tree, "rev-parse", "HEAD") + if commit != args.expected_commit: + raise SystemExit("RIFT commit mismatch: %s != %s" % (commit, args.expected_commit)) + dirty = _git(args.rift_tree, "status", "--porcelain") + if dirty: + raise SystemExit("RIFT tree is dirty:\n" + dirty) + if not args.container.is_file(): + raise SystemExit("missing container: %s" % args.container) + observed_input_hashes = {} + input_hash_mismatches = {} + for path, expected in baseline.get("input_sha256", {}).items(): + observed = _sha256(path) + observed_input_hashes[path] = observed + if observed != expected: + input_hash_mismatches[path] = {"expected": expected, "observed": observed} + if input_hash_mismatches: + raise SystemExit("baseline inputs changed: %s" % json.dumps( + input_hash_mismatches, sort_keys=True)) + container_sha256 = observed_input_hashes.get(str(args.container)) + expected_container_sha256 = baseline.get("input_sha256", {}).get(str(args.container)) + if expected_container_sha256 != container_sha256: + raise SystemExit("container does not match the accepted baseline record") + + cell = "%s_bandlimited_snr%s_seed%s_%s" % ( + baseline["model"], baseline["snr_label"], baseline["seed"], args.control) + out = args.output_root / cell + if out.exists(): + raise SystemExit("refusing to overwrite validation directory: %s" % out) + out.mkdir(parents=True) + output_prefix = out / "output" + _set_option(argv, "--output-file", output_prefix) + + code = args.rift_tree / "MonteCarloMarginalizeCode" / "Code" + ile = code / "bin" / "integrate_likelihood_extrinsic_batchmode" + wrapper = Path(__file__).with_name("telemetry_bandlimited_retained_ile.py") + telemetry = out / "fft_telemetry.json" + env_opts = { + "PYTHONPATH": str(code), + "PATH": str(code / "bin") + ":/usr/local/bin:/usr/bin:/bin", + "OMP_NUM_THREADS": "1", + "OPENBLAS_NUM_THREADS": "1", + "MKL_NUM_THREADS": "1", + "NUMEXPR_NUM_THREADS": "1", + "CUDA_VISIBLE_DEVICES": str(args.gpu), + "RIFT_REAL_ILE": str(ile), + "RIFT_FFT_TELEMETRY_FILE": str(telemetry), + "RIFT_VALIDATION_FORCE_FULL_FFT": "1" if args.control == "full" else "0", + } + launch = ["apptainer", "exec", "--nv"] + for key, value in env_opts.items(): + launch.extend(("--env", key + "=" + value)) + launch.extend((str(args.container), "python3", "-u", str(wrapper))) + timed = ["/usr/bin/time", "-v", "-o", str(out / "resource.txt"), + "taskset", "-c", args.cpuset] + launch + argv + + start = datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z") + (out / "start_utc.txt").write_text(start + "\n") + (out / "rift_commit.txt").write_text(commit + "\n") + (out / "baseline_record.txt").write_text(str(args.baseline_record.resolve()) + "\n") + (out / "argv.nul").write_bytes(b"\0".join(item.encode() for item in argv) + b"\0") + (out / "command.txt").write_text( + " ".join(shlex.quote(item) for item in timed) + "\n") + provenance = { + "schema": 1, + "baseline_record": str(args.baseline_record.resolve()), + "baseline_record_sha256": _sha256(args.baseline_record), + "baseline_status_sha256": baseline.get("status_sha256"), + "input_sha256": observed_input_hashes, + "rift_commit": commit, + "container": str(args.container), + "container_sha256": container_sha256, + "control": args.control, + "physical_gpu": args.gpu, + "cpuset": args.cpuset, + "argv_matches_baseline_except_output": True, + } + (out / "provenance.json").write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\n") + + # One persistent nvidia-smi matches the accepted campaign monitor and avoids + # racing Apptainer's Go runtime with a new helper process every 200 ms on + # login nodes with a tight per-user thread limit. + monitor_path = out / "gpu_usage.csv" + monitor_cmd = [ + "nvidia-smi", + "-i", + str(args.gpu), + "--query-gpu=timestamp,memory.used,utilization.gpu", + "--format=csv,noheader,nounits", + "--loop-ms=200", + "--filename=" + str(monitor_path), + ] + monitor_process = subprocess.Popen( + monitor_cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + launch_env = os.environ.copy() + launch_env["GOMAXPROCS"] = "4" + time.sleep(0.3) + try: + with (out / "run.log").open("w") as log: + process = subprocess.Popen( + timed, stdout=log, stderr=subprocess.STDOUT, env=launch_env) + rc = process.wait() + finally: + monitor_process.terminate() + try: + monitor_process.wait(timeout=5) + except subprocess.TimeoutExpired: + monitor_process.kill() + monitor_process.wait() + + monitor = [] + if monitor_path.exists(): + for line in monitor_path.read_text(errors="replace").splitlines(): + fields = [field.strip() for field in line.split(",")] + if len(fields) != 3: + continue + try: + monitor.append((fields[0], float(fields[1]), float(fields[2]))) + except ValueError: + pass + (out / "exit_code.txt").write_text(str(rc) + "\n") + final_commit = _git(args.rift_tree, "rev-parse", "HEAD") + final_dirty = _git(args.rift_tree, "status", "--porcelain") + (out / "rift_commit_final.txt").write_text(final_commit + "\n") + + status_path = out / "output_0_integrator_status.json" + status = json.loads(status_path.read_text()) if status_path.exists() else {} + resource_text = (out / "resource.txt").read_text(errors="replace") + log_text = (out / "run.log").read_text(errors="replace") + if "CuPy Platform" in log_text and "NVIDIA CUDA" in log_text: + backend = "cuda" + elif "no cupy" in log_text.lower(): + backend = "numpy-cpu" + else: + backend = "unknown" + n_ess = status.get("n_ESS") + khat = status.get("pareto_khat") + run_rejection = [] + if rc: + run_rejection.append("nonzero exit") + if final_commit != commit: + run_rejection.append("RIFT source commit changed during run") + if final_dirty: + run_rejection.append("RIFT source tree became dirty during run") + if status.get("collapsed", False): + run_rejection.append("AV live-volume collapse") + if n_ess is None or not math.isfinite(float(n_ess)) or float(n_ess) < 100: + run_rejection.append("n_ESS below 100") + if khat is None or not math.isfinite(float(khat)) or float(khat) >= 0.7: + run_rejection.append("Pareto k_hat not below 0.7") + if backend != "cuda": + run_rejection.append("requested GPU backend not verified") + if not monitor: + run_rejection.append("GPU memory monitor produced no samples") + + optimization_rejection = [] + if not telemetry.exists(): + optimization_rejection.append("FFT telemetry missing") + fft_telemetry = None + else: + fft_telemetry = json.loads(telemetry.read_text()) + if fft_telemetry.get("failed_calls"): + optimization_rejection.append("band-limited marginalization call failed") + if fft_telemetry.get("full_fft_fallback_rows"): + optimization_rejection.append("retained FFT fell back to full padding") + if args.control == "retained": + if not fft_telemetry.get("retained_fft_rows"): + optimization_rejection.append("retained control used no retained FFT rows") + if fft_telemetry.get("full_fft_selected_rows"): + optimization_rejection.append("retained GPU control selected full padding") + else: + if not fft_telemetry.get("full_fft_selected_rows"): + optimization_rejection.append("full control used no full-padding rows") + if fft_telemetry.get("retained_fft_rows"): + optimization_rejection.append("full control used retained FFT rows") + + memories = [item[1] for item in monitor] + utilizations = [item[2] for item in monitor] + result = { + "schema": 1, + "accepted": not run_rejection and not optimization_rejection, + "sampler_accepted": not run_rejection, + "sampler_rejection_reasons": run_rejection, + "optimization_validated": not optimization_rejection, + "optimization_rejection_reasons": optimization_rejection, + "model": baseline["model"], + "snr_label": baseline["snr_label"], + "seed": baseline["seed"], + "control": args.control, + "exit_code": rc, + "backend_actual": backend, + "rift_commit": commit, + "rift_commit_final": final_commit, + "rift_dirty_final": bool(final_dirty), + "lnZ": status.get("lnL"), + "sigma_lnZ": status.get("sigma_lnL"), + "n_ESS": n_ess, + "pareto_khat": khat, + "ntotal": status.get("ntotal"), + "collapsed": status.get("collapsed", False), + "wall_seconds": _elapsed_seconds(resource_text), + "max_rss_kib": _resource_value(r"Maximum resident set size \(kbytes\):\s*(\d+)", resource_text), + "gpu_peak_mib": max(memories) if memories else None, + "gpu_utilization_median": statistics.median(utilizations) if utilizations else None, + "gpu_monitor_samples": len(monitor), + "fft_telemetry": fft_telemetry, + "baseline": { + key: baseline.get(key) for key in ( + "rift_commit", "lnZ", "sigma_lnZ", "n_ESS", "pareto_khat", + "ntotal", "wall_seconds", "max_rss_kib", "gpu_peak_mib") + }, + } + if result["lnZ"] is not None and baseline.get("lnZ") is not None: + result["delta_lnZ_vs_baseline"] = result["lnZ"] - baseline["lnZ"] + combined = math.hypot(result["sigma_lnZ"], baseline["sigma_lnZ"]) + result["delta_lnZ_over_combined_sigma"] = result["delta_lnZ_vs_baseline"] / combined + (out / "validation_record.json").write_text(json.dumps(result, indent=2, sort_keys=True) + "\n") + if result["accepted"]: + (out / "DONE").touch() + else: + reasons = (["sampler: " + reason for reason in run_rejection] + + ["optimization: " + reason for reason in optimization_rejection]) + (out / "REJECTED").write_text("\n".join(reasons) + "\n") + print(json.dumps(result, indent=2, sort_keys=True)) + return 0 if result["accepted"] else 20 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/telemetry_bandlimited_retained_ile.py b/MonteCarloMarginalizeCode/Code/test/telemetry_bandlimited_retained_ile.py new file mode 100644 index 000000000..1963fcb33 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/telemetry_bandlimited_retained_ile.py @@ -0,0 +1,135 @@ +#!/usr/bin/env python3 +"""Validation-only ILE launcher that aggregates band-limited FFT telemetry. + +The companion harness requires a clean, commit-pinned RIFT tree before importing +this driver. This wrapper never catches or converts likelihood exceptions: +failures are counted, then re-raised. +""" + +import atexit +import json +import os +import runpy +import sys +from collections import Counter +from pathlib import Path + +from RIFT.likelihood import time_marginalization_quadrature as tmq + + +TELEMETRY_PATH = Path(os.environ["RIFT_FFT_TELEMETRY_FILE"]) +REAL_ILE = os.environ["RIFT_REAL_ILE"] +FORCE_FULL = os.environ.get("RIFT_VALIDATION_FORCE_FULL_FFT", "0") == "1" + +_sum_keys = ( + "n_rows", + "n_refined_rows", + "n_wrap_exposed_rows", + "n_unmeasurable_rows", + "n_flat_rows", + "n_refinements", + "retained_fft_batches", + "retained_fft_rows", + "full_fft_selected_batches", + "full_fft_selected_rows", + "full_fft_fallback_batches", + "full_fft_fallback_rows", +) +_max_keys = ( + "upsample_factor", + "max_reflected_period", + "max_dense_factor", + "max_reference_full_fft_length", + "max_retained_fft_length", + "max_retained_grid_length", + "n_retained_fft_plans", +) +_aggregate = { + "schema": 1, + "validation_force_full_fft": FORCE_FULL, + "successful_calls": 0, + "failed_calls": 0, + "failure_types": Counter(), + "backend_calls": Counter(), + "strategy_calls": Counter(), + "factor_rows": Counter(), + "full_fft_selected_reasons": Counter(), + "full_fft_fallback_reasons": Counter(), +} +for _key in _sum_keys + _max_keys: + _aggregate[_key] = 0 + + +if FORCE_FULL: + def _validation_force_full(x, factor, plan_cache, transform_report, xpy=None): + if xpy is None: + xpy = tmq.np + reason = "validation-only explicit full-FFT control" + n_rows = int(x.shape[0]) + reasons = transform_report["full_fft_selected_reasons"] + reasons[reason] = reasons.get(reason, 0) + n_rows + tmq._record_transform( + transform_report, + "full_fft_selected", + n_rows, + 2 * int(x.shape[-1]), + int(factor), + ) + return tmq.reflected_bandlimited_upsample(x, factor, xpy=xpy) + + tmq._reflected_upsample_for_integration = _validation_force_full + + +_original = tmq.time_marginalize_bandlimited + + +def _merge_report(report, xpy): + _aggregate["successful_calls"] += 1 + backend = getattr(xpy, "__name__", type(xpy).__name__) + _aggregate["backend_calls"][backend] += 1 + _aggregate["strategy_calls"][report.get("bandlimited_fft_strategy", "missing")] += 1 + for key in _sum_keys: + _aggregate[key] += int(report.get(key, 0) or 0) + for key in _max_keys: + _aggregate[key] = max(_aggregate[key], int(report.get(key, 0) or 0)) + for factor, rows in report.get("factor_histogram", {}).items(): + _aggregate["factor_rows"][str(factor)] += int(rows) + for key in ("full_fft_selected_reasons", "full_fft_fallback_reasons"): + for reason, rows in report.get(key, {}).items(): + _aggregate[key][reason] += int(rows) + + +def _instrumented(*args, **kwargs): + xpy = kwargs.get("xpy", tmq.np) + try: + result = _original(*args, **kwargs) + except BaseException as exc: + _aggregate["failed_calls"] += 1 + _aggregate["failure_types"][type(exc).__name__] += 1 + raise + _merge_report(tmq.last_report(), xpy) + return result + + +tmq.time_marginalize_bandlimited = _instrumented + + +def _jsonable(): + return { + key: dict(value) if isinstance(value, Counter) else value + for key, value in _aggregate.items() + } + + +def _write_telemetry(): + payload = _jsonable() + TELEMETRY_PATH.parent.mkdir(parents=True, exist_ok=True) + temp = TELEMETRY_PATH.with_suffix(TELEMETRY_PATH.suffix + ".tmp") + temp.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + os.replace(temp, TELEMETRY_PATH) + print("RIFT_FFT_TELEMETRY_JSON=" + json.dumps(payload, sort_keys=True), flush=True) + + +atexit.register(_write_telemetry) +sys.argv[0] = REAL_ILE +runpy.run_path(REAL_ILE, run_name="__main__") diff --git a/MonteCarloMarginalizeCode/Code/test/test_jax_ile_selectable.py b/MonteCarloMarginalizeCode/Code/test/test_jax_ile_selectable.py new file mode 100644 index 000000000..e264eabca --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_jax_ile_selectable.py @@ -0,0 +1,250 @@ +"""Pipeline-selectable ILE executable: --use-jax-ile / --ile-exe. + +RO'S directive (2026-09-08): util_RIFT_pseudo_pipe.py hard-coded +`` `which integrate_likelihood_extrinsic_batchmode` `` as the ILE executable +handed to create_event_parameter_pipeline_BasicIteration, with no way for a +pipeline builder to name bin/integrate_likelihood_extrinsic_jax instead. +pseudo_pipe now exposes --use-jax-ile (resolves to +`which integrate_likelihood_extrinsic_jax`) and --ile-exe (an explicit path), +threaded through the CEPP's existing --ile-exe option to every ILE/ILE_puff/ +ILE_fetch/ILE_extr condor submit file it writes. + +These are full subprocess DAG builds against the same reference ini/coinc +fixtures .travis/test-build.sh uses (.travis/ref_ini/GW150914.ini + coinc.xml), +so what is tested is the actual CLI wiring, not a mock of it. OSG/singularity +are turned off in the ini used here: under --use-osg (without --use-singularity) +write_ILE_sub_simple rewrites the condor "executable" to a fixed OSG wrapper +script (my_wrapper.sh) and carries the real ILE executable inside that +script's body instead, which is orthogonal to the selection wiring under test. +""" + +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest + +CODE = Path(__file__).resolve().parents[1] +BIN = CODE / "bin" +PSEUDO_PIPE = BIN / "util_RIFT_pseudo_pipe.py" +REPO = CODE.parents[1] +REF_INI = REPO / ".travis" / "ref_ini" / "GW150914.ini" +COINC = REPO / ".travis" / "ref_ini" / "coinc.xml" + +BATCHMODE_EXE = str((BIN / "integrate_likelihood_extrinsic_batchmode").resolve()) +JAX_EXE = str((BIN / "integrate_likelihood_extrinsic_jax").resolve()) + +pytestmark = pytest.mark.skipif( + not (REF_INI.exists() and COINC.exists()), + reason="reference ini/coinc fixtures not present in this checkout") + + +def _fast_ini(tmp_path): + """The reference ini with OSG disabled and a tiny initial grid. + + OSG is disabled for the reason in the module docstring. The grid is + shrunk from the production value (5000) to keep this a DAG-BUILD test, + not a several-minute grid-construction benchmark. + """ + text = REF_INI.read_text() + for flag in ("use_osg", "use_osg_file_transfer", "use_osg_cip"): + text = text.replace("{}=True".format(flag), "{}=False".format(flag)) + text = re.sub(r"force-initial-grid-size=\d+", "force-initial-grid-size=4", text) + out = tmp_path / "ref_fast.ini" + out.write_text(text) + return out + + +def _fast_ini_with_osg(tmp_path): + """The reference ini with OSG left ON and a tiny initial grid. + + Unlike _fast_ini, use_osg/use_osg_file_transfer/use_osg_cip are NOT + disabled: this is deliberately the "OSG/singularity deployment this + repository's own reference ini defaults to" that the --use-jax-ile + + --use-osg BLOCKER (util_RIFT_pseudo_pipe.py's OSG/SINGULARITY refusal) + is about. Used only to exercise the refusal itself, which fires before + any of use_osg_file_transfer's downstream branching is reached. + """ + text = REF_INI.read_text() + text = re.sub(r"force-initial-grid-size=\d+", "force-initial-grid-size=4", text) + out = tmp_path / "ref_fast_osg.ini" + out.write_text(text) + return out + + +def _fast_ini_with_osg_cvmfs(tmp_path): + """OSG/singularity ON, but use_osg_file_transfer=False (CVMFS frames). + + With use_osg_file_transfer=True (the reference ini's own default, see + _fast_ini_with_osg), write_ILE_sub_simple additionally wraps the job in a + generated ile_pre.sh that builds local.cache at runtime -- ILE.sub's + "executable" line then names that wrapper, not the ILE driver, and the + driver path only appears inside the wrapper's body. --use-cvmfs-frames + (added by pseudo_pipe when use_osg_file_transfer=False) skips that + wrapper, so this variant isolates exactly the mechanism the BLOCKER + finding is about: write_ILE_sub_simple's SINGULARITY_BASE_EXE_DIR + + basename(exe) rewrite of ILE.sub's own "executable" line. + """ + text = REF_INI.read_text() + text = text.replace("use_osg_file_transfer=True", "use_osg_file_transfer=False") + text = re.sub(r"force-initial-grid-size=\d+", "force-initial-grid-size=4", text) + out = tmp_path / "ref_fast_osg_cvmfs.ini" + out.write_text(text) + return out + + +def _shim_path_dir(tmp_path): + """A directory with 'python' -> this interpreter. + + create_event_parameter_pipeline_BasicIteration is invoked by pseudo_pipe + through `os.system(cmd)` (a bare script name resolved via PATH, run + through its own `#!/usr/bin/env python` shebang) rather than + `sys.executable