| """Evaluate 24-hour rainfall and produce per-step metrics and maps.""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| import yaml |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def categorical(prediction, target, threshold): |
| forecast, observed = prediction >= threshold, target >= threshold |
| hits = np.logical_and(forecast, observed).sum(); false = np.logical_and(forecast, ~observed).sum() |
| misses = np.logical_and(~forecast, observed).sum(); correct_negative = np.logical_and(~forecast, ~observed).sum() |
| total = hits + false + misses + correct_negative; eps = 1e-8 |
| random_hits = (hits + misses) * (hits + false) / max(total, 1) |
| return { |
| "ACC": float((hits + correct_negative) / (total + eps)), |
| "TS": float(hits / (hits + misses + false + eps)), |
| "ETS": float((hits - random_hits) / (hits + misses + false - random_hits + eps)), |
| "BS": float((hits + false) / (hits + misses + eps)), |
| "FAR": float(false / (hits + false + eps)), |
| } |
|
|
|
|
| def fss(prediction, target, threshold, window): |
| forecast = torch.from_numpy((prediction >= threshold).astype(np.float32))[:, None] |
| observed = torch.from_numpy((target >= threshold).astype(np.float32))[:, None] |
| pf = torch.nn.functional.avg_pool2d(forecast, window, stride=1, padding=window // 2) |
| po = torch.nn.functional.avg_pool2d(observed, window, stride=1, padding=window // 2) |
| return float(1 - (pf - po).square().sum() / (pf.square().sum() + po.square().sum() + 1e-8)) |
|
|
|
|
| def main(): |
| config = yaml.safe_load((ROOT / "conf/config.yaml").read_text()) |
| archive = np.load(ROOT / config["paths"]["inference"], allow_pickle=False) |
| steps = int(config["data"]["time_steps"]) |
| lead_hours = [int(value) for value in archive["lead_hours"]] |
| expected_3h = (steps, int(config["data"]["height"]), int(config["data"]["width"])) |
| for key in ("raw_3h", "target_3h", "corrected_3h"): |
| if archive[key].shape[1:] != expected_3h: |
| raise ValueError(f"{key} must have per-day shape {expected_3h}") |
| if lead_hours != config["data"]["lead_hours"]: |
| raise ValueError("inference NPZ lead_hours do not match configuration") |
| thresholds = [float(value) for value in config["evaluation"]["thresholds_mm_24h"]] |
| windows = [int(value) for value in config["evaluation"]["fss_windows"]] |
| metrics = {"units_24h": "mm/24h", "units_per_step": "mm/3h", "per_step_length": steps, |
| "lead_hours": lead_hours, "raw_24h": {}, "corrected_24h": {}, "per_step": {}} |
| for label in ("raw", "corrected"): |
| prediction_24h = archive[f"{label}_24h"] |
| metrics[f"{label}_24h"] = {str(int(t)): categorical(prediction_24h, archive["target_24h"], t) for t in thresholds} |
| metrics[f"{label}_24h"]["FSS@50"] = {str(w): fss(prediction_24h, archive["target_24h"], 50, w) for w in windows} |
| per_step = [] |
| for step in range(steps): |
| item = {"step": step, "lead_hour": lead_hours[step]} |
| item.update({str(int(t)): categorical(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], t) for t in thresholds}) |
| item["FSS@50"] = {str(w): fss(archive[f"{label}_3h"][:, step], archive["target_3h"][:, step], 50, w) for w in windows} |
| per_step.append(item) |
| metrics["per_step"][label] = per_step |
| output = ROOT / config["paths"]["evaluation_dir"] |
| output.mkdir(parents=True, exist_ok=True) |
| (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") |
| figure, axes = plt.subplots(4, steps, figsize=(18, 9), constrained_layout=True) |
| rows = (archive["raw_3h"][0], archive["target_3h"][0], archive["corrected_3h"][0], archive["corrected_3h"][0] - archive["target_3h"][0]) |
| labels = ("ECMWF", "CMPA target", "Corrected", "Error") |
| vmax = max(float(rows[0].max()), float(rows[1].max()), 1) |
| for row, (fields, label) in enumerate(zip(rows, labels)): |
| for step in range(steps): |
| cmap = "RdBu_r" if row == 3 else "Blues" |
| limit = vmax if row < 3 else max(vmax / 2, 1) |
| axes[row, step].imshow(fields[step], cmap=cmap, vmin=-limit if row == 3 else 0, vmax=limit) |
| axes[row, step].axis("off") |
| if row == 0: axes[row, step].set_title(f"+{archive['lead_hours'][step]} h") |
| if step == 0: axes[row, step].set_ylabel(label) |
| figure.savefig(output / "comparison.png", dpi=150); plt.close(figure) |
| print(f"metrics={output / 'metrics.json'} figure={output / 'comparison.png'}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
