script/result.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

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

from model.earthformer import Earthformer
from script.data_loader import make_loader
from script.metrics import metric_sums, metrics_from_sums
from script.utils import clean_state_dict, load_checkpoint_payload, resolve_cli_path, resolve_device


def _squeeze(frame: np.ndarray) -> np.ndarray:
    return np.squeeze(frame)


def plot_lead_time(truth: np.ndarray, prediction: np.ndarray, output_path: Path, title: str) -> Path:
    """Three-panel Truth / Prediction / Difference image, fuxi/fengwu style."""
    truth, prediction = _squeeze(truth), _squeeze(prediction)
    difference = truth - prediction
    rmse = float(np.sqrt(np.mean(difference**2)))
    vmin, vmax = min(truth.min(), prediction.min()), max(truth.max(), prediction.max())
    diff_abs_max = max(float(np.abs(difference).max()), 1e-8)

    panels = [
        {"data": truth, "title": "Truth", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
        {"data": prediction, "title": "Prediction", "cmap": "viridis", "vmin": vmin, "vmax": vmax},
        {"data": difference, "title": f"Difference (RMSE={rmse:.2f})", "cmap": "RdBu_r", "vmin": -diff_abs_max, "vmax": diff_abs_max},
    ]
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for ax, panel in zip(axes, panels):
        image = ax.imshow(panel["data"], cmap=panel["cmap"], vmin=panel["vmin"], vmax=panel["vmax"])
        ax.set_title(panel["title"], fontsize=12, pad=4)
        ax.set_xlabel("Pixel")
        ax.set_ylabel("Pixel")
        plt.colorbar(image, ax=ax, orientation="horizontal")
    fig.suptitle(title, fontsize=14, fontweight="bold", y=0.98)
    fig.savefig(output_path, dpi=300, bbox_inches="tight")
    plt.close(fig)
    return output_path


def visualize_predictions(
    truth: np.ndarray,
    prediction: np.ndarray,
    output_dir: str | Path,
    max_samples: int = 2,
    stride: int = 2,
) -> list[str]:
    """Write per-sample/lead-time PNG images and a per-sample mosaic into output_dir."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    if truth.shape != prediction.shape:
        raise ValueError(f"truth and prediction must have the same shape, got {truth.shape} vs {prediction.shape}")
    samples = min(max_samples, int(truth.shape[0]))
    lead_times = list(range(0, int(truth.shape[1]), stride))
    generated: list[str] = []
    for sample in range(samples):
        for lead_time in lead_times:
            path = output_dir / f"earthformer_sample_{sample:03d}_t{lead_time:02d}.png"
            plot_lead_time(
                truth[sample, lead_time],
                prediction[sample, lead_time],
                path,
                f"Earthformer VIL sample {sample} lead time {lead_time} (+{lead_time * 5} min)",
            )
            generated.append(str(path))
        mosaic = output_dir / f"earthformer_overview_sample_{sample:03d}.png"
        _plot_mosaic(truth[sample], prediction[sample], lead_times, mosaic)
        generated.append(str(mosaic))
    return generated


def _plot_mosaic(truth: np.ndarray, prediction: np.ndarray, lead_times: list[int], output_path: Path) -> Path:
    rows, columns = len(lead_times) * 2, len(lead_times)
    fig, axes = plt.subplots(rows, columns, figsize=(columns * 2.6, rows * 2.2))
    for column, lead_time in enumerate(lead_times):
        truth_frame = _squeeze(truth[lead_time])
        pred_frame = _squeeze(prediction[lead_time])
        vmin, vmax = min(truth_frame.min(), pred_frame.min()), max(truth_frame.max(), pred_frame.max())
        axes[0, column].imshow(truth_frame, cmap="viridis", vmin=vmin, vmax=vmax)
        axes[0, column].set_title(f"+{lead_time * 5} min", fontsize=9)
        axes[len(lead_times), column].imshow(pred_frame, cmap="viridis", vmin=vmin, vmax=vmax)
        for row in range(rows):
            axes[row, column].set_xticks([])
            axes[row, column].set_yticks([])
    axes[0, 0].set_ylabel("Truth", fontsize=10)
    axes[len(lead_times), 0].set_ylabel("Prediction", fontsize=10)
    fig.suptitle("Earthformer VIL sample overview", fontsize=13, fontweight="bold", y=0.99)
    fig.tight_layout(rect=(0, 0, 1, 0.97))
    fig.savefig(output_path, dpi=200, bbox_inches="tight")
    plt.close(fig)
    return output_path


def evaluate(model: torch.nn.Module, loader, device: torch.device) -> dict[str, float]:
    sums = torch.zeros(22, dtype=torch.float64, device=device)
    with torch.no_grad():
        for inputs, targets in loader:
            prediction = model(inputs.to(device)).clamp(0.0, 1.0)
            sums += metric_sums(prediction, targets.to(device))
    result = metrics_from_sums(sums.cpu())
    result["note"] = "Lightweight metrics on configured data; mean CSI is not the official complete SEVIR evaluation"
    return result


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Evaluate Earthformer and visualize predictions (fuxi/fengwu style PNG images)"
    )
    parser.add_argument("--checkpoint", default="data/checkpoint/earthformer.pt")
    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
    parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
    parser.add_argument("--predictions", default="output/predictions.npz", help="Inference arrays for visualization")
    parser.add_argument("--output-dir", default="output/visualization", help="Directory for visualization PNG images")
    parser.add_argument("--max-samples", type=int, default=2)
    parser.add_argument("--stride", type=int, default=2, help="Lead-time stride, aligned with official plot_stride")
    parser.add_argument("--skip-visualization", action="store_true", help="Only evaluate, do not render PNG images")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    device = resolve_device(args.device)
    payload = load_checkpoint_payload(resolve_cli_path(args.checkpoint), device)
    config = payload["config"]
    model = Earthformer(config).to(device)
    model.load_state_dict(clean_state_dict(payload["model"]))
    model.eval()
    loader, _ = make_loader(config, args.split, shuffle=False)
    result = evaluate(model, loader, device)
    report = {"metrics": result}
    if args.skip_visualization:
        print(json.dumps(report, indent=2))
        return
    predictions_path = Path(resolve_cli_path(args.predictions))
    with np.load(predictions_path) as payload:
        truth = payload["targets"]
        prediction = payload["predictions"]
    output_dir = resolve_cli_path(args.output_dir)
    images = visualize_predictions(truth, prediction, output_dir, args.max_samples, args.stride)
    report["output_dir"] = output_dir
    report["images"] = images
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()