giant/giant/config.py

import hashlib
import random
import subprocess
import sys
import tomllib
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch

DEFAULT_CONFIG: dict = {
    "train": {
        "mode": "flow",
        "epochs": 100,
        "batch_size": 4096,
        "lr": 3e-4,
        "weight_decay": 0.01,  # AdamW default — exposed so it can be tuned
        "ema_decay": 0.9999,  # EMA of model weights for sampling; 0 disables
        # per-epoch val loss (not the marginal/KL validate_every pass) is
        # capped to this many batches; 0 = full val set every epoch
        "max_val_batches": 200,
        "val_fraction": 0.1,
        "num_workers": 4,
        "seed": 0,
        "validate_every": 10,
        "validate_steps": 10,
        "warmup_epochs": 5,
        "lambda_nsec": 0.1,
        "lambda_s2": 1.0,
        # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic:
        # critic updates per generator update. gp_weight: gradient-penalty
        # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use
        # `lr`" — not None, since save_config's TOML writer has no null
        # literal to round-trip.
        "n_critic": 5,
        "gp_weight": 10.0,
        "critic_lr": 0.0,
        # Weights & Biases per-epoch metric logging (opt-in; see giant.train).
        # "" for wandb_run_name means "use the checkpoint out_dir name" — not
        # None, since save_config's TOML writer has no null literal to
        # round-trip.
        "wandb": False,
        "wandb_project": "giant",
        "wandb_run_name": "",
        # Batch-granularity metrics (loss/grad_norm/lr) are logged every N
        # optimizer steps, not every batch — a single epoch can be tens of
        # thousands of steps (see steps_per_epoch above), and logging every
        # one of them would flood the run with points the UI has to downsample
        # anyway. Per-epoch metrics (the metrics.csv row) always log in full.
        "wandb_log_every": 50,
    },
    "model": {
        "hidden_dim": 256,
        "n_blocks": 6,
        "emb_dim": 16,
        "dropout": 0.1,
        # WGAN generator noise-vector width (mode == "wgan" only).
        "noise_dim": 64,
        # "physical" conditions on material/particle physical properties via
        # a small MLP (giant.model.network.ConditionEncoder); "embedding"
        # keeps the original learned pdg/material embedding tables — kept
        # available as the generalization-comparison baseline. Checkpoints
        # from before this option existed have no "conditioning" key and
        # load as "embedding" (see giant.model.network.build_models).
        "conditioning": "physical",
        "router": {
            "enabled": False,
            "type": "energy",  # selects the Router impl from ROUTER_REGISTRY
            "n_experts": 4,
            # 0 means "inherit model.hidden_dim/n_blocks" (see
            # resolve_expert_dims below) — not a fixed 128/3, which silently
            # ignored --hidden-dim/--n-blocks whenever routing was enabled.
            # TOML has no null literal to round-trip (same pattern as
            # critic_lr/wandb_run_name above), hence 0 rather than None.
            "expert_hidden_dim": 0,
            "expert_n_blocks": 0,
            "temperature": 0.5,  # energy/pdg-router kwarg
            "learn_centers": True,  # energy/pdg-router kwarg
            "lambda_balance": 0.0,  # optional load-balance aux loss weight
            "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
            "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
            "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
            # (0.0 still trains a working router — the gate gets gradient
            # through the downstream flow loss like EnergyRouter's centers —
            # but only lambda_proc > 0 grounds it in the true `process` label)
            # type = "composed" routes on multiple axes at once (e.g. energy x
            # pdg), each with its own expert count/hyperparameters. Axes are
            # NOT in these defaults (there's no meaningful default axis list)
            # — set them as flat axis{i}_{field} keys instead of "n_experts",
            # e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type =
            # "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See
            # giant.model.network._parse_composed_axes / `--router-axis`.
        },
    },
}


def git_hash() -> str:
    try:
        return (
            subprocess.check_output(
                ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
            )
            .decode()
            .strip()
        )
    except Exception:
        return "unknown"


def auto_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


# Calibration point for estimate_batch_size(training=True): hidden_dim=1024,
# n_blocks=6, batch_size=29696 measured at ~7683 MiB VRAM (post-Phase-2
# architecture, including the Stage-2 secondary decoder and n_sec head).
# Activation memory is assumed to scale linearly with
# batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this
# is a rough estimate rather than a guaranteed bound.
_REF_BYTES = 7683 * 1024**2
_REF_BATCH_SIZE = 29696
_REF_HIDDEN_DIM = 1024
_REF_N_BLOCKS = 6

# Calibration point for estimate_batch_size(training=False): inference has no
# backward graph or optimizer state, so its memory footprint is much smaller
# per sample. hidden_dim=1024, n_blocks=8, batch_size=65536 measured at ~2037
# MiB VRAM.
_REF_BYTES_PREDICT = 2037 * 1024**2
_REF_BATCH_SIZE_PREDICT = 65536
_REF_HIDDEN_DIM_PREDICT = 1024
_REF_N_BLOCKS_PREDICT = 8


def estimate_batch_size(
    hidden_dim: int,
    n_blocks: int,
    device: torch.device,
    safety_factor: float = 0.8,
    min_batch_size: int = 1024,
    training: bool = True,
) -> int:
    """Estimate a batch size that fits in the free memory on `device`.

    Only supported on CUDA devices, which expose a free/total memory query;
    other backends (cpu, mps) raise ValueError. Pass `training=False` for
    inference (e.g. `predict`), which uses a much lower per-sample memory
    calibration since there's no backward graph or optimizer state.
    """
    if device.type != "cuda":
        raise ValueError(
            f"--batch-size auto is only supported on cuda devices, got {device.type!r}"
        )
    device_index = (
        device.index if device.index is not None else torch.cuda.current_device()
    )
    free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index)
    if training:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES,
            _REF_BATCH_SIZE,
            _REF_HIDDEN_DIM,
            _REF_N_BLOCKS,
        )
    else:
        ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = (
            _REF_BYTES_PREDICT,
            _REF_BATCH_SIZE_PREDICT,
            _REF_HIDDEN_DIM_PREDICT,
            _REF_N_BLOCKS_PREDICT,
        )
    bytes_per_unit = ref_bytes / (ref_batch_size * ref_hidden_dim * ref_n_blocks)
    bytes_per_sample = bytes_per_unit * hidden_dim * n_blocks
    batch_size = int(free_bytes * safety_factor / bytes_per_sample)
    batch_size = max(min_batch_size, (batch_size // 1024) * 1024)
    return batch_size


def load_toml(path: Path) -> dict:
    with open(path, "rb") as f:
        return tomllib.load(f)


def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
    """Warn (don't fail) if a config.toml's [meta].git_hash predates the current checkout.

    A config saved by a previous run may have been produced by code that has
    since changed, so its hyperparameters might not mean what they used to —
    surface that as a heads-up rather than blocking the rerun.
    """
    file_hash = file_cfg.get("meta", {}).get("git_hash")
    current_hash = git_hash()
    if not file_hash or file_hash == "unknown" or current_hash == "unknown":
        return
    if file_hash != current_hash:
        print(
            f"warning: {config_path} was generated at git commit {file_hash}, "
            f"but the current checkout is at {current_hash} — hyperparameters "
            "may not match the code that originally produced this config",
            file=sys.stderr,
        )


def load_checkpoint_config(ckpt_path: str | Path) -> dict:
    """Load the full ``[train]``/``[model]``/``[meta]`` config.toml written
    alongside a checkpoint by ``save_config``.

    Returns ``{}`` if no config.toml sits next to the checkpoint (older runs,
    or a checkpoint moved without its sidecar) — this is best-effort
    provenance for threading into a rollout's YAML sidecar, not a hard
    requirement for using the checkpoint itself.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return {}
    return load_toml(config_path)


def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
    """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.

    Training writes config.toml into the same out_dir as its checkpoints, so a
    checkpoint loaded later (for `predict` or `giant.analysis`) can be
    cross-checked the same way `--config` loading is, without the caller having
    to pass the toml path explicitly. Silently does nothing if no config.toml
    is found alongside the checkpoint.
    """
    config_path = Path(ckpt_path).parent / "config.toml"
    if not config_path.exists():
        return
    warn_if_git_hash_mismatch(load_toml(config_path), config_path)


def merge_cli_overrides(
    defaults: dict,
    config_path: Path | None,
    train_overrides: dict,
    model_overrides: dict,
) -> dict:
    """Resolve config as defaults -> TOML file -> explicit CLI flags.

    `model.router` is deep-merged one level (rather than replaced wholesale)
    at each stage, so a TOML file or CLI flag only overriding e.g.
    `router.enabled` doesn't drop the rest of the router defaults.
    """
    cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])}
    cfg["model"]["router"] = dict(defaults["model"]["router"])
    if config_path is not None:
        file_cfg = load_toml(config_path)
        cfg["train"].update(file_cfg.get("train", {}))
        file_model = dict(file_cfg.get("model", {}))
        file_router = file_model.pop("router", None)
        cfg["model"].update(file_model)
        if file_router:
            cfg["model"]["router"].update(file_router)
        warn_if_git_hash_mismatch(file_cfg, config_path)
    model_overrides = dict(model_overrides)
    router_overrides = model_overrides.pop("router", None)
    cfg["train"].update(train_overrides)
    cfg["model"].update(model_overrides)
    if router_overrides:
        cfg["model"]["router"].update(router_overrides)
    return cfg


def resolve_expert_dims(
    router_cfg: dict, hidden_dim: int, n_blocks: int
) -> tuple[int, int]:
    """Resolve a router's expert hidden_dim/n_blocks, inheriting from the
    monolith's when left at the 0 ("unset") sentinel.

    Used by both `giant.pipeline` (to build the checkpoint's `model_config`)
    and `giant.cli`'s batch-size auto-estimate, so `--hidden-dim`/`--n-blocks`
    size the experts the same way in both places unless
    `router.expert_hidden_dim`/`expert_n_blocks` are explicitly overridden.
    """
    expert_hidden_dim = router_cfg.get("expert_hidden_dim") or hidden_dim
    expert_n_blocks = router_cfg.get("expert_n_blocks") or n_blocks
    return expert_hidden_dim, expert_n_blocks


_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb"}


# Priority-ordered candidate fields for default_out_dir_name: (label, getter,
# formatter). `getter(train, model)` returns None when the field is at its
# default (and so should be omitted); otherwise formatter(value) renders the
# name token. The router is a single unit gated on `router.enabled` rather
# than one candidate per router key, since its type/n_experts are meaningless
# while disabled.
def _mode_candidate(train, model):
    return None if train["mode"] == DEFAULT_CONFIG["train"]["mode"] else train["mode"]


def _router_candidate(train, model):
    router = model["router"]
    if router["enabled"] == DEFAULT_CONFIG["model"]["router"]["enabled"]:
        return None
    return f"r-{router['type']}{router['n_experts']}"


def _conditioning_candidate(train, model):
    if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]:
        return None
    code = _CONDITIONING_CODE.get(model["conditioning"], model["conditioning"])
    return f"c{code}"


def _default_field_candidate(section_key, field, prefix):
    def _candidate(train, model):
        section = train if section_key == "train" else model
        value = section[field]
        if value == DEFAULT_CONFIG[section_key][field]:
            return None
        return f"{prefix}{value}"

    return _candidate


_OUT_DIR_NAME_CANDIDATES = [
    ("mode", _mode_candidate),
    ("router", _router_candidate),
    ("conditioning", _conditioning_candidate),
    ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")),
    ("n_blocks", _default_field_candidate("model", "n_blocks", "b")),
    ("emb_dim", _default_field_candidate("model", "emb_dim", "e")),
    ("lr", _default_field_candidate("train", "lr", "lr")),
    ("batch_size", _default_field_candidate("train", "batch_size", "bs")),
    ("seed", _default_field_candidate("train", "seed", "seed")),
    ("epochs", _default_field_candidate("train", "epochs", "ep")),
]

_OUT_DIR_NAME_MAX_FIELDS = 6


def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
    """Build a default checkpoint out_dir name from what's non-default in `cfg`.

    Only fields that differ from DEFAULT_CONFIG are included, so a fully
    default run's name is just its timestamp — see
    `_OUT_DIR_NAME_CANDIDATES` for the fixed, priority-ordered field list.
    Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder
    collapse into a short deterministic hash suffix rather than growing the
    name unboundedly. This name doubles as the run's W&B id (see
    giant.train), which is the reason a timestamp is always included.
    """
    now = now or datetime.now()
    train, model = cfg["train"], cfg["model"]
    tokens = []
    overflow = []
    for label, candidate in _OUT_DIR_NAME_CANDIDATES:
        token = candidate(train, model)
        if token is None:
            continue
        if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS:
            tokens.append(token)
        else:
            overflow.append(f"{label}={token}")

    name = now.strftime("%Y%m%d_%H%M")
    if tokens:
        name += "_" + "_".join(tokens)
    if overflow:
        digest = hashlib.md5("|".join(sorted(overflow)).encode()).hexdigest()[:6]
        name += f"_+{len(overflow)}more-{digest}"
    return name


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _toml_value(v) -> str:
    if isinstance(v, bool):
        return "true" if v else "false"
    if isinstance(v, str):
        return repr(v)
    return str(v)


def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
    lines = []
    # One-level-nested dict values (e.g. model.router) are rendered as their
    # own [section.subsection] table after the parent section, since TOML
    # doesn't accept a bare dict as a `key = value` scalar line.
    nested_sections: list[tuple[str, dict]] = []
    for section, values in cfg.items():
        lines.append(f"[{section}]")
        for k, v in values.items():
            if isinstance(v, dict):
                nested_sections.append((f"{section}.{k}", v))
                continue
            lines.append(f"{k:<14} = {_toml_value(v)}")
        lines.append("")

    for name, values in nested_sections:
        lines.append(f"[{name}]")
        for k, v in values.items():
            lines.append(f"{k:<14} = {_toml_value(v)}")
        lines.append("")

    lines.append("[meta]")
    for k, v in meta.items():
        lines.append(f"{k:<14} = {_toml_value(v)}")

    (out_dir / "config.toml").write_text("\n".join(lines))


def build_run_meta(
    data: Path,
    seed: int,
    n_pdg_codes: int,
    n_materials: int,
    n_train_events: int,
    n_val_events: int,
    n_train_steps: int,
) -> dict:
    return {
        "git_hash": git_hash(),
        "seed": seed,
        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
        "data_path": str(data),
        "n_pdg_codes": n_pdg_codes,
        "n_materials": n_materials,
        "n_train_events": n_train_events,
        "n_val_events": n_val_events,
        "n_train_steps": n_train_steps,
    }