OneScience-Group/FNO

1#!/usr/bin/env python32"""Run strict closed-loop inference for the trained FNO-2D checkpoint."""3 4from __future__ import annotations5 6import argparse7import csv8import hashlib9import json10import os11import sys12import time13from datetime import datetime, timezone14from pathlib import Path15from typing import Any16 17import numpy as np18import torch19 20 21PROJECT_ROOT = Path(__file__).resolve().parents[1]22if str(PROJECT_ROOT) not in sys.path:23    sys.path.insert(0, str(PROJECT_ROOT))24if str(Path(__file__).resolve().parent) not in sys.path:25    sys.path.insert(0, str(Path(__file__).resolve().parent))26 27from models import build_model_from_config  # noqa: E40228from train import (  # noqa: E40229    atomic_write_json,30    build_datasets,31    build_loader,32    data_file_from_config,33    environment_metadata,34    load_config,35    resolve_project_path,36    seed_everything,37    select_device,38    synchronize,39)40 41 42def parse_args() -> argparse.Namespace:43    parser = argparse.ArgumentParser(44        description="Evaluate a trained FNO-2D checkpoint on the fixed test split."45    )46    parser.add_argument(47        "--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml"48    )49    parser.add_argument("--checkpoint", type=Path, default=None)50    parser.add_argument("--output-dir", type=Path, default=None)51    parser.add_argument("--device", default="auto")52    parser.add_argument("--batch-size", type=int, default=None)53    parser.add_argument("--max-test-samples", type=int, default=None, help="Smoke only.")54    parser.add_argument("--rollout-steps", type=int, default=None, help="Smoke only.")55    return parser.parse_args()56 57 58def load_checkpoint(path: Path, device: torch.device) -> dict[str, Any]:59    if not path.is_file():60        raise FileNotFoundError(f"Checkpoint does not exist: {path}")61    try:62        checkpoint = torch.load(path, map_location=device, weights_only=False)63    except TypeError:64        checkpoint = torch.load(path, map_location=device)65    if not isinstance(checkpoint, dict):66        raise TypeError("Checkpoint root must be a mapping")67    required = {68        "model_state_dict",69        "config",70        "epoch",71        "parameter_count",72        "monitor",73        "test_selected",74    }75    missing = sorted(required.difference(checkpoint))76    if missing:77        raise KeyError(f"Checkpoint is missing required keys: {missing}")78    if checkpoint["test_selected"] is not False:79        raise ValueError("This reproduction forbids a test-selected checkpoint")80    if checkpoint["monitor"] != "train_full_relative_l2":81        raise ValueError(f"Unexpected checkpoint monitor: {checkpoint['monitor']}")82    return checkpoint83 84 85def nested_value(mapping: dict[str, Any], dotted_key: str) -> Any:86    value: Any = mapping87    for key in dotted_key.split("."):88        value = value[key]89    return value90 91 92def validate_checkpoint_config(93    current: dict[str, Any], checkpoint_config: dict[str, Any]94) -> None:95    keys = (96        "data.key",97        "data.layout",98        "data.dtype",99        "data.expected_shape",100        "data.resolution",101        "data.ntrain",102        "data.ntest",103        "data.test_start",104        "data.history",105        "data.horizon",106        "data.normalization",107        "model.input_channels",108        "model.output_channels",109        "model.use_grid",110        "model.grid_include_endpoint",111        "model.width",112        "model.modes1",113        "model.modes2",114        "model.num_layers",115        "model.projection_width",116        "model.fft_norm",117        "training.dtype",118        "training.relative_l2_epsilon",119    )120    differences = []121    for key in keys:122        current_value = nested_value(current, key)123        checkpoint_value = nested_value(checkpoint_config, key)124        if current_value != checkpoint_value:125            differences.append(f"{key}: current={current_value!r}, checkpoint={checkpoint_value!r}")126    if differences:127        raise ValueError("Checkpoint/config mismatch:\n" + "\n".join(differences))128 129 130def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:131    digest = hashlib.sha256()132    with path.open("rb") as handle:133        while chunk := handle.read(chunk_size):134            digest.update(chunk)135    return digest.hexdigest()136 137 138def compute_metrics(139    prediction: np.ndarray, target: np.ndarray, epsilon: float140) -> tuple[np.ndarray, np.ndarray]:141    if prediction.shape != target.shape:142        raise ValueError(f"Prediction/target mismatch: {prediction.shape} vs {target.shape}")143    if prediction.ndim != 4:144        raise ValueError(f"Expected [N,H,W,T], received {prediction.shape}")145    if not np.isfinite(prediction).all() or not np.isfinite(target).all():146        raise FloatingPointError("Prediction or target contains NaN/Inf")147 148    difference = prediction.astype(np.float64) - target.astype(np.float64)149    target64 = target.astype(np.float64)150    full_numerator = np.linalg.norm(difference.reshape(prediction.shape[0], -1), axis=1)151    full_denominator = np.linalg.norm(target64.reshape(target.shape[0], -1), axis=1)152    full = full_numerator / (full_denominator + epsilon)153 154    difference_by_time = np.moveaxis(difference, -1, 1).reshape(155        prediction.shape[0], prediction.shape[-1], -1156    )157    target_by_time = np.moveaxis(target64, -1, 1).reshape(158        target.shape[0], target.shape[-1], -1159    )160    per_lead = np.linalg.norm(difference_by_time, axis=2) / (161        np.linalg.norm(target_by_time, axis=2) + epsilon162    )163    return full, per_lead164 165 166def atomic_save_npz(path: Path, **arrays: np.ndarray) -> None:167    path.parent.mkdir(parents=True, exist_ok=True)168    temporary = path.with_suffix(path.suffix + ".tmp")169    with temporary.open("wb") as handle:170        np.savez_compressed(handle, **arrays)171    os.replace(temporary, path)172 173 174def atomic_write_csv(175    path: Path,176    sample_indices: np.ndarray,177    full_metrics: np.ndarray,178    lead_metrics: np.ndarray,179    time_values: np.ndarray,180) -> None:181    path.parent.mkdir(parents=True, exist_ok=True)182    temporary = path.with_suffix(path.suffix + ".tmp")183    header = ["sample_index", "relative_l2_full"] + [184        f"relative_l2_t{int(value)}" for value in time_values185    ]186    with temporary.open("w", encoding="utf-8", newline="") as handle:187        writer = csv.writer(handle)188        writer.writerow(header)189        for row, sample_index in enumerate(sample_indices):190            writer.writerow(191                [int(sample_index), f"{full_metrics[row]:.17g}"]192                + [f"{value:.17g}" for value in lead_metrics[row]]193            )194    os.replace(temporary, path)195 196 197def main() -> None:198    args = parse_args()199    config = load_config(args.config)200    inference = config["inference"]201    training = config["training"]202    seed = int(inference.get("seed", training["seed"]))203    seed_everything(seed, bool(training.get("deterministic", True)))204    device = select_device(args.device)205    checkpoint_path = (206        resolve_project_path(config["paths"]["checkpoint"])207        if args.checkpoint is None208        else args.checkpoint.expanduser().resolve()209    )210    output_dir = (211        resolve_project_path(config["paths"]["results_dir"])212        if args.output_dir is None213        else args.output_dir.expanduser().resolve()214    )215    output_dir.mkdir(parents=True, exist_ok=True)216    checkpoint = load_checkpoint(checkpoint_path, device)217    validate_checkpoint_config(config, checkpoint["config"])218 219    checkpoint_run_type = str(checkpoint.get("run_type", "formal"))220    if checkpoint_run_type == "formal" and (221        args.max_test_samples is not None222        or args.rollout_steps is not None223        or args.output_dir is not None224        or args.checkpoint is not None225    ):226        raise ValueError(227            "Formal inference uses the exact configured checkpoint, test split, horizon, "228            "and results path; overrides are only allowed for smoke checkpoints"229        )230    horizon = int(config["data"]["horizon"])231    rollout_steps = horizon if args.rollout_steps is None else int(args.rollout_steps)232    if not 1 <= rollout_steps <= horizon:233        raise ValueError(f"rollout_steps must be in [1,{horizon}]")234 235    _, test_dataset = build_datasets(236        config,237        max_train_samples=1,238        max_test_samples=args.max_test_samples,239        rollout_steps=rollout_steps,240    )241    batch_size = int(242        inference["batch_size"] if args.batch_size is None else args.batch_size243    )244    test_loader = build_loader(245        test_dataset,246        batch_size=batch_size,247        shuffle=False,248        num_workers=int(training.get("num_workers", 0)),249        pin_memory=bool(training.get("pin_memory", True)) and device.type == "cuda",250        seed=seed,251    )252 253    # Preserve complex64 spectral parameters while moving the model to device.254    model = build_model_from_config(config).to(device=device)255    incompatible = model.load_state_dict(checkpoint["model_state_dict"], strict=True)256    if incompatible.missing_keys or incompatible.unexpected_keys:257        raise RuntimeError(f"Strict state load failed: {incompatible}")258    parameter_count = sum(parameter.numel() for parameter in model.parameters())259    if parameter_count != int(checkpoint["parameter_count"]):260        raise ValueError(261            f"Parameter count mismatch: model={parameter_count}, "262            f"checkpoint={checkpoint['parameter_count']}"263        )264    model.eval()265 266    predictions: list[np.ndarray] = []267    targets: list[np.ndarray] = []268    synchronize(device)269    started = time.perf_counter()270    processed = 0271    with torch.inference_mode():272        for batch_number, (history, target) in enumerate(test_loader, start=1):273            history = history.to(device=device, dtype=torch.float32, non_blocking=True)274            target_device = target.to(device=device, dtype=torch.float32, non_blocking=True)275            window = history276            batch_prediction: list[torch.Tensor] = []277            for step in range(rollout_steps):278                prediction_step = model(window)279                if not torch.isfinite(prediction_step).all():280                    raise FloatingPointError(281                        f"Non-finite prediction at batch {batch_number}, step {step + 1}"282                    )283                batch_prediction.append(prediction_step)284                window = torch.cat((window[..., 1:], prediction_step), dim=-1)285            prediction = torch.cat(batch_prediction, dim=-1)286            predictions.append(prediction.cpu().numpy().astype(np.float32, copy=False))287            targets.append(target_device.cpu().numpy().astype(np.float32, copy=False))288            processed += int(history.shape[0])289            print(290                f"inference_batch={batch_number:03d}/{len(test_loader):03d} "291                f"processed={processed}/{len(test_dataset)}",292                flush=True,293            )294    synchronize(device)295    duration = time.perf_counter() - started296 297    prediction_array = np.concatenate(predictions, axis=0)298    target_array = np.concatenate(targets, axis=0)299    expected_shape = (300        len(test_dataset),301        int(config["data"]["resolution"][0]),302        int(config["data"]["resolution"][1]),303        rollout_steps,304    )305    if prediction_array.shape != expected_shape or target_array.shape != expected_shape:306        raise ValueError(307            f"Unexpected inference arrays: prediction={prediction_array.shape}, "308            f"target={target_array.shape}, expected={expected_shape}"309        )310    epsilon = float(training["relative_l2_epsilon"])311    full_metrics, lead_metrics = compute_metrics(prediction_array, target_array, epsilon)312    test_start = int(config["data"]["test_start"])313    sample_indices = np.arange(314        test_start, test_start + len(test_dataset), dtype=np.int64315    )316    configured_times = np.asarray(config["data"]["future_times"], dtype=np.float32)317    time_values = configured_times[:rollout_steps]318 319    if args.output_dir is None:320        predictions_path = resolve_project_path(config["paths"]["predictions"])321        metrics_path = resolve_project_path(config["paths"]["metrics"])322        csv_path = resolve_project_path(config["paths"]["per_sample_metrics"])323    else:324        predictions_path = output_dir / "predictions.npz"325        metrics_path = output_dir / "metrics.json"326        csv_path = output_dir / "per_sample_metrics.csv"327 328    atomic_save_npz(329        predictions_path,330        prediction=prediction_array,331        target=target_array,332        sample_indices=sample_indices,333        time_values=time_values,334    )335    atomic_write_csv(csv_path, sample_indices, full_metrics, lead_metrics, time_values)336 337    paper_metric = float(config["paper"]["reference_relative_l2"])338    mean_full = float(full_metrics.mean())339    metrics_payload: dict[str, Any] = {340        "schema_version": "fno-ns2d-metrics-v1",341        "created_at": datetime.now(timezone.utc).isoformat(),342        "run_type": checkpoint_run_type,343        "config_path": str(args.config.expanduser().resolve()),344        "checkpoint_path": str(checkpoint_path),345        "checkpoint_sha256": sha256_file(checkpoint_path),346        "checkpoint_epoch": int(checkpoint["epoch"]),347        "checkpoint_monitor": checkpoint["monitor"],348        "test_selected": bool(checkpoint["test_selected"]),349        "data_path": str(data_file_from_config(config)),350        "prediction_path": str(predictions_path),351        "per_sample_metrics_path": str(csv_path),352        "sample_count": int(len(test_dataset)),353        "prediction_shape": list(prediction_array.shape),354        "sample_indices": {"first": int(sample_indices[0]), "last": int(sample_indices[-1])},355        "time_values": [float(value) for value in time_values],356        "metric": {357            "name": "samplewise_relative_l2",358            "formula": "||prediction-target||_2/(||target||_2+epsilon), then arithmetic mean over samples",359            "epsilon": epsilon,360            "full_trajectory_mean": mean_full,361            "full_trajectory_std": float(full_metrics.std(ddof=0)),362            "mean_step_relative_l2": float(lead_metrics.mean()),363            "per_lead_mean": [float(value) for value in lead_metrics.mean(axis=0)],364            "per_lead_std": [float(value) for value in lead_metrics.std(axis=0, ddof=0)],365        },366        "paper_comparison": {367            "paper_relative_l2": paper_metric,368            "signed_difference": mean_full - paper_metric,369            "absolute_difference": abs(mean_full - paper_metric),370        },371        "parameter_count": parameter_count,372        "paper_parameter_count": int(config["paper"]["reference_parameter_count"]),373        "parameter_count_difference": parameter_count374        - int(config["paper"]["reference_parameter_count"]),375        "runtime": {376            **environment_metadata(device),377            "duration_seconds": duration,378            "batch_size": batch_size,379        },380        "assumptions": config.get("assumptions", []),381    }382    atomic_write_json(metrics_path, metrics_payload)383    print(384        f"inference_complete samples={len(test_dataset)} duration={duration:.3f}s "385        f"full_relative_l2={mean_full:.8f} paper={paper_metric:.8f} "386        f"absolute_difference={abs(mean_full-paper_metric):.8f}",387        flush=True,388    )389 390 391if __name__ == "__main__":392    main()393