| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Any |
|
|
| import numpy as np |
| import torch |
| 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 |
| from scripts.train import load_pair, normalize_x, normalize_y, ssim_2d, tensor_to_numpy |
|
|
|
|
| class InferenceDataset(Dataset): |
| def __init__(self, x: torch.Tensor, y: torch.Tensor, stats: dict[str, torch.Tensor]) -> None: |
| self.x = x |
| self.y = y |
| self.stats = stats |
|
|
| def __len__(self) -> int: |
| return int(self.x.shape[0]) |
|
|
| def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: |
| 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 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 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) |
|
|
|
|
| @torch.no_grad() |
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Run compact FNO inference on a NeuralConvection split.") |
| parser.add_argument("--config", default=str(PROJECT_ROOT / "config/config.yaml")) |
| parser.add_argument("--split", default=None) |
| parser.add_argument("--checkpoint", default=None) |
| parser.add_argument("--device", default=None) |
| args = parser.parse_args() |
|
|
| config = load_config(args.config) |
| data_cfg = config["data"] |
| infer_cfg = config["inference"] |
| split = args.split or infer_cfg["split"] |
| device = select_device(args.device or config["experiment"].get("device", "auto")) |
| root = resolve(data_cfg["root_dir"]) |
| stats = torch.load(resolve(data_cfg["stats_path"]), map_location="cpu") |
| x, y = load_pair(root, split, data_cfg["x_file"], data_cfg["y_file"]) |
| loader = DataLoader( |
| InferenceDataset(x, y, stats), |
| batch_size=int(infer_cfg["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) |
| ckpt_path = resolve(args.checkpoint or infer_cfg["checkpoint_path"]) |
| checkpoint = torch.load(ckpt_path, map_location=device) |
| model.load_state_dict(checkpoint["model_state_dict"]) |
| model.eval() |
|
|
| preds: list[np.ndarray] = [] |
| targets: list[np.ndarray] = [] |
| inputs: list[np.ndarray] = [] |
| start = time.perf_counter() |
| for x_batch, y_batch in loader: |
| x_batch = x_batch.to(device, non_blocking=True) |
| pred = tensor_to_numpy(model(x_batch)) |
| preds.append(pred) |
| targets.append(tensor_to_numpy(y_batch)) |
| inputs.append(tensor_to_numpy(x_batch)) |
| elapsed = time.perf_counter() - start |
| pred_all = np.concatenate(preds, axis=0) |
| target_all = np.concatenate(targets, axis=0) |
| input_all = np.concatenate(inputs, axis=0) |
| channel_names = list(data_cfg["channel_names"]) |
| metrics = { |
| "split": split, |
| "checkpoint_path": str(ckpt_path), |
| "num_samples": int(pred_all.shape[0]), |
| "time_seconds_per_sample": elapsed / max(1, pred_all.shape[0]), |
| "per_channel": {}, |
| "paper_table5_compact_fno_2d": config["paper_table5_compact_fno_2d"], |
| } |
| for idx, name in enumerate(channel_names): |
| mse = float(np.mean((pred_all[:, idx] - target_all[:, idx]) ** 2)) |
| 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": mse, "ssim": ssim} |
| paper = config["paper_table5_compact_fno_2d"]["channels"].get(name, {}) |
| print( |
| f"{name}: mse={mse:.8e} ssim={ssim:.6f} paper_mse={paper.get('mse')} paper_ssim={paper.get('ssim')}", |
| flush=True, |
| ) |
|
|
| out_path = resolve(infer_cfg["output_path"]) |
| out_path.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(out_path, pred_norm=pred_all, target_norm=target_all, input_norm=input_all, channel_names=np.array(channel_names)) |
| save_json(resolve(infer_cfg["metrics_path"]), metrics) |
| print(f"saved_predictions={out_path}", flush=True) |
| print(f"saved_metrics={resolve(infer_cfg['metrics_path'])}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
