File size: 2,004 Bytes
f1d3656 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | """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))
|