Add class-balanced secondary particle-type loss (gitea #44)
CI / Lint (ruff check) (push) Successful in 36s
CI / Format (ruff format) (push) Successful in 38s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 38s
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 45s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 47s
CI / Tests (push) Successful in 5m20s
CI / Tests (pull_request) Successful in 4m50s

The v0.3.0 pivot exists because the 2026-08-03 WGAN rollout benchmark
produced zero photon secondaries and ~4M hallucinated antineutrinos —
even with a correctly-sized top-N species vocabulary (gitea #29), plain
cross-entropy over a class distribution spanning orders of magnitude
still under-predicts rare-but-physical species.

stage2_model.particle_type.class_weighting = "none" | "inverse_freq"
(default "none", fully back-compat) weights the stage-2 type head's CE
loss (FlowDDPMStageTrainer._type_loss) by inverse class frequency,
normalized to mean 1 so switching it on doesn't rescale the type loss
against particle_type.lambda / the generator loss it's summed with.

The per-class counts the weighting needs don't already exist despite the
issue's premise: _topn_plus_other_map (giant/data/loader.py) previously
kept counts only for keys folded into "other", dropping the kept classes'
counts on the floor. TopNMap now carries class_counts (index -> count),
round-tripped through the setup-cache sidecar (format version bumped
3->4, since existing sidecars have none) and through checkpoints
(tolerantly — a pre-#44 checkpoint decodes to {}, since only training-time
loss weighting reads it, not inference).

Decisions made during planning (with the user): dropped "effective_num"
from the issue's proposed three-way enum (no beta hyperparameter to
design around) — final domain is "none" | "inverse_freq". Weights are
mean-1-normalized. validate_config rejects class_weighting != "none"
combined with particle_type.target != "onehot" or
stage2_model.generator == "wgan" (both have no class CE to weight),
following the #28/#30 dead-key-must-not-go-silent convention. Branch
fix/issue-44 off master.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-17 23:02:50 +02:00
parent f60af64d00
commit fce47b128c
11 changed files with 360 additions and 22 deletions
+7
View File
@@ -195,6 +195,10 @@ def test_build_topn_map_from_files_keeps_most_frequent(tmp_path):
assert m.class_map["G4_Fe"] == 2 # "other" (n_classes - 1)
assert m.class_map["G4_Pb"] == 2
assert m.other_members == {"G4_Fe": 2, "G4_Pb": 1}
# class_counts (gitea #44): per resulting index, "other" is the sum of
# everything folded into it (2 + 1 = 3), and the total equals row count.
assert m.class_counts == {0: 5, 1: 3, 2: 3}
assert sum(m.class_counts.values()) == len(materials)
def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
@@ -205,6 +209,8 @@ def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
assert m.class_map == {"G4_AIR": 0, "PbWO4": 1}
assert m.other_members == {}
# No "other" bucket ever populated -> no entry for its index either.
assert m.class_counts == {0: 1, 1: 1}
def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
@@ -224,6 +230,7 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
# pooled: 11 -> 5, 22 -> 1 (primary) + 10 (secondary) = 11
assert m.class_map[22] == 0
assert m.class_map[11] == 1
assert m.class_counts == {0: 11, 1: 5}
def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):