| """Show the saved SFNO training and inference summaries.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import torch |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| if __name__ == "__main__": |
| result_dir = ROOT / "result" |
| prediction = torch.load(result_dir / "prediction.pt", map_location="cpu", weights_only=True) |
| target = torch.load(result_dir / "target.pt", map_location="cpu", weights_only=True) |
| error = prediction - target |
| rmse_by_step = torch.sqrt(error.square().mean(dim=(1, 2, 3))) |
| pred_anom = prediction - prediction.mean(dim=(2, 3), keepdim=True) |
| target_anom = target - target.mean(dim=(2, 3), keepdim=True) |
| acc_by_step = (pred_anom * target_anom).sum(dim=(1, 2, 3)) / torch.sqrt( |
| pred_anom.square().sum(dim=(1, 2, 3)) |
| * target_anom.square().sum(dim=(1, 2, 3)) |
| ).clamp_min(1e-12) |
| metrics = { |
| "rmse": float(torch.sqrt(error.square().mean())), |
| "rmse_by_step": rmse_by_step.tolist(), |
| "spatial_acc": float(acc_by_step.mean()), |
| "spatial_acc_by_step": acc_by_step.tolist(), |
| } |
| (result_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") |
| fig, axes = plt.subplots(1, 3, figsize=(12, 3.5)) |
| for ax, data, title in zip( |
| axes, |
| (target[0, 0], prediction[0, 0], error[0, 0]), |
| ("Target", "Prediction", "Error"), |
| ): |
| image = ax.imshow(data, cmap="RdBu_r") |
| ax.set_title(title) |
| plt.colorbar(image, ax=ax, shrink=0.75) |
| plt.tight_layout() |
| plt.savefig(result_dir / "comparison.png", dpi=150) |
| plt.close() |
| print(json.dumps(metrics, indent=2)) |
|
|
