OneScience-Group/FNO

1#!/usr/bin/env python32"""Validate real FNO outputs and render paper-comparison figures."""3 4from __future__ import annotations5 6import argparse7import csv8import hashlib9import json10import os11import sys12from datetime import datetime, timezone13from pathlib import Path14from typing import Any15 16import matplotlib17 18matplotlib.use("Agg")19import matplotlib.pyplot as plt  # noqa: E40220import numpy as np  # noqa: E40221 22 23PROJECT_ROOT = Path(__file__).resolve().parents[1]24if str(PROJECT_ROOT) not in sys.path:25    sys.path.insert(0, str(PROJECT_ROOT))26if str(Path(__file__).resolve().parent) not in sys.path:27    sys.path.insert(0, str(Path(__file__).resolve().parent))28 29from inference import compute_metrics  # noqa: E40230from train import atomic_write_json, load_config, resolve_project_path  # noqa: E40231 32 33def parse_args() -> argparse.Namespace:34    parser = argparse.ArgumentParser(35        description="Validate FNO inference artifacts and generate scientific figures."36    )37    parser.add_argument(38        "--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml"39    )40    parser.add_argument("--output-dir", type=Path, default=None)41    parser.add_argument(42        "--sample-index", type=int, default=0, help="Local test-set index to visualize."43    )44    return parser.parse_args()45 46 47def read_json(path: Path) -> dict[str, Any]:48    if not path.is_file():49        raise FileNotFoundError(f"Required JSON artifact is missing: {path}")50    with path.open("r", encoding="utf-8") as handle:51        payload = json.load(handle)52    if not isinstance(payload, dict):53        raise TypeError(f"Expected a JSON mapping in {path}")54    return payload55 56 57def sha256_file(path: Path, chunk_size: int = 1024 * 1024) -> str:58    digest = hashlib.sha256()59    with path.open("rb") as handle:60        while chunk := handle.read(chunk_size):61            digest.update(chunk)62    return digest.hexdigest()63 64 65def atomic_save_figure(figure: plt.Figure, path: Path, dpi: int = 300) -> None:66    path.parent.mkdir(parents=True, exist_ok=True)67    temporary = path.with_suffix(path.suffix + ".tmp")68    figure.savefig(temporary, format="png", dpi=dpi, bbox_inches="tight")69    plt.close(figure)70    os.replace(temporary, path)71    if path.stat().st_size == 0:72        raise RuntimeError(f"Generated an empty figure: {path}")73 74 75def atomic_write_text(path: Path, content: str) -> None:76    path.parent.mkdir(parents=True, exist_ok=True)77    temporary = path.with_suffix(path.suffix + ".tmp")78    with temporary.open("w", encoding="utf-8") as handle:79        handle.write(content)80    os.replace(temporary, path)81 82 83def verify_csv(84    path: Path,85    sample_indices: np.ndarray,86    full_metrics: np.ndarray,87    lead_metrics: np.ndarray,88    time_values: np.ndarray,89) -> None:90    if not path.is_file():91        raise FileNotFoundError(f"Per-sample metrics CSV is missing: {path}")92    expected_header = ["sample_index", "relative_l2_full"] + [93        f"relative_l2_t{int(value)}" for value in time_values94    ]95    with path.open("r", encoding="utf-8", newline="") as handle:96        rows = list(csv.reader(handle))97    if not rows or rows[0] != expected_header:98        raise ValueError(f"Unexpected CSV header in {path}: {rows[0] if rows else None}")99    if len(rows) - 1 != len(sample_indices):100        raise ValueError(f"Expected {len(sample_indices)} CSV rows, found {len(rows)-1}")101    for row_index, row in enumerate(rows[1:]):102        if int(row[0]) != int(sample_indices[row_index]):103            raise ValueError(f"CSV sample order mismatch at row {row_index + 2}")104        observed = np.asarray([float(value) for value in row[1:]], dtype=np.float64)105        expected = np.concatenate(106            ([full_metrics[row_index]], lead_metrics[row_index].astype(np.float64))107        )108        if not np.allclose(observed, expected, rtol=1e-12, atol=1e-12):109            raise ValueError(f"CSV metric mismatch for sample {sample_indices[row_index]}")110 111 112def validate_history(history_payload: dict[str, Any]) -> list[dict[str, Any]]:113    records = history_payload.get("history")114    if not isinstance(records, list) or not records:115        raise ValueError("Training history contains no epoch records")116    formal = history_payload.get("run_type") == "formal"117    requested = int(history_payload.get("epochs_requested", len(records)))118    if formal and (requested != 500 or len(records) != 500):119        raise ValueError(120            f"Formal paper reproduction requires 500 epochs, got requested={requested}, "121            f"records={len(records)}"122        )123    required = (124        "epoch",125        "learning_rate",126        "duration_seconds",127        "train_step_loss_sum",128        "train_mean_step_relative_l2",129        "train_full_relative_l2",130        "test_mean_step_relative_l2",131        "test_full_relative_l2",132        "best",133    )134    for position, record in enumerate(records, start=1):135        missing = [key for key in required if key not in record]136        if missing:137            raise KeyError(f"Epoch record {position} is missing {missing}")138        if int(record["epoch"]) != position:139            raise ValueError(f"Epoch sequence is not contiguous at record {position}")140        numeric = [float(record[key]) for key in required[1:-1]]141        if not np.isfinite(numeric).all():142            raise FloatingPointError(f"Non-finite training history at epoch {position}")143    return records144 145 146def make_training_figure(147    records: list[dict[str, Any]], paper_metric: float, best_epoch: int148) -> plt.Figure:149    epochs = np.asarray([record["epoch"] for record in records], dtype=np.int64)150    train_full = np.asarray(151        [record["train_full_relative_l2"] for record in records], dtype=np.float64152    )153    test_full = np.asarray(154        [record["test_full_relative_l2"] for record in records], dtype=np.float64155    )156    train_loss = np.asarray(157        [record["train_step_loss_sum"] for record in records], dtype=np.float64158    )159    train_step = np.asarray(160        [record["train_mean_step_relative_l2"] for record in records], dtype=np.float64161    )162    test_step = np.asarray(163        [record["test_mean_step_relative_l2"] for record in records], dtype=np.float64164    )165 166    figure, axes = plt.subplots(1, 2, figsize=(12.5, 4.8), constrained_layout=True)167    left = axes[0]168    left.plot(epochs, train_full, label="Train full relative L2", linewidth=1.6)169    left.plot(epochs, test_full, label="Test full relative L2", linewidth=1.6)170    left.axhline(171        paper_metric,172        color="black",173        linestyle="--",174        linewidth=1.2,175        label=f"Paper benchmark ({paper_metric:.4f})",176    )177    left.axvline(178        best_epoch,179        color="tab:green",180        linestyle=":",181        linewidth=1.2,182        label=f"Best checkpoint epoch ({best_epoch})",183    )184    if np.all(train_full > 0) and np.all(test_full > 0):185        left.set_yscale("log")186    left.set_xlabel("Epoch")187    left.set_ylabel("Full-trajectory relative L2")188    left.set_title("FNO-2D rollout error")189    left.grid(True, alpha=0.25)190    left.legend(fontsize=8)191 192    right = axes[1]193    loss_line = right.plot(194        epochs,195        train_loss,196        color="tab:blue",197        label="Train 10-step loss sum",198        linewidth=1.5,199    )200    right.set_xlabel("Epoch")201    right.set_ylabel("Summed step relative L2", color="tab:blue")202    right.tick_params(axis="y", labelcolor="tab:blue")203    right.grid(True, alpha=0.25)204    diagnostic = right.twinx()205    train_line = diagnostic.plot(206        epochs,207        train_step,208        color="tab:orange",209        label="Train mean-step relative L2",210        linewidth=1.3,211    )212    test_line = diagnostic.plot(213        epochs,214        test_step,215        color="tab:red",216        label="Test mean-step relative L2",217        linewidth=1.3,218    )219    diagnostic.set_ylabel("Mean-step relative L2")220    right.set_title("Training objective and step diagnostics")221    lines = loss_line + train_line + test_line222    right.legend(lines, [line.get_label() for line in lines], fontsize=8, loc="best")223    return figure224 225 226def representative_leads(number_of_steps: int) -> list[int]:227    if number_of_steps <= 0:228        raise ValueError("At least one rollout step is required")229    return sorted({0, number_of_steps // 2, number_of_steps - 1})230 231 232def make_rollout_figure(233    prediction: np.ndarray,234    target: np.ndarray,235    sample_indices: np.ndarray,236    time_values: np.ndarray,237    lead_metrics: np.ndarray,238    local_sample: int,239) -> plt.Figure:240    if not 0 <= local_sample < prediction.shape[0]:241        raise IndexError(242            f"sample-index {local_sample} is outside [0,{prediction.shape[0] - 1}]"243        )244    lead_indices = representative_leads(prediction.shape[-1])245    figure, axes = plt.subplots(246        len(lead_indices),247        3,248        figsize=(11.5, 3.25 * len(lead_indices)),249        squeeze=False,250        constrained_layout=True,251    )252    global_sample = int(sample_indices[local_sample])253    for row, lead_index in enumerate(lead_indices):254        truth = target[local_sample, :, :, lead_index]255        estimate = prediction[local_sample, :, :, lead_index]256        absolute_error = np.abs(estimate - truth)257        shared_limit = max(float(np.max(np.abs(truth))), float(np.max(np.abs(estimate))), 1e-12)258        error_limit = max(float(np.max(absolute_error)), 1e-12)259        time_value = int(time_values[lead_index])260        relative_error = float(lead_metrics[local_sample, lead_index])261 262        fields = (truth, estimate, absolute_error)263        titles = (264            f"Target vorticity w\nsample={global_sample}, t={time_value}",265            f"Predicted vorticity w\nrelative L2={relative_error:.5f}",266            f"Absolute error |prediction-target|\nt={time_value}",267        )268        for column, (field, title) in enumerate(zip(fields, titles)):269            axis = axes[row, column]270            if column < 2:271                image = axis.imshow(272                    field,273                    origin="lower",274                    extent=(0.0, 1.0, 0.0, 1.0),275                    interpolation="nearest",276                    cmap="RdBu_r",277                    vmin=-shared_limit,278                    vmax=shared_limit,279                )280                color_label = "Vorticity w (unit not specified)"281            else:282                image = axis.imshow(283                    field,284                    origin="lower",285                    extent=(0.0, 1.0, 0.0, 1.0),286                    interpolation="nearest",287                    cmap="magma",288                    vmin=0.0,289                    vmax=error_limit,290                )291                color_label = "Absolute error"292            axis.set_aspect("equal")293            axis.set_xlabel("x")294            axis.set_ylabel("y")295            axis.set_title(title, fontsize=9)296            colorbar = figure.colorbar(image, ax=axis, shrink=0.82)297            colorbar.set_label(color_label, fontsize=8)298    return figure299 300 301def main() -> None:302    args = parse_args()303    config = load_config(args.config)304    output_dir = (305        resolve_project_path(config["paths"]["results_dir"])306        if args.output_dir is None307        else args.output_dir.expanduser().resolve()308    )309    if args.output_dir is None:310        history_path = resolve_project_path(config["paths"]["train_history"])311        predictions_path = resolve_project_path(config["paths"]["predictions"])312        metrics_path = resolve_project_path(config["paths"]["metrics"])313        csv_path = resolve_project_path(config["paths"]["per_sample_metrics"])314        training_figure_path = resolve_project_path(config["paths"]["training_curves"])315        rollout_figure_path = resolve_project_path(config["paths"]["rollout_figure"])316        metadata_path = resolve_project_path(config["paths"]["run_metadata"])317        summary_path = resolve_project_path(config["paths"]["summary"])318    else:319        history_path = output_dir / "train_history.json"320        predictions_path = output_dir / "predictions.npz"321        metrics_path = output_dir / "metrics.json"322        csv_path = output_dir / "per_sample_metrics.csv"323        training_figure_path = output_dir / "training_curves.png"324        rollout_figure_path = output_dir / "sample_000_rollout.png"325        metadata_path = output_dir / "run_metadata.json"326        summary_path = output_dir / "summary.md"327 328    output_dir.mkdir(parents=True, exist_ok=True)329    history_payload = read_json(history_path)330    metrics_payload = read_json(metrics_path)331    if metrics_payload.get("run_type") == "formal" and args.output_dir is not None:332        raise ValueError("Formal result generation must use the configured results directory")333    if history_payload.get("run_type") != metrics_payload.get("run_type"):334        raise ValueError("Training history and inference metrics have different run types")335    records = validate_history(history_payload)336    if not predictions_path.is_file():337        raise FileNotFoundError(f"Predictions artifact is missing: {predictions_path}")338    with np.load(predictions_path, allow_pickle=False) as archive:339        required_arrays = {"prediction", "target", "sample_indices", "time_values"}340        missing_arrays = required_arrays.difference(archive.files)341        if missing_arrays:342            raise KeyError(f"Predictions NPZ is missing {sorted(missing_arrays)}")343        prediction = archive["prediction"]344        target = archive["target"]345        sample_indices = archive["sample_indices"]346        time_values = archive["time_values"]347 348    if prediction.dtype != np.float32 or target.dtype != np.float32:349        raise TypeError("Prediction and target arrays must be float32")350    if prediction.shape != target.shape or prediction.ndim != 4:351        raise ValueError(f"Invalid prediction/target shapes: {prediction.shape}, {target.shape}")352    if sample_indices.shape != (prediction.shape[0],):353        raise ValueError("sample_indices shape does not match predictions")354    if time_values.shape != (prediction.shape[-1],):355        raise ValueError("time_values shape does not match rollout horizon")356    if not np.array_equal(sample_indices, np.arange(sample_indices[0], sample_indices[0] + len(sample_indices))):357        raise ValueError("sample_indices must be unique, contiguous, and ordered")358    if not np.all(np.diff(time_values.astype(np.float64)) > 0):359        raise ValueError("time_values must be strictly increasing")360 361    formal = metrics_payload.get("run_type") == "formal"362    if formal:363        expected_shape = (364            int(config["data"]["ntest"]),365            int(config["data"]["resolution"][0]),366            int(config["data"]["resolution"][1]),367            int(config["data"]["horizon"]),368        )369        if prediction.shape != expected_shape:370            raise ValueError(f"Formal prediction shape must be {expected_shape}, got {prediction.shape}")371        expected_indices = np.arange(372            int(config["data"]["test_start"]),373            int(config["data"]["test_start"]) + int(config["data"]["ntest"]),374        )375        if not np.array_equal(sample_indices, expected_indices):376            raise ValueError("Formal sample indices do not match the fixed test split")377 378    epsilon = float(config["training"]["relative_l2_epsilon"])379    full_metrics, lead_metrics = compute_metrics(prediction, target, epsilon)380    observed_mean = float(metrics_payload["metric"]["full_trajectory_mean"])381    if not np.isclose(full_metrics.mean(), observed_mean, rtol=1e-8, atol=1e-8):382        raise ValueError(383            f"metrics.json full relative L2 mismatch: recomputed={full_metrics.mean()}, "384            f"stored={observed_mean}"385        )386    stored_leads = np.asarray(metrics_payload["metric"]["per_lead_mean"], dtype=np.float64)387    if not np.allclose(lead_metrics.mean(axis=0), stored_leads, rtol=1e-8, atol=1e-8):388        raise ValueError("metrics.json per-lead values do not match predictions")389    verify_csv(csv_path, sample_indices, full_metrics, lead_metrics, time_values)390 391    best_epoch = int(history_payload["best_epoch"])392    checkpoint_epoch = int(metrics_payload["checkpoint_epoch"])393    if best_epoch != checkpoint_epoch:394        raise ValueError(395            f"History best epoch {best_epoch} does not match checkpoint epoch {checkpoint_epoch}"396        )397    paper_metric = float(config["paper"]["reference_relative_l2"])398    training_figure = make_training_figure(records, paper_metric, best_epoch)399    atomic_save_figure(training_figure, training_figure_path, dpi=300)400    rollout_figure = make_rollout_figure(401        prediction,402        target,403        sample_indices,404        time_values,405        lead_metrics,406        args.sample_index,407    )408    atomic_save_figure(rollout_figure, rollout_figure_path, dpi=300)409 410    artifact_paths = {411        "train_history": history_path,412        "predictions": predictions_path,413        "metrics": metrics_path,414        "per_sample_metrics": csv_path,415        "training_curves": training_figure_path,416        "rollout_figure": rollout_figure_path,417    }418    artifact_metadata = {419        name: {420            "path": str(path),421            "size_bytes": path.stat().st_size,422            "sha256": sha256_file(path),423        }424        for name, path in artifact_paths.items()425    }426    main_metric = float(full_metrics.mean())427    run_metadata = {428        "schema_version": "fno-ns2d-run-metadata-v1",429        "created_at": datetime.now(timezone.utc).isoformat(),430        "run_type": metrics_payload.get("run_type"),431        "config_path": str(args.config.expanduser().resolve()),432        "data_path": metrics_payload["data_path"],433        "checkpoint_path": metrics_payload["checkpoint_path"],434        "checkpoint_epoch": checkpoint_epoch,435        "test_selected": metrics_payload["test_selected"],436        "split": {437            "train": int(history_payload.get("train_samples", 1000)),438            "validation": 0,439            "test": int(prediction.shape[0]),440        },441        "prediction_shape": list(prediction.shape),442        "normalization": config["data"]["normalization"],443        "metric_formula": metrics_payload["metric"]["formula"],444        "full_trajectory_relative_l2": main_metric,445        "paper_relative_l2": paper_metric,446        "signed_difference": main_metric - paper_metric,447        "absolute_difference": abs(main_metric - paper_metric),448        "parameter_count": metrics_payload["parameter_count"],449        "paper_parameter_count": metrics_payload["paper_parameter_count"],450        "parameter_count_difference": metrics_payload["parameter_count_difference"],451        "runtime": {**metrics_payload["runtime"], "matplotlib": matplotlib.__version__},452        "assumptions": config.get("assumptions", []),453        "conflicts": config.get("conflicts", []),454        "artifacts": artifact_metadata,455        "quality_checks": {456            "all_values_finite": True,457            "metrics_recomputed_from_npz": True,458            "json_metrics_match": True,459            "csv_metrics_match": True,460            "best_epoch_matches_checkpoint": True,461            "figures_nonempty": True,462        },463    }464    atomic_write_json(metadata_path, run_metadata)465 466    summary = f"""# FNO-2D Navier–Stokes reproduction result467 468## Summary469 470- Run type: `{metrics_payload.get('run_type')}`471- Test trajectories: {prediction.shape[0]}472- Forecast shape: `{list(prediction.shape)}`473- Mean full-trajectory relative L2: **{main_metric:.8f}**474- Paper FNO-2D reference (`ν=1e-5`, `T=20`, 1000 train): **{paper_metric:.4f}**475- Absolute difference: **{abs(main_metric-paper_metric):.8f}**476- Checkpoint epoch: {checkpoint_epoch}; selected by train full relative L2 (`test_selected=false`).477 478## Data and method479 480The model uses the fixed first 1000 trajectories for training and the final {prediction.shape[0]} trajectories for testing. Ten observed vorticity frames initialize a closed-loop rollout; every predicted frame updates the next input window. No target frame is used after initialization, and no data normalization, padding, augmentation, or PDE-residual loss is applied.481 482The reported metric is computed per sample as `||prediction-target||₂/(||target||₂+1e-12)` over the full space-time forecast and then averaged. It was recomputed directly from `predictions.npz` and cross-checked against JSON and CSV outputs.483 484## Reproducibility limitations485 486The paper does not specify the exact relative-L2 reduction, batch size, random seed, projection hidden width, coordinate-input choice, block ordering, or checkpoint-selection protocol. These choices are recorded explicitly in `config/config.yaml` and `run_metadata.json`. The paper width 32 also cannot be uniquely reconciled with the reported 414,517 parameters from the published connection details; the actual parameter count is reported rather than hidden.487 488## Artifacts489 490- `{training_figure_path.name}`: train/test rollout errors and step-loss diagnostics.491- `{rollout_figure_path.name}`: target, prediction, and absolute-error vorticity fields.492- `{predictions_path.name}`: full test prediction and target arrays.493- `{metrics_path.name}` and `{csv_path.name}`: aggregate and per-sample metrics.494- `{metadata_path.name}`: provenance, software, assumptions, hashes, and quality checks.495"""496    atomic_write_text(summary_path, summary)497    print(498        f"result_complete full_relative_l2={main_metric:.8f} "499        f"training_figure={training_figure_path} rollout_figure={rollout_figure_path}",500        flush=True,501    )502 503 504if __name__ == "__main__":505    main()506