OneScience-Group/PrecipDD

1"""Compute daily/annual detection metrics and a 7x7 occlusion trend map."""2 3import sys4from pathlib import Path5 6import numpy as np7import torch8 9 10ROOT = Path(__file__).resolve().parents[1]11sys.path.insert(0, str(ROOT))12from model.precipdd import correlation, ensemble_predict, linear_trend, load_config, load_ensemble, write_json13 14 15def main():16    import matplotlib17    matplotlib.use("Agg")18    import matplotlib.pyplot as plt19 20    config = load_config(ROOT / "conf/config.yaml")21    values = np.load(ROOT / config["paths"]["predictions"])22    prediction, target, year = values["prediction"], values["target"], values["year"].astype(float)23    if prediction.shape != target.shape or values["precipitation"].shape[1:] != (1, 55, 160):24        raise ValueError("inference output violates scalar target or [N,1,55,160] input contract")25    years = np.unique(year).astype(int)26    annual_prediction = np.array([prediction[year == current].mean() for current in years])27    annual_target = np.array([target[year == current].mean() for current in years])28    threshold = config["evaluation"]["emergence_threshold_c"]29    em_fraction = np.array([(prediction[year == current] > threshold).mean() for current in years])30    metrics = {31        "daily": {"correlation": correlation(target, prediction), "rmse_c": float(np.sqrt(np.mean((prediction - target) ** 2)))},32        "annual": {"correlation": correlation(annual_target, annual_prediction), "rmse_c": float(np.sqrt(np.mean((annual_prediction - annual_target) ** 2)))},33        "emergence": {"threshold_c": threshold, "fraction_all_days": float((prediction > threshold).mean()),34                      "fraction_trend_per_decade": linear_trend(em_fraction, years.astype(float))},35        "trend_c_per_decade": {"daily_prediction": linear_trend(prediction, year), "annual_prediction": linear_trend(annual_prediction, years.astype(float)),36                               "annual_target": linear_trend(annual_target, years.astype(float))},37        "samples": {"daily": len(prediction), "years": len(years)}38    }39 40    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")41    models, _ = load_ensemble(ROOT / config["paths"]["checkpoint"], device)42    count = min(config["evaluation"]["occlusion_max_days"], len(prediction))43    indices = np.linspace(0, len(prediction) - 1, count, dtype=int)44    selected = torch.from_numpy(values["precipitation"][indices]).float().to(device)45    baseline = ensemble_predict(models, selected).cpu().numpy()46    selected_year = year[indices]47    patch, stride = config["evaluation"]["occlusion_patch"], config["evaluation"]["occlusion_stride"]48    half = patch // 249    sensitivity = np.zeros((55, 160), dtype=np.float32)50    for lat_start in range(0, 55, stride):51        for lon_start in range(0, 160, stride):52            masked = selected.clone()53            lat_stop, lon_stop = min(lat_start + patch, 55), min(lon_start + patch, 160)54            masked[:, :, lat_start:lat_stop, lon_start:lon_stop] = 0.055            delta = baseline - ensemble_predict(models, masked).cpu().numpy()56            score = linear_trend(delta, selected_year) if np.ptp(selected_year) else float(delta.mean())57            sensitivity[lat_start:min(lat_start + stride, 55), lon_start:min(lon_start + stride, 160)] = score58    metrics["occlusion"] = {"patch": [patch, patch], "map_shape": [55, 160], "stride": stride, "sampled_days": count,59                            "quantity": "AGMT occlusion-sensitivity trend in degC per decade"}60    write_json(ROOT / config["paths"]["evaluation_metrics"], metrics)61 62    fig, axes = plt.subplots(3, 1, figsize=(11, 11), constrained_layout=True)63    axes[0].plot(years, annual_target, color="#202020", label="target AGMT")64    axes[0].plot(years, annual_prediction, color="#c84c32", label="DD estimate")65    axes[0].axhline(threshold, color="#777777", linestyle="--", label="0.42 C EM threshold")66    axes[0].set(ylabel="AGMT anomaly (C)", title="Annual mean of daily estimates")67    axes[0].legend(ncol=3)68    axes[1].plot(years, em_fraction, color="#196f82")69    axes[1].set(xlabel="Year", ylabel="Fraction", ylim=(-0.03, 1.03), title="Emergence days (estimated AGMT > 0.42 C)")70    limit = float(np.max(np.abs(sensitivity))) or 1e-671    image = axes[2].imshow(sensitivity, origin="lower", aspect="auto", extent=(0, 400, values["latitude"][0], values["latitude"][-1]),72                           cmap="RdBu_r", vmin=-limit, vmax=limit)73    axes[2].set(xlabel="Longitude (degrees E, extended)", ylabel="Latitude", title="7x7 occlusion-sensitivity trend")74    fig.colorbar(image, ax=axes[2], label="C decade-1")75    figure = ROOT / config["paths"]["comparison_figure"]76    figure.parent.mkdir(parents=True, exist_ok=True)77    fig.savefig(figure, dpi=160)78    plt.close(fig)79    print(f"metrics={config['paths']['evaluation_metrics']} figure={config['paths']['comparison_figure']}")80 81 82if __name__ == "__main__":83    main()84