import copy import numpy as np import pandas as pd import pytest import torch from giant import config as gconfig from giant.constants import COND_DIM from giant.data import setup_cache from giant.data.transforms import Normalizer from giant.pipeline import _seed_energy_router, run_train_job def _unit(v): v = np.asarray(v, dtype=np.float64) n = np.linalg.norm(v) return v / n if n > 1e-9 else np.array([0.0, 0.0, 1.0]) def _make_synthetic_steps(path, n_events=20, seed=0): """A tiny but schema-complete synthetic steps parquet for run_train_job. pdg/material/process are assigned deterministically by row index (not random) so tests that assert on the resulting vocab/proc maps aren't flaky; only continuous quantities (positions/energies/directions) are drawn from `rng`. """ rng = np.random.default_rng(seed) materials = ["G4_AIR", "G4_Fe"] pdgs = [11, 22] processes = ["eIoni", "phot", "compt"] rows = [] row_idx = 0 for event_id in range(n_events): n_steps = int(rng.integers(2, 4)) for s in range(n_steps): pre_E = float(rng.uniform(50.0, 500.0)) n_sec = int(rng.integers(0, 3)) frac_dep = float(rng.uniform(0.05, 0.3)) frac_sec = float(rng.uniform(0.05, 0.2)) if n_sec > 0 else 0.0 frac_post = 1.0 - frac_dep - frac_sec edep = pre_E * frac_dep e_sec = pre_E * frac_sec post_E = pre_E * frac_post pre_pos = rng.uniform(-10, 10, size=3) step_length = float(rng.uniform(0.1, 5.0)) pre_dir = np.array([0.0, 0.0, 1.0]) post_dir = _unit(rng.normal(size=3)) post_pos = pre_pos + step_length * pre_dir sec_energies = list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else [] sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)] sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)] rows.append( { "event_id": event_id, "pdg": pdgs[row_idx % 2], "pre_x": pre_pos[0], "pre_y": pre_pos[1], "pre_z": pre_pos[2], "pre_E": pre_E, "pre_dx": pre_dir[0], "pre_dy": pre_dir[1], "pre_dz": pre_dir[2], "material": materials[row_idx % 2], "layer_id": s, "child_track_ids": list(range(n_sec)), "e_sec": e_sec, "process": processes[row_idx % 3], "step_length": step_length, "post_E": post_E, "edep": edep, "post_dx": post_dir[0], "post_dy": post_dir[1], "post_dz": post_dir[2], "post_x": post_pos[0], "post_y": post_pos[1], "post_z": post_pos[2], "sec_E_list": sec_energies, "sec_pdg_list": sec_pdgs, "sec_dx_list": [d[0] for d in sec_dirs], "sec_dy_list": [d[1] for d in sec_dirs], "sec_dz_list": [d[2] for d in sec_dirs], } ) row_idx += 1 pd.DataFrame(rows).to_parquet(path) return path def _tiny_cfg(**train_overrides): cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG) cfg["train"].update( { "epochs": 1, "batch_size": 8, "val_fraction": 0.2, "seed": 0, "warmup_epochs": 0, "validate_every": 0, "max_val_batches": 1, "wandb": False, } ) cfg["train"].update(train_overrides) cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}) cfg["stage2_model"].update( # decoder="autoregressive" is DEFAULT_CONFIG's default (v0.3.0 step 5) # and left as-is here on purpose, so this pipeline-level fixture # exercises the real default end-to-end against actual data. {"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0} ) cfg["conditioning"]["particle"]["emb_dim"] = 4 cfg["conditioning"]["material"]["emb_dim"] = 4 return cfg def _run(data, out_dir, cfg=None, **kwargs): echoed: list[str] = [] kwargs.setdefault("num_workers", 0) run_train_job( data=data, cfg=cfg or _tiny_cfg(), out_dir=out_dir, device=torch.device("cpu"), shuffle_buffer=64, echo=echoed.append, **kwargs, ) return echoed @pytest.fixture def data(tmp_path): return _make_synthetic_steps(tmp_path / "data.parquet", n_events=20) def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch): echo1 = _run(data, tmp_path / "out1") assert any("fitting normalizer (streaming)" in m for m in echo1) def _forbidden(*a, **k): raise AssertionError("should be served from cache, not recomputed") monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden) monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden) echo2 = _run(data, tmp_path / "out2") joined = "\n".join(echo2) assert "event index: cache hit" in joined assert "vocabulary maps: cache hit" in joined assert "normalizer: cache hit" in joined def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data): """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to "onehot" while conditioning.particle.type stays "physical" — a plain _tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea #29: no longer shared with any conditioning-side onehot map), cache it in the setup-cache sidecar, and persist it into the checkpoint's sec_type_topn_map key, with no extra config needed. pdg_topn_map (conditioning-only) stays unbuilt since conditioning.particle.type is "physical" here.""" echo1 = _run(data, tmp_path / "out1") assert any("building pdg top-N map" in m for m in echo1) loaded = setup_cache.load(data, [data]) assert loaded is not None # stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4 key = setup_cache.topn_key("pdg", 4) assert key in loaded.topn_maps assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22} ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False) assert ckpt.get("pdg_topn_map") is None assert "sec_type_topn_map" in ckpt assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"} echo2 = _run(data, tmp_path / "out2") assert any("pdg top-N map: cache hit" in m for m in echo2) def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data): """conditioning.particle.type="onehot" and stage2_model.particle_type.target="onehot" with different class counts (gitea #29's fix: stage2_model.particle_type.n_classes decouples the two) build two distinct top-N maps, cached under their own (axis, n_classes) key and persisted under two distinct checkpoint keys — no longer forced to share conditioning.particle.emb_dim.""" cfg = _tiny_cfg() cfg["conditioning"]["particle"]["type"] = "onehot" # emb_dim = 4, from _tiny_cfg cfg["stage2_model"]["particle_type"]["n_classes"] = 3 echo = _run(data, tmp_path / "out", cfg=cfg) assert any("mapped to 4 classes" in m for m in echo) assert any("mapped to 3 classes" in m for m in echo) loaded = setup_cache.load(data, [data]) assert loaded is not None cond_key = setup_cache.topn_key("pdg", 4) type_key = setup_cache.topn_key("pdg", 3) assert cond_key in loaded.topn_maps assert type_key in loaded.topn_maps ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) assert ckpt.get("pdg_topn_map") is not None assert ckpt.get("sec_type_topn_map") is not None def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data): """conditioning.material.type="onehot" is an independent axis from the pdg one above, with its own build/cache-hit branch in run_setup_stage — exercise both here the same way the pdg test above does.""" cfg = _tiny_cfg() cfg["conditioning"]["material"]["type"] = "onehot" echo1 = _run(data, tmp_path / "out1", cfg=cfg) assert any("building material top-N map" in m for m in echo1) loaded = setup_cache.load(data, [data]) assert loaded is not None key = setup_cache.topn_key("material", 4) # conditioning.material.emb_dim = 4 assert key in loaded.topn_maps assert set(loaded.topn_maps[key].class_map.keys()) >= {"G4_AIR", "G4_Fe"} ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False) assert "mat_topn_map" in ckpt assert set(ckpt["mat_topn_map"]["class_map"].keys()) >= {"G4_AIR", "G4_Fe"} echo2 = _run(data, tmp_path / "out2", cfg=cfg) assert any("material top-N map: cache hit" in m for m in echo2) 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} echo = _run(data, tmp_path / "out", cfg=cfg) assert not any("top-N map" in m for m in echo) loaded = setup_cache.load(data, [data]) assert loaded is not None assert loaded.topn_maps == {} @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork") def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(tmp_path, data, monkeypatch): # num_workers>0 makes DataLoader actually fork worker subprocesses # (unlike every other test here, which runs with num_workers=0) — pytest # itself is multi-threaded, hence Python's fork-safety warning below. monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2 echo = _run(data, tmp_path / "out", num_workers=3) assert any("num-workers=3" in m and "exceeds" in m for m in echo) @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork") def test_run_train_job_no_warning_when_num_workers_within_shared_quota(tmp_path, data, monkeypatch): monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2 echo = _run(data, tmp_path / "out", num_workers=2) assert not any("exceeds" in m for m in echo) def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data): _run(data, tmp_path / "out", cache_setup=False) assert not setup_cache.sidecar_path(data).exists() def test_run_train_job_rebuild_setup_cache_ignores_existing(tmp_path, data): files = [data] stale = setup_cache.SetupCache.empty(files) stale.vocab = ({999999: 0}, {"G4_AIR": 0}) # deliberately wrong setup_cache.save(data, files, stale) _run(data, tmp_path / "out", rebuild_setup_cache=True) loaded = setup_cache.load(data, files) assert loaded is not None assert loaded.vocab is not None assert set(loaded.vocab[0].keys()) == {11, 22} assert set(loaded.vocab[1].keys()) == {"G4_AIR", "G4_Fe"} def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypatch): _run(data, tmp_path / "out1", cfg=_tiny_cfg(val_fraction=0.1)) def _forbidden(*a, **k): raise AssertionError("vocab should be served from cache") monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden) echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3)) joined = "\n".join(echo2) assert "vocabulary maps: cache hit" in joined assert "fitting normalizer (streaming)" in joined def test_run_train_job_custom_k_max_end_to_end(tmp_path, data): """Regression: stage2_model.k_max other than the K_MAX module constant's default (15) must not produce a shape mismatch between the data pipeline (loader.py/transforms.py padding) and the model (network.py's trunks, sized from this same config value).""" cfg = _tiny_cfg() cfg["stage2_model"]["k_max"] = 3 _run(data, tmp_path / "out", cfg=cfg) ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) assert ckpt["model_config"]["stage2_model"]["k_max"] == 3 def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data): """Regression: conditioning.particle.type and conditioning.material.type are configured independently and may mix freely — e.g. particle "embedding" with material "physical" — end-to-end through the real data pipeline, not just accepted by validate_config.""" cfg = _tiny_cfg() cfg["conditioning"]["particle"]["type"] = "embedding" cfg["conditioning"]["material"]["type"] = "physical" _run(data, tmp_path / "out", cfg=cfg) ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) cond_cfg = ckpt["model_config"]["conditioning"] assert cond_cfg["particle"]["type"] == "embedding" assert cond_cfg["material"]["type"] == "physical" cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) # Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays # unfitted (mean=0/std=1) since "embedding" never computes real values # for it; the material block is fit for real under "physical". from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM assert cond_norm.mean is not None and cond_norm.std is not None np.testing.assert_allclose(cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0) np.testing.assert_allclose(cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0) material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :] assert np.all(material_std > 0) and not np.allclose(material_std, 1.0) def test_run_train_job_share_stages_end_to_end(tmp_path, data): """Regression: conditioning.share_stages = true must actually train (not raise NotImplementedError), and the resulting checkpoint's two stages must reload into a single shared ConditionEncoder instance rather than two independent ones.""" from giant.model.network import build_models cfg = _tiny_cfg() cfg["conditioning"]["share_stages"] = True _run(data, tmp_path / "out", cfg=cfg) ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) assert ckpt["model_config"]["conditioning"]["share_stages"] is True built = build_models(ckpt["model_config"]) stage1, stage2 = built["stage1"], built["stage2"] assert stage1 is not None and stage2 is not None assert stage1.cond_enc is stage2.cond_enc stage1.load_state_dict(ckpt["model"]) stage2.load_state_dict(ckpt["sec_decoder"]) for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()): assert torch.equal(p1, p2) def test_run_train_job_matches_uncached_output(tmp_path, data): _run(data, tmp_path / "uncached", cache_setup=False) _run(data, tmp_path / "cached1", cache_setup=True) _run(data, tmp_path / "cached2", cache_setup=True) # second is a cache hit uncached = torch.load(tmp_path / "uncached" / "last.pt", weights_only=False) cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False) for key in ("cond", "target", "sec_phys"): np.testing.assert_allclose(uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]) np.testing.assert_allclose(uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]) assert uncached["pdg_map"] == cached["pdg_map"] assert uncached["mat_map"] == cached["mat_map"] def _fitted_cond_norm(seed=0): rng = np.random.default_rng(seed) return Normalizer().fit(rng.normal(size=(64, COND_DIM)).astype(np.float32)) @pytest.mark.parametrize( "router_cfg", [ {"enabled": False, "type": "energy", "n_experts": 4}, {"enabled": True, "type": "pdg", "n_experts": 4}, ], ) def test_seed_energy_router_noop_when_not_an_enabled_energy_router(router_cfg): cond_norm = _fitted_cond_norm() echoed = [] _seed_energy_router(router_cfg, cond_norm, np.array([1.0, 2.0]), 3, echoed.append) assert "centers_init" not in router_cfg assert echoed == [] def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples(): router_cfg = {"enabled": True, "type": "energy", "n_experts": 4} cond_norm = _fitted_cond_norm() echoed = [] _seed_energy_router(router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append) assert "centers_init" not in router_cfg assert len(echoed) == 1 assert "falls back to default centers" in echoed[0] def test_seed_energy_router_seeds_centers_from_data_quantiles(): router_cfg = {"enabled": True, "type": "energy", "n_experts": 4} cond_norm = _fitted_cond_norm() energy_idx = 3 # A grid of "raw" quantile values as setup_cache.energy_quantiles_from_sample # would produce them: monotonically increasing, in the same (log-energy) # units as the conditioning column being normalized against. energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32) echoed = [] _seed_energy_router(router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append) assert "centers_init" in router_cfg centers = np.asarray(router_cfg["centers_init"], dtype=np.float32) assert centers.shape == (router_cfg["n_experts"],) levels = np.linspace(0.0, 1.0, router_cfg["n_experts"]) raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels) expected = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx] np.testing.assert_allclose(centers, expected, rtol=1e-5) # Quantile levels are increasing, and the normalizer's std is positive, so # the seeded centers must preserve that order rather than e.g. reversing it. assert np.all(np.diff(centers) > 0) assert len(echoed) == 1 assert "seeded EnergyRouter centers" in echoed[0]