| """Metrics and report plots for Tormenta models.""" |
| import matplotlib.pyplot as plt |
| import numpy as np |
| import pandas as pd |
| from sklearn.inspection import permutation_importance |
| from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score |
|
|
| from src import config |
|
|
|
|
| def metrics(y, pred) -> dict: |
| """Overall RMSE/MAE/R² plus RMSE/MAE on storm hours (Kp ≥ 5-).""" |
| y, pred = np.asarray(y), np.asarray(pred) |
| storm = y >= config.STORM_THRESHOLD |
| return { |
| "rmse": mean_squared_error(y, pred) ** 0.5, |
| "mae": mean_absolute_error(y, pred), |
| "r2": r2_score(y, pred), |
| "storm_rmse": mean_squared_error(y[storm], pred[storm]) ** 0.5 if storm.any() else np.nan, |
| "storm_mae": mean_absolute_error(y[storm], pred[storm]) if storm.any() else np.nan, |
| "n": len(y), |
| } |
|
|
|
|
| def metrics_table(y, preds: dict) -> pd.DataFrame: |
| return pd.DataFrame({name: metrics(y, p) for name, p in preds.items()}).T |
|
|
|
|
| def error_by_level(y, preds: dict) -> pd.DataFrame: |
| """MAE and mean bias (pred − true) per storm level.""" |
| level = pd.cut(pd.Series(np.asarray(y)), bins=config.KP_BINS, labels=config.KP_BIN_LABELS, right=False) |
| rows = {} |
| for name, p in preds.items(): |
| err = pd.Series(np.asarray(p) - np.asarray(y)) |
| rows[(name, "MAE")] = err.abs().groupby(level, observed=False).mean() |
| rows[(name, "bias")] = err.groupby(level, observed=False).mean() |
| return pd.DataFrame(rows).T |
|
|
|
|
| def rmse_by_block_hour(y, preds: dict) -> pd.DataFrame: |
| """RMSE by position inside the 3-h Kp block (0 = first hour), all hours and storm hours.""" |
| y = pd.Series(y) |
| pos = y.index.hour % config.KP_BLOCK_HOURS |
| storm = (y >= config.STORM_THRESHOLD).to_numpy() |
| rows = {} |
| for name, p in preds.items(): |
| sq = pd.Series((np.asarray(p) - y.to_numpy()) ** 2, index=y.index) |
| rows[(name, "all")] = sq.groupby(pos).mean() ** 0.5 |
| rows[(name, "storm")] = sq[storm].groupby(pos[storm]).mean() ** 0.5 |
| out = pd.DataFrame(rows).T |
| out.columns = [f"hour {c + 1} of block" for c in out.columns] |
| return out |
|
|
|
|
| def attach_block_kp(pred: pd.DataFrame, kp_blocks: pd.Series, name: str) -> pd.DataFrame: |
| """Add the Kp of the 3-h block containing each hour (e.g. NOAA estimated or GFZ Kp).""" |
| block = pred.index.floor(f"{config.KP_BLOCK_HOURS}h") |
| return pred.assign(**{name: kp_blocks.reindex(block).to_numpy()}) |
|
|
|
|
| def _save(fig, name): |
| config.FIGURES_DIR.mkdir(parents=True, exist_ok=True) |
| fig.savefig(config.FIGURES_DIR / f"{name}.png", bbox_inches="tight") |
|
|
|
|
| def plot_model_comparison(table: pd.DataFrame, title: str, name: str): |
| fig, axes = plt.subplots(1, 2, figsize=(13, 4)) |
| table["rmse"].plot.barh(ax=axes[0], color="#4C72B0") |
| axes[0].set(title=f"{title}: RMSE (all hours)", xlabel="RMSE (Kp units)") |
| table["storm_rmse"].plot.barh(ax=axes[1], color="#C44E52") |
| axes[1].set(title=f"{title}: RMSE on storm hours (Kp ≥ 5-)", xlabel="RMSE (Kp units)") |
| for ax in axes: |
| ax.invert_yaxis() |
| fig.tight_layout() |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_overfitting(train_rmse: pd.Series, cv_rmse: pd.Series, val_rmse: pd.Series, name: str): |
| df = pd.DataFrame({"train": train_rmse, "CV (train years)": cv_rmse, "validation": val_rmse}) |
| fig, ax = plt.subplots(figsize=(9, 4)) |
| df.plot.bar(ax=ax, rot=0) |
| ax.set(title="Train vs CV vs validation RMSE (gap = overfitting)", ylabel="RMSE") |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_pred_vs_actual(y, pred, title: str, name: str): |
| y, pred = np.asarray(y), np.asarray(pred) |
| fig, axes = plt.subplots(1, 2, figsize=(13, 4.5)) |
| hb = axes[0].hexbin(y, pred, gridsize=40, bins="log", cmap="viridis", extent=(0, 9, 0, 9)) |
| axes[0].plot([0, 9], [0, 9], "r--", lw=1) |
| axes[0].set(title=f"{title}: predicted vs observed", xlabel="observed Kp", ylabel="predicted Kp") |
| fig.colorbar(hb, ax=axes[0], label="hours (log)") |
| resid = pred - y |
| axes[1].hexbin(y, resid, gridsize=40, bins="log", cmap="magma") |
| axes[1].axhline(0, color="w", lw=1) |
| axes[1].set(title="Residual (pred − obs) vs observed", xlabel="observed Kp", ylabel="residual") |
| fig.tight_layout() |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_error_by_level(level_table: pd.DataFrame, name: str): |
| mae = level_table.xs("MAE", level=1) |
| bias = level_table.xs("bias", level=1) |
| fig, axes = plt.subplots(1, 2, figsize=(14, 4)) |
| mae.T.plot.bar(ax=axes[0], rot=0) |
| axes[0].set(title="MAE by storm level (test)", xlabel="observed Kp level", ylabel="MAE") |
| bias.T.plot.bar(ax=axes[1], rot=0, legend=False) |
| axes[1].axhline(0, color="k", lw=1) |
| axes[1].set(title="Mean bias (pred − obs) by level: negative = under-prediction", xlabel="observed Kp level") |
| fig.tight_layout() |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_permutation_importance(model, X, y, name: str, top: int = 20, n_rows: int = 20000): |
| """Permutation importance on a strided sample (model-agnostic, works for HistGB).""" |
| step = max(1, len(X) // n_rows) |
| res = permutation_importance(model, X.iloc[::step], np.asarray(y)[::step], n_repeats=3, |
| scoring="neg_root_mean_squared_error", |
| random_state=config.RANDOM_STATE, n_jobs=1) |
| imp = pd.Series(res.importances_mean, index=X.columns).sort_values().tail(top) |
| fig, ax = plt.subplots(figsize=(8, 6)) |
| imp.plot.barh(ax=ax, color="#55A868") |
| ax.set(title=f"Permutation importance (top {top}): RMSE increase when shuffled", xlabel="ΔRMSE") |
| fig.tight_layout() |
| _save(fig, name) |
| return fig, imp.sort_values(ascending=False) |
|
|
|
|
| def plot_diagnostic_curves(learning: pd.DataFrame, complexity: pd.DataFrame, name: str): |
| """Learning curve (more data) and complexity curve (bigger trees): train vs validation RMSE.""" |
| fig, axes = plt.subplots(1, 2, figsize=(13, 4)) |
| for ax, df, xlabel, title in [ |
| (axes[0], learning, "training hours (log scale)", "Learning curve: more data"), |
| (axes[1], complexity, "max_leaf_nodes per tree (log scale)", "Complexity curve: bigger trees")]: |
| ax.plot(df.index, df["train_rmse"], "o-", label="train") |
| ax.plot(df.index, df["val_rmse"], "s-", label="validation (2017–2019)") |
| ax.set(xscale="log", xlabel=xlabel, ylabel="RMSE (Kp units)", title=title) |
| ax.legend() |
| fig.tight_layout() |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_live(live: pd.DataFrame, name: str | None = None): |
| """Live nowcast with interval vs GFZ / NOAA Kp; hours ahead of the clock are shaded.""" |
| fig, ax = plt.subplots(figsize=(12, 4)) |
| ax.fill_between(live.index, live["kp_lo"], live["kp_hi"], alpha=0.25, label="80% interval") |
| ax.plot(live.index, live["kp_pred"], label="Tormenta nowcast") |
| for col, color, label in [("kp_gfz", "crimson", "GFZ quick-look Kp"), ("kp_swpc", "darkorange", "NOAA estimated Kp")]: |
| if col in live: |
| ax.step(live.index, live[col], where="post", color=color, label=label) |
| if live["ahead"].any(): |
| ax.axvspan(live.index[live["ahead"]].min(), live.index.max(), color="grey", alpha=0.15, |
| label="ahead of now (propagated data)") |
| ax.axhline(config.STORM_THRESHOLD, color="grey", ls="--", lw=1) |
| ax.set(title="Live nowcast vs observed Kp", ylabel="Kp", ylim=(0, 9.3)) |
| ax.legend(loc="upper left") |
| if name: |
| _save(fig, name) |
| return fig |
|
|
|
|
| def plot_storm_case(y, pred, lo, hi, start: str, end: str, title: str, name: str): |
| s = pd.DataFrame({"obs": y, "pred": pred, "lo": lo, "hi": hi}).loc[start:end] |
| fig, ax = plt.subplots(figsize=(12, 4)) |
| ax.fill_between(s.index, s["lo"], s["hi"], alpha=0.25, label="80% interval") |
| ax.plot(s.index, s["pred"], label="predicted") |
| ax.step(s.index, s["obs"], where="post", color="crimson", label="observed (GFZ)") |
| ax.axhline(config.STORM_THRESHOLD, color="grey", ls="--", lw=1) |
| ax.set(title=title, ylabel="Kp", ylim=(0, 9.3)) |
| ax.legend(loc="upper left") |
| _save(fig, name) |
| return fig |
|
|
