monsoon-rl / evaluate_checkpoint_real.py
DHDRL's picture
Update evaluate_checkpoint_real.py
eea3056 verified
Raw
History Blame
42 kB
#!/usr/bin/env python3
"""
evaluate_checkpoint_real.py
===========================
Real-trajectory evaluation for agent decisions vs L1 impact labels.
Closes the gap between:
- scorer product-vs-L1 metrics (backtest_indonesia.py), and
- agent eval that previously used only synthetic _zone_event_flags.
Build order (contracts → code):
entities: EvalDayRecord, EvalResult
modes: SCORER_ORACLE | CHECKPOINT | ALWAYS | NEVER
GT: days inside an L1 event span → gt_source="l1"; days outside every
span for that zone → gt_source="unlabeled" (excluded from P/R/F1).
Sparse catalogs must not manufacture TN/FP from silence.
alert definition: product gate (WARNING+ OR drought≥0.35 OR flood≥0.25)
metrics: product-vs-L1 P/R/F1 on L1 event days only; unlabeled_alert_rate
Does NOT require a trained checkpoint for SCORER_ORACLE / ALWAYS / NEVER —
those baselines prove the harness before GPU time is spent.
Usage examples
--------------
# Scorer oracle on historical cache (no checkpoint)
python evaluate_checkpoint_real.py \\
--pkl historical_continuous_indonesia_v1.pkl \\
--impact-labels impact_labels_java_v1.json \\
--zones karawang_rice,indramayu_rice \\
--start 2023-07-01 --end 2023-11-30 \\
--mode scorer_oracle
# Trained agent (must match env basin_context dim=8)
python evaluate_checkpoint_real.py \\
--pkl historical_continuous_indonesia_v1.pkl \\
--impact-labels impact_labels_java_v1.json \\
--zones karawang_rice \\
--start 2023-07-01 --end 2023-11-30 \\
--mode checkpoint --checkpoint path/to/final.zip \\
--n-zones 1 --max-steps 4
"""
from __future__ import annotations
import argparse
import json
import logging
import pickle
import sys
from collections import Counter
from dataclasses import asdict, dataclass, field
from datetime import date, datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple
import numpy as np
import zone_observation as _zo
assert _zo.SCHEMA_VERSION == 3, (
f"evaluate_checkpoint_real: zone_observation schema mismatch "
f"(expected 3, got {_zo.SCHEMA_VERSION})"
)
from zone_observation import (
BasinContext,
EpisodeContext,
ForecastConfig,
ForecastResult,
RiskScore,
ZoneObs,
)
from crop_risk_scorer import compute_risk_score
from product_alert_service import (
DEFAULT_PRODUCT_GATE,
ProductGateConfig,
is_product_actionable,
)
from weather_forecast_env import make_weather_env
logger = logging.getLogger(__name__)
_PATH_CHECK_N = 0
_PATH_CHECK_DIVERGE = 0
# ---------------------------------------------------------------------------
# Entities
# ---------------------------------------------------------------------------
BELIEF_RAISE_EPS = 0.02
def _path_check_reset() -> None:
global _PATH_CHECK_N, _PATH_CHECK_DIVERGE
_PATH_CHECK_N = 0
_PATH_CHECK_DIVERGE = 0
def _path_check_report() -> None:
if _PATH_CHECK_N <= 0:
return
n_ok = _PATH_CHECK_N - _PATH_CHECK_DIVERGE
print(
f"PATH_CHECK summary: n={_PATH_CHECK_N} ok={n_ok} "
f"argmax_miss_or_hit={_PATH_CHECK_DIVERGE} "
f"rate={_PATH_CHECK_DIVERGE / _PATH_CHECK_N:.3f} "
f"(pack-level argmax [Option A, decision-locked 2026-08-15] is "
f"authoritative for product_actionable; this counts days where "
f"OR-across-zones [Option B] would have differed -- i.e. some "
f"zone's own drought/flood channel independently cleared the gate "
f"while the argmax-selected zone did not. Not a display bug.)",
flush=True,
)
@dataclass
class EvalDayRecord:
date: str
zone_id: str
mode: str
product_alert: bool
elevated: bool
alert_level: str
drought_risk: float
flood_risk: float
event_drought: bool
event_flood: bool
gt_source: str # "l1" | "unlabeled" | "none"
product_emit_code: str = "N/A" # not emitted to bus in this harness
ep_len: int = 0
basin_dim: int = 0
believed_p: float = 0.0
initial_belief: float = 0.0
belief_delta: float = 0.0
belief_raised: bool = False # believed_p > initial_belief + BELIEF_RAISE_EPS
@dataclass
class EvalMetrics:
n_days: int = 0
tp: int = 0
fp: int = 0
fn: int = 0
tn: int = 0
n_l1: int = 0 # days inside an L1 event span (positive labels only in v1)
n_product: int = 0
n_unlabeled: int = 0
n_unlabeled_product: int = 0
n_unlabeled_belief_raised: int = 0
belief_tp: int = 0
belief_fn: int = 0
belief_fp: int = 0
belief_tn: int = 0
n_belief_raised: int = 0
belief_deltas: List[float] = field(default_factory=list)
belief_deltas_l1: List[float] = field(default_factory=list)
belief_deltas_unlabeled: List[float] = field(default_factory=list)
@property
def precision(self) -> Optional[float]:
if self.fp == 0 and self.tn == 0 and (self.tp + self.fn) > 0:
return None
d = self.tp + self.fp
return self.tp / d if d else None
@property
def recall(self) -> Optional[float]:
d = self.tp + self.fn
return self.tp / d if d else None
@property
def f1(self) -> Optional[float]:
p, r = self.precision, self.recall
if p is None or r is None or (p + r) == 0:
return None
return 2 * p * r / (p + r)
@property
def belief_precision(self) -> Optional[float]:
return None
@property
def belief_recall(self) -> Optional[float]:
d = self.belief_tp + self.belief_fn
return self.belief_tp / d if d else None
@property
def belief_f1(self) -> Optional[float]:
return None
def belief_delta_stats(self) -> Dict[str, float]:
xs = self.belief_deltas
if not xs:
return {"n": 0, "mean": 0.0, "std": 0.0, "p10": 0.0, "p50": 0.0, "p90": 0.0,
"frac_pos": 0.0, "frac_neg": 0.0, "frac_near0": 0.0}
arr = sorted(xs)
n = len(arr)
mean = sum(arr) / n
var = sum((x - mean) ** 2 for x in arr) / max(n, 1)
def pct(p: float) -> float:
i = min(n - 1, max(0, int(round(p * (n - 1)))))
return arr[i]
near = sum(1 for x in arr if abs(x) < BELIEF_RAISE_EPS)
return {
"n": n,
"mean": mean,
"std": var ** 0.5,
"p10": pct(0.10),
"p50": pct(0.50),
"p90": pct(0.90),
"frac_pos": sum(1 for x in arr if x > BELIEF_RAISE_EPS) / n,
"frac_neg": sum(1 for x in arr if x < -BELIEF_RAISE_EPS) / n,
"frac_near0": near / n,
}
@property
def unlabeled_alert_rate(self) -> Optional[float]:
if self.n_unlabeled <= 0:
return None
return self.n_unlabeled_product / self.n_unlabeled
@property
def belief_unlabeled_raise_rate(self) -> Optional[float]:
if self.n_unlabeled <= 0:
return None
return self.n_unlabeled_belief_raised / self.n_unlabeled
def to_dict(self) -> Dict[str, Any]:
def _f(x: Optional[float]) -> Optional[float]:
return None if x is None else round(x, 6)
bd = self.belief_delta_stats()
def _subset_stats(xs: List[float]) -> Dict[str, Any]:
saved = self.belief_deltas
self.belief_deltas = xs
out = self.belief_delta_stats()
self.belief_deltas = saved
return {k: (round(v, 6) if isinstance(v, float) else v) for k, v in out.items()}
return {
"n_days": self.n_days,
"n_l1_event_days": self.n_l1,
"n_unlabeled": self.n_unlabeled,
"n_unlabeled_product": self.n_unlabeled_product,
"n_unlabeled_belief_raised": self.n_unlabeled_belief_raised,
"unlabeled_alert_rate": _f(self.unlabeled_alert_rate),
"belief_unlabeled_raise_rate": _f(self.belief_unlabeled_raise_rate),
"n_product_alerts": self.n_product,
"n_belief_raised": self.n_belief_raised,
"tp": self.tp,
"fp": self.fp,
"fn": self.fn,
"tn": self.tn,
"precision": _f(self.precision),
"recall": _f(self.recall),
"f1": _f(self.f1),
"note_metrics": (
"Recall uses only gt_source=l1 (days inside an event span). "
"Unlabeled days are excluded from the confusion matrix and "
"reported via unlabeled_alert_rate / belief_unlabeled_raise_rate. "
"Precision and F1 are null under positive-only L1 (no confirmed "
"negatives → FP/TN stay 0 by construction; a printed 1.0 would "
"be an artifact)."
),
"belief_tp": self.belief_tp,
"belief_fn": self.belief_fn,
"belief_fp": self.belief_fp,
"belief_tn": self.belief_tn,
"belief_precision": _f(self.belief_precision),
"belief_recall": _f(self.belief_recall),
"belief_f1": _f(self.belief_f1),
"note_belief_metrics": (
"belief_recall = belief_tp/(belief_tp+belief_fn) on L1 event "
"days is the legitimate policy-sensitive number. "
"belief_precision and belief_f1 are always null (no FP/TN path)."
),
"belief_delta": {k: (round(v, 6) if isinstance(v, float) else v)
for k, v in bd.items()},
"belief_delta_l1": _subset_stats(self.belief_deltas_l1),
"belief_delta_unlabeled": _subset_stats(self.belief_deltas_unlabeled),
}
# ---------------------------------------------------------------------------
# Historical cache → EpisodeContext
# ---------------------------------------------------------------------------
def _parse_day(s: str) -> date:
return date.fromisoformat(s[:10])
def load_historical_points(
pkl_path: Path,
zone_ids: Sequence[str],
start: date,
end: date,
) -> List[Dict[str, Any]]:
with open(pkl_path, "rb") as f:
cache = pickle.load(f)
trajs = cache.get("trajectories") or []
zone_set = set(zone_ids)
out: List[Dict[str, Any]] = []
for traj in trajs:
meta = traj.get("meta") or {}
zid = meta.get("zone_id")
if zid not in zone_set:
continue
for pt in traj.get("trajectory") or []:
vt = _parse_day(str(pt.get("valid_time", "")))
if vt < start or vt > end:
continue
if pt.get("zone_id") and pt["zone_id"] not in zone_set:
continue
out.append(pt)
out.sort(key=lambda p: (str(p.get("zone_id")), str(p.get("valid_time"))))
return out
def point_to_episode(
pt: Dict[str, Any],
cfg: ForecastConfig,
) -> EpisodeContext:
obs = ZoneObs.from_dict(dict(pt["obs"]))
fc = ForecastResult.from_dict(dict(pt["forecast"]))
basin = None
if pt.get("basin_context"):
try:
basin = BasinContext.from_dict(dict(pt["basin_context"]))
except Exception as e:
logger.warning("basin_context deserialize failed: %s", e)
basin = None
return EpisodeContext(
obs=obs,
forecast=fc,
config=cfg,
zone_ids=[obs.zone_id],
basin_context=basin,
)
def points_to_multi_zone_episode(
pts: Sequence[Dict[str, Any]],
cfg: ForecastConfig,
zone_order: Sequence[str],
) -> EpisodeContext:
by_z = {}
for pt in pts:
obs = ZoneObs.from_dict(dict(pt["obs"]))
by_z[obs.zone_id] = pt
missing = [z for z in zone_order if z not in by_z]
if missing:
raise ValueError(f"points_to_multi_zone_episode missing zones: {missing}")
zone_obs: List[ZoneObs] = []
zone_fc: List[ForecastResult] = []
basin = None
for zid in zone_order:
pt = by_z[zid]
zo = ZoneObs.from_dict(dict(pt["obs"]))
zf = ForecastResult.from_dict(dict(pt["forecast"]))
zone_obs.append(zo)
zone_fc.append(zf)
if basin is None and pt.get("basin_context"):
try:
basin = BasinContext.from_dict(dict(pt["basin_context"]))
except Exception as e:
logger.warning("basin_context deserialize failed: %s", e)
return EpisodeContext(
obs=zone_obs[0],
forecast=zone_fc[0],
config=cfg,
zone_ids=list(zone_order),
basin_context=basin,
zone_obs=zone_obs,
zone_forecasts=zone_fc,
)
def group_points_by_date(
points: Sequence[Dict[str, Any]],
) -> Dict[str, List[Dict[str, Any]]]:
out: Dict[str, List[Dict[str, Any]]] = {}
for pt in points:
obs = pt.get("obs") or {}
vt = str(obs.get("valid_time") or pt.get("valid_time") or "")[:10]
if not vt:
continue
out.setdefault(vt, []).append(pt)
return out
# ---------------------------------------------------------------------------
# Decision policies
# ---------------------------------------------------------------------------
def _multi_zone_risk_score(ctx: EpisodeContext) -> RiskScore:
"""Risk score for a (possibly multi-zone) EpisodeContext, matching
WeatherForecastEnv._compute_multi_zone_risk(): the zone with the
highest supply_shortfall_prob across ALL zones in the episode, not
just ctx.obs/ctx.forecast (which points_to_multi_zone_episode always
sets to zone_ids[0] alone).
ctx.obs/ctx.forecast is only ever a single zone's data by construction
(see points_to_multi_zone_episode / point_to_episode). For single-zone
episodes this is identical to compute_risk_score(ctx.obs, ctx.forecast,
ctx.config); for multi-zone episodes it must aggregate across
ctx.resolved_zone_obs()/ctx.resolved_zone_forecasts() the same way the
env does, or "rs" silently stops representing the same decision as the
env's actual per-step product_actionable flag.
This is Option A (argmax by supply_shortfall_prob) -- the documented,
decision-locked (2026-08-15) rule for the authoritative pack-level
product decision. Used here only to make displayed per-day fields
(alert_level/drought_risk/flood_risk) consistent with whichever zone
actually drove that decision. Do NOT use this for the PATH_CHECK
divergence comparison -- see _or_across_zones_product below for why.
"""
zone_obs = ctx.resolved_zone_obs()
zone_fc = ctx.resolved_zone_forecasts()
best: Optional[RiskScore] = None
for zo, zf in zip(zone_obs, zone_fc):
score = compute_risk_score(zo, zf, ctx.config)
if best is None or score.supply_shortfall_prob > best.supply_shortfall_prob:
best = score
return best if best is not None else compute_risk_score(
ctx.obs, ctx.forecast, ctx.config
)
def _or_across_zones_product(ctx: EpisodeContext, gate: ProductGateConfig) -> bool:
"""Option B from the documented multi-zone product selection decision:
any(is_product_actionable(zone) for zone in zones), independent of
which zone owns the highest supply_shortfall_prob.
Proven property: argmax_product (Option A / loop_product) implies
or_product always, so the only possible divergence is
loop_product=False, or_product=True -- an "argmax miss" where some
zone's own drought/flood channel independently crosses the gate while
the argmax-selected zone (chosen by a supply_shortfall_prob blend that
also includes harvest pressure) does not. This is the real, still-open
signature PATH_CHECK exists to surface (decision recorded 2026-08-15:
keep argmax as the product rule; the divergence itself stays worth
tracking, not silently eliminated).
"""
zone_obs = ctx.resolved_zone_obs()
zone_fc = ctx.resolved_zone_forecasts()
for zo, zf in zip(zone_obs, zone_fc):
score = compute_risk_score(zo, zf, ctx.config)
if is_product_actionable(score, gate):
return True
return False
def decide_scorer_oracle(
obs: ZoneObs,
fc: ForecastResult,
cfg: ForecastConfig,
gate: ProductGateConfig,
) -> Tuple[bool, bool, RiskScore, int, float, float]:
rs = compute_risk_score(obs, fc, cfg)
product = is_product_actionable(rs, gate)
elevated = rs.is_elevated()
return product, elevated, rs, 0, 0.0, 0.0
def decide_always() -> Tuple[bool, bool, None, int, float, float]:
return True, True, None, 0, 1.0, 1.0
def decide_never() -> Tuple[bool, bool, None, int, float, float]:
return False, False, None, 0, 0.0, 0.0
def _initial_belief_from_info(info: Dict[str, Any], cfg: ForecastConfig) -> float:
if "believed_p" in info and info["believed_p"] is not None:
return float(info["believed_p"])
zb = info.get("zone_belief")
if zb is not None:
try:
import numpy as _np
arr = _np.asarray(zb, dtype=float).ravel()
n = int(info.get("n_zones") or cfg.n_zones or 0)
if arr.size:
if n > 0:
return float(_np.max(arr[: min(n, arr.size)]))
return float(_np.max(arr))
except Exception:
pass
if "mean_belief" in info and info["mean_belief"] is not None:
return float(info["mean_belief"])
return float(cfg.prior_belief)
def decide_checkpoint(
model: Any,
env: Any,
ctx: EpisodeContext,
gate: ProductGateConfig,
) -> Tuple[bool, bool, Optional[RiskScore], int, float, float]:
global _PATH_CHECK_N, _PATH_CHECK_DIVERGE
obs, info = env.reset(options={"context": ctx})
initial_belief = _initial_belief_from_info(info, ctx.config)
done = False
ep_len = 0
product = False
elevated = False
believed_p = initial_belief
while not done:
masks = env.action_masks() if hasattr(env, "action_masks") else None
if masks is None and hasattr(env, "env") and hasattr(env.env, "action_masks"):
masks = env.env.action_masks()
action, _ = model.predict(obs, action_masks=masks, deterministic=True)
obs, reward, terminated, truncated, info = env.step(int(action))
ep_len += 1
done = bool(terminated or truncated)
if "product_actionable" in info:
product = bool(info["product_actionable"])
if "elevated" in info:
elevated = bool(info["elevated"])
if "believed_p" in info:
believed_p = float(info["believed_p"])
elif "mean_belief" in info:
believed_p = float(info["mean_belief"])
loop_product, loop_elevated = product, elevated
rs = None
try:
rs = _multi_zone_risk_score(ctx)
except Exception:
pass
try:
or_product = _or_across_zones_product(ctx, gate)
_PATH_CHECK_N += 1
if or_product != loop_product:
_PATH_CHECK_DIVERGE += 1
except Exception:
pass
return loop_product, loop_elevated, rs, ep_len, believed_p, initial_belief
def decide_zero_inspect(
env: Any,
ctx: EpisodeContext,
gate: ProductGateConfig,
) -> Tuple[bool, bool, Optional[RiskScore], int, float, float]:
obs, info = env.reset(options={"context": ctx})
initial_belief = _initial_belief_from_info(info, ctx.config)
base = env.env if hasattr(env, "env") else env
terminate_action = int(getattr(base, "terminate_action", ctx.config.n_zones))
obs, reward, terminated, truncated, info = env.step(terminate_action)
product = bool(info.get("product_actionable", False))
elevated = bool(info.get("elevated", False))
if "believed_p" in info and info["believed_p"] is not None:
believed_p = float(info["believed_p"])
else:
believed_p = _initial_belief_from_info(info, ctx.config)
rs = None
try:
rs = _multi_zone_risk_score(ctx)
except Exception:
pass
return product, elevated, rs, 1, believed_p, initial_belief
# ---------------------------------------------------------------------------
# Core eval loop
# ---------------------------------------------------------------------------
def evaluate_multi_zone_days(
points: Sequence[Dict[str, Any]],
*,
mode: str,
cfg: ForecastConfig,
gate: ProductGateConfig,
zone_order: Sequence[str],
impact_store: Any = None,
model: Any = None,
env: Any = None,
) -> Tuple[List[EvalDayRecord], EvalMetrics]:
if mode not in ("checkpoint", "zero_inspect"):
raise ValueError(
f"evaluate_multi_zone_days only supports checkpoint/zero_inspect "
f"(got {mode!r})"
)
_path_check_reset()
if env is None:
raise RuntimeError("evaluate_multi_zone_days requires env")
if mode == "checkpoint" and model is None:
raise RuntimeError("checkpoint mode requires model")
records: List[EvalDayRecord] = []
m = EvalMetrics()
by_day = group_points_by_date(points)
n_skip_incomplete = 0
zone_set = set(zone_order)
for day_s in sorted(by_day.keys()):
day_pts = []
present = set()
for pt in by_day[day_s]:
zo = ZoneObs.from_dict(dict(pt["obs"]))
if zo.zone_id in zone_set:
day_pts.append(pt)
present.add(zo.zone_id)
if not all(z in present for z in zone_order):
n_skip_incomplete += 1
continue
chosen: List[Dict[str, Any]] = []
for zid in zone_order:
for pt in day_pts:
if ZoneObs.from_dict(dict(pt["obs"])).zone_id == zid:
chosen.append(pt)
break
ctx = points_to_multi_zone_episode(chosen, cfg, zone_order)
day = date.fromisoformat(day_s)
event_d = event_f = False
gt_source = "none"
if impact_store is not None:
any_event = False
for zid in zone_order:
try:
ld, lf = impact_store.labels_for_day(zid, day)
event_d = event_d or bool(ld)
event_f = event_f or bool(lf)
if ld or lf:
any_event = True
except Exception as e:
logger.warning("L1 query failed %s %s: %s", zid, day, e)
gt_source = "l1" if any_event else "unlabeled"
if mode == "checkpoint":
product, elevated, rs, ep_len, believed_p, initial_belief = (
decide_checkpoint(model, env, ctx, gate)
)
else:
product, elevated, rs, ep_len, believed_p, initial_belief = (
decide_zero_inspect(env, ctx, gate)
)
alert_level = (
rs.alert_level.value
if rs is not None
else ("warning" if product else "none")
)
drought_risk = float(rs.drought_risk) if rs is not None else 0.0
flood_risk = float(rs.flood_risk) if rs is not None else 0.0
belief_delta = float(believed_p) - float(initial_belief)
belief_raised = bool(belief_delta > BELIEF_RAISE_EPS)
if gt_source == "l1":
if product:
m.tp += 1
else:
m.fn += 1
if belief_raised:
m.belief_tp += 1
else:
m.belief_fn += 1
m.n_l1 += 1
m.belief_deltas_l1.append(belief_delta)
elif gt_source == "unlabeled":
m.n_unlabeled += 1
if product:
m.n_unlabeled_product += 1
if belief_raised:
m.n_unlabeled_belief_raised += 1
m.belief_deltas_unlabeled.append(belief_delta)
m.n_days += 1
if product:
m.n_product += 1
if belief_raised:
m.n_belief_raised += 1
m.belief_deltas.append(belief_delta)
records.append(
EvalDayRecord(
date=day.isoformat(),
zone_id="+".join(zone_order),
mode=mode,
product_alert=product,
elevated=elevated,
alert_level=alert_level,
drought_risk=drought_risk,
flood_risk=flood_risk,
event_drought=event_d,
event_flood=event_f,
gt_source=gt_source,
ep_len=ep_len,
basin_dim=8,
believed_p=float(believed_p),
initial_belief=float(initial_belief),
belief_delta=belief_delta,
belief_raised=belief_raised,
)
)
if n_skip_incomplete:
print(
f"multi-zone: skipped {n_skip_incomplete} days missing full "
f"zone set {list(zone_order)}"
)
_path_check_report()
return records, m
def evaluate_points(
points: Sequence[Dict[str, Any]],
*,
mode: str,
cfg: ForecastConfig,
gate: ProductGateConfig,
impact_store: Any = None,
model: Any = None,
env: Any = None,
) -> Tuple[List[EvalDayRecord], EvalMetrics]:
_path_check_reset()
records: List[EvalDayRecord] = []
m = EvalMetrics()
for pt in points:
obs = ZoneObs.from_dict(dict(pt["obs"]))
fc = ForecastResult.from_dict(dict(pt["forecast"]))
vt = obs.valid_time
if vt.tzinfo is None:
vt = vt.replace(tzinfo=timezone.utc)
day = vt.date()
zid = obs.zone_id
event_d = event_f = False
gt_source = "none"
if impact_store is not None:
try:
ld, lf = impact_store.labels_for_day(zid, day)
event_d, event_f = bool(ld), bool(lf)
if event_d or event_f:
gt_source = "l1"
else:
gt_source = "unlabeled"
except Exception as e:
logger.warning("L1 query failed %s %s: %s", zid, day, e)
gt_source = "none"
believed_p = 0.0
initial_belief = 0.0
if mode == "scorer_oracle":
product, elevated, rs, ep_len, believed_p, initial_belief = (
decide_scorer_oracle(obs, fc, cfg, gate)
)
elif mode == "always":
product, elevated, rs, ep_len, believed_p, initial_belief = decide_always()
elif mode == "never":
product, elevated, rs, ep_len, believed_p, initial_belief = decide_never()
elif mode == "checkpoint":
if model is None or env is None:
raise RuntimeError("checkpoint mode requires --checkpoint and env")
ctx = point_to_episode(pt, cfg)
product, elevated, rs, ep_len, believed_p, initial_belief = (
decide_checkpoint(model, env, ctx, gate)
)
elif mode == "zero_inspect":
if env is None:
raise RuntimeError("zero_inspect mode requires env")
ctx = point_to_episode(pt, cfg)
product, elevated, rs, ep_len, believed_p, initial_belief = (
decide_zero_inspect(env, ctx, gate)
)
else:
raise ValueError(f"unknown mode: {mode}")
alert_level = (
rs.alert_level.value if rs is not None else ("warning" if product else "none")
)
drought_risk = float(rs.drought_risk) if rs is not None else 0.0
flood_risk = float(rs.flood_risk) if rs is not None else 0.0
belief_delta = float(believed_p) - float(initial_belief)
belief_raised = bool(belief_delta > BELIEF_RAISE_EPS)
if gt_source == "l1":
if product:
m.tp += 1
else:
m.fn += 1
if belief_raised:
m.belief_tp += 1
else:
m.belief_fn += 1
m.n_l1 += 1
m.belief_deltas_l1.append(belief_delta)
elif gt_source == "unlabeled":
m.n_unlabeled += 1
if product:
m.n_unlabeled_product += 1
if belief_raised:
m.n_unlabeled_belief_raised += 1
m.belief_deltas_unlabeled.append(belief_delta)
m.n_days += 1
if product:
m.n_product += 1
if belief_raised:
m.n_belief_raised += 1
m.belief_deltas.append(belief_delta)
records.append(
EvalDayRecord(
date=day.isoformat(),
zone_id=zid,
mode=mode,
product_alert=product,
elevated=elevated,
alert_level=alert_level,
drought_risk=drought_risk,
flood_risk=flood_risk,
event_drought=event_d,
event_flood=event_f,
gt_source=gt_source,
ep_len=ep_len,
basin_dim=8,
believed_p=float(believed_p),
initial_belief=float(initial_belief),
belief_delta=belief_delta,
belief_raised=belief_raised,
)
)
_path_check_report()
return records, m
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _print_metrics(mode: str, m: EvalMetrics) -> None:
def _f(x: Optional[float]) -> str:
return f"{x:.3f}" if x is not None else " - "
print(f"\n=== real eval mode={mode} ===")
print(
f"n={m.n_days} l1_event_days={m.n_l1} unlabeled={m.n_unlabeled} "
f"product_alerts={m.n_product} belief_raised={m.n_belief_raised}"
)
print(
"--- product vs L1 event spans only "
"(unlabeled excluded from confusion matrix) ---"
)
print(f"TP={m.tp} FP={m.fp} FN={m.fn} TN={m.tn}")
print(f"P={_f(m.precision)} R={_f(m.recall)} F1={_f(m.f1)}")
if m.n_unlabeled > 0:
print(
f"--- unlabeled (outside every L1 span; not in P/R/F1) ---"
)
print(
f"n_unlabeled={m.n_unlabeled} "
f"unlabeled_product={m.n_unlabeled_product} "
f"unlabeled_alert_rate={_f(m.unlabeled_alert_rate)}"
)
print(
f"unlabeled_belief_raised={m.n_unlabeled_belief_raised} "
f"belief_unlabeled_raise_rate={_f(m.belief_unlabeled_raise_rate)}"
)
print(
f"--- belief_raised on L1 event days "
f"(delta > {BELIEF_RAISE_EPS} vs episode initial) ---"
)
print(
f"belief_tp={m.belief_tp} belief_fn={m.belief_fn} "
f"(belief_fp/tn unused under positive-only L1)"
)
print(
f"belief_precision={_f(m.belief_precision)} "
f"belief_recall={_f(m.belief_recall)} "
f"belief_f1={_f(m.belief_f1)} "
f"[P/F1 null by construction; R is the real number]"
)
bd = m.belief_delta_stats()
if bd["n"] > 0:
print("--- belief_delta = terminal − initial (all days) ---")
print(
f" n={int(bd['n'])} mean={bd['mean']:+.4f} std={bd['std']:.4f} "
f"p10={bd['p10']:+.4f} p50={bd['p50']:+.4f} p90={bd['p90']:+.4f}"
)
print(
f" frac_pos(>{BELIEF_RAISE_EPS})={bd['frac_pos']:.3f} "
f"frac_neg(<-{BELIEF_RAISE_EPS})={bd['frac_neg']:.3f} "
f"frac_|delta|<{BELIEF_RAISE_EPS}={bd['frac_near0']:.3f}"
)
if m.belief_deltas_l1 or m.belief_deltas_unlabeled:
def _mean(xs: List[float]) -> str:
if not xs:
return " - "
return f"{sum(xs)/len(xs):+.4f}"
print(
f" mean Δ | L1 event {_mean(m.belief_deltas_l1)} "
f"n={len(m.belief_deltas_l1)}"
)
print(
f" mean Δ | unlabeled {_mean(m.belief_deltas_unlabeled)} "
f"n={len(m.belief_deltas_unlabeled)}"
)
def main(argv: Optional[Sequence[str]] = None) -> int:
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
p = argparse.ArgumentParser(description="Real-trajectory product eval (L1)")
p.add_argument("--pkl", required=True, help="historical_continuous_indonesia_v1.pkl")
p.add_argument("--impact-labels", default=None, help="impact_labels_java_v1.json")
p.add_argument("--zones", default="karawang_rice,indramayu_rice")
p.add_argument("--start", required=True, help="YYYY-MM-DD")
p.add_argument("--end", required=True, help="YYYY-MM-DD")
p.add_argument(
"--mode",
choices=("scorer_oracle", "checkpoint", "zero_inspect", "always", "never"),
default="scorer_oracle",
help="checkpoint=agent rollout; zero_inspect=terminate step 1 control; "
"scorer_oracle=deterministic product gate only",
)
p.add_argument("--checkpoint", default=None, help="MaskablePPO .zip (mode=checkpoint)")
p.add_argument("--n-zones", type=int, default=1)
p.add_argument(
"--max-steps",
type=int,
default=4,
help="Eval episode length cap (n_zones=1 only needs ~2; does not "
"need to match train max_steps).",
)
p.add_argument(
"--horizon-days",
type=int,
default=30,
help="Must match the checkpoint's forecast_precip width "
"(train_kaggle / ForecastConfig default is 30).",
)
p.add_argument("--device", default="cpu")
p.add_argument("--out", default=None, help="Write full JSON result")
args = p.parse_args(list(argv) if argv is not None else None)
zone_ids = [z.strip() for z in args.zones.split(",") if z.strip()]
start = _parse_day(args.start)
end = _parse_day(args.end)
pkl_path = Path(args.pkl)
if not pkl_path.is_file():
print(f"FILE_NOT_FOUND: {pkl_path}", file=sys.stderr)
return 2
points = load_historical_points(pkl_path, zone_ids, start, end)
print(f"loaded points: {len(points)} zones={zone_ids} {start}{end}")
if not points:
print("NO_POINTS in window/zones", file=sys.stderr)
return 3
impact_store = None
if args.impact_labels:
from impact_labels import load_impact_events
res = load_impact_events(args.impact_labels)
if not res.success:
print(f"L1 load failed: {res.outcome_code}", file=sys.stderr)
return 4
impact_store = res.data["store"]
print(f"L1: {res.outcome_code} events_loaded={res.data.get('events_loaded')}")
cfg = ForecastConfig(
n_zones=args.n_zones,
max_steps=args.max_steps,
soft_reset=True,
horizon_days=int(args.horizon_days),
)
gate = DEFAULT_PRODUCT_GATE
model = None
env = None
if args.mode in ("checkpoint", "zero_inspect"):
env = make_weather_env(cfg, use_nan_wrapper=True)
obs_space = env.observation_space
if hasattr(env, "env"):
obs_space = env.env.observation_space
bshape = obs_space["basin_context"].shape
precip_shape = obs_space["forecast_precip"].shape
print(f"env basin_context shape: {bshape}")
print(f"env forecast_precip shape: {precip_shape}")
print(f"env horizon_days={cfg.horizon_days} n_zones={cfg.n_zones}")
if args.mode == "checkpoint":
if not args.checkpoint:
print("checkpoint mode requires --checkpoint", file=sys.stderr)
return 5
try:
from sb3_contrib import MaskablePPO
except ImportError:
print("sb3_contrib not installed", file=sys.stderr)
return 6
model = MaskablePPO.load(args.checkpoint, device=args.device)
if int(np.prod(bshape)) != 8:
print(
"WARNING: env basin_context is not 8-dim; "
"checkpoint may be incompatible",
file=sys.stderr,
)
try:
pol_space = model.observation_space
pol_precip = pol_space["forecast_precip"].shape
if tuple(pol_precip) != tuple(precip_shape):
print(
f"SHAPE_MISMATCH: policy forecast_precip {pol_precip} "
f"!= env {precip_shape}. Re-run with "
f"--horizon-days matching training (usually 30).",
file=sys.stderr,
)
return 7
pol_zones = pol_space["zone_belief"].shape
env_zones = obs_space["zone_belief"].shape
if tuple(pol_zones) != tuple(env_zones):
print(
f"SHAPE_MISMATCH: policy zone_belief {pol_zones} "
f"!= env {env_zones}. Use --n-zones matching training "
f"(run_smoke=1, run_nz3=3).",
file=sys.stderr,
)
return 8
except Exception as e:
logger.warning("could not cross-check policy obs space: %s", e)
if (
args.mode in ("checkpoint", "zero_inspect")
and int(args.n_zones) > 1
):
if len(zone_ids) < int(args.n_zones):
print(
f"NEED_ZONES: --n-zones={args.n_zones} but only "
f"{len(zone_ids)} zones listed in --zones",
file=sys.stderr,
)
return 9
zone_order = zone_ids[: int(args.n_zones)]
print(f"multi-zone real eval zone_order={zone_order}")
records, metrics = evaluate_multi_zone_days(
points,
mode=args.mode,
cfg=cfg,
gate=gate,
zone_order=zone_order,
impact_store=impact_store,
model=model,
env=env,
)
else:
records, metrics = evaluate_points(
points,
mode=args.mode,
cfg=cfg,
gate=gate,
impact_store=impact_store,
model=model,
env=env,
)
_print_metrics(args.mode, metrics)
by_zone: Dict[str, EvalMetrics] = {}
for r in records:
zm = by_zone.setdefault(r.zone_id, EvalMetrics())
zm.n_days += 1
if r.gt_source == "l1":
if r.product_alert:
zm.tp += 1
else:
zm.fn += 1
if r.belief_raised:
zm.belief_tp += 1
else:
zm.belief_fn += 1
zm.n_l1 += 1
elif r.gt_source == "unlabeled":
zm.n_unlabeled += 1
if r.product_alert:
zm.n_unlabeled_product += 1
if r.belief_raised:
zm.n_unlabeled_belief_raised += 1
if r.product_alert:
zm.n_product += 1
if r.belief_raised:
zm.n_belief_raised += 1
print("\nper-zone:")
for zid, zm in by_zone.items():
def _f(x: Optional[float]) -> str:
return f"{x:.3f}" if x is not None else " - "
print(
f" {zid:22s} n={zm.n_days:3d} l1={zm.n_l1:3d} "
f"unlab={zm.n_unlabeled:3d} "
f"R={_f(zm.recall)} belief_R={_f(zm.belief_recall)} "
f"TP={zm.tp} FN={zm.fn} "
f"unlab_alert={_f(zm.unlabeled_alert_rate)} "
f"unlab_belief_raise={_f(zm.belief_unlabeled_raise_rate)}"
)
result = {
"mode": args.mode,
"start": start.isoformat(),
"end": end.isoformat(),
"zones": zone_ids,
"metrics": metrics.to_dict(),
"per_zone": {z: m.to_dict() for z, m in by_zone.items()},
"records": [asdict(r) for r in records],
"gate": {
"drought_threshold": gate.drought_threshold,
"flood_threshold": gate.flood_threshold,
"min_alert_level": gate.min_alert_level.value,
},
}
if args.out:
with open(args.out, "w") as f:
json.dump(result, f, indent=2)
print(f"\nWrote {args.out}")
return 0
# ---------------------------------------------------------------------------
# Offline self-test (no pkl required)
# ---------------------------------------------------------------------------
def _self_test() -> None:
print("evaluate_checkpoint_real self-test")
from zone_observation import make_synthetic_zone_obs, make_synthetic_forecast_result
obs = make_synthetic_zone_obs("karawang_rice", drought=True, seed=1)
fc = make_synthetic_forecast_result(
zone_id="karawang_rice", valid_time=obs.valid_time, drought=True, seed=1
)
cfg = ForecastConfig()
gate = DEFAULT_PRODUCT_GATE
product, elevated, rs, _, _bp, _ib = decide_scorer_oracle(obs, fc, cfg, gate)
assert rs is not None
print(f" drought scorer product={product} elevated={elevated} "
f"alert={rs.alert_level.value} drought_risk={rs.drought_risk:.3f}")
m = EvalMetrics(n_days=4, tp=1, fp=1, fn=1, tn=1, n_l1=2, n_product=2)
assert abs((m.precision or 0) - 0.5) < 1e-9
assert abs((m.recall or 0) - 0.5) < 1e-9
print(" metrics arithmetic OK")
m_pos = EvalMetrics(n_days=10, tp=7, fn=3, n_l1=10, fp=0, tn=0)
assert m_pos.precision is None, m_pos.precision
assert m_pos.f1 is None
assert abs((m_pos.recall or 0) - 0.7) < 1e-9
assert m_pos.belief_precision is None
assert m_pos.belief_f1 is None
print(" positive-only null precision OK")
m_u = EvalMetrics(
n_unlabeled=20,
n_unlabeled_product=8,
n_unlabeled_belief_raised=11,
)
assert abs((m_u.unlabeled_alert_rate or 0) - 0.4) < 1e-9
assert abs((m_u.belief_unlabeled_raise_rate or 0) - 0.55) < 1e-9
print(" unlabeled rates OK")
def _f(x):
return f"{x:.3f}" if x is not None else " - "
assert _f(None) == " - "
assert _f(0.0) == "0.000"
assert "0.000" not in _f(None)
print(" None print formatting OK")
print("All evaluate_checkpoint_real self-tests passed.")
if __name__ == "__main__":
if len(sys.argv) == 1:
_self_test()
else:
raise SystemExit(main())