Merge branch 'master' into fix/issue-50

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 -------------------------------------------