stress_operator/train.py

"""Single training entrypoint; ``--config`` selects the stage.

Faithful to Transolver ``exp_elas.py``:
  - AdamW (lr, weight_decay), CosineAnnealingLR(T_max=epochs)
  - batch_size 1, gradient clipping at max_grad_norm (0.1)
  - loss = relative-L2 in physical units: predictions are de-normalized before the loss;
    targets are physical (decode(encode(s)) == s, so storing physical targets is equivalent).

Writes a run-log JSON to ``results/`` per master plan §0.2.
"""
from __future__ import annotations

import argparse
import os
import time
from typing import Any, Dict, Optional

import torch
import yaml

from .data.dataset import build_splits
from .losses.relative_l2 import relative_l2
from .models.transolver import build_model, count_parameters
from .seeds import set_seed
from .utils.logging import MODAL_RATES_PER_SEC, write_run_log


def load_config(path: str) -> Dict[str, Any]:
    with open(path) as f:
        return yaml.safe_load(f)


def run_training(
    config: Dict[str, Any],
    seed: int,
    data_dir: str,
    device: Optional[str] = None,
    gpu_name: str = "CPU",
    results_path: Optional[str] = None,
    ckpt_path: Optional[str] = None,
    log_every: int = 50,
    max_epochs: Optional[int] = None,
    ntrain_override: Optional[int] = None,
    splits=None,
) -> Dict[str, Any]:
    """Train one model for one seed; return final metrics and write a run-log JSON.

    If ``ckpt_path`` is given, also save ``{state_dict, normalizer{mean,std}, config, seed,
    metrics}`` (the normalizer stats are required to de-normalize predictions at inference).
    If ``splits`` (a ``Splits`` from ``build_splits_from_indices``) is given, it overrides the
    default first-1000/last-200 split (used by the OOD evaluation).
    """
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    set_seed(seed)

    data_cfg = config["data"]
    train_cfg = config["train"]
    model_cfg = config["model"]

    if splits is None:
        ntrain = ntrain_override or data_cfg.get("ntrain", 1000)
        ntest = data_cfg.get("ntest", 200)
        splits = build_splits(data_dir, ntrain=ntrain, ntest=ntest)
    ntest = splits.test_coords.shape[0]
    normalizer = splits.normalizer.to(device)

    # GPU-resident dataset: the whole thing is tiny (~10 MB), so we keep it on-device and
    # batch by index. This removes DataLoader + per-iteration host->device + per-iteration
    # .item() sync overhead, which dominates wall-clock at batch_size 1. The math is identical
    # to the DataLoader path (same batch_size, same loss, same seeded shuffle order).
    batch_size = train_cfg.get("batch_size", 1)
    eval_every = int(train_cfg.get("eval_every", 1))

    def _3d(t):
        return (t if t.dim() == 3 else t.unsqueeze(-1)).to(device)

    train_coords = splits.train_coords.to(device)   # (ntrain, 972, 2)
    train_sigma = _3d(splits.train_sigma)            # (ntrain, 972, 1) physical
    test_coords = splits.test_coords.to(device)
    test_sigma = _3d(splits.test_sigma)
    ntrain_eff = train_coords.shape[0]

    base_model = build_model(model_cfg).to(device)
    n_params = count_parameters(base_model)

    # Optional torch.compile (CUDA graphs) to cut per-iteration kernel-launch overhead, which
    # dominates wall-clock at batch_size 1. Same math, static input shape (1, 972, 2). The
    # checkpoint is saved from base_model so its state_dict keys stay clean (no _orig_mod prefix).
    model = base_model
    if bool(train_cfg.get("compile", False)) and device == "cuda":
        try:
            model = torch.compile(base_model, mode="reduce-overhead")
            print(f"[seed {seed}] torch.compile enabled (reduce-overhead)", flush=True)
        except Exception as e:  # pragma: no cover
            print(f"[seed {seed}] torch.compile failed ({e}); falling back to eager", flush=True)
            model = base_model

    lr = float(train_cfg.get("lr", 1e-3))
    wd = float(train_cfg.get("weight_decay", 1e-5))
    betas = tuple(train_cfg.get("betas", (0.9, 0.999)))
    epochs = max_epochs or int(train_cfg.get("epochs", 500))
    max_grad_norm = train_cfg.get("max_grad_norm", None)

    optimizer = torch.optim.AdamW(base_model.parameters(), lr=lr, weight_decay=wd, betas=betas)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

    shuffle_gen = torch.Generator().manual_seed(seed)  # reproducible per-epoch shuffle

    @torch.no_grad()
    def eval_test() -> float:
        model.eval()
        total = 0.0
        for i in range(0, ntest, batch_size):
            out = normalizer.decode(model(test_coords[i:i + batch_size], None))
            total += relative_l2(out, test_sigma[i:i + batch_size], reduction="sum").item()
        return total / ntest

    t0 = time.time()
    best_rel = float("inf")
    test_rel = float("nan")
    history = []
    for ep in range(epochs):
        model.train()
        perm = torch.randperm(ntrain_eff, generator=shuffle_gen).to(device)
        running = torch.zeros((), device=device)
        for s in range(0, ntrain_eff, batch_size):
            idx = perm[s:s + batch_size]
            optimizer.zero_grad()
            out = normalizer.decode(model(train_coords[idx], None))  # -> physical
            loss = relative_l2(out, train_sigma[idx], reduction="sum")
            loss.backward()
            if max_grad_norm is not None:
                torch.nn.utils.clip_grad_norm_(base_model.parameters(), max_grad_norm)
            optimizer.step()
            running += loss.detach()
        scheduler.step()
        train_rel = (running / ntrain_eff).item()

        if (ep % eval_every == 0) or (ep >= epochs - 5):
            test_rel = eval_test()
            best_rel = min(best_rel, test_rel)
        history.append({"epoch": ep, "train_rel": train_rel, "test_rel": test_rel})
        if ep % log_every == 0 or ep == epochs - 1:
            print(
                f"[seed {seed}] epoch {ep:4d}  train_rel={train_rel:.5f}  test_rel={test_rel:.5f}",
                flush=True,
            )

    wall = time.time() - t0
    rate = MODAL_RATES_PER_SEC.get(gpu_name, 0.0)
    est_cost = wall * rate

    final_metrics = {
        "test_rel_l2": round(test_rel, 6),
        "best_test_rel_l2": round(best_rel, 6),
        "train_rel_l2": round(train_rel, 6),
        "n_params": n_params,
        "epochs": epochs,
    }

    if ckpt_path is not None:
        os.makedirs(os.path.dirname(ckpt_path) or ".", exist_ok=True)
        torch.save(
            {
                "state_dict": base_model.state_dict(),
                "normalizer": {
                    "mean": normalizer.mean.detach().cpu(),
                    "std": normalizer.std.detach().cpu(),
                },
                "config": config,
                "seed": seed,
                "metrics": final_metrics,
            },
            ckpt_path,
        )
        print(f"[seed {seed}] saved checkpoint -> {ckpt_path}")

    if results_path is None:
        os.makedirs("results", exist_ok=True)
        results_path = os.path.join("results", f"{config.get('name','run')}_seed{seed}.json")
    write_run_log(
        path=results_path,
        config=config,
        seed=seed,
        final_metrics=final_metrics,
        wall_clock_sec=wall,
        gpu=gpu_name,
        est_cost_usd=est_cost,
        extra={"history_tail": history[-5:]},
    )
    print(
        f"[seed {seed}] DONE  test_rel_l2={test_rel:.6f}  best={best_rel:.6f}  "
        f"params={n_params}  wall={wall:.0f}s  gpu={gpu_name}  est_cost=${est_cost:.4f}"
    )
    return final_metrics


def main() -> int:
    ap = argparse.ArgumentParser(description="Train the stress operator (one seed).")
    ap.add_argument("--config", required=True)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--data-dir", default="data")
    ap.add_argument("--device", default=None)
    ap.add_argument("--gpu-name", default="CPU", help="for cost accounting (A10 / A100-40GB / CPU)")
    ap.add_argument("--results-path", default=None)
    ap.add_argument("--max-epochs", type=int, default=None, help="override epochs (local smoke runs)")
    ap.add_argument("--ntrain", type=int, default=None, help="override ntrain (local smoke runs)")
    args = ap.parse_args()

    config = load_config(args.config)
    run_training(
        config=config,
        seed=args.seed,
        data_dir=args.data_dir,
        device=args.device,
        gpu_name=args.gpu_name,
        results_path=args.results_path,
        max_epochs=args.max_epochs,
        ntrain_override=args.ntrain,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())