#!/usr/bin/env python """Stage 01 — Tx90p climatology, EHD index, and the min()-vs-max() threshold diagnostic. Implements the index EXACTLY as specified (threshold = min(Tx90p, 35 degC)) and, alongside, produces the diagnostic the user asked to see EARLY: how the threshold and resulting EHD frequency differ between min(Tx90p,35) [production] and max(Tx90p,35), by region and month. This stage is a DIAGNOSTIC: the climatology here is fit on the full gridded obs baseline (2002-2023) for illustration. It is NOT a cross-validated skill number; the leave-one-year-out climatology used for verification is fit per-fold later (see CLAUDE.md leakage rule). Outputs: results/diagnostics/tx90p_climatology.nc results/tables/threshold_diagnostic_region_month.csv results/tables/region_centroids.csv results/figures/threshold_diagnostic.png """ from __future__ import annotations import os, sys, logging import numpy as np import pandas as pd import xarray as xr import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from ehd.config import load_config from ehd import climatology as clim logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") log = logging.getLogger("stage01") # Provisional region names (to be verified against centroids below). Codes match the # gridded `reg` mask and station header region codes (SAPA=2, HANOI=3, ALUOI=4, DANANG=5, # PLEIKU=6, CAMAU=7). REGION_NAMES = { 1: "R1 Tay Bac (Northwest)", 2: "R2 Dong Bac (Northeast)", 3: "R3 DB Bac Bo (Red River Delta)", 4: "R4 Bac Trung Bo (North Central)", 5: "R5 Nam Trung Bo (South Central)", 6: "R6 Tay Nguyen (Central Highlands)", 7: "R7 Nam Bo (South)", } def main(): cfg = load_config() res = cfg["paths"]["results"] os.makedirs(os.path.join(res, "diagnostics"), exist_ok=True) os.makedirs(os.path.join(res, "tables"), exist_ok=True) os.makedirs(os.path.join(res, "figures"), exist_ok=True) idx = cfg["index"] abs_c = float(idx["abs_degC"]) pctl = int(idx["percentile"]) win = int(idx["percentile_window_days"]) baseline = list(range(idx["baseline_period"][0], idx["baseline_period"][1] + 1)) # --- load obs + mask --- da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"], cfg["obs"]["gridded_var"], float(cfg["obs"]["gridded_fill"])) log.info("obs loaded: %s %s..%s", da.shape, str(da.time.values[0]), str(da.time.values[-1])) mask = xr.open_dataset(cfg["paths"]["mask"]) reg = mask[cfg["regions"]["mask_var"]].values # (lat, lon) 1..7 or nan land = np.isfinite(reg) years, doy = clim._years_doy(da) months = np.array([t.month for t in da.time.values]) # --- region centroids (verification of names) --- lon2d, lat2d = np.meshgrid(da.lon.values, da.lat.values) cent = [] for r in cfg["regions"]["ids"]: m = reg == r cent.append(dict(region=r, name=REGION_NAMES[r], n_cells=int(m.sum()), lat_centroid=round(float(lat2d[m].mean()), 3), lon_centroid=round(float(lon2d[m].mean()), 3), lat_min=round(float(lat2d[m].min()), 2), lat_max=round(float(lat2d[m].max()), 2))) cent_df = pd.DataFrame(cent).sort_values("lat_centroid", ascending=False) cent_df.to_csv(os.path.join(res, "tables", "region_centroids.csv"), index=False) log.info("region centroids (N->S):\n%s", cent_df.to_string(index=False)) # --- Tx90p climatology (diagnostic baseline = full period) --- log.info("computing Tx90p climatology (pctl=%d, window=%d, baseline=%d-%d)...", pctl, win, baseline[0], baseline[-1]) tx90p = clim.tx90p_climatology(da, baseline, percentile=pctl, window_days=win) tx90p.to_netcdf(os.path.join(res, "diagnostics", "tx90p_climatology.nc")) thr_min = clim.ehd_threshold(tx90p, abs_c, "min") thr_max = clim.ehd_threshold(tx90p, abs_c, "max") # broadcast per-doy thresholds onto every obs time step doy_da = xr.DataArray(doy, dims="time") thr_min_t = thr_min.sel(doy=doy_da).values # (time, lat, lon) thr_max_t = thr_max.sel(doy=doy_da).values tx90p_t = tx90p.sel(doy=doy_da).values obs = da.values ehd_min = (obs >= thr_min_t) ehd_max = (obs >= thr_max_t) pctl_controlled = (tx90p_t < abs_c) # where min() is set by the percentile valid = np.isfinite(obs) # --- aggregate by region x month --- rows = [] for r in cfg["regions"]["ids"]: cellmask = (reg == r) # (lat, lon) for mo in range(1, 13): tsel = months == mo v = valid[tsel][:, cellmask] # (ndays, ncells) if v.sum() == 0: continue def mean_over(a): aa = a[tsel][:, cellmask] return float(np.nanmean(aa[v])) rows.append(dict( region=r, region_name=REGION_NAMES[r], month=mo, n_cells=int(cellmask.sum()), tx90p_mean=round(mean_over(tx90p_t), 2), thr_min_mean=round(mean_over(thr_min_t), 2), thr_max_mean=round(mean_over(thr_max_t), 2), frac_pctl_controlled=round(mean_over(pctl_controlled.astype(float)), 3), ehd_freq_min=round(mean_over(ehd_min.astype(float)), 4), ehd_freq_max=round(mean_over(ehd_max.astype(float)), 4), )) df = pd.DataFrame(rows) df["ehd_freq_ratio_min_over_max"] = ( df["ehd_freq_min"] / df["ehd_freq_max"].replace(0, np.nan)).round(2) out_csv = os.path.join(res, "tables", "threshold_diagnostic_region_month.csv") df.to_csv(out_csv, index=False) log.info("wrote %s (%d rows)", out_csv, len(df)) # annual summary per region ann = df.groupby(["region", "region_name"]).agg( tx90p_mean=("tx90p_mean", "mean"), ehd_freq_min=("ehd_freq_min", "mean"), ehd_freq_max=("ehd_freq_max", "mean"), frac_pctl_controlled=("frac_pctl_controlled", "mean"), ).round(4).reset_index() log.info("annual-mean EHD frequency by region:\n%s", ann.to_string(index=False)) # --- figure: EHD frequency by month, min vs max, per region --- fig, axes = plt.subplots(2, 4, figsize=(16, 7), sharex=True) axes = axes.ravel() cb_blue, cb_orange = "#0072B2", "#E69F00" for ax, r in zip(axes, cfg["regions"]["ids"]): sub = df[df.region == r].sort_values("month") ax.plot(sub.month, sub.ehd_freq_min, "-o", color=cb_blue, label="min(Tx90p,35) [production]") ax.plot(sub.month, sub.ehd_freq_max, "-s", color=cb_orange, label="max(Tx90p,35)") ax.set_title(REGION_NAMES[r], fontsize=9) ax.set_ylim(0, max(0.05, float(sub.ehd_freq_min.max()) * 1.1)) ax.grid(alpha=0.3) ax.set_xticks(range(1, 13)) axes[0].set_ylabel("EHD frequency"); axes[4].set_ylabel("EHD frequency") for ax in axes[4:]: ax.set_xlabel("month") axes[-1].axis("off") h, l = axes[0].get_legend_handles_labels() axes[-1].legend(h, l, loc="center", fontsize=10, frameon=False) fig.suptitle("Observed EHD frequency by region & month: min(Tx90p,35) vs max(Tx90p,35)\n" "gridded obs 2002-2023 (diagnostic climatology, not cross-validated)", fontsize=11) fig.tight_layout(rect=[0, 0, 1, 0.95]) fig_path = os.path.join(res, "figures", "threshold_diagnostic.png") fig.savefig(fig_path, dpi=130) log.info("wrote %s", fig_path) if __name__ == "__main__": main()