File size: 3,648 Bytes
d3e46b7 | 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 | """Compute per-lead latitude-weighted model and persistence RMSE/MBE."""
import json
from pathlib import Path
import sys
import matplotlib.pyplot as plt
import numpy as np
import yaml
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
from model.fuxi_ocean import ocean_channel_depth_indices, weighted_error_sums
def aggregate(prediction, truth, latitude, depth_mask):
channel_depth = ocean_channel_depth_indices()
mask = np.take(depth_mask, channel_depth, axis=1)
squared, bias, weight = weighted_error_sums(prediction, truth, latitude, mask)
return np.sqrt(squared / np.maximum(weight, 1)), bias / np.maximum(weight, 1)
def main():
config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data = np.load(ROOT / config["paths"]["inference"])
model_rmse, model_mbe, persistence_rmse, persistence_mbe = [], [], [], []
for lead in range(data["prediction"].shape[1]):
rmse, mbe = aggregate(data["prediction"][:, lead], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
prmse, pmbe = aggregate(data["initial"], data["truth"][:, lead], data["latitude_deg"], data["depth_mask"])
model_rmse.append(rmse); model_mbe.append(mbe); persistence_rmse.append(prmse); persistence_mbe.append(pmbe)
arrays = [np.asarray(value) for value in (model_rmse, model_mbe, persistence_rmse, persistence_mbe)]
lead_hours = data["lead_hours"].tolist()
groups = json.loads(str(data["variable_groups"]))
def grouped(values, i):
return {name: {"unit": spec["unit"], "rmse": float(np.mean(values[0][i, slice(*spec["channels"])])),
"mbe": float(np.mean(values[1][i, slice(*spec["channels"])]))} for name, spec in groups.items()}
metrics = {"output_kind": str(data["output_kind"]), "format_version": str(data["format_version"]),
"checkpoint_source": str(data["checkpoint_source"]), "sample_count": int(data["sample_count"]),
"coverage_fraction": float(data["coverage_fraction"]), "is_complete_global": bool(data["is_complete_global"]),
"synthetic": bool(data["synthetic"]), "aggregation": "latitude-weighted sampled tiles; never across unit groups",
"output_shape": data["output_shape"].tolist(), "variable_groups": groups,
"per_lead": [{"lead_hours": hour, "model_by_group": grouped(arrays[:2], i),
"persistence_by_group": grouped(arrays[2:], i),
"model_by_channel": {"rmse": arrays[0][i].tolist(), "mbe": arrays[1][i].tolist()},
"persistence_by_channel": {"rmse": arrays[2][i].tolist(), "mbe": arrays[3][i].tolist()}}
for i, hour in enumerate(lead_hours)]}
path = ROOT / config["paths"]["evaluation_metrics"]; path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(metrics, indent=2) + "\n")
figure, axes = plt.subplots(1, 5, figsize=(18, 3.5))
for axis, (name, spec) in zip(axes, groups.items()):
channel_slice = slice(*spec["channels"])
axis.plot(lead_hours, arrays[0][:, channel_slice].mean(1), "o-", label="model")
axis.plot(lead_hours, arrays[2][:, channel_slice].mean(1), "s--", label="persistence")
axis.set(title=name, xlabel="Lead (h)", ylabel=f"RMSE ({spec['unit']})"); axis.legend()
figure.tight_layout()
plot = ROOT / config["paths"]["evaluation_plot"]; figure.savefig(plot, dpi=160); plt.close(figure)
print(f"metrics={path.relative_to(ROOT)} plot={plot.relative_to(ROOT)} leads={lead_hours}")
if __name__ == "__main__":
main()
|