import copy import numpy as np import pandas as pd import pytest import torch from giant import config as gconfig from giant.data import setup_cache from giant.pipeline import 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, } ) cfg["train"].update(train_overrides) cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0}) return cfg def _run(data, out_dir, cfg=None, **kwargs): echoed: list[str] = [] run_train_job( data=data, cfg=cfg or _tiny_cfg(), out_dir=out_dir, device=torch.device("cpu"), shuffle_buffer=64, num_workers=0, 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_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_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"]