DRAFT: jax ILE -- make --save-samples a fair draw, not the raw sampler cloud

[oshaughnessy-junior]

DRAFT — opened for review, not for merge as-is.

What this fixes

bin/integrate_likelihood_extrinsic_jax --save-samples writes whatever cloud the
chosen --mode produced, with no weight column. Every consumer of a
*_samples.dat therefore reads it as an equal-weight draw from the conditional
extrinsic posterior — the contract production ILE's
--fairdraw-extrinsic-output actually provides, by multinomial-resampling
against w = L p / p_s inside RIFT/integrators/mcsampler.py::integrate (and the
identical block in mcsamplerGPU / mcsamplerAdaptiveVolume /
mcsamplerEnsemble / mcsamplerPortfolio).

For several JAX modes that contract was not met, and the correcting weights were
computed and then thrown away:

mode	what theta was at export	weights available?
laplace-is (the default)	draws from the adaptive Gaussian proposal	yes — logw computed in run_laplace_is, discarded
prior-mc	draws from the prior	yes — w = L
flowmc*	chain states at inv_T = --adapt-weight-exponent	yes — post_weight = L^(1-inv_T), returned by samplers.flowmc_sample*, discarded
nuts, nuts-phimarg, multistart-nuts	MCMC chain on the exact target	n/a (correctly nothing to do)

post_weight is uniform only at the default beta = 1. At any other value
the export was a silently tempered — i.e. over-broad — cloud.

The change

Each estimator now returns its per-sample log importance weight (None when the
sampler already targets the posterior), and write_samples fair-draws against it
before writing. The file format is unchanged — equal-weight rows, no weight
column, same header — so downstream tooling is untouched.

Following ILE, the export is capped at 1.5*ESS. Without that a low-ESS cloud is
resampled up to its original length and the file looks like N independent draws
while holding ~ESS distinct points.

--fairdraw-extrinsic-output, --fairdraw-extrinsic-output-n-max and
--n-fairdraw-extrinsic-samples move from the accepted-but-ignored list to
real, typed options that bound the count. --adapt-weight-exponent was also being
reported as ignored while the flowMC modes were in fact using it as the
tempering exponent (samplers.flowmc_sample*: inv_T = 1/temper = beta); that
misreport is fixed.

VERIFIED BY RUNNING CODE

Zero-noise BNS injection (IMRPhenomD, H1L1V1, m1=1.58 m2=1.26, lnL_max≈266),
--mode flowmc-phimarg --distance-marginalization. Ground truth is an
independent defensive importance-sampling reference built on the same
captured JAX likelihood object in the same process: proposal
q = 0.30*prior + 0.70*(400-kernel mixture), exact weights w = L p / q,
multinomial fair draw. Its recovered widths are stable across a 4x proposal
bandwidth range (0.3 / 0.6 / 1.2 → ESS 2700 / 17420 / 16483), so they are not a
proposal artifact.

1. post_weight is exactly uniform at the default beta = 1.0 (min == max == 1/4800, temper = 1.0000) — so the default flowMC path was not
   exporting a mis-weighted cloud, and this change is a no-op there.
   Confirmed end-to-end: patched and unpatched drivers at beta=1.0 produced
   identical exports (4800 rows; sd_ra 0.01407, sd_dec 0.04501, sd_psi
   0.81861, sd_incl 0.25316 in both).

2. At beta = 0.5 the discarded weights were real, and the fix corrects a real
   bias (same event/seed/data, patched vs unpatched, both against the
   reference):

   	rows	sd_ra / ref	sd_dec / ref	sd_psi / ref	sd_incl / ref
   unpatched beta=1.0	4800	0.91	0.86	0.91	1.02
   unpatched beta=0.5	4800	1.40	1.26	0.95	1.08
   patched beta=0.5	2914	0.99	0.93	0.93	0.85

   i.e. tempering silently inflated the exported sky posterior by ~1.5x relative
   to the correctly-drawn arm; the fair draw (ESS 1942 → 2914 rows) removes it.

3. laplace-is, the default mode, is unusable on this posterior regardless.
   200000 samples gave ESS = 97; the exported cloud is 3–60x too narrow
   (sd/ref: ra 0.18, dec 0.37, psi 0.016, incl 0.081) both before and after the
   fix. The fair draw cannot rescue a proposal that missed. This is why the
   1.5*ESS cap and the loud ESS < 200 warning are part of the change, and it
   is a separate bug worth its own issue.

4. Tests: 6 pass under /cvmfs/software.igwn.org/conda/envs/igwn pytest, and
   all 6 were mutation-tested — each was shown to FAIL when the corresponding
   behaviour is reverted:

   • if logw is not None ... → if False (pre-fix export):
     test_export_is_a_fair_draw_of_the_posterior and
     test_fairdraw_count_options_are_live FAIL.
   • remove the uniform-weight early return (always resample):
     test_uniform_weights_are_a_no_op and
     test_tempered_flowmc_weights_are_not_uniform FAIL.
   • remove the 1.5*ESS cap: test_ess_clamp_prevents_manufactured_draws FAILS.
     test_unreweighted_export_would_fail_the_above is a control scored the same
     way as the headline test, so the headline test cannot pass vacuously.

INFERRED, NOT VERIFIED HERE

• The behaviour of --mode prior-mc, flowmc (5-D), flowmc-dpsimarg,
  flowmc-phipsimarg, nuts, nuts-phimarg, multistart-nuts and the
  smc_puffball path under this change is inferred from reading their return
  values. Only flowmc-phimarg and laplace-is were exercised end-to-end.
• smc_puffball_sample returns a uniform post_weight after its own internal
  multinomial resample, so it should be a no-op; not run.
• No claim is made that this change fixes flowMC sampling quality. It does not.

SEPARATE PROBLEM THIS DID NOT FIX (please read)

Even at beta = 1 with uniform weights, the flowmc-phimarg export is measurably
narrower than the reference, and flow re-use across an intrinsic batch makes it
far worse. Measured against the defensive-IS reference on the two actual
production files from a 96-draw batch run as twelve 8-event chunks (flow state
threaded across events unless --no-flow-reuse), reference ESS 22929 and 12788:

production file	sd_ra/ref	sd_dec/ref	sd_psi/ref	sd_incl/ref	incl mean offset
EXT-0_0 (slot 0, fresh flow)	0.93	0.86	0.97	0.90	-0.02 sd
EXT-0_7 (slot 7, flow reused 7x)	0.93	1.00	0.21	0.65	-0.91 sd

The per-slot mean over the 12 intrinsic draws at each slot (the draws are
exchangeable across slots, so this is a pipeline effect, not a physics one) drifts
monotonically:

slot     ra       dec      incl      psi
0     2.09879   0.30243   0.57955   1.58801   <- fresh flow; matches reference 0.551
1     2.09972   0.29867   0.40706   1.52782
4     2.09949   0.29980   0.32189   1.90934
6     2.09907   0.30164   0.31953   1.94863
7     2.09877   0.30242   0.34652   1.90696

Sky (ra, dec) is essentially untouched; inclination and psi are not. That is a
flow-reuse pathology, not a weighting one, and it needs its own fix. Until then,
--no-flow-reuse looks mandatory for any run whose extrinsic samples are used.