GMM: flexible (data-driven) component allocation for the extrinsic integrator

[oshaughnessy-junior]

Replaces the hard-coded per-group GMM component layout (sky=4, dist-incl=2,
phase=wide) with an OPT-IN, data-driven allocation, and fixes two bugs that made a
warm GMM proposal a no-op. Hard-coding remains the default/fallback.

Changes

• O(k³) Hungarian component matching (was O(k!) permutations) — exact same
  optimum, unblocks many-component proposals (k=16: ~6 ms vs infeasible).
• GMM.fit_gmm_adaptive — pick a group's component count by weighted BIC over
  a k-ladder + prune dead components. BIC-at-init, then the stable merge path adapts
  (per-chunk refit-fresh was tried and rejected — the proposal wanders).
• Safety floor — BIC chooses k in [k_min, k_max] with k_min = the stress-tested
  hard-coded count; pruning never drops below it. Opting in can only add
  components (protects broad multi-modal sky posteriors). Fully opt-in:
  --internal-gmm-adaptive-components (default OFF) ⇒ byte-identical hard-coded path.
• Warm-start survival fix — integrate() was silently discarding the models
  fitted by bootstrap_from_samples; now transfers them by dim-group key.
• Portfolio — the flag forwards a scalar cap to the GMM member (honored in
  update_sampling_prior, floored at its n_comp).
• Flags: --internal-gmm-adaptive-components [--internal-gmm-max-components N]
[--internal-gmm-defensive-frac F] [--internal-gmm-inflate X].
• Tests test/integrators/test_gmm_adaptive.py (6/6) + data-free synthetic
  benchmark + RIFT/integrators/DESIGN_flexible_gmm.md.

Key measured finding

Pure GMM importance sampling stalls at n_eff≈1–7 on S250114ax (SNR≈82) regardless
of allocation, warm start, defensive coverage, or inflation (confirmed by ~12 runs

• 5 prior pure-GMM logs); the exp(~1210) dynamic range makes the honest per-chunk
  ESS ≈ 1. AV reaches ≈89 there. On a moderate-SNR synthetic (where importance
  sampling is viable) the flexible allocation is validated: robust, unbiased,
  hands-free. → Use the flexible GMM as a portfolio member with AV, not standalone on
  high-SNR degenerate events. See DESIGN_flexible_gmm.md for the full n_eff-vs-N
  tables.

Base / merge note

Based on rift_O4d_adaptive_integrator before 992486cf (--internal-reparam-dl-incl);
that commit touches the driver in a different region, so this should merge cleanly.

🤖 Generated with Claude Code

[oshaughnessy-junior]

Coordination from the portfolio freeze-policy work (#28, stacked on #26).

Two things there directly unblock the default-vs-adaptive-GMM-component comparison, which was confounded by the AV freeze-out:

1. AV is no longer starved. VARAHA/AV members are now freeze-exempt by default (portfolio_varaha_never_freeze), so in an AV+GMM portfolio the AV member contracts every chunk and becomes the workhorse (measured: its balance weight climbs 0.5 → ~0.65–0.72 on typical events; S250114ax cold n_eff 3.4 → 53). With AV carrying its share, the GMM member's component-allocation effect is no longer masked.

2. The portfolio was silently unusable in the production container. Plugin discovery hard-loaded the NF plugin (import torch), which the cuda118 container doesn't ship, so import mcsamplerPortfolio raised → mcsampler_Portfolio_ok=False → every portfolio run died with a NameError. portfolio: VARAHA never-freeze default + freeze/allocation/clipping policy (gate-validated) #28 wraps plugin loading in try/except (skip unusable plugins). Any container-based GMM-in-portfolio comparison needs this fix to run at all.

So the default-vs-adaptive-component isolation is now a re-run away on top of #28. Harness in test/integrators/bench_multi_event.py (runs a single ILE worker in the event's container with this branch's integrator on PYTHONPATH, swapping only --sampler-method) may be reusable for filling the flexible-GMM table. Happy to run it once #28 + your branch are combined.

[oshaughnessy-junior]

Shape-recovery validation record (pre-merge gate + opt-in feature probe)

1. Differential shape-recovery gate vs PR #26 (safety)

Full expensive_before_merging integrator shape-recovery suite (seeded 3-component mixtures, self-calibrating JS floor, differential PR gate) run on this branch (ccfb087c) vs its base PR #26 tip (8ae375de):

• 0 blocking regressions across the suite.
• Default path (feature OFF) is bitwise-identical to PR Improved adaptive extrinsic integrators (bootstrap / oracle / flow reuse) #26: identical lnI on all 24 targets — the opt-in switch is genuinely inert when not requested.
• Full-suite artifacts: RIFT_roboto_paper repo, analyses/integrator_shape_gate/: pr26_adaptive_integrator.json, pr27_gmm_flexible.json, compare_pr26_vs_pr27.txt.

2. Targeted probe of the opt-in feature (--internal-gmm-adaptive-components → gmm_adaptive kwarg → GMM.fit_gmm_adaptive)

Scenario the feature exists for: the hard-coded per-group component count is wrong (fixed n_comp=1 against a 3-component truth). GMM sampler, correlate_all_dims, nmax=800k, neff=3000, n_chunk=10k, seeded (MixtureTarget(d, 3, 101), run seed 987654). Probe scripts + per-run JSON: analyses/integrator_shape_gate/pr27_optin_probe/.

run	checkout	n_comp	adaptive	n_eff	n_ESS	lnZ bias	max per-dim JS (floor)	max abs mean pull	max width dev	final k
d2 nc1 OFF	PR26	1	off	3064.6	26920	+0.0036	4.9e-5 (3.4e-4)	0.0031	0.34%	1
d2 nc1 OFF	PR27	1	off	3064.6	26920	+0.0036	4.9e-5 (3.4e-4)	0.0031	0.34%	1
d2 nc1 ON	PR27	1	k_max=8	3045.1	26555	-0.0028	6.2e-5 (3.4e-4)	0.0034	0.44%	2 (BIC)
d2 nc3 ctrl	PR27	3	off	3066.5	27239	+0.0042	4.3e-5 (3.5e-4)	0.0010	0.26%	3
d4 nc1 OFF	PR26	1	off	529.7	6149	-0.0079	3.9e-4 (1.2e-3)	0.0169	0.69%	1
d4 nc1 OFF	PR27	1	off	529.7	6149	-0.0079	3.9e-4 (1.2e-3)	0.0169	0.69%	1
d4 nc1 ON	PR27	1	k_max=8	527.0	5960	-0.0131	4.5e-4 (1.3e-3)	0.0198	0.98%	3 (BIC)
d4 nc3 ctrl	PR27	3	off	541.3	5944	+0.0137	8.0e-4 (1.3e-3)	0.0199	1.14%	3

Readings:

• PR26 vs PR27 with feature OFF: identical to full precision (every metric), independently reconfirming the inert default path.
• Feature ON is statistically indistinguishable from OFF and breaks nothing: d2 n_eff 3065→3045, bias +0.004→−0.003; d4 n_eff 530→527, bias −0.008→−0.013. All per-dim JS below the matched-n_ESS floor, mean pulls < 0.02σ, widths within ~1%, zero training errors, zero crashes.
• BIC recovers real structure from a wrong floor: with k_min=1 it allocated k=2 (d2) / k=3 (d4, the true count) — the safety-floor + BIC design behaves as documented.
• The default seed-101 targets are forgiving of wrong n_comp (components overlap; a single tempered Gaussian still covers the support), so the matrix above demonstrates safety and correct component allocation. The rescue itself is demonstrated by the harsh pair below.

Harsh-separation pair (the rescue scenario)

Same d4 3-component target family but well-separated narrow modes (sigma_1d=0.35, scale_x0=9.0), wrong n_comp=1, same budget:

run	n_eff	n_ESS	lnZ bias	max per-dim JS (floor)	max abs mean pull	corr diff	final k
OFF	18.5	116	+0.087	1.6e-2 (5.0e-2)	0.105	0.077	1
ON (k_max=8)	85.8	743	-0.037	4.6e-3 (9.1e-3)	0.047	0.035	4 (BIC)

Feature ON rescues the wrong-n_comp collapse: 4.6x n_eff, 6.4x n_ESS, lnZ |bias| halved, mean pulls 0.105σ → 0.047σ, with BIC allocating 4 components against the 3-mode truth (3 modes + tail capacity). Zero training errors, zero crashes.

• All probes ran CPU-mode with the runtime eigvals workaround for the GPU-port import bug in gaussian_mixture_model.py (_xpy_eigvals/_xpy_eig bound to cupy.linalg at import with no device probe; bisected separately). That bug is now fixed on rift_O4d by PR integrators: probe for a CUDA device before selecting cupy; fail loudly when GMM refit never succeeds #32 (CPU device probe), so the workaround is no longer required on current heads (it is harmless where applied); these probes predate the fix and needed it to measure anything but uniform sampling.

3. Coverage summary

The gate suite exercised the true-n_comp default path (safety: bitwise-identical to PR #26); this probe exercised the wrong-n_comp opt-in path (the feature's intended scenario). Both dimensions of the change are tested.

Suite reference: MonteCarloMarginalizeCode/Code/test/expensive_before_merging/integrators (merged into rift_O4d via PR #31).

🤖 Generated with Claude Code