Update zone_observation.py

zone_observation.py

@@ -10,7 +10,7 @@ import logging
 import math
 import random
 import zlib
-from dataclasses import dataclass, field, asdict
+from dataclasses import dataclass, field, asdict, fields as _dc_fields
 from datetime import datetime, timedelta, timezone
 from enum import Enum, unique
 from typing import Any, ClassVar, Dict, List, Optional, Tuple
@@ -172,11 +172,27 @@ class ForecastConfig:
 
     seed: Optional[int] = None
 
-    real_data_ratio:   float               = 0.7   # fraction of episodes that attempt real data
+    real_data_ratio:   float               = 0.7   # fraction of episodes that attempt real data.
+                                                     # Consumed by era5_data_pipeline.py when
+                                                     # BUILDING a historical cache, AND (new) by
+                                                     # WeatherForecastEnv.reset() when
+                                                     # real_data_pkl_path is set, as the
+                                                     # per-episode probability of sampling a real
+                                                     # historical EpisodeContext instead of a
+                                                     # synthetic one during training.
     era5_ratio:        float               = 0.5   # of real-data attempts, fraction using ERA5
     force_data_source: Optional[DataSource] = None  # pin source for debug/test (overrides above)
-    inject_noise:      bool                = False  # apply stochastic noise after fetch
+    inject_noise:      bool                = False  # apply stochastic noise after fetch, AND (new)
+                                                      # after WeatherForecastEnv real-data sampling
     noise_scale:       float               = 0.05   # noise magnitude (fraction of field range)
+    real_data_pkl_path: Optional[str]      = None   # path to a historical trajectory cache
+                                                      # (see real_episode_sampler.RealEpisodeIndex).
+                                                      # When set, WeatherForecastEnv.reset() may
+                                                      # sample real episodes for training with
+                                                      # probability real_data_ratio, respecting
+                                                      # RealEpisodeIndex's holdout exclusion. When
+                                                      # None (default), training is 100% synthetic,
+                                                      # unchanged from prior behavior.
 
     use_satellite_precip: bool  = False  # prefer IMERG/CHIRPS over ERA5 precip
     use_satellite_soil:   bool  = False  # prefer SMAP over ERA5 soil moisture
@@ -313,6 +329,7 @@ class ForecastConfig:
             ),
             "inject_noise": self.inject_noise,
             "noise_scale": self.noise_scale,
+            "real_data_pkl_path": self.real_data_pkl_path,
             "use_satellite_precip": self.use_satellite_precip,
             "use_satellite_soil": self.use_satellite_soil,
             "include_basin_context": self.include_basin_context,
@@ -947,7 +964,25 @@ class BasinContext:
         _check_schema(sv, "BasinContext")
         d["valid_date"] = datetime.fromisoformat(d["valid_date"])
         d["source"] = DataSource(d.get("source", "synthetic"))
-        return cls(**{k: v for k, v in d.items() if not k.startswith("_")})
+        known = {f.name for f in _dc_fields(cls)}
+        unrecognized = {
+            k: v for k, v in d.items()
+            if not k.startswith("_") and k not in known
+        }
+        if unrecognized:
+            # Forward-compat, not silent data loss: seen in practice with
+            # historical_continuous_indonesia_v1_helio_backfilled_canonical_v2.pkl,
+            # which stamps helio_source_label / helio_backfill_valid_date
+            # directly into basin_context dicts from an out-of-repo backfill
+            # script (surgical_helio_backfill_v2.py) rather than through this
+            # class. Dropping them here is deliberate -- add a real field to
+            # BasinContext if this provenance data should be preserved and
+            # consumed, rather than widening this filter.
+            logger.debug(
+                "BasinContext.from_dict: dropping unrecognized fields %s "
+                "(schema forward-compat)", sorted(unrecognized),
+            )
+        return cls(**{k: v for k, v in d.items() if k in known})
         
 def make_synthetic_basin_context(
     valid_date: Optional[datetime] = None,
@@ -1382,6 +1417,11 @@ if __name__ == "__main__":
     _assert(cfg2.force_data_source is None,      "ForecastConfig force_data_source round-trip")
     _assert(cfg2.inject_noise is False,          "ForecastConfig inject_noise round-trip")
     _assert(cfg2.noise_scale == 0.05,            "ForecastConfig noise_scale round-trip")
+    _assert(cfg2.real_data_pkl_path is None,     "ForecastConfig real_data_pkl_path default round-trip")
+    cfg_pkl = ForecastConfig(real_data_pkl_path="/tmp/example.pkl")
+    cfg_pkl_back = ForecastConfig.from_dict(cfg_pkl.to_dict())
+    _assert(cfg_pkl_back.real_data_pkl_path == "/tmp/example.pkl",
+            "non-default real_data_pkl_path round-trip")
     cfg_era5 = ForecastConfig(force_data_source=DataSource.ERA5_REANALYSIS)
     cfg_era5_back = ForecastConfig.from_dict(cfg_era5.to_dict())
     _assert(
@@ -1514,6 +1554,19 @@ if __name__ == "__main__":
           f"kp={bc.kp_index:.1f}  regime={bc.helio_regime}  "
           f"round-trip OK, EpisodeContext integration OK")
 
+    # 10b. BasinContext.from_dict tolerates unrecognized fields (schema
+    # drift from out-of-repo backfill scripts) instead of crashing.
+    d_drift = bc.to_dict()
+    d_drift["helio_source_label"] = "gfz_kp_daily_max+neutral_sw_goes_v2"
+    d_drift["helio_backfill_valid_date"] = "2023-08-04"
+    bc_drift = BasinContext.from_dict(d_drift)
+    _assert(abs(bc_drift.enso_oni - bc.enso_oni) < 1e-9,
+            "BasinContext.from_dict with unrecognized fields lost a known field")
+    _assert(not hasattr(bc_drift, "helio_source_label"),
+            "BasinContext.from_dict should not silently attach unknown attrs")
+    print("  BasinContext.from_dict schema-drift tolerance OK "
+          "(helio_source_label-style extra keys no longer crash)")
+
     # 11. New optional ZoneObs satellite fields: None-by-default, clipping, round-trip
     obs_sat = ZoneObs(
         zone_id="sat_zone", valid_time=now,