from __future__ import annotations import argparse import csv import json import math import random import sys from pathlib import Path from typing import Any import numpy as np import yaml try: import torch from torch import nn from torch.utils.data import DataLoader, Dataset except ImportError as exc: # pragma: no cover - exercised in environment preflight raise SystemExit(f"PyTorch is required to run training: {exc}") from exc PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from model.FlowBenchFNO import build_model_from_config class FlowBenchTensorDataset(Dataset): def __init__( self, x: np.ndarray, y: np.ndarray, indices: np.ndarray, x_mean: np.ndarray | None = None, x_std: np.ndarray | None = None, y_mean: np.ndarray | None = None, y_std: np.ndarray | None = None, normalize: bool = True, ) -> None: self.x = x self.y = y self.indices = indices.astype(np.int64) self.normalize = normalize self.x_mean = x_mean self.x_std = x_std self.y_mean = y_mean self.y_std = y_std def __len__(self) -> int: return int(self.indices.shape[0]) def __getitem__(self, item: int) -> tuple[torch.Tensor, torch.Tensor]: idx = self.indices[item] x = self.x[idx] y = self.y[idx] if self.normalize: x = (x - self.x_mean) / self.x_std y = (y - self.y_mean) / self.y_std return torch.from_numpy(x.astype(np.float32)), torch.from_numpy(y.astype(np.float32)) 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(path: str | Path, root: Path = PROJECT_ROOT) -> Path: path = Path(path) return path if path.is_absolute() else 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_npz_array(path: Path, key: str) -> np.ndarray: if not path.exists(): raise FileNotFoundError(f"Missing data file: {path}") archive = np.load(path) actual_key = key if key in archive.files else archive.files[0] array = archive[actual_key] if array.ndim != 4: raise ValueError(f"Expected [N,C,H,W] array in {path}, got {array.shape}") return array.astype(np.float32, copy=False) def make_splits(n_samples: int, split_fraction: float, shuffle: bool, seed: int) -> tuple[np.ndarray, np.ndarray]: indices = np.arange(n_samples, dtype=np.int64) if shuffle: rng = np.random.default_rng(seed) rng.shuffle(indices) n_train = int(math.floor(n_samples * split_fraction)) return indices[:n_train], indices[n_train:] def channel_stats(array: np.ndarray, indices: np.ndarray) -> tuple[np.ndarray, np.ndarray]: subset = array[indices] mean = subset.mean(axis=(0, 2, 3), keepdims=False).astype(np.float32)[:, None, None] std = subset.std(axis=(0, 2, 3), keepdims=False).astype(np.float32)[:, None, None] std = np.maximum(std, np.asarray(1.0e-6, dtype=np.float32)) return mean, std def denormalize(tensor: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: return tensor * std + mean def compute_metrics(pred: torch.Tensor, target: torch.Tensor) -> dict[str, float]: diff = pred - target mse = torch.mean(diff.square()).item() mae = torch.mean(torch.abs(diff)).item() per_sample_num = torch.linalg.vector_norm(diff.flatten(1), dim=1) per_sample_den = torch.linalg.vector_norm(target.flatten(1), dim=1).clamp_min(1.0e-12) relative_l2 = torch.mean(per_sample_num / per_sample_den).item() return {"mse": mse, "relative_l2": relative_l2, "mae": mae} @torch.no_grad() def evaluate( model: nn.Module, loader: DataLoader, device: torch.device, y_mean: torch.Tensor, y_std: torch.Tensor, max_batches: int | None = None, ) -> dict[str, float]: model.eval() totals = {"mse": 0.0, "relative_l2": 0.0, "mae": 0.0} count = 0 for batch_idx, (x, y) in enumerate(loader): if max_batches is not None and batch_idx >= max_batches: break x = x.to(device) y = y.to(device) pred = model(x) pred_phys = denormalize(pred, y_mean, y_std) y_phys = denormalize(y, y_mean, y_std) metrics = compute_metrics(pred_phys, y_phys) batch_size = x.shape[0] for key, value in metrics.items(): totals[key] += value * batch_size count += batch_size if count == 0: return {key: float("nan") for key in totals} return {key: value / count for key, value in totals.items()} 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 FlowBench FNO baseline.") parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml")) parser.add_argument("--max-epochs", type=int, default=None, help="Override config training.num_epochs.") parser.add_argument("--sample-limit", type=int, default=None, help="Use the first N samples before splitting.") parser.add_argument("--device", default=None, help="Override experiment.device.") parser.add_argument("--checkpoint", default=None, help="Override best checkpoint path.") args = parser.parse_args() config = load_config(args.config) seed = int(config["experiment"].get("seed", 42)) set_seed(seed) device = select_device(args.device or config["experiment"].get("device", "auto")) data_cfg = config["data"] data_root = resolve_path(data_cfg["root_dir"]) x = load_npz_array(data_root / data_cfg["x_file"], data_cfg.get("x_key", "data")) y = load_npz_array(data_root / data_cfg["y_file"], data_cfg.get("y_key", "data")) if x.shape[0] != y.shape[0]: raise ValueError(f"X/Y sample counts differ: {x.shape[0]} vs {y.shape[0]}") if args.sample_limit is not None: x = x[: args.sample_limit] y = y[: args.sample_limit] model_cfg = config["model"] if x.shape[1] != int(model_cfg["in_channels"]): raise ValueError(f"Config in_channels={model_cfg['in_channels']} but X has {x.shape[1]} channels") if y.shape[1] != int(model_cfg["out_channels"]): raise ValueError(f"Config out_channels={model_cfg['out_channels']} but Y has {y.shape[1]} channels") train_idx, val_idx = make_splits( x.shape[0], float(data_cfg.get("split_fraction", 0.8)), bool(data_cfg.get("shuffle", True)), seed, ) x_mean, x_std = channel_stats(x, train_idx) y_mean, y_std = channel_stats(y, train_idx) normalize = bool(data_cfg.get("normalize", True)) train_ds = FlowBenchTensorDataset(x, y, train_idx, x_mean, x_std, y_mean, y_std, normalize) val_ds = FlowBenchTensorDataset(x, y, val_idx, x_mean, x_std, y_mean, y_std, normalize) 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) optimizer = torch.optim.Adam( model.parameters(), lr=float(config["training"]["learning_rate"]), weight_decay=float(config["training"].get("weight_decay", 0.0)), ) criterion = nn.MSELoss() y_mean_t = torch.from_numpy(y_mean).to(device) y_std_t = torch.from_numpy(y_std).to(device) x_mean_t = torch.from_numpy(x_mean) x_std_t = torch.from_numpy(x_std) num_epochs = args.max_epochs or int(config["training"]["num_epochs"]) best_path = resolve_path(args.checkpoint or config["training"]["checkpoint_path"]) last_path = resolve_path(config["training"]["last_checkpoint_path"]) metrics_path = resolve_path(config["training"]["metrics_path"]) history_path = resolve_path(config["training"]["history_path"]) best_path.parent.mkdir(parents=True, exist_ok=True) last_path.parent.mkdir(parents=True, exist_ok=True) history_path.parent.mkdir(parents=True, exist_ok=True) best_mse = float("inf") with open(history_path, "w", encoding="utf-8", newline="") as handle: writer = csv.DictWriter(handle, fieldnames=["epoch", "train_loss", "val_mse", "val_relative_l2", "val_mae"]) writer.writeheader() for epoch in range(1, num_epochs + 1): model.train() running_loss = 0.0 seen = 0 for x_batch, y_batch in train_loader: x_batch = x_batch.to(device) y_batch = y_batch.to(device) optimizer.zero_grad(set_to_none=True) pred = model(x_batch) loss = criterion(pred, y_batch) loss.backward() clip_norm = config["training"].get("grad_clip_norm") if clip_norm is not None: torch.nn.utils.clip_grad_norm_(model.parameters(), float(clip_norm)) optimizer.step() batch_size = x_batch.shape[0] running_loss += loss.item() * batch_size seen += batch_size train_loss = running_loss / max(1, seen) val_metrics = evaluate(model, val_loader, device, y_mean_t, y_std_t) row = { "epoch": epoch, "train_loss": train_loss, "val_mse": val_metrics["mse"], "val_relative_l2": val_metrics["relative_l2"], "val_mae": val_metrics["mae"], } writer.writerow(row) handle.flush() print( f"epoch={epoch:04d} train_loss={train_loss:.6e} " f"val_mse={val_metrics['mse']:.6e} " f"val_relative_l2={val_metrics['relative_l2']:.6e} " f"val_mae={val_metrics['mae']:.6e}", flush=True, ) checkpoint = { "epoch": epoch, "model_state_dict": model.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "best_val_mse": min(best_mse, val_metrics["mse"]), "config": config, "normalization": { "x_mean": x_mean_t.numpy(), "x_std": x_std_t.numpy(), "y_mean": torch.from_numpy(y_mean), "y_std": torch.from_numpy(y_std), }, "splits": {"train_idx": train_idx, "val_idx": val_idx}, "metrics": {"train_loss": train_loss, **val_metrics}, } torch.save(checkpoint, last_path) if val_metrics["mse"] < best_mse: best_mse = val_metrics["mse"] torch.save(checkpoint, best_path) final_payload = { "best_val_mse": best_mse, "best_checkpoint": str(best_path), "last_checkpoint": str(last_path), "history": str(history_path), "num_samples": int(x.shape[0]), "train_samples": int(train_idx.shape[0]), "val_samples": int(val_idx.shape[0]), "device": str(device), } save_json(metrics_path, final_payload) print(json.dumps(final_payload, indent=2), flush=True) if __name__ == "__main__": main()