Add STROOPWAFEL Adaptive Importance Sampling

[TomWagg]

Okay this is a big one, finally done with the STROOPWAFL implementation.

This PR adds a vectorised implementation of the STROOPWAFEL adaptive importance sampling algorithm (Broekgaarden et al. 2019) for efficiently sampling rare binary-evolution outcomes (e.g. merging BH–BH / NS–NS), where flat Monte Carlo would need millions of evolutions to collect a usable sample.

The sampler explores parameter space, fits a Gaussian-mixture proposal to discovered "hits", concentrates sampling there, and returns importance-weighted results that stay unbiased estimators of the true prior-weighted population. New code lives under cosmic.sample.stroopwafel, with the result/checkpoint types in cosmic.output.

Usage

from cosmic.sample.stroopwafel import AdaptiveSampler, ParameterSpace, Parameter
from cosmic.sample.stroopwafel.presets import any_dco

params = ParameterSpace([
    Parameter('mass_1',      5.0,        150.0,      dist='kroupa'),
    Parameter('q',           0.01,       1.0,        dist='uniform'),
    Parameter('porb',        10**(0.15), 10**(5.5),  dist='sana'),
    Parameter('ecc',         1e-9,       0.9999,     dist='sana_ecc'),
    Parameter('metallicity', 0.0001,     0.03,       dist='flat_in_log'),
    Parameter('natal_kick_1', 0.0, 5000.0, dist='disberg')
])

sampler = AdaptiveSampler(
    parameter_space=params, total_systems=500_000, batch_size=1000,
    BSEDict=BSEDict, SSEDict={'stellar_engine': 'sse'},
    is_interesting=any_dco(kstar_1=[14], kstar_2=[14]),
    derive_params=lambda s: {'mass_2': s['mass_1'] * s['q']},
    reject_systems="default", nproc=8,
)
result = sampler.run()
print(result.hit_rate, result.hit_rate_uncertainty)

What's new

Sampling engine (cosmic.sample.stroopwafel.AdaptiveSampler)

• Three-phase pipeline: exploration → adaptation (EM-fit Gaussian mixture) → refinement, with importance-weight bias correction. mc_only=True gives a plain Monte Carlo baseline. Multiprocessing via nproc.

Parameter space & binary definition

• ParameterSpace / Parameter; columns follow the order parameters are supplied.
• A binary is defined by {mass_1, mass_2, porb, ecc, metallicity}; each is either sampled or supplied by a user derive_params(sampled) -> dict callback (scalars broadcast). Coverage is validated at construction with a clear error.

Composable distributions (cosmic.sample.stroopwafel.distributions)

• Base distributions (Uniform, PowerLaw, BrokenPowerLaw, TruncatedNormal) composed with coordinate transforms (Identity, Log10, Ln, Sin, CosShift). Each bundles sampling, prior pdf, and adaptive kernel width.
• Built-in registry: kroupa (continuous broken power law, α=−1.3 below 0.5 M⊙ and −2.3 above), sana, sana_ecc, flat_in_log, uniform, uniform_in_sine, uniform_in_cosine, disberg. Extend via a Distribution instance, register(...), or subclassing.

Rejection & hit definitions

• rejection.default_reject applies physical cuts (secondary-mass floor, ZAMS contact, Roche-lobe overflow at periastron) using COSMIC's set_reff / a_from_p; it is SSEDict-aware so ZAMS radii use the same stellar engine as evolution (auto-wired into the default reject).
• Hit presets any_dco / merging_dco (presets.py), or any (bpp) -> (n_hits, hit_bin_nums) callable.

Outputs (cosmic.output.COSMICStroopOutput)

• Holds samples, importance weights, hit flags, and full COSMIC tables (bpp/bcm/initC/kick_info). Adds hit_rate / hit_rate_uncertainty, draw_representative_sample(n_samples) (weighted bootstrap returning params + bin numbers), and save / from_file.

Self-contained checkpointing (cosmic.output.STROOPWAFELCheckpoint)

• run_exploration() returns a checkpoint (.save(path)); AdaptiveSampler.from_checkpoint(path) rebuilds the sampler — parameter space, BSEDict/SSEDict, the callables (serialised with dill, so lambdas/closures survive), and the RNG state — and run_refinement() resumes with no re-specification. Any constructor argument can be overridden as a keyword.

New bhflag

• Adds bhflag=4 (src/cosmic/src/kick.f, documented in cosmic-settings.json): BH natal kicks get fallback scaling even when supplied directly via natal_kick_array, which is what enables per-binary BH natal-kick sampling in STROOPWAFEL.

Documentation

New "Adaptive importance sampling" tutorial series under docs/pages/tutorials/adaptive/:

• basics — end-to-end walkthrough (parameter space, binary definition, rejection, hit functions, running, rules of thumb).
• distributions — built-in registry and how to define custom distributions/transforms.
• outputs — interpreting importance weights and drawing a representative sample.
• checkpoint — the self-contained explore/refine workflow.

Wired into the tutorials toctree and the cosmic.sample API reference.

Testing

I got a little carried away on the tests...so these tests now outnumber all other COSMIC tests 😅 (Can you tell weights scare me?)

src/cosmic/tests/test_stroopwafel.py — 60 tests covering transforms/distributions (continuity, normalization, KS goodness-of-fit, broken-power-law slope recovery), the base × transform registry and custom-distribution paths, ParameterSpace (ordering, transforms, prior, sigma), binary-parameter assembly + validation, rejection, the Gaussian mixture, SSEDict wiring, checkpoint save/restore + overrides, and draw_representative_sample (verifying returned bin numbers map to matching initC initial conditions). Validated end-to-end against real COSMIC evolution.

[TomWagg]

Things for me and @katiebreivik to discuss:

• It seems that the RLOF calculation in evolve isn't the same as mine here, and I think mine is more consistently using periastron and the radius calculation from evolvebin
• The docs I've added are static because running stroopwafel stuff would take ~10s of minutes and not sure how to fit that into the CI. Let's think about it

[codecov]

Codecov Report

❌ Patch coverage is 32.85486% with 842 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.46%. Comparing base (8772c07) to head (340328b).
⚠️ Report is 176 commits behind head on develop.

Files with missing lines	Patch %	Lines
src/cosmic/sample/stroopwafel/main.py	0.00%	318 Missing ⚠️
src/cosmic/sample/stroopwafel/distributions.py	0.00%	151 Missing ⚠️
src/cosmic/output.py	0.00%	141 Missing ⚠️
src/cosmic/sample/stroopwafel/mixture_model.py	0.00%	118 Missing ⚠️
src/cosmic/sample/stroopwafel/parameter_space.py	0.00%	59 Missing ⚠️
src/cosmic/sample/stroopwafel/presets.py	0.00%	29 Missing ⚠️
src/cosmic/sample/stroopwafel/rejection.py	0.00%	17 Missing ⚠️
src/cosmic/src/kick.f	0.00%	4 Missing ⚠️
src/cosmic/sample/stroopwafel/__init__.py	0.00%	3 Missing ⚠️
src/cosmic/_version.py	0.00%	1 Missing ⚠️
... and 1 more

Additional details and impacted files

@@             Coverage Diff              @@
##           develop     #810       +/-   ##
============================================
- Coverage    86.91%   74.46%   -12.45%     
============================================
  Files           40       67       +27     
  Lines        25542    29361     +3819     
  Branches         0      986      +986     
============================================
- Hits         22198    21862      -336     
- Misses        3344     7194     +3850     
- Partials         0      305      +305     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.