File size: 3,356 Bytes
7f71cfd | 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 | 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 main():
cfg = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
data = np.load(ROOT / cfg["paths"]["inference"])
pred, target = data["prediction"], data["target"]
if pred.shape != target.shape or pred.shape[1:] != (65, 68):
raise ValueError(f"Expected prediction and target shaped [N,65,68], got {pred.shape} and {target.shape}")
mass = np.concatenate((data["layer_mass"], data["layer_mass"]))
mass = mass / mass.mean()
error = pred - target
mad = (np.abs(error) * mass[None, None]).mean(axis=(0, 2))
bias = error.mean(axis=(0, 2))
flat_target = target[:, 1:].astype(np.float64).reshape(target.shape[0], -1)
flat_pred = pred[:, 1:].astype(np.float64).reshape(pred.shape[0], -1)
source_r2 = []
for truth, estimate in zip(flat_target, flat_pred):
denom = np.sum((truth - truth.mean()) ** 2)
source_r2.append(float(1.0 - np.sum((truth - estimate) ** 2) / max(denom, 1e-12)))
dt = float(data["dt_seconds"])
qt_pred, qt_target = pred[..., 34:], target[..., 34:]
column_water_pred = (qt_pred * data["layer_mass"][None, None]).sum(-1)
column_water_target = (qt_target * data["layer_mass"][None, None]).sum(-1)
lhf_evap = data["surface"][:, :, 1] / 2.5e6
adv_q = (0.5 * (data["advection"][:, :-1, 34:] + data["advection"][:, 1:, 34:]) *
data["layer_mass"][None, None]).sum(-1)
storage = np.diff(column_water_pred, axis=1) / dt
precipitation = np.maximum(0.0, lhf_evap + adv_q - storage)
residual = storage + precipitation - lhf_evap - adv_q
metrics = {"mass_weighted_mad_per_step": mad.tolist(), "bias_per_step": bias.tolist(),
"source_r2": source_r2, "mean_source_r2": float(np.mean(source_r2)),
"water_budget": {"mean_precipitation_kg_m2_s": float(precipitation.mean()),
"mean_abs_residual_kg_m2_s": float(np.abs(residual).mean())},
"per_step_length": int(len(mad)), "step_hours": cfg["data"]["step_hours"]}
out = ROOT / cfg["paths"]["evaluation"]
out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(json.dumps(metrics, indent=2))
hours = data["lead_hours"]
fig, axes = plt.subplots(2, 1, figsize=(9, 7), constrained_layout=True)
axes[0].plot(hours, target[:, :, :34].mean((0, 2)), label="target sL")
axes[0].plot(hours, pred[:, :, :34].mean((0, 2)), "--", label="predicted sL")
ax2 = axes[0].twinx()
ax2.plot(hours, target[:, :, 34:].mean((0, 2)), color="tab:green", label="target qT")
ax2.plot(hours, pred[:, :, 34:].mean((0, 2)), "--", color="tab:red", label="predicted qT")
axes[0].set(xlabel="lead time (h)", title="SCM state rollout")
axes[0].legend(loc="upper left")
ax2.legend(loc="upper right")
axes[1].plot(hours[1:], precipitation.mean(0) * 86400.0, color="navy")
axes[1].set(xlabel="lead time (h)", ylabel="mm day-1", title="Diagnosed precipitation")
fig.savefig(out.parent / "state_precipitation_timeseries.png", dpi=150)
plt.close(fig)
print(f"saved {out}: per_step={len(mad)}, source_R2={metrics['mean_source_r2']:.4f}")
if __name__ == "__main__":
main()
|