Northern-Thailand-Ping-River-Monitor/scripts/backtest_render.py

#!/usr/bin/env python3
"""Regenerate the documented P.1 flood-backtest charts in docs/img/.

For each chart an eval-only model (regression 24 h peak + warning classifier)
is trained on data STRICTLY BEFORE the event, then the event window is walked
hour by hour exactly as the live system would have seen it:

  backtest-2024-p1.png         Oct 2024 record flood, trained < 1 Sep 2024
  backtest-2024-p1-detail.png  22-28 Sep 2024 zoom of the first crossing
  backtest-2025-p1.png         Sep 2025 flood, deployed config (trained <= 2024)

This codifies the previously prose-only acceptance test: the run fails with a
non-zero exit if the model gives less than 12 h of warning before the first
3.70 m crossing of the 2024 event.

Usage:
    python scripts/backtest_render.py                 # uses FLOOD_ML_DB_URL / POSTGRES_*
    python scripts/backtest_render.py --db-url postgresql://...
"""

import argparse
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

import matplotlib

matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd

from src.ml import data, features
from src.ml.train import _make_classifier, _make_regressor

STATION = "P.1"
STAGE1 = 3.70  # official Chiang Mai stage 1 - city flooding begins
STAGE7 = 4.60  # stage 7 - widespread
HORIZON = 24

INK = "#132b35"
BLUE = "#1c6ea4"
AMBER = "#c07d10"
RED = "#d9534f"


def fit_backtest_model(df_long: pd.DataFrame, train_end: str, use_dam: bool = False):
    """Train the 24 h regression + warning heads on rows <= train_end only.

    Mirrors the deployed hgb-v3 pipeline: the regression head learns the RISE
    over the current level, with Open-Meteo catchment-rain features (trailing
    sums + the forward-24h forecast sum); label statistics are bounded to the
    training cutoff. use_dam=True adds the Mae Ngat reservoir columns — an
    ablation-only configuration (2026-08-13 result: costs 1-3 h of lead).
    """
    from src.ml import dam as dam_mod
    from src.ml import rain as rain_mod

    rain_series = rain_mod.catchment_mean(rain_mod.load_history())
    dam_frame = dam_mod.load_history() if use_dam else None
    X, Y, _meta = features.build_matrix(
        df_long,
        STATION,
        (HORIZON,),
        stats_end=train_end,
        rain=rain_series,
        dam=dam_frame,
    )
    train_mask = X.index <= pd.Timestamp(train_end)
    X_train, Y_train = X.loc[train_mask], Y.loc[train_mask]

    max_col, warn_col = f"max_level_{HORIZON}", f"exceed_warn_{HORIZON}"
    reg_rows = Y_train[max_col].notna()
    rise = Y_train.loc[reg_rows, max_col] - X_train.loc[reg_rows, "level"]
    reg = _make_regressor().fit(X_train.loc[reg_rows], rise)
    warn_rows = Y_train[warn_col].notna()
    clf = _make_classifier().fit(
        X_train.loc[warn_rows], Y_train.loc[warn_rows, warn_col].astype(int)
    )
    return X, reg, clf


def event_series(df_long, X, reg, clf, window_start: str, window_end: str):
    """Observed level plus the forecasts the model would have issued hourly."""
    grid = features.make_hourly_grid(df_long)
    # observed has MultiIndex columns (station_code, field)
    observed = grid.observed[(STATION, "water_level")]
    observed = observed.loc[window_start:window_end].dropna().astype(float)

    Xw = X.loc[window_start:window_end]
    forecasts = pd.DataFrame(index=Xw.index)
    # reg predicts the rise; add the current level back (as serving does)
    forecasts["pred_max"] = reg.predict(Xw) + Xw["level"].to_numpy()
    # Belt-and-braces probability: the classifier OR the regression-sigmoid,
    # whichever is more alarmed. The classifier alone proved unreliable on
    # out-of-distribution extremes (silent on the 2024 record flood).
    import numpy as np

    p_clf = clf.predict_proba(Xw)[:, 1]
    p_sig = 1.0 / (1.0 + np.exp(-(forecasts["pred_max"] - STAGE1) / 0.15))
    forecasts["p_flood"] = np.maximum(p_clf, p_sig)

    flood_start = observed[observed >= STAGE1].index.min()
    alerts = forecasts[forecasts["p_flood"] >= 0.5].index
    first_alert = alerts.min() if len(alerts) else None
    return observed, forecasts, flood_start, first_alert


def _style_axes(ax):
    ax.spines[["top", "right"]].set_visible(False)
    ax.tick_params(colors=INK, labelsize=11)
    ax.grid(axis="y", color="#dfe9e7", linewidth=0.8)
    ax.set_axisbelow(True)


def render(
    observed,
    forecasts,
    flood_start,
    first_alert,
    out_path,
    *,
    title,
    subtitle,
    detail=False,
    show_stage7=False,
    peak_note=None,
):
    fig, (ax, axp) = plt.subplots(
        2,
        1,
        figsize=(12.6, 7.6),
        sharex=True,
        gridspec_kw={"height_ratios": [2.2, 1], "hspace": 0.12},
    )
    fig.patch.set_facecolor("white")

    marker = dict(marker="o", markersize=3) if detail else {}
    ax.plot(
        observed.index,
        observed.values,
        color=BLUE,
        linewidth=2.2,
        label="Observed level" + (" (hourly)" if detail else ""),
        **marker,
    )
    marker = dict(marker="s", markersize=3) if detail else {}
    ax.plot(
        forecasts.index,
        forecasts["pred_max"],
        color=AMBER,
        linewidth=2,
        linestyle="--",
        label="Predicted 24 h peak (issued at that hour)",
        **marker,
    )

    ax.axhline(STAGE1, color=RED, linewidth=1, alpha=0.65)
    ax.annotate(
        f"{STAGE1:.2f} m · stage 1 · flooding begins",
        xy=(0.06, STAGE1),
        xycoords=("axes fraction", "data"),
        xytext=(0, 5),
        textcoords="offset points",
        color=RED,
        fontsize=10.5,
    )
    if show_stage7:
        ax.axhline(STAGE7, color=RED, linewidth=1, alpha=0.65)
        ax.annotate(
            f"{STAGE7:.2f} m · stage 7 · widespread",
            xy=(0.06, STAGE7),
            xycoords=("axes fraction", "data"),
            xytext=(0, 5),
            textcoords="offset points",
            color=RED,
            fontsize=10.5,
        )

    if peak_note:
        peak_ts = observed.idxmax()
        ax.annotate(
            peak_note,
            xy=(peak_ts, observed.max()),
            xytext=(12, 10),
            textcoords="offset points",
            color=BLUE,
            fontsize=11.5,
            fontweight="bold",
        )

    ax.set_ylabel("P.1 water level (m)", color=INK, fontsize=11.5)
    ax.legend(loc="upper left", frameon=False, fontsize=10.5)
    _style_axes(ax)

    axp.plot(forecasts.index, forecasts["p_flood"], color=AMBER, linewidth=1.8)
    axp.fill_between(forecasts.index, 0, forecasts["p_flood"], color=AMBER, alpha=0.28)
    axp.axhline(0.5, color=INK, linewidth=0.9, linestyle=":", alpha=0.6)
    axp.set_ylim(-0.02, 1.1)
    axp.set_ylabel(f"P(flooding within {HORIZON} h)", color=INK, fontsize=11.5)
    _style_axes(axp)

    if first_alert is not None:
        lead_h = (
            None
            if flood_start is None
            else int((flood_start - first_alert).total_seconds() // 3600)
        )
        lead_txt = (
            ""
            if lead_h is None
            else (
                f"\n({lead_h} h before flooding began)"
                if lead_h >= 0
                else f"\n({-lead_h} h after flooding began)"
            )
        )
        if detail and flood_start is not None:
            for a in (ax, axp):
                a.axvline(first_alert, color=AMBER, linewidth=1.4, alpha=0.85)
                a.axvline(flood_start, color=BLUE, linewidth=1.4, alpha=0.85)
            # Anchor labels away from each other in chronological order so a
            # late alert (alert AFTER crossing) cannot overprint the labels.
            events = sorted(
                [
                    (first_alert, "model alert", AMBER),
                    (flood_start, "flooding begins", BLUE),
                ]
            )
            for (ts, label, color), (offset, align) in zip(
                events, ((-8, "right"), (8, "left"))
            ):
                ax.annotate(
                    f"{label}\n{ts:%d %b %H:%M}",
                    xy=(ts, observed.min()),
                    xytext=(offset, 18),
                    textcoords="offset points",
                    ha=align,
                    color=color,
                    fontsize=11,
                    fontweight="bold",
                )
            mid_y = observed.min() + (observed.max() - observed.min()) * 0.28
            ax.annotate(
                "",
                xy=(flood_start, mid_y),
                xytext=(first_alert, mid_y),
                arrowprops=dict(arrowstyle="<->", color=INK, lw=1.3),
            )
            arrow_label = (
                f"{lead_h} h warning" if lead_h >= 0 else f"alert {-lead_h} h late"
            )
            ax.annotate(
                arrow_label,
                xy=(first_alert + (flood_start - first_alert) / 2, mid_y),
                xytext=(0, 8),
                textcoords="offset points",
                ha="center",
                color=INK,
                fontsize=11.5,
                fontweight="bold",
            )
        else:
            axp.annotate(
                f"first alert · {first_alert:%d %b %H:%M}{lead_txt}",
                xy=(first_alert, 0.62),
                xytext=(10, 0),
                textcoords="offset points",
                color=RED,
                fontsize=10.5,
                bbox=dict(facecolor="white", alpha=0.75, edgecolor="none"),
            )

    locator = mdates.DayLocator(interval=1 if detail else 3)
    axp.xaxis.set_major_locator(locator)
    axp.xaxis.set_major_formatter(mdates.DateFormatter("%d %b"))
    fig.suptitle(
        f"{title}\n{subtitle}", x=0.07, y=0.985, ha="left", fontsize=15, color=INK
    )
    fig.subplots_adjust(top=0.885, left=0.07, right=0.97, bottom=0.07)
    fig.savefig(out_path, dpi=110)
    plt.close(fig)
    print(f"wrote {out_path}")


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--db-url", default=None)
    parser.add_argument("--out-dir", default=os.path.join("docs", "img"))
    parser.add_argument(
        "--dam",
        action="store_true",
        help="ablation: include Mae Ngat reservoir features "
        "(2026-08 result: costs 1-3 h of alert lead)",
    )
    parser.add_argument(
        "--no-hii-fill",
        action="store_true",
        help="ablation: load without the HII gap-fill merge",
    )
    args = parser.parse_args(argv)

    df = data.load_measurements(db_url=args.db_url, hii_fill=not args.no_hii_fill)
    if df.empty:
        print("no measurement data available", file=sys.stderr)
        return 1
    os.makedirs(args.out_dir, exist_ok=True)

    # --- October 2024 record flood: trained only on data before 1 Sep 2024 ---
    X, reg, clf = fit_backtest_model(df, "2024-08-31", use_dam=args.dam)
    obs, fc, flood_start, first_alert = event_series(
        df, X, reg, clf, "2024-09-10", "2024-10-14 23:00"
    )
    peak = float(obs.max())
    render(
        obs,
        fc,
        flood_start,
        first_alert,
        os.path.join(args.out_dir, "backtest-2024-p1.png"),
        title="October 2024 flood: what the model saw coming",
        subtitle="P.1 Nawarat Bridge — model trained only on data before 1 Sep 2024",
        show_stage7=True,
        peak_note=f"record peak {peak:.2f} m",
    )

    obs_d, fc_d, flood_d, alert_d = event_series(
        df, X, reg, clf, "2024-09-21 18:00", "2024-09-28 06:00"
    )
    lead_h = None
    if alert_d is not None and flood_d is not None:
        lead_h = int((flood_d - alert_d).total_seconds() // 3600)
    render(
        obs_d,
        fc_d,
        flood_d,
        alert_d,
        os.path.join(args.out_dir, "backtest-2024-p1-detail.png"),
        title="Detection in detail: 22–28 September 2024, hour by hour",
        subtitle=(
            f"the model alerts {lead_h} h before the river crosses the flooding line"
            if lead_h is not None and lead_h > 0
            else "model alert vs the river crossing the flooding line"
        ),
        detail=True,
    )

    # --- September 2025 flood: the deployed configuration (trained <= 2024) ---
    X25, reg25, clf25 = fit_backtest_model(df, "2024-12-31", use_dam=args.dam)
    obs25, fc25, flood25, alert25 = event_series(
        df, X25, reg25, clf25, "2025-09-22", "2025-10-02 12:00"
    )
    pred_at_alert = (
        float(fc25.loc[alert25:, "pred_max"].iloc[:24].max())
        if alert25 is not None
        else None
    )
    note = f"peak {float(obs25.max()):.2f} m" + (
        f" (predicted {pred_at_alert:.2f} m)" if pred_at_alert is not None else ""
    )
    render(
        obs25,
        fc25,
        flood25,
        alert25,
        os.path.join(args.out_dir, "backtest-2025-p1.png"),
        title="The September 2025 flood — as forecast by the deployed configuration",
        subtitle="model trained only on data through 2024; this event was never seen in training",
        detail=True,
        peak_note=note,
    )

    print(f"2024: flooding began {flood_start}, first alert {first_alert}")
    print(f"2025: flooding began {flood25}, first alert {alert25}")

    # Acceptance gate: the flagship 2024 event must keep a >= 12 h warning
    if first_alert is None or flood_start is None:
        print("FAIL: 2024 event alert or crossing not found", file=sys.stderr)
        return 1
    lead = (flood_start - first_alert).total_seconds() / 3600
    if lead < 12:
        print(f"FAIL: 2024 first-alert lead {lead:.0f} h < 12 h", file=sys.stderr)
        return 1
    print(f"PASS: 2024 first-alert lead {lead:.0f} h")
    return 0


if __name__ == "__main__":
    sys.exit(main())