""" real_episode_sampler.py ======================== Loads a historical trajectory cache (the same format build_continuous_historical.py produces) and exposes a sampler that WeatherForecastEnv.reset() can draw real multi-zone EpisodeContexts from during training. Import-light by design: weather_forecast_env.py must not become dependent on the era5_data_pipeline stack just to read a cache dict that was already built. -------------------------------------------------------------- DEFAULT_HOLDOUT_RANGES below encodes the exclusion by whole paradigmatic season (matching build_continuous_historical.py's own PARADIGMATIC_SEASONS boundaries), so it unambiguously covers every window already quoted in technical_details.md and this project's evaluate_checkpoint_real.py runs: - 2023-07-01 -> 2023-11-30 (the "dry product freeze" window) falls entirely inside el_nino_2023_24_strong (2023-05-01 -> 2024-04-30). - 2017-06-10 (the PATH_CHECK "argmax miss / OR hit" reference example) falls inside neutral_2017_18 (2017-05-01 -> 2018-04-30). Both whole seasons are held out, not just the specific cited days, so that any future real-eval run inside those seasons remains comparable to what's already on record. """ from __future__ import annotations import logging import pickle import random from datetime import date from pathlib import Path from typing import Any, Dict, List, Optional, Tuple from zone_observation import ( BasinContext, EpisodeContext, ForecastConfig, ForecastResult, ZoneObs, ) logger = logging.getLogger(__name__) DEFAULT_HOLDOUT_RANGES: Tuple[Tuple[str, str], ...] = ( ("2023-05-01", "2024-04-30"), # el_nino_2023_24_strong ("2017-05-01", "2018-04-30"), # neutral_2017_18 ) def _parse_day(s: str) -> date: return date.fromisoformat(str(s)[:10]) def _in_holdout(d: date, holdout: Tuple[Tuple[str, str], ...]) -> bool: for start_s, end_s in holdout: if _parse_day(start_s) <= d <= _parse_day(end_s): return True return False # --------------------------------------------------------------------------- # Precip window perturbation # --------------------------------------------------------------------------- _PRECIP_WINDOW_FIELDS = ( "precip_24h_mm", "precip_7d_mm", "precip_14d_mm", "precip_30d_mm", ) def _perturb_precip_windows(d: Dict[str, Any], scale: float, rng: random.Random) -> None: h24 = float(d.get("precip_24h_mm") or 0.0) d7 = float(d.get("precip_7d_mm") or 0.0) d14 = float(d.get("precip_14d_mm") or 0.0) d30 = float(d.get("precip_30d_mm") or 0.0) bucket_1d = h24 bucket_2_7d = max(0.0, d7 - h24) bucket_8_14d = max(0.0, d14 - d7) bucket_15_30d = max(0.0, d30 - d14) def _jitter(x: float) -> float: return max(0.0, x * (1.0 + rng.gauss(0.0, scale))) bucket_1d = _jitter(bucket_1d) bucket_2_7d = _jitter(bucket_2_7d) bucket_8_14d = _jitter(bucket_8_14d) bucket_15_30d = _jitter(bucket_15_30d) d["precip_24h_mm"] = bucket_1d d["precip_7d_mm"] = d["precip_24h_mm"] + bucket_2_7d d["precip_14d_mm"] = d["precip_7d_mm"] + bucket_8_14d d["precip_30d_mm"] = d["precip_14d_mm"] + bucket_15_30d def _perturb_zone_obs(obs: ZoneObs, scale: float, rng: random.Random) -> ZoneObs: d = obs.to_dict() d.pop("_schema_version", None) _perturb_precip_windows(d, scale, rng) for field_name in ( "temp_mean_c", "temp_max_c", "temp_min_c", "soil_moisture_pct", "rh_mean_pct", "rh_max_pct", "wind_speed_max_ms", "wind_speed_mean_ms", "evapotranspiration_mm", ): val = d.get(field_name) if val is None: continue jitter = 1.0 + rng.gauss(0.0, scale) d[field_name] = max(0.0, float(val) * jitter) for field_name in ( "precip_anomaly_idx", "temp_anomaly_idx", "soil_moisture_anom", "rh_anomaly_idx", ): val = d.get(field_name) if val is None: continue d[field_name] = float(val) + rng.gauss(0.0, scale * 2.0) return ZoneObs.from_dict(d) class RealEpisodeIndex: def __init__( self, pkl_path: str, holdout_ranges: Tuple[Tuple[str, str], ...] = DEFAULT_HOLDOUT_RANGES, ) -> None: self.pkl_path = pkl_path self.holdout_ranges = holdout_ranges self._by_date: Dict[str, Dict[str, Dict[str, Any]]] = {} self._eligible_dates: List[str] = [] self._all_zone_ids: List[str] = [] self._load() def _load(self) -> None: p = Path(self.pkl_path) if not p.is_file(): raise FileNotFoundError( f"RealEpisodeIndex: pkl not found at {self.pkl_path}" ) with open(p, "rb") as f: cache = pickle.load(f) trajs = cache.get("trajectories") or [] n_points = 0 n_excluded = 0 zone_set = set() for traj in trajs: for pt in traj.get("trajectory") or []: obs = pt.get("obs") or {} vt = str(obs.get("valid_time") or pt.get("valid_time") or "")[:10] zid = obs.get("zone_id") or pt.get("zone_id") if not vt or not zid: continue try: d = _parse_day(vt) except Exception: continue if _in_holdout(d, self.holdout_ranges): n_excluded += 1 continue self._by_date.setdefault(vt, {})[zid] = pt zone_set.add(zid) n_points += 1 self._all_zone_ids = sorted(zone_set) self._eligible_dates = sorted(self._by_date.keys()) logger.info( "RealEpisodeIndex: %d training-eligible points (%d excluded by " "holdout=%s) across %d zones, %d dates, from %s", n_points, n_excluded, self.holdout_ranges, len(self._all_zone_ids), len(self._eligible_dates), self.pkl_path, ) if not self._eligible_dates: raise ValueError( f"RealEpisodeIndex: 0 eligible points after holdout filter " f"-- check holdout_ranges={self.holdout_ranges} against " f"the actual date span of {self.pkl_path}" ) @property def n_eligible_dates(self) -> int: return len(self._eligible_dates) @property def all_zone_ids(self) -> List[str]: return list(self._all_zone_ids) def is_holdout(self, day: date) -> bool: return _in_holdout(day, self.holdout_ranges) def sample( self, rng: random.Random, n_zones: int, config: ForecastConfig, inject_noise: bool = False, noise_scale: float = 0.05, max_attempts: int = 20, ) -> Optional[EpisodeContext]: for _ in range(max_attempts): day_s = rng.choice(self._eligible_dates) day_points = self._by_date[day_s] if len(day_points) < n_zones: continue zone_ids = rng.sample(sorted(day_points.keys()), n_zones) try: zone_obs: List[ZoneObs] = [] zone_fc: List[ForecastResult] = [] basin: Optional[BasinContext] = None for zid in zone_ids: pt = day_points[zid] zo = ZoneObs.from_dict(dict(pt["obs"])) zf = ForecastResult.from_dict(dict(pt["forecast"])) if inject_noise: zo = _perturb_zone_obs(zo, noise_scale, rng) 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.debug( "RealEpisodeIndex: basin_context skipped " "for %s/%s (%s)", day_s, zid, e, ) return EpisodeContext( obs=zone_obs[0], forecast=zone_fc[0], config=config, zone_ids=list(zone_ids), basin_context=basin, zone_obs=zone_obs, zone_forecasts=zone_fc, data_source=zone_obs[0].source, ) except Exception as e: logger.debug( "RealEpisodeIndex: sample build failed for %s (%s), " "retrying", day_s, e, ) continue return None # --------------------------------------------------------------------------- # Self-test (python real_episode_sampler.py) -- builds its own tiny # synthetic-shaped pkl on disk so it needs no external data file. # --------------------------------------------------------------------------- def _self_test() -> int: import sys import tempfile from datetime import timezone from zone_observation import make_synthetic_zone_obs, make_synthetic_forecast_result logging.basicConfig(level=logging.WARNING) print("real_episode_sampler.py self-test\n") failures: List[str] = [] def _assert(cond: bool, msg: str) -> None: if not cond: failures.append(msg) print(f" FAIL: {msg}") zone_ids = ["zone_a", "zone_b", "zone_c"] dates_in_holdout = ["2023-08-01", "2023-08-06"] dates_eligible = ["2019-08-01", "2019-08-06", "2019-08-11"] trajectories = [] for zid in zone_ids: points = [] for day_s in dates_in_holdout + dates_eligible: seed = hash((zid, day_s)) % 100000 obs = make_synthetic_zone_obs(zid, seed=seed) d = obs.to_dict() d["valid_time"] = day_s + "T00:00:00+00:00" fc = make_synthetic_forecast_result(zid, valid_time=obs.valid_time, seed=seed) points.append({ "valid_time": day_s, "zone_id": zid, "obs": d, "forecast": fc.to_dict(), "basin_context": None, }) trajectories.append({"meta": {"zone_id": zid}, "trajectory": points}) cache = {"trajectories": trajectories} with tempfile.TemporaryDirectory() as td: pkl_path = str(Path(td) / "test_cache.pkl") with open(pkl_path, "wb") as f: pickle.dump(cache, f) idx = RealEpisodeIndex(pkl_path) # 1. Holdout dates must never be indexed. for day_s in dates_in_holdout: _assert(day_s not in idx._by_date, f"holdout date {day_s} was indexed (must be excluded)") for day_s in dates_eligible: _assert(day_s in idx._by_date, f"eligible date {day_s} was NOT indexed") _assert(idx.n_eligible_dates == len(dates_eligible), f"expected {len(dates_eligible)} eligible dates, " f"got {idx.n_eligible_dates}") print(f" Holdout filtering OK " f"({idx.n_eligible_dates} eligible / " f"{len(dates_in_holdout)} excluded)") # 2. is_holdout() agrees with what actually got indexed. _assert(idx.is_holdout(_parse_day("2023-08-01")), "is_holdout() should be True for a holdout-season date") _assert(not idx.is_holdout(_parse_day("2019-08-01")), "is_holdout() should be False for an eligible date") # 3. Sampling only ever returns eligible-date zones, never holdout. cfg = ForecastConfig(n_zones=3) rng = random.Random(7) seen_dates = set() for _ in range(50): ctx = idx.sample(rng, n_zones=3, config=cfg) _assert(ctx is not None, "sample() unexpectedly returned None") if ctx is not None: d = ctx.obs.valid_time.date() _assert(not idx.is_holdout(d), f"sample() returned a holdout date: {d}") seen_dates.add(d.isoformat()) _assert(len(ctx.zone_ids) == 3, f"expected 3 zones, got {len(ctx.zone_ids)}") _assert(len(ctx.zone_obs) == 3 and len(ctx.zone_forecasts) == 3, "zone_obs/zone_forecasts length mismatch") _assert(len(seen_dates) >= 2, f"sampling looks non-random across dates: only saw {seen_dates}") print(f" Sampling OK (never returned a holdout date across 50 draws, " f"saw {len(seen_dates)}/{len(dates_eligible)} distinct eligible dates)") # 4. Noise injection actually perturbs values (and is reproducible # given the same rng state) without touching zone_id/valid_time, # AND preserves the nested precip-window ordering (the fix under # test: 24h <= 7d <= 14d <= 30d must hold after perturbation, # not just before it). rng_a = random.Random(42) ctx_plain = idx.sample(rng_a, n_zones=3, config=cfg, inject_noise=False) rng_b = random.Random(42) ctx_noisy = idx.sample(rng_b, n_zones=3, config=cfg, inject_noise=True, noise_scale=0.1) _assert(ctx_plain is not None and ctx_noisy is not None, "noise-injection sample() returned None") if ctx_plain is not None and ctx_noisy is not None: differs = any( abs(a.precip_24h_mm - b.precip_24h_mm) > 1e-9 or abs(a.rh_max_pct - b.rh_max_pct) > 1e-9 for a, b in zip(ctx_plain.zone_obs, ctx_noisy.zone_obs) ) _assert(differs, "inject_noise=True produced identical values to inject_noise=False") same_ids = all( a.zone_id == b.zone_id for a, b in zip(ctx_plain.zone_obs, ctx_noisy.zone_obs) ) _assert(same_ids, "noise injection altered zone_id identity") for zo in ctx_noisy.zone_obs: _assert( zo.precip_24h_mm <= zo.precip_7d_mm + 1e-6 and zo.precip_7d_mm <= zo.precip_14d_mm + 1e-6 and zo.precip_14d_mm <= zo.precip_30d_mm + 1e-6, f"noise injection broke precip window nesting for " f"{zo.zone_id}: 24h={zo.precip_24h_mm} 7d={zo.precip_7d_mm} " f"14d={zo.precip_14d_mm} 30d={zo.precip_30d_mm}", ) print(" Noise injection OK (perturbs continuous fields, preserves " "identity fields, preserves nested precip-window ordering)") # 5. n_zones larger than any single day's coverage falls back to None. ctx_toolarge = idx.sample(random.Random(1), n_zones=99, config=ForecastConfig(n_zones=99)) _assert(ctx_toolarge is None, "sample() should return None when no date has enough zone coverage") print(" Graceful None fallback OK (n_zones exceeding available coverage)") # 6. Empty-after-holdout raises clearly rather than silently # returning an unusable index. try: RealEpisodeIndex(pkl_path, holdout_ranges=(("2000-01-01", "2030-01-01"),)) _assert(False, "RealEpisodeIndex should raise when holdout excludes everything") except ValueError: pass print(" Empty-after-holdout raises ValueError OK") print() if failures: print(f"FAILED {len(failures)} test(s):") for f in failures: print(f" - {f}") return 1 print("All 6 test groups passed.") return 0 if __name__ == "__main__": import sys sys.exit(_self_test())