File size: 4,722 Bytes
807a08b | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | """Compute metrics and render ClimODE forecast maps from saved outputs."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import numpy as np
import yaml
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.metrics import evaluate, save_metrics
def _resolve(path: str | Path) -> Path:
value = Path(path)
return value if value.is_absolute() else PROJECT_ROOT / value
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", type=Path, default=PROJECT_ROOT / "conf/config.yaml")
parser.add_argument("--predictions", type=Path, default=None)
parser.add_argument("--targets", type=Path, default=None)
parser.add_argument("--std", type=Path, default=None)
parser.add_argument("--output-dir", type=Path, default=None)
parser.add_argument("--stats-dir", type=Path, default=None)
parser.add_argument("--static-file", type=Path, default=None)
parser.add_argument("--sample", type=int, default=0)
parser.add_argument("--lead", type=int, default=0)
args = parser.parse_args()
with args.config.open("r", encoding="utf-8") as handle:
config = yaml.safe_load(handle)
output_dir = args.output_dir or _resolve(config["data"]["output_dir"])
predictions = np.load(args.predictions or output_dir / "predictions.npy")
targets = np.load(args.targets or output_dir / "targets.npy")
std_path = args.std or output_dir / "std.npy"
std = np.load(std_path) if std_path.is_file() else None
lengths_path = output_dir / "valid_lengths.npy"
valid_lengths = np.load(lengths_path) if lengths_path.is_file() else None
static_file = _resolve(args.static_file or config["data"]["static_file"])
import h5py
with h5py.File(static_file, "r") as handle:
lat2d = handle["lat2d"][:]
stats_dir = _resolve(
args.stats_dir
or config["data"].get("stats_dir", Path(config["data"]["data_dir"]) / "static")
)
minimum = np.load(stats_dir / "min_values.npy").reshape(1, 1, 1, 5, 1, 1)
maximum = np.load(stats_dir / "max_values.npy").reshape(1, 1, 1, 5, 1, 1)
scale = maximum - minimum
metrics = evaluate(
predictions * scale + minimum,
targets * scale + minimum,
lat2d,
std * scale if std is not None else None,
crps_predictions=predictions if std is not None else None,
crps_targets=targets if std is not None else None,
crps_std=std,
valid_lengths=valid_lengths,
)
metrics_path = output_dir.parent / "metrics.json"
save_metrics(metrics, metrics_path)
figure_dir = output_dir / "figures"
figure_dir.mkdir(parents=True, exist_ok=True)
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except ImportError as exc:
raise RuntimeError("Visualization requires matplotlib in the active environment") from exc
if not 0 <= args.sample < predictions.shape[0]:
raise IndexError(f"sample must be in [0,{predictions.shape[0] - 1}]")
if not 0 <= args.lead < predictions.shape[1]:
raise IndexError(f"lead must be in [0,{predictions.shape[1] - 1}]")
if valid_lengths is not None and args.lead >= int(valid_lengths[args.sample]):
raise IndexError(
f"lead {args.lead} is padding for sample {args.sample}; "
f"valid length is {int(valid_lengths[args.sample])}"
)
names = ["z", "t", "t2m", "u10", "v10"]
for channel, name in enumerate(names):
prediction = predictions[args.sample, args.lead, 0, channel]
target = targets[args.sample, args.lead, 0, channel]
difference = prediction - target
figure, axes = plt.subplots(1, 3, figsize=(12, 3.4), constrained_layout=True)
for axis, image, title in zip(
axes,
(prediction, target, difference),
("prediction", "target", "difference"),
):
cmap = "RdBu_r" if title == "difference" else "viridis"
plot = axis.imshow(image, cmap=cmap, origin="upper", aspect="auto")
axis.set_title(title)
axis.set_xlabel("longitude index")
axis.set_ylabel("latitude index")
figure.colorbar(plot, ax=axis, shrink=0.8)
figure.suptitle(f"ClimODE {name}, lead={(args.lead + 1) * 6} h")
figure.savefig(figure_dir / f"{name}_lead_{(args.lead + 1) * 6:03d}h.png", dpi=150)
plt.close(figure)
print(json.dumps({"metrics": str(metrics_path), "figures": str(figure_dir)}))
if __name__ == "__main__":
main()
|