From 000cf79b80318b1f54d70826435618256c7df84b Mon Sep 17 00:00:00 2001 From: Richard O'Shaughnessy Date: Tue, 11 Aug 2026 08:01:45 -0700 Subject: [PATCH 1/2] AV: fix the high-SNR empty-live-volume crash, and report the collapse Backport of rift_O4d PR #63 (d3dd2a2f, 6a828366, f6bfc5d5), reduced to the parts that are safe to carry onto a frozen production line. At high network SNR the extrinsic export died with ===> FAILED ANALYSIS <==== zero-size array to reduction operation maximum which has no identity Probable reasons: SEOB nyquist or starting frequency limit or signal duration The attribution was wrong: nothing about the waveform is involved. Measured on this branch, 12 replicates per arm, zero-noise injections at a fixed intrinsic point, production extrinsic config, rho_net 146.8 -- 11/12 exports crashed. A crashed export writes nothing, so the posterior is silently thinned in an SNR-dependent way: rho 51 -> 0%, rho 72 -> 3%, rho 103 -> 42%, rho 147 -> 92%. ROOT CAUSE. Two defects, stacked. 1. The production likelihood UNDERFLOWS: exp() of a lnL more than ~745 nats below the peak is 0 in float64, so it returns -inf. Over a cold extrinsic prior at rho 147 that is ~99996 of 100000 draws, so AV's live set is born holding a handful of samples -- sometimes one. 2. get_likelihood_threshold could then return a threshold >= max(lkl). With a small live set and one dominant weight BOTH of its terms degenerate: prob_stop_thr saturates at the MAXIMUM (every other weight underflows to exactly 0, so the discard_prob quantile IS the top sample) while lkl_stop_thr falls through the `len > nsel` test to the array MINIMUM. integrate_log applies that as a STRICT `allloglkl > loglkl_thr`, discarding at least one sample per cycle regardless of merit, so the live set ratchets to 1 and then to 0 -- and the empty array reaches `xpy_here.max(allloglkl)`. NOT CUPY-SPECIFIC. numpy raises the identical ValueError from the same line; the upstream traceback names CUPY_CUB_MAX only because production runs on the GPU. The reproduction on this branch raises the numpy flavour, and the new test reproduces it deterministically on the host path. CHANGES RIFT/integrators/mcsamplerAdaptiveVolume.py * get_likelihood_threshold: raise a named LiveVolumeCollapse on empty input, and CLAMP the threshold strictly below max(lkl) so it can never discard the entire live volume. A threshold at/above the max encloses zero probability, contradicting the enc_prob=0.999 the function exists to maintain. With the clamp in place at least the peak always survives, so nrec >= 1 is guaranteed and the empty reduction downstream is unreachable. * live_volume_collapse_verdict() + ess_from_log_weights(): report a degenerate contraction instead of silently exporting one sample. dict_return (which was empty on this branch) gains live_volume_collapsed / collapse_reason / n_live_final / n_ESS, and the run prints an [AV COLLAPSE] block. bin/integrate_likelihood_extrinsic_batchmode * the "Probable reasons: SEOB nyquist ..." hint was printed for EVERY exception in the block. Make it conditional and name the integrator collapse when that is the cause. This misattribution is why the failure went undiagnosed. WHY THE VERDICT IS WORTH CARRYING ONTO A FROZEN LINE. The clamp stops the crash, after which the run that used to crash completes and exports -- sometimes from a single surviving sample. Unreported, that converts a silent THINNING (a crashed export is visibly missing) into a silent CONTAMINATION (a degenerate export looks ordinary to downstream posterior assembly), which is worse. The verdict is purely additive: it reads log_wt after the integration loop and no sampling decision consumes it. WHY THE VERDICT DOES NOT GATE ON k-hat. k-hat exceeds its nominal 0.70 "unresolved tail" threshold even in HEALTHY runs on this problem -- the converged rho=51.4 replicates measure 0.819-1.605 -- so gating on it would false-flag good exports. ESS separates the regimes with nothing in between (16.3-34.3 healthy vs 1.0-2.0 collapsed), so the verdict keys on ESS and live-set size. DELIBERATELY NOT PORTED from upstream, to keep the change small on a production branch: the RIFT_AV_TRACE per-cycle instrumentation, the integrate_log restructuring (finite screening, empty-cycle recovery, previous-state restore), the machine-readable _integrator_status.json sidecar, the ILE --reject-collapsed-live-volume gate, and the warm-seed affine-rank rule. None is needed to stop the crash measured here; all of them change more of the sampling loop or the output contract than this branch should take. TESTS. test/test_av_empty_live_volume.py, a focused backport of the upstream 910-line suite: 24 passed, 2 GPU-only skips. Mutation-tested -- with the clamp disabled, 5 tests fail, including the two that drive integrate_log to the exact production error ("zero-size array to reduction operation maximum which has no identity"). test_clamp_is_inert_when_the_live_set_is_well_populated reproduces the PRE-FIX threshold formula verbatim and asserts the clamped code returns the identical value on a healthy 20000-sample live set, so the clamp cannot silently shift a production lnZ. Co-Authored-By: Claude Opus 5 --- .../integrators/mcsamplerAdaptiveVolume.py | 118 +++++- .../integrate_likelihood_extrinsic_batchmode | 27 +- .../Code/test/test_av_empty_live_volume.py | 372 ++++++++++++++++++ 3 files changed, 515 insertions(+), 2 deletions(-) create mode 100644 MonteCarloMarginalizeCode/Code/test/test_av_empty_live_volume.py diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py index b19687ceb..266b97a7d 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py @@ -94,6 +94,57 @@ def __init__(self, value): def __str__(self): return repr(self.value) +class LiveVolumeCollapse(Exception): + """The adaptive-volume live set is empty (or carries no usable information). + + Raised INSTEAD of the bare numpy/cupy "zero-size array to reduction operation + ... which has no identity" that the empty-live-volume path used to produce, so + callers and logs can tell a degenerate contraction apart from a waveform + generation failure. See the collapse discussion in get_likelihood_threshold. + """ + pass + + +def ess_from_log_weights(log_wt): + """Kish effective sample size (sum w)^2 / sum w^2, from LOG weights.""" + lw = np.asarray(log_wt, dtype=float) + lw = lw[np.isfinite(lw)] + if len(lw) == 0: + return 0.0 + lse = special.logsumexp + return float(np.exp(2 * lse(lw) - lse(2 * lw))) + + +def live_volume_collapse_verdict(n_live, ndim, ess=None): + """Has the adaptive-volume live set degenerated? -> (collapsed, [reasons]) + + A degenerate contraction must be REPORTED rather than silently exported. With the + threshold clamp in place the run no longer crashes, so it now returns a lnZ and a + sample cloud -- but at high SNR both can describe a SINGLE mode, and the cloud is + then not a fair posterior draw. Unreported, that turns a crash (which at least + kept the point out of the posterior) into a silent contamination, which is worse. + + Thresholds, against the separation measured on zero-noise injections at a fixed + intrinsic point (rho_net 51 -> 147): + healthy ESS 16.3-34.3 (rho 51.4, converges cold) + collapsed ESS 1.0-2.0 (rho 103-147) + * n_live <= ndim -- a live set no larger than the dimension cannot span the + space, let alone describe a posterior in it. Geometric, not tuned. + * ESS < 2 -- fewer than two effective samples IS one sample. + The gap between the regimes is an order of magnitude wide, so these sit far from + both sides of it. Pareto k-hat is deliberately NOT used: it exceeds its nominal + 0.70 "unresolved tail" threshold even in the healthy runs on this problem (the 12 + converged rho=51.4 replicates measure 0.819-1.605), so gating on it would have + false-flagged 12 of 12 good exports. + """ + reasons = [] + if n_live <= ndim: + reasons.append("final live volume holds {} sample(s) in {} dimensions".format(n_live, ndim)) + if ess is not None and ess < 2.0: + reasons.append("ESS={:.2f}".format(ess)) + return bool(reasons), reasons + + ### V. Tiwari routines def get_likelihood_threshold(lkl, lkl_thr, nsel, discard_prob,xpy_here=xpy_default): @@ -104,7 +155,18 @@ def get_likelihood_threshold(lkl, lkl_thr, nsel, discard_prob,xpy_here=xpy_defau nsel : integer, has to do with size of array of likelihoods used to evaluate for next array. discard_prob: threshold on CDF to throw away an entire bin. Should be very small """ - + if len(lkl) == 0: + # Caller must not ask for a threshold on an empty live volume: every reduction + # below (max, argsort, [0]) is undefined, and the bare backend error for that + # ("zero-size array to reduction operation CUPY_CUB_MAX which has no identity") + # is what the ILE used to mis-report as a waveform problem. Name the cause. + raise LiveVolumeCollapse( + "adaptive-volume live set is empty: the likelihood returned no finite value " + "inside the sampled volume. At high network SNR this is likelihood UNDERFLOW " + "(exp() of a lnL more than ~745 nats below the peak returns 0), so a cold " + "extrinsic prior yields almost no usable draw. This is NOT a waveform Nyquist/" + "start-frequency/duration problem. Narrow the extrinsic prior or seed the sampler.") + w = xpy_here.exp(lkl - np.max(lkl)) npoints = len(w) sumw = xpy_here.sum(w) @@ -121,6 +183,37 @@ def get_likelihood_threshold(lkl, lkl_thr, nsel, discard_prob,xpy_here=xpy_defau lkl_stop_thr = lkl_stop_thr[-1] lkl_thr = min(lkl_stop_thr, prob_stop_thr) + # CLAMP (backport of rift_O4d PR #63). The threshold is applied downstream as a + # STRICT `lkl > thr`, so a threshold at or above max(lkl) discards the ENTIRE live + # volume; the next cycle then reduces over an empty array and raises + # "zero-size array to reduction operation CUPY_CUB_MAX which has no identity", + # which the ILE driver misreports as a waveform problem. + # + # It happens whenever the live set is small AND one weight dominates -- i.e. at high + # network SNR, where a double-precision exp() of a lnL more than ~745 nats below the + # peak underflows to zero, so almost every cold draw is -inf and the live set starts + # with a handful of members. Then prob_stop_thr saturates at max(lkl) (every other + # weight is 0, so the discard_prob quantile IS the top sample) while lkl_stop_thr + # falls back to the smallest. Each cycle discards at least one point regardless of + # merit and the live set ratchets to zero. + # + # Back the threshold off to the largest value strictly below the maximum, so the peak + # always survives. In a healthy run len(lkl) >> nsel and lkl_stop_thr is the nsel-th + # largest, far below the max, so this clamp never engages and results are unchanged. + # + # Reduce on the ACTIVE backend and move only the scalar: identity_convert(lkl) here + # would copy the whole live set device->host every cycle, on the healthy path too. + lkl_max = float(identity_convert(xpy_here.max(lkl))) + if not (float(identity_convert(lkl_thr)) < lkl_max): + lkl_host = identity_convert(lkl) # rare branch: the live set is tiny by construction + below = lkl_host[lkl_host < lkl_max] + if len(below): + lkl_thr = np.max(below) # keep only the maximum: maximal, but safe, contraction + else: + # every live point is at the maximum: no threshold can separate them, so + # take one below all of them and leave the live volume intact. + lkl_thr = np.nextafter(lkl_max, -np.inf) + truncp = xpy_here.sum(w[lkl < lkl_thr]) / sumw return identity_convert(lkl_thr), identity_convert(truncp) # send both to CPU as needed @@ -686,6 +779,29 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): self._rvs[name] = identity_convert(self._rvs[name]) # this is trivial if xpy_default is numpy, and a conversion otherwise dict_return = {} + + # ------------------------------------------------------------------ + # LIVE-VOLUME COLLAPSE VERDICT. The threshold clamp in get_likelihood_threshold + # stops the empty-reduction CRASH, but the run that used to crash now completes and + # exports -- sometimes from a single surviving sample. Silently exporting that is + # worse than the crash was: a crashed export writes nothing and is visibly missing, + # whereas a degenerate one enters downstream posterior assembly looking ordinary. + # So say so. Thresholds and the measured regimes are documented on + # live_volume_collapse_verdict. Purely additive: no sampling decision reads this. + n_live_final = int(len(log_wt)) + _ess = ess_from_log_weights(log_wt) + collapsed, _reasons = live_volume_collapse_verdict(n_live_final, ndim, ess=_ess) + dict_return['n_ESS'] = _ess + dict_return['n_live_final'] = n_live_final + dict_return['live_volume_collapsed'] = collapsed + if collapsed: + dict_return['collapse_reason'] = "; ".join(_reasons) + print(" [AV COLLAPSE] the live volume degenerated: " + dict_return['collapse_reason'] + ".") + print(" [AV COLLAPSE] lnZ and the exported samples describe a SINGLE mode of the integrand and are") + print(" [AV COLLAPSE] NOT a fair draw from the posterior. Do not use this export unweighted.") + print(" [AV COLLAPSE] At high network SNR this is likelihood underflow over a cold extrinsic prior;") + print(" [AV COLLAPSE] narrow the extrinsic prior (--limit-* zoom box) or seed the sampler.") + return log_int, np.log(rel_var) +2*log_int, eff_samp, dict_return # if outvals: diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index 2c4b83e14..447099fb7 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -9,6 +9,7 @@ Integrate the extrinsic parameters of the prefactored likelihood function. import sys import functools +import traceback from optparse import OptionParser, OptionGroup import numpy @@ -2265,7 +2266,31 @@ for indx in numpy.arange(len(P_list)): # if failure_mode in str(exception_failure): # sys.exit(opts.custom_fail_codes[i]) - print( " Probable reasons: SEOB nyquist or starting frequency limit or signal duration ") + # Attribute the failure to something the traceback actually supports. This line used to + # read "Probable reasons: SEOB nyquist or starting frequency limit or signal duration" + # UNCONDITIONALLY, for every exception raised anywhere in the block above. That pointed + # the high-SNR adaptive-volume live-volume collapse -- the dominant extrinsic-export + # failure at rho_net >~ 100, where >90% of exports died -- at the waveform generation + # code instead, which is why it went undiagnosed for so long. + if ((mcsampler_AV_ok and isinstance(exception_failure, mcsamplerAdaptiveVolume.LiveVolumeCollapse)) + or 'live volume' in str_err or 'live set' in str_err + or ('zero-size array' in str_err and 'no identity' in str_err + and 'mcsamplerAdaptiveVolume' in traceback.format_exc())): + # covers both the named exception and the bare numpy/cupy empty-reduction error that + # older trees (and any other integrator with the same defect) still raise. The string + # fallback additionally requires the traceback to name the integrator, because this + # handler also covers waveform generation, data conditioning and the likelihood stack. + print( " Probable reason: the INTEGRATOR's live volume collapsed -- no samples survived the") + print( " adaptive-volume likelihood threshold. This is NOT a waveform problem: nyquist, start") + print( " frequency and segment duration are all irrelevant to it. At high network SNR the") + print( " likelihood underflows (exp() of a lnL more than ~745 nats below the peak returns 0),") + print( " so a cold extrinsic prior yields almost no finite draw. Narrow the extrinsic prior") + print( " (--limit-right-ascension/--limit-declination/--limit-inclination, tighter --d-min/--d-max).") + elif ('nyquist' in str_err.lower() or 'srate' in str_err.lower() or 'duration' in str_err.lower() + or 'ChooseFDWaveform' in str_err or 'ChooseTDWaveform' in str_err or 'gwsignal' in str_err): + print( " Probable reasons: SEOB nyquist or starting frequency limit or signal duration ") + else: + print( " Cause not classified -- read the traceback above; it names the failing call.") print( " Skipping the following binary! ") # Zero out extrinsic parameters -- these are CUDA-populated / meaningless, but could cause errors if populated P_list[indx].incl = P_list[indx].tref = P_list[indx].dist = P_list[indx].phiref = P_list[indx].psi =P_list[indx].theta = P_list[indx].phi =0 diff --git a/MonteCarloMarginalizeCode/Code/test/test_av_empty_live_volume.py b/MonteCarloMarginalizeCode/Code/test/test_av_empty_live_volume.py new file mode 100644 index 000000000..77138736f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_av_empty_live_volume.py @@ -0,0 +1,372 @@ +#!/usr/bin/env python +""" +Regression tests for the adaptive-volume EMPTY LIVE VOLUME crash +(RIFT/integrators/mcsamplerAdaptiveVolume.py). + +Backport of the focused subset of the rift_O4d suite (PR #63) that covers the parts +of that fix carried onto this branch: the threshold clamp, the named error on an +empty live volume, the collapse verdict, and the ILE's cause attribution. + +Background (the bug these tests lock down). At high network SNR the production +extrinsic likelihood underflows: exp() of a lnL more than ~745 nats below the peak +returns 0, so the likelihood hands back -inf. Over a cold extrinsic prior at +rho_net ~ 147 that is ~99.996% of draws (measured: 99996 of 100000), and the +adaptive-volume live set is therefore born holding a handful of samples, sometimes +one. Two things then went wrong, in sequence: + + 1. get_likelihood_threshold returned a threshold >= max(lkl). With a small live + set the `len(lkl) > nsel` branch falls through to lkl_stop_thr = lkl[-1] (the + array MINIMUM), while prob_stop_thr saturates at the MAXIMUM because every + other weight underflows to zero, so the discard_prob quantile IS the top + sample. min(min, max) is then the live set's own minimum -- and with only one + sample left, its maximum. + 2. integrate_log applies that threshold as a STRICT `allloglkl > loglkl_thr`, so + it discarded at least one sample per cycle regardless of merit, ratcheting the + live set down to 1 and then to 0. + +The empty array then reached `lw = allloglkl - xpy_here.max(allloglkl)` and raised + + ValueError: zero-size array to reduction operation CUPY_CUB_MAX which has no identity + +which bin/integrate_likelihood_extrinsic_batchmode reported as "Probable reasons: +SEOB nyquist or starting frequency limit or signal duration". Measured crash rate +by network SNR on zero-noise injections at a fixed intrinsic point: +rho 51.4 -> 0/12, rho 72.1 -> 3%, rho 102.8 -> 5/12, rho 146.8 -> 11/12. + +NOT CUPY-SPECIFIC. The traceback names CUPY_CUB_MAX only because production runs +on the GPU; numpy raises the identical ValueError ("zero-size array to reduction +operation maximum which has no identity") from the same line. These tests run on +whichever backend the sampler picked, and the cupy path is exercised as well when a +GPU is present. + +The requirement is not merely "does not crash": a degenerate contraction must be +REPORTED (dict_return['live_volume_collapsed']) rather than silently exporting the +one surviving sample as if it were a posterior. +""" + +import os + +import numpy as np +import pytest + +import RIFT.integrators.mcsamplerAdaptiveVolume as mcsamplerAV +from RIFT.integrators.mcsamplerAdaptiveVolume import ( + LiveVolumeCollapse, + ess_from_log_weights, + get_likelihood_threshold, + live_volume_collapse_verdict, +) + +NAMES = ['right_ascension', 'declination', 'phi_orb', 'inclination', 'psi', 'distance'] +NDIM = len(NAMES) + +xpy = mcsamplerAV.xpy_default +to_backend = mcsamplerAV.identity_convert_togpu +to_host = mcsamplerAV.identity_convert + + +def _sampler(n_chunk=10000): + s = mcsamplerAV.MCSampler(n_chunk=n_chunk) + # Bind the sampler to the ACTIVE backend exactly as bin/integrate_likelihood_extrinsic_ + # batchmode does (`sampler.xpy = xpy_default; sampler.identity_convert = ...`). + # MCSampler.__init__ defaults self.xpy to numpy, so on a GPU host an unconfigured + # sampler mixes cupy arrays with numpy calls and dies in the fairdraw block for + # reasons that have nothing to do with what is under test here. + s.xpy = xpy + s.identity_convert = to_host + for name in NAMES: + s.add_parameter(name, pdf=None, left_limit=0.0, right_limit=1.0, + prior_pdf=lambda x: np.ones(np.shape(x)), + adaptive_sampling=True) + return s + + +def _integrate(fn, nmax=200000, neff=8, n_chunk=10000): + s = _sampler(n_chunk) + res = s.integrate_log(fn, *NAMES, nmax=nmax, neff=neff, n=n_chunk, + no_protect_names=True, verbose=False, + igrand_fairdraw_samples=True, + igrand_fairdraw_samples_max=200) + return s, res + + +### +### 1. get_likelihood_threshold must never return a threshold that empties the set +### +# This is the root defect. The threshold is consumed as a STRICT `lkl > thr`, so a +# threshold at or above max(lkl) encloses zero probability -- which contradicts the +# enc_prob = 0.999 the function exists to maintain. + +@pytest.mark.parametrize('n', [1, 2, 3, 10, 999]) +def test_threshold_never_discards_the_entire_live_volume(n): + """The regression: small live set + one dominant weight -> thr was max(lkl).""" + # lnL values spread far enough apart that exp(lkl-max) underflows for all but the top, + # which is exactly the high-SNR condition that saturates prob_stop_thr at the maximum. + lkl_host = 10000.0 + 1000.0 * np.arange(n) + lkl = to_backend(lkl_host) + thr, truncp = get_likelihood_threshold(lkl, -1e15, 1000, 1e-3, xpy_here=xpy) + assert float(thr) < float(lkl_host.max()), \ + "threshold {} >= max {}: strict `>` would empty the live volume".format(thr, lkl_host.max()) + assert int(np.sum(to_host(lkl) > thr)) >= 1 + + +def test_threshold_survives_an_all_equal_live_volume(): + """No contraction is possible when every sample has the same lnL; keep them all.""" + lkl = to_backend(np.full(5, 123.5)) + thr, _ = get_likelihood_threshold(lkl, -1e15, 1000, 1e-3, xpy_here=xpy) + assert float(thr) < 123.5 + assert int(np.sum(to_host(lkl) > thr)) == 5 + + +def test_threshold_on_a_single_sample_keeps_it(): + lkl = to_backend(np.array([42.0])) + thr, _ = get_likelihood_threshold(lkl, -1e15, 1000, 1e-3, xpy_here=xpy) + assert float(thr) < 42.0 + + +def test_threshold_on_an_empty_live_volume_raises_a_named_error(): + """Not a bare 'zero-size array to reduction operation ...' from inside a reduction.""" + with pytest.raises(LiveVolumeCollapse) as excinfo: + get_likelihood_threshold(to_backend(np.array([])), -1e15, 1000, 1e-3, xpy_here=xpy) + msg = str(excinfo.value).lower() + assert 'underflow' in msg or 'no finite value' in msg + # the misattribution this whole investigation chased down must not come back: the + # message may MENTION nyquist, but only to rule it out. + assert 'not a waveform nyquist' in msg + + +@pytest.mark.parametrize('scale', [1.0, 5.0]) +def test_clamp_is_inert_when_the_live_set_is_well_populated(scale): + """The clamp must not move the threshold in the regime production actually runs in. + + Reference value is the PRE-FIX formula, reproduced here verbatim, so this test fails + if the clamp ever starts engaging on a healthy live set (which would silently shift + every production lnZ). + """ + rng = np.random.RandomState(20260810) + lkl_host = rng.normal(0.0, scale, size=20000) + nsel, discard_prob = 1000, 1e-3 + + # --- the original (unclamped) threshold, verbatim from the pre-fix implementation + w = np.exp(lkl_host - np.max(lkl_host)) + prob = w / np.sum(w) + idx = np.argsort(prob) + ecdf = np.cumsum(prob[idx]) + prob_stop_thr = lkl_host[idx][ecdf >= discard_prob][0] + srt = np.sort(lkl_host)[::-1] + lkl_stop_thr = srt[nsel] if len(srt) > nsel else srt[-1] + expected = min(lkl_stop_thr, prob_stop_thr) + # --- + + thr, _ = get_likelihood_threshold(to_backend(lkl_host), -1e15, nsel, discard_prob, xpy_here=xpy) + assert float(thr) == pytest.approx(float(expected)), 'clamp engaged on a healthy live set' + assert int(np.sum(lkl_host > float(thr))) > 1 + + +### +### 2. integrate_log must not crash on the two routes measured in production +### + +def _lone_survivor(*args): + """Exactly one finite draw per chunk: live set of size 1 on cycle 1. + + Signature in production: the crash arrives BEFORE any per-cycle line is printed + (rho_net 146.8). + """ + x = np.array(args).T + out = np.full(len(x), -np.inf) + out[0] = 100.0 + return out + + +def _ratchet_to_one(*args): + """A few distinct finite values, never more, so `>` sheds one sample per cycle. + + Signature in production: int_var 0.7071 (2 samples) -> 0.0 (1 sample) -> crash + (rho_net 72.1). + """ + x = np.array(args).T + out = np.full(len(x), -np.inf) + k = min(3, len(x)) + out[:k] = 100.0 + np.arange(k) + return out + + +@pytest.mark.parametrize('fn,label', [(_lone_survivor, 'lone survivor'), + (_ratchet_to_one, 'ratchet to one')]) +def test_degenerate_live_volume_does_not_raise_an_empty_reduction(fn, label): + try: + s, res = _integrate(fn) + except ValueError as e: # the exact regression + if 'zero-size array' in str(e) or 'no identity' in str(e): + pytest.fail("empty-live-volume crash returned ({}): {}".format(label, e)) + raise + assert np.isfinite(float(res[0])), "lnZ must be a real number, got {}".format(res[0]) + assert len(s._rvs['log_integrand']) >= 1 + + +def test_no_finite_sample_anywhere_raises_a_named_error_not_a_reduction_error(): + """Nothing finite ever: there IS no integral, so fail -- but say why.""" + def all_underflowed(*args): + return np.full(len(np.array(args).T), -np.inf) + + with pytest.raises(LiveVolumeCollapse) as excinfo: + _integrate(all_underflowed, nmax=30000) + msg = str(excinfo.value).lower() + assert 'underflow' in msg or 'no finite value' in msg + assert 'not a waveform nyquist' in msg + + +### +### 3. A degenerate contraction must be REPORTED, not silently exported +### +# The failure mode that survives the crash fix is worse than the crash: an export +# built from one sample, indistinguishable in the output from a converged one. A +# crashed export at least writes nothing and is visibly missing. + +def _peaked(rho, underflow=True): + """6-D Gaussian at lnL scale rho^2/2, with the float64 underflow of the real code.""" + x0 = 0.5 * np.ones(NDIM) + width = 0.5 / rho + lnLmax = 0.5 * rho ** 2 + + def lnL(*args): + x = np.array(args).T + out = lnLmax - 0.5 * np.sum(((x - x0) / width) ** 2, axis=-1) + if underflow: + out = np.where(out > lnLmax - 745.0, out, -np.inf) + return out + return lnL + + +def test_collapse_is_flagged_in_dict_return_at_high_snr(): + """rho ~ 147 with a production-sized chunk: a few finite draws, then a degenerate + contraction. The run COMPLETES -- and must say that its answer is degenerate.""" + np.random.seed(20260810) + s, res = _integrate(_peaked(146.8), nmax=300000, neff=8, n_chunk=100000) + dd = res[3] + assert dd.get('live_volume_collapsed') is True, dd + assert dd.get('collapse_reason') + assert dd.get('n_live_final') is not None + + +@pytest.mark.parametrize('n_live,ess,expect', [ + # (live samples, ESS) -> should the verdict call this collapsed? + (1, 1.0, True), # rho 146.8: one sample exported + (2, 1.7, True), # rho 146.8: slipped an earlier ESS<1.5-only rule + (6, 3.0, True), # live set no larger than the dimension + (4000, 16.3, False), # rho 51.4, measured healthy (bottom of the healthy band) + (4000, 34.3, False), # rho 51.4, measured healthy (top of the healthy band) + (4000, 4.0, False), # hard but not degenerate: low ESS, still many live points +]) +def test_collapse_verdict_matches_the_measured_regimes(n_live, ess, expect): + """Pin the decision boundary against the regimes measured on the real problem.""" + collapsed, reasons = live_volume_collapse_verdict(n_live, NDIM, ess=ess) + assert collapsed is expect, reasons + assert bool(reasons) is expect + + +def test_a_healthy_run_is_not_flagged_as_collapsed(): + """The guard must not cry wolf on the regime that already works (rho ~ 51).""" + np.random.seed(20260810) + s, res = _integrate(_peaked(20.0), nmax=400000, neff=20, n_chunk=20000) + dd = res[3] + assert dd.get('live_volume_collapsed') is False, dd.get('collapse_reason') + assert dd['n_live_final'] > 2 + + +def test_ess_helper_matches_the_kish_definition(): + lw = np.log(np.array([1.0, 1.0, 1.0, 1.0])) + assert ess_from_log_weights(lw) == pytest.approx(4.0) + lw = np.log(np.array([1.0, 1e-12, 1e-12])) + assert ess_from_log_weights(lw) == pytest.approx(1.0, abs=1e-6) + # -inf entries (the high-SNR underflow) must be screened, not poison the reduction + assert ess_from_log_weights(np.array([0.0, -np.inf, -np.inf])) == pytest.approx(1.0) + + +### +### 4. The fix must be inert on well-conditioned integrals +### +# The clamp changes the threshold only when the threshold would have emptied the live +# volume; anything else would silently move production lnZ values. + +def test_known_gaussian_integral_is_recovered(): + """int over [0,1]^6 of exp(-|x-x0|^2/2w^2) = (w sqrt(2pi))^6 for w << 1.""" + w = 0.08 + x0 = 0.5 * np.ones(NDIM) + + def lnL(*args): + x = np.array(args).T + return -0.5 * np.sum(((x - x0) / w) ** 2, axis=-1) + + np.random.seed(20260810) + s, res = _integrate(lnL, nmax=600000, neff=30, n_chunk=20000) + expected = NDIM * np.log(w * np.sqrt(2 * np.pi)) + # Tolerance is loose on purpose: AV carries a known few-tenths-of-a-nat bias on this + # problem which is a property of the estimator, not of the collapse fix. The purpose + # of this test is to catch a gross regression in the integral, not to grade AV. + assert float(res[0]) == pytest.approx(expected, abs=0.35) + assert res[3].get('live_volume_collapsed') is False + + +### +### 5. Backend coverage +### +# The reported traceback is the cupy flavour (CUPY_CUB_MAX). The tests above run on +# whatever backend is active; when a GPU is present, pin the cupy path explicitly so a +# CPU-only CI run can never be mistaken for coverage of the reported configuration. + +@pytest.mark.skipif(not mcsamplerAV.cupy_ok, reason='no cupy/GPU on this host') +def test_threshold_clamp_on_the_cupy_backend(): + import cupy + lkl = cupy.asarray(np.array([10000.0, 11000.0, 12000.0])) + thr, _ = get_likelihood_threshold(lkl, -1e15, 1000, 1e-3, xpy_here=cupy) + assert float(thr) < 12000.0 + assert int(cupy.sum(lkl > thr).get()) >= 1 + + +@pytest.mark.skipif(not mcsamplerAV.cupy_ok, reason='no cupy/GPU on this host') +def test_degenerate_live_volume_on_the_cupy_backend(): + try: + s, res = _integrate(_lone_survivor) + except ValueError as e: + if 'CUPY_CUB_MAX' in str(e) or 'zero-size array' in str(e): + pytest.fail("the reported cupy crash is back: {}".format(e)) + raise + assert np.isfinite(float(res[0])) + + +### +### 6. Wiring: the ILE script must not re-attribute an integrator collapse to the waveform +### + +_ILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', 'bin', 'integrate_likelihood_extrinsic_batchmode') + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_does_not_blame_the_waveform_unconditionally(): + with open(_ILE) as f: + src = f.read() + # match the PRINT, not the prose: the surrounding comment quotes the same text + i_hint = src.find('print( " Probable reasons: SEOB nyquist') + assert i_hint > 0, 'hint text moved; update this test' + + # The hint must live inside a conditional branch, not fire for every exception. + # Check the region between the handler that catches the failure and the hint itself. + i_handler = src.rfind('except Exception as exception_failure:', 0, i_hint) + assert i_handler > 0, 'handler moved; update this test' + handler_body = src[i_handler:i_hint] + assert 'LiveVolumeCollapse' in handler_body, \ + 'the SEOB-nyquist hint is no longer guarded by a cause check: an integrator ' \ + 'collapse would again be reported as a waveform Nyquist/duration failure' + + # ...and the hint must be nested INSIDE that branch, i.e. indented deeper than the + # statements the handler runs unconditionally (such as the FAILED ANALYSIS banner). + def _indent(needle): + i = src.index(needle) + return i - (src.rfind('\n', 0, i) + 1) + + assert _indent('print( " Probable reasons: SEOB nyquist') > _indent('print( " ===> FAILED ANALYSIS'), \ + 'the SEOB-nyquist hint sits at the handler top level again: it would be printed ' \ + 'for every exception, including an integrator live-volume collapse' From ca9c4249f2727b766c3d16c22cdaeeee8ee5fa31 Mon Sep 17 00:00:00 2001 From: Richard O'Shaughnessy Date: Tue, 11 Aug 2026 08:10:03 -0700 Subject: [PATCH 2/2] ILE: extrinsic zoom box, working with the cosine samplers Backport of rift_O4d PR #58 (af68d7ec zoom-box options, f715f4f0 the cosine-sampler fix, 7e6f2ed9 the GPU-safety + CI review round). This branch had no --limit-* options at all, so the port brings the whole feature in its already-fixed form rather than introducing the defect and then repairing it. WHAT THE OPTIONS DO. --limit-right-ascension / --limit-declination / --limit-inclination / --limit-psi restrict the extrinsic sampling AND prior ranges to a truth-centered box, the way --d-min/--d-max already do for distance, so the adaptive sampler can resolve a narrow high-SNR peak it could never find from the full prior. THE DEFECT THIS PORT AVOIDS. Upstream, the parsed limits were stored in param_limits[...] as radians but only the plain (angle) branches consulted them: the --declination-cosine-sampler and --inclination-cosine-sampler branches hardcoded left_limit=-1, right_limit=1 and never looked at param_limits. No error, no warning, no narrowing -- the run simply sampled the whole sphere / whole iota range. The cosine samplers are the production default in the extrinsic-export config, so a requested sky box silently collapsed to RA-only narrowing. THE TRANSFORM. The limits are always given in radians of the PHYSICAL angle and are mapped into the coordinate actually sampled. The conventions are read off the likelihood closures in the same file (dec = pi/2 - arccos(z), iota = arccos(z)): declination: sampled variable is z = sin(dec). sin() is INCREASING on [-pi/2, pi/2], so the limit order is PRESERVED: [lo,hi] -> [sin(lo), sin(hi)]. inclination: sampled variable is z = cos(iota). cos() is DECREASING on [0, pi], so the limit order SWAPS: [lo,hi] -> [cos(hi), cos(lo)]. The naive [cos(lo), cos(hi)] yields an INVERTED, empty sampling interval; the test suite asserts against exactly that. PRIOR SEMANTICS ARE PRESERVED. The cosine branches keep the FULL-RANGE constant prior density 1/2 in z (previously they reused the sampling pdf object as the prior, which was the same thing only because the range was [-1,1]). It is deliberately not renormalized over the box, so the prior mass removed by a box is 0.5*(z_hi - z_lo), identical to the non-cosine branch where prior_pdf is 0.5*cos(dec) resp. 0.5*sin(iota) over the same physical box. The isotropic prior is therefore unchanged and the two samplers give the same posterior and the same lnZ for the same physical box. ALSO IN THIS CHANGE * The non-cosine dec/iota branches use a properly TRUNCATED sampling pdf and cdf_inv when a limit is given: setting left/right_limit while keeping the full-range dec_samp_vector / cos_samp_vector would let mcsampler / mcsamplerGPU (which, unlike the AV sampler, do use pdf and cdf_inv) draw outside the box. With no limit requested the legacy functions are used unchanged. * The truncated closures infer their array module from the argument they are handed (infer_array_module), because mcsamplerGPU.draw_simplified() calls them with a SINGLE positional cupy array; an explicit xpy= still wins. * Limits are validated: a malformed 'LO,HI', an empty or inverted range, or a range that does not overlap the physical domain now fails loudly instead of producing a degenerate sampler. Dec/iota ranges are clipped to [-pi/2,pi/2] resp. [0,pi]. * Fixed an adjacent silent bug in the scalar time-marginalized likelihood closure, which zipped the RAW sampled `inclination` instead of the converted `incl`; under --inclination-cosine-sampler that fed cos(iota) to the waveform as an angle. This branch carried the same bug. KNOWN LIMITATION, CARRIED ACROSS DELIBERATELY. --limit-right-ascension / --limit-declination together with --internal-sky-network-coordinates is now a HARD ERROR: that option samples RA/dec in a rotated, network-aligned frame, so an equatorial sky box would silently select the wrong patch of sky. Failing loudly beats boxing the wrong patch, but it forecloses 2-IFO boxing, where the ring is narrow in the NETWORK frame and that is exactly where a box would help most. Follow-up work, not addressed here. New helpers live in RIFT/integrators/mcsampler.py (infer_array_module, clip_angle_limits, cosine_sampler_limits, ret_dec_samp_vector, ret_dec_samp_cdf_inv_vector, ret_cos_samp_vector, ret_cos_samp_cdf_inv_vector) so they are importable and testable independently of the backend the run selects. The block is byte-identical to upstream and purely additive. TESTS. test/test_limit_cosine_samplers.py, ported unchanged from upstream: 29 passed. Covers the dec order-preserving and the iota order-SWAPPING transform; clipping; empty/inverted/non-finite ranges raising; the truncated samplers reducing to the legacy ones over the full range; identical restricted support and identical prior mass in both samplers; AV-style (volume x prior) equivalence and identical posterior shape; end-to-end MCSampler.integrate agreement between the two parameterizations for both a flat and a peaked integrand; backend inference and an end-to-end mcsamplerGPU.draw_simplified() draw inside the box; and a source-level guard against reintroducing the hardcoded [-1,1] range. The order-swap is mutation-tested: flipping the `reverses` flag on the inclination convention fails 11 of the 29 tests, including test_inclination_limits_map_to_cos_and_SWAP_order. Both new suites are wired into .travis/test-integrate.sh. Co-Authored-By: Claude Opus 5 --- .travis/test-integrate.sh | 8 + .../Code/RIFT/integrators/mcsampler.py | 156 ++++++ .../integrate_likelihood_extrinsic_batchmode | 131 +++-- .../Code/test/test_limit_cosine_samplers.py | 454 ++++++++++++++++++ 4 files changed, 718 insertions(+), 31 deletions(-) create mode 100644 MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py diff --git a/.travis/test-integrate.sh b/.travis/test-integrate.sh index d777cccaf..a215a26dc 100755 --- a/.travis/test-integrate.sh +++ b/.travis/test-integrate.sh @@ -1,5 +1,13 @@ #! /bin/bash +# Unit/regression tests for the sampler helpers themselves (fast, no data needed). +# The extrinsic "zoom box" limits under the cosine samplers live here: they are pure +# coordinate-transform + prior-mass identities, so they belong with the integrator gate +# rather than with the end-to-end run tests. The adaptive-volume empty-live-volume +# regression is here for the same reason (it is a threshold identity, not a run test). +python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py +python -m pytest -q MonteCarloMarginalizeCode/Code/test/test_av_empty_live_volume.py + python MonteCarloMarginalizeCode/Code/test/test_mcsamplerEnsemble_extended.py --as-test --n-max 100000 python MonteCarloMarginalizeCode/Code/test/test_mcsamplerEnsemble_extended.py --as-test --n-max 100000 --use-lnL diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py index e02fbaa52..2588dc270 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py @@ -876,6 +876,162 @@ def dec_samp_cdf_inv_vector(p): return numpy.arccos(2*p-1) - numpy.pi/2 # target from -pi/2 to pi/2 +### +### Truncated isotropic angle samplers ("zoom box" support) +### +# RIFT can sample sky location / orientation either in the ANGLE itself (dec, +# iota) or -- with --declination-cosine-sampler / --inclination-cosine-sampler -- +# in the cosine variable that makes the isotropic prior flat. The two +# conventions, read off the conversions applied in +# bin/integrate_likelihood_extrinsic_batchmode, are +# +# declination: dec = pi/2 - arccos(z) <=> z = sin(dec), z in [-1,1] +# sin() is INCREASING on [-pi/2, pi/2], so a box [lo,hi] maps to +# [sin(lo), sin(hi)]: the order of the limits is PRESERVED. +# inclination: iota = arccos(z) <=> z = cos(iota), z in [-1,1] +# cos() is DECREASING on [0, pi], so a box [lo,hi] maps to +# [cos(hi), cos(lo)]: the order of the limits SWAPS. +# +# In both cases the physical isotropic prior is p(z) dz = dz/2, i.e. a CONSTANT +# density 1/2 in the cosine coordinate. That constant is deliberately NOT +# renormalized over a restricted box, so that restricting the box reduces the +# prior mass (and hence lnZ) by exactly the same factor as in the angle +# coordinate, where the prior density is 0.5*cos(dec) resp. 0.5*sin(iota). + +_COSINE_SAMPLER_CONVENTIONS = { + # name: (angle_min, angle_max, angle->cosine-coordinate map, reverses_order) + 'declination': (-numpy.pi/2, numpy.pi/2, numpy.sin, False), + 'inclination': (0.0, numpy.pi, numpy.cos, True), +} + + +def infer_array_module(x, xpy=None): + """Return the array module (numpy, cupy, ...) that should be used to operate on `x`. + + The truncated samplers below are handed to whichever backend the run selected: + mcsampler and mcsamplerAdaptiveVolume call them with host (numpy) arrays, while + mcsamplerGPU.draw_simplified() calls `self.pdf[param](samples)` / `self.cdf_inv[param](p)` + with a SINGLE positional argument holding a cupy array. Defaulting to numpy would then + hit `numpy.asarray(cupy_array)`, which raises, so the backend is inferred from the + argument instead of assumed. An explicit `xpy=` always wins. + + The lookup is by the array type's top-level module, taken from `sys.modules` (a cupy + array cannot exist unless cupy is already imported), so this file keeps its numpy-only + import list and works for any duck-typed backend exposing the numpy API. + """ + if xpy is not None: + return xpy + mod_name = type(x).__module__.split('.')[0] + if mod_name in ('numpy', 'builtins'): + return numpy + mod = sys.modules.get(mod_name) + if mod is not None and all(hasattr(mod, _attr) for _attr in ('asarray', 'where', 'clip')): + return mod + return numpy + + +def clip_angle_limits(lo, hi, kind): + """Clip an angular range [lo,hi] (radians) to the physical domain of `kind` + ('declination' -> [-pi/2,pi/2], 'inclination' -> [0,pi]). + + Raises ValueError if the requested range is empty/inverted, or if it does + not overlap the physical domain. Returns (lo, hi) as floats with lo < hi. + """ + if kind not in _COSINE_SAMPLER_CONVENTIONS: + raise ValueError("clip_angle_limits: unknown angle '{}' (expected one of {})".format(kind, sorted(_COSINE_SAMPLER_CONVENTIONS))) + angle_min, angle_max, _, _ = _COSINE_SAMPLER_CONVENTIONS[kind] + lo = float(lo) + hi = float(hi) + if not numpy.isfinite(lo) or not numpy.isfinite(hi): + raise ValueError("clip_angle_limits: non-finite {} range [{}, {}]".format(kind, lo, hi)) + if not (hi > lo): + raise ValueError("clip_angle_limits: empty or inverted {} range [{}, {}] (need LO < HI, in radians)".format(kind, lo, hi)) + lo_c = min(max(lo, angle_min), angle_max) + hi_c = min(max(hi, angle_min), angle_max) + if not (hi_c > lo_c): + raise ValueError("clip_angle_limits: {} range [{}, {}] does not overlap the physical domain [{}, {}]".format(kind, lo, hi, angle_min, angle_max)) + return lo_c, hi_c + + +def cosine_sampler_limits(lo, hi, kind): + """Map an angular range [lo,hi] (radians) into the coordinate actually sampled + by RIFT's 'cosine' sky/orientation samplers. + + kind='declination': sampled variable is z = sin(dec); sin is increasing, so + the limit order is preserved: [lo,hi] -> [sin(lo), sin(hi)]. + kind='inclination': sampled variable is z = cos(iota); cos is DECREASING, so + the limit order SWAPS: [lo,hi] -> [cos(hi), cos(lo)]. + + The range is clipped to the physical angular domain first, and the result is + clipped to the sampler domain [-1,1]. Raises ValueError on an empty or + inverted request. Returns (z_lo, z_hi) with z_lo < z_hi. + """ + lo_c, hi_c = clip_angle_limits(lo, hi, kind) + _, _, fn, reverses = _COSINE_SAMPLER_CONVENTIONS[kind] + z_a = float(fn(lo_c)) + z_b = float(fn(hi_c)) + z_lo, z_hi = (z_b, z_a) if reverses else (z_a, z_b) + z_lo = max(z_lo, -1.0) + z_hi = min(z_hi, 1.0) + if not (z_hi > z_lo): + raise ValueError("cosine_sampler_limits: {} range [{}, {}] maps to an empty sampling interval [{}, {}]".format(kind, lo, hi, z_lo, z_hi)) + return z_lo, z_hi + + +def ret_dec_samp_vector(dec_lo, dec_hi): + """Sampling pdf in DECLINATION for a uniform-in-sin(dec) draw truncated to + [dec_lo, dec_hi]. Normalized to unity over that box (the samplers that use + an explicit cdf_inv do not renormalize the pdf themselves). Reduces to + dec_samp_vector for the full range.""" + z_lo, z_hi = cosine_sampler_limits(dec_lo, dec_hi, 'declination') + lo_c, hi_c = clip_angle_limits(dec_lo, dec_hi, 'declination') + norm = z_hi - z_lo + def _pdf(x, xpy=None): + xpy = infer_array_module(x, xpy) + x = xpy.asarray(x, dtype=numpy.float64) + vals = xpy.cos(x)/norm + return xpy.where((x >= lo_c) & (x <= hi_c), vals, xpy.zeros_like(vals)) + return _pdf + + +def ret_dec_samp_cdf_inv_vector(dec_lo, dec_hi): + """Inverse CDF (p in [0,1] -> declination) for uniform-in-sin(dec) truncated + to [dec_lo, dec_hi]. Monotonically increasing in p.""" + z_lo, z_hi = cosine_sampler_limits(dec_lo, dec_hi, 'declination') + def _cdf_inv(p, xpy=None): + xpy = infer_array_module(p, xpy) + p = xpy.asarray(p, dtype=numpy.float64) + return xpy.arcsin(xpy.clip(z_lo + p*(z_hi - z_lo), -1.0, 1.0)) + return _cdf_inv + + +def ret_cos_samp_vector(incl_lo, incl_hi): + """Sampling pdf in INCLINATION for a uniform-in-cos(iota) draw truncated to + [incl_lo, incl_hi]. Normalized to unity over that box. Reduces to + cos_samp_vector for the full range.""" + z_lo, z_hi = cosine_sampler_limits(incl_lo, incl_hi, 'inclination') + lo_c, hi_c = clip_angle_limits(incl_lo, incl_hi, 'inclination') + norm = z_hi - z_lo + def _pdf(x, xpy=None): + xpy = infer_array_module(x, xpy) + x = xpy.asarray(x, dtype=numpy.float64) + vals = xpy.sin(x)/norm + return xpy.where((x >= lo_c) & (x <= hi_c), vals, xpy.zeros_like(vals)) + return _pdf + + +def ret_cos_samp_cdf_inv_vector(incl_lo, incl_hi): + """Inverse CDF (p in [0,1] -> inclination) for uniform-in-cos(iota) + truncated to [incl_lo, incl_hi]. Monotonically increasing in p: p=0 gives + incl_lo (which is arccos of the UPPER cosine limit -- note the swap).""" + z_lo, z_hi = cosine_sampler_limits(incl_lo, incl_hi, 'inclination') + def _cdf_inv(p, xpy=None): + xpy = infer_array_module(p, xpy) + p = xpy.asarray(p, dtype=numpy.float64) + return xpy.arccos(xpy.clip(z_hi - p*(z_hi - z_lo), -1.0, 1.0)) + return _cdf_inv + + def pseudo_dist_samp(r0,r): return r*r*numpy.exp( - (r0/r)*(r0/r)/2. + r0/r)+0.01 # put a floor on probability, so we converge. Note this floor only cuts out NEARBY distances diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index 447099fb7..da03795b9 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -50,6 +50,12 @@ import glue.lal import RIFT.lalsimutils as lalsimutils import RIFT.integrators.mcsampler as mcsampler +# NOTE: the name 'mcsampler' above is REBOUND below to mcsamplerGPU for some --sampler-method +# choices, so the zoom-box helpers are imported under their own names (they are pure-numpy, +# xpy-aware, and identical for every backend). +from RIFT.integrators.mcsampler import (clip_angle_limits, cosine_sampler_limits, + ret_dec_samp_vector, ret_dec_samp_cdf_inv_vector, + ret_cos_samp_vector, ret_cos_samp_cdf_inv_vector) import RIFT.misc.sky_rotations as sky_rotations try: import RIFT.integrators.mcsamplerEnsemble as mcsamplerEnsemble @@ -281,6 +287,10 @@ integration_params.add_option("--d-max", default=10000,type=float,help="Maximum integration_params.add_option("--d-min", default=1,type=float,help="Minimum distance in volume integral. Used to SET THE PRIOR; changing this value changes the numerical answer.") integration_params.add_option("--declination-cosine-sampler",action='store_true',help="If specified, the parameter used for declination is cos(dec), not dec") integration_params.add_option("--inclination-cosine-sampler",action='store_true',help="If specified, the parameter used for inclination is cos(dec), not dec") +integration_params.add_option("--limit-right-ascension",default=None,help="Restrict RA sampling AND prior to 'LO,HI' [rad] (truth-centered zoom box). Narrows the extrinsic prior like --d-min/--d-max do for distance; keep the box large vs the posterior so credible regions are unaffected. Not compatible with --internal-sky-network-coordinates (the sampled sky angles are then in a rotated frame).") +integration_params.add_option("--limit-declination",default=None,help="Restrict declination sampling AND prior to 'LO,HI' [rad]. Always given in radians of DECLINATION: with --declination-cosine-sampler the box is transformed internally to the sampled coordinate sin(dec). Not compatible with --internal-sky-network-coordinates.") +integration_params.add_option("--limit-inclination",default=None,help="Restrict inclination sampling AND prior to 'LO,HI' [rad]. Always given in radians of INCLINATION: with --inclination-cosine-sampler the box is transformed internally to the sampled coordinate cos(iota), which reverses the limit order.") +integration_params.add_option("--limit-psi",default=None,help="Restrict polarization psi sampling AND prior to 'LO,HI' [rad].") integration_params.add_option("--internal-rotate-phase", action='store_true',help="If specified, the integration sampler uses phase_p ==phi+psi and phase_m == phi-psi as sampling coordinates, both ranging from 0 to 4 pi. The prior is twice as large.") integration_params.add_option("--internal-sky-network-coordinates",action='store_true',help="If specified, perform integration in sky coordinates aligned with the first two IFOs provided") integration_params.add_option("--internal-sky-network-coordinates-raw",action='store_true',help="If specified, does not attempt to organize IFO network sensibly, uses them AS PROVIDED IN ORDER.") @@ -761,6 +771,28 @@ param_limits = { "psi": (0, 2*numpy.pi), if opts.internal_rotate_phase: param_limits['psi'] = (0, 4*numpy.pi) param_limits['phi_orb'] = (0, 4*numpy.pi) +# Optional truth-centered "zoom box": narrow the extrinsic sampling AND prior ranges so the +# adaptive sampler can resolve a narrow high-SNR peak it could never find from the full prior. +# Threads through param_limits into every sky/orientation sampler + its pdf/cdf_inv/prior_pdf. +# The limits are ALWAYS specified in radians of the physical angle; the cosine samplers +# (--declination-cosine-sampler / --inclination-cosine-sampler) transform them below into the +# coordinate they actually sample (sin(dec) resp. cos(iota)). +for _optv, _k in [(opts.limit_psi, 'psi'), (opts.limit_right_ascension, 'right_ascension'), + (opts.limit_declination, 'declination'), (opts.limit_inclination, 'inclination')]: + if _optv: + try: + _lo, _hi = [float(_x) for _x in str(_optv).split(',')] + except ValueError: + raise SystemExit(" --limit-{} expects 'LO,HI' in radians, got '{}'".format(_k.replace('_','-'), _optv)) + if _k in ('declination', 'inclination'): + # validates lo _lo): + raise SystemExit(" --limit-{}: empty or inverted range [{}, {}] (need LO < HI)".format(_k.replace('_','-'), _lo, _hi)) + param_limits[_k] = (_lo, _hi) + print(" [limit] restricting {} sampling/prior to [{:.4f}, {:.4f}]".format(_k, _lo, _hi)) +limit_declination_active = bool(opts.limit_declination) +limit_inclination_active = bool(opts.limit_inclination) # # Parameter integral sampling strategy @@ -905,23 +937,40 @@ if (opts.sampler_method == "adaptive_cartesian_gpu" or opts.sampler_method == ' adapt_extra_extrinsic=True if not opts.inclination_cosine_sampler: - incl_sampler = mcsampler.cos_samp_vector # this is NOT dec_samp_vector, because the angular zero point is different! - incl_sampler_cdf_inv = mcsampler.cos_samp_cdf_inv_vector - sampler.add_parameter("inclination", - pdf = incl_sampler, - cdf_inv = incl_sampler_cdf_inv, - left_limit = param_limits["inclination"][0], + if limit_inclination_active: + # truncated uniform-in-cos(iota) draw, expressed in the ANGLE coordinate + incl_sampler = ret_cos_samp_vector(param_limits["inclination"][0], param_limits["inclination"][1]) + incl_sampler_cdf_inv = ret_cos_samp_cdf_inv_vector(param_limits["inclination"][0], param_limits["inclination"][1]) + else: + incl_sampler = mcsampler.cos_samp_vector # this is NOT dec_samp_vector, because the angular zero point is different! + incl_sampler_cdf_inv = mcsampler.cos_samp_cdf_inv_vector + sampler.add_parameter("inclination", + pdf = incl_sampler, + cdf_inv = incl_sampler_cdf_inv, + left_limit = param_limits["inclination"][0], right_limit = param_limits["inclination"][1], prior_pdf = mcsampler.uniform_samp_theta) # do not adapt in parameter going to zero at edge else: - incl_sampler = mcsampler.ret_uniform_samp_vector_alt(-1.0,1.0) - incl_sampler_cdf_inv = lambda x: x*2.0-1. #functools.partial(mcsampler.uniform_samp_cdf_inv_vector,-1,1) - sampler.add_parameter("inclination", - pdf = incl_sampler, - cdf_inv = incl_sampler_cdf_inv, - left_limit = -1, - right_limit = 1, - prior_pdf = incl_sampler, + # Sample uniformly in cos(iota) [=1 face-on, -1 face-off]: the likelihood closure below + # converts back with iota = arccos(z). A --limit-inclination box must therefore be mapped + # into z, and because cos() DECREASES on [0,pi] the limits SWAP: + # [iota_lo, iota_hi] -> [cos(iota_hi), cos(iota_lo)] + # (before this fix the range was hardcoded to [-1,1] and --limit-inclination was silently ignored). + incl_z_lo, incl_z_hi = cosine_sampler_limits(param_limits["inclination"][0], param_limits["inclination"][1], 'inclination') + if limit_inclination_active: + print(" [limit] inclination box [{:.4f}, {:.4f}] rad -> cos(iota) sampling range [{:.6f}, {:.6f}]".format( + param_limits["inclination"][0], param_limits["inclination"][1], incl_z_lo, incl_z_hi)) + incl_sampler = mcsampler.ret_uniform_samp_vector_alt(incl_z_lo, incl_z_hi) + incl_sampler_cdf_inv = lambda x, _a=incl_z_lo, _b=incl_z_hi: _a + x*(_b-_a) # functools.partial(mcsampler.uniform_samp_cdf_inv_vector,_a,_b) + sampler.add_parameter("inclination", + pdf = incl_sampler, + cdf_inv = incl_sampler_cdf_inv, + left_limit = incl_z_lo, + right_limit = incl_z_hi, + # prior density in cos(iota) is the FULL-RANGE constant 1/2, deliberately not renormalized + # to the box, so restricting the box costs exactly the prior mass it should -- matching the + # non-cosine branch, where prior_pdf=0.5*sin(iota) is likewise not renormalized. + prior_pdf = mcsampler.ret_uniform_samp_vector_alt(-1.0,1.0), adaptive_sampling=adapt_extra_extrinsic) # @@ -1017,8 +1066,12 @@ if opts.internal_sky_network_coordinates: else: sky_rotations.assign_sky_frame(ifo_list[0], ifo_list[1], fiducial_epoch) frm = identity_convert_togpu(sky_rotations.frm) - my_rotation = functools.partial(lalsimutils.polar_angles_in_frame_alt,frm,xpy=xpy_default) - my_rotation_cpu = functools.partial(lalsimutils.polar_angles_in_frame_alt,sky_rotations.frm,xpy=np) + my_rotation = functools.partial(lalsimutils.polar_angles_in_frame_alt,frm,xpy=xpy_default) + my_rotation_cpu = functools.partial(lalsimutils.polar_angles_in_frame_alt,sky_rotations.frm,xpy=np) +if opts.internal_sky_network_coordinates and (opts.limit_right_ascension or opts.limit_declination): + # The sampled RA/dec live in the network-aligned frame, so a truth-centered sky box given in + # equatorial coordinates would silently select the wrong patch of sky. Fail loudly. + raise SystemExit(" --limit-right-ascension / --limit-declination are sky boxes in EQUATORIAL coordinates and are not compatible with --internal-sky-network-coordinates (which samples in a rotated, network-aligned frame). Drop --internal-sky-network-coordinates when using a sky zoom box.") # # Intrinsic parameters @@ -1089,26 +1142,42 @@ else: # sky sampler: cos(dec) uniform in [-1, 1), adaptive sampling # if not opts.declination_cosine_sampler: - dec_sampler = mcsampler.dec_samp_vector - dec_sampler_cdf_inv = mcsampler.dec_samp_cdf_inv_vector - sampler.add_parameter("declination", - pdf = dec_sampler, - cdf_inv = dec_sampler_cdf_inv, - left_limit = param_limits["declination"][0], + if limit_declination_active: + # truncated uniform-in-sin(dec) draw, expressed in the ANGLE coordinate + dec_sampler = ret_dec_samp_vector(param_limits["declination"][0], param_limits["declination"][1]) + dec_sampler_cdf_inv = ret_dec_samp_cdf_inv_vector(param_limits["declination"][0], param_limits["declination"][1]) + else: + dec_sampler = mcsampler.dec_samp_vector + dec_sampler_cdf_inv = mcsampler.dec_samp_cdf_inv_vector + sampler.add_parameter("declination", + pdf = dec_sampler, + cdf_inv = dec_sampler_cdf_inv, + left_limit = param_limits["declination"][0], right_limit = param_limits["declination"][1], prior_pdf = mcsampler.uniform_samp_dec, adaptive_sampling = opts.force_adapt_all or (not opts.no_adapt)) else: # Sample uniformly in cos(polar_theta), =1 for north pole, -1 for south pole. # Propagate carefully in conversions: time of flight libraries use RA,DEC - dec_sampler = mcsampler.ret_uniform_samp_vector_alt(-1.0,1.0) - dec_sampler_cdf_inv = lambda x: x*2.0-1. # functools.partial(mcsampler.uniform_samp_cdf_inv_vector,-1,1) - sampler.add_parameter("declination", - pdf = dec_sampler, - cdf_inv = dec_sampler_cdf_inv, - left_limit = -1, - right_limit = 1, - prior_pdf = dec_sampler, + # polar_theta = pi/2 - dec, so the sampled variable is z = sin(dec) (see the likelihood + # closures: dec = pi/2 - arccos(z)). sin() INCREASES on [-pi/2,pi/2], so a + # --limit-declination box maps order-preservingly: [lo,hi] -> [sin(lo), sin(hi)] + # (before this fix the range was hardcoded to [-1,1] and --limit-declination was silently ignored). + dec_z_lo, dec_z_hi = cosine_sampler_limits(param_limits["declination"][0], param_limits["declination"][1], 'declination') + if limit_declination_active: + print(" [limit] declination box [{:.4f}, {:.4f}] rad -> sin(dec) sampling range [{:.6f}, {:.6f}]".format( + param_limits["declination"][0], param_limits["declination"][1], dec_z_lo, dec_z_hi)) + dec_sampler = mcsampler.ret_uniform_samp_vector_alt(dec_z_lo, dec_z_hi) + dec_sampler_cdf_inv = lambda x, _a=dec_z_lo, _b=dec_z_hi: _a + x*(_b-_a) # functools.partial(mcsampler.uniform_samp_cdf_inv_vector,_a,_b) + sampler.add_parameter("declination", + pdf = dec_sampler, + cdf_inv = dec_sampler_cdf_inv, + left_limit = dec_z_lo, + right_limit = dec_z_hi, + # prior density in sin(dec) is the FULL-RANGE constant 1/2, deliberately not renormalized + # to the box, so the prior mass removed by the box matches the non-cosine branch, where + # prior_pdf=uniform_samp_dec=0.5*cos(dec) is likewise not renormalized. + prior_pdf = mcsampler.ret_uniform_samp_vector_alt(-1.0,1.0), adaptive_sampling = opts.force_adapt_all or (not opts.no_adapt)) if not opts.time_marginalization: @@ -1548,7 +1617,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t tvals = numpy.linspace(-t_ref_wind,t_ref_wind,int((t_ref_wind)*2/P.deltaT)) # choose an array at the target sampling rate. P is inherited globally for ph, th, phr, ic, ps, di in zip(right_ascension, dec, - phi_orb, inclination, psi, distance): + phi_orb, incl, psi, distance): # 'incl', NOT the raw sampled 'inclination': under --inclination-cosine-sampler the sampled variable is cos(iota) P.phi = ph # right ascension P.theta = th # declination P.tref = fiducial_epoch # see 'tvals', above diff --git a/MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py b/MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py new file mode 100644 index 000000000..151d63c39 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_limit_cosine_samplers.py @@ -0,0 +1,454 @@ +#!/usr/bin/env python +""" +Regression tests for the extrinsic "zoom box" options +(--limit-declination / --limit-inclination / --limit-right-ascension / --limit-psi) +under the COSINE samplers (--declination-cosine-sampler / --inclination-cosine-sampler). + +Background (the bug these tests lock down): the cosine branches of +bin/integrate_likelihood_extrinsic_batchmode used to hardcode left_limit=-1, +right_limit=1 and never consulted param_limits, so --limit-declination and +--limit-inclination were SILENTLY IGNORED whenever the cosine samplers were on +(which is the production default). No error, no warning, no narrowing. + +Coordinate conventions, read off the likelihood closures in that script +(`dec = pi/2 - arccos(z)`, `iota = arccos(z)`): + + declination: sampled variable z = sin(dec), sin INCREASING on [-pi/2,pi/2] + => [lo,hi] -> [sin(lo), sin(hi)] (order preserved) + inclination: sampled variable z = cos(iota), cos DECREASING on [0,pi] + => [lo,hi] -> [cos(hi), cos(lo)] (order SWAPS) + +The second one is the easy thing to get backwards, so it gets its own test. +""" + +import os +import sys +import types + +import numpy as np +import pytest + +import RIFT.integrators.mcsampler as mcsampler +from RIFT.integrators.mcsampler import ( + clip_angle_limits, + cosine_sampler_limits, + infer_array_module, + ret_cos_samp_cdf_inv_vector, + ret_cos_samp_vector, + ret_dec_samp_cdf_inv_vector, + ret_dec_samp_vector, +) + +# np.trapz was REMOVED in numpy 2.x (renamed np.trapezoid); the modern CI lane runs +# numpy>=2, the legacy lane numpy 1.24.4, so pick whichever exists. +_trapz = getattr(np, 'trapezoid', None) or np.trapz + +# The conversions applied inside the ILE likelihood closures, verbatim +# (the .astype mirrors the numpy.copy(...).astype(numpy.float64) those closures do, +# because mcsampler hands back object arrays). +_dec_from_z = lambda z: np.pi / 2 - np.arccos(np.asarray(z).astype(np.float64)) +_incl_from_z = lambda z: np.arccos(np.asarray(z).astype(np.float64)) + + +### +### 1. Coordinate transform +### + +def test_declination_limits_map_to_sin_and_preserve_order(): + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + assert z_lo == pytest.approx(np.sin(lo)) + assert z_hi == pytest.approx(np.sin(hi)) + assert z_lo < z_hi + # round trip through the conversion the likelihood actually applies + assert _dec_from_z(z_lo) == pytest.approx(lo) + assert _dec_from_z(z_hi) == pytest.approx(hi) + + +def test_inclination_limits_map_to_cos_and_SWAP_order(): + lo, hi = 0.30, 1.20 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + # this is the assertion that fails if someone writes [cos(lo), cos(hi)] + assert z_lo == pytest.approx(np.cos(hi)) + assert z_hi == pytest.approx(np.cos(lo)) + assert z_lo < z_hi + # and the round trip: the LOWER cosine limit is the UPPER angle + assert _incl_from_z(z_lo) == pytest.approx(hi) + assert _incl_from_z(z_hi) == pytest.approx(lo) + # explicit guard against the naive (unswapped) answer + assert (z_lo, z_hi) != pytest.approx((np.cos(lo), np.cos(hi))) + + +def test_inclination_swap_would_produce_empty_or_inverted_interval(): + """A [cos(lo), cos(hi)] implementation is not merely mislabeled: it is inverted.""" + lo, hi = 0.30, 1.20 + naive_lo, naive_hi = np.cos(lo), np.cos(hi) + assert naive_lo > naive_hi # inverted -> would silently give a negative volume + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + assert z_hi - z_lo == pytest.approx(naive_lo - naive_hi) # same width, correct sign + + +def test_full_range_is_a_no_op(): + assert cosine_sampler_limits(-np.pi / 2, np.pi / 2, 'declination') == pytest.approx((-1.0, 1.0)) + assert cosine_sampler_limits(0.0, np.pi, 'inclination') == pytest.approx((-1.0, 1.0)) + + +def test_limits_are_clipped_to_the_physical_domain(): + assert clip_angle_limits(-10.0, 0.1, 'declination') == pytest.approx((-np.pi / 2, 0.1)) + assert clip_angle_limits(0.1, 10.0, 'inclination') == pytest.approx((0.1, np.pi)) + z_lo, z_hi = cosine_sampler_limits(-10.0, 10.0, 'declination') + assert (z_lo, z_hi) == pytest.approx((-1.0, 1.0)) + + +@pytest.mark.parametrize('kind', ['declination', 'inclination']) +def test_empty_or_inverted_range_raises(kind): + with pytest.raises(ValueError): + cosine_sampler_limits(0.5, 0.5, kind) # empty + with pytest.raises(ValueError): + cosine_sampler_limits(0.9, 0.2, kind) # inverted + with pytest.raises(ValueError): + cosine_sampler_limits(np.nan, 0.2, kind) # non-finite + + +def test_range_outside_physical_domain_raises(): + with pytest.raises(ValueError): + cosine_sampler_limits(2.0, 3.0, 'declination') # entirely north of the pole + with pytest.raises(ValueError): + cosine_sampler_limits(-2.0, -1.0, 'inclination') # entirely below iota=0 + + +def test_unknown_angle_raises(): + with pytest.raises(ValueError): + cosine_sampler_limits(0.1, 0.2, 'right_ascension') + + +### +### 2. Support: the box actually restricts, in both samplers +### + +def test_declination_box_restricts_support_in_both_samplers(): + lo, hi = -0.62, -0.41 + p = np.linspace(0.0, 1.0, 4001) + + # plain (angle) sampler, truncated + dec_plain = ret_dec_samp_cdf_inv_vector(lo, hi)(p) + assert dec_plain.min() == pytest.approx(lo) + assert dec_plain.max() == pytest.approx(hi) + + # cosine sampler: uniform in z over the transformed box, then converted back + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + dec_cos = _dec_from_z(z_lo + p * (z_hi - z_lo)) + assert dec_cos.min() == pytest.approx(lo) + assert dec_cos.max() == pytest.approx(hi) + + # ... and the two draws are the SAME map (both are uniform-in-sin(dec) on the box) + assert np.allclose(np.sort(dec_plain), np.sort(dec_cos)) + + +def test_inclination_box_restricts_support_in_both_samplers(): + lo, hi = 0.30, 1.20 + p = np.linspace(0.0, 1.0, 4001) + + incl_plain = ret_cos_samp_cdf_inv_vector(lo, hi)(p) + assert incl_plain.min() == pytest.approx(lo) + assert incl_plain.max() == pytest.approx(hi) + + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + incl_cos = _incl_from_z(z_lo + p * (z_hi - z_lo)) + assert incl_cos.min() == pytest.approx(lo) + assert incl_cos.max() == pytest.approx(hi) + + assert np.allclose(np.sort(incl_plain), np.sort(incl_cos)) + + +def test_truncated_samplers_reduce_to_the_untruncated_ones(): + """Full-range truncated samplers must reproduce the legacy distributions.""" + p = np.linspace(1e-9, 1 - 1e-9, 501) + dec_new = np.sort(ret_dec_samp_cdf_inv_vector(-np.pi / 2, np.pi / 2)(p)) + dec_old = np.sort(mcsampler.dec_samp_cdf_inv_vector(p)) + assert np.allclose(dec_new, dec_old, atol=1e-10) + + incl_new = np.sort(ret_cos_samp_cdf_inv_vector(0.0, np.pi)(p)) + incl_old = np.sort(mcsampler.cos_samp_cdf_inv_vector(p)) + assert np.allclose(incl_new, incl_old, atol=1e-10) + + x = np.linspace(-np.pi / 2 + 1e-6, np.pi / 2 - 1e-6, 257) + assert np.allclose(ret_dec_samp_vector(-np.pi / 2, np.pi / 2)(x), + mcsampler.dec_samp_vector(x)) + y = np.linspace(1e-6, np.pi - 1e-6, 257) + assert np.allclose(ret_cos_samp_vector(0.0, np.pi)(y), mcsampler.cos_samp_vector(y)) + + +### +### 3. Normalization: identical prior mass / lnZ in both samplers +### +# mcsampler / mcsamplerGPU weight each draw by prior_pdf(x)/pdf(x); the expectation of +# that weight over the sampling pdf is the prior MASS inside the box. The two branches +# must agree, otherwise the same physical box would give different lnZ. + +def _weight_plain_dec(lo, hi, dec): + pdf = ret_dec_samp_vector(lo, hi)(dec) + prior = mcsampler.uniform_samp_dec(dec) # 0.5*cos(dec), NOT renormalized + return prior / pdf + + +def _weight_cosine(z_lo, z_hi, z): + pdf = mcsampler.ret_uniform_samp_vector_alt(z_lo, z_hi)(z) + prior = mcsampler.ret_uniform_samp_vector_alt(-1.0, 1.0)(z) # constant 1/2 in z + return prior / pdf + + +def test_declination_box_same_prior_mass_in_both_samplers(): + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + expected = 0.5 * (np.sin(hi) - np.sin(lo)) # isotropic prior mass in the box + + dec = np.linspace(lo + 1e-9, hi - 1e-9, 2001) + w_plain = _weight_plain_dec(lo, hi, dec) + assert np.allclose(w_plain, expected) # constant weight + + z = np.linspace(z_lo + 1e-12, z_hi - 1e-12, 2001) + w_cos = _weight_cosine(z_lo, z_hi, z) + assert np.allclose(w_cos, expected) + + +def test_inclination_box_same_prior_mass_in_both_samplers(): + lo, hi = 0.30, 1.20 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + expected = 0.5 * (np.cos(lo) - np.cos(hi)) + + incl = np.linspace(lo + 1e-9, hi - 1e-9, 2001) + w_plain = mcsampler.uniform_samp_theta(incl) / ret_cos_samp_vector(lo, hi)(incl) + assert np.allclose(w_plain, expected) + + z = np.linspace(z_lo + 1e-12, z_hi - 1e-12, 2001) + assert np.allclose(_weight_cosine(z_lo, z_hi, z), expected) + + +def test_prior_mass_scales_like_the_box_not_the_full_prior(): + """Sanity: narrowing must actually cost prior mass (that is the whole point).""" + full = 0.5 * (np.sin(np.pi / 2) - np.sin(-np.pi / 2)) + lo, hi = -0.62, -0.41 + narrow = 0.5 * (np.sin(hi) - np.sin(lo)) + assert narrow < full + assert full / narrow == pytest.approx(1.0 / (0.5 * (np.sin(hi) - np.sin(lo)))) + + +### +### 4. AV-style estimator equivalence (the production sampler) +### +# mcsamplerAdaptiveVolume draws uniformly in [llim,rlim] and multiplies by the sampling +# volume V_s = prod(rlim-llim), weighting by prior_pdf. Same box => same answer. + +def _av_prior_integral(llim, rlim, prior_pdf, n=200001): + x = np.linspace(llim, rlim, n) + return (rlim - llim) * np.mean(prior_pdf(x)) + + +def test_AV_style_prior_integral_matches_between_samplers_dec(): + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + plain = _av_prior_integral(lo, hi, mcsampler.uniform_samp_dec) + cosine = _av_prior_integral(z_lo, z_hi, lambda x: np.full_like(x, 0.5)) + assert plain == pytest.approx(cosine, rel=1e-6) + assert plain == pytest.approx(0.5 * (np.sin(hi) - np.sin(lo)), rel=1e-6) + + +def test_AV_style_prior_integral_matches_between_samplers_incl(): + lo, hi = 0.30, 1.20 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + plain = _av_prior_integral(lo, hi, mcsampler.uniform_samp_theta) + cosine = _av_prior_integral(z_lo, z_hi, lambda x: np.full_like(x, 0.5)) + assert plain == pytest.approx(cosine, rel=1e-6) + assert plain == pytest.approx(0.5 * (np.cos(lo) - np.cos(hi)), rel=1e-6) + + +def test_AV_style_posterior_shape_matches_between_samplers_dec(): + """Posterior *shape* in declination must be identical (isotropic prior preserved).""" + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + dec_grid = np.linspace(lo, hi, 4001) + + # plain branch: uniform in dec, weight 0.5*cos(dec) + q_plain = mcsampler.uniform_samp_dec(dec_grid) + q_plain = q_plain / _trapz(q_plain, dec_grid) + + # cosine branch: uniform in z=sin(dec), weight 1/2 -> push forward to dec + q_cos = 0.5 * np.cos(dec_grid) # Jacobian dz/ddec = cos(dec) + q_cos = q_cos / _trapz(q_cos, dec_grid) + + assert np.allclose(q_plain, q_cos) + assert z_lo < z_hi + + +### +### 5. End-to-end through MCSampler: same box => same integral +### + +def _integrate_1d(name, pdf, cdf_inv, llim, rlim, prior_pdf, fn, nmax=20000): + s = mcsampler.MCSampler() + s.add_parameter(name, pdf=pdf, cdf_inv=cdf_inv, left_limit=llim, right_limit=rlim, + prior_pdf=prior_pdf) + res = s.integrate(fn, name, nmax=nmax, n=1000, no_protect_names=True, verbose=False) + return res[0] + + +def test_end_to_end_declination_box_gives_same_integral(): + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + expected = 0.5 * (np.sin(hi) - np.sin(lo)) + + # integrand == 1 -> the integral IS the prior mass in the box + unit = lambda declination: np.ones(np.shape(declination)) + + plain = _integrate_1d('declination', + ret_dec_samp_vector(lo, hi), ret_dec_samp_cdf_inv_vector(lo, hi), + lo, hi, mcsampler.uniform_samp_dec, unit) + cosine = _integrate_1d('declination', + mcsampler.ret_uniform_samp_vector_alt(z_lo, z_hi), + lambda x, _a=z_lo, _b=z_hi: _a + x * (_b - _a), + z_lo, z_hi, mcsampler.ret_uniform_samp_vector_alt(-1.0, 1.0), unit) + + assert float(plain) == pytest.approx(expected, rel=1e-6) + assert float(cosine) == pytest.approx(expected, rel=1e-6) + assert float(plain) == pytest.approx(float(cosine), rel=1e-6) + + +def test_end_to_end_inclination_box_gives_same_integral(): + lo, hi = 0.30, 1.20 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'inclination') + expected = 0.5 * (np.cos(lo) - np.cos(hi)) + + unit = lambda inclination: np.ones(np.shape(inclination)) + + plain = _integrate_1d('inclination', + ret_cos_samp_vector(lo, hi), ret_cos_samp_cdf_inv_vector(lo, hi), + lo, hi, mcsampler.uniform_samp_theta, unit) + cosine = _integrate_1d('inclination', + mcsampler.ret_uniform_samp_vector_alt(z_lo, z_hi), + lambda x, _a=z_lo, _b=z_hi: _a + x * (_b - _a), + z_lo, z_hi, mcsampler.ret_uniform_samp_vector_alt(-1.0, 1.0), unit) + + assert float(plain) == pytest.approx(expected, rel=1e-6) + assert float(cosine) == pytest.approx(expected, rel=1e-6) + + +def test_end_to_end_declination_box_same_integral_for_a_peaked_likelihood(): + """Non-constant integrand: the cosine branch must convert z -> dec exactly as ILE does.""" + lo, hi = -0.62, -0.41 + z_lo, z_hi = cosine_sampler_limits(lo, hi, 'declination') + mu, sig = -0.50, 0.03 + like_dec = lambda d: np.exp(-0.5 * ((np.asarray(d, dtype=float) - mu) / sig) ** 2) + + plain = _integrate_1d('declination', + ret_dec_samp_vector(lo, hi), ret_dec_samp_cdf_inv_vector(lo, hi), + lo, hi, mcsampler.uniform_samp_dec, + lambda declination: like_dec(declination), nmax=200000) + cosine = _integrate_1d('declination', + mcsampler.ret_uniform_samp_vector_alt(z_lo, z_hi), + lambda x, _a=z_lo, _b=z_hi: _a + x * (_b - _a), + z_lo, z_hi, mcsampler.ret_uniform_samp_vector_alt(-1.0, 1.0), + lambda declination: like_dec(_dec_from_z(declination)), nmax=200000) + + # analytic reference: int 0.5*cos(dec) * L(dec) ddec over the box + grid = np.linspace(lo, hi, 200001) + ref = _trapz(0.5 * np.cos(grid) * like_dec(grid), grid) + + assert float(plain) == pytest.approx(ref, rel=2e-2) + assert float(cosine) == pytest.approx(ref, rel=2e-2) + + +### +### 6. Backend dispatch: the GPU sampler calls these with ONE positional argument +### +# mcsamplerGPU.draw_simplified() does `self.cdf_inv[param](unif_samples)` and +# `self.pdf[param](param_samples)` with a cupy array and no xpy= keyword. A numpy +# default would then evaluate numpy.asarray(cupy_array), which raises, so the closures +# infer the backend from the argument. There is no GPU in CI, so the cupy module is +# stood in for by a recording shim registered in sys.modules (the inference is a +# sys.modules lookup on the array type's top-level module, exactly as for cupy). + +_XPY_FUNCS = ('asarray', 'where', 'cos', 'sin', 'arcsin', 'arccos', 'clip', 'zeros_like') + + +def _fake_backend(monkeypatch, name='fake_xpy_backend'): + """Return (array_type, calls): a numpy-backed stand-in for cupy.""" + calls = [] + mod = types.ModuleType(name) + for _name in _XPY_FUNCS: + def _record(*args, _f=getattr(np, _name), _n=_name, **kwargs): + calls.append(_n) + return _f(*args, **kwargs) + setattr(mod, _name, _record) + monkeypatch.setitem(sys.modules, name, mod) + + class _FakeArray(np.ndarray): + pass + _FakeArray.__module__ = name # this is what the inference keys on + return _FakeArray, calls + + +def test_infer_array_module_defaults_to_numpy_and_honors_explicit_xpy(): + assert infer_array_module(np.linspace(0, 1, 4)) is np + assert infer_array_module([0.1, 0.2]) is np + assert infer_array_module(0.3) is np + sentinel = object() + assert infer_array_module(np.linspace(0, 1, 4), sentinel) is sentinel + + +@pytest.mark.parametrize('factory,arg', [ + (lambda: ret_dec_samp_vector(-0.62, -0.41), np.linspace(-0.62, -0.41, 32)), + (lambda: ret_cos_samp_vector(0.30, 1.20), np.linspace(0.30, 1.20, 32)), + (lambda: ret_dec_samp_cdf_inv_vector(-0.62, -0.41), np.linspace(0.0, 1.0, 32)), + (lambda: ret_cos_samp_cdf_inv_vector(0.30, 1.20), np.linspace(0.0, 1.0, 32)), +]) +def test_truncated_closures_dispatch_to_the_arrays_own_backend(monkeypatch, factory, arg): + """Called with a single positional non-numpy array, they must use ITS module.""" + fake_array_type, calls = _fake_backend(monkeypatch) + fn = factory() + expected = fn(arg) # host reference, numpy path + del calls[:] + got = fn(arg.view(fake_array_type)) # the draw_simplified() call signature + assert calls, 'closure ignored the backend of its argument (would fail on cupy input)' + assert np.allclose(np.asarray(got), np.asarray(expected)) + + +@pytest.mark.parametrize('angle,factory_pdf,factory_cdf,prior_name,box', [ + ('declination', ret_dec_samp_vector, ret_dec_samp_cdf_inv_vector, 'uniform_samp_dec', (-0.62, -0.41)), + ('inclination', ret_cos_samp_vector, ret_cos_samp_cdf_inv_vector, 'uniform_samp_theta', (0.30, 1.20)), +]) +def test_gpu_sampler_draws_inside_the_box(angle, factory_pdf, factory_cdf, prior_name, box): + """Exercise the actual mcsamplerGPU call path (CPU fallback when cupy is absent).""" + mcsamplerGPU = pytest.importorskip('RIFT.integrators.mcsamplerGPU') + lo, hi = box + s = mcsamplerGPU.MCSampler() + s.add_parameter(angle, pdf=factory_pdf(lo, hi), cdf_inv=factory_cdf(lo, hi), + left_limit=lo, right_limit=hi, + prior_pdf=getattr(mcsamplerGPU, prior_name)) + rv = s.draw_simplified(2000, angle)[-1] + drawn = np.asarray(mcsamplerGPU.identity_convert(rv)).reshape(-1) + assert drawn.min() >= lo - 1e-9 + assert drawn.max() <= hi + 1e-9 + # and it fills the box, i.e. the limits were not merely clipped to a point + assert drawn.max() - drawn.min() > 0.8 * (hi - lo) + + +### +### 7. Wiring: the bin script must not reintroduce the hardcoded [-1,1] range +### + +_ILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), + '..', 'bin', 'integrate_likelihood_extrinsic_batchmode') + + +@pytest.mark.skipif(not os.path.exists(_ILE), reason='ILE executable not in this tree') +def test_ile_cosine_branches_consult_param_limits(): + with open(_ILE) as f: + src = f.read() + for angle in ('declination', 'inclination'): + needle = 'cosine_sampler_limits(param_limits["{}"][0], param_limits["{}"][1], \'{}\')'.format(angle, angle, angle) + assert needle in src, \ + "cosine {} branch no longer transforms param_limits -- --limit-{} would be silently ignored".format(angle, angle) + # the old hardcoded literals must be gone from the sampler setup + assert 'left_limit = -1,' not in src + assert 'right_limit = 1,' not in src