diff --git a/AGENTS.md b/AGENTS.md index fda1117fc..a69653074 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -26,6 +26,16 @@ python MonteCarloMarginalizeCode/Code/test/test_likelihood.py ## Testing Tests use pytest but have no standard runner. Run individual test files directly. +### Merge gate for integrator changes (IMPORTANT) +Any change under `MonteCarloMarginalizeCode/Code/RIFT/integrators/` must pass the +**posterior shape-recovery gate** in +`MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/` +before merging into a production line (run base + candidate with identical seeds, +then `compare_shape_results.py base.json pr.json`; exit 1 = merge-blocking). +The fast CI integral test is NOT sufficient — integrators have shipped confident, +integral-invisible shape failures and silent n_eff~1 degradations that only this +gate catches. See `RIFT/integrators/TESTING.md` for the recipe and caveats. + ## Important CLI tools - `integrate_likelihood_extrinsic_batchmode` - Main PE engine - `create_event_parameter_pipeline_BasicIteration` - Full pipeline diff --git a/MonteCarloMarginalizeCode/Code/RIFT/LISA/lalsimutils_compat.py b/MonteCarloMarginalizeCode/Code/RIFT/LISA/lalsimutils_compat.py index 71441effb..d2d3ef56f 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/LISA/lalsimutils_compat.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/LISA/lalsimutils_compat.py @@ -90,10 +90,19 @@ def hlmoff_for_LISA( ) hlmsdict = lalsimutils.SphHarmFrequencySeries_to_dict(hlms_struct, Lmax) hlmsdict = _filter_modes(hlmsdict, modes) - return { - mode: lal.ResizeCOMPLEX16FrequencySeries(hlm, 0, TDlen) - for mode, hlm in hlmsdict.items() - } + # SimInspiralChooseFDModes returns an ascending two-sided grid + # [-fNyq, ..., 0, ..., +fNyq] of odd length TDlen+1. The resize below + # truncates the +fNyq bin but keeps its -fNyq partner (index 0); zero it + # so the truncation commutes with the conjugate-pair (f -> -f) reflection + # h_{l,-m}(f) = (-1)^l conj(h_{lm}(-f)) for models with support at Nyquist. + out = {} + for mode, hlm in hlmsdict.items(): + truncated = hlm.data.length > TDlen + hlm = lal.ResizeCOMPLEX16FrequencySeries(hlm, 0, TDlen) + if truncated: + hlm.data.data[0] = 0 + out[mode] = hlm + return out if P.approx in {lalsimutils.lalNRHybSur3dq8, lalsimutils.lalIMRPhenomD}: hlms_struct = lalsimutils.hlmoff(P, Lmax=Lmax) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/calmarg/calibration.py b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/calibration.py new file mode 100644 index 000000000..46b74453b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/calmarg/calibration.py @@ -0,0 +1,26 @@ +"""Helpers for matching bilby-pipe calibration conventions.""" + + +def correction_type_for_ifo(setting, ifo_name, parse_dict=None): + """Resolve bilby-pipe's calibration correction type for one detector.""" + if setting is None or setting == "None": + return "template" if ifo_name == "V1" else "data" + + if isinstance(setting, str): + if setting in ("data", "template"): + return setting + if parse_dict is None: + raise ValueError("parse_dict is required for detector-specific settings") + setting = parse_dict(setting) + + try: + correction_type = setting[ifo_name] + except (KeyError, TypeError) as exc: + raise ValueError( + f"No calibration correction type specified for {ifo_name}" + ) from exc + if correction_type not in ("data", "template"): + raise ValueError( + f"Invalid calibration correction type for {ifo_name}: {correction_type}" + ) + return correction_type diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_flexible_gmm.md b/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_flexible_gmm.md new file mode 100644 index 000000000..ddbc77e45 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_flexible_gmm.md @@ -0,0 +1,237 @@ +# Flexible (data-driven) GMM component allocation for the extrinsic integrator + +Status: prototype, benchmarked on S250114ax (SNR~82, real H1L1). See the +measured results and the honest tradeoffs at the bottom -- the headline is that +**this fixed the GMM's hard-coding and two warm-start bugs, but pure GMM +importance sampling is still not the right tool for this particular event**; the +flexible GMM belongs in the portfolio (with AV) or on milder problems. + +## The problem being fixed + +`bin/integrate_likelihood_extrinsic_batchmode` built the GMM (`mcsamplerEnsemble`) +proposal with a **hard-coded** per-group component layout: + +``` +gmm_dict = {(ra,dec):None, (distance,inclination):None, (psi,phi_orb):wide-frozen} +comp_dict = {(ra,dec):4, (distance,inclination):2, (psi,phi_orb):4} +``` + +That pairing (large sky ring = 4 components; a single distance-inclination lobe = 2; +a wide frozen phase-polarization component) targets **quadrupole-dominated, +poorly-localized** binaries. It does not adapt to the actual posterior, and a +product of per-group GMMs cannot represent cross-group correlations. `--internal- +gmm-correlate-all` switches to a single full-dim group but still with a **fixed** +component count. Choosing that count by hand is exactly the "horrible hacky +hard-coding" this work removes. + +## What the extrinsic posterior actually looks like (S250114ax) + +Not a 10^-11 needle -- a **broad, correlated/degenerate 6-D blob** (see +`BREADCRUMB_av_neff_reproduction.md`). Sky is tight; **distance and inclination +are broad and coupled** (the classic distance-inclination degeneracy -- a curved +arc); phase and polarization are similarly degenerate. The prior is huge (all- +sky/all-distance/all-inclination), so the peak is a tiny *fraction* of the prior +even though it is not narrow in absolute terms. + +## Design + +Three orthogonal knobs, prototyped in order of leverage: + +### 1. Scalable component matching (enabling; committed separately) +`gmm._match_components` enumerated all **k! permutations** to align old->new +mixture components in `update()`. Fine for k<=6, impossible for k>=10 (12!~5e8, +16!~2e13). The objective is additive over matched pairs, so it is a linear +assignment problem: `scipy.optimize.linear_sum_assignment` gives the **same +optimum in O(k^3)**. Verified identical to the permutation optimum for k=2..6; +k=16 now matches in ~6 ms. Without this, "wrap the arc with many small Gaussians" +is not even runnable. + +### 2. Warm-start survival (bug fix; high leverage) +`bootstrap_from_samples` fits proposal models and stores them on +`self.integrator`, but `MCSampler.integrate()` **rebuilt a fresh integrator** from +the passed `gmm_dict` (values `None`) and never looked at `self.integrator` -- so +the warm fit was **silently discarded** and a "warm" run started cold (measured: +first chunk n_eff=1.0 despite a `[GMM warm-start] fitted ...` log line). +`integrate()` now transfers any fitted model whose dim-group key matches into the +new integrator (key mismatch -> cold, never biases). + +### 3. Data-driven component count (the flexible allocation) +`fit_gmm_adaptive(samples, bounds, log_weights, k_max, ...)`: + * fit k over a ladder `[1,2,3,4,6,8,...] <= k_max`, + * score each by a **weighted BIC** `-2*wLL + p*ln(N_eff)` (p = free params of + a k-component d-dim mixture, N_eff = Kish effective sample size of the fit + weights), + * keep the best k, then **prune** components whose weight < floor. + +BIC allocates more components only where the importance-weighted cloud is +genuinely non-Gaussian, and stays at k=1 for a single blob; the ln(N_eff) penalty +makes it self-limiting, so it avoids the instability of a fixed over-allocated k. + +**Init-only, then stable merge.** A group with `gmm_adaptive[group]=k_max` picks +its k by BIC at *initialization*, then hands off to the existing, proven-stable +merge adaptation (`model.update()`). A per-chunk BIC refit-fresh was tried and +**rejected**: it makes the proposal wander (n_eff peaks then collapses) because +each fit sees a different elite cloud; the incremental merge smooths that out. + +**Defensive tail coverage (optional, `add_defensive_component`).** A broad box- +covering component with weight `defensive_frac` bounds the importance weights so a +tight fit to the elite cloud cannot blow up n_eff on the broad/degenerate +directions (the AV sampler gets the same guarantee from its cover-fraction floor). +Kept as an option; see the caveat below. + +### Safety: opt-in, and a floor at the stress-tested layout +The flexible allocation is a **refinement layer, never a replacement**: + * It is **opt-in** (`--internal-gmm-adaptive-components`, default OFF). With + the flag off the driver builds the exact hard-coded `gmm_dict`/`comp_dict`/ + `gmm_adapt` as before -- byte-identical behavior for the primary ILE use case. + * When on, BIC chooses k in **[k_min, k_max]** with `k_min` = the group's + stress-tested hard-coded count, and pruning never drops below `k_min`. So + adaptive can only ADD components where the data earns them; it can never + allocate fewer than the validated layout. This protects broad multi-modal + posteriors (e.g. a multi-modal sky keeps its default components even if the + *initial* elite cloud -- fit before the proposal has explored every mode -- + looks single-peaked; the spare capacity lets the merge adaptation grow into + the other modes as they appear). + +### Portfolio +`--internal-gmm-adaptive-components` also works with `--sampler-method portfolio`. +Invoke the members with the flag REPEATED (it is `append`, not comma-split): +`--sampler-portfolio AV --sampler-portfolio GMM` (NOT `"AV,GMM"`, which becomes +one bogus member). When `GMM` is a member the driver sets `sampler_method='GMM'`, +so the full GMM config section runs for the member: it gets the stress-tested +pairing `{sky:4, dist-incl:2, phase:frozen}` AND, with the flag on, the per-group +adaptive allocation on the ADAPTING groups only -- e.g. `gmm_adaptive={(4,5):8, +(3,2):8}` (sky, dist-incl), floored at 4/2, phase excluded. The GMM member honors +this in `update_sampling_prior` (the path the portfolio drives). Default OFF -> +the portfolio's stress-tested GMM member config is unchanged. This is the intended +way to use the flexible GMM: AV carries convergence, the GMM member contributes a +hands-free correlated proposal instead of a hand-tuned component layout. + +**Status (2026-07-21).** After the AV/portfolio fixes on the base branch +(`d44fe486` member setup, `edf775c7` AV empty-selection, `f2d51de0` freeze grace + +revive + NaN-weight guard) the AV+GMM portfolio runs to completion on S250114ax and +the default-vs-adaptive comparison is measurable. Warm, GPU, 4 M cap, +`--sampler-portfolio AV --sampler-portfolio GMM`: + +| warm portfolio, GMM member | peak n_eff (grace=25, default) | peak n_eff (freeze-exempt) | +|------------------------------|-------------------------------:|---------------------------:| +| default (hard-coded pairing) | 3.01 | 1.66 | +| **adaptive (BIC, cap 8)** | **18.2** | 4.16 | + +**The flexible allocation clearly helps the portfolio: the adaptive GMM member +gives ~6x the peak n_eff of the hard-coded member (18.2 vs 3.01)** at the default +grace(25). This is the headline portfolio result -- a hands-free GMM member that +outperforms the hand-tuned layout inside the mix. + +**AV is still not contributing, and freeze-exemption does NOT fix it (a separate +balance-heuristic issue, PR #26 territory).** AV's mixture weight stays pinned at +~0.0099 (= 1/101, a weight floor) for the ENTIRE run whether it is frozen or not -- +so the earlier "freeze-out" was a symptom, not the cause: the balance heuristic +simply never assigns AV meaningful weight when a warm GMM member is present, so the +portfolio always rides the GMM member. Making AV freeze-exempt (grace spanning the +whole run) is actually *worse* (adaptive 4.16 vs 18.2): AV then keeps drawing its +~1% share unproductively (with `nan`s) and dilutes the estimate, instead of being +frozen out of the way. So the open question for the AV/portfolio side is why the +balance heuristic floors AV's weight at 1/101 here -- not the freezing per se. The +GMM-member allocation result above stands regardless. + +### Driver flags +``` +--internal-gmm-adaptive-components # enable BIC allocation (opt-in; OFF by default) +--internal-gmm-max-components N (def 8) # per-group cap (floor = the hard-coded count) +--internal-gmm-defensive-frac F (def 0) # opt-in defensive tail component +--internal-gmm-inflate X (def 1) # covariance (std) inflation +``` + +## Measured results + +### Synthetic (data-free, moderate SNR -- where importance sampling is viable) +6-D target: a curved **banana ridge** in 2 dims (distance-inclination analogue), +a strongly-correlated Gaussian pair (phase-pol analogue), a tight blob (sky). +`test/integrators/synth`-style harness, n_eff vs N (cumulative samples): + +| proposal | n_eff>=100 at N | final n_eff | lnI | +|------------------------------|-----------------|-------------|--------| +| correlate-all, fixed k=1 | 76 k | ~50 | 3.06 | +| correlate-all, fixed k=2 | **28 k** | **312** | 3.05 | +| correlate-all, fixed k=4 | 752 k | ~96 | 3.06 | +| **flexible (BIC, k<=8)** | 220 k | 135 | 3.03 | + +Flexible is **robust and unbiased**: it beats k=1, avoids the k=4 over-allocation +collapse, and lands the same integral -- without any hand-tuning. It does not beat +the *oracle-best* fixed k=2, which is the expected price of a hands-free allocator. +Note fixed k=4 being far worse than k=2 is exactly the "a wrong hard-coded count +hurts" failure the flexible allocation exists to avoid. + +### S250114ax (real, SNR~82) -- n_eff vs N, `--n-max 4e6 --n-eff 100` +Same worker, byte-identical data, only the proposal/config varies: + +| sampler / config | warm | peak n_eff (<=4 M) | +|----------------------------------------------------|------|--------------------| +| **AV (VARAHA), warm (reference)** | yes | **~89 (->100 @~3.4 M)** | +| hard-coded pairing, cold | no | 1.29 | +| correlate-all k=2, warm (warm-start now survives) | yes | 5.7 | +| correlate-all k=8 / k=16, warm + --adapt-adapt | yes | ~1.0 | +| **flexible (BIC k<=8), warm + --adapt-adapt** | yes | 1.03 | +| flexible (BIC k<=8), cold + --adapt-adapt | no | 1.00 | +| (independent) prior pure-GMM logs gmm*/cold* | -- | 1.3 - 3.2 | + +**Pure GMM importance sampling stalls at n_eff ~ 1-7 on this event, regardless of +component allocation (fixed k=2..16, BIC-adaptive), warm start, defensive coverage +(0.05), or covariance inflation (2x-5x).** This is corroborated by five pre- +existing pure-GMM logs in the pipeline (peak 1.3-3.2). The runs that reached ~100 +in that directory are **AV / portfolio** runs, not pure GMM. + +## Why GMM stalls here (and AV does not) + +At SNR~82 the likelihood spans exp(~1210). The honest per-chunk effective sample +size `ESS(lnL + ln p - ln q)` is **exactly 1** every chunk: within any chunk of +proposal draws, the single sample nearest the sharp peak carries ~100+ nats more +log-weight than the rest, so it dominates. For n_eff to exceed 1, the proposal +would have to match `exp(lnL)*prior` to within O(1) *across* the peak -- i.e. be +almost the posterior already. A moving Gaussian-mixture importance proposal does +not get there from a broad start: + * the beta-tempered refit drives the exponent to ~0.005 (nearly flat) to keep + its own ESS up, so it barely uses the likelihood and never concentrates; + * the rank-elite (cross-entropy) refit fits the top-k by lnL, but those elites + are spread along the curved degeneracy ridge, so the fit is a broad Gaussian + over the ridge -- more components do not help because the ESS-1 domination is + set by the *narrow* constrained directions, not the ridge. +AV (VARAHA) is not importance sampling: it contracts an axis-aligned live volume +against a likelihood threshold with a coverage floor, which is the right structure +for a sharply-peaked target (at the cost of not wrapping the diagonal arc, hence +its own ~3e-5 efficiency ceiling). + +## Recommendation / tradeoffs + +* **Keep** the three fixes -- they are correct and independently valuable: + O(k^3) matching, warm-start survival, and BIC allocation remove the hard-coding + and make a warm GMM actually warm. The synthetic shows the allocation is robust + and unbiased. +* **Do not** expect pure GMM to beat AV on high-SNR, strongly-degenerate events. + Use the flexible GMM **inside the portfolio** (`--sampler-method portfolio`), + where the balance-heuristic mixture density keeps a weak member from biasing -- + with a **hands-free** GMM member instead of a hand-tuned component layout. This + is now wired and MEASURED (see Portfolio Status): the adaptive GMM member gives + ~6x the portfolio peak n_eff of the hard-coded member on S250114ax. The one + remaining item is on the AV/portfolio side (PR #26): the balance heuristic floors + the AV member's weight at ~1/101 on this event, so AV never contributes (and + freeze-exemption does not fix it) -- once AV pulls its weight, the portfolio + should combine AV's convergence with the flexible GMM's correlated proposal. +* **Future levers** for making GMM itself competitive here would be *coordinate* + changes that de-curve the arc (a distance-inclination reparametrization, + rotate-phase for phase<->pol) so a low-k axis-aligned mixture fits -- i.e. + attack the correlation in coordinates, then let BIC pick k. See the breadcrumb's + "real levers" section. + +## Code map +* `RIFT/integrators/gaussian_mixture_model.py`: `_match_components` (Hungarian), + `fit_gmm_adaptive`, `gmm.prune_components`, `gmm.num_free_params`, + `add_defensive_component`, `_mixture_log_density_normalized`. +* `RIFT/integrators/MonteCarloEnsemble.py`: `integrator.gmm_adaptive/ + gmm_defensive_frac/gmm_inflate`; BIC-at-init in `_train`. +* `RIFT/integrators/mcsamplerEnsemble.py`: warm-start transfer in `integrate()`; + `gmm_adaptive` threading in `integrate()`/`setup()`. +* `bin/integrate_likelihood_extrinsic_batchmode`: the four flags + `gmm_adaptive` + dict construction in the GMM section. +* `test/integrators/test_gmm_adaptive.py`: unit tests (5/5). diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_portfolio_freeze_policy.md b/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_portfolio_freeze_policy.md new file mode 100644 index 000000000..b860a9c14 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/DESIGN_portfolio_freeze_policy.md @@ -0,0 +1,978 @@ +# Portfolio freeze policy — tuning + benchmarks + +## The problem (starved VARAHA workhorse) + +The portfolio integrator (`mcsamplerPortfolio.py`) mixes several member samplers and reweights +them each chunk by their per-chunk effective sample size `n_ess` (`portfolio_default_weights`). +A member whose weight falls below `portfolio_freeze_wt` (0.05) stops updating its proposal. + +A **VARAHA / AV** member is special: it contracts its live volume **only on the chunk it is +updated** (`update_sampling_prior_selfish`). On its first chunks — before it has contracted — its +`n_ess` is ~1, so the heuristic gives it weight ~0.01 < 0.05 and it is **frozen from chunk 1** and +never contracts. The portfolio then rides a stalling GMM member. On the high-SNR S250114ax event +this produced portfolio cold-peak `n_eff ≈ 1.9` versus standalone AV `≈ 89–100`. + +## Why grace/revive alone cannot fix it + +The earlier fix (commit f2d51de0) added GRACE (never freeze for the first `grace_iters` chunks) +and REVIVE (update a frozen member every `revive_period` chunks). These are **not sufficient** for +a VARAHA member, because of a weight **feedback loop**: AV only earns allocation weight once its +`n_ess` climbs, and its `n_ess` only climbs once it contracts, and it only contracts when updated. +Updating it merely 1/`revive_period` of the time contracts it far too slowly for its `n_ess` to +ever win weight, so it stays starved. The revive sweep below confirms this: revive-every-{8,4,2} +all leave AV at `n_eff` ≈ 1–3, essentially no better than a hard freeze. + +## The fix: VARAHA members are freeze-EXEMPT by default + +Because the portfolio combines members with the **balance-heuristic mixture density** `q_mix` +(`integrate_log`), the estimate is unbiased for **any** member weights — a continuously-updated +VARAHA member can only ever cost a few extra selfish-draw evaluations, never bias the integral. So +the new default (`portfolio_varaha_never_freeze=True`) makes VARAHA/AV members update **every +chunk** past their activation breakpoint, exactly like a standalone AV. Set it False to fall back +to the grace/revive schedule (e.g. to save eval cycles on a VARAHA member known to be a bad fit). + +Also added: a **plateau-aware revive** (`_climbing`) that keeps updating *any* low-weight member +while its own per-chunk `n_ess` is still rising, and a per-member `n_ess` history for diagnostics. + +### Knobs (sampler defaults, all overridable) +| knob | default | meaning | +|------|---------|---------| +| `portfolio_varaha_never_freeze` | **True** | VARAHA members update every chunk past breakpoint | +| `portfolio_grace_iters` | 25 | never freeze ANY member during the first N chunks | +| `portfolio_revive_period` | 8 | update even a frozen member every N chunks (0 disables) | +| `portfolio_freeze_wt` | 0.05 | weight below which a (non-exempt) member stops updating | + +### CLI flags (driver `integrate_likelihood_extrinsic_batchmode`, thread into `sampler.setup`) +`--portfolio-grace-iters N`, `--portfolio-revive-period N`, `--portfolio-freeze-wt X`, +`--portfolio-varaha-never-freeze` (explicit; already the default), `--portfolio-varaha-can-freeze` +(disable the exemption). Also fixed: `--sampler-portfolio AV,GMM` is now comma-split as documented +(it was silently becoming a single bogus member). + +## Benchmark 1 — S250114ax (deliberately HARD: ρ≈82, broad/degenerate extrinsic posterior) + +iteration-0 worker, IMRPhenomD, real H1L1, GPU A100. `N` = sample count at which `n_eff` first +crosses each threshold; `--n-max 4e6 --n-eff 100`. Warm = PE-oracle seed (cover 0.05/inflate 1.3). + +| run | Neff≥5 | ≥10 | ≥20 | ≥50 | ≥100 | final (N, Neff) | +|-----|-------:|----:|----:|----:|-----:|:---------------:| +| **av_warm** (standalone, reference) | 0.70M | 1.37M | 1.69M | 2.27M | **3.64M** | 3.64M, 100 | +| **pf_nf_warm** (portfolio, never-freeze = NEW default) | **0.52M** | **1.10M** | 1.86M | 3.82M | — | 4.0M, **53** | +| pf_cf_warm (portfolio, can-freeze, grace25/revive8) | — | — | — | — | — | 4.0M, 3.4 | +| pf_cf_warm, revive=4 | — | — | — | — | — | 4.0M, 1.3 | +| pf_cf_warm, revive=2 | — | — | — | — | — | 4.0M, 1.6 | +| av_cold (standalone) | 1.75M | — | — | — | — | 4.0M, 3.7 | +| pf_nf_cold (portfolio, never-freeze) | — | — | — | — | — | 0.21M, 1.1 (cold-degenerate, clean stop) | +| GMM (standalone, cold AND warm) | — | — | — | — | — | **NaN on chunk 1 → bails** | + +**Reading it.** +- **Never-freeze rescues the workhorse: `n_eff` 3.4 → 53** (~15×) vs the frozen policy, and it + *tracks or beats* standalone AV through the useful range — it reaches `n_eff`=5 and 10 EARLIER + than standalone AV (0.52M vs 0.70M; 1.10M vs 1.37M), because early on the AV+GMM mixture covers + better than warm-AV alone. +- **grace/revive tuning does NOT rescue AV** (all can-freeze variants stay at `n_eff` 1–3), + confirming the feedback-loop argument above. +- On this *atypical* event the portfolio's deep tail (`n_eff` 50→100) is ~1.5× slower than pure AV + (53 @ 4M vs 100 @ 3.64M): AV alone is optimal here, and the portfolio pays a modest cost for + carrying a weak GMM member. This is the honest limit — on a genuinely AV-optimal, GMM-hostile + event the portfolio cannot beat standalone AV, but with never-freeze it is now in the same + regime instead of starved. +- **Standalone GMM is unusable on this event** (NaN on chunk 1, cold and warm) — it only works + *inside* the portfolio, where AV's coverage via `q_mix` + the NaN-weight guard stabilize it. + +## Benchmark 2 — multi-event robustness (typical events) + +Goal (per reviewer): confirm the AV+GMM portfolio (never-freeze) **replicates the standalone-AV +integral** — same ln Z within MC error — and converges comparably across a spread of *typical* +real events, not just the hard S250114ax target. Method: each event's real iteration-0 ILE worker +(real strain/PSD/intrinsic grid), run in the event's own production container (SEOBNRv5PHM + +gwsignal + cuda118 cupy) with this branch's integrator on `PYTHONPATH`; only `--sampler-method` +differs between the two configs. Single intrinsic point, `--n-eff 30 --n-max 8e5`. + +| event | AV lnZ (n_eff @ N) | portfolio lnZ (n_eff @ N) | ΔlnZ | replicated? | +|-------|:------------------:|:-------------------------:|-----:|:-----------:| +| S231026ab | 17.64 (33 @ 81k) | 17.62 (19 @ 800k) | 0.02 | **yes** | +| S240426s | 29.75 (31 @ 60k) | 29.68 (31 @ 310k) | 0.08 | **yes** | +| S240513ei | 85.73 (9.5 @ 809k)| 85.33 (1.6 @ 800k)| 0.40 | yes, within MC err (both under-converged) | +| S240703ad | 42.67 (11 @ 803k) | 41.70 (2.8 @ 800k)| 0.97 | ~ (pf under-converged, n_eff 2.8) | +| S240601aj | — missing BayesWave glitch-subtracted frame (event-data issue, not integrator) — dropped | + +**Verdict.** The portfolio (never-freeze) **replicates the standalone-AV integral** — every ΔlnZ is +within the (often large) MC error of the lower-n_eff run, and there is no bias. AV is **never +frozen** on any event (0 freeze notices) and, on typical events, becomes the in-portfolio workhorse +(its balance weight climbs from 0.5 to ~0.65–0.72 within a few chunks). So the freeze mechanism is +**NOT fundamentally at odds** with VARAHA's need for continuous contraction — never-freeze gives it +exactly that, unbiased. + +The remaining limit is **efficiency, and it is a DIFFERENT lever than freezing** (addressed next): +with the plain n_ess reweighting the portfolio spent a fixed ~half its budget on the GMM member, so +on AV-favorable events it reached a given n_eff in more evals than standalone AV (S240426s: same +n_eff at 310k vs 60k), and on the hardest events it stayed under-converged. + +## Adaptive-probe draw allocation (the efficiency lever) + +The reason the plain reweighting couldn't concentrate is a **draw catch-22**, structurally identical +to the freeze bug but on draws instead of updates: a member's per-chunk n_ess is *suppressed while +it has few draws* (a VARAHA member contracts slower with fewer samples; any Kish n_ess is noisy on a +small slice), so the member that *should* win is stuck under-observed and never earns more draws. + +Mechanism (`portfolio_adaptive_alloc`, **OPT-IN, default OFF** — see the regression below): keep a +per-member **quality** estimate updated ONLY from chunks where the member had a *fair* allocation; +allocate draws by `quality^exponent` above a small floor; and **round-robin probe** one member per +`probe_period` chunks at a raised share so a suppressed member gets a fair look. `q_mix` keeps every +allocation unbiased. Knobs (`setup` + CLI `--portfolio-adaptive-alloc` to enable): +`portfolio_alloc_exponent` (2.0), `portfolio_alloc_floor` (0.05), `portfolio_quality_decay` (0.5), +`portfolio_probe_period` (4), `portfolio_probe_frac` (0.6). + +### Choosing the quality signal (`portfolio_quality_signal`) + +Three candidates were implemented and measured. Only the third is defensible, and even it cannot +rescue S250114ax — for a reason that turns out **not** to be about allocation at all. + +1. **`ness`** — per-member Kish n_ess. **Fails.** Kish is *scale-invariant* (`(Σw)²/Σw²` is + unchanged if all `w` are scaled), so it cannot see whether a member's samples carry any integral + mass: a self-consistent member sitting off-peak scores as well as one covering the peak. A warm + GMM is instantly self-consistent (n_ess ~120) while a warm AV's per-chunk n_ess is genuinely ~1 + during its slow *cumulative* contraction (value emerges over ~70 chunks). On S250114ax this drove + the true AV workhorse to the floor: **n_eff 8 vs 53** for the legacy allocation — a regression. +2. **mean weight** (per-sample contribution). **Also fails, backwards.** A *well-matched* proposal + correctly has small uniform weights, while a broad proposal's rare huge-weight outlier sets the + maximum. Measured on S250114ax: AV **1e-40** vs GMM **2e-4** — it penalizes the good member. +3. **`global` (default when adaptive is on)** — marginal gain in **pooled** n_eff per sample, + `g_m = 2·mean_w_m/S − mean_w2_m/Q` (`S=Σw`, `Q=Σw²` over all samples). This is the right + objective: it credits weight *mass* and debits weight *variance*. It works on the synthetic + (below), but on S250114ax it still ranks GMM first — and the numbers say exactly why. + +**The S250114ax diagnosis (allocation is not the bottleneck).** With the `global` signal the +measured values are AV `~1e-73` and GMM `1.053e-4`. That GMM value is precisely `1/9500 = 1/n_GMM`, +which is the analytic signature of **one sample owning the entire estimator**: for a member holding +the single dominant outlier, `g = 2/n − 1/n = 1/n`. So the chunk's maximum weight is a catastrophic +GMM outlier ~**10⁷³×** larger than any AV weight — a draw landing where `q_mix ≈ 0` but the target is +nonzero. No allocation signal computed from the current weights can rank AV above that, because the +pooled estimator genuinely *is* dominated by that one sample. + +The consequence: on this event the ceiling is set by the GMM member's **unbounded importance +weights**, not by how draws are split. Even floored at 5% the GMM member still injects outliers, which +is why the legacy allocation reached only 53 (not AV's 100). **The next lever is therefore weight +bounding / member exclusion, not allocation**: a defensive covering component to bound `w`, clipping +or winsorizing member weights, or dropping a member whose weight distribution is unbounded. (PR #27's +`--internal-gmm-defensive-frac` is the related Hesterberg-defensive knob; its help notes it did not +help n_eff on this SNR~82 benchmark, consistent with this being a hard pathology.) + +Because of this, adaptive allocation remains **opt-in**; the default keeps never-freeze + legacy +reweighting. + +## Benchmark 3 — adaptive allocation on synthetic correlated targets + +`test/integrators/test_portfolio_adaptive_alloc.py` (fast, GPU, no ILE). Standalone AV vs standalone +GMM vs AV+GMM portfolio, fixed budget (`nmax 4e5`, ndim 5), reporting the integrator's own eff_samp +and bias vs the analytic ln Z. The GMM member is broad-seeded (a wide peak-covering proposal) so the +test exercises the *allocation* given a functional GMM rather than gambling on cold GMM finding a +thin ridge; AV starts cold. Representative (seed 1234): + +| target | AV n_eff (bias) | GMM n_eff (bias) | **portfolio** n_eff (bias) | portfolio wts (AV,GMM) | +|--------|:---------------:|:----------------:|:--------------------------:|:----------------------:| +| uncorrelated (axis-aligned) | 107 (−0.54) | 398 (−0.01) | **386** (−0.04) | 0.15, 0.85 | +| **correlated (compound-symmetric)** | 23 (−0.67) | 400 (−0.00) | **387** (−0.03) | 0.13, 0.87 | + +- **On the correlated target GMM's full covariance crushes AV's axis-aligned bins (400 vs 23), and + adaptive allocation concentrates on GMM (weight 0.87) so the portfolio BEATS standalone AV** — the + whole point of a portfolio on a correlated problem, and the case the reviewer flagged as the only + regime where beating AV is expected. +- A cold VARAHA/AV under-covers the Gaussian tails and is **biased low** (−0.6 to −0.9); the + portfolio stays **unbiased** because the covering GMM enters `q_mix` — a second reason to prefer + the portfolio. (This is the opposite regime from a warm, cover-frac'd AV on a real ILE likelihood, + where AV is the unbiased workhorse; the point demonstrated is that adaptive allocation follows + whichever member is actually winning — GMM here, AV on a real AV-favorable event.) + +The synthetic result shows the *mechanism* is sound **when the quality signal is right** (there GMM +is genuinely better and adaptive follows it). The S250114ax regression shows the *signal* is wrong on +real high-SNR AV-favorable events. Hence adaptive is shipped **opt-in**, and the default portfolio is +never-freeze + legacy allocation. + +**Overall verdict.** Never-freeze (default) makes the portfolio unbiased and never-starved — it +**replicates standalone AV** and, on typical events, lets AV be the workhorse. Adaptive-probe +allocation (opt-in, with the `global` marginal-pooled-n_eff signal) makes a portfolio **beat AV on a +strongly-correlated target** (synthetic: ~375 vs ~61 n_eff) — the only regime where beating AV is +expected. + +It is still opt-in because of what the global signal *revealed* rather than any deficiency in it: +on S250114ax the pooled estimator is dominated by a single GMM outlier ~10⁷³× the next weight, so +**allocation is not the bottleneck there — unbounded member weights are**. Turning adaptive on +everywhere requires first bounding those weights (defensive component / weight clipping / dropping a +member with an unbounded weight distribution). That is the concrete next lever, and it is a +*member-quality* fix, not an allocation-policy one. + +**Robustness bug fixed along the way (important):** portfolio plugin discovery hard-loaded every +registered plugin at import, and the `NF` plugin does `import torch`, absent in the production GPU +container — so `import mcsamplerPortfolio` raised there, the driver silently set +`mcsampler_Portfolio_ok=False`, and every portfolio run died with a `NameError`. The portfolio +integrator was effectively **unusable in the production container**. Plugin loading is now wrapped +in try/except so a plugin with missing optional deps is skipped, not fatal. + +## Weight clipping (truncated IS) — ADAPTATION-STREAM ONLY; **do not promote to AV yet** + +`portfolio_weight_clip` / `--portfolio-weight-clip C` (OPT-IN, default off) caps weights at +`tau = C*sqrt(n)*mean(w)` (Ionides 2008 truncated IS). Clipping is a **biased** operation, so the +key design decision is *what it is allowed to touch*. + +**⚠ SUPERSEDED — the outliers were an ARTIFACT after all (see PR #33).** This section originally +concluded the 10⁷³× weights were genuine heavy tails, because an instrumented run counted **zero** +`q_mix = max(acc, 1e-300)` underflows. That test was correct but incomplete: it rules out a density +*underflow*, not a density *lie*. PR #33 subsequently found exactly such a lie — AV's +`draw_simplified` head-sliced a **bin-ordered** cloud (returning only ~50–60% of the live-volume +bins) while `sampling_density` claimed uniform coverage of **all** occupied bins, so `q_mix` was +simply wrong, and a member drawing in the region AV never populates gets an arbitrarily inflated +weight. PR #33 also found that default-wired portfolio GMM members **never trained** (`n_comp=None` +silently no-op'd), so the runs below carried an untrained corpse member — which both produced junk +draws and partially masked the density bug. + +**Consequences for everything below:** the S250114ax numbers in this document were taken with a dead +GMM member and a lying AV draw density, so they characterize *those bugs*, not the integrator's real +behavior. The clipping study remains valid as a *methodological* result (which quantities may be +clipped, and why — those arguments are analytic, not event-specific), but every S250114ax +efficiency/ln Z figure needs re-measuring on top of PR #33. The `q_mix UNDERFLOW` counter stays as a +permanent guard, and the lesson generalizes: **when a weight looks impossible, test the density for a +LIE (does the member actually draw where it claims density?), not just for underflow.** + +**First attempt — clip the estimator: a disguised disaster (kept as the cautionary result).** + +| run | Neff≥5 | Neff=100 | final n_eff | **ln Z** | +|-----|-------:|---------:|------------:|---------:| +| standalone AV (reference) | 0.695M | 3.638M | 100.2 | **1191.79** | +| portfolio, no clip | 0.520M | — | 52.6 @4M | 1183.12 | +| portfolio, **estimator**-clip C=1 | **0.030M** | **1.870M** | 100.1 | **1180.25** | + +Clipping the estimator reached n_eff=100 in 1.87M evals — 2× faster than standalone AV — and reported +a perfectly converged run, while biasing ln Z **11.5 nats low** (independent reparam/aniso runs give +~1191.9). **n_eff stops being a validity check the moment the estimator is clipped.** Do not do this. + +**The unbiased redesign — clip the PROPOSAL-FIT INPUT only.** Clipping is both biased *and* +n_ess-distorting, so its scope must be narrow. Final split: +- `log_integrand` → the ESTIMATE (`init_log`/`update_log` → ln Z; `maxval` → n_eff) — UNCLIPPED. +- per-member **n_ess report** and the **allocation signal** — UNCLIPPED (true weights). +- ONLY `member.update_sampling_prior` (the GMM covariance fit) gets the clipped copy + `log_weights_adapt`, so one enormous weight can't make that fit degenerate. + +Two failure modes ruled this scoping. (a) Clipping the **estimator** biases ln Z (the trap above). +(b) Clipping the **n_ess report / allocation** is *also* wrong, and subtly: clipping flattens +weights, which INFLATES a member's Kish n_ess, so the allocation perversely rewards the very member +whose weights had to be clipped. Measured on S250114ax, a first attempt that clipped the report +starved the AV workhorse to the 1% floor and stuck n_eff at ~1 (worse than no-clip's 53). Restricting +clipping to the proposal fit removed that: a short run climbs n_eff normally again. Proposal fitting +only *shapes* the proposal (like warm-starts / oracles / `q_mix`), so it cannot bias ln Z — and this +is strictly better than *dropping* a clipped chunk, which (being conditional on "a big weight +appeared") is data-dependent selection that would bias ln Z low. + +**Synthetic ground truth** (`test/integrators/bench_weight_clip.py`, analytic ln Z, 2 seeds): with the +adaptation-only design the estimator bias is **unchanged at every C** (the falsifiable proof the +estimate is untouched), and where weights are well-behaved clipping is a complete no-op: + +| target | C | n_eff | bias | +|--------|---|------:|-----:| +| uncorrelated | 0 / 1 / 5 | 386 / 385 / 385 | −0.035 / −0.039 / −0.039 | +| correlated | 0 / 1 / 5 | 389 / 387 / 388 | −0.029 / −0.032 / −0.031 | + +**Real S250114ax — unbiased, and (with the correct scope) not harmful.** Proposal-fit clip C=1 +keeps ln Z at the unclipped ~1183 (NOT the biased estimator-clip 1180.3 — the estimator is provably +untouched) and n_eff climbs normally again (a short run tracks the no-clip curve). It does not *speed +up* this event either: the GMM member's proposal is fundamentally heavy-tailed here, and clipping +only stops its covariance fit from going degenerate — it cannot make a bad proposal good. The right +move on this event remains to not carry the GMM member. (An earlier, wrongly-scoped attempt that also +clipped the n_ess report collapsed n_eff to ~1 by starving AV — see the redesign note above; that was +the bug, not a property of clipping.) + +**Side finding (refines the Benchmark-2 claim).** Even *unclipped* the AV+GMM portfolio reads +ln Z = 1183.1 here, 8.7 nats below AV. Heavy-tailed IS is unbiased in expectation but realizes LOW in +almost every run, so in production it behaves like a bias. "Portfolio replicates AV's ln Z" holds on +the four *typical* events (Benchmark 2); it does **not** on S250114ax. + +**Can clipping rescue adaptive allocation on S250114ax?** No. Idea: the global allocation signal was +fooled because one 10⁷³ outlier owned the pooled estimator; since the signal now reads the *clipped* +adaptation weights, a gentle clip (C=20) that removes only that single outlier does flip the +first-chunk signal to correctly favor AV (contrib AV 2e-4 vs GMM 9e-21). But it does not *persist*: +the per-chunk marginal-n_eff signal is too noisy on this event — it oscillates back to GMM within a +few chunks, and both C=1 and C=20 stall the estimator at n_eff ~1 (worse than legacy's 53). So on an +AV-favorable event adaptive allocation fails clipped or not; **legacy allocation stays the right +default there, and adaptive remains a correlated-problem tool.** + +**Typical-event safety validation** (proposal-fit clip C=1, 4 real events in-container, vs no-clip). +Clipping is unbiased everywhere (estimator untouched) and a no-op-to-mild-help on n_eff — and it +recovered the two hard events the broad-scope bug had degraded: + +| event | no-clip n_eff (lnZ) | proposal-fit clip n_eff (lnZ) | +|-------|:-------------------:|:-----------------------------:| +| S231026ab | 19 (17.62) | 25.5 (17.47) | +| S240426s | 31 (29.68) | 30 (29.74) | +| S240513ei | 1.6 (85.33) | 1.3 (83.20) | +| S240703ad | 2.8 (41.70) | **6.9** (42.43) | + +On S240703ad clipping the GMM covariance fit *helped* (n_eff 2.8 → 6.9): protecting the fit from +outliers yielded a better proposal. lnZ differences are within the large MC error at these low n_eff +(the unbiased estimator realizes low on heavy-tailed events; see the side finding below). + +**Verdict.** Proposal-fit clipping is the *correct, unbiased* form of the tool: on well-behaved +weights it is a no-op, and it is a safety valve against a single pathological weight wrecking a +member's covariance fit or the allocation signal. It does **not** manufacture n_eff — where the tail +carries the integral, clipping the adaptation only hurts. The durably valuable artifact is the tracked +withheld-mass fraction, a sharp cheap statement of whether an estimate/fit hangs on a handful of +samples. + +**Guidance before promoting to the individual integrators (AV in particular):** +1. Clip only quantities that feed *adaptation*, never the estimator. If a future design must clip an + estimator, it needs the mass tracking **and** a refuse-to-clip gate (abort once the withheld + fraction exceeds ~1e-3) — otherwise it trades evidence accuracy for a flattering n_eff. +2. Surface the withheld-mass fraction as a first-class run diagnostic regardless — it detects the "a + few samples carry the integral" regime that also invalidates n_eff. +3. Deep weight changes inside AV need the **full LVK PP campaign**, not one-off runs: a −0.1 nat ln Z + bias passes a single-event n_eff check but shows up as PP miscalibration. The −11.5 nat bias hiding + behind a perfect n_eff above is exactly why. + +## POST-#33 RE-MEASUREMENT (S250114ax) — the live-GMM result + +Re-run on top of PR #33 (GMM members actually train; AV's `draw_simplified` no longer lies about its +density), warm, same budget: + +| run | Neff≥5 | Neff≥10 | final n_eff @4M | note | +|-----|-------:|--------:|----------------:|------| +| standalone AV (reference) | 0.695M | 1.374M | **100.2** | unchanged by #33 (standalone AV paths untouched) | +| portfolio, default | — | — | **2.1** | was 52.6 pre-#33 (when the GMM was a corpse) | +| portfolio, adaptive alloc | — | — | 1.1 | | +| portfolio, VARAHA draw floor 0.5 / 0.7 | — | — | ~1–2 | AV got 0.97 / 0.85 of draws | +| portfolio + defensive GMM 0.05 | — | — | 2.5 | defensive mixture alone does not fix it | +| **portfolio + proposal-fit clip** | **0.420M** | **0.590M** | 14.4 | **beats AV to the production target** | + +**The portfolio got *worse* when the GMM member came alive.** Pre-#33 the GMM member was a corpse, so +the portfolio was effectively AV-only and scored 52.6; now that it genuinely trains and draws, the +allocation hands it ~0.84 of the budget and the pooled n_eff collapses to ~2 — against standalone +AV's 100. This is not a freeze problem (never-freeze is working: AV updates every chunk) and not a +clipping problem. It is the **draw-allocation** pathology in its clean form: both allocation rules +score members by per-chunk n_ess, and a VARAHA member's per-chunk n_ess sits at ~1 throughout its +slow *cumulative* contraction, so a member that looks instantly good takes the budget. + +**New opt-in lever: `portfolio_varaha_min_frac` / `--portfolio-varaha-min-frac`** reserves a combined +draw fraction for VARAHA members, applied after either allocation rule (legacy or adaptive). It does +what it says — AV's share went to 0.97 (floor 0.5) and 0.85 (floor 0.7) — but it **does not rescue +this event**: even a 3–15% GMM share still poisons the pooled n_eff, which stayed ~1–2. + +**What actually fixes it: clipping the PROPOSAL-FIT input (the "don't poison the sampling model" +lever).** With the GMM member alive, the thing that goes wrong is its *fit* being corrupted by a few +enormous weights; capping the weights that train it (`--portfolio-weight-clip 1.0`, estimator +untouched) turns the collapse around: + +* **n_eff=5 at 0.420M vs AV's 0.695M (1.7× faster); n_eff=10 at 0.590M vs AV's 1.374M (2.3× faster).** +* **This is the production regime**: the real O4 event configs run `--n-eff 10`, so at the target that + actually ships, the clipped AV+GMM portfolio *beats* standalone AV by ~2.3× on this event. +* It then **plateaus at ~14** rather than climbing to AV's 100, so AV alone still wins the *stress* + target (n_eff 100). The ceiling, not the approach, is what the live GMM member still costs. + +Note this only became visible after #33: pre-#33 the GMM member was a corpse, so there was nothing to +poison and clipping did nothing. Ordering of levers on this event: clip (2.1 → 14.4) ≫ defensive +mixture (2.1 → 2.5) ≈ VARAHA draw floor (no rescue) > adaptive allocation (actively worse, 1.1). + +**Remaining gap.** For the stress target the portfolio is still ceiling-limited by the GMM member, and +down-weighting does not fix that (the draw-floor runs show even a 3–15% share caps the pool). Closing +it needs member *exclusion* (drop a member whose credit is persistently negligible) rather than more +re-weighting. Practical guidance today: **use the portfolio with proposal-fit clipping for production +n_eff targets; use standalone AV if you need n_eff ≫ 10 on an AV-favorable event.** + +## Shape-recovery merge gate (PR #31 requirement) + +Run on a quiet node (pcdev11; pcdev12 at load ~440 kills the suite via RLIMIT_NPROC), base +`rift_O4d @4bac7444` vs this branch, both incl. PR #33. Harness: `~/rift_gate_out/run_gate.sh`. + +**Result: `COMPARE_EXIT=0`, 0 blocking regressions** — base and PR identical in aggregate (strict 8/8, +warn-only 5/5, starved 45/45), and 22 of 23 portfolio rows bitwise identical to base. + +Getting there required fixing a regression the gate caught, which is worth recording because the +cause was the opposite of the obvious one: + +* The first gate run showed ~13 of 20 portfolio rows losing n_eff vs base, several by 2–4×. It did + **not block** (portfolio is warn-only, strict = AV+GMM), but this PR changes the portfolio default + path, so it needed attribution rather than a pass-by-classification. +* **never-freeze — the headline default — was NOT the cause**: toggling it gives ratio 1.00 on those + targets (it only engages where a member would actually be frozen). +* **The plateau-aware `_climbing` revive WAS the sole cause.** Toggling it reproduces base exactly: + `d4_n1_s101 25.9→53.5`, `d4_n3_s202 29.1→83.8`, `d6_n1_s202 64.0→102.1`, `d6_n3_s202 7.2→31.4`, + `d8_n1_s101 37.3→61.9` (base 53.5 / 83.8 / 102.1 / 31.4 / 61.9). Forcing updates of members the + freeze schedule would have parked makes their proposals **worse** — the inverse of the intuition + that motivated it. It now defaults **off** (opt-in only). + +**The one remaining difference** is `mix_d2_n3_s303` (n_eff 736→517, bias −0.0060→−0.0064): the row +where never-freeze genuinely engages. Both PASS comfortably (n_eff ≫ 100, bias unchanged), but it +quantifies never-freeze's cost — **it buys starvation-immunity and pays ~30% n_eff where freezing +would have been harmless.** + +**Method note (a trap worth avoiding).** An isolation that drives `shape_recovery` as a library must +export `PYTHONPATH` (the checkout under test), `CUDA_VISIBLE_DEVICES=""` and `OMP_NUM_THREADS` — only +the wrapper `run_shape_recovery.sh` sets these. My first isolation didn't, silently imported the +**installed** RIFT (where the knob under test does not exist), and confidently reported "no effect" +with n_eff nowhere near the gate's. **A valid isolation reproduces the gate's absolute numbers +row-for-row**; that check is what exposed it. `probe_portfolio_optin_flags.py` now sets this env itself. + +### Flag-ON probe (TESTING.md requirement for opt-in changes) + +`test/expensive_before_merging/integrators/probe_portfolio_optin_flags.py` scores the opt-in features +with the gate's own targets, metrics and `evaluate()`, so a PASS here is a PASS by gate criteria. +**Result: 0 opt-in regressions**, and both features materially help: + +| target | flags OFF | adaptive_alloc ON | weight_clip ON | +|--------|----------:|------------------:|---------------:| +| d2_n1_s303 | 1502 | **3021** | **3037** | +| d2_n3_s303 | 517 | **1058** | **816** | +| d4_n1_s303 | 163 | **415** | **263** | +| d4_n3_s303 | 7 (starved) | 15 | **54** | + +Bias stays small on every PASS row (|lnI−lnZ| ≤ 0.024). This independently corroborates the +S250114ax finding: clipping lifts the worst (starved) row 7 → 54. Caveat: on that starved row the +bias grows (−0.147 → −0.279) — at n_eff ≲ 50 the shape is untestable, so treat clipping's gains in +the starved regime as unvalidated for *shape*, even though n_eff improves. + +Note `adaptive+clip` is identical to `clip` alone on these targets — with clipping active the +allocation rule made no further difference here. + +## Multi-event clip validation, post-#33 (does the S250114ax clip win generalise? — NO) + +4 typical O4 events, warm, `--n-eff 30`, in-container real SEOBNRv5PHM, no-clip vs proposal-fit +`--portfolio-weight-clip 1.0`. Run interactively in the `cuda128` container on idle Blackwell nodes +(pcdev11/13) — which also confirmed SEOBNRv5PHM+cupy run on CC 12.0, matching the A100 result. + +| event | no-clip lnZ (n_eff) | clip lnZ (n_eff) | ΔlnZ | n_eff ratio | +|-------|:-------------------:|:----------------:|-----:|:-----------:| +| S231026ab | 17.54 (28.9) | 17.49 (29.8) | −0.05 | ×1.03 | +| S240426s | 29.74 (31.2) | 29.56 (30.3) | −0.18 | ×0.97 | +| S240513ei | 83.76 (3.1) | 83.83 (1.3) | +0.07 | ×0.42 | +| S240703ad | 41.89 (3.3) | 42.27 (5.0) | +0.38 | ×1.53 | + +**Conclusion.** Clipping's dramatic S250114ax result (n_eff=10 2.3× faster than standalone AV) is +**specific to that event's extreme heavy-tailed pathology and does NOT generalise.** On typical events +it is a near-noop (×0.97–1.03); on the two under-converged hard events it is a wash (one up, one down, +both inside the n_eff≈1–5 scatter). ln Z agrees everywhere (|ΔlnZ| ≤ 0.38, within MC error at these +n_eff) — the portfolio replicates the AV integral with or without clipping. This is exactly why +clipping ships **opt-in, default off**: a targeted tool for a specific failure mode, not a general +speedup to impose on typical runs. It closes the study's last open question. + +## The portfolio's actual purpose: rescuing a high-SNR BEST-FIT evaluation + +Everything above tunes the integrator on a *trial* grid point (`overlap-grid-0.xml.gz --event 0`, +m1/m2 28.29/26.69, lnLmax≈1212) — fine for A/B-ing policy, but it is not the science target. The +target that has to work for a loud event is the **best-fit on-source point** +(`target_params.xml.gz`, m1/m2 37.71/34.03, lnLmax≈3040, ρ≈78): sharply peaked AND +distance-inclination/sky correlated. Measured on GPU, warm, bias-safe cover 0.5, `bench_onsource.sh`: + +| sampler | final n_eff @4M | note | +|---------|----------------:|------| +| AV alone | **1.0** | stalls — axis-aligned bins cannot wrap the correlated peak (also confirms the cardassia CPU result on GPU) | +| GMM alone (adaptive) | **NaN chunk-1** | no coverage floor → weights blow up | +| **AV+GMM portfolio (adaptive)** | **14.7** (lnZ 3016) | ~15× AV, and works where NEITHER member works alone | + +**This is the clearest demonstration of why the portfolio exists.** It is a *GMM-peak + AV-coverage* +event: the GMM member wraps the correlated peak (which AV cannot), and the "dead" AV member — it +reports `nan` per-chunk n_ess in 381/400 chunks and sits at the 1% floor — is NOT wasted, because its +broad warm density still enters `q_mix = frac_AV q_AV + frac_GMM q_GMM`, providing the coverage floor +that keeps GMM's importance weights bounded. Remove AV (run GMM alone) and GMM NaNs; remove GMM (run +AV alone) and it stalls at 1.0. The portfolio is exactly the vehicle that combines a peak-finder with +a coverage member, and the never-freeze/allocation machinery above is what lets it hand the budget to +whichever one is actually working — here, GMM. + +**Adaptive GMM coverage is the lever — but it is NON-MONOTONIC, with a sweet spot:** + +| portfolio config (AV+GMM, warm 0.5) | GMM BIC cap | inflate | n_eff @4M | lnZ | +|-------------------------------------|-----------:|--------:|----------:|-------:| +| baseline | 8 | 1.0 | 14.7 | 3016.13 | +| **sweet spot** | 16 | 1.3 | **56.1** | 3016.08 | +| over-cranked | 24 | 1.5 | **2.3** | 3009.5 ⚠ | + +At the sweet spot: **56× standalone AV** (which stalls at 1.0), ln Z unchanged (3016.08 vs 3016.13) — +real efficiency, not a coverage-shortcut bias. **But more is not better**: cap 24 / inflate 1.5 +collapses to 2.3 AND ln Z drops 6.6 nats (3009.5) — the lnZ shift means over-inflation is biasing, +not just adding variance (an over-wide GMM proposal + a few enormous weights, the same heavy-tail +mode weight clipping was aimed at). So the reviewer's "GMM event with adaptive coverage" framing is +confirmed quantitatively, with the caveat that the coverage knobs need *tuning to a sweet spot*, not +maximizing. Which of the two knobs (BIC cap vs inflation) drives the collapse is under isolation. +Harness: `test/integrators/bench_onsource.sh` (pins the best-fit point; documents it is NOT the trial +point). + +## Files +- `RIFT/integrators/mcsamplerPortfolio.py` — freeze-policy + adaptive-probe allocation, knobs, + n_ess history, plugin-load guard, NaN guard. +- `bin/integrate_likelihood_extrinsic_batchmode` — CLI flags (freeze + allocation) + comma-split fix. +- `test/integrators/bench_portfolio_freeze.sh` — S250114ax single-config runner. +- `test/integrators/bench_multi_event.py` + `run_multi_event.sh` — multi-event robustness suite. +- `test/integrators/parse_neff_traj.py` — trajectory → n_eff-vs-N table parser. +- `test/integrators/test_portfolio_adaptive_alloc.py` — synthetic correlated/uncorrelated test that + the portfolio tracks the winning member and beats AV on a correlated target (Benchmark 3). +- `test/integrators/bench_weight_clip.py` — clipping bias-vs-n_eff sweep against analytic ln Z. + +## The high-SNR rescue is FLAKY — a single run is not a posterior (seed ensemble) + +The 56.1 sweet-spot number above is a **single lucky draw**, not the typical outcome. Repeating +cap 16 / inflate 1.3 across seeds (same best-fit on-source point, warm 0.5, n_max 4M): + +| copy | n_eff @4M | lnZ | +|------|----------:|-------:| +| unseeded | **56.1** | 3016.08 | +| seed 1 | 1.3 | 3011.20 | +| seed 2 | 1.5 | 3017.07 | +| seed 3 | 9.2 | 3015.96 | +| seed 4 | 1.2 | 3010.65 | + +Median n_eff ≈ **1.5**; the distribution is bimodal (mostly collapsed, occasionally lands). lnZ +swings 3010.6 → 3017.1 (6.4 nats) with **no clean sign** — a single dominating outlier can push +evidence high (seed 2: n_eff 1.5 but lnZ 3017.1) or low (seed 4). **One low-n_eff portfolio run on a +high-SNR event is not a usable posterior at any budget.** The operational recipe for such events is +to run MANY independent copies and pool (below), or find a proposal that reliably lands high n_eff. + +**Pooling recovers the answer — but only if pooled by reliability, not naively.** n_eff-weighted mean +of the five copies' lnZ = **3016.0** (the two high-n_eff copies, 3016.08 & 3015.96, agree and +dominate); the *unweighted* mean is 3014.2, biased ~1.8 nats low by the collapsed copies. Naive +concatenation of raw importance samples is WORSE than either: it is dominated by whichever copy owns +the single largest weight — which may be a *collapsed* copy. So "pool many copies" means pool enough +that the **pooled cloud's own n_eff** is high; a handful of copies can still be outlier-dominated. + +**cap-too-high is the failure mode (confirmed).** Bigger BIC cap / inflation → wider GMM proposal → +more prone to the single-enormous-weight collapse: cap 24 median n_eff 2.3–2.8 vs cap 16's occasional +56. The knob buys peak coverage at the cost of tail control; past the sweet spot the tail wins. + +### `--save-samples` is unusable for portfolio shape-checks in this regime (four independent layers) + +Trying to export the extrinsic cloud for a weighted-shape check surfaced that the export path fails +for a peaked (low-to-moderate n_eff) portfolio run at *four* layers — every seed above exported +**0 rows**, even seed 3 at n_eff 9.2: + +1. **Fairdraw** (`--fairdraw-extrinsic-output`) resamples ∝ weight → at n_eff≈1 it returns copies of + the one dominant point or nothing. Useless at low n_eff (per reviewer guidance). +2. **`--save-P` defaults to 0.1** — the export prunes the bottom 10% of *probability*; on a peaked + cloud that discards nearly everything. Raw weighted export needs `--save-P 0`. +3. **`mcsamplerPortfolio` `_rvs` cleanup (draft-inherited, ~line 1135) cumulative-sums the LOG-weights**, + not the weights, and is poisoned by any `-inf` ln_wt entry → 0 rows survive even at n_eff 9.2. + (Pre-existing pattern copied from the ensemble sampler; flagged as a separate fix, not touched here.) +4. **The XML only carries `loglikelihood = log_integrand` (lnL), not the IS weight** + (`log_integrand + log_joint_prior − log_joint_s_prior`). A weighted-posterior check off the XML is + therefore wrong-by-construction (weights by likelihood, not posterior). `shape_extrinsic.py` had + this bug. + +**Correct path for a weight-aware shape/posterior check: `--extrinsic-proposal-output`** — it builds +the TRUE importance log-weights from the raw `_rvs` cloud (driver ~line 2856) and fits a per-group +GMM, bypassing all four failure layers. Pool/compare those GMM fits across copies for the posterior. +(NB: it still needs `--save-P 0`, or the same buggy `_rvs` cleanup prunes the cloud to 0 rows and the +fit dies with "zero-size array to reduction cupy_max". Same root bug as layer 3 above.) + +## The real answer: n_eff is a LOTTERY for every config — pooling is mandatory + +Goal (per reviewer): not max n_eff — *reliable modest* n_eff with a stable, unbiased extrinsic +posterior and no failure mode tied to extrinsic multimodality/degeneracy. Seed ensemble on the +best-fit high-SNR point (warm 0.5, n_max 4M), AV+GMM portfolio, GMM-coverage configs. + +**FIRST, A CORRECTION / METHOD LESSON.** A 3-seed run of cap8 gave {7.0, 10.1, 13.7} and I wrote +"cap8 is reliably modest." That was survivorship bias on 3 draws — the exact trap this document warns +about. Extending cap8 to **10 draws** (GPU runs are non-deterministic even at fixed `--seed`: float +reduction order) gives: + + cap8 n_eff (10 draws): 1.00, 1.00, 1.06, 1.57, 7.0, 10.1, 13.7, 22.0, 55.3, 70.1 + median ~8.5, range 1 -> 70, ~40% collapsed to ~1 + +cap8 is **just as bimodal as cap16** — it is not a reliability fix. n_eff on this high-SNR best-fit +point is a **lottery** for the portfolio regardless of the BIC cap: most runs collapse to ~1, a +minority land 10–70. lnZ tracks the mode (collapsed runs bias lnZ 5–11 nats low). Across configs: + +| config | GMM coverage | n_eff draws | reliability | +|--------|-------------|-------------|-------------| +| cap8 (factored) | cap 8, inflate 1.0 | 1,1,1.06,1.6,7,10,14,22,55,70 (n=10) | bimodal lottery | +| cap16 (factored) | cap 16, inflate 1.3 | 1.2,1.3,1.5,3.5,9.2,13.6,56,59 (n=8) | bimodal lottery (~same) | +| corr (correlate-all) | single 6-D GMM, cap 8 | 1.8,1.9,20.6 (n=3) | strictly WORSE (see below) | + +**Consequence (this is the reviewer's original point, now proven on real data): a single run — any +config — is NOT a posterior on this event. The only robust recipe is to run MANY independent copies +and pool.** Pool by reliability, not naively (see the pooling note above): the pooled *cloud's own* +n_eff must be high. The cap knob changes the odds of a good draw only marginally; it does not remove +the need to pool. + +**The "strongly-correlated problem → correlate-all" hypothesis is REFUTED.** A single full-dimension +(6-D) GMM that *can* represent cross-group (sky–phase, dL–ι) correlation is the WORST here: it +collapses on 2 of 3 seeds and biases lnZ up to 11 nats low (3004.5). Reason: a 6-D mixture needs +~(d+2) effective samples per component; at the modest n_eff these runs produce, its covariances go +near-singular and a few enormous weights dominate. The **factored per-group (2-D) proposal is more +robust** precisely because each low-dimensional fit is cheap and well-conditioned — the correlation +it cannot represent costs less than the fitting variance a full-dim GMM incurs. So *more* proposal +expressiveness is the wrong lever; the lever is **more copies**. + +Harness: `test/integrators/bench_onsource_ensemble.sh` + `compare_extrinsic_breadcrumbs.py`. The +weight-correct extrinsic export needed for the pooled shape check is unblocked by PR #35 (the +`--save-samples`/`_rvs` cleanup fix: linear-weight cumsum + `-inf` guard), cherry-picked here. + +### The failure mode IS an extrinsic-degeneracy collapse — and pooling landed copies is robust + +9-copy cap8 pool (seeds 10–18, `--extrinsic-proposal-output`): 4 landed (n_eff 15,36,39,41), 5 +collapsed (n_eff 1–3.2). The weight-correct per-group GMM fits give a clean picture: + +- **Mode count is a perfect collapse diagnostic, and the collapse is exactly the reviewer's worry.** + Every LANDED copy fits **3–4 modes** in each degeneracy group — (ra,dec) sky **ring**, (distance,ι) + arc, (φ,ψ). Every COLLAPSED copy fits **1 mode** in every group: a single degenerate blob that has + **lost the sky ring / dL–ι arc / phase-pol structure**. So low n_eff ⟺ extrinsic *mode collapse*; + the settings' instability is tied directly to multimodality/degeneracy, and n_eff (or the fitted + mode count) detects it. +- **Landed copies AGREE — when it lands, the posterior is stable and reproducible.** Across the 4 + landers the (ra,dec) mixture mean agrees to ~0.01 (frame units) and (distance,ι) to ~0.02 in ι — + i.e. the recovered extrinsic posterior is *consistent copy-to-copy*, no hidden instability among + good runs. The exception is (φ_orb,ψ): scatter ~1.5 even among landers, because the 2-IFO phase– + polarization degeneracy is genuinely the least-constrained extrinsic direction (expected, not a bug). +- **Reliability-weighted pooling ≈ good-only (correct); naive pooling is biased by the collapsed + copies.** For the well-constrained sky group all three pooling recipes coincide, but for the looser + distance and phase groups naive-unweighted pooling is pulled off the good-only answer (distance: + naive vs good differ ~80 units; phase: −0.74 vs −1.88) while the n_eff-weighted pool tracks good-only. + Reliability-weighted **effective #copies (Kish over n_eff) = 4.1** — the 9-copy pool really rests on + its ~4 landers. **Operational recipe: run ~2–3× as many copies as landers you need, pool weighted by + n_eff (or simply drop n_eff<5 copies).** + +Caveat (honest): the comparator's *physical* un-normalization of the GMM means is in the wrong frame +(the RIFT GMM's internal normalization is not the naive [0,1]-on-bounds I assumed — all means flag +out-of-bounds, so that flag is unreliable). The conclusions above rest only on the frame-INDEPENDENT +signals — mode counts and copy-to-copy agreement (`good-scatter`) — not on absolute mean values. A +correct physical read needs the GMM model's normalization; the mode-collapse / pooling story does not. + +### Coordinates/adaptation help the LANDERS, not the collapse rate — the collapse is an AV peak-lock lottery + +Testing the reviewer's high-SNR recipe (`--force-adapt-all` + rotations). CONFOUND first: the +coordinate-transform flags (`--internal-rotate-phase`, `--internal-sky-network-coordinates`) change +what the sampler's parameter slots MEAN, but `--sampler-warmstart-samples` maps the seed by column +NAME without transforming values -> a PHYSICAL seed poisons the rotated/network proposal. Naive +"add the flags" run: 0/9 landed (every copy collapsed). Fix = a frame-matched seed +(`seed_phi_orb=mod(phi+psi,4pi)`, `seed_psi=mod(phi-psi,4pi)` for rotate-phase; `--force-adapt-all` +is frame-preserving and needs no transform). Now in the lore repo's gotchas. + +With a frame-matched seed (`--force-adapt-all --internal-rotate-phase`, 9 copies): + +| metric | baseline (physical) | +force-adapt-all+rotate-phase | +|--------|--------------------:|------------------------------:| +| landed fraction (n_eff>=5) | 4/9 | **4/9 (unchanged)** | +| landed n_eff | 15,36,39,41 | **41,52,52,41** (higher, tighter) | +| landed sky modes | 3-4 | 3-4 (ring preserved) | + +So phase-decorrelation + full adaptation is a real efficiency win FOR THE LANDERS (n_eff ~50 vs ~30) +but does NOT move the ~55% collapse rate. The lottery is now robust across EVERY config tried (cap8, +cap16, correlate-all, +rotate-phase): same ~50% collapse, same signature (n_eff~1, single mode). The +root cause is therefore not the proposal/coordinates but **AV's contracting box locking onto the +sharp high-SNR peak or contracting around the wrong spot ~50/50** — and the portfolio cannot backstop +better than AV's own contraction reliability, because AV itself is the coin-flip. + +**Targeted fix under test: L0 auto-rescue** (`--sampler-warmstart-retry-neff`). If a pass finishes +n_eff < threshold, re-seed AV from the run's OWN highest-L samples (the peak it did find) and re-run +— same-problem reuse, cannot bias, frame-safe by construction. Was gated to standalone AV; relaxed to +fire for the portfolio too (peak-seed bootstraps into the AV member). This is the in-loop version of +"pool copies": convert each collapsed draw into a land instead of discarding it. Result pending (prr_). + +### THE HIGH-SNR FIX: L0 auto-rescue roughly DOUBLES the landed fraction (4/9 -> 8/9) + +Since the collapse is AV losing the sharp peak ~50/50 (not a proposal/coordinate defect), the fix is +to re-seed a collapsed run from the peak IT DID FIND and re-run: `--sampler-warmstart-retry-neff 5` +(L0 auto-rescue). Bug found + fixed first: the rescue is gated on the sampler being AV or a +portfolio, but `opts.sampler_method` is CLOBBERED to 'GMM' during portfolio member setup (line ~1231, +`opts.sampler_method='GMM'` forces GMM arg-parsing), so an AV+GMM portfolio reports method 'GMM' +everywhere downstream. The gate now detects the portfolio via `opts.sampler_portfolio` (the member +list, which survives the clobber). [Same clobber makes the portfolio-only block at ~1641 dead code -- +harmless, the GMM branch picks up the gmm_adaptive forwarding as a per-group dict -- but a latent +footgun; flagged for cleanup.] + +9-copy pool, cap8 + `--force-adapt-all --internal-rotate-phase` (frame-matched seed) + +`--sampler-warmstart-retry-neff 5`: + +| seed | prior behavior | prr_ result | rescue | +|------|---------------|------------:|:------:| +| s10 | collapse | 4.5 | fired (just under) | +| s11 | chronic ~1 collapse | **33.4** | fired -> LAND | +| s12 | collapse | 33.4 | (landed pass 1) | +| s13 | 35-52 | **47.0** | fired -> LAND | +| s14 | chronic ~1 collapse | **25.2** | fired -> LAND | +| s15 | mixed | **19.6** | fired -> LAND | +| s16 | collapse | **38.1** | fired -> LAND | +| s17 | 15 | 10.0 | (landed pass 1) | +| s18 | 39-41 | 21.4 | (landed pass 1) | + +**LANDED 8/9** (baseline 4/9, pr_ 4/9); 6 rescues fired, 5 converted to clean lands and the 6th to +4.5. Chronic collapsers (s11, s14, both stuck at n_eff~1 across every prior config) now land at 25-33. +Cost: a rescued run does 2 integration passes (~2x). This is the in-loop equivalent of "pool copies", +and it is the real high-SNR lever -- coordinates/adaptation improve the LANDERS, the rescue fixes the +COLLAPSE RATE. + +**Validated high-SNR recipe:** portfolio AV+GMM (cap8, adaptive components) + `--force-adapt-all` ++ `--internal-rotate-phase` (with a phase-frame-matched warm seed) + `--sampler-warmstart-retry-neff 5`. +Even so, for a publication-grade posterior at n_eff this modest, still pool a few landed copies. + +## EVIDENCE AUDIT: which numbers in this document are single draws + +The n_eff lottery (documented above) was discovered LATE, after much of this document was written. +Because a single run on a lottery-prone point is noise-dominated, several earlier claims here rest on +n=1 and must be read as suggestive, not established. Explicit audit: + +**Downgraded to UNPROVEN (single draw on a bimodal quantity):** +- The GMM coverage ladder cap8=14.7 / cap16=56.1 / cap24=2.3 -- all n=1. The cap16 "sweet spot" is + already retracted above; **the companion claim that cap24 over-cranking BIASES lnZ (3009.5, -6.6 + nats) is likewise a single draw and is NOT established.** A collapsed copy shifts lnZ in either + direction (seed 2 of the cap16 ensemble: n_eff 1.5 but lnZ 3017.1, i.e. HIGH). Distinguishing + genuine over-inflation bias from collapse noise needs a seed ensemble per cap, which has not run. +- `--internal-gmm-correlate-all` is worse: n=3 (2/3 collapsed, lnZ up to 11 nats low). Directionally + supported and mechanistically plausible (a 6-D mixture needs ~(d+2) eff-samples/component), but not + firm at n=3. +- Benchmark 1's cold rows (av_cold 3.7, pf_nf_cold 1.1): single draws on the lottery-prone point. + +**Robust (large effect, understood mechanism, and/or well sampled):** +- Never-freeze rescues the workhorse (3.4 -> 53): large, mechanism understood (frozen at chunk 1), + and independently corroborated by zero freeze notices across the multi-event suite. +- Multi-event ln Z replication (Benchmark 2, NON-warm-started): an UNBIASEDNESS claim, structurally + guaranteed by the balance-heuristic q_mix (the estimate is unbiased for any member weights). The + ΔlnZ agreement stands. (The n_eff-efficiency comparisons in that same table are single draws.) +- The lottery itself (cap8 n=10, cap16 n=8), the mode-collapse diagnosis (9 copies, clean 1-mode vs + 3-4-mode split), and the L0 auto-rescue 4/9 -> 8/9 (9 copies + a post-cleanup regression). + +**OPEN: is the lottery high-SNR-only?** Every ensemble here is on the ultra-sharp best-fit point of a +loud event. If typical events are unimodal in n_eff, single-draw comparisons on them (Benchmark 2) are +fine as-is; if not, that table's efficiency numbers need ensembles too. Cheap to settle: one seed +ensemble on a typical event. + +**Not a factor: the sampler_method clobber.** For an AV+GMM portfolio the clobber changed only whether +`return_lnI` was passed, and `mcsamplerPortfolio` never reads it (`use_lnL` was set either way, because +the portfolio branch force-sets `internal_use_lnL=True` before the clobber). Verified by regression: +identical per-group `gmm_adaptive` forwarding and identical rescue behaviour. No result in this +document is invalidated by removing it. + +## COLD-START ensemble: n_eff does NOT certify correctness (the confidently-wrong failure) + +Rerun of the cold (non-warm-started) case after two fixes landed: the pre-existing +`mcsamplerEnsemble` loop-invariant clobber (which had made EVERY cold portfolio start crash at +chunk ~8 with no output at all -- 0/9), and the L0 rescue now firing on degenerate early +termination. Config: portfolio AV+GMM cap8 adaptive, `--force-adapt-all --internal-rotate-phase +--interpolate-time True --sampler-warmstart-retry-neff 5`, 9 seeds, cold. + +| seed | n_eff | lnZ | modes/group | rescue | +|------|------:|--------:|:-----------:|:------:| +| s10 | 1.5 | 3006.16 | 1 | fired | +| s11 | 13.9 | **3012.47** | 1 | fired | +| s12 | 31.0 | **3013.87** | 1 | fired | +| s13 | 1.0 | 3001.19 | 1 | fired | +| s14 | 38.0 | **3013.61** | 4 | fired | +| s15 | 1.1 | 3002.36 | 1 | fired | +| s16 | 18.0 | **3012.54** | 1 | fired | +| s17 | **58.0** | **3001.68** ⚠ | 1 | fired | +| s18 | 3.8 | 3015.53 | 2 | fired | + +(lnZ is only comparable WITHIN this table: `--internal-rotate-phase` doubles the prior, so these +values are offset from the non-rotated benchmarks earlier in this document.) + +**9/9 now produce output (was 0/9 -- the crash), 5/9 land (n_eff>=5).** Cold is materially worse than +warm+rescue (8/9), so a warm seed still earns its keep; but cold now WORKS, which it did not before. + +**The headline result is the lnZ column, not the landed count.** Among the five landed copies lnZ +spans **3001.7 - 3013.9 (12 nats)**, and the single most wrong copy is the one with the **HIGHEST +n_eff**: s17, n_eff 58, lnZ 11 nats below the consensus. Four of five landers agree to within 1.4 +nats (3012.5-3013.9); s17 dissents while looking, by n_eff, like the best run in the ensemble. + +**Consequences (this changes the recommended practice):** +1. **n_eff is NECESSARY BUT NOT SUFFICIENT.** It measures weight concentration, not coverage. A pass + that locks onto one narrow region has low weight variance (high n_eff) while missing posterior + mass (lnZ too low) -- confidently wrong. You CANNOT pick the trustworthy copy by max n_eff, and a + single high-n_eff run is not self-certifying. +2. **Use CONSENSUS across copies, not the best-n_eff copy.** The outlier here is detectable only by + disagreeing with the pool. Prefer the median lnZ over landed copies (median 3012.54 correctly + rejects s17) to an n_eff-argmax or even an n_eff-weighted mean (which s17's weight would drag + down). This is a direct strengthening of the "run MANY copies" recipe: copies are needed not just + to find a good draw, but to DETECT a bad one that looks good. +3. Mode count is a useful but imperfect cross-check here: s14 (4 modes) sits in the consensus, but + s11/s12/s16 are 1-mode and also in the consensus, so a low mode count alone does not condemn a + run at this sample size. Cross-copy agreement remains the strongest signal. + +### AUTO-COLLECTED raw results: AV-backstop / mode-budget sweep (cold, high-SNR best-fit point) + +Config base: portfolio AV+GMM, adaptive components, `--force-adapt-all --internal-rotate-phase`, +`--interpolate-time True`, `--sampler-warmstart-retry-neff 5`, cold (no warm seed). +`bk` = `--portfolio-varaha-min-frac 0.25` (cap 8); `md` = `--internal-gmm-max-components 3` +(no floor); `bkmd` = both. Judged by lnZ CONSENSUS across seeds, not n_eff. + +| config | seed | n_eff | lnZ | AV final frac | +|--------|------|------:|----:|--------------:| +| bk | s10 | 1.0 | 3013.17 | 0.25 | +| bk | s12 | 1.0 | 3009.18 | 0.9900964290627214 | +| bk | s14 | 6.5 | 3013.41 | 0.25 | +| bk | s17 | 11.2 | 3014.23 | 0.25 | +| md | s10 | 12.4 | 3012.42 | 0.009900990099393974 | +| md | s12 | 5.6 | 3006.52 | 0.009900990099649775 | +| md | s14 | 123.6 | 3003.09 | 0.00990112295232892 | +| md | s17 | 22.7 | 3015.00 | 0.009900990099929107 | +| bkmd | s10 | 9.5 | 3010.88 | 0.25 | +| bkmd | s12 | 11.8 | 3011.52 | 0.25 | +| bkmd | s14 | 1.9 | 3011.12 | 0.25 | +| bkmd | s17 | 26.6 | 3013.15 | 0.25 | + +Baseline for the SAME four seeds (no floor, cap 8): s10 3006.16 / s12 3013.87 / s14 3013.61 / +s17 3001.68 -> 12.2 nat spread, with the highest-n_eff copy (s17, n_eff 58) the most wrong. + +lnZ spread per config (max-min over the four seeds): +- `bk`: lnZ = 3013.17 3009.18 3013.41 3014.23 -> spread 5.05 nats +- `md`: lnZ = 3012.42 3006.52 3003.09 3015.00 -> spread 11.91 nats +- `bkmd`: lnZ = 3010.88 3011.52 3011.12 3013.15 -> spread 2.27 nats + +Shape-recovery merge gate: + +### Banded VARAHA share: the share constraint and the mode budget only work TOGETHER + +Adding `--portfolio-varaha-max-frac` (cap) to the existing floor, and crossing it with the GMM BIC +cap. Cold, high-SNR best-fit point, 4 seeds each, judged by lnZ consistency (NOT n_eff): + +| config | VARAHA share | GMM cap | lnZ (s10,s12,s14,s17) | sd | spread | +|--------|--------------|--------:|-----------------------|-----:|------:| +| baseline | unconstrained | 8 | 3006.2 3013.9 3013.6 3001.7 | 5.96 | 12.19 | +| `md` | unconstrained | 3 | 3012.4 3006.5 3003.1 3015.0 | 5.43 | 11.91 | +| `band8` | band .25-.75 | 8 | 3015.0 3011.9 3003.3 3014.6 | 5.42 | 11.65 | +| `bk` | floor .25 | 8 | 3013.2 3009.2 3013.4 3014.2 | 2.26 | 5.05 | +| `bkmd` | floor .25 | 3 | 3010.9 3011.5 3011.1 3013.2 | 1.02 | 2.27 | +| **`band3`** | **band .25-.75** | **3** | **3012.9 3011.2 3012.9 3012.6** | **0.85** | **1.79** | + +**Neither lever works alone.** A reduced mode budget with an unconstrained share (`md`) is no better +than baseline (sd 5.43 vs 5.96) -- and it is the arm that produced the worst confidently-wrong case in +the whole study (n_eff 123.6, lnZ 10 nats low). A share constraint alone helps but inconsistently. +Only the two configurations combining a VARAHA share constraint WITH the modest mode budget are +tight (sd 0.85 / 1.02), and they are the ONLY two with no outlier >=5 nats from their own median. + +Mechanism consistent with the rest of this section: the share constraint keeps a broad backstop in +q_mix so no mode is left uncovered, and the modest mode budget stops the peaked member from splitting +into many narrow components that individually chase structure and collectively lose coverage. + +**Statistical caveat, stated plainly:** n=4 per config. `band8` (5.42) vs `bk` (2.26) differ only by a +cap that bound on one seed, so that gap is almost certainly noise -- an sd on 3 dof swings by ~2x +routinely. Read this table as "the floor+cap3 FAMILY is tight, the rest is not", NOT as a fine +ranking. A confirmation run extending `band3` and `bkmd` to 5 further seeds each (n=9) is under way. + +**Gate status: NOT YET CLEARED.** Every knob here is opt-in, so the default-path merge gate is +bitwise-blind to all of it. `probe_portfolio_optin_flags.py` now carries `varaha floor .25`, +`varaha band .25-.75` and `band + gmm cap3`, scored by the gate's own `evaluate()`. NOTHING here is +proposed as a recommendation or default until that probe passes AND the base-vs-branch gate compares +clean. + +### n=9 like-for-like: the constraint helps, but NOT significantly, and does not make one run trustworthy + +The n=4 table above was a fluke of seeds {10,12,14,17} (band3 sd 0.85 -> 3.11 when extended). Extending +all three configs to the SAME nine seeds {10,12,14,17,20..24}: + +| config | lnZ sd | spread | worst deviation from own median | +|--------|-------:|-------:|--------------------------------:| +| baseline (unconstrained share, cap 8) | 5.04 | 12.71 | 10.4 nats | +| `bkmd` (floor .25, cap 3) | 3.01 | 10.25 | **6.1 nats** | +| `band3` (band .25-.75, cap 3) | 3.08 | 7.79 | 7.2 nats | + +F-test on the variances (n=9 each): +- baseline vs `bkmd` : F=2.81, one-tailed p=0.083 -- **NOT significant at 5%** +- baseline vs `band3` : F=2.69, one-tailed p=0.092 -- **NOT significant at 5%** +- `bkmd` vs `band3` : F=1.05, p=0.48 -- indistinguishable; the CAP adds nothing measurable over the FLOOR + +**Honest reading.** The share constraint cuts lnZ scatter ~40% and roughly halves the worst-case +deviation, which is a real-looking effect with a plausible mechanism (q_mix keeps a broad backstop, so +no mode goes uncovered) -- but at n=9 it does NOT reach significance. Do not present it as an +established improvement. Reaching p<0.05 on a variance ratio this size needs ~20+ seeds per config. + +**What does NOT depend on the significance test:** every catastrophic confidently-wrong case in this +study (n_eff 58 with lnZ 11 nats low; n_eff 123.6 with lnZ 10 nats low) occurred in an +UNCONSTRAINED-share arm, and none occurred in a constrained arm. + +**What is settled regardless:** no configuration makes a SINGLE run trustworthy on this point -- the +best still deviates 6 nats from its own median. Pooling across copies, judged by consensus rather +than by n_eff, remains mandatory. + +### Merge gate: PASSED (PR #34) + +`compare_shape_results.py` over all 96 rows, base `rift_O4d` vs this branch, both arms run +single-process: **0 blocking regressions (strict = AV, GMM), COMPARE_EXIT=0**, no REGRESSION / +BLOCKS-MERGE / ONLY-IN rows. `PREEXISTING-FAIL` rows fail identically on base. +NOTE the gate must be run with `--jobs 1`: its multiprocessing pool DEADLOCKS at higher job counts +(observed on both arms independently) -- see the lore repo's gotchas. + +### Flag-ON gate probe: PASSED (0 opt-in regressions) + +`probe_portfolio_optin_flags.py`, scored by the gate's own `evaluate()` so a PASS here passes by +exactly the gate's criteria. Each configuration compared against the SAME target with flags OFF: + +| configuration | d2_n1_s303 | d2_n3_s303 | d4_n1_s303 | d4_n3_s303 | +|---------------|-----------|-----------|-----------|-----------| +| `varaha floor .25` | PASS | PASS | PASS | STARVED (base STARVED too) | +| `varaha band .25-.75` | PASS | PASS | PASS | STARVED (base STARVED too) | +| `band + gmm cap3` | PASS | PASS | PASS | STARVED (base STARVED too) | + +**opt-in regressions: 0.** The pre-existing opt-in features (adaptive_alloc, weight_clip, and their +combination) also remain at 0 regressions. As expected the constraints are ~no-ops on the gate's +well-behaved targets -- which is the point: they must not COST anything where they are not needed. + +## STATUS SUMMARY (all three bars) + +| bar | result | +|-----|--------| +| merge gate (base vs branch, 96 rows) | **PASS** -- 0 blocking regressions, COMPARE_EXIT=0 | +| flag-ON probe (proposed settings) | **PASS** -- 0 opt-in regressions | +| real-event benefit (n=9, matched seeds) | scatter 5.04 -> ~3.0 sd, worst dev 10.4 -> 6.1 nats, but **p~0.08: NOT significant** | + +So the settings are SAFE (both gates clear) but their benefit is not yet PROVEN. Recommended posture: +keep them opt-in and documented for high-SNR use, do NOT change any default, and either accept the +caveat or spend ~20 seeds/config to settle significance. + +### RETRACTION at n=20: the share constraint does NOT reduce lnZ scatter + +The n=4 result above (`band3` sd 0.85 vs baseline 5.96) does not survive. Extending the same three +configurations on the same point: + +| config | n=4 sd | n=9 sd | **n=20 sd** | n=20 spread | +|--------|-------:|-------:|------------:|------------:| +| baseline (unconstrained, cap 8) | 5.96 | - | **4.02** (n=15) | 13.49 | +| `bkmd` (floor .25, cap 3) | 1.02 | 3.02 | **4.50** | 15.78 | +| `band3` (band .25-.75, cap 3) | 0.85 | 3.11 | **5.02** | 22.03 | + +Variance-ratio F-test vs baseline, one-tailed: `bkmd` F=0.80 p=0.66; `band3` F=0.64 p=0.80. Not +significant, and the point estimates are in the WRONG direction -- the constrained arms have slightly +LARGER scatter than the unconstrained baseline. + +**This is textbook small-sample selection followed by regression to the mean.** `band3` was CHOSEN +because it looked best at n=4; its estimate then decayed 0.85 -> 3.11 -> 5.02 as n grew. The n=4 +caveat recorded above ("read this as a family, not a ranking") was correct but not strong enough: the +honest position is that at n=4 these configurations carried NO usable information about scatter. + +**What this retracts:** any claim that `--portfolio-varaha-min-frac` / `--portfolio-varaha-max-frac` +or a reduced GMM BIC cap improves lnZ consistency on this event. They remain OPT-IN and OFF by +default, and must be described as unproven rather than recommended. + +**What still stands (independent observations, not the remedy):** +- n_eff does not certify correctness: individual copies with the HIGHEST n_eff in their arm were the + most wrong in lnZ (n_eff 58 / 11 nats low; n_eff 123.6 / 10 nats low). Selection must not use n_eff. +- The mixture degenerates to peaked-member-only (VARAHA share -> 0.0099) unless constrained; q_mix + then carries no broad backstop. That is a measured structural fact about the allocation rule. +- The pre-existing cold-start crash (`mcsamplerEnsemble` loop-invariant clobber) and the L0 rescue + gap on degenerate termination were real bugs and are fixed. + +**Gate status: CLEARED.** Default-path merge gate base-vs-branch: `COMPARE_EXIT=0`, 0 blocking +regressions (96/96 rows both arms). Flag-ON probe (`varaha floor .25`, `varaha band .25-.75`, +`band + gmm cap3`, scored by the gate's own evaluate()): **0 opt-in regressions** -- every row PASSes +where the base PASSes and is STARVED where the base is STARVED. So the knobs are SAFE; they are just +not demonstrated to help. + +## L0 rescue REASSESSED after the warm-start fix: most of the 4/9 -> 8/9 was the BROKEN warm start + +The 4/9 -> 8/9 result above was measured while `AV.bootstrap_from_samples` stored the seeded grid in +`self._warm` but nothing installed it on the portfolio draw path (fixed here; see +`_apply_warm_state`). So BOTH arms of that comparison were effectively cold, and the rescue's apparent +benefit was inflated. Re-measured on the same point with the seed now actually applied, 9 seeds each: + +| arm | n_eff by seed | landed (>=5) | +|-----|---------------|-------------:| +| warm, NO rescue | 1, 26, 16, 2, 12, 29, 1, 24, 13 | **6/9** | +| warm + L0 rescue | 14, 41, 15, 2, 14, 7, 15, 13, 15 | **8/9** | + +Fisher two-sided **p = 0.576** -- NOT significant at n=9. + +Reading: +1. **Fixing the warm start is worth more than the rescue.** A single warm-started run now lands 6/9, + versus the 4/9 baseline measured when the seed was silently discarded. That gain is attributable + to the seed actually reaching the sampler, not to any policy. +2. **The rescue's remaining effect is small and unproven** (6/9 -> 8/9, p=0.58). It is cheap (a second + pass only on runs that collapsed) and never hurt in these data, so it stays recommended as + insurance -- but the earlier "roughly doubles the landed fraction" claim is RETRACTED. With a + working warm start there is simply less left for it to rescue. +3. The retry-vs-reseed mechanism question is now moot at this sample size: both arms here are + warm-started, so this measures the rescue as an ADD-ON to a working seed, which is the + configuration anyone would actually run. + +Unchanged by this: the chunk-size and cubic-interpolation results, which do not involve warm start. diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/MonteCarloEnsemble.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/MonteCarloEnsemble.py index 133661090..e6d7887f3 100755 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/MonteCarloEnsemble.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/MonteCarloEnsemble.py @@ -14,12 +14,17 @@ try: import cupy import cupyx.scipy.special + # Probe for an actual device: cupy imports cleanly on GPU-less nodes but + # every kernel launch then dies with cudaErrorNoDevice. getDeviceCount + # raises CUDARuntimeError (not ImportError), hence the broad except. + if cupy.cuda.runtime.getDeviceCount() == 0: + raise ImportError("cupy installed but no CUDA device available") xpy_default = cupy xpy_special_default = cupyx.scipy.special identity_convert = cupy.asnumpy identity_convert_togpu = cupy.asarray cupy_ok = True -except ImportError: +except Exception: xpy_default = np xpy_special_default = None identity_convert = lambda x: x @@ -96,13 +101,26 @@ class integrator: def __init__(self, d, bounds, gmm_dict, n_comp, n=None, prior=None, user_func=None, proc_count=None, L_cutoff=None, use_lnL=False,return_lnI=False,gmm_adapt=None,gmm_epsilon=None,tempering_exp=1,temper_log=False,lnw_failure_cut=None, - tempering_adapt=False, ess_target=None, ess_floor=None): + tempering_adapt=False, ess_target=None, ess_floor=None, gmm_adaptive=None, + gmm_defensive_frac=0.05, gmm_inflate=1.0): # if 'return_lnI' is active, 'integral' holds the *logarithm* of the integral. # user-specified parameters self.d = d self.bounds = bounds self.gmm_dict = gmm_dict self.gmm_adapt = gmm_adapt + # gmm_adaptive: {dim_group: k_max}. Groups listed here choose their + # component count from the data by BIC (GMM.fit_gmm_adaptive) at + # initialization, then adapt via the stable merge path, instead of using + # a fixed n_comp -- see _train. + self.gmm_adaptive = gmm_adaptive + # defensive tail coverage + covariance inflation for adaptive groups + self.gmm_defensive_frac = gmm_defensive_frac + # Opt-in: install the defensive component on the FIXED-COMPONENT fit paths too. + # Off by default because it costs n_eff at d>=6; a portfolio that relies on this + # member for coverage turns it on (mcsamplerPortfolio.setup). + self.gmm_defensive_all_paths = False + self.gmm_inflate = gmm_inflate self.gmm_epsilon= gmm_epsilon self.n_comp = n_comp self.user_func=user_func @@ -146,6 +164,10 @@ def __init__(self, d, bounds, gmm_dict, n_comp, n=None, prior=None, if self.return_lnI: self.total_value = None self.n_max = float('inf') + # set to a descriptive string when integrate() exits abnormally (error + # budget exhausted); None means a clean run. Callers that cannot catch + # the consecutive-refit-failure RuntimeError can inspect this instead. + self.integration_error = None # saved values self.cumulative_samples = self.xpy.empty((0, d)) self.cumulative_values = self.xpy.empty(0) @@ -336,13 +358,74 @@ def _train(self): for dim in dim_group: temp_samples[:,index] = sample_array[:,dim] index += 1 + # gmm_adaptive may be a dict {group:k_max} (per-group opt-in) or a + # scalar/bool (apply to every adapting group -- used by the portfolio, + # whose GMM member's grouping is not known here). + adaptive_kmax = None + if self.gmm_adaptive: + if isinstance(self.gmm_adaptive, dict): + adaptive_kmax = self.gmm_adaptive.get(dim_group) + elif isinstance(self.gmm_adaptive, bool): + adaptive_kmax = 8 # default cap when enabled globally + else: + adaptive_kmax = int(self.gmm_adaptive) if model is None: - if isinstance(self.n_comp, int) and self.n_comp != 0: + if adaptive_kmax: + # FLEXIBLE allocation: choose this group's component count + # from the data by BIC at INITIALIZATION, then hand off to the + # proven-stable merge adaptation below (model.update()). We + # deliberately do NOT re-fit fresh every chunk: a per-chunk + # BIC refit makes the proposal wander (measured: n_eff peaks + # then collapses) because each fit sees a different elite + # cloud; the incremental merge smooths that out. + # SAFETY FLOOR: never fewer components than the stress-tested + # hard-coded count for this group (self.n_comp) -- adaptive is + # a REFINEMENT that only adds capacity, e.g. a broad multi-modal + # sky keeps its default components. + if isinstance(self.n_comp, dict): + k_floor = self.n_comp.get(dim_group, 1) + else: + k_floor = self.n_comp + k_floor = int(k_floor) if isinstance(k_floor, int) and k_floor > 0 else 1 + model = GMM.fit_gmm_adaptive(temp_samples, new_bounds, + log_sample_weights=log_weights, + k_max=max(int(adaptive_kmax), k_floor), + k_min=k_floor, + epsilon=self.gmm_epsilon, + defensive_frac=self.gmm_defensive_frac, + inflate=self.gmm_inflate) + elif isinstance(self.n_comp, int) and self.n_comp != 0: model = GMM.gmm(self.n_comp, new_bounds,epsilon=self.gmm_epsilon) model.fit(temp_samples, log_sample_weights=log_weights) + # The defensive component is the ONLY thing that actually guarantees this member + # has support across the box -- gmm.score() merely FLOORS at 1e-300, which is a + # numerical guard, not coverage (a sample there would carry weight ~1e300). + # fit_gmm_adaptive adds it; the fixed-component path did not. OPT-IN, because + # measured on the shape gate a 5% broad component costs real n_eff in + # higher dimensions (d6_n3_s303 119->75, d8_n1_s303 448->210): it spends + # 5% of draws where the likelihood is negligible. Only a consumer that + # NEEDS this member as its coverage guarantee should pay -- so a + # portfolio sets gmm_defensive_all_paths on its members, and a standalone + # GMM user is unaffected. + GMM.add_defensive_component(model, defensive_frac=( + getattr(self,'gmm_defensive_frac',0.0) + if getattr(self,'gmm_defensive_all_paths',False) else 0.0)) elif isinstance(self.n_comp, dict) and self.n_comp[dim_group] != 0: model = GMM.gmm(self.n_comp[dim_group], new_bounds,epsilon=self.gmm_epsilon) model.fit(temp_samples, log_sample_weights=log_weights) + # The defensive component is the ONLY thing that actually guarantees this member + # has support across the box -- gmm.score() merely FLOORS at 1e-300, which is a + # numerical guard, not coverage (a sample there would carry weight ~1e300). + # fit_gmm_adaptive adds it; the fixed-component path did not. OPT-IN, because + # measured on the shape gate a 5% broad component costs real n_eff in + # higher dimensions (d6_n3_s303 119->75, d8_n1_s303 448->210): it spends + # 5% of draws where the likelihood is negligible. Only a consumer that + # NEEDS this member as its coverage guarantee should pay -- so a + # portfolio sets gmm_defensive_all_paths on its members, and a standalone + # GMM user is unaffected. + GMM.add_defensive_component(model, defensive_frac=( + getattr(self,'gmm_defensive_frac',0.0) + if getattr(self,'gmm_defensive_all_paths',False) else 0.0)) else: model.update(temp_samples, log_sample_weights=log_weights) try: @@ -432,6 +515,14 @@ def integrate(self, func, min_iter=10, max_iter=20, var_thresh=0.0, max_err=10, self._verbose_diag = verbose # per-chunk adaptation diagnostics in _train err_count = 0 + # Consecutive-refit-failure budget: if the proposal refit fails this + # many chunks IN A ROW the proposal has never adapted and the returned + # integral/eff_samp are meaningless (the cupy-without-GPU regression + # produced exactly this: every refit raised, 'Error training, + # resetting...' each chunk, and integrate() returned eff_samp~1 with no + # error signal). Fail loudly instead. + max_train_fail = int(kwargs["max_consecutive_train_failures"]) if "max_consecutive_train_failures" in kwargs else 5 + consec_train_fail = 0 cumulative_eval_time = 0 adapting=True if nmax is None: @@ -445,6 +536,7 @@ def integrate(self, func, min_iter=10, max_iter=20, var_thresh=0.0, max_err=10, adapting=False if err_count >= max_err: print('Exiting due to errors...') + self.integration_error = 'exited after {} sampling/results/training errors'.format(err_count) break try: self._sample() @@ -490,6 +582,7 @@ def integrate(self, func, min_iter=10, max_iter=20, var_thresh=0.0, max_err=10, try: if adapting: self._train() + consec_train_fail = 0 except KeyboardInterrupt: print('KeyboardInterrupt, exiting...') break @@ -497,7 +590,11 @@ def integrate(self, func, min_iter=10, max_iter=20, var_thresh=0.0, max_err=10, print(traceback.format_exc()) print('Error training, resetting...') err_count += 1 + consec_train_fail += 1 self._reset() + if consec_train_fail >= max_train_fail: + self.integration_error = 'proposal refit failed {} consecutive times; proposal never adapted'.format(consec_train_fail) + raise RuntimeError('GMM ' + self.integration_error) from e if self.user_func is not None: self.user_func(self) if progress: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/TESTING.md b/MonteCarloMarginalizeCode/Code/RIFT/integrators/TESTING.md new file mode 100644 index 000000000..7e5e909f8 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/TESTING.md @@ -0,0 +1,37 @@ +# Before merging changes to this directory + +**Any PR that touches the integrators must pass the posterior SHAPE-recovery +merge gate**, not just the fast CI integral test: + + MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/ + +Integrals are easy — importance-sampling estimates of Z are unbiased under +weak conditions — while the recovered posterior *shape* (the weighted sample +cloud that CIP/fairdraws consume) can be confidently, silently wrong. Real +examples caught by this gate: a GMM run with evidence correct to +0.009 nats +whose marginals had JS=0.29 and widths 2.8-3.8x too broad; and a GPU-port +change whose swallowed per-refit exceptions returned n_eff~1 with no error +flag (rift_O4c -> rift_O4d GMM regression, bisected 2026-07). + +Quick recipe (see the suite README for details; ~10 min per branch on a +quiet head node): + + source ~/RIFT_develUWM/bin/activate # or equivalent env + cd .../test/expensive_before_merging/integrators + SHAPE_JOBS=12 OMP_NUM_THREADS=1 ./run_shape_recovery.sh base.json + SHAPE_JOBS=12 OMP_NUM_THREADS=1 ./run_shape_recovery.sh pr.json + python compare_shape_results.py base.json pr.json # exit 1 = merge-blocking + +Notes for agents: +- The suite is self-contained: it runs against ANY checkout via PYTHONPATH + (the two runs above use the SAME suite files against different checkouts). +- Run one suite at a time: LDG head nodes have RLIMIT_NPROC=500. +- CPU-only by design (CUDA_VISIBLE_DEVICES="") — this also exercises the + cupy-installed-but-no-GPU worker configuration that has repeatedly bitten + production (module-level cupy selection without a device probe). +- "STARVED" rows (n_eff < 100) are not absolute failures — high-D mixtures + legitimately exhaust production budgets — but base-healthy -> starved IS a + blocking regression. +- If your change is behind an opt-in flag, the default-path gate will show + bitwise-identical results; you must ALSO probe the flag ON (use + shape_recovery.py as a library; see its docstring). diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/gaussian_mixture_model.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/gaussian_mixture_model.py index 05fdf1874..83216613e 100755 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/gaussian_mixture_model.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/gaussian_mixture_model.py @@ -17,12 +17,19 @@ try: import cupy import cupyx.scipy.special + # cupy imports cleanly on GPU-less nodes (shared install, or a GPU node + # with CUDA_VISIBLE_DEVICES masked); probe for an actual device before + # selecting the GPU backend, else every cupy kernel launch dies at call + # time with cudaErrorNoDevice. getDeviceCount raises CUDARuntimeError + # (not ImportError) in that case, hence the broad except below. + if cupy.cuda.runtime.getDeviceCount() == 0: + raise ImportError("cupy installed but no CUDA device available") xpy_default = cupy xpy_special_default = cupyx.scipy.special identity_convert = cupy.asnumpy identity_convert_togpu = cupy.asarray cupy_ok = True -except ImportError: +except Exception: xpy_default = np xpy_special_default = None # scipy.special is used via scipy if needed identity_convert = lambda x: x @@ -66,6 +73,7 @@ def mvnun(lower, upper, mean, cov, maxpts=None, abseps=1e-5, releps=1e-5): from scipy.special import logsumexp from . import multivariate_truncnorm as truncnorm import itertools +import math def _xpy_logsumexp(a, axis=None): @@ -87,21 +95,6 @@ def _xpy_logsumexp(a, axis=None): return logsumexp(a, axis=axis) -# Symmetric (Hermitian) eigen-routines. cupy.linalg only provides the Hermitian -# variants (eigh/eigvalsh), not the general eig/eigvals. The matrices fed to -# _near_psd below are covariance/correlation matrices and hence symmetric, so -# the Hermitian routines are both correct and the only ones available on GPU. -if cupy_ok: - _xpy_eigvals = cupy.linalg.eigvalsh - _xpy_eig = cupy.linalg.eigh -else: - # Symmetric routines on CPU as well: the inputs are covariance/correlation - # matrices. eigvalsh/eigh are faster, return real eigenvalues (no spurious - # complex output from round-off asymmetry), and match the GPU path. - _xpy_eigvals = np.linalg.eigvalsh - _xpy_eig = np.linalg.eigh - - def _near_psd_impl(x, epsilon, xpy): ''' Shared, hardened nearest-PSD projection for covariance matrices. @@ -124,12 +117,17 @@ def _near_psd_impl(x, epsilon, xpy): floor = xpy.maximum(diag, epsilon) x = x + xpy.diag(floor - diag) x = 0.5 * (x + x.T) # symmetrize: eigh assumes it, round-off breaks it + # Symmetric (Hermitian) eigen-routines, resolved through the CALLER's xpy: + # cupy.linalg only provides eigh/eigvalsh (not general eig/eigvals), and the + # inputs here are covariance/correlation matrices, so the Hermitian variants + # are correct on both backends. Do not bind these at import time -- that is + # how a cupy install without a GPU broke every CPU refit (cudaErrorNoDevice). for _ in range(10): # bounded: the legacy `while True` could spin forever var_list = xpy.sqrt(xpy.diag(x)) y = x / (var_list[:, None] * var_list[None, :]) - if bool(xpy.min(_xpy_eigvals(y)) > epsilon): + if bool(xpy.min(xpy.linalg.eigvalsh(y)) > epsilon): return x - eigval, eigvec = _xpy_eig(y) + eigval, eigvec = xpy.linalg.eigh(y) val_psd = xpy.maximum(eigval, epsilon) near_corr = eigvec @ xpy.diag(val_psd) @ eigvec.T near_cov = near_corr * (var_list[:, None] * var_list[None, :]) @@ -426,28 +424,81 @@ def fit(self, sample_array, log_sample_weights=None): self.p_nk = model.p_nk self.log_prob = model.log_prob + def num_free_params(self): + '''Number of free parameters of this k-component d-dim mixture: + k means (k*d) + k covariances (k*d*(d+1)/2) + (k-1) mixture weights.''' + d = self.d + return self.k*d + self.k*(d*(d+1))//2 + (self.k - 1) + + def prune_components(self, weight_floor=1e-3, min_keep=1): + '''Drop mixture components whose weight falls below weight_floor and + renormalize. Over-allocated components collapse to ~zero weight under + EM; removing them (a) prevents a spurious sharp component from dominating + the importance weights and (b) cuts score() cost, which is O(k) in the + per-component mvnun box normalization. Keeps at least max(1, min_keep) + components (the highest-weight ones) -- pass min_keep to preserve a safety + floor. No-op if nothing is below the floor.''' + min_keep = max(1, int(min_keep)) + w = np.asarray(self.identity_convert(self.weights), dtype=float) + keep = np.where(w >= weight_floor)[0] + if len(keep) < min_keep: + # keep the min_keep highest-weight components + keep = np.argsort(w)[::-1][:min(min_keep, self.k)] + if len(keep) == self.k: + return + keep = np.sort(keep) + self.means = [self.means[i] for i in keep] + self.covariances = [self.covariances[i] for i in keep] + w_keep = w[keep] + w_keep = w_keep / w_keep.sum() + self.weights = self.identity_convert_togpu(w_keep) + self.adapt = [self.adapt[i] for i in keep] if isinstance(self.adapt, list) else self.adapt + self.k = len(keep) + def _match_components(self, new_model): ''' Match components in new model to those in current model by minimizing the - net Mahalanobis between all pairs of components + net Mahalanobis between all pairs of components. + + The objective is a SUM of per-pair distances, so the optimal old->new + assignment is a linear assignment problem, solved exactly in O(k^3) by + the Hungarian algorithm. The legacy implementation enumerated all k! + permutations (itertools.permutations), which is fine for k<=6 but + explodes (8!=40320, 12!~5e8, 16!~2e13) -- it made any many-component + proposal (e.g. a chain of small Gaussians wrapping a curved degeneracy + arc) impossible to refit through update(). linear_sum_assignment + returns the SAME optimum (identical additive objective); only tie-break + ordering can differ. Returns a tuple `order` with order[i]=j meaning + old component i is matched to new component j. ''' - orders = list(itertools.permutations(list(range(self.k)), self.k)) - distances = np.empty(len(orders)) - index = 0 - for order in orders: - dist = 0 - i = 0 - for j in order: - # These are likely small vectors, stay on CPU - diff = self.identity_convert(new_model.means[j]) - self.identity_convert(self.means[i]) - cov_inv = np.linalg.inv(self.identity_convert(self.covariances[i])) - temp_cov_inv = np.linalg.inv(self.identity_convert(new_model.covariances[j])) - dist += np.sqrt(np.dot(np.dot(diff, cov_inv), diff)) - dist += np.sqrt(np.dot(np.dot(diff, temp_cov_inv), diff)) - i += 1 - distances[index] = dist - index += 1 - return orders[np.argmin(distances)] + k = self.k + # cost[i,j] = mahalanobis(new_j - old_i) under old_i cov + under new_j cov + cost = np.empty((k, k)) + old_means = [self.identity_convert(m) for m in self.means] + new_means = [self.identity_convert(m) for m in new_model.means] + old_cov_inv = [np.linalg.inv(self.identity_convert(c)) for c in self.covariances] + new_cov_inv = [np.linalg.inv(self.identity_convert(c)) for c in new_model.covariances] + for i in range(k): + for j in range(k): + diff = new_means[j] - old_means[i] + cost[i, j] = np.sqrt(np.dot(np.dot(diff, old_cov_inv[i]), diff)) \ + + np.sqrt(np.dot(np.dot(diff, new_cov_inv[j]), diff)) + try: + from scipy.optimize import linear_sum_assignment + row_ind, col_ind = linear_sum_assignment(cost) + # row_ind is sorted 0..k-1, so col_ind[i] is the new index for old i + return tuple(int(j) for j in col_ind) + except Exception: + # Defensive fallback (should not trigger: scipy.optimize is a hard + # RIFT dependency). Greedy nearest assignment, O(k^2 log k). + order = [None] * k + used = set() + for i in np.argsort(cost.min(axis=1)): + j = int(min((jj for jj in range(k) if jj not in used), + key=lambda jj: cost[i, jj])) + order[i] = j + used.add(j) + return tuple(order) def _merge(self, new_model, M): ''' @@ -494,10 +545,38 @@ def _near_psd(self, x): ''' return _near_psd_impl(x, self.epsilon, self.xpy) + def _strip_defensive_component(self): + """Detach the defensive component (always appended last) and renormalize the rest.""" + if self.k <= 1: + return 0.0 + dfrac = float(getattr(self, 'defensive_frac', 0.0) or 0.0) + w = np.asarray(self.identity_convert(self.weights), dtype=float)[:-1] + means = [self.identity_convert(m) for m in self.means][:-1] + covs = [self.identity_convert(c) for c in self.covariances][:-1] + s = w.sum() + w = w / s if s > 0 else np.ones(len(w)) / max(len(w), 1) + self.means = [self.identity_convert_togpu(m) for m in means] + self.covariances = [self.identity_convert_togpu(c) for c in covs] + self.weights = self.identity_convert_togpu(w) + if isinstance(self.adapt, list): + self.adapt = list(self.adapt)[:-1] + self.k = len(means) + self.defensive_frac = 0.0 + return dfrac + def update(self, sample_array, log_sample_weights=None): ''' Updates the model with new data without doing a full retraining. ''' + # PROTECT THE DEFENSIVE COMPONENT. _merge() blends component i of this model with + # component order[i] of the freshly fitted one for every i in range(self.k); it does NOT + # consult self.adapt. So the broad box-covering component -- marked adapt=False by + # add_defensive_component precisely so it would be left alone -- was dragged toward the + # fitted cloud on every update: its mean, covariance and weight drifted while + # defensive_frac stayed set, so has_unbounded_support kept reporting coverage that no + # longer existed. Detach it, update the real components, then reinstate it. + _dfrac = self._strip_defensive_component() if ( + getattr(self, 'defensive_frac', 0.0) or 0.0) > 0 else 0.0 # halve the covariance regularizer but FLOOR it: an unbounded decay # (the legacy behavior) eventually leaves sharp refits unregularized self.tempering_coeff = max(self.tempering_coeff / 2, 1e-12) @@ -516,6 +595,8 @@ def update(self, sample_array, log_sample_weights=None): M, _ = sample_array.shape self._merge(new_model, M) self.N += M + if _dfrac > 0: + add_defensive_component(self, defensive_frac=_dfrac) def score(self, sample_array,assume_normalized=True): ''' @@ -641,3 +722,175 @@ def print_params(self): print(weight, '\n') else: print(i, weight, self._unnormalize(np.array([mean]))[0,0], mean[0], np.sqrt(cov[0,0])) + + +def _mixture_log_density_normalized(model, Xn): + '''Log mixture density (n,) of a fitted `gmm` at NORMALIZED samples Xn (n,d), + in the model's normalized [-1,1] coordinate frame. Backend-portable.''' + xpy = model.xpy + n = Xn.shape[0] + logk = xpy.empty((n, model.k)) + for j in range(model.k): + mean = model.means[j] + cov = model.covariances[j] + if cupy_ok: + lp = gpu_logpdf(Xn, mean, cov, xpy) + else: + lp = multivariate_normal.logpdf(x=model.identity_convert(Xn), + mean=model.identity_convert(mean), + cov=model.identity_convert(cov), + allow_singular=True) + logk[:, j] = lp + xpy.log(model.weights[j]) + return _xpy_logsumexp(logk, axis=1) + + +def add_defensive_component(model, defensive_frac=0.05, width_norm=1.0): + '''Append a broad, box-covering "defensive" component to a fitted mixture so + the proposal has heavy enough tails for importance sampling. + + This is the single most important fix for the SNR~82 extrinsic posterior: a + mixture fit to the (tight) elite cloud UNDER-COVERS the broad, degenerate + directions (distance-inclination), so the importance weight L*p/q blows up on + the rare draw that lands in a poorly-covered high-likelihood pocket and the + effective sample size collapses to ~1. A defensive component (Hesterberg + 1995) with weight `defensive_frac`, wide in the model's normalized [-1,1] + frame, bounds the weights: q >= defensive_frac * q_broad everywhere, so no + single sample can dominate. The AV sampler gets the same guarantee from its + cover-fraction floor; the fitted GMM had none. + + width_norm is the std of the defensive Gaussian in normalized coords (1.0 ~ + covers the whole [-1,1] box; truncated to the box it is near-uniform). + ''' + if not defensive_frac or defensive_frac <= 0: + model.defensive_frac = 0.0 + return model + xpy = model.xpy + d = model.d + w = np.asarray(model.identity_convert(model.weights), dtype=float) + means = [model.identity_convert(m) for m in model.means] + covs = [model.identity_convert(c) for c in model.covariances] + means.append(np.zeros(d)) # box center (normalized) + covs.append((width_norm ** 2) * np.eye(d)) # broad, box-covering + w = np.concatenate([w * (1.0 - defensive_frac), [defensive_frac]]) + model.means = [model.identity_convert_togpu(m) for m in means] + model.covariances = [model.identity_convert_togpu(c) for c in covs] + model.weights = model.identity_convert_togpu(w / w.sum()) + model.adapt = list(model.adapt) + [False] if isinstance(model.adapt, list) else model.adapt + model.k = len(means) + # MARKER: the portfolio must be able to VERIFY this component is installed rather than + # infer it from a config value -- gmm_defensive_frac>0 was being read as a guarantee + # while the fixed-component fit paths never called this function. + model.defensive_frac = float(defensive_frac) + return model + + +def fit_gmm_adaptive(sample_array, bounds, log_sample_weights=None, k_max=8, + k_min=1, k_candidates=None, epsilon=None, tempering_coeff=1e-8, + prune_weight_floor=1e-3, defensive_frac=0.05, inflate=1.0): + '''Fit a GMM whose COMPONENT COUNT is chosen from the data by BIC, then + prune near-zero-weight components. Data-driven replacement for a hard-coded + per-group component count. + + Rationale (measured on the S250114ax extrinsic posterior, SNR~82): + * A fixed SMALL k (e.g. the correlate-all default of 2) cannot wrap a + curved distance-inclination degeneracy arc: the elite fit is one broad + Gaussian over the ridge, the proposal never locks onto the peak, and + the honest effective sample size stays ~1. + * A fixed LARGE k is both statistically fragile (a spurious sharp + component collapses onto ~1 elite sample and dominates the importance + weights) and computationally costly (score() does an O(k) per-component + mvnun box normalization on the CPU). + BIC threads between the two: fit k over a ladder, penalize free parameters + by ln(N_eff), keep the best, and drop dead components. It allocates more + components only where the (importance-weighted) cloud is genuinely + non-Gaussian and stays at k=1 for a single blob. + + SAFETY FLOOR: k is chosen in [k_min, k_max], and pruning never drops below + k_min. Pass k_min = the stress-tested hard-coded per-group count so opting + into adaptive can only ADD components where the data earns them, never fewer + than the layout that was validated for the primary ILE use case (e.g. a broad + multi-modal sky keeps its default components even if the INITIAL elite cloud + -- fit before the proposal has explored every mode -- looks single-peaked). + + Parameters + ---------- + sample_array : (N, d) array in ORIGINAL coordinates. + bounds : (d, 2) array of [llim, rlim] per dimension (as gmm expects). + log_sample_weights : (N,) importance/elite log-weights (default: equal). + k_max : cap on the number of components. + k_min : floor on the number of components (default 1); the stress- + tested hard-coded count when used as a refinement layer. + k_candidates : explicit ladder (overrides k_max/k_min-derived ladder). + prune_weight_floor : components below this mixture weight are removed (but + never below k_min). + + Returns a fitted `gmm`. + ''' + xpy = xpy_default + N, d = sample_array.shape + if log_sample_weights is None: + log_sample_weights = xpy.zeros(N) + # Kish effective sample size of the fit weights (drives both the BIC penalty + # and the per-component sample-count cap). + lw = xpy.where(xpy.isfinite(log_sample_weights), log_sample_weights, + -xpy.inf * xpy.ones(N)) + lw_max = xpy.max(lw) + if not bool(xpy.isfinite(lw_max)): + wn = xpy.ones(N) + else: + wn = xpy.exp(lw - lw_max) + wn = xpy.where(xpy.isfinite(wn), wn, xpy.zeros(N)) + sw = xpy.sum(wn) + if not bool(sw > 0): + wn = xpy.ones(N); sw = float(N) + wn = wn / sw + N_eff = float(1.0 / xpy.sum(wn ** 2)) + + k_min = max(1, int(k_min)) + k_max = max(k_min, int(k_max)) + if k_candidates is None: + base = [1, 2, 3, 4, 6, 8, 12, 16, 24, 32] + k_candidates = [k for k in base if k_min <= k <= k_max] + for kk in (k_min, k_max): # always evaluate the endpoints + if int(kk) not in k_candidates: + k_candidates.append(int(kk)) + # cap k so each component retains >~ (d+2) effective samples (an EM stability + # floor mirroring estimator._m_step's ESS>=d+1 guard) -- but never below the + # safety floor k_min (the stress-tested count), even if the initial elite + # cloud is small. + k_cap = max(k_min, int(N_eff // max(d + 2, 4))) + k_candidates = sorted(set(int(k) for k in k_candidates if k_min <= k <= max(k_min, k_cap))) + if not k_candidates: + k_candidates = [k_min] + + ln_Neff = math.log(max(N_eff, 2.0)) + wn_scaled = N_eff * wn # effective-count weights (sum to N_eff) + best, best_bic = None, None + for k in k_candidates: + try: + model = gmm(k, bounds, epsilon=epsilon, tempering_coeff=tempering_coeff) + model.fit(sample_array, log_sample_weights=log_sample_weights) + logmix = _mixture_log_density_normalized(model, model._normalize(sample_array)) + wll = float(xpy.sum(wn_scaled * logmix)) # weighted log-likelihood + bic = -2.0 * wll + model.num_free_params() * ln_Neff + except Exception: + continue + if best_bic is None or bic < best_bic: + best, best_bic = model, bic + if best is None: # every candidate failed: fall back to the floor count + best = gmm(k_min, bounds, epsilon=epsilon, tempering_coeff=tempering_coeff) + best.fit(sample_array, log_sample_weights=log_sample_weights) + if prune_weight_floor: + # never prune below the safety floor: the extra components carry the + # capacity to capture modes the merge adaptation discovers later. + best.prune_components(prune_weight_floor, min_keep=k_min) + if inflate and inflate != 1.0: + # widen every fitted component so the proposal has heavier tails than the + # (tight) elite cloud it was fit to -- a basic importance-sampling + # requirement the raw EM fit violates on a peaked/degenerate posterior. + fac = float(inflate) ** 2 + best.covariances = [best.identity_convert_togpu(fac * best.identity_convert(c)) + for c in best.covariances] + if defensive_frac: + add_defensive_component(best, defensive_frac=defensive_frac) + return best diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py index 822605d7b..323c12533 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsampler.py @@ -8,7 +8,7 @@ import numpy from RIFT.precision import RiftFloat # platform-portable replacement for np.float128 from scipy import integrate, interpolate -from ..integrators.statutils import cumvar, welford, update, finalize +from ..integrators.statutils import cumvar, welford, update, finalize, pareto_khat_from_log, ess_from_log_weights, block_scatter_sigma, bootstrap_lnZ_quantiles import itertools import functools @@ -470,6 +470,9 @@ def integrate(self, func, *args, **kwargs): maxlnL = -float("Inf") eff_samp = 0 mean, var = None, RiftFloat(0) # to prevent infinite variance due to overflow + # per-chunk lnZ record for the between-chunk error floor (each chunk used a + # different adapted proposal; the pooled variance cannot see their scatter) + lnZ_chunk_list = []; n_chunk_list = [] if bShowEvaluationLog: print("iteration Neff sqrt(2*lnLmax) sqrt(2*lnLmarg) ln(Z/Lmax) int_var") @@ -565,7 +568,7 @@ def integrate(self, func, *args, **kwargs): self._rvs["integrand"] = numpy.hstack( (self._rvs["integrand"], fval) ) self._rvs["joint_prior"] = numpy.hstack( (self._rvs["joint_prior"], joint_p_prior) ) self._rvs["joint_s_prior"] = numpy.hstack( (self._rvs["joint_s_prior"], joint_p_s) ) - self._rvs["weights"] = numpy.hstack( (self._rvs["joint_s_prior"], fval*joint_p_prior/joint_p_s) ) + self._rvs["weights"] = numpy.hstack( (self._rvs["weights"], fval*joint_p_prior/joint_p_s) ) # BUGFIX: was appending onto joint_s_prior, corrupting the weights record else: self._rvs["integrand"] = fval self._rvs["joint_prior"] = joint_p_prior @@ -607,6 +610,11 @@ def integrate(self, func, *args, **kwargs): var = outvals[-1] # running integral (note also in current_aggregate) int_val1 += int_val.sum() + # per-chunk lnZ for the between-chunk error floor (log first: RiftFloat-safe) + try: + lnZ_chunk_list.append(float(numpy.log(int_val.sum())) - numpy.log(n)); n_chunk_list.append(n) + except Exception: + pass # running number of evaluations self.ntotal += n # FIXME: Likely redundant with int_val1 @@ -743,6 +751,32 @@ def integrate(self, func, *args, **kwargs): else: self._rvs[key] = self._rvs[key][indx_list] + # MC-error diagnostics (before the fairdraw resampling rewrites _rvs). + # The pooled weight variance is 1/ESS restated and tail-blind; disclose the + # tail (Pareto k-hat), the between-chunk scatter, and -- when the naive + # relative error is already large -- bootstrap lnZ quantiles. + mc_diag = {} + try: + _sb = block_scatter_sigma(lnZ_chunk_list, n_chunk_list) + if _sb is not None: + mc_diag['sigma_lnZ_block'] = _sb + if "integrand" in self._rvs and len(self._rvs["integrand"]) > 0: + # log first (RiftFloat-safe), cast after + _lw_diag = numpy.asarray(numpy.log(self._rvs["integrand"]) + numpy.log(self._rvs["joint_prior"]) - numpy.log(self._rvs["joint_s_prior"]), dtype=float) + _kh = pareto_khat_from_log(_lw_diag) + if _kh is not None: + mc_diag['pareto_khat'] = _kh + mc_diag['n_ESS'] = ess_from_log_weights(_lw_diag) + _sig_rel_naive = numpy.inf + if int_val1 > 0 and self.ntotal > 1: + _sig_rel_naive = float(numpy.sqrt(var*self.ntotal)/int_val1) + if _sig_rel_naive > 0.3 or mc_diag.get('sigma_lnZ_block', 0) > 0.3: + _q = bootstrap_lnZ_quantiles(_lw_diag, n_total=self.ntotal) + if _q is not None: + mc_diag['lnZ_ci90'] = _q + except Exception as _e_diag: + print(" mcsampler: MC-error diagnostics failed ({}); continuing.".format(_e_diag), file=sys.stderr) + # Do a fair draw of points, if option is set if bFairdraw and not(n_extr is None): n_extr = int(numpy.min([n_extr,1.5*eff_samp,1.5*neff])) @@ -762,6 +796,7 @@ def integrate(self, func, *args, **kwargs): dict_return ={} if convergence_tests is not None: dict_return["convergence_test_results"] = last_convergence_test + dict_return.update(mc_diag) # MC-error diagnostics (pareto_khat, n_ESS, sigma_lnZ_block, lnZ_ci90) return int_val1/self.ntotal, var/self.ntotal, eff_samp, dict_return diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py index ec1a9da73..e582dac55 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerAdaptiveVolume.py @@ -91,7 +91,7 @@ def profile(fn): except: print(" - No healpy - ") -from RIFT.integrators.statutils import update,finalize, init_log,update_log,finalize_log +from RIFT.integrators.statutils import update,finalize, init_log,update_log,finalize_log, pareto_khat_from_log, ess_from_log_weights, bootstrap_lnZ_quantiles #from multiprocessing import Pool @@ -139,10 +139,24 @@ def get_likelihood_threshold(lkl, lkl_thr, nsel, discard_prob,xpy_here=xpy_defau return identity_convert(lkl_thr), identity_convert(truncp) # send both to CPU as needed def sample_from_bins(xrange, dx, bu, ninbin, reject_out_of_range=False): - + # Draw uniformly within each occupied hypercube bin. VECTORIZED: the old + # implementation looped over bins in Python (a list comprehension + vstack + # over one entry per bin), which is O(n_bins) per chunk and becomes the + # bottleneck once the live volume is finely resolved -- e.g. a concentrated + # warm start from a full PE posterior can seed thousands of bins. Here we + # instead repeat each bin's lower corner ninbin[k] times and add a single + # (N, ndim) uniform draw, so cost is O(N) with no Python-level bin loop. ndim = xrange.shape[0] - xlo, xhi = xrange.T[0] + dx * bu, xrange.T[0] + dx * (bu+1) - x = xpy_default.vstack([xpy_default.random.uniform(xlo[kk], xhi[kk], size = (npb, ndim)) for kk, npb in enumerate(ninbin)]) + # per-bin lower corners + the point->bin expansion are done on the HOST + # (np.repeat with an int-array of counts is reliable everywhere; cupy.repeat + # with array repeats is version-fragile), then the single uniform draw is on + # the active backend so the output matches the previous cupy behaviour. + bu_h = identity_convert(bu); dx_h = np.asarray(identity_convert(dx)) + xlo_h = np.asarray(identity_convert(xrange)).T[0] + dx_h * np.asarray(bu_h) # (n_bins, ndim) + reps = np.asarray(identity_convert(ninbin)).astype(int) + lo_per_point = np.repeat(xlo_h, reps, axis=0) # host (N, ndim) + N = lo_per_point.shape[0] + x = xpy_default.asarray(lo_per_point) + xpy_default.asarray(dx_h) * xpy_default.random.uniform(0.0, 1.0, size=(N, ndim)) # remove points that are out of range. Due to rounding issues etc, the sampler above can generate points out of range! # Note this rejection will bias the integral, because volumes are calculated assuming a regular grid. We *should* fix the grid sizes to integers if reject_out_of_range: @@ -153,6 +167,11 @@ def sample_from_bins(xrange, dx, bu, ninbin, reject_out_of_range=False): class MCSampler(object): + # COMPACT SUPPORT: this sampler's density is EXACTLY ZERO outside its contracted live volume, + # so once seeded or contracted it cannot serve as the mixture's coverage guarantee. + # mcsamplerPortfolio reads this to decide whether it must hold one member cold. + has_unbounded_support = False + """ Class to define a set of parameter names, limits, and probability densities. """ @@ -218,7 +237,14 @@ def __init__(self,n_chunk=400000,**kwargs): # sampling tool self.V=None # fractional volume self.delta_V=None # fractional volume - + self._warm=None # bootstrap/warm-start live-volume state (see bootstrap_from_*) + self._warm_applied=False # has _warm been installed into the ACTIVE grid? (see _apply_warm_state) + # Opt-in ANISOTROPIC bin allocation: give each axis a different number of bins + # (fine where the live points cluster tightly -- phase/pol/sky; coarse where they are + # broad -- distance/inclination), instead of the default equal split. Keeps the same + # total bin budget (prod(nbins)=1/delta_V) so the estimator is unchanged. Default off. + self.anisotropic_bins = False + def setup(self, **kwargs): ndim = len(self.params) @@ -312,12 +338,65 @@ def prior_prod(self, x): return p_out + def _apply_warm_state(self): + """Install a seeded live-volume grid (self._warm) into the ACTIVE sampling state. + + `bootstrap_from_samples` / `bootstrap_from_oracle` / `load_state` only STORE the seeded grid + in `self._warm`; historically it was installed only inside `integrate_log`. Anything that + drives this sampler WITHOUT calling its own integrate_log -- above all a PORTFOLIO, which + calls draw_simplified()/update_sampling_prior() directly, and the driver's L0 auto-rescue, + which re-seeds a portfolio then re-runs -- therefore kept drawing from the COLD grid while + reporting that it had been warm-started. Idempotent; safe to call on every draw. + """ + warm = getattr(self, '_warm', None) + if warm is None or getattr(self, '_warm_applied', False): + return + try: + self.binunique = np.array(warm['binunique']) + self.dx = np.array(warm['dx']) + self.nbins = np.array(warm['nbins']) + self.ninbin = ((self.n_chunk // self.binunique.shape[0] + 1) + * np.ones(self.binunique.shape[0])).astype(int) + if 'V' in warm: + self.V = float(warm['V']) + if 'loglkl_thr' in warm: + self.lnL_thresh = float(warm['loglkl_thr']) + self._warm_applied = True + print(" [AV warm-start] seeded grid APPLIED to the active draw path: " + "live bins={}".format(self.binunique.shape[0])) + except Exception as e: + # never let a malformed seed break sampling: fall back to the cold grid + print(" [AV warm-start] could not apply seeded grid ({}); continuing cold".format(e)) + self._warm_applied = True + def draw_simplified(self,n_to_get, *args, **kwargs): + # Self-contained cold start. A PORTFOLIO (mcsamplerPortfolio) drives draw_simplified on + # its members directly, WITHOUT running each member's own integrate()/setup(), so a cold + # AV member may not have its live-volume grid (my_ranges/dx/binunique/ninbin) built yet + # -> AttributeError on self.my_ranges. + # Build the cold full-box grid on first use so AV works as a portfolio member cold or warm. + if getattr(self, 'my_ranges', None) is None: + self.setup() + # ... and if a seed was supplied, INSTALL it: setup() above (and the driver, which calls + # setup BEFORE bootstrap_from_samples) leaves the active grid cold, so without this a + # warm-started portfolio member draws from the cold grid. + self._apply_warm_state() rv, log_p = self.draw_simple() - p = np.exp(log_p)[:n_to_get] + # Subsample RANDOMLY, never a head slice: sample_from_bins emits points + # grouped in lexicographic bin order (binunique from np.unique), so + # rv[:n_to_get] returns only the first ~n_to_get/ninbin bins of the live + # volume while sampling_density (hence the portfolio's q_mix) claims + # uniform coverage of ALL occupied bins. In any multi-member portfolio + # that mismatch systematically biased the recovered shape (pulls up to + # 0.6 sigma at d2; random subsample collapses them to ~1e-3). + n_have = len(rv) + if n_to_get < n_have: + keep = np.random.choice(n_have, size=int(n_to_get), replace=False) + rv = rv[keep] + log_p = log_p[keep] + p = np.exp(log_p) ps = self.xpy.ones(len(p))*self.V_s/self.V # sampling prior, full hypercube normalized to 1 - ps = ps[:n_to_get] - rv = rv[:n_to_get].T + rv = rv.T return ps, p, rv def draw_simple(self): @@ -329,11 +408,99 @@ def draw_simple(self): indx_p = self.params_ordered.index(p) x[:,indx_p] = self.params_pinned_vals[p] - # probabilities at these points. + # probabilities at these points. log_p = np.log(self.prior_prod(x)) # Not including any sampling prior factors, since it is de facto uniform right now (just discarding 'irrelevant' regions) return x, log_p + def sampling_density(self, X): + """Pointwise sampling density q(theta) of THIS member, evaluated at + ARBITRARY points X (shape (N, ndim), columns in self.params_ordered + order). Returns a host (numpy) array of length N, or None if the + live-volume state has not been set up yet. + + VARAHA draws uniformly over its live volume -- the union of the + currently-occupied hypercubes (self.binunique), each of width self.dx. + The density is therefore the SAME constant this sampler reports in + integrate_log's log_joint_s_prior, + + q_live = 1 / (n_occupied_bins * prod(dx)) (== 1/(V*prod(dx0)) + for VARAHA's geometric V), + + inside the live volume and 0 outside it. We use the geometric form + 1/(n_bins*prod(dx)) directly: it is *exactly* the density of the points + draw_simple() produces (equal draws per occupied bin, uniform within a + bin), so a multiple-importance-sampling denominator built from it is + unbiased regardless of any drift between the tracked scalar V and the + actual bin grid. + + This method is READ-ONLY -- it does not touch any sampler state and does + not affect this sampler's own integrate_log. It exists so the portfolio + can form the balance-heuristic mixture density q_mix = sum_m w_m q_m. + """ + binunique = getattr(self, 'binunique', None) + dx = getattr(self, 'dx', None) + if binunique is None or dx is None or not hasattr(self, 'my_ranges'): + return None + X = np.atleast_2d(np.asarray(identity_convert(X), dtype=float)) + ndim = len(self.params_ordered) + if X.shape[1] != ndim and X.shape[0] == ndim: + X = X.T # tolerate (ndim, N) + box_lo = self.my_ranges.T[0] + box_hi = self.my_ranges.T[1] + dx = np.asarray(identity_convert(dx), dtype=float) + bins = np.asarray(identity_convert(binunique)).astype(np.int64) + n_bins = bins.shape[0] + if n_bins == 0: + return np.zeros(X.shape[0], dtype=float) + # bin index of each point (same floor((x-lo)/dx) mapping the sampler uses + # in integrate_log to build binidx), then test membership in the occupied + # set. Points outside the box floor out of range and are excluded below. + binidx = np.floor((X - box_lo) / dx).astype(np.int64) + binset = set(map(tuple, bins.tolist())) + inside = np.array([tuple(row) in binset for row in binidx], dtype=bool) + inside &= np.all((X >= box_lo) & (X <= box_hi), axis=1) + q_live = 1.0 / (float(n_bins) * float(np.prod(dx))) + q = np.zeros(X.shape[0], dtype=float) + q[inside] = q_live + return q + + def _allocate_nbins(self, live_pts, delta_V, ndim): + """Per-axis bin counts whose product over adaptive dims equals 1/delta_V (the same + total resolution the isotropic split uses, so the volume V=n_bins*prod(dx) and hence + the estimator are unchanged). + + Default (self.anisotropic_bins False): equal split -- nbins_i = (1/delta_V)**(1/d). + Anisotropic: redistribute that SAME total bin budget by each axis's *compressibility* + c_i = log(range_i / spread_i), where spread_i is the std of the live points on axis i. + Axes whose points fill only a small fraction of their range (tight: phase, polarization, + sky) get many bins (fine); broad axes (distance, inclination) get few (coarse). This + lets the live hypercube wrap a correlated/degenerate posterior far more tightly than an + isotropic grid, which must use one resolution for both the narrow and the broad axes.""" + nbins = np.ones(ndim) + if self.d_adaptive <= 0: + return nbins + adaptive = np.ones(ndim, dtype=bool) + if len(self.indx_not_adaptive): + adaptive[np.array(self.indx_not_adaptive, dtype=int)] = False + total_log = -np.log(max(float(delta_V), 1e-300)) # log(1/delta_V): total log-bins to spread + n_live = 0 if live_pts is None else len(live_pts) + if (not self.anisotropic_bins) or n_live < max(8, 2 * self.d_adaptive): + nbins[adaptive] = np.exp(total_log / self.d_adaptive) # isotropic fallback + return nbins + lp = np.asarray(identity_convert(live_pts)) + rng = np.diff(self.my_ranges, axis=1).flatten() # range per axis + spread = lp.std(axis=0) + spread = np.maximum(spread, 1e-6 * np.maximum(rng, 1e-30)) + c = np.clip(np.log(np.maximum(rng, 1e-30) / spread), 0.0, None) # compressibility + c[~adaptive] = 0.0 + csum = c[adaptive].sum() + if csum <= 0: + nbins[adaptive] = np.exp(total_log / self.d_adaptive) # degenerate -> isotropic + else: + nbins[adaptive] = np.exp(c[adaptive] / csum * total_log) # prod(adaptive)=1/delta_V + return nbins + def update_sampling_prior_selfish(self, lnF, *args, xpy=xpy_default,no_protect_names=True,**kwargs): """ update_sampling_prior @@ -386,15 +553,29 @@ def update_sampling_prior_selfish(self, lnF, *args, xpy=xpy_default,no_protect_n # For now: no prior, just duplicate VT algorithm log_integrand =lnL + log_joint_p_prior - + loglkl = log_integrand # note we are putting the prior in here - idxsel = xpy_here.where(loglkl > loglkl_thr) + # admit only FINITE samples above threshold: a cold portfolio member can draw points + # whose loglkl is -inf/NaN (out-of-support / degenerate extrinsic config). With the + # initial threshold -1e15 the plain "> thr" test then passes NaN (breaking later maxes) + # or, if ALL are non-finite, yields an empty set -> the reported crash chain + # (get_likelihood_threshold max of empty array; then this method's max at line ~532). + idxsel = xpy_here.where(xpy_here.logical_and(loglkl > loglkl_thr, xpy_here.isfinite(loglkl))) #only admit samples that lie inside the live volume, i.e. one that cross likelihood threshold allx = xpy_here.append(allx, rv[idxsel], axis = 0) allloglkl = xpy_here.append(allloglkl, loglkl[idxsel]) allp = xpy_here.append(allp, log_joint_p_prior[idxsel]) ninj = len(allloglkl) + if ninj == 0: + # Nothing finite in the live volume this step (cold portfolio member / degenerate + # draw). Leave V and the grid UNCHANGED rather than crashing on empty-array + # reductions downstream. The portfolio's other members carry this step; a later + # draw with finite samples lets AV resume training. (This method is a SINGLE + # selfish step, so an early return is correct -- unlike integrate_log's loop.) + print(" [AV selfish-update] no finite in-volume samples this step; live volume unchanged") + self.V = V + return #just some test to verify if we dont discard more than 1 - Pthr probability @@ -410,6 +591,11 @@ def update_sampling_prior_selfish(self, lnF, *args, xpy=xpy_default,no_protect_n allp = allp[idxsel] allx = allx[idxsel] nrec = len(allloglkl) # recovered size of active volume at present, after selection + if nrec == 0: + # threshold selected nothing (degenerate all-equal finite draws): leave the + # live volume unchanged instead of crashing on max()/divide-by-zero below. + self.V = V + return # Weights lw = allloglkl - xpy_here.max(allloglkl) @@ -424,7 +610,8 @@ def update_sampling_prior_selfish(self, lnF, *args, xpy=xpy_default,no_protect_n # Redefine bin sizes, reassign points to redefined hypercube set. [Asymptotically this becomes stationary] # Note hypercube calculation is on CPU at present, always if self.d_adaptive > 0: - self.nbins = np.ones(ndim)*(1/delta_V) ** (1/self.d_adaptive) # uniform split in each dimension is normal, but we have array - can be irregular + # per-axis (anisotropic) or equal (default) split; same total bin budget either way + self.nbins = self._allocate_nbins(allx, delta_V, ndim) self.nbins[self.indx_not_adaptive] = 1 # reset to 1 bin for non-adaptive dimensions else: self.nbins = np.ones(ndim) # why are we even doing this! @@ -443,7 +630,311 @@ def update_sampling_prior_selfish(self, lnF, *args, xpy=xpy_default,no_protect_n self.V = V self.delta_V = delta_V - + + + ### + ### BOOTSTRAP / WARM-START SUPPORT + ### + # The VARAHA algorithm normally starts every integrate_log() call cold: one + # bin spanning the whole box, threshold -1e15, fractional volume V=1, and + # spends its first several chunks carving the live volume down from the full + # prior. In production (repeated ILE instances, successive CIP iterations, + # or events with a known Fisher matrix) we already know roughly where the + # posterior lives, so that carving is wasted work -- worst in high dimension. + # + # These methods seed the live-volume state (`self._warm`) from prior + # information; integrate_log() then starts from that concentrated grid. The + # seeded fractional volume is set GEOMETRICALLY (n_occupied_bins / prod(nbins)) + # so the final integral normalization (log_joint_s_prior = log(1/V) - sum log dx0) + # stays unbiased regardless of how the state was produced. + + def _order_columns(self, samples, params=None): + """Return samples as an (M, ndim) array whose columns are in + self.params_ordered order. `params` names the columns of `samples`; + if None the caller guarantees they are already in order.""" + X = np.atleast_2d(np.asarray(samples, dtype=float)) + if X.shape[1] != len(self.params_ordered) and X.shape[0] == len(self.params_ordered): + X = X.T # tolerate (ndim, M) + if params is None: + return X + out = np.empty((X.shape[0], len(self.params_ordered))) + for j, p in enumerate(self.params_ordered): + out[:, j] = X[:, list(params).index(p)] + return out + + def _build_grid_from_points(self, pts, loglkl=None, enc_prob=0.999, dilate=1, + resolution_pts=None): + """Build a VARAHA live-volume grid (binunique, dx, nbins) and a + geometrically-consistent fractional volume V from points that populate + the high-likelihood region. Mirrors the bin-refinement block of + integrate_log() so a warm start lands on the same kind of grid the cold + algorithm would have converged to. + + `dilate` (>=0): grow the occupied-bin set by this many axis-neighbor + layers along the adaptive dimensions. This is a SAFETY margin: VARAHA's + live volume only ever *contracts*, so a warm start that seeded a grid + tighter than the true support could never recover the missing region and + would bias the integral low. Dilating guarantees the seed is a superset + of the sampled support (at a small efficiency cost the first few chunks + then trim away).""" + ndim = len(self.params_ordered) + pts = np.atleast_2d(np.asarray(pts, dtype=float)) + box_lo = self.my_ranges.T[0] + box_hi = self.my_ranges.T[1] + inside = np.all((pts >= box_lo) & (pts <= box_hi), axis=1) + pts = pts[inside] + if loglkl is not None: + loglkl = np.asarray(loglkl, dtype=float)[inside] + nrec = len(pts) + if nrec < 2: + raise ValueError("AV bootstrap needs >=2 in-box reference points (got {})".format(nrec)) + box = box_hi - box_lo + # Bin RESOLUTION (nbins) is set from the CONCENTRATED core, not the full + # cloud: when a coverage floor (cover_frac) adds uniform full-box points, + # they must not coarsen the grid to a single bin per dim (which collapses + # V to 1 and throws away the seed's concentration). resolution_pts is the + # core (the actual proposal, without the uniform floor); coverage points + # then land in scattered fine bins that still guarantee coverage. + res_pts = pts if resolution_pts is None else np.atleast_2d(np.asarray(resolution_pts, dtype=float)) + n_res = max(len(res_pts), 2) + lo = np.quantile(res_pts, 0.5 * (1 - enc_prob), axis=0) + hi = np.quantile(res_pts, 1 - 0.5 * (1 - enc_prob), axis=0) + ext = np.clip(hi - lo, box * 1e-6, None) + V_extent = float(np.prod(ext / box)) + # VARAHA bin count: nbins = (1/delta_V)^(1/d_adaptive), delta_V = V/sqrt(nrec) + delta_V = V_extent / np.sqrt(n_res) + if self.d_adaptive > 0: + # per-axis (anisotropic) or equal (default) split of the warm-seed grid + nbins = self._allocate_nbins(res_pts, delta_V, ndim) + nbins[self.indx_not_adaptive] = 1 + else: + nbins = np.ones(ndim) + nbins = np.maximum(np.floor(nbins), 1) + dx = box / nbins + # CLIP bin indices to [0, nbins-1]: a point exactly on the upper box edge + # maps to binidx == nbins (out of range), which would put out-of-range bins + # in binunique -> V = n_bins/prod(nbins) can exceed 1 (an invalid fractional + # volume) and draw_simple would sample outside the box. This bites hardest + # for a WIDE seed (e.g. a full PE posterior + cover_frac spanning the box). + nb_int = np.maximum(nbins.astype(np.int64), 1) + binidx = np.clip(((pts - box_lo) / dx).astype(np.int64), 0, nb_int - 1) + binunique = np.unique(binidx, axis=0) + # SAFETY dilation: grow occupied bins by axis-neighbor layers along the + # adaptive dims, clipped to [0, nbins-1]. Uses a bounded 2*d_adaptive + # neighborhood per layer (not the full 3^d) so the volume grows linearly. + if dilate and self.d_adaptive > 0: + bins = set(map(tuple, binunique.tolist())) + nb_max = nbins.astype(int) + for _ in range(int(dilate)): + grown = set(bins) + for b in bins: + for ax in self.indx_adaptive: + for step in (-1, 1): + nb = list(b); nb[ax] += step + if 0 <= nb[ax] < nb_max[ax]: + grown.add(tuple(nb)) + bins = grown + binunique = np.array(sorted(bins)) + # fractional volume ACTUALLY sampled = occupied bins / total bins + V = float(binunique.shape[0] / np.prod(nbins)) + # seed the threshold just below the reference support so the first chunk + # keeps the seeded region; if no lnL given, let integrate_log recompute it + # (the concentrated grid already delivers the efficiency win). + loglkl_thr = -1e15 if loglkl is None else float(np.min(loglkl)) + return dict(binunique=binunique, dx=dx, nbins=nbins, V=V, + loglkl_thr=loglkl_thr, trunc_p=1e-10) + + def bootstrap_from_samples(self, samples, params=None, loglkl=None, enc_prob=0.999, + cover_frac=0.0, dilate=1, inflate=1.0, seed=None): + """Warm-start from an explicit set of reference points populating the + high-likelihood region (e.g. a previous run's posterior draws, a puff of + an earlier MAP point, or fair-draw samples from a prior ILE instance). + `loglkl` (optional) is L*prior at those points, used to seed the threshold. + + `cover_frac` (0..1) mixes this fraction of uniform full-box points into the + seed cloud, widening the seeded live volume. Leave 0 when reusing a proposal + for the SAME problem (e.g. an in-run second pass). + + IT IS NOT A COVERAGE GUARANTEE, despite what this docstring claimed until + 2026-08. A FINITE set of uniform points occupies only the bins it lands in, + so the seeded grid is NOT a superset of a cold (uniform) start -- and the + shortfall grows fast with dimension. Measured, fraction of the [-5,5]^d prior + box covered by the seeded grid (a cold start is 1.0 by construction): + + cover_frac: 0.0 0.2 0.5 0.9 + d=2 0.027 0.634 0.982 1.000 + d=4 0.0015 0.028 0.104 0.620 + d=6 6.3e-05 0.00087 0.0033 0.0287 + + At d=6 even cover_frac=0.9 leaves 97% of the box unsampled. So the claim that + "a warm-started integral can never be MORE biased than a cold one" was false; + do not rely on it. + + WHAT ACTUALLY PROTECTS YOU is having a component with support everywhere. In + the default AV+GMM portfolio that is the GMM member: a Gaussian mixture has + nonzero density over the whole box, so q_mix never vanishes and a badly-seeded + AV member costs efficiency rather than bias (measured with a deliberately + displaced seed at d=4 and d=6: |lnZ bias| <= 0.05 in every run). An ALL-AV + portfolio has no such member -- every component is a hard-edged box -- and the + same displaced seed gave lnZ bias -1.0 to -6.8 nats. + + And note the limit of any coverage fix: in that all-AV test, keeping one member + fully cold (V=1) still gave -1.1 to -4.2 nats, because a uniform member at d=6 + finds a sharp peak too rarely to carry the integral within the budget (n_eff + 3-9). Coverage in principle is necessary, not sufficient. A badly mismatched + seed is an efficiency catastrophe no knob repairs -- detect it (the L0 rescue) + or do not warm-start across dissimilar points. + + `inflate` (>=1) is the HANDOFF SAFETY MARGIN: widen the seed cloud by this + factor about its own mean before building the grid. When importing a proposal + that came from a *neighbouring* intrinsic point, the true peak is shifted (and + often slightly broader) at this point, so an un-inflated seed may sit just off + it; inflate>1 (e.g. 1.5-2) gives margin for that shift while staying far tighter + than a cold start. cover_frac is the coarse safety net for gross mismatch; + inflate is the fine margin for a modest shift.""" + if not hasattr(self, 'my_ranges'): + self.setup() + X = self._order_columns(samples, params) + inflate = float(max(inflate, 1.0)) + if inflate > 1.0 and len(X) >= 2: + _m = np.mean(X, axis=0) + X = _m + inflate * (X - _m) + X = np.clip(X, self.my_ranges.T[0], self.my_ranges.T[1]) + loglkl = None # inflated points no longer carry their original lnL + cover_frac = float(np.clip(cover_frac, 0.0, 1.0)) + _core = X # the concentrated proposal; sets the grid RESOLUTION + if cover_frac > 0: + rng = np.random.RandomState(seed) + n_cover = max(int(cover_frac / (1.0 - cover_frac) * len(X)), 1) + Xc = rng.uniform(self.my_ranges.T[0], self.my_ranges.T[1], + size=(n_cover, len(self.params_ordered))) + X = np.vstack([X, Xc]) + # the cover points are not part of the high-L region, so drop the + # lnL-threshold seed (let integrate_log recompute it from the data) + loglkl = None + # resolution from the core (not the uniform cover), so the coverage floor + # cannot coarsen away the proposal's concentration (a wide PE + cover_frac + # would otherwise collapse the grid to one bin per dim, V->1) + self._warm = self._build_grid_from_points(X, loglkl=loglkl, enc_prob=enc_prob, + dilate=dilate, resolution_pts=_core) + self._warm_applied = False # a NEW seed must be re-installed (L0 rescue re-seeds mid-run) + return self._warm + + def bootstrap_from_gaussian(self, mean, cov, n=None, params=None, enc_prob=0.999, + seed=None, cover_frac=0.0, dilate=1): + """Warm-start from a single Gaussian proposal N(mean, cov) -- the + Fisher-oracle entry point. Draws `n` points from the (box-clipped) + Gaussian and builds the live-volume grid from them. + + IMPORTANT -- this is a UNIMODAL seed. A single Gaussian covers only one + mode, and because VARAHA's live volume only ever contracts, any mode the + seed misses is lost forever and biases the integral low. Use this only + when the target is (locally) unimodal -- e.g. a Fisher matrix at the MAP. + For known multimodal structure use bootstrap_from_gaussian_mixture(); for + an empirical proposal use bootstrap_from_samples() (both cover the full + support and stay unbiased). + + `cover_frac` (0..1): optional safety valve -- fraction of the seed cloud + drawn uniformly from the full box, trading efficiency for coverage on a + possibly-misspecified seed. Default 0.""" + if not hasattr(self, 'my_ranges'): + self.setup() + rng = np.random.RandomState(seed) + mean = np.asarray(mean, dtype=float) + cov = np.atleast_2d(np.asarray(cov, dtype=float)) + if params is not None: + order = [list(params).index(p) for p in self.params_ordered] + mean = mean[order] + cov = cov[np.ix_(order, order)] + n = int(n or self.n_chunk) + n_cover = int(np.clip(cover_frac, 0.0, 1.0) * n) + n_gauss = n - n_cover + X = rng.multivariate_normal(mean, cov, size=n_gauss) + X = np.clip(X, self.my_ranges.T[0], self.my_ranges.T[1]) + if n_cover > 0: + Xc = rng.uniform(self.my_ranges.T[0], self.my_ranges.T[1], + size=(n_cover, len(self.params_ordered))) + X = np.vstack([X, Xc]) + self._warm = self._build_grid_from_points(X, enc_prob=enc_prob, dilate=dilate) + self._warm_applied = False # a NEW seed must be re-installed (L0 rescue re-seeds mid-run) + return self._warm + + def bootstrap_from_fisher(self, mean, fisher, **kwargs): + """Warm-start from a Fisher matrix (mean, Gamma): cov = Gamma^{-1}. + This is the 'Fisher-matrix oracle' -- an essentially free substitute for + an expensively-trained flow, giving the integrator a correct-to-2nd-order + starting proposal.""" + cov = np.linalg.inv(np.atleast_2d(np.asarray(fisher, dtype=float))) + return self.bootstrap_from_gaussian(mean, cov, **kwargs) + + def bootstrap_from_gaussian_mixture(self, means, covs, weights=None, n=None, + params=None, enc_prob=0.999, seed=None, dilate=1): + """Warm-start from a MIXTURE of Gaussians -- the general oracle seed for + multimodal targets. This is what a flow oracle, a GMM fit of a previous + posterior, or a set of known degenerate modes (e.g. sky reflections) + provides. Because the seed cloud covers every component, the resulting + live volume is a superset of the support and the integral stays + unbiased.""" + if not hasattr(self, 'my_ranges'): + self.setup() + rng = np.random.RandomState(seed) + means = [np.asarray(m, dtype=float) for m in means] + covs = [np.atleast_2d(np.asarray(c, dtype=float)) for c in covs] + k = len(means) + weights = np.ones(k) / k if weights is None else np.asarray(weights, float) / np.sum(weights) + if params is not None: + order = [list(params).index(p) for p in self.params_ordered] + means = [m[order] for m in means] + covs = [c[np.ix_(order, order)] for c in covs] + n = int(n or self.n_chunk) + counts = rng.multinomial(n, weights) + chunks = [] + for c, m, cov in zip(counts, means, covs): + if c > 0: + chunks.append(rng.multivariate_normal(m, cov, size=c)) + X = np.vstack(chunks) + X = np.clip(X, self.my_ranges.T[0], self.my_ranges.T[1]) + self._warm = self._build_grid_from_points(X, enc_prob=enc_prob, dilate=dilate) + self._warm_applied = False # a NEW seed must be re-installed (L0 rescue re-seeds mid-run) + return self._warm + + def save_state(self, path): + """Serialize the compact live-volume state (occupied bins + widths + + volume + threshold) to a lightweight .npz. This is RIFT's cheap + alternative to persisting a trained flow: the entire adapted proposal is + just an integer bin-index array plus a few scalars.""" + warm = getattr(self, '_warm', None) + if warm is None: + thr = float(self.lnL_thresh) if np.isfinite(self.lnL_thresh) else -1e15 + warm = dict(binunique=self.binunique, dx=self.dx, nbins=self.nbins, + V=float(self.V), loglkl_thr=thr, trunc_p=1e-10) + np.savez(path, + params=np.array([str(p) for p in self.params_ordered]), + llim=self.my_ranges.T[0], rlim=self.my_ranges.T[1], + binunique=warm['binunique'], dx=warm['dx'], nbins=warm['nbins'], + V=warm['V'], loglkl_thr=warm['loglkl_thr'], + trunc_p=warm.get('trunc_p', 1e-10)) + return path + + def load_state(self, path): + """Restore a live-volume state saved by save_state(). Verifies the + parameter names and box match this sampler before warm-starting.""" + if not hasattr(self, 'my_ranges'): + self.setup() + d = np.load(path, allow_pickle=True) + saved_params = [str(p) for p in d['params']] + if saved_params != [str(p) for p in self.params_ordered]: + raise ValueError("saved state params {} != sampler params {}".format( + saved_params, [str(p) for p in self.params_ordered])) + if not (np.allclose(d['llim'], self.my_ranges.T[0]) and + np.allclose(d['rlim'], self.my_ranges.T[1])): + raise ValueError("saved state box does not match sampler box") + self._warm_applied = False # new seed -> must be re-installed + self._warm = dict(binunique=np.array(d['binunique']), dx=np.array(d['dx']), + nbins=np.array(d['nbins']), V=float(d['V']), + loglkl_thr=float(d['loglkl_thr']), trunc_p=float(d['trunc_p'])) + return self._warm + @profile def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): @@ -514,6 +1005,9 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): save_intg = kwargs["save_intg"] if "save_intg" in kwargs else False + # opt-in anisotropic (per-axis) bin allocation; also settable as a sampler attribute + if "anisotropic_bins" in kwargs: + self.anisotropic_bins = bool(kwargs["anisotropic_bins"]) # FIXME: The adaptive step relies on the _rvs cache, so this has to be # on in order to work if n_adapt > 0 and tempering_exp > 0.0: @@ -562,6 +1056,27 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): trunc_p = 1e-10 #How much probability analysis removes with evolution nsel = 1000# number of largest log-likelihood samples selected to estimate lkl_thr for the next cycle. nsel = np.min([nsel, int(0.1*self.n_chunk)]) # if chunk size is small, don't pick too many points + + # WARM START: if this sampler was bootstrapped (bootstrap_from_* / + # load_state), override the cold single-bin grid, fractional volume and + # threshold with the seeded live-volume state. self.setup() above has + # already reset these to cold defaults, so we re-apply the seed here. + warm = getattr(self, '_warm', None) + if warm is not None: + self.binunique = np.array(warm['binunique']) + self.dx = np.array(warm['dx']) + self.nbins = np.array(warm['nbins']) + self.ninbin = ((self.n_chunk // self.binunique.shape[0] + 1) * np.ones(self.binunique.shape[0])).astype(int) + V = float(warm['V']) + loglkl_thr = float(warm['loglkl_thr']) + trunc_p = float(warm.get('trunc_p', 1e-10)) + if bShowEvaluationLog: + print(" [AV warm-start] live bins={} V={:.3e} loglkl_thr={:.3g}".format( + self.binunique.shape[0], V, loglkl_thr)) + + var_lnV = 0.0 # accumulated variance of ln(V): V is a stochastic product of per-cycle + # binomial survival fractions, and Z ~ V*mean(w), so Var(lnV) is a + # component of the lnZ error the weight variance is structurally blind to if cupy_ok: allx = identity_convert_togpu(allx) allloglkl = identity_convert_togpu(allloglkl) @@ -575,17 +1090,26 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): rv = identity_convert_togpu(rv) # send random numbers to GPU : ugh log_joint_p_prior = identity_convert_togpu(log_joint_p_prior) # send to GPU if required. Don't waste memory reassignment otherwise - # Evaluate function, protecting argument order. The user integrand is - # a host function in general, so feed it CPU samples; lnL is pushed - # back to the active backend just below. identity_convert is a no-op - # without cupy. - rv_cpu = identity_convert(rv) - if 'no_protect_names' in kwargs: - unpacked0 = rv_cpu.T - lnL = lnF(*unpacked0) # do not protect order + # Evaluate the integrand. Two contracts exist: the production + # GPU/vectorized ILE likelihood is DEVICE-native (wants cupy arrays), + # while synthetic/host integrands (CI tests, benchmarks) want numpy. + # Feed the native (device) array -- matching production -- but if a + # host-only integrand chokes on a cupy array, fall back to a host copy + # and remember the choice for the rest of the run. (The previous + # version always fed a host copy, which silently broke the real GPU + # ILE likelihood with 'Unsupported type numpy.ndarray'.) + def _eval_integrand(samples): + if 'no_protect_names' in kwargs: + return lnF(*samples.T) + return lnF(**dict(list(zip(self.params_ordered, samples.T)))) + if getattr(self, '_integrand_wants_host', False): + lnL = _eval_integrand(identity_convert(rv)) else: - unpacked = dict(list(zip(self.params_ordered,rv_cpu.T))) - lnL= lnF(**unpacked) # protect order using dictionary + try: + lnL = _eval_integrand(rv) + except (TypeError, ValueError): + self._integrand_wants_host = True + lnL = _eval_integrand(identity_convert(rv)) # take log if we are NOT using lnL if cupy_ok: if not(isinstance(lnL,cupy.ndarray)): @@ -638,7 +1162,8 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # Redefine bin sizes, reassign points to redefined hypercube set. [Asymptotically this becomes stationary] # Note hypercube calculation is on CPU at present, always if self.d_adaptive > 0: - self.nbins = np.ones(ndim)*(1/delta_V) ** (1/self.d_adaptive) # uniform split in each dimension is normal, but we have array - can be irregular + # per-axis (anisotropic) or equal (default) split; same total bin budget either way + self.nbins = self._allocate_nbins(allx, delta_V, ndim) self.nbins[self.indx_not_adaptive] = 1 # reset to 1 bin for non-adaptive dimensions else: self.nbins = np.ones(ndim) # why are we even doing this! @@ -653,6 +1178,12 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): self.binunique = np.unique(binidx, axis = 0) self.ninbin = ((self.n_chunk // self.binunique.shape[0] + 1) * np.ones(self.binunique.shape[0])).astype(int) self.ntotal = current_log_aggregate[0] + # accumulate the binomial variance of this cycle's ln(V) update: + # Var(ln p_hat) ~= (1-p_hat)/(n p_hat) = (1-nrec/ninj)/nrec. Cycles reuse + # surviving samples, so this is an approximate (disclosed) budget rather + # than a rigorous iid propagation; it vanishes as the volume stabilizes. + if nrec > 0 and ninj > 0: + var_lnV += (1.0 - nrec/ninj)/nrec if super_verbose: print(ntotal_true,eff_samp, np.round(neff_varaha), np.round(np.max(allloglkl), 1), len(allloglkl), np.mean(self.nbins), V, len(self.binunique), np.round(loglkl_thr, 1), trunc_p) @@ -681,7 +1212,13 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): log_wt = self._rvs["log_integrand"] + self._rvs["log_joint_prior"] - self._rvs["log_joint_s_prior"] log_wt = identity_convert(log_wt) # convert to CPU log_int = special.logsumexp( log_wt) - np.log(len(log_wt)) # mean value - rel_var = np.var( np.exp(log_wt - log_int))/len(log_wt) # error in integral, estimated: just taking int = , so error is V(w_k)/N (sample mean/variance) + rel_var_mc = np.var( np.exp(log_wt - log_int))/len(log_wt) # error in integral, estimated: just taking int = , so error is V(w_k)/N (sample mean/variance) + # Total DISCLOSED relative variance: the naive weight-variance term above is + # structurally blind to (a) the stochasticity of the live volume V itself + # (Z ~ V*mean(w); var_lnV accumulated per cycle) and (b) the probability + # deliberately truncated by the likelihood threshold (trunc_p, a one-sided + # systematic entered here as a variance in quadrature). Add them. + rel_var = rel_var_mc + var_lnV + trunc_p**2 eff_samp = np.sum(np.exp(log_wt - np.max(log_wt))) maxval = np.max(allloglkl) # max of log @@ -718,6 +1255,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 = {} + # MC-error diagnostics: disclose the components and the weight-tail state. + # NOTE the AV estimator assigns the surviving (threshold-selected) samples a + # pretend-uniform density on the final live volume, so the naive term is if + # anything MORE optimistic than for the other samplers -- k-hat matters here. + try: + mc_diag = {'sigma_lnZ_mc': float(np.sqrt(rel_var_mc)), + 'sigma_lnV': float(np.sqrt(var_lnV)), + 'trunc_p': float(trunc_p)} + _kh = pareto_khat_from_log(log_wt) + if _kh is not None: + mc_diag['pareto_khat'] = _kh + mc_diag['n_ESS'] = ess_from_log_weights(log_wt) + if np.sqrt(rel_var) > 0.3: + _q = bootstrap_lnZ_quantiles(log_wt, n_total=len(log_wt)) + if _q is not None: + mc_diag['lnZ_ci90'] = _q + dict_return.update(mc_diag) + print(" [AV mc diag] sigma_mc={:.4f} sigma_lnV={:.4f} trunc_p={:.2e} khat={} ESS={}".format( + mc_diag['sigma_lnZ_mc'], mc_diag['sigma_lnV'], mc_diag['trunc_p'], + round(mc_diag['pareto_khat'],3) if 'pareto_khat' in mc_diag else None, + round(mc_diag['n_ESS'],1) if 'n_ESS' in mc_diag else None)) + except Exception as _e_diag: + print(" mcsamplerAdaptiveVolume: MC-error diagnostics failed ({}); continuing.".format(_e_diag)) return log_int, np.log(rel_var) +2*log_int, eff_samp, dict_return # if outvals: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerEnsemble.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerEnsemble.py index ac433f43a..780cbd1b6 100755 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerEnsemble.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerEnsemble.py @@ -9,12 +9,17 @@ try: import cupy import cupyx.scipy.special + # Probe for an actual device: cupy imports cleanly on GPU-less nodes but + # every kernel launch then dies with cudaErrorNoDevice. getDeviceCount + # raises CUDARuntimeError (not ImportError), hence the broad except. + if cupy.cuda.runtime.getDeviceCount() == 0: + raise ImportError("cupy installed but no CUDA device available") xpy_default = cupy xpy_special_default = cupyx.scipy.special identity_convert = cupy.asnumpy identity_convert_togpu = cupy.asarray cupy_ok = True -except ImportError: +except Exception: xpy_default = np xpy_special_default = None identity_convert = lambda x: x @@ -47,6 +52,51 @@ def __str__(self): return repr(self.value) class MCSampler(object): + + @property + def has_unbounded_support(self): + """Does this member's proposal genuinely have support across the WHOLE prior box? + + mcsamplerPortfolio uses this to decide whether it must hold a member cold on a warm start. + Getting it wrong costs coverage silently, so it is answered from the INSTALLED MODELS, not + from a configuration value. + + Two traps this avoids: + * `gmm_defensive_frac > 0` is only a REQUEST. add_defensive_component() is called by + fit_gmm_adaptive, but the fixed-component fit paths did not call it, and gmm_adaptive + defaults to None (off) -- so the default configuration asked for a defensive component + and never installed one. + * gmm.score() FLOORS its return at 1e-300, so a member always looks like it has nonzero + density everywhere. That is a numerical guard against log(0), not coverage: a sample + landing there carries weight L*p/q ~ 1e300 and would wreck the estimate rather than + support it. Measured, a fixed-component fit to a tight cloud returns exactly that + floor at the far corner for every d >= 4. + + Reports False whenever it cannot be verified -- before the integrator exists, before any + group has been trained, or if ANY trained group lacks the component. + """ + integ = getattr(self, 'integrator', None) + if integ is None: + return False + if not (getattr(integ, 'gmm_defensive_frac', 0.0) or 0.0) > 0: + return False + models = [m for m in getattr(integ, 'gmm_dict', {}).values() if m is not None] + if not models: + if not getattr(integ, 'gmm_defensive_all_paths', False) and not getattr( + integ, 'gmm_adaptive', None): + # Neither path that installs the component is active: the request in + # gmm_defensive_frac will not be honoured, so do not promise coverage. + return False + # UNTRAINED. The portfolio has to decide about warm-starting before any group is + # fitted, so there is nothing to inspect yet. Trusting the config is justified only + # because EVERY fit path now installs the component (fit_gmm_adaptive did already; + # the fixed-component paths in this file and in MonteCarloEnsemble were fixed at the + # same time as this check). test_every_fit_path_installs_the_defensive_component + # pins that invariant -- if a new fit path is added without it, that test fails rather + # than this property silently over-promising again. + return True + return all((getattr(m, 'defensive_frac', 0.0) or 0.0) > 0 for m in models) + """ Class to define a set of parameter names, limits, and probability densities. """ @@ -169,7 +219,46 @@ def calc_pdf(self, samples): temp_ret *= pdf_vals.reshape( temp_ret.shape) return temp_ret - def setup(self,n_comp=None,**kwargs): + def setup(self, n_comp=None, **kwargs): + """Build the integrator. REMEMBERS its arguments and re-applies them on later calls. + + setup() is not called once: bootstrap_from_samples() re-runs it to rebuild the proposal as + a single full-dim group, and mcsamplerPortfolio replays it to reset a member between + points. Each of those rebuilt the integrator from ONLY the kwargs of that call, so every + option the caller set originally was silently dropped. That has now bitten three separate + settings -- gmm_dict (the dimension grouping and any seeded models), gmm_defensive_frac, + and gmm_defensive_all_paths -- each found as its own P1, each patched individually, and a + fourth would have followed. + + So: merge this call's kwargs OVER the remembered ones, and remember the result. An + explicit argument still wins; an omitted one keeps whatever it was configured to be + instead of reverting to a library default. Pass `setup_forget=True` to start clean. + """ + _prev = dict(getattr(self, '_setup_kwargs_seen', {}) or {}) + if kwargs.pop('setup_forget', False): + _prev = {} + if n_comp is None: + n_comp = _prev.get('n_comp', None) + merged = dict(_prev) + merged.update(kwargs) + merged['n_comp'] = n_comp + self._setup_kwargs_seen = dict(merged) + kwargs = dict(merged) + kwargs.pop('n_comp', None) + return self._setup_impl(n_comp=n_comp, **kwargs) + + def _setup_impl(self, n_comp=None, **kwargs): + # n_comp=None silently disabled ALL training downstream: the integrator + # stores it verbatim and update_sampling_prior only builds a model for + # int!=0 or dict n_comp, so every gmm_dict entry stayed None forever. In + # default portfolio wiring (setup() forwarded without GMM args) the GMM + # member therefore never trained and "portfolio" ran as AV-only, with no + # error (2026-07-22 shape-gate probe). Default to a single component and + # say so; n_comp=0 remains the explicit off-switch. + if n_comp is None: + print(" mcsamplerEnsemble: setup() called without n_comp; defaulting n_comp=1 " + "(n_comp=None previously disabled GMM training silently; pass n_comp=0 to disable adaptation)") + n_comp = 1 integrator_func = kwargs['integrator_func'] if "integrator_func" in kwargs else None mcsamp_func = kwargs['mcsamp_func'] if "mcsamp_func" in kwargs else None proc_count = kwargs['proc_count'] if "proc_count" in kwargs else None @@ -180,6 +269,10 @@ def setup(self,n_comp=None,**kwargs): write_to_file = kwargs['write_to_file'] if "write_to_file" in kwargs else False correlate_all_dims = kwargs['correlate_all_dims'] if "correlate_all_dims" in kwargs else False gmm_adapt = kwargs['gmm_adapt'] if "gmm_adapt" in kwargs else None + gmm_adaptive = kwargs['gmm_adaptive'] if "gmm_adaptive" in kwargs else None + gmm_defensive_frac = kwargs['gmm_defensive_frac'] if "gmm_defensive_frac" in kwargs else 0.05 + _defensive_all = kwargs['gmm_defensive_all_paths'] if "gmm_defensive_all_paths" in kwargs else False + gmm_inflate = kwargs['gmm_inflate'] if "gmm_inflate" in kwargs else 1.0 gmm_epsilon = kwargs['gmm_epsilon'] if "gmm_epsilon" in kwargs else None L_cutoff = kwargs["L_cutoff"] if "L_cutoff" in kwargs else None tempering_exp = kwargs["tempering_exp"] if "tempering_exp" in kwargs else 1.0 @@ -229,7 +322,9 @@ def setup(self,n_comp=None,**kwargs): self.integrator = monte_carlo.integrator(dim, bounds, gmm_dict, n_comp, n=self.n, prior=self.calc_pdf, user_func=integrator_func, proc_count=proc_count,L_cutoff=L_cutoff,gmm_adapt=gmm_adapt,gmm_epsilon=gmm_epsilon,tempering_exp=tempering_exp, - tempering_adapt=tempering_adapt, ess_target=ess_target, ess_floor=ess_floor) + tempering_adapt=tempering_adapt, ess_target=ess_target, ess_floor=ess_floor, gmm_adaptive=gmm_adaptive, + gmm_defensive_frac=gmm_defensive_frac, gmm_inflate=gmm_inflate) + self.integrator.gmm_defensive_all_paths = bool(_defensive_all) def update_sampling_prior(self,ln_weights, n_history,tempering_exp=1,log_scale_weights=True,floor_integrated_probability=0,external_rvs=None,**kwargs): rvs_here = self._rvs @@ -241,11 +336,20 @@ def update_sampling_prior(self,ln_weights, n_history,tempering_exp=1,log_scale_w gmm_dict = self.integrator.gmm_dict - n_history_to_use = self.xpy.min([n_history, len(ln_weights), len(rvs_here[self.params_ordered[0]])] ) - + # These are all host ints; use the Python builtin min (self.xpy.min([list]) + # crashes on cupy -- "'list' object has no attribute 'min'" -- the same + # backend-min-of-a-list bug fixed in integrate()'s fairdraw block). A host + # int is also required for the [-n_history_to_use:] slices just below. + n_history_to_use = int(min(n_history, len(ln_weights), len(rvs_here[self.params_ordered[0]]))) + + # external_rvs (e.g. the portfolio's host history) may be host numpy while + # sample_array lives on the active backend (cupy on GPU); assigning a host + # slice into a cupy row raises "non-scalar numpy.ndarray cannot be used for + # fill". Convert each slice to the backend first so this method is + # backend-consistent (previously it only worked on CPU). sample_array = self.xpy.empty( (len(self.params_ordered), n_history_to_use)) for indx, p in enumerate(self.params_ordered): - sample_array[indx] = rvs_here[p][-n_history_to_use:] + sample_array[indx] = self.identity_convert_togpu(rvs_here[p][-n_history_to_use:]) sample_array = sample_array.T for dim_group in gmm_dict: @@ -255,29 +359,161 @@ def update_sampling_prior(self,ln_weights, n_history,tempering_exp=1,log_scale_w continue new_bounds = self.xpy.empty((len(dim_group), 2)) new_bounds = self.integrator.bounds[dim_group] + # per-dimension (uncorrelated) groups: setup() hands the integrator raw + # (dim,2) array bounds, so bounds[(i,)] is a bare (2,) row; GMM.fit + # needs (n_dims,2). Same up-shape guard as _sample()/q-scoring. + # (Latent until now: the n_comp=None bug meant this line was never + # reached in the default portfolio configuration.) + if len(new_bounds.shape) < 2: + new_bounds = self.xpy.array([new_bounds]) model = self.integrator.gmm_dict[dim_group] temp_samples = self.xpy.empty((n_history_to_use, len(dim_group))) index = 0 for dim in dim_group: - temp_samples[:,index] = self.identity_convert(sample_array[:,dim]) + # keep on the active backend: temp_samples and sample_array are + # both self.xpy arrays, and the GMM model.fit/update below runs on + # self.xpy. (The old identity_convert here forced a host array + # into a cupy column -> the same fill error as above on GPU.) + temp_samples[:,index] = sample_array[:,dim] index += 1 + # Drop NaN-weight samples before fitting. NOTE: filter into LOOP-LOCAL names. This used + # to reassign `ln_weights` itself, which is loop-INVARIANT (built once, before the loop + # over dim_groups): the first group with any NaN shrank it (e.g. 10000 -> 8686), and every + # LATER group then rebuilt temp_samples at full n_history_to_use but reused the stale, + # shorter weights -> "boolean index did not match indexed array" inside GMM.update / + # GMM.fit. Only reachable when weights actually contain NaN, i.e. a degenerate/cold pass, + # which is why warm runs never hit it and cold portfolio starts died on chunk ~8. if self.xpy.any(self.xpy.isnan(ln_weights)): ok_indx = ~self.xpy.isnan(ln_weights) - temp_samples = temp_samples[ok_indx] - ln_weights = ln_weights[ok_indx] + temp_samples = temp_samples[ok_indx] # rebuilt each iteration: safe to filter + ln_weights_group = ln_weights[ok_indx] # loop-LOCAL: never touch ln_weights itself + else: + ln_weights_group = ln_weights + # Data-driven component count (matches integrator._train): scalar or + # per-group gmm_adaptive picks k by BIC at init, floored at the + # stress-tested n_comp, then the merge path below adapts. This is the + # path the PORTFOLIO drives its GMM member through (update_sampling_prior). + adaptive_kmax = None + _ga = getattr(self.integrator, 'gmm_adaptive', None) + if _ga: + if isinstance(_ga, dict): + adaptive_kmax = _ga.get(dim_group) + elif isinstance(_ga, bool): + adaptive_kmax = 8 + else: + adaptive_kmax = int(_ga) if model is None: - if isinstance(self.integrator.n_comp, int) and self.integrator.n_comp != 0: + if adaptive_kmax: + if isinstance(self.integrator.n_comp, dict): + k_floor = self.integrator.n_comp.get(dim_group, 1) + else: + k_floor = self.integrator.n_comp + k_floor = int(k_floor) if isinstance(k_floor, int) and k_floor > 0 else 1 + model = GMM.fit_gmm_adaptive(temp_samples, new_bounds, + log_sample_weights=ln_weights_group, + k_max=max(int(adaptive_kmax), k_floor), + k_min=k_floor, + epsilon=self.integrator.gmm_epsilon, + defensive_frac=getattr(self.integrator,'gmm_defensive_frac',0.0), + inflate=getattr(self.integrator,'gmm_inflate',1.0)) + elif isinstance(self.integrator.n_comp, int) and self.integrator.n_comp != 0: model = GMM.gmm(self.integrator.n_comp, new_bounds,epsilon=self.integrator.gmm_epsilon) - model.fit(temp_samples, log_sample_weights=ln_weights) + model.fit(temp_samples, log_sample_weights=ln_weights_group) + # The defensive component is the ONLY thing that actually guarantees this member + # has support across the box -- gmm.score() merely FLOORS at 1e-300, which is a + # numerical guard, not coverage (a sample there would carry weight ~1e300). + # fit_gmm_adaptive adds it; the fixed-component path did not. OPT-IN, because + # measured on the shape gate a 5% broad component costs real n_eff in + # higher dimensions (d6_n3_s303 119->75, d8_n1_s303 448->210): it spends + # 5% of draws where the likelihood is negligible. Only a consumer that + # NEEDS this member as its coverage guarantee should pay -- so a + # portfolio sets gmm_defensive_all_paths on its members, and a standalone + # GMM user is unaffected. + GMM.add_defensive_component(model, defensive_frac=( + getattr(self.integrator,'gmm_defensive_frac',0.0) + if getattr(self.integrator,'gmm_defensive_all_paths',False) else 0.0)) elif isinstance(self.integrator.n_comp, dict) and self.integrator.n_comp[dim_group] != 0: model = GMM.gmm(self.integrator.n_comp[dim_group], new_bounds,epsilon=self.integrator.gmm_epsilon) - model.fit(temp_samples, log_sample_weights=ln_weights) + model.fit(temp_samples, log_sample_weights=ln_weights_group) + # The defensive component is the ONLY thing that actually guarantees this member + # has support across the box -- gmm.score() merely FLOORS at 1e-300, which is a + # numerical guard, not coverage (a sample there would carry weight ~1e300). + # fit_gmm_adaptive adds it; the fixed-component path did not. OPT-IN, because + # measured on the shape gate a 5% broad component costs real n_eff in + # higher dimensions (d6_n3_s303 119->75, d8_n1_s303 448->210): it spends + # 5% of draws where the likelihood is negligible. Only a consumer that + # NEEDS this member as its coverage guarantee should pay -- so a + # portfolio sets gmm_defensive_all_paths on its members, and a standalone + # GMM user is unaffected. + GMM.add_defensive_component(model, defensive_frac=( + getattr(self.integrator,'gmm_defensive_frac',0.0) + if getattr(self.integrator,'gmm_defensive_all_paths',False) else 0.0)) + elif not (self.integrator.n_comp == 0 or + (isinstance(self.integrator.n_comp, dict) and self.integrator.n_comp.get(dim_group) == 0)): + # invalid n_comp (e.g. None from an integrator built outside + # setup()): never no-op silently -- that hid a dead GMM + # portfolio member in production. n_comp==0 is the only + # sanctioned way to skip training. + if not getattr(self, '_warned_invalid_n_comp', False): + self._warned_invalid_n_comp = True + print(" mcsamplerEnsemble: update_sampling_prior SKIPPING training for dim_group {}: " + "invalid n_comp {!r} (use n_comp=0 to disable adaptation intentionally)".format( + dim_group, self.integrator.n_comp)) else: - model.update(temp_samples, log_sample_weights=ln_weights) + model.update(temp_samples, log_sample_weights=ln_weights_group) self.integrator.gmm_dict[dim_group] = model + def bootstrap_from_samples(self, samples, params=None, n_comp_warm=2, **kwargs): + """Warm-start: fit this GMM's proposal to a seed cloud so it samples AT the peak from + the first draw, instead of having to discover the peak location from scratch. For a + needle-in-a-haystack extrinsic posterior (peak ~ 10^-11 of the prior box) this is the + difference between converging and never finding the peak by cold draws. + + Builds the integrator if setup() has not run yet, then (re)fits one GMM per dim-group + from the seed with EQUAL weights -- this only shapes the proposal, carries no lnL + information, so it cannot bias the estimate (the importance weights still use the true + likelihood). AV-only kwargs (cover_frac / inflate) are accepted and ignored, so a + portfolio can forward a single seed to every member uniformly. + + `samples`: (N, ndim) array, columns in self.params_ordered order.""" + samples = np.asarray(self.identity_convert(samples)) + ndim = len(self.params_ordered) + if samples.ndim != 2 or samples.shape[1] != ndim: + raise ValueError("GMM warm-start expects (N,{}) samples in params_ordered order".format(ndim)) + # (Re)build the integrator as a SINGLE full-dimensional Gaussian group. Two reasons: + # * modeling -- a localized high-SNR peak has strong cross-parameter correlations + # (sky<->phase<->distance); one full-dim mixture captures them, whereas the default + # per-dimension factored proposal cannot and "can stall at the prior". + # * robustness -- the default gmm_dict=None path leaves integrator.bounds as a raw + # array (not a per-group dict), which breaks the per-group fit; the correlate-all + # path builds proper dict bounds. + n_comp = n_comp_warm + if (self.integrator is not None and isinstance(self.integrator.n_comp, int) + and self.integrator.n_comp > 0): + n_comp = self.integrator.n_comp + # CARRY THE COVERAGE CONFIG THROUGH. setup() rebuilds the integrator from its kwargs, + # so calling it bare here reset gmm_defensive_all_paths to False and refitted the warm + # GMM with NO defensive component -- after the portfolio had already decided, on the + # strength of that flag, that it was safe to contract its AV member. The guarantee has + # to survive the sampler's own lifecycle, not just its initial setup. + _prev = getattr(self, 'integrator', None) + # gmm_dict=None EXPLICITLY. setup() now remembers its kwargs, so a remembered explicit + # grouping would survive this call and correlate_all_dims=True would have no effect -- + # defeating the single full-dimensional group this path exists to build (the whole point + # is to capture sky<->phase<->distance correlations at high SNR). Passing None overrides + # the remembered value; the coverage settings below are still carried forward. + self.setup(n_comp=int(n_comp), correlate_all_dims=True, gmm_dict=None, + gmm_defensive_frac=getattr(_prev, 'gmm_defensive_frac', 0.05), + gmm_defensive_all_paths=getattr(_prev, 'gmm_defensive_all_paths', False)) + rvs = {p: samples[:, j] for j, p in enumerate(self.params_ordered)} + # equal weights == "put proposal mass at these seed locations" (no lnL info) + self.update_sampling_prior(self.xpy.zeros(len(samples)), len(samples), + external_rvs=rvs, log_scale_weights=True) + print(" [GMM warm-start] fitted full-dim proposal to {} seed samples (n_comp={})".format( + len(samples), n_comp)) + return True def draw_simplified(self,n,*args,**kwargs): n_samples = int(n) @@ -306,6 +542,46 @@ def draw_simplified(self,n,*args,**kwargs): return joint_p_s, joint_p_prior, rv + def sampling_density(self, X): + """Pointwise sampling density q(theta) of THIS GMM member, evaluated at + ARBITRARY points X (shape (N, ndim), columns in self.params_ordered + order). Returns a host (numpy) array of length N, or None if the + integrator/GMM has not been built yet. + + This is exactly the per-sample product MonteCarloEnsemble._sample stores + as sampling_prior_array, but evaluated at supplied points rather than at + the member's own draws: for each grouped set of dimensions it is the + fitted mixture density gmm.score(...) (already normalized to integrate to + 1 over the box, in ORIGINAL coordinates), or the uniform density 1/vol + for a not-yet-fitted (None) group. READ-ONLY; does not affect this + sampler's own integrate(). Used by the portfolio balance heuristic. + """ + integrator = getattr(self, 'integrator', None) + if integrator is None: + return None + Xc = np.atleast_2d(np.asarray(self.identity_convert(X), dtype=float)) + ndim = len(self.params_ordered) + if Xc.shape[1] != ndim and Xc.shape[0] == ndim: + Xc = Xc.T # tolerate (ndim, N) + Xg = self.identity_convert_togpu(Xc) + q = self.xpy.ones(Xg.shape[0]) + for dim_group in integrator.gmm_dict: + new_bounds = integrator.bounds[dim_group] + if len(new_bounds.shape) < 2: + new_bounds = self.xpy.array([new_bounds]) + model = integrator.gmm_dict[dim_group] + cols = self.xpy.empty((Xg.shape[0], len(dim_group))) + for index, dim in enumerate(dim_group): + cols[:, index] = Xg[:, dim] + if model is None: + llim = new_bounds[:, 0] + rlim = new_bounds[:, 1] + vol = self.xpy.prod(rlim - llim) + q *= 1.0 / vol + else: + q *= model.score(cols) + return self.identity_convert(q) + def integrate_log(self, func, *args,**kwargs): args_passed = {} @@ -334,6 +610,9 @@ def integrate(self, func, *args,**kwargs): write_to_file = kwargs['write_to_file'] if "write_to_file" in kwargs else False correlate_all_dims = kwargs['correlate_all_dims'] if "correlate_all_dims" in kwargs else False gmm_adapt = kwargs['gmm_adapt'] if "gmm_adapt" in kwargs else None + gmm_adaptive = kwargs['gmm_adaptive'] if "gmm_adaptive" in kwargs else None + gmm_defensive_frac = kwargs['gmm_defensive_frac'] if "gmm_defensive_frac" in kwargs else 0.05 + gmm_inflate = kwargs['gmm_inflate'] if "gmm_inflate" in kwargs else 1.0 gmm_epsilon = kwargs['gmm_epsilon'] if "gmm_epsilon" in kwargs else None L_cutoff = kwargs["L_cutoff"] if "L_cutoff" in kwargs else None tempering_exp = kwargs["tempering_exp"] if "tempering_exp" in kwargs else 1.0 @@ -397,14 +676,33 @@ def integrate(self, func, *args,**kwargs): integrator = monte_carlo.integrator(dim, bounds, gmm_dict, n_comp, n=n, prior=self.calc_pdf, user_func=integrator_func, proc_count=proc_count,L_cutoff=L_cutoff,gmm_adapt=gmm_adapt,gmm_epsilon=gmm_epsilon,tempering_exp=tempering_exp, - tempering_adapt=tempering_adapt, ess_target=ess_target, ess_floor=ess_floor) + tempering_adapt=tempering_adapt, ess_target=ess_target, ess_floor=ess_floor, gmm_adaptive=gmm_adaptive, + gmm_defensive_frac=gmm_defensive_frac, gmm_inflate=gmm_inflate) + # Warm-start survival: a prior setup()/bootstrap_from_samples fits proposal + # models and stores them on self.integrator, but integrate() rebuilds a fresh + # integrator from the passed gmm_dict (values None) -- so without this the + # bootstrapped fit is SILENTLY DISCARDED and a "warm" run starts cold + # (measured: warm correlate-all began at n_eff=1.0, climbed to only ~7 @4M). + # Transfer any fitted model whose dim-group key matches; a key mismatch + # (e.g. bootstrap built correlate-all but the run uses the factored pairing) + # simply falls back to cold, so this can never bias or crash. + prev = getattr(self, 'integrator', None) + if prev is not None and prev is not integrator and getattr(prev, 'gmm_dict', None): + n_xfer = 0 + for key, model in prev.gmm_dict.items(): + if model is not None and key in integrator.gmm_dict and integrator.gmm_dict[key] is None: + integrator.gmm_dict[key] = model + n_xfer += 1 + if n_xfer: + print(" [GMM warm-start] transferred {} fitted proposal group(s) into the integrator".format(n_xfer)) + self.integrator = integrator if not direct_eval: func = self.evaluate if use_lnL: print(" ==> input assumed as lnL ") if return_lnI: print(" ==> internal calculations and return values are lnI ") - integrator.integrate(func, min_iter=min_iter, max_iter=max_iter, var_thresh=var_thresh, neff=neff, nmax=nmax,max_err=max_err,verbose=verbose,progress=super_verbose,tripwire_fraction=tripwire_fraction,tripwire_epsion=tripwire_epsilon,use_lnL=use_lnL,return_lnI=return_lnI,lnw_failure_cut=lnw_failure_cut) + integrator.integrate(func, min_iter=min_iter, max_iter=max_iter, var_thresh=var_thresh, neff=neff, nmax=nmax,max_err=max_err,verbose=verbose,progress=super_verbose,tripwire_fraction=tripwire_fraction,tripwire_epsilon=tripwire_epsilon,use_lnL=use_lnL,return_lnI=return_lnI,lnw_failure_cut=lnw_failure_cut) self.n = int(integrator.n) self.ntotal = int(integrator.ntotal) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py index 20687d233..bff6f14fd 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerGPU.py @@ -83,7 +83,7 @@ def profile(fn): except: print(" - No healpy - ") -from ..integrators.statutils import update,finalize, init_log,update_log,finalize_log +from ..integrators.statutils import update,finalize, init_log,update_log,finalize_log, pareto_khat_from_log, ess_from_log_weights, block_scatter_sigma, bootstrap_lnZ_quantiles #from multiprocessing import Pool @@ -93,6 +93,12 @@ def profile(fn): rosDebugMessages = True +# Minimum uniform-mixture fraction applied to every adapted histogram (see +# compute_hist): guarantees no bin has exactly zero sampling probability, since a +# zero bin can never be re-drawn (absorbing state) and silently truncates the +# integration domain. Override at module level for controlled experiments. +HIST_FLOOR_LEVEL_MIN = 1e-2 + class NanOrInf(Exception): def __init__(self, value): self.value = value @@ -297,7 +303,12 @@ def compute_hist(self, x_samples, param,weights=None,floor_level=0): # Smooth the histogram # kernel_size =3 # histogram_values = self.xpy.convolve( histogram_values, self.xpy.ones(kernel_size)/kernel_size,mode='same') - # Mix with a uniform sampling + # Mix with a uniform sampling. A bin with exactly zero probability is an + # absorbing state: it can never be drawn again, so the sampled support is + # permanently truncated and the integral is systematically biased LOW by the + # mass outside the support -- a bias no within-run error estimate can see. + # Enforce a minimal floor so every bin stays reachable. + floor_level = max(floor_level, HIST_FLOOR_LEVEL_MIN) histogram_values = histogram_values*(1-floor_level)+floor_level*self.xpy.ones(len(histogram_values))/len(histogram_values) # Evaluate the CDF by taking a cumulative sum of the histogram. @@ -344,7 +355,7 @@ def pdf_from_hist(self, x, param): y = (x - self.x_min[param]) / self.x_max_minus_min[param] # Compute the indices of the histogram bins that `x` falls into. indices = self.xpy.trunc(y / self.dx[param], out=y).astype(np.int32) - indices = self.xpy.minimum(indices,self.n_bins[param]) # prevent being out of range due to rounding ! + indices = self.xpy.minimum(indices,self.n_bins[param]-1) # prevent being out of range due to rounding (x == right edge maps to last bin) # Return the value of the histogram. return self.histogram_values[param][indices] @@ -645,7 +656,7 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): print(" Note: cannot adapt, no history ") tempering_exp = kwargs["tempering_exp"] if "tempering_exp" in kwargs else 0.0 - n_adapt = int(kwargs["n_adapt"]*n) if "n_adapt" in kwargs else 1000 # default to adapt to 1000 chunks, then freeze + n_adapt = int(kwargs["n_adapt"]*n) if "n_adapt" in kwargs else 1000*n # default to adapt to 1000 chunks, then freeze. NOTE: scaled by n, matching integrate() floor_integrated_probability = kwargs["floor_level"] if "floor_level" in kwargs else 0 temper_log = kwargs["tempering_log"] if "tempering_log" in kwargs else False tempering_adapt = kwargs["tempering_adapt"] if "tempering_adapt" in kwargs else False @@ -674,6 +685,9 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): maxval=0 # max weight outvals=None # define in top level scope self.ntotal = 0 + # per-chunk lnZ record: each chunk used a (different) adapted proposal, so the + # between-chunk scatter is an error floor the pooled variance cannot see + lnZ_chunk_list = []; n_chunk_list = [] if bShowEvaluationLog: print("iteration Neff sqrt(2*lnLmax) sqrt(2*lnLmarg) ln(Z/Lmax) int_var") @@ -770,6 +784,14 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): else: current_log_aggregate = update_log(current_log_aggregate, log_integrand,xpy=xpy,special=xpy_special_default) outvals = finalize_log(current_log_aggregate,xpy=xpy) + # per-chunk lnZ for the between-chunk error floor (init_log returns + # (n, log_mean, log_M2, log_ref) so lnZ_chunk = log_mean + log_ref) + try: + _chunk_agg = init_log(log_integrand,xpy=xpy,special=xpy_special_default) + lnZ_chunk_list.append(float(identity_convert(_chunk_agg[1])) + float(identity_convert(_chunk_agg[3]))) + n_chunk_list.append(int(_chunk_agg[0])) + except Exception: + pass self.ntotal = current_log_aggregate[0] # effective samples maxval = max(maxval, identity_convert(self.xpy.max(log_integrand) )) @@ -804,8 +826,7 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # The total number of adaptive steps is reached # # FIXME: We need a better stopping condition here - if self.ntotal > n_adapt*n: - print(n_adapt,self.ntotal) + if self.ntotal > n_adapt: # n_adapt already scaled by n above; the old test (n_adapt*n) double-counted n and never froze continue # @@ -817,8 +838,14 @@ def inner(arg): return f(arg, p) return inner - weights_alt = self._rvs["log_integrand"][-n_history:]+np.max([maxlnL, 200]) # try to make sure we have some dynamic range here - weights_alt = self.xpy.maximum(weights_alt, 1e-5) # prevent negative weights. NOTE THIS IS IMPORTANT: if you are integrating a function with lnL<0, use an offset! + # Tempered importance weights exp(tempering_exp*lnL + ln p - ln p_s), so the + # weighted histogram of draws estimates the FIXED target L^tempering_exp * prior. + # (The old lnL + max(maxlnL,200) weights ignored tempering_exp and the 1/p_s + # correction: near-flat weights made each histogram replay the previous + # proposal's sampling noise, a multiplicative random walk that collapses the + # proposal onto a comb of surviving bins.) + weights_alt = self._rvs["log_weights"][-n_history:] + weights_alt = self.xpy.exp(weights_alt - self.xpy.max(weights_alt)) weights_alt = weights_alt/(weights_alt.sum()) if weights_alt.dtype == RiftFloat: weights_alt = weights_alt.astype(numpy.float64,copy=False) @@ -874,6 +901,36 @@ def inner(arg): else: self._rvs[key] = self._rvs[key][indx_list] + # MC-error diagnostics (must run BEFORE the fairdraw resampling below rewrites + # _rvs). See statutils: the pooled weight variance is 1/ESS restated and + # tail-blind; disclose the tail (k-hat), the between-chunk scatter, and -- + # when the naive relative error is already large -- bootstrap lnZ quantiles. + mc_diag = {} + try: + _sb = block_scatter_sigma(lnZ_chunk_list, n_chunk_list) + if _sb is not None: + mc_diag['sigma_lnZ_block'] = _sb + if "log_integrand" in self._rvs: + _lw_diag = numpy.asarray(identity_convert(self._rvs["log_integrand"] + self._rvs["log_joint_prior"] - self._rvs["log_joint_s_prior"]), dtype=float) + _kh = pareto_khat_from_log(_lw_diag) + if _kh is not None: + mc_diag['pareto_khat'] = _kh + mc_diag['n_ESS'] = ess_from_log_weights(_lw_diag) + _sig_rel_naive = np.inf + if outvals is not None: + _sig_rel_naive = float(np.exp(identity_convert(outvals[1])/2 - identity_convert(outvals[0]) - np.log(self.ntotal)/2)) + if _sig_rel_naive > 0.3 or mc_diag.get('sigma_lnZ_block', 0) > 0.3: + _q = bootstrap_lnZ_quantiles(_lw_diag, n_total=self.ntotal) + if _q is not None: + mc_diag['lnZ_ci90'] = _q + print(" [mc diag] khat={} ESS={} sigma_block={} (chunks={})".format( + round(mc_diag['pareto_khat'],3) if 'pareto_khat' in mc_diag else None, + round(mc_diag['n_ESS'],1) if 'n_ESS' in mc_diag else None, + round(mc_diag['sigma_lnZ_block'],4) if 'sigma_lnZ_block' in mc_diag else None, + len(lnZ_chunk_list))) + except Exception as _e_diag: + print(" mcsamplerGPU: MC-error diagnostics failed ({}); continuing.".format(_e_diag)) + # Do a fair draw of points, if option is set. CAST POINTS BACK TO NUMPY, IDEALLY if bFairdraw and not(n_extr is None): n_extr = int(numpy.min([n_extr,1.5*identity_convert(eff_samp),1.5*neff])) @@ -896,6 +953,7 @@ def inner(arg): dict_return ={} if convergence_tests is not None: dict_return["convergence_test_results"] = last_convergence_test + dict_return.update(mc_diag) # MC-error diagnostics (pareto_khat, n_ESS, sigma_lnZ_block, lnZ_ci90) # perform type conversion of all stored variables if cupy_ok: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerNFlow.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerNFlow.py index 480cf7310..f91991740 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerNFlow.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerNFlow.py @@ -439,6 +439,8 @@ def __init__(self,n_chunk=400000,**kwargs): self.nf_epoch = 0 self.mean_affine_set = False + # pre-trained flow to warm-load (see load_flow); applied inside integrate_log + self._preloaded_state = None def setup(self, nf_method='default',**kwargs): @@ -697,6 +699,48 @@ def update_sampling_prior(self, lnw, *args, xpy=xpy_default,no_protect_names=Tru self.nf_epoch +=1 + ### + ### FLOW STORAGE / REUSE + ### + # Training a normalizing flow is slow and expensive, and only pays off if the + # trained flow can be RE-USED across the many ILE instances that share similar + # posterior structure. These helpers persist a trained flow to a small file + # and warm-load it into a fresh sampler, which can then either sample from it + # directly (n_adapt=0) or do a few cheap 'polish' epochs (small n_adapt) to + # adapt it to the new instance -- amortizing the training cost. + + def save_flow(self, path): + """Serialize the trained flow (torch state_dict + the architecture + metadata needed to rebuild it) to `path`.""" + if self.nf_flow is None: + raise Exception("mcsamplerNFlow.save_flow: no trained flow to save (run integrate first)") + payload = dict( + state_dict=self.nf_flow.state_dict(), + params_ordered=[str(p) for p in self.params_ordered], + bounds=[[float(self.llim[p]), float(self.rlim[p])] for p in self.params_ordered], + nf_method=getattr(self, 'nf_method', 'default'), + num_layers=int(getattr(self, 'num_layers', max(1, len(self.params_ordered) // 2))), + nf_epoch=int(self.nf_epoch), + ) + torch.save(payload, path) + return path + + def load_flow(self, path): + """Stage a pre-trained flow saved by save_flow(). The weights are applied + inside integrate_log() (after the architecture is rebuilt), so this must be + called before integrate/integrate_log. Verifies the parameters and box + match this sampler.""" + payload = torch.load(path, map_location='cpu') + if [str(p) for p in payload['params_ordered']] != [str(p) for p in self.params_ordered]: + raise ValueError("saved flow params {} != sampler params {}".format( + payload['params_ordered'], [str(p) for p in self.params_ordered])) + saved_bounds = np.array(payload['bounds'], dtype=float) + my_bounds = np.array([[self.llim[p], self.rlim[p]] for p in self.params_ordered], dtype=float) + if not np.allclose(saved_bounds, my_bounds): + raise ValueError("saved flow box does not match sampler box") + self._preloaded_state = payload + return payload + @profile def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): """ @@ -722,7 +766,13 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): xpy_here = self.xpy - + # The normalizing flow (nflows/torch) samples and scores on the host, and + # the integrand is a host function, so the NF integrator AGGREGATES on the + # host. Force the running estimate onto numpy/scipy regardless of the + # xpy=cupy default -- mixing a cupy xpy with the flow's host arrays is what + # crashed the GPU path (cupy rv handed to a numpy integrand). + xpy = self.xpy # = numpy + special_here = special # scipy.special (host) # # Determine stopping conditions # @@ -748,10 +798,12 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): save_intg = kwargs["save_intg"] if "save_intg" in kwargs else False - # FIXME: The adaptive step relies on the _rvs cache, so this has to be - # on in order to work - if n_adapt > 0 and tempering_exp > 0.0: - save_intg = True + # The NF's final integral estimate reads log_integrand/log_joint_prior/ + # log_joint_s_prior back out of self._rvs, so those MUST be accumulated + # regardless of adaptation. (Previously save_intg was only turned on when + # n_adapt>0 and tempering_exp>0, so pure flow REUSE with n_adapt=0 hit a + # KeyError('log_integrand') at the end.) + save_intg = True deltalnL = kwargs['igrand_threshold_deltalnL'] if 'igrand_threshold_deltalnL' in kwargs else float("Inf") # default is to return all deltaP = kwargs["igrand_threshold_p"] if 'igrand_threshold_p' in kwargs else 0 # default is to omit 1e-7 of probability @@ -783,8 +835,25 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): print("iteration Neff sqrt(2*lnLmax) sqrt(2*lnLmarg) ln(Z/Lmax) int_var") self.n_chunk = n - self.setup(nf_method=nf_method) - + # a warm-loaded flow dictates the architecture; use its method so the + # rebuilt transform matches the saved weights + if self._preloaded_state is not None: + nf_method = self._preloaded_state.get('nf_method', nf_method) + self.setup(nf_method=nf_method) + + # WARM-LOAD: rebuild the flow object on the freshly-created architecture + # and load the pre-trained weights, so this run starts from a trained flow + # (n_adapt=0 -> pure reuse; small n_adapt -> a few polish epochs). + if self._preloaded_state is not None: + flow = Flow(self.nf_model, StandardNormal(shape=[len(self.params_ordered)])) + flow.load_state_dict(self._preloaded_state['state_dict']) + self.nf_flow = flow + self.nf_trainer.flow = flow + self.mean_affine_set = True # affine layer is part of the loaded weights + self.nf_epoch = int(self._preloaded_state.get('nf_epoch', 0)) + if bShowEvaluationLog: + print(" [NF warm-load] restored pre-trained flow ({} layers)".format(self.num_layers)) + ntotal_true = 0 max_epochs_requested =300 while (eff_samp < neff and ntotal_true < nmax ): # and (not bConvergenceTests): @@ -798,28 +867,33 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): log_joint_p_s = np.log(joint_p_s) log_joint_p_prior = np.log(joint_p_prior) ntotal_true += len(joint_p_s) - if cupy_ok: - rv = identity_convert_togpu(rv) # send random numbers to GPU : ugh - log_joint_p_prior = identity_convert_togpu(log_joint_p_prior) # send to GPU if required. Don't waste memory reassignment otherwise - - # Evaluate function, protecting argument order + # rv is a host array from the flow (nflows/torch samples on the host). + rv = identity_convert(rv) params = [] for item in self.params_ordered: # USE IN ORDER if isinstance(item, tuple): params.extend(item) else: params.append(item) - unpacked = unpacked0 = rv #numpy.hstack([r.flatten() for r in rv]).reshape(len(args), -1) - unpacked = dict(list(zip(params, unpacked))) - if 'no_protect_names' in kwargs: - lnL = lnF(*unpacked0) # do not protect order + # Evaluate the integrand. The real GPU ILE likelihood is DEVICE-native + # (wants cupy); synthetic/CI integrands are host-native. Feed + # device-first, fall back to host on a type error, and remember the + # choice (same contract as AV / the portfolio). The flow's own math + # stays on the host, so lnL is brought back to host afterwards. + def _eval_integrand(cols): + if 'no_protect_names' in kwargs: + return lnF(*cols) + return lnF(**dict(list(zip(params, cols)))) + if getattr(self, '_integrand_wants_host', False) or not cupy_ok: + lnL = _eval_integrand(rv) else: - unpacked = dict(list(zip(self.params_ordered,rv.T))) - lnL= lnF(**unpacked) # protect order using dictionary - # take log if we are NOT using lnL - if cupy_ok: - if not(isinstance(lnL,cupy.ndarray)): - lnL = identity_convert_togpu(lnL) # send to GPU, if not already there + try: + lnL = _eval_integrand(identity_convert_togpu(rv)) + except (TypeError, ValueError): + self._integrand_wants_host = True + lnL = _eval_integrand(rv) + # bring lnL back to the host for the flow-side / aggregation math + lnL = identity_convert(lnL) # For now: no prior, just duplicate VT algorithm @@ -828,9 +902,9 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # log_weights = tempering_exp*lnL + log_joint_p_prior # log aggregate: NOT USED at present, remember the threshold is floating if current_log_aggregate is None: - current_log_aggregate = init_log(log_integrand,xpy=xpy,special=xpy_special_default) + current_log_aggregate = init_log(log_integrand,xpy=xpy,special=special_here) else: - current_log_aggregate = update_log(current_log_aggregate, log_integrand,xpy=xpy,special=xpy_special_default) + current_log_aggregate = update_log(current_log_aggregate, log_integrand,xpy=xpy,special=special_here) # Monitoring for i/o outvals = finalize_log(current_log_aggregate,xpy=xpy) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerPortfolio.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerPortfolio.py index 016acce45..2985e17a6 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerPortfolio.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/mcsamplerPortfolio.py @@ -2,7 +2,7 @@ import math #import bisect from collections import defaultdict -from types import ModuleType +from types import ModuleType, FunctionType, BuiltinFunctionType import numpy np=numpy #import numpy as np @@ -11,7 +11,7 @@ import itertools import functools -from copy import deepcopy +from copy import deepcopy, copy as shallow_copy import os @@ -88,8 +88,13 @@ def profile(fn): known_pipelines = {} for pipeline in discovered_plugins: print(" Portfolio discovery: loading ", pipeline.name) - known_pipelines[pipeline.name] = pipeline.load() -print('RIFT portfolio plugins:', [ep.name for ep in discovered_plugins]) + try: + known_pipelines[pipeline.name] = pipeline.load() + except Exception as e: + # optional plugins (e.g. the NF pipeline needing torch) must not make + # importing mcsamplerPortfolio itself fail on torch-free installations + print(" Portfolio discovery: SKIPPING {} (unavailable: {})".format(pipeline.name, e)) +print('RIFT portfolio plugins:', sorted(known_pipelines)) class NanOrInf(Exception): @@ -111,6 +116,12 @@ def portfolio_default_weights(n_ess_list, wt_previous, portfolio_probability_flo # don't update if we have insane answers if any(np.isnan(rewt)): return wt_previous + # if every member is degenerate (n_ess ~ 1, e.g. a very hard target's first + # chunk found nothing), the normalization below would divide by zero and + # produce nan weights -> negative per-member sample counts downstream. Keep + # the previous (typically uniform) weights instead. + if np.sum(rewt) <= 0: + return wt_previous rewt = np.ones(len(rewt))*portfolio_probability_floor + (rewt/np.sum(rewt)) * (1-portfolio_probability_floor) net = (rewt * history_factor + wt_previous*(1-history_factor)) return net/np.sum(net) # make SURE normalized correctly @@ -150,8 +161,129 @@ def __init__(self,portfolio=None,portfolio_weights=None,oracle_realizations =Non if not(self.portfolio_weights ): self.portfolio_weights = np.ones(len(self.portfolio))/(1.0*len(self.portfolio)) - self.portfolio_adapt = np.ones(len(self.portfolio),dtype=bool) # default : everything adapts. + self.portfolio_adapt = np.ones(len(self.portfolio),dtype=bool) # default : everything adapts. self.portfolio_freeze_wt =portfolio_freeze_wt # if weight is below this number, the portfolio member's distribution will NOT update. SCALAR + # Freeze protection. A member (esp. a VARAHA/AV workhorse) contributes little on its + # first chunks -- before it has contracted -- so a plain weight OFF, vs gate base): + # d4_n1_s101 25.9 -> 53.5 (base 53.5) d4_n3_s202 29.1 -> 83.8 (base 83.8) + # d6_n1_s202 64.0 -> 102.1 (base 102.1) d6_n3_s202 7.2 -> 31.4 (base 31.4) + # d8_n1_s101 37.3 -> 61.9 (base 61.9) + # i.e. OFF reproduces base EXACTLY, so this was the sole default-path regression in + # PR #28 (never-freeze was measured to be a no-op on these targets, ratio 1.00). + # Kept available for experimentation; DO NOT default it on without re-running the gate. + self.portfolio_plateau_revive = kwargs.get('portfolio_plateau_revive', False) + # Diagnostic: per-member n_ess history (one list per portfolio member), appended each + # chunk in the report block. Enables plateau-aware policies and post-hoc analysis. + self.portfolio_member_ness_history = [[] for _ in range(len(self.portfolio))] + + # ADAPTIVE-PROBE DRAW ALLOCATION -- OPT-IN (default OFF). Idea: keep a per-member QUALITY + # estimate updated only from fair-allocation chunks, allocate draws by quality^exponent, and + # round-robin PROBE each member at a raised share so a suppressed member can prove itself. + # q_mix keeps this unbiased for ANY allocation. On strongly-correlated SYNTHETIC targets it + # works well (a full-covariance GMM wins the probe and the portfolio beats AV -- + # test_portfolio_adaptive_alloc.py). + # + # BUT it is NOT a safe default: the quality signal is each member's per-chunk Kish n_ess, + # which rewards SELF-CONSISTENCY, not integral coverage. A warm GMM is instantly + # self-consistent (per-chunk n_ess ~120) while a warm VARAHA/AV member's per-chunk n_ess is + # genuinely ~1 during its slow, CUMULATIVE contraction (its value emerges over ~70 chunks). + # So on a real high-SNR AV-favorable event (S250114ax) adaptive drives the true AV workhorse + # to the floor and rides the self-consistent-but-worse GMM: measured n_eff 8 vs 53 for the + # legacy allocation -- a regression. The probe cannot rescue AV because AV still looks bad + # at high allocation until fully contracted. A correct default needs a GLOBAL-impact quality + # signal (how much a member improves the pooled q_mix n_eff), not per-member self-n_ess; that + # is future work. Until then the DEFAULT keeps the legacy n_ess reweighting. + self.portfolio_adaptive_alloc = kwargs.get('portfolio_adaptive_alloc', False) + # QUALITY SIGNAL for the allocation: + # 'global' (default) -- each member's MARGINAL GAIN IN POOLED n_eff PER SAMPLE, + # g_m = 2*mean_w_m/S - mean_w2_m/Q (S=sum w, Q=sum w^2 over ALL samples; see the + # derivation where it is computed). This directly optimizes the quantity we care + # about: it credits a member for the weight MASS it contributes and debits it for the + # weight VARIANCE it injects. The two simpler candidates both fail: + # * Kish n_ess is SCALE-INVARIANT ((sum w)^2/sum w^2 is unchanged if all w are + # scaled), so it cannot see whether a member carries any integral mass at all -- a + # self-consistent member sitting off-peak scores as well as one covering the peak. + # * mean weight alone REWARDS badly-matched proposals: a well-matched contracted AV + # correctly has small uniform weights, while a broad GMM's rare huge-weight outlier + # sets the max (measured on S250114ax: AV 1e-40 vs GMM 2e-4 -- backwards). + # g_m is also self-correcting at low allocation: a starved peak-covering member sees + # inflated weights (q_mix is small there) so it earns share, and as its share grows + # q_mix rises and the weights fall -- an equilibrium, with no under-observation trap. + # 'ness' -- legacy per-member Kish n_ess (kept for comparison; see the S250114ax regression). + self.portfolio_quality_signal = kwargs.get('portfolio_quality_signal', 'global') + self.portfolio_alloc_exponent = kwargs.get('portfolio_alloc_exponent', 1.0) # weights ~ quality^p + self.portfolio_alloc_floor = kwargs.get('portfolio_alloc_floor', 0.05) # min share (coverage+probe) + self.portfolio_quality_decay = kwargs.get('portfolio_quality_decay', 0.5) # EMA alpha for quality + self.portfolio_probe_period = kwargs.get('portfolio_probe_period', 4) # probe one member every N chunks + self.portfolio_probe_frac = kwargs.get('portfolio_probe_frac', 0.6) # raise probed member to >= this + # WEIGHT CLIPPING (truncated importance sampling) -- OPT-IN, default off, PROPOSAL-FIT INPUT + # ONLY (see the clipping block in integrate_log). Cap w at + # tau = portfolio_weight_clip * sqrt(n) * mean(w) (Ionides 2008, truncated IS) + # Clipping is BIASED and distorts n_ess, so the clipped copy feeds ONLY + # member.update_sampling_prior (the GMM covariance fit) -- one enormous weight can't make + # that fit degenerate. The estimator (ln Z, n_eff), the n_ess report, and the allocation + # signal all use the TRUE weights, so they stay exactly unbiased/undistorted. The withheld + # mass is accumulated in log space and reported as a TAIL DIAGNOSTIC, not a bias. + self.portfolio_weight_clip = kwargs.get('portfolio_weight_clip', 0.0) # 0 = off + self.portfolio_clip_log_removed = -np.inf + self.portfolio_clip_log_total = -np.inf + self.portfolio_clip_n = 0 + # diagnostic: samples whose mixture density UNDERFLOWED to 0 and hit the 1e-300 floor (those + # produce spurious ~1/1e-300 weights -- a numerical artifact, not real tail mass) + self.portfolio_qmix_underflow = 0 + # VARAHA DRAW FLOOR (opt-in, default 0 = off). never-freeze guarantees a VARAHA/AV member + # keeps UPDATING (contracting), but nothing guarantees it keeps DRAWING: both allocation + # rules score members by per-chunk n_ess, and a VARAHA member's per-chunk n_ess sits at ~1 + # during its slow CUMULATIVE contraction, so a member that looks instantly good (a live GMM) + # can take almost the whole budget. Measured on S250114ax after the PR #33 fixes made the + # GMM member genuinely live: the allocation gave GMM ~0.84 and the portfolio collapsed to + # n_eff ~2 at 4M, versus ~100 for standalone AV. Setting this to f reserves a combined + # fraction f of the draws for VARAHA members (applied AFTER whichever allocation rule runs, + # so it protects the legacy and adaptive paths alike). q_mix keeps any allocation unbiased, + # so this only trades efficiency. + self.portfolio_varaha_min_frac = kwargs.get('portfolio_varaha_min_frac', 0.0) + self.portfolio_varaha_max_frac = kwargs.get('portfolio_varaha_max_frac', 0.0) # 0 = no cap + # Range restriction (see setup(): portfolio_restrict_ranges). Default OFF, so every code + # path guarded by these is inert unless a member is explicitly narrowed. + self._has_restricted_member = False + self._full_support_members = [] + self._pending_range_overrides = set() # (member, param) awaiting add_parameter; see setup() + # Keep the full-support backstop COLD when warm-starting. Seeding every member removes + # the mixture's coverage of the prior box (cover_frac cannot restore it -- see + # bootstrap_from_samples), which turns a mis-placed seed from an efficiency cost into a + # silent low bias. Set False to restore the old seed-everything behaviour. + self.portfolio_warmstart_backstop_cold = True + self.portfolio_quality = np.ones(len(self.portfolio)) # per-member quality (EMA of the signal) + self.portfolio_quality_nobs = np.zeros(len(self.portfolio), dtype=int) # #updates per member + self.portfolio_probe_ptr = 0 # round-robin probe pointer + + # ---- SUPPORT-MISMATCH (warm-start) DIAGNOSTIC state. OFF-PATH: read-only reduction over + # densities integrate_log already evaluates to build q_mix; nothing here feeds the estimate. + # See _update_support_diagnostics() for the definitions and the measured caveats. + self._reset_support_diagnostics() # Total number of samples drawn self.ntotal = 0 @@ -179,23 +311,603 @@ def __init__(self,portfolio=None,portfolio_weights=None,oracle_realizations =Non # extra args, created during setup self.extra_args = {} + # PER-MEMBER RANGE OVERRIDES for interval narrowing: {member_index: {param: (lo, hi)}}. + # Populate with restrict_member_range() BEFORE add_parameter()/setup(). Member 0 is the + # designated full-support backstop and must never be narrowed. + self.member_range_overrides = {} + + def restrict_member_range(self, member_index, param, lo, hi): + """Narrow ONE portfolio member's sampling range for `param` to [lo, hi] (interval narrowing). + + This is a PROPOSAL-only change: the member's prior callables are untouched, so it keeps + reporting the true global prior, and the balance-heuristic mixture density q_mix keeps the + estimate unbiased with no renormalization -- PROVIDED at least one member retains full + support. Member 0 is that backstop by convention and may not be narrowed. + + Why the backstop is not optional: measured on a truth-known ladder, a WRONG sub-box costs a + STANDALONE sampler up to -1949 nats (while still reporting a healthy n_eff of 220-840, i.e. + confidently wrong), but only ~1 nat inside a portfolio whose full-box member keeps q_mix + covering the complement. Restriction without a backstop converts a rare pathology into a + systematic one. + + Call before add_parameter(); narrowing is applied there, and setup() then builds every + derived quantity from the narrowed range. + """ + n_members = len(self.portfolio_realizations) + member_index = int(member_index) + # Validate STRICTLY: a negative index or a misspelled parameter used to be accepted here and + # then silently fail to match the positive enumerate()/params checks in add_parameter, so the + # call "succeeded" while applying no restriction at all. + if member_index == 0: + raise ValueError( + "mcsamplerPortfolio.restrict_member_range: member 0 is the full-support backstop and " + "must not be narrowed -- q_mix would then have no component covering the complement, " + "and a mode outside every sub-box becomes uncoverable rather than merely under-covered.") + if not (1 <= member_index < n_members): + raise ValueError("restrict_member_range: member_index must satisfy 1 <= i < {} (got {}); " + "negative indices are NOT accepted -- they never match the positive " + "enumerate() in add_parameter and would silently be a no-op.".format( + n_members, member_index)) + if not (hi > lo): + raise ValueError("restrict_member_range: need hi > lo, got [{}, {}]".format(lo, hi)) + self.member_range_overrides.setdefault(member_index, {})[param] = (float(lo), float(hi)) + # Track for the consumed-check in setup(): a parameter name that never arrives via + # add_parameter must be an ERROR, not a silent no-op. + self._pending_range_overrides = getattr(self, '_pending_range_overrides', set()) + self._pending_range_overrides.add((member_index, param)) + + # CENTRALISE the coverage invariants: these are the SAME flags the + # setup(portfolio_restrict_ranges=...) path establishes. Setting them only there meant this + # public API disabled the full-support draw floor, the restricted-only active-member guard and + # the q_mix fallback guard -- so member 0 could be allocated zero draws and the mixture could + # silently lose full support, which is precisely the failure restriction is supposed to avoid. + restricted = set(self.member_range_overrides) + if len(restricted) >= n_members: + raise ValueError( + "restrict_member_range: that would restrict EVERY member, leaving no component with " + "full support. The mixture would not cover L*p outside the sub-boxes and the integral " + "would be biased low with no diagnostic. Leave at least one member unrestricted.") + self._has_restricted_member = True + self._full_support_members = [i for i in range(n_members) if i not in restricted] + def add_parameter(self, params, pdf, **kwargs): """ Add one (or more) parameters to sample dimensions. params is either a string describing the parameter, or a tuple of strings. The tuple will indicate to the sampler that these parameters must be sampled together. left_limit and right_limit are on the infinite interval by default, but can and probably should be specified. If several params are given, left_limit, and right_limit must be a set of tuples with corresponding length. Sampling PDF is required, and if not provided, the cdf inverse function will be determined numerically from the sampling PDF. """ self.params.add(params) # does NOT preserve order in which parameters are provided self.params_ordered.append(params) - for member in self.portfolio_realizations + self.oracle_realizations: + _all_members = self.portfolio_realizations + self.oracle_realizations + for indx, member in enumerate(_all_members): member.add_parameter(params, pdf, **kwargs) - # update dictionary limits, yes this is super-redundant, but we have a scoping issue and this is easier to code - self.llim.update( member.llim) - self.rlim.update(member.rlim) - # set master list of adaptive parameters + # The PORTFOLIO's own limits must always describe the FULL prior range, never a + # restricted member's sub-box: they are the reference range used downstream (L0-rescue + # puff width, breadcrumb bounds, distance-marginalization bounds). Take them from + # member 0, which is the designated FULL-SUPPORT member by convention (see + # restrict_member_range), and take them BEFORE any narrowing is applied below. + if indx == 0: + self.llim.update( member.llim) + self.rlim.update(member.rlim) + # set master list of adaptive parameters self.adaptive = member.adaptive # top level list of adaptive coordinates + # PER-MEMBER RANGE RESTRICTION (interval narrowing). + # At high SNR the posterior can occupy a vanishing fraction of the prior box, so a member + # confined to a well-chosen sub-box resolves it far better (measured: n_eff 2495 vs 1629, and + # SNR-INDEPENDENT, on the truth-known ladder). We do this by narrowing ONE member's limits + # rather than clipping the prior, which is what makes it safe: + # * The estimator weight is L*p_prior/q_mix with p_prior the TRUE prior. A member's range + # is purely a PROPOSAL choice -- proposals need not cover the prior, only the MIXTURE + # must cover the support of L*p. So NO prior renormalization and NO clipped-volume + # correction are required, PROVIDED a full-support member remains (see _full_support_members). + # * We must NOT rebuild the prior callables for the narrowed member: `prior_prod` evaluates + # the callables handed to add_parameter, and those are absolute densities normalized over + # the ORIGINAL range. Sharing them is what keeps every member reporting the SAME true + # prior -- the portfolio takes joint_p_prior from whichever member drew each sample, so + # a member that renormalized its prior over its sub-box would silently bias the integral. + # Hence we only overwrite llim/rlim here, and only AFTER add_parameter has installed the + # shared callables. + # Narrowing happens before setup(), so every derived AV quantity (my_ranges, dx, dx0, V, + # binunique, ninbin) is built from the narrowed range and nothing is left stale. + for indx, member in enumerate(self.portfolio_realizations): + _ov = self.member_range_overrides.get(indx) + if not _ov or params not in _ov: + continue + lo, hi = _ov[params] + # NARROW ONLY. Widening past the member's own limits would sample where the SHARED + # prior callables (normalized over the ORIGINAL range) are not normalized, so the + # member would report a prior density that is wrong outside the original box -- a + # biased integral, silently. The name says "restrict"; refuse rather than quietly + # clip, so a caller who meant to widen finds out instead of getting a no-op. + if lo < member.llim[params] or hi > member.rlim[params]: + raise ValueError( + "restrict_member_range: requested [{}, {}] for {!r} on member {} is NOT contained " + "in that member's range [{}, {}]. This API can only narrow: the prior callables " + "are absolute densities normalized over the original range, so sampling outside " + "it would bias the integral.".format(lo, hi, params, indx, + member.llim[params], member.rlim[params])) + member.llim[params] = lo + member.rlim[params] = hi + getattr(self, '_pending_range_overrides', set()).discard((indx, params)) + print(" [portfolio] member {} range for {} narrowed to [{}, {}] (proposal only; " + "prior callables untouched)".format(indx, params, lo, hi)) + + + @staticmethod + def _snapshot_setup_args(args): + """Copy the mutable containers in a setup-argument dict; pass everything else by reference. + + REQUIRED for correctness, not tidiness. Setup arguments are not inert: production supplies + `gmm_dict` as a grouping spec ({(0,1,2): None, ...}), mcsamplerEnsemble hands that very + object to monte_carlo.integrator, which stores it WITHOUT copying (MonteCarloEnsemble.py:110) + and then writes trained models into it (`self.gmm_dict[dim_group] = model`, :403). Keeping a + reference and replaying it would hand the next point the PREVIOUS point's trained proposal -- + reintroducing, through the reset itself, exactly the state leak the reset exists to remove. + + Objects NESTED INSIDE a spec container are cloned too, not just the container. A seeded + GMM model supplied via `--extrinsic-proposal-breadcrumb` lives as a VALUE in gmm_dict, and + with `--extrinsic-proposal-adapt` it keeps adapting: `model.update()` mutates it in place + (gaussian_mixture_model.py:548). Copying only the dict would leave the stored "baseline" + pointing at the live model, so it would drift during point 1 and be replayed into point 2 -- + the same leak one level down. (With adapt OFF, the default, `_train` skips seeded groups + and nothing mutates, so this path was previously harmless.) + + TOP-LEVEL non-container arguments are still passed by reference: those are callables, + modules and sampler objects, which a deepcopy would try to clone and can fail on or spend + real time duplicating. If a nested clone fails, the original is kept and the failure is + REPORTED -- a silently shared object is how this class of bug survives. + """ + _unclonable = [] + + def _cp(v, depth=0): + if depth > 6: # runaway guard; setup specs are shallow + return v + if isinstance(v, dict): + return dict((k, _cp(x, depth + 1)) for k, x in v.items()) + if isinstance(v, list): + return [_cp(x, depth + 1) for x in v] + if isinstance(v, tuple): + return tuple(_cp(x, depth + 1) for x in v) + if isinstance(v, set): + return set(v) + if isinstance(v, np.ndarray): + return v.copy() + if depth == 0 or v is None or isinstance(v, (bool, int, float, complex, str, bytes)): + return v + if isinstance(v, (ModuleType, FunctionType, BuiltinFunctionType, type)): + return v # stateless: sharing these is safe and cloning them is not + # A nested object with mutable state -- e.g. a seeded GMM `estimator`. deepcopy is + # NOT usable: a real estimator holds a module reference (`xpy`) and deepcopy raises + # "cannot pickle 'module' object", which would send us down the fallback and leave the + # model SHARED -- i.e. not fixed at all. Shallow-copy the object (which never + # pickles) and then clone its mutable attributes, leaving module/function refs shared. + try: + new_obj = shallow_copy(v) + d = getattr(new_obj, '__dict__', None) + if d is None: + raise TypeError("no __dict__ (__slots__?), cannot clone attribute state") + for k, val in list(d.items()): + d[k] = _cp(val, depth + 1) + return new_obj + except Exception as e: + _unclonable.append("{} ({})".format(type(v).__name__, e)) + return v + + out = dict((k, _cp(v)) for k, v in args.items()) + if _unclonable: + print(" [portfolio] WARNING: could not clone {} nested setup object(s): {}. These are " + "SHARED with the live sampler, so if anything mutates them in place their state " + "will persist across a reset.".format(len(_unclonable), "; ".join(_unclonable))) + return out + + def clear_warm_state(self): + """Clear any warm-start seed AND the installed active grid on every member. + + Setting `portfolio._warm = None` does NOT do this: `_warm` and the contracted AV grid live on + the MEMBERS, not on the portfolio object. Now that a seed is actually installed on the draw + path (AV._apply_warm_state), failing to clear it between points would let the next point reuse + the PREVIOUS point's contracted live volume -- which can exclude the new point's support and + bias it low with no diagnostic. Called by the driver wherever it used to do + `sampler._warm = None`, including the seed-capture failure and exception paths. + + Failures PROPAGATE. A reset that quietly did not happen leaves the next point drawing from + the previous point's grid, which is the exact silent-wrong-answer this method exists to + prevent -- so it must not be reducible to a log line. + """ + self._warm = None + _groups = [(list(getattr(self, 'portfolio_realizations', [])), + getattr(self, '_member_setup_args', None)), + (list(getattr(self, 'oracle_realizations', [])), + getattr(self, '_oracle_setup_args', None))] + for members, saved_args in _groups: + for indx, member in enumerate(members): + member._warm = None + member._warm_applied = False + if not hasattr(member, 'setup'): + continue + # REPLAY the member's original setup arguments. A bare setup() restores the cold + # grid but DISCARDS the configuration: mcsamplerEnsemble.setup() rebuilds its + # dimension grouping and re-reads n_comp / gmm_adapt / correlate_all_dims from + # kwargs, so a configured (0,1) GMM with n_comp=3 and adaptation off comes back as + # separate (0,), (1,) groups with n_comp defaulted and gmm_adapt=None -- a quietly + # different sampler for every point after the first. + args_here = None + if saved_args is not None and indx < len(saved_args): + args_here = saved_args[indx] + if args_here is None: + # setup() was never run through the portfolio: nothing to replay, nothing to + # lose. AV.setup() ignores kwargs and rebuilds from the member's own llim/rlim + # (so a narrowed member stays narrowed across the reset). + member.setup() + else: + # a FRESH copy per replay: passing the stored dict itself would let the + # rebuilt integrator train into our snapshot, so the reset after next would + # replay a polluted spec and the leak would return one point later. + member.setup(**self._snapshot_setup_args(args_here)) + + ### + ### SUPPORT-MISMATCH (WARM-START) DIAGNOSTIC -- strictly OFF-PATH + ### + # WHAT IT IS. A warm-started AV/VARAHA member draws UNIFORMLY over its contracted live volume, + # so its density is EXACTLY ZERO outside that volume (a hard-edged union of boxes). If the seed + # was built at a different point -- a neighbouring intrinsic grid point, a stale breadcrumb, a + # displaced posterior -- the true peak can lie outside the seeded box entirely. Define, for + # member m, over ALL samples of the run: + # + # escaped_mass[m] = sum_{i : q_m(x_i) == 0} w_i / sum_i w_i (w = L p_prior / q_mix) + # + # i.e. the fraction of the total posterior weight carried by samples member m COULD NOT HAVE + # DRAWN. Matched seed -> small; a seed displaced off the peak -> ->1. Its power is meant to + # come from lnL AMPLITUDE, not sample count: one draw at a peak the warm member misses carries + # weight ~e^{Delta lnL} times everything inside it, so the statistic can fire while n_eff is + # still 3-9 and estimation is hopeless. This is exactly what n_eff and the Pareto k-hat cannot + # do -- both are functions only of the weights actually drawn. + # + # COMPANION (soft) STATISTIC. weight_share[m] = sum of w over samples DRAWN BY m / sum of all w. + # Far less invasive (no density comparison at all), so it is reported alongside; see the study + # notes for whether it is actually weaker. + # + # WHY IT IS FREE. integrate_log already evaluates EVERY member's density at EVERY pooled sample + # to build q_mix = sum_m frac_m q_m, and keeps frac_m*q_m per member in self._chunk_mix_parts. + # This reduction is one comparison and one masked sum per member per chunk -- the same order as + # forming q_mix, and it touches neither log_integrand nor the weights. + # + # LIMITS, MEASURED, NOT ASSUMED (do not read the number without these): + # * A member that drew ZERO samples in a chunk has no entry in _chunk_mix_parts, so that chunk + # is excluded from BOTH its numerator and its denominator. The per-member denominator is + # therefore the total weight of the chunks that member participated in, not of the run. + # * "q_m == 0" is a FLOATING-POINT test. A genuinely soft member (a GMM) evaluated far into + # its tail can UNDERFLOW to 0 in double precision and be scored as hard-edged. The + # hard_edged flag below records only that a zero was ever seen; it does not distinguish + # "bounded support" from "underflowed tail", and for the detector's purpose it need not -- + # a density that underflows to zero cannot be sampled from either. + # * It measures where the WEIGHT the portfolio FOUND is, so it is blind to a peak NO member + # ever sampled. A cold uniform member in high dimension may land on a narrow true peak so + # rarely that a genuinely mismatched seed still reads escaped_mass ~ 0 (starvation false + # negative). The statistic is a lower bound on mismatch, never an upper bound. + # + # MEASURED VERDICT -- IT IS A WARM-START QUALITY MONITOR, NOT A lnZ-BIAS ALARM. Full ROC in + # test/expensive_before_merging/integrators/escaped_mass_study.py (1440 truth-known runs, + # 20 independent target seeds per cell, d=4 and d=6). In brief: + # * Use escaped_mass_EARLY (the first chunk), NOT the cumulative number. A correctly-placed + # seed's converged live volume legitimately excludes median 0.51 (d=4) / 0.80 (d=6) of the + # posterior WEIGHT, so the cumulative statistic has no usable floor. The first-chunk floor + # is 6e-6 / 2.5e-7 (max over 20 seeds 2.3e-4 / 6.8e-3). + # * It only works when SOME member is left COLD/broad. If every member is warm-started from + # the same cloud -- which is what bootstrap_from_samples does by default, mcsamplerEnsemble + # having its own bootstrap_from_samples -- the first chunk draws nothing outside the seeded + # volume and the statistic reads exactly 0.000 no matter how wrong the seed is (320/320 + # runs). A monitor that wants this signal must keep an unseeded probe member. + # * RECOMMENDED THRESHOLD, valid ONLY under those two conditions (first-chunk statistic, at + # least one cold/broad member): escaped_mass_early > 1e-2. Measured 0/40 false positives + # on matched seeds at d=4 and d=6 (Wilson/rule-of-three 95% upper bound ~0.07), with + # true-positive rate 0.20/0.75/1.00/1.00 at d=4 and 0.10/0.50/0.80/1.00 at d=6 for seed + # displacements of 1/1.5/2/3 prior-box units. + # * And in exactly that configuration lnZ is already protected by the balance heuristic + # (|bias| <= 0.31 nat at d=4 over every displacement tested). Where displacement DOES bias + # lnZ -- an all-AV portfolio, no soft backstop -- nothing can be seen escaping, because the + # mixture never samples outside the union of the AV volumes. Do not build a lnZ-bias gate + # on this number. + def _reset_support_diagnostics(self): + """(Re)initialize the support-mismatch accumulators. Log-space, because the per-chunk + weight totals of a peaked target span many orders of magnitude and a linear running sum + would be dominated by whichever chunk found the peak (or underflow to 0 before it does).""" + m = len(getattr(self, 'portfolio_realizations', []) or getattr(self, 'portfolio', []) or []) + self.portfolio_escape_log_num = np.full(m, -np.inf) # log sum w over q_m == 0 samples + self.portfolio_escape_log_den = np.full(m, -np.inf) # log sum w over chunks m took part in + self.portfolio_share_log_num = np.full(m, -np.inf) # log sum w over samples m DREW + self.portfolio_weight_log_total = -np.inf # log sum w over the whole run + self.portfolio_escape_n = np.zeros(m, dtype=np.int64) # count of zero-density samples + self.portfolio_escape_nsamp = np.zeros(m, dtype=np.int64) # samples m was evaluated at + self.portfolio_escape_hard = np.zeros(m, dtype=bool) # ever saw an exact zero + # PER-CHUNK history (list of length-m arrays). The cumulative fraction mixes regimes: a + # VARAHA member's live volume CONTRACTS as the run proceeds, so escaped mass grows even for + # a perfectly-placed seed, while a MISPLACED seed is already fully escaped in chunk 1. + # Keeping the history lets a caller read the early (seed-state) signal, which is the one + # that is actually about the warm start. m floats per chunk: free. + self.portfolio_escape_history = [] + self._member_index = {} + + def _update_support_diagnostics(self, log_weights, q_mix): + """Accumulate escaped_mass / weight_share for THIS chunk. Pure reduction; no state used + by the estimator is read or written. Any failure is swallowed: a diagnostic must never be + able to take down an integral.""" + try: + parts = getattr(self, '_chunk_mix_parts', None) + if not parts or q_mix is None: + return + n_mem = len(self.portfolio_realizations) + if len(self.portfolio_escape_log_num) != n_mem: + self._reset_support_diagnostics() + if not self._member_index: + self._member_index = dict((id(mem), i) + for i, mem in enumerate(self.portfolio_realizations)) + lw = numpy.asarray(self.identity_convert(log_weights), dtype=float) + fin = numpy.isfinite(lw) + if not bool(numpy.any(fin)): + return + mx = float(numpy.max(lw[fin])) + # max-subtracted linear weights; the offset mx is carried back in log space so chunks + # with wildly different scales combine exactly. + u = numpy.where(fin, numpy.exp(lw - mx), 0.0) + tot = float(numpy.sum(u)) + if not (tot > 0): + return + n_here = len(u) + ln_tot = float(numpy.log(tot) + mx) + self.portfolio_weight_log_total = numpy.logaddexp(self.portfolio_weight_log_total, ln_tot) + + # --- soft comparator: weight share by DRAWING member. draw() lays the pooled chunk out + # as contiguous per-member blocks in _chunk_members order, with the counts recorded in + # _chunk_fractions, so the blocks are recoverable exactly (do NOT re-derive them from + # self.portfolio_weights -- those have already been updated for the NEXT chunk). + fracs = getattr(self, '_chunk_fractions', None) + members = getattr(self, '_chunk_members', []) + if fracs is not None and len(members) == len(fracs): + counts = numpy.rint(numpy.asarray(fracs, dtype=float) * n_here).astype(int) + start = 0 + for cnt, mem in zip(counts, members): + if cnt <= 0: + continue + end = min(start + int(cnt), n_here) + idx = self._member_index.get(id(mem)) + if idx is not None and end > start: + s_here = float(numpy.sum(u[start:end])) + if s_here > 0: + self.portfolio_share_log_num[idx] = numpy.logaddexp( + self.portfolio_share_log_num[idx], numpy.log(s_here) + mx) + start = end + + # --- escaped mass, per member that HAS a density this chunk. parts[id(m)] is + # frac_m*q_m with frac_m > 0 by construction, so it vanishes exactly where q_m does. + _this_chunk = np.full(n_mem, np.nan) + for idx, mem in enumerate(self.portfolio_realizations): + pc = parts.get(id(mem)) + if pc is None: + continue # member drew nothing this chunk: no density evaluated, no evidence + self.portfolio_escape_log_den[idx] = numpy.logaddexp( + self.portfolio_escape_log_den[idx], ln_tot) + self.portfolio_escape_nsamp[idx] += n_here + zero = numpy.asarray(pc, dtype=float) <= 0.0 + n_zero = int(numpy.count_nonzero(zero)) + _this_chunk[idx] = 0.0 + if n_zero: + self.portfolio_escape_n[idx] += n_zero + self.portfolio_escape_hard[idx] = True + s_esc = float(numpy.sum(u[zero])) + _this_chunk[idx] = min(1.0, s_esc / tot) + if s_esc > 0: + self.portfolio_escape_log_num[idx] = numpy.logaddexp( + self.portfolio_escape_log_num[idx], numpy.log(s_esc) + mx) + self.portfolio_escape_history.append(_this_chunk) + except Exception as e: + print(" [portfolio] support diagnostic skipped this chunk ({}: {})".format( + type(e).__name__, e)) + + def support_diagnostics(self): + """Finalize the support-mismatch statistics. Returns a dict; see + _update_support_diagnostics for definitions and limits.""" + n_mem = len(getattr(self, 'portfolio_escape_log_num', [])) + esc = np.full(n_mem, np.nan) + shr = np.full(n_mem, np.nan) + for m in range(n_mem): + if np.isfinite(self.portfolio_escape_log_den[m]): + esc[m] = float(np.exp(min(0.0, self.portfolio_escape_log_num[m] + - self.portfolio_escape_log_den[m]))) + if np.isfinite(self.portfolio_weight_log_total): + shr[m] = float(np.exp(min(0.0, self.portfolio_share_log_num[m] + - self.portfolio_weight_log_total))) + hard = np.asarray(getattr(self, 'portfolio_escape_hard', np.zeros(n_mem, dtype=bool))) + # HEADLINE: the worst hard-edged member. A member whose density is nowhere exactly zero + # (a live GMM) trivially scores 0 and must not dilute the maximum -- reporting a mean over + # all members would let a soft member hide a fully-escaped warm AV. + esc_max = float(np.nanmax(esc[hard])) if bool(np.any(hard)) else 0.0 + hist = np.asarray(getattr(self, 'portfolio_escape_history', []) or + np.zeros((0, n_mem)), dtype=float) + # EARLY signal: the first chunk in which the member had a density. This is the state the + # WARM SEED put it in, before any contraction, so it isolates seed misplacement from the + # ordinary contraction that inflates the cumulative number. + early = np.full(n_mem, np.nan) + if hist.size: + for m in range(n_mem): + col = hist[:, m] + ok = np.flatnonzero(np.isfinite(col)) + if len(ok): + early[m] = float(col[ok[0]]) + early_max = float(np.nanmax(early[hard])) if bool(np.any(hard)) and np.any( + np.isfinite(early[hard])) else 0.0 + return dict(escaped_mass=esc, weight_share=shr, hard_edged=hard, + escaped_mass_max=esc_max, escaped_mass_history=hist, + escaped_mass_early=early, escaped_mass_early_max=early_max, + escape_n_zero=np.asarray(getattr(self, 'portfolio_escape_n', np.zeros(n_mem))), + escape_n_eval=np.asarray(getattr(self, 'portfolio_escape_nsamp', + np.zeros(n_mem)))) + + + def reset_adaptation(self): + """FULL reset: member proposals AND the portfolio's own adaptive bookkeeping. + + clear_warm_state() rebuilds the MEMBERS, but the portfolio itself also learns during a run + -- draw allocation, per-member quality EMAs and their observation counts, the round-robin + probe pointer, the iteration counter, breakpoint progression and the per-member n_ess + histories. MC-error replicas that inherit those start with scheduling learned from the + earlier replicas, so they are not adaptation-independent and the between-replica scatter + still understates the true error -- which is the entire quantity the replicas exist to + measure. Restores every field to its post-setup value. + """ + self.clear_warm_state() + n = len(self.portfolio) + # draw allocation: back to uniform (or the caller's explicit initial weights) + w0 = getattr(self, '_portfolio_weights_initial', None) + self.portfolio_weights = np.array(w0) if w0 is not None else np.ones(n) / (1.0 * n) + self.portfolio_quality = np.ones(n) + self.portfolio_quality_nobs = np.zeros(n, dtype=int) + self.portfolio_probe_ptr = 0 + self.portfolio_draw_iteration = 0 + self.portfolio_member_ness_history = [[] for _ in range(n)] + # breakpoints are a SCHEDULE (set at setup), not learned state: restore the schedule that + # setup() installed rather than zeroing it, or a replica would activate members on a + # different iteration than the first run did. + bp0 = getattr(self, '_portfolio_breakpoints_initial', None) + if bp0 is not None: + self.portfolio_breakpoints = np.array(bp0) + for _attr in ('portfolio_frozen', 'portfolio_grace_left', 'portfolio_last_revive'): + _v0 = getattr(self, '_' + _attr + '_initial', None) + if _v0 is not None: + setattr(self, _attr, np.array(_v0) if hasattr(_v0, '__len__') else _v0) + + def bootstrap_from_samples(self, samples, params=None, keep_backstop_cold=None, **kwargs): + """Warm-start: forward a seed cloud to every member that supports it (e.g. the + AV/VARAHA member's live volume), EXCEPT the full-support backstop (see below). + Members without bootstrap_from_samples are left cold. A warm start only shapes a + member's proposal, and the portfolio combines members with the balance-heuristic + mixture density (q_mix), so a mis-seeded member costs efficiency rather than bias -- + BUT ONLY WHILE SOME MEMBER STILL COVERS THE SUPPORT. Column + order matches self.params_ordered, which every member shares (add_parameter forwards + to all members in the same order), so no per-member remapping is needed. + + Only VARAHA/AV-style members (those exposing bootstrap_from_samples) are seeded + directly here. GMM / adaptive-Gaussian members cannot be seeded pre-integration + (their internal integrator, hence gmm_dict, does not exist until the first + integrate() call), but they still warm up STRUCTURALLY during the run: the portfolio + lets the cold GMM member adapt its proposal from the warm AV member's high-likelihood + draws, which is what gives the mixture a faster early n_eff than warm-AV alone. An + explicit GMM pre-seed (build an initial gmm_dict from the sample cloud) is a future + enhancement. q_mix keeps any cold/mis-seeded member from biasing the estimate. + + THE BACKSTOP MUST STAY COLD. Seeding EVERY member destroys the coverage invariant the + whole design rests on. `cover_frac` does not save it: it mixes a FINITE number of uniform + points into the seed, and a finite point set occupies only the bins it lands in, so the + seeded live volume is NOT a superset of a cold start. Measured on a tight seed in the + [-5,5]^d box -- fraction of the prior box covered, against 1.0 cold: + + d=2: cover_frac 0.0 / 0.2 / 0.5 / 0.9 -> 0.027 / 0.634 / 0.982 / 1.000 + d=4: -> 0.0015 / 0.028 / 0.104 / 0.620 + d=6: -> 6.3e-05 / 0.00087 / 0.0033 / 0.0287 + + so at d=6 even cover_frac=0.9 leaves 97% of the box unsampled. Before this change a + portfolio warm start narrowed member 0 to V=0.0033 along with everyone else. Keeping + member 0 (the designated backstop, which restrict_member_range also refuses to narrow) + cold makes the mixture-level guarantee TRUE instead of merely asserted. + + BE PRECISE ABOUT WHAT THIS BUYS, because the measurements are not what one expects: + + * With a CORRECT seed it is nearly free -- n_eff 3630/3466/5695 cold-backstop vs + 3876/3461/5621 seeded at d=4, and 5078/5689/5163 vs 5857/5387/5266 at d=6. Within + run-to-run scatter, so it is cheap insurance. That is the case for the default. + * It is NOT what protects the DEFAULT AV+GMM portfolio. The GMM member is a Gaussian + mixture with nonzero density over the whole box, so q_mix never vanishes there + regardless of this setting. With a deliberately displaced seed, |lnZ bias| stayed + <= 0.05 in every d=4 and d=6 run in BOTH arms. That unbounded support is the real + (previously undocumented) reason production has not been biased by warm starts. + * It does NOT rescue a badly mismatched seed. In an ALL-AV portfolio (every component + a hard-edged box) with a displaced seed at d=6, the cold backstop still gave lnZ bias + -1.1 to -4.2 nats with n_eff 3-9, versus -1.0 to -6.8 seeded. A uniform member finds + a sharp 6-D peak too rarely to carry the integral within budget. Coverage in + principle is necessary, not sufficient -- for a mismatched seed the mitigations are + detection (L0 rescue) and not warm-starting across dissimilar points. + + `keep_backstop_cold`: None (default) uses self.portfolio_warmstart_backstop_cold, itself + True by default; pass False to restore the old seed-everything behaviour. + + Returns the number of members warm-started (0 is fine; the portfolio still runs).""" + samples = np.asarray(samples) + if keep_backstop_cold is None: + keep_backstop_cold = bool(getattr(self, 'portfolio_warmstart_backstop_cold', True)) + # Which members must retain full support? If range restriction is in play it already + # computed them; otherwise it is member 0 by the same convention. + # The invariant is NOT "member 0 must be cold" -- it is "SOME member must have support + # everywhere". A GMM/ensemble member satisfies that inherently: it carries an explicit + # uniform defensive component (gmm_defensive_frac, default 0.05) plus Gaussian tails, so + # q_mix never vanishes however it is seeded. Measured: in the default [AV, GMM] portfolio + # a deliberately displaced seed left |lnZ bias| <= 0.05 whether or not a member was held + # cold, while an ALL-AV portfolio gave -1.0 to -6.8 nats. + # So hold a member cold ONLY when EVERY member has compact support. Doing it + # unconditionally disables the AV warm start in [AV, GMM] -- member 0 IS the AV member -- + # to buy a guarantee the GMM member already provides. The merge gate caught exactly that. + # Default FALSE: a sampler must DECLARE full support to be counted. Defaulting to True + # meant any member that simply had not been annotated was treated as the coverage + # guarantee -- the safe default is to assume compact and keep a cold backstop. + # And a nominally broad member does NOT count if it has been RANGE-RESTRICTED: its + # proposal is confined to a sub-box, so it no longer covers the prior. Without this, + # [unrestricted AV, restricted GMM] reported _full_support_members == [0] and then + # warm-started and contracted member 0 anyway, leaving nothing covering the prior box -- + # the silent low bias this whole mechanism exists to prevent. + if getattr(self, '_has_restricted_member', False): + _unrestricted = set(getattr(self, '_full_support_members', []) or []) + else: + _unrestricted = set(range(len(self.portfolio_realizations))) + _has_broad = any(getattr(m, 'has_unbounded_support', False) and (i in _unrestricted) + for i, m in enumerate(self.portfolio_realizations)) + if keep_backstop_cold and _has_broad: + keep_backstop_cold = False + print(" [portfolio] warm-starting all members: a full-support member is present " + "(defensive mixture), so no cold backstop is needed") + _backstop = set(getattr(self, '_full_support_members', None) or [0]) if keep_backstop_cold else set() + if len(_backstop) >= len(self.portfolio_realizations): + _backstop = set([0]) # never refuse to warm-start EVERY member + self._warmstart_backstop_cold = sorted(_backstop) + n_warmed = 0 + for indx, member in enumerate(self.portfolio_realizations): + if indx in _backstop: + print(" [portfolio] member {} kept COLD as the full-support backstop " + "(cover_frac cannot make a seeded grid cover the prior box; see " + "bootstrap_from_samples docstring)".format(indx)) + continue + if not hasattr(member, 'bootstrap_from_samples'): + continue + try: + member.bootstrap_from_samples(samples, params=params, **kwargs) + n_warmed += 1 + except Exception as e: + print(" [portfolio] member {} warm-start skipped ( {} )".format(indx, e)) + print(" [portfolio] warm-started {}/{} members directly (others warm structurally)".format( + n_warmed, len(self.portfolio_realizations))) + return n_warmed def setup(self, **kwargs): self.extra_args =kwargs # may need to pass/use during the 'update' step + # allow the driver/CLI to tune the freeze-protection knobs. A None means "not set on + # the CLI" (optparse default), so keep the current value rather than clobbering it. + def _kw_keep(name): + v = kwargs.get(name, None) + if v is not None: + setattr(self, name, v) + _kw_keep('portfolio_freeze_wt') + _kw_keep('portfolio_grace_iters') + _kw_keep('portfolio_revive_period') + _kw_keep('portfolio_varaha_never_freeze') + _kw_keep('portfolio_plateau_revive') + _kw_keep('portfolio_adaptive_alloc') + _kw_keep('portfolio_quality_signal') + _kw_keep('portfolio_alloc_exponent') + _kw_keep('portfolio_alloc_floor') + _kw_keep('portfolio_quality_decay') + _kw_keep('portfolio_probe_period') + _kw_keep('portfolio_probe_frac') + _kw_keep('portfolio_weight_clip') + _kw_keep('portfolio_varaha_min_frac') + _kw_keep('portfolio_varaha_max_frac') + _kw_keep('portfolio_warmstart_backstop_cold') if 'oracle_realizations' in kwargs: if kwargs['oracle_realizations']: self.oracle_realizations = kwargs['oracle_realizations'] # might not have been initialized earlier @@ -214,12 +926,110 @@ def setup(self, **kwargs): portfolio_extra_args = kwargs['portfolio_args'] else: print(" PORTFOLIO - format ERROR ", kwargs['portfolio_args']) - for indx, member in enumerate(self.portfolio): + # RANGE RESTRICTION (opt-in): narrow ONE OR MORE members to a sub-box, so a member can put + # its fixed bin budget where the posterior actually is. Applied HERE -- after add_parameter + # (which forwards identical limits AND the shared prior callables to every member) and + # BEFORE member.setup() (which rebuilds my_ranges/dx/dx0/V_s/binunique/ninbin/V from + # llim/rlim) -- so no derived state can go stale. + # + # WHY THIS NEEDS NO NORMALIZATION CORRECTION. The estimator weights are + # lnL + log(joint_p_prior) - log(q_mix). `joint_p_prior` comes from each member's STORED + # prior callables (AV.prior_prod), which are absolute densities over the ORIGINAL physical + # ranges and never consult llim/rlim -- so a narrowed member still reports the TRUE global + # prior. A member's range is therefore purely a PROPOSAL choice: proposals need not cover + # the prior, only the MIXTURE must cover the support of L*p. Hence no prior renormalization + # and no ln(V_sub/V_full) term. The invariant that makes this true is enforced below. + _restrict = kwargs.get('portfolio_restrict_ranges', None) + if _restrict: + if len(_restrict) != len(self.portfolio_realizations): + raise Exception("portfolio_restrict_ranges must align with the member list " + "({} entries for {} members)".format(len(_restrict), + len(self.portfolio_realizations))) + _n_restricted = 0 + for indx, member in enumerate(self.portfolio_realizations): + spec = _restrict[indx] + if not spec: + continue + for p, (lo, hi) in dict(spec).items(): + if p not in member.llim: + raise Exception("portfolio_restrict_ranges: member {} has no parameter {!r}".format(indx, p)) + # NARROW ONLY -- clip into the existing range. Widening a member beyond the + # prior's support would sample where the prior callable is not normalized. + lo_new = max(float(lo), float(member.llim[p])) + hi_new = min(float(hi), float(member.rlim[p])) + if not (hi_new > lo_new): + raise Exception("portfolio_restrict_ranges: empty sub-range for {!r} on member {}".format(p, indx)) + member.llim[p] = lo_new + member.rlim[p] = hi_new + _n_restricted += 1 + print(" PORTFOLIO: member {} RESTRICTED to sub-box {}".format(indx, dict(spec))) + # COVERAGE INVARIANT: at least one member must keep FULL support, otherwise the mixture + # no longer covers L*p outside the union of sub-boxes and the integral is biased low by + # the missing mass (silently -- n_eff can even look BETTER). Refuse rather than bias. + # UNION with any restrictions registered through the public restrict_member_range() API. + # Both entry points must feed the SAME bookkeeping: computing _full_support_members from + # _restrict alone would declare an API-restricted member "full support" and hand it the + # draw floor that is meant to protect a genuinely unrestricted component. + _restricted_set = set(i for i in range(len(self.portfolio_realizations)) if _restrict[i]) + _restricted_set |= set(getattr(self, 'member_range_overrides', {})) + if len(_restricted_set) >= len(self.portfolio_realizations): + raise Exception( + "portfolio_restrict_ranges: every member is restricted, so no member retains " + "full support. The mixture would not cover L*p outside the sub-boxes and the " + "integral would be biased low with no diagnostic. Leave at least one member " + "unrestricted (it is the defensive component).") + self._has_restricted_member = bool(_restricted_set) + # index of a full-support member: the per-member draw floor below protects it + self._full_support_members = [i for i in range(len(self.portfolio_realizations)) + if i not in _restricted_set] + + # Snapshot the post-setup values of everything reset_adaptation() restores, so a replica + # returns to THIS state rather than to a hard-coded guess. + self._portfolio_weights_initial = np.array(self.portfolio_weights) + self._portfolio_breakpoints_initial = np.array(self.portfolio_breakpoints) + for _attr in ('portfolio_frozen', 'portfolio_grace_left', 'portfolio_last_revive'): + if hasattr(self, _attr): + _v = getattr(self, _attr) + setattr(self, '_' + _attr + '_initial', + np.array(_v) if hasattr(_v, '__len__') else _v) + + # CONSUMED CHECK. restrict_member_range() only takes effect if it was called BEFORE + # add_parameter forwarded that parameter to the members. A restriction naming a parameter + # that never arrives (typo, or the call came too late) used to be a SILENT no-op: the caller + # believes a member is focused on the posterior while it still samples the full box. Fail + # loudly instead -- a narrowing that quietly did nothing is a wasted member, not a safe one. + _pending = getattr(self, '_pending_range_overrides', set()) + if _pending: + raise Exception( + "restrict_member_range: {} restriction(s) were never applied: {}. Either the " + "parameter name does not exist on that member, or restrict_member_range() was " + "called AFTER add_parameter() -- it must be called before.".format( + len(_pending), sorted(_pending))) + + # Iterate the INSTANTIATED samplers (portfolio_realizations), NOT self.portfolio: the + # latter may hold modules/names (see __init__), which lack .setup(), so member setup was + # silently skipped -> a cold member's internal state (AV my_ranges, GMM integrator) was + # never built and draw_simplified failed. Setting up the realizations fixes AV+GMM cold. + # REMEMBER each member's setup arguments. clear_warm_state() has to re-run setup() to + # restore a member's cold grid, and calling it bare would silently DISCARD the member's + # configuration: mcsamplerEnsemble.setup() rebuilds its dimension grouping and re-reads + # n_comp / gmm_adapt / correlate_all_dims etc from kwargs, so a configured (0,1) GMM with + # n_comp=3 and adaptation off comes back as separate (0,), (1,) groups with n_comp + # defaulted and gmm_adapt=None. Store the exact args and replay them. + self._member_setup_args = [None] * len(self.portfolio_realizations) + self._oracle_setup_args = [None] * len(self.oracle_realizations) + for indx, member in enumerate(self.portfolio_realizations): if hasattr(member, 'setup'): print(" PORTFOLIO setup ", member, portfolio_extra_args[indx]) args_here = {} args_here.update(kwargs) args_here.update(portfolio_extra_args[indx]) + # snapshot BEFORE setup: the member (or its integrator) may mutate these in place + # A portfolio member may be the mixture's ONLY full-support component, so ask for + # the defensive component on every fit path for OUR members. Standalone users of + # the same sampler are unaffected -- measured, it costs real n_eff at d>=6. + args_here.setdefault('gmm_defensive_all_paths', True) + self._member_setup_args[indx] = self._snapshot_setup_args(args_here) member.setup(**args_here) for indx, member in enumerate(self.oracle_realizations): if hasattr(member, 'setup'): @@ -227,9 +1037,64 @@ def setup(self, **kwargs): args_here = {} args_here.update(kwargs) args_here.update(portfolio_extra_args[indx]) + self._oracle_setup_args[indx] = self._snapshot_setup_args(args_here) member.setup(**args_here) member.params_ordered = list(self.params_ordered) # enforce parameters for oracle being sane + def _adaptive_allocation(self, ness_now, frac_now, iteration): + """Adaptive-probe draw allocation (see __init__). Returns the next chunk's per-member + draw weights. Decouples a QUALITY estimate (EMA of n_ess, updated only from chunks where + the member had a fair allocation) from the ALLOCATION (quality^exponent, floored), and + round-robin PROBES one member per `probe_period` chunks at a raised share so a suppressed + member can prove itself. Unbiased for any allocation (q_mix handles correctness).""" + m = len(self.portfolio) + if m <= 1: + return np.ones(m) + # 'global' (marginal pooled n_eff) and 'credit' (MIS credit) are both ZERO-based + # contribution measures; only the legacy Kish n_ess signal is floored at 1. + _global = self.portfolio_quality_signal in ('global', 'credit') + _floor_obs = 0.0 if _global else 1.0 # contribution is zero-based; Kish n_ess is >= 1 + obs = np.asarray(ness_now, dtype=float) + obs = np.where(np.isfinite(obs), obs, _floor_obs) + frac = np.asarray(frac_now, dtype=float) + # 1) update QUALITY. With the 'ness' signal a member drawn at the floor has too few/too + # noisy samples to trust, so only fair-allocation chunks count. The 'global' + # contribution signal is SELF-CORRECTING at low allocation (a starved peak-covering + # member shows inflated weights), so every chunk is informative -- no gating needed. + fair = 0.0 if _global else 0.9 / m + a = self.portfolio_quality_decay + for k in range(m): + if frac[k] > 0 and frac[k] >= fair: + _o = max(obs[k], _floor_obs) + if self.portfolio_quality_nobs[k] == 0: + # first real observation: adopt it outright. The 'global' contribution signal + # has an arbitrary scale, so EMA-ing from the placeholder 1.0 would bias it. + self.portfolio_quality[k] = _o + else: + self.portfolio_quality[k] = (1 - a) * self.portfolio_quality[k] + a * _o + self.portfolio_quality_nobs[k] += 1 + # 2) base allocation ~ quality^exponent above a floor. For 'ness' we use the EXCESS over the + # degenerate n_ess=1 (a member at n_ess 1 contributes nothing); the 'global' contribution + # is already zero-based. + q = np.maximum(self.portfolio_quality - _floor_obs, 0.0) + if np.sum(q) <= 0: + base = np.ones(m) / m + else: + w = (q / np.sum(q)) ** self.portfolio_alloc_exponent + base = w / np.sum(w) + base = self.portfolio_alloc_floor + base * (1.0 - m * self.portfolio_alloc_floor) + base = base / np.sum(base) + # 3) round-robin probe: raise ONE member to >= probe_frac every probe_period chunks so an + # under-observed member gets a fair look next chunk (breaks the under-observation trap). + if self.portfolio_probe_period > 0 and (iteration % self.portfolio_probe_period == 0): + k = self.portfolio_probe_ptr % m + self.portfolio_probe_ptr += 1 + if base[k] < self.portfolio_probe_frac: + base = base * (1.0 - self.portfolio_probe_frac) / max(1e-12, 1.0 - base[k]) + base[k] = self.portfolio_probe_frac + base = base / np.sum(base) + return base + def draw(self,n_samples, *args, **kwargs): """ draw n_samples @@ -260,7 +1125,27 @@ def draw(self,n_samples, *args, **kwargs): # if only one method is active, just call the low-level function if len(indx_active) == 1: - joint_p_s, joint_p_prior, rv = self.portfolio[indx_active[0]].draw_simplified(n_samples, *self.params_ordered, **kwargs) + # Single-member fast path: q_mix degenerates to this member's own density. If that lone + # member is a RESTRICTED one, nothing covers L*p outside its sub-box for this chunk and + # the estimate is biased low with no diagnostic. (Reachable when activation breakpoints + # delay the full-support member.) Refuse rather than silently bias. + if (getattr(self, '_has_restricted_member', False) + and int(indx_active[0]) not in set(getattr(self, '_full_support_members', []))): + raise Exception( + "mcsamplerPortfolio: the only ACTIVE member (realization {}) has a RESTRICTED " + "range, so this chunk has no full-support component and the integral would be " + "biased low. Give the full-support member an activation breakpoint of 0." + .format(int(indx_active[0]))) + only_member = self.portfolio_realizations[indx_active[0]] + joint_p_s, joint_p_prior, rv = only_member.draw_simplified(n_samples, *self.params_ordered, **kwargs) + # The portfolio aggregates on the host (self.xpy is numpy); members + # may be GPU-backed (cupy), so bring their draws to the host. + joint_p_s = identity_convert(joint_p_s); joint_p_prior = identity_convert(joint_p_prior); rv = identity_convert(rv) + # Record which members produced this chunk, and each member's SAMPLING + # FRACTION (n_from_member / n_total). integrate_log uses these to form + # the balance-heuristic mixture density q_mix = sum_m frac_m * q_m. + self._chunk_members = [only_member] + self._chunk_fractions = np.array([1.0]) else: # Identify number of samples per member of the portfolio. Can be zero. n_samples_per_member = ((np.array(weights_active))*n_samples).astype(int) @@ -273,6 +1158,30 @@ def draw(self,n_samples, *args, **kwargs): n_samples_per_member[-1] = 0 n_samples_per_member[-2] = n_samples - np.sum(n_samples_per_member[0:-2]) + # PER-MEMBER DRAW FLOOR for full-support members when some member is RESTRICTED. + # A member that draws 0 samples this chunk contributes NOTHING to q_mix (the mixture loop + # skips frac_m <= 0), so its coverage vanishes for that chunk. That is harmless when all + # members share a support, but FATAL once a member has been narrowed: the full-support + # member is the only thing covering L*p outside the sub-box, and a chunk where it is + # rounded to zero draws is a chunk with an uncovered region -- the silent low-bias failure + # this whole design exists to avoid. The VARAHA share band is a GROUP constraint over all + # is_varaha members, so it does NOT protect an individual full-box AV against a restricted + # sibling absorbing the group's share. Enforce a per-member minimum of 1 draw here. + if getattr(self, '_has_restricted_member', False): + # CAREFUL: n_samples_per_member is indexed by POSITION WITHIN portfolio_active (the + # breakpoint-filtered subset), while _full_support_members holds REALIZATION indices. + # Map one to the other; indexing directly by realization index protects the wrong member + # as soon as any member is still behind its activation breakpoint. + _full_set = set(getattr(self, '_full_support_members', [])) + _full = [pos for pos, ridx in enumerate(indx_active) if int(ridx) in _full_set] + for i in _full: + if n_samples_per_member[i] < 1: + # take the deficit from the largest member so the total is preserved exactly + j = int(np.argmax(n_samples_per_member)) + if j != i and n_samples_per_member[j] > 1: + n_samples_per_member[j] -= 1 + n_samples_per_member[i] = 1 + n_index_start_per_member = np.zeros(len(portfolio_active),dtype=int) n_index_start_per_member[1:] = np.cumsum(n_samples_per_member)[:-1] @@ -282,17 +1191,27 @@ def draw(self,n_samples, *args, **kwargs): joint_p_s_here, joint_p_prior_here, rv_here = member.draw_simplified( n_samples_per_member[indx_member], *self.params_ordered, **kwargs ) - # type convert as needed, to GPU - if not(isinstance( type(joint_p_s_here), type(joint_p_s))): - joint_p_s_here = self.identity_convert_togpu(joint_p_s_here) - joint_p_prior_here = self.identity_convert_togpu(joint_p_prior_here) - rv_here = self.identity_convert_togpu(rv_here) + # Bring member draws to the host backend the portfolio aggregates in + # (self.xpy is numpy). identity_convert is cupy.asnumpy when a member + # is GPU-backed, else a no-op. (The previous isinstance(type(x),..) + # guard never fired, leaving cupy arrays to collide with numpy ones.) + joint_p_s_here = identity_convert(joint_p_s_here) + joint_p_prior_here = identity_convert(joint_p_prior_here) + rv_here = identity_convert(rv_here) indx_start = int(n_index_start_per_member[indx_member]) indx_end = indx_start + int(n_samples_per_member[indx_member]) joint_p_s[indx_start:indx_end] = joint_p_s_here joint_p_prior[indx_start:indx_end] = joint_p_prior_here rv[:,indx_start:indx_end] = rv_here - + + # Record the ACTUAL per-member sampling fractions for this chunk (the + # counts actually drawn, not the raw portfolio_weights). These are the + # w_m in the balance-heuristic mixture density q_mix = sum_m w_m q_m + # that integrate_log builds; using the true drawn fractions is what + # keeps the deterministic-mixture estimator exactly unbiased. + self._chunk_members = list(portfolio_active) + self._chunk_fractions = np.array(n_samples_per_member, dtype=float) / float(n_samples) + # # Cache the samples we chose. REQUIRED # @@ -317,8 +1236,18 @@ def integrate(self, lnF, *args, xpy=xpy_default,**kwargs): def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): - xpy_here = self.xpy - + # The portfolio AGGREGATES on the host: draw() brings member draws to the + # host, so force the running-estimate math onto numpy/scipy regardless of + # what the driver set self.xpy to (it sets cupy for GPU members). Members + # still do their own heavy sampling/adaptation on their own backend. The + # INTEGRAND, however, may be device-native (the real vectorized GPU ILE + # likelihood) or host-native (synthetic/CI): _eval_integrand() below feeds + # it device-first and falls back to host, so both work. + self.xpy = numpy + xpy_here = numpy + xpy = numpy + special_here = special # scipy.special (host); statutils uses this + # # Determine stopping conditions # @@ -328,6 +1257,12 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): convergence_tests = kwargs["convergence_tests"] if "convergence_tests" in kwargs else None save_no_samples = kwargs["save_no_samples"] if "save_no_samples" in kwargs else None portfolio_wt_func = kwargs['portfolio_schedule'] if 'portfolio_schedule' in kwargs else portfolio_default_weights + # allow a per-integration override of the adaptive-probe allocation (else use the instance + # default set at init/setup); falling back to the legacy n_ess reweighting when off. + use_adaptive_alloc = kwargs.get('portfolio_adaptive_alloc', self.portfolio_adaptive_alloc) + # per-integration override of the weight clip (else the instance default from init/setup) + if 'portfolio_weight_clip' in kwargs: + self.portfolio_weight_clip = kwargs['portfolio_weight_clip'] # # Adaptive sampling parameters @@ -384,6 +1319,9 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): n_zero_prior =0 it_max_oracle = 7 it_now =0 + # the support-mismatch diagnostic describes THIS integral: a second point reusing the same + # sampler object must not inherit the previous point's escaped mass. + self._reset_support_diagnostics() if 'integrand' in self._rvs: # remove conflict del self._rvs['integrand'] @@ -413,22 +1351,174 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): params.extend(item) else: params.append(item) - unpacked = unpacked0 = rv #numpy.hstack([r.flatten() for r in rv]).reshape(len(args), -1) - unpacked = dict(list(zip(params, unpacked))) - - # Evaluate function, protecting argument order - if 'no_protect_names' in kwargs: - lnL = lnF(*unpacked0) # do not protect order + # Evaluate the integrand. rv is on the host (see draw()). The real + # GPU ILE likelihood is DEVICE-native (wants cupy); synthetic/CI + # integrands are host-native. Feed device-first, fall back to host on + # a type error, and remember the choice (same contract as the AV + # integrator). lnL is brought back to the host for aggregation. + def _eval_integrand(cols): + if 'no_protect_names' in kwargs: + return lnF(*cols) + return lnF(**dict(list(zip(params, cols)))) + if getattr(self, '_integrand_wants_host', False) or not cupy_ok: + lnL = _eval_integrand(rv) else: - lnL= lnF(**unpacked) # protect order using dictionary - # take log if we are NOT using lnL - if cupy_ok: - if not(isinstance(lnL,cupy.ndarray)): - lnL = identity_convert_togpu(lnL) # send to GPU, if not already there + try: + lnL = _eval_integrand(identity_convert_togpu(rv)) + except (TypeError, ValueError): + self._integrand_wants_host = True + lnL = _eval_integrand(rv) + # bring lnL back to the host for the host-side aggregation + lnL = identity_convert(lnL) + + # ---- BALANCE-HEURISTIC (deterministic-mixture) sampling density ---- + # The pooled draw is, by construction, a sample from the MIXTURE + # q_mix(theta) = sum_m frac_m * q_m(theta), + # where frac_m = (# samples member m contributed)/n for THIS chunk and + # q_m is member m's own sampling density evaluated at theta. Using + # q_mix (rather than each sample's own member density -- the previous + # STRATIFIED estimator) makes the estimate unbiased for ANY member + # weights, provided the mixture covers the peak (Veach & Guibas MIS + # balance heuristic). A broad member with even a small weight then + # guarantees coverage, so a wrongly-contracted member can no longer + # drive the integral low. See portfolio_default_weights: n_ess-based + # weighting made the old stratified denominator UNSAFE. + # + # We require EVERY active member (that drew >0 samples) to expose a + # sampling_density; if any does not (e.g. an AC-histogram member with + # no pointwise density yet), we fall back to the legacy per-member + # joint_p_s so those portfolios keep running unchanged. + use_mixture = kwargs['portfolio_use_mixture_density'] if 'portfolio_use_mixture_density' in kwargs else True + q_mix = None + # drop last chunk's per-member densities: keeping them would let a chunk in which the + # mixture could not be formed be scored against the PREVIOUS chunk's supports. + self._chunk_mix_parts = None + if use_mixture: + X_all = numpy.asarray(identity_convert(rv), dtype=float) + if X_all.shape[0] == len(self.params_ordered): + X_all = X_all.T # -> (N, ndim), columns in params_ordered order + members_here = getattr(self, '_chunk_members', []) + fracs_here = getattr(self, '_chunk_fractions', None) + if len(members_here) > 0 and fracs_here is not None: + acc = numpy.zeros(X_all.shape[0], dtype=float) + _mix_parts = {} + all_ok = True + any_active = False + for frac_m, member_m in zip(fracs_here, members_here): + if frac_m <= 0: + continue # member drew nothing this chunk + dens_fn = getattr(member_m, 'sampling_density', None) + q_m = dens_fn(X_all) if dens_fn is not None else None + if q_m is None: + all_ok = False + break + _contrib_m = float(frac_m) * numpy.asarray(identity_convert(q_m), dtype=float) + acc = acc + _contrib_m + # retain frac_m*q_m per member: the balance-heuristic credit + # frac_m q_m / q_mix is the MIS share of each sample owed to member m + _mix_parts[id(member_m)] = _contrib_m + any_active = True + if all_ok and any_active: + # every pooled sample was drawn by some active member, so + # q_mix >= frac*q_m(own) > 0 there; floor only guards FP. + # UNDERFLOW DIAGNOSTIC: mathematically acc>0 for every drawn sample, so any + # acc==0 is a floating-point UNDERFLOW of the linear-space density sum. The + # 1e-300 floor then turns it into a spurious ~target/1e-300 weight, which can + # single-handedly dominate the pooled estimator. Count these so a numerical + # artifact can be told apart from genuine heavy-tailed weights (the former + # wants a log-space q_mix / member fix, the latter wants weight clipping). + _n_uf = int(numpy.sum(acc <= 0)) + if _n_uf > 0: + self.portfolio_qmix_underflow += _n_uf + print(" PORTFOLIO: q_mix UNDERFLOW on {}/{} samples this chunk" + " (density summed to 0 -> floored 1e-300 -> spurious huge weight;" + " cumulative {})".format(_n_uf, len(acc), self.portfolio_qmix_underflow)) + q_mix = numpy.maximum(acc, 1e-300) + self._chunk_mix_parts = _mix_parts + if q_mix is not None: + joint_p_s = q_mix # deterministic-mixture denominator + else: + # The legacy stratified per-member density is only valid when every member shares the + # SAME support. If any member has been given a RESTRICTED range (see + # portfolio_restricted_members) the stratified estimator is silently WRONG -- it does + # not form the true mixture denominator -- so refuse rather than return a biased + # number. Unequal supports are exactly the configuration the fallback cannot handle. + if getattr(self, '_has_restricted_member', False): + raise Exception( + "mcsamplerPortfolio: a member has a RESTRICTED sampling range, but the " + "balance-heuristic q_mix could not be formed (some active member lacks " + "sampling_density). The legacy stratified density is invalid for members " + "with unequal support and would bias the integral; refusing to continue.") + if use_mixture and getattr(self, '_warned_no_mixture', False) is False: + print(" PORTFOLIO: some active member lacks sampling_density; " + "falling back to legacy stratified per-member density.") + self._warned_no_mixture = True log_integrand =lnL + self.xpy.log(joint_p_prior) - self.xpy.log(joint_p_s) # tempering_exp done inside the update proposal, NOT here log_weights = lnL + self.xpy.log(joint_p_prior) - self.xpy.log(joint_p_s) + # NaN guard: a frozen/degenerate member (e.g. an un-contracted VARAHA) can emit NaN + # samples -> NaN lnL / joint_p_s -> NaN weights, which otherwise propagate into the + # aggregation and the reported crash ("boolean index did not match ... 10000 vs 9882" + # when a NaN mask is applied downstream). Map any non-finite weight to -inf (zero + # weight) IN PLACE, keeping the array length fixed so no mask-size mismatch can arise. + _bad = ~self.xpy.isfinite(log_integrand) + if bool(self.identity_convert(self.xpy.any(_bad))): + log_integrand = self.xpy.where(_bad, -self.xpy.inf, log_integrand) + log_weights = self.xpy.where(_bad, -self.xpy.inf, log_weights) + + # SUPPORT-MISMATCH DIAGNOSTIC (off-path; see _update_support_diagnostics). Placed on + # the TRUE, unclipped weights and before any adaptation, so it describes the estimator's + # own weights. Reads only self._chunk_mix_parts, which q_mix construction already built. + self._update_support_diagnostics(log_weights, q_mix) + + # WEIGHT CLIPPING (truncated IS; OPT-IN) -- PROPOSAL-TRAINING INPUT ONLY. + # Clipping is BIASED and also distorts n_ess, so its scope is deliberately narrow: it + # produces log_weights_adapt, which is fed ONLY to member.update_sampling_prior (the GMM + # covariance fit), so that a single enormous weight cannot make that fit degenerate. + # Everything else uses the TRUE weights: + # * log_integrand -> the ESTIMATE (ln Z, eff_samp): UNCLIPPED -> exactly unbiased. + # * the per-member n_ess REPORT and the ALLOCATION signal: UNCLIPPED -> undistorted + # (clipping flattens weights and would INFLATE the clipped member's Kish n_ess, + # perversely rewarding the very member whose weights had to be clipped). + # This is also strictly better than DROPPING a chunk that clipped: dropping conditional + # on "a big weight appeared" is data-dependent selection and would bias ln Z low. The + # tracked withheld mass is a TAIL DIAGNOSTIC (how much weight the proposal fit ignored). + log_weights_adapt = log_weights + if self.portfolio_weight_clip and self.portfolio_weight_clip > 0: + _lw = numpy.asarray(self.identity_convert(log_weights), dtype=float) + _fin = numpy.isfinite(_lw) + if bool(numpy.any(_fin)): + _mx = float(numpy.max(_lw[_fin])) + _u = numpy.where(_fin, numpy.exp(_lw - _mx), 0.0) + _n_here = max(1, len(_u)) + _total = float(numpy.sum(_u)) + _tau = self.portfolio_weight_clip * numpy.sqrt(_n_here) * (_total / _n_here) + if _total > 0: + self.portfolio_clip_log_total = numpy.logaddexp( + self.portfolio_clip_log_total, numpy.log(_total) + _mx) + _over = _u > _tau + _n_over = int(numpy.sum(_over)) + if _n_over > 0 and _tau > 0: + _removed = float(numpy.sum(_u[_over] - _tau)) + if _removed > 0: + self.portfolio_clip_log_removed = numpy.logaddexp( + self.portfolio_clip_log_removed, numpy.log(_removed) + _mx) + self.portfolio_clip_n += _n_over + _u = numpy.minimum(_u, _tau) + # ADAPTATION copy only -- log_integrand (the estimator) is deliberately + # untouched, so ln Z and n_eff remain exactly unbiased. + log_weights_adapt = numpy.where(_u > 0, + numpy.log(numpy.maximum(_u, 1e-300)) + _mx, + -numpy.inf) + _frac = float(numpy.exp(self.portfolio_clip_log_removed + - self.portfolio_clip_log_total)) \ + if numpy.isfinite(self.portfolio_clip_log_removed) else 0.0 + _frac = min(max(_frac, 0.0), 1.0 - 1e-15) + print(" PORTFOLIO: proposal-fit weight-clip tau={:.3e}(rel max) clipped {} " + "this chunk ({} total); cumulative tail mass withheld from PROPOSAL FIT" + " ={:.3e} (estimator + n_ess report + allocation all unclipped)" + .format(_tau, _n_over, self.portfolio_clip_n, _frac)) if save_intg: # FIXME: See warning at beginning of function. The prior values @@ -455,9 +1545,9 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # n, Mean, error tracked by statutils structure if current_log_aggregate is None: - current_log_aggregate = init_log(log_integrand,xpy=xpy,special=xpy_special_default) + current_log_aggregate = init_log(log_integrand,xpy=xpy,special=special_here) else: - current_log_aggregate = update_log(current_log_aggregate, log_integrand,xpy=xpy,special=xpy_special_default) + current_log_aggregate = update_log(current_log_aggregate, log_integrand,xpy=xpy,special=special_here) outvals = finalize_log(current_log_aggregate,xpy=xpy) self.ntotal = current_log_aggregate[0] # effective samples @@ -499,44 +1589,186 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): n_index_start_per_member = np.zeros(len(self.portfolio_realizations),dtype=int) n_index_start_per_member[1:] = np.cumsum(n_samples_per_member)[:-1] + # GLOBAL-IMPACT signal: each member's MARGINAL GAIN IN POOLED n_eff PER SAMPLE. + # Pooled Kish n_eff = S^2/Q with S = sum(w), Q = sum(w^2) over ALL members' samples. + # One extra sample from member m adds (in expectation) mean_w_m to S and mean_w2_m to Q, + # so d(n_eff)/dn_m divided by n_eff gives the relative per-sample gain + # g_m = 2*mean_w_m/S - mean_w2_m/Q . + # This is the quantity the allocation should maximize: it credits a member for the + # weight MASS it supplies but debits it for the weight VARIANCE it injects, so an + # outlier-heavy broad member (a few enormous weights) scores LOW or negative -- those + # outliers are precisely what destroys pooled n_eff. Note both simpler candidates fail: + # Kish n_ess is SCALE-INVARIANT (blind to whether a member carries any integral mass), + # and mean weight alone REWARDS badly-matched proposals (a well-matched, contracted AV + # correctly has small uniform weights, while a broad GMM's rare huge-weight outlier sets + # the maximum) -- measured on S250114ax, mean weight ranked AV at 1e-40 vs GMM 2e-4. + # A single global normalization (the chunk's max log-weight) keeps members comparable. + # NB: the report n_ess and the allocation signal use the TRUE (unclipped) weights. + # Feeding them the clipped copy is WRONG: clipping flattens weights, which INFLATES a + # member's Kish n_ess, so the allocation would perversely favor exactly the member whose + # weights had to be clipped (measured: on S250114ax it starved the AV workhorse to the + # 1% floor and collapsed n_eff to ~1). Clipping's ONLY job is to protect proposal FITS. + _lw_all = numpy.asarray(self.identity_convert(log_weights), dtype=float) + _finite = numpy.isfinite(_lw_all) + _lw_max = float(numpy.max(_lw_all[_finite])) if bool(numpy.any(_finite)) else 0.0 + _u_all = numpy.where(_finite, numpy.exp(_lw_all - _lw_max), 0.0) + _S_tot = float(numpy.sum(_u_all)); _Q_tot = float(numpy.sum(_u_all * _u_all)) + contrib_per_sample = numpy.zeros(len(self.portfolio)) + # MIS CREDIT ASSIGNMENT (q_mix-native, the 'credit' quality signal). Under the + # balance heuristic each sample's contribution is owed to members in proportion to + # their share of the mixture density there, so member m's credit is + # credit_m = sum_i [ frac_m q_m(x_i) / q_mix(x_i) ] * w_i . + # Unlike Kish n_ess (scale-invariant, hence blind to whether a member carries any + # integral mass) this credits a member for COVERING WHERE THE INTEGRAND IS, even if + # it drew few samples there -- exactly the signal a slow-contracting VARAHA member + # needs. Normalized per drawn sample so members are comparable at unequal shares. + credit_per_sample = numpy.zeros(len(self.portfolio)) + _parts = getattr(self, '_chunk_mix_parts', None) + if _parts and q_mix is not None: + _qm = numpy.asarray(self.identity_convert(q_mix), dtype=float) + _w_all = numpy.where(numpy.isfinite(_lw_all), numpy.exp(_lw_all - _lw_max), 0.0) + for _im, _mem in enumerate(self.portfolio_realizations): + _pc = _parts.get(id(_mem)) + if _pc is None: + continue + _share = numpy.where(_qm > 0, _pc / _qm, 0.0) + # DIVIDE OUT THE MEMBER'S OWN ALLOCATION. The raw share frac_m*q_m/q_mix scales + # with frac_m, so a member accrues credit simply BECAUSE it is dominant -- the + # same circularity the n_ess signal has (measured: a 0.95-share GMM scored 6e-4 + # vs AV's 9e-10 and starved AV to the floor). Normalizing by frac_m turns this + # into "integral explained PER UNIT ALLOCATION", which is allocation-invariant + # and is what an allocation rule must compare. + _frac_m = float(n_samples_per_member[_im]) / float(max(1, n_samples)) + if _frac_m > 0: + credit_per_sample[_im] = float(numpy.sum(_share * _w_all)) / (_frac_m * max(1, n_samples)) + portfolio_report = {} for indx_member, member in enumerate(self.portfolio): indx_start = int(n_index_start_per_member[indx_member]) - indx_end = indx_start + int(n_samples_per_member[indx_member]) - ln_wt_here = log_weights[indx_start:indx_end] + indx_end = indx_start + int(n_samples_per_member[indx_member]) + _n_here = max(1, indx_end - indx_start) + _u_here = _u_all[indx_start:indx_end] + if _S_tot > 0 and _Q_tot > 0 and indx_end > indx_start: + _mean_w = float(numpy.sum(_u_here)) / _n_here + _mean_w2 = float(numpy.sum(_u_here * _u_here)) / _n_here + contrib_per_sample[indx_member] = 2.0 * _mean_w / _S_tot - _mean_w2 / _Q_tot + ln_wt_here = log_weights[indx_start:indx_end] # TRUE weights (see note above); not the clipped copy ln_wt_here += - np.max(ln_wt_here) # evaluate n_ess, n_eff for this set of samples in batch specifically, portfolio_report[indx_member] = [ self.portfolio_weights[indx_member], self.identity_convert(self.xpy.sum(self.xpy.exp(ln_wt_here))**2/self.xpy.sum(self.xpy.exp(ln_wt_here*2))), identity_convert(self.xpy.sum(self.xpy.exp(ln_wt_here)))] print("\t",portfolio_report) + if use_adaptive_alloc and len(self.portfolio) > 1: + print("\t credit/sample (MIS credit):", numpy.array2string(credit_per_sample, precision=3)) + print("\t contrib/sample (global-impact signal):", numpy.array2string(contrib_per_sample, precision=3), + " quality:", numpy.array2string(np.asarray(self.portfolio_quality, dtype=float), precision=3)) + # Record each member's per-chunk n_ess so freeze policies (and post-hoc analysis) + # can tell a member that is still CLIMBING from one that has PLATEAUED. + for indx_member in range(len(self.portfolio)): + self.portfolio_member_ness_history[indx_member].append(float(portfolio_report[indx_member][1])) # Weight based on n_ESS from batch. remember these are >=1, so no negatives or 0 will happen dat =np.array([ portfolio_report[k][1] for k in range(len(self.portfolio))]) - self.portfolio_weights = portfolio_wt_func(dat, self.portfolio_weights, xpy=self.xpy, identity_convert=self.identity_convert) # call weighting function + if use_adaptive_alloc and len(self.portfolio) > 1: + # adaptive-probe allocation: quality-EMA + round-robin probe (see _adaptive_allocation). + # frac_now = the fraction each member actually drew THIS chunk (n_samples_per_member is + # derived from self.portfolio_weights just above). The quality OBSERVABLE is either the + # global-impact contribution (default) or the legacy per-member Kish n_ess. + frac_now = np.array(n_samples_per_member, dtype=float) / float(max(1, n_samples)) + if self.portfolio_quality_signal == 'credit': + _obs = credit_per_sample + elif self.portfolio_quality_signal == 'global': + _obs = contrib_per_sample + else: + _obs = dat + self.portfolio_weights = self._adaptive_allocation(_obs, frac_now, self.portfolio_draw_iteration) + else: + self.portfolio_weights = portfolio_wt_func(dat, self.portfolio_weights, xpy=self.xpy, identity_convert=self.identity_convert) # call weighting function + # VARAHA DRAW FLOOR (see __init__): reserve a combined fraction for VARAHA members, so a + # slow-contracting workhorse cannot be starved of DRAWS by a member that merely looks + # good per-chunk. Applied after either allocation rule; unbiased (q_mix). + # BANDED: a floor alone is not enough. Measured on a loud-event best-fit point: with no + # floor the mixture degenerates to GMM-alone (VARAHA share -> 0.0099), q_mix loses its + # broad backstop, and a mode the peaked member misses is uncovered -> lnZ silently low + # while n_eff looks GOOD (the confidently-wrong failure). With a floor but no cap, one + # seed ran away the OTHER way (VARAHA -> 0.99) and was the outlier of its arm. Both are + # mixture degeneration. Constraining the VARAHA share to a BAND keeps q_mix genuinely + # mixed -- a broad backstop AND a peaked component -- by construction. Unbiased either + # way (q_mix balance heuristic), so this costs at most draws, never correctness. + _vmin = float(self.portfolio_varaha_min_frac) + _vmax = float(self.portfolio_varaha_max_frac) # <=0 or >=1 => no cap (back-compatible) + _cap_on = (0.0 < _vmax < 1.0) + if (_vmin > 0 or _cap_on) and len(self.portfolio) > 1: + _is_v = np.array([hasattr(m, 'is_varaha') for m in self.portfolio_realizations]) + if _is_v.any() and not _is_v.all(): + _w = np.asarray(self.portfolio_weights, dtype=float) + _w = np.where(np.isfinite(_w) & (_w > 0), _w, 0.0) + _sv = _w[_is_v].sum(); _so = _w[~_is_v].sum() + _target = None + if _vmin > 0 and _sv < _vmin: + _target = _vmin + elif _cap_on and _sv > _vmax: + _target = _vmax + if _target is not None and (_sv > 0 or _so > 0): + # put the VARAHA group at _target and the rest at (1-_target), each preserving its + # own internal split; if a group is all-zero, spread its share evenly within it. + if _sv > 0: + _w[_is_v] *= _target / _sv + else: + _w[_is_v] = _target / max(int(_is_v.sum()), 1) + if _so > 0: + _w[~_is_v] *= (1.0 - _target) / _so + else: + _w[~_is_v] = (1.0 - _target) / max(int((~_is_v).sum()), 1) + _tot = _w.sum() + if _tot > 0: + self.portfolio_weights = _w / _tot ### ### ORACLE BLOCK ### + # Oracles PROPOSE points (hill-climb hotspots, a Fisher/Gaussian, a + # previous posterior). We evaluate the true likelihood there and + # APPEND those (point, weight) pairs to the training data the other + # portfolio members adapt from, so they learn about regions the plain + # sampling missed. Oracles never enter the integral estimate itself, + # so they cannot bias it -- at worst they cost a few evaluations. rvs_train = self._rvs + log_weights_train = log_weights_adapt # weights aligned with rvs_train tail (clipped copy) if it_now < it_max_oracle and len(self.oracle_realizations )>0: rvs_train = deepcopy(self._rvs) # duplicate deeply, since we will append to it n_samples_per_oracle = int(n*0.1/len(self.oracle_realizations)) # try to minimize oracle effort - print(" ORACLE: attempting updates ") - # update each oracle - for member in self.oracle_realizations: - member.update_sampling_prior(log_weights, n_history, external_rvs=rvs_train, log_scale_weights=True) - # generate samples from oracles - rv_oracle = self.xpy.empty((n_samples_per_oracle*len(self.oracle_realizations), len(self.params_ordered))) - base_now = 0 - for member in self.oracle_realizations: - _, _, rv_here = member.draw_simplified(n_samples_per_oracle) - rv_oracle[base_now:base_now+n_samples_per_oracle] = rv_here - base_now += n_samples_per_oracle - # evaluate lnL for each, - lnL_oracles = lnF(*rv_oracle.T) - # put into weights and rvs, for use in training other samples - self.xpy.append(log_weights, lnL_oracles) - for indx, p in enumerate(self.params_ordered): - self.xpy.append(rvs_train[p], rv_oracle[:,indx]) + if n_samples_per_oracle > 0: + print(" ORACLE: attempting updates ") + # update each oracle from the current (host) history + for member in self.oracle_realizations: + member.update_sampling_prior(log_weights_adapt, n_history, external_rvs=rvs_train, log_scale_weights=True) + # generate proposals from oracles (oracles are host/numpy) + rv_list = [] + for member in self.oracle_realizations: + _, _, rv_here = member.draw_simplified(n_samples_per_oracle) + rv_list.append(numpy.asarray(identity_convert(rv_here))) # (n, ndim) host + rv_oracle = numpy.vstack(rv_list) # host, (n_oracle_total, ndim) + # evaluate the true integrand at the proposals; feed it device-first + # (real GPU ILE likelihood) with host fallback, mirroring the main loop + _cols = rv_oracle.T + if getattr(self, '_integrand_wants_host', False) or not cupy_ok: + _lnLo = lnF(*_cols) if 'no_protect_names' in kwargs else lnF(**dict(zip(self.params_ordered, _cols))) + else: + try: + _colsg = identity_convert_togpu(_cols) + _lnLo = lnF(*_colsg) if 'no_protect_names' in kwargs else lnF(**dict(zip(self.params_ordered, _colsg))) + except (TypeError, ValueError): + self._integrand_wants_host = True + _lnLo = lnF(*_cols) if 'no_protect_names' in kwargs else lnF(**dict(zip(self.params_ordered, _cols))) + lnL_oracles = numpy.asarray(identity_convert(_lnLo)) + # training weight for a proposal = its lnL (same log scale as + # log_weights up to the shared normalization the members remove) + log_w_oracle = lnL_oracles + # ACTUALLY append (numpy.append is not in-place -- must reassign) + for indx, p in enumerate(self.params_ordered): + base = identity_convert(rvs_train[p]) + rvs_train[p] = numpy.append(base, rv_oracle[:, indx]) + log_weights_train = numpy.append(identity_convert(log_weights_adapt), log_w_oracle) ### @@ -549,12 +1781,34 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # update sampling prior, using ALL past data # Don't update samples which are not being drawn # always update if we have an oracle - don't freeze out out oracle, UNLESS we have explicitly frozen it with a breakpoint + _is_varaha = hasattr(member, 'is_varaha') + # VARAHA EXEMPTION: a VARAHA/AV member contracts its live volume ONLY on the chunk + # it is updated, so it must update EVERY chunk (like standalone AV) to become the + # workhorse. q_mix keeps this unbiased regardless of weight, so exempt it from the + # freeze schedule entirely by default (see portfolio_varaha_never_freeze). + _varaha_exempt = _is_varaha and self.portfolio_varaha_never_freeze + # GRACE: don't freeze anyone during the first grace_iters iterations (let a slow + # starter like a VARAHA member contract before its weight is judged). + _in_grace = (self.portfolio_draw_iteration <= self.portfolio_grace_iters) + # REVIVE: periodically update even a frozen member so it gets a chance to recover + # instead of being starved forever. + _revive = (self.portfolio_revive_period > 0 + and (self.portfolio_draw_iteration % self.portfolio_revive_period == 0)) + # PLATEAU-AWARE revive: also update a low-weight member while its OWN per-chunk + # n_ess is still climbing (it is still learning); only let the freeze schedule + # govern a member that has plateaued. Uses the n_ess history recorded above. + _climbing = False + _hist = self.portfolio_member_ness_history[indx] + if self.portfolio_plateau_revive and len(_hist) >= 3: + _recent = _hist[-1]; _older = np.median(_hist[-3:-1]) + _climbing = (_recent > 1.05*max(_older, 1.0)) if self.portfolio_draw_iteration < self.portfolio_breakpoints[indx]: print(" - before activation breakpoint for member {} ".format( indx)) pass - elif (len(self.oracle_realizations) > 0 and it_now self.portfolio_freeze_wt): - if not(hasattr(member, 'is_varaha')): - member.update_sampling_prior(log_weights, n_history,external_rvs=rvs_train,log_scale_weights=True, **update_dict) + elif (len(self.oracle_realizations) > 0 and it_now self.portfolio_freeze_wt) or _in_grace or _revive or _varaha_exempt or _climbing: + if not(_is_varaha): + # log_weights_train / rvs_train include any oracle proposals appended above + member.update_sampling_prior(log_weights_train, n_history,external_rvs=rvs_train,log_scale_weights=True, **update_dict) else: # just do a single VARAHA step, independent of others member.update_sampling_prior_selfish(lnF) @@ -590,10 +1844,23 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): ln_wt = self._rvs["log_integrand"] + self._rvs["log_joint_prior"] - self._rvs["log_joint_s_prior"] # Convert to CPU as needed ln_wt = identity_convert(ln_wt) - ln_wt += - np.max(ln_wt) # remove maximum value, irrelevant - wt = np.exp(ln_wt) # exponentiate. Danger underflow + ln_wt = numpy.asarray(ln_wt, dtype=float) + # Guard: rejected/underflowed samples (-inf) or bad priors (nan) must not poison the + # cumulative sum below. Map any non-finite log-weight to -inf so it exponentiates to a + # zero linear weight instead of corrupting cumsum/normalization (which would drop ALL rows). + ln_wt[~numpy.isfinite(ln_wt)] = -numpy.inf + ln_wt_max = numpy.max(ln_wt) + if numpy.isfinite(ln_wt_max): + ln_wt = ln_wt - ln_wt_max # remove maximum value, irrelevant to the normalized cumulative prob + wt = numpy.exp(ln_wt) # exponentiate to LINEAR weights (max-subtracted). Underflow -> 0, which is fine + else: + # degenerate: no finite-weight sample survived Step 1 -- keep everything rather than drop all rows + wt = numpy.ones(len(ln_wt)) idx_sorted_index = numpy.lexsort((numpy.arange(len(wt)), wt)) # Sort the array of weights, recovering index values - indx_list = numpy.array( [[k, ln_wt[k]] for k in idx_sorted_index]) # pair up with the weights again. NOTE NOT INTEGER TYPE ANY MORE + # Pair the sorted index with the LINEAR weight wt[k] (NOT the log-weight ln_wt[k]): the cumulative + # sum below must be a cumulative PROBABILITY, matching mcsampler/mcsamplerEnsemble. Cumsumming the + # log-weights (<=0, and -inf for rejects) is not a probability threshold and kept 0 rows for peaked runs. + indx_list = numpy.array( [[k, wt[k]] for k in idx_sorted_index]) # pair up with the LINEAR weights again. NOTE NOT INTEGER TYPE ANY MORE cum_sum = numpy.cumsum(indx_list[:,1]) # find the cumulative sum cum_sum = cum_sum/cum_sum[-1] # normalize the cumulative sum indx_list = [int(indx_list[k, 0]) for k, value in enumerate(cum_sum > deltaP) if value] # find the indices that preserve > 1e-7 of total probability. RECAST TO INTEGER @@ -627,6 +1894,30 @@ def integrate_log(self, lnF, *args, xpy=xpy_default,**kwargs): # if convergence_tests is not None: # dict_return["convergence_test_results"] = None # last_convergence_test + # SUPPORT-MISMATCH (warm-start) DIAGNOSTIC -- reported, never acted on. Nothing above this + # line consumes it, so the estimate returned below is bit-identical with and without it. + try: + _sd = self.support_diagnostics() + dict_return['portfolio_escaped_mass'] = _sd['escaped_mass'] + dict_return['portfolio_member_weight_share'] = _sd['weight_share'] + dict_return['portfolio_member_hard_edged'] = _sd['hard_edged'] + dict_return['portfolio_escaped_mass_max'] = _sd['escaped_mass_max'] + dict_return['portfolio_escaped_mass_early'] = _sd['escaped_mass_early'] + dict_return['portfolio_escaped_mass_early_max'] = _sd['escaped_mass_early_max'] + dict_return['portfolio_escaped_mass_history'] = _sd['escaped_mass_history'] + dict_return['portfolio_escape_n_zero'] = _sd['escape_n_zero'] + dict_return['portfolio_escape_n_eval'] = _sd['escape_n_eval'] + print(" PORTFOLIO support: escaped_mass={} early={} (hard-edged members {}; " + "max {:.3e}, early max {:.3e}) weight_share={}".format( + numpy.array2string(np.asarray(_sd['escaped_mass'], dtype=float), precision=3), + numpy.array2string(np.asarray(_sd['escaped_mass_early'], dtype=float), precision=3), + list(numpy.flatnonzero(_sd['hard_edged'])), _sd['escaped_mass_max'], + _sd['escaped_mass_early_max'], + numpy.array2string(np.asarray(_sd['weight_share'], dtype=float), precision=3))) + except Exception as _e_sd: + print(" PORTFOLIO support diagnostic unavailable ({}: {})".format( + type(_e_sd).__name__, _e_sd)) + # perform type conversion of all stored variables if cupy_ok: for name in self._rvs: diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/proposal_field.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/proposal_field.py new file mode 100644 index 000000000..21bb729e3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/proposal_field.py @@ -0,0 +1,134 @@ +""" +proposal_field.py -- L3 substrate: an intrinsic-space field of extrinsic proposals. + +Scaffolding for iteration-to-iteration warm-start reuse (the "breadcrumb" strategy). +After an ILE iteration, each converged point contributes a compact extrinsic proposal +(its high-likelihood samples, or an AV live-volume state). Those are aggregated, keyed +by the intrinsic parameters lambda, into a ProposalField. The next iteration's ILE +workers query the field at their own lambda and warm-start from the nearest proposal. + +Design constraints (see DESIGN_warmstart_threading.md): + * A warm start only shapes p_s, never the estimator -> a stale/mismatched field entry + can only cost efficiency, never bias. Callers should still pass cover_frac>0 to + bootstrap_from_samples for cross-problem reuse as a belt-and-suspenders floor. + * Entries are small (samples ~KB, AV state ~8KB) -> the whole field is a small file, + fine for Condor transfer_input_files. + * Nearest-neighbour lookup uses a whitened intrinsic metric so "nearby" respects the + very different scales of chirp mass, mass ratio, and spins. + +This module intentionally does NOT touch the DAG. Pipeline wiring (build the field as a +post-iteration node, pass it forward, query it in the ILE driver via a new +--extrinsic-proposal-field hook) is the next step and mirrors the existing calmarg +extrinsic-breadcrumb plumbing. +""" +import numpy as np + + +class ProposalField(object): + """A set of (lambda, extrinsic-proposal) entries with nearest-lambda lookup. + + lambda vectors are intrinsic-parameter coordinates (e.g. [mc, q, chi1z, chi2z, ...]); + proposals are (M_k, d_extrinsic) arrays of high-likelihood extrinsic samples in the + sampler's coordinate convention. `extrinsic_params` names the proposal columns.""" + + def __init__(self, intrinsic_params=None, extrinsic_params=None): + self.intrinsic_params = list(intrinsic_params) if intrinsic_params else None + self.extrinsic_params = list(extrinsic_params) if extrinsic_params else None + self._lambdas = [] # list of (d_intrinsic,) arrays + self._proposals = [] # list of (M_k, d_extrinsic) arrays + self._scale = None # per-intrinsic-dim scale for the whitened metric + + def add(self, lam, proposal): + """Add one point's proposal. lam: (d_intrinsic,); proposal: (M, d_extrinsic).""" + lam = np.asarray(lam, dtype=float).ravel() + proposal = np.atleast_2d(np.asarray(proposal, dtype=float)) + if proposal.shape[0] < 2: + return # too few points to seed a live volume + self._lambdas.append(lam) + self._proposals.append(proposal) + self._scale = None # invalidate cached metric + + def _metric_scale(self): + if self._scale is None and self._lambdas: + L = np.vstack(self._lambdas) + s = np.std(L, axis=0) + s[s <= 0] = 1.0 + self._scale = s + return self._scale + + def nearest(self, lam, k=1): + """Return the proposal(s) from the k intrinsic points nearest `lam` (whitened + Euclidean). For k=1 returns a single (M, d_extrinsic) array; for k>1 the + vertically-stacked union (a broader, safer seed).""" + if not self._lambdas: + return None + lam = np.asarray(lam, dtype=float).ravel() + s = self._metric_scale() + d = np.array([np.sum(((lam - lj) / s) ** 2) for lj in self._lambdas]) + order = np.argsort(d)[:max(1, int(k))] + if len(order) == 1: + return self._proposals[order[0]] + return np.vstack([self._proposals[i] for i in order]) + + def warm_seed_for(self, lam, k=1): + """Convenience: the seed array to hand to AV.bootstrap_from_samples(..., + params=self.extrinsic_params, cover_frac=...) for intrinsic point `lam`.""" + return self.nearest(lam, k=k) + + # --- serialization (small npz; fine for Condor transfer) --- + def save(self, path): + if not self._lambdas: + raise ValueError("ProposalField is empty") + arr = {} + arr['lambdas'] = np.vstack(self._lambdas) + arr['sizes'] = np.array([p.shape[0] for p in self._proposals]) + arr['proposals'] = np.vstack(self._proposals) # concatenated; split by sizes + arr['intrinsic_params'] = np.array(self.intrinsic_params or []) + arr['extrinsic_params'] = np.array(self.extrinsic_params or []) + np.savez_compressed(path, **arr) + return path + + @classmethod + def load(cls, path): + d = np.load(path, allow_pickle=True) + pf = cls(intrinsic_params=[str(x) for x in d['intrinsic_params']] or None, + extrinsic_params=[str(x) for x in d['extrinsic_params']] or None) + L = d['lambdas']; sizes = d['sizes']; P = d['proposals'] + off = 0 + for i in range(len(L)): + n = int(sizes[i]) + pf._lambdas.append(np.asarray(L[i], dtype=float)) + pf._proposals.append(np.asarray(P[off:off + n], dtype=float)) + off += n + return pf + + def __len__(self): + return len(self._lambdas) + + +#: canonical intrinsic key: masses (Msun) + the six spin components +LAMBDA_INTRINSIC_PARAMS = ["m1", "m2", "s1x", "s1y", "s1z", "s2x", "s2y", "s2z"] + + +def lambda_from_P(P): + """Canonical intrinsic coordinate for a RIFT ChooseWaveformParams P: + [m1, m2 (Msun), s1x, s1y, s1z, s2x, s2y, s2z]. Used as the field key so + 'nearby' is measured in the physical intrinsic parameters (the whitened metric + in ProposalField then handles their different scales).""" + try: + import lal + msun = lal.MSUN_SI + except Exception: + msun = 1.98892e30 + return np.array([P.m1 / msun, P.m2 / msun, + P.s1x, P.s1y, P.s1z, P.s2x, P.s2y, P.s2z], dtype=float) + + +def build_field_from_run_outputs(entries, intrinsic_params=None, extrinsic_params=None): + """Aggregate a list of (lam, proposal) pairs (one per converged ILE point in an + iteration) into a ProposalField. Intended to be called by a small post-iteration + node that scans the iteration's ILE outputs.""" + pf = ProposalField(intrinsic_params=intrinsic_params, extrinsic_params=extrinsic_params) + for lam, proposal in entries: + pf.add(lam, proposal) + return pf diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/statutils.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/statutils.py index 166d84562..3e98db873 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/statutils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/statutils.py @@ -187,9 +187,145 @@ def finalize_log(existingAggregate,xpy=numpy): """ (count, log_mean_orig, log_M2, log_ref) = existingAggregate - (log_mean, log_sampleVariance) = (log_mean_orig+log_ref, log_M2 + 2*log_ref - xpy.log((count - 1))) + (log_mean, log_sampleVariance) = (log_mean_orig+log_ref, log_M2 + 2*log_ref - xpy.log((count - 1))) # print( log_mean, log_sampleVariance) if count < 2: return float('nan') else: return (log_mean, log_sampleVariance) + + +# +# MC-error stabilization helpers (host-side, numpy only). +# +# Motivation: the sample variance of the importance weights, computed from the +# SAME draws as the integral, is algebraically 1/ESS_hat - 1/n restated. It is +# tail-blind: a run that has not sampled the dominant weight region reports BOTH +# a low integral AND a small error bar, so the naive estimate fails conditionally +# on the run being wrong. These helpers provide (a) a generalized-Pareto tail +# diagnostic (PSIS k-hat, Vehtari et al. JMLR 2024), (b) an ESS statistic, +# (c) a between-chunk jackknife scatter that sees adaptation nonstationarity, +# and (d) bootstrap quantiles of lnZ for honest (asymmetric) intervals when the +# relative error is large. All are cheap relative to likelihood evaluations and +# must be called on CPU (numpy) arrays. +# + +def pareto_khat_from_log(log_wt, tail_frac=0.2, min_tail=20): + """Generalized-Pareto tail index k of the importance-weight distribution, + fit to the largest weights (Zhang & Stephens 2009 posterior-mean estimator, + as used by Pareto-smoothed importance sampling). Input is LOG weights on + any scale (k is invariant under overall rescaling). Interpretation: + k < 0.5 : weight variance finite, the naive error estimate is meaningful; + 0.5-0.7 : variance marginal, treat the naive sigma as optimistic; + k > 0.7 : the weight tail is unresolved -- the naive sigma is a lower bound + and the integral itself may be dominated by unseen tail mass. + Returns float k, or None if there are too few finite weights to fit.""" + lw = numpy.asarray(log_wt, dtype=float) + lw = lw[numpy.isfinite(lw)] + n = len(lw) + if n < 5 * min_tail: + return None + lw = numpy.sort(lw) + w = numpy.exp(lw - lw[-1]) # rescale by max: tail values are O(1), rest may underflow harmlessly + M = int(min(tail_frac * n, numpy.ceil(3 * numpy.sqrt(n)))) + M = max(M, min_tail) + if M >= n: + M = n - 1 + tail = w[-M:] + mu = w[-M - 1] # threshold = largest non-tail weight + x = tail - mu + if x[-1] <= 0: + return 0.0 # massive ties at the top: no resolvable tail + x = x[x > 0] + nt = len(x) + if nt < min_tail: + return 0.0 + xstar = x[int(nt / 4 + 0.5) - 1] + if xstar <= 0: + return 0.0 + m = 30 + int(numpy.sqrt(nt)) + jj = numpy.arange(1, m + 1) + theta = 1.0 / x[-1] + (1 - numpy.sqrt(m / (jj - 0.5))) / (3.0 * xstar) + theta[theta == 0] = 1e-12 # avoid the (measure-zero) singular point + # profile log-likelihood of the GPD for each candidate theta + k_of = -numpy.mean(numpy.log1p(-numpy.outer(theta, x)), axis=1) + k_of[numpy.abs(k_of) < 1e-12] = 1e-12 + with numpy.errstate(divide='ignore', invalid='ignore'): + lp = nt * (numpy.log(theta / k_of) + k_of - 1) + lp[~numpy.isfinite(lp)] = -numpy.inf + lp -= lp.max() + wts = numpy.exp(lp) + s = wts.sum() + if not numpy.isfinite(s) or s <= 0: + return None + theta_hat = numpy.sum(theta * wts) / s + if theta_hat == 0: + return 0.0 + # Zhang & Stephens parameterize the GPD with k_ZS = -xi (their k>0 is a + # BOUNDED tail); return the PSIS/Vehtari tail index xi = -k_ZS, so that + # heavy tails give POSITIVE k-hat and the 0.5/0.7 thresholds apply. + return float(numpy.mean(numpy.log1p(-theta_hat * x))) + + +def ess_from_log_weights(log_wt): + """Kish effective sample size (sum w)^2 / sum w^2 from LOG weights.""" + lw = numpy.asarray(log_wt, dtype=float) + lw = lw[numpy.isfinite(lw)] + if len(lw) == 0: + return 0.0 + lse = scipy.special.logsumexp + return float(numpy.exp(2 * lse(lw) - lse(2 * lw))) + + +def block_scatter_sigma(lnZ_blocks, n_blocks): + """sigma(lnZ) from the between-chunk scatter of per-chunk mean estimates, + via a delete-one jackknife of the n-weighted pooled mean. Each chunk of an + adaptive run used a different proposal, so this sees the nonstationarity the + pooled within-run variance averages away. (It still cannot see modes that + EVERY chunk missed -- only independent replicas can.) Returns float sigma + or None if fewer than 2 usable chunks.""" + lnZ = numpy.asarray(lnZ_blocks, dtype=float) + nb = numpy.asarray(n_blocks, dtype=float) + good = numpy.isfinite(lnZ) & (nb > 0) + lnZ = lnZ[good] + nb = nb[good] + K = len(lnZ) + if K < 2: + return None + ref = lnZ.max() + Z = numpy.exp(lnZ - ref) + tot = numpy.sum(nb * Z) + N = nb.sum() + loo = (tot - nb * Z) / (N - nb) # leave-one-out pooled means (relative to ref) + if tot <= 0 or numpy.any(loo <= 0): + return None + ln_loo = numpy.log(loo) + var_jk = (K - 1) / K * numpy.sum((ln_loo - ln_loo.mean()) ** 2) + return float(numpy.sqrt(var_jk)) + + +def bootstrap_lnZ_quantiles(log_wt, n_total=None, n_boot=200, quantiles=(0.05, 0.5, 0.95), rng_seed=None): + """Bootstrap quantiles of lnZ_hat = ln( sum_i w_i / n_total ) by resampling + the stored LOG weights with replacement. When the relative error is O(1) + the delta-method +-sigma interval on lnZ is meaningless (the distribution is + strongly skewed); these quantiles are an honest same-sample interval. They + remain blind to tail mass never sampled -- pair with pareto_khat_from_log. + n_total: divisor if the stored weights are a pruned subset of a larger run + (the pruned-away weights contribute negligibly to the sum). Returns a + numpy array of lnZ quantiles, or None if too few weights.""" + lw = numpy.asarray(log_wt, dtype=float) + lw = lw[numpy.isfinite(lw)] + n = len(lw) + if n < 10: + return None + if n_total is None: + n_total = n + rng = numpy.random.default_rng(rng_seed) + ref = lw.max() + w = numpy.exp(lw - ref) + out = numpy.empty(n_boot) + for b in range(n_boot): + idx = rng.integers(0, n, n) + out[b] = numpy.log(numpy.sum(w[idx])) + out += ref - numpy.log(n_total) + return numpy.quantile(out, quantiles) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/.gitignore b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/.gitignore new file mode 100644 index 000000000..7a60b85e1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/__init__.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/__init__.py new file mode 100644 index 000000000..b8b34dfdd --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/__init__.py @@ -0,0 +1,22 @@ +""" +unreliable_oracle -- proposal 'oracles' for the RIFT portfolio integrator. + +An oracle only PROPOSES candidate points; the portfolio evaluates the true +likelihood there and folds them into the training data used to adapt the other +integrators. Because oracles never contribute to the integral estimate +directly, an inaccurate oracle cannot bias the result -- it can only waste a few +likelihood evaluations. That makes them a safe channel for injecting cheap, +approximate posterior knowledge (a Fisher matrix, a hill-climbed hotspot, or a +previous run's posterior samples). +""" +from .resampling import ResamplingOracle +from .puffball import PuffballOracle +from .hill_climber import ClimbingOracle +from .fisher_gaussian import FisherGaussianOracle + +__all__ = [ + "ResamplingOracle", + "PuffballOracle", + "ClimbingOracle", + "FisherGaussianOracle", +] diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/fisher_gaussian.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/fisher_gaussian.py new file mode 100644 index 000000000..89ca2f260 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/fisher_gaussian.py @@ -0,0 +1,150 @@ +""" +fisher_gaussian.py + +A Fisher-matrix / Gaussian oracle for the RIFT portfolio integrator. + +The paper's remark is that a normalizing-flow oracle, trained per-inference, in +effect just supplies a fast local approximation to the posterior -- and that a +Fisher matrix gives the same thing essentially for free. This oracle is that +cheap substitute: given a mean and a covariance (or a Fisher matrix Gamma, cov = +Gamma^{-1}), it proposes points from the box-truncated Gaussian N(mean, cov), and +can optionally refresh (mean, cov) from the weighted history it is shown. + +Oracles only PROPOSE points; the portfolio evaluates the true likelihood at them +and folds them into the training data for the other integrators. A proposal can +therefore never bias the integral -- at worst it wastes a few evaluations -- so +an approximate Fisher is a safe, robust way to inject known posterior shape. + +Interface matches MCSamplerGeneric: setup(), update_sampling_prior(), draw_simplified(). +Backend-agnostic: all math is host numpy (proposals are cheap and small); the +portfolio moves arrays to the active backend as needed. +""" +import numpy as np + +from RIFT.integrators.mcsampler_generic import MCSamplerGeneric + + +class FisherGaussianOracle(MCSamplerGeneric): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.reference_mean = None + self.reference_cov = None + self._chol = None + # if True, refit (mean,cov) from weighted history on each update + self.adapt = False + # inflate the covariance when proposing, to stay broad (avoid collapse) + self.cov_inflate = 1.0 + + def add_parameter(self, params, pdf, **kwargs): + super().add_parameter(params, pdf, **kwargs) + + def setup(self, mean=None, cov=None, fisher=None, adapt=False, + cov_inflate=1.0, params=None, **kwargs): + """mean/cov describe the proposal; alternatively pass a Fisher matrix + (cov = fisher^{-1}). `params` names the columns of mean/cov if they are + not already in this oracle's params_ordered order. adapt=True refreshes + (mean,cov) from weighted history on each update_sampling_prior call.""" + super().setup(**kwargs) + self.adapt = adapt + self.cov_inflate = float(cov_inflate) + if cov is None and fisher is not None: + cov = np.linalg.inv(np.atleast_2d(np.asarray(fisher, dtype=float))) + if mean is not None and cov is not None: + mean = np.asarray(mean, dtype=float) + cov = np.atleast_2d(np.asarray(cov, dtype=float)) + if params is not None and self.params_ordered: + order = [list(params).index(p) for p in self.params_ordered] + mean = mean[order] + cov = cov[np.ix_(order, order)] + self._set_gaussian(mean, cov) + + def _set_gaussian(self, mean, cov): + cov = 0.5 * (cov + cov.T) + # regularize to SPD + w, V = np.linalg.eigh(cov) + w = np.clip(w, 1e-12 * max(np.max(w), 1e-300), None) + cov = (V * w) @ V.T + self.reference_mean = mean + self.reference_cov = cov * self.cov_inflate + self._chol = np.linalg.cholesky(self.reference_cov) + + def update_sampling_prior(self, ln_weights, n_history, lnw_cut=-10, + external_rvs=None, verbose=False, **kwargs): + """Optionally refit (mean, cov) from the recent weighted history. If + adapt is False (default), the oracle keeps the supplied Fisher/Gaussian + and this is a no-op -- the robust behaviour when a trustworthy Fisher is + available. With adapt=True it becomes a self-refining Gaussian proposal.""" + if not self.adapt: + return + rvs_here = external_rvs if external_rvs else self._rvs + if rvs_here is None or len(self.params_ordered) == 0: + return + p0 = self.params_ordered[0] + n_avail = len(rvs_here[p0]) + n_use = int(min(n_history, n_avail)) + if ln_weights is not None: + n_use = int(min(n_use, len(ln_weights))) + if n_use < len(self.params_ordered) + 2: + return + X = np.empty((n_use, len(self.params_ordered))) + for j, p in enumerate(self.params_ordered): + X[:, j] = np.asarray(rvs_here[p])[-n_use:] + w = None + if ln_weights is not None: + lw = np.asarray(ln_weights)[-n_use:].astype(float) + lw = lw - np.max(lw) + if lnw_cut is not None: + keep = lw > lnw_cut + if np.sum(keep) >= len(self.params_ordered) + 2: + X = X[keep] + lw = lw[keep] + w = np.exp(lw) + w = w / np.sum(w) + mean = np.average(X, axis=0, weights=w) + cov = np.cov(X.T, aweights=w) + cov = np.atleast_2d(cov) + self._set_gaussian(mean, cov) + if verbose: + print(" oracle - fisher_gaussian - refit mean", mean) + + def _bounds(self): + lo = np.array([self.llim[p] for p in self.params_ordered], dtype=float) + hi = np.array([self.rlim[p] for p in self.params_ordered], dtype=float) + return lo, hi + + def draw_simplified(self, n_samples, *args, **kwargs): + """Draw n_samples from the box-truncated Gaussian. Returns + (p_s, p_prior, rv) with rv shape (n_samples, ndim); p_s is the Gaussian + proposal density (so the draws can be used with correct importance + weights if desired), p_prior is left None (the portfolio supplies it).""" + if self.reference_mean is None: + raise Exception("FisherGaussianOracle: setup(mean=,cov=/fisher=) required before draw") + d = len(self.params_ordered) + lo, hi = self._bounds() + out = np.empty((n_samples, d)) + logdens = np.empty(n_samples) + n_out = 0 + cov = self.reference_cov + inv = np.linalg.inv(cov) + logdet = np.linalg.slogdet(cov)[1] + lognorm = -0.5 * (d * np.log(2 * np.pi) + logdet) + guard = 0 + while n_out < n_samples and guard < 1000: + guard += 1 + batch = np.random.multivariate_normal(self.reference_mean, cov, size=2 * n_samples) + inside = np.all((batch >= lo) & (batch <= hi), axis=1) + batch = batch[inside] + if len(batch) == 0: + continue + take = min(len(batch), n_samples - n_out) + b = batch[:take] + out[n_out:n_out + take] = b + dx = b - self.reference_mean + logdens[n_out:n_out + take] = lognorm - 0.5 * np.einsum('ij,jk,ik->i', dx, inv, dx) + n_out += take + if n_out < n_samples: + # fall back to filling the remainder uniformly in-box (keeps coverage) + rem = n_samples - n_out + out[n_out:] = np.random.uniform(lo, hi, size=(rem, d)) + logdens[n_out:] = -np.sum(np.log(hi - lo)) + return np.exp(logdens), None, out diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/hill_climber.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/hill_climber.py index 0395b60f6..806f44ca1 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/hill_climber.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/hill_climber.py @@ -51,18 +51,23 @@ def update_sampling_prior(self,ln_weights, n_history,lnw_cut = -10, external_rvs n_history_to_use = np.min([n_history, len(rvs_here[self.params_ordered[0]])] ) if not(ln_weights is None): - n_history_to_use = np.min([n_history,len(ln_weights)]) - + n_history_to_use = np.min([n_history,len(ln_weights), len(rvs_here[self.params_ordered[0]])]) + + # sample_array is (ndim, n_history_to_use): rows params, columns samples sample_array = self.xpy.empty( (len(self.params_ordered), n_history_to_use)) for indx, p in enumerate(self.params_ordered): sample_array[indx] = rvs_here[p][-n_history_to_use:] - if lnw_cut and not(ln_weights is None): # we can override and trainon all data - sample_array = sample_array[:, ln_weights > np.max(ln_weights) + lnw_cut ] # training range + if lnw_cut and not(ln_weights is None): # start climbs from the high-weight region + lnw = np.asarray(ln_weights)[-n_history_to_use:] + keep = lnw > np.max(lnw) + lnw_cut + if np.sum(keep) > 0: + sample_array = sample_array[:, keep] # select COLUMNS (samples) - # Pick points to climb - drawn_indx = np.random.choice(range(len(sample_array)), replace=True,size=self.n_climbers) # random samples + # Pick starting points to climb from (index over SAMPLES = columns) + n_avail = sample_array.shape[1] + drawn_indx = np.random.choice(range(n_avail), replace=True,size=self.n_climbers) # random samples sample_array = sample_array[:,drawn_indx].T # Climb diff --git a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/puffball.py b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/puffball.py index cb62ac4ef..a73c678d2 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/puffball.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/integrators/unreliable_oracle/puffball.py @@ -40,20 +40,24 @@ def update_sampling_prior(self,ln_weights, n_history,lnw_cut = -10, external_rvs raise Exception(" oracle - puffball -update_sampling_prior requires initial history") n_history_to_use = np.min([n_history, len(rvs_here[self.params_ordered[0]])] ) - if ln_weights: - n_history_to_use = np.min([n_history,len(ln_weights)]) - + if ln_weights is not None: + n_history_to_use = np.min([n_history,len(ln_weights), len(rvs_here[self.params_ordered[0]])]) + + # sample_array is (ndim, n_history_to_use): rows are parameters, columns samples sample_array = self.xpy.empty( (len(self.params_ordered), n_history_to_use)) for indx, p in enumerate(self.params_ordered): sample_array[indx] = rvs_here[p][-n_history_to_use:] - if lnw_cut and ln_weights: # we can override and trainon all data - sample_array = sample_array[ln_weights > np.max(ln_weights) + lnw_cut ] # training range + if lnw_cut is not None and ln_weights is not None: # restrict to high-weight training range + lnw = np.asarray(ln_weights)[-n_history_to_use:] + keep = lnw > np.max(lnw) + lnw_cut + if np.sum(keep) > len(self.params_ordered) + 1: # need enough for a covariance + sample_array = sample_array[:, keep] # select COLUMNS (samples), not rows - # compute mean and cov + # compute mean and cov (over samples = columns) self.reference_mean = np.mean(sample_array, axis=-1) - self.reference_cov = np.cov(sample_array) + self.reference_cov = np.atleast_2d(np.cov(sample_array)) if verbose: print(" oracle - puffball - update_sampling_prior ", self.reference_mean, self.reference_cov) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py index 287e17fc4..53b9f8b6e 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/lalsimutils.py @@ -3313,7 +3313,17 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil # but that is very difficult to do because modes of different 'm' and thus typical frequency generally mix # Note also that unless the segment length is large, this is often surprisingly few frequency bins for tapering if not(no_condition): - our_fvals = evaluate_fvals(hlmsdict[(2,2)]) + # SimInspiralChooseFDModes returns modes on an ascending two-sided grid + # [-fNyq, ..., 0, ..., +fNyq] with DC at index npts//2 (odd length TDlen+1). + # Do NOT use evaluate_fvals here: it assumes RIFT's reversed packing and, for + # odd length, assigns f_assumed = -f_true + deltaF/2. Since (l,m) and (l,-m) + # modes occupy opposite signs of f, that offset shifts the high-pass window by + # one bin between the members of a pair, violating + # h_{l,-m}(f) = (-1)^l conj(h_{lm}(-f)) — and hence the TD conjugate-pair + # identity — at the percent level. The window must be exactly even in the + # true frequency to commute with complex conjugation. + npts_fd = hlmsdict[(2,2)].data.length + our_fvals = P.deltaF*(np.arange(npts_fd) - npts_fd//2) vectaper_symmetric = np.ones(len(our_fvals)) indx_below = np.logical_and(np.abs(our_fvals)=P.fmin*fd_standoff_factor) vectaper_symmetric[indx_below] = 0.5 + 0.5*np.cos(np.pi* (np.abs(our_fvals[indx_below])/P.fmin - 1)/(1-fd_standoff_factor)) @@ -3339,6 +3349,11 @@ def hlmoft(P, Lmax=2,nr_polarization_convention=False, fixed_tapering=False, sil for mode in hlmsdict: hlmsdict[mode] = lal.ResizeCOMPLEX16FrequencySeries(hlmsdict[mode],0, TDlen) + if not(no_condition): + # The resize above truncated the +fNyq bin from the two-sided grid but kept + # its -fNyq partner (index 0). Zero it so the truncation commutes with the + # conjugate-pair (f -> -f) reflection for models with support at Nyquist. + hlmsdict[mode].data.data[0] = 0 hlmsT[mode] = DataInverseFourier(hlmsdict[mode]) # Phase factors: see crazy conventions in https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalsimulation/lib/LALSimInspiral.c if True: #P.approx == lalIMRPhenomXHM or P.approx == lalIMRPhenomHM: @@ -4634,6 +4649,16 @@ def frame_data_to_hoft(fname, channel, start=None, stop=None, window_shape=0., with open(fname) as cfile: cachef = Cache.fromfile(cfile) cachef=cachef.sieve(ifos=channel[:1]) + # ET's three detectors (E1,E2,E3) share the site letter, so the channel[:1] sieve above is + # ambiguous and would read the wrong frame (e.g. E2:... served from the E1 frame -> "channel + # not found"). When the cache carries full-IFO observatories (E1/E2/E3/C1 from -*.gwf + # frames), narrow to the exact IFO. Standard single-letter observatories ('H','L','V') never + # equal the 2-char IFO, so this leaves ordinary H1/L1/V1 caches untouched. + ifo_full = channel.split(':')[0] + if len(ifo_full) > 1: + exact = [e for e in cachef if getattr(e, 'observatory', None) == ifo_full] + if exact: + cachef = cachef.__class__(exact) for name in cachef: print(name) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py index 5668e8292..ef02fe145 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/__init__.py @@ -45,8 +45,11 @@ JAXDistPhiMargLikelihood, JAXDistPsiMargLikelihood, build_data_from_precompute, + build_rotation_data_from_precompute, + build_freqresponse_data_from_precompute, EXTRINSIC_PARAM_ORDER, ) +from .banded import build_rotation_data, build_freqresponse_data from .coordinates import ( build_network_frame, equatorial_to_network, @@ -65,6 +68,10 @@ "JAXDistPhiMargLikelihood", "JAXDistPsiMargLikelihood", "build_data_from_precompute", + "build_rotation_data_from_precompute", + "build_freqresponse_data_from_precompute", + "build_rotation_data", + "build_freqresponse_data", "EXTRINSIC_PARAM_ORDER", "build_network_frame", "equatorial_to_network", diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py new file mode 100644 index 000000000..b6f87e5b1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/banded.py @@ -0,0 +1,149 @@ +""" +Builders for the multi-band (slow-rotation / finite-size) JAX likelihood data. + +These take the SAME packed precompute banks the cupy NoLoop path uses -- the +outputs of ``factored_likelihood_with_rotation.pack_rotation_arrays`` (Path A/B) +and ``factored_likelihood_freqresponse.pack_freqresponse_arrays`` (Path D) -- and +wrap them into a :class:`~RIFT.likelihood.jax_ile.core.JAXLikelihoodData` tagged +with a ``feature`` so that :func:`core._accumulate_unit` routes through the +multi-band accumulator. Because that accumulator returns the identical +``(kappa_unit, rho_sq_unit)`` contract as the baseline, every marginalization +variant (distance / phi_ref / psi) works with the feature unchanged. + +The heavy, data-touching precompute (frame reading, ````, U/V, packing) +is reused verbatim -- only the cheap extrinsic->lnL contraction is JAX. +""" + +import numpy as np +import jax.numpy as jnp + +import lal +import lalsimulation as lalsim + +from .core import build_likelihood_data, DIST_MPC_REF +from . import response_slowrot as _rs +from . import response_freqresponse as _rf + + +def _stack_bank(by_key, keys, det): + """Stack ``by_key[det][k]`` over the ordered ``keys`` -> leading axis = band.""" + return np.stack([np.asarray(by_key[det][k], dtype=np.complex128) for k in keys], + axis=0) + + +def _base_data(packed_scalar, deltaT, tref, tvals, distMpcRef): + """Build the JAXLikelihoodData scaffold (scalars + minimal detector dict). + + We reuse the baseline container for the time grid / Simpson weights / gmst / + epoch bookkeeping, then attach the banded arrays and geometry below. + """ + return build_likelihood_data(packed_scalar, deltaT, tref, tvals, distMpcRef) + + +def build_rotation_data(meta, lookupNKDict, rho_by_a, U_by_aa, V_by_aa, epochDict, + deltaT, tvals, distMpcRef=DIST_MPC_REF): + """Banded JAXLikelihoodData for the slow-rotation (Path A/B) likelihood. + + Parameters mirror ``pack_rotation_arrays`` outputs plus the time grid. + ``meta`` carries ``a_list`` (ordered ``(p,n)``), ``harmonics``, ``p_max`` and + ``event_time_geo`` (the fiducial epoch / sidereal reference ``tref``). + """ + a_list = [(int(p), int(n)) for (p, n) in meta["a_list"]] + tref = float(meta["event_time_geo"]) + detectors = list(rho_by_a.keys()) + + # Minimal baseline-shaped packed dict (rholmArray of the FIRST band as a + # stand-in) so build_likelihood_data can set up lms/epoch/location/response. + a0 = a_list[0] + packed_scalar = {} + for det in detectors: + packed_scalar[det] = dict( + lms=np.asarray(lookupNKDict[det]), + rholmArray=np.asarray(rho_by_a[det][a0], dtype=np.complex128), + U=np.asarray(U_by_aa[det][(a0, a0)], dtype=np.complex128), + V=np.asarray(V_by_aa[det][(a0, a0)], dtype=np.complex128), + epoch=float(epochDict[det])) + data = _base_data(packed_scalar, deltaT, tref, tvals, distMpcRef) + + # Attach the full band banks + geometry. + pairs = [(a, ap) for a in a_list for ap in a_list] # unused; kept for clarity + for det in detectors: + dd = data.detectors[det] + Q_bank = _stack_bank(rho_by_a, a_list, det) # (A, K, npts_full) + dd["Q_bank"] = jnp.asarray(np.ascontiguousarray( + np.transpose(Q_bank, (0, 2, 1)))) # (A, npts_full, K) + A = len(a_list) + K = len(dd["lms"]) + U = np.empty((A, A, K, K), dtype=np.complex128) + V = np.empty((A, A, K, K), dtype=np.complex128) + for i, a in enumerate(a_list): + for j, ap in enumerate(a_list): + U[i, j] = np.asarray(U_by_aa[det][(a, ap)], dtype=np.complex128) + V[i, j] = np.asarray(V_by_aa[det][(a, ap)], dtype=np.complex128) + dd["U_bank"] = jnp.asarray(U) + dd["V_bank"] = jnp.asarray(V) + + data.feature = "rotation" + data.band = dict( + a_list=a_list, + p_max=int(meta["p_max"]), + harmonics=tuple(int(h) for h in meta["harmonics"]), + refl_idx=np.asarray(_rs.reflection_index(a_list), dtype=np.int64), + ) + return data + + +def build_freqresponse_data(meta, lookupNKDict, rho_by_p, U_by_pp, V_by_pp, + epochDict, deltaT, tvals, det_geom, + distMpcRef=DIST_MPC_REF): + """Banded JAXLikelihoodData for the finite-size (Path D) likelihood. + + Parameters mirror ``pack_freqresponse_arrays`` outputs plus the time grid. + ``meta`` carries ``p_list`` (0..Qmax+1), ``Qmax`` and ``event_time_geo``. + ``det_geom`` maps ``det -> (response, x_arm, y_arm, L)`` (from + ``slowrot_freqresponse.detector_geometry``); the arm unit vectors enter the + finite-size coefficients ``beta_q`` -- they are NOT recoverable from the LAL + response tensor alone. + """ + p_list = [int(p) for p in meta["p_list"]] + tref = float(meta["event_time_geo"]) + detectors = list(rho_by_p.keys()) + + p0 = p_list[0] + packed_scalar = {} + for det in detectors: + packed_scalar[det] = dict( + lms=np.asarray(lookupNKDict[det]), + rholmArray=np.asarray(rho_by_p[det][p0], dtype=np.complex128), + U=np.asarray(U_by_pp[det][(p0, p0)], dtype=np.complex128), + V=np.asarray(V_by_pp[det][(p0, p0)], dtype=np.complex128), + epoch=float(epochDict[det])) + data = _base_data(packed_scalar, deltaT, tref, tvals, distMpcRef) + + for det in detectors: + dd = data.detectors[det] + Q_bank = _stack_bank(rho_by_p, p_list, det) # (A, K, npts_full) + dd["Q_bank"] = jnp.asarray(np.ascontiguousarray( + np.transpose(Q_bank, (0, 2, 1)))) # (A, npts_full, K) + A = len(p_list) + K = len(dd["lms"]) + U = np.empty((A, A, K, K), dtype=np.complex128) + V = np.empty((A, A, K, K), dtype=np.complex128) + for i, p in enumerate(p_list): + for j, pp in enumerate(p_list): + U[i, j] = np.asarray(U_by_pp[det][(p, pp)], dtype=np.complex128) + V[i, j] = np.asarray(V_by_pp[det][(p, pp)], dtype=np.complex128) + dd["U_bank"] = jnp.asarray(U) + dd["V_bank"] = jnp.asarray(V) + resp, x_arm, y_arm, L = det_geom[det] + dd["x_arm"] = jnp.asarray(np.asarray(x_arm, dtype=np.float64)) + dd["y_arm"] = jnp.asarray(np.asarray(y_arm, dtype=np.float64)) + dd["L_arm"] = float(L) + + data.feature = "freqresponse" + data.band = dict( + p_list=p_list, + Qmax=int(meta["Qmax"]), + refl_idx=np.asarray(_rf.reflection_index(p_list), dtype=np.int64), + ) + return data diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py index 239e6925e..f5f4a32af 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/core.py @@ -70,6 +70,8 @@ from .detector import compute_detamresponse, time_delay_from_earth_center from .spherical import spherical_harmonics_vectorized +from . import response_slowrot as _rs +from . import response_freqresponse as _rf # Fiducial template distance (Mpc); identical to factored_likelihood.distMpcRef. DIST_MPC_REF = 1000.0 @@ -199,7 +201,38 @@ def _gather_linear(Q_col, pos): return jnp.where(valid, val, 0.0 + 0.0j) -_GATHERERS = {"nearest": _gather_nearest, "linear": _gather_linear} +def _gather_cubic(Q_col, pos): + """Four-point cubic-Lagrange interpolation of Q_col at continuous ``pos``. + + Mirrors the production ``factored_likelihood._cubic_Q_window_numpy`` / + ``Q_inner_product_cubic`` stencil EXACTLY: with ``i0 = floor(pos)`` and + ``u = pos - i0`` the value is ``sum_{k=-1}^{2} w_k(u) Q[i0+k]`` with the cubic + Lagrange weights below (at integer ``pos`` it reproduces the sample). Unlike + linear, cubic captures the curvature of the razor-sharp high-frequency rholm + peak; for 3G/high-SNR signals linear *undershoots* that peak (worse than + nearest) and biases the recovered arrival time -- hence the sky -- so this is + the interpolation the maintained likelihood uses. Still differentiable in + ``pos`` (a polynomial in ``u``), so it drives gradient sampling. Stencil + points outside the buffer contribute ZERO (per-point zero extension, matching + the reference), so an over-running window falls off to zero. + """ + n = Q_col.shape[0] + i0 = jnp.floor(pos).astype(jnp.int32) + u = pos - jnp.floor(pos) + w = (-u * (u - 1.0) * (u - 2.0) / 6.0, + (u + 1.0) * (u - 1.0) * (u - 2.0) / 2.0, + -(u + 1.0) * u * (u - 2.0) / 2.0, + (u + 1.0) * u * (u - 1.0) / 6.0) + out = jnp.zeros(pos.shape, dtype=jnp.complex128) + for off, wk in zip((-1, 0, 1, 2), w): + idx = i0 + off + valid = (idx >= 0) & (idx < n) + out = out + wk * jnp.where(valid, Q_col[jnp.clip(idx, 0, n - 1)], 0.0 + 0.0j) + return out + + +_GATHERERS = {"nearest": _gather_nearest, "linear": _gather_linear, + "cubic": _gather_cubic} def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, @@ -210,7 +243,16 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, is complex; the distance factor and the Re/abs reduction are applied by the caller (so the same accumulation feeds both the fixed-distance and the distance-marginalized paths). + + When ``data`` carries a slow-rotation / finite-size ``feature`` (built by + :func:`banded.build_rotation_data` / :func:`banded.build_freqresponse_data`), + the multi-band accumulator is used instead. It returns the *identical* + ``(kappa_unit, rho_sq_unit)`` contract, so every downstream marginalization + variant (distance, phi_ref, psi, ...) inherits the feature for free. """ + if getattr(data, "feature", None) is not None: + return _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, + phase_marginalization) ra = jnp.asarray(ra, dtype=jnp.float64) dec = jnp.asarray(dec, dtype=jnp.float64) psi = jnp.asarray(psi, dtype=jnp.float64) @@ -270,6 +312,117 @@ def _accumulate_unit(data, ra, dec, psi, incl, phiref, interp, return kappa_unit, rho_sq_unit +def _banded_coefficients(data, det, ra, dec, psi): + """Per-sample response coefficients ``C`` of shape (A, S) for detector ``det``. + + Dispatches on ``data.feature``: ``"rotation"`` -> the sidereal-harmonic + coefficients ``C_{(p,n)}`` (Path A/B), ``"freqresponse"`` -> the finite-size + basis coefficients ``b_p`` (Path D). ``data.gmst`` (= GMST(tref), a host + constant) is the sidereal reference, exactly as in the numpy NoLoop path. + """ + dd = data.detectors[det] + b = data.band + if data.feature == "rotation": + return _rs.rotation_coefficients_packed( + dd["response"], dd["location"], ra, dec, psi, data.gmst, + b["p_max"], b["a_list"]) + if data.feature == "freqresponse": + return _rf.response_coefficients_packed( + dd["response"], dd["x_arm"], dd["y_arm"], ra, dec, psi, data.gmst, + b["Qmax"], b["p_list"]) + raise ValueError("unknown banded feature %r" % (data.feature,)) + + +def _accumulate_unit_banded(data, ra, dec, psi, incl, phiref, interp, + phase_marginalization): + """Multi-band (slow-rotation / finite-size) network kappa and rho^2. + + Generalizes :func:`_accumulate_unit` by an extra summed "band" index + ``a`` (sidereal harmonic ``(p,n)`` / finite-size basis weight ``p``), sized + ``A``, contracted with the per-sample coefficient vector ``C_a`` from + :func:`_banded_coefficients`. The baseline is the ``A==1``, ``C==[F]`` case. + + Mirrors the numpy NoLoop references + ``DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation`` and + ``DiscreteFactoredLogLikelihoodFreqResponseNoLoop``: + + kappa_unit = sum_a conj(C_a) sum_lm conj(Y_lm) Q^a_lm(t) + rho_sq_unit = 0.5 Re[ sum_{a,a'} conj(C_a) C_a' (Ybar U^{a,a'} Y) + + C_{aR} C_a' (Y V^{a,a'} Y) ] + + with the caller applying ``invDist`` / ``invDist^2`` and the ``Re / -1/2`` + reduction (so ``-0.5*rho_sq_unit == -0.25 Re[...]`` matches the reference + ``term2``). ``aR`` is the V-term reflection (``(p,-n)`` for rotation, the + identity for finite-size), supplied as ``data.band['refl_idx']``. + + ``phase_marginalization`` is not supported for banded features. + """ + if phase_marginalization: + raise NotImplementedError( + "phase marginalization is not supported for slow-rotation / " + "finite-size (banded) likelihoods") + + ra = jnp.asarray(ra, dtype=jnp.float64) + dec = jnp.asarray(dec, dtype=jnp.float64) + psi = jnp.asarray(psi, dtype=jnp.float64) + incl = jnp.asarray(incl, dtype=jnp.float64) + phiref = jnp.asarray(phiref, dtype=jnp.float64) + + gather = _GATHERERS[interp] + gmst = data.gmst + inv_deltaT = 1.0 / data.deltaT + S = ra.shape[0] + npts = data.npts + t_offsets = jnp.arange(npts, dtype=jnp.float64) + refl_idx = data.band["refl_idx"] # (A,) int, static + + kappa_unit = jnp.zeros((S, npts), dtype=jnp.complex128) + rho_sq_unit = jnp.zeros((S, npts), dtype=jnp.float64) + + for det in data.detector_names: + dd = data.detectors[det] + lms = dd["lms"] + Q_bank = dd["Q_bank"] # (A, npts_full, K) + U_bank = dd["U_bank"] # (A, A, K, K) + V_bank = dd["V_bank"] # (A, A, K, K) + A = Q_bank.shape[0] + K = len(lms) + + Y = spherical_harmonics_vectorized(lms, incl, -phiref, l_max=dd["l_max"]) + conjY = jnp.conj(Y) # (S, K) + + C = _banded_coefficients(data, det, ra, dec, psi) # (A, S) complex + C_refl = C[refl_idx] # (A, S) + + t_det = (data.tref_minus_epoch(det) + + time_delay_from_earth_center(dd["location"], ra, dec, gmst)) + p0 = (t_det + data.tval0) * inv_deltaT + pos = p0[:, None] + t_offsets[None, :] # (S, npts) + + # --- term1: sum_a conj(C_a) * ( sum_lm conj(Y_lm) Q^a_lm(t) ) --- + kappa_det = jnp.zeros((S, npts), dtype=jnp.complex128) + for a in range(A): + inner_a = jnp.zeros((S, npts), dtype=jnp.complex128) + Qa = Q_bank[a] # (npts_full, K) + for k in range(K): + inner_a = inner_a + conjY[:, k][:, None] * gather(Qa[:, k], pos) + kappa_det = kappa_det + jnp.conj(C[a])[:, None] * inner_a + kappa_unit = kappa_unit + kappa_det + + # --- term2: 0.5 Re[ sum_{a,a'} conj(C_a)C_a' YbarUY + C_aR C_a' YVY ] --- + # YUY[a,a'] = einsum(conjY, Y, U_bank[a,a']); YVY[a,a'] = einsum(Y, Y, V) + YUY = jnp.einsum("si,sj,abij->abs", conjY, Y, U_bank) # (A,A,S) + YVY = jnp.einsum("si,sj,abij->abs", Y, Y, V_bank) # (A,A,S) + # conj(C_a) C_a' and C_aR C_a' contracted over (a,a') + CC_U = jnp.einsum("as,bs->abs", jnp.conj(C), C) # (A,A,S) + CC_V = jnp.einsum("as,bs->abs", C_refl, C) # (A,A,S) + term2_c = jnp.sum(CC_U * YUY + CC_V * YVY, axis=(0, 1)) # (S,) complex + rho_sq_det = 0.5 * term2_c.real # (S,) + rho_sq_unit = rho_sq_unit + rho_sq_det[:, None] + + return kappa_unit, rho_sq_unit + + def _time_marginalize(lnL_t, w_t): """log integral_t exp(lnL_t) dt via constant Simpson weights, log-sum-exp stable.""" m = jnp.max(lnL_t, axis=-1, keepdims=True) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_freqresponse.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_freqresponse.py new file mode 100644 index 000000000..d5c5a7f80 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_freqresponse.py @@ -0,0 +1,127 @@ +""" +JAX port of the finite-size (Path D) response coefficients. + +The extrinsic layer of the frequency-dependent-response likelihood +(``factored_likelihood_freqresponse``) needs, per detector, the complex +coefficients ``b_p`` of the response basis ``p = 0 .. Qmax+1`` -- these multiply +the (sky-independent) precompute banks ``Q^p_lm``, ``U^{p,p'}``, ``V^{p,p'}``. +The frequency basis ``c_q(f)`` is already folded into the precompute weights +``W_p(f)``, so the extrinsic layer only needs the *scalar* sky/pol coefficients + + b_0 = F0 (exact long-wavelength lal.ComputeDetAMResponse) + b_{1+q} = beta_q = (1/2)[ zx^2 a_x^q - zy^2 a_y^q ] , q = 0 .. Qmax + +which are closed-form analytic functions of ``(RA, DEC, psi)`` (with ``GMST(tref)`` +a host constant). Ported to ``jax.numpy`` (differentiable in the sky/pol angles), +mirroring ``slowrot_freqresponse.finite_size_geometry`` / ``finite_size_beta`` and +``factored_likelihood_freqresponse.response_coefficients``. Validated against the +numpy reference to ~1e-12 by ``test/jax/test_jax_freqresponse.py``. + +Unlike the sidereal-rotation case, every basis weight ``W_p`` is Hermitian, so the +V cross term needs NO harmonic reflection: the reflection index is the identity. + +Detector geometry (response tensor ``D``, arm unit vectors ``x_arm``, ``y_arm``, +arm length ``L``) is supplied by the caller as host constants (from +``slowrot_freqresponse.detector_geometry``); only ``RA, DEC, psi`` are JAX leaves. +""" + +import numpy as np +import jax.numpy as jnp + + +def _triad_jax(dec, psi, g): + """Polarization triad X, Y and source direction nhat at hour angle g=GMST-RA. + + JAX port of ``slowrot_freqresponse._triad``; vectors carry the 3-component on + the last axis. dec, psi, g are (S,) JAX arrays. + """ + cd, sd = jnp.cos(dec), jnp.sin(dec) + cp, sp = jnp.cos(psi), jnp.sin(psi) + cg, sg = jnp.cos(g), jnp.sin(g) + ones = jnp.ones_like(sg) + X = jnp.stack([-cp * sg - sp * cg * sd, + -cp * cg + sp * sg * sd, + sp * cd * ones], axis=-1) + Y = jnp.stack([sp * sg - cp * cg * sd, + sp * cg + cp * sg * sd, + cp * cd * ones], axis=-1) + nhat = jnp.stack([cd * cg, -cd * sg, sd * ones], axis=-1) + return X, Y, nhat + + +def _lwl_response_jax(D, X, Y): + """Long-wavelength F_+, F_x (== ComputeDetAMResponse). D host, X/Y JAX (S,3).""" + XDX = jnp.einsum('...i,ij,...j->...', X, D, X) + YDY = jnp.einsum('...i,ij,...j->...', Y, D, Y) + XDY = jnp.einsum('...i,ij,...j->...', X, D, Y) + YDX = jnp.einsum('...i,ij,...j->...', Y, D, X) + return XDX - YDY, XDY + YDX + + +def response_coefficients_dict(response, x_arm, y_arm, RA, DEC, psi, gmst_tref, + Qmax): + """JAX analogue of ``response_coefficients``: ``{p: (S,) complex}``. + + Parameters + ---------- + response : (3,3) host array detector response tensor. + x_arm, y_arm : (3,) host arrays Earth-fixed arm unit vectors. + RA, DEC, psi : (S,) JAX arrays. + gmst_tref : float GMST(tref) [rad], host constant. + Qmax : int highest arm-projection power retained. + + b_0 = F0 (exact lal baseline), b_{1+q} = beta_q. The arm length L does NOT + enter here (it lives in the precompute's W_p weights); only the arm *unit + vectors* enter, through the projections a_x, a_y and zx, zy. + """ + RA = jnp.asarray(RA, dtype=jnp.float64) + DEC = jnp.asarray(DEC, dtype=jnp.float64) + psi = jnp.asarray(psi, dtype=jnp.float64) + D = jnp.asarray(response, dtype=jnp.float64) + xa = jnp.asarray(np.asarray(x_arm, dtype=float)) + ya = jnp.asarray(np.asarray(y_arm, dtype=float)) + + g = gmst_tref - RA + X, Y, nhat = _triad_jax(DEC, psi, g) + Fp_lwl, Fc_lwl = _lwl_response_jax(D, X, Y) + F0 = Fp_lwl + 1j * Fc_lwl # (S,), exact lal baseline + + Xx = jnp.einsum('...i,i->...', X, xa) + Yx = jnp.einsum('...i,i->...', Y, xa) + Xy = jnp.einsum('...i,i->...', X, ya) + Yy = jnp.einsum('...i,i->...', Y, ya) + zx = Xx + 1j * Yx + zy = Xy + 1j * Yy + ax = jnp.einsum('...i,i->...', nhat, xa) + ay = jnp.einsum('...i,i->...', nhat, ya) + zx2, zy2 = zx ** 2, zy ** 2 + + b = {0: F0} + for q in range(Qmax + 1): + b[1 + q] = 0.5 * (zx2 * ax ** q - zy2 * ay ** q) + return b + + +def pack_coefficients(coeff_dict, p_list, S): + """Align a ``{p: (S,)}`` coefficient dict to the fixed ``p_list`` order. + + Returns a ``(A, S)`` complex JAX array, row ``i`` = ``coeff_dict[p_list[i]]``. + """ + rows = [] + for p in p_list: + rows.append(jnp.broadcast_to(coeff_dict[int(p)], (S,)).astype(jnp.complex128)) + return jnp.stack(rows, axis=0) + + +def reflection_index(p_list): + """Identity map (A,) -- the finite-size V term needs no reflection (W_p Hermitian).""" + return np.arange(len(p_list), dtype=np.int64) + + +def response_coefficients_packed(response, x_arm, y_arm, RA, DEC, psi, gmst_tref, + Qmax, p_list): + """Convenience: ``response_coefficients_dict`` + ``pack_coefficients`` -> (A,S).""" + S = int(jnp.asarray(RA).shape[0]) + cdict = response_coefficients_dict(response, x_arm, y_arm, RA, DEC, psi, + gmst_tref, Qmax) + return pack_coefficients(cdict, p_list, S) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py new file mode 100644 index 000000000..bd926603f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/response_slowrot.py @@ -0,0 +1,202 @@ +""" +JAX ports of the slow-rotation (Path A / Path B) response coefficients. + +The extrinsic layer of the rotation-aware likelihood +(``factored_likelihood_with_rotation``) needs, per detector, the complex +coefficients ``C_a`` of each elementary modulated template ``a = (p, n)`` -- +these multiply the (sky-independent) precompute banks ``Q^a_lm``, ``U^{a,a'}``, +``V^{a,a'}``. ``C_a`` is a closed-form analytic function of ``(RA, DEC, psi)`` +(with ``GMST(tref)`` a host constant), so it ports cleanly to ``jax.numpy`` and +is differentiable in the sky/polarization angles. + +This mirrors, term for term, the numpy reference +``factored_likelihood_with_rotation.rotation_coefficients_vector`` (which in turn +builds on ``slowrot_response.antenna_harmonics_vector`` / +``delay_harmonics_vector``). Validated against it to ~1e-12 by +``test/jax/test_jax_slowrot.py``. + +The detector-fixed inputs (``response`` tensor, ``location`` vector) are host +constants supplied by the caller (from ``lalsimulation.DetectorPrefixToLALDetector``); +only ``DEC, psi, RA`` are JAX (differentiable) leaves and ``gmst_tref`` a host float. +""" + +import math + +import numpy as np +import jax.numpy as jnp + +# Sidereal angular rate [rad/s]; identical to +# factored_likelihood_with_rotation.OMEGA_EARTH (used only for meta consistency). +OMEGA_EARTH = 7.292115e-5 +C_SI = 299792458.0 # m/s, matches slowrot_response.C_SI + + +def _antenna_harmonics_jax(D, dec, psi): + """JAX port of ``slowrot_response.antenna_harmonics_vector``. + + Returns ``{n: (S,) complex}`` for ``n in (-2,-1,0,1,2)`` -- the complex + antenna-pattern harmonics ``A_n`` with ``F(t)=sum_n A_n exp(i n g)``, + ``g = GMST(t) - RA``. Depends only on the (host) response tensor ``D`` and + the (JAX) declination/polarization. + """ + D = jnp.asarray(D, dtype=jnp.float64) + dec = jnp.asarray(dec, dtype=jnp.float64) + psi = jnp.asarray(psi, dtype=jnp.float64) + cd, sd = jnp.cos(dec), jnp.sin(dec) + cp, sp = jnp.cos(psi), jnp.sin(psi) + z = jnp.zeros_like(dec) + + def vec(a, b, c): + return jnp.stack(jnp.broadcast_arrays(a, b, c), axis=-1) # (S,3) + + Xc = vec(-sp * sd, -cp, z) + Xs = vec(-cp, sp * sd, z) + X0 = vec(z, z, sp * cd) + Yc = vec(-cp * sd, sp, z) + Ys = vec(sp, cp * sd, z) + Y0 = vec(z, z, cp * cd) + + def B(u, v): + return jnp.einsum('...i,ij,...j->...', u, D, v) + + Pp0 = 0.5 * (B(Xc, Xc) + B(Xs, Xs)) + B(X0, X0) - (0.5 * (B(Yc, Yc) + B(Ys, Ys)) + B(Y0, Y0)) + Pp1 = 2.0 * (B(Xc, X0) - B(Yc, Y0)) + Qp1 = 2.0 * (B(Xs, X0) - B(Ys, Y0)) + Pp2 = 0.5 * ((B(Xc, Xc) - B(Xs, Xs)) - (B(Yc, Yc) - B(Ys, Ys))) + Qp2 = B(Xc, Xs) - B(Yc, Ys) + + Pc0 = B(Xc, Yc) + B(Xs, Ys) + 2.0 * B(X0, Y0) + Pc1 = 2.0 * (B(Xc, Y0) + B(X0, Yc)) + Qc1 = 2.0 * (B(Xs, Y0) + B(X0, Ys)) + Pc2 = B(Xc, Yc) - B(Xs, Ys) + Qc2 = B(Xc, Ys) + B(Xs, Yc) + + P0 = Pp0 + 1j * Pc0 + P1 = Pp1 + 1j * Pc1 + Q1 = Qp1 + 1j * Qc1 + P2 = Pp2 + 1j * Pc2 + Q2 = Qp2 + 1j * Qc2 + return { + 0: P0, + 1: 0.5 * (P1 - 1j * Q1), + -1: 0.5 * (P1 + 1j * Q1), + 2: 0.5 * (P2 - 1j * Q2), + -2: 0.5 * (P2 + 1j * Q2), + } + + +def _delay_harmonics_jax(location, dec): + """JAX port of ``slowrot_response.delay_harmonics_vector``. + + Returns ``{m: (S,) complex}`` for ``m in (-1,0,1)`` -- the geometric-delay + harmonics ``B_m`` [s], ``tau(t)=sum_m B_m exp(i m g)``. + """ + r = np.asarray(location, dtype=float) + dec = jnp.asarray(dec, dtype=jnp.float64) + cd, sd = jnp.cos(dec), jnp.sin(dec) + T0 = -(r[2] * sd) / C_SI + T1c = -(cd * r[0]) / C_SI + T1s = (cd * r[1]) / C_SI + return {0: T0 + 0j, 1: 0.5 * (T1c - 1j * T1s), -1: 0.5 * (T1c + 1j * T1s)} + + +def _convolve_harmonics(a, b): + """Convolve two harmonic sequences (dicts {m: coef}) -> dict {m: coef}. + + Same as ``factored_likelihood_with_rotation._convolve_harmonics`` but the + coefficients are JAX arrays. + """ + out = {} + for m1, c1 in a.items(): + for m2, c2 in b.items(): + out[m1 + m2] = out.get(m1 + m2, 0.0) + c1 * c2 + return out + + +def rotation_coefficients_dict(response, location, RA, DEC, psi, gmst_tref, + p_max): + """JAX analogue of ``rotation_coefficients_vector``: ``{(p,n): (S,) complex}``. + + Parameters + ---------- + response : (3,3) host array detector response tensor. + location : (3,) host array detector location [m]. + RA, DEC, psi : (S,) JAX arrays (differentiable leaves). + gmst_tref : float GMST(tref) [rad], host constant. + p_max : int 0 = Path A (amplitude only); >=1 = Path B. + + Same algebra as the numpy reference; ``g_ev = gmst_tref - RA``. + """ + RA = jnp.asarray(RA, dtype=jnp.float64) + g_ev = gmst_tref - RA + A = _antenna_harmonics_jax(response, DEC, psi) + Atil = {n: A[n] * jnp.exp(1j * n * g_ev) for n in A} + if p_max == 0: + return {(0, n): Atil[n] for n in Atil} + Bd = _delay_harmonics_jax(location, DEC) + Btil = {m: Bd[m] * jnp.exp(1j * m * g_ev) for m in Bd} + tau0 = jnp.real(sum(Btil.values())) + D = {m: Btil[m] for m in Btil} + D[0] = D[0] - tau0 + negD = {m: -D[m] for m in D} + C = {} + E = {0: jnp.ones_like(g_ev, dtype=jnp.complex128)} + for p in range(p_max + 1): + if p > 0: + E = _convolve_harmonics(E, negD) + inv = 1.0 / math.factorial(p) + for n, an in Atil.items(): + for m, em in E.items(): + key = (p, n + m) + C[key] = C.get(key, 0.0) + inv * an * em + return C + + +def pack_coefficients(coeff_dict, a_list, S): + """Align a ``{(p,n): (S,)}`` coefficient dict to the fixed ``a_list`` order. + + Returns a ``(A, S)`` complex JAX array with row ``i`` = ``coeff_dict[a_list[i]]`` + (zeros where the dict lacks that key), exactly matching the numpy NoLoop + ``Cg`` behaviour (keys of the dict outside ``a_list`` are dropped; ``a_list`` + entries absent from the dict contribute zero). + """ + rows = [] + for a in a_list: + a = (int(a[0]), int(a[1])) + if a in coeff_dict: + rows.append(jnp.broadcast_to(coeff_dict[a], (S,)).astype(jnp.complex128)) + else: + rows.append(jnp.zeros((S,), dtype=jnp.complex128)) + return jnp.stack(rows, axis=0) # (A, S) + + +def reflection_index(a_list): + """Static index map ``i -> j`` with ``a_list[j] = (p, -n)`` for ``a_list[i]=(p,n)``. + + The rotation V cross term contracts ``C_{(p,-n)}`` (harmonic reflection). The + harmonic set is symmetric (the reference asserts this), so ``(p,-n)`` is always + present in ``a_list`` and the map is total. Returns an ``(A,)`` int numpy array. + """ + a_list = [(int(p), int(n)) for (p, n) in a_list] + pos = {a: i for i, a in enumerate(a_list)} + refl = [] + for (p, n) in a_list: + key = (p, -n) + if key not in pos: + raise ValueError( + "reflection partner (p,-n)=%r absent from a_list -- harmonic set " + "must be symmetric for the V term" % (key,)) + refl.append(pos[key]) + return np.asarray(refl, dtype=np.int64) + + +def rotation_coefficients_packed(response, location, RA, DEC, psi, gmst_tref, + p_max, a_list): + """Convenience: ``rotation_coefficients_dict`` + ``pack_coefficients``. + + Returns a ``(A, S)`` complex array aligned to ``a_list``. + """ + S = int(jnp.asarray(RA).shape[0]) + cdict = rotation_coefficients_dict(response, location, RA, DEC, psi, + gmst_tref, p_max) + return pack_coefficients(cdict, a_list, S) diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py index 42258535e..90a365f01 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/samplers.py @@ -288,6 +288,8 @@ def multistart_nuts(like, d_min, d_max, n_starts=8, num_warmup=300, num_samples=500, n_prior_pilot=8000, seed=0, target_accept=0.8, min_sep=0.3, proposal_inflate=2.0, n_is=40000, sky_coords="equatorial", + dense_mass=True, max_tree_depth=10, rotate_phase=False, + polish_seeds=True, extra_seeds=None, verbose=False, chain_progress_bar=False): """Multimodal posterior sampling by multi-start gradient-based NUTS. @@ -316,6 +318,49 @@ def multistart_nuts(like, d_min, d_max, n_starts=8, num_warmup=300, Base PRNG seed (numpy + JAX). target_accept : float NUTS target acceptance probability. + dense_mass : bool + Adapt a FULL (dense) mass matrix during warmup instead of the default + diagonal one. The distance-marginalized angular posterior is strongly + correlated -- at high SNR it is a thin, curved sky ring entangled with + the psi/incl/phiref degeneracies -- so a diagonal mass matrix leaves the + Hamiltonian geometry wildly anisotropic and NUTS hits ``max_tree_depth`` + (~2^depth leapfrog steps) on essentially every sample, stalling the run. + A dense mass matrix ≈ the inverse posterior covariance whitens the + geometry so trajectories are short and acceptance is high. ``True`` is + the sane production default for this problem; only set ``False`` for a + deliberately cheap low-SNR run where the posterior is broad and round. + max_tree_depth : int or (int, int) + NUTS maximum tree depth (numpyro passthrough). Bounds the worst-case + leapfrog steps per sample (``2^depth``) so an ill-conditioned *early* + warmup window -- before the mass matrix has adapted -- cannot blow up + wall-clock. A ``(warmup_depth, sampling_depth)`` tuple caps warmup more + tightly than sampling; a scalar applies to both. + rotate_phase : bool + Sample the rotated "polarization-phase" coordinates + ``phase_p = phiref + psi`` and ``phase_m = phiref - psi`` (each over + ``[0, 4pi)``) instead of ``(psi, phiref)`` directly, then map back + ``psi = (phase_p - phase_m)/2``, ``phiref = (phase_p + phase_m)/2``. + This is the JAX mirror of production RIFT's ``--internal-rotate-phase``: + the quadrupole-dominated likelihood depends on ``2psi +/- 2phiref``, so + the curved psi/phiref degeneracy ridge becomes AXIS-ALIGNED in + ``(phase_p, phase_m)`` -- the sampler's (dense) mass matrix is then + near-diagonal and NUTS keeps a healthy step at high SNR. The map is a + constant-Jacobian rotation, so the flat prior is preserved (exactly, in + the enlarged periodic domain). Combine with ``sky_coords="network"`` + (which similarly straightens the sky time-delay ring) for the full + high-SNR reparameterization. Exact for the (2,+/-2) quadrupole; still a + valid (just less-perfectly-decorrelating) reparameterization with higher + modes. + polish_seeds : bool + Gradient-ascend (+ Newton) each pilot seed to its local MAP before + running NUTS. At very high SNR the sky posterior is a ~1/SNR-thin ring + that a finite prior pilot cannot land ON -- the best raw pilot draw sits + many nats below the peak (e.g. SNR=1000: seed lnL ~24500 below + 0.5), so a chain started there samples the wrong arc (MAP degrees + off truth). A few hundred AD-gradient steps + Fisher-inverse Newton + steps climb from the broad basin onto the true peak, so NUTS starts AT + the needle -- the whole point of having exact gradients. Cheap + (a few hundred grad evals per seed); default on. min_sep : float Minimum angular separation (radians, in the combined sky+angle metric) between seeds. @@ -356,6 +401,33 @@ def multistart_nuts(like, d_min, d_max, n_starts=8, num_warmup=300, print(" chose %d seeds (lnL): %s" % (len(seeds), np.array2string(seed_lnL, precision=1))) + # Optional caller-supplied seeds (each a length-5 (ra,dec,psi,incl,phiref)), + # PREPENDED to the pilot seeds before the polish. At very high SNR the true + # peak is thinner than 1 pilot draw can resolve (~(1/SNR)^2 of the sky), so a + # blind pilot + gradient polish can settle on a secondary mode nats below the + # global peak; a known seed near the true basin (in production: the intrinsic + # grid + coarse extrinsic pass; here: the injected truth) guarantees one chain + # characterizes the injected mode. Still polished, so it snaps to the exact MAP. + if extra_seeds is not None: + ex = np.atleast_2d(np.asarray(extra_seeds, dtype=float)) + seeds = np.vstack([ex, seeds]) + seed_lnL = np.concatenate([eval_lnL(like, ex), seed_lnL]) + if verbose: + print(" + %d caller seed(s) (lnL): %s" % + (len(ex), np.array2string(eval_lnL(like, ex), precision=1))) + + # Gradient MAP-polish: climb each raw pilot seed onto the true (1/SNR-thin) + # peak so NUTS starts AT the needle rather than degrees off it on the wrong + # arc. Uses the exact JAX gradient (+ Fisher-inverse Newton); cheap. Keeps + # each seed at its OWN local MAP (preserves the multi-start mode coverage). + if polish_seeds: + _, _, _pol = _map_polish_4(like, seeds, bounds=_BOUNDS5) + seeds = np.array([p[0] for p in _pol]) + seed_lnL = np.array([p[1] for p in _pol]) + if verbose: + print(" polished seeds to MAP (lnL): %s" + % np.array2string(seed_lnL, precision=1)) + # Optional: sample the sky in NETWORK-frame coordinates (polar axis = the # baseline of the first two detectors), which folds the time-delay ring onto # a constant-polar-angle line. Falls back to equatorial if <2 detectors. @@ -379,57 +451,89 @@ def multistart_nuts(like, d_min, d_max, n_starts=8, num_warmup=300, # well-conditioned space with the prior Jacobians handled automatically; # the uniform sky prior is uniform in (cos_theta_n, phi_n) too, since the # rotation preserves the sphere measure. + # Shared phase parameterization: either sample (psi, phiref) directly, or + # the rotated (phase_p, phase_m) = (phiref+psi, phiref-psi) that decorrelate + # the 2psi+/-2phiref degeneracy (production --internal-rotate-phase). + _4PI = 4.0 * _PI + + def _sample_phase(): + if rotate_phase: + pp = numpyro.sample("phase_p", dist.Uniform(0.0, _4PI)) + pm = numpyro.sample("phase_m", dist.Uniform(0.0, _4PI)) + return (pp - pm) * 0.5, (pp + pm) * 0.5 # psi, phiref + psi = numpyro.sample("psi", dist.Uniform(0.0, _PI)) + phiref = numpyro.sample("phiref", dist.Uniform(0.0, _TWO_PI)) + return psi, phiref + + def _init_phase(th0): + psi0, phi0 = float(th0[2]), float(th0[4]) + if rotate_phase: + return {"phase_p": (phi0 + psi0) % _4PI, + "phase_m": (phi0 - psi0) % _4PI} + return {"psi": psi0, "phiref": phi0} + + def _extract_phase(s): + if rotate_phase: + pp = np.asarray(s["phase_p"]); pm = np.asarray(s["phase_m"]) + return np.mod((pp - pm) * 0.5, _PI), np.mod((pp + pm) * 0.5, _TWO_PI) + return np.mod(np.asarray(s["psi"]), _PI), np.mod(np.asarray(s["phiref"]), _TWO_PI) + if net is None: def model(): ra = numpyro.sample("ra", dist.Uniform(0.0, _TWO_PI)) sin_dec = numpyro.sample("sin_dec", dist.Uniform(-1.0, 1.0)) - psi = numpyro.sample("psi", dist.Uniform(0.0, _PI)) cos_incl = numpyro.sample("cos_incl", dist.Uniform(-1.0, 1.0)) - phiref = numpyro.sample("phiref", dist.Uniform(0.0, _TWO_PI)) + psi, phiref = _sample_phase() lnL = like._scalar(jnp.stack( [ra, jnp.arcsin(sin_dec), psi, jnp.arccos(cos_incl), phiref])) numpyro.factor("loglike", lnL) def make_init(th0): - return {"ra": float(th0[0]), "sin_dec": float(np.sin(th0[1])), - "psi": float(th0[2]), "cos_incl": float(np.cos(th0[3])), - "phiref": float(th0[4])} + d = {"ra": float(th0[0]), "sin_dec": float(np.sin(th0[1])), + "cos_incl": float(np.cos(th0[3]))} + d.update(_init_phase(th0)) + return d def extract(s): + psi, phiref = _extract_phase(s) return np.stack([np.asarray(s["ra"]), np.arcsin(np.asarray(s["sin_dec"])), - np.asarray(s["psi"]), np.arccos(np.asarray(s["cos_incl"])), - np.asarray(s["phiref"])], axis=-1) + psi, np.arccos(np.asarray(s["cos_incl"])), + phiref], axis=-1) else: _C, R, gmst = net def model(): cos_tn = numpyro.sample("cos_theta_n", dist.Uniform(-1.0, 1.0)) phi_n = numpyro.sample("phi_n", dist.Uniform(0.0, _TWO_PI)) - psi = numpyro.sample("psi", dist.Uniform(0.0, _PI)) cos_incl = numpyro.sample("cos_incl", dist.Uniform(-1.0, 1.0)) - phiref = numpyro.sample("phiref", dist.Uniform(0.0, _TWO_PI)) + psi, phiref = _sample_phase() ra, dec = _C.network_to_equatorial(jnp.arccos(cos_tn), phi_n, R, gmst) lnL = like._scalar(jnp.stack([ra, dec, psi, jnp.arccos(cos_incl), phiref])) numpyro.factor("loglike", lnL) def make_init(th0): tn, pn = _C.equatorial_to_network(float(th0[0]), float(th0[1]), R, gmst) - return {"cos_theta_n": float(np.cos(float(tn))), - "phi_n": float(float(pn) % _TWO_PI), - "psi": float(th0[2]), "cos_incl": float(np.cos(th0[3])), - "phiref": float(th0[4])} + d = {"cos_theta_n": float(np.cos(float(tn))), + "phi_n": float(float(pn) % _TWO_PI), + "cos_incl": float(np.cos(th0[3]))} + d.update(_init_phase(th0)) + return d def extract(s): tn = np.arccos(np.asarray(s["cos_theta_n"])) ra, dec = _C.network_to_equatorial(tn, np.asarray(s["phi_n"]), R, gmst) - return np.stack([np.asarray(ra), np.asarray(dec), np.asarray(s["psi"]), + psi, phiref = _extract_phase(s) + return np.stack([np.asarray(ra), np.asarray(dec), psi, np.arccos(np.asarray(s["cos_incl"])), - np.asarray(s["phiref"])], axis=-1) + phiref], axis=-1) # -- 2. one NUTS chain per seed, pooled -------------------------------- + # (len(seeds) may exceed n_starts when the caller supplies extra_seeds) per_chain = [] # (num_samples, 5) per seed, in equatorial theta5 - for k in range(n_starts): + n_chains = len(seeds) + for k in range(n_chains): kernel = NUTS(model, target_accept_prob=target_accept, + dense_mass=dense_mass, max_tree_depth=max_tree_depth, init_strategy=init_to_value(values=make_init(seeds[k]))) mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples, num_chains=1, progress_bar=chain_progress_bar) @@ -437,11 +541,46 @@ def extract(s): per_chain.append(extract(mcmc.get_samples())) if verbose: print(" chain %d/%d done (seed lnL=%.2f)" % - (k + 1, n_starts, seed_lnL[k])) + (k + 1, n_chains, seed_lnL[k])) theta = np.concatenate(per_chain, axis=0) lnL = eval_lnL(like, theta) + # Evidence-weighted pooling. Multi-start places one chain per mode, but the + # modes carry vastly different posterior mass -- at high SNR the sub-dominant + # time-delay-ring images and amplitude-degeneracy branches sit many nats below + # the true peak. Pooling the chains with EQUAL weight over-represents those + # negligible modes (they dominate a naive credible-region or corner plot). + # Weight each chain by a LAPLACE estimate of its mode evidence, + # log Z_k ~= peak_lnL_k + 1/2 log det Sigma^sky_k , + # i.e. the mode's peak likelihood times its (sky) width. Peak alone is wrong: + # at LOW SNR the chains sample one broad, overlapping posterior with similar + # peaks, and a tiny peak difference would spuriously collapse it -- the width + # term keeps broad modes comparable there, while at HIGH SNR the sub-dominant + # modes are suppressed by their far-lower peak regardless of width. ``theta`` + # stays the raw pooled draws; ``post_weight`` is the per-sample posterior weight + # callers use for credible regions, sky areas, and corner plots. + n_per = [len(c) for c in per_chain] + _off = np.cumsum([0] + n_per) + logev = np.full(len(per_chain), -np.inf) + for k in range(len(per_chain)): + if n_per[k] < 3: + continue + thk = per_chain[k] + ra_k, dec_k = thk[:, 0], thk[:, 1] + ra0 = np.angle(np.mean(np.exp(1j * ra_k))) + x = ((ra_k - ra0 + np.pi) % (2 * np.pi) - np.pi) * np.cos(np.median(dec_k)) + y = dec_k - np.median(dec_k) + cov = np.cov(np.vstack([x, y])) + 1e-8 * np.eye(2) + logdet = float(np.log(max(np.linalg.det(cov), 1e-30))) + peak_k = float(lnL[_off[k]:_off[k + 1]].max()) + logev[k] = peak_k + 0.5 * logdet + cw = np.exp(logev - np.max(logev)) + post_weight = np.concatenate([ + np.full(n_per[k], cw[k] / max(n_per[k], 1)) for k in range(len(per_chain))]) + sw = post_weight.sum() + post_weight = post_weight / sw if sw > 0 else np.full(len(theta), 1.0 / max(len(theta), 1)) + # -- 3. Gaussian-mixture importance evidence (one comp per seed) ------- mus, covs = [], [] for th in per_chain: @@ -477,8 +616,14 @@ def extract(s): logZ, sigma_over_Z, neff = _finalize_evidence( logZ, sigma_over_Z, neff, float(np.max(lnL)) if len(lnL) else np.nan) + # Per-chain posterior draws (stacked) so callers can compute a POSTERIOR + # effective-sample-size / R-hat -- the right "did we resolve the posterior" + # diagnostic, distinct from the importance-sampling evidence ``neff`` above + # (which is limited by the Gaussian-mixture proposal's fit to the target). + theta_per_chain = np.stack(per_chain, axis=0) # (n_starts, num_samples, 5) return dict(theta=theta, lnL=lnL, seeds=seeds, seed_lnL=seed_lnL, - logZ=logZ, sigma_over_Z=sigma_over_Z, neff=neff) + logZ=logZ, sigma_over_Z=sigma_over_Z, neff=neff, + theta_per_chain=theta_per_chain, post_weight=post_weight) # --------------------------------------------------------------------------- @@ -635,6 +780,9 @@ def logpdf(theta5, data): _BOUNDS4 = [(0.0, _TWO_PI), (-_PI / 2 + 1e-3, _PI / 2 - 1e-3), (0.0, _PI), (1e-3, _PI - 1e-3)] +# 5-D support (ra, dec, psi, incl, phiref) for MAP-polishing the 5-D +# distance-marginalized seeds in multistart_nuts. +_BOUNDS5 = _BOUNDS4 + [(0.0, _TWO_PI)] def sample_prior_4(n, rng): @@ -1894,10 +2042,12 @@ def model(): "logZ=%.3f neff(IS)=%.1f mode mass=%s" % (len(theta), K, logZ, neff, np.array2string(mass, precision=3))) + theta_per_chain = per_chain # list of (num_samples, 4) arrays, one per mode return dict(theta=theta, lnL=lnL, post_weight=post_weight, logZ=logZ, sigma_over_Z=sigma_over_Z, neff=neff, theta_map=modes[0], modes=modes, mode_lnL=mode_lnL, - mode_logZ=mode_logZ, flow_state=None) + mode_logZ=mode_logZ, theta_per_chain=theta_per_chain, + flow_state=None) # --------------------------------------------------------------------------- diff --git a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py index 474491619..ee51cb4d0 100644 --- a/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py +++ b/MonteCarloMarginalizeCode/Code/RIFT/likelihood/jax_ile/wrapper.py @@ -36,6 +36,102 @@ EXTRINSIC_PARAM_ORDER = ("ra", "dec", "psi", "incl", "phiref", "distMpc") +def build_rotation_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, + integration_window_half, Lmax, fMax, + t_window=0.1, harmonics=(-2, -1, 0, 1, 2), + p_max=0, analyticPSD_Q=False, + inv_spec_trunc_Q=False, T_spec=0.0, + tvals=None, verbose=False, + **precompute_kwargs): + """One-call builder for the slow-rotation (Path A/B) banded JAX likelihood. + + Runs the production ``PrecomputeLikelihoodTermsWithRotation`` + + ``pack_rotation_arrays`` (heavy, data-touching -- reused verbatim) and wraps + the packed banks into a banded :class:`JAXLikelihoodData`. The returned + object flows through every ``fused_log_likelihood*`` / marginalization + variant and the samplers exactly like the baseline data. + + ``t_window`` is the rholm-buffer half width for the rotation precompute (it + builds its own buffer, unlike the baseline two-window driver); ``tvals`` is + the marginalization grid (defaults to ``linspace(-iwh, iwh, 2*iwh/deltaT)``). + """ + import RIFT.likelihood.factored_likelihood_with_rotation as flwr + from .banded import build_rotation_data + + ri, ct, ctV, rho, meta = flwr.PrecomputeLikelihoodTermsWithRotation( + fiducial_epoch, t_window, P, data_dict, psd_dict, Lmax, fMax, + harmonics=harmonics, p_max=p_max, f_sidereal=flwr.F_SIDEREAL, + analyticPSD_Q=analyticPSD_Q, inv_spec_trunc_Q=inv_spec_trunc_Q, + T_spec=T_spec, verbose=verbose, quiet=not verbose, + skip_interpolation=True, **precompute_kwargs) + lk, rbn, ubn, vbn, ep = flwr.pack_rotation_arrays(meta, rho, ct, ctV) + + deltaT = float(P.deltaT) + if tvals is None: + # tvals spaced EXACTLY by deltaT (arange, not linspace) so the grid matches + # the pos<->sample mapping and Simpson weights the likelihood assumes; the + # maintained NoLoop path uses this same arange(-Nw,Nw)*deltaT convention. + # (A linspace grid is spaced deltaT*npts/(npts-1) and shifts the time + # reference by a fraction of a sample -> a sky bias that only shows up at + # high SNR, where cubic interpolation resolves the razor-sharp peak.) + Nw = int(integration_window_half / deltaT) + tvals = np.arange(-Nw, Nw) * deltaT + data = build_rotation_data(meta, lk, rbn, ubn, vbn, ep, deltaT, tvals) + extras = dict(meta=meta, rho_by_a=rbn, U_by_aa=ubn, V_by_aa=vbn, + epochDict=ep, lookupNKDict=lk) + return data, extras + + +def build_freqresponse_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, + integration_window_half, Lmax, fMax, + t_window=0.1, Qmax=4, L_arm=None, + analyticPSD_Q=False, + inv_spec_trunc_Q=False, T_spec=0.0, + tvals=None, verbose=False, + **precompute_kwargs): + """One-call builder for the finite-size (Path D) banded JAX likelihood. + + Runs ``PrecomputeLikelihoodTermsFreqResponse`` + ``pack_freqresponse_arrays`` + (reused verbatim) and wraps the packed banks into a banded + :class:`JAXLikelihoodData`. ``L_arm`` overrides the arm length (e.g. 40000. + for a 40-km CE arm; ``None`` = native LAL arm lengths). The finite-size + coefficients need the arm *unit vectors*, so the detector geometry is + recomputed via ``slowrot_freqresponse.detector_geometry`` and attached. + """ + import RIFT.likelihood.factored_likelihood_freqresponse as flfr + import RIFT.likelihood.slowrot_freqresponse as sfr + from .banded import build_freqresponse_data + + bk = flfr.PrecomputeLikelihoodTermsFreqResponse( + fiducial_epoch, t_window, P, data_dict, psd_dict, Lmax, fMax, + Qmax=Qmax, L_arm=L_arm, analyticPSD_Q=analyticPSD_Q, + inv_spec_trunc_Q=inv_spec_trunc_Q, T_spec=T_spec, verbose=verbose, + quiet=not verbose, skip_interpolation=True, **precompute_kwargs) + meta = bk[4] + lk, rbp, ubp, vbp, ep = flfr.pack_freqresponse_arrays(bk[4], bk[3], bk[1], bk[2]) + + def _L_of(det): + return L_arm.get(det, None) if isinstance(L_arm, dict) else L_arm + det_geom = {det: sfr.detector_geometry(det, L_arm=_L_of(det)) + for det in data_dict.keys()} + + deltaT = float(P.deltaT) + if tvals is None: + # tvals spaced EXACTLY by deltaT (arange, not linspace) so the grid matches + # the pos<->sample mapping and Simpson weights the likelihood assumes; the + # maintained NoLoop path uses this same arange(-Nw,Nw)*deltaT convention. + # (A linspace grid is spaced deltaT*npts/(npts-1) and shifts the time + # reference by a fraction of a sample -> a sky bias that only shows up at + # high SNR, where cubic interpolation resolves the razor-sharp peak.) + Nw = int(integration_window_half / deltaT) + tvals = np.arange(-Nw, Nw) * deltaT + data = build_freqresponse_data(meta, lk, rbp, ubp, vbp, ep, deltaT, tvals, + det_geom) + extras = dict(meta=meta, rho_by_p=rbp, U_by_pp=ubp, V_by_pp=vbp, + epochDict=ep, lookupNKDict=lk, det_geom=det_geom) + return data, extras + + def build_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, storage_window_half, integration_window_half, Lmax, fMax, @@ -83,9 +179,10 @@ def build_data_from_precompute(P, data_dict, psd_dict, fiducial_epoch, deltaT = float(P.deltaT) if tvals is None: - npts = int(2 * integration_window_half / deltaT) - tvals = np.linspace(-integration_window_half, - integration_window_half, npts) + # arange(-Nw,Nw)*deltaT: spacing exactly deltaT (see the freqresponse + # builder) -- matches the maintained NoLoop tvals convention. + Nw = int(integration_window_half / deltaT) + tvals = np.arange(-Nw, Nw) * deltaT data = build_likelihood_data(packed, deltaT, float(fiducial_epoch), tvals) extras = dict(rholms=rholms, cross_terms=cross_terms, diff --git a/MonteCarloMarginalizeCode/Code/bin/calibration_reweighting.py b/MonteCarloMarginalizeCode/Code/bin/calibration_reweighting.py index 63b61a4e7..96886668e 100755 --- a/MonteCarloMarginalizeCode/Code/bin/calibration_reweighting.py +++ b/MonteCarloMarginalizeCode/Code/bin/calibration_reweighting.py @@ -58,6 +58,7 @@ # TODO this should not be a hardcoded path! import RIFT.calmarg.rift_source as rift_source +from RIFT.calmarg.calibration import correction_type_for_ifo from bilby.core.utils import logger @@ -316,6 +317,8 @@ def alt_reweight(result, label=None, new_likelihood=None, new_prior=None, spline_calibration_envelope_dict = bilby_pipe.utils.convert_string_to_dict( data.meta_data['command_line_args']['spline_calibration_envelope_dict']) +calibration_correction_type = data.meta_data['command_line_args'].get( + 'calibration_correction_type') ifos_for_reweighting = deepcopy(ifos) for ifo in ifos: # removes any model for the calibration that was set up in the file ifo.calibration_model = bilby.gw.calibration.Recalibrate() @@ -395,7 +398,10 @@ def alt_reweight(result, label=None, new_likelihood=None, new_prior=None, if args.use_local_cal_files: calibration_file_path = './cal_envelopes/' + os.path.basename(calibration_file_path) # force local, specific name. Copied in place earlier ifo_calibration_priors = bilby.gw.prior.CalibrationPriorDict.from_envelope_file( - calibration_file_path, ifo.minimum_frequency, ifo.maximum_frequency, 10, ifo.name) + calibration_file_path, ifo.minimum_frequency, ifo.maximum_frequency, 10, + ifo.name, correction_type=correction_type_for_ifo( + calibration_correction_type, ifo.name, + parse_dict=bilby_pipe.utils.convert_string_to_dict)) # TODO FOR DEBUGGING PURPOSES # for key in ifo_calibration_priors.keys(): diff --git a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py index aeec1ae4a..e300aefd3 100755 --- a/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py +++ b/MonteCarloMarginalizeCode/Code/bin/helper_LDG_Events.py @@ -217,6 +217,8 @@ def get_observing_run(t): parser.add_argument("--internal-ile-rotate-phase", action='store_true') parser.add_argument("--internal-ile-auto-logarithm-offset",action='store_true',help="Passthrough to ILE") parser.add_argument("--internal-ile-use-lnL",action='store_true',help="Passthrough to ILE. Will DISABLE auto-logarithm-offset and manual-logarithm-offset") +parser.add_argument("--internal-ile-n-chunk",default=None,type=int,help="Override the extrinsic chunk size (--n-chunk) passed to ILE. Default: 40000, scaled linearly with SNR above 40 and capped at 160000. Rationale: at high SNR the posterior is a vanishing fraction of the prior volume, so a small chunk gives few informative samples per adaptation step; measured collapse on a truth-known SNR ladder falls 88%%->50%% (SNR160) and 69%%->25%% (SNR80) going 1e4->1.6e5, and the gain survives at fixed budget. Larger chunks cost GPU memory, so raise the ILE memory request if you raise this a lot.") +parser.add_argument("--internal-ile-interpolate-time",action='store_true',help="Evaluate Q_lm at FRACTIONAL detector times by cubic interpolation instead of snapping to the nearest sample bin (passes --interpolate-time True). Requires the maintained NoLoop likelihood, i.e. the --vectorized --gpu --force-xpy combination. Nearest-bin evaluation injects a time-quantization non-smoothness into the extrinsic likelihood surface that is a discretization artifact, not physics; removing it makes convergence more robust. Default off for backward compatibility.") parser.add_argument("--internal-cip-use-lnL",action='store_true') parser.add_argument("--ile-n-eff",default=50,type=int,help="Target n_eff passed to ILE. Try to keep above 2") parser.add_argument("--test-convergence",action='store_true',help="If present, the code will terminate if the convergence test passes. WARNING: if you are using a low-dimensional model the code may terminate during the low-dimensional model!") @@ -1081,6 +1083,34 @@ def crit_m2(delta): else: helper_cip_args_extra = " --internal-use-lnL " # always use lnL scaling for loud signals at late times, overflow issue for most integrators rescaled_base_ile = True + +# EXTRINSIC CHUNK SIZE, SCALED WITH SNR. +# At high SNR the extrinsic posterior is a vanishing fraction of the prior volume, so a small chunk +# carries very few informative samples per adaptation step and the sampler adapts on noise. Measured +# on a truth-known synthetic SNR ladder (demos/integrator_snr_lottery; AV, d=4, collapse = failure to +# recover the known lnZ), at EQUAL adaptation steps: +# SNR 40: 31% collapse @1e4 -> 19% @4e4 -> 0% @1.6e5 +# SNR 80: 69% collapse @1e4 -> 38% @4e4 -> 25% @1.6e5 +# SNR 160: 88% collapse @1e4 -> 62% @4e4 -> 50% @1.6e5 +# and the gain SURVIVES at fixed total budget (pooled Fisher p=0.014 for SNR>=80, 10k vs 160k), +# despite the larger chunk taking 16x fewer adaptation steps. The driver default (1e4) is therefore +# too small for loud events. 40k is the new baseline; scale up with SNR, capped, because GPU memory +# grows with the chunk and an over-large request matches fewer slots (held jobs / idle capacity). +# Override with --internal-ile-n-chunk. NOTE the skymap branch later sets --n-chunk 500 deliberately +# and must keep winning: optparse takes the LAST occurrence, and that append happens after this one. +if opts.internal_ile_n_chunk: + n_chunk_ile = int(opts.internal_ile_n_chunk) +else: + n_chunk_ile = 40000 + if "SNR" in event_dict.keys(): + # 40k up to SNR 40, then linear in SNR, capped at 160k (the largest size measured) + n_chunk_ile = int(40000 * np.max([1.0, event_dict["SNR"] / 40.0])) + n_chunk_ile = int(np.min([n_chunk_ile, 160000])) +helper_ile_args += " --n-chunk " + str(n_chunk_ile) + " " +if opts.internal_ile_interpolate_time: + # cubic Q_lm time interpolation; needs the NoLoop path (--vectorized --gpu --force-xpy) + helper_ile_args += " --interpolate-time True " + if opts.internal_ile_auto_logarithm_offset and not opts.internal_ile_use_lnL: helper_ile_args += " --auto-logarithm-offset " rescaled_base_ile = True diff --git a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode index 774c05700..379a6f4c0 100755 --- a/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode +++ b/MonteCarloMarginalizeCode/Code/bin/integrate_likelihood_extrinsic_batchmode @@ -243,6 +243,10 @@ optp.add_option("--calibration-proposal-breadcrumb",default=None, help="Opt-in ( optp.add_option("--calibration-dump-responsibilities",default=None, help="Opt-in (Option C / adaptive pilot): path to write per-cal-realization log-responsibilities (length n_cal), accumulated over the evaluated grid, plus the cal node draws. This is the pilot's output, fitted into a proposal by util_CalPilotFit.py. No effect on the returned likelihood.") optp.add_option("--calibration-pilot-extrinsic",default=256,type=int, help="Pilot only: number of uniform-prior extrinsic samples used to extrinsic-marginalize the per-realization cal responsibility at each intrinsic point. Cal is ~extrinsic-independent, so a modest batch suffices.") optp.add_option("--calibration-mc-error-extrinsic",default=8192,type=int, help="Calmarg error budget: CAP on the number of extrinsic-prior samples used to estimate the calibration Monte-Carlo contribution to the lnL error (per-realization responsibilities a_c -> Var(lnZ) ~= n_cal*Var_c(a_c)), added IN QUADRATURE to the reported sigma column. The batch is ADAPTIVE: it starts small and doubles until the estimate stabilizes or this cap is reached. Distance is drawn from the RUN'S distance prior (sampler prior / --d-prior; with a PINNED distance the probe runs at that fixed value and warns that the estimate is conservative). The extrinsic sampler's variance cannot see the spread over the (fixed) cal draw set, so without this term the reported error badly understates the truth whenever the cal n_eff is small. Set 0 to disable (restores the old, extrinsic-only sigma).") +optp.add_option("--mc-error-replicas",default=0,type=int, help="MC-error stabilization: when the reported lnL error is untrustworthy (see the trigger options below), re-run the extrinsic integration this many EXTRA times as cold replicas (adaptation reset, sample cache dropped, fresh RNG draws) and report lnL from the LINEAR mean of the replica integrals with sigma from the max of the propagated error and the between-replica scatter (t-distributed, K-1 dof). The naive per-run sigma is computed from the SAME weights as the integral, so it is small exactly when the run silently missed the peak; only independent replicas can see that. NEVER combine replicas by inverse-variance weighting -- that overweights the worst replica. The posterior/fairdraw export POOLS the replicas (weights renormalized so each contributes Z_k/K; fairdraw blocks contribute equal within-block weights, since those samples already carry their weights once), so the exported samples represent the same mixture as the reported evidence. Default 0 = off (production behavior unchanged).") +optp.add_option("--mc-error-sigma-trigger",default=0.4,type=float, help="Replicate (see --mc-error-replicas) when the reported sigma_lnZ exceeds this value.") +optp.add_option("--mc-error-khat-trigger",default=0.7,type=float, help="Replicate when the Pareto k-hat weight-tail diagnostic exceeds this value (0.7 = the PSIS reliability threshold: above it the weight variance is effectively unresolved and the naive sigma is a lower bound).") +optp.add_option("--mc-error-ess-trigger",default=30.,type=float, help="Replicate when the Kish effective sample size (sum w)^2/sum w^2 of the run's weights falls below this value.") optp.add_option("--calibration-neff-cal-target",default=10,type=float, help="Calmarg ADAPTIVE draw count: after the cal-block precompute, probe the effective number of contributing cal draws (neff_cal) at this intrinsic point; while it is below this target, DOUBLE the cal draw set (drawing fresh independent realizations and appending their precomputed blocks) up to --calibration-n-realizations-max. Set 0 to disable (fixed --calibration-n-realizations).") optp.add_option("--calibration-n-realizations-max",default=0,type=int, help="Cap for the adaptive cal draw count (see --calibration-neff-cal-target). Default 0 = 8x --calibration-n-realizations.") optp.add_option("--calibration-burn-in-neff",default=None,type=float, help="Opt-in: before the production cal-marginalized integration, BURN IN the extrinsic sampler on the cheap ZERO-CAL (n_cal=1) likelihood until this effective sample count, then switch to the full cal-marginalized likelihood. The extrinsic posterior is ~cal-independent. CAVEAT: the AV sampler RESETS between integrate() calls (no seedable AV yet), so this gives AV no speedup (correctness-safe only). It can warm-start GMM/portfolio (model reuse). Awaiting a seedable / boundary-shifting AV; see DESIGN_adaptive_driver.md. No effect unless calmarg is active.") @@ -263,7 +267,7 @@ optp.add_option("--freqresponse-arm-length", default=None, help="Arm-length over optp.add_option("--force-xpy", action="store_true", help="Use the xpy code path. Use with --vectorized --gpu to use the fallback CPU-based code path. Useful for debugging.") optp.add_option("-o", "--output-file", help="Save result to this file.") optp.add_option("-O", "--output-format", default='xml', help="[xml|hdf5]") -optp.add_option("-S", "--save-samples", action="store_true", help="Save sample points to output-file. Requires --output-file to be defined.") +optp.add_option("-S", "--save-samples", action="store_true", help="Save sample points to output-file (sparse sim_inspiral XML). Requires --output-file to be defined. NOTE: the XML carries lnL only, not the importance weight -- it does NOT persist the joint prior / sampling prior, so it must not be reweighted by likelihood for a weighted-posterior/shape check. For that, use the ASCII per-sample outputs --extrinsic-proposal-output (full log-weight) or --calibration-export-posterior.") optp.add_option("--save-samples-process-params", action="store_true", help="XML output retains process_params table, Default is not to do this") optp.add_option("-L", "--save-deltalnL", type=float, default=float("Inf"), help="Threshold on deltalnL for points preserved in output file. Requires --output-file to be defined") optp.add_option("-P", "--save-P", type=float,default=0.1, help="Threshold on cumulative probability for points preserved in output file. Requires --output-file to be defined") @@ -303,6 +307,10 @@ integration_params.add_option("--adapt-log",action='store_true',help="Use a loga integration_params.add_option("--internal-gmm-correlate-all",action='store_true',help="GMM sampler: use a SINGLE full-dimension GMM group instead of the default (sky)(distance,inclination)(psi,phi) pairing. The default pairing targets quadrupole-dominated binaries with a large sky ring; a product of per-group GMMs cannot represent cross-group correlations (e.g. sky-phase), and for a strongly-localized single-peak source the factored proposal can stall at the prior. Component count from --internal-gmm-sky-components (default 2 in this mode).") integration_params.add_option("--internal-gmm-sky-components",type=int,default=None,help="GMM sampler: number of mixture components for the (ra,dec) group (default 4, sized for a large sky ring; use 1-2 for a well-localized single peak, e.g. 3+ IFOs / high SNR). With --internal-gmm-correlate-all, sets the single full-dimension group's component count.") integration_params.add_option("--internal-gmm-phase-components",type=int,default=None,help="GMM sampler: number of mixture components for the (psi,phi_orb) group (default 4; use 1-2 for a single dominant phase peak).") +integration_params.add_option("--internal-gmm-adaptive-components",action='store_true',help="GMM sampler (FLEXIBLE allocation): choose each adaptive group's component count from the DATA by BIC each chunk (GMM.fit_gmm_adaptive), instead of the hard-coded per-group counts (sky=4,dist-incl=2,...) that target quadrupole/large-sky-ring binaries. BIC allocates more components only where the importance-weighted cloud is genuinely non-Gaussian (e.g. a curved distance-inclination arc) and stays at k=1 for a single blob; a defensive tail component (see --internal-gmm-defensive-frac) keeps the importance weights bounded. Cap per group via --internal-gmm-max-components.") +integration_params.add_option("--internal-gmm-max-components",type=int,default=8,help="Cap on the per-group component count for --internal-gmm-adaptive-components (default 8).") +integration_params.add_option("--internal-gmm-defensive-frac",type=float,default=0.0,help="Weight of the broad box-covering 'defensive' mixture component added to each adaptive GMM group (default 0 = OFF). Intended to bound the importance weights (Hesterberg defensive IS), but on a wide extrinsic prior the broad component draws physically-extreme points where the likelihood is NaN and it did not improve n_eff on the SNR~82 benchmark -- opt-in only.") +integration_params.add_option("--internal-gmm-inflate",type=float,default=1.0,help="Covariance inflation factor (std multiplier) applied to each adaptive GMM component (default 1.0 = none). A value >1 widens the proposal relative to the elite cloud it was fit to; complements --internal-gmm-defensive-frac.") integration_params.add_option("--interpolate-time", default=False,help="If using the maintained NoLoop likelihood, evaluate Q_lm at fractional detector times using cubic interpolation instead of nearest sample bins. Accepts truthy values such as True/1/yes. (Default=false)") 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") @@ -323,7 +331,47 @@ integration_params.add_option("--internal-use-lnL",action='store_true',help="lik integration_params.add_option("--sampler-method",default="adaptive_cartesian_gpu",help="adaptive_cartesian|GMM|adaptive_cartesian_gpu") integration_params.add_option("--sampler-portfolio",default=None,action='append',type=str,help="comma-separated strings, matching sampler methods other than portfolio") integration_params.add_option("--sampler-portfolio-args",default=None, action='append', type=str, help='eval-able dictionaryo to be passed to that sampler') +# Portfolio freeze-policy knobs (only meaningful with --sampler-method portfolio). A member +# whose balance weight drops below --portfolio-freeze-wt normally stops updating its proposal; +# these control that. Defaults (None here) mean "use the sampler's built-in default". +integration_params.add_option("--portfolio-grace-iters",default=None,type=int,help="Portfolio: never freeze ANY member during the first N integration chunks (let slow starters contract). Sampler default 25.") +integration_params.add_option("--portfolio-revive-period",default=None,type=int,help="Portfolio: every N chunks, update even a frozen member one step so it can recover. 0 disables. Sampler default 8.") +integration_params.add_option("--portfolio-freeze-wt",default=None,type=float,help="Portfolio: a member whose balance weight is below this stops updating its proposal (subject to grace/revive/VARAHA-exemption). Sampler default 0.05.") +integration_params.add_option("--portfolio-varaha-never-freeze",action='store_true',default=False,help="Portfolio: VARAHA/AV members always update every chunk past their breakpoint (freeze-exempt). This is the sampler default; the flag is here for explicitness/pipe pass-through.") +integration_params.add_option("--portfolio-varaha-can-freeze",action='store_true',default=False,help="Portfolio: DISABLE the VARAHA freeze-exemption, so VARAHA/AV members obey the grace/revive/weight freeze schedule like other members. Use only if a VARAHA member is a known-bad fit and you want to save its selfish-draw eval cycles.") +# Portfolio DRAW-ALLOCATION policy (adaptive-probe): OPT-IN (default off). Concentrates the draw +# budget on the member with the highest per-chunk n_ess, with round-robin probing. Unbiased for +# any allocation (q_mix). Helps on strongly-correlated targets, but the n_ess signal rewards +# self-consistency and STARVES a slow-contracting VARAHA/AV member on real high-SNR events, so it +# is not the default -- see DESIGN_portfolio_freeze_policy.md. +integration_params.add_option("--portfolio-adaptive-alloc",action='store_true',default=False,help="Portfolio: ENABLE (opt-in) adaptive-probe draw allocation -- concentrate draws on the best per-chunk-n_ess member. Good on strongly-correlated targets; NOT recommended for AV-favorable high-SNR events (it starves the slow-contracting AV workhorse). Off by default (legacy n_ess reweighting).") +integration_params.add_option("--portfolio-varaha-min-frac",default=None,type=float,help="Portfolio: reserve this combined DRAW fraction for VARAHA/AV members (0/unset = off). never-freeze keeps a VARAHA member UPDATING, but both allocation rules score by per-chunk n_ess, which sits at ~1 during VARAHA's slow cumulative contraction -- so a member that looks instantly good can take nearly the whole budget (measured on S250114ax post-#33: GMM took ~0.84 and the portfolio collapsed to n_eff ~2 vs ~100 for standalone AV). Unbiased for any allocation (q_mix); trades efficiency only.") +integration_params.add_option("--portfolio-varaha-max-frac",default=None,type=float,help="Portfolio: CAP the combined DRAW fraction of VARAHA/AV members (0/unset = no cap). Use WITH --portfolio-varaha-min-frac to constrain the VARAHA share to a BAND. Rationale: a floor alone stops the mixture degenerating to peaked-member-only (which strips q_mix of its broad backstop, so a missed mode goes uncovered and lnZ is silently low while n_eff looks GOOD), but the share can then run away the OTHER way to ~1 and the mixture degenerates to VARAHA-only instead. A band (e.g. 0.25/0.75) keeps q_mix genuinely mixed by construction. Unbiased either way (balance heuristic), so it costs at most draws, never correctness.") +integration_params.add_option("--portfolio-weight-clip",default=None,type=float,help="Portfolio: OPT-IN truncated importance sampling applied to the PROPOSAL-FIT INPUT ONLY. Caps the weights fed to member.update_sampling_prior (the GMM covariance fit) at tau = C*sqrt(n)*mean(w) (0/unset = off; C~1 is the standard Ionides choice), so one enormous weight cannot make that fit degenerate. The estimator (ln Z, n_eff), the n_ess report, and the allocation signal all use the TRUE unclipped weights, so they stay exactly unbiased and undistorted. Do NOT clip the estimator (measured on S250114ax: n_eff=100 2x faster than AV but ln Z biased -11.5 nats) or the n_ess report (clipping inflates the clipped member's n_ess and starves the AV workhorse). The withheld tail mass is tracked and reported as a diagnostic. NOTE: if huge weights come from q_mix UNDERFLOW (watch for the warning) they are a numerical artifact, not tail mass.") +integration_params.add_option("--portfolio-quality-signal",default=None,type=str,help="Portfolio adaptive allocation: which per-member quality signal to rank members by. 'global' (default) = marginal gain in POOLED n_eff per sample (credits weight mass, debits weight variance); 'credit' = q_mix-native MIS credit assignment, sum_i [frac_m q_m/q_mix]_i * w_i per drawn sample (credits a member for COVERING where the integrand is, even if it drew few samples there); 'ness' = legacy per-member Kish n_ess (scale-invariant, misranks a slow-contracting AV -- see DESIGN_portfolio_freeze_policy.md).") +integration_params.add_option("--portfolio-alloc-exponent",default=None,type=float,help="Portfolio: adaptive allocation ~ member_quality^exponent. Higher concentrates harder on the winner. Sampler default 1.0.") +integration_params.add_option("--portfolio-probe-period",default=None,type=int,help="Portfolio: round-robin probe one member at a raised draw share every N chunks (breaks the under-observation trap). 0 disables probing. Sampler default 4.") integration_params.add_option("--sampler-xpy",default=None,help="numpy|cupy if the adaptive_cartesian_gpu sampler is active, use that.") +# Integrator warm-start / reuse (bootstrap AV, persist/reuse a trained NF flow). All +# default off. A warm start only affects the initial PROPOSAL, never the integral. +integration_params.add_option("--sampler-warmstart-samples",default=None,help="AV only: ASCII file (named columns) of prior extrinsic samples used to warm-start the adaptive-volume live region. Intended for the CHERRY-PICKED-PILOT workflow: after iteration 0, run ONE ILE at the best (highest-lnL / CIP-MAP) point with --save-samples (~tens of KB for a single point), then warm-start every subsequent point from it. Do NOT --save-samples the whole grid (disk) and do NOT pick the pilot at random (a poor fit endangers the grid) -- pick the best point. Columns matched to the sampler's extrinsic parameters by name; the coverage-floor + inflation margins below keep a shifted peak from biasing.") +integration_params.add_option("--sampler-warmstart-cover-frac",type=float,default=0.5,help="Coverage floor for --sampler-warmstart-samples (default 0.5): fraction of full-prior coverage mixed in so a mismatched pilot degrades to cold rather than biasing. 0.5 is the MEASURED-safe floor (test_AV_warmstart_safety.py, 20 calibration seeds: max |bias| 0.164, max degradation vs cold 0.113); 0.1 is genuinely under-covered (1.1-1.7 in log bias across seeds) and raising --n-max does not rescue it, because the runs terminate on n_eff first. Lower it only for SAME-problem reuse, where the peak is already in the seed.") +integration_params.add_option("--sampler-warmstart-inflate",type=float,default=1.5,help="Handoff safety margin for --sampler-warmstart-samples (default 1.5): widen the pilot seed about its mean to cover the peak shift between the pilot point and this one.") +integration_params.add_option("--sampler-load-state",default=None,help="AV only: load a saved live-volume state (.npz from --sampler-save-state) to warm-start this integration. Overrides --sampler-warmstart-samples.") +integration_params.add_option("--sampler-save-state",default=None,help="AV only: after integration, write the adapted live-volume state (.npz) for reuse by later instances/iterations. Point --sampler-load-state at the same file across a grid to warm-start each point from the previous one.") +integration_params.add_option("--nf-flow-load",default=None,help="NF only: load a pre-trained normalizing flow (.pt from --nf-flow-save); with --n-adapt 0 this reuses it directly (skips training), otherwise it is polished.") +integration_params.add_option("--nf-flow-save",default=None,help="NF only: after integration, serialize the trained normalizing flow (.pt) for reuse across ILE instances.") +integration_params.add_option("--sampler-sequential-warmstart",action='store_true',help="AV only: when a worker analyzes several intrinsic points (--n-events-to-analyze>1), warm-start each point's extrinsic integral from the previous point's converged high-likelihood samples. Points are processed in their given order (NOT reordered), so a truncated/failed worker still drops a spatially-unbiased subset. A coverage floor (see --sampler-sequential-warmstart-cover-frac) keeps a poorly-matched transfer from ever biasing the result.") +integration_params.add_option("--sampler-sequential-warmstart-cover-frac",type=float,default=0.5,help="Coverage floor for --sampler-sequential-warmstart: fraction of full-prior coverage mixed into the seed so the warm live volume always contains a cold start (a mis-matched proposal then only costs efficiency, never bias). Default 0.5, the measured-safe floor (see --sampler-warmstart-cover-frac); 0.1 is under-covered.") +integration_params.add_option("--sampler-sequential-warmstart-deltalnL",type=float,default=15.0,help="Keep previous-point samples within this lnL of the max as the warm seed for the next point. Default 15.") +integration_params.add_option("--sampler-l0-rescue-accept-truncated", action='store_true', default=False, help="Report the L0 rescue's warm pass even when it lands well below the full-support cold pass (see --sampler-l0-rescue-reject-dlnZ). Default OFF: on that evidence the cold result is kept instead, since the warm pass is confined to the seeded peak and may be missing a mode. The rescue itself still runs either way.") +integration_params.add_option("--sampler-l0-rescue-reject-dlnZ", type=float, default=0.5, help="Evidence threshold (nats) for rejecting the L0 rescue's warm pass: reject when the full-support cold pass reports lnZ this much HIGHER, which indicates the seed missed mass. Larger = more permissive.") +integration_params.add_option("--sampler-warmstart-retry-neff",type=float,default=None,help="AV or portfolio (L0 auto-rescue): if a pass finishes below this n_eff (i.e. it stalled on a very sharp / high-amplitude peak), automatically re-run a second pass warm-started from THIS point's own highest-likelihood samples. Same-problem reuse in the sense that the seed provably contains the peak the cold pass found -- but NOT that every mode is represented, so the warm pass can be biased low if the seed missed one. The rescue still runs as before; its result is rejected in favour of the cold pass only on positive evidence of lost mass (see --sampler-l0-rescue-reject-dlnZ). A portfolio is unaffected: its GMM member carries a defensive component. Directly targets the high-SNR n_eff LOTTERY (a large fraction of independent runs collapse to n_eff~1 by contracting onto the wrong spot); the rescue re-seeds a collapsed run from the peak it did find. Recommended for high-SNR events; e.g. 5.") +integration_params.add_option("--sampler-anisotropic-bins",action="store_true",help="AV only: give each extrinsic axis a DIFFERENT number of bins during contraction -- fine where the live points cluster tightly (phase/polarization/sky), coarse where they are broad (distance/inclination) -- instead of the default equal split. Keeps the same total bin budget, so the estimator is unchanged; helps AV wrap a correlated/degenerate posterior more tightly.") +integration_params.add_option("--internal-reparam-dl-incl",action="store_true",help="Sample the DISTANCE axis as an effective distance D_eff = d_L / A(iota), with A(iota)=sqrt(((1+cos^2 i)/2)^2 + cos^2 i) the leading (l=|m|=2) inclination amplitude. This axis-aligns the distance<->inclination degeneracy (L depends mostly on A(iota)/d_L), decorrelating the two broad directions so the sampler wraps them efficiently. The likelihood reconstructs physical d_L=D_eff*A(iota); the measure correction is PRIOR-AGNOSTIC -- ln p(d_L) - ln p(D_eff) + ln A(iota), using the ACTUAL --d-prior (dist_prior_pdf), so it is correct for Euclidean, cosmo, cosmo_sourceframe, pseudo_cosmo alike (normalization cancels in the ratio; reduces to +3 ln A only for Euclidean). The physical d_L bound is enforced. NOT compatible with --d-prior-redshift (errors out). Estimator stays unbiased (validate vs baseline posterior).") +integration_params.add_option("--extrinsic-proposal-field",default=None,help="AV only (L3): path to a ProposalField (.npz built by util_BuildProposalField.py from a previous ILE iteration). Each intrinsic point warm-starts its extrinsic integral from the field's nearest entry. Cross-problem reuse, so a coverage floor + an inflation margin are applied (see the two options below); a stale/mismatched field can only cost efficiency, never bias.") +integration_params.add_option("--extrinsic-proposal-field-cover-frac",type=float,default=0.5,help="Coverage floor for --extrinsic-proposal-field handoff (default 0.5, the measured-safe floor -- see --sampler-warmstart-cover-frac; 0.1 is under-covered).") +integration_params.add_option("--extrinsic-proposal-field-inflate",type=float,default=1.5,help="Handoff safety margin for --extrinsic-proposal-field: widen the imported seed by this factor about its mean to cover the peak shift between the neighbouring intrinsic point and this one (default 1.5).") integration_params.add_option("--supplementary-likelihood-factor-code", default=None,type=str,help="Import a module (in your pythonpath!) containing a supplementary factor for the likelihood. Used to impose supplementary external priors of arbitrary complexity and external dependence (e.g., EM observations). EXPERTS-ONLY") integration_params.add_option("--supplementary-likelihood-factor-function", default=None,type=str,help="With above option, specifies the specific function used as an external prior. EXPERTS ONLY") integration_params.add_option("--supplementary-likelihood-factor-ini", default=None,type=str,help="With above option, specifies an ini file that is parsed (here) and passed to the preparation code, called when the module is first loaded, to configure the module. EXPERTS ONLY") @@ -486,6 +534,17 @@ if opts.resample_time_marginalization and not(opts.fairdraw_extrinsic_output): # raise Exception(" Fairdraw not available for this sampler") +def _reparam_A_of_incl(incl_rad, xpy=numpy): + """Leading (l=|m|=2) inclination amplitude A(iota)=sqrt(((1+cos^2 i)/2)^2 + cos^2 i). + Used by --internal-reparam-dl-incl to map the effective distance D_eff <-> physical d_L + (d_L = D_eff * A(iota)). A in [0.5 (edge-on), sqrt(2) (face-on)].""" + ci = xpy.cos(incl_rad) + return xpy.sqrt(((1.0 + ci*ci)/2.0)**2 + ci*ci) + +_REPARAM_A_MIN = 0.5 # A(iota=pi/2), edge-on +_REPARAM_A_MAX = numpy.sqrt(2.0) # A(iota=0), face-on +_REPARAM_LNF = 0.0 # ln(physical-range prior mass fraction); set at setup for --internal-reparam-dl-incl + supplemental_ln_likelihood= None supplemental_ln_likelhood_prep=None supplemental_ln_likelhood_parsed_ini=None @@ -1081,6 +1140,19 @@ 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) +if opts.internal_reparam_dl_incl: + # The reparam operates in DISTANCE units (D_eff = d_L/A in Mpc). --d-prior-redshift samples + # in redshift and converts, which does not compose with the amplitude relation -- refuse it + # rather than silently bias. All distance-space priors (Euclidean/cosmo/cosmo_sourceframe/ + # pseudo_cosmo) ARE supported: the measure correction below uses dist_prior_pdf directly. + if getattr(opts, 'd_prior_redshift', False): + raise SystemExit(" --internal-reparam-dl-incl is not compatible with --d-prior-redshift (redshift-space sampling). Use a distance-space --d-prior.") + # sample the distance axis as D_eff = d_L / A(iota); widen it so physical d_L = D_eff*A(iota) + # can cover [dmin,dmax] for all iota (A in [0.5, sqrt(2)]). dist_prior_pdf auto-normalizes + # over this range (=> a constant lnZ offset vs baseline, computable; posterior unaffected); + # the physical d_L bound and the prior-agnostic measure term are applied in the likelihood closure. + param_limits['distance'] = (dmin/_REPARAM_A_MAX, dmax/_REPARAM_A_MIN) + print(" [reparam d_L<->D_eff] distance axis is D_eff; sampling range {:.1f}..{:.1f}".format(*param_limits['distance'])) # 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. @@ -1100,6 +1172,7 @@ for _optv, _k in [(opts.limit_psi, 'psi'), (opts.limit_right_ascension, 'right_a # Portfolio use_portfolio=False +use_gmm_member=False # set when a portfolio carries a GMM member (see the portfolio setup loop) params = {} sampler = mcsampler.MCSampler() xpy_asarray_already = functools.partial(xpy_default.asarray,dtype=np.float64) @@ -1140,7 +1213,11 @@ elif opts.sampler_method == "portfolio" and mcsampler_Portfolio_ok: use_portfolio=True opts.internal_use_lnL=True # required, we only implement those scenarios right now sampler_list = [] - sampler_types = opts.sampler_portfolio + # --sampler-portfolio is action='append' AND documented as comma-separated, so honor BOTH: + # flatten the appended list and split every element on ',' (e.g. ['AV,GMM'] -> ['AV','GMM'], + # ['AV','GMM'] -> ['AV','GMM']). Without this, 'AV,GMM' was one bogus member name that matched + # no branch, silently yielding a single-member portfolio. + sampler_types = [s.strip() for item in opts.sampler_portfolio for s in str(item).split(',') if s.strip()] # prep xpy, etc my_xpy = xpy_default @@ -1158,8 +1235,14 @@ elif opts.sampler_method == "portfolio" and mcsampler_Portfolio_ok: sampler = mcsamplerAdaptiveVolume.MCSampler(n_chunk=opts.n_chunk) # enforce now, so provided for setup phase elif name =='GMM': sampler = mcsamplerEnsemble.MCSampler() - # following override means sampler_method is CHANGED, so THIS MUST BE LAST, and can't condition on portfolio - opts.sampler_method = 'GMM' # this will force the creation/parsing of GMM-specific arguments below, so they are properly passed + # A GMM member needs the GMM-specific argument blocks below to run so its config is + # forwarded. Historically this was done by CLOBBERING opts.sampler_method='GMM', which + # silently broke every downstream `sampler_method == "portfolio"` test (the portfolio + # setup block became dead code, and the L0 auto-rescue gate never fired for a portfolio) + # and made a portfolio take GMM-only branches (e.g. return_lnI). Instead flag it + # non-destructively: sampler_method stays 'portfolio', and the GMM blocks below key off + # `use_gmm_args` = standalone GMM OR a portfolio carrying a GMM member. + use_gmm_member = True elif name == "adaptive_cartesian_gpu" or name == 'AC': sampler = mcsamplerGPU.MCSampler() mcsampler = mcsamplerGPU # force use of routines in that file, for properly configured GPU-accelerated code as needed @@ -1320,8 +1403,19 @@ elif not opts.distance_marginalization: elif opts.d_prior != 'Euclidean': print(" ==== WARNING UNKNOWN DISTANCE PRIOR === ") raise Exception('distance prior') + if opts.internal_reparam_dl_incl: + # dist_prior_pdf is normalized over the WIDENED D_eff range; the physical prior must be + # normalized over [dmin,dmax]. ln F = ln(mass of dist_prior_pdf in [dmin,dmax]); the closure + # subtracts it so the reported lnZ matches the baseline physical-range normalization exactly. + _xg = numpy.linspace(dmin, dmax, 40000) + try: + _pg = numpy.asarray(dist_prior_pdf(_xg), dtype=float) + except Exception: + _pg = numpy.array([float(dist_prior_pdf(numpy.array([_x]))[0]) for _x in _xg]) + _REPARAM_LNF = float(numpy.log(numpy.trapz(_pg, _xg))) + print(" [reparam] physical-range prior-mass fraction F={:.4f} (lnF={:.3f}); lnZ normalization matched".format(numpy.exp(_REPARAM_LNF), _REPARAM_LNF)) #dist_sampler_cdf_inv=None - sampler.add_parameter("distance", + sampler.add_parameter("distance", pdf = dist_sampler, cdf_inv = dist_sampler_cdf_inv, left_limit = param_limits["distance"][0], @@ -1544,9 +1638,16 @@ pinned_params.update({ "igrand_fairdraw_samples_max": np.min([opts.fairdraw_extrinsic_output_n_max,opts.n_eff]) }) if opts.sampler_method == "adaptive_cartesian_gpu": - pinned_params.update({"save_no_samples":True}) # do not exhaust GPU memory with MC samples! + pinned_params.update({"save_no_samples":True}) # do not exhaust GPU memory with MC samples! +# GMM-specific argument blocks must run for a STANDALONE GMM *or* for a portfolio carrying a GMM +# member (whose config still has to be forwarded). This used to be achieved by clobbering +# opts.sampler_method='GMM' during portfolio setup, which broke every downstream 'portfolio' test; +# key off this explicit flag instead so sampler_method keeps meaning what the user asked for. +use_gmm_args = (opts.sampler_method == "GMM") or use_gmm_member return_lnL=False -if opts.sampler_method=="GMM" and opts.internal_use_lnL: +if (opts.sampler_method=="GMM") and opts.internal_use_lnL: + # standalone GMM only: return_lnI is an mcsamplerEnsemble kwarg; the portfolio does not consume it + # (the portfolio's own use_lnL wiring is in the portfolio block below). return_lnL=True pinned_params.update({"use_lnL":True,"return_lnI":True}) if opts.sampler_method =="adaptive_cartesian_gpu" and opts.internal_use_lnL: @@ -1561,7 +1662,18 @@ if opts.sampler_method =="AV" and opts.internal_use_lnL: if opts.sampler_method =="portfolio": return_lnL=True pinned_params.update({"use_lnL":True}) -if opts.sampler_method == "GMM": + # FLEXIBLE allocation for the portfolio's GMM member. NOTE: when the portfolio HAS a GMM member + # (use_gmm_member), the GMM block below already forwards gmm_adaptive as a per-group DICT (keyed by + # the member's actual parameter groups) via extra_args -- which is richer than the scalar cap here, + # and is the path every portfolio benchmark on this branch actually exercised (it ran because + # sampler_method used to be clobbered to 'GMM'). Only fall back to the scalar form for a portfolio + # WITHOUT a GMM member, where that block does not run. Setting both would double-specify it. + if opts.internal_gmm_adaptive_components and not use_gmm_member: + pinned_params.update({'gmm_adaptive': int(opts.internal_gmm_max_components), + 'gmm_defensive_frac': float(opts.internal_gmm_defensive_frac), + 'gmm_inflate': float(opts.internal_gmm_inflate)}) + print(" Portfolio: GMM member adaptive components enabled (BIC, cap {})".format(opts.internal_gmm_max_components)) +if use_gmm_args: # standalone GMM, or a portfolio carrying a GMM member (see use_gmm_args above) n_step =pinned_params["n"] n_max_blocks = ((1.0*int(opts.n_max))/n_step) # pairing coordinates for adaptive integration: see definition of order below @@ -1639,6 +1751,21 @@ if opts.sampler_method == "GMM": else: comp_dict = {pair_ra_dec:n_sky,pair_d_incl:n_d,pair_phi_psi:n_phase} extra_args = {'n_comp':comp_dict,'max_iter':n_max_blocks,'gmm_dict':gmm_dict, 'gmm_adapt':gmm_adapt} # made up for now, should adjust + # FLEXIBLE allocation: each ADAPTING group chooses its component count from + # the data by BIC (GMM.fit_gmm_adaptive), replacing the hard-coded per-group + # counts above. A defensive tail component + optional inflation keep the + # importance weights bounded (the fix that gets GMM n_eff off ~1 at high SNR). + if opts.internal_gmm_adaptive_components: + k_cap = int(opts.internal_gmm_max_components) + gmm_adaptive = {} + for _g in gmm_dict: + if (gmm_adapt is None) or gmm_adapt.get(_g, True): # only groups that adapt + gmm_adaptive[_g] = k_cap + extra_args['gmm_adaptive'] = gmm_adaptive + extra_args['gmm_defensive_frac'] = float(opts.internal_gmm_defensive_frac) + extra_args['gmm_inflate'] = float(opts.internal_gmm_inflate) + print("GMM adaptive components (BIC, cap {}, defensive {}, inflate {}): groups {}".format( + k_cap, opts.internal_gmm_defensive_frac, opts.internal_gmm_inflate, list(gmm_adaptive.keys()))) print("GMM:",extra_args) print("GMM:",sampler.params_ordered) # if opts.distance_marginalization: @@ -1685,7 +1812,44 @@ if use_portfolio: if not(isinstance(opts.sampler_portfolio_args[indx], dict)): print(indx,opts.sampler_portfolio_args[indx]) print(" ARGS ", opts.sampler_portfolio_args) - sampler.setup(portfolio_args=opts.sampler_portfolio_args, **pinned_params) # directly pass all parameters set above to low-level portfolios. In particular, GMM setup + # Assemble freeze-policy overrides from the CLI. Only include options the user actually + # set (None = unset) so the sampler keeps its built-in defaults otherwise. The two VARAHA + # flags are mutually exclusive; --portfolio-varaha-can-freeze wins if both are given. + _freeze_policy_kwargs = {} + if opts.portfolio_grace_iters is not None: + _freeze_policy_kwargs['portfolio_grace_iters'] = opts.portfolio_grace_iters + if opts.portfolio_revive_period is not None: + _freeze_policy_kwargs['portfolio_revive_period'] = opts.portfolio_revive_period + if opts.portfolio_freeze_wt is not None: + _freeze_policy_kwargs['portfolio_freeze_wt'] = opts.portfolio_freeze_wt + if opts.portfolio_varaha_can_freeze: + _freeze_policy_kwargs['portfolio_varaha_never_freeze'] = False + elif opts.portfolio_varaha_never_freeze: + _freeze_policy_kwargs['portfolio_varaha_never_freeze'] = True + # adaptive-probe draw allocation (OPT-IN; off by default in the sampler) + if opts.portfolio_adaptive_alloc: + _freeze_policy_kwargs['portfolio_adaptive_alloc'] = True + if opts.portfolio_varaha_min_frac is not None: + _freeze_policy_kwargs['portfolio_varaha_min_frac'] = opts.portfolio_varaha_min_frac + if opts.portfolio_varaha_max_frac is not None: + _freeze_policy_kwargs['portfolio_varaha_max_frac'] = opts.portfolio_varaha_max_frac + if opts.portfolio_weight_clip is not None: + _freeze_policy_kwargs['portfolio_weight_clip'] = opts.portfolio_weight_clip + if opts.portfolio_quality_signal is not None: + _freeze_policy_kwargs['portfolio_quality_signal'] = opts.portfolio_quality_signal + if opts.portfolio_alloc_exponent is not None: + _freeze_policy_kwargs['portfolio_alloc_exponent'] = opts.portfolio_alloc_exponent + if opts.portfolio_probe_period is not None: + _freeze_policy_kwargs['portfolio_probe_period'] = opts.portfolio_probe_period + print(" PORTFOLIO freeze-policy overrides: ", _freeze_policy_kwargs) + sampler.setup(portfolio_args=opts.sampler_portfolio_args, **_freeze_policy_kwargs, **pinned_params) # directly pass all parameters set above to low-level portfolios. In particular, GMM setup + # NOTE: the portfolio oracle MECHANISM is fixed (proposals now actually enter + # member training; see mcsamplerPortfolio), and a FisherGaussianOracle can be + # attached via sampler.oracle_realizations. It is not auto-wired here because + # (a) ILE has no natural Fisher source for the full extrinsic space, and + # (b) MCSampler.setup() re-runs each oracle's setup(), which would reset a + # pre-configured skymap ResamplingOracle. Wire deliberately when a proposal + # source is available. # initialize sampler, before we call integrate, so we can seed it if opts.sampler_method == 'adaptive_cartesian_gpu' and opts.skymap_file: @@ -1739,11 +1903,25 @@ def resample_samples(my_samples, lnL_out = np.zeros(n_samples) # IF UPSAMPLING, PERFORM NOW. (Currently on if opts.srate_resample_time_marginalization and opts.srate_resample_time_marginalization > fSample: - deltaT_orig = tvals[1]-tvals[0] - tvals_denser = tvals[0] + deltaT_orig/2 * np.arange(2*len(tvals)) + # Resample the marginalization-time grid to EXACTLY the requested rate, so + # the exported geocenter time is quantized at 1/srate_resample seconds. We + # step by exactly 1/srate_resample rather than by an integer subdivision of + # the internal grid: that internal grid is a closed-interval linspace whose + # spacing is ~1/fSample but NOT exactly (here ~4086.7 Hz vs 4096), so an + # integer-factor upsample would land at ~n/deltaT_orig, tens of percent off + # the request. For the usual power-of-two rates 1/srate is exactly + # representable in float64, so consecutive output times differ by exactly + # that step. + dt_target = 1.0/opts.srate_resample_time_marginalization + # floor(): stay within [tvals[0], tvals[-1]] so the spline never + # extrapolates. At most one step (<1/srate s, tens of us) is dropped at the + # far edge of the +-75 ms window, where the time-marginalized likelihood is + # negligible. + n_dense = int(np.floor((tvals[-1]-tvals[0])/dt_target)) + 1 + tvals_denser = tvals[0] + dt_target * np.arange(n_dense) from scipy.interpolate import RegularGridInterpolator, CubicSpline # cubic spline at first, easiest - generally not exporting too many events - lnLt_new = np.zeros( (lnLt.shape[0], lnLt.shape[1]*2) ) + lnLt_new = np.zeros( (lnLt.shape[0], n_dense) ) for indx_here in np.arange(n_samples): cs = CubicSpline(tvals, lnLt[indx_here]) lnLt_new[indx_here] = cs(tvals_denser) @@ -1764,6 +1942,211 @@ def resample_samples(my_samples, +def _rvs_len(rvs): + for v in rvs.values(): + try: + return len(numpy.atleast_1d(numpy.asarray(v)).ravel()) + except Exception: + continue + return 0 + + +def _pool_replica_rvs(rep_rvs, sampler, rep_lnZ=None, already_resampled=False): + """Concatenate the replicas' samples into one correctly-weighted set. + + Each replica k is an independent importance-sampling estimate with weights w_ki and its own + sample count n_k, and the reported evidence is the linear mean (1/K) sum_k Z_k. The posterior + that matches THAT estimator is the concatenation with weights w_ki/(K n_k) -- equivalently the + importance weight against the pooled proposal q'_ki = q_ki * K * n_k, which is the actual + density of "pick a replica uniformly, then one of its n_k draws". So the K*n_k factor goes + into the sampling prior, where every downstream weight computation already accounts for it. + + Falls back to the first replica if the record shape is unexpected: a degraded export is + recoverable, a silently mis-weighted one is not. + """ + rep_rvs = [r for r in rep_rvs if r] + if len(rep_rvs) <= 1: + return rep_rvs[0] if rep_rvs else {} + # `already_resampled` -- the records are FAIRDRAW output. Those samples were already drawn in + # proportion to their own importance weights, so reusing those weights applies them a second + # time and the pooled block follows w^2 instead of w. Renormalizing to Z_k/K fixes the block's + # SCALE but not its SHAPE, so it does not help here. A fairdraw block is an equal-weight draw + # from its own posterior, so that is what it must contribute: constant weights within the + # block, summing to Z_k/K. + # + # DO NOT assume the records are raw importance samples. integrate() may have thresholded or + # fairdraw-resampled _rvs before we see it, in which case sum_i w_ki over the RETAINED rows is + # no longer Z_k * n_k and a 1/n_k rescale would mis-weight the replica (a fairdraw record is + # already posterior-resampled, so scaling it by its retained length weights it twice). When + # the reported per-replica lnZ is available, renormalize each block so it contributes exactly + # Z_k/K -- correct whether the rows are raw, pruned or resampled, since only their RELATIVE + # weights need be right. + keys = set(rep_rvs[0]) + for r in rep_rvs[1:]: + keys &= set(r) + log_key = 'log_joint_s_prior' if 'log_joint_s_prior' in keys else None + lin_key = 'joint_s_prior' if (log_key is None and 'joint_s_prior' in keys) else None + if log_key is None and lin_key is None: + print(" [mc error] pooling skipped: no sampling-prior column in the replica records; " + "exporting the FIRST replica (consistent weights, fewer samples)") + return rep_rvs[0] + K = len(rep_rvs) + out = {} + try: + cols = {k: [] for k in keys} + for _i, r in enumerate(rep_rvs): + n_k = _rvs_len(r) + if n_k <= 0: + continue + _flat_block = False + if already_resampled and rep_lnZ is not None and _i < len(rep_lnZ) \ + and numpy.isfinite(rep_lnZ[_i]): + # equal weights within the block, summing to Z_k/K + _flat_block = True + _target_lw = float(rep_lnZ[_i]) - numpy.log(float(K)) - numpy.log(float(n_k)) + scale = 0.0 + elif rep_lnZ is not None and _i < len(rep_lnZ) and numpy.isfinite(rep_lnZ[_i]): + # target: this block's weights sum to Z_k/K + _cur = _lnZ_of_rvs(r, already_pooled=True) + if _cur is None or not numpy.isfinite(_cur): + scale = numpy.log(float(K) * float(n_k)) + else: + scale = _cur - (float(rep_lnZ[_i]) - numpy.log(float(K))) + else: + scale = numpy.log(float(K) * float(n_k)) + if _flat_block and log_key is not None: + # force lw_i = log_integrand + log_joint_prior - log_joint_s_prior == _target_lw + _li = numpy.atleast_1d(numpy.asarray( + sampler.identity_convert(r['log_integrand']), dtype=float)).ravel() + _lp = numpy.atleast_1d(numpy.asarray( + sampler.identity_convert(r['log_joint_prior']), dtype=float)).ravel() + _forced = _li + _lp - _target_lw + for k in keys: + v = numpy.atleast_1d(numpy.asarray(sampler.identity_convert(r[k]))).ravel() + if _flat_block and log_key is not None and k == log_key: + v = _forced + elif _flat_block and lin_key is not None and k == lin_key: + _ig = numpy.atleast_1d(numpy.asarray( + sampler.identity_convert(r['integrand']), dtype=float)).ravel() + _jp = numpy.atleast_1d(numpy.asarray( + sampler.identity_convert(r['joint_prior']), dtype=float)).ravel() + v = _ig * _jp / numpy.exp(_target_lw) + elif k == log_key: + v = v + scale + elif k == lin_key: + v = v * (float(K) * float(n_k)) + cols[k].append(v) + for k in keys: + out[k] = numpy.concatenate(cols[k]) if cols[k] else numpy.array([]) + # CACHED WEIGHTS MUST FOLLOW THE COMPONENTS. _rvs may carry a precomputed 'log_weights' + # (mcsamplerPortfolio writes one), and the .dgrid and calibration-posterior exporters + # PREFER it -- they only fall back to log_integrand + log_joint_prior - log_joint_s_prior + # when it is absent. Concatenating the per-replica caches unchanged would hand those + # scientific outputs the ORIGINAL weights while the estimate used the corrected ones: + # replica rebalancing ignored, and fairdraw blocks double-weighted again in exactly the + # products this pooling exists to make consistent. Recompute from the canonical columns. + _lw_pooled = None + if all(k in out for k in ('log_integrand', 'log_joint_prior', 'log_joint_s_prior')): + _lw_pooled = (numpy.asarray(out['log_integrand'], dtype=float) + + numpy.asarray(out['log_joint_prior'], dtype=float) + - numpy.asarray(out['log_joint_s_prior'], dtype=float)) + elif all(k in out for k in ('integrand', 'joint_prior', 'joint_s_prior')): + _ig = numpy.asarray(out['integrand'], dtype=float) + _jp = numpy.asarray(out['joint_prior'], dtype=float) + _js = numpy.asarray(out['joint_s_prior'], dtype=float) + _ok = (_ig > 0) & (_jp > 0) & (_js > 0) + _lw_pooled = numpy.full(len(_ig), -numpy.inf) + _lw_pooled[_ok] = numpy.log(_ig[_ok]) + numpy.log(_jp[_ok]) - numpy.log(_js[_ok]) + if _lw_pooled is not None: + if 'log_weights' in out: + out['log_weights'] = _lw_pooled + if 'weights' in out: + out['weights'] = numpy.exp(_lw_pooled - numpy.max(_lw_pooled[numpy.isfinite(_lw_pooled)])) + elif 'log_weights' in out or 'weights' in out: + # cannot rebuild them -> DROP, so consumers fall through to whatever components exist + # rather than silently trusting a stale cache. + out.pop('log_weights', None) + out.pop('weights', None) + print(" [mc error] pooled record: dropped stale cached weights (components unavailable" + " to rebuild them); consumers will reconstruct from what remains") + except Exception as e: + print(" [mc error] pooling failed ({}); exporting the FIRST replica".format(e)) + return rep_rvs[0] + return out + + +def _lnZ_of_rvs(rvs, already_pooled=True): + """log of the evidence implied by an _rvs record. + + For a POOLED record the weights already carry their 1/(K n_k) factor, so the estimate is the + plain sum; for a single run it is the mean. Returns None when the weights cannot be rebuilt. + """ + try: + if 'log_integrand' in rvs and 'log_joint_prior' in rvs and 'log_joint_s_prior' in rvs: + lw = numpy.asarray(rvs['log_integrand'], dtype=float) \ + + numpy.asarray(rvs['log_joint_prior'], dtype=float) \ + - numpy.asarray(rvs['log_joint_s_prior'], dtype=float) + elif 'integrand' in rvs and 'joint_prior' in rvs and 'joint_s_prior' in rvs: + w = (numpy.asarray(rvs['integrand'], dtype=float) + * numpy.asarray(rvs['joint_prior'], dtype=float) + / numpy.asarray(rvs['joint_s_prior'], dtype=float)) + lw = numpy.log(numpy.where(w > 0, w, numpy.nan)) + else: + return None + lw = lw[numpy.isfinite(lw)] + if lw.size == 0: + return None + m = numpy.max(lw) + tot = m + numpy.log(numpy.sum(numpy.exp(lw - m))) + return float(tot if already_pooled else tot - numpy.log(lw.size)) + except Exception: + return None + + +def _kish_neff_of_rvs(rvs): + """Kish effective sample size of an _rvs record, or None if the weights are not reconstructible.""" + try: + if 'log_integrand' in rvs and 'log_joint_prior' in rvs and 'log_joint_s_prior' in rvs: + lw = numpy.asarray(rvs['log_integrand'], dtype=float) \ + + numpy.asarray(rvs['log_joint_prior'], dtype=float) \ + - numpy.asarray(rvs['log_joint_s_prior'], dtype=float) + elif 'integrand' in rvs and 'joint_prior' in rvs and 'joint_s_prior' in rvs: + w = (numpy.asarray(rvs['integrand'], dtype=float) + * numpy.asarray(rvs['joint_prior'], dtype=float) + / numpy.asarray(rvs['joint_s_prior'], dtype=float)) + lw = numpy.log(numpy.where(w > 0, w, numpy.nan)) + else: + return None + lw = lw[numpy.isfinite(lw)] + if lw.size == 0: + return None + lw = lw - numpy.max(lw) + w = numpy.exp(lw) + return float(numpy.sum(w) ** 2 / numpy.sum(w ** 2)) + except Exception: + return None + + +def _clear_warm_state(sampler): + """Clear a warm-start seed AND any grid it installed, reaching PORTFOLIO MEMBERS too. + + `sampler._warm = None` alone is not enough for mcsamplerPortfolio: `_warm` and the contracted + AV grid live on each MEMBER, and portfolio.integrate_log() does not rerun each member's setup(), + so the next point would silently draw from the PREVIOUS point's contracted live volume. If the + new point's support falls outside it, lnZ is biased low with a healthy-looking n_eff and no + error. Portfolio exposes clear_warm_state(); everything else keeps the old behaviour. + """ + # Deliberately NOT wrapped in try/except. A reset that quietly did not happen leaves the next + # point drawing from the previous point's contracted grid -- the exact silent bias this guards + # against -- so a failure must abort the point (the per-point handler below reports it) rather + # than degrade to a log line nobody reads. + if hasattr(sampler, 'clear_warm_state'): + sampler.clear_warm_state() + else: + sampler._warm = None + sampler._warm_applied = False + + def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_trunc_Q=inv_spec_trunc_Q, T_spec=T_spec): nEvals=0 P = P_list[indx_event] @@ -2244,6 +2627,9 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t distance = redshift_to_distance(distance) P.psi = xpy_default.asarray(psi, dtype=np.float64) + if opts.internal_reparam_dl_incl: + # distance axis holds D_eff; reconstruct physical d_L = D_eff * A(iota) + distance = distance * _reparam_A_of_incl(P.incl, xpy=xpy_default) P.dist = xpy_default.asarray(distance* 1.e6 * lalsimutils.lsu_PC, dtype=np.float64) # luminosity distance # rotate sky if needed @@ -2280,6 +2666,19 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t cal_method=('fused' if use_fused_calmarg and opts._noloop_time_interp == 'nearest' else 'loop'), cal_log_weights=calibration_log_weights, time_interp=opts._noloop_time_interp) # non-distmarg: default-helper fused kernel (cal_distmarg=None) # nEvals +=len(right_ascension) + if opts.internal_reparam_dl_incl: + # PRIOR-AGNOSTIC measure for the D_eff<->d_L reparam: the integrand needs + # p_prior(d_L) * |dd_L/dD_eff| but the sampler applied p_prior(D_eff), + # so add ln p(d_L) - ln p(D_eff) + ln A , using the ACTUAL --d-prior + # (dist_prior_pdf). Its normalization cancels in the ratio, so this is correct + # for Euclidean / cosmo / cosmo_sourceframe / pseudo_cosmo (reduces to +3 ln A + # only for Euclidean). Then enforce the physical d_L in [dmin,dmax]. + _A = _reparam_A_of_incl(P.incl, xpy=xpy_default) + _dl = distance # physical d_L (reconstructed above) + _deff = _dl / _A # the sampled D_eff + # physical prior normalized over [dmin,dmax] (=> - _REPARAM_LNF); sampler prior at D_eff; Jacobian A + lnL = lnL + (xpy_default.log(dist_prior_pdf(_dl)) - _REPARAM_LNF) - xpy_default.log(dist_prior_pdf(_deff)) + xpy_default.log(_A) + lnL = xpy_default.where((_dl >= dmin) & (_dl <= dmax), lnL, -1e300) if supplemental_ln_likelihood: lnL += supplemental_ln_likelihood(P.phi, P.theta, P.phiref ,P.incl, P.psi, P.dist,xpy=xpy_default) # use these variables so they are already float-type if return_lnL: @@ -2503,7 +2902,7 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t if 'distance' in sampler.params: sampler.reset_sampling('distance') sampler.reset_sampling('inclination') - elif opts.sampler_method == "GMM": + elif use_gmm_args: # standalone GMM or a portfolio with a GMM member (gmm_dict exists in both) if 'distance' in sampler.params: pair_d_incl = sampler_param_tuple(sampler, ['distance','inclination']) if pair_d_incl in gmm_dict: @@ -2534,6 +2933,44 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t lnL_oracles = np.zeros(opts.n_chunk) sampler.update_sampling_prior(lnL_oracles, opts.n_chunk, external_rvs=rvs_train,log_scale_weights=True,floor_integrated_probability=opts.adapt_floor_level) + # Warm-start / bootstrap. The AV (VARAHA) sampler has no update_sampling_prior, so the + # skymap-oracle seeding above is skipped for it; instead seed its live volume directly + # (bootstrap_from_samples). A PORTFOLIO forwards the seed to its warm-startable members + # (e.g. its AV member) via the same method; members without it stay cold, and the + # balance-heuristic mixture density (q_mix) keeps a cold/mis-seeded member from biasing. + # A seed only shapes the initial proposal, never the integral, so this cannot bias. + # Fires for any sampler exposing bootstrap_from_samples (AV directly, or a portfolio). + if hasattr(sampler, 'bootstrap_from_samples'): + try: + if opts.sampler_load_state and hasattr(sampler, 'load_state'): + print(" warm-start: loading saved sampler state from", opts.sampler_load_state) + sampler.load_state(opts.sampler_load_state) + elif opts.sampler_warmstart_samples: + _dat = np.genfromtxt(opts.sampler_warmstart_samples, names=True) + _cols = np.vstack([np.asarray(_dat[p], dtype=float) for p in sampler.params_ordered]).T + print(" warm-start: bootstrapping from", opts.sampler_warmstart_samples, _cols.shape, + "(cover_frac={}, inflate={})".format(opts.sampler_warmstart_cover_frac, opts.sampler_warmstart_inflate)) + # cover_frac + inflate are the handoff safety margins: this seed usually + # comes from a DIFFERENT (cherry-picked pilot) point, so it must not be + # able to bias if the peak has shifted. Default them >0 in this path. + sampler.bootstrap_from_samples(_cols, + cover_frac=opts.sampler_warmstart_cover_frac, + inflate=opts.sampler_warmstart_inflate) + elif oracleRS: + _, _, _rv_oracle = oracleRS.draw_simplified(opts.n_chunk) + print(" AV warm-start: bootstrapping live volume from skymap oracle") + sampler.bootstrap_from_samples(_rv_oracle, params=list(sampler.params_ordered)) + except Exception as _e_ws: + print(" AV warm-start skipped (", _e_ws, ")") + + # NF flow reuse: warm-load a pre-trained flow so this instance skips/shortens training. + if opts.nf_flow_load and hasattr(sampler, 'load_flow'): + try: + print(" NF: loading pre-trained flow from", opts.nf_flow_load) + sampler.load_flow(opts.nf_flow_load) + except Exception as _e_nf: + print(" NF flow load skipped (", _e_nf, ")") + # Optional ZERO-CAL burn-in (generally useful; see RIFT/calmarg/DESIGN_adaptive_driver.md). # Adapt the extrinsic sampler cheaply on the n_cal=1 baseline first, then run the full # cal-marginalized integration reusing the adapted proposal. The likelihood closures @@ -2555,8 +2992,131 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t print(" [calmarg burn-in] integrate failed ({}); proceeding to production".format(_eb)) n_cal_for_likelihood = _ncal_full # restore: production uses the full cal set + # opt-in anisotropic per-axis bin allocation: set on the AV sampler and any AV portfolio members + if getattr(opts, 'sampler_anisotropic_bins', False): + _aniso_targets = [sampler] + list(getattr(sampler, 'portfolio_realizations', [])) + for _t in _aniso_targets: + if hasattr(_t, 'anisotropic_bins'): + _t.anisotropic_bins = True + print(" AV: anisotropic per-axis bin allocation ENABLED") + res, var, neff, dict_return = sampler.integrate(like_to_integrate, *unpinned_params, **pinned_params) + # L0 auto-rescue: on a very sharply-peaked (high-amplitude) point a cold AV can + # stall at n_eff ~ 1 because it never draws near the tiny peak. If so, seed a + # SECOND pass from this same point's own highest-likelihood samples and re-run. + # This is SAME-problem reuse (the peak provably lies in the seed, since the cold + # pass found it), so no coverage floor is needed (cover_frac=0) and it cannot + # bias the result. AV or a portfolio carrying an AV member; opt-in via + # --sampler-warmstart-retry-neff. For a portfolio the peak-seed is bootstrapped into its + # warm-startable members (the AV live volume) and the whole mixture is re-run; the seed comes + # from the run's OWN _rvs (already in the sampling frame), so it is coordinate-safe by construction. + # A DEGENERATE EARLY TERMINATION (neff is None) is the strongest possible rescue trigger, not a + # reason to skip: mcsamplerPortfolio/AV return (None,None,None,None) from their "terminate early" + # branch when the live volume never finds finite in-volume samples -- i.e. exactly the cold, very + # sharp peak this rescue exists for. Such a pass still populates _rvs (it DID sample the peak, + # it just could not build a volume around it), so the peak-seed below is available. Treat + # neff=None as "below threshold". + _neff_val = None if neff is None else float(sampler.identity_convert(neff)) + _needs_l0_rescue = (_neff_val is None) or (_neff_val < float(opts.sampler_warmstart_retry_neff or 0)) + if (opts.sampler_method in ('AV', 'portfolio') and opts.sampler_warmstart_retry_neff + and hasattr(sampler, 'bootstrap_from_samples') + and _needs_l0_rescue): + try: + _lnkey = 'log_integrand' if 'log_integrand' in sampler._rvs else ('integrand' if 'integrand' in sampler._rvs else None) + _lnv = np.asarray(sampler.identity_convert(sampler._rvs[_lnkey]), dtype=float).ravel() if _lnkey else np.array([]) + if _lnv.size >= 1 and np.any(np.isfinite(_lnv)): + _cols = np.vstack([np.asarray(sampler.identity_convert(sampler._rvs[p]), dtype=float).ravel() for p in sampler.params_ordered]).T + _kthr = np.nanmax(_lnv) - opts.sampler_sequential_warmstart_deltalnL + _seed = _cols[_lnv > _kthr] + if len(_seed) < 2: # peak found with too few points -> puff the single best point + _best = _cols[int(np.nanargmax(_lnv))] + _wid = (np.array([sampler.rlim[p] for p in sampler.params_ordered]) - np.array([sampler.llim[p] for p in sampler.params_ordered])) / 200.0 + _seed = np.random.RandomState(0).normal(_best, _wid, size=(2000, len(_best))) + print(" [L0 auto-rescue] cold n_eff {} < {}; re-running warm from this point's peak ({} pts)".format( + "DEGENERATE (early termination)" if _neff_val is None else "{:.1f}".format(_neff_val), + opts.sampler_warmstart_retry_neff, len(_seed))) + # The warm pass is an estimate over TRUNCATED support: the seeded box provably + # contains the peak the cold pass found, and says nothing about what that pass did + # not reach, so it is biased low by any missed mode. Three things that look like + # fixes are not: + # * cover_frac is a finite sprinkle of points into the grid (2.9% of the box at + # d=6), not a mixture, so it gives no coverage guarantee; + # * pooling cold+warm propagates the bias in diluted form, because averaging Z is + # unbiased only when EVERY term is. Measured on a bimodal target whose seed + # caught one mode: true -4.6052, cold -4.6265, warm -5.3009 (-log 2, the missed + # mode), pooled -4.9079 (log 0.75). Pinned in test_replica_pooling.py; + # * a warning alone does not correct the number that gets reported. + # AND THE OBVIOUS "REAL FIX" IS A TRAP. Giving AV a defensive component with + # support everywhere requires per-sample sampling densities (integrate_log applies + # ONE scalar log_joint_s_prior to every sample), i.e. a proposal that is a weighted + # mixture of components whose densities are evaluated per sample and combined. + # That is mcsamplerPortfolio: q_mix, defensive members, and the coverage + # bookkeeping around them already exist there, tested. Rebuilding it inside AV + # would leave two implementations of the same mathematics to keep in step, and the + # bugs found in this review -- a capability flag that lied, a defensive component + # absorbed by an update, config silently dropped on re-setup -- are precisely the + # kind that appear when one of two parallel paths is updated and the other is not. + # So this is NOT scoped as future work on AV. The architectural answer is that a + # run needing a coverage guarantee uses the portfolio; standalone AV stays a fast + # single-proposal sampler with this limitation documented at its call site. + # + # WHAT THIS DOES, DELIBERATELY NARROWLY. The rescue still runs, because it exists + # to fix the high-SNR n_eff lottery and removing it by default would be a certain + # production regression traded against a possible bias. What changes is only the + # case where we have POSITIVE EVIDENCE of lost mass: the cold pass had full + # support, so if the warm evidence lands well below it, the seed missed something + # the cold pass reached. There we keep the cold result rather than report the + # precise-but-truncated one. Detection is imperfect -- a missed mode need not + # produce this ordering -- so this narrows the failure, it does not close it. + _cold_rvs = sampler._rvs + _cold_lnZ = _lnZ_of_rvs(_cold_rvs, already_pooled=False) + # dict_return too: khat, block scatter, ESS, the confidence interval and the + # replica trigger downstream all read it, so keeping the warm pass's diagnostics + # beside a restored cold result would describe a run we did not report. + _cold_res, _cold_var, _cold_neff, _cold_dict = res, var, neff, dict_return + sampler.bootstrap_from_samples(_seed, cover_frac=0.0) + res, var, neff, dict_return = sampler.integrate(like_to_integrate, *unpinned_params, **pinned_params) + _warm_lnZ = _lnZ_of_rvs(sampler._rvs, already_pooled=False) + _evidence_of_loss = ( + (_cold_lnZ is not None) and (_warm_lnZ is not None) + and numpy.isfinite(_cold_lnZ) and numpy.isfinite(_warm_lnZ) + and (_cold_lnZ - _warm_lnZ) > float(opts.sampler_l0_rescue_reject_dlnZ)) + if _evidence_of_loss: + print(" [L0 auto-rescue] *** REJECTING the warm pass *** its lnZ {:.3f} is" + " {:.3f} nats BELOW the full-support cold pass ({:.3f}), which is evidence" + " the seed missed mass the cold pass reached.".format( + _warm_lnZ + manual_avoid_overflow_logarithm, + _cold_lnZ - _warm_lnZ, + _cold_lnZ + manual_avoid_overflow_logarithm)) + if opts.sampler_l0_rescue_accept_truncated: + print(" [L0 auto-rescue] --sampler-l0-rescue-accept-truncated set:" + " reporting the warm pass anyway (may be biased LOW).") + else: + print(" [L0 auto-rescue] keeping the COLD (full-support) result; its n_eff" + " is lower but it is not missing mass. A portfolio avoids this" + " trade entirely -- its GMM member carries a defensive component.") + sampler._rvs = _cold_rvs + res, var, neff, dict_return = _cold_res, _cold_var, _cold_neff, _cold_dict + _clear_warm_state(sampler) + except Exception as _e_l0: + print(" [L0 auto-rescue] skipped (", _e_l0, ")") + _clear_warm_state(sampler) + + # Persist adapted state / trained flow for reuse by later instances. + if opts.sampler_method == 'AV' and opts.sampler_save_state and hasattr(sampler, 'save_state'): + try: + sampler.save_state(opts.sampler_save_state) + print(" AV: saved live-volume state to", opts.sampler_save_state) + except Exception as _e_ss: + print(" AV: could not save state (", _e_ss, ")") + if opts.nf_flow_save and hasattr(sampler, 'save_flow'): + try: + sampler.save_flow(opts.nf_flow_save) + print(" NF: saved trained flow to", opts.nf_flow_save) + except Exception as _e_fs: + print(" NF: could not save flow (", _e_fs, ")") + if not(res): # no resut raise ValueError(" No integral result returned") @@ -2567,6 +3127,138 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t log_res = res sqrt_var_over_res = numpy.exp(var/2 - log_res) + # ------------------------------------------------------------------ + # MC-error stabilization (see RIFT/integrators/statutils.py helpers). + # The sampler's naive sigma is (1/ESS_hat - 1/n)^{1/2} computed from the SAME + # weights as the integral: tail-blind, and small exactly when the run silently + # missed the peak. Three disclosed defenses: + # (1) floor sigma at the between-chunk lnZ scatter (adaptation nonstationarity); + # (2) print the Pareto k-hat tail diagnostic (k>0.7: sigma is a LOWER BOUND); + # (3) if triggered and --mc-error-replicas>0, re-run cold replicas and combine + # by the LINEAR mean with scatter-based error (never inverse-variance). + # ------------------------------------------------------------------ + def _extract_mc_diag(dd): + dd = dd if isinstance(dd, dict) else {} + return dd.get('pareto_khat', None), dd.get('sigma_lnZ_block', None), dd.get('n_ESS', None), dd.get('lnZ_ci90', None) + _khat, _sig_block, _n_ess, _ci90 = _extract_mc_diag(dict_return) + if _sig_block is not None and numpy.isfinite(_sig_block) and _sig_block > sqrt_var_over_res: + print(" [mc error] sigma_lnZ raised to the between-chunk scatter: {:.4f} -> {:.4f}".format(float(sqrt_var_over_res), float(_sig_block))) + sqrt_var_over_res = float(_sig_block) + if _khat is not None: + print(" [mc error] Pareto k-hat = {:.3f}{}".format(float(_khat), " (> {:.2f}: weight tail unresolved; the reported sigma is a LOWER BOUND)".format(opts.mc_error_khat_trigger) if _khat > opts.mc_error_khat_trigger else "")) + if _ci90 is not None: + print(" [mc error] bootstrap lnZ 5/50/95 quantiles: {}".format(numpy.array2string(numpy.asarray(_ci90) + manual_avoid_overflow_logarithm, precision=4))) + + _trigger_reasons = [] + if opts.mc_error_replicas > 0: + _neff_target = pinned_params.get('neff', None) + if sqrt_var_over_res > opts.mc_error_sigma_trigger: + _trigger_reasons.append('sigma={:.3f}>{:.2f}'.format(float(sqrt_var_over_res), opts.mc_error_sigma_trigger)) + if _khat is not None and _khat > opts.mc_error_khat_trigger: + _trigger_reasons.append('khat={:.2f}>{:.2f}'.format(float(_khat), opts.mc_error_khat_trigger)) + if _n_ess is not None and _n_ess < opts.mc_error_ess_trigger: + _trigger_reasons.append('ESS={:.1f}<{:g}'.format(float(_n_ess), opts.mc_error_ess_trigger)) + if _neff_target is not None and float(neff) < float(_neff_target): + _trigger_reasons.append('neff={:.1f} pooled weight w_ki / (K n_k) + # which is exactly the importance weight against the POOLED proposal density + # q'_ki = q_ki * K * n_k (pick a replica uniformly, then draw one of its n_k samples). + # Folding the factor into log_joint_s_prior is therefore a statement of the real pooled + # sampling density, not a fudge -- and it leaves every downstream weight computation + # (which all form log_integrand + log_joint_prior - log_joint_s_prior) correct untouched. + sampler._rvs = _pool_replica_rvs(_rep_rvs, sampler, rep_lnZ=_rep_lnZ, + already_resampled=bool(opts.fairdraw_extrinsic_output)) + if len(_rep_lnZ) > 1: + _K = len(_rep_lnZ) + _l = numpy.array(_rep_lnZ); _s = numpy.array(_rep_sig) + _lref = numpy.max(_l) + _Z = numpy.exp(_l - _lref) + _Zbar = numpy.mean(_Z) + _lnZ_comb = numpy.log(_Zbar) + _lref # linear mean over replicas: unbiased in Z + _sig_prop = float(numpy.sqrt(numpy.sum((_s*_Z)**2))/(_K*_Zbar)) + _sig_scatter = float(numpy.std(_l, ddof=1)/numpy.sqrt(_K)) # t_{K-1}: small-K quantiles are wider than Gaussian, hence the max() below + _sig_comb = max(_sig_prop, _sig_scatter) + print(" [mc error] combined {} replicas: lnZ {} -> {:.4f} (shift {:+.3f} vs first); sigma propagated {:.3f} / scatter {:.3f} -> {:.3f}; neff {} -> {:.1f}".format( + _K, numpy.array2string(_l + manual_avoid_overflow_logarithm, precision=3), float(_lnZ_comb + manual_avoid_overflow_logarithm), float(_lnZ_comb - _rep_lnZ[0]), + _sig_prop, _sig_scatter, _sig_comb, numpy.array2string(numpy.asarray(_rep_neff), precision=1), float(numpy.sum(_rep_neff)))) + log_res = float(_lnZ_comb) + sqrt_var_over_res = _sig_comb + # Report the POOLED n_eff, not the sum. The sum claims the posterior carries the + # combined effective sample size of K independent runs, which is only true if they + # agree; when they disagree -- the case these replicas exist to detect -- the pooled + # Kish n_eff is smaller, and that disagreement is exactly what should show up here. + _neff_pooled = _kish_neff_of_rvs(sampler._rvs) + neff = float(_neff_pooled) if _neff_pooled is not None else float(numpy.sum(_rep_neff)) + if _neff_pooled is not None: + print(" [mc error] pooled posterior: {} samples, Kish n_eff {:.1f} (sum over replicas was {:.1f})".format( + len(numpy.atleast_1d(list(sampler._rvs.values())[0])) if sampler._rvs else 0, + float(_neff_pooled), float(numpy.sum(_rep_neff)))) + # keep the (res, var) pair consistent for any downstream reader + if not(opts.internal_use_lnL): + res = numpy.exp(log_res); var = (sqrt_var_over_res*res)**2 + else: + res = log_res; var = 2*numpy.log(sqrt_var_over_res) + 2*log_res + # Calibration MC error budget. The sampler's `var` is the EXTRINSIC sampling # variance with the cal draw set held FIXED -- it is structurally blind to the # Monte-Carlo error of the (1/n_cal) sum over realizations, which dominates @@ -3004,6 +3696,12 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t else: samples["latitude"] = samples["declination"] samples["longitude"] = samples["right_ascension"] + # NOTE: the sim_inspiral XML schema is deliberately sparse -- it carries lnL only (below, mapped + # to alpha1), NOT the importance weight. It does NOT persist log_joint_prior/log_joint_s_prior, + # so a downstream consumer CANNOT reconstruct the true weight (log_integrand + log_joint_prior - + # log_joint_s_prior) from this file and must not reweight it by likelihood for a weighted-posterior + # or shape check. For richer, ASCII, per-sample output that carries the full log-weight, use + # --extrinsic-proposal-output (writes lnL + ln(prior) - ln(s_prior)) or --calibration-export-posterior. if "log_integrand" in samples: samples["loglikelihood"] = samples["log_integrand"] + manual_avoid_overflow_logarithm else: @@ -3192,6 +3890,25 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t # integrate_log() fairdraw does self._rvs[key][indx_list] and raises # "list indices must be integers or slices, not ndarray". This silently dropped # every binary after the first in each ILE batch whenever fairdraw export was on. + # SEQUENTIAL WARM-START SEED must be captured BEFORE the _rvs wipe below. The caller's capture + # block runs after this function returns, by which point _rvs is empty -- so + # --sampler-sequential-warmstart was silently inert (it always saw no samples and never seeded + # the next point). Stash the seed here instead; the caller consumes _SEQ_WS_PENDING. + global _SEQ_WS_PENDING + _SEQ_WS_PENDING = None + if getattr(opts, 'sampler_sequential_warmstart', False) and hasattr(sampler, 'bootstrap_from_samples'): + try: + _lnkey = 'log_integrand' if 'log_integrand' in sampler._rvs else ('integrand' if 'integrand' in sampler._rvs else None) + if _lnkey is not None and all(p in sampler._rvs for p in sampler.params_ordered): + _lnv = np.asarray(sampler.identity_convert(sampler._rvs[_lnkey]), dtype=float).ravel() + if _lnv.size >= 2 and np.any(np.isfinite(_lnv)): + _cols = np.vstack([np.asarray(sampler.identity_convert(sampler._rvs[p]), dtype=float).ravel() + for p in sampler.params_ordered]).T + _keep = _lnv > (np.nanmax(_lnv) - opts.sampler_sequential_warmstart_deltalnL) + _SEQ_WS_PENDING = _cols[_keep] if np.sum(_keep) >= 2 else (_cols if _cols.shape[0] >= 2 else None) + except Exception as _e_cap: + print(" [seq warm-start] could not capture seed ({})".format(_e_cap)) + sampler._rvs = {} return res @@ -3200,16 +3917,76 @@ def analyze_event(P_list, indx_event, data_dict, psd_dict, fmax, opts,inv_spec_t lnL_sofar = -np.inf no_adapt_sky = False +# L3: load a proposal field (from a previous ILE iteration) once, if provided +_proposal_field = None +if opts.extrinsic_proposal_field and hasattr(sampler, 'bootstrap_from_samples'): + try: + from RIFT.integrators.proposal_field import ProposalField, lambda_from_P + _proposal_field = ProposalField.load(opts.extrinsic_proposal_field) + print(" [proposal-field] loaded {} entries from {}".format(len(_proposal_field), opts.extrinsic_proposal_field)) + except Exception as _e_pf: + print(" [proposal-field] could not load {} ({}); ignoring".format(opts.extrinsic_proposal_field, _e_pf)) + _proposal_field = None + +_seq_ws_proposal = None # L1 sequential hot-feed: previous point's extrinsic seed (in-memory) +_SEQ_WS_PENDING = None # seed captured INSIDE analyze_event, before it wipes sampler._rvs for indx in numpy.arange(len(P_list)): try: + # UNCONDITIONAL per-point reset. mcsamplerPortfolio.integrate_log does NOT call self.setup() + # (it is commented out at mcsamplerPortfolio.py:874), so member state -- including the AV live + # volume contracted around the PREVIOUS point -- survives into the next integral. With + # --n-events-to-analyze > 1 the second point then draws from a box shaped by the first, and if + # its support falls outside, lnZ is biased low with a healthy-looking n_eff and no error. This + # is independent of any warm-start feature: it bites users who never enable one. Reset FIRST, + # so an intentional seed installed just below survives. + if indx > 0: + _clear_warm_state(sampler) + # Warm-start this point's extrinsic integral. Order is preserved (no grid + # reordering), so a truncated worker still drops a spatially-unbiased subset; a + # coverage floor + inflation margin make a poorly-matched transfer degrade to + # cold rather than bias. AV only. Priority: L3 proposal field (cross-iteration) + # over L1 sequential (same worker). + if _proposal_field is not None and len(_proposal_field) > 0: + try: + from RIFT.integrators.proposal_field import lambda_from_P + _seed = _proposal_field.warm_seed_for(lambda_from_P(P_list[indx]), k=1) + if _seed is not None and len(_seed) >= 2: + sampler.bootstrap_from_samples(_seed, + params=_proposal_field.extrinsic_params, + cover_frac=opts.extrinsic_proposal_field_cover_frac, + inflate=opts.extrinsic_proposal_field_inflate) + print(" [proposal-field] point {} seeded from nearest entry ({} pts, cover_frac={}, inflate={})".format( + indx, len(_seed), opts.extrinsic_proposal_field_cover_frac, opts.extrinsic_proposal_field_inflate)) + except Exception as _e_pfq: + print(" [proposal-field] seed skipped for point {} ({})".format(indx, _e_pfq)) + _clear_warm_state(sampler) + elif opts.sampler_sequential_warmstart and (_seq_ws_proposal is not None) and hasattr(sampler, 'bootstrap_from_samples'): + try: + sampler.bootstrap_from_samples(_seq_ws_proposal, + cover_frac=opts.sampler_sequential_warmstart_cover_frac) + print(" [seq warm-start] point {} seeded from previous point ({} pts, cover_frac={})".format( + indx, len(_seq_ws_proposal), opts.sampler_sequential_warmstart_cover_frac)) + except Exception as _e_sw: + print(" [seq warm-start] skipped for point {} ({})".format(indx, _e_sw)) + _clear_warm_state(sampler) res = analyze_event(P_list, indx, data_dict, psd_dict, fmax, opts) + # capture this point's converged high-likelihood extrinsic samples as the seed + # for the next point (in-memory; no files leave the job) + if opts.sampler_sequential_warmstart and hasattr(sampler, 'bootstrap_from_samples'): + # Consume the seed captured inside analyze_event BEFORE it wiped sampler._rvs. Reading + # sampler._rvs here would always find it empty (that wipe is required: it fixes a fairdraw + # export bug that silently dropped every binary after the first). + _seq_ws_proposal = _SEQ_WS_PENDING + if _seq_ws_proposal is None: + print(" [seq warm-start] no seed captured from point {}".format(indx)) + _clear_warm_state(sampler) # clear before the next point (re-seeded above if enabled) # abort if horrible (nan event) - done with 'raise' lnL_sofar = np.max([lnL_sofar,res]) if opts.force_reset_all: # depends on integrator! May not always be availble if opts.sampler_method == "adaptive_cartesian_gpu": for name in sampler.params: sampler.reset_sampling(name) - elif opts.sampler_method == "GMM": + elif use_gmm_args: # standalone GMM or a portfolio with a GMM member # reset the GMM dictionary for component in gmm_dict: gmm_dict[component] = None @@ -3230,6 +4007,11 @@ for indx in numpy.arange(len(P_list)): except Exception as exception_failure: print( " ===> FAILED ANALYSIS <==== ") print( exception_failure) + # The message alone ("boolean index did not match...", "index out of range", ...) is rarely enough + # to locate a failure inside the sampler stack, and this handler is often the ONLY record a batch + # job leaves behind. Print the traceback too -- it costs nothing on the success path. + import traceback as _tb_mod + _tb_mod.print_exc() if opts.internal_make_empty_file_on_error: fname_output_txt = opts.output_file +"_"+str(indx)+"_" + ".dat" open(fname_output_txt,'a').close() # create empty file diff --git a/MonteCarloMarginalizeCode/Code/bin/util_BuildProposalField.py b/MonteCarloMarginalizeCode/Code/bin/util_BuildProposalField.py new file mode 100644 index 000000000..350f0f5ad --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/bin/util_BuildProposalField.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python +""" +util_BuildProposalField.py + +L3 producer (MULTI-pilot generalization): aggregate the extrinsic outputs of a FEW +cherry-picked points into a ProposalField keyed by intrinsic parameters, consumed by +integrate_likelihood_extrinsic_batchmode --extrinsic-proposal-field. + +IMPORTANT -- do NOT point this at a whole grid analyzed with --save-samples: per-point +sample dumps blow up disk fast, and a full field is rarely needed. The standard, +disk-safe path is a SINGLE cherry-picked pilot (util_PickPilotPoint.py picks the best +point by lnL after a cheap iteration 0; run ONLY that point with --save-samples; +warm-start the rest with --sampler-warmstart-samples). Use this multi-pilot field only +when a few (2-4) well-separated pilots are genuinely warranted (e.g. a known multimodal +source) -- run save-samples on THOSE few points only. + +For each provided point it reads that point's saved extrinsic samples (--save-samples +.xml.gz), keeps the high-likelihood subset, converts to the sampler's coordinate +convention (matching --declination-cosine-sampler / --inclination-cosine-sampler), and +records it against the point's intrinsic lambda = [m1, m2, s1x..s2z]. + +A proposal only ever shapes p_s, so a stale/partial field can only cost efficiency, never +bias -- missing points are simply skipped. + +Usage: + util_BuildProposalField.py --grid overlap-grid-5.xml.gz --output-prefix ile_5 \ + --out proposal_field_5.npz [--deltalnL 15] [--max-per-point 4000] \ + [--no-cosine-dec] [--no-cosine-incl] +""" +from __future__ import print_function +import argparse +import glob +import os +import numpy as np + +# sampler extrinsic coordinate order used by the AV extrinsic integrator +EXTRINSIC_PARAMS = ["right_ascension", "declination", "phi_orb", "inclination", "psi", "distance"] + + +def _read_extrinsic_xml(path): + """Return (samples Nx6 in sampler coords, lnL N) from an ILE --save-samples xml.gz, + or (None, None). Converts physical (dec, incl) to the cosine-sampler variables.""" + try: + from igwn_ligolw import ligolw, lsctables, utils + xd = utils.load_filename(path, contenthandler=lsctables.use_in(ligolw.LIGOLWContentHandler)) + t = lsctables.SimInspiralTable.get_table(xd) + except Exception as e: + print(" (could not read {}: {})".format(path, e)) + return None, None + if len(t) == 0: + return None, None + ra = np.array([r.longitude for r in t]); lat = np.array([r.latitude for r in t]) + dist = np.array([r.distance for r in t]); incl = np.array([r.inclination for r in t]) + psi = np.array([r.polarization for r in t]); phi = np.array([r.coa_phase for r in t]) + lnL = np.array([getattr(r, 'alpha1', 0.0) for r in t]) + return ra, lat, dist, incl, psi, phi, lnL + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--grid", required=True, help="sim xml with the iteration's intrinsic grid") + ap.add_argument("--output-prefix", required=True, help="ILE --output-file prefix; reads __.xml.gz") + ap.add_argument("--out", required=True, help="output ProposalField .npz") + ap.add_argument("--deltalnL", type=float, default=15.0, help="keep samples within this lnL of each point's max") + ap.add_argument("--max-per-point", type=int, default=4000) + ap.add_argument("--no-cosine-dec", action="store_true", help="grid did NOT use --declination-cosine-sampler") + ap.add_argument("--no-cosine-incl", action="store_true", help="grid did NOT use --inclination-cosine-sampler") + args = ap.parse_args() + + from RIFT.integrators.proposal_field import ProposalField, lambda_from_P, LAMBDA_INTRINSIC_PARAMS + import RIFT.lalsimutils as lalsimutils + + P_list = lalsimutils.xml_to_ChooseWaveformParams_array(args.grid) + pf = ProposalField(intrinsic_params=LAMBDA_INTRINSIC_PARAMS, extrinsic_params=EXTRINSIC_PARAMS) + + n_added = 0 + for i, P in enumerate(P_list): + path = "{}_{}_.xml.gz".format(args.output_prefix, i) + if not os.path.exists(path): + continue + got = _read_extrinsic_xml(path) + if got[0] is None: + continue + ra, lat, dist, incl, psi, phi, lnL = got + dec_s = np.sin(lat) if not args.no_cosine_dec else lat + incl_s = np.cos(incl) if not args.no_cosine_incl else incl + cols = np.vstack([ra, dec_s, phi, incl_s, psi, dist]).T + keep = lnL > (np.nanmax(lnL) - args.deltalnL) + if np.sum(keep) < 2: + keep = np.ones(len(cols), dtype=bool) # fall back to all saved samples + sub = cols[keep] + if len(sub) > args.max_per_point: + sub = sub[np.random.RandomState(0).choice(len(sub), args.max_per_point, replace=False)] + pf.add(lambda_from_P(P), sub) + n_added += 1 + + if len(pf) == 0: + print("WARNING: no usable per-point outputs found; writing nothing.") + return + pf.save(args.out) + print("Built ProposalField with {} entries (of {} grid points) -> {} ({} bytes)".format( + n_added, len(P_list), args.out, os.path.getsize(args.out))) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py b/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py index 57683fe3d..a2c88d6f1 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_CleanILE.py @@ -89,12 +89,33 @@ for key in data_at_intrinsic: lnL, sigmaOverL, ntot,neff = np.transpose(data_at_intrinsic[key]) + lnL = np.atleast_1d(lnL); sigmaOverL = np.atleast_1d(sigmaOverL); ntot = np.atleast_1d(ntot); neff = np.atleast_1d(neff) sigmaOverL = np.maximum(sigmaOverL, 1e-7*np.ones(len(lnL))) # prevent accidental underflow during debugging/using synthetic data with no error lnLmax = np.max(lnL) - sigma = sigmaOverL*np.exp(lnL-lnLmax) # remove overall Lmax factor, which factors out from the weights constructed from \sigma - wts = weight_simulations.AverageSimulationWeights(None, None,sigma) - lnLmeanMinusLmax = np.log(np.sum(np.exp(lnL - lnLmax)*wts)) - sigmaNetOverL = (np.sqrt(1./np.sum(1./sigma/sigma)))/np.exp(lnLmeanMinusLmax) + L = np.exp(lnL - lnLmax) # remove overall Lmax factor, which factors out of the combination + K = len(lnL) + # Combine repeated evaluations by their SAMPLE-COUNT-weighted LINEAR mean. + # DO NOT inverse-variance weight with the reported sigmas: each sigma is + # computed from the same importance weights as its lnL, so a replica that + # silently missed the likelihood peak reports BOTH a low lnL AND a small + # sigma -- 1/sigma^2 weighting then overweights the worst replica, giving a + # systematically low combined lnL with an overconfident combined error. + # The pooled (ntot-weighted) linear mean is unbiased in L regardless. + wts = np.asarray(ntot, dtype=float) + if np.any(wts <= 0) or not np.all(np.isfinite(wts)): + wts = np.ones(K) + wts = wts/np.sum(wts) + Lbar = np.sum(wts*L) + lnLmeanMinusLmax = np.log(Lbar) + # Error: max(propagated per-run sigmas, between-replica scatter). Only the + # scatter term can see the replica lottery (correlated underreporting); with + # K replicas it has K-1 dof, so treat the result as a t-interval downstream. + sigma_prop = np.sqrt(np.sum((wts*sigmaOverL*L)**2))/Lbar + if K > 1: + sigma_scatter = np.sqrt( np.sum(wts**2 * (L - Lbar)**2) * K/(K-1.) )/Lbar + else: + sigma_scatter = 0. + sigmaNetOverL = max(sigma_prop, sigma_scatter) if opts.eccentricity: diff --git a/MonteCarloMarginalizeCode/Code/bin/util_PickPilotPoint.py b/MonteCarloMarginalizeCode/Code/bin/util_PickPilotPoint.py new file mode 100644 index 000000000..24c54e2aa --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/bin/util_PickPilotPoint.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python +""" +util_PickPilotPoint.py + +Cherry-pick the best intrinsic grid point(s) after iteration 0, to run as warm-start +PILOT(s). This is the disk- and fit-safe way to seed later ILE evaluations: + + * We must NOT --save-samples the whole grid (disk blows up fast), and + * we must NOT pick pilots at random (a poor point makes a poor extrinsic proposal, and + -- more importantly -- the CIP fit that builds the next grid needs good coverage, so + random subsetting is dangerous). + +So: after a cheap iteration-0 ILE (per-point marginal lnL only, no sample dump), pick the +TOP-k points by lnL and emit them as a small sub-grid. A follow-up ILE runs only those +few points WITH --save-samples (~tens of KB each) to produce the pilot extrinsic +proposal, which then warm-starts the rest of the grid (via +integrate_likelihood_extrinsic_batchmode --sampler-warmstart-samples). + +Usage: + util_PickPilotPoint.py --grid overlap-grid-0.xml.gz --output-prefix ile_0 \ + --top-k 1 --out pilot-grid.xml.gz + # then run ILE on pilot-grid.xml.gz with --save-samples, and warm-start iteration 1+ + # from the pilot's saved extrinsic samples. + +lnL is read from the per-point ILE .dat outputs (__.dat, column 10 = marginal +lnL), or from a single --net composite file with an explicit --lnL-column. +""" +from __future__ import print_function +import argparse +import os +import numpy as np + + +def _lnL_from_dat(prefix, i): + path = "{}_{}_.dat".format(prefix, i) + if not os.path.exists(path): + return None + try: + row = np.atleast_2d(np.loadtxt(path)) + # ILE .dat layout: idx m1 m2 s1x s1y s1z s2x s2y s2z lnL sigma ntotal neff + return float(row[0, 9]) + except Exception: + return None + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--grid", required=True, help="iteration-0 intrinsic grid xml") + ap.add_argument("--output-prefix", default=None, help="ILE --output-file prefix; reads __.dat for lnL") + ap.add_argument("--net", default=None, help="alternative: a composite/net file with one row per grid point") + ap.add_argument("--lnL-column", type=int, default=9, help="0-based lnL column in --net (default 9)") + ap.add_argument("--top-k", type=int, default=1) + ap.add_argument("--out", required=True, help="output pilot sub-grid xml") + args = ap.parse_args() + + import RIFT.lalsimutils as lalsimutils + P_list = lalsimutils.xml_to_ChooseWaveformParams_array(args.grid) + + lnL = np.full(len(P_list), -np.inf) + if args.net: + dat = np.atleast_2d(np.loadtxt(args.net)) + m = min(len(P_list), dat.shape[0]) + lnL[:m] = dat[:m, args.lnL_column] + elif args.output_prefix: + for i in range(len(P_list)): + v = _lnL_from_dat(args.output_prefix, i) + if v is not None: + lnL[i] = v + else: + raise SystemExit("provide --output-prefix or --net for lnL values") + + finite = np.isfinite(lnL) + if not np.any(finite): + raise SystemExit("no finite lnL found; cannot pick a pilot") + order = np.argsort(lnL)[::-1] + order = [i for i in order if np.isfinite(lnL[i])][:max(1, args.top_k)] + print("Picked {} pilot point(s) by lnL: {}".format( + len(order), [(int(i), round(float(lnL[i]), 1)) for i in order])) + + P_out = [P_list[i] for i in order] + lalsimutils.ChooseWaveformParams_array_to_xml(P_out, fname=args.out.replace('.xml.gz', '').replace('.xml', '')) + print("Wrote pilot sub-grid ->", args.out) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py index daa9e9ebd..daa64cf77 100755 --- a/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py +++ b/MonteCarloMarginalizeCode/Code/bin/util_RIFT_pseudo_pipe.py @@ -468,6 +468,7 @@ def run_lisa_known_sky_surface(opts): parser.add_argument("--internal-ile-srate-time-resampling",default=None, help=" Adds --srate-resample-time-marginalization to ILE for output, to provide higher-resolution time output ") parser.add_argument("--internal-ile-srate-internal",default=None, help=" Adds --srate-internal to ILE, modifying how calculations are performed internally to use a higher sampling rate ") parser.add_argument("--internal-ile-interpolate-time",action='store_true',help="Pass --interpolate-time True to ILE, enabling cubic interpolation of Q_lm at fractional detector arrival times in the maintained NoLoop likelihood.") +parser.add_argument("--internal-ile-n-chunk",default=None,type=int,help="Override the extrinsic chunk size (--n-chunk) passed to ILE, via the helper. Default behaviour (helper): 40000, scaled linearly with SNR above 40 and capped at 160000, because at high SNR the posterior is a vanishing fraction of the prior volume and a small chunk gives few informative samples per adaptation step. Larger chunks cost GPU memory but measured HOST memory (what RequestMemory governs) is flat, so no memory-request change is normally needed. EXPERTS ONLY.") parser.add_argument("--batch-extrinsic",action='store_true') parser.add_argument("--fmin",default=20,type=int,help="Mininum frequency for integration. template minimum frequency (we hope) so all modes resolved at this frequency") # should be 23 for the BNS parser.add_argument("--fmin-template",default=None,type=float,help="Mininum frequency for template. If provided, then overrides automated settings for fmin-template = fmin/Lmax") # should be 23 for the BNS @@ -1160,6 +1161,13 @@ def run_lisa_known_sky_surface(opts): cmd += " --internal-ile-auto-logarithm-offset " if opts.internal_ile_rotate_phase: cmd += " --internal-ile-rotate-phase " +if opts.internal_ile_interpolate_time: + # HELPER passthrough (not a raw ILE arg): the helper owns ILE argument construction, and it + # also knows whether the NoLoop path (--vectorized --gpu --force-xpy) that --interpolate-time + # requires is actually in use. + cmd += " --internal-ile-interpolate-time " +if not(opts.internal_ile_n_chunk is None): + cmd += " --internal-ile-n-chunk {} ".format(int(opts.internal_ile_n_chunk)) # If user provides ini file *and* ini file has fake-cache field, generate a local.cache file, and pass it as argument if opts.use_ini: # config = ConfigParser.ConfigParser() @@ -1274,8 +1282,6 @@ def run_lisa_known_sky_surface(opts): # - requested or # - AC + not freezeadapt line += " --force-reset-all " -if opts.internal_ile_interpolate_time: - line += " --interpolate-time True " if not(opts.manual_extra_ile_args is None): line += " {} ".format(opts.manual_extra_ile_args) # embed with space on each side, avoid collisions if '--declination ' in opts.manual_extra_ile_args: # if we are pinning dec, we aren't using a cosine coordinate. Don't mess up. diff --git a/MonteCarloMarginalizeCode/Code/test/README.md b/MonteCarloMarginalizeCode/Code/test/README.md index e77cc4266..c5340fc67 100644 --- a/MonteCarloMarginalizeCode/Code/test/README.md +++ b/MonteCarloMarginalizeCode/Code/test/README.md @@ -12,3 +12,5 @@ See pp * ``test_mcsamplerEnsemble_extended.py`` : best single-contact test. 3d gaussian integration, with plot of recovered CDF. * ``test_mcsampler_rosenbrock``: Simple 2d test + +* ``expensive_before_merging/integrators``: **posterior shape-recovery merge gate** — REQUIRED before merging any integrator change into a production line; much stronger than the integral tests above (catches integral-invisible shape failures and silent n_eff collapse). See RIFT/integrators/TESTING.md. diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/README.md b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/README.md new file mode 100644 index 000000000..304ef15aa --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/README.md @@ -0,0 +1,38 @@ +# Expensive pre-merge validation suites + +Tests in this tree are **NOT** run in per-commit CI. They are the strong, +slow checks run **before confirming a merge** into a production line +(`rift_O4d`, future `rift_O4e`, ...), and when back-checking a production line +against its predecessor (e.g. `rift_O4c` -> `rift_O4d`). + +Rationale: the fast CI gate (`.travis/test-integrate.sh`) validates the +*integral* on a single 3-D Gaussian. Integrals are easy — importance-sampling +estimates of Z are unbiased under weak conditions — while the recovered +*posterior shape* (the weighted sample cloud consumed by CIP and the fairdraw +machinery) can be subtly wrong: clipped tails, wrong widths, missing mixture +components, distorted correlations. Production merges must pass the shape +test, not just the integral test. + +## Suites + +* `integrators/` — posterior shape-recovery gate for the MC integrators + (AV, GMM, NF, portfolio; optionally AC/default). Random seeded Gaussian + mixtures across dimensions, following RIFT-FinerNet + `demos/integrators/multigauss_direct` (Wagner et al). See + `integrators/shape_recovery.py` docstring for method and thresholds, and + `integrators/run_shape_recovery.sh` for the standard invocation. + +## Merge workflow + +1. Run the suite on the **base** branch: `--json base.json`. +2. Run the suite on the **candidate** branch (same preset/seeds): `--json pr.json`. +3. `python integrators/compare_shape_results.py base.json pr.json` + - Merge-blocking: any strict-sampler run that regresses PASS -> FAIL, or a + metric regression beyond tolerance (see script). + - Pre-existing failures (FAIL on both) do not block, but should be ticketed. +4. Attach both JSON files + the comparison output to the PR before confirming. + +Policy: AV is the gold-standard production sampler and is always strict. +GMM is strict by default. NF and portfolio are warn-only by default (known +weaker in older lines, e.g. rift_O4c); tighten with `--strict-samplers` as +they harden. diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/PROPOSAL_bootstrap_gate_cases.md b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/PROPOSAL_bootstrap_gate_cases.md new file mode 100644 index 000000000..ff9fcfafe --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/PROPOSAL_bootstrap_gate_cases.md @@ -0,0 +1,423 @@ +# PROPOSAL: warm-start / bootstrap cases for the shape-recovery merge gate + +Status: **proposal only** — nothing in this directory has been modified. Every number below was +measured on this branch (`rift_O4d_portfolio_freeze_tuning`) on CPU +(`CUDA_VISIBLE_DEVICES=""`, `OMP_NUM_THREADS=1`, +`/cvmfs/software.igwn.org/conda/envs/igwn-py310`, numpy 1.24.4), using the gate's own +`MixtureTarget`, `build_sampler`, `shape_metrics` and `evaluate`. + +(Note on location: the task brief called this directory `demos/`. The gate actually lives +at `MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/`; `demos/` only holds +`integrator_snr_lottery`. This file is placed with the gate it modifies.) + +--- + +## 1. How the existing gate specifies and scores a case + +### Case specification +`shape_recovery.py` builds a pure product matrix. A "case" is a 4-tuple + + (sampler kind, ndim, ncomp, target_seed) + +* `PRESETS[preset]` gives `dims x ncomps x seeds x nmax_per_dim x neff` + (`quick`: dims 2,4 / ncomps 2 / seeds 101 / nmax_per_dim 50000 / neff 2000; + `standard`: dims 2,4,6,8 / ncomps 1,3 / seeds 101,202,303 / nmax_per_dim 200000 / neff 3000). +* `main()` expands the product over `--samplers` into jobs + `(kind, (d, nc, ts), nmax = nmax_per_dim*d, neff, run_seed)` and dispatches them through + `_worker` (spawn `multiprocessing.Pool`, so anything monkey-patched in the parent is lost — + see `probe_portfolio_optin_flags.py`). +* `MixtureTarget(ndim, ncomp, seed)` is fully determined by the seed: weights `U(0.1,1.1)` + normalised, means `U(-3/sqrt(d), 3/sqrt(d))`, Wishart covariances around `sigma_1d=0.7`, all on + the box `[-5,5]^d`. `true_lnZ = LNL_OFFSET + ln(in-box mass) - sum ln(width)` from a 10^6-point + rejection-sampled truth pool. +* `run_one` drives the sampler through its production API with + `n=n_chunk, n_adapt=100, floor_level=0, tempering_exp=0.1, neff, nmax, save_intg=True`, + reads the weighted cloud back out of `sampler._rvs`, and **never raises** — exceptions land in + `record["error"]`. + +### Scoring — `evaluate(record)` returns one of `PASS / FAIL / STARVED / ERROR` +1. `error` set -> `ERROR` +2. `n_eff < MIN_NEFF_FOR_SHAPE (100)` -> `STARVED` **and nothing else is checked** +3. otherwise all of, per dimension `d`: + * `JS[d] < JS_MULT*floor[d] + JS_ABS_MIN` (3.0 x self-calibrated floor + 0.004), where the floor + is JS(truth subsample at this run's own n_ESS, truth pool), mean + 2 sd over 5 subsamples; + * `|mean_pull[d]| <= max(5/sqrt(n_ESS), 0.05)`; + * `|width_ratio[d] - 1| <= max(5/sqrt(2 n_ESS), 0.05)`; + * `corr_diff_max <= max(8/sqrt(n_ESS), 0.08)`; + * `|lnZ_hat - true_lnZ| <= max(4*rel_err, 0.10)`. +4. `main()` exits 1 iff a **strict** sampler (`--strict-samplers`, default `AV,GMM`) scored `FAIL`. + `STARVED` never sets the exit code; `WARN` (non-strict FAIL) never sets it. + +### Base-vs-candidate — `compare_shape_results.py` +Records are paired on `(kind, target)`. Blocking verdicts, for strict kinds only: +* `PASS` on base -> anything else on candidate = `REGRESSION(pass->...)` — this *includes* + `PASS -> STARVED`; +* both `PASS` but a summary metric worsens by more than `TOL_WORSE` + (`js .005, mean_pull .05, width_dev .05, corr .05, bias_ln .10`) or + `n_eff_cand < 0.5 * n_eff_base` = `REGRESSION(metrics)`. +`STARVED -> STARVED` is `BOTH-STARVED` and **never blocks**; `STARVED -> PASS` is `IMPROVED`. +The script needs no code change to accept new kinds — it keys on `record["kind"]` generically. + +--- + +## 2. What the gate does not exercise, and what actually catches it + +Two independent defects, both of which produce *no exception and no obviously bad diagnostic*: + +**(a) warm seed stored but never installed on the draw path.** `bootstrap_from_*` only writes +`self._warm`; historically only `mcsamplerAdaptiveVolume.integrate_log` consumed it. A PORTFOLIO +calls `member.draw_simplified()` directly and never runs the member's `integrate_log`, so the AV +member drew from the cold single-bin grid. `mcsamplerAdaptiveVolume._apply_warm_state()` (called +from `draw_simplified`) is the fix. + +**(b) stale contracted live volume leaking between sequential points.** +`mcsamplerPortfolio.integrate_log` has `self.setup()` **commented out** (mcsamplerPortfolio.py +~L874), so member state survives the call. Point 2 therefore inherits point 1's contracted AV +live volume. `mcsamplerPortfolio.clear_warm_state()` (and the driver's `_clear_warm_state()` +helper, `bin/integrate_likelihood_extrinsic_batchmode` L1943) is the fix. +Standalone AV is immune: `AV.integrate_log` calls `self.setup()` on entry. + +### Measurement 1 — bug (a) is bit-for-bit invisible +Emulating the pre-fix code (`member._warm_applied = True` right after seeding, so +`_apply_warm_state()` is a no-op) and seeding **only the AV member**, the run reproduces the cold +run to the last printed digit. 15/15 pairs identical, e.g. d4 nc1 ts101 rs987654: + +| mode | n_eff | lnZ bias | n_eval | +|---|---|---|---| +| cold | 46.3 | -0.0209 | 200000 | +| warm_av_INERT | 46.3 | -0.0209 | 200000 | +| warm_av (fixed)| 45.3 | -0.1398 | 200000 | + +### Measurement 2 — a black-box "warm beats cold" assertion does NOT test the AV path +The gate's portfolio is AV + GMM, and `mcsamplerEnsemble` *also* has `bootstrap_from_samples` +(added in 8c29876b). `portfolio.bootstrap_from_samples` seeds both, and the GMM member captures +essentially the whole win. d=6, ncomp=1, `nmax=6e5`, `neff=3000`, 3 target seeds x 4 run seeds, +scored with the gate's own `evaluate()`: + +| target seed | fixed n_eff | INERT n_eff (bug (a) present) | verdict, both | +|---|---|---|---| +| 101 | 5504 – 5941 | 5078 – 5972 | PASS | +| 202 | 3240 – 3439 | 5310 – 5623 | PASS | +| 303 | 5360 – 5787 | 4692 – 5540 | PASS | +| **all 12** | **3240 – 5941** | **4692 – 5972** | **12/12 PASS both** | + +The buggy configuration is sometimes *better*. **No n_eff / lnZ / shape assertion can separate +these.** What does separate them, exactly and with zero RNG dependence, is the AV member's live +volume after a single `portfolio.draw()`: + +| | AV member `V` | AV member live bins | +|---|---|---| +| fixed | 0.0313 – 0.129 (d6), 0.0432 (d2), 0.0862 – 0.159 (d4) | 47 – 656 | +| INERT | **1.000 exactly** | **1 exactly** | + +Behavioural form of the same statement — fraction of AV-member draws landing inside the seed +cloud's bounding box, versus the uniform-box expectation `V_unif` (10000 draws, 2 target seeds x +2 run seeds each): + +| d | V_unif | fixed fraction | ratio | INERT ratio | +|---|---|---|---|---| +| 2 | 0.336 / 0.187 | 0.729 – 0.764 | 2.17 / 4.09 | 0.93 – 1.01 | +| 4 | 0.0489 / 0.0311 | 0.240 – 0.289 | 4.9 / 9.2 | 0.95 – 1.06 | +| 6 | 0.00350 / 0.0210 | 0.072 – 0.139 | 6.1 – 22.3 | 0.86 – 1.09 | + +Run-seed scatter of the fraction is < 0.005 absolute (binomial, n=10^4). + +### Measurement 3 — bug (b) is enormous, and does not even need a warm start +Two displaced targets A (`offset=-2`) and B (`offset=+2`), both `scale_x0=1.0`, integrated +sequentially on ONE portfolio, d=2, ncomp=1, `nmax=1e5`, `neff=2000`, 3 target seeds x 5 run seeds: + +| between-point handling | B n_eff (min–max) | \|B lnZ bias\| (median, max) | +|---|---|---| +| `sampler._warm = None` only (pre-fix driver) | **0.0 – 3.9** | 16.8, **337** | +| `clear_warm_state()` (post-fix) | 799.8 – 2089.1 | 0.0038, 0.0248 | +| fresh sampler for B (reference) | 819.1 – 2082.3 | 0.0054, 0.0124 | + +With **no bootstrap at all** (pure sequential reuse, same displaced pair): leak B n_eff 0.0 – 9.9, +|bias| up to 24.2; `clear` 1587 – 2094, |bias| <= 0.019. The leak is the run-contracted grid, not +the seed — so this case guards `--n-events-to-analyze > 1` even for users who never warm-start. +Standalone `AV` shows leak == cold bit-for-bit, confirming the defect is portfolio-only. + +Scored with the gate's own `evaluate()` (4 run seeds x 3 target seeds): + +| mode (12 runs each) | verdict | n_eff | \|bias\| | JSmax | width dev | max pull | +|---|---|---|---|---|---|---| +| clear | 12/12 PASS | 815.8 – 2103.4 | 0.0004 – 0.0188 | <= 0.0003 | <= 0.010 | <= 0.017 | +| LEAK | 12/12 STARVED | 0.0 – 12.8 | 0.20 – 318.5 | <= 0.69 | <= 1.00 | <= 6.78 | + +--- + +## 3. Proposed cases + +Two new sampler kinds, so `compare_shape_results.py` pairs them automatically. + +### Case W — `portfolio_warm` (guards bug (a) + "the warm-start feature went inert") +* target: `MixtureTarget(ndim=6, ncomp=1, seed in {101, 202, 303})`, unmodified. +* seed cloud: 3000 fair draws from the target's own truth pool, `RandomState(target.seed+13)`. +* budget: `nmax = 6*100000`, `neff = 3000`, `n_chunk = 10000`. +* procedure: + 1. probe sampler: build portfolio, `setup()`, `bootstrap_from_samples(cloud, cover_frac=0.0)`, + one `draw(n_chunk)`; record `warm_V`, `warm_bins` from `portfolio_realizations[0]`, and + `warm_box_frac` / `warm_box_frac_uniform`. + 2. fresh sampler, same seeding, `integrate_log` -> the record scored by the normal metrics. +* assertions (all hard; this kind is STRICT and STARVED is promoted to FAIL): + * **A1 (install)** `warm_V < 0.9` and `warm_bins > 1`. Measured over 12 runs: fixed + `V in [0.0313, 0.1292]`, `bins in [489, 656]`; buggy `V = 1.000` and `bins = 1` in 12/12. + No RNG enters `V` or `binunique` — the margin is categorical, not statistical. + * **A2 (behavioural install)** `warm_box_frac >= 3 * warm_box_frac_uniform`. Measured at d6: + 6.1x – 22.3x; buggy: 0.86 – 1.09. Margin >= 2x on the worst measured seed. + * **A3 (feature-level)** `n_eff >= 1000`. Measured warm 2426 – 5941 (5 target seeds x 5 run + seeds at neff=2000, plus 3 x 4 at neff=3000); cold at the same budget 3.8 – 91.7. 2.4x margin + below the warm minimum, 11x above the cold maximum. + * normal JS / pull / width / corr / lnZ checks (measured warm: JSmax <= 0.0008, + width dev <= 0.026, |bias| <= 0.0038; 12/12 PASS). + + A3 does **not** isolate the AV path (Measurement 2); it catches "all warm-start channels went + inert", which is a real and separate regression. A1/A2 are what catch bug (a). A1 is + deliberately white-box: bug (a) has no statistical signature, so a purely black-box gate cannot + see it, and pretending otherwise would give a case that never fires. + +### Case S — `portfolio_seq` (guards bug (b)) +* targets: `A = MixtureTarget(2, 1, ts, offset=-2.0, scale_x0=1.0)`, + `B = MixtureTarget(2, 1, ts, offset=+2.0, scale_x0=1.0)`, `ts in {101, 202, 303}`. + Mean separation 4.0 with `sigma_1d = 0.7` -> B's mass is far outside A's contracted volume, and + both stay inside the `[-5,5]^2` box (`scale_x0=1.0` keeps the random means within +-1). +* budget: `nmax = 1e5`, `neff = 2000`, `n_chunk = 10000`, one portfolio reused. +* procedure: seed from A's truth pool, integrate A, `sampler._rvs = {}`, then + `sampler.clear_warm_state()` **if present else `sampler._warm = None`** (so the case also RUNS on + a base branch that lacks the API and correctly fails there), then integrate B. The record is + point B, scored against `B.true_lnZ` and B's truth pool. +* assertions (hard; STRICT, STARVED promoted to FAIL): + * **B1** `n_eff >= 100`. Measured clear 799.8 – 2103.4 over 27 runs (>=8x margin); leak + 0.0 – 12.8. No overlap. + * **B2** `|lnZ_hat - true_lnZ_B| <= 0.10` (the gate's existing floor). Measured clear <= 0.0248 + (4x margin); leak 0.20 – 318.5 in the 12 scored runs (median 16.8 over the wider 15-run set). + * normal shape checks (measured clear: JSmax <= 0.0003, width dev <= 0.010, pull <= 0.017). +* optional companion row `portfolio_seq_nobs` — identical but with no bootstrap at all, which + isolates the run-contraction leak from the seed. Measured clear 1587 – 2094; leak 0.0 – 9.9. + Cheap (same cost) and strictly more informative; recommended. + +### Negative control (recommended, ~free) +`AV_seq`: the same sequential construction with the standalone AV sampler at d=2, which must be +unaffected. Measured leak == cold **bit-for-bit** (n_eff 2616 – 3028, |bias| 0.0075 – 0.043 in +both). Keep it warn-only: it documents that the defect is portfolio-specific. +Do **not** extend it to d=4 — AV alone on the displaced pair has |bias| 0.27 – 0.47 there and would +fail its own lnZ tolerance for reasons unrelated to this feature. + +--- + +## 4. Cases considered and REJECTED as too flaky + +**R1. "warm-from-the-correct-target beats cold" as the test for bug (a), on the AV member only.** +Seeding only the AV member does isolate the path (the inert variant reproduces cold bit-for-bit), +but the effect is not reliably positive. 15 runs per cell (3 target seeds x 5 run seeds): + +| d | cold n_eff (med, min–max) | warm_av n_eff | cold \|bias\| med | warm_av \|bias\| med | +|---|---|---|---|---| +| 2 | 1659 (817 – 2062) | 2250 (2021 – 2334) | 0.0024 | 0.0035 | +| 4 | **59.5** (19.8 – 117) | **47.6** (16.7 – 111) | 0.022 | **0.043** | +| 6 | 21.6 (3.8 – 91.7) | 35.8 (9.2 – 86.9) | 0.041 | **0.098** | + +At d=4 the warm run is *worse* on both n_eff and bias; at d=6 n_eff improves but bias degrades. +Any threshold that passes d=2 fails d=4 on some seeds. Rejected. + +**R2. Case S at d=4.** `clear` gives B n_eff 23.7 – 125.2, straddling the `STARVED` floor of 100, +so the verdict flips on the target seed alone. Worse, the discriminant collapses: at ts101 the +*leak* run gives n_eff 16.5 – 25.1 with |bias| <= 0.16, which overlaps the *clear* run at ts202 +(23.7 – 38.6, |bias| <= 0.09). Rejected. + +**R3. Case S at d=6.** `clear` gives B n_eff 6.1 – 91.6 — always STARVED even when the code is +correct — while leak gives 1.0 – 16.6 with |bias| as low as 0.058. The case would report +`BOTH-STARVED` (non-blocking) on every branch. Rejected. + +**R4. A pure lnZ-bias assertion for case S.** The leak's bias is heavy-tailed, not uniformly +large: at d=2 ts101 rs987654 the leaked run had |bias| = 0.025, inside the gate's 0.10 tolerance, +while n_eff was 2.1. `n_eff` is the reliable discriminant; bias is the corroborating one. Keep +both, but do not rely on bias alone. + +**R5. Warm-vs-cold n_eff ratio (instead of the absolute floor A3).** Cold n_eff at d6 ncomp=1 +ranges 3.8 – 91.7 across seeds, a factor 24 — a ratio threshold inherits that scatter and doubles +the runtime by requiring a paired cold run. The absolute floor (`n_eff >= 1000`, warm min 2426, +cold max 91.7) is both cheaper and tighter. + +--- + +## 5. Diff-sized changes + +### 5.1 `shape_recovery.py` + +**(i) displaced targets — the only change to `MixtureTarget` (6 lines).** `shape_recovery.py` +cannot currently express a displaced pair: `MixtureTarget` exposes `sigma_1d` and `scale_x0` but no +translation, and two different seeds give random, typically overlapping mean offsets (|mean| <= +3/sqrt(d) with sigma ~ 0.9 at d=2), so mode displacement cannot be guaranteed from seeds alone. + +```python + def __init__(self, ndim, ncomp, seed, sigma_1d=0.7, scale_x0=3.0, offset=0.0): + ... + self.offset = np.zeros(ndim) + np.asarray(offset, dtype=float) + if np.any(self.offset): + self.name += "_o{:+.2f}".format(float(np.mean(self.offset))) + ... + for k in range(ncomp): + x0 = rng.uniform(-scale_x0/np.sqrt(ndim), scale_x0/np.sqrt(ndim), ndim) + self.offset +``` + +Adding `offset` AFTER the `rng.uniform` draw keeps the RNG stream identical, so `offset=+a` and +`offset=-a` are the same mixture translated — exactly the "same shape, displaced support" +construction the case needs. `pool` and `true_lnZ` follow automatically (they are derived from +`means`/`covs`). Default `0.0` leaves every existing target and every existing `name` bit-identical. + +**(ii) new kinds + explicit case list (~90 lines, additive).** + +```python +WARM_KINDS = ("portfolio_warm", "portfolio_seq", "portfolio_seq_nobs", "AV_seq") +STARVE_IS_FAIL = ("portfolio_warm", "portfolio_seq", "portfolio_seq_nobs") +WARM_NEFF_FLOOR = 1000.0 # case W A3; measured warm 2426-5941, cold 3.8-91.7 +WARM_V_MAX = 0.9 # case W A1; measured fixed 0.031-0.129, buggy exactly 1.0 +WARM_BOX_MULT = 3.0 # case W A2; measured 6.1-22.3x at d6, buggy ~1.0x + +WARM_CASES = [ # (kind, ndim, ncomp, tseed, nmax, neff, extra) + ("portfolio_warm", 6, 1, 101, 600000, 3000, {}), + ("portfolio_warm", 6, 1, 202, 600000, 3000, {}), + ("portfolio_warm", 6, 1, 303, 600000, 3000, {}), + ("portfolio_seq", 2, 1, 101, 100000, 2000, dict(offset=2.0, scale_x0=1.0)), + ("portfolio_seq", 2, 1, 202, 100000, 2000, dict(offset=2.0, scale_x0=1.0)), + ("portfolio_seq", 2, 1, 303, 100000, 2000, dict(offset=2.0, scale_x0=1.0)), + ("portfolio_seq_nobs", 2, 1, 101, 100000, 2000, dict(offset=2.0, scale_x0=1.0)), + ("AV_seq", 2, 1, 101, 100000, 2000, dict(offset=2.0, scale_x0=1.0)), +] +``` + +* `_warm_seed_cloud(target, n=3000)` -> `target.pool[RandomState(target.seed+13).choice(...)]`. +* `run_warm_case(...)`: probe sampler (build / `setup` / `bootstrap_from_samples(cloud, + cover_frac=0.0)` / one `draw(n_chunk)`) recording `warm_V`, `warm_bins`, `warm_box_frac`, + `warm_box_frac_uniform`; then a fresh sampler, same seeding, `integrate_log`, then the existing + `shape_metrics` + record assembly. +* `run_seq_case(...)`: build `A` (offset `-o`) and `B` (offset `+o`); integrate A; + `sampler._rvs = {}`; `sampler.clear_warm_state()` if present else `sampler._warm = None`; + integrate B; record B. `AV_seq` uses `build_sampler("AV", ...)`; `portfolio_seq_nobs` skips the + bootstrap. +* `run_one` dispatches on `kind in WARM_KINDS` before its existing branch chain; both helpers keep + the never-raise contract (wrap in the same `try/except` that fills `record["error"]`). +* `main()`: `--warm-cases {auto,on,off}` (default `auto` = on for `--preset standard`, off for + `quick`), appending `WARM_CASES` jobs to `jobs` after the product expansion. + +**(iii) `evaluate()` — 3 additive blocks, no change to existing behaviour.** + +```python + if r["kind"] in STARVE_IS_FAIL and r["n_eff"] < MIN_NEFF_FOR_SHAPE: + return "FAIL", ["n_eff={:.0f} < {:.0f}: warm/sequential case must not starve" + .format(r["n_eff"], MIN_NEFF_FOR_SHAPE)] + if r["n_eff"] < MIN_NEFF_FOR_SHAPE: # unchanged + return "STARVED", [...] + ... + if r["kind"] == "portfolio_warm": + if not (r.get("warm_V", 1.0) < WARM_V_MAX and r.get("warm_bins", 1) > 1): + reasons.append("warm seed NOT installed on the draw path: AV member V={:.3f}, " + "live bins={} (cold state)".format(r.get("warm_V"), r.get("warm_bins"))) + if r.get("warm_box_frac", 0) < WARM_BOX_MULT * r.get("warm_box_frac_uniform", 1.0): + reasons.append("warm draws not concentrated in the seed box: {:.3f} < {:.1f}x{:.4f}" + .format(r["warm_box_frac"], WARM_BOX_MULT, r["warm_box_frac_uniform"])) + if r["n_eff"] < WARM_NEFF_FLOOR: + reasons.append("warm n_eff {:.0f} < {:.0f}".format(r["n_eff"], WARM_NEFF_FLOOR)) +``` + +### 5.2 `compare_shape_results.py` +No code change required (it keys on `record["kind"]`). One default change: + +```python +- ap.add_argument("--strict-samplers", default="AV,GMM") ++ ap.add_argument("--strict-samplers", ++ default="AV,GMM,portfolio_warm,portfolio_seq,portfolio_seq_nobs") +``` + +Note the intended base-vs-candidate behaviour: `portfolio_seq` FAILs on a base branch without +`clear_warm_state` and PASSes here, i.e. `IMPROVED(fail->pass)` — non-blocking, as intended. Its +value is forward-looking: once this branch is the base, any change that re-breaks the clearing +gives `REGRESSION(pass->fail)` and blocks. + +### 5.3 `run_shape_recovery.sh` +```sh +-exec python "${HERE}/shape_recovery.py" --preset standard --jobs "${SHAPE_JOBS:-8}" \ +- --json "${OUT}" "$@" ++exec "${PYTHON:-python3}" "${HERE}/shape_recovery.py" --preset standard \ ++ --jobs "${SHAPE_JOBS:-8}" --warm-cases auto --json "${OUT}" "$@" +``` +The `python` -> `python3` change is unrelated to this proposal but is a live foot-gun: this host has +no `python` on PATH, so the wrapper fails immediately. + +### 5.4 `test_shape_recovery.py` +Optional: add a second parametrisation over `WARM_CASES` so the new cases also appear under pytest. +Keep them out of the existing `_MATRIX` (which is a strict-sampler x preset product). + +--- + +## 6. Runtime + +Measured single-threaded on this host (`OMP_NUM_THREADS=1`): + +| item | cost | +|---|---| +| truth pool (10^6 draws) | 5.0 s at d=2, 2.3 s at d=6 | +| case W (probe draw + warm integrate, terminates in 1-2 chunks) | 3.0 – 5.3 s per (seed) after pool | +| case S (2 pools + 2 integrations) | 7.1 – 7.7 s per (seed) after the first | + +Full proposed set (3 W + 3 S + 1 S_nobs + 1 AV_seq = 8 cases): **~70 s of serial CPU**, ~15 s wall +at `--jobs 8`. The `standard` preset is 96 runs, many at 200k – 1.6M evaluations, so the added +cost is well under 2% of the gate. Nothing here needs a GPU; both cases are pure-CPU deterministic +in the same sense as the rest of the suite. + +--- + +## 7. Honest summary + +* Bug (b) is cheap, deterministic and worth gating: a 3-orders-of-magnitude n_eff separation with + no seed overlap in 15/15 runs at d=2. Take case S. +* Bug (a) **cannot** be caught by any statistical assertion on the gate's AV+GMM portfolio, because + the GMM member's warm start supplies nearly the whole win and the buggy configuration sometimes + scores better. The only reliable detector is the direct one: after seeding, the AV member's live + volume must actually be contracted on the draw path (A1/A2). Take that, and do not dress it up + as a statistical test. +* The obvious-looking "warm must beat cold" assertion is fit for A3 only (all-channels-inert), not + for bug (a), and is unusable in the AV-isolated form (R1). +* Case S must be pinned at d=2. At d=4 and d=6 the correct behaviour is itself starved and the + verdict becomes seed lottery (R2, R3). + +--- + +## 8. Corrections found when implementing this (2026-08-05) + +The proposal above was implemented as written and then each bug was **reintroduced** to check the +cases actually fire. Two of the proposed assertions did not survive that check. Both are corrected +in `shape_recovery.py`; this section records what was measured, not what was expected. + +**C1. `portfolio_seq` does NOT catch the leak — `portfolio_seq_nobs` does.** +With `clear_warm_state()` no-op'd, `portfolio_seq` (which re-bootstraps on point B) measured +**PASS, n_eff 5979, bias -0.008**: the fresh B seed simply overwrites the stale contracted grid, so +the leak never manifests. `portfolio_seq_nobs` with the same injection: + +| target seed | leak n_eff | leak lnZ bias | correct n_eff | correct bias | +|---|---|---|---|---| +| 101 | 9.9 | -0.559 | 2094 | +0.019 | +| 202 | 1.0 | -22.811 | 1593 | -0.000 | +| 303 | 1.0 | -59.667 | 822 | +0.002 | + +So the case list now runs `portfolio_seq_nobs` at **all three** target seeds and keeps a single +`portfolio_seq` row, whose only job is to cover the reseed-after-reset path. + +**C2. A2 measured on the portfolio mixture is not a discriminant; it must be measured on the AV +member's own draws.** With the AV install disabled, the *mixture* still concentrated **28x** in the +seed box, because the GMM member is warm-started through a separate channel. Re-measuring A2 from +`av.draw_simplified(n_chunk)` gives a clean separation: + +| | ts101 | ts202 | ts303 | +|---|---|---|---| +| installed (V, bins, box ratio) | 0.042, 656, **21.0x** | 0.032, 498, **27.2x** | 0.129, 529, **6.3x** | +| inert (V, bins, box ratio) | 1.000, 1, **1.0x** | 1.000, 1, **0.9x** | 1.000, 1, **0.9x** | +| n_eff installed / inert | 3159 / **4857** | 3368 / **5718** | 5707 / **5532** | + +The threshold `WARM_BOX_MULT = 3.0` sits between 1.0x and 6.3x. The n_eff row is the important +one: **the broken code scores HIGHER n_eff in 2 of 3 seeds**, which is the direct confirmation of +the proposal's Measurement 2 — no statistical assertion can catch this bug, only A1/A2. + +**Measured added runtime:** all 8 warm cases complete inside a 26 s wall-clock run at `--jobs 4` +(quick preset, single-threaded BLAS). diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/compare_shape_results.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/compare_shape_results.py new file mode 100755 index 000000000..1d2bf33d2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/compare_shape_results.py @@ -0,0 +1,166 @@ +#!/usr/bin/env python +""" +compare_shape_results.py BASE.json CANDIDATE.json [--strict-samplers ...] + +Compare two shape_recovery.py --json outputs (same preset/seeds!) run on a +base branch and a candidate branch. Exit 1 iff a strict-sampler run +REGRESSES: PASS on base -> FAIL on candidate, or a shape metric worsens +beyond tolerance. Pre-existing failures (FAIL on both) are reported but do +not block; improvements are celebrated. +""" +from __future__ import print_function + +import argparse +import json +import sys + +import numpy as np + +from shape_recovery import evaluate + +# metric-worsening tolerances (candidate - base), applied only when both pass +TOL_WORSE = dict(js=0.005, mean_pull=0.05, width_dev=0.05, corr=0.05, + bias_ln=0.10, neff_frac=0.5) + + +def _key(r): + return (r["kind"], r["target"]) + + +def _summ(r): + if r.get("error"): + return None + return dict(js=max(r["js"]), + mean_pull=max(abs(p) for p in r["mean_pull"]), + width_dev=max(abs(w - 1.0) for w in r["width_ratio"]), + corr=r["corr_diff_max"], + bias_ln=abs(r["bias_ln"]), + neff=r["n_eff"]) + + +def classify(b, c): + """Return (verdict, note) for one base/candidate record pair. + + SINGLE SOURCE OF TRUTH for what counts as a regression. confirm_regressions.py + imports this: it previously reimplemented only the PASS->non-PASS case and was blind to + REGRESSION(metrics), so a real metric regression (measured: n_eff 448->210) produced + "no blocking regressions to confirm" and exited 0. Two copies of this logic will always + drift; there is now one. + """ + if c is None and b is not None: + # The candidate produced NO record for a row the base did. That is a regression, not a + # bookkeeping curiosity: a candidate that crashes before emitting a result would otherwise + # be classified ONLY-IN-BASE, never reach confirmation, and exit the gate successfully -- + # bypassing the fail-closed rerun logic entirely. + return "REGRESSION(missing-in-candidate)", "candidate produced no record for this row" + if b is None: + return "ONLY-IN-CANDIDATE", "" + st_b, _ = evaluate(b) + st_c, why_c = evaluate(c) + sb, sc = _summ(b), _summ(c) + verdict, note = "OK", "" + if st_b == "PASS" and st_c != "PASS": + # includes healthy->STARVED: candidate lost the efficiency the + # base had on this target -> regression + verdict = "REGRESSION(pass->{})".format(st_c.lower()) + note = "; ".join(why_c) + elif st_b != "PASS" and st_c == "PASS": + verdict = "IMPROVED({}->pass)".format(st_b.lower()) + elif st_b == "STARVED" and st_c == "STARVED": + verdict = "BOTH-STARVED" + elif st_b == "STARVED" and st_c in ("FAIL", "ERROR"): + # base gave no shape information here; candidate at least reaches + # testability (or crashes) -- flag, don't block + verdict = "NEWLY-TESTABLE-" + st_c + note = "; ".join(why_c) + elif st_b in ("FAIL", "ERROR") and st_c != "PASS": + verdict = "PREEXISTING-FAIL" + elif sb and sc: + worse = [] + for m, tol in TOL_WORSE.items(): + if m == "neff_frac": + if sc["neff"] < TOL_WORSE["neff_frac"] * sb["neff"]: + worse.append("n_eff {:.0f}->{:.0f}".format(sb["neff"], sc["neff"])) + elif sc[m] - sb[m] > tol: + worse.append("{} {:.3f}->{:.3f}".format(m, sb[m], sc[m])) + if worse: + verdict = "REGRESSION(metrics)" + note = "; ".join(worse) + return verdict, note + + +def is_blocking(verdict, kind, strict): + return verdict.startswith("REGRESSION") and kind in strict + + +def blocking_keys(base, cand, strict): + """Every (kind, target) the gate would BLOCK on -- both regression flavours.""" + out = [] + for k in sorted(set(base) | set(cand)): + v, _ = classify(base.get(k), cand.get(k)) + if is_blocking(v, k[0], strict): + out.append(k) + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("base") + ap.add_argument("candidate") + # The warm/sequential kinds are STRICT: they exist to catch silent wrong answers, so a + # regression there must block. Note the intended asymmetry on first merge -- portfolio_seq + # FAILs on a base without clear_warm_state and PASSes here, i.e. IMPROVED (non-blocking). + # Its value is forward-looking: once this is the base, re-breaking the reset blocks. + # ENFORCEMENT. Given both checkouts, a blocking regression is re-tested at fresh seeds + # before it is allowed to fail the gate, and THIS script's exit code reflects the confirmed + # verdict. Without these the script only reports, and confirmation is advisory -- which is + # how the first version shipped: documented in the runner but never actually invoked. + ap.add_argument("--confirm-base-checkout", default=None, + help="with --confirm-cand-checkout: re-test blocking rows at fresh seeds") + ap.add_argument("--confirm-cand-checkout", default=None) + ap.add_argument("--confirm-repeats", type=int, default=5) + ap.add_argument("--confirm-jobs", type=int, default=4) + ap.add_argument("--strict-samplers", + default="AV,GMM,portfolio_warm,portfolio_seq,portfolio_seq_nobs") + opts = ap.parse_args() + strict = set(x.strip() for x in opts.strict_samplers.split(",")) + + with open(opts.base) as fh: + base = {_key(r): r for r in json.load(fh)} + with open(opts.candidate) as fh: + cand = {_key(r): r for r in json.load(fh)} + + n_block = 0 + rows = [] + for k in sorted(set(base) | set(cand)): + verdict, note = classify(base.get(k), cand.get(k)) + blocking = is_blocking(verdict, k[0], strict) + if blocking: + n_block += 1 + rows.append((k, verdict + (" <-- BLOCKS MERGE" if blocking else ""), note)) + + for (kind, tgt), verdict, note in rows: + print("{:<10s} {:<16s} {} {}".format(kind, tgt, verdict, + ("[" + note + "]") if note else "")) + print("# blocking regressions (strict={}): {}".format(sorted(strict), n_block)) + if not n_block: + return 0 + if not (opts.confirm_base_checkout and opts.confirm_cand_checkout): + print("# NOT CONFIRMED AT FRESH SEEDS: pass --confirm-base-checkout/--confirm-cand-checkout\n" + "# to re-test these rows before treating them as real. Every threshold here is a\n" + "# hard cut on a stochastic quantity, so a single blocking row is a hypothesis.") + return 1 + import confirm_regressions + print("\n# re-testing {} blocking row(s) at {} fresh seeds per arm".format( + n_block, opts.confirm_repeats)) + return confirm_regressions.main([ + opts.base, opts.candidate, + "--base-checkout", opts.confirm_base_checkout, + "--cand-checkout", opts.confirm_cand_checkout, + "--repeats", str(opts.confirm_repeats), + "--jobs", str(opts.confirm_jobs), + "--strict-samplers", opts.strict_samplers]) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/confirm_regressions.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/confirm_regressions.py new file mode 100644 index 000000000..15224582e --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/confirm_regressions.py @@ -0,0 +1,181 @@ +#!/usr/bin/env python +"""Re-test the blocking regressions a merge-gate comparison reported, at NEW random seeds. + +WHY THIS EXISTS. Every gate verdict is a hard threshold (n_eff >= 100, JS < 3*floor + 0.004, ...) +applied to a stochastic quantity, so a cell sitting near a threshold flips on realization alone. +Observed: `GMM mix_d6_n3_s303` read n_eff 66 / 119 / 104 across runs of the SAME unchanged +checkout -- straddling the 100 floor -- purely from where it landed in the worker pool. Reported +as a REGRESSION once, it would have blocked a merge that changed nothing about that sampler. + +The fix is NOT to make the samplers deterministic. Independent copies that localize differently +are our main detector for support/mode-collapse failures; pinning every fit to one seed would +silence it, and would make N copies of a production run no better than one. The fix is to ask the +question again, properly: re-run the disputed cell in BOTH arms at several fresh run seeds and see +whether the candidate is really worse. + +Usage: + confirm_regressions.py base.json cand.json --base-checkout DIR --cand-checkout DIR \\ + [--repeats 3] [--jobs 4] [--seeds 11,22,33] + +Exit 0 if no regression is CONFIRMED; 1 if any is. A regression is confirmed when the candidate +is worse than the base in a MAJORITY of the fresh seeds (ties count as not-worse: the burden of +proof is on the claim that the candidate broke something). +""" +import argparse +import json +import os +import subprocess +import sys +import tempfile + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +# The comparator OWNS the definition of "blocking". Importing it -- rather than reimplementing +# the PASS->non-PASS case, as this script first did -- is what keeps the two in step: the local +# copy was blind to REGRESSION(metrics), so a real metric regression (measured: n_eff 448->210) +# reported "no blocking regressions to confirm" and exited 0. +from compare_shape_results import classify, is_blocking, blocking_keys # noqa: E402 + + +def _key(r): + return (r["kind"], r["target"]) + + +def _blocking(base_path, cand_path, strict): + with open(base_path) as fh: + base = {_key(r): r for r in json.load(fh)} + with open(cand_path) as fh: + cand = {_key(r): r for r in json.load(fh)} + return [(k, base.get(k), cand.get(k)) for k in blocking_keys(base, cand, strict)] + + +def _rerun(checkout, rec, seed, jobs, tag): + """Re-run ONE cell of the matrix at a given run seed; return its record or None.""" + fd, path = tempfile.mkstemp(suffix=".json", prefix="confirm_%s_" % tag) + os.close(fd) + cmd = [os.environ.get("PYTHON", "python3"), os.path.join(HERE, "shape_recovery.py"), + "--preset", "standard", "--json", path, "--jobs", str(jobs), + "--samplers", rec["kind"] if not rec["kind"].startswith(("portfolio_warm", + "portfolio_seq", "AV_seq")) + else "AV", + "--dims", str(rec["ndim"]), "--ncomps", str(rec["ncomp"]), + "--target-seeds", str(rec["target_seed"]), "--run-seed", str(seed), + "--warm-cases", "on" if rec["kind"] in ("portfolio_warm", "portfolio_seq", + "portfolio_seq_nobs", "AV_seq") else "off"] + env = dict(os.environ) + env["PYTHONPATH"] = os.path.join(checkout, "MonteCarloMarginalizeCode", "Code") + \ + os.pathsep + env.get("PYTHONPATH", "") + env["CUDA_VISIBLE_DEVICES"] = "" + try: + subprocess.run(cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + check=False) + with open(path) as fh: + for r in json.load(fh): + if _key(r) == _key(rec): + return r + except Exception: + return None + finally: + try: + os.unlink(path) + except OSError: + pass + return None + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[1]) + ap.add_argument("base") + ap.add_argument("candidate") + ap.add_argument("--base-checkout", required=True) + ap.add_argument("--cand-checkout", required=True) + ap.add_argument("--repeats", type=int, default=3, + help="fresh run seeds per arm (default 3; use more for a near-threshold cell)") + ap.add_argument("--seeds", default=None, help="explicit comma list, overrides --repeats") + ap.add_argument("--jobs", type=int, default=4) + ap.add_argument("--min-valid", type=int, default=None, + help="usable base/candidate pairs required for a verdict (default: all " + "seeds). Fewer -> INCONCLUSIVE and exit 1, never a silent clear.") + ap.add_argument("--strict-samplers", + default="AV,GMM,portfolio_warm,portfolio_seq,portfolio_seq_nobs") + opts = ap.parse_args(argv) + + strict = set(x.strip() for x in opts.strict_samplers.split(",") if x.strip()) + seeds = ([int(x) for x in opts.seeds.split(",")] if opts.seeds + else [987654 + 1000 * (i + 1) for i in range(opts.repeats)]) + if opts.min_valid is None: + opts.min_valid = len(seeds) + + disputed = _blocking(opts.base, opts.candidate, strict) + if not disputed: + print("# no blocking regressions to confirm") + return 0 + print("# confirming {} blocking regression(s) at {} fresh seed(s): {}".format( + len(disputed), len(seeds), seeds)) + + n_confirmed = 0 + n_inconclusive = 0 + for k, brec, crec in disputed: + worse = same = 0 + detail = [] + for s in seeds: + # The cell to re-run is defined by whichever record exists -- for a + # REGRESSION(missing-in-candidate) row the candidate has no record, but the base + # record still tells us which (kind, dim, ncomp, seed) to run, so the candidate CAN + # and must be re-tested rather than written off. + spec = brec if brec is not None else crec + rb = _rerun(opts.base_checkout, spec, s, opts.jobs, "base") + rc = _rerun(opts.cand_checkout, spec, s, opts.jobs, "cand") + if rc is None and rb is None: + detail.append("seed {}: BOTH reruns produced no record (no evidence either way)" + .format(s)) + continue + if rc is None: + # The CANDIDATE failed where the base did not. That is not missing evidence, it + # IS the regression: crashing or emitting no record is worse than passing. + # Discarding it -- as this script first did -- let a candidate that failed on + # every seed be declared "not confirmed". + worse += 1 + detail.append("seed {}: CANDIDATE PRODUCED NO RECORD (counts against candidate)" + .format(s)) + continue + if rb is None: + detail.append("seed {}: base rerun produced no record; pair unusable".format(s)) + continue + # the SAME classifier the gate uses, so a metrics-only regression is judged here + # exactly as it was there + verdict, note = classify(rb, rc) + if is_blocking(verdict, k[0], strict): + worse += 1 + else: + same += 1 + detail.append("seed {}: {} (n_eff {:.0f} vs {:.0f}){}".format( + s, verdict, rb.get("n_eff", float("nan")), rc.get("n_eff", float("nan")), + " [" + note + "]" if note else "")) + + valid = worse + same + if valid < opts.min_valid: + status = ("INCONCLUSIVE -- {}/{} valid pairs, need {}: NOT cleared" + .format(valid, len(seeds), opts.min_valid)) + n_inconclusive += 1 + elif worse > same: + status = "CONFIRMED REGRESSION -- BLOCKS ({} worse / {} not-worse)".format(worse, same) + n_confirmed += 1 + else: + status = ("NOT CONFIRMED (realization noise; does not block) ({} worse / {} not-worse)" + .format(worse, same)) + print("\n{} {}".format(k[0], k[1])) + for d in detail: + print(" " + d) + print(" -> " + status) + + print("\n# confirmed blocking regressions: {}".format(n_confirmed)) + if n_inconclusive: + print("# INCONCLUSIVE rows (too few valid reruns): {}".format(n_inconclusive)) + # Inconclusive must NOT read as success: we failed to obtain the evidence that would clear + # the row, so the gate stays red until a human looks. + return 1 if (n_confirmed or n_inconclusive) else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_report.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_report.py new file mode 100644 index 000000000..fa8845e7b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_report.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python +"""escaped_mass_report.py -- tabulate / ROC the output of escaped_mass_study.py. + + python escaped_mass_report.py run1.json [run2.json ...] + +Candidate statistics compared (all for the WARM-STARTED member, index 0): + esc_cum cumulative escaped_mass over the whole run + esc_early escaped_mass in the FIRST chunk the member had a density in + esc_first3 max over the first 3 chunks + lo_share 1 - weight_share (the cheap comparator: a useless warm member stops + attracting weight, so a LOW share is the alarm) + inv_neff 1000/n_eff (null comparator: an efficiency signal, not a support signal) + +Two questions are scored separately, because they have different answers: + * SEED-MISMATCH ROC -- can the statistic tell offset>0 from offset=0? + * BIAS ROC -- can it tell |lnZ bias| > BIAS_MATERIAL from |bias| <= BIAS_MATERIAL? + This is the one that matters operationally: a detector that fires only on runs that were + going to be fine is a nuisance alarm, and one that stays quiet on biased runs is worthless. +""" +from __future__ import print_function + +import json +import sys + +import numpy as np + +BIAS_MATERIAL = 0.5 # nats; the scale at which a lnZ error starts to matter downstream + + +def _firstk(r, k, m=0): + h = r.get("esc_hist") or [] + col = [row[m] for row in h[:k] if len(row) > m and np.isfinite(row[m])] + return float(np.max(col)) if col else np.nan + + +STATS = { + "esc_cum": lambda r: r["esc_warm"], + "esc_early": lambda r: r["esc_early_warm"], + "esc_first3": lambda r: _firstk(r, 3), + "lo_share": lambda r: 1.0 - r["share_warm"], + "inv_neff": lambda r: 1000.0 / max(r["n_eff"], 1e-9), +} + + +def auc(pos, neg): + """Mann-Whitney AUC: P(stat_pos > stat_neg), ties counted as 1/2.""" + pos = np.asarray([x for x in pos if np.isfinite(x)], dtype=float) + neg = np.asarray([x for x in neg if np.isfinite(x)], dtype=float) + if not len(pos) or not len(neg): + return np.nan + gt = np.sum(pos[:, None] > neg[None, :]) + eq = np.sum(pos[:, None] == neg[None, :]) + return float((gt + 0.5 * eq) / (len(pos) * len(neg))) + + +def qs(a): + a = np.asarray([x for x in a if np.isfinite(x)], dtype=float) + if not len(a): + return (np.nan,) * 4 + return (float(np.median(a)), float(np.percentile(a, 10)), + float(np.percentile(a, 90)), float(np.max(a))) + + +def main(paths): + rs = [] + for p in paths: + rs += json.load(open(p)) + n_err = sum(1 for r in rs if r.get("error")) + rs = [r for r in rs if not r.get("error")] + dims = sorted(set(r["ndim"] for r in rs)) + arms = sorted(set(r["arm"] for r in rs)) + offs = sorted(set(r["offset"] for r in rs)) + print("# {} records ({} errors dropped); dims {} arms {} offsets {}".format( + len(rs), n_err, dims, arms, offs)) + + # ---------------- 1. sensitivity table ---------------- + print("\n== SENSITIVITY: statistic and |lnZ bias| vs seed displacement ==") + hdr = ("%-11s %2s %5s %3s | %8s %8s %8s | %9s %9s | %9s %9s | %9s | %7s" % + ("arm", "d", "off", "N", "med bias", "|b|max", "med neff", + "esc_cum", "(p10)", "esc_early", "(p10)", "esc_f3", "share0")) + print(hdr); print("-" * len(hdr)) + for d in dims: + for arm in arms: + for off in offs: + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm and r["offset"] == off] + if not sel: + continue + b = np.asarray([r["bias_ln"] for r in sel]) + print("%-11s %2d %5.1f %3d | %8.3f %8.2f %8.0f | %9.2e %9.2e | %9.2e %9.2e | %9.2e | %7.3f" % ( + arm, d, off, len(sel), np.median(b), np.max(np.abs(b)), + np.median([r["n_eff"] for r in sel]), + qs([r["esc_warm"] for r in sel])[0], qs([r["esc_warm"] for r in sel])[1], + qs([r["esc_early_warm"] for r in sel])[0], qs([r["esc_early_warm"] for r in sel])[1], + qs([_firstk(r, 3) for r in sel])[0], + np.median([r["share_warm"] for r in sel]))) + print() + + # ---------------- 2. false-positive floor at offset 0 ---------------- + print("== FALSE-POSITIVE FLOOR at offset=0 (across independent target seeds) ==") + hdr = "%-11s %2s %3s | %-10s %10s %10s %10s %10s" % ( + "arm", "d", "N", "stat", "median", "p90", "max", "frac>1e-3") + print(hdr); print("-" * len(hdr)) + floors = {} + for d in dims: + for arm in arms: + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm and r["offset"] == 0.0] + if not sel: + continue + for name in ("esc_cum", "esc_early", "esc_first3", "lo_share"): + v = np.asarray([STATS[name](r) for r in sel], dtype=float) + v = v[np.isfinite(v)] + m, p10, p90, mx = qs(v) + floors[(d, arm, name)] = mx + print("%-11s %2d %3d | %-10s %10.3e %10.3e %10.3e %10.2f" % ( + arm, d, len(v), name, m, p90, mx, float(np.mean(v > 1e-3)))) + print() + + # ---------------- 3. seed-mismatch ROC (AUC vs the offset=0 control) ---------------- + print("== SEED-MISMATCH AUC: P(stat[offset] > stat[offset=0]), 0.5 = useless ==") + names = ["esc_cum", "esc_early", "esc_first3", "lo_share", "inv_neff"] + hdr = "%-11s %2s %5s | " % ("arm", "d", "off") + " ".join("%10s" % n for n in names) + print(hdr); print("-" * len(hdr)) + for d in dims: + for arm in arms: + ctl = [r for r in rs if r["ndim"] == d and r["arm"] == arm and r["offset"] == 0.0] + for off in offs: + if off == 0.0: + continue + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm and r["offset"] == off] + if not sel or not ctl: + continue + row = [auc([STATS[n](r) for r in sel], [STATS[n](r) for r in ctl]) for n in names] + print("%-11s %2d %5.1f | " % (arm, d, off) + + " ".join("%10.3f" % x for x in row)) + print() + + # ---------------- 4. BIAS ROC: does it flag the runs that are actually WRONG? ---------- + print("== BIAS DETECTION: positives = |lnZ bias| > {} nat, pooled over offsets ==".format( + BIAS_MATERIAL)) + hdr = "%-11s %2s | %5s %5s | " % ("arm", "d", "Npos", "Nneg") + " ".join("%10s" % n for n in names) + print(hdr); print("-" * len(hdr)) + for d in dims: + for arm in arms: + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm] + pos = [r for r in sel if abs(r["bias_ln"]) > BIAS_MATERIAL] + neg = [r for r in sel if abs(r["bias_ln"]) <= BIAS_MATERIAL] + if not pos or not neg: + print("%-11s %2d | %5d %5d | (degenerate: one class empty)" % ( + arm, d, len(pos), len(neg))) + continue + row = [auc([STATS[n](r) for r in pos], [STATS[n](r) for r in neg]) for n in names] + print("%-11s %2d | %5d %5d | " % (arm, d, len(pos), len(neg)) + + " ".join("%10.3f" % x for x in row)) + print() + + # ---------------- 5. operating point: threshold = the measured offset=0 max ---------- + print("== OPERATING POINT: threshold = max over the 20 offset=0 controls of the SAME cell ==") + hdr = "%-11s %2s %-10s | %10s | " % ("arm", "d", "stat", "thresh") + \ + " ".join("%6.1f" % o for o in offs if o > 0) + print(hdr); print("-" * len(hdr)) + for d in dims: + for arm in arms: + for name in ("esc_cum", "esc_early", "esc_first3", "lo_share"): + thr = floors.get((d, arm, name)) + if thr is None or not np.isfinite(thr): + continue + cells = [] + for off in offs: + if off == 0.0: + continue + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm + and r["offset"] == off] + v = np.asarray([STATS[name](r) for r in sel], dtype=float) + v = v[np.isfinite(v)] + cells.append(float(np.mean(v > thr)) if len(v) else np.nan) + print("%-11s %2d %-10s | %10.3e | " % (arm, d, name, thr) + + " ".join("%6.2f" % c for c in cells)) + print() + print("# (numbers are TRUE-POSITIVE RATE at a threshold with 0/20 false positives by " + "construction; 1-sided 95% upper bound on that FP rate is ~0.14)") + + # ---------------- 6. BIAS operating point: the question that actually matters ---------- + # Threshold set on the BENIGN runs (|bias| <= BIAS_MATERIAL) at a 10% false-alarm rate; then + # report what fraction of the genuinely-wrong runs it catches, and how wrong they were. + print("\n== BIAS OPERATING POINT: threshold = p90 of the statistic over runs with " + "|bias| <= {} nat ==".format(BIAS_MATERIAL)) + hdr = ("%-11s %2s %-10s | %10s | %5s %5s | %6s | %10s %10s" % + ("arm", "d", "stat", "thr(FP=.10)", "Npos", "Nneg", "TPR", "med|b| hit", "med|b| miss")) + print(hdr); print("-" * len(hdr)) + for d in dims: + for arm in arms: + sel = [r for r in rs if r["ndim"] == d and r["arm"] == arm] + pos = [r for r in sel if abs(r["bias_ln"]) > BIAS_MATERIAL] + neg = [r for r in sel if abs(r["bias_ln"]) <= BIAS_MATERIAL] + if len(pos) < 5 or len(neg) < 5: + print("%-11s %2d %-10s | (too few in one class: %d pos / %d neg)" % ( + arm, d, "-", len(pos), len(neg))) + continue + for name in ("esc_cum", "esc_early", "esc_first3", "lo_share", "inv_neff"): + vneg = np.asarray([STATS[name](r) for r in neg], dtype=float) + vneg = vneg[np.isfinite(vneg)] + thr = float(np.percentile(vneg, 90)) + vpos = np.asarray([STATS[name](r) for r in pos], dtype=float) + bpos = np.abs([r["bias_ln"] for r in pos]) + hit = vpos > thr + print("%-11s %2d %-10s | %11.3e | %5d %5d | %6.2f | %10.2f %10.2f" % ( + arm, d, name, thr, len(pos), len(neg), float(np.mean(hit)), + float(np.median(bpos[hit])) if np.any(hit) else float("nan"), + float(np.median(bpos[~hit])) if np.any(~hit) else float("nan"))) + print() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_study.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_study.py new file mode 100644 index 000000000..69da80131 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/escaped_mass_study.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python +"""escaped_mass_study.py -- ROC / sensitivity study for the portfolio's warm-start +SUPPORT-MISMATCH detector (mcsamplerPortfolio.support_diagnostics). + +QUESTION. A RIFT extrinsic point is often warm-started from a proposal built at a DIFFERENT +point (a neighbouring intrinsic grid point, a stale breadcrumb, a recovered posterior). If the +seed is misplaced, an AV/VARAHA member's live volume -- whose density is EXACTLY ZERO outside it +-- can exclude the true peak. n_eff and the Pareto k-hat cannot see this: both are functions +only of the weights actually drawn (k-hat has been measured at 0.435 on a -1949-nat run). The +proposed detector is + + escaped_mass[m] = sum_{i : q_m(x_i)==0} w_i / sum_i w_i + +-- the fraction of total posterior weight carried by samples member m could not have drawn -- and +its cheaper comparator, weight_share[m] (fraction of weight from samples m DREW). + +METHOD. Truth-known testbed (shape_recovery.MixtureTarget). Warm-start from the DISPLACED +target's own truth pool while integrating the TRUE (offset=0) target; sweep the displacement. +Three portfolio arms, because the answer depends entirely on what ELSE is in the mixture: + avgmm_cold : [AV warm-seeded, GMM COLD] -- an independent broad member + avgmm_warm : [AV warm-seeded, GMM warm-seeded from the SAME displaced cloud] -- partly blind + avav : [AV warm-seeded, AV warm-seeded] -- no soft component at all + +Usage (CPU, deterministic): + export PYTHONPATH=/MonteCarloMarginalizeCode/Code + export CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 + python escaped_mass_study.py --dims 4,6 --offsets 0,0.5,1,1.5,2,2.5,3,4 --copies 20 \ + --arms avgmm_cold,avgmm_warm,avav --jobs 8 --json out.json + python escaped_mass_report.py out.json + +MEASURED VERDICT (960 early-stopping runs + 480 fixed-budget runs; 20 independent target seeds +per cell, d=4 and d=6, ncomp=1, nmax 120000, n_chunk 5000). Read this before using the number. + + 1. THE CUMULATIVE STATISTIC HAS NO USABLE FLOOR. The hypothesis "matched seed -> escaped_mass + ~0" is FALSE. With a PERFECT seed (offset=0) the cumulative escaped mass of the warm AV + member is median 0.51 at d=4 and 0.80 at d=6 (20 seeds; max 0.84 / 0.88). A correctly-placed + VARAHA member contracts to a likelihood contour that legitimately excludes half to four-fifths + of the posterior WEIGHT, and the statistic cannot tell that from a misplaced seed: AUC vs the + offset=0 control is only 0.48-0.63 out to offset 2. + + 2. THE FIRST-CHUNK STATISTIC IS SHARP -- IN ONE ARM ONLY. Measured at the seed's own live + volume, before any contraction, the offset=0 floor is 6e-6 (max 2.3e-4) at d=4 and 2.5e-7 + (max 6.8e-3) at d=6, four to six decades below the signal. In the avgmm_cold arm, at a + threshold equal to the measured offset=0 MAXIMUM (0/20 false positives): + d=4 TPR 0.15 / 0.70 / 0.95 / 1.00 / 1.00 / 1.00 / 1.00 at offset 0.5/1/1.5/2/2.5/3/4 + d=6 TPR 0.00 / 0.10 / 0.50 / 0.80 / 0.75 / 1.00 / 1.00 + The d=6 degradation is the predicted STARVATION false negative: a cold uniform member has to + land on the true peak for the escape to be observable at all. + It beats the soft comparator decisively: 1-weight_share at the same 0-FP threshold reaches + only TPR 0.20 (d=4) / 0.05 (d=6) at offset 2, where escaped_mass_early is at 1.00 / 0.80. + + 3. IT IS BLIND IN BOTH OTHER ARMS -- EXACTLY 0.000 IN 320/320 RUNS. If every member is seeded + from the same displaced cloud (avgmm_warm -- which is what portfolio.bootstrap_from_samples + does by DEFAULT, since mcsamplerEnsemble also implements bootstrap_from_samples), or if the + portfolio is all-AV, then in the first chunk NOTHING is drawn outside the seeded volume and + there is no escaping weight to measure. AUC 0.500 / 0.45 at every offset, both dims. + + 4. THE APPARENT PERFECT DETECTOR IN avgmm_warm IS A RUN-LENGTH ARTIFACT. Under production-style + early stopping the cumulative statistic separates offset 0 from every offset >= 0.5 at + AUC 1.000 -- because the matched run reaches neff in ~3 chunks and its AV never contracts, + while a mismatched run burns 24. With the budget FIXED (neff disabled) the offset=0 floor + moves from 4.9e-5 to 0.524 (d=4) and 1.5e-4 to 0.783 (d=6) and the AUC collapses to 0.47/0.58 + at offset 1/2. 1000/n_eff scores AUC 1.000 in the same cells: the statistic was measuring + n_eff, not support. ALWAYS compare at matched budget. + + 5. STRUCTURAL LIMIT (the reason 2 and 3 cannot both be fixed). escaped_mass is a reduction over + samples that were DRAWN. Weight can only be seen escaping member m if some OTHER member + covers the complement of m's support -- which is precisely the configuration in which the + balance heuristic already keeps lnZ unbiased (measured: |bias| <= 0.31 over all 160 avgmm_cold + runs at d=4 and <= 1.29 at d=6, at every displacement out to 4). In the arms where + displacement DOES bias lnZ (all-AV: median -0.40 nat at d=4 / -1.27 at d=6 even at offset 0, + reaching -22.6 median and -26000 worst at offset 4) no variant reaches a usable threshold: + esc_cum has TPR 0.85 (d=4) / 0.60 (d=6) at offset 3, where the median bias is already -4.2 / + -2.0 nats, and at d=6 the runs it MISSES are the worse ones (median |bias| 2.27 vs 1.42 for + the hits). The detector fires when it does not matter and is quiet when it does. +""" +from __future__ import print_function + +import argparse +import json +import os +import sys +import time + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import shape_recovery as sr # noqa: E402 + + +# The truth pool only supplies (a) the in-box mass for true_lnZ and (b) the warm seed cloud. +# 2e5 draws give ~0.2% on the box mass = 0.002 nat, far below the 0.05-nat bias scale we resolve, +# and each study run builds TWO pools (true + displaced), so this is the dominant fixed cost. +DEFAULT_POOL_N = 200000 + + +def _build_portfolio(arm, target, n_chunk): + """Assemble the portfolio for one arm and return (sampler, members).""" + if not sr._gpu_available(): + sr._force_cpu_modules() + from RIFT.integrators import (mcsamplerPortfolio, mcsamplerAdaptiveVolume, + mcsamplerEnsemble) + + def _av(): + try: + return mcsamplerAdaptiveVolume.MCSampler(n_chunk=n_chunk) + except TypeError: + return mcsamplerAdaptiveVolume.MCSampler() + + if arm == "avav": + members = [_av(), _av()] + else: + members = [_av(), mcsamplerEnsemble.MCSampler()] + if not sr._gpu_available(): + for m in members: + sr._force_cpu(m) + s = mcsamplerPortfolio.MCSampler(portfolio=list(members)) + if not sr._gpu_available(): + sr._force_cpu(s) + + def uniform_pdf(d): + wdt = target.rlim[d] - target.llim[d] + return np.vectorize(lambda x, wdt=wdt: 1.0 / wdt) + + for d, p in enumerate(target.params): + s.add_parameter(p, uniform_pdf(d), prior_pdf=uniform_pdf(d), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + return s, members + + +def run_case(ndim, ncomp, target_seed, offset, arm, run_seed, + nmax, neff, n_chunk, ncomp_gmm=2, verbose=False): + """One (target, displacement, arm) run. Never raises: errors are recorded.""" + t0 = time.time() + out = dict(ndim=ndim, ncomp=ncomp, target_seed=target_seed, offset=float(offset), + arm=arm, run_seed=run_seed, nmax=int(nmax), n_chunk=int(n_chunk)) + try: + np.random.seed(run_seed) + true_t = sr.MixtureTarget(ndim, ncomp, target_seed) # what we integrate + seed_t = sr.MixtureTarget(ndim, ncomp, target_seed, offset=offset) # where the seed came from + cloud = sr._warm_seed_cloud(seed_t) + + s, members = _build_portfolio(arm, true_t, n_chunk) + _dims = tuple(range(ndim)) + setup_kw = {} + if arm != "avav": + # PRODUCTION SHAPE: a real run always supplies a grouping spec, and that is the + # configuration in which the GMM member is a genuine full-dim mixture. + setup_kw = dict(n_comp={_dims: ncomp_gmm}, gmm_dict={_dims: None}, + correlate_all_dims=True) + try: + s.setup(**setup_kw) + except TypeError: + s.setup() + + # SEEDING. avgmm_cold seeds ONLY the AV member; the other two arms seed every member that + # exposes bootstrap_from_samples (which is what portfolio.bootstrap_from_samples does, and + # what production therefore does by default). + if arm == "avgmm_cold": + members[0].bootstrap_from_samples(cloud, cover_frac=0.0) + out["n_warmed"] = 1 + else: + out["n_warmed"] = int(s.bootstrap_from_samples(cloud, cover_frac=0.0)) + + extra = dict(n=n_chunk, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=neff, nmax=int(nmax), save_intg=True, verbose=verbose) + lnI, logvar, eff, dret = s.integrate_log(true_t.as_lnfunc(), *true_t.params, + no_protect_names=True, **extra) + lnI = float(sr._asnumpy(lnI)) + logvar = float(sr._asnumpy(logvar)) + ln_wt = sr.log_weights_from_rvs(s._rvs) + + esc = np.asarray(dret.get("portfolio_escaped_mass", []), dtype=float) + early = np.asarray(dret.get("portfolio_escaped_mass_early", []), dtype=float) + share = np.asarray(dret.get("portfolio_member_weight_share", []), dtype=float) + hard = np.asarray(dret.get("portfolio_member_hard_edged", []), dtype=bool) + hist = np.asarray(dret.get("portfolio_escaped_mass_history", []), dtype=float) + out.update( + lnI=lnI, true_lnZ=float(true_t.true_lnZ), bias_ln=lnI - float(true_t.true_lnZ), + n_eff=float(sr._asnumpy(eff)), n_ess=float(sr.n_ess_kish(ln_wt)), + n_eval=int(getattr(s, "ntotal", 0)), + rel_err=float(np.exp(0.5 * logvar - lnI)) if np.isfinite(logvar) else float("nan"), + escaped_mass=[float(x) for x in esc], + escaped_mass_early=[float(x) for x in early], + weight_share=[float(x) for x in share], + hard_edged=[bool(x) for x in hard], + # HEADLINE STATISTICS, as a monitor would read them: + # esc_warm -- the warm-started member (index 0), the one under test; + # esc_max -- worst over members observed to be hard-edged (what production reads, + # since it does not know which member was warm-started); + # share_warm-- the soft comparator for the same member. + esc_warm=float(esc[0]) if len(esc) else float("nan"), + esc_early_warm=float(early[0]) if len(early) else float("nan"), + esc_max=float(dret.get("portfolio_escaped_mass_max", np.nan)), + esc_early_max=float(dret.get("portfolio_escaped_mass_early_max", np.nan)), + share_warm=float(share[0]) if len(share) else float("nan"), + n_chunks=int(hist.shape[0]) if hist.ndim == 2 else 0, + # full per-chunk history (n_chunks x n_members) so any "first K chunks" variant of the + # statistic can be evaluated post hoc without re-running: the cumulative and the + # first-chunk numbers are two points on this curve, and which one is the detector is + # exactly what the study has to decide. + esc_hist=hist.tolist() if hist.ndim == 2 else [], + wallclock=time.time() - t0, error=None) + except Exception as e: + import traceback + out.update(error="{}: {}".format(type(e).__name__, e), + traceback=traceback.format_exc(), wallclock=time.time() - t0) + return out + + +def _worker(job): + sr.TRUTH_POOL_N = job.pop("pool_n", DEFAULT_POOL_N) + return run_case(**job) + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[1]) + ap.add_argument("--dims", default="4,6") + ap.add_argument("--ncomp", type=int, default=1) + ap.add_argument("--ncomp-gmm", type=int, default=2) + ap.add_argument("--offsets", default="0,0.5,1,1.5,2,3") + ap.add_argument("--arms", default="avgmm_cold,avgmm_warm,avav") + ap.add_argument("--copies", type=int, default=20, + help="independent (target seed, run seed) pairs per cell") + ap.add_argument("--seed0", type=int, default=1000) + ap.add_argument("--nmax", type=int, default=120000) + ap.add_argument("--neff", type=int, default=3000) + ap.add_argument("--n-chunk", type=int, default=5000) + ap.add_argument("--pool-n", type=int, default=DEFAULT_POOL_N) + ap.add_argument("--jobs", type=int, default=1) + ap.add_argument("--json", default=None) + ap.add_argument("--verbose", action="store_true") + opts = ap.parse_args(argv) + + dims = [int(x) for x in opts.dims.split(",") if x.strip()] + offsets = [float(x) for x in opts.offsets.split(",") if x.strip()] + arms = [x.strip() for x in opts.arms.split(",") if x.strip()] + + jobs = [] + for d in dims: + for off in offsets: + for arm in arms: + for c in range(opts.copies): + ts = opts.seed0 + c + jobs.append(dict(ndim=d, ncomp=opts.ncomp, target_seed=ts, offset=off, + arm=arm, run_seed=900000 + 37 * ts + 11 * d, + nmax=opts.nmax, neff=opts.neff, n_chunk=opts.n_chunk, + ncomp_gmm=opts.ncomp_gmm, verbose=opts.verbose, + pool_n=opts.pool_n)) + print("# escaped_mass_study: {} runs ({} dims x {} offsets x {} arms x {} copies)".format( + len(jobs), len(dims), len(offsets), len(arms), opts.copies)) + sys.stdout.flush() + + t0 = time.time() + if opts.jobs > 1: + import multiprocessing as mp + with mp.get_context("spawn").Pool(opts.jobs) as pool: + results = pool.map(_worker, jobs, chunksize=1) + else: + results = [_worker(j) for j in jobs] + print("# wallclock {:.1f}s".format(time.time() - t0)) + + if opts.json: + with open(opts.json, "w") as fh: + json.dump(results, fh, indent=1) + print("# wrote", opts.json) + n_err = sum(1 for r in results if r.get("error")) + print("# errors:", n_err) + for r in results[:3]: + if r.get("error"): + print(r.get("traceback", r["error"])) + return 1 if n_err == len(results) else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/probe_portfolio_optin_flags.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/probe_portfolio_optin_flags.py new file mode 100644 index 000000000..c8ea13495 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/probe_portfolio_optin_flags.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python +""" +probe_portfolio_optin_flags.py -- shape-gate probe for the portfolio OPT-IN flags. + +`RIFT/integrators/TESTING.md` requires that a change behind an opt-in flag ALSO be probed with the +flag ON: the default-path merge gate necessarily shows bitwise-identical results for opt-in code, so +it proves nothing about that code. This probe covers the two opt-in portfolio features: + + portfolio_adaptive_alloc (adaptive-probe draw allocation) + portfolio_weight_clip (truncated IS on the GMM proposal-fit input) + +Method: reuse the merge-gate suite as a library (per shape_recovery.py's docstring) so the targets, +truth pools, metrics and pass thresholds are IDENTICAL to the gate. We monkey-patch +`build_sampler` to switch the flags on for portfolio runs, and run IN-PROCESS (jobs=1): the gate's +multiprocessing path uses spawn, which would not carry the patch into workers. + +Each configuration is scored with the gate's own `evaluate()`, so a row that PASSes here passes by +exactly the gate's criteria. + +Usage (CPU, like the gate): + export PYTHONPATH=/MonteCarloMarginalizeCode/Code:$PYTHONPATH + export CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 + python probe_portfolio_optin_flags.py [--dims 2,4] [--ncomps 1,3] [--seeds 303] +""" +from __future__ import print_function +import argparse +import os +import sys + +# The merge-gate WRAPPER (run_shape_recovery.sh) exports these; library mode does NOT. Without +# them you silently import the INSTALLED RIFT (not the checkout under test) and/or hit the +# cupy-without-a-device path -- both yield confident, meaningless numbers. A valid probe +# reproduces the gate's ABSOLUTE values row-for-row. +os.environ.setdefault('CUDA_VISIBLE_DEVICES', '') +os.environ.setdefault('OMP_NUM_THREADS', '1') +_CODE = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', '..')) +if _CODE not in sys.path: + sys.path.insert(0, _CODE) + +import shape_recovery as SR + + +def patched_build(flags): + """Return a build_sampler that switches the opt-in flags on for portfolio samplers. + + Most knobs are plain attributes on the portfolio object, so setattr suffices. The GMM + component cap is NOT: it lives on the GMM MEMBER's integrator (`gmm_adaptive`), which the + portfolio forwards through setup(). Since the probe patches AFTER build, reach into the + realized members for that one. Use the reserved key '_gmm_adaptive_cap'. + """ + orig = SR.build_sampler + + def build(kind, target, n_chunk): + s = orig(kind, target, n_chunk) + if kind == "portfolio": + for k, v in flags.items(): + if k == "_gmm_adaptive_cap": + # per-group BIC cap on the portfolio's GMM member(s) + for m in list(getattr(s, "portfolio_realizations", [])): + integ = getattr(m, "integrator", None) + if integ is not None and hasattr(integ, "gmm_dict"): + setattr(integ, "gmm_adaptive", + {g: int(v) for g in integ.gmm_dict}) + continue + setattr(s, k, v) + return s + return build + + +def run_config(label, flags, jobs_spec, nmax_per_dim, neff, run_seed): + """Run the portfolio rows for one flag configuration; return list of (job, record).""" + SR.build_sampler = patched_build(flags) + out = [] + for (d, nc, ts) in jobs_spec: + target = SR.MixtureTarget(d, nc, ts) + rec = SR.run_one("portfolio", target, nmax_per_dim * d, neff, seed=run_seed) + verdict = SR.evaluate(rec) + out.append(((d, nc, ts), rec, verdict)) + print(" {:22s} d{}_n{}_s{} n_eff={:8.0f} lnI-lnZ={:+.4f} {}".format( + label, d, nc, ts, + float(rec.get("n_eff", float("nan"))), + float(rec.get("bias_ln", float("nan"))), + verdict if isinstance(verdict, str) else verdict[0])) + sys.stdout.flush() + return out + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--dims", default="2,4") + ap.add_argument("--ncomps", default="1,3") + ap.add_argument("--seeds", default="303") + ap.add_argument("--nmax-per-dim", type=int, default=None) + ap.add_argument("--neff", type=int, default=None) + ap.add_argument("--run-seed", type=int, default=987654) + args = ap.parse_args() + + cfg = dict(SR.PRESETS["standard"]) + nmax_per_dim = args.nmax_per_dim or cfg["nmax_per_dim"] + neff = args.neff or cfg["neff"] + jobs_spec = [(int(d), int(nc), int(ts)) + for d in args.dims.split(",") + for nc in args.ncomps.split(",") + for ts in args.seeds.split(",")] + + configs = [ + ("flags OFF (default)", {}), + ("adaptive_alloc ON", {"portfolio_adaptive_alloc": True}), + ("weight_clip ON", {"portfolio_weight_clip": 1.0}), + ("adaptive+clip ON", {"portfolio_adaptive_alloc": True, "portfolio_weight_clip": 1.0}), + # VARAHA draw-share constraints (see DESIGN_portfolio_freeze_policy.md). Motivation: on a + # sharp high-SNR target the mixture degenerates to peaked-member-only (VARAHA share -> ~0.01), + # q_mix loses its broad backstop, and a missed mode goes uncovered -> lnZ silently low while + # n_eff looks GOOD. A floor blocks that; a floor WITHOUT a cap lets the share run away to ~1 + # (VARAHA-only), which is the same degeneracy mirrored. These rows check the constraints do + # not damage shape recovery on the gate's own targets, which are NOT pathological -- the + # constraint should be close to a no-op there, and must not regress it. + ("varaha floor .25", {"portfolio_varaha_min_frac": 0.25}), + ("varaha band .25-.75", {"portfolio_varaha_min_frac": 0.25, + "portfolio_varaha_max_frac": 0.75}), + ("band + gmm cap3", {"portfolio_varaha_min_frac": 0.25, + "portfolio_varaha_max_frac": 0.75, + "_gmm_adaptive_cap": 3}), + ] + print("# portfolio opt-in flag probe: {} targets x {} configs " + "(nmax_per_dim={}, neff={})".format(len(jobs_spec), len(configs), nmax_per_dim, neff)) + + results = {} + for label, flags in configs: + print("== {} ==".format(label)) + results[label] = run_config(label, flags, jobs_spec, nmax_per_dim, neff, args.run_seed) + + # Summary: the opt-in paths must not be WORSE than the default path on the gate's own verdict. + print("\n# SUMMARY (verdict per target; opt-in must not regress vs flags OFF)") + base = {k: (v, d) for k, v, d in results["flags OFF (default)"]} + bad = 0 + for label, _ in configs[1:]: + for key, rec, verdict in results[label]: + b_rec, b_verdict = base[key] + vs = verdict if isinstance(verdict, str) else verdict[0] + bs = b_verdict if isinstance(b_verdict, str) else b_verdict[0] + flag = "" + if bs == "PASS" and vs not in ("PASS", "STARVED"): + flag = " <-- REGRESSION (base PASS -> {})".format(vs); bad += 1 + print(" {:22s} d{}_n{}_s{} base={:8s} flag={:8s}{}".format( + label, key[0], key[1], key[2], bs, vs, flag)) + print("\n# opt-in regressions: {}".format(bad)) + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/run_shape_recovery.sh b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/run_shape_recovery.sh new file mode 100755 index 000000000..95c870fef --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/run_shape_recovery.sh @@ -0,0 +1,49 @@ +#!/bin/bash +# Standard merge-gate invocation of the shape-recovery suite. +# +# A FAIL FROM THIS SCRIPT IS NOT A VERDICT. Every gate threshold (n_eff >= 100, JS, pull, width) +# is a hard cut on a stochastic quantity, so any cell sitting near a threshold flips on +# realization alone. Before treating a blocking regression as real, re-test it at fresh seeds: +# +# compare_shape_results.py base.json cand.json \ +# --confirm-base-checkout DIR --confirm-cand-checkout DIR --confirm-repeats 5 +# +# With those flags the comparison ENFORCES confirmation: blocking rows are re-tested and the exit +# code is the confirmed verdict. Without them it still exits 1 on a blocking row, but says so +# explicitly rather than implying the row was confirmed. (confirm_regressions.py also runs +# standalone against an existing pair of JSONs.) +# +# It re-runs only the disputed cells, in BOTH arms, at several new run seeds, and blocks only if +# the candidate is worse in a majority. A candidate that produces NO record where the base did +# counts against the candidate, and too few usable pairs is INCONCLUSIVE (exit 1), never a clear. Worked example: `GMM mix_d6_n3_s303` was reported as a +# blocking REGRESSION in two consecutive full runs (base 119, candidate 66) and looked +# reproducible -- but at 5 fresh seeds the two arms were BIT-IDENTICAL (93/93, 80/80, 119/119, +# 95/95, 96/96) and 4 of the 5 starved. The cell simply sits on the n_eff=100 floor; its PASS at +# the default seed was the lucky draw, and the apparent regression was an artifact of where the +# job landed in the worker pool. +# +# Do NOT "fix" this by seeding the samplers deterministically. Independent copies that localize +# differently are the working detector for support/mode-collapse failures; pinning every fit to +# one seed silences it, and makes N production copies no better than one. +# +# ./run_shape_recovery.sh /path/to/checkout results.json [extra args...] +# +# Runs CPU-only (deterministic; also exercises the cupy-installed-but-no-GPU +# worker configuration that has repeatedly bitten production). Use --jobs to +# parallelize across cores. +set -e +HERE="$(cd "$(dirname "$0")" && pwd)" +CHECKOUT=${1:?usage: run_shape_recovery.sh /path/to/checkout results.json [extra args]} +OUT=${2:?need output json path} +shift 2 + +export PYTHONPATH="${CHECKOUT}/MonteCarloMarginalizeCode/Code:${PYTHONPATH}" +export CUDA_VISIBLE_DEVICES="" +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} +export MKL_NUM_THREADS=${OMP_NUM_THREADS} +export OPENBLAS_NUM_THREADS=${OMP_NUM_THREADS} + +# NOT `python`: several IGWN/conda environments (and this submit host) provide only python3, +# where a bare `python` makes the whole gate exit 127 before it starts. +exec "${PYTHON:-python3}" "${HERE}/shape_recovery.py" --preset standard \ + --jobs "${SHAPE_JOBS:-8}" --warm-cases auto --json "${OUT}" "$@" diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/shape_recovery.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/shape_recovery.py new file mode 100755 index 000000000..5f7a7b0c9 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/shape_recovery.py @@ -0,0 +1,863 @@ +#!/usr/bin/env python +""" +shape_recovery.py -- posterior SHAPE-recovery merge gate for RIFT MC integrators. + +Motivation +---------- +The fast CI gate (.travis/test-integrate.sh -> test/test_mcsamplerEnsemble_extended.py) +checks only the *integral* on a single 3-D correlated Gaussian. Integrals are +easy: importance-sampling estimators of Z are unbiased under very weak +conditions, while the recovered *posterior shape* (the weighted sample cloud +used downstream by CIP / fairdraws) can be subtly biased -- wrong widths, +clipped tails, missing mixture components, distorted correlations -- without +the integral moving outside its error bars. This suite is the strong, +expensive check run before confirming a merge (see ../README.md). + +Method (follows RIFT-FinerNet demos/integrators/multigauss_direct, Wagner et al): + * Targets: random Gaussian mixtures over a range of dimensions. Weights + ~U(0.1,1.1) normalized, means uniform in a central sub-box, covariances + Wishart-drawn (random orientation + condition number). Mixture recipe is + seeded, so every branch under test sees the *identical* targets. + * Truth: 10^6 exact fair draws per target by rejection sampling inside the + integration box (also yields the in-box mass for the true evidence). + * Recovery: each integrator runs through its production API + (add_parameter / setup / integrate[_log]) with save_intg=True; the weighted + posterior cloud is read back from sampler._rvs exactly as ILE/CIP consume it. + * Shape metrics, each dimension: + - JS divergence (nats) of the weighted 1-D marginal histogram vs the + truth-pool histogram; + - mean pull (weighted mean - true mean)/true sigma; + - width ratio weighted sigma / true sigma; + plus max |Delta corr(i,j)| over dimension pairs, evidence bias + lnI - lnZ_true, n_eff and Kish n_ESS. + * Self-calibrating JS pass threshold: the sampling floor for each dim is + measured by subsampling the truth pool down to the run's own n_ESS and + computing JS(subsample, pool) -- i.e. the JS a *perfect* sampler with the + same effective sample count would score. PASS requires + JS < JS_MULT * floor + JS_ABS_MIN. + This keeps one threshold meaningful across samplers/dimensions/branches. + +Policy +------ +Strict (hard-fail) samplers default to AV + GMM; NF and portfolio default to +warn-only (they are known-weaker in older code lines, e.g. rift_O4c). Override +with --strict-samplers / --samplers. Exit code 1 iff any strict run fails. + +Usage +----- + # environment: any venv with numpy/scipy (torch+nflows only needed for NF), + # PYTHONPATH pointing at the checkout under test: + export PYTHONPATH=/path/to/checkout/MonteCarloMarginalizeCode/Code:$PYTHONPATH + export CUDA_VISIBLE_DEVICES="" # CPU: deterministic merge-gate default + + python shape_recovery.py --preset quick # ~minutes, smoke + python shape_recovery.py --preset standard --jobs 8 --json results.json + +This file is self-contained on purpose: it must run unmodified against ANY +branch (including historical ones that lack test/integrators helpers). +""" +from __future__ import print_function + +import argparse +import json +import os +import sys +import time + +import numpy as np +from scipy.special import logsumexp +from scipy.stats import multivariate_normal + +# production-scale constant lnL offset (a modest-SNR detection), as in the +# FinerNet multigauss demo: keeps us honest about lnL-vs-L overflow handling. +LNL_OFFSET = 100.0 +BOX_HALF_WIDTH = 5.0 +TRUTH_POOL_N = 1000000 +JS_NBINS = 50 +JS_MULT = 3.0 # pass if JS < JS_MULT*floor + JS_ABS_MIN +JS_ABS_MIN = 0.004 +MIN_NEFF_FOR_SHAPE = 100.0 +NF_NMAX_CAP = 400000 # NF trains a flow per chunk; cap its budget (warn-only sampler) + + +# ---------------------------------------------------------------------------- +# Target: seeded random Gaussian mixture with exact truth +# ---------------------------------------------------------------------------- +class MixtureTarget(object): + """Random `ncomp`-component Gaussian mixture in `ndim` dimensions on the + box [-BOX_HALF_WIDTH, BOX_HALF_WIDTH]^ndim, FinerNet multigauss recipe.""" + + def __init__(self, ndim, ncomp, seed, sigma_1d=0.7, scale_x0=3.0, offset=0.0): + self.ndim = int(ndim) + self.ncomp = int(ncomp) + self.seed = int(seed) + self.name = "mix_d{}_n{}_s{}".format(ndim, ncomp, seed) + # `offset` TRANSLATES the whole mixture. Two targets built with the same seed and + # offset=-a / +a are the SAME shape displaced -- which is what the sequential cases need + # and what different seeds cannot guarantee (random means typically overlap). It is + # applied AFTER the rng.uniform draw below, so the RNG stream and every existing target + # are bit-identical at the default offset=0. + self.offset = np.zeros(int(ndim)) + np.asarray(offset, dtype=float) + if np.any(self.offset): + self.name += "_o{:+.2f}".format(float(np.mean(self.offset))) + self.params = ["x{}".format(i) for i in range(ndim)] + self.llim = -BOX_HALF_WIDTH * np.ones(ndim) + self.rlim = BOX_HALF_WIDTH * np.ones(ndim) + rng = np.random.RandomState(seed) + wt = rng.uniform(size=ncomp) + 0.1 + self.wt = wt / np.sum(wt) + import scipy.stats as ss + self.means, self.covs, self._mvns = [], [], [] + for k in range(ncomp): + x0 = rng.uniform(-scale_x0 / np.sqrt(ndim), scale_x0 / np.sqrt(ndim), ndim) + self.offset + Sig = (sigma_1d ** 2) * np.diag(rng.uniform(1.0, 2.0, ndim)) + Sig = ss.wishart.rvs(df=ndim, scale=Sig / ndim, random_state=rng) / 1.25 + Sig = np.atleast_2d(Sig) + self.means.append(x0) + self.covs.append(Sig) + self._mvns.append(multivariate_normal(x0, Sig, allow_singular=True)) + self._pool = None + self._box_mass = None + + def lnL(self, X): + X = np.atleast_2d(X) + terms = np.empty((len(X), self.ncomp)) + for k in range(self.ncomp): + terms[:, k] = self._mvns[k].logpdf(X) + np.log(self.wt[k]) + return LNL_OFFSET + logsumexp(terms, axis=1) + + def as_lnfunc(self): + def ln_f(*cols): + return self.lnL(np.array([np.asarray(c, dtype=float) for c in cols]).T) + return ln_f + + def as_func(self): + ln_f = self.as_lnfunc() + return lambda *cols: np.exp(ln_f(*cols)) + + def _build_pool(self): + """Exact fair draws of the box-truncated posterior, by rejection.""" + rng = np.random.RandomState(self.seed + 7) + kept, n_tot, n_in = [], 0, 0 + while n_in < TRUTH_POOL_N: + counts = rng.multinomial(200000, self.wt) + chunks = [rng.multivariate_normal(self.means[k], self.covs[k], counts[k]) + for k in range(self.ncomp) if counts[k] > 0] + draw = np.vstack(chunks) + rng.shuffle(draw) # multinomial blocks are ordered by component + n_tot += len(draw) + inside = np.all((draw > self.llim) & (draw < self.rlim), axis=1) + draw = draw[inside] + n_in += len(draw) + kept.append(draw) + self._pool = np.vstack(kept)[:TRUTH_POOL_N] + self._box_mass = float(n_in) / n_tot + + @property + def pool(self): + if self._pool is None: + self._build_pool() + return self._pool + + @property + def true_lnZ(self): + """Truth for the sampler-returned integral \\int L p_prior dx with + normalized uniform prior: LNL_OFFSET + ln(in-box mass) - sum ln(width).""" + if self._box_mass is None: + self._build_pool() + return (LNL_OFFSET + np.log(self._box_mass) + - float(np.sum(np.log(self.rlim - self.llim)))) + + +# ---------------------------------------------------------------------------- +# Reading back the weighted posterior cloud (tolerant of _rvs conventions) +# ---------------------------------------------------------------------------- +def _asnumpy(a): + try: + import cupy + if isinstance(a, cupy.ndarray): + return cupy.asnumpy(a) + except Exception: + pass + return np.asarray(a) + + +def log_weights_from_rvs(rvs): + """ln(weight) = lnL + ln p_prior - ln p_sampling per stored sample, across + the heterogeneous _rvs conventions (log-keyed AV/NF/portfolio, linear-keyed + default/AC, GMM storing lnL under 'integrand' when return_lnI is set).""" + if "log_integrand" in rvs: + lnL = _asnumpy(rvs["log_integrand"]).astype(float) + elif "integrand" in rvs: + L = _asnumpy(rvs["integrand"]).astype(float) + lnL = L if np.nanmin(L) < 0 else np.log(L + 1e-300) + else: + raise KeyError("no integrand in _rvs; run with save_intg=True") + + def _get_log(logkey, linkey, n): + if logkey in rvs: + return _asnumpy(rvs[logkey]).astype(float) + if linkey in rvs: + return np.log(_asnumpy(rvs[linkey]).astype(float) + 1e-300) + return np.zeros(n) + + n = len(lnL) + lnp = _get_log("log_joint_prior", "joint_prior", n) + lnps = _get_log("log_joint_s_prior", "joint_s_prior", n) + return lnL + lnp - lnps + + +def n_ess_kish(ln_wt): + ln_wt = ln_wt - np.max(ln_wt) + w = np.exp(ln_wt) + return float(np.sum(w) ** 2 / np.sum(w ** 2)) + + +# ---------------------------------------------------------------------------- +# Shape metrics +# ---------------------------------------------------------------------------- +def _js_from_hists(p, q): + p = p / p.sum() + q = q / q.sum() + m = 0.5 * (p + q) + + def _kl(a, b): + mask = a > 0 + return float(np.sum(a[mask] * np.log(a[mask] / (b[mask] + 1e-300)))) + + return 0.5 * _kl(p, m) + 0.5 * _kl(q, m) + + +def shape_metrics(target, X, ln_wt, rng): + """Compare weighted cloud (X, ln_wt) against the target truth pool. + + Returns dict with per-dim JS + matched-n_ESS JS floors, mean pulls, width + ratios, and max correlation-coefficient discrepancy.""" + pool = target.pool + w = np.exp(ln_wt - np.max(ln_wt)) + w = w / np.sum(w) + ness = n_ess_kish(ln_wt) + + mu_true = pool.mean(axis=0) + sd_true = pool.std(axis=0) + mu_w = np.sum(w[:, None] * X, axis=0) + var_w = np.sum(w[:, None] * (X - mu_w) ** 2, axis=0) + sd_w = np.sqrt(var_w) + + js, js_floor = [], [] + n_sub = int(min(max(ness, 50), len(pool) // 10)) + for d in range(target.ndim): + edges = np.linspace(target.llim[d], target.rlim[d], JS_NBINS + 1) + h_pool, _ = np.histogram(pool[:, d], bins=edges) + h_run, _ = np.histogram(X[:, d], bins=edges, weights=w) + js.append(_js_from_hists(h_run.astype(float), h_pool.astype(float))) + # JS floor: perfect sampler at the same effective sample size + f = [] + for _ in range(5): + sub = pool[rng.choice(len(pool), size=n_sub, replace=False), d] + h_sub, _ = np.histogram(sub, bins=edges) + f.append(_js_from_hists(h_sub.astype(float), h_pool.astype(float))) + js_floor.append(float(np.mean(f) + 2.0 * np.std(f))) + + # correlation matrices (guard zero-width dims) + corr_diff = 0.0 + if target.ndim > 1: + cov_w = np.einsum("i,ij,ik->jk", w, X - mu_w, X - mu_w) + corr_w = cov_w / np.outer(sd_w, sd_w) + corr_t = np.corrcoef(pool.T) + corr_diff = float(np.max(np.abs(corr_w - corr_t) + [np.triu_indices(target.ndim, 1)])) + + return dict( + n_ess=ness, + js=[float(x) for x in js], + js_floor=js_floor, + mean_pull=[float(x) for x in (mu_w - mu_true) / sd_true], + width_ratio=[float(x) for x in sd_w / sd_true], + corr_diff_max=corr_diff, + ) + + +# ---------------------------------------------------------------------------- +# Sampler adapter (production API, tolerant of older branches) +# ---------------------------------------------------------------------------- +KNOWN_SAMPLERS = ("AV", "GMM", "NF", "portfolio", "AC", "default") + + +def _gpu_available(): + try: + import cupy + return cupy.cuda.runtime.getDeviceCount() > 0 + except Exception: + return False + + +def _force_cpu(s): + """Mirror production's --sampler-xpy numpy instance override. Needed for + GMM: mcsamplerEnsemble sets cupy_ok on *import* success without a device + probe, so on a cupy-installed GPU-less node it crashes in setup().""" + s.xpy = np + s.identity_convert = lambda x: x + s.identity_convert_togpu = lambda x: x + return s + + +_CPU_PATCHED = False + + +def _force_cpu_modules(): + """The GMM stack (mcsamplerEnsemble -> MonteCarloEnsemble -> + gaussian_mixture_model) selects cupy at *module* level on import success + with no device probe, and the inner integrator has no xpy argument; on a + cupy-installed GPU-less node the only recourse is patching the module + globals to numpy (what production sees when cupy is absent).""" + global _CPU_PATCHED + if _CPU_PATCHED: + return + import importlib + import scipy.special as _sp + for name in ("mcsampler", "mcsamplerEnsemble", "MonteCarloEnsemble", + "gaussian_mixture_model", "mcsamplerGPU", + "mcsamplerAdaptiveVolume", "mcsamplerPortfolio", + "mcsamplerNFlow"): + try: + mod = importlib.import_module("RIFT.integrators." + name) + except Exception: + continue + for attr, val in (("xpy_default", np), ("cupy_ok", False), + ("xpy_special_default", _sp), + ("identity_convert", lambda x: x), + ("identity_convert_togpu", lambda x: x)): + if hasattr(mod, attr): + setattr(mod, attr, val) + _CPU_PATCHED = True + + +def build_sampler(kind, target, n_chunk): + if not _gpu_available(): + _force_cpu_modules() + from RIFT.integrators import mcsampler + + def uniform_pdf(d): + wdt = target.rlim[d] - target.llim[d] + return np.vectorize(lambda x, wdt=wdt: 1.0 / wdt) + + if kind == "default": + s = mcsampler.MCSampler() + elif kind == "AC": + from RIFT.integrators import mcsamplerGPU + s = mcsamplerGPU.MCSampler() + elif kind == "GMM": + from RIFT.integrators import mcsamplerEnsemble + s = mcsamplerEnsemble.MCSampler() + if not _gpu_available(): + _force_cpu(s) + elif kind == "AV": + from RIFT.integrators import mcsamplerAdaptiveVolume + try: + s = mcsamplerAdaptiveVolume.MCSampler(n_chunk=n_chunk) + except TypeError: # older signature + s = mcsamplerAdaptiveVolume.MCSampler() + elif kind == "NF": + from RIFT.integrators import mcsamplerNFlow + s = mcsamplerNFlow.MCSampler() + elif kind == "portfolio": + from RIFT.integrators import (mcsamplerPortfolio, + mcsamplerAdaptiveVolume, mcsamplerEnsemble) + try: + m1 = mcsamplerAdaptiveVolume.MCSampler(n_chunk=n_chunk) + except TypeError: + m1 = mcsamplerAdaptiveVolume.MCSampler() + m2 = mcsamplerEnsemble.MCSampler() + if not _gpu_available(): + _force_cpu(m1) + _force_cpu(m2) + s = mcsamplerPortfolio.MCSampler(portfolio=[m1, m2]) + if not _gpu_available(): + _force_cpu(s) + else: + raise ValueError("unknown sampler kind %r" % kind) + + for d, p in enumerate(target.params): + s.add_parameter(p, uniform_pdf(d), prior_pdf=uniform_pdf(d), + left_limit=float(target.llim[d]), + right_limit=float(target.rlim[d]), + adaptive_sampling=True) + return s + + +def run_one(kind, target, nmax, neff, n_chunk=10000, seed=987654, verbose=False): + """Run one sampler on one target; return metrics dict (never raises).""" + t0 = time.time() + if kind == "NF": + nmax = min(nmax, NF_NMAX_CAP) + out = dict(kind=kind, target=target.name, ndim=target.ndim, + ncomp=target.ncomp, target_seed=target.seed, nmax=int(nmax)) + try: + np.random.seed(seed) + try: + import torch + torch.manual_seed(seed) + torch.set_num_threads(max(1, int(os.environ.get("OMP_NUM_THREADS", "4")))) + except Exception: + pass + s = build_sampler(kind, target, n_chunk) + ln_f = target.as_lnfunc() + params = target.params + extra = dict(n=n_chunk, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=neff, nmax=int(nmax), save_intg=True, verbose=verbose) + if hasattr(s, "setup"): + try: + s.setup() + except TypeError: + pass + + if kind == "default": + f = target.as_func() + I, var, eff, _ = s.integrate(f, *params, no_protect_names=True, **extra) + lnI = float(np.log(I)) + relerr = float(np.sqrt(var) / I) + elif kind == "AC": + lnI, logvar, eff, _ = s.integrate(ln_f, *params, no_protect_names=True, + use_lnL=True, **extra) + lnI = float(_asnumpy(lnI)) + relerr = float(np.exp(float(_asnumpy(logvar)) / 2 - lnI)) + elif kind == "GMM": + n_iters = max(2, int(nmax / n_chunk)) + lnI, logvar, eff, _ = s.integrate(ln_f, *params, min_iter=n_iters, + max_iter=n_iters, correlate_all_dims=True, + n_comp=max(1, target.ncomp), + use_lnL=True, return_lnI=True, **extra) + lnI = float(_asnumpy(lnI)) + relerr = float(np.exp(float(_asnumpy(logvar)) / 2 - lnI)) + else: # AV, NF, portfolio: integrate_log + lnI, logvar, eff, _ = s.integrate_log(ln_f, *params, + no_protect_names=True, **extra) + lnI = float(_asnumpy(lnI)) + logvar = float(_asnumpy(logvar)) + relerr = float(np.exp(0.5 * logvar - lnI)) if np.isfinite(logvar) else float("nan") + + eff = float(_asnumpy(eff)) + ln_wt = log_weights_from_rvs(s._rvs) + X = np.column_stack([_asnumpy(s._rvs[p]).astype(float).flatten() + for p in params]) + rng = np.random.RandomState(seed + 1) + out.update(shape_metrics(target, X, ln_wt, rng)) + out.update(lnI=lnI, true_lnZ=float(target.true_lnZ), + bias_ln=lnI - float(target.true_lnZ), rel_err=relerr, + n_eff=eff, n_eval=int(getattr(s, "ntotal", 0)) or int(nmax), + wallclock=time.time() - t0, error=None) + except Exception as e: + import traceback + out.update(error="{}: {}".format(type(e).__name__, e), + traceback=traceback.format_exc(), wallclock=time.time() - t0) + return out + + +# ---------------------------------------------------------------------------- +# Warm-start / sequential-reuse cases +# ---------------------------------------------------------------------------- +# These guard two portfolio defects that are INVISIBLE to the ordinary matrix, because both +# produce a wrong answer with no exception and (for the first) no statistical signature at all: +# +# (a) a warm-start seed that is accepted but never installed on the draw path. Measured: with +# the AV install disabled, 12/12 runs still PASS and n_eff is 4692-5972 versus 3240-5941 +# for the correct code -- the broken config often scores BETTER, because the portfolio's +# GMM member supplies nearly the whole warm-start win. No black-box assertion can see +# this; only a direct check that the AV member's live volume actually contracted. +# +# (b) state leaking between sequential points. mcsamplerPortfolio.integrate_log does not call +# self.setup(), so a member's contracted live volume survives into the next integral. If +# the next point's support lies outside it, lnZ is biased low with a healthy-looking n_eff. +# This needs no warm-start feature at all -- it bites any --n-events-to-analyze > 1 run. +WARM_KINDS = ("portfolio_warm", "portfolio_seq", "portfolio_seq_nobs", "AV_seq") +# For these kinds a starved run is a FAILURE, not "untestable": the whole point of the case is +# that correct code comfortably clears the floor (measured margins >= 8x). +STARVE_IS_FAIL = ("portfolio_warm", "portfolio_seq", "portfolio_seq_nobs") +WARM_NEFF_FLOOR = 1000.0 # case W A3; measured warm 2426-5941, cold 3.8-91.7 +WARM_V_MAX = 0.9 # case W A1; measured installed 0.031-0.129, inert exactly 1.000 +WARM_BOX_MULT = 3.0 # case W A2; measured on AV-member draws only -- see run_warm_case +SEQ_OFFSET = 2.0 # +-2 with sigma_1d=0.7 -> mean separation 4.0, both inside [-5,5] + +# CASE-LIST NOTE, from directly reintroducing each bug and re-running (not from reasoning): +# * portfolio_seq_nobs is the DISCRIMINANT for the leak. With clear_warm_state() no-op'd it +# gives n_eff 1.0 / 1.0 / 9.9 and lnZ bias -22.8 / -59.7 / -0.56 at ts 101/202/303, versus +# n_eff ~2094 and bias +0.019 when the reset works. +# * portfolio_seq (which re-bootstraps on point B) does NOT catch it: measured PASS with +# n_eff 5979 while the bug was active, because the fresh B seed overwrites the stale grid. +# It is kept as ONE row only, and only because it covers the reseed-after-reset path. +WARM_CASES = [ # (kind, ndim, ncomp, target_seed, nmax, neff, extra) + ("portfolio_warm", 6, 1, 101, 600000, 3000, {}), + ("portfolio_warm", 6, 1, 202, 600000, 3000, {}), + ("portfolio_warm", 6, 1, 303, 600000, 3000, {}), + ("portfolio_seq_nobs", 2, 1, 101, 100000, 2000, dict(scale_x0=1.0)), + ("portfolio_seq_nobs", 2, 1, 202, 100000, 2000, dict(scale_x0=1.0)), + ("portfolio_seq_nobs", 2, 1, 303, 100000, 2000, dict(scale_x0=1.0)), + # covers reseed-after-reset; NOT a leak discriminant (see note above) + ("portfolio_seq", 2, 1, 101, 100000, 2000, dict(scale_x0=1.0)), + # negative control: standalone AV must be unaffected (it reruns its own setup). Warn-only. + ("AV_seq", 2, 1, 101, 100000, 2000, dict(scale_x0=1.0)), +] + + +def _warm_seed_cloud(target, n=3000): + """Fair draws from the target's own truth pool -- a PERFECT seed, so any shortfall is the + warm-start machinery, not a bad proposal.""" + pool = target.pool + rng = np.random.RandomState(target.seed + 13) + idx = rng.choice(len(pool), size=min(int(n), len(pool)), replace=False) + return np.asarray(pool[idx], dtype=float) + + +def _finish_record(out, target, s, lnI, logvar, eff, nmax, seed, t0): + """Shared tail of run_one: shape metrics + lnZ bookkeeping from a finished sampler.""" + eff = float(_asnumpy(eff)) + ln_wt = log_weights_from_rvs(s._rvs) + X = np.column_stack([_asnumpy(s._rvs[p]).astype(float).flatten() for p in target.params]) + rng = np.random.RandomState(seed + 1) + out.update(shape_metrics(target, X, ln_wt, rng)) + lnI = float(_asnumpy(lnI)) + logvar = float(_asnumpy(logvar)) + relerr = float(np.exp(0.5 * logvar - lnI)) if np.isfinite(logvar) else float("nan") + out.update(lnI=lnI, true_lnZ=float(target.true_lnZ), bias_ln=lnI - float(target.true_lnZ), + rel_err=relerr, n_eff=eff, n_eval=int(getattr(s, "ntotal", 0)) or int(nmax), + wallclock=time.time() - t0, error=None) + return out + + +def run_warm_case(target, nmax, neff, n_chunk=10000, seed=987654, verbose=False): + """Case W: does a warm-start seed reach the AV member's ACTIVE draw state? + + Two samplers: a cheap PROBE that is seeded and drawn once (white-box: reads the AV member's + live volume and bin count), then a fresh one that is seeded and integrated (black-box). + The probe is what catches the inert-seed bug; the integral catches "all warm channels died".""" + t0 = time.time() + out = dict(kind="portfolio_warm", target=target.name, ndim=target.ndim, + ncomp=target.ncomp, target_seed=target.seed, nmax=int(nmax)) + try: + np.random.seed(seed) + cloud = _warm_seed_cloud(target) + lo, hi = cloud.min(axis=0), cloud.max(axis=0) + + probe = build_sampler("portfolio", target, n_chunk) + try: + probe.setup() + except TypeError: + pass + probe.bootstrap_from_samples(cloud, cover_frac=0.0) + probe.draw(n_chunk) + av = probe.portfolio_realizations[0] + out["warm_V"] = float(_asnumpy(getattr(av, "V", 1.0))) + out["warm_bins"] = int(len(getattr(av, "binunique", [0]))) + # Measure concentration on the AV MEMBER'S OWN draws, not the portfolio mixture. Measured: + # with the AV install disabled the MIXTURE still concentrates 28x in the seed box, because + # the GMM member is separately warm-started -- so a mixture-level ratio is not a + # discriminant for this bug at all. Drawing from the member isolates the path under test. + _ps, _p, rv_av = av.draw_simplified(n_chunk) # rv_av is (ndim, n) + Xp = np.asarray(_asnumpy(rv_av), dtype=float).T + out["warm_box_frac"] = float(np.mean(np.all((Xp >= lo) & (Xp <= hi), axis=1))) + # what a UNIFORM (cold) proposal would put in the same box -- the null this must beat + out["warm_box_frac_uniform"] = float(np.prod((hi - lo) / (target.rlim - target.llim))) + + s = build_sampler("portfolio", target, n_chunk) + try: + s.setup() + except TypeError: + pass + s.bootstrap_from_samples(cloud, cover_frac=0.0) + extra = dict(n=n_chunk, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=neff, nmax=int(nmax), save_intg=True, verbose=verbose) + lnI, logvar, eff, _ = s.integrate_log(target.as_lnfunc(), *target.params, + no_protect_names=True, **extra) + _finish_record(out, target, s, lnI, logvar, eff, nmax, seed, t0) + except Exception as e: + import traceback + out.update(error="{}: {}".format(type(e).__name__, e), + traceback=traceback.format_exc(), wallclock=time.time() - t0) + return out + + +def run_seq_case(kind, target_b, target_a, nmax, neff, n_chunk=10000, seed=987654, verbose=False): + """Cases S / S-nobs / AV_seq: integrate DISPLACED target A then target B on ONE sampler. + + Scores point B only. If the sampler carries A's contracted live volume into B, B's mass sits + outside it and lnZ collapses. Uses clear_warm_state() when present and falls back to the old + `_warm = None` otherwise, so the case RUNS on a base branch without the API -- and fails there, + which is the point.""" + t0 = time.time() + out = dict(kind=kind, target=target_b.name, ndim=target_b.ndim, ncomp=target_b.ncomp, + target_seed=target_b.seed, nmax=int(nmax)) + try: + np.random.seed(seed) + s = build_sampler("AV" if kind == "AV_seq" else "portfolio", target_a, n_chunk) + # PRODUCTION SHAPE: supply gmm_dict. Production always does, and it is not an inert spec -- + # monte_carlo.integrator stores the caller's dict without copying and writes trained models + # into it, so this is the configuration in which stale-proposal aliasing can occur. A gate + # that only ever ran the gmm_dict=None branch could not see that class of defect at all. + _setup_kw = {} + if kind != "AV_seq": + _dims = tuple(range(len(target_a.params))) + _setup_kw = dict(n_comp={_dims: 2}, gmm_dict={_dims: None}, correlate_all_dims=True) + try: + s.setup(**_setup_kw) + except TypeError: + s.setup() + extra = dict(n=n_chunk, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=neff, nmax=int(nmax), save_intg=True, verbose=verbose) + if kind != "portfolio_seq_nobs": + s.bootstrap_from_samples(_warm_seed_cloud(target_a), cover_frac=0.0) + s.integrate_log(target_a.as_lnfunc(), *target_a.params, + no_protect_names=True, **extra) + # record that point A ACTUALLY trained something -- otherwise "cleared" could pass + # vacuously on a run where the GMM never fitted at all + _trained_before_reset = False + if kind != "AV_seq": + try: + _trained_before_reset = any( + v is not None for v in s.portfolio_realizations[1].integrator.gmm_dict.values()) + except Exception as _e_tr: + out["seq_gmm_error"] = "pre-reset inspection failed: {}: {}".format( + type(_e_tr).__name__, _e_tr) + + # ---- the transition the driver performs between two points ---- + s._rvs = {} + if hasattr(s, "clear_warm_state"): + s.clear_warm_state() + else: + s._warm = None + out["used_clear_api"] = bool(hasattr(s, "clear_warm_state")) + # WHITE-BOX: did the reset actually clear the GMM member's TRAINED PROPOSAL? Measured, the + # statistical rows cannot answer this: with the trained proposal leaking, n_eff is 196-1700 + # and |lnZ bias| <= 0.010 -- degraded but passing, because a stale GMM proposal is merely a + # bad proposal (the AV member still covers the support), unlike the AV grid leak which + # removes support and does bias. So the leak must be observed directly. + if kind != "AV_seq": + try: + _gd = s.portfolio_realizations[1].integrator.gmm_dict + out["seq_gmm_trained_before_reset"] = bool(_trained_before_reset) + out["seq_gmm_cleared"] = all(v is None for v in _gd.values()) + except Exception as _e_gd: + out["seq_gmm_cleared"] = None + out["seq_gmm_error"] = "{}: {}".format(type(_e_gd).__name__, _e_gd) + + if kind != "portfolio_seq_nobs": + s.bootstrap_from_samples(_warm_seed_cloud(target_b), cover_frac=0.0) + lnI, logvar, eff, _ = s.integrate_log(target_b.as_lnfunc(), *target_b.params, + no_protect_names=True, **extra) + _finish_record(out, target_b, s, lnI, logvar, eff, nmax, seed, t0) + except Exception as e: + import traceback + out.update(error="{}: {}".format(type(e).__name__, e), + traceback=traceback.format_exc(), wallclock=time.time() - t0) + return out + + +# ---------------------------------------------------------------------------- +# Pass/fail policy +# ---------------------------------------------------------------------------- +def evaluate(r): + """Return (status, reasons) for one run record. + + status: "PASS" | "FAIL" | "STARVED" | "ERROR". + STARVED (n_eff below the shape-testability floor at this budget) is NOT an + absolute failure: high-dimensional mixtures legitimately exhaust production + budgets (the FinerNet high-D degradation), so starved rows gate only + *differentially* -- a candidate that starves where its base was healthy is + a regression (see compare_shape_results.py).""" + if r.get("error"): + return "ERROR", ["ERROR " + r["error"]] + if r["kind"] in STARVE_IS_FAIL and r["n_eff"] < MIN_NEFF_FOR_SHAPE: + # NOT "untestable": correct code clears this floor by >= 8x on these cases (measured + # 799-2103 vs 0.0-12.8 when state leaks), so starvation here IS the defect. + return "FAIL", ["n_eff={:.0f} < {:.0f}: warm/sequential case must not starve".format( + r["n_eff"], MIN_NEFF_FOR_SHAPE)] + if r["n_eff"] < MIN_NEFF_FOR_SHAPE: + return "STARVED", ["n_eff={:.0f} < {:.0f}: shape untestable at this budget".format( + r["n_eff"], MIN_NEFF_FOR_SHAPE)] + reasons = [] + if r["kind"] == "portfolio_warm": + # A1: white-box, and deliberately so. An inert seed leaves V at exactly 1.000 with a + # single live bin; a working one contracted to 0.031-0.129 with 489-656 bins in every + # measured run. No RNG enters either quantity, so the margin is categorical. + if not (r.get("warm_V", 1.0) < WARM_V_MAX and r.get("warm_bins", 1) > 1): + reasons.append("warm seed NOT installed on the draw path: AV member V={:.3f}, " + "live bins={} (cold state)".format(r.get("warm_V", float("nan")), + r.get("warm_bins", -1))) + # A2: behavioural confirmation -- draws must actually concentrate in the seed box. + if r.get("warm_box_frac", 0.0) < WARM_BOX_MULT * r.get("warm_box_frac_uniform", 1.0): + reasons.append("warm draws not concentrated in the seed box: {:.3f} < {:.1f}x{:.4f}" + .format(r.get("warm_box_frac", float("nan")), WARM_BOX_MULT, + r.get("warm_box_frac_uniform", float("nan")))) + # A3: feature level -- catches "every warm-start channel went inert", which A1/A2 (AV + # member only) would miss. + if r["n_eff"] < WARM_NEFF_FLOOR: + reasons.append("warm n_eff {:.0f} < {:.0f}".format(r["n_eff"], WARM_NEFF_FLOOR)) + if r["kind"] in ("portfolio_seq", "portfolio_seq_nobs"): + # Require BOTH observations to be affirmatively True. Testing only for `cleared is False` + # let three distinct failures read as success: training never happened (so "cleared" is + # trivially true and proves nothing), the inspection raised (cleared is None), or the + # fields were absent entirely. A check that cannot run is a failed check, not a pass. + _trained = r.get("seq_gmm_trained_before_reset") + _cleared = r.get("seq_gmm_cleared") + if _trained is not True: + reasons.append("GMM clearing check could not run: trained_before_reset={!r} -- the " + "member never trained, so a 'cleared' verdict would be vacuous{}".format( + _trained, + " [" + r["seq_gmm_error"] + "]" if r.get("seq_gmm_error") else "")) + elif _cleared is not True: + reasons.append("reset did NOT clear the GMM member's trained proposal (cleared={!r}): " + "point B inherits point A's fitted components (setup-arg aliasing){}".format( + _cleared, + " [" + r["seq_gmm_error"] + "]" if r.get("seq_gmm_error") else "")) + ness = max(r["n_ess"], 1.0) + for d, (js, floor) in enumerate(zip(r["js"], r["js_floor"])): + thresh = JS_MULT * floor + JS_ABS_MIN + if js > thresh: + reasons.append("JS[{}]={:.4f} > {:.4f} (floor {:.4f})".format( + d, js, thresh, floor)) + tol_mean = max(5.0 / np.sqrt(ness), 0.05) + for d, pull in enumerate(r["mean_pull"]): + if abs(pull) > tol_mean: + reasons.append("mean_pull[{}]={:+.3f} > {:.3f}".format(d, pull, tol_mean)) + tol_wid = max(5.0 / np.sqrt(2.0 * ness), 0.05) + for d, wr in enumerate(r["width_ratio"]): + if abs(wr - 1.0) > tol_wid: + reasons.append("width_ratio[{}]={:.3f} (tol {:.3f})".format(d, wr, tol_wid)) + tol_corr = max(8.0 / np.sqrt(ness), 0.08) + if r["corr_diff_max"] > tol_corr: + reasons.append("corr_diff_max={:.3f} > {:.3f}".format( + r["corr_diff_max"], tol_corr)) + relerr = r["rel_err"] if np.isfinite(r.get("rel_err", float("nan"))) else 0.05 + tol_lnZ = max(4.0 * relerr, 0.10) + if abs(r["bias_ln"]) > tol_lnZ: + reasons.append("lnZ bias {:+.3f} > {:.3f}".format(r["bias_ln"], tol_lnZ)) + return ("FAIL" if reasons else "PASS"), reasons + + +# ---------------------------------------------------------------------------- +# Matrix presets + CLI +# ---------------------------------------------------------------------------- +PRESETS = { + # (dims, ncomps, target_seeds, nmax_per_dim, neff) + "quick": (dict(dims=[2, 4], ncomps=[2], seeds=[101], nmax_per_dim=50000, neff=2000)), + "standard": (dict(dims=[2, 4, 6, 8], ncomps=[1, 3], seeds=[101, 202, 303], + nmax_per_dim=200000, neff=3000)), +} + + +def _worker(job): + kind, tgt_args, nmax, neff, seed = job[:5] + extra = job[5] if len(job) > 5 else {} + if kind == "portfolio_warm": + return run_warm_case(MixtureTarget(*tgt_args, **extra), nmax, neff, seed=seed) + if kind in ("portfolio_seq", "portfolio_seq_nobs", "AV_seq"): + # SAME mixture, translated: A at -SEQ_OFFSET, B at +SEQ_OFFSET. Different seeds would + # give random, typically overlapping means and would not test displacement at all. + a = MixtureTarget(*tgt_args, offset=-SEQ_OFFSET, **extra) + b = MixtureTarget(*tgt_args, offset=+SEQ_OFFSET, **extra) + return run_seq_case(kind, b, a, nmax, neff, seed=seed) + return run_one(kind, MixtureTarget(*tgt_args, **extra), nmax, neff, seed=seed) + + +def main(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[1]) + ap.add_argument("--preset", default="standard", choices=sorted(PRESETS)) + ap.add_argument("--samplers", default="AV,GMM,NF,portfolio", + help="comma list from: " + ",".join(KNOWN_SAMPLERS)) + ap.add_argument("--warm-cases", default="auto", choices=("auto", "on", "off"), + help="run the warm-start/sequential-reuse cases " + "(auto = on for --preset standard, off for quick)") + ap.add_argument("--strict-samplers", default="AV,GMM", + help="samplers whose failures set exit code 1 (others warn)") + ap.add_argument("--dims", default=None, help="override preset, e.g. 2,4,8") + ap.add_argument("--ncomps", default=None) + ap.add_argument("--target-seeds", default=None) + ap.add_argument("--nmax-per-dim", type=int, default=None, + help="nmax = this * ndim") + ap.add_argument("--neff", type=int, default=None) + ap.add_argument("--run-seed", type=int, default=987654) + ap.add_argument("--jobs", type=int, default=1) + ap.add_argument("--json", default=None, help="write full records here") + ap.add_argument("--verbose", action="store_true") + opts = ap.parse_args(argv) + + cfg = dict(PRESETS[opts.preset]) + if opts.dims: + cfg["dims"] = [int(x) for x in opts.dims.split(",")] + if opts.ncomps: + cfg["ncomps"] = [int(x) for x in opts.ncomps.split(",")] + if opts.target_seeds: + cfg["seeds"] = [int(x) for x in opts.target_seeds.split(",")] + if opts.nmax_per_dim: + cfg["nmax_per_dim"] = opts.nmax_per_dim + if opts.neff: + cfg["neff"] = opts.neff + + samplers = [x.strip() for x in opts.samplers.split(",") if x.strip()] + strict = set(x.strip() for x in opts.strict_samplers.split(",") if x.strip()) + + jobs = [] + for d in cfg["dims"]: + for nc in cfg["ncomps"]: + for ts in cfg["seeds"]: + for kind in samplers: + jobs.append((kind, (d, nc, ts), + cfg["nmax_per_dim"] * d, cfg["neff"], opts.run_seed)) + n_matrix = len(jobs) + want_warm = (opts.warm_cases == "on" or + (opts.warm_cases == "auto" and opts.preset == "standard")) + if want_warm: + for kind, d, nc, ts, nmax, neff, extra in WARM_CASES: + jobs.append((kind, (d, nc, ts), nmax, neff, opts.run_seed, dict(extra))) + print("# shape_recovery: {} runs ({} targets x {} samplers){}, preset={}".format( + len(jobs), n_matrix // len(samplers), len(samplers), + " + {} warm/sequential cases".format(len(jobs) - n_matrix) if want_warm else "", + opts.preset)) + sys.stdout.flush() + + if opts.jobs > 1: + import multiprocessing as mp + with mp.get_context("spawn").Pool(opts.jobs) as pool: + results = pool.map(_worker, jobs) + else: + results = [_worker(j) for j in jobs] + + n_fail_strict, n_fail_warn, n_starved = 0, 0, 0 + print("\n{:<10s} {:<16s} {:>9s} {:>9s} {:>7s} {:>7s} {:>8s} {:>7s} {}".format( + "sampler", "target", "n_eff", "n_ESS", "JSmax", "|pull|", "widthdev", + "lnZbias", "verdict")) + for r in results: + status, reasons = evaluate(r) + if status == "PASS": + tag = "PASS" + elif status == "STARVED": + tag = "STARVED" # non-blocking; gates differentially vs base + n_starved += 1 + elif r["kind"] in strict: + tag = "FAIL" + n_fail_strict += 1 + else: + tag = "WARN" + n_fail_warn += 1 + if r.get("error"): + print("{:<10s} {:<16s} {:>9s} {:>9s} {:>7s} {:>7s} {:>8s} {:>7s} {} {}".format( + r["kind"], r["target"], "-", "-", "-", "-", "-", "-", tag, r["error"])) + continue + print("{:<10s} {:<16s} {:>9.0f} {:>9.0f} {:>7.4f} {:>7.3f} {:>8.3f} {:>+7.3f} {}{}".format( + r["kind"], r["target"], r["n_eff"], r["n_ess"], max(r["js"]), + max(abs(p) for p in r["mean_pull"]), + max(abs(w - 1) for w in r["width_ratio"]), r["bias_ln"], tag, + (" [" + "; ".join(reasons) + "]") if reasons else "")) + sys.stdout.flush() + + if opts.json: + with open(opts.json, "w") as fh: + json.dump(results, fh, indent=1) + print("# wrote", opts.json) + print("# strict failures: {} warn-only failures: {} starved (non-blocking): {}".format( + n_fail_strict, n_fail_warn, n_starved)) + return 1 if n_fail_strict else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_confirm_regressions.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_confirm_regressions.py new file mode 100644 index 000000000..7f985b603 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_confirm_regressions.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python +"""Unit tests for the confirm-on-fail accounting. + +These cover the ways a confirmation step can WRONGLY CLEAR a real regression, which is the only +dangerous direction: a false block costs a rerun, a false clear ships a bug. + +Run: python test_confirm_regressions.py +""" +import sys + +import confirm_regressions as CR +from compare_shape_results import classify, is_blocking + +STRICT = {"GMM", "AV"} + + +def _rec(kind="GMM", target="t", n_eff=3000.0, js=0.0001, bias=0.001): + return dict(kind=kind, target=target, ndim=4, ncomp=1, target_seed=101, n_eff=n_eff, + n_ess=n_eff * 3, js=[js, js], js_floor=[0.0005, 0.0005], + mean_pull=[0.005, 0.005], width_ratio=[1.001, 1.001], corr_diff_max=0.005, + rel_err=0.01, bias_ln=bias, error=None) + + +def test_metrics_only_regression_is_recognised(): + """The comparator blocks on REGRESSION(metrics) too. A confirm step that only knew about + PASS->non-PASS reported 'nothing to confirm' and exited 0 on a real n_eff collapse.""" + b, c = _rec(n_eff=4000.0), _rec(n_eff=400.0) # 10x n_eff drop, both still PASS + verdict, _ = classify(b, c) + assert verdict == "REGRESSION(metrics)", verdict + assert is_blocking(verdict, "GMM", STRICT) + + +def _run_with(monkey_results, seeds=(1, 2, 3), min_valid=None): + """Drive main() with _rerun stubbed to a scripted sequence of (base, cand) records.""" + calls = {"i": 0} + + def fake_rerun(checkout, rec, seed, jobs, tag): + pair = monkey_results[calls["i"] // 2] + out = pair[0] if tag == "base" else pair[1] + calls["i"] += 1 + return out + + orig = CR._rerun + CR._rerun = fake_rerun + try: + import json, tempfile, os + b, c = _rec(n_eff=4000.0), _rec(n_eff=400.0) + paths = [] + for recs in ([b], [c]): + fd, p = tempfile.mkstemp(suffix=".json") + os.close(fd) + json.dump(recs, open(p, "w")) + paths.append(p) + argv = [paths[0], paths[1], "--base-checkout", "/b", "--cand-checkout", "/c", + "--seeds", ",".join(str(s) for s in seeds)] + if min_valid is not None: + argv += ["--min-valid", str(min_valid)] + return CR.main(argv) + finally: + CR._rerun = orig + + +def test_candidate_crash_counts_against_the_candidate(): + """If the candidate produces no record where the base does, that IS the regression. + Discarding those pairs let a candidate that failed on every seed be 'not confirmed'.""" + good = _rec(n_eff=4000.0) + rc = _run_with([(good, None), (good, None), (good, None)]) + assert rc == 1, "candidate produced no record on every seed but was cleared (rc={})".format(rc) + + +def test_insufficient_valid_pairs_is_inconclusive_not_a_pass(): + """Missing evidence must not read as 'cleared'.""" + good = _rec(n_eff=4000.0) + rc = _run_with([(None, None), (None, None), (None, None)]) + assert rc == 1, "zero valid pairs was reported as success (rc={})".format(rc) + + +def test_genuine_noise_clears(): + """A row that is equivalent at fresh seeds must clear, or the step is useless.""" + good = _rec(n_eff=4000.0) + rc = _run_with([(good, good), (good, good), (good, good)]) + assert rc == 0, "equivalent arms were reported as a confirmed regression (rc={})".format(rc) + + +def test_real_regression_is_confirmed(): + good, bad = _rec(n_eff=4000.0), _rec(n_eff=200.0) + rc = _run_with([(good, bad), (good, bad), (good, bad)]) + assert rc == 1, "a reproducible 20x n_eff drop was not confirmed (rc={})".format(rc) + + + + +def test_missing_candidate_record_is_a_blocking_regression(): + """A candidate that emits no record for a strict row must BLOCK. + + Classified as ONLY-IN-BASE it was not a regression, so it never reached confirmation and the + gate exited 0 -- a candidate crashing before its first result would bypass the fail-closed + rerun logic entirely.""" + b = _rec() + verdict, note = classify(b, None) + assert verdict.startswith("REGRESSION"), verdict + assert is_blocking(verdict, "GMM", STRICT), "missing candidate record did not block" + # and it must be picked up as a row to confirm + from compare_shape_results import blocking_keys + keys = blocking_keys({("GMM", "t"): b}, {}, STRICT) + assert keys == [("GMM", "t")], keys + + +def test_extra_candidate_record_is_not_a_regression(): + """The reverse direction is not a defect: a NEW row in the candidate must not block.""" + verdict, _ = classify(None, _rec()) + assert not verdict.startswith("REGRESSION"), verdict + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print("PASS", name) + print("confirm-on-fail accounting holds") diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_escaped_mass_diagnostic.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_escaped_mass_diagnostic.py new file mode 100644 index 000000000..08ad2f914 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_escaped_mass_diagnostic.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python +"""Correctness tests for the portfolio's SUPPORT-MISMATCH diagnostic +(mcsamplerPortfolio.support_diagnostics / escaped_mass). + +These are CORRECTNESS tests, not a claim that the statistic is a useful detector -- that is what +escaped_mass_study.py measures, and its verdict is arm-dependent (see the study docstring). What +is asserted here is only what must hold for the number to mean anything at all: + + 1. OFF-PATH. With the diagnostic reduced to a no-op, the returned lnZ / var / n_eff are + BIT-IDENTICAL. A diagnostic that perturbs the estimator is worse than no diagnostic. + MEASURED CAVEAT: only the ALL-AV portfolio is bit-reproducible run to run. An [AV, GMM] + portfolio is NOT -- repeating the identical configuration with the identical np.random seed + and the diagnostic ON BOTH TIMES gives lnZ 90.75572176844321 vs 90.75570550321379 (the + sklearn mixture fit inside mcsamplerEnsemble does not reproduce). That is a pre-existing + property of the GMM member, not of this diagnostic, so the bit-identity assertion is made on + the arm where it is meaningful, and the [AV, GMM] arm is covered by the structural test + instead. + 2. NON-INVASIVE. A direct call to _update_support_diagnostics changes no sampler attribute + outside its own accumulator namespace and does not modify its inputs. + 3. IT FIRES. Warm-start the AV member from a cloud placed entirely off the true peak, in a + portfolio whose other member is a COLD (broad) GMM, and escaped_mass for the AV member must + go to ~1 while the matched-seed control stays at ~0 on the first-chunk statistic. + 4. IT IS READING THE SUPPORT, not a proxy: make the AV member's density strictly positive + everywhere (so nothing CAN escape) and the same misplaced seed must score 0 / not hard-edged. + 5. A member whose density is nowhere exactly zero is reported hard_edged=False and does not + contribute to escaped_mass_max, so a soft member cannot mask an escaped hard-edged one. + +Run: RIFT_RUN_EXPENSIVE=1 pytest -v test_escaped_mass_diagnostic.py + (or: python test_escaped_mass_diagnostic.py) +""" +from __future__ import print_function + +import os +import sys + +import numpy as np + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +import shape_recovery as sr # noqa: E402 +from escaped_mass_study import _build_portfolio # noqa: E402 + +try: + import pytest + pytestmark = pytest.mark.skipif( + not os.environ.get("RIFT_RUN_EXPENSIVE"), + reason="expensive merge-gate suite; set RIFT_RUN_EXPENSIVE=1") +except ImportError: # allow bare `python test_escaped_mass_diagnostic.py` + pytest = None + +NDIM, NCOMP, TSEED = 4, 1, 4242 +N_CHUNK, NMAX, NEFF = 4000, 40000, 2000 +sr.TRUTH_POOL_N = 100000 # only feeds the box mass + the seed cloud; see the study + + +def _run(arm, offset, run_seed=13579, disable_diag=False, soften_av=False): + """Integrate the TRUE target after warm-starting from a target displaced by `offset`.""" + np.random.seed(run_seed) + true_t = sr.MixtureTarget(NDIM, NCOMP, TSEED) + seed_t = sr.MixtureTarget(NDIM, NCOMP, TSEED, offset=offset) + s, members = _build_portfolio(arm, true_t, N_CHUNK) + kw = {} + if arm != "avav": + _dims = tuple(range(NDIM)) + kw = dict(n_comp={_dims: 2}, gmm_dict={_dims: None}, correlate_all_dims=True) + s.setup(**kw) + if disable_diag: + # break the diagnostic, not the sampler: everything else must be untouched + s._update_support_diagnostics = lambda *a, **k: None + members[0].bootstrap_from_samples(sr._warm_seed_cloud(seed_t), cover_frac=0.0) + if soften_av: + # BREAK THE THING THE STATISTIC TARGETS: keep the identical misplaced seed but make the AV + # member's reported density strictly positive everywhere, so no sample can be OUTSIDE its + # support. If escaped_mass were reading anything other than the support (n_eff, weight + # concentration, the offset itself) it would still fire; it must not. + _orig = members[0].sampling_density + + def _soft(X, _o=_orig): + q = _o(X) + return None if q is None else np.maximum(np.asarray(q, dtype=float), 1e-12) + members[0].sampling_density = _soft + lnI, logvar, eff, dret = s.integrate_log( + true_t.as_lnfunc(), *true_t.params, no_protect_names=True, + n=N_CHUNK, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=NEFF, nmax=NMAX, save_intg=True) + return float(sr._asnumpy(lnI)), float(sr._asnumpy(logvar)), float(sr._asnumpy(eff)), dret + + +def test_diagnostic_is_off_path(): + """The estimate must be BIT-identical when the diagnostic is disabled. + + Asserted on the ALL-AV portfolio, the arm that is bit-reproducible at all (see the module + docstring: an [AV, GMM] portfolio does not reproduce itself run to run, with or without this + code, so a bit-identity assertion there would be testing sklearn, not the diagnostic).""" + a = _run("avav", 0.0, disable_diag=False) + b = _run("avav", 0.0, disable_diag=True) + assert a[0] == b[0], "lnZ changed with the diagnostic on: {!r} vs {!r}".format(a[0], b[0]) + assert a[1] == b[1], "log-variance changed: {!r} vs {!r}".format(a[1], b[1]) + assert a[2] == b[2], "n_eff changed: {!r} vs {!r}".format(a[2], b[2]) + # and the disabled run must genuinely report nothing, otherwise the comparison is vacuous + assert float(b[3]["portfolio_escaped_mass_max"]) == 0.0 + assert int(np.sum(b[3]["portfolio_escape_n_eval"])) == 0, \ + "diagnostic still accumulated after being disabled -- the off-path check is vacuous" + assert int(np.sum(a[3]["portfolio_escape_n_eval"])) > 0, \ + "diagnostic never accumulated in the ENABLED run -- the off-path check is vacuous" + + +def test_diagnostic_mutates_nothing_outside_its_namespace(): + """Structural off-path check, valid for EVERY arm including the nondeterministic [AV, GMM]. + + Runs one chunk, then calls _update_support_diagnostics a second time by hand and verifies that + the only attributes whose value changed are the diagnostic's own accumulators, and that the + inputs it is handed come back unmodified.""" + from RIFT.integrators import mcsamplerPortfolio # noqa: F401 (import check) + np.random.seed(24680) + t = sr.MixtureTarget(NDIM, NCOMP, TSEED) + s, members = _build_portfolio("avgmm_cold", t, N_CHUNK) + _dims = tuple(range(NDIM)) + s.setup(n_comp={_dims: 2}, gmm_dict={_dims: None}, correlate_all_dims=True) + members[0].bootstrap_from_samples(sr._warm_seed_cloud(t), cover_frac=0.0) + s.integrate_log(t.as_lnfunc(), *t.params, no_protect_names=True, + n=N_CHUNK, n_adapt=100, floor_level=0.0, tempering_exp=0.1, + neff=NEFF, nmax=2 * N_CHUNK, save_intg=True) + assert s._chunk_mix_parts, "no per-member densities retained; the check would be vacuous" + + own = set(k for k in vars(s) if k.startswith("portfolio_escape") or + k in ("portfolio_weight_log_total", "portfolio_share_log_num", "_member_index")) + assert own, "diagnostic namespace not found" + before = {} + for k, v in vars(s).items(): + before[k] = v.copy() if isinstance(v, np.ndarray) else v + + n = len(next(iter(s._chunk_mix_parts.values()))) + lw = np.linspace(-3.0, 1.0, n) + qm = np.ones(n) + lw_in, qm_in = lw.copy(), qm.copy() + s._update_support_diagnostics(lw, qm) + + assert np.array_equal(lw, lw_in), "log_weights were modified in place" + assert np.array_equal(qm, qm_in), "q_mix was modified in place" + changed = [] + for k, v in vars(s).items(): + if k in own: + continue + old = before.get(k, "") + same = (np.array_equal(v, old) if isinstance(v, np.ndarray) + else (v is old or v == old if not isinstance(old, np.ndarray) else False)) + if not same: + changed.append(k) + assert not changed, "diagnostic mutated sampler state outside its namespace: {}".format(changed) + + # NON-VACUITY: the same comparison must SEE a deliberate mutation, otherwise "nothing changed" + # proves only that the comparison is blind. + s.ntotal = s.ntotal + 1 + seen = [k for k, v in vars(s).items() + if k not in own and not isinstance(v, np.ndarray) and + not (v is before.get(k, "") or v == before.get(k, ""))] + assert "ntotal" in seen, "the mutation check cannot detect a change; it is vacuous" + + +def test_escaped_mass_fires_on_a_misplaced_seed(): + """Matched seed -> first-chunk escaped mass ~0; seed displaced clear of the peak -> ~1. + + Uses the FIRST-CHUNK statistic, which is the one that isolates the seed's own live volume: + the cumulative statistic also absorbs the ordinary contraction of a correctly-placed member + and has a large, target-dependent floor (measured median 0.35 at d=4 / 0.82 at d=6).""" + _, _, _, ok = _run("avgmm_cold", 0.0) + _, _, _, bad = _run("avgmm_cold", 4.0) + e_ok = float(ok["portfolio_escaped_mass_early"][0]) + e_bad = float(bad["portfolio_escaped_mass_early"][0]) + assert 0.0 <= e_ok <= 1.0 and 0.0 <= e_bad <= 1.0, "escaped_mass out of [0,1]" + assert e_bad > 0.9, "misplaced seed did NOT fire the detector: early escaped_mass={}".format(e_bad) + assert e_ok < 1e-2, "matched seed produced a false positive: early escaped_mass={}".format(e_ok) + assert bool(ok["portfolio_member_hard_edged"][0]) or e_ok == 0.0 + + +def test_statistic_reads_the_support_and_not_a_proxy(): + """Same misplaced seed, but the AV member's density is floored strictly positive: with no + region outside its support, escaped_mass MUST read 0 and the member must not be hard-edged. + This is the break-it check for the "it fires" assertion above -- everything else about the + run (the displaced seed, the low n_eff, the concentrated weights) is unchanged.""" + _, _, _, bad = _run("avgmm_cold", 4.0) + _, _, _, soft = _run("avgmm_cold", 4.0, soften_av=True) + assert float(bad["portfolio_escaped_mass_early"][0]) > 0.9, \ + "control did not fire; the comparison would be vacuous" + assert float(soft["portfolio_escaped_mass"][0]) == 0.0, \ + "escaped_mass nonzero for a member with strictly positive density: it is not reading support" + assert not bool(soft["portfolio_member_hard_edged"][0]) + assert float(soft["portfolio_escaped_mass_max"]) == 0.0 + + +def test_soft_member_is_not_scored_as_hard_edged(): + """A live GMM's density is nowhere exactly zero on these targets, so it must be reported + hard_edged=False and must not enter escaped_mass_max (a soft member reading 0 would otherwise + drag a max/mean down and mask a fully-escaped AV member).""" + _, _, _, bad = _run("avgmm_cold", 4.0) + hard = np.asarray(bad["portfolio_member_hard_edged"], dtype=bool) + esc = np.asarray(bad["portfolio_escaped_mass"], dtype=float) + assert hard[0], "the AV member must be observed hard-edged" + if not hard[1]: + assert esc[1] == 0.0 + assert float(bad["portfolio_escaped_mass_max"]) == esc[0], \ + "escaped_mass_max must ignore the soft member" + + +if __name__ == "__main__": + for fn in (test_diagnostic_is_off_path, + test_diagnostic_mutates_nothing_outside_its_namespace, + test_escaped_mass_fires_on_a_misplaced_seed, + test_statistic_reads_the_support_and_not_a_proxy, + test_soft_member_is_not_scored_as_hard_edged): + fn() + print("PASS", fn.__name__) diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_seq_gmm_check.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_seq_gmm_check.py new file mode 100644 index 000000000..4716d45eb --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_seq_gmm_check.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python +"""The sequential-case GMM clearing check must not pass VACUOUSLY. + +A check keyed only on `cleared is False` treats three distinct failures as success: the member +never trained (so "cleared" is trivially true and proves nothing), the inspection raised (cleared +is None), or the fields are absent entirely. A check that could not run is a failed check. + +Run: python test_seq_gmm_check.py +""" +import sys + +import shape_recovery as SR + + +def _rec(**kw): + """A record that passes every ordinary metric, so only the GMM check can move the verdict.""" + r = dict(kind="portfolio_seq_nobs", n_eff=2000.0, n_ess=9000.0, + js=[0.0001, 0.0001], js_floor=[0.0005, 0.0005], + mean_pull=[0.005, 0.005], width_ratio=[1.001, 1.001], + corr_diff_max=0.005, rel_err=0.01, bias_ln=0.002, error=None) + r.update(kw) + return r + + +def test_healthy_case_passes(): + st, why = SR.evaluate(_rec(seq_gmm_trained_before_reset=True, seq_gmm_cleared=True)) + assert st == "PASS", (st, why) + + +def test_leak_fails(): + st, why = SR.evaluate(_rec(seq_gmm_trained_before_reset=True, seq_gmm_cleared=False)) + assert st == "FAIL", (st, why) + assert any("did NOT clear" in w for w in why), why + + +def test_never_trained_is_not_a_pass(): + """"Cleared" is trivially true if nothing was ever trained -- that must not read as success.""" + st, why = SR.evaluate(_rec(seq_gmm_trained_before_reset=False, seq_gmm_cleared=True)) + assert st == "FAIL", "a vacuous 'cleared' verdict passed: {} {}".format(st, why) + assert any("vacuous" in w for w in why), why + + +def test_inspection_failure_is_not_a_pass(): + st, why = SR.evaluate(_rec(seq_gmm_trained_before_reset=True, seq_gmm_cleared=None, + seq_gmm_error="AttributeError: no integrator")) + assert st == "FAIL", "an unreadable GMM state passed: {} {}".format(st, why) + assert any("AttributeError" in w for w in why), "the underlying error was not surfaced: {}".format(why) + + +def test_missing_fields_are_not_a_pass(): + st, why = SR.evaluate(_rec()) + assert st == "FAIL", "absent instrumentation passed: {} {}".format(st, why) + + +def test_av_seq_is_exempt(): + """AV_seq has no GMM member, so the check must not fire on it.""" + st, why = SR.evaluate(_rec(kind="AV_seq")) + assert st == "PASS", (st, why) + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print("PASS", name) + print("sequential GMM clearing check cannot pass vacuously") diff --git a/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_shape_recovery.py b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_shape_recovery.py new file mode 100755 index 000000000..cfbc6f559 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators/test_shape_recovery.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python +"""pytest wrapper for the shape-recovery merge gate. + +Guarded by an env var so ordinary `pytest` sweeps stay fast: + + RIFT_RUN_EXPENSIVE=1 pytest -v test_shape_recovery.py # quick matrix + RIFT_RUN_EXPENSIVE=1 RIFT_SHAPE_PRESET=standard pytest -v ... # full gate + +The canonical merge-gate invocation is run_shape_recovery.sh (JSON output + +base-vs-candidate comparison); this wrapper exists so the suite also shows up +in standard pytest tooling. +""" +import os + +import pytest + +from shape_recovery import MixtureTarget, PRESETS, evaluate, run_one + +pytestmark = pytest.mark.skipif( + not os.environ.get("RIFT_RUN_EXPENSIVE"), + reason="expensive merge-gate suite; set RIFT_RUN_EXPENSIVE=1") + +_PRESET = PRESETS[os.environ.get("RIFT_SHAPE_PRESET", "quick")] +_STRICT = os.environ.get("RIFT_SHAPE_STRICT", "AV,GMM").split(",") + +_MATRIX = [(kind, d, nc, ts) + for kind in _STRICT + for d in _PRESET["dims"] + for nc in _PRESET["ncomps"] + for ts in _PRESET["seeds"]] + + +@pytest.mark.parametrize("kind,ndim,ncomp,tseed", _MATRIX) +def test_shape_recovery(kind, ndim, ncomp, tseed): + target = MixtureTarget(ndim, ncomp, tseed) + r = run_one(kind, target, _PRESET["nmax_per_dim"] * ndim, _PRESET["neff"]) + ok, reasons = evaluate(r) + assert ok, "{} on {}: {}".format(kind, target.name, "; ".join(reasons)) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/README_benchmark.md b/MonteCarloMarginalizeCode/Code/test/integrators/README_benchmark.md new file mode 100644 index 000000000..86e384b0c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/README_benchmark.md @@ -0,0 +1,35 @@ +# Integrator benchmark harness (`benchmark_integrators.py`) + +A quantitative, API-matching benchmark for RIFT Monte-Carlo integrators, giving a +*testable* definition of "a better integrator". + +## Targets (known truth) +- `corrgauss{3,5,8}` — scaled correlated Gaussian, one narrow dim (generalizes the + CI target `test_mcsamplerEnsemble_extended.py` to arbitrary D). +- `rosenbrock` — 2-D Rosenbrock, true log-evidence -5.804. +- `gaussmix{4,8}` — superposition of Gaussians (the high-D stress test where AV + degrades; cf. FinerNet `multigauss_direct`). + +## Metrics +`bias_ln = lnI - lnZ_true`, fractional MC error `sqrt(var)/I`, RIFT `n_eff = Σp/max p`, +Kish `n_ESS = (Σw)²/Σw²`, **efficiency `n_eff/N_eval`** (headline), `N_eval`/wallclock to +reach the target `neff`, and Jensen–Shannon divergence (nats) of recovered vs analytic +1-D marginals. + +## Backend +Selected by the caller's environment, exactly as production ILE: `CUDA_VISIBLE_DEVICES=""` +→ CPU/numpy; set to an idle GPU index → cupy. Each row reports the backend actually used. + +## Run +``` +source ~/RIFT_develUWM/bin/activate +export PYTHONPATH=/MonteCarloMarginalizeCode/Code:$PYTHONPATH +export CUDA_VISIBLE_DEVICES=1 +python benchmark_integrators.py --target gaussmix4 --samplers default,AC,GMM,AV --nmax 200000 --neff 1000 --json out.json +``` + +## Cold-vs-warm +`run(..., warm_start=callable(sampler,target))` seeds prior information before +`integrate()`, for measuring bootstrap gains (see the bootstrappable-AV work). + +NOTE: the benchmarks here measure efficiency/accuracy interactively; the pre-merge REQUIREMENT is the shape-recovery gate in ../expensive_before_merging/integrators/ (see RIFT/integrators/TESTING.md). diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_coverfrac.sh b/MonteCarloMarginalizeCode/Code/test/integrators/bench_coverfrac.sh new file mode 100755 index 000000000..70646edea --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_coverfrac.sh @@ -0,0 +1,46 @@ +#!/bin/bash +# +# bench_coverfrac.sh -- answer the cardassia cover_frac ask on the REFERENCE point/backend. +# +# Runs the S250114ax **iteration-0 worker** point (overlap-grid-0.xml.gz, --event 0) -- the point the +# reference matrix in BREADCRUMB_av_neff_reproduction.md actually used (lnLmax~1212, +# sqrt(2 lnLmax)~49.2) -- on GPU, varying ONLY the warm-start coverage floor, with the auto-rescue +# DISABLED so nothing masks the seed's own behaviour. +# +# NB this is deliberately NOT the "loud on-source coinc point" (lnLmax~3020): that is a different, +# much louder target, and mixing the two is what made the reference look unreproducible. +# +# Usage: NAME= COVER= [GPU=n] bench_coverfrac.sh [extra ILE flags...] +set -u +WT=/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c +CODE=$WT/MonteCarloMarginalizeCode/Code +PIPE=/home/richard.oshaughnessy/RIFT_roboto_paper/analyses/integrator_demos/S250114ax_pipeline +RUNPE=$PIPE/run_PE +SEED=$PIPE/pe_warm_seed.dat +OUTDIR=$RUNPE/iteration_0_ile +NAME=${NAME:?set NAME}; COVER=${COVER:?set COVER}; GPU=${GPU:-0} +NCHUNK=${NCHUNK:-10000}; NMAX=${NMAX:-4000000}; NEFF=${NEFF:-999} + +export PYTHONPATH=$CODE:${PYTHONPATH:-} +export PATH=$CODE/bin:$PATH +export PYTHONUNBUFFERED=1 CUDA_VISIBLE_DEVICES=$GPU OMP_NUM_THREADS=2 + +LOG=$OUTDIR/cf_${NAME}.log +cd $OUTDIR +echo "# COVERFRAC bench NAME=$NAME cover=$COVER gpu=$GPU nchunk=$NCHUNK nmax=$NMAX (rescue OFF) $(date)" > $LOG + +/home/richard.oshaughnessy/RIFT_develUWM/bin/python -u $CODE/bin/integrate_likelihood_extrinsic_batchmode \ + --save-P 0.1 --fmax 1792.0 --cache $PIPE/local.cache --event-time 1420878141.22266 \ + --channel-name H1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file H1=$RUNPE/H1-psd.xml.gz --fmin-ifo H1=20 \ + --channel-name L1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file L1=$RUNPE/L1-psd.xml.gz --fmin-ifo L1=20 \ + --fmin-template 20.0 --reference-freq 20 --d-max 10000 \ + --data-start-time 1420878135.222656 --data-end-time 1420878143.222656 --inv-spec-trunc-time 0 \ + --window-shape 0.1 --time-marginalization --inclination-cosine-sampler --declination-cosine-sampler \ + --n-max $NMAX --n-eff $NEFF --n-chunk $NCHUNK --vectorized --gpu --srate 4096 \ + --adapt-weight-exponent 0.1 --l-max 2 --approx IMRPhenomD --force-xpy \ + --internal-waveform-fd-L-frame --n-events-to-analyze 1 \ + --sim-xml $RUNPE/overlap-grid-0.xml.gz --event 0 \ + --sampler-method AV \ + --sampler-warmstart-samples $SEED --sampler-warmstart-cover-frac $COVER --sampler-warmstart-inflate 1.3 \ + "$@" --output-file $OUTDIR/cf_${NAME}.xml >> $LOG 2>&1 +echo "# EXIT $? $(date)" >> $LOG diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_multi_event.py b/MonteCarloMarginalizeCode/Code/test/integrators/bench_multi_event.py new file mode 100644 index 000000000..9e802edf3 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_multi_event.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python3 +""" +Multi-event robustness check for the portfolio freeze-policy change. + +For each event we take its REAL iteration-0 ILE worker args (real strain/PSD/grid, from the +event's ILE.sub), strip the gwsignal/SEOBNR *container-only* flags, and swap the waveform to +the bare-venv-native IMRPhenomXPHM (l-max 4, precessing). The waveform is only the integrand; +the point of this test is the INTEGRATOR: does the AV+GMM portfolio (VARAHA never-freeze, the +new default) REPLICATE the standalone-AV integral (same ln Z within MC error) and converge to a +comparable n_eff, across a spread of real events? + +Usage: + bench_multi_event.py smoke # tiny-budget single AV run (plumbing check) + bench_multi_event.py run # one run; sampler... e.g. AV OR portfolio --sampler-portfolio AV,GMM +Env: GPU (default 2), NEFF (default 40), NMAX (default 2000000) + +ln Z / n_eff are read back from the ILE output .xml_0_.dat (lnZ = field[-4], neff=field[-1]). +""" +import re, shlex, subprocess, sys, os + +EVENTS_BASE = "/home/richard.oshaughnessy/unixhome/Projects/LIGO-ScienceMode/O4_era/RIFT_roboto_paper/analyses/rerun_o4ab_distance_export/project/working" +WT = "/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c" +CODE = WT + "/MonteCarloMarginalizeCode/Code" +BIN = CODE + "/bin/integrate_likelihood_extrinsic_batchmode" +PY = "/home/richard.oshaughnessy/RIFT_develUWM/bin/python" + +def event_dir(event): + return "{}/{}/rift-distexport-nocal".format(EVENTS_BASE, event) + +# exact production container for these events (from the event ILE.sub SingularityImage line) +SIF = "/home/richard.oshaughnessy/rift_cit_build_container_family/built_containers/rift_o4d-calmarg_in_loop_cc60-90_cuda118_20260615b.sif" + +def build_args(event, container=False): + # Parse the pipeline-generated command-single.sh: it holds the fully-formed, correctly-quoted + # single-worker ILE command (in particular the --internal-waveform-extra-kwargs dict), so we + # avoid re-deriving condor's '""'/''-escaping from ILE.sub (which mangled the nested quotes). + edir = event_dir(event) + txt = open(edir + "/command-single.sh").read() + m = re.search(r'^\S*integrate_likelihood_extrinsic_batchmode\s+(.*)$', txt, re.M) + argv = shlex.split(m.group(1)) + if not container: + # BARE-VENV fallback (unused for the robustness suite): swap to a native waveform and drop + # gwsignal/cosmo-prior flags. FAILS on grids with transverse spins (RIFT cannot recover + # modes from H+/Hx for precessing configs) -- that is why the suite uses the container. + drop, out, skip = {"--use-gwsignal", "--force-gpu-only"}, [], 0 + for tok in argv: + if skip: skip = 0; continue + if tok == "--internal-waveform-extra-kwargs": skip = 1; continue + if tok in drop: continue + out.append("IMRPhenomXHM" if tok == "SEOBNRv5PHM" else tok) + argv = out + return edir, argv + +def filtered(argv, container=False): + """Drop flags we override: sampler-method, n-eff, n-max, output-file, n-events-to-analyze, + event (we re-add --event 0 --n-events-to-analyze 1 to integrate a single intrinsic point). + In BARE mode also drop --d-prior (cosmo_* needs cupyx.scipy.interpolate, absent in the old + venv cupy 10.6); in CONTAINER mode the cosmo prior works, so keep it (faithful integrand).""" + drop_val = {"--sampler-method", "--n-eff", "--n-max", "--n-events-to-analyze", "--event"} + if not container: + drop_val.add("--d-prior") + out, skip = [], 0 + for tok in argv: + if skip: + skip = 0; continue + if tok in drop_val: + skip = 1; continue + if tok.startswith("--output-file") or tok.startswith("--event="): + continue + out.append(tok) + return out + +def run(event, tag, sampler_extra, neff, nmax, gpu, container=False, wrap=True): + """container=True builds the FULL production args (gwsignal/SEOBNRv5PHM/cosmo prior). + wrap=False runs python directly instead of nesting singularity -- use when the job is ALREADY + inside the container (e.g. condor supplied it via MY.SingularityImage).""" + edir, argv = build_args(event, container=container) + argv = filtered(argv, container=container) + out = "{}/mev_{}.xml".format(edir, tag) + argv += ["--sampler-method"] + sampler_extra + ["--n-eff", str(int(neff)), "--n-max", str(int(nmax)), + "--n-events-to-analyze", "1", "--event", "0", + "--output-file", out] + log = "{}/mev_{}.log".format(edir, tag) + if container and wrap: + # Run inside the event's production container (real SEOBNRv5PHM + gwsignal + cuda118 cupy) + # but force MY worktree RIFT onto PYTHONPATH so the CONTAINER supplies the waveform/cupy + # stack while the INTEGRATOR code under test is this branch's. Bind ceph frames + the + # worktree; $HOME is auto-mounted so edir/local.cache/PSDs resolve. + inner = ("cd {edir} && PYTHONPATH={code}:$PYTHONPATH PATH={code}/bin:$PATH " + "CUDA_VISIBLE_DEVICES={gpu} OMP_NUM_THREADS=2 PYTHONUNBUFFERED=1 " + "python -u {bin} {args}").format( + edir=edir, code=CODE, gpu=gpu, bin=BIN, + args=" ".join(shlex.quote(a) for a in argv)) + cmd = ["singularity", "exec", "--nv", "--bind", "/ceph", "--bind", "/cvmfs", + "--bind", WT, SIF, "bash", "-c", inner] + env = dict(os.environ) + else: + # direct exec: either the bare-venv fallback, or we are already inside the container + _py = "python" if (container and not wrap) else PY + cmd = [_py, "-u", BIN] + argv + env = dict(os.environ) + env["PYTHONPATH"] = CODE + ":" + env.get("PYTHONPATH", "") + env["PATH"] = CODE + "/bin:" + env["PATH"] + env["PYTHONUNBUFFERED"] = "1"; env["CUDA_VISIBLE_DEVICES"] = str(gpu); env["OMP_NUM_THREADS"] = "2" + with open(log, "w") as lf: + lf.write("# MULTIEVENT {} tag={} gpu={} neff={} nmax={} container={}\n# {}\n".format( + event, tag, gpu, neff, nmax, container, " ".join(argv))) + lf.flush() + rc = subprocess.call(cmd, cwd=edir, stdout=lf, stderr=subprocess.STDOUT, env=env) + lf.write("\n# EXIT {}\n".format(rc)) + return log + +def read_result(event, tag): + dat = "{}/mev_{}.xml_0_.dat".format(event_dir(event), tag) + try: + with open(dat) as f: + line = f.readline().split() + vals = [float(x) for x in line] + return {"lnZ": vals[-4], "sigOverL": vals[-3], "ntot": vals[-2], "neff": vals[-1]} + except Exception as e: + return {"error": str(e)} + +if __name__ == "__main__": + cmd = sys.argv[1] + gpu = int(os.environ.get("GPU", 2)); neff = float(os.environ.get("NEFF", 40)); nmax = int(os.environ.get("NMAX", 2000000)) + container = os.environ.get("CONTAINER", "1") == "1" # default: faithful container path + wrap = os.environ.get("NO_SINGULARITY", "0") != "1" # 0 => nest singularity; 1 => already inside + if cmd == "smoke": + log = run(sys.argv[2], "smoke", ["AV"], neff=999, nmax=60000, gpu=gpu, container=container, wrap=wrap) + print("smoke log:", log); print(read_result(sys.argv[2], "smoke")) + elif cmd == "run": + event, tag = sys.argv[2], sys.argv[3]; sampler_extra = sys.argv[4:] + log = run(event, tag, sampler_extra, neff, nmax, gpu, container=container, wrap=wrap) + print(tag, read_result(event, tag), "log:", log) + elif cmd == "read": + print(sys.argv[2], sys.argv[3], read_result(sys.argv[2], sys.argv[3])) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource.sh b/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource.sh new file mode 100755 index 000000000..54f1101a2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# +# bench_onsource.sh -- the HIGH-SNR BEST-FIT (on-source) point, not a trial grid point. +# +# Two DIFFERENT problems have been conflated in this study; this script exists to keep them apart: +# * integrator tuning on a TRIAL point : overlap-grid-0.xml.gz --event 0 (m1/m2 28.29/26.69, +# lnLmax~1212). Fine for A/B-ing integrator policy; says nothing about the science target. +# * rescuing the BEST-FIT evaluation : target_params.xml.gz (m1/m2 37.71/34.03, lnLmax~3020, +# rho~78). THIS is the one that has to work for a high-SNR event, and where AV stalls at +# n_eff~1 (measured independently on cardassia CPU and reproduced here). +# +# The PE warm seed is the production extrinsic posterior FOR THIS EVENT, so warm-starting is +# physically appropriate here in a way it is not for an off-source trial point. Coverage floor +# defaults to the bias-safe 0.5 (do NOT use 0.05/0: under-covered, biases ln Z low). +# +# Usage: NAME= [COVER=0.5] [GPU=n] bench_onsource.sh +set -u +# WT = the checkout whose RIFT code is exercised. Override via env to A/B this branch against a +# base worktree (e.g. WT=/path/to/rift_O4d_worktree) with an otherwise identical command line. +WT=${WT:-/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c} +CODE=$WT/MonteCarloMarginalizeCode/Code +PIPE=/home/richard.oshaughnessy/RIFT_roboto_paper/analyses/integrator_demos/S250114ax_pipeline +RUNPE=$PIPE/run_PE +SEED=$PIPE/pe_warm_seed.dat +OUTDIR=$RUNPE/iteration_0_ile +NAME=${NAME:?set NAME}; COVER=${COVER:-0.5}; GPU=${GPU:-0} +NCHUNK=${NCHUNK:-10000}; NMAX=${NMAX:-4000000}; NEFF=${NEFF:-999} +WARM=${WARM:-1} + +export PYTHONPATH=$CODE:${PYTHONPATH:-} +export PATH=$CODE/bin:$PATH +export PYTHONUNBUFFERED=1 CUDA_VISIBLE_DEVICES=$GPU OMP_NUM_THREADS=2 + +WARMFLAGS=() +[ "$WARM" = "1" ] && WARMFLAGS=( --sampler-warmstart-samples $SEED \ + --sampler-warmstart-cover-frac $COVER --sampler-warmstart-inflate 1.3 ) + +LOG=$OUTDIR/os_${NAME}.log +cd $OUTDIR +echo "# ONSOURCE bench NAME=$NAME cover=$COVER warm=$WARM gpu=$GPU nmax=$NMAX $(date)" > $LOG +echo "# point: target_params.xml.gz (best-fit, m1/m2 37.71/34.03) -- NOT the trial grid point" >> $LOG + +/home/richard.oshaughnessy/RIFT_develUWM/bin/python -u $CODE/bin/integrate_likelihood_extrinsic_batchmode \ + --save-P 0.1 --fmax 1792.0 --cache $PIPE/local.cache --event-time 1420878141.22266 \ + --channel-name H1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file H1=$RUNPE/H1-psd.xml.gz --fmin-ifo H1=20 \ + --channel-name L1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file L1=$RUNPE/L1-psd.xml.gz --fmin-ifo L1=20 \ + --fmin-template 20.0 --reference-freq 20 --d-max 10000 \ + --data-start-time 1420878135.222656 --data-end-time 1420878143.222656 --inv-spec-trunc-time 0 \ + --window-shape 0.1 --time-marginalization --inclination-cosine-sampler --declination-cosine-sampler \ + --n-max $NMAX --n-eff $NEFF --n-chunk $NCHUNK --vectorized --gpu --srate 4096 \ + --adapt-weight-exponent 0.1 --l-max 2 --approx IMRPhenomD --force-xpy \ + `# CUBIC Q_lm time interpolation instead of nearest-sample-bin. Requires the maintained NoLoop` \ + `# likelihood, i.e. the --vectorized --gpu --force-xpy combo set above. Removes a superfluous` \ + `# extrinsic non-smoothness (time quantization), which makes convergence more robust.` \ + --interpolate-time True \ + --internal-waveform-fd-L-frame --n-events-to-analyze 1 \ + --sim-xml $RUNPE/target_params.xml.gz --event 0 \ + "${WARMFLAGS[@]}" "$@" --output-file $OUTDIR/os_${NAME}.xml >> $LOG 2>&1 +echo "# EXIT $? $(date)" >> $LOG diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource_ensemble.sh b/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource_ensemble.sh new file mode 100755 index 000000000..475a773db --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource_ensemble.sh @@ -0,0 +1,51 @@ +#!/usr/bin/env bash +# Seed ensemble on the S250114ax best-fit on-source point, comparing GMM-coverage configs for +# RELIABILITY (n_eff distribution) and EXTRINSIC-POSTERIOR STABILITY. The goal is NOT max n_eff: +# it is MODEST, RELIABLE n_eff with a stable, unbiased extrinsic posterior that preserves the real +# degeneracy structure (sky ring / dL-inclination / psi-phi), i.e. no copy silently collapsing a mode. +# +# Uses --extrinsic-proposal-output (weight-correct GMM fit of the run's TRUE-weighted extrinsic +# posterior) instead of --save-samples, which is unusable here (fairdraw/save-P/log-cumsum/lnL-only, +# see DESIGN_portfolio_freeze_policy.md). Compare the breadcrumbs with compare_extrinsic_breadcrumbs.py. +# +# Usage: bench_onsource_ensemble.sh [GPU] [SEEDS...] (default GPU 2, seeds 1 2 3) +set -u +HERE="$(cd "$(dirname "$0")" && pwd)" +BENCH="$HERE/bench_onsource.sh" +GPU="${1:-2}"; shift || true +SEEDS=("$@"); [ ${#SEEDS[@]} -eq 0 ] && SEEDS=(1 2 3) +RUNPE=~/RIFT_roboto_paper/analyses/integrator_demos/S250114ax_pipeline/run_PE/iteration_0_ile +MAXCONC=3 + +# The bench does NOT select a sampler (driver default is adaptive_cartesian_gpu, the classic GPU +# Cartesian sampler -- NOT VARAHA/AV, NOT the portfolio). Every config must explicitly request the +# AV(=mcsamplerAdaptiveVolume)+GMM(=mcsamplerEnsemble) portfolio, or the --internal-gmm-* flags are inert. +# canonical form: --sampler-portfolio is action='append', so repeat the flag (one member each). +# (comma-separated AV,GMM also works via a recent driver split, but repeated flags are the robust form.) +PORTFOLIO="--sampler-method portfolio --sampler-portfolio AV --sampler-portfolio GMM" + +# config name -> extra GMM-coverage flags (portfolio selection prepended in run_one). +declare -A CFG +CFG[cap8]="--internal-gmm-adaptive-components --internal-gmm-max-components 8 --internal-gmm-inflate 1.0" +CFG[cap16]="--internal-gmm-adaptive-components --internal-gmm-max-components 16 --internal-gmm-inflate 1.3" +CFG[corr]="--internal-gmm-correlate-all --internal-gmm-adaptive-components --internal-gmm-max-components 8" + +run_one() { + local cfg="$1" + local seed="$2" + local name="e_${cfg}_s${seed}" + GPU=$GPU NAME="$name" bash "$BENCH" $PORTFOLIO ${CFG[$cfg]} \ + --seed "$seed" --extrinsic-proposal-output "$RUNPE/ext_${name}.npz" +} + +# simple 3-wide job pool +running=0 +for cfg in cap8 cap16 corr; do + for s in "${SEEDS[@]}"; do + run_one "$cfg" "$s" & + running=$((running+1)) + if [ $running -ge $MAXCONC ]; then wait -n 2>/dev/null || wait; running=$((running-1)); fi + done +done +wait +echo "ENSEMBLE DONE" diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_portfolio_freeze.sh b/MonteCarloMarginalizeCode/Code/test/integrators/bench_portfolio_freeze.sh new file mode 100755 index 000000000..d4ea76f52 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_portfolio_freeze.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# +# bench_portfolio_freeze.sh -- S250114ax iteration-0 ILE benchmark for the portfolio +# freeze-policy tuning (grace / revive / VARAHA-never-freeze). Runs ONE integrator +# configuration to the n_eff target or the n-max cap and tees the trajectory to a log. +# +# The core args are the iteration-0 S250114ax worker args (from run_PE/ILE.sub, macros +# substituted: macroiteration->0 macroevent->0 macrongroup->1). The sampler / warm-start / +# freeze-policy flags are supplied by the caller so one script drives every variant. +# +# Usage: +# NAME= [GPU=2] [WARM=1] bench_portfolio_freeze.sh +# WARM=1 -> append the PE-oracle warm-start flags (cover 0.05 / inflate 1.3 / retry 5) +# WARM=0 -> cold (no warm start) [default] +# Writes $OUTDIR/frz_.log and $OUTDIR/frz_.xml_0_.dat +# +# Examples: +# NAME=av_warm WARM=1 bench_portfolio_freeze.sh --sampler-method AV +# NAME=pf_neverfreeze_warm WARM=1 bench_portfolio_freeze.sh --sampler-method portfolio --sampler-portfolio AV,GMM +# NAME=pf_canfreeze_warm WARM=1 bench_portfolio_freeze.sh --sampler-method portfolio --sampler-portfolio AV,GMM --portfolio-varaha-can-freeze +# +set -u + +WT=/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c +CODE=$WT/MonteCarloMarginalizeCode/Code +PIPE=/home/richard.oshaughnessy/RIFT_roboto_paper/analyses/integrator_demos/S250114ax_pipeline +RUNPE=$PIPE/run_PE +SEED=$PIPE/pe_warm_seed.dat +OUTDIR=$RUNPE/iteration_0_ile +GPU=${GPU:-2} +WARM=${WARM:-0} +NAME=${NAME:?set NAME=} + +# use MY worktree's source tree (has the freeze-policy code + CLI flags); the installed venv bin is stale +export PYTHONPATH=$CODE:${PYTHONPATH:-} +export PATH=$CODE/bin:$PATH +export PYTHONUNBUFFERED=1 +export CUDA_VISIBLE_DEVICES=$GPU +export OMP_NUM_THREADS=2 + +BIN=$CODE/bin/integrate_likelihood_extrinsic_batchmode +LOG=$OUTDIR/frz_${NAME}.log +OUT=$OUTDIR/frz_${NAME}.xml + +CORE=( --save-P 0.1 --fmax 1792.0 --cache $PIPE/local.cache --event-time 1420878141.22266 \ + --channel-name H1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file H1=$RUNPE/H1-psd.xml.gz --fmin-ifo H1=20 \ + --channel-name L1=DCS-CALIB_STRAIN_CLEAN_AR01 --psd-file L1=$RUNPE/L1-psd.xml.gz --fmin-ifo L1=20 \ + --fmin-template 20.0 --reference-freq 20 --d-max 10000 \ + --data-start-time 1420878135.222656 --data-end-time 1420878143.222656 --inv-spec-trunc-time 0 \ + --window-shape 0.1 --time-marginalization --inclination-cosine-sampler --declination-cosine-sampler \ + --n-max 4000000 --n-eff 100 --vectorized --gpu --srate 4096 --adapt-weight-exponent 0.1 --l-max 2 \ + --approx IMRPhenomD --force-xpy --internal-waveform-fd-L-frame --n-events-to-analyze 1 \ + --sim-xml $RUNPE/overlap-grid-0.xml.gz --event 0 ) + +WARMFLAGS=() +if [ "$WARM" = "1" ]; then + WARMFLAGS=( --sampler-warmstart-samples $SEED --sampler-warmstart-cover-frac 0.05 \ + --sampler-warmstart-inflate 1.3 --sampler-warmstart-retry-neff 5 ) +fi + +cd $OUTDIR +echo "# BENCH $NAME GPU=$GPU WARM=$WARM $(date)" > $LOG +echo "# extra flags: $*" >> $LOG +echo "# BIN=$BIN" >> $LOG +/home/richard.oshaughnessy/RIFT_develUWM/bin/python -u "$BIN" \ + "${CORE[@]}" "${WARMFLAGS[@]}" "$@" --output-file "$OUT" >> $LOG 2>&1 +echo "# EXIT $? $(date)" >> $LOG diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/bench_weight_clip.py b/MonteCarloMarginalizeCode/Code/test/integrators/bench_weight_clip.py new file mode 100644 index 000000000..36679670f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/bench_weight_clip.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python +""" +bench_weight_clip.py -- quantify the BIAS vs n_eff trade of portfolio weight clipping +(truncated importance sampling) against targets with an ANALYTIC true ln Z. + +Weight clipping caps each importance weight at tau = C*sqrt(n)*mean(w) (Ionides 2008). A single +enormous weight crushes pooled n_eff = (sum w)^2/sum w^2, so clipping can buy a large variance +reduction -- but it is a BIASED estimator (it discards the clipped mass), and the whole question is +whether the bias is small enough to be worth it. Because these targets have a known true ln Z we +can measure BOTH sides directly, and the sampler also reports the removed-mass fraction, so the +predicted bias log1p(-frac) can be checked against the measured bias. + +Sweeps C over several values (C=0 is clipping OFF, the unbiased reference) on the correlated and +uncorrelated Gaussians from test_portfolio_adaptive_alloc.py, over several seeds. + +Usage: + CUDA_VISIBLE_DEVICES=2 OMP_NUM_THREADS=2 PYTHONPATH= python bench_weight_clip.py + options: --seeds 3 --nmax 400000 --n-chunk 10000 --ndim 5 +""" +from __future__ import print_function +import argparse +import numpy as np + +import benchmark_integrators as B +import test_portfolio_adaptive_alloc as T + + +def run_clip(target, clip, n_chunk, nmax, seed, adaptive=False): + np.random.seed(seed) + port = T.build(target, ['AV', 'GMM'], n_chunk) + lnI, _, eff, _ = port.integrate_log( + T._host_lnfunc(target), *target.params, no_protect_names=True, + nmax=nmax, neff=10 ** 9, n=n_chunk, n_adapt=100, tempering_exp=0.3, + floor_level=0.0, use_lnL=True, save_intg=True, verbose=False, + portfolio_adaptive_alloc=adaptive, portfolio_weight_clip=clip) + lnI = float(B._asnumpy(lnI)) + # removed-mass fraction the sampler tracked (0 if nothing clipped) + frac = 0.0 + if np.isfinite(port.portfolio_clip_log_removed) and np.isfinite(port.portfolio_clip_log_total): + frac = float(np.exp(port.portfolio_clip_log_removed - port.portfolio_clip_log_total)) + frac = min(max(frac, 0.0), 1.0 - 1e-15) + return dict(lnI=lnI, bias=lnI - float(target.true_lnZ), n_eff=float(B._asnumpy(eff)), + clip_frac=frac, predicted_bias=float(np.log1p(-frac)) if frac > 0 else 0.0, + n_clipped=int(port.portfolio_clip_n)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ndim", type=int, default=5) + ap.add_argument("--nmax", type=int, default=400000) + ap.add_argument("--n-chunk", type=int, default=10000) + ap.add_argument("--seeds", type=int, default=3) + ap.add_argument("--clips", type=str, default="0,0.5,1,2,5,20") + args = ap.parse_args() + + clips = [float(c) for c in args.clips.split(',')] + targets = [("uncorrelated", B.CorrelatedGaussian(ndim=args.ndim, rho=0.0, narrow=0.1)), + ("correlated", T.CompoundCorrelatedGaussian(ndim=args.ndim))] + seeds = [1234 + 101 * i for i in range(args.seeds)] + + print("# weight-clip sweep: nmax={} n_chunk={} ndim={} seeds={}".format( + args.nmax, args.n_chunk, args.ndim, seeds)) + print("# clip C=0 is OFF (unbiased reference). bias = lnI - true_lnZ (mean +/- std over seeds)") + for name, tgt in targets: + print("\n== {} true_lnZ={:.4f} ==".format(name, tgt.true_lnZ)) + print("{:>6} {:>12} {:>18} {:>12} {:>12}".format( + "C", "n_eff", "bias", "clip_frac", "pred_bias")) + for c in clips: + rows = [run_clip(tgt, c, args.n_chunk, args.nmax, s) for s in seeds] + ne = np.array([r["n_eff"] for r in rows]) + bi = np.array([r["bias"] for r in rows]) + cf = np.mean([r["clip_frac"] for r in rows]) + pb = np.mean([r["predicted_bias"] for r in rows]) + print("{:>6.2f} {:>6.0f}+/-{:<5.0f} {:>+8.3f}+/-{:<7.3f} {:>12.2e} {:>+12.3f}".format( + c, ne.mean(), ne.std(), bi.mean(), bi.std(), cf, pb)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_gmm_flexible_synth.py b/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_gmm_flexible_synth.py new file mode 100644 index 000000000..63aaa1394 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_gmm_flexible_synth.py @@ -0,0 +1,104 @@ +"""Data-free synthetic benchmark for the flexible GMM component allocation +(RIFT/integrators/DESIGN_flexible_gmm.md). Mimics the S250114ax extrinsic +posterior's HARD feature: a CURVED degeneracy arc that no axis-aligned binning +(and no few-component Gaussian) can wrap. + +Reference results (n_eff vs cumulative N; GPU, seconds): + corrall k=1 -> n_eff>=100 @76k, final ~50 ; corrall k=2 -> @28k, final ~312 + corrall k=4 -> @752k (over-allocation collapse) + adaptive (BIC, k<=8) -> @220k, final ~135 (robust, unbiased, hands-free) + +6D target on a broad box (needle-ish: peak is a small fraction of the prior): + dims (2,3): a parabolic BANANA ridge <-> distance-inclination arc (curved) + dims (0,1): a strongly-correlated Gaussian <-> phase-polarization degeneracy + dims (4,5): a tight isotropic Gaussian <-> well-localized sky + +We measure n_eff vs N (cumulative samples) for correlate-all GMM proposals with +varying component count, plus (later) the flexible-k prototype. + +Usage: + python synth_bench.py [k] + mode = corrall -> single full-dim GMM, fixed k components (k from argv) + mode = adaptive -> flexible-k prototype (n_comp='adaptive') + mode = pairing -> factored pairing {(0,1),(2,3),(4,5)} k each +""" +import sys, numpy as np +np.random.seed(1234) +from RIFT.integrators import mcsamplerEnsemble + +# ---- box (broad prior) ---- +LO = np.array([-6.,-6., -6.,-30., -6.,-6.]) +HI = np.array([ 6., 6., 6., 30., 6., 6.]) + +# ---- target lnL ---- +# banana in (x2,x3): ridge v = C*(u^2 - M); narrow across ridge, broad along it +C, M, SU, SV = 3.0, 3.0, 1.2, 1.5 +# correlated pair (x0,x1) +RHO, SP = 0.92, 1.0 +_cov = np.array([[SP**2, RHO*SP*SP],[RHO*SP*SP, SP**2]]) +_covinv = np.linalg.inv(_cov) +# tight sky (x4,x5) +SK = 0.35 +MU4, MU5 = 1.0, -1.0 + +def lnL_np(x): + x = np.asarray(x) + u, v = x[:,2], x[:,3] + ban = -0.5*(u/SU)**2 - 0.5*((v - C*(u**2 - M))/SV)**2 + d0, d1 = x[:,0], x[:,1] + q = _covinv[0,0]*d0*d0 + 2*_covinv[0,1]*d0*d1 + _covinv[1,1]*d1*d1 + corr = -0.5*q + sky = -0.5*((x[:,4]-MU4)**2 + (x[:,5]-MU5)**2)/SK**2 + return ban + corr + sky + +def like(*args): + # args: one array per param, in params_ordered order + X = np.array(args).T + return lnL_np(X) # returns lnL (we run with use_lnL) + +def build_sampler(): + s = mcsamplerEnsemble.MCSampler() + for i in range(6): + s.add_parameter(str(i), left_limit=float(LO[i]), right_limit=float(HI[i]), + adaptive_sampling=True) + return s + +def run(mode, k): + s = build_sampler() + params = [str(i) for i in range(6)] + traj = [] + def hook(integrator): + traj.append((int(integrator.ntotal), float(integrator.identity_convert(integrator.eff_samp)))) + kw = dict(min_iter=5, max_iter=300, n=4000, nmax=400_000, neff=5000, + use_lnL=True, return_lnI=True, integrator_func=hook, + verbose=False, super_verbose=False) + if mode == 'corrall': + kw.update(correlate_all_dims=True, n_comp=int(k)) + elif mode == 'adaptive': + # correlate-all single group with data-driven k (BIC), cap = k + g = tuple(range(6)) + kw.update(gmm_dict={g:None}, n_comp={g:2}, gmm_adaptive={g:int(k)}) + elif mode == 'adaptpair': + gd = {(0,1):None,(2,3):None,(4,5):None} + kw.update(gmm_dict=gd, n_comp={(0,1):2,(2,3):2,(4,5):2}, + gmm_adaptive={(0,1):int(k),(2,3):int(k),(4,5):int(k)}) + elif mode == 'pairing': + gd = {(0,1):None,(2,3):None,(4,5):None} + kw.update(gmm_dict=gd, n_comp={(0,1):int(k),(2,3):int(k),(4,5):int(k)}) + integral, err2, eff, _ = s.integrate(like, *params, **kw) + return traj, float(s.identity_convert(eff)), float(s.identity_convert(integral)) + +if __name__ == '__main__': + mode = sys.argv[1] if len(sys.argv)>1 else 'corrall' + k = int(sys.argv[2]) if len(sys.argv)>2 else 2 + traj, eff, integral = run(mode, k) + label = "{}{}".format(mode, k if mode!='adaptive' else '') + print("MODE", label, "final_eff", eff, "lnI", integral) + # print n_eff-vs-N crossings + for target in [5,10,20,50,100,200,500,1000]: + cross = next((N for (N,e) in traj if e>=target), None) + print(" neff>={:<5} at N= {}".format(target, cross)) + # dump full trajectory sparsely + for i,(N,e) in enumerate(traj): + if i%10==0 or i==len(traj)-1: + print(" traj N={:>8} eff={:.1f}".format(N,e)) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_integrators.py b/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_integrators.py new file mode 100644 index 000000000..ec252dc8a --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/benchmark_integrators.py @@ -0,0 +1,442 @@ +#!/usr/bin/env python +""" +benchmark_integrators.py + +A reusable, quantitative benchmark harness for RIFT Monte-Carlo integrators. + +Goal +---- +Provide a *testable* definition of "a better integrator", matching the existing +RIFT sampler API (add_parameter / setup / integrate[_log]), on analytic targets +with known truth. Supports: + + * multiple analytic targets with known integral & known 1-D marginals: + - CorrelatedGaussian(ndim) (the CI 3-D target, generalized to any D) + - Rosenbrock2D (true log-evidence -5.804) + - GaussianMixture(ndim,ncomp) (the FinerNet multigauss high-D stress test) + * a uniform adapter over the heterogeneous sampler call/return conventions + (default mcsampler, AC/mcsamplerGPU, GMM/mcsamplerEnsemble, AV, NF, portfolio) + * the paper's quality metrics: + - integral bias ln(I_hat) - ln(Z_true) + - fractional MC error sqrt(var)/I + - n_eff (RIFT: sum p / max p) and n_ESS = (sum w)^2 / sum w^2 (Kish) + - EFFICIENCY eff = n_eff / N_eval <-- headline scaling metric + - N_eval and wallclock consumed to reach a target n_eff + - Jensen-Shannon divergence of recovered vs true 1-D marginals (nats) + * cold-vs-warm comparison: an optional `warm_start` hook that seeds a sampler + with prior information before integrate() (used by the bootstrappable-AV work). + +Backend (CPU vs GPU / xpy) is selected by the *caller's* environment exactly as +in production ILE: set CUDA_VISIBLE_DEVICES="" for CPU/numpy, or to an idle GPU +index for cupy. The harness records which backend each sampler actually used. + +This file is import-safe (no side effects) and has a CLI at the bottom. +""" +from __future__ import print_function + +import sys +import time +import json + +import numpy as np +from scipy.stats import multivariate_normal, norm +from scipy.special import logsumexp, erf + + +# ---------------------------------------------------------------------------- +# Targets +# ---------------------------------------------------------------------------- +class Target(object): + """Analytic integrand with known truth. + + Contract: + name : str + ndim : int + params : list[str] (ordered parameter names) + llim, rlim : arrays length ndim (integration box) + lnL(X) : X is (N, ndim) -> (N,) log-integrand (the 'likelihood') + true_lnZ : ln of the true integral of exp(lnL) over the box + true_marginal_pdf(dim, x) : analytic 1-D marginal density of the + *normalized posterior* (integrand/Z), or None + """ + name = "abstract" + + def lnL(self, X): + raise NotImplementedError + + # convenience: a callable of positional scalars/arrays, matching RIFT's + # `no_protect_names=True` integrand signature f(x0, x1, ...). + def as_lnfunc(self): + def ln_f(*cols): + X = np.array(cols, dtype=float).T + return self.lnL(np.atleast_2d(X)) + return ln_f + + def as_func(self): + ln_f = self.as_lnfunc() + def f(*cols): + return np.exp(ln_f(*cols)) + return f + + def true_marginal_pdf(self, dim, x): + return None + + +class CorrelatedGaussian(Target): + """Scaled multivariate normal with one narrow dimension and off-diagonal + correlation. Generalizes the CI test (test_mcsamplerEnsemble_extended.py) + to arbitrary dimension. With mu in the middle of the box and small widths, + the box captures essentially all the mass, so the true integral over the box + equals `scale` (the Gaussian integrates to 1).""" + def __init__(self, ndim=3, width=10.0, scale=100.0, seed=123456, rho=-0.1, + narrow=0.05): + self.name = "corrgauss_d{}".format(ndim) + self.ndim = ndim + self.width = width + self.scale = scale + self.params = [str(i) for i in range(ndim)] + self.llim = -0.5 * width * np.ones(ndim) + self.rlim = 0.5 * width * np.ones(ndim) + rng = np.random.RandomState(seed) + self.mu = rng.uniform(-width / 4.0, width / 4.0, ndim) + cov = np.identity(ndim) + cov[ndim - 1][ndim - 1] = narrow # one narrow dimension + cov[0][ndim - 1] = rho # a correlation + cov[ndim - 1][0] = rho + # keep SPD in high-D: shrink off-diagonals if needed + while np.min(np.linalg.eigvalsh(cov)) <= 1e-6: + cov[0][ndim - 1] *= 0.5 + cov[ndim - 1][0] *= 0.5 + self.cov = cov + self._mvn = multivariate_normal(self.mu, self.cov) + # The samplers return I = \int L * p_prior dx , with p_prior the + # normalized uniform density (1/width per dim). Over the box the + # Gaussian integrates to `scale`, so the returned quantity's truth is + # ln(scale) minus the uniform-prior log-normalization. + self.true_lnZ = np.log(scale) - np.sum(np.log(self.rlim - self.llim)) + + def lnL(self, X): + return np.atleast_1d(np.log(self.scale * self._mvn.pdf(X) + 1e-300)) + + def true_marginal_pdf(self, dim, x): + return norm.pdf(x, loc=self.mu[dim], scale=np.sqrt(self.cov[dim][dim])) + + +class Rosenbrock2D(Target): + """2-D Rosenbrock likelihood distributed with RIFT. Known true log-evidence + over the box [-5,5]^2 is -5.804 (see FinerNet neff_linear demo).""" + def __init__(self, box=5.0, lnL_offset=0.0): + self.name = "rosenbrock2d" + self.ndim = 2 + self.params = ["0", "1"] + self.llim = -box * np.ones(2) + self.rlim = box * np.ones(2) + self.lnL_offset = lnL_offset + self.true_lnZ = -5.804 + lnL_offset + + def lnL(self, X): + x1 = X[:, 0]; x2 = X[:, 1] + minus = (1.0 - x1) ** 2 + 100.0 * (x2 - x1 ** 2) ** 2 + return np.atleast_1d(self.lnL_offset - minus) + + def true_marginal_pdf(self, dim, x): + if dim == 0: + # exact 1-D marginal in x1 (unnormalized then normalized by exp(true_lnZ)) + box = self.rlim[1] + m = (1.0 / 20.0) * np.sqrt(np.pi) * (erf(10 * (box - x ** 2)) + + erf(10 * (box + x ** 2))) * np.exp(-(1 - x) ** 2) + return m / np.exp(self.true_lnZ - self.lnL_offset) + return None + + +class GaussianMixture(Target): + """Superposition of `ncomp` multivariate normals with random weights, means + and Wishart-drawn covariances (the FinerNet multigauss stress test that + exposes AV's high-dimensional degradation). Normalized so the true integral + over the box equals `scale`.""" + def __init__(self, ndim=4, ncomp=3, width=10.0, scale=100.0, seed=42, + sigma_1d=0.7, scale_x0=3.0): + self.name = "gaussmix_d{}_n{}".format(ndim, ncomp) + self.ndim = ndim + self.ncomp = ncomp + self.width = width + self.scale = scale + self.params = [str(i) for i in range(ndim)] + self.llim = -0.5 * width * np.ones(ndim) + self.rlim = 0.5 * width * np.ones(ndim) + import scipy.stats as ss + rng = np.random.RandomState(seed) + wt = rng.uniform(size=ncomp) + 0.1 + self.wt = wt / np.sum(wt) + self.means = [] + self.covs = [] + self._rvs = [] + for k in range(ncomp): + x0 = rng.uniform(-scale_x0 / np.sqrt(ndim), scale_x0 / np.sqrt(ndim), ndim) + Sig = (sigma_1d ** 2) * np.diag(rng.uniform(1.0, 2.0, ndim)) + Sig = ss.wishart.rvs(df=ndim, scale=Sig / ndim, random_state=rng) / 1.25 + Sig = np.atleast_2d(Sig) + self.means.append(x0) + self.covs.append(Sig) + self._rvs.append(multivariate_normal(x0, Sig, allow_singular=True)) + # returned I = \int L p_prior dx ; mixture integrates to `scale` over box + self.true_lnZ = np.log(scale) - np.sum(np.log(self.rlim - self.llim)) + + def lnL(self, X): + val = np.zeros(len(X)) + for k in range(self.ncomp): + val += self.wt[k] * self._rvs[k].pdf(X) + return np.atleast_1d(np.log(self.scale * val + 1e-300)) + + def true_marginal_pdf(self, dim, x): + p = np.zeros_like(x, dtype=float) + for k in range(self.ncomp): + p += self.wt[k] * norm.pdf(x, loc=self.means[k][dim], + scale=np.sqrt(self.covs[k][dim][dim])) + return p + + def sample_truth(self, n, seed=0): + rng = np.random.RandomState(seed) + counts = rng.multinomial(n, self.wt) + out = [] + for k in range(self.ncomp): + out.append(rng.multivariate_normal(self.means[k], self.covs[k], counts[k])) + return np.vstack(out) + + +# ---------------------------------------------------------------------------- +# Metric helpers +# ---------------------------------------------------------------------------- +def _asnumpy(a): + try: + import cupy + if isinstance(a, cupy.ndarray): + return cupy.asnumpy(a) + except Exception: + pass + return np.asarray(a) + + +def log_weights_from_rvs(rvs): + """Return per-sample ln(weight) = ln L + ln p_prior - ln p_sampling, from a + sampler's _rvs cache, tolerating the heterogeneous storage conventions: + log-keyed (AV/NF/portfolio) vs linear-keyed (default/AC), and GMM which + stores 'integrand' as *log* L (negative values) alongside linear priors.""" + # integrand -> ln L + if "log_integrand" in rvs: + lnL = _asnumpy(rvs["log_integrand"]).astype(float) + elif "integrand" in rvs: + L = _asnumpy(rvs["integrand"]).astype(float) + # a genuine integrand/density is >=0; negative values mean it is + # already stored as ln L (GMM under return_lnI). + lnL = L if np.nanmin(L) < 0 else np.log(L + 1e-300) + else: + raise KeyError("no integrand in _rvs; run with save_intg=True") + # prior and sampling prior, each possibly log- or linear-keyed + def _get_log(logkey, linkey, n): + if logkey in rvs: + return _asnumpy(rvs[logkey]).astype(float) + if linkey in rvs: + return np.log(_asnumpy(rvs[linkey]).astype(float) + 1e-300) + return np.zeros(n) + n = len(lnL) + lnp = _get_log("log_joint_prior", "joint_prior", n) + lnps = _get_log("log_joint_s_prior", "joint_s_prior", n) + return lnL + lnp - lnps + + +def n_ess_kish(ln_wt): + ln_wt = ln_wt - np.max(ln_wt) + w = np.exp(ln_wt) + return float(np.sum(w) ** 2 / np.sum(w ** 2)) + + +def marginal_js(target, rvs, ln_wt, dim, nbins=60): + """Jensen-Shannon divergence (nats) between the weighted-sample 1-D marginal + and the analytic true marginal, over the box for `dim`.""" + tp = target.true_marginal_pdf(dim, np.array([0.0])) + if tp is None: + return float("nan") + name = target.params[dim] + x = _asnumpy(rvs[name]).astype(float).flatten() + w = np.exp(ln_wt - np.max(ln_wt)) + lo, hi = target.llim[dim], target.rlim[dim] + edges = np.linspace(lo, hi, nbins + 1) + centers = 0.5 * (edges[:-1] + edges[1:]) + hist, _ = np.histogram(x, bins=edges, weights=w) + if hist.sum() <= 0: + return float("nan") + q = hist / hist.sum() + tpdf = target.true_marginal_pdf(dim, centers) + p = tpdf / tpdf.sum() + m = 0.5 * (p + q) + def _kl(a, b): + mask = a > 0 + return np.sum(a[mask] * np.log(a[mask] / (b[mask] + 1e-300))) + return float(0.5 * _kl(p, m) + 0.5 * _kl(q, m)) + + +# ---------------------------------------------------------------------------- +# Sampler adapter +# ---------------------------------------------------------------------------- +def build_sampler(kind, target, n_chunk=10000): + """Instantiate + add_parameter a sampler of the requested kind for `target`. + Returns (sampler, backend_str).""" + from RIFT.integrators import (mcsampler, mcsamplerEnsemble, mcsamplerGPU, + mcsamplerAdaptiveVolume) + + def uniform_pdf(d): + w = target.rlim[d] - target.llim[d] + return np.vectorize(lambda x, w=w: 1.0 / w) + + if kind == "default": + s = mcsampler.MCSampler(); backend = "cpu" + elif kind in ("AC", "adaptive_cartesian_gpu"): + s = mcsamplerGPU.MCSampler(); backend = "gpu" if mcsamplerGPU.cupy_ok else "cpu" + elif kind in ("GMM", "gmm"): + s = mcsamplerEnsemble.MCSampler(); backend = "cpu" + elif kind == "AV": + s = mcsamplerAdaptiveVolume.MCSampler(n_chunk=n_chunk) + backend = "gpu" if mcsamplerAdaptiveVolume.cupy_ok else "cpu" + elif kind == "NF": + from RIFT.integrators import mcsamplerNFlow + s = mcsamplerNFlow.MCSampler(); backend = "cpu(torch)" + else: + raise ValueError("unknown sampler kind %r" % kind) + + for d, p in enumerate(target.params): + s.add_parameter(p, uniform_pdf(d), prior_pdf=uniform_pdf(d), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + return s, backend + + +def run(kind, target, nmax=200000, neff=1000, n_chunk=10000, tempering_exp=0.1, + n_adapt=100, warm_start=None, verbose=False, seed=None): + """Run one sampler on one target and return a metrics dict. + + warm_start : optional callable(sampler, target) invoked after setup() and + before integrate(), used to seed prior information (cold-vs-warm). + """ + if seed is not None: + np.random.seed(seed) + s, backend = build_sampler(kind, target, n_chunk=n_chunk) + ln_f = target.as_lnfunc() + f = target.as_func() + params = target.params + extra = dict(n=n_chunk, n_adapt=n_adapt, floor_level=0.0, + tempering_exp=tempering_exp, neff=neff, nmax=nmax, + save_intg=True, verbose=verbose) + + # setup + optional warm start + if hasattr(s, "setup"): + try: + s.setup() + except TypeError: + pass + if warm_start is not None: + warm_start(s, target) + + t0 = time.time() + if kind == "default": + I, var, eff, _ = s.integrate(f, *params, no_protect_names=True, **extra) + lnI = np.log(I); ln_relerr = np.log(np.sqrt(var) / I) + elif kind in ("AC", "adaptive_cartesian_gpu"): + lnI, logvar, eff, _ = s.integrate(ln_f, *params, no_protect_names=True, + use_lnL=True, **extra) + lnI = float(_asnumpy(lnI)); ln_relerr = float(_asnumpy(logvar)) / 2 - lnI + elif kind in ("GMM", "gmm"): + n_iters = int(nmax / n_chunk) + lnI, logvar, eff, _ = s.integrate(ln_f, *params, min_iter=n_iters, + max_iter=n_iters, correlate_all_dims=True, + n_comp=1, use_lnL=True, return_lnI=True, **extra) + lnI = float(_asnumpy(lnI)); ln_relerr = float(_asnumpy(logvar)) / 2 - lnI + elif kind == "AV": + lnI, logvar, eff, _ = s.integrate_log(ln_f, *params, no_protect_names=True, **extra) + lnI = float(_asnumpy(lnI)); ln_relerr = 0.5 * (float(_asnumpy(logvar)) - 2 * lnI) + elif kind == "NF": + lnI, logvar, eff, _ = s.integrate_log(ln_f, *params, no_protect_names=True, **extra) + lnI = float(_asnumpy(lnI)); ln_relerr = 0.5 * (float(_asnumpy(logvar)) - 2 * lnI) + else: + raise ValueError(kind) + wall = time.time() - t0 + eff = float(_asnumpy(eff)) + n_eval = int(getattr(s, "ntotal", 0)) or int(nmax) + + # sample-based metrics + ln_wt = log_weights_from_rvs(s._rvs) + ness = n_ess_kish(ln_wt) + js = [] + for d in range(target.ndim): + js.append(marginal_js(target, s._rvs, ln_wt, d)) + js = [x for x in js if x == x] # drop nan + js_mean = float(np.mean(js)) if js else float("nan") + + return dict( + kind=kind, target=target.name, backend=backend, + ndim=target.ndim, n_eval=n_eval, wallclock=wall, + lnI=lnI, true_lnZ=float(target.true_lnZ), + bias_ln=lnI - float(target.true_lnZ), + rel_err=float(np.exp(ln_relerr)), + n_eff=eff, n_ess=ness, + efficiency=eff / max(n_eval, 1), + js_marginal=js_mean, + ) + + +# ---------------------------------------------------------------------------- +# CLI +# ---------------------------------------------------------------------------- +_TARGETS = { + "corrgauss3": lambda: CorrelatedGaussian(ndim=3), + "corrgauss5": lambda: CorrelatedGaussian(ndim=5), + "corrgauss8": lambda: CorrelatedGaussian(ndim=8), + "rosenbrock": lambda: Rosenbrock2D(), + "gaussmix4": lambda: GaussianMixture(ndim=4, ncomp=3), + "gaussmix8": lambda: GaussianMixture(ndim=8, ncomp=3), +} + + +def _fmt(r): + return ("{kind:>8s} {target:>14s} [{backend:>7s}] N={n_eval:>8d} " + "t={wallclock:6.1f}s lnI-lnZ={bias_ln:+7.3f} relerr={rel_err:7.4f} " + "neff={n_eff:9.1f} nESS={n_ess:10.1f} eff={efficiency:.2e} " + "JS={js_marginal:.4f}").format(**r) + + +def main(): + import optparse + p = optparse.OptionParser() + p.add_option("--target", default="corrgauss3") + p.add_option("--samplers", default="default,AC,GMM,AV") + p.add_option("--nmax", type=int, default=200000) + p.add_option("--neff", type=int, default=1000) + p.add_option("--n-chunk", type=int, default=10000) + p.add_option("--seed", type=int, default=123456) + p.add_option("--json", default=None, help="write results as JSON to this path") + p.add_option("--verbose", action="store_true") + opts, _ = p.parse_args() + + tgt = _TARGETS[opts.target]() + print("# target: {} ndim={} true_lnZ={:.4f}".format(tgt.name, tgt.ndim, tgt.true_lnZ)) + results = [] + for kind in opts.samplers.split(","): + kind = kind.strip() + try: + r = run(kind, tgt, nmax=opts.nmax, neff=opts.neff, n_chunk=opts.n_chunk, + verbose=opts.verbose, seed=opts.seed) + results.append(r) + print(_fmt(r)) + except Exception as e: + import traceback + print(" {:>8s} FAILED: {}".format(kind, e)) + if opts.verbose: + traceback.print_exc() + if opts.json: + with open(opts.json, "w") as fh: + json.dump(results, fh, indent=2) + print("# wrote", opts.json) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/compare_extrinsic_breadcrumbs.py b/MonteCarloMarginalizeCode/Code/test/integrators/compare_extrinsic_breadcrumbs.py new file mode 100755 index 000000000..438277fbf --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/compare_extrinsic_breadcrumbs.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +""" +compare_extrinsic_breadcrumbs.py -- cross-copy EXTRINSIC-POSTERIOR stability check. + +Reads the per-run breadcrumbs written by `--extrinsic-proposal-output` (a weight-correct GMM fit of +each run's TRUE-weighted extrinsic posterior; see RIFT/calmarg/extrinsic_handoff.py). For a seed +ensemble of one config, it answers the question that n_eff alone cannot: + + Across independent copies, is the recovered extrinsic posterior STABLE, and does it preserve the + real degeneracy structure -- sky ring (ra,dec), dL-inclination arc, psi-phi -- or does some copy + silently COLLAPSE a group (fewer modes / a shifted blob)? A collapse is a failure mode even when + n_eff looks acceptable. + +Per group we report, per copy: the number of effective mixture modes (weight > MODE_WT) and the +mixture mean + spread in the model's NORMALIZED frame (all copies share the same bounds, so the +normalized frame is directly comparable -- we un-normalize the summary to physical units too). Then +per group across the ensemble: the cross-copy scatter of the mode count and of the group mean. A +stable config has consistent mode counts and small cross-copy mean scatter. + +Usage: compare_extrinsic_breadcrumbs.py [ ...] +Groups files by config prefix (ext__s.npz). +""" +from __future__ import print_function +import sys, os, re +import numpy as np + +# repo import: RIFT/calmarg/breadcrumbs.py +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "RIFT", "..")) +try: + from RIFT.calmarg import breadcrumbs +except Exception as e: + print("cannot import RIFT.calmarg.breadcrumbs (%s); set PYTHONPATH to the repo Code dir" % e) + sys.exit(2) + +MODE_WT = 0.10 # a mixture component counts as a 'mode' if its weight exceeds this +NAME_RE = re.compile(r"ext_(?P.+?)_s(?P\d+)\.npz$") + + +def _phys_mean(g): + """Mixture mean per param, un-normalized to physical units via stored bounds. + RIFT GMM works in a [0,1]-normalized frame per dimension: x_phys = lo + x_norm*(hi-lo).""" + means = np.asarray(g["means"], dtype=float) # (K,d) normalized + w = np.asarray(g["weights"], dtype=float); w = w / w.sum() + b = np.asarray(g["bounds"], dtype=float) # (d,2) + lo, hi = b[:, 0], b[:, 1] + mu_norm = (w[:, None] * means).sum(axis=0) # (d,) + return lo + mu_norm * (hi - lo) # (d,) physical + + +def load_group_summaries(paths): + """-> dict cfg -> list of per-copy dicts {seed, neff, nsamp, groups:{gname:{K,modes,mean_phys}}}""" + out = {} + for p in paths: + m = NAME_RE.search(os.path.basename(p)) + cfg = m.group("cfg") if m else "?" + seed = int(m.group("seed")) if m else -1 + try: + bc = breadcrumbs.load(p) + except Exception as e: + print(" skip %s (%s)" % (p, str(e)[:60])); continue + meta = bc.get("meta", {}) + ext = bc.get("extrinsic") + rec = dict(seed=seed, neff=float(meta.get("neff", np.nan)), + nsamp=int(meta.get("n_samples", 0)), groups={}) + if ext: + for g in ext["groups"]: + gname = ",".join(g["params"]) + w = np.asarray(g["weights"], dtype=float) + rec["groups"][gname] = dict(K=len(w), modes=int((w > MODE_WT).sum()), + mean_phys=_phys_mean(g), + bounds=np.asarray(g["bounds"], dtype=float)) + out.setdefault(cfg, []).append(rec) + for cfg in out: + out[cfg].sort(key=lambda r: r["seed"]) + return out + + +GOOD_NEFF = 5.0 # a copy counts as 'landed' (usable posterior) above this n_eff + + +def _in_bounds(rec): + """True if the group's phys-mean lies inside its bounds -- a collapsed run's degenerate GMM + fit drifts a component out of range, so out-of-bounds is a collapse signature.""" + b = rec.get("bounds") + m = rec.get("mean_phys") + if b is None or m is None: + return True + lo, hi = np.asarray(b)[:, 0], np.asarray(b)[:, 1] + return bool(np.all(m >= lo - 1e-6) and np.all(m <= hi + 1e-6)) + + +def report(summaries): + for cfg in sorted(summaries): + copies = summaries[cfg] + neffs = np.array([c["neff"] for c in copies], dtype=float) + ngood = int(np.sum(neffs >= GOOD_NEFF)) + # Kish effective #copies over the reliability weights: how many copies the pool really rests on + w = np.where(np.isfinite(neffs), neffs, 0.0) + kish_copies = (w.sum() ** 2 / np.sum(w * w)) if np.sum(w * w) > 0 else 0.0 + print("=" * 78) + print("CONFIG %s (%d copies, %d landed n_eff>=%.0f) n_eff: %s" % ( + cfg, len(copies), ngood, GOOD_NEFF, " ".join("%.1f" % x for x in neffs))) + print(" reliability-weighted effective #copies (Kish over n_eff) = %.1f" % kish_copies) + + gnames = [] + for c in copies: + for gn in c["groups"]: + if gn not in gnames: + gnames.append(gn) + + # per-copy: n_eff + per-group mode count + in-bounds (collapse detector) + print(" per-copy structure (seed: n_eff | group->modes,inbounds):") + for c in copies: + parts = [] + for gn in gnames: + r = c["groups"].get(gn) + if r is None: + parts.append("%s:-" % gn.split(",")[0]); continue + ib = "ok" if _in_bounds(r) else "OOB" + parts.append("%s:m%d/%s" % (gn.split(",")[0], r["modes"], ib)) + tag = "" if c["neff"] >= GOOD_NEFF else " (collapsed)" + print(" s%-3d n_eff=%6.1f | %s%s" % (c["seed"], c["neff"], " ".join(parts), tag)) + + # POOLING: reliability-weighted vs naive vs good-only, per group mean + print(" POOLED group mean [reliability-weighted (all) | naive-unweighted | good-only]:") + for gn in gnames: + recs = [(c["neff"], c["groups"][gn]) for c in copies if gn in c["groups"]] + if not recs: + continue + means = np.array([r["mean_phys"] for _, r in recs]) # (n,d) + ne = np.array([max(0.0, n) for n, _ in recs]) + wpool = ne / ne.sum() if ne.sum() > 0 else np.ones(len(ne)) / len(ne) + m_relw = (wpool[:, None] * means).sum(axis=0) # reliability-weighted + m_naive = means.mean(axis=0) # naive (corrupted by collapses) + good = ne >= GOOD_NEFF + m_good = means[good].mean(axis=0) if good.any() else np.full(means.shape[1], np.nan) + # do the good copies AGREE? scatter of good-only means + good_scatter = np.nanstd(means[good], axis=0) if good.sum() >= 2 else np.full(means.shape[1], np.nan) + print(" [%s]" % gn) + print(" relw =%s naive=%s good =%s good-scatter=%s" % ( + np.array2string(m_relw, precision=2, suppress_small=True), + np.array2string(m_naive, precision=2, suppress_small=True), + np.array2string(m_good, precision=2, suppress_small=True), + np.array2string(good_scatter, precision=2, suppress_small=True))) + print() + + +if __name__ == "__main__": + paths = sys.argv[1:] + if not paths: + print(__doc__); sys.exit(1) + report(load_group_summaries(paths)) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/parse_neff_traj.py b/MonteCarloMarginalizeCode/Code/test/integrators/parse_neff_traj.py new file mode 100644 index 000000000..80cd2eb85 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/parse_neff_traj.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Parse portfolio/AV ILE trajectory logs (lines beginning ' : N Neff ...') and report +the N (sample count) at which Neff first crosses a set of thresholds, plus the final +(N, Neff, sqrt(2 lnLmax)). Usage: parse_neff_traj.py [ ...]""" +import re, sys + +THRESH = [5, 10, 20, 50, 100] +# Two trajectory formats: +# portfolio (mcsamplerPortfolio): " : ..." +# standalone AV (mcsamplerAdaptiveVolume): " - ..." +# Match an optional leading ' :', then N (integer-ish sample count), Neff, sqrt(2lnLmax). +row = re.compile(r"^\s*(?::\s*)?(\d[\d]*)\s+(nan|inf|[0-9.eE+-]+)\s+(nan|inf|-|[0-9.eE+-]+)\b") + +def _f(tok): + try: + return float(tok) + except ValueError: + return float('nan') + +def parse(path): + Ns, Neffs, lmax = [], [], [] + with open(path, errors='replace') as f: + for line in f: + m = row.match(line) + if not m: + continue + N = _f(m.group(1)) + if N < 1000: # skip N=0 header/degenerate lines and non-trajectory numeric lines + continue + Ns.append(N); Neffs.append(_f(m.group(2))); lmax.append(_f(m.group(3))) + return Ns, Neffs, lmax + +def main(): + print(f"{'run':32s} {'Neff>=5':>9} {'>=10':>9} {'>=20':>9} {'>=50':>9} {'>=100':>9} {'finalN':>10} {'finalNeff':>9} {'sq2lnLmax':>9}") + for path in sys.argv[1:]: + Ns, Neffs, lmax = parse(path) + name = path.split('/')[-1].replace('frz_', '').replace('.log', '') + if not Ns: + print(f"{name:32s} (no trajectory lines)") + continue + cross = {} + for t in THRESH: + hit = next((Ns[i] for i in range(len(Ns)) if Neffs[i] >= t), None) + cross[t] = hit + def fmt(v): + return f"{v/1e6:.3f}M" if v is not None else " -- " + print(f"{name:32s} {fmt(cross[5]):>9} {fmt(cross[10]):>9} {fmt(cross[20]):>9} " + f"{fmt(cross[50]):>9} {fmt(cross[100]):>9} {Ns[-1]/1e6:>9.3f}M {Neffs[-1]:>9.1f} {lmax[-1]:>9.2f}") + +if __name__ == '__main__': + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event.sh b/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event.sh new file mode 100755 index 000000000..c1a9cb505 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Orchestrate the multi-event robustness suite: for each event run standalone AV and the +# never-freeze AV+GMM portfolio (in the event's container), then dump an lnZ/n_eff comparison. +set -u +CODE=/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c/MonteCarloMarginalizeCode/Code +PY=/home/richard.oshaughnessy/RIFT_develUWM/bin/python +MEV=$CODE/test/integrators/bench_multi_event.py +EVENTS="S231026ab S240426s S240513ei S240601aj S240703ad" +export CONTAINER=1 NEFF=${NEFF:-30} NMAX=${NMAX:-800000} +GPUS=(1 3) +MAXJOBS=4 + +launch() { # event tag gpu sampler... + local ev=$1 tag=$2 gpu=$3; shift 3 + GPU=$gpu $PY $MEV run $ev $tag "$@" >/dev/null 2>&1 & +} + +i=0 +for ev in $EVENTS; do + g=${GPUS[$((i % ${#GPUS[@]}))]} + launch $ev av_$ev $g AV + launch $ev pf_$ev $g portfolio --sampler-portfolio AV,GMM + i=$((i+1)) + # throttle: wait if too many background jobs + while [ "$(jobs -rp | wc -l)" -ge "$MAXJOBS" ]; do sleep 15; done +done +wait +echo "=== ALL MULTI-EVENT RUNS DONE ===" +for ev in $EVENTS; do + $PY $MEV read $ev av_$ev + $PY $MEV read $ev pf_$ev +done diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event_clip.sh b/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event_clip.sh new file mode 100644 index 000000000..7ffdd78de --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/run_multi_event_clip.sh @@ -0,0 +1,22 @@ +#!/bin/bash +# Safety check: does ADAPTATION-only weight clipping (C=1) hurt TYPICAL events? Runs the portfolio +# with --portfolio-weight-clip 1.0 on each event (in its container) and prints lnZ/n_eff, to compare +# against the earlier no-clip portfolio results (clipping must leave lnZ unbiased and not hurt n_eff). +set -u +CODE=/home/richard.oshaughnessy/RIFT_develUWM/src/research-projects-RIT/.claude/worktrees/rift-adaptive-integrator/.claude/worktrees/gifted-herschel-caf99c/MonteCarloMarginalizeCode/Code +PY=/home/richard.oshaughnessy/RIFT_develUWM/bin/python +MEV=$CODE/test/integrators/bench_multi_event.py +EVENTS="S231026ab S240426s S240513ei S240703ad" +export CONTAINER=1 NEFF=${NEFF:-30} NMAX=${NMAX:-800000} +GPUS=(0 1) +MAXJOBS=2 +i=0 +for ev in $EVENTS; do + g=${GPUS[$((i % ${#GPUS[@]}))]} + GPU=$g $PY $MEV run $ev pfclip_$ev portfolio --sampler-portfolio AV,GMM --portfolio-weight-clip 1.0 >/dev/null 2>&1 & + i=$((i+1)) + while [ "$(jobs -rp | wc -l)" -ge "$MAXJOBS" ]; do sleep 15; done +done +wait +echo "=== CLIP MULTI-EVENT DONE ===" +for ev in $EVENTS; do $PY $MEV read $ev pfclip_$ev; done diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/shape_extrinsic.py b/MonteCarloMarginalizeCode/Code/test/integrators/shape_extrinsic.py new file mode 100644 index 000000000..ad51dcba7 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/shape_extrinsic.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +""" +shape_extrinsic.py -- weighted-sample SHAPE check on an ILE extrinsic export. + +n_eff alone can hide a wrong posterior SHAPE: a run can report a healthy n_eff while a handful of +samples own the estimate (over-broad GMM proposal at too-high a component cap) or the recovered +marginals are degenerate. This reads the extrinsic sample cloud saved by `--save-samples` and +reports, per extrinsic parameter, the WEIGHTED marginal (mean, std) plus global diagnostics: + * n_eff (Kish) of the weighted cloud, + * max single-sample weight FRACTION -- the outlier-dominance signal (the cap-too-high failure), + * effective vs raw sample count. + +Usage: shape_extrinsic.py [ ...] +Auto-detects the format. Compare several caps side by side to see the failure mode emerge. +""" +from __future__ import print_function +import sys +import numpy as np + +# extrinsic columns we summarize (name -> unit label) +PARAMS = [("distance", "Mpc"), ("inclination", "rad"), ("right_ascension", "rad"), + ("declination", "rad"), ("psi", "rad"), ("phi_orb", "rad")] +# aliases as written by --save-samples (longitude/latitude/polarization/coa_phase) +ALIAS = {"right_ascension": ["longitude", "ra"], "declination": ["latitude", "dec"], + "psi": ["polarization"], "phi_orb": ["coa_phase"]} + + +def _from_xml(path): + """Read the sim_inspiral-style table --save-samples writes; return dict of arrays incl weights.""" + from igwn_ligolw import ligolw, lsctables, utils as ligolw_utils + xmldoc = ligolw_utils.load_filename(path, contenthandler=lsctables.use_in(ligolw.LIGOLWContentHandler)) + tbl = lsctables.SimInspiralTable.get_table(xmldoc) + cols = {} + # standard extrinsic mapping (see the save-samples block in the driver) + getters = {"distance": "distance", "inclination": "inclination", + "right_ascension": "longitude", "declination": "latitude", + "psi": "polarization", "phi_orb": "coa_phase", "loglikelihood": "alpha1"} + for k in ("distance", "inclination", "longitude", "latitude", "polarization", "coa_phase"): + try: + cols[k] = np.array([getattr(r, k) for r in tbl], dtype=float) + except Exception: + pass + # weight: RIFT stores the sampling info in alpha columns; loglikelihood via alpha1 typically. + for wk in ("alpha1", "alpha", "snr"): + try: + cols["loglikelihood"] = np.array([getattr(r, wk) for r in tbl], dtype=float) + break + except Exception: + continue + return cols + + +def _from_dat(path): + """The .dat ASCII form: RIFT writes extrinsic columns + weights when --save-samples is on. + Column order is the ILE convention; we detect it by width and pull the log-weight column.""" + arr = np.loadtxt(path) + if arr.ndim == 1: + arr = arr[None, :] + # A single-row .dat is the marginalized point (no cloud) -- not a shape export. + if arr.shape[0] < 5: + return None + return arr + + +def summarize(path): + name = path.split("/")[-1] + try: + if path.endswith(".xml") or path.endswith(".xml.gz"): + cols = _from_xml(path) + ll = cols.get("loglikelihood") + if ll is None: + print("%-28s (no weight column found in XML)" % name); return + w = np.exp(ll - np.max(ll)) + data = {"distance": cols.get("distance"), "inclination": cols.get("inclination"), + "right_ascension": cols.get("longitude"), "declination": cols.get("latitude"), + "psi": cols.get("polarization"), "phi_orb": cols.get("coa_phase")} + else: + arr = _from_dat(path) + if arr is None: + print("%-28s (single-row .dat: marginalized point, not a cloud)" % name); return + # heuristic: last col = neff-ish, cols mapped by the standard ILE .dat layout is + # fragile, so require the XML path for full shape; here just report n_eff proxy. + print("%-28s (.dat cloud reader needs the XML; run made %d rows)" % (name, arr.shape[0])); return + except Exception as e: + print("%-28s ERROR %s" % (name, str(e)[:70])); return + + w = np.asarray(w, dtype=float) + w = np.where(np.isfinite(w), w, 0.0) + W = w.sum() + if W <= 0: + print("%-28s (all-zero weights)" % name); return + neff = W * W / np.sum(w * w) + maxfrac = float(np.max(w) / W) + print("== %s ==" % name) + print(" n_eff(Kish)=%.1f raw=%d max-weight-frac=%.3e%s" % ( + neff, len(w), maxfrac, " <-- ONE SAMPLE DOMINATES" if maxfrac > 0.05 else "")) + for p, unit in PARAMS: + x = data.get(p) + if x is None or len(x) != len(w): + continue + m = np.sum(w * x) / W + s = np.sqrt(max(0.0, np.sum(w * (x - m) ** 2) / W)) + print(" %-16s mean=%9.3f std=%9.3f %s" % (p, m, s, unit)) + + +if __name__ == "__main__": + for p in sys.argv[1:]: + summarize(p) + print() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_bootstrap.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_bootstrap.py new file mode 100644 index 000000000..e94961869 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_bootstrap.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python +""" +test_AV_bootstrap.py + +Quantitative test of the bootstrappable AdaptiveVolume integrator. + +Demonstrates and validates that warm-starting AV from prior information reaches a +target effective sample size in far fewer likelihood evaluations than a cold +start, WITHOUT biasing the integral (must stay within the CI 4-sigma gate). + +Three warm-start channels are exercised, all matching the production sampler API: + 1. Fisher matrix (bootstrap_from_fisher) -- the free "Fisher oracle" + 2. reference samples (bootstrap_from_samples) -- e.g. a previous ILE posterior + 3. serialized state round-trip (save_state/load_state) -- reuse across instances + +Usage: + python test_AV_bootstrap.py # corrgauss5, GPU if visible + python test_AV_bootstrap.py --target gaussmix8 --as-test +""" +from __future__ import print_function +import argparse +import os +import tempfile +import numpy as np + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerAdaptiveVolume as AVmod + + +def _numeric_fisher(target, at=None, eps=1e-3): + """Central-difference Fisher (negative Hessian of lnL) at the mode `at`.""" + d = target.ndim + if at is None: + # crude mode: densest of a coarse random scan + rng = np.random.RandomState(0) + X = rng.uniform(target.llim, target.rlim, size=(20000, d)) + at = X[np.argmax(target.lnL(X))] + H = np.zeros((d, d)) + scale = (target.rlim - target.llim) * eps + f0 = target.lnL(np.atleast_2d(at))[0] + for i in range(d): + for j in range(i, d): + ei = np.zeros(d); ei[i] = scale[i] + ej = np.zeros(d); ej[j] = scale[j] + fpp = target.lnL(np.atleast_2d(at + ei + ej))[0] + fpm = target.lnL(np.atleast_2d(at + ei - ej))[0] + fmp = target.lnL(np.atleast_2d(at - ei + ej))[0] + fmm = target.lnL(np.atleast_2d(at - ei - ej))[0] + H[i, j] = H[j, i] = (fpp - fpm - fmp + fmm) / (4 * scale[i] * scale[j]) + fisher = -0.5 * (H + H.T) + # regularize to SPD + w, Vv = np.linalg.eigh(fisher) + w = np.clip(w, 1e-6, None) + return at, Vv @ np.diag(w) @ Vv.T + + +def run_cold(target, **kw): + return B.run("AV", target, **kw) + + +def run_warm(target, warm_kind, seed_info, **kw): + def warm_start(sampler, tgt): + if warm_kind == "fisher": + mean, fisher = seed_info + sampler.bootstrap_from_fisher(mean, fisher, n=sampler.n_chunk, seed=1) + elif warm_kind == "mixture": + means, covs, weights = seed_info + sampler.bootstrap_from_gaussian_mixture(means, covs, weights, + n=sampler.n_chunk, seed=1) + elif warm_kind == "samples": + sampler.bootstrap_from_samples(seed_info) + elif warm_kind == "state": + sampler.load_state(seed_info) + return B.run("AV", target, warm_start=warm_start, **kw) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--target", default="corrgauss5") + ap.add_argument("--nmax", type=int, default=300000) + ap.add_argument("--neff", type=int, default=1000) + ap.add_argument("--n-chunk", type=int, default=10000) + ap.add_argument("--as-test", action="store_true") + ap.add_argument("--seed", type=int, default=123456) + args = ap.parse_args() + + target = B._TARGETS[args.target]() + kw = dict(nmax=args.nmax, neff=args.neff, n_chunk=args.n_chunk, seed=args.seed) + multimodal = hasattr(target, "means") and len(getattr(target, "means")) > 1 + print("# target {} ndim={} true_lnZ={:.4f} {}".format( + target.name, target.ndim, target.true_lnZ, + "MULTIMODAL" if multimodal else "unimodal")) + + # --- cold baseline --- + cold = run_cold(target, **kw) + print("COLD ", B._fmt(cold)) + + results = [] # (label, result) + + at, fisher = _numeric_fisher(target) + if multimodal: + # a single Fisher cannot cover multiple modes; use a mixture oracle + # (stand-in for a GMM/flow oracle fit to a previous posterior) and also + # empirical samples -- both cover the full support. + seed_mix = (target.means, target.covs, target.wt) + results.append(("mixture", run_warm(target, "mixture", seed_mix, **kw))) + else: + results.append(("fisher", run_warm(target, "fisher", (at, fisher), **kw))) + + # --- warm: reference samples (analytic-truth draws stand in for a previous + # ILE posterior; the mixture can be sampled exactly) --- + if hasattr(target, "sample_truth"): + ref = target.sample_truth(5000, seed=7) + else: + ref = np.random.RandomState(3).multivariate_normal(at, np.linalg.inv(fisher), 5000) + results.append(("samples", run_warm(target, "samples", ref, **kw))) + + # --- warm: serialized-state round trip (reuse across sampler instances) --- + s0, _ = B.build_sampler("AV", target, n_chunk=args.n_chunk) + s0.setup() + if multimodal: + s0.bootstrap_from_gaussian_mixture(target.means, target.covs, target.wt, + n=args.n_chunk, seed=1) + else: + s0.bootstrap_from_fisher(at, fisher, n=args.n_chunk, seed=1) + tmp = os.path.join(tempfile.gettempdir(), "av_state_%s.npz" % target.name) + s0.save_state(tmp) + results.append(("state", run_warm(target, "state", tmp, **kw))) + sz = os.path.getsize(tmp) + + for label, r in results: + print("WARM({:<8s}".format(label + ")"), B._fmt(r)) + print("# serialized state size: {} bytes ({} live bins)".format(sz, len(np.load(tmp)['binunique']))) + + # --- summary: efficiency and samples-to-neff speedups --- + print("\n# --- speedup (warm vs cold) ---") + for name, r in results: + eff_ratio = r["efficiency"] / cold["efficiency"] + n_ratio = cold["n_eval"] / max(r["n_eval"], 1) + print(" {:>8s}: efficiency x{:.2f} N_eval-to-neff x{:.2f} fewer " + "bias_ln={:+.3f} (cold {:+.3f})".format(name, eff_ratio, n_ratio, + r["bias_ln"], cold["bias_ln"])) + + if args.as_test: + # Correctness criterion: a warm start must not make the integral MORE + # biased than a cold start (AV can have its own intrinsic bias on hard + # high-D targets; the bootstrap must not worsen it) and must not lose + # efficiency. Gate = cold's own |bias| plus a margin. + tol = max(0.10, abs(cold["bias_ln"]) + 3 * cold["rel_err"]) + ok = True + for name, r in results: + if abs(r["bias_ln"]) > tol: + print(" FAIL: warm({}) bias_ln {:+.3f} exceeds tol {:.3f} (cold {:+.3f})".format( + name, r["bias_ln"], tol, cold["bias_ln"])) + ok = False + if r["efficiency"] < 0.95 * cold["efficiency"]: + print(" WARN: warm({}) efficiency {:.2e} < cold {:.2e}".format( + name, r["efficiency"], cold["efficiency"])) + if not ok: + raise SystemExit(1) + print(" PASS: all warm starts no more biased than cold (tol {:.3f}) and no less efficient".format(tol)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_warmstart_safety.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_warmstart_safety.py new file mode 100644 index 000000000..2588d0ac2 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_AV_warmstart_safety.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python +""" +test_AV_warmstart_safety.py + +Guards the anti-bias property required for reusing a warm-start proposal ACROSS +different problems (a neighbouring intrinsic point, a stale breadcrumb, a proposal +carried between pipeline iterations): a mis-placed proposal must never make the +warm-started integral MORE biased than a cold one. + +Because VARAHA's live volume only contracts, a warm start seeded at the WRONG +location silently contracts there, misses the true peak, and returns a +catastrophically biased integral that nonetheless LOOKS converged (a healthy +n_eff) -- the worst possible failure. The coverage floor (cover_frac) mixes a +fraction of full-prior coverage into the seed, so the warm live volume always +contains a cold start: a wrong proposal then only costs efficiency, never bias. + +This test deliberately seeds AV at a decoy far from the true mode and asserts: + * NO floor -> badly biased (demonstrates the danger), and + * cover_frac>0 -> unbiased, comparable to cold. +""" +from __future__ import print_function +import argparse +import numpy as np + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerAdaptiveVolume as AV + + +def _run(target, warm=None, nmax=200000, neff=1500, n_chunk=10000, seed=1234): + np.random.seed(seed) + s = AV.MCSampler(n_chunk=n_chunk) + for i, p in enumerate(target.params): + w = target.rlim[i] - target.llim[i] + s.add_parameter(p, np.vectorize(lambda x, w=w: 1.0 / w), + prior_pdf=np.vectorize(lambda x, w=w: 1.0 / w), + left_limit=float(target.llim[i]), right_limit=float(target.rlim[i]), + adaptive_sampling=True) + s.setup() + if warm is not None: + warm(s) + r, v, eff, _ = s.integrate_log(target.as_lnfunc(), *target.params, no_protect_names=True, + nmax=nmax, n=n_chunk, n_adapt=100, neff=neff, tempering_exp=0.1) + return float(B._asnumpy(r)) - float(target.true_lnZ), float(B._asnumpy(eff)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--as-test", action="store_true") + ap.add_argument("--cover-frac", type=float, default=0.5, + help="Coverage floor to exercise. Default 0.5 = the production default: " + "measured-safe across 20 calibration seeds (max |bias| 0.164, max " + "degradation vs cold 0.113). 0.1 is under-covered (1.1-1.7 log bias " + "across seeds), and raising nmax does NOT rescue it because the runs " + "terminate on n_eff=1500 first.") + args = ap.parse_args() + + target = B.CorrelatedGaussian(ndim=3) # cold AV converges here (unbiased control) + # a WRONG proposal: a tight cloud far from the true mode (a stale/neighbour seed) + decoy = np.clip(np.random.RandomState(1).normal([-4.0, 4.0, -4.0], 0.3, size=(3000, 3)), + target.llim + 1e-3, target.rlim - 1e-3) + + cold_b, cold_n = _run(target) + bad_b, bad_n = _run(target, warm=lambda s: s.bootstrap_from_samples(decoy)) + safe_b, safe_n = _run(target, warm=lambda s: s.bootstrap_from_samples(decoy, cover_frac=args.cover_frac)) + + print("COLD bias_ln=%+.3f neff=%.0f" % (cold_b, cold_n)) + print("WARM wrong seed, NO floor bias_ln=%+.3f neff=%.0f (danger: biased but 'converged')" % (bad_b, bad_n)) + print("WARM wrong seed, cover_frac bias_ln=%+.3f neff=%.0f (safe: ~cold)" % (safe_b, safe_n)) + + if args.as_test: + ok = True + # the no-floor case MUST demonstrate the danger (else the test is not exercising it) + if abs(bad_b) < 1.0: + print(" WARN: no-floor decoy did not bias strongly (%.3f); test may be too easy" % bad_b) + # the safety requirement: covered warm start is no more biased than cold + margin + tol = max(0.30, 4 * abs(cold_b) + 0.15) + if abs(safe_b) > tol: + print(" FAIL: cover_frac warm start biased %+.3f > tol %.3f (cold %+.3f)" % (safe_b, tol, cold_b)) + ok = False + if not ok: + raise SystemExit(1) + print(" PASS: coverage floor keeps a mis-placed warm start unbiased (tol %.3f)" % tol) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_NF_reuse.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_NF_reuse.py new file mode 100644 index 000000000..a30aaca45 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_NF_reuse.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python +""" +test_NF_reuse.py + +Demonstrates + validates the normalizing-flow storage/reuse framework. + +NF training is slow and only pays off if the trained flow is re-used across the +many ILE instances that share similar posterior structure. This test: + 1. trains a flow once on a target and saves it (save_flow); + 2. reuses it in FRESH samplers via load_flow, either + - pure reuse (n_adapt=0: sample straight from the trained flow), or + - a few 'polish' epochs (small n_adapt) to adapt to the instance; + 3. checks the reused runs match the cold-trained integral (unbiased) while + spending far less wallclock on training. + +Usage: + python test_NF_reuse.py --as-test +Run with thread caps, e.g. OMP_NUM_THREADS=2, to avoid torch oversubscription. +""" +from __future__ import print_function +import argparse +import os +import tempfile +import time +import numpy as np + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerNFlow + + +def _build(target, n_chunk=10000): + s = mcsamplerNFlow.MCSampler(n_chunk=n_chunk) + for d, p in enumerate(target.params): + w = target.rlim[d] - target.llim[d] + s.add_parameter(p, np.vectorize(lambda x, w=w: 1.0 / w), + prior_pdf=np.vectorize(lambda x, w=w: 1.0 / w), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + return s + + +def _integrate(s, target, nmax, neff, n_chunk, n_adapt, load_path=None): + if load_path is not None: + s.load_flow(load_path) + t0 = time.time() + lnI, logvar, eff, _ = s.integrate_log(target.as_lnfunc(), *target.params, + no_protect_names=True, nmax=nmax, neff=neff, + n=n_chunk, n_adapt=n_adapt, tempering_exp=1.0, + verbose=False) + wall = time.time() - t0 + lnI = float(B._asnumpy(lnI)); eff = float(B._asnumpy(eff)) + ln_wt = B.log_weights_from_rvs(s._rvs) + return dict(bias_ln=lnI - float(target.true_lnZ), n_eff=eff, + n_ess=B.n_ess_kish(ln_wt), n_eval=int(getattr(s, "ntotal", 0)) or nmax, + wall=wall) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ndim", type=int, default=3) + ap.add_argument("--n-chunk", type=int, default=10000) + ap.add_argument("--train-nmax", type=int, default=80000) + ap.add_argument("--nmax", type=int, default=40000) + ap.add_argument("--neff", type=int, default=300) + ap.add_argument("--as-test", action="store_true") + args = ap.parse_args() + + target = B.CorrelatedGaussian(ndim=args.ndim) + print("# NF reuse on {} true_lnZ={:.4f}".format(target.name, target.true_lnZ)) + + # --- phase 1: train once and save (the expensive, one-time cost) --- + trainer = _build(target, args.n_chunk) + r_train = _integrate(trainer, target, args.train_nmax, args.neff, args.n_chunk, n_adapt=8) + path = os.path.join(tempfile.gettempdir(), "nf_flow_%s.pt" % target.name) + trainer.save_flow(path) + sz = os.path.getsize(path) + print("TRAIN+SAVE bias={bias_ln:+.3f} neff={n_eff:.1f} nESS={n_ess:.1f} " + "t={wall:.1f}s -> saved {sz} bytes".format(sz=sz, **r_train)) + + # --- phase 2a: cold (fresh sampler, train from scratch) --- + cold = _integrate(_build(target, args.n_chunk), target, args.nmax, args.neff, args.n_chunk, n_adapt=8) + print("COLD bias={bias_ln:+.3f} neff={n_eff:.1f} nESS={n_ess:.1f} t={wall:.1f}s".format(**cold)) + + # --- phase 2b: warm reuse (load flow, NO training) --- + reuse = _integrate(_build(target, args.n_chunk), target, args.nmax, args.neff, args.n_chunk, + n_adapt=0, load_path=path) + print("WARM(reuse) bias={bias_ln:+.3f} neff={n_eff:.1f} nESS={n_ess:.1f} t={wall:.1f}s".format(**reuse)) + + # --- phase 2c: warm + polish (load flow, a couple of epochs) --- + polish = _integrate(_build(target, args.n_chunk), target, args.nmax, args.neff, args.n_chunk, + n_adapt=2, load_path=path) + print("WARM(polish)bias={bias_ln:+.3f} neff={n_eff:.1f} nESS={n_ess:.1f} t={wall:.1f}s".format(**polish)) + + print("\n# training-cost saving: cold {:.1f}s vs warm-reuse {:.1f}s (x{:.1f} faster)".format( + cold["wall"], reuse["wall"], cold["wall"] / max(reuse["wall"], 1e-3))) + + if args.as_test: + ok = True + tol = max(0.20, 4 * abs(cold["bias_ln"]) + 0.1) + for name, r in [("reuse", reuse), ("polish", polish)]: + if abs(r["bias_ln"]) > tol: + print(" FAIL: warm({}) biased {:+.3f} > {:.3f}".format(name, r["bias_ln"], tol)); ok = False + # reuse must be materially cheaper than cold (no training loop) + if not (reuse["wall"] < 0.6 * cold["wall"]): + print(" FAIL: warm reuse not faster than cold ({:.1f}s vs {:.1f}s)".format(reuse["wall"], cold["wall"])); ok = False + # reused flow must actually sample the mode (not degenerate) + if not (reuse["n_ess"] > 20): + print(" FAIL: reused flow n_ESS too low ({:.1f})".format(reuse["n_ess"])); ok = False + if not ok: + raise SystemExit(1) + print(" PASS: flow reuse is unbiased and skips training (tol {:.3f})".format(tol)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_gmm_adaptive.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_gmm_adaptive.py new file mode 100644 index 000000000..c97e2d630 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_gmm_adaptive.py @@ -0,0 +1,148 @@ +"""Unit tests for the flexible / data-driven GMM component allocation added to +RIFT.integrators.gaussian_mixture_model: + + * fit_gmm_adaptive -- choose k by BIC, then prune dead components + * gmm.prune_components + * gmm._match_components -- O(k^3) Hungarian == old O(k!) permutation optimum + +Backend note: gaussian_mixture_model uses cupy when a GPU is visible, numpy +otherwise. These tests are backend-agnostic; run with CUDA_VISIBLE_DEVICES set +to a GPU, or with a numpy-only build. No ILE data required (seconds to run). + +Run: + CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 \ + python test/integrators/test_gmm_adaptive.py +""" +from __future__ import print_function +import sys, itertools +import numpy as np +from RIFT.integrators import gaussian_mixture_model as GMM + +cvt = GMM.identity_convert +rng = np.random.RandomState(20250721) + + +def _bounds(d, lo=-8., hi=8.): + b = np.empty((d, 2)); b[:, 0] = lo; b[:, 1] = hi + return GMM.xpy_default.array(b) + + +def test_bic_picks_one_for_single_gaussian(): + """A single Gaussian blob should be modeled with k=1 (BIC penalizes extra + components that do not improve the weighted likelihood).""" + d = 3 + X = rng.normal(0.0, 0.7, size=(4000, d)) + model = GMM.fit_gmm_adaptive(GMM.xpy_default.array(X), _bounds(d), k_max=8, + defensive_frac=0.0) + print(" single-gaussian -> k =", model.k) + assert model.k == 1, "expected k=1 for a single blob, got %d" % model.k + + +def test_bic_grows_for_separated_modes(): + """Three well-separated blobs should earn more than one component.""" + d = 2 + centers = np.array([[-5., -5.], [0., 5.], [5., -5.]]) + X = np.vstack([rng.normal(c, 0.4, size=(1500, d)) for c in centers]) + model = GMM.fit_gmm_adaptive(GMM.xpy_default.array(X), _bounds(d), k_max=8, + defensive_frac=0.0) + print(" three-modes -> k =", model.k) + assert model.k >= 3, "expected k>=3 for three separated modes, got %d" % model.k + + +def test_bic_respects_weights(): + """With importance weights that select one of two blobs, BIC should prefer + fewer components (only the up-weighted blob carries effective mass).""" + d = 2 + A = rng.normal([-4, 0], 0.4, size=(2000, d)) + B = rng.normal([4, 0], 0.4, size=(2000, d)) + X = np.vstack([A, B]) + # up-weight only blob A + lw = np.concatenate([np.zeros(len(A)), -50.0 * np.ones(len(B))]) + model = GMM.fit_gmm_adaptive(GMM.xpy_default.array(X), _bounds(d), + log_sample_weights=GMM.xpy_default.array(lw), + k_max=8, defensive_frac=0.0) + print(" weighted-one-of-two -> k =", model.k) + assert model.k <= 2, "expected small k when weights select one blob, got %d" % model.k + # the fitted mass should sit near blob A (-4,0), not the midpoint + means = np.array([cvt(m) for m in model.means]) + mean_un = model._unnormalize(GMM.xpy_default.array(means)) + mx = float(cvt(mean_un)[:, 0].mean()) + print(" weighted mean x =", mx) + assert mx < -1.0, "weighted fit should sit on the up-weighted blob" + + +def test_k_min_safety_floor(): + """A single blob would BIC-select k=1, but k_min must floor the count so a + stress-tested hard-coded allocation is never reduced (multi-modal-sky + safety).""" + d = 2 + X = rng.normal(0.0, 0.7, size=(4000, d)) + model = GMM.fit_gmm_adaptive(GMM.xpy_default.array(X), _bounds(d), k_max=8, + k_min=4, defensive_frac=0.0) + print(" single-gaussian, k_min=4 -> k =", model.k) + assert model.k >= 4, "k_min floor violated: got %d < 4" % model.k + + +def test_prune_removes_dead_components(): + d = 2 + model = GMM.gmm(4, _bounds(d)) + # fit to a single blob so 3 of 4 components collapse to ~zero weight + X = rng.normal(0.0, 0.5, size=(3000, d)) + model.fit(GMM.xpy_default.array(X)) + k_before = model.k + model.prune_components(weight_floor=1e-2) + print(" prune: k %d -> %d" % (k_before, model.k)) + assert model.k <= k_before + w = np.asarray(cvt(model.weights), dtype=float) + assert abs(w.sum() - 1.0) < 1e-6, "weights must renormalize to 1" + assert len(model.means) == model.k and len(model.covariances) == model.k + + +def test_matching_matches_permutation_optimum(): + """Hungarian _match_components must reproduce the exact permutation optimum.""" + def objective(order, om, oc, nm, nc): + val = 0.0 + for i, j in enumerate(order): + diff = nm[j] - om[i] + val += np.sqrt(diff @ np.linalg.inv(oc[i]) @ diff) + val += np.sqrt(diff @ np.linalg.inv(nc[j]) @ diff) + return val + for k in [2, 3, 4, 5]: + d = 3 + model = GMM.gmm(k, _bounds(d)) + new = GMM.gmm(k, _bounds(d)) + model.d = new.d = d + model.means = [GMM.xpy_default.array(rng.randn(d)) for _ in range(k)] + new.means = [GMM.xpy_default.array(rng.randn(d)) for _ in range(k)] + mk = lambda: (lambda A: GMM.xpy_default.array(A @ A.T + np.eye(d)))(rng.randn(d, d)) + model.covariances = [mk() for _ in range(k)] + new.covariances = [mk() for _ in range(k)] + om = [cvt(m) for m in model.means]; nm = [cvt(m) for m in new.means] + oc = [cvt(c) for c in model.covariances]; nc = [cvt(c) for c in new.covariances] + # brute-force optimum + best = min(itertools.permutations(range(k)), + key=lambda o: objective(o, om, oc, nm, nc)) + got = model._match_components(new) + assert abs(objective(got, om, oc, nm, nc) - objective(best, om, oc, nm, nc)) < 1e-9, \ + "k=%d: Hungarian objective != permutation optimum" % k + print(" matching == permutation optimum for k in 2..5") + + +if __name__ == "__main__": + tests = [test_bic_picks_one_for_single_gaussian, + test_bic_grows_for_separated_modes, + test_bic_respects_weights, + test_k_min_safety_floor, + test_prune_removes_dead_components, + test_matching_matches_permutation_optimum] + nfail = 0 + for t in tests: + try: + print("[RUN]", t.__name__) + t() + print("[PASS]", t.__name__) + except AssertionError as e: + nfail += 1 + print("[FAIL]", t.__name__, "->", e) + print("\n%d/%d passed" % (len(tests) - nfail, len(tests))) + sys.exit(1 if nfail else 0) diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_adaptive_alloc.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_adaptive_alloc.py new file mode 100644 index 000000000..4b697e7a6 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_adaptive_alloc.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python +""" +test_portfolio_adaptive_alloc.py + +Validates mcsamplerPortfolio's ADAPTIVE-PROBE draw allocation: the portfolio should +automatically concentrate its draw budget on whichever member is actually winning, so it +TRACKS the best single member on both a weakly- and a strongly-correlated target -- and +therefore BEATS standalone AV on the correlated one (where a full-covariance GMM wraps the +degeneracy that AV's axis-aligned bins cannot). + +Two targets, both scaled Gaussians with known true_lnZ: + * UNCORRELATED : axis-aligned anisotropic Gaussian. + * CORRELATED : a COMPOUND-SYMMETRIC Gaussian (every coordinate pair correlated) -- its narrow + eigen-directions are OFF the coordinate axes, so AV's axis-aligned bins cannot + wrap the tilted ridge, while a single full-covariance GMM component captures it. + +For each target we run standalone AV, standalone GMM, and the AV+GMM portfolio to a FIXED sample +budget and compare n_eff (efficiency) and bias = lnI - true_lnZ (correctness). The GMM member is +broad-seeded (a wide peak-covering proposal) so the test deterministically exercises the ALLOCATION +policy given a member that CAN model the correlation, rather than gambling on cold GMM finding a +thin ridge; AV starts cold (axis-aligned bins cannot wrap the correlation, seed or not). Observed: + * On the CORRELATED target GMM's n_eff is several-fold AV's, and adaptive allocation concentrates + on GMM, so the portfolio BEATS standalone AV (the whole point of a portfolio on a correlated + problem). On the uncorrelated target the portfolio matches/beats the better single member too. + * A cold VARAHA/AV only CONTRACTS, so it under-covers the Gaussian tails and is BIASED LOW on + both targets; the portfolio stays UNBIASED because the covering GMM member enters q_mix -- a + second reason to prefer the portfolio over AV alone here. (This is the opposite regime from a + warm, cover-frac'd AV on a real ILE likelihood, where AV is the unbiased workhorse; the point + tested here is that adaptive allocation follows whichever member is actually winning.) + +Usage: + CUDA_VISIBLE_DEVICES=2 OMP_NUM_THREADS=2 OPENBLAS_NUM_THREADS=2 \\ + python test_portfolio_adaptive_alloc.py --as-test +""" +from __future__ import print_function +import argparse +import numpy as np +from scipy.stats import multivariate_normal + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerAdaptiveVolume as AVmod +from RIFT.integrators import mcsamplerEnsemble as Emod +from RIFT.integrators import mcsamplerPortfolio as Pmod + + +class CompoundCorrelatedGaussian(B.CorrelatedGaussian): + """A COMPOUND-SYMMETRIC correlated Gaussian: every pair of coordinates is correlated (cov_ij=c + for i!=j). Its eigen-structure is one wide direction along (1,1,...,1) and (ndim-1) narrow + directions that are OFF the coordinate axes -- exactly the correlated/degenerate geometry that + a full-covariance GMM captures in one component but that AV's axis-aligned bins cannot wrap + (they must staircase the tilted ridge, wasting most of their bounding box). `base` scales the + whole covariance so the wide direction is comfortably contained in the box AND the narrow + directions stay findable cold (std ~0.3, not a needle).""" + def __init__(self, ndim=5, c=0.85, base=0.5, width=10.0, scale=100.0, seed=7): + super(CompoundCorrelatedGaussian, self).__init__(ndim=ndim, width=width, scale=scale, + seed=seed, rho=0.0) + self.cov = base * ((1.0 - c) * np.eye(ndim) + c * np.ones((ndim, ndim))) + self.mu = np.zeros(ndim) # centered -> contained in the box + self._mvn = multivariate_normal(self.mu, self.cov) + self.name = "compound_d{}".format(ndim) + self.true_lnZ = np.log(scale) - np.sum(np.log(self.rlim - self.llim)) + + +def _host_lnfunc(target): + """cupy-tolerant wrapper: AV's selfish self-update evaluates on device-native draws, but the + synthetic integrand is host/numpy -- move any device args to the host first.""" + base = target.as_lnfunc() + def ln_f(*cols): + cols = [Emod.identity_convert(c) for c in cols] + return base(*cols) + return ln_f + + +def _seed_gmm_broad(gmm, target, broad=3.0, n=8000, seed=7): + """Give the GMM member a BROAD but peak-covering full-covariance proposal (fit to a wide cloud + N(mu, broad^2 cov) around the mode). This removes the cold-start LOTTERY -- cold GMM only + sometimes finds a thin correlated ridge from a uniform start -- so the test deterministically + exercises the ALLOCATION policy given a member that *can* model the correlation (AV cannot, + seed or not). The member still adapts/tightens during the run.""" + rng = np.random.RandomState(seed) + cloud = rng.multivariate_normal(target.mu, broad ** 2 * np.atleast_2d(target.cov), n) + cloud = np.clip(cloud, target.llim + 1e-3, target.rlim - 1e-3) + gmm.update_sampling_prior(np.zeros(len(cloud)), 2 * len(cloud), + external_rvs={p: cloud[:, i] for i, p in enumerate(gmm.params_ordered)}, + log_scale_weights=True) + + +def build(target, members, n_chunk): + """Build a portfolio of the requested members ('AV', 'GMM' or both). GMM members are seeded + with a broad peak-covering proposal (see _seed_gmm_broad); AV members start cold.""" + objs, gmms = [], [] + for name in members: + if name == 'AV': + objs.append(AVmod.MCSampler(n_chunk=n_chunk)) + else: + g = Emod.MCSampler(); objs.append(g); gmms.append(g) + port = Pmod.MCSampler(portfolio=objs, n_chunk=n_chunk) + for d, p in enumerate(target.params): + w = target.rlim[d] - target.llim[d] + port.add_parameter(p, np.vectorize(lambda x, w=w: 1.0 / w), + prior_pdf=np.vectorize(lambda x, w=w: 1.0 / w), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + # GMM: single full-covariance component (captures a correlated ridge in one component) + port.setup(portfolio_breakpoints=None, n_comp=1, correlate_all_dims=True, n=n_chunk) + for g in gmms: + _seed_gmm_broad(g, target) + return port + + +def run(target, members, n_chunk, nmax, seed=1234): + np.random.seed(seed) + port = build(target, members, n_chunk) + lnI, _, eff, _ = port.integrate_log( + _host_lnfunc(target), *target.params, no_protect_names=True, + nmax=nmax, neff=10**9, n=n_chunk, n_adapt=100, tempering_exp=0.3, + floor_level=0.0, use_lnL=True, save_intg=True, verbose=False, + portfolio_adaptive_alloc=True) # opt-in: this test exercises adaptive allocation + lnI = float(B._asnumpy(lnI)) + # use the integrator's OWN reported effective-sample count (the q_mix-based pooled eff_samp), + # the quantity it actually targets -- comparable across standalone AV/GMM and the portfolio. + n_eff = float(B._asnumpy(eff)) + return dict(lnI=lnI, bias=lnI - float(target.true_lnZ), n_eff=n_eff, + wts=np.round(np.array(port.portfolio_weights), 3)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ndim", type=int, default=5) + ap.add_argument("--nmax", type=int, default=400000) + ap.add_argument("--n-chunk", type=int, default=10000) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--as-test", action="store_true") + args = ap.parse_args() + + uncorr = B.CorrelatedGaussian(ndim=args.ndim, rho=0.0, narrow=0.1) + corr = CompoundCorrelatedGaussian(ndim=args.ndim) + kw = dict(n_chunk=args.n_chunk, nmax=args.nmax, seed=args.seed) + + print("# fixed budget nmax={} n_chunk={} ndim={}\n".format(args.nmax, args.n_chunk, args.ndim)) + rows = {} + for label, tgt in [("UNCORRELATED (axis-aligned)", uncorr), ("CORRELATED (compound-symmetric)", corr)]: + print("== {} true_lnZ={:.3f} ==".format(label, tgt.true_lnZ)) + av = run(tgt, ['AV'], **kw) + gm = run(tgt, ['GMM'], **kw) + pf = run(tgt, ['AV', 'GMM'], **kw) + rows[label] = (av, gm, pf) + for nm, r in [("AV ", av), ("GMM ", gm), ("PORT ", pf)]: + extra = " final wts(AV,GMM)={}".format(r["wts"]) if nm == "PORT " else "" + print(" {}: n_eff={:9.1f} bias={:+.3f}{}".format(nm, r["n_eff"], r["bias"], extra)) + print() + + if args.as_test: + # Only the ROBUST claims are gated (cold GMM's absolute n_eff on the correlated target is + # stochastic run-to-run; the portfolio also legitimately carries the biased AV member, so it + # is not always >= GMM-alone). The durable, seed-insensitive facts are: the portfolio is + # UNBIASED, and on the CORRELATED target adaptive allocation concentrates on the full-cov GMM + # and the portfolio clearly BEATS standalone AV (the correlated-problem win). + ok = True + for label, (av, gm, pf) in rows.items(): + for nm, r in [("GMM", gm), ("PORT", pf)]: + if abs(r["bias"]) > 0.2: + print(" FAIL[{}]: {} biased ({:+.3f})".format(label, nm, r["bias"])); ok = False + av_c, gm_c, pf_c = rows["CORRELATED (compound-symmetric)"] + if not (pf_c["n_eff"] > 1.5 * av_c["n_eff"]): + print(" FAIL: portfolio did not clearly beat standalone AV on the correlated target " + "(PORT {:.1f} vs AV {:.1f})".format(pf_c["n_eff"], av_c["n_eff"])); ok = False + if not (pf_c["wts"][1] > 0.6): + print(" FAIL: adaptive allocation did not concentrate on GMM on the correlated target " + "(GMM weight {:.2f})".format(pf_c["wts"][1])); ok = False + if not ok: + raise SystemExit(1) + print("\n PASS: portfolio unbiased on both targets, and on the correlated target adaptive " + "allocation concentrates on the full-cov GMM so the portfolio beats standalone AV " + "({:.0f} vs {:.0f} n_eff).".format(pf_c["n_eff"], av_c["n_eff"])) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_balance_heuristic.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_balance_heuristic.py new file mode 100644 index 000000000..0384196be --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_balance_heuristic.py @@ -0,0 +1,255 @@ +#!/usr/bin/env python +""" +test_portfolio_balance_heuristic.py + +Correctness test for mcsamplerPortfolio's SAFETY under a wrong ("decoy") member. + +Setup (the failure mode the fix targets) +---------------------------------------- +The portfolio pools draws from several member samplers and estimates + I = \\int L(theta) prior(theta) dtheta. +Members are re-weighted by their per-member effective sample size (n_ess). A +warm-started AdaptiveVolume (VARAHA) member seeded at a DECOY -- a wrong location +far from the true mode -- draws a tight, self-consistent cloud of LOW-likelihood +points. Those points have nearly-uniform weights, so the decoy member reports a +HIGH per-member n_ess and gets driven to weight ~1, starving the broad covering +member (a GMM/mcsamplerEnsemble) down to the ~1% floor. + + * OLD (STRATIFIED) estimator: each pooled sample keeps its OWN member's + sampling density p_s in L*prior/p_s. Then + E[I_hat] = sum_m w_m * Z_m, + where Z_m is the true integral over member m's support. The decoy member + covers only the (empty) decoy region (Z_decoy ~ 0) and has w ~ 1, so the + estimate is biased LOW by ~ ln(w_covering_floor). A broad member CANNOT + rescue it. This is a real bias, not just variance. + + * NEW (BALANCE-HEURISTIC / deterministic-mixture) estimator: every pooled + sample is weighted by the MIXTURE density + q_mix(theta) = sum_m frac_m * q_m(theta), frac_m = n_drawn_m / n, + evaluated at that sample. Then E[I_hat] = \\int q_mix * L*prior/q_mix = I, + UNBIASED for any member weights, provided the mixture covers the peak. The + broad member's small-but-positive weight guarantees q_mix>0 at the true mode, + so the wrongly-contracted decoy member can no longer bias the result. + +This test builds AV(decoy) + GMM(broad) on a correlated-Gaussian target where a +cold AV converges, and checks: + * OLD estimator -> badly biased low, + * NEW estimator -> unbiased (matches true integral within a few percent), + * and a no-regression control: a NORMAL portfolio (cold AV + GMM, both sane) + stays unbiased under the NEW estimator. + +Usage: + CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=2 OPENBLAS_NUM_THREADS=2 \\ + python test_portfolio_balance_heuristic.py --as-test +""" +from __future__ import print_function +import argparse +import numpy as np + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerAdaptiveVolume as AVmod +from RIFT.integrators import mcsamplerEnsemble as Emod +from RIFT.integrators import mcsamplerPortfolio as Pmod + + +class PeakPlusPlateau(B.CorrelatedGaussian): + """CorrelatedGaussian peak + a small CONSTANT likelihood floor over the whole + box. The floor makes any far-off-peak region a genuine FLAT plateau: an AV + member that contracts onto a tight cloud out there sees an essentially + constant likelihood, so its per-member weights L*prior/p_s are ~uniform and + its Kish n_ess is near-maximal (n). That is exactly the pathology the + portfolio's n_ess re-weighting rewards -- it drives such a decoy member to + weight ~1 and starves the broad covering member -- even though the plateau + carries almost none of the integral. The floor is set to contribute a tiny + fraction of the total evidence so the true integral is essentially unchanged.""" + def __init__(self, floor_frac=1e-3, **kw): + super(PeakPlusPlateau, self).__init__(**kw) + self.name = "peakplateau_d{}".format(self.ndim) + Vbox = float(np.prod(self.rlim - self.llim)) + # floor * Vbox = floor_frac * scale ==> floor contributes floor_frac of Z + self.floor = floor_frac * self.scale / Vbox + self.true_lnZ = np.log(self.scale + self.floor * Vbox) - np.sum(np.log(self.rlim - self.llim)) + + def lnL(self, X): + return np.atleast_1d(np.log(self.scale * self._mvn.pdf(X) + self.floor)) + + +def _host_lnfunc(target): + """A cupy-tolerant wrapper around the benchmark's host integrand. + + The AV member's VARAHA self-update (update_sampling_prior_selfish) evaluates + the integrand on its own DEVICE-native draws (cupy on GPU); the synthetic + benchmark integrand is host/numpy-only and would choke on a cupy array. In + production the ILE likelihood is device-native so this never arises; for the + synthetic target we simply move any device args to the host first, so the + same test exercises the portfolio identically on CPU and GPU.""" + base = target.as_lnfunc() + + def ln_f(*cols): + cols = [Emod.identity_convert(c) for c in cols] + return base(*cols) + return ln_f + + +def _seed_av_decoy(av, decoy): + """Seed the AV member's live-volume state (binunique/dx/V) at the decoy cloud + and apply it to the LIVE attributes draw_simplified() reads. bootstrap_from_* + only stashes self._warm (integrate_log applies it); here the member is driven + through draw_simplified() by the portfolio, so we apply it ourselves. VARAHA's + live volume only ever contracts, so a seed at the decoy stays stuck there.""" + warm = av.bootstrap_from_samples(decoy) # no cover_frac: deliberately wrong + av.binunique = np.array(warm['binunique']) + av.dx = np.array(warm['dx']) + av.nbins = np.array(warm['nbins']) + av.V = float(warm['V']) + av.ninbin = ((av.n_chunk // av.binunique.shape[0] + 1) + * np.ones(av.binunique.shape[0])).astype(int) + + +def _seed_gmm_broad(gmm, target, broad_factor=3.0, n=8000, seed=7): + """Make the GMM member a BROAD but peak-covering proposal: fit it (uniform + weights) to a wide cloud N(mu, broad_factor^2 * cov) around the true mode. + + This is the covering member's job -- a reasonable, deliberately-wider-than-the + -peak proposal (e.g. from a Fisher matrix or a previous posterior). It matters + for the TEST because when the flawed n_ess re-weighting starves this member to + the ~1% floor, its few samples must still land near the peak for the q_mix + estimate to have usable variance; a member left uniform over the whole box + would be unbiased only in expectation but astronomically noisy (the peak is a + ~1e-4 volume fraction). The member still ADAPTS during the run and tightens + further; the point being tested is the ESTIMATOR, given a sane covering member.""" + rng = np.random.RandomState(seed) + cov = broad_factor ** 2 * np.atleast_2d(target.cov) + cloud = rng.multivariate_normal(target.mu, cov, n) + cloud = np.clip(cloud, target.llim + 1e-3, target.rlim - 1e-3) + gmm.update_sampling_prior(np.zeros(len(cloud)), 2 * len(cloud), + external_rvs={p: cloud[:, i] for i, p in enumerate(gmm.params_ordered)}, + log_scale_weights=True) + + +def build_portfolio(target, n_chunk, decoy=None, broad_gmm=True): + """AV + GMM portfolio. If `decoy` is given the AV member is seeded there + (the failure case); otherwise AV starts cold (the no-regression control). + `broad_gmm` pre-fits the GMM as a broad peak-covering proposal.""" + av = AVmod.MCSampler(n_chunk=n_chunk) + gmm = Emod.MCSampler() + members = [av, gmm] + port = Pmod.MCSampler(portfolio=members, portfolio_freeze_wt=0.1, n_chunk=n_chunk) + for d, p in enumerate(target.params): + w = target.rlim[d] - target.llim[d] + port.add_parameter(p, np.vectorize(lambda x, w=w: 1.0 / w), + prior_pdf=np.vectorize(lambda x, w=w: 1.0 / w), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + # propagate GMM configuration (full-covariance single component) through setup + port.setup(portfolio_breakpoints=None, n_comp=1, correlate_all_dims=True, n=n_chunk) + if broad_gmm: + _seed_gmm_broad(gmm, target) + if decoy is not None: + _seed_av_decoy(av, decoy) + # FREEZE the seeded AV member at the decoy: no-op its VARAHA self-update so + # it keeps drawing from the seeded decoy grid every chunk. This is a + # faithful stand-in for "contracted and stuck" -- VARAHA's live volume only + # ever contracts, so a real run seeded here stays near the decoy -- and it + # sidesteps an unrelated numerical edge case (VARAHA's threshold search + # empties the live set when the likelihood is perfectly flat). The point + # of the test is the ESTIMATOR (stratified vs q_mix), not VARAHA dynamics. + av.update_sampling_prior_selfish = (lambda *a, **k: None) + return port, members + + +def run(target, n_chunk, nmax, neff, use_mixture, decoy=None, seed=1234, + tempering_exp=0.3, verbose=False): + np.random.seed(seed) + port, members = build_portfolio(target, n_chunk, decoy=decoy) + ln_f = _host_lnfunc(target) + lnI, logvar, eff, _ = port.integrate_log( + ln_f, *target.params, no_protect_names=True, + nmax=nmax, neff=neff, n=n_chunk, n_adapt=100, + tempering_exp=tempering_exp, floor_level=0.0, use_lnL=True, + save_intg=True, verbose=verbose, + portfolio_use_mixture_density=use_mixture, + # This test isolates the q_mix ESTIMATOR under a PINNED pathological allocation (the decoy + # AV is frozen and, in the stratified case, dominates). Adaptive-probe allocation would + # dynamically re-allocate away from the decoy and change the scenario, so pin it off here; + # the adaptive policy itself is exercised in test_portfolio_adaptive_alloc.py. + portfolio_adaptive_alloc=False) + lnI = float(B._asnumpy(lnI)) + ln_wt = B.log_weights_from_rvs(port._rvs) + return dict(lnI=lnI, bias=lnI - float(target.true_lnZ), + n_eval=int(getattr(port, "ntotal", 0)) or nmax, + n_ess=B.n_ess_kish(ln_wt), + final_weights=np.array(port.portfolio_weights)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ndim", type=int, default=3) + ap.add_argument("--nmax", type=int, default=400000) + ap.add_argument("--neff", type=int, default=2000) + ap.add_argument("--n-chunk", type=int, default=20000) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--verbose", action="store_true") + ap.add_argument("--as-test", action="store_true") + args = ap.parse_args() + + target = PeakPlusPlateau(ndim=args.ndim) + # a TIGHT decoy cloud on the far side of the box from the true mode. On the + # flat plateau its likelihood is ~constant -> maximal per-member n_ess. + span = target.rlim - target.llim + decoy_center = np.clip(target.mu - 0.55 * span * np.sign(target.mu + 1e-9), + target.llim + 0.08 * span, target.rlim - 0.08 * span) + rng = np.random.RandomState(1) + decoy = np.clip(rng.normal(decoy_center, 0.015 * span, size=(4000, args.ndim)), + target.llim + 1e-3, target.rlim - 1e-3) + dist = np.linalg.norm(decoy_center - target.mu) + print("# corrgauss ndim={} true_lnZ={:.4f} mode={} decoy={} |decoy-mode|={:.2f}".format( + args.ndim, target.true_lnZ, np.round(target.mu, 2), + np.round(decoy_center, 2), dist)) + + kw = dict(n_chunk=args.n_chunk, nmax=args.nmax, neff=args.neff, + seed=args.seed, verbose=args.verbose) + + old = run(target, use_mixture=False, decoy=decoy, **kw) # legacy stratified + new = run(target, use_mixture=True, decoy=decoy, **kw) # balance heuristic + ctl = run(target, use_mixture=True, decoy=None, **kw) # no-regression control + + print("\nDECOY portfolio (AV seeded at decoy + broad GMM):") + print(" OLD stratified estimator : lnI-lnZ = {:+.3f} n_ess={:8.1f} wts={}".format( + old["bias"], old["n_ess"], np.round(old["final_weights"], 3))) + print(" NEW q_mix estimator : lnI-lnZ = {:+.3f} n_ess={:8.1f} wts={}".format( + new["bias"], new["n_ess"], np.round(new["final_weights"], 3))) + print("NORMAL portfolio (cold AV + GMM), NEW q_mix estimator:") + print(" control : lnI-lnZ = {:+.3f} n_ess={:8.1f}".format( + ctl["bias"], ctl["n_ess"])) + print("\n# decoy bias improvement: old {:+.3f} -> new {:+.3f} " + "(factor exp = {:.1f}x closer to truth)".format( + old["bias"], new["bias"], + np.exp(abs(old["bias"]) - abs(new["bias"])))) + + if args.as_test: + ok = True + # 1. the OLD estimator MUST demonstrate the danger (biased low) + if not (old["bias"] < -0.7): + print(" FAIL: old stratified estimator not badly biased low ({:+.3f}); " + "decoy not exercised".format(old["bias"])); ok = False + # 2. the NEW estimator must be unbiased within a few percent (a few % in + # the integral is ~0.03-0.20 in ln); allow a modest gate + if abs(new["bias"]) > 0.20: + print(" FAIL: new q_mix estimator biased ({:+.3f} > 0.20)".format(new["bias"])); ok = False + # 3. the new estimator must be dramatically better than the old + if not (abs(new["bias"]) < abs(old["bias"]) - 0.5): + print(" FAIL: q_mix did not fix the decoy bias"); ok = False + # 4. no regression: normal portfolio stays unbiased under q_mix + if abs(ctl["bias"]) > 0.20: + print(" FAIL: normal-portfolio control biased under q_mix " + "({:+.3f})".format(ctl["bias"])); ok = False + if not ok: + raise SystemExit(1) + print("\n PASS: q_mix balance heuristic keeps the portfolio unbiased with a " + "decoy member (old {:+.3f} -> new {:+.3f}); control unbiased " + "({:+.3f}).".format(old["bias"], new["bias"], ctl["bias"])) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_gmm_member_trains.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_gmm_member_trains.py new file mode 100644 index 000000000..d17af207f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_gmm_member_trains.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +"""Regression test: the GMM member of a default-configured portfolio must +actually TRAIN. + +Bug (2026-07-22, found by the PR#28 freeze-territory probe): portfolio setup() +forwarded kwargs that lack n_comp, mcsamplerEnsemble.setup() defaulted +n_comp=None, and update_sampling_prior silently no-opped for n_comp=None -- +so in default wiring the GMM member never trained and every 'portfolio' was +effectively AV-only, with no error or warning. n_comp=0 remains the +intentional off-switch and must stay off. + +Run: python test_portfolio_gmm_member_trains.py (exit 0 = pass) +""" +from __future__ import print_function + +import numpy as np +from scipy.stats import multivariate_normal + +# CPU nodes with cupy installed but no GPU: work around the import-time +# cupy binding in gaussian_mixture_model (fixed separately); harmless once +# that fix lands. +import RIFT.integrators.gaussian_mixture_model as _gmmmod +if hasattr(_gmmmod, "_xpy_eigvals"): + _gmmmod._xpy_eigvals = np.linalg.eigvalsh +if hasattr(_gmmmod, "_xpy_eig"): + _gmmmod._xpy_eig = np.linalg.eig + +from RIFT.integrators import (mcsamplerAdaptiveVolume, mcsamplerEnsemble, + mcsamplerPortfolio) + +rng = np.random.RandomState(31415) +NDIM = 2 +LLIM, RLIM = -5.0, 5.0 +MU = rng.uniform(-1.5, 1.5, NDIM) +COV = 0.3 * np.identity(NDIM) +_mvn = multivariate_normal(MU, COV) + + +def ln_f(*cols): + X = np.array([np.asarray(c, dtype=float) for c in cols]).T + return 100.0 + np.log(_mvn.pdf(np.atleast_2d(X)) + 1e-300) + + +def _gmm_trained(gmm_sampler): + integ = getattr(gmm_sampler, "integrator", None) + if integ is None: + return False + return any(m is not None for m in integ.gmm_dict.values()) + + +def _build_portfolio(portfolio_args=None): + try: + av = mcsamplerAdaptiveVolume.MCSampler(n_chunk=5000) + except TypeError: + av = mcsamplerAdaptiveVolume.MCSampler() + gmm = mcsamplerEnsemble.MCSampler() + port = mcsamplerPortfolio.MCSampler(portfolio=[av, gmm]) + params = ["x{}".format(i) for i in range(NDIM)] + for p in params: + pdf = np.vectorize(lambda x: 1.0 / (RLIM - LLIM)) + port.add_parameter(p, pdf, prior_pdf=pdf, left_limit=LLIM, + right_limit=RLIM, adaptive_sampling=True) + if portfolio_args is None: + port.setup() + else: + port.setup(portfolio_args=portfolio_args) + return port, gmm, params + + +def test_default_portfolio_gmm_member_trains(): + # NO explicit n_comp anywhere: this is exactly the default production wiring. + port, gmm, params = _build_portfolio() + port.integrate_log(ln_f, *params, no_protect_names=True, nmax=60000, + n=5000, neff=50000, n_adapt=100, tempering_exp=0.1, + save_intg=True, verbose=False) + assert _gmm_trained(gmm), ( + "portfolio GMM member never trained (n_comp default regression): " + "gmm_dict models are all None") + print("PASS: portfolio GMM member trained under default configuration") + + +def test_n_comp_zero_remains_off_switch(): + port, gmm, params = _build_portfolio(portfolio_args=[{}, dict(n_comp=0)]) + port.integrate_log(ln_f, *params, no_protect_names=True, nmax=30000, + n=5000, neff=50000, n_adapt=100, tempering_exp=0.1, + save_intg=True, verbose=False) + assert not _gmm_trained(gmm), "n_comp=0 must remain the off-switch" + print("PASS: n_comp=0 off-switch still honored") + + +if __name__ == "__main__": + test_default_portfolio_gmm_member_trains() + test_n_comp_zero_remains_off_switch() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_oracle.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_oracle.py new file mode 100644 index 000000000..1365bffbf --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_oracle.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +""" +test_portfolio_oracle.py + +Quantitative test of the portfolio oracle mechanism, focused on the case where +oracles are supposed to help: a NEEDLE -- a narrow likelihood mode far off-centre +in a large prior box, which uniform sampling almost never finds, but a +Fisher/Gaussian oracle points straight at. + +Compares mcsamplerPortfolio (two adaptive-cartesian members) WITH and WITHOUT a +FisherGaussianOracle, on the same needle, and checks that: + * the oracle-seeded run reaches a substantially higher n_eff / n_ESS, and + * both runs stay unbiased (oracles only propose; they cannot bias the integral). + +Usage: + python test_portfolio_oracle.py # GPU if visible + python test_portfolio_oracle.py --as-test +""" +from __future__ import print_function +import argparse +import numpy as np + +import benchmark_integrators as B +from RIFT.integrators import mcsamplerGPU, mcsamplerPortfolio +from RIFT.integrators.unreliable_oracle.fisher_gaussian import FisherGaussianOracle + + +class Needle(B.Target): + """Narrow correlated Gaussian mode placed far off-centre in a big box. The + posterior occupies a tiny fraction of the prior volume, so uniform sampling + has a very low hit rate -- the regime where a proposal oracle pays off.""" + def __init__(self, ndim=4, width=20.0, scale=100.0, sigma=0.15, seed=3): + self.name = "needle_d{}".format(ndim) + self.ndim = ndim + self.width = width + self.scale = scale + self.params = [str(i) for i in range(ndim)] + self.llim = -0.5 * width * np.ones(ndim) + self.rlim = 0.5 * width * np.ones(ndim) + rng = np.random.RandomState(seed) + self.mu = rng.uniform(0.30 * 0.5 * width, 0.42 * 0.5 * width, ndim) * rng.choice([-1, 1], ndim) + A = rng.normal(size=(ndim, ndim)) + cov = A @ A.T + d = np.sqrt(np.diag(cov)) + cov = cov / np.outer(d, d) * (sigma ** 2) # correlated, ~sigma per dim + self.cov = cov + from scipy.stats import multivariate_normal + self._mvn = multivariate_normal(self.mu, self.cov) + self.true_lnZ = np.log(scale) - np.sum(np.log(self.rlim - self.llim)) + + def lnL(self, X): + return np.atleast_1d(np.log(self.scale * self._mvn.pdf(X) + 1e-300)) + + +def build_portfolio(target, with_oracle, n_chunk): + members = [mcsamplerGPU.MCSampler(), mcsamplerGPU.MCSampler()] + oracles = [] + if with_oracle: + oracles = [FisherGaussianOracle()] + port = mcsamplerPortfolio.MCSampler(portfolio=members, portfolio_freeze_wt=0.1, + oracle_realizations=oracles, n_chunk=n_chunk) + for d, p in enumerate(target.params): + w = target.rlim[d] - target.llim[d] + port.add_parameter(p, np.vectorize(lambda x, w=w: 1.0 / w), + prior_pdf=np.vectorize(lambda x, w=w: 1.0 / w), + left_limit=float(target.llim[d]), right_limit=float(target.rlim[d]), + adaptive_sampling=True) + port.setup(portfolio_breakpoints=None) + for m in members: + m.setup() + if with_oracle: + # seed the Fisher oracle with the (known) mode shape -- in production this + # is a Fisher matrix at the MAP; here we use the true mean/cov. + oracles[0].setup(mean=target.mu, cov=target.cov) + return port + + +def run(target, with_oracle, nmax, neff, n_chunk, seed): + np.random.seed(seed) + port = build_portfolio(target, with_oracle, n_chunk) + ln_f = target.as_lnfunc() + import time + t0 = time.time() + lnI, logvar, eff, _ = port.integrate_log(ln_f, *target.params, no_protect_names=True, + nmax=nmax, neff=neff, n=n_chunk, n_adapt=100, + tempering_exp=0.1, floor_level=0.0, use_lnL=True, + save_intg=True, verbose=False) + wall = time.time() - t0 + lnI = float(B._asnumpy(lnI)); eff = float(B._asnumpy(eff)) + ln_wt = B.log_weights_from_rvs(port._rvs) + ness = B.n_ess_kish(ln_wt) + n_eval = int(getattr(port, "ntotal", 0)) or nmax + return dict(kind="portfolio" + ("+oracle" if with_oracle else ""), + target=target.name, backend="gpu" if mcsamplerGPU.cupy_ok else "cpu", + ndim=target.ndim, n_eval=n_eval, wallclock=wall, lnI=lnI, + true_lnZ=float(target.true_lnZ), bias_ln=lnI - float(target.true_lnZ), + rel_err=float(np.exp(0.5 * (float(B._asnumpy(logvar)) - 2 * lnI))), + n_eff=eff, n_ess=ness, efficiency=eff / max(n_eval, 1), js_marginal=float("nan")) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--ndim", type=int, default=4) + ap.add_argument("--nmax", type=int, default=400000) + ap.add_argument("--neff", type=int, default=500) + ap.add_argument("--n-chunk", type=int, default=20000) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--as-test", action="store_true") + args = ap.parse_args() + + target = Needle(ndim=args.ndim) + print("# needle target ndim={} true_lnZ={:.4f} mode at {}".format( + target.ndim, target.true_lnZ, np.round(target.mu, 2))) + + base = run(target, False, args.nmax, args.neff, args.n_chunk, args.seed) + print("NO-ORACLE ", B._fmt(base)) + orc = run(target, True, args.nmax, args.neff, args.n_chunk, args.seed) + print("ORACLE ", B._fmt(orc)) + + ness_gain = orc["n_ess"] / max(base["n_ess"], 1e-9) + neff_gain = orc["n_eff"] / max(base["n_eff"], 1e-9) + print("\n# oracle vs none: n_eff x{:.2f} n_ESS x{:.2f} " + "bias(no-oracle)={:+.3f} bias(oracle)={:+.3f}".format( + neff_gain, ness_gain, base["bias_ln"], orc["bias_ln"])) + + if args.as_test: + ok = True + tol = max(0.15, 3 * max(base["rel_err"], orc["rel_err"])) + if abs(orc["bias_ln"]) > tol: + print(" FAIL: oracle run biased ({:+.3f} > {:.3f})".format(orc["bias_ln"], tol)); ok = False + if not (orc["n_eff"] >= base["n_eff"] * 1.2 or orc["n_ess"] >= base["n_ess"] * 1.2): + print(" FAIL: oracle did not improve n_eff/n_ESS by >=1.2x"); ok = False + if not ok: + raise SystemExit(1) + print(" PASS: oracle improved sampling and stayed unbiased (tol {:.3f})".format(tol)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_restrict_and_warm.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_restrict_and_warm.py new file mode 100644 index 000000000..80e24d61b --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_portfolio_restrict_and_warm.py @@ -0,0 +1,564 @@ +#!/usr/bin/env python +"""Cheap CPU-only regression tests for two portfolio invariants that fail SILENTLY. + +Both bugs these cover produce a WRONG ANSWER with no exception and a healthy-looking n_eff, so +there is no runtime signal to catch them -- they can only be caught here. + + 1. Range restriction must register coverage bookkeeping. `restrict_member_range()` narrows a + member so it can spend its fixed bin budget where the posterior is. That is safe ONLY while + some member keeps full support: proposals need not cover the prior, but the MIXTURE must cover + the support of L*p. The per-member draw floor, the restricted-only active-member guard and + the q_mix fallback guard are all keyed off `_has_restricted_member`/`_full_support_members`. + If the public API narrows a member without setting them, those guards go dark, member 0 can be + allocated zero draws, and the mixture loses full support -- biasing the integral LOW. + + 2. Warm state must be cleared on the MEMBERS. `_warm` and the contracted AV grid live on each + member; portfolio.integrate_log() does not rerun member setup(). A driver that clears only + `portfolio._warm` leaves the previous point's CONTRACTED live volume installed, so the next + point draws from a box that may exclude its own support. + +Run: python test_portfolio_restrict_and_warm.py +""" +import numpy as np + +import RIFT.integrators.mcsamplerPortfolio as mcsP +import RIFT.integrators.mcsamplerAdaptiveVolume as mcsAV +import RIFT.integrators.mcsamplerEnsemble as mcsGMM + + +def _mk(n=3): + return mcsP.MCSampler(portfolio=[mcsAV] * n) + + +def _mk_av_gmm(): + """AV + GMM, the production portfolio. An AV-only portfolio cannot detect configuration loss + across a reset -- AV.setup() ignores kwargs, so a bare setup() looks identical to a replayed + one. Only a member that CONSUMES its setup arguments can show the difference.""" + return mcsP.MCSampler(portfolio=[mcsAV, mcsGMM]) + + +def _flat(x): + return np.vectorize(lambda z: 0.1) + + +def test_restrict_rejects_invalid_member(): + s = _mk() + for bad in (-1, 0, 99): + try: + s.restrict_member_range(bad, 'x', 0., 1.) + except ValueError: + continue + raise AssertionError( + "restrict_member_range accepted member_index={}; it would never match the positive " + "enumerate() in add_parameter and would be a silent no-op".format(bad)) + + +def test_restrict_sets_coverage_bookkeeping(): + s = _mk() + s.restrict_member_range(1, 'x', -1., 1.) + assert s._has_restricted_member is True + assert s._full_support_members == [0, 2], s._full_support_members + + +def test_restrict_refuses_to_restrict_every_member(): + s = _mk(3) + s.restrict_member_range(1, 'x', -1., 1.) + s.restrict_member_range(2, 'x', -1., 1.) + # members 1 and 2 restricted, member 0 is the backstop -> still fine + assert s._full_support_members == [0] + # and member 0 can never be restricted, so full coverage cannot be lost through this API + try: + s.restrict_member_range(0, 'x', -1., 1.) + except ValueError: + return + raise AssertionError("restrict_member_range narrowed the full-support backstop") + + +def test_unconsumed_restriction_raises_at_setup(): + """A restriction naming a parameter that never arrives must FAIL, not silently do nothing.""" + s = _mk() + s.restrict_member_range(1, 'typo_param', -1., 1.) + s.add_parameter('x', _flat('x'), left_limit=-5., right_limit=5.) + try: + s.setup() + except Exception as e: + assert 'never applied' in str(e), str(e) + return + raise AssertionError("an unapplied range restriction survived setup() as a silent no-op") + + +def test_restriction_narrows_only_that_member(): + s = _mk() + s.restrict_member_range(1, 'x', -1., 1.) + s.add_parameter('x', _flat('x'), left_limit=-5., right_limit=5.) + s.setup() + lims = [(m.llim['x'], m.rlim['x']) for m in s.portfolio_realizations] + assert lims == [(-5., 5.), (-1., 1.), (-5., 5.)], lims + # the PORTFOLIO's own reference limits must stay the full physical range: they are taken from + # member 0 before narrowing, and downstream code uses them as the prior's extent. + assert (s.llim['x'], s.rlim['x']) == (-5., 5.) + + +def test_clear_warm_state_reaches_members(): + s = _mk() + s.add_parameter('x', _flat('x'), left_limit=-5., right_limit=5.) + s.setup() + m = s.portfolio_realizations[1] + m._warm = {'binunique': np.array([[0]]), 'dx': np.array([1.0]), + 'nbins': np.array([1]), 'ninbin': [10], 'V': 0.001} + m._warm_applied = True + m.V = 1e-6 # pretend a heavily contracted live volume + m.dx = np.array([1e-3]) + s.clear_warm_state() + assert m._warm is None + assert m._warm_applied is False + assert m.V == 1, "clear_warm_state left the contracted grid installed (V={})".format(m.V) + assert np.allclose(m.dx, [10.]), m.dx + + +def test_restrict_refuses_to_widen(): + """`restrict` must not silently WIDEN. The prior callables are absolute densities normalized + over the ORIGINAL range, so a member sampling outside it reports a wrong prior and biases the + integral -- exactly the failure the API is supposed to prevent.""" + for lo, hi, what in [(-9., 9., "two-sided"), (-1., 7., "upper-only"), (-7., 1., "lower-only")]: + s = _mk() + s.restrict_member_range(1, 'x', lo, hi) # not contained in the [-5,5] added below + try: + s.add_parameter('x', _flat('x'), left_limit=-5., right_limit=5.) + except ValueError: + continue + lims = (s.portfolio_realizations[1].llim['x'], s.portfolio_realizations[1].rlim['x']) + raise AssertionError("{} widening to [{}, {}] was accepted: member range is now {}".format( + what, lo, hi, lims)) + + +def test_clear_warm_state_preserves_member_configuration(): + """The reset must REPLAY each member's setup arguments. + + A bare setup() restores the cold grid but discards configuration: mcsamplerEnsemble.setup() + rebuilds its dimension grouping and re-reads n_comp/gmm_adapt from kwargs, so a configured + (0,1) GMM would come back as separate (0,), (1,) groups with n_comp defaulted -- a quietly + different sampler for every point after the first. An AV-only portfolio cannot see this.""" + s = _mk_av_gmm() + for p in ('x', 'y'): + s.add_parameter(p, _flat(p), left_limit=-5., right_limit=5.) + # supply gmm_dict, as production does -- the no-gmm_dict path takes a different branch in + # mcsamplerEnsemble.setup() and would not exercise what production actually runs + cfg = dict(n_comp={(0, 1): 3}, gmm_adapt={(0, 1): False}, correlate_all_dims=True, + gmm_dict={(0, 1): None}) + s.setup(**cfg) + gmm = s.portfolio_realizations[1] + + def _snapshot(): + # repr(), not dict(): a bare setup() collapses n_comp from {(0,1): 3} to the scalar + # default, and dict() on an int raises TypeError instead of reporting the defect. + i = gmm.integrator + return (sorted(i.gmm_dict), repr(i.n_comp), repr(i.gmm_adapt)) + + before = _snapshot() + s.clear_warm_state() + after = _snapshot() + assert before[0] == after[0], \ + "reset changed the GMM dimension grouping: {} -> {}".format(before[0], after[0]) + assert before[1] == after[1], "reset lost n_comp: {} -> {}".format(before[1], after[1]) + assert before[2] == after[2], "reset lost gmm_adapt: {} -> {}".format(before[2], after[2]) + + +class _FakeModel(object): + """Stand-in for a trained GMM component; identity is all this test needs.""" + def __repr__(self): + return "" + + +def test_clear_warm_state_clears_trained_proposal_not_just_config(): + """The reset must clear the TRAINED PROPOSAL, not merely restore the grouping. + + `gmm_dict` is not an inert spec. mcsamplerEnsemble hands the caller's dict straight to + monte_carlo.integrator, which stores it WITHOUT copying (MonteCarloEnsemble.py:110) and then + writes trained models into it (`self.gmm_dict[dim_group] = model`, :403). So a reset that + replays a *reference* to those setup arguments hands the next point the previous point's + trained proposal -- reintroducing, through the reset itself, the leak the reset exists to + remove. Checking grouping / n_comp / gmm_adapt alone cannot see this: all three survive.""" + s = _mk_av_gmm() + for p in ('x', 'y'): + s.add_parameter(p, _flat(p), left_limit=-5., right_limit=5.) + caller_spec = {(0, 1): None} + s.setup(n_comp={(0, 1): 3}, gmm_adapt={(0, 1): False}, correlate_all_dims=True, + gmm_dict=caller_spec) + gmm = s.portfolio_realizations[1] + + # the stored args must not alias the caller's dict, or training pollutes them + assert s._member_setup_args[1]['gmm_dict'] is not caller_spec, \ + "stored setup args alias the caller's gmm_dict" + + # point 1 trains: the integrator writes a model into its gmm_dict + gmm.integrator.gmm_dict[(0, 1)] = _FakeModel() + s.clear_warm_state() + assert gmm.integrator.gmm_dict.get((0, 1)) is None, \ + "reset left point 1's trained proposal installed: {}".format(gmm.integrator.gmm_dict) + assert sorted(gmm.integrator.gmm_dict) == [(0, 1)], "reset lost the grouping" + + # point 2 trains, and resets again: the stored snapshot must not have been polluted by the + # first replay (passing the stored dict itself would let the rebuilt integrator train into it, + # so the leak would simply return one point later) + gmm2 = s.portfolio_realizations[1] + gmm2.integrator.gmm_dict[(0, 1)] = _FakeModel() + s.clear_warm_state() + assert gmm2.integrator.gmm_dict.get((0, 1)) is None, \ + "second reset leaked: the stored snapshot was polluted by the first replay" + + +class _AdaptiveSeed(object): + """A seeded GMM whose update() mutates in place, as gaussian_mixture_model.gmm.update does.""" + def __init__(self): + self.tempering_coeff = 1.0 + self.n_updates = 0 + self.means = np.zeros(2) + + def update(self, *a, **kw): + self.tempering_coeff /= 2.0 + self.n_updates += 1 + self.means += 1.0 + + +def test_snapshot_clones_a_real_gmm_model(): + """The clone path must work on the ACTUAL model class, not just a stand-in. + + A real `gmm` holds a module reference (`xpy`) and bound functions, so copy.deepcopy raises + "cannot pickle 'module' object". If the snapshot fell back to sharing on that failure, seeded + models would not be isolated at all -- the fix would be inert on exactly the configuration it + exists for. Hence the shallow-copy-then-clone-attributes path.""" + from RIFT.integrators.gaussian_mixture_model import gmm + rng = np.random.RandomState(0) + m = gmm(2, np.array([[-5., 5.], [-5., 5.]])) + m.fit(rng.normal(size=(400, 2)), log_sample_weights=np.zeros(400)) + clone = mcsP.MCSampler._snapshot_setup_args({'gmm_dict': {(0, 1): m}})['gmm_dict'][(0, 1)] + assert clone is not m, "the real gmm model was SHARED, not cloned" + t0 = m.tempering_coeff + mu0 = None if m.means is None else np.array(m.means) + clone.update(rng.normal(size=(200, 2)), log_sample_weights=np.zeros(200)) + assert m.tempering_coeff == t0 and (mu0 is None or np.array_equal(m.means, mu0)), \ + "mutating the clone changed the original: attribute state is still shared" + + +def test_adaptive_seeded_model_does_not_drift_across_reset(): + """A seeded-and-ADAPTING GMM must not carry point 1's adaptation into point 2. + + Production seeds per-group GMMs from a breadcrumb; with --extrinsic-proposal-adapt those + seeded groups keep re-fitting, and update() mutates the model object in place. Snapshotting + only the containing dict would leave the stored baseline pointing at the live model, so the + baseline drifts during point 1 and is replayed into point 2 -- the same leak one level down. + (With adapt OFF, the default, _train skips seeded groups, so nothing mutates.)""" + s = _mk_av_gmm() + for p in ('x', 'y'): + s.add_parameter(p, _flat(p), left_limit=-5., right_limit=5.) + seed = _AdaptiveSeed() + s.setup(n_comp={(0, 1): 3}, gmm_adapt={(0, 1): True}, correlate_all_dims=True, + gmm_dict={(0, 1): seed}) + gmm_member = s.portfolio_realizations[1] + stored = s._member_setup_args[1]['gmm_dict'][(0, 1)] + assert stored is not seed, "stored baseline aliases the live seeded model" + + # point 1 adapts the live model in place + gmm_member.integrator.gmm_dict[(0, 1)].update(None) + gmm_member.integrator.gmm_dict[(0, 1)].update(None) + assert s._member_setup_args[1]['gmm_dict'][(0, 1)].n_updates == 0, \ + "the stored baseline drifted while point 1 adapted" + + s.clear_warm_state() + restored = s.portfolio_realizations[1].integrator.gmm_dict[(0, 1)] + assert restored is not None, "reset cleared the seeded model entirely" + assert restored.n_updates == 0, \ + "point 2 inherited point 1's adaptation (n_updates={})".format(restored.n_updates) + assert restored.tempering_coeff == 1.0, "point 2 inherited point 1's tempering state" + + +def test_warm_start_keeps_a_backstop_when_every_member_is_compact(): + """With no full-support member, a warm start must not narrow ALL of them. + + `cover_frac` is not a coverage guarantee: a FINITE set of uniform points occupies only the + bins it lands in, so a seeded grid is not a superset of a cold start (measured at d=6, even + cover_frac=0.9 covers 2.9% of the box). In an ALL-AV portfolio every component is a + hard-edged box, so seeding all of them removes the mixture's coverage of the prior box. + Member 0 -- the backstop restrict_member_range also refuses to narrow -- stays cold.""" + d = 4 + s = mcsP.MCSampler(portfolio=[mcsAV, mcsAV]) + for i in range(d): + p = "x%d" % i + s.add_parameter(p, _flat(p), prior_pdf=_flat(p), left_limit=-5., right_limit=5., + adaptive_sampling=True) + s.setup() + rng = np.random.RandomState(1) + cloud = rng.normal(0, 0.2, size=(1500, d)) # a tight seed, as at high SNR + s.bootstrap_from_samples(cloud, cover_frac=0.5) + for i in (0, 1): + s.portfolio_realizations[i].draw_simplified(500) # forces the seed onto the draw path + v0 = float(s.portfolio_realizations[0].V) + v1 = float(s.portfolio_realizations[1].V) + assert v0 >= 1.0, "the full-support backstop was narrowed by the warm start (V={})".format(v0) + assert v1 < 0.5, "member 1 was not actually warm-started (V={}); the test proves nothing".format(v1) + + +def test_warm_start_backstop_opt_out_still_works(): + """The old seed-everything behaviour must remain reachable, for A/B and for callers who + know their seed is right.""" + d = 4 + s = mcsP.MCSampler(portfolio=[mcsAV, mcsAV]) + for i in range(d): + p = "x%d" % i + s.add_parameter(p, _flat(p), prior_pdf=_flat(p), left_limit=-5., right_limit=5., + adaptive_sampling=True) + s.setup() + rng = np.random.RandomState(1) + s.bootstrap_from_samples(rng.normal(0, 0.2, size=(1500, d)), cover_frac=0.5, + keep_backstop_cold=False) + for i in (0, 1): + s.portfolio_realizations[i].draw_simplified(500) + assert float(s.portfolio_realizations[0].V) < 0.5, "opt-out did not seed member 0" + + +def test_warm_start_does_not_sacrifice_av_when_a_gmm_is_present(): + """The rule is "SOME member has support everywhere", not "member 0 is cold". + + A GMM member carries an explicit uniform defensive component (gmm_defensive_frac, default + 0.05) plus Gaussian tails, so q_mix never vanishes however it is seeded -- measured, a + displaced seed in [AV, GMM] left |lnZ bias| <= 0.05 either way. Holding member 0 cold there + would disable the AV warm start entirely (member 0 IS the AV member) to buy a guarantee that + already exists. The merge gate caught exactly that regression, so it is pinned here.""" + d = 4 + s = mcsP.MCSampler(portfolio=[mcsAV, mcsGMM]) + for i in range(d): + p = "x%d" % i + s.add_parameter(p, _flat(p), prior_pdf=_flat(p), left_limit=-5., right_limit=5., + adaptive_sampling=True) + s.setup() + rng = np.random.RandomState(1) + s.bootstrap_from_samples(rng.normal(0, 0.2, size=(1500, d)), cover_frac=0.5) + s.portfolio_realizations[0].draw_simplified(500) + v0 = float(s.portfolio_realizations[0].V) + assert v0 < 0.5, ("the AV member was left cold even though a full-support GMM member is " + "present: the warm start is disabled for no benefit (V={})".format(v0)) + +def test_clear_warm_state_propagates_failures(): + """A reset that quietly did not happen leaves the next point on the previous point's grid. + That must abort, not become a log line.""" + s = _mk() + s.add_parameter('x', _flat('x'), left_limit=-5., right_limit=5.) + s.setup() + + def _boom(**kwargs): + raise RuntimeError("member reset failed") + s.portfolio_realizations[1].setup = _boom + try: + s.clear_warm_state() + except RuntimeError: + return + raise AssertionError("clear_warm_state swallowed a failed member reset") + + + + +def _mk_dim_portfolio(members, d=4, restrict_member=None, **setup_kw): + s = mcsP.MCSampler(portfolio=list(members)) + if restrict_member is not None: + for i in range(d): + s.restrict_member_range(restrict_member, "x%d" % i, -1., 1.) + for i in range(d): + p = "x%d" % i + s.add_parameter(p, _flat(p), prior_pdf=_flat(p), left_limit=-5., right_limit=5., + adaptive_sampling=True) + s.setup(**setup_kw) + return s + + +def _warm_and_measure_member0(s, d=4): + rng = np.random.RandomState(1) + s.bootstrap_from_samples(rng.normal(0, 0.2, size=(1500, d)), cover_frac=0.5) + s.portfolio_realizations[0].draw_simplified(500) + return float(s.portfolio_realizations[0].V) + + +def test_restricted_broad_member_does_not_count_as_the_backstop(): + """A nominally full-support member that has been RANGE-RESTRICTED is confined to a sub-box, + so it no longer covers the prior and must not license contracting everyone else. + + Without this, [unrestricted AV, restricted GMM] reported _full_support_members == [0] and then + warm-started member 0 anyway (measured V=0.095), leaving nothing covering the prior box.""" + s = _mk_dim_portfolio([mcsAV, mcsGMM], restrict_member=1) + assert s._full_support_members == [0], s._full_support_members + v0 = _warm_and_measure_member0(s) + assert v0 >= 1.0, ("member 0 was contracted even though the only other member is " + "range-restricted: nothing covers the prior box (V={})".format(v0)) + + +def test_full_support_capability_must_be_declared(): + """Default FALSE. Treating un-annotated samplers as full-support made any member that simply + had not been marked act as the coverage guarantee.""" + class _Unannotated(object): + pass + assert getattr(_Unannotated(), 'has_unbounded_support', False) is False + assert mcsAV.MCSampler.has_unbounded_support is False + + +def test_defensive_frac_zero_is_not_full_support(): + """The GMM's guarantee is the UNIFORM DEFENSIVE COMPONENT, not Gaussian tails (which underflow + to exactly zero far from the mode). With gmm_defensive_frac=0 it must not be counted.""" + s = _mk_dim_portfolio([mcsAV, mcsGMM], gmm_defensive_frac=0.0) + assert s.portfolio_realizations[1].has_unbounded_support is False + v0 = _warm_and_measure_member0(s) + assert v0 >= 1.0, "member 0 contracted although no member guarantees coverage (V={})".format(v0) + # and the normal case still warm-starts everything + s2 = _mk_dim_portfolio([mcsAV, mcsGMM]) + assert s2.portfolio_realizations[1].has_unbounded_support is True + assert _warm_and_measure_member0(s2) < 0.5 + + +_PORTFOLIO_ADAPTIVE_STATE = ('portfolio_weights', 'portfolio_quality', 'portfolio_quality_nobs', + 'portfolio_probe_ptr', 'portfolio_draw_iteration', + 'portfolio_breakpoints', 'portfolio_member_ness_history') + + +def _adaptive_snapshot(s): + out = {} + for a in _PORTFOLIO_ADAPTIVE_STATE: + v = getattr(s, a, None) + # repr(), not np.array(): the n_ess histories are RAGGED (one list per member, different + # lengths), and np.array on a ragged nested list raises rather than comparing. + out[a] = repr(np.asarray(v).tolist()) if isinstance(v, np.ndarray) else repr(v) + return out + + +def test_reset_adaptation_restores_all_portfolio_state(): + """MC-error replicas must be adaptation-independent at the PORTFOLIO level too. + + clear_warm_state() rebuilds the members, but the portfolio itself learns draw allocation, + per-member quality EMAs and their counts, the probe pointer, the iteration counter and the + n_ess histories. A replica inheriting those starts with scheduling learned from earlier + replicas, so the between-replica scatter -- the very quantity the replicas exist to measure -- + still understates the error.""" + s = _mk_dim_portfolio([mcsAV, mcsGMM]) + before = _adaptive_snapshot(s) + + # simulate a run having adapted the portfolio-level bookkeeping + s.portfolio_weights = np.array([0.9, 0.1]) + s.portfolio_quality = np.array([3.0, 0.2]) + s.portfolio_quality_nobs = np.array([7, 4]) + s.portfolio_probe_ptr = 5 + s.portfolio_draw_iteration = 42 + s.portfolio_member_ness_history = [[1.0, 2.0], [3.0]] + + s.reset_adaptation() + after = _adaptive_snapshot(s) + diffs = [a for a in _PORTFOLIO_ADAPTIVE_STATE if before[a] != after[a]] + assert not diffs, "reset_adaptation left portfolio state carried over: {}".format( + {a: (before[a], after[a]) for a in diffs}) + + + + +def test_every_fit_path_installs_the_defensive_component(): + """The portfolio's coverage guarantee rests on this, so it must hold on EVERY fit path. + + `gmm_defensive_frac > 0` is only a request: add_defensive_component() was called by + fit_gmm_adaptive but NOT by the fixed-component paths, and gmm_adaptive defaults to off -- so + the default configuration asked for a defensive component and never installed one. Worse, + gmm.score() floors at 1e-300, so the member still LOOKED like it had support everywhere; a + sample landing there would carry weight ~1e300. Measured: a fixed-component fit to a tight + cloud returned exactly that floor at the far corner for every d >= 4. + + has_unbounded_support trusts the config while untrained, which is only sound while this holds. + """ + import RIFT.integrators.gaussian_mixture_model as _GMM + rng = np.random.RandomState(0) + for d in (2, 6): + bounds = np.repeat([[-5., 5.]], d, axis=0) + X = rng.normal(0, 0.2, size=(1500, d)) + far = np.full((1, d), 4.9) + + m = _GMM.gmm(2, bounds) + m.fit(X, log_sample_weights=np.zeros(len(X))) + floored = float(np.asarray(m.score(far)).flatten()[0]) + _GMM.add_defensive_component(m, defensive_frac=0.05) + real = float(np.asarray(m.score(far)).flatten()[0]) + assert getattr(m, 'defensive_frac', 0.0) > 0, "marker not set at d={}".format(d) + assert real > 1e6 * max(floored, 1e-300), ( + "defensive component gave no real far-field density at d={}: {:.3g} -> {:.3g}".format( + d, floored, real)) + + # and a member trained through the DEFAULT portfolio path must report the capability honestly + s = _mk_dim_portfolio([mcsAV, mcsGMM], d=2) + gmm = s.portfolio_realizations[1] + assert gmm.has_unbounded_support is True, "untrained default member should trust the config" + s2 = _mk_dim_portfolio([mcsAV, mcsGMM], d=2, gmm_defensive_frac=0.0) + assert s2.portfolio_realizations[1].has_unbounded_support is False, \ + "defensive_frac=0 must never report coverage" + + +def _far_density(model, d): + return float(np.asarray(model.score(np.full((1, d), 4.9))).flatten()[0]) + + +def test_defensive_component_survives_repeated_updates(): + """Coverage must hold through the LIFECYCLE, not just at installation. + + gmm._merge() blends component i with the freshly fitted component order[i] for every i and + never consults self.adapt -- so the broad component, marked adapt=False precisely so it would + be left alone, was dragged toward the fitted cloud on every update. Its mean, covariance and + weight drifted while `defensive_frac` stayed set, so has_unbounded_support kept reporting + coverage that no longer existed. Assert on the actual far-field density and on the + component's own parameters, not on the marker.""" + import RIFT.integrators.gaussian_mixture_model as _GMM + rng = np.random.RandomState(0) + d = 4 + bounds = np.repeat([[-5., 5.]], d, axis=0) + m = _GMM.gmm(2, bounds) + m.fit(rng.normal(0, 0.2, size=(1500, d)), log_sample_weights=np.zeros(1500)) + _GMM.add_defensive_component(m, defensive_frac=0.05) + + def _defensive(): + w = np.asarray(m.identity_convert(m.weights)).flatten()[-1] + mu = np.asarray(m.identity_convert(m.means[-1])).flatten() + cov = np.asarray(m.identity_convert(m.covariances[-1])) + return float(w), mu, cov + + d0 = _far_density(m, d) + w0, mu0, cov0 = _defensive() + for _ in range(4): + m.update(rng.normal(0, 0.2, size=(800, d)), log_sample_weights=np.zeros(800)) + assert _far_density(m, d) > 0.2 * d0, ( + "far-field density collapsed after an update: {:.3g} -> {:.3g}".format( + d0, _far_density(m, d))) + w1, mu1, cov1 = _defensive() + assert abs(w1 - w0) < 1e-9, "defensive weight drifted {:.4f} -> {:.4f}".format(w0, w1) + assert np.allclose(mu1, mu0), "defensive mean drifted toward the fitted cloud" + assert np.allclose(cov1, cov0), "defensive covariance drifted toward the fitted cloud" + + +def test_warm_bootstrap_preserves_the_defensive_opt_in(): + """bootstrap_from_samples() re-runs setup(), which rebuilds the integrator from its kwargs. + Calling it bare reset gmm_defensive_all_paths to False and refitted the warm GMM with no + defensive component -- AFTER the portfolio had already decided, on the strength of that flag, + that it was safe to contract its AV member.""" + d = 4 + s = mcsP.MCSampler(portfolio=[mcsAV, mcsGMM]) + for i in range(d): + p = "x%d" % i + s.add_parameter(p, _flat(p), prior_pdf=_flat(p), left_limit=-5., right_limit=5., + adaptive_sampling=True) + s.setup() + gmm = s.portfolio_realizations[1] + assert gmm.integrator.gmm_defensive_all_paths is True, "portfolio did not opt its member in" + rng = np.random.RandomState(1) + s.bootstrap_from_samples(rng.normal(0, 0.2, size=(1200, d)), cover_frac=0.5) + assert gmm.integrator.gmm_defensive_all_paths is True, \ + "warm bootstrap reset the defensive opt-in" + assert gmm.has_unbounded_support is True, \ + "member stopped guaranteeing coverage after a warm bootstrap, while AV stays contracted" + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith('test_'): + fn() + print("PASS", name) + print("all portfolio restrict/warm invariants hold") diff --git a/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py b/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py new file mode 100644 index 000000000..3794a7d3c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/integrators/test_replica_pooling.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python +"""Tests for MC-error replica pooling and the L0 coverage warning. + +Both concern the same trap: an estimate that is PRECISE but computed over truncated support looks +better by every efficiency metric than a noisy estimate over full support. + +Run: python test_replica_pooling.py +""" +import os +import re +import types + +import numpy + + +def _load_driver_helpers(): + """Import the helpers out of the driver script without executing it.""" + 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 + for fn in ("_rvs_len", "_pool_replica_rvs", "_lnZ_of_rvs", "_kish_neff_of_rvs"): + 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__) + return mod + + +DRV = _load_driver_helpers() + + +class _S(object): + identity_convert = staticmethod(lambda x: x) + + +def _replica(rng, n, lnZ, spread): + """A record whose importance weights average to exp(lnZ).""" + lw = rng.normal(0, spread, size=n) + lw = lw - numpy.log(numpy.mean(numpy.exp(lw))) + lnZ + return dict(log_integrand=lw, log_joint_prior=numpy.zeros(n), + log_joint_s_prior=numpy.zeros(n), x=rng.normal(size=n)) + + +def test_pooling_reproduces_the_combined_evidence(): + """The exported posterior must be a draw from the SAME mixture the reported lnZ describes.""" + rng = numpy.random.RandomState(0) + reps = [_replica(rng, 4000, 0.0, 1.2), _replica(rng, 3000, 0.05, 1.2), + _replica(rng, 5000, -0.03, 1.0)] + Zk = [numpy.mean(numpy.exp(r['log_integrand'])) for r in reps] + lnZ_comb = numpy.log(numpy.mean(Zk)) # the reported combination + pooled = DRV._pool_replica_rvs(reps, _S()) + assert abs(DRV._lnZ_of_rvs(pooled) - lnZ_comb) < 1e-9, ( + "pooled samples imply lnZ {} but the reported combination is {}".format( + DRV._lnZ_of_rvs(pooled), lnZ_comb)) + # unequal replica sizes must still be handled: pooling is 1/(K n_k), not a plain concatenation + assert len(pooled['x']) == 12000 + + +def test_max_neff_selection_would_export_the_collapsed_replica(): + """Why selection by n_eff is the wrong rule: n_eff measures CONCENTRATION, not coverage, so a + mode-collapsed replica scores highest and would be the one exported alongside a combined + evidence it does not represent.""" + rng = numpy.random.RandomState(1) + broad_a = _replica(rng, 4000, 0.0, 1.2) + broad_b = _replica(rng, 4000, 0.05, 1.2) + narrow = _replica(rng, 4000, -0.9, 0.05) # collapsed: low Z, tiny weight spread + neffs = [DRV._kish_neff_of_rvs(r) for r in (broad_a, broad_b, narrow)] + assert int(numpy.argmax(neffs)) == 2, neffs + pooled = DRV._pool_replica_rvs([broad_a, broad_b, narrow], _S()) + # the pooled n_eff must be honest: smaller than the naive sum over replicas + assert DRV._kish_neff_of_rvs(pooled) < sum(neffs) + + +def test_pooling_does_not_repair_a_truncated_support_estimate(): + """The reason the L0 rescue WARNS rather than pooling. + + Averaging Z is unbiased only when every term is unbiased. A warm pass over truncated support + is biased low, and pooling it with a full-support pass yields (Z + Z/2)/2 = 0.75 Z -- better + than warm-only, still wrong. Pinned here so nobody 'improves' the rescue by pooling it.""" + rng = numpy.random.RandomState(2) + full = _replica(rng, 4000, 0.0, 1.0) # unbiased + truncated = _replica(rng, 4000, numpy.log(0.5), 0.2) # missed half the mass + pooled = DRV._pool_replica_rvs([full, truncated], _S()) + bias = DRV._lnZ_of_rvs(pooled) - 0.0 + assert abs(bias - numpy.log(0.75)) < 0.05, ( + "expected pooling to inherit log(0.75) of bias, got {:+.3f}".format(bias)) + + +def test_lnZ_of_rvs_handles_a_single_run_and_a_pooled_record(): + rng = numpy.random.RandomState(3) + r = _replica(rng, 2000, 0.25, 0.8) + assert abs(DRV._lnZ_of_rvs(r, already_pooled=False) - 0.25) < 1e-9 + pooled = DRV._pool_replica_rvs([r, r], _S()) + assert abs(DRV._lnZ_of_rvs(pooled) - 0.25) < 1e-9 + + +def test_missing_columns_degrade_to_the_first_replica(): + """A degraded export is recoverable; a silently mis-weighted one is not.""" + rng = numpy.random.RandomState(4) + a = dict(x=rng.normal(size=10)); b = dict(x=rng.normal(size=10)) + out = DRV._pool_replica_rvs([a, b], _S()) + assert out is a + assert DRV._kish_neff_of_rvs(a) is None + + + + +def test_pooling_uses_reported_lnZ_when_records_are_not_raw(): + """Production _rvs may already be thresholded or fairdraw-resampled. + + Then sum_i w_ki over the RETAINED rows is no longer Z_k * n_k, so a 1/n_k rescale mis-weights + the replica -- a fairdraw record is already posterior-resampled and would be weighted twice. + Given the reported per-replica lnZ, each block is renormalized to contribute exactly Z_k/K, + which is correct whether the rows are raw, pruned or resampled. + """ + rng = numpy.random.RandomState(7) + raw_a = _replica(rng, 4000, 0.0, 1.0) + raw_b = _replica(rng, 4000, 0.3, 1.0) + lnZ = [0.0, 0.3] + # simulate pruning: keep only the top-weight half of replica b + lw_b = raw_b['log_integrand'] + keep = numpy.argsort(lw_b)[len(lw_b) // 2:] + pruned_b = {k: numpy.asarray(v)[keep] for k, v in raw_b.items()} + + target = numpy.log(numpy.mean(numpy.exp(numpy.array(lnZ)))) + pooled = DRV._pool_replica_rvs([raw_a, pruned_b], _S(), rep_lnZ=lnZ) + got = DRV._lnZ_of_rvs(pooled) + assert abs(got - target) < 1e-9, ( + "pooled lnZ {} != reported combination {} for a pruned replica".format(got, target)) + + # without the reported lnZ the naive 1/n_k rescale gets the pruned replica wrong -- which is + # exactly the failure mode this guards against + naive = DRV._lnZ_of_rvs(DRV._pool_replica_rvs([raw_a, pruned_b], _S())) + assert abs(naive - target) > 0.05, ( + "expected the naive rescale to mis-weight a pruned replica; it did not, so this test " + "no longer demonstrates the hazard") + + + + +def _fairdraw(rng, rec): + """Resample a record in proportion to its own weights, as the samplers do for fairdraw output. + + The returned rows keep their ORIGINAL weight columns -- which is exactly the trap. + """ + lw = (numpy.asarray(rec['log_integrand']) + numpy.asarray(rec['log_joint_prior']) + - numpy.asarray(rec['log_joint_s_prior'])) + w = numpy.exp(lw - numpy.max(lw)) + w = w / w.sum() + idx = rng.choice(len(w), size=len(w), replace=True, p=w) + return {k: numpy.asarray(v)[idx] for k, v in rec.items()} + + +def test_fairdraw_blocks_are_not_weighted_twice(): + """Fairdraw samples were already drawn in proportion to their weights. + + Reusing those weights applies them a second time, so the block follows w^2 rather than w. + Renormalizing to Z_k/K fixes the block's SCALE but not its SHAPE, which is why + already_resampled needs its own handling: a fairdraw block is an equal-weight draw from its own + posterior and must contribute constant weights within the block. + """ + rng = numpy.random.RandomState(11) + raw = _replica(rng, 6000, 0.0, 1.3) + fd = _fairdraw(rng, raw) + lnZ = [0.0, 0.0] + + pooled = DRV._pool_replica_rvs([fd, fd], _S(), rep_lnZ=lnZ, already_resampled=True) + lw = (numpy.asarray(pooled['log_integrand']) + numpy.asarray(pooled['log_joint_prior']) + - numpy.asarray(pooled['log_joint_s_prior'])) + # within-block weights must be CONSTANT -- that is what "already carries its weights" means + assert numpy.ptp(lw) < 1e-9, "fairdraw block did not get equal within-block weights" + # and the pooled evidence must still be the reported combination + assert abs(DRV._lnZ_of_rvs(pooled) - 0.0) < 1e-9 + + # the un-handled path leaves a w^2 spread: strictly wider than the w spread it came from + naive = DRV._pool_replica_rvs([fd, fd], _S(), rep_lnZ=lnZ) + lw_naive = (numpy.asarray(naive['log_integrand']) + numpy.asarray(naive['log_joint_prior']) + - numpy.asarray(naive['log_joint_s_prior'])) + assert numpy.ptp(lw_naive) > 1.0, ( + "expected the naive path to retain a spread of weights on an already-resampled block; " + "this test no longer demonstrates the hazard") + # concretely: n_eff of the naive pooling is much worse, because w^2 concentrates + assert DRV._kish_neff_of_rvs(naive) < 0.5 * DRV._kish_neff_of_rvs(pooled) + + + + +def test_cached_log_weights_follow_the_pooled_components(): + """The .dgrid and calibration-posterior exporters PREFER a cached `log_weights` column and only + fall back to the components. mcsamplerPortfolio writes that column, so concatenating the + per-replica caches unchanged would hand those scientific outputs the ORIGINAL weights while the + evidence used the corrected ones -- replica rebalancing ignored, fairdraw blocks double-weighted + again, in exactly the products this pooling exists to make consistent.""" + rng = numpy.random.RandomState(21) + a = _replica(rng, 3000, 0.0, 1.1) + b = _replica(rng, 3000, 0.4, 1.1) + for r in (a, b): # a stale cache, as the portfolio would leave behind + r['log_weights'] = (numpy.asarray(r['log_integrand']) + numpy.asarray(r['log_joint_prior']) + - numpy.asarray(r['log_joint_s_prior'])) + stale_a = numpy.array(a['log_weights']) + + pooled = DRV._pool_replica_rvs([a, b], _S(), rep_lnZ=[0.0, 0.4]) + comp = (numpy.asarray(pooled['log_integrand']) + numpy.asarray(pooled['log_joint_prior']) + - numpy.asarray(pooled['log_joint_s_prior'])) + assert numpy.allclose(pooled['log_weights'], comp), \ + "cached log_weights disagree with the pooled components the estimate used" + # and it must actually have CHANGED -- otherwise the test would pass on the buggy code + assert not numpy.allclose(pooled['log_weights'][:len(stale_a)], stale_a), \ + "pooled log_weights are the stale per-replica values; the cache was not rebuilt" + + +def test_cached_weights_follow_a_flat_fairdraw_block(): + """The case that matters most: a fairdraw block's cache must become constant too, or the + exporters reapply the weights the resampling already used.""" + rng = numpy.random.RandomState(22) + raw = _replica(rng, 4000, 0.0, 1.3) + fd = _fairdraw(rng, raw) + fd['log_weights'] = (numpy.asarray(fd['log_integrand']) + numpy.asarray(fd['log_joint_prior']) + - numpy.asarray(fd['log_joint_s_prior'])) + pooled = DRV._pool_replica_rvs([fd, fd], _S(), rep_lnZ=[0.0, 0.0], already_resampled=True) + assert numpy.ptp(pooled['log_weights']) < 1e-9, \ + "fairdraw block's cached log_weights are not constant: exporters would double-weight it" + + +if __name__ == "__main__": + for name, fn in sorted(globals().items()): + if name.startswith("test_"): + fn() + print("PASS", name) + print("replica pooling / coverage-warning invariants hold") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_fs_consistency.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_fs_consistency.py new file mode 100644 index 000000000..499e722d6 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_fs_consistency.py @@ -0,0 +1,94 @@ +""" +Debug the injection<->recovery consistency for the finite-size (Path D) path. + +Injects a finite-size signal with slowrot_fs_lib and evaluates the JAX +freqresponse likelihood at the EXACT injected extrinsic parameters, comparing to +the Cauchy-Schwarz bound half. If lnL(truth) ~= half and the sky scan +peaks at truth, the likelihood is consistent (any offset is a sampler artifact); +if lnL(truth) < half and a shifted sky scores higher, the injection and +recovery models are inconsistent (a real convention bug). + +Also runs the SAME check on a POINT-response (L_arm->0) injection to isolate +whether the offset is specific to the finite-size path or a general sky/time +convention issue. + +Run in the JAX container: + apptainer exec --nv env ... PYTHONPATH=: \ + python test/jax/debug_fs_consistency.py +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB = os.environ.get("SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB) +import slowrot_fs_lib as fslib + +from RIFT.likelihood.jax_ile.wrapper import build_freqresponse_data_from_precompute +from RIFT.likelihood.jax_ile.core import fused_log_likelihood + +NET = os.environ.get("SLOWROT_NET", "CE+ET") +SNR = float(os.environ.get("SLOWROT_SNR", "300")) +QMAX = int(os.environ.get("SLOWROT_QMAX", "4")) +IWH, TBUF = 0.03, 0.12 + + +def _eval(data, ra, dec, psi, incl, phiref, dist, interp): + return float(np.asarray(fused_log_likelihood( + data, np.array([ra]), np.array([dec]), np.array([psi]), + np.array([incl]), np.array([phiref]), np.array([dist]), interp=interp))[0]) + + +def run(net_name, arm_scale=1.0, tag=""): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=0.4, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD") + net = fslib.network(net_name) + if arm_scale != 1.0: # shrink arms -> point (LWL) response + net = {d: (psd, L * arm_scale) for d, (psd, L) in net.items()} + dist = fslib.distance_for_snr(src, net, SNR) + dd, pd, arm, meta = fslib.build_finite_size_data(src, net, dist) + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + data, _ = build_freqresponse_data_from_precompute( + P0, dd, pd, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm, analyticPSD_Q=True, verbose=False) + half_dd = meta["half_dd"] + rt, dt_, pt, it, ft = src.ra, src.dec, src.psi, src.incl, src.phiref + print("\n=== %s%s SNR=%.0f half=%.1f (arm_scale=%g) ===" % + (net_name, tag, meta["snr"], half_dd, arm_scale)) + for interp in ("nearest", "linear"): + lnL_t = _eval(data, rt, dt_, pt, it, ft, dist, interp) + print(" lnL(truth, %s) = %.1f half-lnL = %.1f (%.2f%%)" % + (interp, lnL_t, half_dd - lnL_t, 100 * (half_dd - lnL_t) / half_dd)) + # fine sky scan about truth (fixed true dist/incl/psi/phiref), CUBIC interp + best = (-np.inf, 0, 0) + for ddec in np.linspace(-2.5, 2.5, 41): + for dra in np.linspace(-2.5, 2.5, 41): + ra = rt + np.radians(dra) / np.cos(dt_) + dec = dt_ + np.radians(ddec) + v = _eval(data, ra, dec, pt, it, ft, dist, "cubic") + if v > best[0]: + best = (v, dra, ddec) + lnL_t = _eval(data, rt, dt_, pt, it, ft, dist, "cubic") + print(" sky scan peak: lnL=%.1f at (dRA*cosd,dDec)=(%+.2f,%+.2f) deg vs lnL(truth)=%.1f " + "=> peak-offset=%.2f deg" % (best[0], best[1], best[2], lnL_t, + np.hypot(best[1], best[2]))) + return best + + +def main(): + print("QMAX=%d" % QMAX) + # 1) finite-size injection + finite-size recovery (the case that showed 1.6 deg) + run(NET, arm_scale=1.0, tag=" [finite-size]") + # 2) near-point injection (arms x0.001 -> LWL response) + same recovery: + # isolates whether the offset is the finite-size response or a general + # sky/time convention (a point injection must peak at truth). + run(NET, arm_scale=1e-3, tag=" [near-point/LWL]") + print("\nDEBUG DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py new file mode 100644 index 000000000..3ac880125 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_jax_vs_cupy_inj.py @@ -0,0 +1,79 @@ +""" +Decisive consistency test: on ONE finite-size injection, evaluate lnL(truth) +with (a) the cupy/numpy freqresponse NoLoop and (b) the JAX banded likelihood, +from the SAME packed precompute, for several t_window / tvals configs. + +If JAX==cupy at truth, the JAX port is consistent with the validated reference +and any lnL(truth) deficit is a precompute/config effect (fixable by matching the +validated t_window/tvals). If JAX!=cupy, the JAX build has a convention bug. +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB = os.environ.get("SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB) +import slowrot_fs_lib as fslib +import RIFT.likelihood.factored_likelihood_freqresponse as flfr +import RIFT.likelihood.slowrot_freqresponse as sfr +import RIFT.lalsimutils as lsu +import lal + +from RIFT.likelihood.jax_ile.banded import build_freqresponse_data +from RIFT.likelihood.jax_ile.core import fused_log_likelihood + +NET = os.environ.get("SLOWROT_NET", "CE+ET") +SNR = float(os.environ.get("SLOWROT_SNR", "300")) +QMAX = int(os.environ.get("SLOWROT_QMAX", "4")) + + +def main(): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=0.4, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD") + net = fslib.network(NET) + dist = fslib.distance_for_snr(src, net, SNR) + dd, pd, arm, meta = fslib.build_finite_size_data(src, net, dist) + deltaT, deltaF = meta["deltaT"], meta["deltaF"] + half_dd = meta["half_dd"] + print("NET=%s SNR=%.0f half=%.1f deltaT=%.3e" % (NET, meta["snr"], half_dd, deltaT)) + + rt, dt_, pt, it, ft = src.ra, src.dec, src.psi, src.incl, src.phiref + det_geom = {d: sfr.detector_geometry(d, L_arm=arm.get(d)) for d in dd} + + for t_window in (0.06, 0.10): + Psig = fslib._base_params(src, dist, deltaT, deltaF) + pk = fslib._pack_finite(fslib.EVENT_TIME, t_window, Psig, dd, pd, arm, src.fmax, QMAX) + for iwh in (0.03, 0.06): + Nw = int(iwh / deltaT); tvals = np.arange(-Nw, Nw) * deltaT + # cupy/numpy NoLoop at truth (nearest + cubic) + Pv = Psig.manual_copy() + Pv.phi = np.array([rt]); Pv.theta = np.array([dt_]); Pv.psi = np.array([pt]) + Pv.incl = np.array([it]); Pv.phiref = np.array([ft]) + Pv.dist = np.array([dist]) * 1e6 * lsu.lsu_PC + Pv.tref = float(fslib.EVENT_TIME); Pv.deltaT = deltaT + cu = {} + for ti in ("nearest", "cubic"): + cu[ti] = float(np.asarray(flfr.DiscreteFactoredLogLikelihoodFreqResponseNoLoop( + tvals, Pv, pk["meta"], pk["lk"], pk["rbp"], pk["ubp"], pk["vbp"], pk["ep"], + Lmax=fslib.LMAX, time_interp=ti, xpy=np))[0]) + # JAX banded from the SAME packed data + data = build_freqresponse_data(pk["meta"], pk["lk"], pk["rbp"], pk["ubp"], + pk["vbp"], pk["ep"], deltaT, tvals, det_geom) + jx = {ti: float(np.asarray(fused_log_likelihood( + data, np.array([rt]), np.array([dt_]), np.array([pt]), + np.array([it]), np.array([ft]), np.array([dist]), + interp=ti))[0]) + for ti in ("nearest", "cubic")} + print(" t_win=%.2f tvals=+/-%.2f : cupy(near/cubic)=%.1f/%.1f " + "JAX(near/lin)=%.1f/%.1f half-cupy_near=%.1f (%.2f%%)" + % (t_window, iwh, cu["nearest"], cu["cubic"], jx["nearest"], jx["cubic"], + half_dd - cu["nearest"], 100 * (half_dd - cu["nearest"]) / half_dd)) + print("DEBUG2 DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_orientation_degeneracy.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_orientation_degeneracy.py new file mode 100644 index 000000000..3c0965c9f --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_orientation_degeneracy.py @@ -0,0 +1,70 @@ +"""Is the SNR~600 orientation offset a genuine (2,2)-mode degeneracy or a sampler miss? + +Rebuilds the representative finite-size event and evaluates the distance- +marginalized lnL at (truth sky, TRUTH orientation) vs (truth sky, RECOVERED +orientation from samples_snr600.npz) and a small scan over (incl, phiref). If +lnL(truth) ~= lnL(recovered), the two orientations are observationally +degenerate for the dominant quadrupole (expected; HM would break it); if +lnL(truth) >> lnL(recovered) the sampler missed the injected mode. +""" +import os, sys +import numpy as np +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB = os.environ.get("SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB) +import slowrot_fs_lib as fslib +from RIFT.likelihood.jax_ile.wrapper import ( + build_freqresponse_data_from_precompute, JAXDistanceMarginalizedLikelihood) + +NET = os.environ.get("SLOWROT_NET", "CE+ET+K") +SNR = float(os.environ.get("SLOWROT_SNR_REP", "600")) +QMAX = int(os.environ.get("SLOWROT_QMAX", "4")) +IWH, TBUF = 0.03, 0.12 +FIGDIR = os.environ.get("SLOWROT_FIG_DIR", + os.path.join(_FSLIB, "3g", "figdata_3site")) + + +def main(): + INCL = float(os.environ.get("SLOWROT_INCL", "0.4")) + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=INCL, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD") + net = fslib.network(NET) + dist = fslib.distance_for_snr(src, net, SNR) + dd, pd, arm, meta = fslib.build_finite_size_data(src, net, dist) + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + half_dd = meta["half_dd"] + rt, dt = src.ra, src.dec + print("=== %s SNR=%.0f half=%.1f ===" % (NET, meta["snr"], half_dd)) + + # Qmax truncation sweep: does lnL(truth)->half and does the injected + # orientation become the MAP as the finite-size basis is resolved? + qsweep = [int(x) for x in os.environ.get("SLOWROT_QSWEEP", "4,8,12,16").split(",")] + for q in qsweep: + data, _ = build_freqresponse_data_from_precompute( + P0, dd, pd, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=q, L_arm=arm, analyticPSD_Q=True, verbose=False) + d_min = max(1.0, dist * 0.3); d_max = dist * 2.5 + like = JAXDistanceMarginalizedLikelihood(data, d_min, d_max, n_grid=256, interp="cubic") + + def L(ra, dec, psi, incl, phiref): + return float(np.asarray(like.log_likelihood( + np.array([ra]), np.array([dec]), np.array([psi]), + np.array([incl]), np.array([phiref]))[0])) + + lnL_truth = L(rt, dt, src.psi, src.incl, src.phiref) + # best-incl at truth sky/psi/phiref (map the degeneracy displacement) + incs = np.radians(np.linspace(2, 80, 40)) + li = np.array([L(rt, dt, src.psi, ic, src.phiref) for ic in incs]) + ibest = int(np.argmax(li)) + print("Qmax=%2d: lnL(truth)=%.1f deficit=half-lnL=%.1f (%.3f%%) " + "best-incl=%.1f deg (truth 22.9), lnL@best-truth=%+.1f" + % (q, lnL_truth, half_dd - lnL_truth, + 100 * (half_dd - lnL_truth) / half_dd, + np.degrees(incs[ibest]), li[ibest] - lnL_truth)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py b/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py new file mode 100644 index 000000000..7985c0073 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/debug_timeshift.py @@ -0,0 +1,88 @@ +"""Is the ~1% lnL(truth) deficit a fractional-sample time misalignment between the +fslib injection and the RIFT precompute t=0? + +Applies a pure time shift exp(2*pi*i*f*dt_shift) to the injected data (per detector, +in FD) over a fine grid of dt_shift within +-1 sample, rebuilds the freqresponse +recovery, and evaluates lnL(truth). If the deficit half-lnL(truth) is +minimized (->~0) at some dt_shift != 0, the injection<->precompute time reference +is misaligned by that fraction of a sample (a fixable convention), and that offset +is the self-consistent injection shift. If the minimum sits at dt_shift=0 with the +deficit intact, the ~1% is a genuine response-model gap, not a time reference. +""" +import os, sys +import numpy as np +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB = os.environ.get("SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB) +import slowrot_fs_lib as fslib +import RIFT.likelihood.factored_likelihood_freqresponse as flfr +from RIFT.likelihood.jax_ile.wrapper import ( + build_freqresponse_data_from_precompute, JAXDistanceMarginalizedLikelihood) + +NET = os.environ.get("SLOWROT_NET", "CE+ET+K") +SNR = float(os.environ.get("SLOWROT_SNR_REP", "600")) +QMAX = int(os.environ.get("SLOWROT_QMAX", "4")) +INCL = float(os.environ.get("SLOWROT_INCL", "1.05")) +IWH, TBUF = 0.03, 0.12 + + +def _shift_data(data_dict, dt_shift): + """Return a copy of data_dict with each series multiplied by exp(2 pi i f dt).""" + out = {} + for det, d in data_dict.items(): + n = d.data.length + fvals = flfr.evaluate_fvals_from_length(n, d.deltaF) + nd = lal_copy(d) + nd.data.data[:] = d.data.data * np.exp(2j * np.pi * fvals * dt_shift) + out[det] = nd + return out + + +def lal_copy(d): + import lal + nd = lal.CreateCOMPLEX16FrequencySeries(d.name, d.epoch, d.f0, d.deltaF, + d.sampleUnits, d.data.length) + nd.data.data[:] = d.data.data + return nd + + +def main(): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=INCL, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, approx="IMRPhenomD") + net = fslib.network(NET) + dist = fslib.distance_for_snr(src, net, SNR) + dd, pd, arm, meta = fslib.build_finite_size_data(src, net, dist) + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + half_dd = meta["half_dd"]; deltaT = meta["deltaT"] + rt, dt = src.ra, src.dec + print("=== %s SNR=%.0f incl=%.1f deg half=%.1f deltaT=%.3e ===" % + (NET, meta["snr"], np.degrees(INCL), half_dd, deltaT)) + + def eval_truth(data_dict): + data, _ = build_freqresponse_data_from_precompute( + P0, data_dict, pd, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm, analyticPSD_Q=True, verbose=False) + d_min = max(1.0, dist * 0.3); d_max = dist * 2.5 + like = JAXDistanceMarginalizedLikelihood(data, d_min, d_max, n_grid=256, interp="cubic") + return float(np.asarray(like.log_likelihood( + np.array([rt]), np.array([dt]), np.array([src.psi]), + np.array([src.incl]), np.array([src.phiref]))[0])), data + + fracs = np.linspace(-1.0, 1.0, 21) + best = (-np.inf, 0.0) + for fr in fracs: + dts = fr * deltaT + lnL_t, _ = eval_truth(_shift_data(dd, dts) if fr != 0 else dd) + print(" dt_shift=%+.3f samples (%+.3e s): lnL(truth)=%.1f deficit=%.1f (%.3f%%)" % + (fr, dts, lnL_t, half_dd - lnL_t, 100 * (half_dd - lnL_t) / half_dd)) + if lnL_t > best[0]: + best = (lnL_t, fr) + print("\nBEST dt_shift = %+.3f samples, lnL(truth)=%.1f deficit=%.1f (%.3f%%)" % + (best[1], best[0], half_dd - best[0], 100 * (half_dd - best[0]) / half_dd)) + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_flowmc.py b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_flowmc.py new file mode 100644 index 000000000..7d0cd73f1 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_flowmc.py @@ -0,0 +1,104 @@ +""" +flowMC on the finite-size (Path D) high-SNR extrinsic posterior — run in parallel +with the coordinate-rotation NUTS (demo_slowrot_highsnr_nuts.py). + +flowMC interleaves local MALA with a normalizing-flow global proposal that LEARNS +the curved multimodal geometry (the sky ring + phase/polarization structure) that +a constant-metric NUTS cannot whiten. This tests whether flowMC holds up where +naive dense-mass NUTS lost ESS at high SNR. + +Requires flowMC importable (installed to ~/flowmc_libs; put it on PYTHONPATH). +Run on GPU in the JAX container (pin an idle GPU): + apptainer exec --nv \ + env PYTHON_JULIAPKG_OFFLINE=yes JAX_ENABLE_X64=1 JAX_ILE_DISTMARG_GH=64 \ + CUDA_VISIBLE_DEVICES= XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 \ + PYTHONPATH=::/analyses/slowrot_finite-size \ + python test/jax/demo_slowrot_flowmc.py +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB_DIR = os.environ.get( + "SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB_DIR) +import slowrot_fs_lib as fslib + +from RIFT.likelihood.jax_ile.wrapper import build_freqresponse_data_from_precompute +from RIFT.likelihood.jax_ile.wrapper import JAXDistanceMarginalizedLikelihood +from RIFT.likelihood.jax_ile import samplers + +NETWORK = os.environ.get("SLOWROT_NET", "CE+ET") +QMAX = 4 +IWH = 0.03 +TBUF = 0.12 +SNRS = [float(x) for x in os.environ.get("SLOWROT_SNRS", "100,300,1000").split(",")] + + +def _gc_dist(ra, dec, ra0, dec0): + c = (np.sin(dec) * np.sin(dec0) + + np.cos(dec) * np.cos(dec0) * np.cos(ra - ra0)) + return np.degrees(np.arccos(np.clip(c, -1.0, 1.0))) + + +def run_one(src, net, target_snr): + dist = fslib.distance_for_snr(src, net, target_snr) + data_dict, psd_dict, arm_dict, meta = fslib.build_finite_size_data(src, net, dist) + print("\n=== target SNR %.0f -> dist=%.2f Mpc actual SNR=%.1f half=%.3e ===" + % (target_snr, dist, meta["snr"], meta["half_dd"])) + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + data, _ = build_freqresponse_data_from_precompute( + P0, data_dict, psd_dict, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm_dict, analyticPSD_Q=True, verbose=False) + + d_min = max(1.0, dist * 0.3) + d_max = dist * 2.5 + like = JAXDistanceMarginalizedLikelihood(data, d_min, d_max, n_grid=256) + + res = samplers.flowmc_sample( + like, d_min, d_max, n_chains=20, n_local_steps=20, n_global_steps=20, + n_training_loops=4, n_production_loops=4, n_epochs=10, + n_prior_pilot=int(max(2e4, 50.0 * target_snr)), seed=1, verbose=True) + + th = np.asarray(res["theta"]); lnLs = np.asarray(res["lnL"]) + ra, dec = th[:, 0], th[:, 1] + imap = int(np.argmax(lnLs)) if len(lnLs) else 0 + d_map = float(_gc_dist(np.array([ra[imap]]), np.array([dec[imap]]), + src.ra, src.dec)[0]) if len(ra) else float("nan") + nb = int(np.clip(np.sqrt(max(len(ra), 1)) / 2.0, 64, 256)) + area = fslib.sky_area_90(ra, dec, np.ones_like(ra), nside_bins=nb) if len(ra) else float("nan") + print(" flowMC: n_draws=%d evidence_neff=%.1f logZ=%.2f 90%% area=%.3e deg^2 " + "MAP d=%.2f deg truth=(%.3f,%.3f)" + % (len(ra), res["neff"], res["logZ"], area, d_map, src.ra, src.dec)) + return dict(target_snr=target_snr, snr=meta["snr"], n=len(ra), + neff=float(res["neff"]), area=float(area), d_map=d_map) + + +def main(): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=0.4, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, + approx="IMRPhenomD") + print("network=%s Qmax=%d distmarg_gh=%s (flowMC)" % + (NETWORK, QMAX, os.environ.get("JAX_ILE_DISTMARG_GH", "0"))) + net = fslib.network(NETWORK) + rows = [] + for snr in SNRS: + try: + rows.append(run_one(src, net, snr)) + except Exception as e: + import traceback; traceback.print_exc() + print(" SNR %.0f FAILED: %s" % (snr, e)) + print("\n==== SUMMARY flowMC (finite-size, network=%s) ====" % NETWORK) + print(" target_snr actual_snr n_draws evid_neff 90%_area_deg2 MAP_deg") + for r in rows: + print(" %8.0f %8.1f %7d %8.1f %12.3e %7.2f" + % (r["target_snr"], r["snr"], r["n"], r["neff"], r["area"], r["d_map"])) + print("FLOWMC DEMO DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_highsnr_nuts.py b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_highsnr_nuts.py new file mode 100644 index 000000000..090d84164 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_highsnr_nuts.py @@ -0,0 +1,233 @@ +""" +High-SNR demonstration: gradient-based NUTS resolves the finite-size (Path D) +extrinsic sky posterior where the production AdaptiveVolume (AV) Monte-Carlo +sampler collapses to n_eff≈1. + +Injects a zero-spin BNS into a 3G network (CE + ET) WITH the frequency-dependent +finite-size detector response (reusing the validated injection machinery in +``~/RIFT_roboto_paper/analyses/slowrot_finite-size/slowrot_fs_lib.py``), builds +the differentiable banded JAX finite-size likelihood (this branch), and runs +``multistart_nuts`` at network SNR 100 / 300 / 1000. Reports, per SNR: + + * n_eff of the evidence estimator (AV gives ~1 at SNR≳100; the whole point), + * 90% credible sky area [deg^2], + * recovered sky (circular mean) vs the injected truth. + +Runs on GPU inside the JAX container: + apptainer exec --nv \ + env PYTHON_JULIAPKG_OFFLINE=yes JAX_ENABLE_X64=1 JAX_ILE_DISTMARG_GH=64 \ + PYTHONPATH=/Code:/analyses/slowrot_finite-size \ + python test/jax/demo_slowrot_highsnr_nuts.py +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +# slowrot_fs_lib lives in the paper repo; allow an env override for its dir. +_FSLIB_DIR = os.environ.get( + "SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB_DIR) +import slowrot_fs_lib as fslib + +from RIFT.likelihood.jax_ile.wrapper import ( + build_freqresponse_data_from_precompute, JAXDistanceMarginalizedLikelihood) +from RIFT.likelihood.jax_ile import samplers + +NETWORK = os.environ.get("SLOWROT_NET", "CE+ET") +QMAX = 4 +IWH = 0.03 # marginalization window half-width [s] +TBUF = 0.12 # rholm buffer half-width [s] (covers CE<->ET delay excursions) +SNRS = [float(x) for x in os.environ.get("SLOWROT_SNRS", "100,300,1000").split(",")] + + +def _posterior_ess(theta_per_chain, dims=(0, 1, 2, 3, 4)): + """Total posterior effective sample size: within-chain ESS summed over chains. + + Uses numpyro.diagnostics.effective_sample_size per chain (shape (1, nsamp)) + and sums, giving the number of effectively-independent posterior draws NUTS + produced. Reduction dims (default all 5) map to physical proxies: + 0 = sin(dec), 1 = sky-x (cos dec cos ra), 2 = psi, 3 = incl, 4 = phiref. + Returns the MIN ESS over the selected dims (worst-mixing = the honest one). + Pass ``dims=(0,1)`` for a sky-only ESS. + """ + if theta_per_chain is None: + return float("nan") + from numpyro.diagnostics import effective_sample_size + tpc = np.asarray(theta_per_chain) # (n_chain, nsamp, 5) + proxy = np.stack([np.sin(tpc[..., 1]), # 0 sin dec + np.cos(tpc[..., 1]) * np.cos(tpc[..., 0]), # 1 sky x + tpc[..., 2], tpc[..., 3], tpc[..., 4]], axis=-1) + per_dim = [] + for j in dims: + tot = 0.0 + for c in range(proxy.shape[0]): + x = proxy[c, :, j][None, :] # (1, nsamp) + try: + tot += float(effective_sample_size(x)) + except Exception: + tot += float(proxy.shape[1]) + per_dim.append(tot) + return float(np.min(per_dim)) + + +def _cov_sky_area_90(ra, dec): + """90% credible sky area [deg^2] from the sample COVARIANCE (Gaussian proxy). + + The histogram sky_area_90 is floored by the bin size (~5 deg^2/cell), so it + cannot show a compact high-SNR peak shrinking. For a compact (near-Gaussian) + sky blob the 90% area is pi * chi2_0.9(2df) * sqrt(det Cov) in a local + East-North tangent plane (x=(ra-ra0)cos dec0, y=dec-dec0), which resolves + sub-deg^2 localization from a few thousand samples. (Overestimates for a + genuinely curved ring arc -- read alongside the histogram value.) + """ + ra = np.asarray(ra); dec = np.asarray(dec) + ra0 = np.angle(np.mean(np.exp(1j * ra))) + dec0 = float(np.mean(dec)) + x = ((ra - ra0 + np.pi) % (2 * np.pi) - np.pi) * np.cos(dec0) + y = dec - dec0 + cov = np.cov(np.vstack([x, y])) + det = float(np.linalg.det(cov)) + if not np.isfinite(det) or det <= 0: + return float("nan") + chi2_90 = 4.60517 # chi2.ppf(0.9, df=2) + area_sr = np.pi * chi2_90 * np.sqrt(det) + return float(area_sr * (180.0 / np.pi) ** 2) + + +def _gc_dist(ra, dec, ra0, dec0): + """Great-circle distance [deg] from each (ra,dec) to (ra0,dec0).""" + c = (np.sin(dec) * np.sin(dec0) + + np.cos(dec) * np.cos(dec0) * np.cos(ra - ra0)) + return np.degrees(np.arccos(np.clip(c, -1.0, 1.0))) + + +def _truth_in_cred(ra, dec, ra0, dec0, cred=0.9, nside_bins=64): + """Is (ra0,dec0) inside the `cred` credible sky region of the samples? + + Uses the same equal-solid-angle (ra x sin dec) binning as + slowrot_fs_lib.sky_area_90; the truth is "in" if its cell is among the + highest-density cells accumulating to `cred` of the mass. + """ + Nra = 2 * nside_bins; Nsd = nside_bins + def _cell(r, d): + ir = int(np.clip((r % (2*np.pi)) / (2*np.pi) * Nra, 0, Nra - 1e-9)) + isd = int(np.clip((np.sin(d) + 1) / 2 * Nsd, 0, Nsd - 1e-9)) + return ir * Nsd + isd + flat = np.array([_cell(r, d) for r, d in zip(ra, dec)]) + H = np.bincount(flat, minlength=Nra * Nsd).astype(float) + H /= H.sum() + order = np.argsort(H)[::-1] + csum = np.cumsum(H[order]) + ncell = int(np.searchsorted(csum, cred) + 1) + keep = set(order[:ncell].tolist()) + return _cell(ra0, dec0) in keep + + +def run_one(src, net, target_snr): + dist = fslib.distance_for_snr(src, net, target_snr) + data_dict, psd_dict, arm_dict, meta = fslib.build_finite_size_data(src, net, dist) + print("\n=== target SNR %.0f -> dist=%.2f Mpc actual SNR=%.1f half=%.3e ===" + % (target_snr, dist, meta["snr"], meta["half_dd"])) + + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + data, extras = build_freqresponse_data_from_precompute( + P0, data_dict, psd_dict, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm_dict, analyticPSD_Q=True, verbose=False) + + # distance bracket around the (narrow, ~d/SNR) posterior + d_min = max(1.0, dist * 0.3) + d_max = dist * 2.5 + like = JAXDistanceMarginalizedLikelihood(data, d_min, d_max, n_grid=256) + + # Pilot must land on the (~1/SNR-thin) time-delay ring to seed NUTS, so scale + # the prior scan with SNR (cheap: the lnL eval is vectorized on GPU). + n_pilot = int(max(2e4, 50.0 * target_snr)) + # Full high-SNR reparameterization (mirrors production RIFT): + # * sky_coords="network" -> baseline-frame sky, straightens the time-delay ring; + # * rotate_phase=True -> (phase_p,phase_m)=(phiref+/-psi), axis-aligns the + # 2psi+/-2phiref degeneracy so the dense mass matrix + # is near-diagonal. + # dense_mass=True mops up the residual; (7,10) caps the pre-adaptation warmup. + res = samplers.multistart_nuts( + like, d_min, d_max, n_starts=6, num_warmup=300, num_samples=500, + n_prior_pilot=n_pilot, seed=1, sky_coords="network", rotate_phase=True, + dense_mass=True, max_tree_depth=(7, 10), verbose=True) + + ra = np.asarray(res["theta"][:, 0]); dec = np.asarray(res["theta"][:, 1]) + lnLs = np.asarray(res["lnL"]) + w = np.ones_like(ra) + # Finer binning at high SNR (the ring is << the default 2.8 deg cells); the + # nside is capped by the pooled sample count so cells stay populated. + nb = int(np.clip(np.sqrt(len(ra)) / 2.0, 64, 256)) + area = fslib.sky_area_90(ra, dec, w, nside_bins=nb) + area_cov = _cov_sky_area_90(ra, dec) # sample-efficient compact-peak area + + # POSTERIOR effective sample size: within-chain ESS summed over chains -- the + # honest "how many effective posterior draws did NUTS get" (contrast: AV gets + # ~1). Distinct from the evidence-estimator neff (Gaussian-mixture IS). + ess = _posterior_ess(res.get("theta_per_chain")) # min over all 5 dims + ess_sky = _posterior_ess(res.get("theta_per_chain"), dims=(0, 1)) # sky only + ess_by = {nm: _posterior_ess(res.get("theta_per_chain"), dims=(j,)) + for j, nm in enumerate(("sindec", "skyx", "psi", "incl", "phiref"))} + print(" per-dim ESS:", {k: int(v) for k, v in ess_by.items()}) + + # Ring-aware sky diagnostics (CE+ET is a 2-SITE timing net -> ring posterior, + # so a circular mean is meaningless). Report: great-circle distance from the + # truth to the NEAREST posterior sample, the MAP (highest-lnL) sample's sky, + # and whether the truth falls inside the 90% credible sky region. + d_ang = _gc_dist(ra, dec, src.ra, src.dec) # (N,) deg + d_near = float(np.min(d_ang)) + imap = int(np.argmax(lnLs)) + ra_map, dec_map = float(ra[imap]), float(dec[imap]) + d_map = float(_gc_dist(np.array([ra_map]), np.array([dec_map]), src.ra, src.dec)[0]) + truth_in90 = _truth_in_cred(ra, dec, src.ra, src.dec, cred=0.9) + + print(" NUTS: posterior_ESS=%.0f sky_ESS=%.0f (of %d pooled draws) " + "evidence_neff=%.1f logZ=%.2f" + % (ess, ess_sky, len(ra), res["neff"], res["logZ"])) + print(" sky: 90%% area hist=%.3e cov=%.3e deg^2 nearest=%.2f deg " + "MAP=(%.3f,%.3f) d=%.2f deg truth=(%.3f,%.3f)" % + (area, area_cov, d_near, ra_map, dec_map, d_map, src.ra, src.dec)) + return dict(target_snr=target_snr, snr=meta["snr"], ess=float(ess), + ess_sky=float(ess_sky), n_pool=int(len(ra)), neff=float(res["neff"]), + logZ=float(res["logZ"]), area=float(area), area_cov=float(area_cov), + d_near=d_near, d_map=d_map, truth_in90=bool(truth_in90)) + + +def main(): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=0.4, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, + approx="IMRPhenomD") + print("network=%s Qmax=%d distmarg_gh=%s" % + (NETWORK, QMAX, os.environ.get("JAX_ILE_DISTMARG_GH", "0"))) + net = fslib.network(NETWORK) + rows = [] + for snr in SNRS: + try: + rows.append(run_one(src, net, snr)) + except Exception as e: + import traceback; traceback.print_exc() + print(" SNR %.0f FAILED: %s" % (snr, e)) + print("\n==== SUMMARY (finite-size, network=%s) ====" % NETWORK) + print(" target_snr actual_snr sky_ESS area_hist area_cov_deg2 MAP_deg") + for r in rows: + print(" %8.0f %8.1f %7.0f %.3e %.3e %7.2f" + % (r["target_snr"], r["snr"], r["ess_sky"], + r["area"], r["area_cov"], r["d_map"])) + print(" (area_cov = covariance-ellipse 90%% area, resolves the compact peak the" + " bin-floored histogram cannot; MAP_deg = recovered sky vs truth.)") + print("\n post_ESS = effective independent POSTERIOR draws from NUTS (the sampling" + " win; AV gives ~1 -- it never lands on the peak).") + print(" evid_neff = Gaussian-mixture importance EVIDENCE estimator quality; it" + " degrades at high SNR because a Gaussian mixture cannot wrap the thin CURVED") + print(" sky ring -- a known jax_ile limitation (ring-aware evidence is" + " future work), NOT a sampling failure.") + print("HIGH-SNR NUTS DEMO DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_reparam.py b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_reparam.py new file mode 100644 index 000000000..46b4cfa7d --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/demo_slowrot_reparam.py @@ -0,0 +1,146 @@ +""" +High-SNR reparameterized sampling: phase-marginalized + Fisher-whitened NUTS on +the finite-size (Path D) extrinsic posterior, vs the naive 5-D multistart NUTS. + +The naive 5-D sampler (demo_slowrot_highsnr_nuts.py) samples (ra,dec,psi,incl, +phiref) directly. Even with a dense mass matrix its posterior ESS collapses at +high SNR (536 -> 34 -> 18 at SNR 100/300/1000) because the sky posterior is a +thin CURVED ring entangled with the psi/phiref degeneracy, which a global +constant metric cannot whiten. + +This driver uses the reparameterization instead (samplers.fisher_nuts_sample_phimarg): + * phi_ref (phase) is MARGINALIZED analytically (JAXDistPhiMargLikelihood), + removing the curved psi/phi_ref ridge -> a 4-D (ra,dec,psi,incl) target; + * each discrete sky mode is Fisher-WHITENED (theta = MAP + A y, A from the + inverse-Fisher), so the ~1/SNR-narrow ring is O(1) scale in y and NUTS keeps + a healthy step at any SNR. + +Reports the POSTERIOR effective sample size (the honest "resolved the posterior" +metric) and sky recovery, to compare against the naive numbers. + +Run on GPU in the JAX container (pin an idle GPU on a shared box): + apptainer exec --nv \ + env PYTHON_JULIAPKG_OFFLINE=yes JAX_ENABLE_X64=1 JAX_ILE_DISTMARG_GH=64 \ + CUDA_VISIBLE_DEVICES= XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 \ + PYTHONPATH=:/analyses/slowrot_finite-size \ + python test/jax/demo_slowrot_reparam.py +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB_DIR = os.environ.get( + "SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB_DIR) +import slowrot_fs_lib as fslib + +from RIFT.likelihood.jax_ile.wrapper import build_freqresponse_data_from_precompute +from RIFT.likelihood.jax_ile.wrapper import JAXDistPhiMargLikelihood +from RIFT.likelihood.jax_ile import samplers + +NETWORK = os.environ.get("SLOWROT_NET", "CE+ET") +QMAX = 4 +IWH = 0.03 +TBUF = 0.12 +NPHI = int(os.environ.get("SLOWROT_NPHI", "32")) +SNRS = [float(x) for x in os.environ.get("SLOWROT_SNRS", "100,300,1000").split(",")] + + +def _gc_dist(ra, dec, ra0, dec0): + c = (np.sin(dec) * np.sin(dec0) + + np.cos(dec) * np.cos(dec0) * np.cos(ra - ra0)) + return np.degrees(np.arccos(np.clip(c, -1.0, 1.0))) + + +def _posterior_ess_4(theta_per_chain): + """Total posterior ESS (within-chain, summed over chains) for 4-D (ra,dec,psi,incl). + + MIN over the sampled dims (worst-mixing direction = the honest number). + """ + if not theta_per_chain: + return float("nan") + from numpyro.diagnostics import effective_sample_size + per_dim = [] + ndim = theta_per_chain[0].shape[-1] + for j in range(ndim): + tot = 0.0 + for th in theta_per_chain: + x = np.asarray(th)[:, j][None, :] + try: + tot += float(effective_sample_size(x)) + except Exception: + tot += float(np.asarray(th).shape[0]) + per_dim.append(tot) + return float(np.min(per_dim)) + + +def run_one(src, net, target_snr): + dist = fslib.distance_for_snr(src, net, target_snr) + data_dict, psd_dict, arm_dict, meta = fslib.build_finite_size_data(src, net, dist) + print("\n=== target SNR %.0f -> dist=%.2f Mpc actual SNR=%.1f half=%.3e ===" + % (target_snr, dist, meta["snr"], meta["half_dd"])) + + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + data, _ = build_freqresponse_data_from_precompute( + P0, data_dict, psd_dict, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm_dict, analyticPSD_Q=True, verbose=False) + + d_min = max(1.0, dist * 0.3) + d_max = dist * 2.5 + # phi_ref-marginalized 4-D target (ra,dec,psi,incl); works on banded data. + like4 = JAXDistPhiMargLikelihood(data, d_min, d_max, nphi=NPHI, n_grid=256) + + res = samplers.fisher_nuts_sample_phimarg( + like4, num_warmup=300, num_samples=500, n_starts=12, n_modes=4, + n_prior_pilot=int(max(2e4, 50.0 * target_snr)), seed=1, verbose=True) + + th = np.asarray(res["theta"]) # (N,4): ra,dec,psi,incl + lnLs = np.asarray(res["lnL"]) + ess = _posterior_ess_4(res.get("theta_per_chain")) + ra, dec = th[:, 0], th[:, 1] + imap = int(np.argmax(lnLs)) + d_map = float(_gc_dist(np.array([ra[imap]]), np.array([dec[imap]]), + src.ra, src.dec)[0]) + d_near = float(np.min(_gc_dist(ra, dec, src.ra, src.dec))) + nb = int(np.clip(np.sqrt(len(ra)) / 2.0, 64, 256)) + area = fslib.sky_area_90(ra, dec, np.asarray(res["post_weight"]), nside_bins=nb) + print(" REPARAM(phimarg+Fisher-whiten): posterior_ESS=%.0f (of %d draws) " + "evidence_neff=%.1f logZ=%.2f" % (ess, len(ra), res["neff"], res["logZ"])) + print(" sky: MAP d=%.2f deg nearest=%.2f deg 90%% area=%.3e deg^2 (nbin=%d) " + "modes=%d truth=(%.3f,%.3f)" + % (d_map, d_near, area, nb, len(res["modes"]), src.ra, src.dec)) + return dict(target_snr=target_snr, snr=meta["snr"], ess=float(ess), + n=len(ra), neff=float(res["neff"]), area=float(area), + d_map=d_map, d_near=d_near, n_modes=len(res["modes"])) + + +def main(): + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=0.4, + phiref=0.0, fmin=50.0, fmax=1024.0, seglen=32.0, + approx="IMRPhenomD") + print("network=%s Qmax=%d nphi=%d distmarg_gh=%s (REPARAMETERIZED sampler)" % + (NETWORK, QMAX, NPHI, os.environ.get("JAX_ILE_DISTMARG_GH", "0"))) + net = fslib.network(NETWORK) + rows = [] + for snr in SNRS: + try: + rows.append(run_one(src, net, snr)) + except Exception as e: + import traceback; traceback.print_exc() + print(" SNR %.0f FAILED: %s" % (snr, e)) + print("\n==== SUMMARY reparam (phimarg+Fisher-whiten), network=%s ====" % NETWORK) + print(" target_snr actual_snr post_ESS evid_neff 90%_area_deg2 MAP_deg modes") + for r in rows: + print(" %8.0f %8.1f %8.0f %8.1f %12.3e %7.2f %4d" + % (r["target_snr"], r["snr"], r["ess"], r["neff"], r["area"], + r["d_map"], r["n_modes"])) + print("\n vs naive 5-D dense-mass NUTS post_ESS: 536 / 34 / 18 at SNR 100/300/1000.") + print("REPARAM DEMO DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py b/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py new file mode 100644 index 000000000..420fe2435 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/make_3g_figdata.py @@ -0,0 +1,196 @@ +""" +Generate figure DATA for the paper's 3G subsection (8.B): + + 1. sky area vs network SNR (Fig-10 analog) -- one row per SNR: + snr, area_cov_deg2, area_hist_deg2, sky_ESS, map_dist_deg + 2. full posterior SAMPLES at a representative SNR (Fig-2/Fig-4 analog pair): + ra, dec, psi, incl, phiref, distMpc (distance drawn from its per-sample + conditional), plus the injected truth. + +Finite-size (Path D) CE+ET injection via slowrot_fs_lib, sampled with the full +high-SNR reparameterization stack (network sky coords + phase rotation + dense +mass + gradient seed-polish). Writes .npz files to SLOWROT_FIG_DIR. + +Run on GPU in the JAX container (pin an idle GPU): + apptainer exec --nv env PYTHON_JULIAPKG_OFFLINE=yes JAX_ENABLE_X64=1 \ + JAX_ILE_DISTMARG_GH=64 CUDA_VISIBLE_DEVICES= \ + XLA_PYTHON_CLIENT_MEM_FRACTION=0.4 SLOWROT_FIG_DIR= \ + PYTHONPATH=:/analyses/slowrot_finite-size \ + python test/jax/make_3g_figdata.py +""" +import os +import sys +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +_FSLIB_DIR = os.environ.get( + "SLOWROT_FS_LIB_DIR", + os.path.expanduser("~/RIFT_roboto_paper/analyses/slowrot_finite-size")) +sys.path.insert(0, _FSLIB_DIR) +import slowrot_fs_lib as fslib + +from RIFT.likelihood.jax_ile.wrapper import ( + build_freqresponse_data_from_precompute, JAXDistanceMarginalizedLikelihood) +from RIFT.likelihood.jax_ile import samplers +from RIFT.likelihood.jax_ile import core as _core + +OUT = os.environ.get("SLOWROT_FIG_DIR", "/tmp/slowrot_3g_fig") +os.makedirs(OUT, exist_ok=True) +NETWORK = os.environ.get("SLOWROT_NET", "CE+ET") +QMAX = int(os.environ.get("SLOWROT_QMAX", "4")) +IWH = 0.03 +TBUF = 0.12 +# SNR ladder for the area-vs-SNR curve; SNR of the representative event. +SNRS = [float(x) for x in os.environ.get("SLOWROT_SNRS", "40,100,200,400,700,1000").split(",")] +SNR_REP = float(os.environ.get("SLOWROT_SNR_REP", "600")) + + +def _cov_sky_area_90(ra, dec): + ra0 = np.angle(np.mean(np.exp(1j * ra))); dec0 = float(np.mean(dec)) + x = ((ra - ra0 + np.pi) % (2 * np.pi) - np.pi) * np.cos(dec0) + y = dec - dec0 + cov = np.cov(np.vstack([x, y])); det = float(np.linalg.det(cov)) + if not np.isfinite(det) or det <= 0: + return float("nan") + return float(np.pi * 4.60517 * np.sqrt(det) * (180.0 / np.pi) ** 2) + + +def _gc_dist(ra, dec, ra0, dec0): + c = np.sin(dec) * np.sin(dec0) + np.cos(dec) * np.cos(dec0) * np.cos(ra - ra0) + return np.degrees(np.arccos(np.clip(c, -1.0, 1.0))) + + +def _draw_distance(data, ra, dec, psi, incl, phiref, d_min, d_max, seed=0): + """Draw a luminosity distance per angular sample from its conditional. + + The per-(sample,time) distance integrand is exp(K x - 0.5 R x^2) with + x = dref/d, K = Re(kappa_unit), R = rho_sq_unit (from _accumulate_unit). At + the matched-filter time bin (max K^2/R, K>0) we draw x from the normalized + p(x) ∝ exp(K x - 0.5 R x^2) x^{-4} (the d^2 volumetric prior in x) on a grid, + by inverse-CDF. Good enough for the figure's distance posterior. + """ + rng = np.random.default_rng(seed) + K, R = _core._accumulate_unit(data, ra, dec, psi, incl, phiref, "cubic", False) + K = np.asarray(K.real); R = np.maximum(np.asarray(R), 1e-30) # (S, npts) + snr2 = np.where(K > 0, K * K / R, -np.inf) + tb = np.argmax(snr2, axis=1) # best time bin + S = ra.shape[0] + Kb = K[np.arange(S), tb]; Rb = R[np.arange(S), tb] + dref = float(data.distMpcRef) + x_lo, x_hi = dref / d_max, dref / d_min + xg = np.linspace(x_lo, x_hi, 512) + d = np.empty(S) + for i in range(S): + lg = Kb[i] * xg - 0.5 * Rb[i] * xg ** 2 - 4.0 * np.log(xg) + lg -= lg.max() + w = np.exp(lg); c = np.cumsum(w); c /= c[-1] + u = rng.random() + xi = np.interp(u, c, xg) + d[i] = dref / xi + return d + + +def run_one(src, net, target_snr, want_samples=False): + dist = fslib.distance_for_snr(src, net, target_snr) + # SELFCONSISTENT (int Qmax): render the injection with the recovery's own b_p*W_p + # response so truth is the exact global maximum -- combined with a finely-sampled + # rholm (fmax>=2048, deltaT<=1/4096) this removes the ~0.16 deg cubic-interpolation + # timing systematic that otherwise displaces the razor-sharp high-SNR sky posterior. + sc = os.environ.get("SLOWROT_SELFCONSISTENT") + data_dict, psd_dict, arm_dict, meta = fslib.build_finite_size_data( + src, net, dist, selfconsistent_Qmax=(int(sc) if sc else None)) + P0 = fslib._base_params(src, dist, meta["deltaT"], meta["deltaF"]) + data, _ = build_freqresponse_data_from_precompute( + P0, data_dict, psd_dict, fslib.EVENT_TIME, IWH, fslib.LMAX, src.fmax, + t_window=TBUF, Qmax=QMAX, L_arm=arm_dict, analyticPSD_Q=True, verbose=False) + d_min = max(1.0, dist * 0.3); d_max = dist * 2.5 + like = JAXDistanceMarginalizedLikelihood(data, d_min, d_max, n_grid=256, interp="cubic") + n_pilot = int(max(2e4, 50.0 * target_snr)) + # At very high SNR + fine deltaT the true peak is thinner than the pilot can + # resolve, so seed one chain at the injected truth (production: intrinsic grid + + # coarse extrinsic pass provides this). Gated on SLOWROT_SEED_TRUTH so the + # area-vs-SNR sweep, where the pilot already finds the (broader) modes, is untouched. + extra = None + if os.environ.get("SLOWROT_SEED_TRUTH"): + extra = np.array([[src.ra, src.dec, src.psi, src.incl, src.phiref]]) + res = samplers.multistart_nuts( + like, d_min, d_max, n_starts=6, num_warmup=300, num_samples=500, + n_prior_pilot=n_pilot, seed=1, sky_coords="network", rotate_phase=True, + dense_mass=True, max_tree_depth=(7, 10), polish_seeds=True, + extra_seeds=extra, verbose=True) + th = np.asarray(res["theta"]); lnL_all = np.asarray(res["lnL"]) + ra, dec, psi, incl, phiref = (th[:, 0], th[:, 1], th[:, 2], th[:, 3], th[:, 4]) + # Dominant-mode mask: the credible region is the region carrying the posterior + # mass; sub-dominant multi-start modes many nats below the peak carry none. + # Keep draws with lnL within DTHR of the peak -- broad at low SNR (one wide + # mode), compact at high SNR (secondary modes dropped). (Raw pooled draws are + # saved; the mask is applied for the area and, in plot_3g, the recovery figure.) + DTHR = 40.0 + dom = lnL_all > (lnL_all.max() - DTHR) + tpc = res.get("theta_per_chain") + from numpyro.diagnostics import effective_sample_size + sky_ess = 0.0 + for c in np.asarray(tpc): + sky_ess += float(effective_sample_size(np.sin(c[:, 1])[None, :])) + # areas over the DOMINANT mode (dom mask) + area_cov = _cov_sky_area_90(ra[dom], dec[dom]) + area_hist = fslib.sky_area_90(ra[dom], dec[dom], np.ones(dom.sum()), nside_bins=64) + try: + area_kde = float(fslib.sky_area_90_kde(ra[dom], dec[dom], np.ones(dom.sum()))) + except Exception: + area_kde = float("nan") + imap = int(np.argmax(lnL_all)) + map_d = float(_gc_dist(np.array([ra[imap]]), np.array([dec[imap]]), src.ra, src.dec)[0]) + row = dict(snr=meta["snr"], area_cov=area_cov, area_hist=area_hist, + area_kde=area_kde, sky_ess=sky_ess, map_dist=map_d, dist_true=dist, + frac_dom=float(dom.mean())) + print(" SNR %.0f: sky_ESS=%.0f area_kde=%.3e area_cov=%.3e MAP=%.2f deg dom=%.0f%%" % + (meta["snr"], sky_ess, area_kde, area_cov, map_d, 100 * dom.mean())) + # save the sky samples at EVERY SNR (with lnL) so areas can be recomputed/plotted + np.savez(os.path.join(OUT, "sky_snr%d.npz" % int(round(target_snr))), + ra=ra, dec=dec, lnL=lnL_all, snr=meta["snr"], + truth=np.array([src.ra, src.dec])) + if want_samples: + dist_s = _draw_distance(data, ra, dec, psi, incl, phiref, d_min, d_max) + np.savez(os.path.join(OUT, "samples_snr%d.npz" % int(round(target_snr))), + ra=ra, dec=dec, psi=psi, incl=incl, phiref=phiref, distMpc=dist_s, + lnL=lnL_all, + truth=np.array([src.ra, src.dec, src.psi, src.incl, src.phiref, dist]), + snr=meta["snr"]) + print(" saved samples_snr%d.npz (%d draws)" % (int(round(target_snr)), len(ra))) + return row + + +def main(): + # Representative-event inclination: the area-vs-SNR sweep is orientation- + # independent (sky localization), but the recovery corner needs an INCLINED + # source (default 60 deg) so the distance-inclination-polarization degeneracy + # of a near-face-on dominant-quadrupole source is broken and the orientation + # sector recovers on truth. Override with SLOWROT_INCL. + incl = float(os.environ.get("SLOWROT_INCL", "0.4")) + # fmax sets both the waveform bandlimit AND the rholm sampling deltaT=1/(2 fmax); + # 2048 (deltaT=1/4096) finely samples the rholm so the recovery's cubic time + # interpolation reproduces the per-detector fractional-sample delays -> no sky bias. + fmax = float(os.environ.get("SLOWROT_FMAX", "1024.0")) + src = fslib.Source(m1=1.6, m2=1.4, ra=1.2, dec=0.3, psi=0.5, incl=incl, + phiref=0.0, fmin=50.0, fmax=fmax, seglen=32.0, + approx="IMRPhenomD") + net = fslib.network(NETWORK) + print("3G FIGDATA network=%s rep_snr=%.0f snrs=%s" % (NETWORK, SNR_REP, SNRS)) + rows = [] + snr_set = sorted(set(SNRS) | {SNR_REP}) + for snr in snr_set: + try: + rows.append(run_one(src, net, snr, want_samples=(snr == SNR_REP))) + except Exception as e: + import traceback; traceback.print_exc(); print(" SNR %.0f FAILED: %s" % (snr, e)) + arr = {k: np.array([r[k] for r in rows]) for k in rows[0]} + np.savez(os.path.join(OUT, "area_vs_snr.npz"), **arr) + print("\nsaved area_vs_snr.npz (%d SNRs) to %s" % (len(rows), OUT)) + print("3G FIGDATA DONE") + + +if __name__ == "__main__": + main() diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py new file mode 100644 index 000000000..4448c5704 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot.py @@ -0,0 +1,171 @@ +""" +Validation ladder for the slow-rotation (Path A/B) and finite-size (Path D) +JAX likelihoods, mirroring test/jax/test_jax_endtoend.py but for the banded +features. + +Gates: + (a) JAX interp="nearest" reproduces the cupy/numpy NoLoop references + DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation (rotation) + DiscreteFactoredLogLikelihoodFreqResponseNoLoop (freqresponse) + on the SAME packed data, to ~1e-13. + (b) interp="linear" gradient (distance-marginalized, smooth) vs finite diff ~1e-6. + (c) jit / vmap / grad / hessian all execute and stay finite. + +Run: + PYTHONPATH=<...>/Code taskset -c 0-3 python test/jax/test_jax_slowrot.py +""" +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) +import jax.numpy as jnp + +import lal +import lalsimulation as lalsim + +import RIFT.lalsimutils as lsu +import RIFT.likelihood.factored_likelihood as fl +import RIFT.likelihood.factored_likelihood_with_rotation as flwr +import RIFT.likelihood.factored_likelihood_freqresponse as flfr +import RIFT.likelihood.slowrot_freqresponse as sfr + +from RIFT.likelihood.jax_ile.core import fused_log_likelihood +from RIFT.likelihood.jax_ile.banded import (build_rotation_data, + build_freqresponse_data) +from RIFT.likelihood.jax_ile.wrapper import JAXDistanceMarginalizedLikelihood + +if not getattr(fl, "numba_on", True): + fl.lalylm = np.vectorize(lal.SpinWeightedSphericalHarmonic, otypes=[complex]) + +fSample = 4096.0; fmin = 30.0; fmax = 1700.0; event_time = 1e9 +t_window = 0.1; Lmax = 2; deltaT = 1.0 / fSample; deltaF = 1.0 / 4.0 +HARM = (-2, -1, 0, 1, 2) +L_CE = 40000.0; Qmax = 4 +PC = lal.PC_SI + +Psig = lsu.ChooseWaveformParams( + fmin=fmin, radec=True, incl=0.3, phiref=0.0, theta=0.2, phi=1.0, psi=0.4, + m1=30 * lal.MSUN_SI, m2=25 * lal.MSUN_SI, detector='H1', + dist=200e6 * lal.PC_SI, deltaT=deltaT, tref=event_time, deltaF=deltaF) +DETS = ("H1", "L1", "V1") +data_dict = {} +for d in DETS: + _p = Psig.manual_copy(); _p.detector = d + data_dict[d] = lsu.non_herm_hoff(_p) +psd_dict = {d: lalsim.SimNoisePSDaLIGOZeroDetHighPower for d in data_dict} +TVALS = np.arange(int(2 * 0.03 / deltaT)) * deltaT - 0.03 + + +def _P_vec(K=48, seed=71): + rng = np.random.RandomState(seed) + Pv = Psig.manual_copy() + Pv.phi = rng.uniform(0, 2 * np.pi, K) + Pv.theta = np.arcsin(rng.uniform(-1, 1, K)) + Pv.psi = rng.uniform(0, np.pi, K) + Pv.incl = np.arccos(rng.uniform(-1, 1, K)) + Pv.phiref = rng.uniform(0, 2 * np.pi, K) + Pv.dist = (rng.uniform(100, 800, K) * 1e6 * lsu.lsu_PC) + Pv.tref = float(event_time); Pv.deltaT = deltaT + return Pv + + +def _distMpc(Pv): + return np.asarray(Pv.dist) / (PC * 1e6) + + +def _finite_diff_grad(fn, x0, h=1e-4): + """Central-difference gradient of a scalar fn at vector x0.""" + g = np.zeros_like(x0) + for i in range(len(x0)): + xp = x0.copy(); xp[i] += h + xm = x0.copy(); xm[i] -= h + g[i] = (fn(xp) - fn(xm)) / (2 * h) + return g + + +def check_rotation(): + print("\n=== ROTATION (Path A, p_max=0) ===") + ri, ct, ctV, rho, meta = flwr.PrecomputeLikelihoodTermsWithRotation( + event_time, t_window, Psig, data_dict, psd_dict, Lmax, fmax, + harmonics=HARM, p_max=0, f_sidereal=flwr.F_SIDEREAL, analyticPSD_Q=True, + verbose=False, quiet=True, skip_interpolation=True) + lk, rbn, ubn, vbn, ep = flwr.pack_rotation_arrays(meta, rho, ct, ctV) + Pv = _P_vec() + lnL_ref = flwr.DiscreteFactoredLogLikelihoodViaArrayVectorNoLoopWithRotation( + TVALS, Pv, meta, lk, rbn, ubn, vbn, ep, Lmax=Lmax, time_interp='nearest', + xpy=np) + data = build_rotation_data(meta, lk, rbn, ubn, vbn, ep, deltaT, TVALS) + lnL_jax = np.asarray(fused_log_likelihood( + data, Pv.phi, Pv.theta, Pv.psi, Pv.incl, Pv.phiref, _distMpc(Pv), + interp="nearest")) + fin = np.isfinite(lnL_ref) & np.isfinite(lnL_jax) + err = np.max(np.abs(lnL_ref[fin] - lnL_jax[fin])) + rel = np.max(np.abs(lnL_ref[fin] - lnL_jax[fin]) / (1 + np.abs(lnL_ref[fin]))) + print("(a) nearest vs numpy NoLoop-with-rotation: max|abs| = %.3e max|rel| = %.3e" + " (%d samples)" % (err, rel, fin.sum())) + assert rel < 1e-10, "rotation nearest mismatch (rel) %g" % rel + return data + + +def check_freqresponse(): + print("\n=== FREQRESPONSE (Path D, Qmax=%d, L=%.0f m) ===" % (Qmax, L_CE)) + bk = flfr.PrecomputeLikelihoodTermsFreqResponse( + event_time, t_window, Psig, data_dict, psd_dict, Lmax, fmax, + Qmax=Qmax, L_arm=L_CE, analyticPSD_Q=True, verbose=False, quiet=True, + skip_interpolation=True) + meta = bk[4] + lk, rbp, ubp, vbp, ep = flfr.pack_freqresponse_arrays(bk[4], bk[3], bk[1], bk[2]) + Pv = _P_vec() + lnL_ref = flfr.DiscreteFactoredLogLikelihoodFreqResponseNoLoop( + TVALS, Pv, meta, lk, rbp, ubp, vbp, ep, Lmax=Lmax, time_interp='nearest', + xpy=np) + det_geom = {d: sfr.detector_geometry(d, L_arm=L_CE) for d in DETS} + data = build_freqresponse_data(meta, lk, rbp, ubp, vbp, ep, deltaT, TVALS, + det_geom) + lnL_jax = np.asarray(fused_log_likelihood( + data, Pv.phi, Pv.theta, Pv.psi, Pv.incl, Pv.phiref, _distMpc(Pv), + interp="nearest")) + fin = np.isfinite(lnL_ref) & np.isfinite(lnL_jax) + err = np.max(np.abs(lnL_ref[fin] - lnL_jax[fin])) + rel = np.max(np.abs(lnL_ref[fin] - lnL_jax[fin]) / (1 + np.abs(lnL_ref[fin]))) + print("(a) nearest vs numpy FreqResponse NoLoop: max|abs| = %.3e max|rel| = %.3e" + " (%d samples)" % (err, rel, fin.sum())) + assert rel < 1e-10, "freqresponse nearest mismatch (rel) %g" % rel + return data + + +def check_ad(data, tag): + print("--- AD checks (%s) ---" % tag) + # (c) jit + vmap of the fixed-distance likelihood + f = jax.jit(lambda ra, dec, psi, incl, phiref, d: fused_log_likelihood( + data, ra, dec, psi, incl, phiref, d, interp="linear")) + th = (jnp.array([1.0]), jnp.array([0.2]), jnp.array([0.4]), + jnp.array([0.9]), jnp.array([1.1]), jnp.array([300.0])) + v = np.asarray(f(*th)) + assert np.all(np.isfinite(v)), "jit likelihood non-finite" + print("(c) jit fused_log_likelihood finite: lnL = %.4f" % v[0]) + + # (b,c) distance-marginalized: grad vs finite diff, hessian finite + dlike = JAXDistanceMarginalizedLikelihood(data, 5.0, 3000.0, n_grid=128, + interp="linear") + x0 = np.array([1.0, 0.2, 0.4, 0.9, 1.1]) + val, grad = dlike.value_and_grad(x0) + fd = _finite_diff_grad(lambda x: dlike.value(x), x0, h=1e-4) + rel = np.max(np.abs(grad - fd) / (1 + np.abs(fd))) + print("(b) distmarg grad vs finite-diff: max|rel| = %.3e" % rel) + print(" grad =", np.array2string(np.asarray(grad), precision=4)) + print(" fin-diff =", np.array2string(fd, precision=4)) + assert np.all(np.isfinite(grad)), "distmarg grad non-finite" + assert rel < 1e-4, "distmarg grad disagrees with finite diff: %g" % rel + H = dlike.fisher(x0) + assert np.all(np.isfinite(H)), "hessian non-finite" + print("(c) hessian finite, Fisher diag =", + np.array2string(np.diag(H), precision=2)) + + +if __name__ == "__main__": + d_rot = check_rotation() + check_ad(d_rot, "rotation") + d_fr = check_freqresponse() + check_ad(d_fr, "freqresponse") + print("\nSLOWROT + FREQRESPONSE JAX VALIDATION PASSED") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py new file mode 100644 index 000000000..c2512e6cd --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_coeffs.py @@ -0,0 +1,95 @@ +""" +Gate 1a: JAX response-coefficient ports vs the numpy references, to ~1e-12. + +Validates (no heavy precompute needed -- pure analytic algebra): + * response_slowrot.rotation_coefficients_dict vs + factored_likelihood_with_rotation.rotation_coefficients_vector (Path A & B) + * response_freqresponse.response_coefficients_dict vs + factored_likelihood_freqresponse.response_coefficients (Path D) + +over H1/L1/V1 and random (RA,DEC,psi). + +Run: + PYTHONPATH=<...>/Code taskset -c 0-3 python test/jax/test_jax_slowrot_coeffs.py +""" +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +import lal +import lalsimulation as lalsim + +import RIFT.likelihood.factored_likelihood_with_rotation as flwr +import RIFT.likelihood.factored_likelihood_freqresponse as ffr +import RIFT.likelihood.slowrot_freqresponse as sfr +from RIFT.likelihood.jax_ile import response_slowrot as rs +from RIFT.likelihood.jax_ile import response_freqresponse as rf + +DETS = ["H1", "L1", "V1"] +TREF = 1126259462.0 +HARM = (-2, -1, 0, 1, 2) + + +def _gmst(tref): + return float(lal.GreenwichMeanSiderealTime(lal.LIGOTimeGPS(float(tref)))) + + +def test_rotation_coefficients(): + rng = np.random.default_rng(3) + S = 64 + RA = rng.uniform(0, 2 * np.pi, S) + DEC = np.arcsin(rng.uniform(-1, 1, S)) + psi = rng.uniform(0, np.pi, S) + gmst = _gmst(TREF) + worst = 0.0 + for p_max in (0, 1, 2): + for det in DETS: + lald = lalsim.DetectorPrefixToLALDetector(det) + resp = np.asarray(lald.response, dtype=float) + loc = np.asarray(lald.location, dtype=float) + C_np = flwr.rotation_coefficients_vector(det, RA, DEC, psi, TREF, p_max) + C_jx = rs.rotation_coefficients_dict(resp, loc, RA, DEC, psi, gmst, p_max) + # compare the union of keys + keys = set(C_np) | set(C_jx) + for k in keys: + a = np.asarray(C_np.get(k, np.zeros(S, complex))) + b = np.asarray(C_jx.get(k, np.zeros(S, complex))) + d = np.max(np.abs(a - b)) + worst = max(worst, d) + print("[rotation coeff] max|jax-np| over dets/p_max = %.3e" % worst) + assert worst < 1e-11, "rotation coefficient mismatch %g" % worst + + +def test_freqresponse_coefficients(): + rng = np.random.default_rng(7) + S = 64 + RA = rng.uniform(0, 2 * np.pi, S) + DEC = np.arcsin(rng.uniform(-1, 1, S)) + psi = rng.uniform(0, np.pi, S) + gmst = _gmst(TREF) + Qmax = 4 + worst = 0.0 + for L_arm in (None, 40000.0): # native LIGO arm and a 40-km CE arm + for det in DETS: + resp, x_arm, y_arm, L = sfr.detector_geometry(det, L_arm=L_arm) + b_jx = rf.response_coefficients_dict(resp, x_arm, y_arm, RA, DEC, psi, + gmst, Qmax) + # numpy reference is scalar -> loop the S samples + b_np = {p: np.empty(S, complex) for p in range(Qmax + 2)} + for i in range(S): + bi = ffr.response_coefficients(det, float(RA[i]), float(DEC[i]), + float(psi[i]), TREF, Qmax, L_arm=L_arm) + for p in range(Qmax + 2): + b_np[p][i] = bi[p] + for p in range(Qmax + 2): + d = np.max(np.abs(np.asarray(b_jx[p]) - b_np[p])) + worst = max(worst, d) + print("[freqresponse coeff] max|jax-np| over dets/L = %.3e" % worst) + assert worst < 1e-11, "freqresponse coefficient mismatch %g" % worst + + +if __name__ == "__main__": + test_rotation_coefficients() + test_freqresponse_coefficients() + print("COEFFICIENT PORTS VALIDATED (Gate 1a PASSED)") diff --git a/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py new file mode 100644 index 000000000..5b3b8a14d --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/jax/test_jax_slowrot_wrapper.py @@ -0,0 +1,76 @@ +""" +Smoke test for the one-call banded builders (build_*_data_from_precompute) and +package import. Confirms they reproduce the manual precompute+pack+build path +and yield a finite, differentiable likelihood. + +Run: + PYTHONPATH=<...>/Code taskset -c 0-3 python test/jax/test_jax_slowrot_wrapper.py +""" +import numpy as np + +import jax +jax.config.update("jax_enable_x64", True) + +import lal +import lalsimulation as lalsim + +import RIFT.lalsimutils as lsu +import RIFT.likelihood.factored_likelihood as fl + +# exercise the package public API +from RIFT.likelihood.jax_ile import ( + build_rotation_data_from_precompute, + build_freqresponse_data_from_precompute, +) +from RIFT.likelihood.jax_ile.wrapper import JAXDistanceMarginalizedLikelihood +from RIFT.likelihood.jax_ile.core import fused_log_likelihood + +if not getattr(fl, "numba_on", True): + fl.lalylm = np.vectorize(lal.SpinWeightedSphericalHarmonic, otypes=[complex]) + +fSample = 4096.0; fmin = 30.0; fmax = 1700.0; event_time = 1e9 +Lmax = 2; deltaT = 1.0 / fSample; deltaF = 1.0 / 4.0 +IWH = 0.03 +PC = lal.PC_SI + +P = lsu.ChooseWaveformParams( + fmin=fmin, radec=True, incl=0.3, phiref=0.0, theta=0.2, phi=1.0, psi=0.4, + m1=30 * lal.MSUN_SI, m2=25 * lal.MSUN_SI, detector='H1', + dist=200e6 * lal.PC_SI, deltaT=deltaT, tref=event_time, deltaF=deltaF) +DETS = ("H1", "L1") +data_dict = {} +for d in DETS: + _p = P.manual_copy(); _p.detector = d + data_dict[d] = lsu.non_herm_hoff(_p) +psd_dict = {d: lalsim.SimNoisePSDaLIGOZeroDetHighPower for d in data_dict} + +rng = np.random.RandomState(5) +S = 16 +ra = rng.uniform(0, 2 * np.pi, S) +dec = np.arcsin(rng.uniform(-1, 1, S)) +psi = rng.uniform(0, np.pi, S) +incl = np.arccos(rng.uniform(-1, 1, S)) +phiref = rng.uniform(0, 2 * np.pi, S) +distMpc = rng.uniform(100, 800, S) + + +def _run(builder, tag, **kw): + data, extras = builder(P.manual_copy(), data_dict, psd_dict, event_time, + IWH, Lmax, fmax, analyticPSD_Q=True, verbose=False, **kw) + lnL = np.asarray(fused_log_likelihood(data, ra, dec, psi, incl, phiref, + distMpc, interp="nearest")) + assert np.all(np.isfinite(lnL)), "%s produced non-finite lnL" % tag + # differentiable distmarg path + dlike = JAXDistanceMarginalizedLikelihood(data, 5.0, 3000.0, n_grid=64) + v, g = dlike.value_and_grad([ra[0], dec[0], psi[0], incl[0], phiref[0]]) + assert np.isfinite(v) and np.all(np.isfinite(g)), "%s distmarg AD non-finite" % tag + print("[%s] one-call build OK: lnL[0]=%.3f distmarg lnL=%.3f |grad|=%.2f" + % (tag, lnL[0], v, np.linalg.norm(g))) + return data + + +if __name__ == "__main__": + _run(build_rotation_data_from_precompute, "rotation", p_max=0) + _run(build_freqresponse_data_from_precompute, "freqresponse", Qmax=4, + L_arm=40000.0) + print("ONE-CALL BUILDER SMOKE TEST PASSED") diff --git a/MonteCarloMarginalizeCode/Code/test/test_calmarg_calibration.py b/MonteCarloMarginalizeCode/Code/test/test_calmarg_calibration.py new file mode 100644 index 000000000..43cdb0a4c --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_calmarg_calibration.py @@ -0,0 +1,37 @@ +import pytest + +from RIFT.calmarg.calibration import correction_type_for_ifo + + +@pytest.mark.parametrize( + "ifo_name, expected", + [("H1", "data"), ("L1", "data"), ("K1", "data"), ("V1", "template")], +) +def test_bilby_pipe_default_correction_types(ifo_name, expected): + assert correction_type_for_ifo(None, ifo_name) == expected + + +@pytest.mark.parametrize("setting", ["data", "template"]) +def test_global_correction_type(setting): + assert correction_type_for_ifo(setting, "V1") == setting + + +def test_detector_specific_correction_types(): + setting = {"H1": "template", "V1": "data"} + assert correction_type_for_ifo(setting, "H1") == "template" + assert correction_type_for_ifo(setting, "V1") == "data" + + +def test_string_detector_specific_correction_types(): + def parse_dict(value): + assert value == "{H1: data, V1: template}" + return {"H1": "data", "V1": "template"} + + assert correction_type_for_ifo( + "{H1: data, V1: template}", "V1", parse_dict=parse_dict + ) == "template" + + +def test_missing_detector_is_rejected(): + with pytest.raises(ValueError, match="No calibration correction type"): + correction_type_for_ifo({"H1": "data"}, "V1") diff --git a/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py b/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py new file mode 100644 index 000000000..8bf495a3d --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/test_srate_resample_time_marginalization.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +""" +Unit tests for the time-marginalisation upsampling used by +``--srate-resample-time-marginalization`` in +``bin/integrate_likelihood_extrinsic_batchmode``. + +Before the fix, the option was effectively a boolean: whenever the requested +rate exceeded --srate, the internal time grid was refined by a hardcoded factor +of two and the requested value was discarded. With the O4c production settings +(--srate 4096, --data-integration-window-half 0.075) asking for 16384 Hz +delivered ~8173 Hz, and the exported geocentre times inherited that resolution. + +These tests exercise a transcription of the shipped block, kept in sync by +``test_source_matches_reference_implementation`` below. +""" + +import os +import re + +import numpy as np +import pytest + +# RIFT defaults exercised by the O4c production configuration. +SRATE = 4096.0 +WINDOW_HALF = 75e-3 # --data-integration-window-half default +REQUESTED = 16384 + +ILE_SCRIPT = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "bin", + "integrate_likelihood_extrinsic_batchmode", +) + + +def rift_tvals(srate=SRATE, window_half=WINDOW_HALF): + """The internal time grid built by analyze_event_extrinsic_export.""" + n_points = int(2 * window_half / (1.0 / srate)) + return np.linspace(-window_half, window_half, n_points) + + +def upsample(tvals, lnLt, requested, fsample=SRATE): + """Reference implementation, mirroring the shipped code.""" + from scipy.interpolate import CubicSpline + + if not (requested and requested > fsample): + return tvals, lnLt + dt_target = 1.0 / requested + n_dense = int(np.floor((tvals[-1] - tvals[0]) / dt_target)) + 1 + tvals_denser = tvals[0] + dt_target * np.arange(n_dense) + lnLt_new = np.zeros((lnLt.shape[0], n_dense)) + for index in range(lnLt.shape[0]): + lnLt_new[index] = CubicSpline(tvals, lnLt[index])(tvals_denser) + return tvals_denser, lnLt_new + + +def output_spacing(tvals): + """The one spacing of the (uniform) output grid.""" + diffs = np.diff(tvals) + assert np.allclose(diffs, diffs[0], rtol=0, atol=1e-15), "grid is not uniform" + return diffs[0] + + +def effective_rate(tvals): + return 1.0 / output_spacing(tvals) + + +@pytest.fixture +def toy_lnl(): + """A smooth, sharply peaked lnL(t): Gaussians of ~1 ms width.""" + tvals = rift_tvals() + peak = np.array([[-0.7e-3], [0.0], [1.3e-3]]) + return tvals, -0.5 * ((tvals[None, :] - peak) / 1.0e-3) ** 2 + + +def test_internal_grid_is_slightly_coarser_than_srate(): + """ + linspace(-W, W, N) with N = int(2*W*fS) spans the closed interval with N + points, so the spacing is deltaT*N/(N-1) - about 0.2% coarser than 1/fS. + The refinement factor must therefore be derived from the grid spacing, not + from fSample, or the result lands just short of the requested rate. + """ + tvals = rift_tvals() + assert len(tvals) == 614 + assert tvals[1] - tvals[0] > 1.0 / SRATE + assert effective_rate(tvals) == pytest.approx(4086.67, rel=1e-4) + + +@pytest.mark.parametrize("requested", [8192, 16384, 32768, 65536]) +def test_recovers_the_exact_requested_rate(toy_lnl, requested): + """ + The whole point of the fix: the output rate must equal the requested rate, + not merely reach or exceed it. The requested rates are powers of two, so + 1/requested is exactly representable in float64 and consecutive output + times differ by exactly that step, to the bit. + """ + tvals, lnl = toy_lnl + dense, _ = upsample(tvals, lnl, requested) + spacing = output_spacing(dense) + assert spacing == 1.0 / requested # bit-exact, not approx + assert effective_rate(dense) == float(requested) + + +@pytest.mark.parametrize("requested", [16384, 32768]) +def test_output_times_lie_on_the_requested_grid(toy_lnl, requested): + """ + Every output time is tvals[0] + k/requested for integer k, i.e. the + exported geocenter time is quantized at exactly 1/requested seconds. + """ + tvals, lnl = toy_lnl + dense, _ = upsample(tvals, lnl, requested) + k = (dense - dense[0]) * requested + np.testing.assert_allclose(k, np.round(k), rtol=0, atol=1e-9) + + +def test_scales_with_the_request(toy_lnl): + """Doubling the request exactly halves the output spacing.""" + tvals, lnl = toy_lnl + s16 = output_spacing(upsample(tvals, lnl, 16384)[0]) + s32 = output_spacing(upsample(tvals, lnl, 32768)[0]) + assert s16 == 2.0 * s32 + + +def test_does_not_extrapolate_outside_the_original_grid(toy_lnl): + """ + The dense grid must stay within [tvals[0], tvals[-1]] so the cubic spline + never extrapolates (the old grid ran half a sample past tvals[-1]). We + floor the point count, so the far edge is left short by < 1/requested s. + """ + tvals, lnl = toy_lnl + dense, _ = upsample(tvals, lnl, REQUESTED) + assert dense[0] == tvals[0] + assert dense[-1] <= tvals[-1] + assert (tvals[-1] - dense[-1]) < 1.0 / REQUESTED + + +def test_recovers_the_peak_to_the_requested_resolution(toy_lnl): + """ + The exported geocentre time is drawn from this grid, so the grid spacing + floors the achievable time resolution. + """ + tvals, lnl = toy_lnl + truth = np.array([-0.7e-3, 0.0, 1.3e-3]) + dense, lnl_dense = upsample(tvals, lnl, REQUESTED) + error = np.abs(dense[np.argmax(lnl_dense, axis=1)] - truth).max() + assert error < 1.0 / REQUESTED + + +def test_switch_is_off_at_or_below_fsample(toy_lnl): + tvals, lnl = toy_lnl + for requested in (None, 0, 2048, int(SRATE)): + dense, lnl_dense = upsample(tvals, lnl, requested) + np.testing.assert_array_equal(dense, tvals) + np.testing.assert_array_equal(lnl_dense, lnl) + + +def test_source_matches_reference_implementation(): + """ + Guard against the shipped block and this reference drifting apart - the + tests above are only meaningful if they describe the real code. + """ + if not os.path.exists(ILE_SCRIPT): + pytest.skip("ILE script not found next to the test directory") + with open(ILE_SCRIPT) as handle: + source = handle.read() + + block = re.search( + r"if opts\.srate_resample_time_marginalization and .*?lnLt_norm = " + r"scipy\.special\.logsumexp\(lnLt,axis=-1\)", + source, + re.S, + ) + assert block, "could not locate the upsampling block" + text = block.group(0) + + # The output step must be exactly 1/requested, so the requested rate is + # recovered exactly rather than snapped to a multiple of the grid. + assert "dt_target = 1.0/opts.srate_resample_time_marginalization" in text + assert "dt_target * np.arange(n_dense)" in text + # ...and the old hardcoded doubling and the integer-factor upsample are gone. + assert "np.arange(2*len(tvals))" not in text + assert "lnLt.shape[1]*2" not in text + assert "n_upsample" not in text + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/MonteCarloMarginalizeCode/Code/test/waveforms/README.md b/MonteCarloMarginalizeCode/Code/test/waveforms/README.md new file mode 100644 index 000000000..7bbb9edbc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/waveforms/README.md @@ -0,0 +1,105 @@ +# Waveform-level symmetry tests + +Checks on the `U` and `V` mode-cross-term matrices built by +`RIFT.likelihood.factored_likelihood` (`ComputeModeCrossTermIP`): + +``` +U[(A,B)] = < h_A | h_B > (crossTerms) +V[(A,B)] = < h_A^* | h_B > (crossTermsV) +``` + +with `A=(l,m)`, `B=(l',m')`. + +## `test_uv_symmetry.py` + +Verifies, over semi-random parameters looped across an active waveform list: + +1. **`U` Hermitian** — `U[(A,B)] = conj(U[(B,A)])` (definitional) +2. **`U` diagonal real & positive** — `U[(A,A)] > 0`, real (definitional) +3. **`V` complex-symmetric** — `V[(A,B)] = V[(B,A)]` (definitional) +4. **Reflection / parity (aligned-spin only)** — + `V[((l,m),B)] = (-1)^l U[((l,-m),B)]`, from + `h_{l,-m} = (-1)^l conj(h_{l,m})`. + +Every matrix element is recomputed independently (`same_waveform_Q=False`), so +the code's internal symmetrization shortcut is bypassed and the checks are real. + +Checks 1–3 are exercised on both aligned-spin (`ACTIVE_WAVEFORMS`) and +precessing (`PRECESSING_WAVEFORMS`) models; check 4 only on aligned-spin models, +since precessing approximants do not obey the simple reflection relation even at +zero in-plane spin. Models the local `lalsuite` build cannot generate are +skipped with a reason, not failed. + +```bash +pytest -v test_uv_symmetry.py +python test_uv_symmetry.py --approximant IMRPhenomXHM --Lmax 3 --seed 42 +python test_uv_symmetry.py --list +``` + +### Placeholder (expected-fail) check + +`test_full_nonlinear_reflection_symmetry_left_as_exercise` is a placeholder for +the full non-linear reflection algebra that is not yet implemented. It is marked +`@pytest.mark.xfail(strict=True)` so it stays visible (reported `XFAIL`) without +reddening the suite; if the algebra is ever implemented and it starts passing, +strict xfail turns the `XPASS` into a failure so the placeholder gets removed. +Deselect with `-k 'not left_as_exercise'`, or run the script with +`--skip-ludicrous`. + +## Parity (orbital-plane reflection) diagnostics — NOT collected by pytest + +Two script-only diagnostics implement the physics the placeholder above points +at: the exact parity identity of GR, + +``` +h_lm[(s_x,s_y,s_z) -> (-s_x,-s_y,s_z)](t) = (-1)^l conj( h_{l,-m}(t) ), +``` + +with no time- or phase-shift freedom. They are deliberately named so pytest +does NOT collect them: run against currently released precessing models they +FAIL, correctly (NRSur7dq4 at the percent level generically and tens of +percent in superkick subdominant-mode amplitudes; SEOBNRv5PHM with +antisymmetric modes at few x 1e-4; SEOBNRv4PHM at 0.3-1% with a frame origin). +Their role is rapid assessment of upstream "developer leakage" when adopting a +model version or interface — not CI gating. + +- `parity_check_hlm.py [models...]` — mode-level check of the identity above + over superkick-like / generic-precessing / nonprecessing configurations. + Use *perturbed* superkicks: the exact degenerate point is a convention + branch point for several models. +- `uv_parity_diagnostics.py [models...]` — the same physics on the U/V + cross-term matrices ILE builds: + D1 (any waveform; failure = code bug): U = U^dagger, V = V^T; + D2 (single generation, nonprecessing points, including nonprecessing limits + of precessing models): V_{(l,m),B} = (-1)^l U_{(l,-m),B}; + D3 (two generations, any configuration): reflected-pair relations + U'_{(lm),(l'm')} = (-1)^{l+l'} U_{(l',-m'),(l,-m)}, + V'_{(lm),(l'm')} = (-1)^{l+l'} conj(V_{(l,-m),(l',-m')}). + Clean models: <= 1e-10 relative Frobenius residual; violating models: + >= 1e-4. Suggested tolerance: 1e-8. (This extends check (4) of + `test_uv_symmetry.py` to precessing models, where the known violations + live; note that check (4)'s 3e-2 tolerance would pass NRSur7dq4's 4e-4 + aligned-spin violation.) + +NRSur7dq4 needs `LAL_DATA_PATH` pointing at a directory containing +`NRSur7dq4_v1.0.h5`. A marginal-likelihood impact demonstration (at what SNR +a failed check biases PE) lives in the RIFT_roboto_paper repository under +`demos/waveform_symmetry/`. + +### X-family (ChooseFDModes) models: two caveats + +1. **Frame convention**: raw `ChooseFDModes` output for IMRPhenomXPHM / + IMRPhenomXPNR satisfies the parity identity only after a global rotation + by exactly pi about z (their mode frame rotates under reflection of the + in-plane spins). Degenerate with phi_ref — cancels in marginalized PE — + but naive complex mode-level residuals report O(1) "violations". The + amplitude residual is the convention-robust column. After removing this + rotation: XPNR raw modes are parity-clean to 1e-5 (amplitudes 1e-14); + XPHM likewise, apart from a small genuine 2e-3 (2,+-1) equatorial + asymmetry in its aligned-spin limit. +2. **RIFT interface artifact**: raw IMRPhenomXHM satisfies the equatorial + identity exactly (residual 0.0), but through `RIFT.lalsimutils.hlmoft`'s + ChooseFDModes->TD conditioning acquires ~1% (2,+-2) spurious amplitude + asymmetry. This affects every ChooseFDModes-consumed model and dominates + the through-RIFT parity residuals for the X family (D2 for XHM sits at + ~1e-2 instead of <=1e-10 until the conditioning is fixed). diff --git a/MonteCarloMarginalizeCode/Code/test/waveforms/parity_check_hlm.py b/MonteCarloMarginalizeCode/Code/test/waveforms/parity_check_hlm.py new file mode 100644 index 000000000..605b4f8b6 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/waveforms/parity_check_hlm.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python +""" +parity_check_hlm.py : waveform-level parity (orbital-plane reflection) check. + +Exact GR requirement: reflecting the binary through the orbital plane at fref +maps (s_ix, s_iy, s_iz) -> (-s_ix, -s_iy, s_iz) for both spins (spins are +pseudovectors; L is preserved; positions/orbital phase unchanged), and the +spin-weighted spherical harmonic modes must satisfy + + h_lm[reflected params](t) = (-1)^l conj( h_{l,-m}[original params](t) ) + +with NO time or phase shift freedom (same fmin/fref/phiref). Nonprecessing +configurations are fixed points of the reflection, recovering the usual +equatorial identity h_{l,-m} = (-1)^l conj(h_lm). + +We test this mode-by-mode for superkick-like and generic precessing configs. +Metrics per (l,m): + amp_resid : relative L2 difference of |h^B_lm| vs |h^A_{l,-m}| (phase-convention robust) + cplx_resid: relative L2 difference of h^B_lm vs (-1)^l conj(h^A_{l,-m}) + asym : physical (+m,-m) asymmetry content of config A itself, + ||h^A_lm - (-1)^l conj(h^A_{l,-m})|| / ||h^A_lm|| (scale reference) +""" +import os, sys, json +import numpy as np +import lal +import lalsimulation as lalsim +import RIFT.lalsimutils as lalsimutils + +DELTA_T = 1./4096 +LMAX = 4 + +def make_P(m1_msun, m2_msun, s1, s2, fmin=20., fref=20., approx_str=None): + P = lalsimutils.ChooseWaveformParams() + P.m1 = m1_msun*lal.MSUN_SI + P.m2 = m2_msun*lal.MSUN_SI + P.s1x, P.s1y, P.s1z = s1 + P.s2x, P.s2y, P.s2z = s2 + P.fmin = fmin; P.fref = fref + P.deltaT = DELTA_T + P.dist = 400.*1e6*lal.PC_SI + P.phiref = 0.0; P.incl = 0.0; P.psi = 0.0 + P.tref = 0.0 + if approx_str is not None: + P.approx = lalsim.GetApproximantFromString(approx_str) + return P + +def reflected(P): + Pr = P.copy() + Pr.s1x, Pr.s1y = -P.s1x, -P.s1y + Pr.s2x, Pr.s2y = -P.s2x, -P.s2y + return Pr + +def get_modes(approx_str, P, Lmax=LMAX): + """Return dict {(l,m): (epoch_float, complex ndarray)}""" + P = P.copy() + if approx_str.startswith("SEOBNRv5"): + import RIFT.physics.GWSignal as rgws + hlmT = rgws.hlmoft(P, Lmax=Lmax, approx_string=approx_str) + else: + P.approx = lalsim.GetApproximantFromString(approx_str) + if approx_str.startswith("IMRPhenomX"): + P.deltaF = 1./16 # ChooseFDModes path needs an explicit segment length + hlmT = lalsimutils.hlmoft(P, Lmax=Lmax) + if not isinstance(hlmT, dict): + hlmT = lalsimutils.SphHarmTimeSeries_to_dict(hlmT, Lmax) + out = {} + for k, v in hlmT.items(): + out[k] = (float(v.epoch), np.array(v.data.data, dtype=complex)) + return out + +def l2(x): + return np.sqrt(np.sum(np.abs(x)**2)) + +def aligned_pair(eA, hA, eB, hB, deltaT=DELTA_T): + """Trim two series (epochs eA,eB) to their common time support, nearest-sample.""" + off = (eB - eA)/deltaT + n = int(round(off)) + if abs(off - n) > 1e-3: + # non-integer offset: interpolate B onto A's grid + tA = eA + deltaT*np.arange(len(hA)) + tB = eB + deltaT*np.arange(len(hB)) + re = np.interp(tA, tB, hB.real, left=0, right=0) + im = np.interp(tA, tB, hB.imag, left=0, right=0) + return hA, re + 1j*im, ("interp", off) + # integer offset: shift + if n >= 0: + a = hA[n:]; b = hB + else: + a = hA; b = hB[-n:] + m = min(len(a), len(b)) + return a[:m], b[:m], ("shift", n) + +def compare(modesA, modesB, label=""): + rows = [] + for (l, m) in sorted(modesA.keys()): + if (l, -m) not in modesA or (l, m) not in modesB: + continue + eB, hB = modesB[(l, m)] + eA, hA = modesA[(l, -m)] + target = (-1)**l * np.conj(hA) # prediction for h^B_lm from config A + hB_al, tgt_al, how = aligned_pair(eB, hB, eA, target) + nrm = max(l2(hB_al), l2(tgt_al)) + if nrm == 0: + continue + cplx_resid = l2(hB_al - tgt_al)/nrm + amp_resid = l2(np.abs(hB_al) - np.abs(tgt_al))/nrm + # physical asymmetry content of config A + eA2, hA2 = modesA[(l, m)] + a1, a2, _ = aligned_pair(eA2, hA2, eA, target) + nrma = max(l2(a1), l2(a2)) + asym = l2(a1 - a2)/nrma if nrma > 0 else 0. + rows.append(dict(l=l, m=m, cplx_resid=float(cplx_resid), + amp_resid=float(amp_resid), asym=float(asym))) + return rows + +CONFIGS = { + # superkick: q=1, antiparallel in-plane spins, generic azimuth + "superkick": dict(m1=40., m2=40., + s1=(0.8*np.cos(0.4), 0.8*np.sin(0.4), 0.), + s2=(-0.8*np.cos(0.4), -0.8*np.sin(0.4), 0.)), + # hangup-kick-like: in-plane antiparallel + aligned component + "superkick_tilted": dict(m1=40., m2=40., + s1=(0.6*np.cos(0.4), 0.6*np.sin(0.4), 0.5), + s2=(-0.6*np.cos(0.4), -0.6*np.sin(0.4), 0.5)), + # superkick broken slightly: unequal masses + azimuth offset -> not a fixed point + # of exchange symmetry, and total in-plane spin no longer exactly zero + "superkick_perturbed": dict(m1=40.8, m2=39.2, + s1=(0.8*np.cos(0.4), 0.8*np.sin(0.4), 0.), + s2=(-0.75*np.cos(0.45), -0.75*np.sin(0.45), 0.)), + # generic precessing, unequal mass + "generic_prec": dict(m1=48., m2=32., s1=(0.5, 0.2, 0.3), s2=(-0.1, 0.4, -0.2)), + # nonprecessing control: reflection is the identity + "nonprec_control": dict(m1=44., m2=36., s1=(0., 0., 0.5), s2=(0., 0., -0.3)), +} + +MODELS = sys.argv[1:] if len(sys.argv) > 1 else \ + ["IMRPhenomTPHM", "SEOBNRv4PHM", "NRSur7dq4", "SEOBNRv5PHM", + "IMRPhenomXPHM", "IMRPhenomXPNR"] + +only = os.environ.get("PARITY_CONFIGS") +if only: + CONFIGS = {k: v for k, v in CONFIGS.items() if k in only.split(",")} + +results = {} +for model in MODELS: + results[model] = {} + for cname, c in CONFIGS.items(): + try: + PA = make_P(c["m1"], c["m2"], c["s1"], c["s2"]) + PB = reflected(PA) + mA = get_modes(model, PA) + mB = get_modes(model, PB) + rows = compare(mA, mB) + results[model][cname] = rows + worst = max(rows, key=lambda r: r["cplx_resid"]) + print(f"[{model:14s}] {cname:18s} worst mode ({worst['l']},{worst['m']:+d}): " + f"cplx={worst['cplx_resid']:.3e} amp={worst['amp_resid']:.3e} " + f"(physical asym scale {worst['asym']:.3e})", flush=True) + for r in rows: + print(f" ({r['l']},{r['m']:+d}) cplx={r['cplx_resid']:.3e} " + f"amp={r['amp_resid']:.3e} asym={r['asym']:.3e}", flush=True) + except Exception as e: + results[model][cname] = f"ERROR: {e}" + print(f"[{model:14s}] {cname:18s} ERROR: {e}", flush=True) + +out = os.path.join(os.path.dirname(os.path.abspath(__file__)), "parity_check_results.json") +with open(out, "w") as f: + json.dump(results, f, indent=1) +print("wrote", out) diff --git a/MonteCarloMarginalizeCode/Code/test/waveforms/test_uv_symmetry.py b/MonteCarloMarginalizeCode/Code/test/waveforms/test_uv_symmetry.py new file mode 100644 index 000000000..44b529026 --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/waveforms/test_uv_symmetry.py @@ -0,0 +1,458 @@ +#! /usr/bin/env python +# +# test_uv_symmetry.py +# +# Waveform-level symmetry checks on the U and V mode-cross-term matrices +# built by RIFT.likelihood.factored_likelihood. +# +# WHAT IS BEING TESTED +# The factored likelihood pre-computes two dictionaries of PSD-weighted mode +# inner products (see factored_likelihood.ComputeModeCrossTermIP): +# +# U[(A,B)] = < h_A | h_B > (crossTerms) +# V[(A,B)] = < h_A^* | h_B > (crossTermsV) +# +# where A=(l,m), B=(l',m'), h_A^* is the time-domain complex conjugate mode, +# and < a | b > = 2 \int a^*(f) b(f) / S_n(f) df (RIFT.lalsimutils.ComplexIP). +# +# Three of the properties tested below follow purely from the DEFINITION of +# the inner product and hold for every waveform (a genuine numerical check, +# because we recompute every matrix element independently rather than relying +# on the code's own symmetrization shortcut): +# +# (1) U is Hermitian: U[(A,B)] = conj(U[(B,A)]) +# (2) U has real, positive diag: U[(A,A)] real and > 0 +# (3) V is complex-symmetric: V[(A,B)] = V[(B,A)] +# +# The fourth property is PHYSICS, and only holds for non-precessing +# (aligned-spin) binaries, which obey the reflection / parity relation +# +# h_{l,-m}(t) = (-1)^l conj(h_{l,m}(t)). +# +# Because V is built from the conjugated modes, this implies the cross-matrix +# identity +# +# (4) V[((l,m),B)] = (-1)^l U[((l,-m),B)] (aligned-spin binaries only). +# +# We exercise all four over semi-random parameters, looped over an active list +# of waveform approximants. +# +# EXAMPLES +# pytest -v test_uv_symmetry.py +# python test_uv_symmetry.py --approximant IMRPhenomXHM --Lmax 3 --seed 42 +# python test_uv_symmetry.py --list # show the active waveform list +# +# NOTE +# This module also carries one placeholder check +# (test_full_nonlinear_reflection_symmetry_left_as_exercise), marked +# xfail(strict) so it stays visible without reddening the suite; see its +# docstring. Deselect it with `-k 'not left_as_exercise'` or run the script +# with `--skip-ludicrous`. + +from __future__ import print_function + +import argparse +import itertools +import sys + +import numpy as np + +import lal +import lalsimulation as lalsim + +import RIFT +import RIFT.lalsimutils as lalsimutils +import RIFT.likelihood.factored_likelihood as factored_likelihood + +try: + import pytest + _HAVE_PYTEST = True +except ImportError: # allow running as a bare script without pytest installed + _HAVE_PYTEST = False + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + +# The "active waveform list": aligned-spin, multi-mode-capable approximants, +# which obey the reflection relation (4). Anything that fails to generate in the +# local lalsuite build is skipped (with a reason) rather than treated as a +# symmetry failure. +ACTIVE_WAVEFORMS = [ + "IMRPhenomXHM", + "IMRPhenomHM", + "SEOBNRv4HM", # skipped on builds whose SimIMRSpinAlignedEOBModes signature differs + "SEOBNRv5HM", # skipped on builds that reject the current call +] + +# Precessing models. Their inertial-frame modes do NOT obey the simple +# reflection relation (4) even at zero in-plane spin (frame/phase conventions), +# so they are used only for the definitional checks (1)-(3). +PRECESSING_WAVEFORMS = [ + "IMRPhenomXPHM", +] + +# Analytic PSD used to weight the inner products. Any physical (finite, > 0 in +# band) PSD works: the symmetry relations are independent of its choice. +PSD_FUNC = lalsim.SimNoisePSDaLIGOZeroDetHighPower + +# Fixed grid, chosen so every mode series shares (deltaT, deltaF) and stays well +# away from ISCO/wraparound for a heavy-ish system. +FMIN = 20.0 +FMAX = 1700.0 +DELTA_T = 1.0 / 4096.0 +DELTA_F = 1.0 / 8.0 +MTOT_MSUN = 60.0 +LMAX_DEFAULT = 3 + +# Tolerances, expressed relative to the geometric mean of the two diagonal +# norms so they are dimensionless. +TOL_DEFINITIONAL = 1e-6 # (1)-(3): exact up to floating-point round-off +TOL_REFLECTION = 3e-2 # (4): physics + finite-length / tapering noise + + +# --------------------------------------------------------------------------- +# Core: build the U and V matrices at the waveform level +# --------------------------------------------------------------------------- + +class _WaveformUnavailable(Exception): + """Raised when a model cannot be resolved or generated in this build. + + Only this exception is treated as a legitimate reason to *skip*; any other + exception (in the inner products or symmetry algebra) is an analysis + regression and is allowed to propagate as a real failure. + """ + + +def _make_params(approximant, seed, aligned=True, mtot=MTOT_MSUN): + """Semi-random ChooseWaveformParams on a fixed, well-behaved grid. + + `seed` makes each trial reproducible; `aligned` zeroes the in-plane spins so + the reflection relation (4) applies. + """ + rng = np.random.RandomState(seed) + + P = lalsimutils.ChooseWaveformParams() + P.ampO = -1 # keep all higher modes the model offers + P.phaseO = 7 + P.taper = lalsimutils.lsu_TAPER_START + P.deltaT = DELTA_T + P.deltaF = DELTA_F + P.fmin = FMIN + P.fref = 20.0 + + # Semi-random intrinsic parameters (reproducible via seed). + q = rng.uniform(0.5, 1.0) # m2/m1 + m1 = mtot / (1.0 + q) + m2 = mtot - m1 + P.m1 = m1 * lal.MSUN_SI + P.m2 = m2 * lal.MSUN_SI + + s1z = rng.uniform(-0.6, 0.6) + s2z = rng.uniform(-0.6, 0.6) + P.s1x = P.s1y = P.s2x = P.s2y = 0.0 + P.s1z = s1z + P.s2z = s2z + if not aligned: + # Only used for the definitional checks (1)-(3), which do not require + # reflection symmetry. + P.s1x = rng.uniform(-0.4, 0.4) + P.s1y = rng.uniform(-0.4, 0.4) + P.s2x = rng.uniform(-0.4, 0.4) + + # Extrinsic angles are irrelevant to U/V (they act on the Ylm sum, not the + # mode inner products), but set them to something non-trivial anyway. + P.incl = rng.uniform(0.0, np.pi) + P.phiref = rng.uniform(0.0, 2 * np.pi) + P.psi = rng.uniform(0.0, np.pi) + P.dist = factored_likelihood.distMpcRef * 1e6 * lal.PC_SI + + try: + P.approx = lalsim.GetApproximantFromString(approximant) + except Exception as e: # model name not known to this lalsuite build + raise _WaveformUnavailable("approximant {} unavailable: {}".format(approximant, e)) + return P + + +def _generate_modes(P, Lmax): + """Generate (hlms, hlms_conj), raising _WaveformUnavailable on failure. + + Only *waveform generation* is treated as skippable (a model may not be + compiled into the local lalsuite build). Everything downstream -- the inner + products and the symmetry algebra -- is an analysis step that must surface + its errors, not be swallowed into a skip. + """ + try: + return factored_likelihood.internal_hlm_generator( + P, Lmax, verbose=False, quiet=True) + except Exception as e: + raise _WaveformUnavailable("cannot generate {}: {}".format(P.approx, e)) + + +def build_uv(P, Lmax, psd_func=PSD_FUNC, fmin=FMIN, fmax=FMAX, verbose=False): + """Generate the modes for P and return (hlms, U, V). + + U and V are computed with same_waveform_Q=False so that *every* matrix + element is an independent inner product -- the code's internal symmetrized + fast path is intentionally bypassed so the symmetry tests below are real. + + Generation errors raise _WaveformUnavailable (skippable); inner-product / + analysis errors propagate unchanged (so regressions fail, not skip). + """ + hlms, hlms_conj = _generate_modes(P, Lmax) + + fNyq = 1.0 / (2.0 * P.deltaT) + U = factored_likelihood.ComputeModeCrossTermIP( + hlms, hlms, psd_func, fmin, fmax, fNyq, P.deltaF, + analyticPSD_Q=True, verbose=False, prefix="U", same_waveform_Q=False) + V = factored_likelihood.ComputeModeCrossTermIP( + hlms_conj, hlms, psd_func, fmin, fmax, fNyq, P.deltaF, + analyticPSD_Q=True, verbose=False, prefix="V", same_waveform_Q=False) + return hlms, U, V + + +# --------------------------------------------------------------------------- +# Symmetry checks. Each returns a list of human-readable violation strings. +# --------------------------------------------------------------------------- + +def _scale(U, A, B): + """Geometric mean of the diagonal norms, used to non-dimensionalize.""" + dA = abs(U[(A, A)]) + dB = abs(U[(B, B)]) + s = np.sqrt(dA * dB) + return s if s > 0 else 1.0 + + +def check_U_hermitian(U, tol=TOL_DEFINITIONAL): + """(1) U[(A,B)] = conj(U[(B,A)]).""" + viol = [] + modes = sorted({A for (A, _) in U.keys()}) + for A, B in itertools.combinations(modes, 2): + lhs = U[(A, B)] + rhs = np.conj(U[(B, A)]) + rel = abs(lhs - rhs) / _scale(U, A, B) + if rel > tol: + viol.append("U not Hermitian for {},{}: |dU|/scale={:.3e}".format(A, B, rel)) + return viol + + +def check_U_diagonal_real_positive(U, tol=TOL_DEFINITIONAL): + """(2) U[(A,A)] is real and positive.""" + viol = [] + modes = sorted({A for (A, _) in U.keys()}) + for A in modes: + d = U[(A, A)] + if abs(d) == 0: + # A zero diagonal is a positivity failure: < h_A | h_A > must be > 0 + # for any mode with power. Flag it rather than skipping. + viol.append("U[{0},{0}] is exactly zero (mode has no power)".format(A)) + continue + imag_frac = abs(np.imag(d)) / abs(d) + if imag_frac > tol: + viol.append("U[{0},{0}] not real: Im/|.|={1:.3e}".format(A, imag_frac)) + if np.real(d) <= 0: + viol.append("U[{0},{0}] not positive: Re={1:.3e}".format(A, np.real(d))) + return viol + + +def check_V_symmetric(V, U, tol=TOL_DEFINITIONAL): + """(3) V[(A,B)] = V[(B,A)].""" + viol = [] + modes = sorted({A for (A, _) in V.keys()}) + for A, B in itertools.combinations(modes, 2): + rel = abs(V[(A, B)] - V[(B, A)]) / _scale(U, A, B) + if rel > tol: + viol.append("V not symmetric for {},{}: |dV|/scale={:.3e}".format(A, B, rel)) + return viol + + +def check_reflection_aligned(U, V, tol=TOL_REFLECTION): + """(4) V[((l,m),B)] = (-1)^l U[((l,-m),B)] (aligned-spin binaries). + + Only pairs for which the reflected mode (l,-m) is present are tested. + """ + viol = [] + n_tested = 0 + modes = sorted({A for (A, _) in U.keys()}) + mode_set = set(modes) + for A in modes: + (l, m) = A + A_refl = (l, -m) + if A_refl not in mode_set: + continue + for B in modes: + n_tested += 1 + lhs = V[(A, B)] + rhs = ((-1) ** l) * U[(A_refl, B)] + rel = abs(lhs - rhs) / _scale(U, A, B) + if rel > tol: + viol.append( + "reflection broken for A={},B={}: " + "|V - (-1)^l U_refl|/scale={:.3e}".format(A, B, rel)) + if n_tested == 0: + viol.append("reflection check exercised no mode pairs (no +/-m partners found)") + return viol + + +def run_all_checks(P, Lmax, aligned, verbose=False): + """Build U,V for P and return the concatenated violation list.""" + hlms, U, V = build_uv(P, Lmax, verbose=verbose) + if verbose: + print(" modes:", sorted(hlms.keys())) + viol = [] + viol += check_U_hermitian(U) + viol += check_U_diagonal_real_positive(U) + viol += check_V_symmetric(V, U) + if aligned: + viol += check_reflection_aligned(U, V) + return viol + + +# --------------------------------------------------------------------------- +# Waveform generation guard: skip (don't fail) ONLY if a model is unavailable +# --------------------------------------------------------------------------- + +def _try_build(approximant, seed, Lmax, aligned=True): + """Return (P, violations). + + Raises _WaveformUnavailable if the model cannot be resolved/generated in + this build (caller may skip). Analysis errors are NOT caught here -- they + propagate so a regression fails loudly instead of masquerading as a skip. + """ + P = _make_params(approximant, seed, aligned=aligned) + viol = run_all_checks(P, Lmax, aligned=aligned) + return P, viol + + +# --------------------------------------------------------------------------- +# pytest entry points +# --------------------------------------------------------------------------- + +if _HAVE_PYTEST: + + @pytest.mark.parametrize("approximant", ACTIVE_WAVEFORMS) + def test_uv_symmetry_aligned(approximant): + """Definitional + reflection symmetry for aligned-spin binaries.""" + try: + _P, viol = _try_build(approximant, seed=1234, Lmax=LMAX_DEFAULT, aligned=True) + except _WaveformUnavailable as e: + pytest.skip(str(e)) + assert not viol, "symmetry violations for {}:\n {}".format( + approximant, "\n ".join(viol)) + + @pytest.mark.parametrize("approximant", PRECESSING_WAVEFORMS) + def test_uv_definitional_precessing(approximant): + """Definitional checks (1)-(3) must also hold for precessing systems.""" + try: + P = _make_params(approximant, seed=99, aligned=False) + _hlms, U, V = build_uv(P, LMAX_DEFAULT) + except _WaveformUnavailable as e: + pytest.skip(str(e)) + viol = (check_U_hermitian(U) + + check_U_diagonal_real_positive(U) + + check_V_symmetric(V, U)) + assert not viol, "definitional violations for {}:\n {}".format( + approximant, "\n ".join(viol)) + + @pytest.mark.xfail(strict=True, + reason="full non-linear reflection algebra not yet " + "implemented; placeholder left as an exercise") + def test_full_nonlinear_reflection_symmetry_left_as_exercise(): + """LUDICROUS / INTENTIONAL FAILURE (marked xfail). + + A complete waveform-symmetry test would also verify the higher-order, + fully non-linear reflection identities relating U and V across *all* + (l, m) sectors simultaneously (the closed algebra of parity, time- + reversal and mode-mixing operators), not just the pairwise relation (4). + + That verification is not implemented here. It is marked xfail(strict) + so it does not redden the suite, yet stays visible as an unfinished + item: if someone ever implements the algebra and this starts passing, + strict xfail turns the XPASS into a failure, forcing the marker (and + this placeholder) to be removed. + """ + assert False, "this failure is left as a test" + + +# --------------------------------------------------------------------------- +# Script runner (mirrors the style of test/waveform/check_waveform_random.py) +# --------------------------------------------------------------------------- + +def _main(argv=None): + parser = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--approximant", type=str, default=None, + help="single approximant to test (default: loop over the active list)") + parser.add_argument("--Lmax", type=int, default=LMAX_DEFAULT) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--n-trials", type=int, default=1, + help="number of semi-random trials per approximant") + parser.add_argument("--precessing", action="store_true", + help="use precessing spins and, unless --approximant is " + "given, loop over PRECESSING_WAVEFORMS " + "(disables reflection check (4))") + parser.add_argument("--list", action="store_true", + help="print the active waveform lists and exit") + parser.add_argument("--skip-ludicrous", action="store_true", + help="do not run the placeholder (expected-fail) check") + parser.add_argument("--verbose", action="store_true") + opts = parser.parse_args(argv) + + aligned = not opts.precessing + + if opts.list: + print("Active waveform list (aligned, reflection-symmetric):") + for a in ACTIVE_WAVEFORMS: + print(" ", a) + print("Precessing list (definitional checks only):") + for a in PRECESSING_WAVEFORMS: + print(" ", a) + return 0 + + # Default loop: aligned -> ACTIVE_WAVEFORMS; --precessing -> PRECESSING_WAVEFORMS. + if opts.approximant: + approximants = [opts.approximant] + else: + approximants = PRECESSING_WAVEFORMS if opts.precessing else ACTIVE_WAVEFORMS + + n_fail = 0 + n_skip = 0 + for approximant in approximants: + for trial in range(opts.n_trials): + seed = opts.seed + trial + label = "{} (seed={}, {})".format( + approximant, seed, "aligned" if aligned else "precessing") + try: + _P, viol = _try_build(approximant, seed, opts.Lmax, aligned=aligned) + except _WaveformUnavailable as e: + print("SKIP {}: {}".format(label, e)) + n_skip += 1 + continue + except Exception as e: # analysis regression -> real failure, keep going + n_fail += 1 + print("FAIL {}: analysis error: {}".format(label, e)) + continue + if viol: + n_fail += 1 + print("FAIL {}".format(label)) + for v in viol: + print(" - {}".format(v)) + else: + print("PASS {}".format(label)) + + if not opts.skip_ludicrous: + # Expected (xfail-style) failure: a placeholder for the not-yet-implemented + # full non-linear reflection algebra. It does NOT count as a real failure. + print("\n--- placeholder check (expected to fail) ---") + try: + assert False, "this failure is left as a test" + except AssertionError as e: + print("XFAIL full_nonlinear_reflection_symmetry_left_as_exercise: {}".format(e)) + + print("\nSummary: {} failing, {} skipped".format(n_fail, n_skip)) + return 1 if n_fail else 0 + + +if __name__ == "__main__": + sys.exit(_main()) diff --git a/MonteCarloMarginalizeCode/Code/test/waveforms/uv_parity_diagnostics.py b/MonteCarloMarginalizeCode/Code/test/waveforms/uv_parity_diagnostics.py new file mode 100644 index 000000000..c4e24bbdc --- /dev/null +++ b/MonteCarloMarginalizeCode/Code/test/waveforms/uv_parity_diagnostics.py @@ -0,0 +1,117 @@ +#!/usr/bin/env python +""" +uv_parity_diagnostics.py : prototype waveform-QA checks expressed purely at +the level of RIFT's U/V cross-term matrices (the objects ILE already builds +in PrecomputeLikelihoodTerms). Goal: catch waveform parity violations at +precompute time, before they bias PE. + +Definitions (RIFT conventions, ComplexIP over two-sided band, even PSD): + U_{(lm),(l'm')} = < h_lm | h_l'm' > + V_{(lm),(l'm')} = < conj(h_lm) | h_l'm' > + +Diagnostics: + D1 (exact identities; failure = code bug, ANY waveform): + U = U^dagger , V = V^T + D2 (single config; applies whenever the CONFIG is reflection-symmetric, + i.e. nonprecessing -- even if the model is a precessing model): + V_{(l,m),(l',m')} = (-1)^l U_{(l,-m),(l',m')} + Failure = spurious (+m,-m) asymmetry (parity violation) in the model. + D3 (reflected pair; applies to ANY config; two waveform generations): + with primes = quantities of the reflected config (s_xy -> -s_xy): + U'_{(lm),(l'm')} = (-1)^{l+l'} U_{(l',-m'),(l,-m)} + V'_{(lm),(l'm')} = (-1)^{l+l'} conj( V_{(l,-m),(l',-m')} ) + Failure = parity violation, full precessing case. + +Each check reports a relative Frobenius-norm residual. +""" +import numpy as np +import lal +import lalsimulation as lalsim +import RIFT.lalsimutils as lsu +import RIFT.likelihood.factored_likelihood as fl + +LMAX = 4 + +def make_P(m1, m2, s1, s2, approx_str): + P = lsu.ChooseWaveformParams() + P.m1 = m1*lal.MSUN_SI; P.m2 = m2*lal.MSUN_SI + P.s1x, P.s1y, P.s1z = s1 + P.s2x, P.s2y, P.s2z = s2 + P.fmin = 20.; P.fref = 20. + P.deltaT = 1./4096 + P.deltaF = 1./16 + P.dist = 400.*1e6*lal.PC_SI + P.phiref = 0.; P.incl = 0.; P.psi = 0.; P.tref = 0. + P.approx = lalsim.GetApproximantFromString(approx_str) + return P + +def uv_matrices(P): + hlmF = lsu.hlmoff(P.copy(), LMAX) + hlmF_conj = lsu.conj_hlmoff(P.copy(), LMAX) + if not isinstance(hlmF, dict): + hlmF = lsu.SphHarmFrequencySeries_to_dict(hlmF, LMAX) + hlmF_conj = lsu.SphHarmFrequencySeries_to_dict(hlmF_conj, LMAX) + fNyq = 0.5/P.deltaT + U = fl.ComputeModeCrossTermIP(hlmF, hlmF, lal.LIGOIPsd, P.fmin, fNyq, fNyq, + P.deltaF, analyticPSD_Q=True, verbose=False) + V = fl.ComputeModeCrossTermIP(hlmF_conj, hlmF, lal.LIGOIPsd, P.fmin, fNyq, fNyq, + P.deltaF, analyticPSD_Q=True, verbose=False, prefix="V") + return U, V + +def frob(d, keys): + return np.sqrt(sum(abs(d[k])**2 for k in keys)) + +def rel(dA, dB, keys): + return np.sqrt(sum(abs(dA[k]-dB[k])**2 for k in keys))/max(frob(dA, keys), 1e-300) + +def check_D1(U, V): + keys = list(U.keys()) + Udag = {(p1, p2): np.conj(U[(p2, p1)]) for (p1, p2) in keys} + Vt = {(p1, p2): V[(p2, p1)] for (p1, p2) in keys} + return rel(U, Udag, keys), rel(V, Vt, keys) + +def check_D2(U, V): + keys = list(U.keys()) + Vpred = {} + for (p1, p2) in keys: + Vpred[(p1, p2)] = (-1)**p1[0]*U[((p1[0], -p1[1]), p2)] + return rel(V, Vpred, keys) + +def check_D3(U, V, Up, Vp): + keys = list(U.keys()) + Upred, Vpred = {}, {} + for (p1, p2) in keys: + f1 = (p1[0], -p1[1]); f2 = (p2[0], -p2[1]) + s = (-1)**(p1[0]+p2[0]) + Upred[(p1, p2)] = s*U[(f2, f1)] + Vpred[(p1, p2)] = s*np.conj(V[(f1, f2)]) + return rel(Up, Upred, keys), rel(Vp, Vpred, keys) + +CONFIGS = { + "nonprec": dict(m1=44., m2=36., s1=(0., 0., 0.5), s2=(0., 0., -0.3)), + "superkick_perturbed": dict(m1=40.8, m2=39.2, + s1=(0.8*np.cos(0.4), 0.8*np.sin(0.4), 0.), + s2=(-0.75*np.cos(0.45), -0.75*np.sin(0.45), 0.)), + "generic_prec": dict(m1=48., m2=32., s1=(0.5, 0.2, 0.3), s2=(-0.1, 0.4, -0.2)), +} + +import sys +MODELS = sys.argv[1:] if len(sys.argv) > 1 else \ + ["IMRPhenomTPHM", "NRSur7dq4", "IMRPhenomXPHM", "IMRPhenomXPNR"] + +for model in MODELS: + for cname, c in CONFIGS.items(): + try: + P = make_P(c["m1"], c["m2"], c["s1"], c["s2"], model) + U, V = uv_matrices(P) + e1u, e1v = check_D1(U, V) + e2 = check_D2(U, V) + Pr = P.copy() + Pr.s1x, Pr.s1y = -P.s1x, -P.s1y + Pr.s2x, Pr.s2y = -P.s2x, -P.s2y + Up, Vp = uv_matrices(Pr) + e3u, e3v = check_D3(U, V, Up, Vp) + print(f"[{model:14s}] {cname:20s} D1(U)={e1u:.1e} D1(V)={e1v:.1e} | " + f"D2={e2:.3e} | D3(U)={e3u:.3e} D3(V)={e3v:.3e}", flush=True) + except Exception as e: + print(f"[{model:14s}] {cname:20s} ERROR: {e}", flush=True) diff --git a/containers/PRECOMPILE_ASSESSMENT.md b/containers/PRECOMPILE_ASSESSMENT.md new file mode 100644 index 000000000..e41a955b5 --- /dev/null +++ b/containers/PRECOMPILE_ASSESSMENT.md @@ -0,0 +1,194 @@ +# RIFT Container Pre-Compilation Assessment + +## Summary + +Pre-building parts of the RIFT container is feasible, but there are two distinct +classes of startup cost: + +1. **Install/build cost**: dependency resolution, source checkout, editable + install, Python bytecode generation, and optional packages. This is fully + image-build-time work and should be moved out of job startup wherever possible. +2. **Runtime compiler cost**: CuPy `RawKernel`/`ElementwiseKernel` compilation and + JAX/XLA compilation. This can be reduced with persistent caches and warmup + scripts, but cache hits depend on CUDA version, GPU architecture, driver/JAX + versions, backend, function shape, dtype, and selected likelihood mode. + +The most practical near-term path is to add a survey-driven container warmup +phase that materializes persistent CuPy and JAX caches under a known in-image or +job-local cache path, plus runtime environment defaults that cap JAX/XLA thread +creation. For production on heterogeneous OSG/LDG GPUs, this should be paired +with the existing container-family mechanism so each image targets a bounded +CUDA/GPU capability range. A concrete design for that survey/warmup layer is in +`containers/SURVEY_SCAN_PROPOSAL.md`. + +## Current State + +The current Apptainer template (`containers/rift_container.def.in`) clones RIFT, +runs `pip3 install -e .`, installs the selected `cupy-cuda11x`/`cupy-cuda12x` +wheel, then installs the shared requirements. The top-level `rift_container.def` +does the same pattern inline. This means jobs receive source and dependencies, +but no application-specific GPU/JAX compilation has been warmed. + +Core CuPy compilation sites: + +- `RIFT.likelihood.Q_inner_product.Q_inner_product_cupy`: reads + `cuda_Q_inner_product.cu` and constructs `cupy.RawKernel`. +- `RIFT.likelihood.Q_fused_calmarg`: lazily constructs `RawKernel` objects for + `cuda_Q_fused_calmarg.cu` and `cuda_Q_fused_calmarg_distmarg.cu`. +- `RIFT.interpolators.interp_gpu`: memoizes a CuPy `ElementwiseKernel`. + +Core JAX compilation sites: + +- `RIFT.likelihood.jax_ile.wrapper` constructs jitted likelihood closures per + likelihood object. +- `RIFT.likelihood.jax_ile.samplers._warmup_compile` already forces a small + first-call compile and reports the latency, because some modes can spend + tens of seconds in XLA compile. +- Additional `jax.jit`, `jax.grad`, and `jax.hessian` calls appear in polishing, + Fisher, NUTS, and flowMC paths. + +## Feasibility + +### CuPy + +CuPy cache warming is feasible and likely worth doing. The current kernels are +small, have stable source, and are concentrated in a few modules. A warmup script +can import CuPy, allocate representative arrays, call each kernel once, and +leave artifacts in `CUPY_CACHE_DIR`. + +Important constraints: + +- CuPy compiles for the active CUDA toolkit/driver/GPU architecture. A cache + baked on one architecture may not serve another. +- Building a container usually has no GPU unless the builder is a GPU node and + Apptainer is run with NVIDIA support. Without a GPU, build-time warmup cannot + compile device-specific SASS. +- If the image must run across old and modern GPUs, the existing container-family + split is the right place to bound compatibility and avoid one cache trying to + serve every target. + +Recommended first implementation: + +- Add `RIFT/likelihood/warmup_gpu_kernels.py` or a `bin/rift_warmup_gpu_kernels` + script. +- Set `CUPY_CACHE_DIR` to a stable path, for example + `/opt/rift-cache/cupy` in the image or `${_CONDOR_SCRATCH_DIR}/.cupy/kernel_cache` + at runtime. +- During container build, run the script only when a GPU is available; otherwise + install the script and rely on a one-time prolog/warmup job on each GPU class. +- Stop using `CUPY_CACHE_IN_MEMORY=1` as the only default for production caching; + it avoids disk writes but also prevents reuse across processes. + +### JAX/XLA + +JAX pre-compilation is feasible only as persistent-cache warming, not as a +single universal binary baked into the source install. The likelihood closes +over event data and compiles by argument shape and static branch choices: + +- number of detectors; +- number of modes; +- `npts` / time grid length; +- distance grid size; +- phi/psi marginalization grid sizes; +- interpolation and phase-marginalization choices; +- sampler mode and use of value/grad/Hessian. + +This means a generic warmup can populate common shapes, but real events with +different shapes may still compile. The most valuable cache targets are the +standard O4 production settings and the expensive modes already using +`_warmup_compile`. + +Recommended first implementation: + +- Add a JAX warmup command that builds a synthetic `JAXLikelihoodData` matching + standard production shapes and instantiates the production wrappers. +- Enable JAX persistent compilation cache via environment variables before JAX + import. For modern JAX this can be done with + `JAX_COMPILATION_CACHE_DIR=/opt/rift-cache/jax` or the corresponding + `jax.config` calls in the warmup command. +- Warm the common wrapper modes: fixed-distance, distance-marginalized, + phi-marginalized, phi+psi-marginalized, and the value/grad/Hessian paths that + samplers invoke. +- Keep runtime fallback behavior: if a shape misses the cache, it should compile + once and continue. + +Threading should be addressed independently. Container defaults should set a +conservative CPU-thread policy for JAX/XLA jobs, especially in Condor slots: + +```sh +XLA_FLAGS=--xla_cpu_multi_thread_eigen=false +OMP_NUM_THREADS=1 +OPENBLAS_NUM_THREADS=1 +MKL_NUM_THREADS=1 +NUMEXPR_NUM_THREADS=1 +``` + +These should be opt-out or mode-aware if CPU-only JAX performance is important. + +## Required Modifications + +1. **Container recipes** + - Add cache directories such as `/opt/rift-cache/cupy` and + `/opt/rift-cache/jax`. + - Export stable cache/thread defaults in `%environment`. + - Optionally run warmup scripts during `%post` when a GPU is visible. + - Add JAX/numpyro/flowMC dependencies to a separate JAX-enabled image flavor; + they should not silently enter the minimal production image. + +2. **Build family** + - Extend `containers/build_family.sh` matrix with optional feature columns: + CUDA family, target capability band, and JAX-enabled vs non-JAX image. + - Keep a CPU-safe fallback image. + - Publish warmed caches per image, not shared across CUDA major versions. + +3. **Warmup scripts** + - CuPy script: allocate tiny representative arrays and invoke + `Q_inner_product_cupy`, `Q_fused_calmarg_cupy`, + `Q_fused_calmarg_distmarg_cupy`, and `interp_gpu.interp`. + - JAX script: construct synthetic likelihood data with configurable + detectors/modes/time-grid/distance-grid sizes and call each wrapper's + warmup path. + - Scripts should report cache directory, backend, device, CUDA/JAX/CuPy + versions, and elapsed compile time. + +4. **Runtime pipeline** + - Add optional Condor prolog or first-node warmup job per container/GPU class. + - Ensure writable caches are available when the image cache is read-only. + For Apptainer on CVMFS, a job-local cache path may be more reliable than + trying to update an in-image cache. + - Thread through environment for `CUPY_CACHE_DIR`, + `JAX_COMPILATION_CACHE_DIR`, and XLA thread flags. + +5. **Packaging** + - Prefer a normal wheel install for release images instead of `pip install -e .` + when the image should be immutable. Editable source installs are convenient + for development but do not give stronger startup guarantees. + - Keep `.cu` files as installed package data, or move them into package data + explicitly rather than relying only on `data_files`, so warmup scripts and + installed modules resolve the same paths. + +## Risks and Limitations + +- A cache warmed on one GPU architecture may miss on another; use the container + family to limit variation. +- Build environments frequently lack GPUs, so some warming must happen as a + deployment/prolog step rather than in `%post`. +- JAX cache portability is version-sensitive. Pinning JAX/JAXLIB/CUDA/Python is + more important for JAX images than for the current unpinned container canary. +- JAX event-shape variation prevents complete elimination of first-call compile. +- Baked caches increase image size and should be measured against CVMFS/OSDF + transfer cost. + +## Suggested Sequence + +1. Add `survey_scan survey` to record the target cluster's GPU/driver classes. +2. Add CuPy warmup script and persistent `CUPY_CACHE_DIR`; validate on one GPU + per dominant image band. +3. Add a JAX-enabled container-family entry with pinned JAX/JAXLIB/numpyro/flowMC. +4. Add synthetic JAX warmup for the standard O4 shape/mode set. +5. Measure cold vs warm startup for: + - NoLoop CuPy standard path; + - fused calmarg path; + - `integrate_likelihood_extrinsic_jax` distance/phi-marginalized modes. +6. Promote warmup to `%post` only for builders with GPU access; otherwise use a + site prolog or one-time cache seeding job per published image/GPU class. diff --git a/containers/README.md b/containers/README.md index 56eaa9e81..988621dda 100644 --- a/containers/README.md +++ b/containers/README.md @@ -1,7 +1,7 @@ # RIFT containers This directory holds the multi-architecture container build and the "container -family" deployment mechanism. It has two related but independent pieces: +family" deployment mechanism. It has three related pieces: 1. **Multi-target build** — build a *family* of RIFT containers (different base image + cupy/CUDA variant, targeting different GPU compute capabilities) from @@ -9,6 +9,9 @@ family" deployment mechanism. It has two related but independent pieces: 2. **Family deployment** — let `SINGULARITY_RIFT_IMAGE` point at a YAML *manifest* describing that family, so each Condor job picks the right image for the machine it lands on. +3. **Survey + warmup scans** — survey a target Condor GPU pool and emit + representative CuPy/JAX warmup jobs for the image bands that pool actually + uses. The top-level [`rift_container.def`](../rift_container.def) is unchanged and remains the default single-image build. @@ -188,7 +191,44 @@ wrapper to exercise it): that the pilot evaluates the expression-valued --- -## 3. CI dependency-resolution canary +## 3. Survey + Warmup Scans + +`survey_scan.sh` records the target pool's GPU classes, emits one Condor warmup +job per container/profile combination, and collects JSON timing/cache reports. +The submit-side tools use only the Python standard library; the CuPy/JAX imports +happen inside the container on the execute node. + +```console +containers/survey_scan.sh survey \ + --out survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml +containers/survey_scan.sh emit-jobs \ + --survey survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml +cd survey/cit-YYYYMMDD/jobs +./submit_all.sh +containers/survey_scan.sh collect --survey survey/cit-YYYYMMDD +``` + +Profiles: + +- `cupy` warms common NoLoop/fused-calmarg CuPy kernels: + `Q_inner_product_cupy`, `Q_fused_calmarg_cupy`, + `Q_fused_calmarg_distmarg_cupy`, and `interp_gpu.interp`. +- `jax` warms synthetic JAX ILE wrapper shapes. Use this only for JAX-enabled + images, for example `--profiles cupy,jax` on a JAX image manifest. + +Generated job wrappers set `CUPY_CACHE_DIR`, `JAX_COMPILATION_CACHE_DIR`, and +conservative thread defaults before `apptainer exec --nv`. If an image is listed +as `osdf://...`, the wrapper fetches only that selected image with `stashcp` or +`pelican`. + +See [`survey_scan/README.md`](survey_scan/README.md) and +[`SURVEY_SCAN_PROPOSAL.md`](SURVEY_SCAN_PROPOSAL.md) for details. + +--- + +## 4. CI dependency-resolution canary The default container build uses *unpinned* deps, so a fresh upstream release (e.g. `swig>=4.4.0`, see issue #136) can silently break RIFT and we only find out diff --git a/containers/SURVEY_SCAN_PROPOSAL.md b/containers/SURVEY_SCAN_PROPOSAL.md new file mode 100644 index 000000000..ce1fa978e --- /dev/null +++ b/containers/SURVEY_SCAN_PROPOSAL.md @@ -0,0 +1,254 @@ +# `survey_scan` Proposal for RIFT Container Builds + +## Goal + +Add a `survey_scan` tool to the RIFT container build framework that surveys the +actual target cluster, classifies the small number of GPU/driver/container +combinations we care about, and runs representative warmup probes so a published +container family carries, or can seed, the most common CuPy and JAX startup +caches. + +This is intentionally a "cover the common cases" tool. It does not need to find +every kernel, every RIFT executable, or every possible event shape. It should +reduce cold-start cost for the dominant NoLoop/CuPy and JAX ILE modes on the GPU +classes we actually schedule onto. + +## Concrete Starting Point: CIT + +The current CIT build kit on `ldas-grid-alt` already has the right shape: + +- `~/rift_cit_build_container_family/build_cit_family.sh` + builds the `cc60-90` CUDA 11.8 image and `cc90-120` CUDA 12.8 image. +- `~/rift_cit_build_container_family/build_jax_container.sh` + builds a separate JAX GPU image. +- `built_containers/rift_container_family.cit.yaml` + records the deployable family. +- `built_containers/rift_container_select.sh` + selects the image at runtime on OSG-like pools when `MY.SingularityImage` + expressions are not evaluated. +- `validate/gpu_check.py` and `validate_jax/` already prove how to submit + real-GPU validation jobs. + +A live CIT GPU census currently looks like: + +```text +554 GeForce GTX 1050 Ti cc 6.1 4040 MB +164 NVIDIA RTX PRO 4000 Blackwell SFF Edition cc 12.0 24027 MB +3 NVIDIA A30 cc 8.0 24188 MB +1 GeForce GTX 1650 cc 7.5 3912 MB +3190 undefined/undefined/undefined +``` + +So for CIT, two non-JAX images remain the right default bands: + +- `cc60-90`: CUDA 11.8 runtime, `cupy-cuda11x`, fallback / old-GPU image. +- `cc90-120`: CUDA 12.8 devel, `cupy-cuda12x`, Blackwell/Hopper image. The devel + base matters because Blackwell may need NVRTC headers for first-use CuPy JIT. + +The JAX GPU image can initially target `cc90-120`, where the CUDA 12.8 stack is +already validated. + +## What `survey_scan` Should Do + +### 1. Survey + +Run on a build/login host with Condor tools available: + +```sh +containers/survey_scan.sh survey --pool cit --out survey/cit-YYYYMMDD +``` + +Collect: + +- `condor_status` grouped by GPU name, `GPUs_Capability`, memory, driver-ish + attributes when advertised, and slot count. +- A normalized JSON summary with recommended image bands. +- The current build matrix and manifest labels, so the survey can say which + observed classes are covered, uncovered, or only fallback-covered. + +Suggested output: + +```text +survey/cit-YYYYMMDD/ + gpu_inventory.tsv + gpu_inventory.json + recommended_matrix.json + coverage.md +``` + +### 2. Generate Scan Jobs + +Create Condor submit files that run a container-specific probe on one machine per +dominant GPU class: + +```sh +containers/survey_scan.sh emit-jobs \ + --survey survey/cit-YYYYMMDD \ + --manifest built_containers/rift_container_family.cit.yaml \ + --out survey/cit-YYYYMMDD/jobs +``` + +Each job should: + +- constrain to one GPU class or image band; +- run with `apptainer exec --nv`; +- set persistent cache directories; +- run a standard warmup probe; +- archive cache metadata and timing logs. + +### 3. Warm CuPy Common Paths + +The CuPy probe should run inside the selected image and exercise: + +- `RIFT.likelihood.Q_inner_product.Q_inner_product_cupy`; +- `RIFT.likelihood.Q_fused_calmarg.Q_fused_calmarg_cupy`; +- `RIFT.likelihood.Q_fused_calmarg.Q_fused_calmarg_distmarg_cupy`; +- `RIFT.interpolators.interp_gpu.interp`. + +It should use tiny, deterministic arrays with representative dtypes and shapes. +The important thing is to trigger compilation and confirm the cache path is +populated, not to benchmark production throughput. + +Recommended environment: + +```sh +CUPY_CACHE_DIR=/rift_cache/cupy/${image_label}/${gpu_capability} +CUPY_CACHE_IN_MEMORY=0 +CUDA_PATH=/usr/local/cuda +``` + +For read-only published images, the runtime wrapper can instead seed or reuse a +job-local cache: + +```sh +CUPY_CACHE_DIR=${_CONDOR_SCRATCH_DIR}/.rift_cache/cupy +``` + +### 4. Warm JAX Common Modes + +The JAX probe should be separate and run only in JAX-enabled images. It should +build synthetic `JAXLikelihoodData` matching standard production shapes and +compile the expensive wrapper modes: + +- fixed-distance `JAXExtrinsicLikelihood`; +- distance-marginalized `JAXDistanceMarginalizedLikelihood`; +- phi-marginalized and phi+psi-marginalized wrappers when present; +- `value_and_grad` and Hessian/Fisher paths used by samplers; +- the existing sampler `_warmup_compile` path. + +Recommended environment: + +```sh +JAX_COMPILATION_CACHE_DIR=/rift_cache/jax/${image_label}/${gpu_capability} +JAX_ENABLE_X64=1 +XLA_FLAGS=--xla_cpu_multi_thread_eigen=false +OMP_NUM_THREADS=1 +OPENBLAS_NUM_THREADS=1 +MKL_NUM_THREADS=1 +NUMEXPR_NUM_THREADS=1 +``` + +The first version can target one standard O4 shape. Later versions can read a +small profile file, for example: + +```yaml +jax_profiles: + - name: o4_default_lmax4_hl + detectors: [H1, L1] + l_max: 4 + npts: 614 + distance_grid: 256 + phi_grid: 32 +``` + +### 5. Report + +After jobs finish: + +```sh +containers/survey_scan.sh collect --survey survey/cit-YYYYMMDD +``` + +Produce: + +- GPU classes observed; +- image selected for each class; +- CuPy/JAX versions; +- cold compile time; +- warm second-call time; +- cache size and file count; +- failures by GPU class; +- recommended manifest/build-matrix changes. + +## Where This Should Live + +Start in this repo under `containers/`, not in a separate repo. + +Reasons: + +- The probes need RIFT-specific imports, CLI names, and expected shapes. +- The container manifest and build-family code already live here. +- Keeping the first version local makes it easier for pipeline changes to evolve + with the warmup profiles. + +Split into a separate repository only after the tool has stable boundaries, for +example if it becomes a general IGWN GPU-container survey/warmup kit. A clean +future split would keep generic Condor/GPU inventory and cache-collection logic +outside RIFT, while RIFT keeps its own warmup profile scripts. + +## Proposed File Layout + +```text +containers/ + survey_scan.sh + survey_scan/ + README.md + gpu_inventory.py + emit_condor_jobs.py + collect_results.py + profiles/ + rift_cupy_common.py + rift_jax_ile_common.py + o4_default.yaml +``` + +The Python files should be dependency-light: standard library plus optional +`PyYAML` when parsing manifests/profiles. They should not require JAX or CuPy on +the submit host; those imports happen inside the container probe. + +## Integration with the Current CIT Kit + +For the remote CIT kit, `survey_scan` can be used in place before or after image +builds: + +1. `survey` before a build to confirm the matrix still covers the pool. +2. `emit-jobs` after images are built/staged to run one warmup job per class. +3. `collect` to decide whether caches should be baked into the next image or + distributed as a cache tarball beside each image. + +The existing `rift_container_select.sh` wrapper is a good runtime integration +point. It already detects compute capability and selects the image. It could also +set: + +```sh +export RIFT_GPU_CAPABILITY="$cap" +export RIFT_CONTAINER_LABEL="${LABELS[$sel]}" +export CUPY_CACHE_DIR="${RIFT_CACHE_ROOT:-${_CONDOR_SCRATCH_DIR}/.rift_cache}/cupy/${LABELS[$sel]}" +export JAX_COMPILATION_CACHE_DIR="${RIFT_CACHE_ROOT:-${_CONDOR_SCRATCH_DIR}/.rift_cache}/jax/${LABELS[$sel]}" +``` + +For CVMFS deployments, if preseeded caches are published next to the SIF, the +wrapper can copy or bind the matching cache directory into the job scratch area +before running the real command. + +## First Implementation Milestone + +Do not start with full automatic cache baking. Start with observability and +repeatable probes: + +1. Add `survey_scan survey` and commit its CIT inventory output format. +2. Add `rift_cupy_common.py` warmup and a Condor job emitter for one GPU class. +3. Run on `cc60-90` and `cc90-120`; record cold/warm timings and cache sizes. +4. Add the JAX warmup only for the JAX image after the CuPy path is stable. +5. Decide whether the next image build should bake caches in `%post`, publish a + sidecar cache tarball, or simply rely on per-slot first-use warming. diff --git a/containers/survey_scan.sh b/containers/survey_scan.sh new file mode 100755 index 000000000..ddc63aa04 --- /dev/null +++ b/containers/survey_scan.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# Survey target GPU pools and emit representative RIFT container warmup probes. +set -euo pipefail + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PYTHON="${PYTHON:-python3}" + +usage() { + cat <&2 + usage >&2 + exit 2 + ;; +esac diff --git a/containers/survey_scan/README.md b/containers/survey_scan/README.md new file mode 100644 index 000000000..1c2d1d41f --- /dev/null +++ b/containers/survey_scan/README.md @@ -0,0 +1,43 @@ +# RIFT `survey_scan` + +`survey_scan` is a lightweight build/deployment companion for RIFT container +families. It surveys the target Condor GPU pool, emits one warmup job per +container image band, and collects JSON timing/cache reports from common CuPy +and JAX startup probes. + +It is deliberately RIFT-specific: the useful probes exercise the NoLoop CuPy +kernels, fused calmarg kernels, and JAX ILE wrapper shapes that dominate startup +cost. + +## Commands + +```sh +containers/survey_scan.sh survey \ + --out survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml +containers/survey_scan.sh emit-jobs \ + --survey survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml +containers/survey_scan.sh collect --survey survey/cit-YYYYMMDD +``` + +The submit-side commands use only the Python standard library. `PyYAML` is used +when available for manifest parsing; otherwise a small parser handles the simple +RIFT container-family YAML schema. + +## Profiles + +- `rift_cupy_common.py`: warms `Q_inner_product_cupy`, + `Q_fused_calmarg_cupy`, `Q_fused_calmarg_distmarg_cupy`, and + `interp_gpu.interp`. +- `rift_jax_ile_common.py`: warms synthetic JAX ILE wrapper modes. Use this only + for JAX-enabled images. + +The generated jobs run the profile inside the chosen container with: + +```sh +apptainer exec --nv python3 --json-out .json +``` + +If the manifest image is an `osdf://` URL, the generated wrapper fetches only +that image with `stashcp` or `pelican`. diff --git a/containers/survey_scan/collect_results.py b/containers/survey_scan/collect_results.py new file mode 100755 index 000000000..65137fbae --- /dev/null +++ b/containers/survey_scan/collect_results.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +"""Collect completed survey_scan warmup JSON outputs.""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +from common import read_json, write_json + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--survey", required=True, help="Survey directory.") + ap.add_argument("--out", default=None, help="Summary JSON path.") + args = ap.parse_args(argv) + + survey = Path(args.survey) + jobs = survey / "jobs" + result_files = sorted(jobs.glob("*.json")) + results = [] + for path in result_files: + try: + item = read_json(path) + item["_path"] = str(path) + results.append(item) + except Exception as exc: # noqa: BLE001 + results.append({"_path": str(path), "error": str(exc)}) + + summary = { + "survey": str(survey), + "n_results": len(results), + "results": results, + } + out = Path(args.out) if args.out else survey / "warmup_summary.json" + write_json(out, summary) + + md = out.with_suffix(".md") + with md.open("w", encoding="utf-8") as f: + f.write("# Warmup Summary\n\n") + f.write("| profile | status | device | elapsed s | cache bytes | path |\n") + f.write("|---|---|---|---:|---:|---|\n") + for item in results: + profile = item.get("profile", "?") + status = "PASS" if item.get("ok") else "FAIL" + device = item.get("device", {}).get("name", "?") if isinstance(item.get("device"), dict) else "?" + elapsed = item.get("elapsed_s", "") + cache = item.get("cache", {}).get("bytes", "") if isinstance(item.get("cache"), dict) else "" + f.write(f"| {profile} | {status} | {device} | {elapsed} | {cache} | {item.get('_path')} |\n") + print(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/containers/survey_scan/common.py b/containers/survey_scan/common.py new file mode 100644 index 000000000..d2ac57722 --- /dev/null +++ b/containers/survey_scan/common.py @@ -0,0 +1,147 @@ +"""Shared helpers for container survey_scan tooling.""" + +from __future__ import annotations + +import json +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass +class ContainerEntry: + label: str + image: str + cuda_capability_min: float | None + cuda_capability_max: float | None + note: str = "" + + +def write_json(path: Path, obj: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as f: + json.dump(obj, f, indent=2, sort_keys=True) + f.write("\n") + + +def read_json(path: Path) -> Any: + with path.open(encoding="utf-8") as f: + return json.load(f) + + +def _coerce_scalar(value: str) -> Any: + value = value.strip() + if value in ("", "null", "None", "~"): + return None + if (value.startswith('"') and value.endswith('"')) or ( + value.startswith("'") and value.endswith("'") + ): + return value[1:-1] + try: + return int(value) + except ValueError: + pass + try: + return float(value) + except ValueError: + return value + + +def _parse_simple_yaml_manifest(text: str) -> dict[str, Any]: + """Parse the simple RIFT container-family YAML schema without dependencies. + + This is not a general YAML parser. It handles the schema emitted by + containers/build_family.sh and the CIT build kit: top-level scalars plus a + `containers:` list of scalar mappings. + """ + + result: dict[str, Any] = {"containers": []} + in_containers = False + current: dict[str, Any] | None = None + for raw in text.splitlines(): + line = raw.split("#", 1)[0].rstrip() + if not line.strip(): + continue + if line.strip() == "containers:": + in_containers = True + continue + if not in_containers: + if ":" in line: + key, value = line.split(":", 1) + result[key.strip()] = _coerce_scalar(value) + continue + stripped = line.strip() + if stripped.startswith("- "): + if current is not None: + result["containers"].append(current) + current = {} + stripped = stripped[2:] + if stripped and ":" in stripped: + key, value = stripped.split(":", 1) + current[key.strip()] = _coerce_scalar(value) + elif current is not None and ":" in stripped: + key, value = stripped.split(":", 1) + current[key.strip()] = _coerce_scalar(value) + if current is not None: + result["containers"].append(current) + return result + + +def load_manifest(path: Path) -> dict[str, Any]: + text = path.read_text(encoding="utf-8") + try: + import yaml # type: ignore + + loaded = yaml.safe_load(text) + if isinstance(loaded, dict): + return loaded + except Exception: + pass + return _parse_simple_yaml_manifest(text) + + +def manifest_entries(path: Path) -> list[ContainerEntry]: + manifest = load_manifest(path) + entries = [] + for item in manifest.get("containers", []): + entries.append( + ContainerEntry( + label=str(item.get("label", "")), + image=str(item.get("image", "")), + cuda_capability_min=_as_float_or_none( + item.get("cuda_capability_min") + ), + cuda_capability_max=_as_float_or_none( + item.get("cuda_capability_max") + ), + note=str(item.get("note", "") or ""), + ) + ) + return entries + + +def _as_float_or_none(value: Any) -> float | None: + if value is None or value == "": + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def safe_name(value: str) -> str: + value = value.strip() or "unknown" + return re.sub(r"[^A-Za-z0-9_.-]+", "_", value) + + +def repo_root_from_here() -> Path: + return Path(__file__).resolve().parents[2] + + +def rel_or_abs(path: Path) -> str: + try: + return os.path.relpath(path, Path.cwd()) + except ValueError: + return str(path) diff --git a/containers/survey_scan/emit_condor_jobs.py b/containers/survey_scan/emit_condor_jobs.py new file mode 100755 index 000000000..59e4aab38 --- /dev/null +++ b/containers/survey_scan/emit_condor_jobs.py @@ -0,0 +1,175 @@ +#!/usr/bin/env python3 +"""Emit Condor jobs that run RIFT container warmup profiles.""" + +from __future__ import annotations + +import argparse +import os +import stat +import sys +from pathlib import Path + +from common import manifest_entries, rel_or_abs, repo_root_from_here, safe_name + + +PROFILE_MAP = { + "cupy": "rift_cupy_common.py", + "jax": "rift_jax_ile_common.py", +} + + +def _constraint(min_cap: float | None, max_cap: float | None) -> str: + parts = ["(Capability =!= undefined)"] + if min_cap is not None: + parts.append(f"(Capability >= {min_cap})") + if max_cap is not None: + parts.append(f"(Capability <= {max_cap})") + return " && ".join(parts) + + +def _write_runner(path: Path, image: str, profile: str, result: str) -> None: + path.write_text( + f"""#!/usr/bin/env bash +set -euo pipefail +log() {{ echo "[survey_scan] $*" >&2; }} + +image={image!r} +profile={profile!r} +result={result!r} +sif="$image" + +if [[ "$image" == osdf://* ]]; then + base="$(basename "$image")" + if [ -e "$base" ]; then + sif="./$base" + else + log "fetching $image" + if command -v stashcp >/dev/null 2>&1; then + stashcp "$image" "$base" + elif command -v pelican >/dev/null 2>&1; then + pelican object get "$image" "$base" + else + log "FATAL: no stashcp or pelican available for $image" + exit 4 + fi + sif="./$base" + fi +fi + +cache_root="${{RIFT_SURVEY_CACHE_ROOT:-${{_CONDOR_SCRATCH_DIR:-$PWD}}/.rift_cache}}" +mkdir -p "$cache_root" +export CUPY_CACHE_DIR="${{CUPY_CACHE_DIR:-$cache_root/cupy}}" +export CUPY_CACHE_IN_MEMORY="${{CUPY_CACHE_IN_MEMORY:-0}}" +export JAX_COMPILATION_CACHE_DIR="${{JAX_COMPILATION_CACHE_DIR:-$cache_root/jax}}" +export JAX_ENABLE_X64="${{JAX_ENABLE_X64:-1}}" +export XLA_FLAGS="${{XLA_FLAGS:---xla_cpu_multi_thread_eigen=false}}" +export OMP_NUM_THREADS="${{OMP_NUM_THREADS:-1}}" +export OPENBLAS_NUM_THREADS="${{OPENBLAS_NUM_THREADS:-1}}" +export MKL_NUM_THREADS="${{MKL_NUM_THREADS:-1}}" +export NUMEXPR_NUM_THREADS="${{NUMEXPR_NUM_THREADS:-1}}" + +log "image=$sif" +log "profile=$profile" +log "result=$result" +apptainer exec --nv "$sif" python3 "$profile" --json-out "$result" +""", + encoding="utf-8", + ) + path.chmod(path.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + + +def _write_submit( + path: Path, + runner: Path, + profile_path: Path, + result_name: str, + min_cap: float | None, + max_cap: float | None, + request_disk: str, +) -> None: + path.write_text( + f"""universe = vanilla +executable = {runner.name} +arguments = +request_GPUs = 1 +request_disk = {request_disk} +require_gpus = {_constraint(min_cap, max_cap)} +transfer_input_files = {profile_path} +transfer_output_files = {result_name} +output = $(Cluster).$(Process).out +error = $(Cluster).$(Process).err +log = $(Cluster).log +queue 1 +""", + encoding="utf-8", + ) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--survey", required=True, help="Survey directory.") + ap.add_argument("--manifest", required=True, help="Container-family manifest.") + ap.add_argument("--out", default=None, help="Output jobs directory.") + ap.add_argument( + "--profiles", + default="cupy", + help="Comma-separated profiles: cupy,jax. Default: cupy.", + ) + ap.add_argument("--request-disk", default="16000M") + args = ap.parse_args(argv) + + survey = Path(args.survey) + out = Path(args.out) if args.out else survey / "jobs" + out.mkdir(parents=True, exist_ok=True) + profile_dir = repo_root_from_here() / "containers" / "survey_scan" / "profiles" + selected_profiles = [x.strip() for x in args.profiles.split(",") if x.strip()] + + manifest = Path(args.manifest) + entries = manifest_entries(manifest) + if not entries: + raise SystemExit(f"No container entries found in {manifest}") + + generated = [] + for entry in entries: + for profile_key in selected_profiles: + profile_name = PROFILE_MAP.get(profile_key, profile_key) + profile_path = profile_dir / profile_name + if not profile_path.exists(): + raise SystemExit(f"Profile not found: {profile_path}") + stem = safe_name(f"{entry.label}_{Path(profile_name).stem}") + runner = out / f"run_{stem}.sh" + result = f"{stem}.json" + submit = out / f"{stem}.sub" + _write_runner(runner, entry.image, profile_path.name, result) + _write_submit( + submit, + runner, + Path(profile_path.name), + result, + entry.cuda_capability_min, + entry.cuda_capability_max, + args.request_disk, + ) + generated.append(submit) + + # Copy profile scripts next to the jobs so condor_submit can run from out/. + for profile_key in selected_profiles: + profile_name = PROFILE_MAP.get(profile_key, profile_key) + src = profile_dir / profile_name + dst = out / profile_name + dst.write_text(src.read_text(encoding="utf-8"), encoding="utf-8") + dst.chmod(dst.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) + + with (out / "submit_all.sh").open("w", encoding="utf-8") as f: + f.write("#!/usr/bin/env bash\nset -euo pipefail\n") + f.write("cd \"$(dirname \"$0\")\"\n") + for sub in generated: + f.write(f"condor_submit {sub.name}\n") + (out / "submit_all.sh").chmod(0o755) + + print(rel_or_abs(out)) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/containers/survey_scan/gpu_inventory.py b/containers/survey_scan/gpu_inventory.py new file mode 100755 index 000000000..ffb6209f2 --- /dev/null +++ b/containers/survey_scan/gpu_inventory.py @@ -0,0 +1,193 @@ +#!/usr/bin/env python3 +"""Survey Condor GPU inventory for RIFT container-family planning.""" + +from __future__ import annotations + +import argparse +import datetime as _dt +import shutil +import socket +import subprocess +import sys +from collections import Counter +from pathlib import Path + +from common import manifest_entries, write_json + + +FIELDS = [ + "Name", + "GPUs_DeviceName", + "GPUs_Capability", + "GPUs_GlobalMemoryMb", + "CUDACapability", + "CUDADeviceName", + "CUDADeviceGlobalMemoryMb", +] + + +def _run_condor_status(constraint: str) -> list[dict[str, str]]: + if shutil.which("condor_status") is None: + raise SystemExit("condor_status not found on PATH") + cmd = ["condor_status", "-constraint", constraint, "-af", *FIELDS] + proc = subprocess.run(cmd, text=True, capture_output=True, check=False) + if proc.returncode != 0: + raise SystemExit(proc.stderr.strip() or "condor_status failed") + rows = [] + for line in proc.stdout.splitlines(): + parts = line.split() + if len(parts) < len(FIELDS): + parts = parts + ["undefined"] * (len(FIELDS) - len(parts)) + rows.append(dict(zip(FIELDS, parts[: len(FIELDS)]))) + return rows + + +def _norm(row: dict[str, str]) -> tuple[str, str, str]: + name = row.get("GPUs_DeviceName") or row.get("CUDADeviceName") or "undefined" + cap = row.get("GPUs_Capability") or row.get("CUDACapability") or "undefined" + mem = row.get("GPUs_GlobalMemoryMb") or row.get("CUDADeviceGlobalMemoryMb") or "undefined" + return name, cap, mem + + +def _recommend_bands(summary: Counter[tuple[str, str, str]]) -> list[dict[str, object]]: + caps = [] + for (_name, cap, _mem), count in summary.items(): + try: + caps.append((float(cap), count)) + except ValueError: + continue + if not caps: + return [] + min_cap = min(c for c, _ in caps) + max_cap = max(c for c, _ in caps) + bands = [] + if min_cap < 9.0: + bands.append( + { + "label": "cc60-90" if min_cap >= 6.0 else "default", + "cuda_capability_min": max(3.5, min_cap), + "cuda_capability_max": 9.0, + "reason": "Observed pre-Blackwell CUDA 11-compatible GPUs.", + } + ) + if max_cap >= 9.0: + bands.append( + { + "label": "cc90-120", + "cuda_capability_min": 9.0, + "cuda_capability_max": max(12.0, max_cap), + "reason": "Observed Hopper/Blackwell-class GPUs; use CUDA 12 devel when NVRTC headers are needed.", + } + ) + return bands + + +def _coverage(summary: Counter[tuple[str, str, str]], manifest: Path | None) -> list[dict[str, object]]: + if manifest is None: + return [] + entries = manifest_entries(manifest) + rows = [] + for (device, cap, mem), count in summary.most_common(): + matches = [] + try: + cap_f = float(cap) + except ValueError: + cap_f = None + if cap_f is not None: + for entry in entries: + lo = entry.cuda_capability_min + hi = entry.cuda_capability_max + if lo is not None and cap_f < lo: + continue + if hi is not None and cap_f > hi: + continue + matches.append(entry.label) + rows.append( + { + "device": device, + "capability": cap, + "memory_mb": mem, + "slots": count, + "manifest_matches": matches, + "covered": bool(matches), + } + ) + return rows + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--out", default=None, help="Output directory.") + ap.add_argument( + "--constraint", + default="TotalGPUs > 0", + help="condor_status constraint for GPU inventory.", + ) + ap.add_argument("--manifest", default=None, help="Optional container-family manifest to check coverage.") + args = ap.parse_args(argv) + + stamp = _dt.datetime.utcnow().strftime("%Y%m%dT%H%M%SZ") + out = Path(args.out or f"survey/{socket.gethostname()}-{stamp}") + out.mkdir(parents=True, exist_ok=True) + + rows = _run_condor_status(args.constraint) + summary = Counter(_norm(row) for row in rows) + + manifest_path = Path(args.manifest) if args.manifest else None + coverage = _coverage(summary, manifest_path) + + write_json(out / "gpu_inventory.json", { + "created_utc": stamp, + "host": socket.gethostname(), + "constraint": args.constraint, + "manifest": str(manifest_path) if manifest_path else None, + "fields": FIELDS, + "rows": rows, + "summary": [ + {"device": k[0], "capability": k[1], "memory_mb": k[2], "slots": v} + for k, v in summary.most_common() + ], + "coverage": coverage, + }) + write_json(out / "recommended_matrix.json", { + "created_utc": stamp, + "bands": _recommend_bands(summary), + }) + + with (out / "gpu_inventory.tsv").open("w", encoding="utf-8") as f: + f.write("slots\tdevice\tcapability\tmemory_mb\n") + for (device, cap, mem), count in summary.most_common(): + f.write(f"{count}\t{device}\t{cap}\t{mem}\n") + + with (out / "coverage.md").open("w", encoding="utf-8") as f: + f.write("# GPU Survey\n\n") + f.write(f"- Created UTC: `{stamp}`\n") + f.write(f"- Host: `{socket.gethostname()}`\n") + f.write(f"- Constraint: `{args.constraint}`\n\n") + f.write("| slots | device | capability | memory MB |\n") + f.write("|---:|---|---:|---:|\n") + for (device, cap, mem), count in summary.most_common(): + f.write(f"| {count} | {device} | {cap} | {mem} |\n") + f.write("\n## Suggested Bands\n\n") + for band in _recommend_bands(summary): + f.write( + f"- `{band['label']}`: cc {band['cuda_capability_min']} - " + f"{band['cuda_capability_max']} ({band['reason']})\n" + ) + if coverage: + f.write("\n## Manifest Coverage\n\n") + f.write("| slots | device | capability | matching labels |\n") + f.write("|---:|---|---:|---|\n") + for row in coverage: + labels = ", ".join(row["manifest_matches"]) or "UNCOVERED" + f.write( + f"| {row['slots']} | {row['device']} | " + f"{row['capability']} | {labels} |\n" + ) + + print(out) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/containers/survey_scan/profiles/rift_cupy_common.py b/containers/survey_scan/profiles/rift_cupy_common.py new file mode 100755 index 000000000..433d913ff --- /dev/null +++ b/containers/survey_scan/profiles/rift_cupy_common.py @@ -0,0 +1,148 @@ +#!/usr/bin/env python3 +"""Warm common RIFT CuPy kernels inside a GPU-enabled container.""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import sys +import time +from pathlib import Path + + +def _cache_stats(path: str | None) -> dict[str, object]: + if not path: + return {"path": None, "files": 0, "bytes": 0} + root = Path(path) + files = 0 + total = 0 + if root.exists(): + for p in root.rglob("*"): + if p.is_file(): + files += 1 + total += p.stat().st_size + return {"path": str(root), "files": files, "bytes": total} + + +def _write(path: str | None, obj: dict[str, object]) -> None: + text = json.dumps(obj, indent=2, sort_keys=True) + "\n" + if path: + Path(path).write_text(text, encoding="utf-8") + print(text) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--json-out", default=None) + args = ap.parse_args(argv) + + t0 = time.perf_counter() + steps: list[dict[str, object]] = [] + ok = True + err = None + device: dict[str, object] = {} + try: + import cupy as cp + + props = cp.cuda.runtime.getDeviceProperties(0) + name = props["name"].decode() if isinstance(props["name"], bytes) else props["name"] + device = { + "name": name, + "compute_capability": f"{props['major']}.{props['minor']}", + "runtime_version": cp.cuda.runtime.runtimeGetVersion(), + "cupy_version": cp.__version__, + } + + # 1. Standard NoLoop Q inner-product RawKernel. + from RIFT.likelihood.Q_inner_product import Q_inner_product_cupy + + Q = cp.ascontiguousarray( + (cp.random.random((64, 4)) + 1j * cp.random.random((64, 4))).astype(cp.complex128) + ) + A = cp.ascontiguousarray( + (cp.random.random((8, 4)) + 1j * cp.random.random((8, 4))).astype(cp.complex128) + ) + starts = cp.asarray([0, 2, 4, 6, 8, 10, 12, 14], dtype=cp.int32) + s0 = time.perf_counter() + out = Q_inner_product_cupy(Q, A, starts, 16) + cp.cuda.Stream.null.synchronize() + steps.append({"name": "Q_inner_product_cupy", "elapsed_s": time.perf_counter() - s0, "shape": list(out.shape)}) + + # 2. Fused calmarg kernels. + from RIFT.likelihood.Q_fused_calmarg import ( + Q_fused_calmarg_cupy, + Q_fused_calmarg_distmarg_cupy, + ) + + n_det, n_cal, n_window, n_lms, n_ext, npts = 2, 3, 32, 4, 8, 16 + Qf = cp.ascontiguousarray( + (cp.random.random((n_det, n_cal * n_window, n_lms)) + + 1j * cp.random.random((n_det, n_cal * n_window, n_lms))).astype(cp.complex128) + ) + Af = cp.ascontiguousarray( + (cp.random.random((n_det, n_ext, n_lms)) + + 1j * cp.random.random((n_det, n_ext, n_lms))).astype(cp.complex128) + ) + ifirst = cp.ascontiguousarray(cp.tile(cp.arange(n_ext, dtype=cp.int32), (n_det, 1))) + inv_dist = cp.ones(n_ext, dtype=cp.float64) + rho_sq = cp.ones((n_ext, npts), dtype=cp.float64) + w_t = cp.ones(npts, dtype=cp.float64) / npts + s0 = time.perf_counter() + y = Q_fused_calmarg_cupy(Qf, Af, ifirst, inv_dist, rho_sq, w_t, n_cal, n_window) + cp.cuda.Stream.null.synchronize() + steps.append({"name": "Q_fused_calmarg_cupy", "elapsed_s": time.perf_counter() - s0, "shape": list(y.shape)}) + + lnI = cp.zeros((8, 8), dtype=cp.float64) + distmarg = { + "lnI_array": lnI, + "s0": -4.0, + "ds": 1.0, + "smin": -4.0, + "smax": 3.0, + "t0": -4.0, + "dt": 1.0, + "tmax": 3.0, + "xmin": 0.001, + "xmax": 10.0, + "sqrt_bmax": 1.0, + "bref": 1.0, + } + s0 = time.perf_counter() + yd = Q_fused_calmarg_distmarg_cupy(Qf, Af, ifirst, inv_dist, rho_sq, w_t, n_cal, n_window, distmarg) + cp.cuda.Stream.null.synchronize() + steps.append({"name": "Q_fused_calmarg_distmarg_cupy", "elapsed_s": time.perf_counter() - s0, "shape": list(yd.shape)}) + + # 3. RIFT's temporary cupy interp ElementwiseKernel. + from RIFT.interpolators.interp_gpu import interp + + xp = cp.linspace(0.0, 1.0, 32, dtype=cp.float64) + fp = cp.sin(xp) + x = cp.linspace(-0.1, 1.1, 128, dtype=cp.float64) + s0 = time.perf_counter() + zi = interp(x, xp, fp) + cp.cuda.Stream.null.synchronize() + steps.append({"name": "interp_gpu.interp", "elapsed_s": time.perf_counter() - s0, "shape": list(zi.shape)}) + + except Exception as exc: # noqa: BLE001 + ok = False + err = f"{type(exc).__name__}: {exc}" + + elapsed = time.perf_counter() - t0 + result = { + "profile": "rift_cupy_common", + "ok": ok, + "error": err, + "host": socket.gethostname(), + "elapsed_s": elapsed, + "device": device, + "cache": _cache_stats(os.environ.get("CUPY_CACHE_DIR")), + "steps": steps, + } + _write(args.json_out, result) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/containers/survey_scan/profiles/rift_jax_ile_common.py b/containers/survey_scan/profiles/rift_jax_ile_common.py new file mode 100755 index 000000000..25bfbd6f8 --- /dev/null +++ b/containers/survey_scan/profiles/rift_jax_ile_common.py @@ -0,0 +1,169 @@ +#!/usr/bin/env python3 +"""Warm common synthetic RIFT JAX ILE wrappers inside a JAX-enabled container.""" + +from __future__ import annotations + +import argparse +import json +import os +import socket +import sys +import time +from pathlib import Path + + +def _cache_stats(path: str | None) -> dict[str, object]: + if not path: + return {"path": None, "files": 0, "bytes": 0} + root = Path(path) + files = 0 + total = 0 + if root.exists(): + for p in root.rglob("*"): + if p.is_file(): + files += 1 + total += p.stat().st_size + return {"path": str(root), "files": files, "bytes": total} + + +def _write(path: str | None, obj: dict[str, object]) -> None: + text = json.dumps(obj, indent=2, sort_keys=True) + "\n" + if path: + Path(path).write_text(text, encoding="utf-8") + print(text) + + +def _synthetic_data(npts: int, n_full: int, l_max: int): + import numpy as np + import jax.numpy as jnp + from RIFT.likelihood.jax_ile.core import JAXLikelihoodData + + lms = [(2, -2), (2, 2)] + if l_max >= 3: + lms += [(3, -3), (3, 3)] + if l_max >= 4: + lms += [(4, -4), (4, 4)] + k = len(lms) + rng = np.random.default_rng(1234) + detectors = {} + for i, det in enumerate(("H1", "L1")): + q = rng.normal(size=(n_full, k)) + 1j * rng.normal(size=(n_full, k)) + u = np.eye(k, dtype=np.complex128) + v = 0.05 * np.eye(k, dtype=np.complex128) + detectors[det] = { + "lms": lms, + "Q": jnp.asarray(q, dtype=jnp.complex128), + "U": jnp.asarray(u, dtype=jnp.complex128), + "V": jnp.asarray(v, dtype=jnp.complex128), + "epoch": 1000000000.0 + i * 0.002, + "location": jnp.asarray([3000.0 + i, 4000.0 - i, 5000.0 + 2 * i], dtype=jnp.float64), + "response": jnp.asarray(np.eye(3), dtype=jnp.float64), + "npts_full": n_full, + "l_max": l_max, + } + tvals = np.linspace(-0.05, 0.05, npts) + return JAXLikelihoodData(detectors, deltaT=1.0 / 4096.0, gmst=1.0, tvals=tvals, tref=1000000000.0) + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--json-out", default=None) + ap.add_argument("--npts", type=int, default=64) + ap.add_argument("--n-full", type=int, default=512) + ap.add_argument("--l-max", type=int, default=4) + ap.add_argument("--distance-grid", type=int, default=64) + ap.add_argument("--phi-grid", type=int, default=16) + ap.add_argument("--psi-grid", type=int, default=8) + args = ap.parse_args(argv) + + # Must happen before first substantial JAX use. + os.environ.setdefault("JAX_ENABLE_X64", "1") + os.environ.setdefault("XLA_FLAGS", "--xla_cpu_multi_thread_eigen=false") + + t0 = time.perf_counter() + steps: list[dict[str, object]] = [] + ok = True + err = None + device: dict[str, object] = {} + try: + import numpy as np + import jax + import jax.numpy as jnp + from RIFT.likelihood.jax_ile.wrapper import ( + JAXDistanceMarginalizedLikelihood, + JAXDistPhiMargLikelihood, + JAXDistPhiPsiMargLikelihood, + JAXExtrinsicLikelihood, + ) + + device = { + "jax_version": jax.__version__, + "backend": jax.default_backend(), + "devices": [str(d) for d in jax.devices()], + } + data = _synthetic_data(args.npts, args.n_full, args.l_max) + batch = { + "ra": jnp.asarray([0.1, 1.0]), + "dec": jnp.asarray([0.2, -0.1]), + "psi": jnp.asarray([0.3, 0.7]), + "incl": jnp.asarray([0.8, 1.1]), + "phiref": jnp.asarray([0.4, 1.2]), + "dist": jnp.asarray([500.0, 800.0]), + } + + s0 = time.perf_counter() + like6 = JAXExtrinsicLikelihood(data) + y = like6.log_likelihood(batch["ra"], batch["dec"], batch["psi"], batch["incl"], batch["phiref"], batch["dist"]) + np.asarray(y).tolist() + v, g = like6.value_and_grad(np.array([0.1, 0.2, 0.3, 0.8, 0.4, 500.0])) + steps.append({"name": "JAXExtrinsicLikelihood", "elapsed_s": time.perf_counter() - s0, "value": float(v), "grad_norm": float(np.linalg.norm(g))}) + + s0 = time.perf_counter() + like5 = JAXDistanceMarginalizedLikelihood(data, 100.0, 2000.0, n_grid=args.distance_grid) + y = like5.log_likelihood(batch["ra"], batch["dec"], batch["psi"], batch["incl"], batch["phiref"]) + np.asarray(y).tolist() + v, g = like5.value_and_grad(np.array([0.1, 0.2, 0.3, 0.8, 0.4])) + steps.append({"name": "JAXDistanceMarginalizedLikelihood", "elapsed_s": time.perf_counter() - s0, "value": float(v), "grad_norm": float(np.linalg.norm(g))}) + + s0 = time.perf_counter() + like4 = JAXDistPhiMargLikelihood(data, 100.0, 2000.0, nphi=args.phi_grid, n_grid=args.distance_grid) + y = like4.log_likelihood(batch["ra"], batch["dec"], batch["psi"], batch["incl"]) + np.asarray(y).tolist() + v, g = like4.value_and_grad(np.array([0.1, 0.2, 0.3, 0.8])) + steps.append({"name": "JAXDistPhiMargLikelihood", "elapsed_s": time.perf_counter() - s0, "value": float(v), "grad_norm": float(np.linalg.norm(g))}) + + s0 = time.perf_counter() + like3 = JAXDistPhiPsiMargLikelihood( + data, + 100.0, + 2000.0, + nphi=args.phi_grid, + npsi=args.psi_grid, + n_grid=args.distance_grid, + ) + y = like3.log_likelihood(batch["ra"], batch["dec"], batch["incl"]) + np.asarray(y).tolist() + v, g = like3.value_and_grad(np.array([0.1, 0.2, 0.8])) + steps.append({"name": "JAXDistPhiPsiMargLikelihood", "elapsed_s": time.perf_counter() - s0, "value": float(v), "grad_norm": float(np.linalg.norm(g))}) + + except Exception as exc: # noqa: BLE001 + ok = False + err = f"{type(exc).__name__}: {exc}" + + elapsed = time.perf_counter() - t0 + result = { + "profile": "rift_jax_ile_common", + "ok": ok, + "error": err, + "host": socket.gethostname(), + "elapsed_s": elapsed, + "device": device, + "cache": _cache_stats(os.environ.get("JAX_COMPILATION_CACHE_DIR")), + "steps": steps, + } + _write(args.json_out, result) + return 0 if ok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/containers/survey_scan/test_survey_scan.py b/containers/survey_scan/test_survey_scan.py new file mode 100644 index 000000000..a2cba546a --- /dev/null +++ b/containers/survey_scan/test_survey_scan.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python3 +"""Filesystem-level regression tests for the survey_scan tooling.""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from collections import Counter +from pathlib import Path + +import collect_results +import emit_condor_jobs +import gpu_inventory +from common import manifest_entries + + +MANIFEST = """\ +schema_version: 1 +containers: + - label: legacy + image: osdf://example.org/rift-legacy.sif + cuda_capability_min: 6.0 + cuda_capability_max: 8.9 + - label: modern + image: /cvmfs/example/rift-modern.sif + cuda_capability_min: 9.0 + cuda_capability_max: 12.0 +""" + + +class SurveyScanTests(unittest.TestCase): + def test_manifest_entries_and_inventory_coverage(self): + with tempfile.TemporaryDirectory() as tmp: + manifest = Path(tmp) / "manifest.yaml" + manifest.write_text(MANIFEST, encoding="utf-8") + + entries = manifest_entries(manifest) + self.assertEqual([entry.label for entry in entries], ["legacy", "modern"]) + self.assertEqual(entries[0].cuda_capability_max, 8.9) + + summary = Counter( + { + ("NVIDIA_A100", "8.0", "40960"): 3, + ("NVIDIA_H100", "9.0", "81920"): 2, + ("unknown", "undefined", "undefined"): 1, + } + ) + coverage = gpu_inventory._coverage(summary, manifest) + by_device = {row["device"]: row for row in coverage} + self.assertEqual(by_device["NVIDIA_A100"]["manifest_matches"], ["legacy"]) + self.assertEqual(by_device["NVIDIA_H100"]["manifest_matches"], ["modern"]) + self.assertFalse(by_device["unknown"]["covered"]) + + bands = gpu_inventory._recommend_bands(summary) + self.assertEqual([band["label"] for band in bands], ["cc60-90", "cc90-120"]) + + def test_emit_jobs_builds_constraints_and_osdf_runner(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + manifest = root / "manifest.yaml" + manifest.write_text(MANIFEST, encoding="utf-8") + survey = root / "survey" + out = survey / "jobs" + + rc = emit_condor_jobs.main( + [ + "--survey", + str(survey), + "--manifest", + str(manifest), + "--out", + str(out), + "--profiles", + "cupy", + ] + ) + self.assertEqual(rc, 0) + + legacy_submit = (out / "legacy_rift_cupy_common.sub").read_text(encoding="utf-8") + self.assertIn("(Capability >= 6.0)", legacy_submit) + self.assertIn("(Capability <= 8.9)", legacy_submit) + self.assertIn("transfer_output_files = legacy_rift_cupy_common.json", legacy_submit) + + legacy_runner = (out / "run_legacy_rift_cupy_common.sh").read_text(encoding="utf-8") + self.assertIn("stashcp", legacy_runner) + self.assertIn("pelican object get", legacy_runner) + self.assertIn("apptainer exec --nv", legacy_runner) + self.assertIn("JAX_COMPILATION_CACHE_DIR", legacy_runner) + self.assertTrue((out / "rift_cupy_common.py").exists()) + self.assertTrue((out / "submit_all.sh").exists()) + + def test_collect_results_records_success_and_malformed_output(self): + with tempfile.TemporaryDirectory() as tmp: + survey = Path(tmp) / "survey" + jobs = survey / "jobs" + jobs.mkdir(parents=True) + (jobs / "good.json").write_text( + json.dumps( + { + "profile": "cupy", + "ok": True, + "elapsed_s": 1.25, + "device": {"name": "A100"}, + "cache": {"bytes": 4096}, + } + ), + encoding="utf-8", + ) + (jobs / "bad.json").write_text("{not-json", encoding="utf-8") + + out = survey / "summary.json" + rc = collect_results.main(["--survey", str(survey), "--out", str(out)]) + self.assertEqual(rc, 0) + + summary = json.loads(out.read_text(encoding="utf-8")) + self.assertEqual(summary["n_results"], 2) + self.assertEqual(sum("error" in row for row in summary["results"]), 1) + markdown = out.with_suffix(".md").read_text(encoding="utf-8") + self.assertIn("| cupy | PASS | A100 | 1.25 | 4096 |", markdown) + self.assertIn("| ? | FAIL | ? |", markdown) + + +if __name__ == "__main__": + unittest.main() diff --git a/demos/integrator_snr_lottery/.gitignore b/demos/integrator_snr_lottery/.gitignore new file mode 100644 index 000000000..32eb564ef --- /dev/null +++ b/demos/integrator_snr_lottery/.gitignore @@ -0,0 +1,2 @@ +# raw sampler stdout: large, machine-specific, and superseded by the summarised results/*.txt +*.log diff --git a/demos/integrator_snr_lottery/README.md b/demos/integrator_snr_lottery/README.md new file mode 100644 index 000000000..5318fbc83 --- /dev/null +++ b/demos/integrator_snr_lottery/README.md @@ -0,0 +1,95 @@ +# Integrator SNR-lottery demos + +Truth-known (`shape_recovery.MixtureTarget.true_lnZ`), CPU-only studies of high-SNR extrinsic +integration failure. SNR ladder = peak width: `sigma_1d = 0.7*(20/SNR)` on the box [-5,5]^d. + +## results/restrict_ladder.txt -- restricted-range AV (n=20 per case per SNR) + +| case | SNR 40 bias / n_eff / collapse | SNR 80 | SNR 160 | +|------|-------------------------------|--------|---------| +| A_full standalone AV, full box | +0.026 / 2029 / 0% | +0.097 / 1802 / 0% | +0.102 / 1629 / 0% | +| A_sub standalone AV, CORRECT sub-box | -0.029 / 2495 / 0% | -0.029 / 2495 / 0% | -0.029 / 2495 / 0% | +| A_sub_wrong standalone, WRONG sub-box | **-35.9** / 840 / 0% | **-266.8** / 356 / 0% | **-1949.3** / 220 / 0% | +| P_full portfolio AV+AV, both full | -0.456 / 21 / 30% | -0.545 / 16 / 30% | -0.678 / 30 / 30% | +| P_mix portfolio full + CORRECT sub | -0.427 / 196 / 10% | -0.428 / 182 / 10% | -0.455 / 186 / 10% | +| P_mix_wrong portfolio full + WRONG sub| -0.699 / 9 / 55% | -1.076 / 12 / 45% | -1.183 / 35 / 25% | + +Readings: +1. **A_sub is SNR-INDEPENDENT** -- identical bias/n_eff at every SNR, because restricting the box to + the posterior makes the problem self-similar. A_full meanwhile degrades with SNR (n_eff 2029 -> + 1629, bias +0.026 -> +0.102). This is the mechanism working as intended. +2. **A_sub_wrong is the confidently-wrong failure in its purest form**: bias -36 / -267 / -1949 nats + while `collapse% = 0` and n_eff stays 220-840. n_eff reports a HEALTHY run. This is the strongest + argument in this whole study that n_eff cannot certify correctness, and it is the ideal generator + for validating a tail diagnostic (see tools/khat_validation.py). +3. **Fail-safe CONFIRMED**: the same wrong sub-box inside a portfolio (P_mix_wrong) costs ~1 nat and + some efficiency, not ~1949 nats -- the full-box member keeps q_mix covering. This is the reason to + prefer the multi-AV variant over mutating a single AV's range. +4. **P_mix improves on P_full**: collapse 30% -> 10%, n_eff ~20 -> ~185, at equal budget. + +HONEST CAVEAT: the portfolio rows carry a persistent ~-0.4 nat bias even in the GOOD case (P_mix +-0.43), but P_full shows the same (-0.456), so it is NOT caused by the restriction -- it is the known +downward skew of IS estimates at low n_eff (heavy-tailed weights). The portfolio path is therefore not +demonstrated unbiased at these n_eff; only the restriction's SAFETY and EFFICIENCY are established. + +## results/khat_decisive.txt -- does Pareto k-hat catch the confidently-wrong runs? NO (for this failure mode) + +Scored against KNOWN truth at SNR 160, n=20 per case. k-hat computed from the true importance +log-weights (log_integrand + log_joint_prior - log_joint_s_prior) via statutils.pareto_khat_from_log. + +| case | bias_med | n_eff_med | khat_med | % copies khat>0.7 | +|------|---------:|----------:|---------:|------------------:| +| A_full (accurate) | +0.102 | 1629 | -0.268 | 0% | +| A_sub (accurate) | -0.029 | 2495 | -0.302 | 0% | +| **A_sub_wrong (CATASTROPHIC)**| **-1949** | 220 | **0.435** | **10%** | +| P_full | -0.678 | 30 | 0.706 | 50% | +| **P_mix (BEST portfolio)** | -0.455 | 186 | **0.766** | **80%** | +| P_mix_wrong | -1.183 | 35 | 0.691 | 45% | + +**k-hat misses the catastrophe and fires on the good run.** The run biased by -1949 nats has +khat 0.435 -- BELOW the 0.7 "unresolved tail" threshold -- and only 10% of its copies trip the +threshold, so ~90% of catastrophically-wrong runs pass the check. Meanwhile the most ACCURATE +portfolio configuration (P_mix) has the HIGHEST khat (0.766, 80% firing). Ranking by khat is +anti-correlated with actual error here. + +**Mechanism (why this is not a bug in k-hat).** k-hat estimates the tail index of the weights you +ACTUALLY DREW. A sampler confined to a wrong sub-box never draws from the true peak at all, so its +observed weights are narrow and self-consistent -- the tail genuinely IS resolved, for the region it +sampled. The failure is total support non-overlap, not a heavy tail. **k-hat can detect mass whose +tail you have begun to sample; it cannot detect mass you have never touched.** Conversely P_mix has a +legitimately heavy tail (its full-box member occasionally lands a huge-weight point near the peak) -- +k-hat correctly flags that, but the run is accurate, so a k-hat gate would reject the best config. + +**Implication.** k-hat is a sound tail diagnostic and worth keeping, but it must NOT be used as a +pass/fail correctness gate for proposal/support-mismatch (mode-collapse) failures -- the dominant +high-SNR failure in this study. For that class the working detector remains CROSS-COPY DISAGREEMENT +(replicas / bootstrap quantiles): independent copies that localize differently disagree, and that is +observable, whereas a single run's own weights are not. Note this is one generator, d=4, n=20 -- +the mechanism is principled but the numbers are one configuration. + +## results/chunk_memory.txt -- the resource price of a larger chunk + +Real ILE likelihood, on-source point, warm start, nmax 4e5, one run per chunk size: + +| n_chunk | host peak (MiB) | wall (s) | +|--------:|----------------:|---------:| +| 10,000 | 2409 | 109 | +| 40,000 | 2369 | 51 | +| 160,000 | 2506 | 33 | + +**HOST memory is FLAT across a 16x chunk range** (2369-2506 MiB, ~3% spread = noise). Since condor's +`RequestMemory` governs the HOST, the well-motivated action is **change nothing**: raising the chunk +does not require a memory-request bump, and therefore does not restrict which slots a job can match. +This is consistent with the prior observation that RIFT extrinsic jobs already request 35-105x the +host RAM they actually use. + +**Wall time falls 3.3x** (109 -> 33 s) as the chunk grows -- fewer, larger GPU kernel launches. So the +SNR-scaled chunk is not merely free on host memory, it is faster. + +HONEST GAPS: +* **GPU memory was NOT successfully measured** (column read 0): the per-PID `nvidia-smi` filter did + not match the process that owns the CUDA context. The GPU-side cost of a larger chunk therefore + remains UNQUANTIFIED. It is the surface that plausibly does scale with chunk, so this gap matters + if a site is GPU-memory constrained -- it just is not the surface `RequestMemory` controls. +* n_eff in this table (7.3 / 5.5 / 4.8) is ONE cold run per chunk on the pathological point; that is + lottery noise, not a trend. The collapse-rate evidence is the multi-copy study above, not this. diff --git a/demos/integrator_snr_lottery/results/chunk_fixed_budget.json.gz b/demos/integrator_snr_lottery/results/chunk_fixed_budget.json.gz new file mode 100644 index 000000000..4d5ecefef Binary files /dev/null and b/demos/integrator_snr_lottery/results/chunk_fixed_budget.json.gz differ diff --git a/demos/integrator_snr_lottery/results/chunk_fixed_steps.json.gz b/demos/integrator_snr_lottery/results/chunk_fixed_steps.json.gz new file mode 100644 index 000000000..b2113b71f Binary files /dev/null and b/demos/integrator_snr_lottery/results/chunk_fixed_steps.json.gz differ diff --git a/demos/integrator_snr_lottery/results/chunk_memory.txt b/demos/integrator_snr_lottery/results/chunk_memory.txt new file mode 100644 index 000000000..92d6a081d --- /dev/null +++ b/demos/integrator_snr_lottery/results/chunk_memory.txt @@ -0,0 +1,8 @@ +# chunk-size resource cost, REAL ILE likelihood, on-source point +# GPU=2 nmax=400000 warm=1 Thu Jul 30 18:08:16 PDT 2026 +# host_MiB = peak VmHWM over the job's process tree (what RequestMemory governs) +# gpu_MiB = peak nvidia-smi used_memory for THIS job's pids only +chunk host_MiB gpu_MiB wall_s n_eff +10000 2409 0 109 7.3 +40000 2369 0 51 5.5 +160000 2506 0 33 4.8 diff --git a/demos/integrator_snr_lottery/results/khat_decisive.txt b/demos/integrator_snr_lottery/results/khat_decisive.txt new file mode 100644 index 000000000..7cebe21f3 --- /dev/null +++ b/demos/integrator_snr_lottery/results/khat_decisive.txt @@ -0,0 +1,7 @@ +# case bias_med bias_sd neff_med collapse% n_ok khat_med khat>0.7% + A_full +0.102 0.177 1629 0% 20 -0.268 0% + A_sub -0.029 0.047 2495 0% 20 -0.302 0% + A_sub_wrong -1949.282 80690.958 220 0% 20 0.435 10% + P_full -0.678 3.338 30 30% 20 0.706 50% + P_mix -0.455 0.679 186 10% 20 0.766 80% + P_mix_wrong -1.183 1.205 35 25% 20 0.691 45% diff --git a/demos/integrator_snr_lottery/results/khat_validation.json.gz b/demos/integrator_snr_lottery/results/khat_validation.json.gz new file mode 100644 index 000000000..a8ca8253d Binary files /dev/null and b/demos/integrator_snr_lottery/results/khat_validation.json.gz differ diff --git a/demos/integrator_snr_lottery/results/restrict_ladder.txt b/demos/integrator_snr_lottery/results/restrict_ladder.txt new file mode 100644 index 000000000..48a170ab0 --- /dev/null +++ b/demos/integrator_snr_lottery/results/restrict_ladder.txt @@ -0,0 +1,40 @@ +########## SNR 40 ########## + Adapting x3 + PORTFOLIO setup {} + PORTFOLIO setup {} + trial 19 P_mix_wrong bias=-2.906 n_eff=12 (10s) + +# case bias_med bias_sd neff_med collapse% n_ok + A_full +0.026 0.121 2029 0% 20 + A_sub -0.029 0.047 2495 0% 20 + A_sub_wrong -35.881 741.736 840 0% 20 + P_full -0.456 0.964 21 30% 20 + P_mix -0.427 0.653 196 10% 20 + P_mix_wrong -0.699 1.237 9 55% 20 +########## SNR 80 ########## + Adapting x3 + PORTFOLIO setup {} + PORTFOLIO setup {} + trial 19 P_mix_wrong bias=-3.499 n_eff=13 (5s) + +# case bias_med bias_sd neff_med collapse% n_ok + A_full +0.097 0.140 1802 0% 20 + A_sub -0.029 0.047 2495 0% 20 + A_sub_wrong -266.759 6039.169 356 0% 20 + P_full -0.545 1.767 16 30% 20 + P_mix -0.428 0.676 182 10% 20 + P_mix_wrong -1.076 1.194 12 45% 20 +########## SNR 160 ########## + Adapting x3 + PORTFOLIO setup {} + PORTFOLIO setup {} + trial 19 P_mix_wrong bias=-2.951 n_eff=4 (5s) + +# case bias_med bias_sd neff_med collapse% n_ok + A_full +0.102 0.177 1629 0% 20 + A_sub -0.029 0.047 2495 0% 20 + A_sub_wrong -1949.282 80690.958 220 0% 20 + P_full -0.678 3.338 30 30% 20 + P_mix -0.455 0.679 186 10% 20 + P_mix_wrong -1.183 1.205 35 25% 20 +LADDER DONE diff --git a/demos/integrator_snr_lottery/tools/chunk_study.py b/demos/integrator_snr_lottery/tools/chunk_study.py new file mode 100644 index 000000000..c5f34be43 --- /dev/null +++ b/demos/integrator_snr_lottery/tools/chunk_study.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python +""" +chunk_study.py -- is CHUNK SIZE a stability knob for high-SNR extrinsic integration? + +HYPOTHESIS (reviewer): part of the high-SNR collapse is that the chunk is too small. At high SNR the +posterior occupies a vanishing fraction of the prior volume, so a small chunk contains almost no +informative samples per adaptation step and the sampler adapts on noise. A larger chunk may raise +the SNR at which extrinsic integration collapses. + +COUNTERWEIGHT: GPU memory scales with chunk size. Bigger chunks restrict which resources a job can +run on and would force per-job memory tuning that production does not normally do (held jobs, idle +capacity). So the deliverable is not "bigger is better" but WHERE the trade sits. + +WHY THIS HARNESS RATHER THAN A REAL EVENT + * TRUTH IS KNOWN. MixtureTarget exposes `true_lnZ`, so we measure real BIAS (lnI - true_lnZ), not + scatter about an unknown answer. The real-event study could only ever measure scatter, and that + is what made it so easy to fool ourselves with small samples. + * It is CPU-ONLY, so copies are cheap. The real-event work was GPU-bound and ran 4-9 copies per + cell; at that size a variance estimate is nearly worthless (we retracted a result for exactly + this reason). Here we can afford tens of copies per cell. + +SNR LADDER. Posterior width scales as 1/SNR, and MixtureTarget's `sigma_1d` IS the peak width on a +fixed box [-5,5]^d. So we set + + sigma_1d = SIGMA_REF * (SNR_REF / SNR), SIGMA_REF=0.7 at SNR_REF=20 + +i.e. SNR 20 -> sigma 0.7 (the gate's own default), SNR 160 -> sigma 0.0875. The peak's volume +fraction falls like (sigma/10)^d, which is the mechanism we care about. + +FAIR COMPARISON. Total budget `nmax` is held FIXED across chunk sizes, so this is a same-cost +comparison. Note the coupling that makes it interesting: n_steps = nmax/n_chunk, so a larger chunk +buys better per-step statistics at the price of FEWER adaptation steps. That is the real trade. + +Usage: + export PYTHONPATH=/MonteCarloMarginalizeCode/Code + export CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 + python chunk_study.py --copies 24 --jobs 32 --json results/chunk_study.json +""" +from __future__ import print_function +import argparse, json, os, sys, time +import numpy as np +from multiprocessing import Pool + +HERE = os.path.dirname(os.path.abspath(__file__)) +GATE = os.path.abspath(os.path.join( + HERE, "..", "..", "..", "MonteCarloMarginalizeCode", "Code", + "test", "expensive_before_merging", "integrators")) +sys.path.insert(0, GATE) +import shape_recovery as SR # reuse the gate's targets, runner, metrics and verdicts + +SIGMA_REF, SNR_REF = 0.7, 20.0 + + +def sigma_for_snr(snr): + """Peak width for an 'SNR': posterior width ~ 1/SNR, anchored at the gate's default.""" + return SIGMA_REF * (SNR_REF / float(snr)) + + +def _one(job): + (snr, n_chunk, ndim, ncomp, nmax, neff, seed, kind) = job + # nmax is resolved by the caller: FIXED-BUDGET mode passes the same nmax to every chunk size + # (same cost, but steps = nmax/n_chunk falls as the chunk grows), while FIXED-STEPS mode passes + # nmax = n_chunk*steps (equal adaptation opportunities, but cost grows with the chunk). + sigma = sigma_for_snr(snr) + t0 = time.time() + try: + target = SR.MixtureTarget(ndim, ncomp, seed, sigma_1d=sigma) + rec = SR.run_one(kind, target, nmax, neff, n_chunk=n_chunk, seed=seed) + verdict = SR.evaluate(rec) + status = verdict if isinstance(verdict, str) else verdict[0] + except Exception as e: # never let one cell kill the sweep + rec, status = {"error": str(e)[:200]}, "ERROR" + return dict(snr=snr, sigma=sigma, n_chunk=n_chunk, ndim=ndim, ncomp=ncomp, + nmax=nmax, seed=seed, kind=kind, status=status, + n_steps=int(nmax // n_chunk), + bias_ln=float(rec.get("bias_ln", np.nan)), + n_eff=float(rec.get("n_eff", np.nan)), + wall=time.time() - t0) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--snrs", default="20,40,80,160") + ap.add_argument("--chunks", default="10000,40000,160000") + ap.add_argument("--ndim", type=int, default=4) + ap.add_argument("--ncomp", type=int, default=3) + ap.add_argument("--nmax", type=int, default=1000000, help="FIXED total budget (same cost)") + ap.add_argument("--steps", type=int, default=None, + help="FIXED-STEPS mode: set nmax = n_chunk*steps per cell instead of --nmax. " + "Equal adaptation opportunities across chunk sizes, so it isolates whether " + "richer per-step statistics help; NOT a same-cost comparison.") + ap.add_argument("--neff", type=int, default=3000) + ap.add_argument("--copies", type=int, default=24) + ap.add_argument("--kinds", default="AV,portfolio") + ap.add_argument("--seed0", type=int, default=4000) + ap.add_argument("--jobs", type=int, default=16) + ap.add_argument("--json", default=None) + a = ap.parse_args() + + snrs = [float(x) for x in a.snrs.split(",")] + chunks = [int(x) for x in a.chunks.split(",")] + kinds = [k.strip() for k in a.kinds.split(",") if k.strip()] + + jobs = [] + for kind in kinds: + for snr in snrs: + for nc in chunks: + nmax_here = nc * a.steps if a.steps else a.nmax + for c in range(a.copies): + jobs.append((snr, nc, a.ndim, a.ncomp, nmax_here, a.neff, + a.seed0 + c, kind)) + print("# chunk study: {} kinds x {} SNR x {} chunks x {} copies = {} runs" + .format(len(kinds), len(snrs), len(chunks), a.copies, len(jobs))) + if a.steps: + print("# FIXED-STEPS mode: steps={} so nmax = n_chunk*steps (cost GROWS with chunk); " + "isolates per-step statistics".format(a.steps)) + else: + print("# FIXED-BUDGET mode: nmax={} for every chunk (same cost; steps = nmax/n_chunk " + "FALLS as the chunk grows)".format(a.nmax)) + print("# ndim={} ncomp={} neff={}".format(a.ndim, a.ncomp, a.neff)) + print("# sigma ladder: " + ", ".join("SNR%g->%.4f" % (s, sigma_for_snr(s)) for s in snrs)) + sys.stdout.flush() + + t0 = time.time() + pool = Pool(a.jobs) + recs = pool.map(_one, jobs) + pool.close(); pool.join() + print("# done in {:.1f} min".format((time.time() - t0) / 60.0)) + + if a.json: + os.makedirs(os.path.dirname(os.path.abspath(a.json)), exist_ok=True) + json.dump(recs, open(a.json, "w"), indent=1) + print("# wrote", a.json) + + summarize(recs) + + +def summarize(recs): + """Per cell: collapse fraction (status != PASS) and median |bias| among PASSing copies.""" + import collections + cells = collections.OrderedDict() + for r in recs: + cells.setdefault((r["kind"], r["snr"], r["n_chunk"]), []).append(r) + print("\n%-10s %6s %9s %7s %8s %12s %10s" % + ("kind", "SNR", "n_chunk", "steps", "collapse", "med|bias|", "med n_eff")) + for (kind, snr, nc), rs in cells.items(): + bad = [r for r in rs if r["status"] != "PASS"] + ok = [r for r in rs if r["status"] == "PASS"] + mb = np.median([abs(r["bias_ln"]) for r in ok]) if ok else float("nan") + mn = np.median([r["n_eff"] for r in rs if np.isfinite(r["n_eff"])]) if rs else float("nan") + print("%-10s %6g %9d %7d %7.0f%% %12.4f %10.0f" % + (kind, snr, nc, rs[0]["n_steps"], 100.0 * len(bad) / len(rs), mb, mn)) + + +if __name__ == "__main__": + main() diff --git a/demos/integrator_snr_lottery/tools/khat_validation.py b/demos/integrator_snr_lottery/tools/khat_validation.py new file mode 100644 index 000000000..bf30fa3b5 --- /dev/null +++ b/demos/integrator_snr_lottery/tools/khat_validation.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python +""" +khat_validation.py -- does the Pareto-k tail diagnostic actually CATCH the confidently-wrong runs? + +MOTIVATION. On a sharp high-SNR target we measured runs that were confidently wrong: the copy with +the HIGHEST n_eff in its arm was the MOST biased in lnZ (n_eff 58 -> 11 nats low; n_eff 123.6 -- the +highest of the whole study -- -> 10 nats low). n_eff (Kish) measures weight CONCENTRATION, not +COVERAGE, so a proposal that has missed mass looks confident. Any error estimate keyed on n_eff +inherits that, and it fails in the dangerous direction. + +`RIFT.integrators.statutils.pareto_khat_from_log` (added with the MC-error-estimate work) is the +proposed instrument: the generalized-Pareto tail index of the importance weights, with + k < 0.5 variance finite, naive sigma meaningful + 0.5-0.7 variance marginal, naive sigma optimistic + k > 0.7 tail unresolved -- naive sigma is a LOWER BOUND and the integral may be dominated by + unseen tail mass. +That is exactly the failure above, so the question is empirical: on REAL failures (not synthetic GPD +draws), does k-hat fire when n_eff does not? + +WHY THIS HARNESS CAN ANSWER IT. We reuse the merge gate's MixtureTarget, which exposes `true_lnZ`. +So every run has a KNOWN bias, and we can score k-hat as a binary classifier of "this run is wrong" +-- sensitivity, specificity, and a head-to-head against n_eff. Sampling on a real event could never +do this: there is no truth to score against. + +k-hat is computed from the SAME importance log-weights the gate uses for its shape metrics, captured +by wrapping `shape_recovery.shape_metrics` (which receives ln_wt). We do not modify the gate: today +`run_one` discards `dict_return`, so the sampler-emitted `mc_diag['pareto_khat']` never reaches the +record. + +Usage (CPU; keep jobs small -- RLIMIT_NPROC counts THREADS on this cluster): + OMP_NUM_THREADS=1 python khat_validation.py --snrs 80,160 --copies 40 --jobs 6 +""" +from __future__ import print_function +import argparse, json, os, sys, time +import numpy as np +from multiprocessing import Pool + +HERE = os.path.dirname(os.path.abspath(__file__)) +GATE = os.path.abspath(os.path.join( + HERE, "..", "..", "..", "MonteCarloMarginalizeCode", "Code", + "test", "expensive_before_merging", "integrators")) +sys.path.insert(0, GATE) +import shape_recovery as SR +from RIFT.integrators.statutils import pareto_khat_from_log + +SIGMA_REF, SNR_REF = 0.7, 20.0 +_CAP = {} + + +def _install_capture(): + """Wrap the gate's shape_metrics to stash the importance log-weights it is handed. + Idempotent, and applied inside the worker so it survives fork/spawn.""" + if getattr(SR, "_khat_capture_installed", False): + return + _orig = SR.shape_metrics + + def _wrapped(target, X, ln_wt, rng, *a, **kw): + try: + _CAP["ln_wt"] = np.asarray(ln_wt, dtype=float).ravel().copy() + except Exception: + _CAP["ln_wt"] = None + return _orig(target, X, ln_wt, rng, *a, **kw) + SR.shape_metrics = _wrapped + SR._khat_capture_installed = True + + +def _one(job): + (snr, ndim, ncomp, nmax, neff, n_chunk, seed, kind) = job + _install_capture() + _CAP.pop("ln_wt", None) + sigma = SIGMA_REF * (SNR_REF / float(snr)) + try: + target = SR.MixtureTarget(ndim, ncomp, seed, sigma_1d=sigma) + rec = SR.run_one(kind, target, nmax, neff, n_chunk=n_chunk, seed=seed) + status = SR.evaluate(rec) + status = status if isinstance(status, str) else status[0] + except Exception as e: + return dict(snr=snr, seed=seed, kind=kind, status="ERROR", err=str(e)[:120], + bias=float("nan"), n_eff=float("nan"), khat=None) + lw = _CAP.get("ln_wt") + khat = None + if lw is not None and len(lw): + try: + khat = pareto_khat_from_log(lw) + except Exception: + khat = None + return dict(snr=snr, seed=seed, kind=kind, status=status, + bias=float(rec.get("bias_ln", float("nan"))), + n_eff=float(rec.get("n_eff", float("nan"))), + khat=(float(khat) if khat is not None else None)) + + +def score(recs, bias_tol, khat_cut, neff_cut): + """Score k-hat and n_eff as binary detectors of 'this run is materially wrong'.""" + use = [r for r in recs if np.isfinite(r["bias"]) and r["khat"] is not None] + wrong = [r for r in use if abs(r["bias"]) > bias_tol] + right = [r for r in use if abs(r["bias"]) <= bias_tol] + def rate(rs, pred): + return (100.0 * sum(1 for r in rs if pred(r)) / len(rs)) if rs else float("nan") + k_flag = lambda r: r["khat"] > khat_cut + n_flag = lambda r: r["n_eff"] < neff_cut + print("\nscored on %d runs with finite bias and a k-hat (%d wrong, %d accurate; |bias|>%.2f = wrong)" + % (len(use), len(wrong), len(right), bias_tol)) + print(" %-28s %12s %12s" % ("detector", "sensitivity", "false alarm")) + print(" %-28s %11.0f%% %11.0f%%" % ("k-hat > %.2f" % khat_cut, rate(wrong, k_flag), rate(right, k_flag))) + print(" %-28s %11.0f%% %11.0f%%" % ("n_eff < %g" % neff_cut, rate(wrong, n_flag), rate(right, n_flag))) + # the decisive subset: runs n_eff would have PASSED but that are actually wrong + sneaky = [r for r in wrong if r["n_eff"] >= neff_cut] + if sneaky: + caught = sum(1 for r in sneaky if k_flag(r)) + print(" CONFIDENTLY WRONG (n_eff>=%g yet |bias|>%.2f): %d runs, k-hat catches %d (%.0f%%)" + % (neff_cut, bias_tol, len(sneaky), caught, 100.0 * caught / len(sneaky))) + print(" their k-hat: %s" % " ".join("%.2f" % r["khat"] for r in sorted(sneaky, key=lambda r: -abs(r["bias"]))[:12])) + print(" their bias : %s" % " ".join("%+.2f" % r["bias"] for r in sorted(sneaky, key=lambda r: -abs(r["bias"]))[:12])) + else: + print(" (no confidently-wrong runs in this sample -- raise SNR or copies)") + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--snrs", default="80,160") + ap.add_argument("--ndim", type=int, default=4) + ap.add_argument("--ncomp", type=int, default=3) + ap.add_argument("--nmax", type=int, default=2000000) + ap.add_argument("--neff", type=int, default=3000) + ap.add_argument("--n-chunk", type=int, default=10000) + ap.add_argument("--copies", type=int, default=40) + ap.add_argument("--kinds", default="AV") + ap.add_argument("--seed0", type=int, default=7000) + ap.add_argument("--jobs", type=int, default=6) + ap.add_argument("--bias-tol", type=float, default=0.10, help="gate's lnZ tolerance") + ap.add_argument("--khat-cut", type=float, default=0.7) + ap.add_argument("--neff-cut", type=float, default=100.0) + ap.add_argument("--json", default=None) + a = ap.parse_args() + + jobs = [(float(s), a.ndim, a.ncomp, a.nmax, a.neff, a.n_chunk, a.seed0 + c, k) + for k in a.kinds.split(",") for s in a.snrs.split(",") for c in range(a.copies)] + print("# k-hat validation: %d runs (%s, SNR %s, %d copies, nmax=%d, chunk=%d)" + % (len(jobs), a.kinds, a.snrs, a.copies, a.nmax, a.n_chunk)) + sys.stdout.flush() + t0 = time.time() + pool = Pool(a.jobs); recs = pool.map(_one, jobs); pool.close(); pool.join() + print("# done in %.1f min" % ((time.time() - t0) / 60.0)) + if a.json: + json.dump(recs, open(a.json, "w"), indent=1); print("# wrote", a.json) + ok = [r for r in recs if r["khat"] is not None] + print("# k-hat available for %d/%d runs" % (len(ok), len(recs))) + score(recs, a.bias_tol, a.khat_cut, a.neff_cut) + + +if __name__ == "__main__": + main() diff --git a/demos/integrator_snr_lottery/tools/measure_chunk_memory.sh b/demos/integrator_snr_lottery/tools/measure_chunk_memory.sh new file mode 100755 index 000000000..5c4930576 --- /dev/null +++ b/demos/integrator_snr_lottery/tools/measure_chunk_memory.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# measure_chunk_memory.sh -- price the SNR-scaled chunk size, on BOTH memory surfaces. +# +# WHY TWO NUMBERS. GPU memory and host memory are DIFFERENT control surfaces, and condor's +# RequestMemory governs the HOST, not the GPU. Enlarging --n-chunk grows the device-side sample +# arrays; whether it moves host RSS at all is a separate question. Reporting only one (as an earlier +# version of this script did) cannot answer "must I raise RequestMemory?". +# +# Relevant prior: RIFT extrinsic jobs have been measured requesting 35-105x the host RAM they +# actually use. So the likely correct conclusion is "raise the chunk, leave RequestMemory alone" -- +# but that must be MEASURED, not assumed, and changing a production resource setting without evidence +# is exactly the kind of unmotivated churn to avoid. +# +# METHOD. Run the pinned on-source ILE bench per chunk size, sample every 2 s: +# host : max over the job's process TREE of VmHWM (peak RSS) from /proc//status [KiB -> MiB] +# gpu : nvidia-smi --query-compute-apps, FILTERED TO THIS JOB'S PIDs (the earlier version took the +# max over all compute apps and so reported other users' jobs -- identical 37946 MiB for +# every chunk, which is what exposed the bug) +# +# Usage: [GPU=2] [NMAX=400000] [WARM=1] measure_chunk_memory.sh [chunk ...] (default 10000 40000 160000) +set -u +HERE="$(cd "$(dirname "$0")" && pwd)" +WT="$(cd "$HERE/../../.." && pwd)" +BENCH="$WT/MonteCarloMarginalizeCode/Code/test/integrators/bench_onsource.sh" +RUNPE=~/RIFT_roboto_paper/analyses/integrator_demos/S250114ax_pipeline/run_PE/iteration_0_ile +GPU=${GPU:-2} +NMAX=${NMAX:-400000} +WARM=${WARM:-1} # warm so the run does real work rather than bailing early +CHUNKS=${@:-"10000 40000 160000"} +OUT="$HERE/../results/chunk_memory.txt" +mkdir -p "$(dirname "$OUT")" + +{ + echo "# chunk-size resource cost, REAL ILE likelihood, on-source point" + echo "# GPU=$GPU nmax=$NMAX warm=$WARM $(date)" + echo "# host_MiB = peak VmHWM over the job's process tree (what RequestMemory governs)" + echo "# gpu_MiB = peak nvidia-smi used_memory for THIS job's pids only" + printf "%-9s %10s %10s %9s %9s\n" chunk host_MiB gpu_MiB wall_s n_eff +} | tee "$OUT" + +descendants() { # pid -> pid and all descendants + local p=$1; echo "$p" + local kids; kids=$(pgrep -P "$p" 2>/dev/null) + local k; for k in $kids; do descendants "$k"; done +} + +for nc in $CHUNKS; do + name="mem_c${nc}" + t0=$(date +%s) + env GPU=$GPU WARM=$WARM NAME=$name NMAX=$NMAX NCHUNK=$nc bash "$BENCH" \ + --sampler-method portfolio --sampler-portfolio AV --sampler-portfolio GMM \ + --internal-gmm-adaptive-components --internal-gmm-max-components 8 \ + --force-adapt-all --internal-rotate-phase --seed 10 >/dev/null 2>&1 & + JOB=$! + host_peak=0; gpu_peak=0 + while kill -0 $JOB 2>/dev/null; do + pids=$(descendants $JOB 2>/dev/null | sort -u) + # host: peak RSS over the tree + for p in $pids; do + v=$(awk '/VmHWM/{print $2}' /proc/$p/status 2>/dev/null) + [ -n "${v:-}" ] && [ "$v" -gt "$host_peak" ] 2>/dev/null && host_peak=$v + done + # gpu: only rows whose pid is in our tree + while read -r gp gm; do + case " $pids " in *" $gp "*) + [ -n "$gm" ] && [ "$gm" -gt "$gpu_peak" ] 2>/dev/null && gpu_peak=$gm ;; + esac + done < <(nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader,nounits 2>/dev/null | tr -d ',') + sleep 2 + done + wait $JOB 2>/dev/null + t1=$(date +%s) + d=$RUNPE/os_${name}.xml_0_.dat + ne=$([ -e "$d" ] && awk 'END{printf "%.1f", $NF}' "$d" || echo "-") + printf "%-9s %10s %10s %9s %9s\n" "$nc" "$((host_peak/1024))" "${gpu_peak:-0}" "$((t1-t0))" "$ne" | tee -a "$OUT" +done +echo "# wrote $OUT" diff --git a/demos/integrator_snr_lottery/tools/run_chunk_study.sh b/demos/integrator_snr_lottery/tools/run_chunk_study.sh new file mode 100755 index 000000000..626218b32 --- /dev/null +++ b/demos/integrator_snr_lottery/tools/run_chunk_study.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# Chunk-size study driver. Runs BOTH designs, because they answer different questions: +# A) FIXED BUDGET -- same cost per run, so steps = nmax/n_chunk falls as the chunk grows. +# This is the PRODUCTION question: "at the budget I already pay, should the chunk be bigger?" +# B) FIXED STEPS -- nmax = n_chunk*steps, equal adaptation opportunities, cost grows with chunk. +# This is the MECHANISM question: "do richer per-step statistics help, independent of steps?" +# CPU-only and truth-known (MixtureTarget.true_lnZ), so copies are cheap and we measure real BIAS. +set -u +HERE="$(cd "$(dirname "$0")" && pwd)" +WT="$(cd "$HERE/../../.." && pwd)" +OUT="$HERE/../results" +export PATH=/home/richard.oshaughnessy/RIFT_develUWM/bin:$PATH +export PYTHONPATH="$WT/MonteCarloMarginalizeCode/Code" +export CUDA_VISIBLE_DEVICES="" OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 +JOBS=${JOBS:-32} +COPIES=${COPIES:-24} +mkdir -p "$OUT" +echo "############ A) FIXED BUDGET (same cost) ############" +python -u "$HERE/chunk_study.py" --snrs 20,40,80,160 --chunks 10000,40000,160000 \ + --nmax 2000000 --copies $COPIES --jobs $JOBS --kinds AV,portfolio \ + --json "$OUT/chunk_fixed_budget.json" 2>&1 | grep -vE "Adding parameter|Adapting|^ *[0-9]+ " +echo +echo "############ B) FIXED STEPS (isolate per-step statistics) ############" +python -u "$HERE/chunk_study.py" --snrs 20,40,80,160 --chunks 10000,40000,160000 \ + --steps 50 --copies $COPIES --jobs $JOBS --kinds AV,portfolio \ + --json "$OUT/chunk_fixed_steps.json" 2>&1 | grep -vE "Adding parameter|Adapting|^ *[0-9]+ " +echo "CHUNK STUDY DONE" diff --git a/docs/source/containers.rst b/docs/source/containers.rst index 59f5cf4ac..f6d811309 100644 --- a/docs/source/containers.rst +++ b/docs/source/containers.rst @@ -15,6 +15,9 @@ That still works exactly as before. This page documents two additions: capabilities, and let HTCondor pick the right one per matched machine; and * a **multi-target build** that produces such a family from one template. +For an operator workflow that inventories a target GPU pool and creates +container-cache warmup jobs, see :doc:`executables/survey_scan`. + .. note:: If ``SINGULARITY_RIFT_IMAGE`` is a plain ``.sif`` path or a single diff --git a/docs/source/demos.rst b/docs/source/demos.rst new file mode 100644 index 000000000..21c0307d5 --- /dev/null +++ b/docs/source/demos.rst @@ -0,0 +1,299 @@ +===== +Demos +===== + +RIFT ships several demonstration directories under +``MonteCarloMarginalizeCode/Code/demo``. Use this page as a map from the +problem you want to understand to the smallest existing example that exercises +that workflow. + +The demos are not all the same kind of artifact. Some are fast local smoke +tests, some build Condor DAGs without submitting them, and some are advanced +operator tutorials that assume LIGO cluster credentials or external software. +Each entry below states what it is useful for before pointing you at the source +files. + +.. contents:: Demo catalog + :local: + :depth: 2 + +Pipeline-builder smoke tests +============================ + +Path + ``MonteCarloMarginalizeCode/Code/demo/pipeline`` + +Use this when + You want a fast, submission-free check that ``util_RIFT_pseudo_pipe.py`` can + build a complete RIFT run directory and thread command-line options through + CEPP into the generated Condor submit files. + +What it demonstrates + The demo uses fake data and reference ``.ini`` / ``coinc.xml`` inputs to + create run directories without submitting jobs or requiring real frames, + PSDs, GPUs, or a Condor pool. It is especially useful for regression tests + of argument plumbing: a flag supplied to ``util_RIFT_pseudo_pipe.py`` should + survive into the correct ``args_*.txt`` file and ``*.sub`` file. + +Primary files + * ``README.md`` — target descriptions and expected assertions. + * ``Makefile`` — ``baseline``, ``grid``, ``slices``, ``all``, and ``clean`` + targets. + +Typical command + From the demo directory, in a configured RIFT environment:: + + make all + + or, using the repository Pixi environment from the README:: + + pixi run --manifest-path ../../../../../pixi.toml make all + +Notes + The ``grid`` and ``slices`` targets exercise last-iteration extrinsic export + behavior. They check that distance-grid or distance-slice flags land on the + final extrinsic ILE stage without leaking into the intrinsic ILE jobs. + +Zero-spin IMRPhenomD distance-grid validation +============================================ + +Path + ``MonteCarloMarginalizeCode/Code/demo/pipeline/zero_spin_phenomD`` + +Use this when + You need a compact end-to-end validation of per-distance likelihood export, + consolidation, and posterior reconstruction on a laptop-scale example. + +What it demonstrates + This demo uses zero-spin IMRPhenomD, the AV sampler, fake zero-noise BBH + inputs, and a small mass grid. It bypasses Condor for the local execution + step but uses the same production code paths for building the pipeline, + running ``integrate_likelihood_extrinsic_batchmode``, consolidating + ``.dgrid`` files, and reconstructing a joint intrinsic-plus-distance + posterior. + +Primary files + * ``README.md`` — full four-stage validation walkthrough. + * ``Makefile`` — ``build``, ``run-extr``, ``consolidate``, ``posterior``, + ``all``, and ``clean`` targets. + * ``zero_spin_phenomD.ini`` — minimal pseudo-pipe configuration. + +Typical command + From the demo directory:: + + make all + +Notes + The default settings intentionally evaluate only a few events so the test is + fast. This is a code-path validation, not a scientific accuracy benchmark. + Increase ``N_EVENTS`` if you want a more meaningful posterior check. + +HyperPipe demos +=============== + +Path + ``MonteCarloMarginalizeCode/Code/demo/hyperpipe`` + +Use this when + You want to learn HyperPipe, test generalized likelihood drivers, compare + baseline posterior resampling to tracer placement, or adapt a toy + coordinate-free workflow into a real one. + +What it demonstrates + The directory contains runnable YAML configurations for the same 3-D + Gaussian toy likelihood. The variants exercise the iterative + ``MARG -> CON -> UNIFY -> EOS_POST -> PUFF/placement -> TEST`` loop, OSG + submit-host settings, coordinate transformation, and parsimonious/tracer + placement. + +Primary files + * ``README.md`` — detailed description of every configuration. + * ``technical_doc.txt`` — pedagogical implementation notes. + * ``hyperpipe_conf.yaml`` — baseline posterior-resampling workflow. + * ``hyperpipe_conf_tracer.yaml`` — tracer/parsimonious-placement workflow. + * ``hyperpipe_conf_osg.yaml`` — OSG/IGWN-oriented submit configuration. + * ``hyperpipe_conf_linear_uvw.yaml`` — fit in transformed coordinates while + sampling in the original coordinates. + * ``example_gaussian*.py`` — toy likelihood drivers. + * ``Makefile`` — convenience targets such as ``rundir`` and + ``rundir_tracer``. + +Typical commands + Baseline demo:: + + util_RIFT_hyperpipe.py --config ./hyperpipe_conf.yaml + + Tracer-placement demo:: + + util_RIFT_hyperpipe.py --config ./hyperpipe_conf_tracer.yaml + +Notes + Start here before writing a new HyperPipe configuration from scratch. The + YAML files show the expected schema and the generated run directories expose + the exact executable arguments in ``args_*.txt`` and Condor ``*.sub`` files. + +Population-study demo +===================== + +Path + ``MonteCarloMarginalizeCode/Code/demo/populations`` + +Use this when + You want a worked outline for generating mock compact-binary populations + with GWKokab and producing RIFT parameter estimates for those injections. + +What it demonstrates + The README describes a multi-environment workflow: generate injections with + GWKokab, validate the population inference setup, switch to a separate RIFT + environment, prepare injections, generate MDC files, create RIFT run + directories, submit PE jobs, and produce diagnostics. + +Primary files + * ``README.md`` — full tutorial and environment notes. + * ``Makefile`` — workflow automation points. + * ``pop-example.ini`` — example RIFT configuration for the population run. + * ``injections.dat`` — example injection table. + * ``write_mdc.py`` and ``gwk_pop_conversion.py`` — conversion/setup helpers. + * ``plot_all.sh`` and ``collect_all.sh`` — post-processing helpers. + +Typical command + This is an advanced, environment-dependent workflow. Read and edit the + Makefile variables and ``pop-example.ini`` before running targets. The + README starts with GWKokab setup and then moves into RIFT setup. + +Notes + Keep GWKokab and RIFT in separate environments. The prior ranges in + ``pop-example.ini`` must match the population used to generate + ``injections.dat``; otherwise the resulting PE runs are not meaningful. + +Distance-grid export demo +========================= + +Path + ``MonteCarloMarginalizeCode/Code/demo/rift/add_distance_grids`` + +Use this when + You need to understand or validate the ``--export-marginal-distance-grid`` + path and the generated ``.dgrid`` likelihood-density files. + +What it demonstrates + The demo builds a small zero-spin RIFT workflow with distance-grid export + enabled for ILE jobs. It reuses fake zero-noise CI assets and verifies that + the generated ILE arguments include ``--export-marginal-distance-grid`` and + ``--internal-use-lnL``. + +Primary files + * ``README.md`` — build/submit instructions and environment warning notes. + * ``PLAN_B_DESIGN.md`` — design notes for fixed-distance slice export and + re-marginalization. + * ``Makefile`` — ``dag`` and ``submit`` targets. + * ``add_distance_grids.ini`` — zero-spin/fake-data configuration. + * ``validate_distance_grid.py`` and ``validate_distance_slices.py`` — helper + validation scripts. + +Typical commands + Build the DAG without submitting:: + + make dag + + Submit the generated workflow after inspection:: + + make submit + +Notes + LALSuite/SWIG compatibility warnings may appear in some environments. The + README explains how to distinguish those warnings from distance-grid + failures. + +Numerical relativity with RIFT +============================== + +Path + ``MonteCarloMarginalizeCode/Code/demo/nr_w_rift`` + +Use this when + You want an advanced tutorial for comparing gravitational-wave data to + numerical-relativity simulations rather than analytic waveform models. + +What it demonstrates + The workflow obtains event data, constructs an NR simulation grid, builds a + RIFT run directory through NR-specific pipeline tools, and runs a refine + stage before the final CIP posterior construction. + +Primary files + * ``README.md`` — cluster-oriented tutorial and required manual settings. + * ``Makefile`` — data, grid, and run-directory construction targets. + +Typical commands + This tutorial assumes LIGO computing access and event-specific manual edits. + Read the README first, then configure the event identifiers, channels, NR + group, mass range, and event time in the Makefile before running targets such + as ``make data``, ``make grid``, and ``make rundir``. + +Notes + This is not a quickstart. It assumes a working RIFT environment, LIGO data + access, NR catalog access, and familiarity with production RIFT runs. + +Internal and test-oriented demos +================================ + +Some demo-like directories are primarily regression or development harnesses. +They are useful for developers, but should not be presented as first-stop user +quickstarts until their assumptions are documented. + +Known examples include: + +* ``MonteCarloMarginalizeCode/Code/demo/rift/test_frameworks/zero_likelihood`` + — zero-likelihood HyperPipe/Condor smoke-test material. + +When promoting one of these into a user-facing tutorial, first document: + +* whether it requires Condor, GPUs, GraceDB, LIGO credentials, or external data; +* whether it submits jobs or only builds run directories; +* expected runtime and expected outputs; +* cleanup commands; and +* which scientific result, if any, should be trusted. + +Choosing a starting point +========================= + +.. list-table:: Demo selection guide + :header-rows: 1 + :widths: 24 34 42 + + * - If you want to... + - Start with... + - Why + * - Check pseudo-pipe argument plumbing quickly + - ``demo/pipeline`` + - Fast fake-data DAG construction without submission. + * - Validate distance-grid export end to end + - ``demo/pipeline/zero_spin_phenomD`` + - Exercises build, extrinsic likelihood, consolidation, and posterior + reconstruction. + * - Learn HyperPipe + - ``demo/hyperpipe`` + - Small Gaussian likelihood with baseline, tracer, OSG, and coordinate + transform variants. + * - Try tracer placement + - ``demo/hyperpipe/hyperpipe_conf_tracer.yaml`` + - Existing parsimonious-placement example with generated DAG output. + * - Explore population-study workflows + - ``demo/populations`` + - End-to-end GWKokab-to-RIFT outline, with environment caveats. + * - Inspect distance-grid export internals + - ``demo/rift/add_distance_grids`` + - Focused DAG build and validation helpers for ``.dgrid`` output. + * - Work with numerical-relativity simulations + - ``demo/nr_w_rift`` + - Advanced cluster-oriented NR workflow. + +Related pages +============= + +* :doc:`hyperpipe` +* :doc:`using-pipeline` +* :doc:`examples-ini` +* :doc:`osg` +* :doc:`plotting` +* :doc:`troubleshooting` diff --git a/docs/source/executables/index.rst b/docs/source/executables/index.rst index c2941a579..cb6f19a4a 100644 --- a/docs/source/executables/index.rst +++ b/docs/source/executables/index.rst @@ -16,6 +16,7 @@ This section documents the core user-facing command-line executables in RIFT. util_ManualOverlapGrid convergence_test_samples util_ParameterPuffball + survey_scan Core Executables Overview ========================= @@ -28,4 +29,4 @@ The following executables form the core user interface for RIFT: Related Documentation ==================== -- :doc:`../api_reference/index` - API Reference \ No newline at end of file +- :doc:`../api_reference/index` - API Reference diff --git a/docs/source/executables/survey_scan.rst b/docs/source/executables/survey_scan.rst new file mode 100644 index 000000000..6761c9117 --- /dev/null +++ b/docs/source/executables/survey_scan.rst @@ -0,0 +1,123 @@ +############### +``survey_scan`` +############### + +``containers/survey_scan.sh`` is an operator-facing companion for RIFT +container families. It surveys a target HTCondor GPU pool, generates one +warmup job for each selected container/profile combination, and summarizes the +JSON reports returned by completed jobs. It does not run an analysis or submit +the generated jobs itself. + +For the container-family manifest and deployment model, see :doc:`../containers`. + +Prerequisites and boundaries +============================= + +The submit-side commands use Python's standard library. Manifest parsing uses +PyYAML when it is installed and otherwise supports the simple RIFT +container-family YAML schema. ``survey`` needs ``condor_status`` on the host. +To execute the jobs, the target environment needs HTCondor, a compatible GPU, +Apptainer, and an image containing the requested CuPy or JAX dependencies. + +This command is an operator-run pool inventory and cache-warmup workflow. The +reference documents the generated workflow; it is not evidence that a given +pool, container runtime, or image has been exercised successfully. + +Survey a pool +============= + +Run ``survey`` before selecting image bands or generating jobs:: + + containers/survey_scan.sh survey \ + --out survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml + +The exact interface is:: + + containers/survey_scan.sh survey [--out DIR] [--constraint EXPR] [--manifest FILE] + +``--constraint`` is passed to ``condor_status`` and defaults to +``TotalGPUs > 0``. ``--manifest`` is optional; when supplied, the inventory +also records which manifest labels cover each observed capability. If ``--out`` +is omitted, the command creates a timestamped directory under ``survey/``. + +The survey directory contains: + +* ``gpu_inventory.json`` — raw ClassAd fields, normalized summary, and optional + manifest coverage; +* ``gpu_inventory.tsv`` — the summarized slot/device/capability/memory table; +* ``recommended_matrix.json`` — suggested capability bands; and +* ``coverage.md`` — a readable inventory, suggested bands, and manifest + coverage when a manifest was supplied. + +Generate warmup jobs +==================== + +Generate Condor submit files from a survey directory and a container-family +manifest:: + + containers/survey_scan.sh emit-jobs \ + --survey survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml + +The exact interface is:: + + containers/survey_scan.sh emit-jobs --survey DIR --manifest FILE [--out DIR] [--profiles LIST] [--request-disk REQUEST_DISK] + +``--survey`` and ``--manifest`` are required. ``--out`` defaults to +``DIR/jobs``; ``--profiles`` is a comma-separated list and defaults to +``cupy``; and ``--request-disk`` defaults to ``16000M``. Supported profile +names are ``cupy`` and ``jax``. The command uses ``--survey`` only to choose +that default output location: it emits every manifest entry times every +selected profile, rather than selecting jobs from the survey inventory. +Consequently, a manifest band absent from the surveyed pool can still yield a +submit file that remains unmatched. For example, request both profiles only +for a JAX-enabled image:: + + containers/survey_scan.sh emit-jobs \ + --survey survey/cit-YYYYMMDD \ + --manifest container_family/rift_container_family.generated.yaml \ + --profiles cupy,jax + +The ``cupy`` profile warms common NoLoop and fused-calmarg kernels. The +``jax`` profile warms synthetic JAX ILE-wrapper shapes. The generated directory +has a ``.sub`` and executable ``run_*.sh`` wrapper for each selected +container/profile pair, copied profile scripts, and ``submit_all.sh``. Submit +them deliberately from that directory:: + + cd survey/cit-YYYYMMDD/jobs + ./submit_all.sh + +Each wrapper sets cache locations such as ``CUPY_CACHE_DIR`` and +``JAX_COMPILATION_CACHE_DIR`` before running ``apptainer exec --nv``. For an +``osdf://`` image URL, it fetches only that selected image, using ``stashcp`` or +``pelican``; one of those tools must therefore be available on the execute +node. Size ``--request-disk`` for the selected image and its work area. + +Collect results +=============== + +After jobs have returned their JSON outputs to the jobs directory, collect a +single summary:: + + containers/survey_scan.sh collect --survey survey/cit-YYYYMMDD + +The exact interface is:: + + containers/survey_scan.sh collect --survey DIR [--out FILE] + +``--survey`` is required. By default, the command writes +``warmup_summary.json`` and its Markdown counterpart +``warmup_summary.md`` in the survey directory. ``--out`` selects a different +JSON summary path; the Markdown report uses the same basename with a ``.md`` +suffix. The collector globs only JSON files already present under +``jobs/*.json``; it cannot identify jobs that never produced a result. Invalid +present JSON is retained as an error entry so the summary can report malformed +output. + +See also +======== + +* :doc:`../containers` for building and deploying a container family. +* ``containers/survey_scan/README.md`` in the source tree for a concise + operator-oriented overview of the warmup profiles. diff --git a/docs/source/hyperpipe.rst b/docs/source/hyperpipe.rst index 1313dfa22..acab9ac9c 100644 --- a/docs/source/hyperpipe.rst +++ b/docs/source/hyperpipe.rst @@ -19,6 +19,8 @@ Depending on your goal, choose the appropriate guide below: and legacy support. * **Troubleshooting**: See :doc:`hyperpipe/troubleshooting` for diagnostic commands and RIFT-specific caveats. +* **Distance grids**: See :doc:`hyperpipe/add_distance_grids` to export and + validate per-point luminosity-distance likelihood grids. .. toctree:: :maxdepth: 2 @@ -29,3 +31,4 @@ Depending on your goal, choose the appropriate guide below: hyperpipe/driver_dev hyperpipe/reference hyperpipe/troubleshooting + hyperpipe/add_distance_grids diff --git a/docs/source/hyperpipe/add_distance_grids.rst b/docs/source/hyperpipe/add_distance_grids.rst new file mode 100644 index 000000000..968b35ad0 --- /dev/null +++ b/docs/source/hyperpipe/add_distance_grids.rst @@ -0,0 +1,100 @@ +=========================== +``add_distance_grids`` demo +=========================== + +The ``add_distance_grids`` demo provides a small zero-spin RIFT workflow to +configure for luminosity-distance likelihood-grid export. Use it to verify the +distance-grid export path before adapting the same options to another RIFT +workflow. + +The runnable sources are in +``MonteCarloMarginalizeCode/Code/demo/rift/add_distance_grids``. The demo uses +the fake-data inputs in ``.travis/ILE-GPU-Paper/demos`` and its +``add_distance_grids.ini`` records the corresponding configuration. + +Generate and configure the DAG +============================== + +From the demo directory, first check that the CI-style inputs are present: + +.. code-block:: console + + $ make inputs + +Then generate the baseline DAG: + +.. code-block:: console + + $ make dag + +.. warning:: + + ``make dag`` removes and recreates ``rundir/`` before generating the DAG. + Copy any results you need from that directory before rerunning it. + +The current ``make dag`` recipe is a baseline generator, not a completed +distance-grid run: its ``validate-args`` check expects an export flag that the +recipe does not put in ``rundir/args_ile.txt``. Consequently, stock ``make +dag`` can stop at that check and does not by itself create a dgrid-producing +``ILE_extr`` stage. + +Before submitting, configure the pipeline's final extrinsic ILE stage and its +distance-grid export option. For ``create_event_parameter_pipeline_BasicIteration``, +this means enabling ``--last-iteration-extrinsic`` together with +``--last-iteration-export-marginal-distance-grid`` (and supplying the required +pipeline configuration, including compatible conversion arguments). Inspect +the resulting ``ILE_extr`` submit arguments to confirm both +``--export-marginal-distance-grid`` and ``--internal-use-lnL`` are present. +Submit only that correctly configured DAG; the demo never submits +automatically. + +Validate an output grid +======================= + +After a correctly configured ILE-extrinsic run completes, locate a producer +output under the completed run results. Each ILE evaluation writes +``__.dgrid``. The exact ILE-output prefix and +result-directory nesting are assigned by the DAG, so from the generated run +directory discover completed outputs with: + +.. code-block:: console + + $ find rundir -type f -name '*_*.dgrid' -print + +Use one returned path with the loader before reconstruction: the +``reconstruct_marginal_lnL`` API accepts the parsed grid table, not a filename. +For example: + +.. code-block:: python + + from RIFT.misc.distance_grid import ( + load_distance_grid, + reconstruct_marginal_lnL, + ) + + grid = load_distance_grid("") + reconstructed = reconstruct_marginal_lnL(grid) + +With the default argument, reconstruction uses the stored sampling-distance +prior when it is present. Compare the result to the ordinary marginalized +likelihood for the same intrinsic point, allowing for that run's Monte Carlo +uncertainty. For a controlled synthetic check, run +``validate_distance_grid.py`` from the demo directory; it exercises the +Plan-A ``.dgrid`` table-level reconstruction path directly. The neighboring +``validate_distance_slices.py`` script covers the separate Plan-B ``.dslice`` +behavior and is out of scope for this demo. + +Environment note +================ + +LALSuite SWIG/Python memory-leak messages can indicate an incompatible local +binding build rather than a distance-grid failure. The demo README describes a +known-good environment constraint: changing the local SWIG executable does not +alter already-built LAL Python bindings. + +Review checklist +================ + +Before requesting human review, verify that this page is reachable from the +HyperPipe landing-page toctree, retain the Sphinx error-delta result, and have +an independent reviewer check the rendered guide against the demo sources. diff --git a/docs/source/index.rst b/docs/source/index.rst index d396c8f61..91f567304 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -24,6 +24,7 @@ Rapid inference via Iterative FiTting: this algorithm provides a framework for e containers injections plotting + demos hyperpipe troubleshooting api_samples_utils