scripts/result.py

"""Evaluate ACE rollouts and render forecast figures."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import torch

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from ACE.model.paths import GENERATED_DATA_PATH, INFER_PATH, PIC_DIR, configured_path
from ACE.model.physics import forecast_metrics
from ACE.model.variables import DIAGNOSTIC_CHANNELS, PROGNOSTIC_CHANNELS


CHANNELS = PROGNOSTIC_CHANNELS + DIAGNOSTIC_CHANNELS
MAP_CHANNELS = ("T_0", "T_7", "Ts_land_or_seaice", "ps", "P", "LHF")
SERIES_CHANNELS = ("T_0", "T_7", "qT_7", "ps", "P", "LHF", "SHF")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", type=Path, default=Path(__file__).resolve().parents[1] / "conf" / "config.yaml")
    parser.add_argument("--prediction-path", type=Path, default=None)
    parser.add_argument("--truth-path", type=Path, default=None)
    parser.add_argument("--output-dir", type=Path, default=None)
    parser.add_argument("--area-weights", type=Path, default=None, help="NPY [H,W] or [H]")
    parser.add_argument("--dpi", type=int, default=160)
    return parser.parse_args()


def _rollout(values: np.ndarray, name: str) -> np.ndarray:
    values = np.asarray(values)
    if values.ndim == 5:
        values = values[0]
    elif values.ndim == 4:
        if values.shape[1] != len(CHANNELS):
            raise ValueError(f"{name} must have 44 channels, got {values.shape}")
        values = values if name == "prediction" else values[:1]
    elif values.ndim == 3:
        values = values[None]
    else:
        raise ValueError(f"{name} must be [B,T,C,H,W], [T,C,H,W], or [C,H,W], got {values.shape}")
    if values.ndim != 4 or values.shape[1] != len(CHANNELS):
        raise ValueError(f"{name} must resolve to [T,44,H,W], got {values.shape}")
    return values.astype(np.float32, copy=False)


def _truth(pred_data, truth_path, config):
    if truth_path is not None:
        data = np.load(truth_path)
        return data["predictions"] if "predictions" in data else data["targets"], data
    if "targets" in pred_data:
        return pred_data["targets"], pred_data
    data = np.load(configured_path(config, "data_path", GENERATED_DATA_PATH))
    return data["targets"], data


def _coordinates(*sources, height: int, width: int):
    lat = lon = None
    for source in sources:
        if source is None:
            continue
        if lat is None and "lat" in source:
            lat = np.asarray(source["lat"], dtype=np.float32)
        if lon is None and "lon" in source:
            lon = np.asarray(source["lon"], dtype=np.float32)
    lat = lat if lat is not None else np.linspace(-90, 90, height, dtype=np.float32)
    lon = lon if lon is not None else np.linspace(0, 360, width, endpoint=False, dtype=np.float32)
    if lat.size != height or lon.size != width:
        raise ValueError(f"coordinate shape mismatch: lat={lat.shape}, lon={lon.shape}, field={(height, width)}")
    return lat, lon


def _area_grid(path, lat, width):
    if path is not None:
        area = np.asarray(np.load(path), dtype=np.float64)
        if area.ndim == 1:
            area = area[:, None]
        if area.shape != (lat.size, width):
            raise ValueError(f"area weights must match {(lat.size, width)}, got {area.shape}")
        return area
    return np.broadcast_to(np.cos(np.deg2rad(lat))[:, None], (lat.size, width)).copy()


def _global_mean(fields, area):
    weights = area / max(float(area.sum()), np.finfo(np.float64).eps)
    return (fields * weights[None, None]).sum(axis=(-2, -1))


def _limits(values, symmetric=False):
    finite = np.asarray(values, dtype=np.float64)
    finite = finite[np.isfinite(finite)]
    if finite.size == 0:
        return -1.0, 1.0
    low, high = np.percentile(finite, (2, 98))
    if not np.isfinite(low) or not np.isfinite(high) or low == high:
        low, high = float(finite.min()), float(finite.max())
    if symmetric:
        bound = max(abs(float(low)), abs(float(high)), 1e-12)
        return -bound, bound
    return float(low), float(high if high > low else low + 1e-12)


def _plt():
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
        return plt
    except ImportError as exc:
        raise RuntimeError("PNG visualization requires matplotlib; it is available in develop_base") from exc


def _save_maps(path, pred, lat, lon, dpi):
    plt = _plt()
    selected = [(name, CHANNELS.index(name)) for name in MAP_CHANNELS if name in CHANNELS]
    steps = np.unique(np.linspace(0, pred.shape[0] - 1, min(pred.shape[0], 4), dtype=int))
    fig, axes = plt.subplots(len(selected), len(steps), figsize=(3.7 * len(steps), 2.8 * len(selected)), squeeze=False, constrained_layout=True)
    extent = (float(lon.min()), float(lon.max()), float(lat.min()), float(lat.max()))
    for row, (name, channel) in enumerate(selected):
        low, high = _limits(pred[:, channel])
        for col, step in enumerate(steps):
            image = axes[row, col].imshow(pred[step, channel], origin="lower", extent=extent, aspect="auto", cmap="viridis", vmin=low, vmax=high)
            axes[row, col].set_title(f"{name} | lead {(step + 1) * 6} h")
            axes[row, col].set_xlabel("longitude (deg)")
            axes[row, col].set_ylabel("latitude (deg)")
            fig.colorbar(image, ax=axes[row, col], shrink=0.82)
    fig.suptitle("ACE rollout fields (first sample)", fontsize=15)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def _save_comparison(path, pred, truth, lat, lon, dpi):
    plt = _plt()
    selected = [(name, CHANNELS.index(name)) for name in ("T_0", "T_7", "ps", "P", "LHF") if name in CHANNELS]
    fig, axes = plt.subplots(len(selected), 3, figsize=(11, 2.9 * len(selected)), squeeze=False, constrained_layout=True)
    extent = (float(lon.min()), float(lon.max()), float(lat.min()), float(lat.max()))
    for row, (name, channel) in enumerate(selected):
        predicted, reference = pred[0, channel], truth[0, channel]
        error = predicted - reference
        low, high = _limits(np.stack([predicted, reference]))
        bound = _limits(error, symmetric=True)[1]
        for col, (field, title, cmap, vmin, vmax) in enumerate(((predicted, "prediction", "viridis", low, high), (reference, "truth", "viridis", low, high), (error, "prediction - truth", "coolwarm", -bound, bound))):
            image = axes[row, col].imshow(field, origin="lower", extent=extent, aspect="auto", cmap=cmap, vmin=vmin, vmax=vmax)
            axes[row, col].set_title(f"{name}: {title}")
            axes[row, col].set_xlabel("longitude (deg)")
            axes[row, col].set_ylabel("latitude (deg)")
            fig.colorbar(image, ax=axes[row, col], shrink=0.82)
    fig.suptitle("First-step forecast comparison", fontsize=15)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def _save_series(path, means, truth_means, dpi):
    plt = _plt()
    selected = [name for name in SERIES_CHANNELS if name in CHANNELS]
    fig, axes = plt.subplots(len(selected), 1, figsize=(10, 2.1 * len(selected)), squeeze=False, sharex=True, constrained_layout=True)
    time = np.arange(means.shape[0]) * 6 / 24.0
    for row, name in enumerate(selected):
        channel = CHANNELS.index(name)
        axes[row, 0].plot(time, means[:, channel], marker="o", linewidth=1.6, label="prediction")
        if truth_means is not None:
            axes[row, 0].plot(np.arange(truth_means.shape[0]) * 6 / 24.0, truth_means[:, channel], "x--", label="truth")
        axes[row, 0].set_ylabel(name)
        axes[row, 0].grid(alpha=0.25)
        axes[row, 0].legend(loc="best", fontsize=8)
    axes[-1, 0].set_xlabel("forecast lead (days)")
    fig.suptitle("Area-weighted global means", fontsize=15)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def _save_heatmap(path, means, dpi):
    plt = _plt()
    fig, axis = plt.subplots(figsize=(11, 13), constrained_layout=True)
    image = axis.imshow(means.T, aspect="auto", interpolation="nearest", cmap="RdBu_r")
    axis.set_yticks(np.arange(len(CHANNELS)))
    axis.set_yticklabels(CHANNELS, fontsize=7)
    axis.set_xticks(np.arange(means.shape[0]))
    axis.set_xticklabels([f"+{(step + 1) * 6}h" for step in range(means.shape[0])], rotation=45, ha="right")
    axis.set_xlabel("forecast lead")
    axis.set_title("Global mean of all 44 output channels")
    fig.colorbar(image, ax=axis, label="area-weighted mean")
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)


def _save_rmse(path, pred, truth, dpi):
    plt = _plt()
    rmse = np.sqrt(np.mean(np.square(pred - truth), axis=(0, 2, 3)))
    fig, axis = plt.subplots(figsize=(15, 5), constrained_layout=True)
    colors = ["#2b6cb0" if index < len(PROGNOSTIC_CHANNELS) else "#c05621" for index in range(len(CHANNELS))]
    axis.bar(np.arange(len(CHANNELS)), rmse, color=colors)
    axis.set_xticks(np.arange(len(CHANNELS)))
    axis.set_xticklabels(CHANNELS, rotation=75, ha="right", fontsize=7)
    axis.set_ylabel("RMSE")
    axis.set_title("Forecast RMSE by output channel")
    axis.grid(axis="y", alpha=0.25)
    fig.savefig(path, dpi=dpi, bbox_inches="tight")
    plt.close(fig)
    return rmse


def main() -> int:
    args = parse_args()
    import yaml
    with args.config.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle) or {}
    prediction_path = args.prediction_path or configured_path(config, "infer_path", INFER_PATH)
    output_dir = args.output_dir or configured_path(config, "pic_dir", PIC_DIR)
    if not prediction_path.exists():
        raise SystemExit(f"prediction not found: {prediction_path}; run 'python scripts/inference.py' first")
    pred_data = np.load(prediction_path)
    if "predictions" not in pred_data:
        raise KeyError("inference NPZ must contain predictions")
    truth_raw, truth_data = _truth(pred_data, args.truth_path, config)
    prediction = _rollout(pred_data["predictions"], "prediction")
    truth = _rollout(truth_raw, "truth")
    if truth.shape[0] == prediction.shape[0]:
        eval_pred, eval_truth = prediction, truth
    else:
        eval_pred, eval_truth = prediction[:1], truth[:1]
    height, width = prediction.shape[-2:]
    lat, lon = _coordinates(pred_data, truth_data, height=height, width=width)
    area = _area_grid(args.area_weights, lat, width)
    metrics = forecast_metrics(torch.from_numpy(eval_pred).float(), torch.from_numpy(eval_truth).float(), torch.from_numpy(area).float())
    means = _global_mean(prediction, area)
    truth_means = _global_mean(truth, area) if truth.shape[0] == prediction.shape[0] else None
    output_dir.mkdir(parents=True, exist_ok=True)
    (output_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    _save_maps(output_dir / "rollout_maps.png", prediction, lat, lon, args.dpi)
    _save_comparison(output_dir / "first_step_comparison.png", eval_pred, eval_truth, lat, lon, args.dpi)
    _save_series(output_dir / "global_mean_timeseries.png", means, truth_means, args.dpi)
    _save_heatmap(output_dir / "all_channel_global_means.png", means, args.dpi)
    rmse = _save_rmse(output_dir / "channel_rmse.png", eval_pred, eval_truth, args.dpi)
    manifest = {"prediction_path": str(prediction_path), "shape": list(prediction.shape), "channels": list(CHANNELS), "figures": ["rollout_maps.png", "first_step_comparison.png", "global_mean_timeseries.png", "all_channel_global_means.png", "channel_rmse.png"]}
    (output_dir / "visualization_manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(json.dumps({"status": "success", "metrics": metrics, "output_dir": str(output_dir), "figures": 5, "shape": list(prediction.shape), "max_channel_rmse": float(rmse.max())}))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())