giant/giant/pipeline.py

from dataclasses import dataclass
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from giant import config
from giant.constants import (
    COND_DIM,
    EMB_DIM,
    K_MAX,
    PARTICLE_PHYS_DIM,
    SEC_SLOT_DIM,
    X_DIM,
)
from giant.data import setup_cache
from giant.data.loader import (
    find_parquet_files,
    iter_file_chunks,
    build_index_maps_from_files,
    build_process_map_from_files,
)
from giant.data.transforms import (
    Normalizer,
    build_features,
    _WelfordAccumulator,
    _ReservoirSampler,
    sorted_membership,
)
from giant.data.dataset import make_event_split, StreamingStepsDataset
from giant.model.network import build_models, build_critics
from giant.train import train as run_training


@dataclass
class SetupStageResult:
    """Everything `run_train_job`'s pre-epoch setup stage derives from `data`.

    Also returned standalone by `run_setup_stage` for callers (e.g. `dwarf
    warm-cache`) that only want to populate/refresh the setup cache sidecar
    without actually training a model.
    """

    files: list[Path]
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
    proc_map: dict[str, int] | None
    cond_norm: Normalizer
    tgt_norm: Normalizer
    sec_phys_norm: Normalizer
    train_events: set
    val_events: set
    n_train_steps: int


def run_setup_stage(
    data: str | Path,
    val_fraction: float,
    seed: int,
    conditioning: str,
    router_cfg: dict,
    cache_setup: bool = True,
    rebuild_setup_cache: bool = False,
    echo=print,
) -> SetupStageResult:
    """Scan `data` for everything training needs before the epoch loop: the
    train/val event split, pdg/material vocab maps, an optional process map
    (`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers.

    Reads from and writes to the `giant.data.setup_cache` sidecar when
    `cache_setup` is set (`rebuild_setup_cache` ignores — but still
    refreshes — any existing sidecar content). `router_cfg` may be mutated
    in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`)
    gets its `centers_init` seeded from real data quantiles here.
    """
    files = find_parquet_files(data)
    echo(f"found {len(files)} parquet file(s)")

    cache: setup_cache.SetupCache | None = None
    if cache_setup:
        if rebuild_setup_cache:
            echo("setup cache: --rebuild-setup-cache given, recomputing all sections")
            loaded = None
        else:
            loaded = setup_cache.load(data, files, echo=echo)
        cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)

    if cache is not None and cache.event_index is not None:
        unique_ids, counts = cache.event_index
        echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
    else:
        echo("scanning event IDs …")
        unique_ids, counts = setup_cache.compute_event_index_from_files(files)
        if cache is not None:
            cache.event_index = (unique_ids, counts)

    train_events, val_events = make_event_split(
        unique_ids, val_fraction=val_fraction, seed=seed
    )
    events_arr = np.array(sorted(train_events))
    n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
    echo(
        f"  {int(counts.sum()):,} steps  |  "
        f"{len(train_events)} train events  |  "
        f"{len(val_events)} val events"
    )

    if cache is not None and cache.vocab is not None:
        pdg_map, mat_map = cache.vocab
        echo(
            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
            f"{len(mat_map)} materials)"
        )
    else:
        echo("building vocabulary maps …")
        pdg_map, mat_map = build_index_maps_from_files(files)
        echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
        if cache is not None:
            cache.vocab = (pdg_map, mat_map)

    proc_map: dict[str, int] | None = None
    if router_cfg.get("enabled") and router_cfg.get("type") == "process":
        n_experts = router_cfg["n_experts"]
        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
        if cached_proc_map is not None:
            proc_map = cached_proc_map
            echo(
                f"process vocabulary: cache hit ({len(proc_map)} labels, "
                f"{n_experts} experts)"
            )
        else:
            echo("building process vocabulary …")
            proc_map = build_process_map_from_files(files, n_experts=n_experts)
            echo(f"  {len(proc_map)} process labels mapped to {n_experts} experts")
            if cache is not None:
                cache.proc_maps[n_experts] = proc_map

    energy_router_active = (
        router_cfg.get("enabled") and router_cfg.get("type") == "energy"
    )
    energy_idx = router_cfg.get("energy_idx", 3)
    norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
    entry = cache.normalizers.get(norm_key) if cache is not None else None

    if entry is not None:
        echo(f"normalizer: cache hit (key={norm_key!r})")
        cond_norm = entry.cond_norm
        tgt_norm = entry.tgt_norm
        sec_phys_norm = entry.sec_phys_norm
        energy_sample = entry.energy_reservoir_sample
    else:
        echo("fitting normalizer (streaming) …")
        cond_acc = _WelfordAccumulator(COND_DIM)
        tgt_acc = _WelfordAccumulator(X_DIM)
        sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
        # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
        # the z-normalized energy column is roughly uniform, which real energy
        # spectra rarely are — collect a reservoir sample here (reusing this
        # same pass, not a second scan) so centers can instead be seeded from
        # actual data quantiles below. Collected whenever the setup cache is
        # being populated, not only when *this* run's router is
        # energy-typed, so a later run enabling --router-type energy against
        # this same (val_fraction, seed, conditioning) key never needs to
        # rescan just to seed centers.
        collect_energy_sample = energy_router_active or cache is not None
        energy_sampler = (
            _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
        )
        for path in files:
            for chunk in iter_file_chunks(path):
                mask = sorted_membership(chunk["event_id"], events_arr)
                if not mask.any():
                    continue
                chunk_tr = {k: v[mask] for k, v in chunk.items()}
                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
                    chunk_tr,
                    pdg_map,
                    mat_map,
                    proc_map=proc_map,
                    require_secondaries=True,
                    conditioning=conditioning,
                    sec_phys_only=True,
                )
                cond_acc.update(cond_cont)
                tgt_acc.update(target_s1)
                if energy_sampler is not None:
                    energy_sampler.update(cond_cont[:, energy_idx])
                sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
                sec_phys = sec_cont[:, :, 4:6][sec_valid]
                if len(sec_phys) > 0:
                    sec_phys_acc.update(sec_phys)
        cond_norm = cond_acc.to_normalizer()
        tgt_norm = tgt_acc.to_normalizer()
        sec_phys_norm = sec_phys_acc.to_normalizer()
        energy_sample = (
            energy_sampler.sample
            if energy_sampler is not None
            else np.empty(0, dtype=np.float32)
        )
        if cache is not None:
            cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_sample
            )

    if energy_router_active and energy_sample.size > 0:
        assert cond_norm.mean is not None and cond_norm.std is not None
        normalized_sample = (
            energy_sample - cond_norm.mean[energy_idx]
        ) / cond_norm.std[energy_idx]
        quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
        centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
        router_cfg["centers_init"] = centers_init.tolist()
        echo(
            f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
        )
    elif energy_router_active:
        echo(
            "  warning: no energy samples collected — EnergyRouter falls back to "
            "default centers"
        )

    if cache is not None:
        setup_cache.save(data, files, cache, echo=echo)

    return SetupStageResult(
        files=files,
        pdg_map=pdg_map,
        mat_map=mat_map,
        proc_map=proc_map,
        cond_norm=cond_norm,
        tgt_norm=tgt_norm,
        sec_phys_norm=sec_phys_norm,
        train_events=train_events,
        val_events=val_events,
        n_train_steps=n_train_steps,
    )


def run_train_job(
    data: Path,
    cfg: dict,
    out_dir: Path,
    device: torch.device,
    shuffle_buffer: int,
    num_workers: int,
    resume: Path | None = None,
    cache_setup: bool = True,
    rebuild_setup_cache: bool = False,
    echo=print,
) -> None:
    t, m = cfg["train"], cfg["model"]
    config.seed_everything(t["seed"])

    out_dir = Path(out_dir)

    router_cfg = m["router"]
    if t["mode"] == "wgan" and router_cfg.get("enabled"):
        raise ValueError(
            "--mode wgan does not support --router (no routed WGAN generator/"
            "critic exists) — disable one or the other"
        )

    conditioning = m["conditioning"]
    setup = run_setup_stage(
        data,
        val_fraction=t["val_fraction"],
        seed=t["seed"],
        conditioning=conditioning,
        router_cfg=router_cfg,
        cache_setup=cache_setup,
        rebuild_setup_cache=rebuild_setup_cache,
        echo=echo,
    )
    files = setup.files
    pdg_map, mat_map, proc_map = setup.pdg_map, setup.mat_map, setup.proc_map
    cond_norm, tgt_norm, sec_phys_norm = (
        setup.cond_norm,
        setup.tgt_norm,
        setup.sec_phys_norm,
    )
    train_events, val_events, n_train_steps = (
        setup.train_events,
        setup.val_events,
        setup.n_train_steps,
    )

    total_train_batches = n_train_steps // t["batch_size"]
    echo(f"  ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")

    train_ds = StreamingStepsDataset(
        files=files,
        split_events=train_events,
        pdg_map=pdg_map,
        mat_map=mat_map,
        cond_normalizer=cond_norm,
        target_normalizer=tgt_norm,
        batch_size=t["batch_size"],
        shuffle_buffer=shuffle_buffer,
        shuffle=True,
        proc_map=proc_map,
        conditioning=conditioning,
        sec_phys_normalizer=sec_phys_norm,
    )
    val_ds = StreamingStepsDataset(
        files=files,
        split_events=val_events,
        pdg_map=pdg_map,
        mat_map=mat_map,
        cond_normalizer=cond_norm,
        target_normalizer=tgt_norm,
        batch_size=t["batch_size"],
        shuffle=False,
        proc_map=proc_map,
        conditioning=conditioning,
        sec_phys_normalizer=sec_phys_norm,
    )

    pin = device.type == "cuda"
    train_loader = DataLoader(
        train_ds,
        batch_size=None,
        num_workers=num_workers,
        pin_memory=pin,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=None,
        num_workers=num_workers,
        pin_memory=pin,
    )

    emb_dim = m.get("emb_dim", EMB_DIM)
    expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims(
        router_cfg, m["hidden_dim"], m["n_blocks"]
    )
    if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != (
        m["hidden_dim"],
        m["n_blocks"],
    ):
        # Only reachable via an explicit router.expert_hidden_dim/n_blocks
        # override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/
        # ["n_blocks"] — see resolve_expert_dims), so this is never a false
        # positive from inheritance, only a deliberate narrow/wide-experts
        # config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't
        # reflect.
        echo(
            f"  warning: experts are {expert_hidden_dim}x{expert_n_blocks}, "
            f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x"
            f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, "
            "not the experts actually being trained"
        )

    model_config = {
        "pdg_vocab": len(pdg_map),
        "mat_vocab": len(mat_map),
        "hidden_dim": m["hidden_dim"],
        "n_blocks": m["n_blocks"],
        "emb_dim": emb_dim,
        "dropout": m["dropout"],
        "k_max": K_MAX,
        "sec_slot_dim": SEC_SLOT_DIM,
        "conditioning": conditioning,
        "router": dict(router_cfg),
        "expert_hidden_dim": expert_hidden_dim,
        "expert_n_blocks": expert_n_blocks,
        # Read by `predict`/`rollout` (which never receive their own --mode
        # flag) to auto-detect which sampler a checkpoint needs.
        "mode": t["mode"],
        "noise_dim": m.get("noise_dim", 64),
    }

    stage1_model, sec_decoder = build_models(model_config)
    echo(
        f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters  |  "
        f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
    )

    critic = None
    sec_critic = None
    if t["mode"] == "wgan":
        critic, sec_critic = build_critics(model_config)
        echo(
            f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters  |  "
            f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters"
        )

    out_dir.mkdir(parents=True, exist_ok=True)
    meta = config.build_run_meta(
        data=data,
        seed=t["seed"],
        n_pdg_codes=len(pdg_map),
        n_materials=len(mat_map),
        n_train_events=len(train_events),
        n_val_events=len(val_events),
        n_train_steps=n_train_steps,
    )
    config.save_config(cfg, out_dir, meta)

    run_training(
        stage1_model=stage1_model,
        sec_decoder=sec_decoder,
        train_loader=train_loader,
        val_loader=val_loader,
        mode=t["mode"],
        epochs=t["epochs"],
        lr=t["lr"],
        weight_decay=t["weight_decay"],
        ema_decay=t["ema_decay"],
        warmup_epochs=t["warmup_epochs"],
        device=device,
        out_dir=out_dir,
        lambda_nsec=t.get("lambda_nsec", 0.1),
        lambda_s2=t.get("lambda_s2", 1.0),
        lambda_balance=router_cfg.get("lambda_balance", 0.0),
        lambda_proc=router_cfg.get("lambda_proc", 0.0),
        normalizer_dict={
            "cond": cond_norm.to_dict(),
            "target": tgt_norm.to_dict(),
            "sec_phys": sec_phys_norm.to_dict(),
        },
        pdg_map={str(k): v for k, v in pdg_map.items()},
        mat_map={str(k): v for k, v in mat_map.items()},
        proc_map=proc_map,
        model_config=model_config,
        resume_path=resume,
        validate_every=t["validate_every"],
        validate_steps=t["validate_steps"],
        max_val_batches=t["max_val_batches"],
        total_train_batches=total_train_batches,
        critic=critic,
        sec_critic=sec_critic,
        n_critic=t.get("n_critic", 5),
        gp_weight=t.get("gp_weight", 10.0),
        critic_lr=t.get("critic_lr") or None,
        use_wandb=t.get("wandb", False),
        wandb_project=t.get("wandb_project", "giant"),
        wandb_run_name=t.get("wandb_run_name", ""),
        wandb_log_every=t.get("wandb_log_every", 50),
    )