madspace: fix overweight tail after MadNIS training (truncation budget seeded with full target)

[oliviermattelaer]

Problem

With MadNIS on, unweighted samples carry about 3× the intended overweight tail. With the weight histogram from #190, p p > t t~ j j (13.6 TeV, 100k events, default card) has 0.26–0.30% of the cross section above w/⟨w⟩ = 1, reaching 3.2. max_overweight_truncation = 0.001 is meant to cap that at ~0.1%.

It looked multiplicity-dependent, but it isn't. p p > t t~ j j with [madnis] enable = false has the smallest tail of the four runs:

run (100k events)	channels	unw. eff.	tail above 1	above 1.05	reach
t t~ j, MadNIS off (default)	21	0.082	0.087%	0.055%	2.4
t t~ j, MadNIS on	8	0.388	0.163%	0.092%	2.7
t t~ j j, MadNIS on (default)	18	0.117	0.26–0.30%	0.224%	3.2
t t~ j j, MadNIS off	88	0.0086	0.084%	0.067%	3.1

The two processes differ only because set_madnis_auto_settings always enables MadNIS for n_out >= 4, but for n_out == 3 only when rsd > 10, events > 1M, or a gridpack is written. So a 2→3 process in those cases is affected too.

Cause: introduced by ae5a7f2

ae5a7f2 ("max weight determination bugfix", 3 Sep 2026, merged to main on 7 Sep via #51) added this at the start of EventGenerator::generate():

for (auto& channel : _channels)
    if (channel->status().count_opt == 0)
        channel->set_target_count(_config.target_count);

The truncation budget in update_max_weight is max_overweight_truncation × min(count_target, freeze_max_weight_after). Seeding the full run target makes each seeded channel's first-batch budget 100 (for 100k events) instead of 0.001 × its share. Instrumenting update_max_weight shows what happens to channel 4.1 of t t~ j j (final target 291, real budget 0.29): one 1000-event batch at budget 100 moved its cap from ~160 to 1.53. After that:

• _max_weight only ever rises;
• the channel met its real target immediately (acceptance is ~100% with that cap), so generate() skipped it for the rest of the run and update_max_weight was never called on it again.

That cap was therefore permanent. Four such small channels (0.74% of the events) produced 61% of the whole overweight tail.

With the seed removed (the logic before ae5a7f2), the same MadNIS run gives 0.093%, so this commit accounts for the entire regression. I could not reproduce the case ae5a7f2 was fixing: without the seed the efficiency was 0.124 against 0.117 with it. This PR therefore keeps the seed and fixes what it leaves behind.

Why MadNIS off was not affected

The seed only fires for channels with count_opt == 0 when generate() starts.

• MadNIS off: generate() runs on the same EventGenerator that ran survey(). Every channel already has count_opt > 0, and survey() ended with update_integral_fractions(), which set real per-channel targets. The seed never fires. Budget excursions during the survey do not carry over either: each VEGAS batch calls clear_events(), which resets _max_weight and _large_weights.
• MadNIS on: build_madnis() builds new Channel objects (event_generator = None), and train_madnis() then calls build_event_generator() without a survey. So every channel is a fresh ChannelEventGenerator with count_opt == 0, and the seed fires for all of them. Flows also have no VEGAS/discrete optimizer, so needs_optimization() is false and clear_events() never runs to reset the cap.

This rebuild-without-survey has been there since be0fe9f (2025) and was harmless before ae5a7f2.

Fix

ChannelEventGenerator::set_target_count() moves out of the header. When the target decreases, it re-applies the truncation budget (the existing loop, factored into apply_truncation_budget()). This can only raise _max_weight. Raising it resets _unweighted_count, so unweight_all() re-unweights the whole weight file against the final cap, and the per-channel invariant (excess ≤ max_overweight_truncation × count_target) holds whatever the seed was. The freeze is respected (no update once count_unweighted > freeze_max_weight_after).

I also tried seeding with target_count / n_channels instead. It gives 0.104%, but small channels still end up above their budget (up to 14×), so I did not include it. Combining the two gives 0.094%, which adds nothing over this change alone.

Validation

• p p > t t~ j j, defaults (MadNIS on): tail 0.26–0.30% → 0.096%, reach 3.2 → 2.x. Unweighting efficiency 0.117 → 0.120, σ unchanged (256.06 ± 0.19 pb). With the instrumented build, no channel exceeds its budget (worst 0.43×, was 70×).
• p p > t t~ j, defaults (MadNIS off): 0.087% → 0.093%, efficiency 0.082 → 0.084, σ 411.98 pb: unchanged within run-to-run noise.
• madspace tests: 1584 passed, 11 skipped.

freeze_max_weight_after and the default max_overweight_truncation are not involved (no channel reached the freeze in any run). Neither needs to scale with multiplicity.

Large batches (GPU): no return of the bug ae5a7f2 fixed

Before ae5a7f2, a channel entering generate() unsurveyed kept count_target = 1. Its first batch therefore ran with budget ≈ 0, and its cap became the maximum weight of that whole batch. Since the cap only rises, that choice was permanent, and with a 64k GPU batch it is very conservative. This PR keeps the seed and only ever raises caps, so I checked that it does not drift back to that behaviour. The first batch after MadNIS is one device batch, so cpu_batch_size = 64000 reproduces the GPU case on CPU. Three modes, all at 64k:

p p > t t~ j j, cpu_batch_size = 64000	unw. eff.	samples	tail above 1	budget used
main (seed with the full target)	0.064	1.57M	0.43% (reach 4.8)	103%, badly spread: even the large channels at 1.3–2.2× their budget
before ae5a7f2 (no seed)	0.032	3.11M	0.011%	3%: caps 2–5× too high, which is the old bug
this PR	0.050–0.053	1.89–2.02M	0.083–0.085%	26%; large channels use 17–59%, none above 1×

• With the PR, the final caps of the large channels (0.4: 348, 0.1: 372, 0.5: 778) are 2–2.5× below the old-bug caps (814, 830, 1821). Their budget use (17–59%) matches the healthy surveyed path (t t~ j, MadNIS off: 10–30%). So the caps come from each channel's real budget, and nothing is left of the conservative first-batch maximum.
• The PR's efficiency sits below main at 64k (0.053 against 0.064) only because main spends 4× its truncation budget. At 64k main is also worse than at 1k (0.43% against 0.30%), because the budget-100 first batch spans 64× more events.
• Some efficiency loss at 64k is plain granularity: round one alone is 18 channels × 64k = 1.15M samples. The no-MadNIS t t~ j path shows the same drop (0.084 → 0.067), identically in all three modes. That path never applies the seed, so this PR does not change it.

🤖 Generated with Claude Code

[theoheimel]

We should also add a test to check that the truncation stays below the limit to prevent this from regressing again

[theoheimel]

Looks good, let's merge this