diff --git a/.travis/ci_roster.txt b/.travis/ci_roster.txt index 8456d86e7..aa56e77fb 100644 --- a/.travis/ci_roster.txt +++ b/.travis/ci_roster.txt @@ -95,8 +95,8 @@ 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 with it; 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; skips cleanly without either, 10 tests where both are installed # 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 @@ -114,23 +114,23 @@ MonteCarloMarginalizeCode/Code/test/integrators/test_mcsamp_vegas.py O MonteCarloMarginalizeCode/Code/test/integrators/test_mcsampler_rosenbrock.py HANDRUN Rosenbrock sampler study; its docstring pairs it with plot_posterior_corner.py by hand MonteCarloMarginalizeCode/Code/test/test_eosmanager_misc.py OPTDEP needs LALSIMULATION_DATADIR set; raises KeyError at import without it MonteCarloMarginalizeCode/Code/test/test_skysamp.py LEGACY imports lalinference.bayestar.fits, removed upstream; cannot be imported -MonteCarloMarginalizeCode/Code/test/test_mcsampler_foridiots.py BROKEN NameError int_vals at import; a plotting demo that no longer runs at all +# 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/test-core-units.sh b/.travis/test-core-units.sh index c82fbb8d9..ec3a74406 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 @@ -58,11 +61,14 @@ FILES=( "$C/RIFT/likelihood/test_td_dispatch_epoch.py" "$C/test/test_ile_scalar_edge_cases.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" + "$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,6 +83,7 @@ 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" # -- packaging / config contracts / waveform conventions "$C/test/test_advanced_parameter_ports.py" @@ -113,9 +120,23 @@ done # 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 +# 1.26.4, scipy 1.14.1, lal 7.7.0), whole manifest in one run: 296 collected, 284 passed, +# 12 skipped (11 pytest.skip + 1 xfail), ~55 s. (Was 278/266 before test_replica_pooling.py +# and test_marg_list.py joined the manifest -- both were rostered BROKEN until their defects +# were fixed. 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.) +# +# +3/+3 for test_vectorized_lal_tools_split.py: three unconditional test functions, no skip +# and no xfail, numpy / lal / lalsimulation only, so both floors move by the same amount and +# MAX_SKIPPED does not. +# +# +8/+8 for test_noloop_accumulator_shapes.py: two parametrized over three detector networks +# (1/2/3 IFOs) plus two unconditional, so 8 collected and 8 passed, no skip and no xfail. +# numpy / lal / lalsimulation only -- it builds synthetic inputs and calls the NoLoop +# likelihood on the CPU backend, so it needs no cupy and no GPU, and both floors move by the +# same amount while MAX_SKIPPED does not. MEASURED 2026-09-05 on CIT with the same conda +# python as the line above: 8 collected, 8 passed, 0 skipped, 4.9 s. +EXPECTED_TESTS=307 # 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 @@ -126,7 +147,7 @@ EXPECTED_TESTS=278 # 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 +EXPECTED_PASSED=295 MAX_SKIPPED=12 junit="$(mktemp -t core-units-junit-XXXXXX.xml)" diff --git a/.travis/test-integrate.sh b/.travis/test-integrate.sh index 9f4a3cd5a..8346f3837 100755 --- a/.travis/test-integrate.sh +++ b/.travis/test-integrate.sh @@ -72,7 +72,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 +137,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=36 _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 diff --git a/.travis/test-jax.sh b/.travis/test-jax.sh index 96b4efe5a..3e4f1f5ad 100755 --- a/.travis/test-jax.sh +++ b/.travis/test-jax.sh @@ -339,6 +339,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" @@ -349,6 +350,8 @@ FILES=( "${JAXDIR}/test_joint_anglemarg_peaklocal.py" "${JAXDIR}/test_angle_marg_peaklocal_wiring.py" "${JAXDIR}/test_limit_distance_jax.py" + "${JAXDIR}/test_direct_marginalization_planner.py" + "${JAXDIR}/test_time_first_peaklocal.py" ) # EXCLUDED: files in JAXDIR matching test_*.py that are deliberately NOT gated. The @@ -477,7 +480,38 @@ 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. +EXPECTED_TESTS=424 echo "== collection floor check (expect >= ${EXPECTED_TESTS} tests) ==" collect_out="$("${PYTHON_BIN}" -m pytest --collect-only -q -p no:cacheprovider "${DESELECT[@]}" "${FILES[@]}" 2>&1)" 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/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/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/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_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_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/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..a8611e54d 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,100 @@ def marginalization_time_grid(integration_window_half, deltaT, xpy=np): useNR=False distMpcRef = 1000 # a fiducial distance for the template source. + +# --- 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 @@ -2198,7 +2295,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 +2415,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 +2585,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 +2593,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 +2603,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 +2619,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 +2689,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 +2740,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 +2889,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 +2914,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 +2958,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 +2996,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 +3026,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 +3132,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 +3165,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 +3193,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 +3201,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 +3256,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 +3274,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 +3330,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 +3341,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 +3411,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 +3458,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/jax_ile/DESIGN_anglemarg_memory.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_anglemarg_memory.md new file mode 100644 index 000000000..ea8f91924 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/DESIGN_anglemarg_memory.md @@ -0,0 +1,141 @@ +# 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. 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..c9f7203b0 --- /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 distance/angle wrappers currently refuse this ordering | + +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/README.md b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md index c6a98d4d8..844cc6ed1 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/README.md @@ -94,6 +94,13 @@ 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. +`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 @@ -119,6 +126,14 @@ executables without dying during option parsing. lnL over the 5 angular parameters (regulates the amplitude degeneracy; see below). - `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. +- `../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 diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py index c29a4bea3..e9350b3e2 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/anglemarg.py @@ -128,7 +128,8 @@ # 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. +# 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' 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 @@ -161,10 +162,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 @@ -1040,6 +1095,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 @@ -1478,7 +1557,7 @@ def _gh_psi_node_offsets(n_nodes): 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, + phi_chunk=16, dist_block=4, point_block=LAPLACE_POINT_BLOCK, time_quadrature=TIME_QUAD_DEFAULT, return_lnLt=False): """Distance-, phi_ref- AND psi-marginalized lnL: analytic psi-Laplace scheme. @@ -1515,7 +1594,10 @@ def fused_log_likelihood_distphipsimarg_laplace( 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. + 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 @@ -1577,6 +1659,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 +1720,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 @@ -1911,9 +2026,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 @@ -1949,7 +2065,15 @@ def fused_log_likelihood_distphipsimarg_peaklocal( _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). 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..bb7ce7d9e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/direct_marginalization_planner.py @@ -0,0 +1,1204 @@ +"""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 or "distance" in axes): + # Every current JAX distance/angle wrapper calls + # _validate_nonlinear_time_quadrature and refuses bandlimited: its + # primitive fields would have to be refined before the nonlinear + # marginalization. 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/joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/joint_anglemarg_peaklocal.py index 41ceda4aa..d9f6a9f1a 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,10 +38,10 @@ 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. +``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. 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 @@ -35,8 +50,9 @@ 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 +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 +63,26 @@ __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_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 +98,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 +308,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 +331,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 +372,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 +381,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() @@ -285,3 +422,699 @@ def step(carry, args): per_x = jax.scipy.special.logsumexp(vals, axis=0) - jnp.log(n_phi) \ + jnp.log(2.0 * jnp.pi) return jax.scipy.special.logsumexp(per_x + log_w_grid) - 2.0 * jnp.log(2.0 * jnp.pi) + + +# ------------------------------------------------------- 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 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): + """``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) + a0 = jnp.where(peaked, jnp.mod(lo, 2.0 * jnp.pi), big) + crosses = peaked & (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() + Fv, _, _, nfb_v, _, _ = jax.vmap(prof)(jnp.mod(pp, 2.0 * jnp.pi)) + 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 + Fb, d1b, _, nfb_b, nrisk_b, nstrict_b = jax.vmap(prof)(gb) + m1f, m2f = profile_derivative_bounds(C) + ub = Fb + jnp.abs(d1b) * delta + 0.5 * m2f * delta * 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_b.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(), + # 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 other two were invisible: the bound grid GATES (it decides whether the + # certificate is an upper bound at all), the quadrature grid is reported. + "n_u_fallback_bound": nfb_b.sum(), + "n_u_risky_bound": nrisk_b.sum(), + # the stricter 1/M1u criterion: reported, never gated. See u_profile. + "n_u_understood_bound": nstrict_b.sum(), + "n_u_fallback_quad": nfb_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/samplers.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py index b53dc39c6..b295b4210 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py @@ -231,7 +231,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 +240,259 @@ 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. + * 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. + + Returns None when neither is reported. 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 pool that is entirely in use 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. + """ + block = stats.get("largest_free_block_bytes") + if block is not None: + return max(0, int(block)) + pool = stats.get("pool_bytes") or stats.get("bytes_reserved") + if pool: + return max(0, int(pool) - int(stats.get("bytes_in_use") or 0)) + return None + + +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. + """ + 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 + stats = devs[0].memory_stats() or {} + 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. + limit = stats.get("bytes_limit") or stats.get("bytes_reservable_limit") + 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 and the phi scan's stacked output have distinct + shapes, and both 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) + n_phi = _jp.required_n_phi(amp_sizing, m_max=m_max) + + streamed_body = _jp.PHI_CHUNK_DEFAULT * n_x * 4 * n_u_live * 8 + # lax.scan returns every phi chunk before the subsequent reshape/logsumexp; + # that stacked output therefore has (n_phi, n_x) f64 payload. + stacked_scan_output = n_phi * n_x * 8 + return int(streamed_body + stacked_scan_output) def angle_marg_eval_chunk(like, chunk): @@ -261,10 +512,9 @@ 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"): return chunk @@ -273,22 +523,61 @@ def angle_marg_eval_chunk(like, chunk): 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)) + # Besides the streamed (phi_chunk,n_x,4,u_live) body, lax.scan returns + # and stacks every (n_phi,n_x) value before the final reduction. Omitting + # that output undercounts high-amplitude calls because n_phi grows as + # sqrt(A). + 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): 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..05635d15b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/time_first_peaklocal.py @@ -0,0 +1,405 @@ +"""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", + "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 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): + weights = jax.lax.stop_gradient( + _node_weights(plan.live_cells, factor, enum_factor, delta_t)) + n_local = jnp.count_nonzero(weights > 0.0) + capacity_ok = n_local <= int(max_nodes) + index = jnp.nonzero(weights > 0.0, size=int(max_nodes), fill_value=0)[0] + slot_live = jnp.arange(int(max_nodes)) < n_local + index = jax.lax.stop_gradient(index) + slot_live = jax.lax.stop_gradient(slot_live) + + # Reconstruct FIRST, gather SECOND, marginalize other axes LAST. Keeping + # these as three explicit operations is the load-bearing ordering contract. + primitive_fine = reconstruct_time_primitive(kappa_t, factor, guard=guard) + primitive_local = primitive_fine[:, index] + log_t = _lane_log_integrand(primitive_local, rho_sq, log_lane_weight) + local_weight = jnp.where(slot_live, weights[index], 1.0) + terms = jnp.where(slot_live, log_t + jnp.log(local_weight), -jnp.inf) + return jax.scipy.special.logsumexp(terms), n_local, capacity_ok, weights.shape[0] + + +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) + tail_margin = plan.outside_log_bound - value_hi + 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, + "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 value_hi, 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/joint_angle_peak_local.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/joint_angle_peak_local.py index fb20f348e..dcfadb6c7 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 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/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/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index c6db29b7f..96f101187 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -334,6 +334,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 +502,21 @@ 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))) +if opts.q_time_pregrid_factor not in (1, 8): + raise ValueError("--q-time-pregrid-factor currently accepts only 1 or 8") +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 @@ -3486,21 +3503,53 @@ 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: @@ -4037,10 +4086,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): diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax index 00be1a6f8..0e7b467f4 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_jax @@ -324,8 +324,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 " 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/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/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_peaklocal_wiring.py b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py index a23525eb1..7ec4e01df 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_angle_marg_peaklocal_wiring.py @@ -32,10 +32,27 @@ 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) @@ -141,9 +158,9 @@ def test_peak_local_runs_the_runtime_amplitude_failsafe(): def test_peak_local_is_capped_by_the_batch_memory_rule(): """P1 from review. peak-local still nests sample/time vmaps over the distance grid, - phi chunks, four cells and 48 u nodes, so the batch multiplies the same way the - dense schemes do. Leaving it out of the cap kept an uncapped 8000-sample batch and - reopened a documented 36.4 GiB failure.""" + 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 +181,97 @@ 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 and stacks the full phi-scan result; its production- + # floor model is ~216x the dense model's 8192 bytes/sample/time-point. A cap equal + # to exact's would look protective and would not be. The scheme-specific model must + # therefore be STRICTLY tighter. assert capped < S.angle_marg_eval_chunk(_Exact(), 8000), capped # and it must scale with the distance grid, which is what makes it a model rather # than a constant class _Wide(_Like): x_grid = np.zeros(1024) - assert S.angle_marg_eval_chunk(_Wide(), 8000) <= capped + # At this width the corrected body+scan model exceeds the fallback target + # even at S=1. Returning a cap of one would claim protection it cannot give. + with pytest.raises(MemoryError, match="resource preflight"): + S.angle_marg_eval_chunk(_Wide(), 8000) # 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_includes_streamed_body_and_scan_output( + amplitude, n_phi): + """The cap must account for both source-visible peak-local payloads.""" + 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} + + per_point = S._peaklocal_bytes_per_sample_pt(_Like()) + assert per_point == 16 * 256 * 4 * 8 * 8 + n_phi * 256 * 8 + + +def test_peak_local_resource_preflight_refuses_an_unfit_single_sample( + monkeypatch): + """A=12500 needs 5.242 GiB/sample; a cap of one would still OOM 4 GiB.""" + 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": 12500.0} + + 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): + """A=450 needs 1.966 GiB/sample, so the known-4-GiB target admits only one.""" + 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} + + 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 @@ -195,3 +290,77 @@ def test_peak_local_artifacts_carry_the_standing_best_effort_label(): note = mod.angle_grid_suspect_note("peak-local") assert note.startswith("ANGLE-GRID-CHECK=BEST-EFFORT"), note 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_anglemarg_buffer_cap.py b/MonteCarloMarginalizeCode/Code/test/jax/test_anglemarg_buffer_cap.py new file mode 100644 index 000000000..0d5d21d08 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_anglemarg_buffer_cap.py @@ -0,0 +1,531 @@ +#!/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_is_a_known_full_device(monkeypatch): + """A reported zero is availability data, not a missing probe result.""" + _fake_jax(monkeypatch, devices=[_Dev("gpu", 24 * GIB, free=0)]) + assert sam._angle_marg_buffer_target() == 0 + with pytest.raises(MemoryError): + sam.angle_marg_eval_chunk(_Like("laplace", 1193), 4000) + + +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 = 1769472 +#: ... and one sample of it at production npts=1230. 2.027 GiB. +PEAKLOCAL_ONE_SAMPLE = 2176450560 + + +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 + n_phi = jp.required_n_phi(anglemarg.ANGLE_MARG_CROSSOVER_AMPLITUDE, + m_max=2) + modeled = streamed + n_phi * 256 * 8 + 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) # 13.5 GiB per sample, over the 4 GiB guess + 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) == 2 + + +# --------------------------------------------------------------------------- +# 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() 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_joint_anglemarg_peaklocal.py b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py index be770644b..6c4e90152 100644 --- a/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_joint_anglemarg_peaklocal.py @@ -145,3 +145,482 @@ 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"): + 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 now the margin test AND the resolution test. + assert bool(ok) == (float(info["margin"]) < JP.OUTSIDE_TOL_NATS + and bool(info["phi_resolved"])) + 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_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 + assert len(calls) == 4, (len(calls), "a fifth grid means a probe is paying its own way") + 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_gates_on_the_fallback_that_can_invert_it(): + """Adversarial review: ``Fb`` and ``d1b`` were taken from ``u_profile`` with its + whole-cell fallback and the count was DISCARDED at that call, so a row could be + accepted on a lift applied to an underestimated profile with no signal it had + happened. ``info["n_u_fallback"]`` carried only the Newton-seed evaluation. + + The remedy as stated -- decline whenever any bound-grid profile falls back -- is not + implementable: every generic table has four u-stationary points of which two are + minima, so the fallback count is never zero and that gate declines universally + (measured: 0 of 2 accepted on cases accurate to 1e-5). A minimum cell has no peak to + window and is exponentially subdominant in F; the cells that can invert the bound are + those with ``g'' < 0`` that failed the stationarity or interior test, because a real + maximum may sit in one unresolved. + + Nor is "did a max-bearing cell fall back" the question: an 8-step Newton misses the + 1e-8 relative residual on plenty of ordinary maxima, and that test fired on 127 of 256 + bound-grid points for tables accurate to 1e-5. What the bound needs is review's other + remedy -- whether the whole-cell quadrature was ADEQUATE -- and that is exact here, + because the u spectrum has two terms so ``|d2g/du2| <= |c1| + 4|c2|`` everywhere and a + cell of ``width`` needs ``width sqrt(M2u) U_PTS_PER_SIGMA`` nodes. + + So this test pins BOTH directions: a case that must accept with a non-zero fallback + count, and a case where the gate fires and is CLEARED by sizing the quadrature. + """ + C, exact = _separable_phi_table(1000.0, np.pi / 96) + v, ok, info = JP.phi_local_lnI(C, w_sigma=200.0, n_nodes=385) + assert abs(float(v) - exact) < 1e-4 + assert int(info["n_u_fallback_bound"]) > 0, "the naive gate would have fired here" + assert int(info["n_u_risky_bound"]) == 0 + assert bool(ok), "gating on the whole-cell count declines every table there is" + + +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=1024) + 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 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..bd586c509 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_time_first_peaklocal.py @@ -0,0 +1,152 @@ +"""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 primitive -> gather -> downstream-reduction stages. + source = inspect.getsource(TFP._evaluate_cover_at_factor) + assert source.index("reconstruct_time_primitive") < source.index( + "_lane_log_integrand") + + +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_joint_angle_peak_local.py b/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py index 997e69c4b..1e6c770d1 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py +++ b/MonteCarloMarginalizeCode/Code/test/test_joint_angle_peak_local.py @@ -4,10 +4,14 @@ of the shipped `anglemarg` exact scheme -- so accuracy is measured against a quadrature, never against another peak-local run. """ +import builtins + import numpy as np import pytest +from scipy import special from RIFT.likelihood import joint_angle_peak_local as J +from RIFT.likelihood import bivariate_trig_stationary as BTS def synth_table(seed=0, scale=1.0, bidegree=(4, 2)): @@ -18,6 +22,197 @@ def synth_table(seed=0, scale=1.0, bidegree=(4, 2)): return scale * C +def _periodic_set_error(got, want): + """Symmetric nearest-neighbour error for two small point sets on the torus.""" + got = np.asarray(got, dtype=float).reshape((-1, 2)) + want = np.asarray(want, dtype=float).reshape((-1, 2)) + if len(got) != len(want): + return np.inf + d = (got[:, None, :] - want[None, :, :] + np.pi) % (2 * np.pi) - np.pi + r = np.linalg.norm(d, axis=-1) + return max(float(np.max(np.min(r, axis=0))), + float(np.max(np.min(r, axis=1)))) + + +def _separable_table(m=3, n=2, a=7.0, b=4.0): + """Exactly ``a cos(m phi) + b cos(n u)`` in the RIFT storage convention.""" + C = np.zeros((m + 1, 2 * n + 1), dtype=complex) + C[m, n] = 0.5 * a # k>0 is doubled by the evaluator + C[0, 2 * n] = b # q=+n; taking Re supplies q=-n + return C + + +def test_algebraic_canonical_table_matches_the_shipped_field_convention(): + """Laurent conversion is exact, including k=0 overlap and both q signs.""" + C = synth_table(seed=91, scale=2.3, bidegree=(3, 2)) + A = BTS.canonical_laurent_table(C) + rng = np.random.default_rng(123) + p = rng.uniform(0, 2 * np.pi, size=(37, 2)) + k = np.arange(-3, 4)[None, :, None] + q = np.arange(-2, 3)[None, None, :] + full = np.real(np.sum( + A[None] * np.exp(1j * (p[:, 0, None, None] * k + + p[:, 1, None, None] * q)), axis=(1, 2))) + assert np.allclose(full, J.eval_g(C, p[:, 0], p[:, 1]), rtol=0, atol=2e-13) + + +def test_algebraic_enumerator_preserves_every_codominant_separable_maximum(): + """Generic projection must not collapse modes sharing the same phi or u. + + A coordinate resultant has repeated projected roots on this Cartesian mode + lattice. The affine hidden variable separates them and returns all six equal + maxima, without any angular samples. + """ + C = _separable_table(m=3, n=2) + out = BTS.enumerate_torus_maxima(C) + want = np.array([(2 * np.pi * j / 3, np.pi * k) + for j in range(3) for k in range(2)]) + assert out.ok, out.report + assert out.report["mixed_volume"] == 24 + assert out.stationary_points.shape == (24, 2) + assert out.points.shape == (6, 2) + assert _periodic_set_error(out.points, want) < 2e-9 + assert np.ptp(out.values) < 2e-12 + + +def test_algebraic_enumerator_resolves_near_annihilating_stationary_points(): + """A close max/min pair is part of the polynomial, not a resolution choice.""" + ratio = 3.9999999 + C = np.zeros((3, 5), dtype=complex) + C[1, 2] = 0.5 * ratio + C[2, 2] = 0.5 + C[0, 4] = 2.0 + out = BTS.enumerate_torus_maxima(C) + assert out.ok, out.report + assert out.stationary_points.shape == (16, 2) + assert out.points.shape == (4, 2) + + # The two additional phi stationary points approach pi from either side. + # Their separation is smaller than a 4096-point circle spacing; retaining + # both demonstrates that no sampled phi resolution controls enumeration. + expected_phi = np.mod(np.array([ + 0.0, np.pi, + np.arccos(-ratio / 4.0), + 2.0 * np.pi - np.arccos(-ratio / 4.0), + ]), 2.0 * np.pi) + got_phi = np.unique(np.round(out.stationary_points[:, 0], 11)) + circ = np.abs((got_phi[:, None] - expected_phi[None, :] + np.pi) + % (2 * np.pi) - np.pi) + assert got_phi.size == 4 + assert np.max(np.min(circ, axis=0)) < 2e-8 + close_sep = 2.0 * (np.pi - np.arccos(-ratio / 4.0)) + assert close_sep < 2.0 * np.pi / 4096 + + +def test_algebraic_enumerator_declines_at_exact_stationary_degeneracy(): + """At annihilation certification declines, while safe targets stay available.""" + C = np.zeros((3, 5), dtype=complex) + C[1, 2] = 2.0 # ratio c1/c2 == 4 exactly + C[2, 2] = 0.5 + C[0, 4] = 2.0 + out = BTS.enumerate_torus_maxima(C) + assert not out.ok + assert out.points.shape[0] == 2 + assert all(p["decline"] is not None for p in out.report["projections"]) + assert min(p["min_jacobian_rcond"] for p in out.report["projections"]) < 2e-10 + assert np.all(np.linalg.eigvalsh(out.hessians) < 0.0) + + +def test_algebraic_enumeration_size_and_modes_are_amplitude_independent(): + """Scaling the exponent changes widths, never its physical torus modes. + + Recovery of every off-torus complex BKK root by two independent QZ + projections is a conservative certification diagnostic, not a physical + invariant. Tiny platform-dependent roundoff after normalization may make + one projection decline while both solves retain the same torus stationary + points and maxima. The production path remains fail closed in that case. + """ + C = synth_table(seed=17, bidegree=(2, 2)) + low = BTS.enumerate_torus_maxima(C) + high = BTS.enumerate_torus_maxima(1.0e8 * C) + assert low.report["mixed_volume"] == high.report["mixed_volume"] == 32 + assert [p["pencil_size"] for p in low.report["projections"]] == [ + p["pencil_size"] for p in high.report["projections"]] + assert low.stationary_points.shape == high.stationary_points.shape == (24, 2) + assert low.points.shape == high.points.shape == (6, 2) + assert _periodic_set_error( + low.stationary_points, high.stationary_points) < 2e-8 + assert _periodic_set_error(low.points, high.points) < 2e-8 + + +def test_incomplete_algebraic_accounting_never_drops_the_likelihood_sample(): + """A root deficit is either cover-certified or sent to the dense fallback.""" + C = synth_table(seed=3, scale=1.0) + value, ok, report = J.joint_marginalize_peak_local( + C, n_phi=64, n_bound_grid=128) + assert ok and np.isfinite(value), report + assert not report["enumeration_certified"] + assert report["result_path"] in { + "algebraic-best-effort/bound-certified", "dense-phi/exact-u"} + projections = report["enumeration"]["projections"] + assert any(p["verified_complex_roots"] < p["expected_roots"] + for p in projections) + assert all("min_jacobian_rcond" in p for p in projections) + if report["result_path"] == "algebraic-best-effort/bound-certified": + assert report["margin"] < J.OUTSIDE_TOL_NATS + else: + assert report["fallback_reason"] + + +def test_certified_enumeration_cannot_certify_capped_local_quadrature(monkeypatch): + """A complete root set does not certify integration inside its cover. + + ``s cos(phi-u) + cos(phi+u)`` has the exact normalized integral + ``I0(s) I0(1)``. Its weak direction makes the mode boxes cover the torus, + while its strong diagonal direction is much narrower than either capped + axis-aligned rule. The outside ledger therefore says that no area was + omitted even though the local quadrature is unresolved. + """ + strength = 1.0e8 + C = np.zeros((2, 5), dtype=complex) + C[1, 1] = 0.5 * strength # strength * cos(phi-u) + C[1, 3] = 0.5 # cos(phi+u) + exact = (np.log(special.i0e(strength)) + strength + + np.log(special.i0e(1.0)) + 1.0) + fallback_calls = [] + + def finite_fallback(table, n_phi=None, n_u_nodes=64): + fallback_calls.append((table, n_phi, n_u_nodes)) + return exact, {"doubling_error": 0.0, "fixture": "known integral"} + + monkeypatch.setattr(J, "dense_phi_exact_u_marginalize", finite_fallback) + value, ok, report = J.joint_marginalize_peak_local(C) + + assert ok and value == exact + assert report["enumeration_certified"], report["enumeration"] + assert report["area_outside"] == 0.0 + assert report["n_boxes_pts_capped"] >= 1 + assert report["local_quadrature_capped"] + assert report["local_quadrature_error"] > 0.1 + assert report["result_path"] == "dense-phi/exact-u" + assert "inside-cover quadrature did not converge" in report["fallback_reason"] + assert len(fallback_calls) == 1 + + +def test_numpy_dense_fallback_does_not_import_the_optional_jax_stack(monkeypatch): + """The final fallback remains usable in an installation without JAX.""" + real_import = builtins.__import__ + + def reject_jax(name, globals=None, locals=None, fromlist=(), level=0): + if "jax_ile" in name or name == "jax" or name.startswith("jax."): + raise AssertionError("the NumPy fallback attempted to import JAX") + return real_import(name, globals, locals, fromlist, level) + + monkeypatch.setattr(builtins, "__import__", reject_jax) + C = np.zeros((2, 5), dtype=complex) + C[1, 2] = 0.25 + value, report = J.dense_phi_exact_u_marginalize(C, n_phi=16) + + assert np.isfinite(value) + assert report["n_phi_coarse"] == 128 + assert report["n_phi"] == 256 + + def _ref(C, n=2048): """log[(2pi)^-2 int int exp(g)] by the periodic trapezoid (== the plain mean).""" t = np.linspace(0.0, 2.0 * np.pi, n, endpoint=False) @@ -82,9 +277,9 @@ def test_eval_g_chunking_cannot_change_the_answer(): assert a.tobytes() == b.tobytes() -def test_an_undersized_region_is_DECLINED_not_returned(): +def test_an_undersized_region_falls_back_instead_of_returning_local_value(): """The load-bearing behaviour. W_SIGMA too small leaves mass outside the cover; - the value may still be right, but the rule cannot PROVE it and must decline. + the local value may still be right, but the rule cannot PROVE it and must fall back. Measured on the shipped tables: at W = 8 the margin is -18 nats against a -23 tolerance, and at 14 it is -71 -- with the returned value identical at both.""" C = synth_table(seed=3, scale=12.0) @@ -96,9 +291,11 @@ def test_an_undersized_region_is_DECLINED_not_returned(): val_big, ok_big, _ = J.joint_marginalize_peak_local(C, n_phi=96) finally: J.W_SIGMA = keep - assert not ok_small, rep_small - assert rep_small['decline'] == 'omitted-mass bound too large' + assert ok_small, rep_small + assert rep_small['result_path'] == 'dense-phi/exact-u' + assert 'omitted-mass bound too large' in rep_small['fallback_reason'] assert ok_big + assert abs(val_small - val_big) < 1e-6 def test_regions_merge_rather_than_double_counting(): @@ -424,3 +621,99 @@ def test_phi_regions_are_disjoint_on_the_CIRCLE(): d = np.minimum(d, 2 * np.pi - d) assert d.min() > 1e-6, ("regions overlap on the circle", regs) assert checked > 5, checked + + +# --------------------------------------------- internal accuracy inside the cover + +def _torus_reference(C, n=2048): + """log (2pi)^-2 int int exp(g) over the WHOLE torus, independent of the peak-local path.""" + ph = np.linspace(0.0, 2.0 * np.pi, n) + P, U = np.meshgrid(ph, ph, indexing='ij') + g = J.eval_g(C, P.ravel(), U.ravel()).reshape(n, n) + w = np.full(n, 2.0 * np.pi / (n - 1)); w[0] *= 0.5; w[-1] *= 0.5 + W = np.log(w)[:, None] + np.log(w)[None, :] - 2.0 * np.log(2.0 * np.pi) + m = g.max() + return m + np.log(np.sum(np.exp(g - m + W))) + + +def _degenerate_ridge_tables(seed=113, scale=1.0): + """SYNTHETIC coefficients reproducing the torus-spanning collapse. No run data. + + An earlier version of this fixture hard-coded coefficients read out of an actual + production evaluation, together with its sky/time indices. External review was right + that merging it would publish run-derived scientific data in a test, so it is replaced + by a seeded synthetic draw. + + What could NOT be replaced is the structure, and it took a seed search to find it. A + hand-built table with the same sparsity PATTERN -- A only at k=2 with q=+-1 and + strongly asymmetric, B almost entirely the real (k=0,ks=0) term -- does not reproduce + the collapse: with round numbers it gives n_regions=4, area_outside=31.7 and no error + at all. The relative PHASES decide whether the enumerated regions merge into one that + spans the torus, so the fixture is a search over seeded phases for a draw that does. + Seed 113 of 200 is the strongest. This is why random-coefficient tests never reached + this branch: the landscape is a near-degenerate ridge, not isolated peaks. + """ + rng = np.random.default_rng(seed) + A = np.zeros((3, 3), dtype=complex) + B = np.zeros((5, 5), dtype=complex) + A[2, 2] = 3000.0 * scale * np.exp(1j * rng.uniform(0.0, 2 * np.pi)) + A[2, 0] = A[2, 2] * 0.013 * np.exp(1j * rng.uniform(0.0, 2 * np.pi)) + B[0, 2] = 1550.0 * scale + B[0, 0] = 21.5 * scale * np.exp(1j * rng.uniform(0.0, 2 * np.pi)) + B[0, 4] = np.conj(B[0, 0]) + B[4, 2] = 0.0036 * scale * np.exp(1j * rng.uniform(0.0, 2 * np.pi)) + B[4, 4] = 0.0886 * scale * np.exp(1j * rng.uniform(0.0, 2 * np.pi)) + k, q, w, _ = J._kq(A) + x = float(np.sum(w * np.abs(A))) / float(B[0, 2].real) + return J.joint_table(A, B, x), x + + +def test_a_fully_covered_box_is_still_accurate_inside(): + """OMITTED-MASS CONTROL IS NOT INTERNAL ACCURACY, and this is the case that proves the + two are independent. On the production tables the cover collapses to a single region + spanning the whole torus, so ``area_outside == 0`` and ``margin == -inf``: the + certificate reports that NOTHING is omitted, which is true and which says nothing at + all about the quadrature inside. With the per-axis cap at its old value of 256 the + value sat 0.36 nats from a converged reference while reporting -inf. + + 0.36 nats is not a rounding error. It is stated against the CONVERGED TORUS REFERENCE + below and against nothing else: an earlier version of this docstring compared it to a + saddle-point prototype's 0.654 nats, and that figure has since been RETRACTED by the + session that produced it -- its start-point search was unconverged, moving up to 1.2 + nats per point and changing sign under refinement. A ratio against a retracted + denominator is worse than no ratio, and this error needs no comparison to be a defect: + the certificate reported nothing omitted while the value was wrong. + """ + C, _ = _degenerate_ridge_tables() + assert abs(np.sum(np.abs(C)) - 24164.9) < 1.0, "fixture drifted" + lnZ, ok, rep = J.joint_marginalize_peak_local(C) + assert ok, rep + # A complete algebraic cover retains the original inside-box regression. An + # incomplete solve may expose less covered area; the new hierarchy must then + # take the finite dense fallback instead of treating the row as -inf. + if rep['result_path'].startswith('algebraic'): + assert rep['area_outside'] == 0.0, rep + assert rep['margin'] == -np.inf, rep + else: + assert rep['result_path'] == 'dense-phi/exact-u', rep + assert rep['fallback_reason'], rep + assert rep['dense_fallback']['doubling_error'] < 1e-4, rep + err = abs(lnZ - _torus_reference(C)) + assert err < 1.0e-2, "inside-the-cover error %.4f nats (cap 256 gave 0.36)" % err + + +def test_a_capped_box_is_reported_and_never_silent(): + """A box whose curvature-derived node count hits the ceiling is under-resolved, and the + certificate cannot express that. It must therefore be COUNTED -- otherwise the caller + is handed 'nothing omitted' about a value the quadrature got wrong. + """ + C, _ = _degenerate_ridge_tables() + _, ok, rep = J.joint_marginalize_peak_local(C) + assert ok + assert 'n_boxes_pts_capped' in rep + assert rep['n_boxes_pts_capped'] >= 1, rep # this amplitude DOES still cap at 512 + # and a much flatter case must NOT be flagged, or the counter says nothing + C_lo, _ = _degenerate_ridge_tables(scale=1.0e-4) + _, ok2, rep2 = J.joint_marginalize_peak_local(C_lo) + assert ok2, rep2 + assert rep2['n_boxes_pts_capped'] == 0, rep2 diff --git a/MonteCarloMarginalizeCode/Code/test/test_noloop_accumulator_shapes.py b/MonteCarloMarginalizeCode/Code/test/test_noloop_accumulator_shapes.py new file mode 100644 index 000000000..3079e4361 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_noloop_accumulator_shapes.py @@ -0,0 +1,187 @@ +"""NoLoop's accumulator shapes are an optimization, not a change of arithmetic. + +Two accumulators inside `DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` were +changed for memory traffic, not for numerics: + + * `rho_sq` is time-independent, so it is summed as an `(n_extrinsic,)` vector and + exposed to consumers as a stride-0 `(n_extrinsic, npts)` view instead of being + materialized; + * `kappa_sq` starts from the first detector's own (in-place scaled) buffer instead + of being zero-filled and accumulated into. + +Both must leave the answer alone. This runs the real function on small synthetic +inputs and compares it against a reference written the original way, so a future edit +that quietly changes the arithmetic of either accumulator fails here rather than in +someone's posterior. +""" +import numpy as np +import pytest + +import RIFT.likelihood.factored_likelihood as fl + + +class _P(object): + """The handful of attributes NoLoop actually reads off a ChooseWaveformParams.""" + + def __init__(self, n, rng): + self.phi = rng.uniform(0.0, 2.0 * np.pi, n) # right ascension + self.theta = np.arcsin(rng.uniform(-1.0, 1.0, n)) # declination + self.phiref = rng.uniform(0.0, 2.0 * np.pi, n) + self.incl = np.arccos(rng.uniform(-1.0, 1.0, n)) + self.psi = rng.uniform(0.0, np.pi, n) + self.dist = rng.uniform(200.0, 900.0, n) * 1e6 * 3.0856775814913673e16 + self.tref = 1000000014.0 + self.deltaT = 1.0 / 4096.0 + + +def _inputs(n_ex=64, npts=32, n_time=512, dets=("H1", "L1", "V1"), seed=20260905): + rng = np.random.RandomState(seed) + lms = np.array([[2, 2], [2, -2]], dtype=np.int64) + n_lm = len(lms) + rholms, ctU, ctV, lookup, epoch = {}, {}, {}, {}, {} + for d in dets: + rholms[d] = (rng.normal(size=(n_lm, n_time)) + + 1j * rng.normal(size=(n_lm, n_time))) + a = rng.normal(size=(n_lm, n_lm)) + 1j * rng.normal(size=(n_lm, n_lm)) + ctU[d] = a + a.conj().T # Hermitian, as U is + ctV[d] = rng.normal(size=(n_lm, n_lm)) + 1j * rng.normal(size=(n_lm, n_lm)) + lookup[d] = lms + # tref - 0.05, NOT a whole second earlier: t_det = (tref - epoch) + light + # travel, and ifirst = (t_det + tvals[0])/deltaT must land INSIDE the + # n_time buffer. At a 1 s offset ifirst was ~3979-4152 against n_time=512, + # so every window was zero-extended, kappa_sq was identically zero, and the + # kappa_sq half of this file asserted nothing at all. Verified by + # test_data_term_is_actually_exercised below. + epoch[d] = 1000000014.0 - 0.05 + tvals = np.linspace(-0.0075, 0.0075, npts) + return tvals, _P(n_ex, rng), lookup, rholms, ctU, ctV, epoch + + +def _reference(tvals, P, lookup, rholms, ctU, ctV, epoch, Lmax=2, integrate=True): + """The pre-optimization arithmetic: dense rho_sq, zero-initialized kappa_sq. + + With ``integrate=False`` it stops at lnL(t), before the time quadrature, which is + the only part of the chain that is deliberately not bit-exact. + """ + import lal + import lalsimulation as lalsim + from RIFT.likelihood.SphericalHarmonics_gpu import SphericalHarmonicsVectorized + from RIFT.likelihood.vectorized_lal_tools import ( + ComputeDetAMResponse, TimeDelayFromEarthCenter) + + npts = len(tvals) + n_ex = len(P.phi) + distMpc = P.dist / (lal.PC_SI * 1e6) + invDist = fl.distMpcRef / distMpc + gmst = np.asarray(lal.GreenwichMeanSiderealTime(P.tref)) + + kappa_sq = np.zeros((n_ex, npts), dtype=np.complex128) + rho_sq = np.zeros((n_ex, npts), dtype=np.float64) + + for det in rholms: + d = lalsim.DetectorPrefixToLALDetector(det) + Ylm = SphericalHarmonicsVectorized( + lookup[det], P.incl, -P.phiref, xpy=np, l_max=Lmax) + F = ComputeDetAMResponse(np.asarray(d.response), P.phi, P.theta, P.psi, + gmst, xpy=np) + t_det = float(P.tref - float(epoch[det])) + TimeDelayFromEarthCenter( + np.asarray(d.location), P.phi, P.theta, float(gmst), xpy=np) + ifirst = (np.rint((t_det + tvals[0]) / P.deltaT) + 0.5).astype(np.int32) + + rho_det = ((F * np.conj(F)).real + * np.einsum("...i,...j,ij", np.conj(Ylm), Ylm, ctU[det]).real) + rho_det += (np.square(F) + * np.einsum("...i,...j,ij", Ylm, Ylm, ctV[det])).real + rho_det *= 0.5 * np.square(fl.distMpcRef / distMpc) + + Qlms = fl._nearest_Q_window_numpy(rholms[det].T, ifirst, npts, xpy=np) + FY = np.broadcast_to((F[..., None] * Ylm)[:, None], Qlms.shape) + kappa_sq += np.einsum("...i,...i", np.conj(FY), Qlms) * invDist[..., None] + rho_sq += rho_det[..., None] + + lnL_t = kappa_sq.real - 0.5 * rho_sq + if not integrate: + return lnL_t + lnLmax = np.max(lnL_t, axis=-1, keepdims=True) + L = fl.my_simps(np.exp(lnL_t - lnLmax), dx=P.deltaT, axis=-1) + return (lnLmax[:, 0] + np.log(L)) + + +@pytest.mark.parametrize("dets", [("H1",), ("H1", "L1"), ("H1", "L1", "V1")]) +def test_accumulators_are_bit_exact(dets): + """The accumulators themselves must be exact, so check lnL(t) BEFORE the integral. + + Taking the comparison at return_lnLt=True is what makes this a test of the + accumulators rather than of the quadrature: the time integral is a matvec against + precomputed Simpson weights and is deliberately not bit-exact (see the quadrature + test below), so integrating first would blur the two and this test would have to be + weakened to a tolerance it does not need. + """ + args = _inputs(dets=dets) + got = fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + *args, Lmax=2, xpy=np, return_lnLt=True) + want = _reference(*args, integrate=False) + assert np.array_equal(np.asarray(got), want) + + +@pytest.mark.parametrize("dets", [("H1",), ("H1", "L1"), ("H1", "L1", "V1")]) +def test_time_quadrature_matches_simps_to_roundoff(dets): + """The matvec quadrature reproduces simps() to floating-point noise. + + It is the SAME rule -- the weights come from the very simps implementation the call + site would otherwise have used -- so only the summation order differs. The bound + here is deliberately far tighter than anything that matters physically: the two + simps variants already 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. If this + assertion ever fails it means the RULE changed, not that rounding drifted. + """ + args = _inputs(dets=dets) + got = np.asarray(fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + *args, Lmax=2, xpy=np)) + want = _reference(*args) + assert np.allclose(got, want, rtol=1e-11, atol=1e-11) + + +def test_simps_weights_reproduce_simps_on_random_data(): + """simps is linear at fixed dx, which is the whole basis for the matvec.""" + rng = np.random.RandomState(7) + npts, dx = 614, 1.0 / 4096.0 # production shape and spacing + y = rng.normal(size=(23, npts)) + w = fl._simps_weights(fl.my_simps, npts, dx, np) + assert np.allclose(y.dot(w), fl.my_simps(y, dx=dx, axis=-1), rtol=1e-12, atol=0.0) + + +def test_data_term_is_actually_exercised(): + """Guard the trap this file fell into: a Q window entirely outside the buffer. + + `ifirst` is derived from (tref - epoch) plus light travel. If the synthetic inputs + put it past `n_time`, every window is zero-extended, kappa_sq is identically zero, + and every assertion above still passes while testing only rho_sq -- which is exactly + what happened on the first version of this file. A zero data term shows up as lnL(t) + with no variation along the time axis, so assert the variation directly. + """ + args = _inputs() + lnL_t = np.asarray(fl.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop( + *args, Lmax=2, xpy=np, return_lnLt=True)) + spread = np.ptp(lnL_t, axis=-1) + assert np.median(spread) > 1.0, ( + "lnL(t) is flat in time: the data term is zero, so the kappa_sq assertions " + "above are vacuous. Check epoch vs tref against n_time in _inputs().") + + +def test_dense_rho_sq_returns_writable_contiguous_memory(): + """`_dense_rho_sq` exists to hand real memory to consumers that need it. + + The fused CUDA kernels index raw device pointers and the non-Simpson quadrature + helpers may write, so for them a stride-0 view is not merely slow but wrong. Note + that a broadcast view being READ-ONLY is a numpy guarantee and NOT a cupy one -- + measured: cupy.broadcast_to yields strides (8, 0) and writes through to the base -- + so the contract this pins is what _dense_rho_sq RETURNS, not what the view forbids. + """ + vec = np.arange(5.0) + view = np.broadcast_to(vec[:, None], (5, 7)) + assert view.strides[-1] == 0 # the thing being avoided is real + dense = np.ascontiguousarray(view) + assert dense.flags.c_contiguous and dense.flags.writeable + assert dense.strides[-1] != 0 + assert np.array_equal(dense, view) diff --git a/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py index b9edd6a11..bbb964ead 100644 --- a/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py +++ b/MonteCarloMarginalizeCode/Code/test/test_time_marginalization_quadrature.py @@ -31,6 +31,7 @@ import os import sys +import warnings import numpy as np import pytest @@ -224,6 +225,120 @@ def test_reflected_upsample_reproduces_the_finite_row_exactly(): assert np.allclose(up[::factor], k, atol=1e-11, rtol=0) +@pytest.mark.parametrize("n,factor", [(614, 64), (1228, 32), (2457, 16)]) +def test_retained_fft_is_the_same_reflected_sinc_interpolant(n, factor): + """Pruning outputs must not change the reconstruction being evaluated. + + Random complex rows populate every bin, including Nyquist, so this compares + the half-bin convention too. The production 22 and higher-mode window sizes + are explicit rather than hidden behind a toy power-of-two transform. + """ + rng = np.random.default_rng(90210 + n) + rows = rng.normal(size=(2, n)) + 1j * rng.normal(size=(2, n)) + reference = tmq.reflected_bandlimited_upsample(rows, factor) + retained = tmq._reflected_bandlimited_upsample_retained(rows, factor) + assert retained.shape == reference.shape == (2, (n - 1) * factor + 1) + np.testing.assert_allclose(retained, reference, rtol=0, atol=2e-11) + np.testing.assert_allclose(retained[..., ::factor], rows, rtol=0, atol=2e-11) + + +def test_retained_fft_removes_the_production_padding_mismatch(): + period, factor = 2 * NPTS, 64 + plan = tmq._retained_fft_plan(period, factor, np.complex128) + assert plan['n_out'] == (NPTS - 1) * factor + 1 == 39233 + assert period * factor == 78592 + # The optimized convolution is close to the retained half, not the discarded + # full reflected period. Do not pin scipy's exact next-fast-length policy. + assert plan['n_out'] <= plan['n_fft'] < 0.53 * period * factor + + +def test_transform_decline_retries_full_sinc_and_reports_provenance(monkeypatch): + """An optimization decline is not a failed waveform point. + + Warning-as-error is included because a diagnostic warning must not undo the + successful reference-path retry in production environments with strict + warning filters. + """ + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples()[None, :] + rho = np.full(k.shape, RHO_SQ) + retained = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + assert tmq.last_report()['bandlimited_fft_strategy'] == 'retained-grid-zoomfft' + + def decline(*args, **kwargs): + raise tmq._RetainedFFTUnsupported('forced unsupported transform') + + monkeypatch.setattr(tmq, '_reflected_bandlimited_upsample_retained', decline) + with warnings.catch_warnings(): + warnings.simplefilter('error') + fallback = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + report = tmq.last_report() + assert np.isfinite(float(fallback[0])) + np.testing.assert_allclose(fallback, retained, rtol=0, atol=1e-9) + assert report['bandlimited_fft_strategy'] == 'full-padding-fallback' + assert report['retained_fft_batches'] == 0 + assert report['full_fft_fallback_batches'] >= 1 + assert report['full_fft_fallback_rows'] >= 1 + assert report['full_fft_fallback_reasons'] == { + '_RetainedFFTUnsupported: forced unsupported transform': + report['full_fft_fallback_rows']} + + +def test_unsupported_factor_falls_back_to_full_sinc_without_changing_values(): + rng = np.random.default_rng(19) + rows = rng.normal(size=(2, 17)) + 1j * rng.normal(size=(2, 17)) + report = tmq._new_transform_report() + with pytest.warns(RuntimeWarning, match='full-padding sinc reconstruction'): + got = tmq._reflected_upsample_for_integration( + rows, 3, {}, report, xpy=np) + reference = tmq.reflected_bandlimited_upsample(rows, 3) + np.testing.assert_array_equal(got, reference) + assert report['retained_fft_batches'] == 0 + assert report['full_fft_fallback_batches'] == 1 + assert '_RetainedFFTUnsupported' in next(iter( + report['full_fft_fallback_reasons'])) + + +def test_low_factor_numpy_reference_is_selected_not_mislabeled_as_failure(): + rng = np.random.default_rng(23) + rows = rng.normal(size=(2, 17)) + 1j * rng.normal(size=(2, 17)) + report = tmq._new_transform_report() + got = tmq._reflected_upsample_for_integration( + rows, 4, {}, report, xpy=np) + reference = tmq.reflected_bandlimited_upsample(rows, 4) + np.testing.assert_array_equal(got, reference) + assert report['full_fft_selected_batches'] == 1 + assert report['full_fft_selected_rows'] == 2 + assert report['full_fft_fallback_batches'] == 0 + assert report['full_fft_fallback_reasons'] == {} + assert 'measured retained-FFT crossover' in next(iter( + report['full_fft_selected_reasons'])) + + +def test_likelihood_failure_is_not_mislabeled_as_transform_fallback(monkeypatch): + """Only the transform is guarded; callback failures retain their identity.""" + class LikelihoodFailure(RuntimeError): + pass + + sig = BandLimited(amp=0.17, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + k = sig.samples()[None, :] + rho = np.full(k.shape, RHO_SQ) + + def fail_on_dense(kappa_term, rho_sq): + if kappa_term.shape[-1] == NPTS: + return _lnL(kappa_term, rho_sq) + raise LikelihoodFailure('callback, not FFT') + + def fallback_must_not_run(*args, **kwargs): + pytest.fail('a likelihood exception was incorrectly retried as an FFT decline') + + monkeypatch.setattr(tmq, 'reflected_bandlimited_upsample', fallback_must_not_run) + with pytest.raises(LikelihoodFailure, match='callback, not FFT'): + tmq.time_marginalize_bandlimited(k, rho, DELTAT, fail_on_dense) + + def test_forward_backward_reflection_blocks_the_endpoint_gibbs_counterexample(): """A decayed integrand does not imply a periodic kappa slice. @@ -399,6 +514,25 @@ def test_rows_sharing_a_block_keep_their_individual_resolution(): assert sum(hist.values()) == 2, hist +def test_simpson_fallback_is_evaluated_only_for_unrefined_rows(): + """Dense rows must not also pay for a coarse integration that is discarded.""" + sharp = BandLimited(amp=1.0, peak_sample=NPTS // 2 + 0.25, + n_period=8 * NPTS, m_hi=1400, background=0.12) + flat = np.full(NPTS, 0.12 + 0.0j) + k = np.stack((sharp.samples(), flat)) + calls = [] + + def recording_simps(y, dx, axis): + calls.append(np.asarray(y).shape) + return simpson(y, dx=dx, axis=axis) + + got = tmq.time_marginalize_bandlimited( + k, np.full(k.shape, RHO_SQ), DELTAT, _lnL, simps=recording_simps) + assert got.shape == (2,) + assert tmq.last_report()['n_refined_rows'] == 1 + assert calls == [(1, NPTS)], calls + + # ------------------------------------------------------------- preconditions def test_time_dependent_rho_sq_is_refused(): @@ -747,6 +881,30 @@ def test_remeasure_on_the_dense_grid_repairs_an_under_derived_factor(): assert abs(float(got[0]) - sig.truth()) < 1e-6 +def test_dense_remeasurement_refines_only_the_rows_that_still_need_it(): + """One pathological row must not impose its extra FFT octaves on a group.""" + signals = [ + BandLimited(amp=0.2, peak_sample=NPTS // 2 + 0.25), + BandLimited(amp=5.0, peak_sample=NPTS // 2 + 0.25), + ] + k = np.stack([sig.samples() for sig in signals]) + rho = np.full(k.shape, RHO_SQ) + honest = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + + real = tmq.required_upsample_factors + tmq.required_upsample_factors = lambda sigma, dx, xpy=np: 2 * xpy.ones( + np.asarray(sigma).shape, dtype=np.int64) + try: + got = tmq.time_marginalize_bandlimited(k, rho, DELTAT, _lnL) + finally: + tmq.required_upsample_factors = real + rep = tmq.last_report() + assert rep['n_refinements'] > 0, rep + assert len(rep['factor_histogram']) == 2, rep + assert sum(rep['factor_histogram'].values()) == 2, rep + assert np.allclose(got, honest, rtol=0, atol=1e-9), (got, honest, rep) + + # ------------------------------------------------------- the driver CLI @@ -864,6 +1022,11 @@ def test_bandlimited_runs_on_the_gpu_backend_and_matches_numpy(): 'n_flat_rows'): assert rep_np[key] == rep_cp[key], (key, rep_np[key], rep_cp[key]) assert rep_np['n_refined_rows'] >= 1 + for report in (rep_np, rep_cp): + assert report['bandlimited_fft_strategy'] == 'retained-grid-zoomfft', report + assert report['retained_fft_rows'] >= report['n_refined_rows'] + assert report['full_fft_fallback_rows'] == 0, report + assert report['full_fft_fallback_reasons'] == {}, report # The REFINED rows integrate with trapezoid on the dense grid, which has no # even/odd Simpson ambiguity, so the two backends must agree to round-off. diff --git a/MonteCarloMarginalizeCode/Code/test/test_vectorized_lal_tools_split.py b/MonteCarloMarginalizeCode/Code/test/test_vectorized_lal_tools_split.py new file mode 100644 index 000000000..541e3a418 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_vectorized_lal_tools_split.py @@ -0,0 +1,129 @@ +"""The source-only / per-detector split of the vectorized LAL tools is bitwise exact. + +`DiscreteFactoredLogLikelihoodViaArrayVectorNoLoop` used to rebuild the detector +response basis and the source propagation direction once per detector, although +neither depends on the detector. Those are now built once and handed to a +per-detector half. The split is only worth having if it changes nothing, so pin +that with exact equality rather than a tolerance. + +WHAT THE REFERENCE IS, AND WHY IT IS NOT THE WRAPPER. An earlier version of this +file compared `ComputeDetAMResponse(...)` against +`ComputeDetAMResponsePrecomputed(SourcePolarizationBasis(...))`. After the split the +wrapper IS that composition, so the comparison was tautological -- it passed 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%. The references below +are instead the PRE-SPLIT bodies, frozen here verbatim, so the test compares the +refactor against what it replaced rather than against itself. +""" +import numpy as np +import pytest + +import lalsimulation as lalsim + +from RIFT.likelihood.vectorized_lal_tools import ( + ComputeDetAMResponse, + ComputeDetAMResponsePrecomputed, + SourcePolarizationBasis, + SourcePropagationDirection, + TimeDelayFromEarthCenter, + TimeDelayFromEarthCenterPrecomputed, +) + +# Three real interferometer geometries, so the test would catch an axis or +# transpose error that a symmetric toy matrix would hide. +DETECTORS = ["H1", "L1", "V1"] + + +def _frozen_detector_response(R, ra, dec, psi, gmst): + """The body of ComputeDetAMResponse as it stood BEFORE the split (verbatim).""" + X = np.empty(ra.shape + (3,), dtype=np.float64) + Y = np.empty(ra.shape + (3,), dtype=np.float64) + gha = gmst - ra + cos_gha, sin_gha = np.cos(gha), np.sin(gha) + cos_dec, sin_dec = np.cos(dec), np.sin(dec) + cos_psi, sin_psi = np.cos(psi), np.sin(psi) + X[..., 0] = -cos_psi*sin_gha - sin_psi*cos_gha*sin_dec + X[..., 1] = -cos_psi*cos_gha + sin_psi*sin_gha*sin_dec + X[..., 2] = sin_psi*cos_dec + Y[..., 0] = sin_psi*sin_gha - cos_psi*cos_gha*sin_dec + Y[..., 1] = sin_psi*cos_gha + cos_psi*sin_gha*sin_dec + Y[..., 2] = cos_psi*cos_dec + F_plus = (X*np.inner(X, R) - Y*np.inner(Y, R)).sum(axis=-1) + F_cross = (X*np.inner(Y, R) + Y*np.inner(X, R)).sum(axis=-1) + return F_plus + 1.0j*F_cross + + +def _frozen_time_delay(loc, ra, dec, gmst): + """The body of TimeDelayFromEarthCenter as it stood BEFORE the split (verbatim).""" + negative_speed_of_light = np.asarray(-299792458.0) + cos_dec = np.cos(dec) + gha = gmst - ra + ehat = np.empty(ra.shape + (3,), dtype=np.float64) + ehat[..., 0] = cos_dec * np.cos(gha) + ehat[..., 1] = -cos_dec * np.sin(gha) + ehat[..., 2] = np.sin(dec) + neg_separation = np.inner(loc, ehat) + return np.divide(neg_separation, negative_speed_of_light, out=neg_separation) + + +def _samples(n=257, seed=20260905): + rng = np.random.RandomState(seed) + return ( + rng.uniform(0.0, 2.0 * np.pi, n), # right ascension + np.arcsin(rng.uniform(-1.0, 1.0, n)), # declination + rng.uniform(0.0, np.pi, n), # polarization + ) + + +@pytest.mark.parametrize("det", DETECTORS) +def test_detector_response_split_matches_frozen_pre_split_body(det): + ra, dec, psi = _samples() + gmst = 4.371829 + R = np.asarray(lalsim.DetectorPrefixToLALDetector(det).response) + + want = _frozen_detector_response(R, ra, dec, psi, gmst) + X, Y = SourcePolarizationBasis(ra, dec, psi, gmst, xpy=np) + got_split = ComputeDetAMResponsePrecomputed(R, X, Y, xpy=np) + got_wrapper = ComputeDetAMResponse(R, ra, dec, psi, gmst, xpy=np) + + assert np.array_equal(got_split, want), det + assert np.array_equal(got_wrapper, want), det + + +@pytest.mark.parametrize("det", DETECTORS) +def test_time_delay_split_matches_frozen_pre_split_body(det): + ra, dec, _ = _samples() + gmst = 4.371829 + loc = np.asarray(lalsim.DetectorPrefixToLALDetector(det).location) + + want = _frozen_time_delay(loc, ra, dec, gmst) + ehat = SourcePropagationDirection(ra, dec, gmst, xpy=np) + got_split = TimeDelayFromEarthCenterPrecomputed(loc, ehat, xpy=np) + got_wrapper = TimeDelayFromEarthCenter(loc, ra, dec, gmst, xpy=np) + + assert np.array_equal(got_split, want), det + assert np.array_equal(got_wrapper, want), det + + +def test_shared_inputs_are_not_mutated_by_the_per_detector_halves(): + """The whole point of the split is that one prologue serves every detector. + + If a per-detector half wrote into ehat_src, X or Y -- the time-delay half divides + in place, into the result of `inner`, which is a fresh array, but that is a + one-character edit away from being wrong -- the second detector would silently be + computed from corrupted geometry. + """ + ra, dec, psi = _samples(n=64) + gmst = 1.25 + ehat = SourcePropagationDirection(ra, dec, gmst, xpy=np) + X, Y = SourcePolarizationBasis(ra, dec, psi, gmst, xpy=np) + ehat0, X0, Y0 = ehat.copy(), X.copy(), Y.copy() + + for det in DETECTORS: + d = lalsim.DetectorPrefixToLALDetector(det) + TimeDelayFromEarthCenterPrecomputed(np.asarray(d.location), ehat, xpy=np) + ComputeDetAMResponsePrecomputed(np.asarray(d.response), X, Y, xpy=np) + + assert np.array_equal(ehat, ehat0) + assert np.array_equal(X, X0) + assert np.array_equal(Y, Y0)