Stratified ("quiet") resampling

[Nathan-Majernik]

Adds the option for stratified ("quiet") resampling (in addition to the current, fully random downsampling) to resample(). Includes a new demo notebook contrasting random vs stratified downsampling and their effects on the bunching factor of a beam

This higher noise power from random downsampling can cause collective effects like CSR to be incorrectly represented in simulation, leading to major, macroscopic divergences from high fidelity simulations; stratified downsampling can mitigate this.

Stratified downsampling can cause distortions if the new population is not much smaller than the initial population. By default, if the ratio is not more than 5, stratified_resample_particles() will fall back to random. The user may override this behavior by setting allow_bad_sampling_ratio to True.

Claude made some tests for the new function; not sure how useful they are.

[ChristopherMayes]

Fable 5 review:

Review: PR #149 — Stratified ("quiet") resampling

Overview

Adds stratified_resample_particles() to statistics.py, exposes it via a new ParticleGroup.stratified_resample() method and a method= kwarg on resample(), plus a demo notebook wired into mkdocs and three new tests. I verified the core math (strata edges via linspace(0, m, n+1).astype(int) always yield width ≥ 1 when n ≤ m; offsets stay in-stratum; charge is preserved; id is carried through when present), and the new tests pass in beamphysics-dev (3 passed, 1.16s). The implementation is clean, well-documented, and backward compatible. No correctness bugs found — the items below are design and polish.

Main design point: the weight check happens after the ratio fallback

In stratified_resample_particles (beamphysics/statistics.py:592-601), the fallback to random runs before the variable-weight rejection. So whether variable weights are accepted depends on n:

• Pvar.stratified_resample(n_alive // 2) → silently succeeds (random fallback tolerates variable weights)
• Pvar.stratified_resample(n_alive // 10) → raises ValueError

Asking for more output particles succeeds while asking for fewer errors, which is backwards from what a user would guess, and there's no signal that the fallback occurred. Since Ken's feedback removed the warnings, the silent fallback is presumably intentional for the ratio itself — but I'd still move the weight check above the fallback so the "stratified requires constant weights" contract holds unconditionally. If the current behavior is deliberate, it should at least be stated in the docstring.

Other suggestions

• Reproducibility: rng = np.random.default_rng() is unseeded and unseedable, while the random path responds to np.random.seed(). The demo notebook even seeds the global RNG "for a reproducible figure" — but the stratified curve in that same figure is not reproducible, so the docs plot changes on every rebuild. An optional rng=None parameter (rng = np.random.default_rng(rng)) would fix both.
• Simplification: offsets = (rng.random(n) * (edges[1:] - edges[:-1])).astype(int) can be pick = order[rng.integers(edges[:-1], edges[1:])] — one line, exactly uniform, no float truncation.
• min_ratio drift: it's a local variable, but the statistics.py docstring refers to min_ratio (which the user can't see) while the particles.py docstring hard-codes "5". A module-level constant (e.g. STRATIFIED_MIN_RATIO = 5) referenced from both docstrings keeps them from drifting.
• Notebook: there's a trailing empty code cell (cell 10). Otherwise the demo is nice — the bunching-spectrum comparison with the analytic Gaussian form factor and the 1/√n noise floor makes the motivation very clear.

On the tests ("not sure how useful they are")

They're genuinely useful — good coverage of the happy path (count, status, charge preservation, equal weights, picks-from-source), dispatch through resample(method=...), all three error branches, and the ratio fallback. Two gaps:

1. Nothing tests that the output is actually stratified. A pure-random sample would pass every current assertion. The defining property is cheap to check: sort the source by key, compute the strata edges, and assert exactly one picked particle falls in each stratum (or more loosely, that sorted Q.t values land in monotonically increasing strata). That's the test that would catch a real regression in this feature.
2. The fallback test is indirect — it detects which code path ran by whether variable weights are rejected. Clever, but it's coupled to the check-ordering issue above: if you move the weight check before the fallback (as I suggest), this test breaks by design. A more direct probe would be monkeypatching or comparing statistical noise, but honestly, if you reorder the checks, "bad ratio + variable weights raises" becomes the correct expectation and the test simplifies.

Minor: set(np.asarray(Q.t)).issubset(...) works (exact copies, so float equality is safe) but np.isin(Q.t, alive.t).all() is the idiomatic form.

Verdict

Solid PR — mergeable as-is from a correctness standpoint. The one thing I'd change before merge is the weight-check ordering (or an explicit docstring note if the current behavior is intended); the RNG seedability and a stratification-property test are worthwhile follow-ups.

[Nathan-Majernik]

@ChristopherMayes All suggestions implemented