scripts/train.py

from __future__ import annotations

import argparse
import csv
import json
import random
import sys
import time
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset
import yaml

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

from model.neural_convection_fno import build_model_from_config


class TensorPairDataset(Dataset):
    def __init__(self, x: torch.Tensor, y: torch.Tensor, indices: np.ndarray, stats: dict[str, torch.Tensor]) -> None:
        self.x = x
        self.y = y
        self.indices = indices.astype(np.int64)
        self.stats = stats

    def __len__(self) -> int:
        return int(self.indices.shape[0])

    def __getitem__(self, item: int) -> tuple[torch.Tensor, torch.Tensor]:
        idx = int(self.indices[item])
        return normalize_x(self.x[idx], self.stats), normalize_y(self.y[idx], self.stats)


def load_config(path: str | Path) -> dict[str, Any]:
    with open(path, "r", encoding="utf-8") as handle:
        return yaml.safe_load(handle)


def resolve(path: str | Path) -> Path:
    path = Path(path)
    return path if path.is_absolute() else PROJECT_ROOT / path


def set_seed(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 select_device(name: str) -> torch.device:
    if name == "auto":
        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
    return torch.device(name)


def load_pair(root: Path, split: str, x_file: str, y_file: str) -> tuple[torch.Tensor, torch.Tensor]:
    x_path = root / split / x_file
    y_path = root / split / y_file
    if not x_path.exists() or not y_path.exists():
        raise FileNotFoundError(f"Missing tensors for split {split}: {x_path}, {y_path}")
    x = torch.load(x_path, map_location="cpu")
    y = torch.load(y_path, map_location="cpu")
    if not isinstance(x, torch.Tensor) or not isinstance(y, torch.Tensor):
        raise TypeError(f"Expected torch.Tensor objects in {x_path} and {y_path}")
    x = x.float().contiguous()
    y = y.float().contiguous()
    if tuple(x.shape) != tuple(y.shape):
        raise ValueError(f"X/Y shape mismatch for {split}: {tuple(x.shape)} vs {tuple(y.shape)}")
    if x.ndim != 4 or x.shape[1:] != (4, 64, 64):
        raise ValueError(f"Expected [N,4,64,64] tensors for {split}, got {tuple(x.shape)}")
    return x, y


def make_train_val_indices(n: int, val_fraction: float, seed: int, shuffle: bool) -> tuple[np.ndarray, np.ndarray]:
    indices = np.arange(n, dtype=np.int64)
    if shuffle:
        rng = np.random.default_rng(seed)
        rng.shuffle(indices)
    n_val = int(round(n * val_fraction))
    n_train = n - n_val
    return indices[:n_train], indices[n_train:]


def channel_minmax(tensor: torch.Tensor, indices: np.ndarray) -> tuple[torch.Tensor, torch.Tensor]:
    subset = tensor[torch.as_tensor(indices, dtype=torch.long)]
    mins = subset.amin(dim=(0, 2, 3)).view(1, -1, 1, 1)
    maxs = subset.amax(dim=(0, 2, 3)).view(1, -1, 1, 1)
    return mins, maxs


def build_stats(x: torch.Tensor, y: torch.Tensor, train_idx: np.ndarray) -> dict[str, torch.Tensor]:
    x_min, x_max = channel_minmax(x, train_idx)
    y_min, y_max = channel_minmax(y, train_idx)
    return {
        "x_min": x_min,
        "x_max": x_max,
        "x_range": (x_max - x_min).clamp_min(1.0e-8),
        "y_min": y_min,
        "y_max": y_max,
        "y_range": (y_max - y_min).clamp_min(1.0e-8),
        "train_indices": torch.as_tensor(train_idx, dtype=torch.long),
    }


def normalize_x(x: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
    return (x - stats["x_min"].squeeze(0)) / stats["x_range"].squeeze(0)


def normalize_y(y: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
    return (y - stats["y_min"].squeeze(0)) / stats["y_range"].squeeze(0)


def stats_to_device(stats: dict[str, torch.Tensor], device: torch.device) -> dict[str, torch.Tensor]:
    return {key: value.to(device) if isinstance(value, torch.Tensor) else value for key, value in stats.items()}


def count_parameters(model: nn.Module) -> int:
    return sum(param.numel() for param in model.parameters() if param.requires_grad)


def tensor_to_numpy(tensor: torch.Tensor, dtype: np.dtype = np.float32) -> np.ndarray:
    """Convert tensors without relying on torch's NumPy ABI bridge."""
    return np.asarray(tensor.detach().cpu().tolist(), dtype=dtype)


def ssim_2d(pred: np.ndarray, target: np.ndarray) -> float:
    try:
        from skimage.metrics import structural_similarity

        return float(structural_similarity(target, pred, data_range=1.0))
    except Exception:
        mu_x = float(pred.mean())
        mu_y = float(target.mean())
        sigma_x = float(pred.var())
        sigma_y = float(target.var())
        sigma_xy = float(((pred - mu_x) * (target - mu_y)).mean())
        c1 = 0.01**2
        c2 = 0.03**2
        return float(((2 * mu_x * mu_y + c1) * (2 * sigma_xy + c2)) / ((mu_x**2 + mu_y**2 + c1) * (sigma_x + sigma_y + c2)))


@torch.no_grad()
def evaluate_loader(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    channel_names: list[str],
    collect: bool = False,
) -> tuple[dict[str, Any], dict[str, np.ndarray] | None]:
    model.eval()
    total_loss = 0.0
    total_count = 0
    sqerr_sum = torch.zeros(len(channel_names), dtype=torch.float64)
    elem_count = 0
    preds: list[np.ndarray] = []
    targets: list[np.ndarray] = []
    inputs: list[np.ndarray] = []
    start = time.perf_counter()
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
        pred = model(x)
        loss = torch.mean((pred - y) ** 2)
        batch = x.shape[0]
        total_loss += float(loss.item()) * batch
        total_count += batch
        sqerr_sum += (pred.detach().cpu().double() - y.detach().cpu().double()).square().sum(dim=(0, 2, 3))
        elem_count += batch * y.shape[-1] * y.shape[-2]
        if collect:
            preds.append(tensor_to_numpy(pred))
            targets.append(tensor_to_numpy(y))
            inputs.append(tensor_to_numpy(x))
    elapsed = time.perf_counter() - start
    mse = np.asarray((sqerr_sum / max(1, elem_count)).tolist(), dtype=np.float64)
    metrics: dict[str, Any] = {
        "loss": total_loss / max(1, total_count),
        "per_channel": {},
        "time_seconds_per_sample": elapsed / max(1, total_count),
    }
    pred_all = target_all = input_all = None
    if collect:
        pred_all = np.concatenate(preds, axis=0)
        target_all = np.concatenate(targets, axis=0)
        input_all = np.concatenate(inputs, axis=0)
    for idx, name in enumerate(channel_names):
        ssim = None
        if collect and pred_all is not None and target_all is not None:
            ssim = float(np.mean([ssim_2d(pred_all[i, idx], target_all[i, idx]) for i in range(pred_all.shape[0])]))
        metrics["per_channel"][name] = {"mse": float(mse[idx]), "ssim": ssim}
    arrays = None
    if collect and pred_all is not None and target_all is not None and input_all is not None:
        arrays = {"pred_norm": pred_all, "target_norm": target_all, "input_norm": input_all}
    return metrics, arrays


def save_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as handle:
        json.dump(payload, handle, indent=2)


def main() -> None:
    parser = argparse.ArgumentParser(description="Train compact FNO for 2D NeuralConvection.")
    parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml"))
    parser.add_argument("--epochs", type=int, default=None)
    parser.add_argument("--device", default=None)
    args = parser.parse_args()

    config = load_config(args.config)
    seed = int(config["experiment"].get("seed", 2026))
    set_seed(seed)
    device = select_device(args.device or config["experiment"].get("device", "auto"))
    print(f"device={device}", flush=True)

    data_cfg = config["data"]
    root = resolve(data_cfg["root_dir"])
    channel_names = list(data_cfg["channel_names"])
    x_train, y_train = load_pair(root, data_cfg["train_split"], data_cfg["x_file"], data_cfg["y_file"])
    train_idx, val_idx = make_train_val_indices(
        x_train.shape[0],
        float(data_cfg["validation_fraction"]),
        seed,
        bool(data_cfg.get("shuffle_train", True)),
    )
    stats = build_stats(x_train, y_train, train_idx)
    stats_path = resolve(data_cfg["stats_path"])
    stats_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save({**stats, "val_indices": torch.as_tensor(val_idx, dtype=torch.long)}, stats_path)
    print(
        f"loaded train tensors X/Y={tuple(x_train.shape)}; train={len(train_idx)} val={len(val_idx)}; stats={stats_path}",
        flush=True,
    )

    train_ds = TensorPairDataset(x_train, y_train, train_idx, stats)
    val_ds = TensorPairDataset(x_train, y_train, val_idx, stats)
    train_loader = DataLoader(
        train_ds,
        batch_size=int(config["training"]["batch_size"]),
        shuffle=True,
        num_workers=int(data_cfg.get("num_workers", 0)),
        pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda",
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=int(config["training"]["batch_size"]),
        shuffle=False,
        num_workers=int(data_cfg.get("num_workers", 0)),
        pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda",
    )

    model = build_model_from_config(config).to(device)
    print(f"model_parameters={count_parameters(model)}", flush=True)
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=float(config["training"]["learning_rate"]),
        weight_decay=float(config["training"]["weight_decay"]),
    )
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer,
        mode=config["training"]["scheduler"].get("mode", "min"),
        factor=float(config["training"]["scheduler"].get("factor", 0.5)),
        patience=int(config["training"]["scheduler"].get("patience", 20)),
    )
    criterion = nn.MSELoss()
    epochs = int(args.epochs or config["training"]["epochs"])
    grad_clip = float(config["training"]["grad_clip_norm"])
    best_val = float("inf")
    best_epoch = 0
    ckpt_path = resolve(config["training"]["checkpoint_path"])
    last_path = resolve(config["training"]["last_checkpoint_path"])
    ckpt_path.parent.mkdir(parents=True, exist_ok=True)
    history_path = resolve(config["training"]["history_path"])
    history_path.parent.mkdir(parents=True, exist_ok=True)

    with open(history_path, "w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "val_loss", "lr"])
        writer.writeheader()
        for epoch in range(1, epochs + 1):
            model.train()
            train_total = 0.0
            train_count = 0
            for x, y in train_loader:
                x = x.to(device, non_blocking=True)
                y = y.to(device, non_blocking=True)
                optimizer.zero_grad(set_to_none=True)
                pred = model(x)
                loss = criterion(pred, y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                optimizer.step()
                batch = x.shape[0]
                train_total += float(loss.item()) * batch
                train_count += batch
            train_loss = train_total / max(1, train_count)
            val_metrics, _ = evaluate_loader(model, val_loader, device, channel_names, collect=False)
            val_loss = float(val_metrics["loss"])
            scheduler.step(val_loss)
            lr = float(optimizer.param_groups[0]["lr"])
            writer.writerow({"epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "lr": lr})
            handle.flush()
            print(f"epoch={epoch:03d} train_loss={train_loss:.8e} val_loss={val_loss:.8e} lr={lr:.6e}", flush=True)
            payload = {
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "config": config,
                "best_val_loss": min(best_val, val_loss),
                "normalization_stats_path": str(stats_path),
            }
            torch.save(payload, last_path)
            if val_loss < best_val:
                best_val = val_loss
                best_epoch = epoch
                torch.save(payload, ckpt_path)
                print(f"saved_best_checkpoint={ckpt_path} best_val_loss={best_val:.8e}", flush=True)

    checkpoint = torch.load(ckpt_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    stats_device = stats_to_device(stats, device)
    del stats_device

    metrics_out: dict[str, Any] = {
        "best_epoch": best_epoch,
        "best_val_loss": best_val,
        "checkpoint_path": str(ckpt_path),
        "stats_path": str(stats_path),
        "splits": {"train": int(len(train_idx)), "val": int(len(val_idx))},
        "tests": {},
        "paper_table5_compact_fno_2d": config["paper_table5_compact_fno_2d"],
    }
    for label, split in data_cfg["test_splits"].items():
        x_test, y_test = load_pair(root, split, data_cfg["x_file"], data_cfg["y_file"])
        test_idx = np.arange(x_test.shape[0], dtype=np.int64)
        test_ds = TensorPairDataset(x_test, y_test, test_idx, stats)
        test_loader = DataLoader(
            test_ds,
            batch_size=int(config["inference"]["batch_size"]),
            shuffle=False,
            num_workers=int(data_cfg.get("num_workers", 0)),
            pin_memory=bool(data_cfg.get("pin_memory", True)) and device.type == "cuda",
        )
        collect = label == "in_distribution"
        test_metrics, arrays = evaluate_loader(model, test_loader, device, channel_names, collect=collect)
        metrics_out["tests"][label] = test_metrics
        print(f"test_split={label} loss={test_metrics['loss']:.8e} time_s_per_sample={test_metrics['time_seconds_per_sample']:.6e}", flush=True)
        for channel, vals in test_metrics["per_channel"].items():
            print(f"  {channel}: mse={vals['mse']:.8e} ssim={vals['ssim']}", flush=True)
        if collect and arrays is not None:
            np.savez_compressed(resolve(config["inference"]["output_path"]), **arrays, channel_names=np.array(channel_names))
    save_json(resolve(config["training"]["metrics_path"]), metrics_out)
    print(f"metrics_path={resolve(config['training']['metrics_path'])}", flush=True)


if __name__ == "__main__":
    main()