scripts/inference.py

#!/usr/bin/env python3
"""Evaluate the best FAST-DeepONet checkpoint on all ten test fields."""

from __future__ import annotations

import argparse
import csv
import json
import math
import os
import sys
from pathlib import Path
from typing import Any

import numpy as np
import torch

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

from model.FASTDeepONet import FASTDeepONet
from scripts.common import (
    atomic_write_json,
    load_config,
    load_split,
    make_coordinate_features,
    relative_l2,
    select_device,
)


def load_checkpoint(path: Path, device: torch.device) -> dict[str, Any]:
    try:
        checkpoint = torch.load(path, map_location=device, weights_only=False)
    except TypeError:
        checkpoint = torch.load(path, map_location=device)
    if not isinstance(checkpoint, dict) or "model_state_dict" not in checkpoint:
        raise ValueError(f"Invalid checkpoint payload: {path}")
    return checkpoint


def build_from_checkpoint(checkpoint: dict[str, Any], device: torch.device) -> FASTDeepONet:
    arguments = dict(checkpoint["model_arguments"])
    recorded_effective_rank = int(arguments.pop("effective_factor_rank"))
    basis = checkpoint["spectral_basis"]
    model = FASTDeepONet(spectral_basis=basis, **arguments).to(device)
    if model.effective_factor_rank != recorded_effective_rank:
        raise ValueError("Checkpoint effective rank is inconsistent with its spectral basis")
    model.load_state_dict(checkpoint["model_state_dict"], strict=True)
    model.eval()
    return model


def atomic_save_predictions(path: Path, **arrays: np.ndarray) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    with temporary.open("wb") as handle:
        np.savez_compressed(handle, **arrays)
    os.replace(temporary, path)


def write_metrics_csv(path: Path, filenames: list[str], metrics: list[float], mean_metric: float) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    with temporary.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(["sample_index", "filename", "relative_l2"])
        for index, (filename, metric) in enumerate(zip(filenames, metrics)):
            writer.writerow([index, filename, f"{metric:.12e}"])
        writer.writerow(["mean", "", f"{mean_metric:.12e}"])
    os.replace(temporary, path)


def run_inference(config: dict[str, Any], force_cpu: bool) -> None:
    paths = config["paths"]
    data_config = config["data"]
    model_config = config["model"]
    training = config["training"]
    device = torch.device("cpu") if force_cpu else select_device(require_accelerator=False)
    checkpoint_path = Path(paths["checkpoint"])
    if not checkpoint_path.is_file():
        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
    checkpoint = load_checkpoint(checkpoint_path, device)
    model = build_from_checkpoint(checkpoint, device)
    filenames, inputs_np, targets_np, masks_np = load_split(
        paths["data_root"],
        "test",
        data_config["split_counts"]["test"],
        data_config["points_per_field"],
    )
    coordinates_np = make_coordinate_features(
        data_config["grid_size"], model_config["fourier_order"]
    )
    coordinates = torch.from_numpy(coordinates_np).to(device)
    target_mean = np.asarray(checkpoint["target_mean"].cpu(), dtype=np.float32)
    target_std = float(checkpoint["target_std"])
    predictions: list[np.ndarray] = []
    metrics: list[float] = []
    batch_size = int(training["batch_size"])

    print(
        f"inference_start checkpoint={checkpoint_path} best_epoch={checkpoint['best_epoch']} "
        f"best_tune_loss={float(checkpoint['best_tune_mse']):.8e} device={device}",
        flush=True,
    )
    with torch.inference_mode():
        for sample_index, (filename, input_field, target_field) in enumerate(
            zip(filenames, inputs_np, targets_np)
        ):
            branch = torch.from_numpy(input_field[None, :]).to(device)
            branch_latent = model.encode_branch(branch)
            normalized_chunks: list[np.ndarray] = []
            for start in range(0, coordinates.shape[0], batch_size):
                coordinate_batch = coordinates[start : start + batch_size]
                expanded_branch = branch_latent.expand(coordinate_batch.shape[0], -1)
                normalized = model.predict_from_branch_encoding(expanded_branch, coordinate_batch)
                normalized_chunks.append(normalized.cpu().numpy().astype(np.float32, copy=False))
            normalized_prediction = np.concatenate(normalized_chunks)
            prediction = normalized_prediction * target_std + target_mean
            metric = relative_l2(target_field, prediction)
            if not math.isfinite(metric):
                raise FloatingPointError(f"Non-finite relative L2 for {filename}")
            predictions.append(prediction.astype(np.float32, copy=False))
            metrics.append(metric)
            print(
                f"test_sample={sample_index:02d} file={filename} relative_l2={metric:.10e}",
                flush=True,
            )

    mean_metric = float(np.mean(metrics))
    print(f"test_mean_relative_l2={mean_metric:.10e} samples={len(metrics)}", flush=True)
    results_dir = Path(paths["results_dir"])
    results_dir.mkdir(parents=True, exist_ok=True)
    metrics_payload = {
        "metric": "relative_l2_per_field",
        "equation": "||y-y_hat||_2 / (||y||_2 + 1e-12)",
        "checkpoint": str(checkpoint_path),
        "best_epoch": int(checkpoint["best_epoch"]),
        "best_tune_mse": float(checkpoint["best_tune_mse"]),
        "sample_count": len(metrics),
        "samples": [
            {"sample_index": index, "filename": filename, "relative_l2": metric}
            for index, (filename, metric) in enumerate(zip(filenames, metrics))
        ],
        "mean_relative_l2": mean_metric,
        "paper_comparison_note": (
            "This fixed run uses 10 test fields and one model seed; the paper's Table 2 "
            "uses 100 independent test fields and summarizes five fits."
        ),
    }
    atomic_write_json(results_dir / "test_metrics.json", metrics_payload)
    write_metrics_csv(results_dir / "test_metrics.csv", filenames, metrics, mean_metric)
    atomic_save_predictions(
        results_dir / "test_predictions.npz",
        filenames=np.asarray(filenames),
        inputs=inputs_np.astype(np.float32, copy=False),
        masks=masks_np.astype(np.float32, copy=False),
        targets=targets_np.astype(np.float32, copy=False),
        predictions=np.stack(predictions).astype(np.float32, copy=False),
        relative_l2=np.asarray(metrics, dtype=np.float64),
    )
    print(f"inference_complete results_dir={results_dir}", flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "config" / "config.yaml")
    parser.add_argument("--cpu", action="store_true", help="Force CPU inference")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    run_inference(load_config(args.config), force_cpu=args.cpu)


if __name__ == "__main__":
    main()