Merge branch 'master' into fix/issue-50

giant/config.py

@@ -468,6 +468,18 @@ class ParticleTypeConfig:
     # silently the same number). 0 = inherit conditioning.particle.emb_dim,
     # preserving pre-#29 behavior.
     n_classes: int = 0
+    # Class-balances the target = "onehot" cross-entropy loss against the
+    # secondary-species long tail (gitea #44: the failure mode motivating the
+    # v0.3.0 pivot was specifically a species collapse — zero photon
+    # secondaries, hallucinated antineutrinos). "none": plain CE (pre-#44
+    # behavior). "inverse_freq": CE weighted by 1/count per class,
+    # normalized to mean 1 across classes so lambda_weight doesn't need
+    # retuning when this is switched on. validate_config requires target =
+    # "onehot" and stage2_model.generator != "wgan" whenever this isn't
+    # "none" — "embedding"/"physical" have no class CE to weight, and the
+    # WGAN stage-2 path feeds its type slice to the critic via a
+    # straight-through Gumbel relaxation instead of a CE loss.
+    class_weighting: str = "none"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
@@ -477,6 +489,7 @@ class ParticleTypeConfig:
             lambda_weight=d.get("lambda", 1.0),
             other_policy=d.get("other_policy", "sample"),
             n_classes=d.get("n_classes", 0),
+            class_weighting=d.get("class_weighting", "none"),
         )
 
     def to_dict(self) -> dict:
@@ -485,6 +498,7 @@ class ParticleTypeConfig:
             "lambda": self.lambda_weight,
             "other_policy": self.other_policy,
             "n_classes": self.n_classes,
+            "class_weighting": self.class_weighting,
         }
 
 
@@ -1438,6 +1452,26 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
             f"{particle_type!r})"
         )
 
+    class_weighting = _get_path(cfg, "stage2_model.particle_type.class_weighting")
+    if class_weighting not in ("none", "inverse_freq"):
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {class_weighting!r} — must be 'none' or 'inverse_freq'"
+        )
+    if class_weighting != "none" and pt_target != "onehot":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' requires "
+            f"stage2_model.particle_type.target = 'onehot' (there is no class "
+            f"cross-entropy to weight under target = {pt_target!r})"
+        )
+    if class_weighting != "none" and _get_path(cfg, "stage2_model.generator") == "wgan":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' is "
+            "incompatible with stage2_model.generator = 'wgan' — that path "
+            "feeds the type slice to the critic via a straight-through "
+            "Gumbel relaxation instead of a class cross-entropy, so there is "
+            "nothing to weight"
+        )
+
     for stage_name in ("stage1_model", "stage2_model"):
         if _get_path(cfg, f"{stage_name}.freeze") and not _get_path(cfg, f"{stage_name}.init_from") and not resume:
             raise ValueError(

giant/data/loader.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Iterator
 
@@ -256,24 +256,32 @@ def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
     return counts
 
 
-def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
+def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
     """Frequency-capped value->index map: the `n_classes - 1` most frequent
     keys get their own index; every rarer key is bucketed into a shared
     "other" index (`n_classes - 1`).
 
-    Returns `(class_map, other_members)` — `other_members` is `{key: count}`
-    for every key bucketed into "other" (the empirical within-bucket
-    distribution, for `other_policy = "sample"` at rollout).
+    Returns `(class_map, other_members, class_counts)` — `other_members` is
+    `{key: count}` for every key bucketed into "other" (the empirical
+    within-bucket distribution, for `other_policy = "sample"` at rollout);
+    `class_counts` is `{index: total_count}` for every resulting class index
+    (0-indexed; the "other" index's count is the sum of `other_members`),
+    the per-class frequencies `stage2_model.particle_type.class_weighting`
+    (gitea #44) needs and that would otherwise be dropped once `counts` is
+    collapsed into `class_map`.
     """
     ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
     keep = ranked[: max(n_classes - 1, 0)]
     class_map = {k: i for i, k in enumerate(keep)}
+    class_counts = {i: counts[k] for i, k in enumerate(keep)}
     other_idx = n_classes - 1
     other_members: dict = {}
     for k in ranked[len(keep) :]:
         class_map[k] = other_idx
         other_members[k] = counts[k]
-    return class_map, other_members
+    if other_members:
+        class_counts[other_idx] = sum(other_members.values())
+    return class_map, other_members, class_counts
 
 
 def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
@@ -287,7 +295,7 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
     fixed-width n_sec_head classifier.
     """
     counts = _rank_by_frequency_from_files(files, "process", str)
-    class_map, _ = _topn_plus_other_map(counts, n_experts)
+    class_map, _, _ = _topn_plus_other_map(counts, n_experts)
     return class_map
 
 
@@ -299,6 +307,11 @@ class TopNMap:
 
     class_map: dict
     other_members: dict
+    # {class_index: total_count} — see _topn_plus_other_map. Empty for a
+    # TopNMap decoded from a checkpoint/sidecar predating gitea #44; only
+    # stage2_model.particle_type.class_weighting reads it, and it raises
+    # loudly if it needs counts that aren't there (giant/training/trainers.py).
+    class_counts: dict = field(default_factory=dict)
 
 
 def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
@@ -315,8 +328,8 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
     free during this same scan.
     """
     counts = _rank_by_frequency_from_files(files, column, cast)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
 
 
 def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
@@ -347,5 +360,5 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
         if has_sec:
             exploded = df["sec_pdg_list"].explode().dropna()
             _accumulate_value_counts(counts, exploded, int)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)

giant/data/setup_cache.py

@@ -35,7 +35,10 @@ from giant.data.transforms import Normalizer, sorted_membership
 # v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
 # energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
 # v2 sidecar has no such grid to fall back on, so it must be recomputed.
-_CACHE_FORMAT_VERSION = 3
+# v4: TopNMap gained class_counts (gitea #44, stage2_model.particle_type.
+# class_weighting) — a v3 sidecar's cached topn_maps have no counts, so they
+# must be rebuilt rather than silently cached with class_counts={}.
+_CACHE_FORMAT_VERSION = 4
 
 _DIMS = {
     "COND_DIM": COND_DIM,
@@ -131,6 +134,7 @@ def topnmap_to_json(m: TopNMap) -> dict:
     return {
         "class_map": {str(k): v for k, v in m.class_map.items()},
         "other_members": {str(k): v for k, v in m.other_members.items()},
+        "class_counts": {str(k): v for k, v in m.class_counts.items()},
     }
 
 
@@ -139,6 +143,11 @@ def topnmap_from_json(d: dict, axis: str) -> TopNMap:
     return TopNMap(
         class_map={cast(k): v for k, v in d["class_map"].items()},
         other_members={cast(k): v for k, v in d["other_members"].items()},
+        # Missing for a checkpoint's topn maps predating gitea #44 — {} is
+        # the correct decode there (inference never reads class_counts; only
+        # stage2_model.particle_type.class_weighting does, at train time, and
+        # it raises loudly if it needs counts a checkpoint doesn't have).
+        class_counts={int(k): v for k, v in d.get("class_counts", {}).items()},
     )
 
 

giant/model/summary.py

@@ -74,6 +74,7 @@ _NOT_BUILD_TIME: dict[str, str] = {
     "stage2_model.n_sec.lambda": "trainers.py: StageSpec.n_sec_lambda, the n_sec-head loss weight",
     "stage2_model.particle_type.lambda": "trainers.py: Stage2Trainer.particle_type_lambda, the type-head loss weight",
     "stage2_model.particle_type.other_policy": "giant/rollout.py: resolves an 'other'-bucket secondary's PDG code at inference",
+    "stage2_model.particle_type.class_weighting": "trainers.py: FlowDDPMStageTrainer.type_class_weights, shapes the type-head loss, not the built graph (gitea #44)",
     "stage2_model.autoregressive.teacher_forcing": "giant/training/stage2_inputs.py's training-time input assembly",
     "stage2_model.autoregressive.tf_p_start": "trainers.py's teacher-forcing schedule",
     "stage2_model.autoregressive.tf_p_end": "trainers.py's teacher-forcing schedule",

giant/training/loop.py

@@ -132,7 +132,8 @@ def train(
     validate_steps = t.get("validate_steps", 10)
     max_val_batches = t.get("max_val_batches", 0)
 
-    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
+    sec_type_class_counts = sec_type_topn_map.class_counts if sec_type_topn_map is not None else None
+    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches, sec_type_class_counts)
     if not trainers:
         raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
     for line in init_stages_from_checkpoints(trainers):

giant/training/trainers.py

@@ -76,6 +76,41 @@ def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
     return type(batch)(*(t.to(device) for t in batch))
 
 
+def _type_class_weight_vector(class_counts: dict[int, int], n_classes: int, scheme: str) -> list[float] | None:
+    """Per-class `F.cross_entropy(weight=...)` vector for the stage-2 type
+    head's `class_weighting` (gitea #44), or `None` under `"none"` (the
+    pre-#44 unweighted-CE behavior — the caller must pass that through as
+    `weight=None`, not a vector of ones, so old runs stay bit-identical).
+
+    `"inverse_freq"`: `1 / count` per class, normalized to mean 1 over
+    `n_classes` so switching this on doesn't rescale the type loss against
+    `particle_type.lambda` / the generator loss it's summed with. A class
+    with zero training examples (fewer distinct species than `n_classes - 1`
+    slots) clamps its count to 1 — its weight is otherwise undefined, and
+    since it never appears in a batch's labels the value is inert anyway.
+
+    Raises if `scheme != "none"` and `class_counts` is empty: that means the
+    `TopNMap` behind this run predates gitea #44 (a stale checkpoint's decode
+    map, or a not-yet-rebuilt setup-cache sidecar) and truly has no
+    frequency information to weight by — silently falling back to uniform
+    weights would look like the feature is active when it isn't.
+    """
+    if scheme == "none":
+        return None
+    if not class_counts:
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {scheme!r} requires "
+            "per-class counts, but this run's sec_type_topn_map has none "
+            "(class_counts={}) — it was built before gitea #44 or loaded "
+            "from a stale setup-cache sidecar/checkpoint; rebuild the setup "
+            "cache (giant train --rebuild-setup-cache) or retrain."
+        )
+    counts = [max(class_counts.get(i, 0), 1) for i in range(n_classes)]
+    inv = [1.0 / c for c in counts]
+    mean_inv = sum(inv) / len(inv)
+    return [w / mean_inv for w in inv]
+
+
 @dataclass(frozen=True)
 class StageSpec:
     """One stage's resolved training configuration.
@@ -102,6 +137,10 @@ class StageSpec:
     # particle-type target (stage 2 only)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
     particle_type_n_classes: int = 16
+    # Resolved by from_config from sec_type_class_counts (dataset-derived,
+    # not itself a cfg value — see _type_class_weight_vector) crossed with
+    # particle_type.class_weighting (gitea #44). None under "none".
+    type_class_weights: list[float] | None = None
 
     # optimization
     lr: float = 3e-4
@@ -139,7 +178,18 @@ class StageSpec:
     type_gumbel_tau_end: float = 0.1
 
     @classmethod
-    def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
+    def from_config(
+        cls,
+        cfg: dict,
+        name: str,
+        is_stage2: bool,
+        steps_per_epoch: int,
+        sec_type_class_counts: dict[int, int] | None = None,
+    ) -> "StageSpec":
+        """`sec_type_class_counts` is dataset-derived (`sec_type_topn_map.class_counts`,
+        gitea #44), not a `cfg` value — it's the one input to `StageSpec` that
+        doesn't come from `cfg`, kept separate from the "only place that
+        reads `cfg`" invariant below on purpose."""
         t = TrainConfig.from_dict(cfg["train"])
         # n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are
         # stage-2-only concepts, always read off s2_spec (guarded by
@@ -152,6 +202,9 @@ class StageSpec:
         # stage 1's).
         s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
         stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"])
+        particle_type_n_classes = resolve_type_n_classes(
+            s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
+        )
         return cls(
             name=name,
             is_stage2=is_stage2,
@@ -163,8 +216,9 @@ class StageSpec:
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
             n_sec_mode=s2_spec.n_sec.mode,
             particle_type=s2_spec.particle_type,
-            particle_type_n_classes=resolve_type_n_classes(
-                s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
+            particle_type_n_classes=particle_type_n_classes,
+            type_class_weights=_type_class_weight_vector(
+                sec_type_class_counts or {}, particle_type_n_classes, s2_spec.particle_type.class_weighting
             ),
             # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
             # (giant/config.py), so TrainConfig.from_dict never has to fall
@@ -619,6 +673,12 @@ class FlowDDPMStageTrainer(StageTrainer):
         # "onehot"/"embedding" pull it out into model.type_head instead (0
         # here).
         self._flow_type_dim = None if self.particle_type_cfg.target == "physical" else 0
+        # gitea #44: None under class_weighting = "none" (the default),
+        # matching F.cross_entropy's own unweighted default — a real tensor
+        # only materializes when the config asked for one.
+        self.type_class_weights = (
+            None if spec.type_class_weights is None else torch.tensor(spec.type_class_weights, device=device)
+        )
 
         self.params = list(self.model.parameters())
         self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
@@ -712,8 +772,12 @@ class FlowDDPMStageTrainer(StageTrainer):
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
         if self.particle_type_cfg.target == "onehot":
-            ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
+            ce = F.cross_entropy(
+                type_out.transpose(1, 2), sec_type_idx, weight=self.type_class_weights, reduction="none"
+            )
             l_type = (ce * mask).sum() / denom
+            # Unweighted, deliberately — type_acc is a diagnostic of raw
+            # per-slot correctness, not the (possibly class-weighted) loss.
             type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
         else:  # "embedding"
             target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
@@ -1120,6 +1184,7 @@ def build_stage_trainers(
     critics: dict[str, torch.nn.Module | None],
     device: torch.device,
     total_train_batches: int,
+    sec_type_class_counts: dict[int, int] | None = None,
 ) -> dict[str, StageTrainer]:
     """One trainer per active stage — `models[name] is None` means that stage
     is `active = false` and is simply never constructed.
@@ -1128,13 +1193,18 @@ def build_stage_trainers(
     stage-2 trainer to the stage-1 one (`StageTrainer.attach_stage1`) so it
     can draw a real stage-1 sample instead of only ever seeing the
     ground-truth stage-1 outcome — `validate_config` already guarantees both
-    stages are active whenever that config value is set."""
+    stages are active whenever that config value is set.
+
+    `sec_type_class_counts` (`sec_type_topn_map.class_counts`, gitea #44) is
+    the one dataset-derived input `StageSpec.from_config` needs beyond `cfg`
+    — `None`/absent whenever `stage2_model.particle_type.class_weighting =
+    "none"` (the default), which never reads it."""
     trainers: dict[str, StageTrainer] = {}
     for name, is_stage2 in (("stage1", False), ("stage2", True)):
         model = models.get(name)
         if model is None:
             continue
-        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
+        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1), sec_type_class_counts)
         if build_objective(spec.generator).is_adversarial:
             critic = critics.get(name)
             assert critic is not None, (

tests/test_config.py

@@ -132,7 +132,15 @@ def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
     assert gconfig.ParticleTypeConfig().n_classes == 0
     spec = gconfig.ParticleTypeConfig.from_dict({"n_classes": 32})
     assert spec.n_classes == 32
-    assert spec.to_dict()["n_classes"] == 32
+
+
+def test_particle_type_config_class_weighting_defaults_to_none_and_round_trips():
+    """gitea #44: an existing config.toml with no
+    stage2_model.particle_type.class_weighting key must reproduce the
+    pre-#44 unweighted-CE behavior exactly."""
+    assert gconfig.ParticleTypeConfig().class_weighting == "none"
+    spec = gconfig.ParticleTypeConfig.from_dict({"class_weighting": "inverse_freq"})
+    assert spec.class_weighting == "inverse_freq"
 
 
 def test_router_config_extra_round_trips_composed_axis_keys():
@@ -704,6 +712,42 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
     gconfig.validate_config(cfg)  # must not raise
 
 
+def test_validate_config_bad_class_weighting_rejected():
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "effective_num"})
+    with pytest.raises(ValueError, match="class_weighting"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_requires_onehot_target():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "physical",
+        }
+    )
+    with pytest.raises(ValueError, match="onehot"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_incompatible_with_wgan_generator():
+    # stage2_model.generator defaults to "wgan" and particle_type.target
+    # defaults to "onehot", so only class_weighting needs overriding here.
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "inverse_freq"})
+    with pytest.raises(ValueError, match="wgan"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_passes_with_onehot_and_flow():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "onehot",
+            "stage2_model.generator": "flow",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 def test_validate_config_mixed_particle_material_conditioning_is_valid():
     """The particle and material conditioning axes are configured
     independently and may mix freely — e.g. material

tests/test_loader.py

@@ -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):

tests/test_pipeline.py

@@ -231,7 +231,7 @@ def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, da
 
 def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
     cfg = _tiny_cfg()
-    cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
+    cfg["stage2_model"]["particle_type"].update({"target": "physical", "lambda": 1.0})
     echo = _run(data, tmp_path / "out", cfg=cfg)
     assert not any("top-N map" in m for m in echo)
 

tests/test_setup_cache.py

@@ -99,7 +99,7 @@ def test_save_load_round_trip_topn_maps(tmp_path):
 
     cache = SetupCache.empty(files)
     cache.topn_maps[setup_cache.topn_key("pdg", 3)] = TopNMap(
-        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}
+        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}, class_counts={0: 100, 1: 50, 2: 5}
     )
     cache.topn_maps[setup_cache.topn_key("material", 2)] = TopNMap(
         class_map={"G4_AIR": 0, "PbWO4": 1}, other_members={}
@@ -114,9 +114,21 @@ def test_save_load_round_trip_topn_maps(tmp_path):
     assert pdg_m.other_members == {2212: 5}
     # key type is int (matches pdg_map's own key type), not str
     assert all(isinstance(k, int) for k in pdg_m.class_map)
+    # class_counts (gitea #44) round-trips too, keyed by class index (always
+    # int, independent of the pdg/material axis's own key type).
+    assert pdg_m.class_counts == {0: 100, 1: 50, 2: 5}
+    assert all(isinstance(k, int) for k in pdg_m.class_counts)
 
     mat_m = loaded.topn_maps[setup_cache.topn_key("material", 2)]
     assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}
+    assert mat_m.class_counts == {}
+
+
+def test_topnmap_from_json_missing_class_counts_defaults_empty():
+    """A checkpoint's topn map predating gitea #44 has no class_counts key at
+    all — must decode to {}, not raise, since inference never reads it."""
+    m = setup_cache.topnmap_from_json({"class_map": {"11": 0}, "other_members": {}}, axis="pdg")
+    assert m.class_counts == {}
 
 
 def test_topn_key_unknown_axis_raises():

tests/test_train.py

@@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch
 import numpy as np
 import pytest
 import torch
+import torch.nn.functional as F
 
 from giant.config import ParticleTypeConfig
 from giant.constants import (
@@ -35,6 +36,7 @@ from giant.training import (
     train,
 )
 from giant.training.metrics import _wandb_run_config
+from giant.training.trainers import _type_class_weight_vector
 from giant.training.stage2_inputs import (
     _ar_has_prev,
     _assemble_stage2_ar_inputs,
@@ -599,6 +601,151 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
 
+# --- gitea #44: class-balanced secondary particle-type loss -----------------
+
+
+def test_type_class_weight_vector_none_scheme_returns_none():
+    assert _type_class_weight_vector({0: 100, 1: 5}, n_classes=2, scheme="none") is None
+
+
+def test_type_class_weight_vector_raises_without_counts():
+    with pytest.raises(ValueError, match="class_counts"):
+        _type_class_weight_vector({}, n_classes=4, scheme="inverse_freq")
+
+
+def test_type_class_weight_vector_inverse_freq_favors_rare_class_and_has_mean_one():
+    weights = _type_class_weight_vector({0: 1000, 1: 10, 2: 1, 3: 1}, n_classes=4, scheme="inverse_freq")
+    assert weights is not None
+    assert len(weights) == 4
+    assert weights[1] > weights[0]  # rarer class -> larger weight
+    assert math.isclose(sum(weights) / len(weights), 1.0, rel_tol=1e-9)
+
+
+def test_type_class_weight_vector_missing_index_clamps_to_count_one():
+    # n_classes=3 but only index 0 was ever observed (e.g. a tiny dataset) —
+    # indices 1/2 must not divide by zero.
+    weights = _type_class_weight_vector({0: 10}, n_classes=3, scheme="inverse_freq")
+    assert weights is not None
+    assert all(math.isfinite(w) for w in weights)
+
+
+def _onehot_flow_stage2_setup():
+    """A built stage-2 model + a batch, under target='onehot' + generator='flow'
+    (mirrors the 'stage2_onehot_target_flow' case in test_train_end_to_end)."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
+    model_config = _model_config(cfg)
+    model = build_models(model_config)["stage2"]
+    assert model is not None
+    batch = _fake_batches(1, 8)[0]
+    device = torch.device("cpu")
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+    # Mostly class 0 (common), a few slot 1's set to class 1 (rare) —
+    # PARTICLE_CFG's emb_dim=8, n_classes=0 (inherit) -> 8 type classes.
+    sec_type_idx = torch.zeros(8, K_MAX, dtype=torch.long)
+    sec_type_idx[:, :2] = 1
+    sec_mask = torch.ones(8, K_MAX, dtype=torch.bool)
+    return model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device
+
+
+def test_flow_ddpm_trainer_type_loss_none_leaves_weight_unset():
+    model, *_ = _onehot_flow_stage2_setup()
+    spec = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"),
+        particle_type_n_classes=8,
+        ema_decay=0.0,
+    )
+    trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu"))
+    assert trainer.type_class_weights is None
+
+
+def test_flow_ddpm_trainer_type_loss_matches_manual_weighted_cross_entropy():
+    model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device = _onehot_flow_stage2_setup()
+    class_counts = {0: 1000, 1: 10, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1, 7: 1}
+    weights = _type_class_weight_vector(class_counts, n_classes=8, scheme="inverse_freq")
+    spec = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="inverse_freq"),
+        particle_type_n_classes=8,
+        type_class_weights=weights,
+        ema_decay=0.0,
+    )
+    trainer = FlowDDPMStageTrainer(spec, model, device)
+    assert trainer.type_class_weights is not None
+    stage1_ctx = trainer._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None)
+
+    with torch.no_grad():
+        type_out = model.predict_type(cond_cont, cond_cat, stage1_ctx)
+        weight_t = torch.tensor(weights)
+        ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, weight=weight_t, reduction="none")
+        expected = (ce * sec_mask.float()).sum() / sec_mask.float().sum().clamp(min=1)
+
+        l_type, _ = trainer._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device)
+
+    assert torch.allclose(l_type, expected, atol=1e-6)
+
+    # Unweighted trainer, same model/batch — the two losses must differ
+    # (the batch mixes the common and rare classes, so weighting changes the
+    # per-slot contributions), confirming the weight is actually plumbed in.
+    spec_none = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"),
+        particle_type_n_classes=8,
+        ema_decay=0.0,
+    )
+    trainer_none = FlowDDPMStageTrainer(spec_none, model, device)
+    with torch.no_grad():
+        l_type_none, _ = trainer_none._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device)
+    assert not torch.allclose(l_type, l_type_none)
+
+
+def test_build_stage_trainers_threads_sec_type_class_counts_into_weights():
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["particle_type"] = {
+        "target": "onehot",
+        "lambda": 1.0,
+        "class_weighting": "inverse_freq",
+    }
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    class_counts = {i: 100 for i in range(8)}
+    class_counts[1] = 1  # one rare class
+    trainers = build_stage_trainers(
+        cfg, models, critics, torch.device("cpu"), total_train_batches=4, sec_type_class_counts=class_counts
+    )
+    stage2_trainer = trainers["stage2"]
+    assert isinstance(stage2_trainer, FlowDDPMStageTrainer)
+    weights = stage2_trainer.type_class_weights
+    assert weights is not None
+    assert weights[1] > weights[0]
+
+
+def test_build_stage_trainers_no_class_counts_with_none_weighting_is_fine():
+    """The overwhelmingly common case (class_weighting = 'none', the
+    default): build_stage_trainers must not require sec_type_class_counts at
+    all."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0})
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+    stage2_trainer = trainers["stage2"]
+    assert isinstance(stage2_trainer, FlowDDPMStageTrainer)
+    assert stage2_trainer.type_class_weights is None
+
+
 # --- gitea #42: freeze / init_from -------------------------------------------