| """Show the saved tiny AtmoRep training and inference summaries.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| if __name__ == "__main__": |
| history_path = ROOT / "weight" / "training" / "history.json" |
| train_history = json.loads(history_path.read_text()) if history_path.exists() else [] |
| prediction = torch.load(ROOT / "result" / "prediction.pt", map_location="cpu", weights_only=True) |
| target = torch.load(ROOT / "result" / "target.pt", map_location="cpu", weights_only=True) |
| ensemble = prediction["ensemble"] |
| target = target["target"] if isinstance(target, dict) else target |
| mean = ensemble.mean(dim=1) |
| std = ensemble.std(dim=1, unbiased=False) |
| ensemble_rmse = torch.sqrt((ensemble - target[:, None]).square().mean()) |
| mean_rmse = torch.sqrt((mean - target).square().mean()) |
| target_std = target.std(dim=-1, unbiased=False, keepdim=True).expand_as(target) |
| spread_rmse = torch.sqrt((std - target_std).square().mean()) |
| summary = { |
| "train_history": train_history, |
| "ensemble_shape": list(ensemble.shape), |
| "target_shape": list(target.shape), |
| "ensemble_rmse": float(ensemble_rmse), |
| "mean_rmse": float(mean_rmse), |
| "spread_rmse": float(spread_rmse), |
| "finite": bool(torch.isfinite(ensemble).all()), |
| } |
| (ROOT / "result" / "metrics.json").write_text(json.dumps(summary, indent=2) + "\n") |
| fig, axes = plt.subplots(1, 4, figsize=(12, 3)) |
| for ax, data, title in zip(axes, (target[0, 0], mean[0, 0], (mean - target)[0, 0].abs(), std[0, 0]), ("Target", "Mean", "Abs error", "Spread")): |
| image = ax.imshow(data.reshape(4, 4).numpy(), cmap="viridis") |
| ax.set_title(title) |
| plt.colorbar(image, ax=ax, shrink=0.75) |
| plt.tight_layout() |
| plt.savefig(ROOT / "result" / "comparison.png", dpi=150) |
| plt.close() |
| print(json.dumps(summary, indent=2)) |
|
|
