File size: 4,192 Bytes
2e913c2 | 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 | """Compute ClimateBench spatial/global NRMSE and plot all four targets."""
import json
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
def area_mean(field: np.ndarray, latitude: np.ndarray) -> np.ndarray:
weights = np.cos(np.deg2rad(latitude)).clip(0)[:, None]
return np.sum(field * weights, axis=(-2, -1)) / (weights.sum() * field.shape[-1])
def climatebench_metrics(prediction: np.ndarray, target: np.ndarray, latitude: np.ndarray, alpha: float) -> dict:
normalization = float(area_mean(target, latitude).mean())
if abs(normalization) <= 1e-12:
raise ValueError("area-weighted target mean must be nonzero for NRMSE normalization")
temporal_rmse = np.sqrt(np.mean((prediction - target) ** 2, axis=0))
spatial = float(area_mean(temporal_rmse, latitude) / normalization)
predicted_global = area_mean(prediction, latitude)
target_global = area_mean(target, latitude)
global_nrmse = float(np.sqrt(np.mean((predicted_global - target_global) ** 2)) / normalization)
return {"spatial_nrmse": spatial, "global_nrmse": global_nrmse,
"total_nrmse": spatial + alpha * global_nrmse}
def main() -> None:
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data = np.load(ROOT / config["paths"]["inference"])
if str(data["format_version"]) != config["data"]["format_version"]:
raise ValueError("incompatible prediction format")
start, end = int(data["evaluation_start_year"]), int(data["evaluation_end_year"])
if (start, end) != (int(config["evaluation"]["start_year"]), int(config["evaluation"]["end_year"])):
raise ValueError("evaluation period metadata must be 2080-2100")
names = data["target_names"].tolist()
alpha = float(config["evaluation"]["global_weight"])
metrics = {name: climatebench_metrics(data["predictions"][:, index], data["targets"][:, index],
data["latitude"], alpha) for index, name in enumerate(names)}
report = {"variables": metrics, "mean_total_nrmse": float(np.mean([value["total_nrmse"] for value in metrics.values()])),
"evaluation_protocol": {"scenario": str(data["scenario"]), "period": f"{start}-{end}",
"target_aggregation": str(data["target_aggregation"]),
"spatial_metric": "cosine-area mean of grid-cell temporal RMSE",
"global_weighting": "cosine latitude cell area",
"total": "spatial_nrmse + 5 * global_nrmse",
"normalization": "cosine-area and sample mean of target field"}}
numeric = [number for values in metrics.values() for number in values.values()]
if not np.isfinite(numeric).all():
raise FloatingPointError("evaluation contains NaN or Inf")
output = ROOT / config["paths"]["evaluation_dir"]
output.mkdir(parents=True, exist_ok=True)
(output / "metrics.json").write_text(json.dumps(report, indent=2) + "\n")
figure, axes = plt.subplots(4, 3, figsize=(12, 12), constrained_layout=True)
longitude, latitude = data["longitude"], data["latitude"]
for index, name in enumerate(names):
fields = (data["targets"][:, index].mean(0), data["predictions"][:, index].mean(0),
(data["predictions"][:, index] - data["targets"][:, index]).mean(0))
for column, (field, title) in enumerate(zip(fields, ("Target", "Prediction", "Error"))):
image = axes[index, column].pcolormesh(longitude, latitude, field, shading="auto", cmap="coolwarm")
axes[index, column].set(title=f"{name}: {title}", xlabel="Longitude", ylabel="Latitude")
figure.colorbar(image, ax=axes[index, column], shrink=0.75)
figure.suptitle(f"ClimateBench {str(data['scenario'])} {start}-{end}")
figure.savefig(output / "four_targets.png", dpi=150)
plt.close(figure)
print(f"evaluation={output.relative_to(ROOT)} variables={','.join(names)}")
if __name__ == "__main__":
main()
|