scripts/train.py

from __future__ import annotations

import argparse
import csv
import os
import sys
from pathlib import Path
from time import time

import numpy as np
import torch

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from model.ContinuousFlowOperator import (  # noqa: E402
    CFOBatch,
    ContinuousFlowOperator,
    quintic_spline_coefficients_np,
    rmse,
    trajectory_relative_l2,
)
from scripts.common import ensure_dirs, load_burgers_splits, load_config, write_json  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train Quintic CFO on 1D Burgers")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config" / "config.yaml"))
    parser.add_argument("--steps", type=int, default=None)
    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--eval-interval", type=int, default=None)
    parser.add_argument("--print-interval", type=int, default=None)
    parser.add_argument("--smoke", action="store_true")
    return parser.parse_args()


def device_from_runtime() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def prepare_splines(cfg: dict, train_data: np.ndarray, tspan: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    results_dir = Path(cfg["paths"]["results_dir"])
    cache_path = results_dir / "quintic_spline_train_seed43_ratio100.npz"
    if cfg["training"].get("spline_cache", True) and cache_path.exists():
        cached = np.load(cache_path)
        return cached["coeffs"], cached["starts"], cached["ends"]
    print("Preparing quintic spline coefficients for train trajectories...", flush=True)
    coeffs, starts, ends = quintic_spline_coefficients_np(train_data, tspan)
    if cfg["training"].get("spline_cache", True):
        np.savez_compressed(cache_path, coeffs=coeffs, starts=starts, ends=ends)
    return coeffs, starts.astype(np.float32), ends.astype(np.float32)


def make_batch(
    coeffs: np.ndarray,
    starts: np.ndarray,
    ends: np.ndarray,
    batch_size: int,
    device: torch.device,
    generator: torch.Generator,
) -> CFOBatch:
    idx = torch.randint(0, len(coeffs), (batch_size,), generator=generator)
    idx_np = idx.cpu().numpy()
    spline_coef = torch.from_numpy(coeffs[idx_np]).to(device=device, dtype=torch.float32)
    t_start = torch.from_numpy(starts[idx_np]).to(device=device, dtype=torch.float32)
    t_end = torch.from_numpy(ends[idx_np]).to(device=device, dtype=torch.float32)
    dt = t_end - t_start
    delta_t = torch.rand(batch_size, generator=generator).to(device=device, dtype=torch.float32) * dt
    eps = torch.randn(tuple(spline_coef[:, 0].shape), generator=generator).to(device=device, dtype=torch.float32)
    return CFOBatch(spline_coef=spline_coef, t_start=t_start, t_end=t_end, delta_t=delta_t, eps=eps)


@torch.no_grad()
def evaluate(model: ContinuousFlowOperator, data: np.ndarray, device: torch.device, steps_per_segment: int, batch_size: int) -> dict:
    preds = []
    for start in range(0, len(data), batch_size):
        chunk = torch.from_numpy(data[start : start + batch_size, 0]).to(device=device, dtype=torch.float32)
        pred = model.rollout(chunk, trajectory_points=data.shape[1], steps_per_segment=steps_per_segment)
        preds.append(pred.numpy())
    pred_arr = np.concatenate(preds, axis=0)
    return {
        "trajectory_relative_l2": trajectory_relative_l2(data, pred_arr),
        "rmse": rmse(data, pred_arr),
    }


def save_checkpoint(path: Path, model: ContinuousFlowOperator, optimizer: torch.optim.Optimizer, step: int, metrics: dict, cfg: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".tmp")
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "step": step,
            "metrics": metrics,
            "config": cfg,
        },
        tmp,
    )
    tmp.replace(path)


def main() -> None:
    args = parse_args()
    cfg = load_config(args.config)
    ensure_dirs(cfg)
    train_cfg = cfg["training"]
    infer_cfg = cfg["inference"]
    if args.steps is not None:
        train_cfg["steps"] = args.steps
    if args.batch_size is not None:
        train_cfg["batch_size"] = args.batch_size
    if args.eval_interval is not None:
        train_cfg["eval_interval"] = args.eval_interval
    if args.print_interval is not None:
        train_cfg["print_interval"] = args.print_interval
    if args.smoke:
        train_cfg["steps"] = min(int(train_cfg["steps"]), 2)
        train_cfg["batch_size"] = min(int(train_cfg["batch_size"]), 4)
        train_cfg["eval_interval"] = 1
        train_cfg["print_interval"] = 1

    np.random.seed(int(train_cfg["seed"]))
    torch.manual_seed(int(train_cfg["seed"]))
    device = device_from_runtime()
    print(f"Runtime device={device} torch={torch.__version__} cuda_available={torch.cuda.is_available()}", flush=True)

    splits, tspan = load_burgers_splits(cfg)
    train_data, val_data, test_data = splits["train"], splits["val"], splits["test"]
    print(f"Data split shapes train={train_data.shape} val={val_data.shape} test={test_data.shape}", flush=True)
    coeffs, starts, ends = prepare_splines(cfg, train_data, tspan)
    print(f"Spline dataset coeffs={coeffs.shape} starts={starts.shape} ends={ends.shape}", flush=True)

    model = ContinuousFlowOperator(input_size=cfg["model"]["input_size"], gamma=cfg["model"]["gamma"]).to(device)
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=float(train_cfg["learning_rate"]),
        betas=tuple(float(x) for x in train_cfg["betas"]),
    )
    gen = torch.Generator(device="cpu")
    gen.manual_seed(int(train_cfg["seed"]))

    history_path = Path(cfg["paths"]["results_dir"]) / "train_history.csv"
    best_ckpt = Path(train_cfg["best_checkpoint"])
    latest_ckpt = Path(train_cfg["latest_checkpoint"])
    best_val = float("inf")
    best_step = -1
    start_time = time()
    with history_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["step", "loss", "val_rel_l2", "val_rmse", "elapsed_sec"])
        writer.writeheader()
        for step in range(1, int(train_cfg["steps"]) + 1):
            model.train()
            batch = make_batch(coeffs, starts, ends, int(train_cfg["batch_size"]), device, gen)
            optimizer.zero_grad(set_to_none=True)
            loss = model.loss_from_batch(batch)
            loss.backward()
            optimizer.step()

            row = {
                "step": step,
                "loss": float(loss.detach().cpu()),
                "val_rel_l2": "",
                "val_rmse": "",
                "elapsed_sec": time() - start_time,
            }
            do_eval = (step % int(train_cfg["eval_interval"]) == 0) or args.smoke
            if do_eval:
                metrics = evaluate(
                    model,
                    val_data,
                    device,
                    steps_per_segment=int(infer_cfg["steps_per_segment"]),
                    batch_size=int(infer_cfg["eval_batch_size"]),
                )
                row["val_rel_l2"] = metrics["trajectory_relative_l2"]
                row["val_rmse"] = metrics["rmse"]
                if metrics["trajectory_relative_l2"] < best_val:
                    best_val = metrics["trajectory_relative_l2"]
                    best_step = step
                    save_checkpoint(best_ckpt, model, optimizer, step, metrics, cfg)
                print(
                    f"step={step} loss={row['loss']:.8e} val_rel_l2={metrics['trajectory_relative_l2']:.8e} "
                    f"val_rmse={metrics['rmse']:.8e} best_step={best_step} best_val={best_val:.8e}",
                    flush=True,
                )
            elif step % int(train_cfg["print_interval"]) == 0 or step == 1:
                print(f"step={step} loss={row['loss']:.8e} elapsed_sec={row['elapsed_sec']:.1f}", flush=True)
            writer.writerow(row)
            f.flush()
    final_metrics = evaluate(
        model,
        test_data,
        device,
        steps_per_segment=int(infer_cfg["steps_per_segment"]),
        batch_size=int(infer_cfg["eval_batch_size"]),
    )
    save_checkpoint(latest_ckpt, model, optimizer, int(train_cfg["steps"]), final_metrics, cfg)
    summary = {
        "status": "success",
        "steps": int(train_cfg["steps"]),
        "best_step": best_step,
        "best_val_trajectory_relative_l2": best_val,
        "latest_test_trajectory_relative_l2": final_metrics["trajectory_relative_l2"],
        "latest_test_rmse": final_metrics["rmse"],
        "best_checkpoint": str(best_ckpt),
        "latest_checkpoint": str(latest_ckpt),
    }
    write_json(Path(cfg["paths"]["results_dir"]) / "train_summary.json", summary)
    print("Training complete.", flush=True)
    print(summary, flush=True)


if __name__ == "__main__":
    main()