giant/tests/test_pipeline.py

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.data.transforms import Normalizer
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,
            "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_pdg_topn_map(tmp_path, data):
    """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
    "onehot" (docs/v0.3.0-design.md §3.3/§8) — a plain _tiny_cfg() run must
    build the shared pdg top-N map, cache it in the setup-cache sidecar, and
    persist it into the checkpoint, with no extra config needed."""
    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
    key = setup_cache.topn_key("pdg", 4)  # conditioning.particle.emb_dim = 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 "pdg_topn_map" in ckpt
    assert set(ckpt["pdg_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_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 == {}


def test_run_train_job_warns_when_num_workers_exceeds_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=3)
    assert any("num-workers=3" in m and "exceeds" in m for m in echo)


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):
    """docs/v0.3.0-followups.md item 3 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):
    """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type
    and conditioning.material.type are configured independently and may mix
    freely (docs/v0.3.0-design.md §3.1) — 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):
    """docs/v0.3.0-followups.md item 5 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"]